⚡️ Speed up method AutoformerLayernorm.forward by 27%

[codeflash-ai]

📄 27% (0.27x) speedup for AutoformerLayernorm.forward in src/transformers/models/autoformer/modeling_autoformer.py

⏱️ Runtime : 4.43 milliseconds → 3.48 milliseconds (best of 43 runs)

📝 Explanation and details

The optimization replaces an inefficient tensor manipulation sequence with PyTorch's native broadcasting mechanism.

Key Change: The bias calculation was changed from torch.mean(x_hat, dim=1).unsqueeze(1).repeat(1, x.shape[1], 1) to torch.mean(x_hat, dim=1, keepdim=True).

Why This is Faster:

• Eliminates redundant memory allocation: The original code explicitly creates a full-sized tensor through .repeat(), copying the mean values across the entire sequence dimension
• Leverages PyTorch's optimized broadcasting: Using keepdim=True maintains the dimension structure, allowing PyTorch to broadcast the subtraction operation without creating intermediate tensors
• Reduces memory bandwidth: Broadcasting operations are handled at the kernel level, avoiding the memory overhead of creating and copying full-sized bias tensors

Performance Impact: The line profiler shows the bias calculation time dropped from 3.22ms (39.8% of total time) to 1.34ms (21.8% of total time) - a ~58% reduction in that operation's cost.

Test Case Performance: The optimization is particularly effective for larger tensors, showing 35-47% speedups on most test cases, with the largest improvements on constant/simple inputs where the broadcasting advantage is most pronounced. Even edge cases with small tensors see 20-40% improvements, demonstrating the optimization's broad applicability.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 85 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests and Runtime

import pytest  # used for our unit tests
import torch
from transformers.models.autoformer.modeling_autoformer import \
    AutoformerLayernorm


# function to test
class DummyConfig:
    """Minimal config to mimic AutoformerConfig for testing purposes."""
    def __init__(self, d_model):
        self.d_model = d_model
from transformers.models.autoformer.modeling_autoformer import \
    AutoformerLayernorm

# unit tests

# ----------- Basic Test Cases -----------

def test_forward_basic_identity():
    # Test that output shape matches input shape for basic input
    config = DummyConfig(d_model=4)
    layernorm = AutoformerLayernorm(config)
    x = torch.rand(2, 3, 4)  # batch=2, seq=3, features=4
    codeflash_output = layernorm.forward(x); y = codeflash_output # 85.7μs -> 62.8μs (36.5% faster)

def test_forward_basic_zero_input():
    # Test that zero input returns zero output
    config = DummyConfig(d_model=5)
    layernorm = AutoformerLayernorm(config)
    x = torch.zeros(1, 2, 5)
    codeflash_output = layernorm.forward(x); y = codeflash_output # 82.2μs -> 58.9μs (39.5% faster)

def test_forward_basic_constant_input():
    # Test that constant input returns zero output
    config = DummyConfig(d_model=6)
    layernorm = AutoformerLayernorm(config)
    x = torch.ones(2, 4, 6) * 7.0
    codeflash_output = layernorm.forward(x); y = codeflash_output # 72.7μs -> 50.0μs (45.4% faster)

def test_forward_basic_varying_input():
    # Test that output mean along sequence is zero
    config = DummyConfig(d_model=3)
    layernorm = AutoformerLayernorm(config)
    x = torch.tensor([[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]])  # shape [1,2,3]
    codeflash_output = layernorm.forward(x); y = codeflash_output # 79.9μs -> 58.2μs (37.2% faster)
    # Output mean along sequence dimension should be zero for each batch and feature
    mean = torch.mean(y, dim=1)

# ----------- Edge Test Cases -----------

def test_forward_edge_single_element():
    # Test with batch=1, seq=1, d_model=1
    config = DummyConfig(d_model=1)
    layernorm = AutoformerLayernorm(config)
    x = torch.tensor([[[42.0]]])
    codeflash_output = layernorm.forward(x); y = codeflash_output # 79.2μs -> 58.6μs (35.1% faster)

def test_forward_edge_large_d_model():
    # Test with large d_model
    config = DummyConfig(d_model=512)
    layernorm = AutoformerLayernorm(config)
    x = torch.rand(2, 3, 512)
    codeflash_output = layernorm.forward(x); y = codeflash_output # 83.9μs -> 61.8μs (35.7% faster)
    # Check mean along sequence is close to zero
    mean = torch.mean(y, dim=1)

def test_forward_edge_empty_sequence():
    # Test with sequence length 0
    config = DummyConfig(d_model=8)
    layernorm = AutoformerLayernorm(config)
    x = torch.rand(2, 0, 8)
    codeflash_output = layernorm.forward(x); y = codeflash_output # 67.0μs -> 54.5μs (22.9% faster)

def test_forward_edge_empty_batch():
    # Test with batch size 0
    config = DummyConfig(d_model=8)
    layernorm = AutoformerLayernorm(config)
    x = torch.rand(0, 5, 8)
    codeflash_output = layernorm.forward(x); y = codeflash_output # 74.4μs -> 54.4μs (36.8% faster)

def test_forward_edge_negative_input():
    # Test with negative values
    config = DummyConfig(d_model=4)
    layernorm = AutoformerLayernorm(config)
    x = -torch.ones(3, 2, 4) * 5.0
    codeflash_output = layernorm.forward(x); y = codeflash_output # 73.1μs -> 49.7μs (47.0% faster)

def test_forward_edge_nan_input():
    # Test with NaN values
    config = DummyConfig(d_model=4)
    layernorm = AutoformerLayernorm(config)
    x = torch.tensor([[[float('nan'), 1.0, 2.0, 3.0]]])
    codeflash_output = layernorm.forward(x); y = codeflash_output # 79.0μs -> 57.3μs (37.8% faster)

def test_forward_edge_inf_input():
    # Test with Inf values
    config = DummyConfig(d_model=4)
    layernorm = AutoformerLayernorm(config)
    x = torch.tensor([[[float('inf'), 1.0, 2.0, 3.0]]])
    codeflash_output = layernorm.forward(x); y = codeflash_output # 78.7μs -> 56.7μs (38.8% faster)

# ----------- Large Scale Test Cases -----------

def test_forward_large_batch_and_seq():
    # Test with large batch and sequence, but under 100MB
    config = DummyConfig(d_model=16)
    batch_size = 128
    seq_len = 64
    x = torch.rand(batch_size, seq_len, config.d_model)
    layernorm = AutoformerLayernorm(config)
    codeflash_output = layernorm.forward(x); y = codeflash_output # 331μs -> 292μs (13.3% faster)
    # Output mean along sequence should be zero
    mean = torch.mean(y, dim=1)

def test_forward_large_d_model():
    # Test with large d_model, but under 100MB
    config = DummyConfig(d_model=1024)
    batch_size = 8
    seq_len = 8
    x = torch.rand(batch_size, seq_len, config.d_model)
    layernorm = AutoformerLayernorm(config)
    codeflash_output = layernorm.forward(x); y = codeflash_output # 129μs -> 95.3μs (35.8% faster)
    mean = torch.mean(y, dim=1)

def test_forward_large_random_values():
    # Test with large random values to check numerical stability
    config = DummyConfig(d_model=32)
    batch_size = 64
    seq_len = 32
    x = torch.rand(batch_size, seq_len, config.d_model) * 1e6
    layernorm = AutoformerLayernorm(config)
    codeflash_output = layernorm.forward(x); y = codeflash_output # 154μs -> 128μs (20.0% faster)
    mean = torch.mean(y, dim=1)

def test_forward_large_extreme_values():
    # Test with extreme values (very high and low)
    config = DummyConfig(d_model=16)
    batch_size = 32
    seq_len = 16
    x = torch.cat([
        torch.full((batch_size // 2, seq_len, config.d_model), 1e9),
        torch.full((batch_size // 2, seq_len, config.d_model), -1e9)
    ])
    layernorm = AutoformerLayernorm(config)
    codeflash_output = layernorm.forward(x); y = codeflash_output # 96.6μs -> 74.5μs (29.6% faster)
    mean = torch.mean(y, dim=1)

def test_forward_large_shape_limit():
    # Test with tensor size close to 100MB limit
    config = DummyConfig(d_model=128)
    batch_size = 64
    seq_len = 64
    x = torch.rand(batch_size, seq_len, config.d_model)
    layernorm = AutoformerLayernorm(config)
    codeflash_output = layernorm.forward(x); y = codeflash_output # 674μs -> 574μs (17.5% faster)
    mean = torch.mean(y, dim=1)

# ----------- Mutation Testing Robustness -----------

def test_forward_mutation_behavior():
    # If the function is mutated to not subtract the mean, output mean along sequence will not be zero
    config = DummyConfig(d_model=4)
    layernorm = AutoformerLayernorm(config)
    x = torch.rand(2, 5, 4)
    codeflash_output = layernorm.forward(x); y = codeflash_output # 82.1μs -> 59.8μs (37.2% faster)
    mean = torch.mean(y, dim=1)

# ----------- Miscellaneous Test Cases -----------

def test_forward_gradients():
    # Check that gradients can be computed through the layer
    config = DummyConfig(d_model=3)
    layernorm = AutoformerLayernorm(config)
    x = torch.rand(2, 2, 3, requires_grad=True)
    codeflash_output = layernorm.forward(x); y = codeflash_output # 79.8μs -> 57.5μs (38.7% faster)
    loss = y.sum()
    loss.backward()

def test_forward_dtype_float32():
    # Test with float32 input
    config = DummyConfig(d_model=4)
    layernorm = AutoformerLayernorm(config)
    x = torch.rand(2, 3, 4, dtype=torch.float32)
    codeflash_output = layernorm.forward(x); y = codeflash_output # 81.2μs -> 59.3μs (36.9% faster)

def test_forward_dtype_float64():
    # Test with float64 input
    config = DummyConfig(d_model=4)
    layernorm = AutoformerLayernorm(config)
    x = torch.rand(2, 3, 4, dtype=torch.float64)
    codeflash_output = layernorm.forward(x); y = codeflash_output

def test_forward_different_shapes():
    # Test with multiple shape combinations
    config = DummyConfig(d_model=2)
    layernorm = AutoformerLayernorm(config)
    for batch in [1, 5]:
        for seq in [1, 7]:
            x = torch.rand(batch, seq, 2)
            codeflash_output = layernorm.forward(x); y = codeflash_output
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import pytest  # used for our unit tests
import torch  # used for tensor operations
from transformers.models.autoformer.modeling_autoformer import \
    AutoformerLayernorm


# function to test
# (copied from above)
class DummyConfig:
    """Minimal config for testing purposes."""
    def __init__(self, d_model):
        self.d_model = d_model
from transformers.models.autoformer.modeling_autoformer import \
    AutoformerLayernorm

# unit tests

# ---- Basic Test Cases ----

def test_forward_basic_identity():
    # Test that the output is zero when input is constant
    config = DummyConfig(d_model=4)
    model = AutoformerLayernorm(config)
    # batch_size=2, seq_len=3, d_model=4
    x = torch.ones(2, 3, 4)
    codeflash_output = model.forward(x); out = codeflash_output # 77.6μs -> 56.6μs (37.1% faster)

def test_forward_basic_random():
    # Test with random input, check shape and mean subtraction
    config = DummyConfig(d_model=5)
    model = AutoformerLayernorm(config)
    x = torch.randn(2, 4, 5)
    codeflash_output = model.forward(x); out = codeflash_output # 79.7μs -> 58.8μs (35.5% faster)
    # For each batch, mean along seq_len should be zero
    means = torch.mean(out, dim=1)

def test_forward_basic_batch_size_one():
    # Test with batch size 1
    config = DummyConfig(d_model=3)
    model = AutoformerLayernorm(config)
    x = torch.tensor([[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]])  # shape (1, 2, 3)
    codeflash_output = model.forward(x); out = codeflash_output # 79.0μs -> 58.0μs (36.3% faster)
    means = torch.mean(out, dim=1)

def test_forward_basic_seq_len_one():
    # Test with sequence length 1
    config = DummyConfig(d_model=2)
    model = AutoformerLayernorm(config)
    x = torch.tensor([[[1.0, 2.0]], [[3.0, 4.0]]])  # shape (2, 1, 2)
    codeflash_output = model.forward(x); out = codeflash_output # 78.0μs -> 56.7μs (37.4% faster)
    means = torch.mean(out, dim=1)

# ---- Edge Test Cases ----

def test_forward_edge_d_model_one():
    # Test with d_model=1 (single feature)
    config = DummyConfig(d_model=1)
    model = AutoformerLayernorm(config)
    x = torch.tensor([[[1.0], [2.0], [3.0]]])  # shape (1, 3, 1)
    codeflash_output = model.forward(x); out = codeflash_output # 81.4μs -> 59.8μs (36.3% faster)
    means = torch.mean(out, dim=1)

def test_forward_edge_large_values():
    # Test with large values to check for numerical stability
    config = DummyConfig(d_model=2)
    model = AutoformerLayernorm(config)
    x = torch.tensor([[[1e10, -1e10], [1e10, -1e10]]])  # shape (1, 2, 2)
    codeflash_output = model.forward(x); out = codeflash_output # 78.5μs -> 57.5μs (36.4% faster)
    means = torch.mean(out, dim=1)

def test_forward_edge_small_values():
    # Test with very small values
    config = DummyConfig(d_model=2)
    model = AutoformerLayernorm(config)
    x = torch.tensor([[[1e-10, -1e-10], [1e-10, -1e-10]]])  # shape (1, 2, 2)
    codeflash_output = model.forward(x); out = codeflash_output # 79.0μs -> 58.2μs (35.7% faster)
    means = torch.mean(out, dim=1)

def test_forward_edge_zero_input():
    # Test with all zeros input
    config = DummyConfig(d_model=3)
    model = AutoformerLayernorm(config)
    x = torch.zeros(2, 4, 3)
    codeflash_output = model.forward(x); out = codeflash_output # 80.8μs -> 58.5μs (38.2% faster)

def test_forward_edge_negative_values():
    # Test with all negative values
    config = DummyConfig(d_model=3)
    model = AutoformerLayernorm(config)
    x = -torch.ones(2, 4, 3)
    codeflash_output = model.forward(x); out = codeflash_output # 79.1μs -> 56.6μs (39.7% faster)

def test_forward_edge_single_element():
    # Test with a single element tensor
    config = DummyConfig(d_model=1)
    model = AutoformerLayernorm(config)
    x = torch.tensor([[[42.0]]])  # shape (1, 1, 1)
    codeflash_output = model.forward(x); out = codeflash_output # 80.8μs -> 59.3μs (36.2% faster)

def test_forward_edge_non_contiguous_input():
    # Test with non-contiguous input
    config = DummyConfig(d_model=3)
    model = AutoformerLayernorm(config)
    x = torch.randn(2, 4, 3)
    x_t = x.transpose(0, 1)  # make non-contiguous
    codeflash_output = model.forward(x_t.transpose(0, 1)); out = codeflash_output # 79.0μs -> 59.6μs (32.6% faster)
    means = torch.mean(out, dim=1)

def test_forward_edge_requires_grad():
    # Test with requires_grad input
    config = DummyConfig(d_model=3)
    model = AutoformerLayernorm(config)
    x = torch.randn(2, 4, 3, requires_grad=True)
    codeflash_output = model.forward(x); out = codeflash_output # 80.4μs -> 58.9μs (36.6% faster)
    means = torch.mean(out, dim=1)

def test_forward_edge_dtype_float16():
    # Test with float16 input
    config = DummyConfig(d_model=3)
    model = AutoformerLayernorm(config)
    x = torch.randn(2, 4, 3, dtype=torch.float16)
    codeflash_output = model.forward(x); out = codeflash_output # 87.2μs -> 65.0μs (34.1% faster)
    means = torch.mean(out, dim=1)

def test_forward_edge_dtype_float64():
    # Test with float64 input
    config = DummyConfig(d_model=3)
    model = AutoformerLayernorm(config)
    x = torch.randn(2, 4, 3, dtype=torch.float64)
    codeflash_output = model.forward(x); out = codeflash_output
    means = torch.mean(out, dim=1)

# ---- Large Scale Test Cases ----

def test_forward_large_batch():
    # Test with large batch size
    config = DummyConfig(d_model=8)
    model = AutoformerLayernorm(config)
    batch_size = 512
    seq_len = 10
    x = torch.randn(batch_size, seq_len, config.d_model)
    codeflash_output = model.forward(x); out = codeflash_output # 235μs -> 211μs (11.1% faster)
    means = torch.mean(out, dim=1)

def test_forward_large_seq_len():
    # Test with large sequence length
    config = DummyConfig(d_model=8)
    model = AutoformerLayernorm(config)
    batch_size = 4
    seq_len = 512
    x = torch.randn(batch_size, seq_len, config.d_model)
    codeflash_output = model.forward(x); out = codeflash_output # 137μs -> 113μs (21.2% faster)
    means = torch.mean(out, dim=1)

def test_forward_large_d_model():
    # Test with large d_model
    config = DummyConfig(d_model=512)
    model = AutoformerLayernorm(config)
    batch_size = 2
    seq_len = 4
    x = torch.randn(batch_size, seq_len, config.d_model)
    codeflash_output = model.forward(x); out = codeflash_output # 85.5μs -> 62.8μs (36.2% faster)
    means = torch.mean(out, dim=1)

def test_forward_large_total_elements():
    # Test with large total number of elements, but under 100MB
    config = DummyConfig(d_model=64)
    model = AutoformerLayernorm(config)
    batch_size = 16
    seq_len = 64
    x = torch.randn(batch_size, seq_len, config.d_model)
    codeflash_output = model.forward(x); out = codeflash_output # 143μs -> 113μs (26.9% faster)
    means = torch.mean(out, dim=1)

# ---- Error Handling Test Cases ----

def test_forward_error_wrong_shape():
    # Test with wrong input shape (missing d_model)
    config = DummyConfig(d_model=3)
    model = AutoformerLayernorm(config)
    x = torch.randn(2, 4)  # should be (batch, seq, d_model)
    with pytest.raises(RuntimeError):
        model.forward(x) # 72.3μs -> 72.0μs (0.442% faster)

def test_forward_error_mismatched_d_model():
    # Test with d_model mismatch
    config = DummyConfig(d_model=3)
    model = AutoformerLayernorm(config)
    x = torch.randn(2, 4, 2)  # last dim != d_model
    with pytest.raises(RuntimeError):
        model.forward(x) # 66.4μs -> 65.8μs (0.813% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-AutoformerLayernorm.forward-mha36s10 and push.