⚡️ Speed up method MraSparseDenseMatMul.forward by 77%

[codeflash-ai]

📄 77% (0.77x) speedup for MraSparseDenseMatMul.forward in src/transformers/models/mra/modeling_mra.py

⏱️ Runtime : 99.9 microseconds → 56.5 microseconds (best of 12 runs)

📝 Explanation and details

The optimization achieves a 76% speedup by reducing redundant attribute access and improving tensor operations:

Key Optimizations:

1. Shape Caching: Instead of calling .size() multiple times on tensors (which creates new tuples each time), the code caches tensor shapes once with .shape and reuses them. This eliminates repeated method calls and tuple creation overhead.

2. Efficient Tensor Reshaping: Replaced .reshape() with .view() for the dense_key transformation. Since the tensor is already made contiguous later, .view() is faster as it avoids unnecessary memory allocations when possible.

3. Streamlined Validation: Consolidated dimension checks using cached shapes (e.g., sq_shape[2] instead of sparse_query.size(2)), reducing method call overhead in validation logic.

4. Optimized Method Chaining: Combined indices.int() and .contiguous() into a single chained operation, reducing intermediate tensor creation.

Performance Impact by Test Case:

• Error handling tests show 11-461% improvements due to faster early validation
• The most dramatic improvements (455-461%) occur in tests that trigger the indices dimension error, where shape caching provides maximum benefit
• Even basic validation cases see 12-32% improvements from reduced attribute access

The optimizations are particularly effective for validation-heavy code paths and functions called frequently in ML pipelines, where minimizing method call overhead compounds significantly.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import pytest  # used for our unit tests
import torch
from transformers.models.mra.modeling_mra import MraSparseDenseMatMul

# --- Basic Test Cases ---



def test_forward_block_size_mismatch_raises():
    # block_size mismatch on sparse_query
    batch_size = 1
    key_size = 32
    dim = 32
    block_size = 32
    query_num_block = 1

    # Wrong block_size in sparse_query
    sparse_query = torch.rand(batch_size, query_num_block, block_size + 1, block_size)
    indices = torch.randint(0, key_size // block_size, (batch_size, query_num_block))
    dense_key = torch.rand(batch_size, key_size, dim)

    with pytest.raises(ValueError, match="The size of the second dimension of sparse_query must be equal to the block_size."):
        MraSparseDenseMatMul.forward(None, sparse_query, indices, dense_key, query_num_block) # 4.96μs -> 3.83μs (29.5% faster)

def test_forward_block_size_mismatch_third_dim_raises():
    # block_size mismatch on third dim of sparse_query
    batch_size = 1
    key_size = 32
    dim = 32
    block_size = 32
    query_num_block = 1

    sparse_query = torch.rand(batch_size, query_num_block, block_size, block_size + 1)
    indices = torch.randint(0, key_size // block_size, (batch_size, query_num_block))
    dense_key = torch.rand(batch_size, key_size, dim)

    with pytest.raises(ValueError, match="The size of the third dimension of sparse_query must be equal to the block_size."):
        MraSparseDenseMatMul.forward(None, sparse_query, indices, dense_key, query_num_block) # 4.59μs -> 3.48μs (31.8% faster)

def test_forward_key_size_not_divisible_by_block_size_raises():
    batch_size = 1
    key_size = 33  # Not divisible by block_size
    dim = 32
    block_size = 32
    query_num_block = 1

    sparse_query = torch.rand(batch_size, query_num_block, block_size, block_size)
    indices = torch.randint(0, key_size // block_size, (batch_size, query_num_block))
    dense_key = torch.rand(batch_size, key_size, dim)

    with pytest.raises(ValueError, match="key_size .* must be divisible by block_size."):
        MraSparseDenseMatMul.forward(None, sparse_query, indices, dense_key, query_num_block) # 3.59μs -> 3.21μs (11.9% faster)



def test_forward_indices_not_2d_raises():
    batch_size = 1
    key_size = 32
    dim = 32
    block_size = 32
    query_num_block = 1

    sparse_query = torch.rand(batch_size, query_num_block, block_size, block_size)
    # Make indices 1D
    indices = torch.randint(0, key_size // block_size, (batch_size * query_num_block,))
    dense_key = torch.rand(batch_size, key_size, dim)

    with pytest.raises(ValueError, match="indices must be a 2-dimensional tensor."):
        MraSparseDenseMatMul.forward(None, sparse_query, indices, dense_key, query_num_block) # 23.1μs -> 4.11μs (461% faster)








#------------------------------------------------
import pytest  # used for our unit tests
import torch  # used for tensor creation and manipulation
from transformers.models.mra.modeling_mra import MraSparseDenseMatMul

# --- Unit Tests ---

# Basic Test Cases




def test_forward_key_size_not_divisible_by_block_size():
    """
    Test that ValueError is raised if key_size is not divisible by block_size.
    """
    batch_size = 1
    key_size = 35  # Not divisible by 32
    dim = 32
    query_num_block = 1
    block_size = 32

    sparse_query = torch.randn(batch_size, query_num_block, block_size, block_size)
    indices = torch.randint(0, key_size // block_size, (batch_size, query_num_block))
    dense_key = torch.randn(batch_size, key_size, dim)

    with pytest.raises(ValueError, match="key_size.*divisible.*block_size"):
        MraSparseDenseMatMul.forward(None, sparse_query, indices, dense_key, query_num_block) # 4.10μs -> 3.65μs (12.3% faster)



def test_forward_indices_wrong_shape():
    """
    Test that ValueError is raised if indices does not have 2 dimensions.
    """
    batch_size = 1
    key_size = 32
    dim = 32
    query_num_block = 1
    block_size = 32

    sparse_query = torch.randn(batch_size, query_num_block, block_size, block_size)
    # Make indices 1D
    indices = torch.randint(0, key_size // block_size, (query_num_block,))
    dense_key = torch.randn(batch_size, key_size, dim)

    with pytest.raises(ValueError, match="indices must be a 2-dimensional tensor"):
        MraSparseDenseMatMul.forward(None, sparse_query, indices, dense_key, query_num_block) # 22.7μs -> 4.09μs (455% faster)

def test_forward_sparse_query_block_size_mismatch():
    """
    Test that ValueError is raised if sparse_query's last two dimensions are not equal to block_size.
    """
    batch_size = 1
    key_size = 32
    dim = 32
    query_num_block = 1
    block_size = 32

    # Wrong last dimension
    sparse_query = torch.randn(batch_size, query_num_block, block_size, block_size + 1)
    indices = torch.randint(0, key_size // block_size, (batch_size, query_num_block))
    dense_key = torch.randn(batch_size, key_size, dim)

    with pytest.raises(ValueError, match="third dimension of sparse_query must be equal to the block_size"):
        MraSparseDenseMatMul.forward(None, sparse_query, indices, dense_key, query_num_block) # 4.33μs -> 3.43μs (26.0% faster)


def test_forward_empty_tensors():
    """
    Test that the function raises an error for empty tensors.
    """
    batch_size = 0
    key_size = 32
    dim = 32
    query_num_block = 0
    block_size = 32

    sparse_query = torch.randn(batch_size, query_num_block, block_size, block_size)
    indices = torch.randint(0, key_size // block_size, (batch_size, query_num_block))
    dense_key = torch.randn(batch_size, key_size, dim)

    # Should raise an error due to empty tensors
    with pytest.raises(Exception):
        MraSparseDenseMatMul.forward(None, sparse_query, indices, dense_key, query_num_block) # 32.7μs -> 30.7μs (6.34% faster)

To edit these changes git checkout codeflash/optimize-MraSparseDenseMatMul.forward-mhjwkggk and push.