⚡️ Speed up function compute_dp_attention_local_info by 21%

[codeflash-ai]

📄 21% (0.21x) speedup for compute_dp_attention_local_info in python/sglang/srt/layers/dp_attention.py

⏱️ Runtime : 613 microseconds → 509 microseconds (best of 193 runs)

📝 Explanation and details

The optimization improves performance by eliminating redundant integer divisions and reducing temporary expression evaluations.

Key Changes:

1. Precomputed divisor: tp_size // local_tp_size is calculated once and stored in divisor, avoiding recalculation in the max() expression
2. Simplified max logic: Replaced max(1, dp_size // divisor) with explicit comparison 1 if quotient < 1 else quotient, which is faster than the max() builtin for this binary case
3. Intermediate quotient: Store dp_size // divisor result to avoid redundant division

Why This is Faster:

• Integer division is expensive, especially for large numbers. The original code performed tp_size // local_tp_size twice - once inside the max() call and implicitly again
• The max() builtin has function call overhead compared to a simple conditional
• Fewer temporary objects are created during expression evaluation

Impact on Workloads:
Based on the function reference, this function is called during initialize_dp_attention() - a setup phase for distributed attention mechanisms. While not in a tight loop, the 20% speedup is beneficial because:

• Initialization time affects model startup latency
• The function handles tensor parallelism and data parallelism coordination, which is critical for multi-GPU setups
• Large-scale deployments (as shown in test cases with tp_size=1000) benefit most from the optimization

Test Case Performance:
The optimization shows consistent 30-40% improvements across most enabled attention test cases, with particularly strong gains on large-scale scenarios (35-42% faster for 1000+ parameter cases), indicating the optimization scales well with input size.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from __future__ import annotations

# imports
import pytest  # used for our unit tests
from sglang.srt.layers.dp_attention import compute_dp_attention_local_info

# unit tests

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

def test_basic_disable_attention():
    # If attention is disabled, should return (tp_rank, tp_size, 0)
    codeflash_output = compute_dp_attention_local_info(False, 2, 4, 8, 0) # 403ns -> 385ns (4.68% faster)
    codeflash_output = compute_dp_attention_local_info(False, 0, 1, 1, 0) # 196ns -> 215ns (8.84% slower)
    codeflash_output = compute_dp_attention_local_info(False, 5, 10, 20, 3) # 128ns -> 130ns (1.54% slower)

def test_basic_enable_attention_no_moe():
    # Basic enabled attention, no moe_dense_tp_size
    # tp_size = 4, dp_size = 2, tp_rank = 1
    # local_tp_size = tp_size = 4
    # local_tp_rank = 1 % 4 = 1
    # local_dp_size = max(1, 2 // (4 // 4)) = 2 // 1 = 2
    # local_attn_tp_size = 4 // 2 = 2
    # local_attn_dp_rank = 1 // 2 = 0
    # local_attn_tp_rank = 1 % 2 = 1
    codeflash_output = compute_dp_attention_local_info(True, 1, 4, 2, 0) # 1.17μs -> 817ns (42.6% faster)

def test_basic_enable_attention_with_moe():
    # moe_dense_tp_size overrides tp_size
    # tp_size = 4, moe_dense_tp_size = 2, dp_size = 2, tp_rank = 3
    # local_tp_size = 2
    # local_tp_rank = 3 % 2 = 1
    # local_dp_size = max(1, 2 // (4 // 2)) = 2 // 2 = 1
    # local_attn_tp_size = 2 // 1 = 2
    # local_attn_dp_rank = 1 // 2 = 0
    # local_attn_tp_rank = 1 % 2 = 1
    codeflash_output = compute_dp_attention_local_info(True, 3, 4, 2, 2) # 1.00μs -> 723ns (38.9% faster)

def test_basic_enable_attention_moe_zero():
    # moe_dense_tp_size = 0, so fallback to tp_size
    codeflash_output = compute_dp_attention_local_info(True, 2, 4, 2, 0) # 948ns -> 705ns (34.5% faster)

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

def test_edge_tp_rank_out_of_bounds():
    # tp_rank >= tp_size, should wrap around by modulo
    codeflash_output = compute_dp_attention_local_info(True, 5, 4, 2, 0) # 954ns -> 682ns (39.9% faster)
    codeflash_output = compute_dp_attention_local_info(True, 8, 4, 2, 0) # 637ns -> 575ns (10.8% faster)


def test_edge_dp_size_zero():
    # dp_size = 0, should fallback to max(1, ...)
    codeflash_output = compute_dp_attention_local_info(True, 0, 4, 0, 0) # 1.55μs -> 1.16μs (33.9% faster)
    codeflash_output = compute_dp_attention_local_info(True, 2, 4, 0, 0) # 572ns -> 406ns (40.9% faster)


def test_edge_local_tp_size_divisor_zero():
    # local_tp_size = tp_size, dp_size = 0, so divisor is zero, but max(1, ...) ensures no division by zero
    # Should not raise ZeroDivisionError
    try:
        codeflash_output = compute_dp_attention_local_info(True, 0, 4, 0, 0); result = codeflash_output
    except ZeroDivisionError:
        pytest.fail("ZeroDivisionError raised unexpectedly")

def test_edge_local_attn_tp_size_zero():
    # If local_tp_size < local_dp_size, local_attn_tp_size = 0, which would cause ZeroDivisionError
    # But max(1, ...) prevents local_dp_size from being zero
    # Let's force local_tp_size = 1, local_dp_size = 2, so local_attn_tp_size = 0
    # Should raise ZeroDivisionError
    with pytest.raises(ZeroDivisionError):
        compute_dp_attention_local_info(True, 0, 2, 4, 1) # 1.61μs -> 1.22μs (32.2% faster)

def test_edge_all_zero_inputs():
    # All inputs zero, should not crash, but may return (0, 0, 0)
    codeflash_output = compute_dp_attention_local_info(False, 0, 0, 0, 0)
    # If enabled, should fallback to max(1, ...)
    codeflash_output = compute_dp_attention_local_info(True, 0, 0, 0, 0)

def test_edge_negative_inputs():
    # Negative values: Should handle gracefully, but may produce negative modulo results
    codeflash_output = compute_dp_attention_local_info(False, -1, 4, 2, 0) # 489ns -> 467ns (4.71% faster)
    # If enabled, negative tp_rank modulo local_tp_size
    codeflash_output = compute_dp_attention_local_info(True, -1, 4, 2, 0) # 1.13μs -> 740ns (53.0% faster)
    # Negative tp_size or dp_size
    codeflash_output = compute_dp_attention_local_info(True, 1, -4, 2, 0) # 363ns -> 295ns (23.1% faster)
    codeflash_output = compute_dp_attention_local_info(True, 1, 4, -2, 0) # 455ns -> 369ns (23.3% faster)

def test_edge_moe_dense_tp_size_zero():
    # Explicitly test moe_dense_tp_size = 0, should fallback to tp_size
    codeflash_output = compute_dp_attention_local_info(True, 1, 4, 2, 0) # 1.01μs -> 804ns (25.9% faster)

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

def test_large_scale_tp_size():
    # Large tp_size, dp_size, tp_rank
    tp_size = 999
    dp_size = 997
    tp_rank = 998
    moe_dense_tp_size = 0
    # local_tp_size = 999
    # local_tp_rank = 998 % 999 = 998
    # local_dp_size = max(1, 997 // (999 // 999)) = 997 // 1 = 997
    # local_attn_tp_size = 999 // 997 = 1
    # local_attn_dp_rank = 998 // 1 = 998
    # local_attn_tp_rank = 998 % 1 = 0
    codeflash_output = compute_dp_attention_local_info(True, tp_rank, tp_size, dp_size, moe_dense_tp_size) # 1.13μs -> 835ns (35.8% faster)

def test_large_scale_moe_dense_tp_size():
    # Large moe_dense_tp_size
    tp_size = 500
    dp_size = 100
    tp_rank = 499
    moe_dense_tp_size = 999
    # local_tp_size = 999
    # local_tp_rank = 499 % 999 = 499
    # local_dp_size = max(1, 100 // (500 // 999)) = 100 // 0 = ZeroDivisionError
    # Should raise ZeroDivisionError
    with pytest.raises(ZeroDivisionError):
        compute_dp_attention_local_info(True, tp_rank, tp_size, dp_size, moe_dense_tp_size) # 1.19μs -> 1.05μs (13.4% faster)

def test_large_scale_all_inputs():
    # All inputs large, but valid
    tp_size = 999
    dp_size = 999
    tp_rank = 998
    moe_dense_tp_size = 999
    # local_tp_size = 999
    # local_tp_rank = 998 % 999 = 998
    # local_dp_size = max(1, 999 // (999 // 999)) = 999 // 1 = 999
    # local_attn_tp_size = 999 // 999 = 1
    # local_attn_dp_rank = 998 // 1 = 998
    # local_attn_tp_rank = 998 % 1 = 0
    codeflash_output = compute_dp_attention_local_info(True, tp_rank, tp_size, dp_size, moe_dense_tp_size) # 1.28μs -> 899ns (42.6% faster)

def test_large_scale_loop_over_tp_rank():
    # Test all tp_rank values for large tp_size
    tp_size = 1000
    dp_size = 100
    moe_dense_tp_size = 0
    for tp_rank in range(tp_size):
        codeflash_output = compute_dp_attention_local_info(True, tp_rank, tp_size, dp_size, moe_dense_tp_size); result = codeflash_output # 267μs -> 227μs (17.7% faster)
        # Should not raise, and local_attn_tp_rank in [0, local_attn_tp_size-1]
        local_tp_size = tp_size
        local_tp_rank = tp_rank % local_tp_size
        local_dp_size = max(1, dp_size // (tp_size // local_tp_size))
        local_attn_tp_size = local_tp_size // local_dp_size
        local_attn_tp_rank = local_tp_rank % local_attn_tp_size if local_attn_tp_size > 0 else 0

def test_large_scale_loop_over_dp_size():
    # Test all dp_size values for fixed tp_size
    tp_size = 100
    tp_rank = 10
    moe_dense_tp_size = 0
    for dp_size in range(1, 100):
        codeflash_output = compute_dp_attention_local_info(True, tp_rank, tp_size, dp_size, moe_dense_tp_size); result = codeflash_output # 27.9μs -> 23.4μs (19.4% faster)
        local_tp_size = tp_size
        local_tp_rank = tp_rank % local_tp_size
        local_dp_size = max(1, dp_size // (tp_size // local_tp_size))
        local_attn_tp_size = local_tp_size // local_dp_size
        local_attn_tp_rank = local_tp_rank % local_attn_tp_size if local_attn_tp_size > 0 else 0
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from __future__ import annotations

# imports
import pytest  # used for our unit tests
from sglang.srt.layers.dp_attention import compute_dp_attention_local_info

# unit tests

# --- Basic Test Cases ---

def test_basic_no_attention():
    # If attention is disabled, should return (tp_rank, tp_size, 0)
    codeflash_output = compute_dp_attention_local_info(False, 1, 4, 2, 0) # 404ns -> 383ns (5.48% faster)
    codeflash_output = compute_dp_attention_local_info(False, 0, 8, 4, 2) # 200ns -> 197ns (1.52% faster)

def test_basic_attention_default_moe():
    # Attention enabled, moe_dense_tp_size is 0 (so use tp_size)
    # tp_size=4, dp_size=2, tp_rank=1
    # local_tp_size = 4, local_tp_rank = 1, local_dp_size = max(1, 2 // (4//4)) = 2
    # local_attn_tp_size = 4 // 2 = 2
    # local_attn_dp_rank = 1 // 2 = 0
    # local_attn_tp_rank = 1 % 2 = 1
    codeflash_output = compute_dp_attention_local_info(True, 1, 4, 2, 0) # 1.01μs -> 734ns (37.2% faster)

def test_basic_attention_with_moe():
    # Attention enabled, moe_dense_tp_size is set
    # tp_size=8, moe_dense_tp_size=4, dp_size=2, tp_rank=5
    # local_tp_size = 4, local_tp_rank = 5 % 4 = 1
    # local_dp_size = max(1, 2 // (8//4)) = max(1, 2 // 2) = 1
    # local_attn_tp_size = 4 // 1 = 4
    # local_attn_dp_rank = 1 // 4 = 0
    # local_attn_tp_rank = 1 % 4 = 1
    codeflash_output = compute_dp_attention_local_info(True, 5, 8, 2, 4) # 984ns -> 730ns (34.8% faster)

def test_basic_attention_tp_rank_wrap():
    # tp_rank > local_tp_size, so modulo applies
    # tp_size=6, moe_dense_tp_size=3, tp_rank=5
    # local_tp_size = 3, local_tp_rank = 5 % 3 = 2
    # local_dp_size = max(1, dp_size // (tp_size//local_tp_size)) = max(1, 2 // (6//3)) = max(1, 2//2) = 1
    # local_attn_tp_size = 3 // 1 = 3
    # local_attn_dp_rank = 2 // 3 = 0
    # local_attn_tp_rank = 2 % 3 = 2
    codeflash_output = compute_dp_attention_local_info(True, 5, 6, 2, 3) # 940ns -> 676ns (39.1% faster)

# --- Edge Test Cases ---


def test_edge_dp_size_zero():
    # dp_size=0, should use max(1, ...)
    # tp_size=4, moe_dense_tp_size=0
    # local_tp_size = 4
    # local_dp_size = max(1, 0 // (4//4)) = max(1, 0 // 1) = 1
    # local_attn_tp_size = 4 // 1 = 4
    codeflash_output = compute_dp_attention_local_info(True, 2, 4, 0, 0) # 1.06μs -> 785ns (34.5% faster)

def test_edge_moe_dense_tp_size_zero():
    # Should fall back to tp_size
    codeflash_output = compute_dp_attention_local_info(True, 3, 4, 2, 0) # 994ns -> 737ns (34.9% faster)

def test_edge_tp_rank_zero():
    # tp_rank=0
    codeflash_output = compute_dp_attention_local_info(True, 0, 4, 2, 0) # 1.08μs -> 790ns (36.8% faster)

def test_edge_tp_rank_equals_tp_size():
    # tp_rank = tp_size, so modulo wraps to 0
    codeflash_output = compute_dp_attention_local_info(True, 4, 4, 2, 0) # 1.02μs -> 755ns (35.2% faster)

def test_edge_tp_size_not_divisible_by_local_tp_size():
    # tp_size=7, moe_dense_tp_size=3, dp_size=2, tp_rank=6
    # local_tp_size = 3, local_tp_rank = 6 % 3 = 0
    # local_dp_size = max(1, 2 // (7//3)) = max(1, 2 // 2) = 1
    # local_attn_tp_size = 3 // 1 = 3
    # local_attn_dp_rank = 0 // 3 = 0
    # local_attn_tp_rank = 0 % 3 = 0
    codeflash_output = compute_dp_attention_local_info(True, 6, 7, 2, 3) # 1.00μs -> 783ns (28.2% faster)

def test_edge_local_tp_size_larger_than_tp_size():
    # Unusual: moe_dense_tp_size > tp_size
    # tp_size=4, moe_dense_tp_size=8, tp_rank=3, dp_size=2
    # local_tp_size = 8, local_tp_rank = 3 % 8 = 3
    # local_dp_size = max(1, 2 // (4//8)) = max(1, 2 // 0) => division by zero
    # Should raise ZeroDivisionError
    with pytest.raises(ZeroDivisionError):
        compute_dp_attention_local_info(True, 3, 4, 2, 8) # 1.12μs -> 1.03μs (8.45% faster)

def test_edge_local_tp_size_equals_tp_size():
    # Should behave as if moe_dense_tp_size is not set
    codeflash_output = compute_dp_attention_local_info(True, 2, 4, 2, 4) # 1.14μs -> 845ns (35.4% faster)

def test_edge_all_zeros():
    # All inputs zero except enable_dp_attention
    # tp_rank=0, tp_size=0, dp_size=0, moe_dense_tp_size=0
    # local_tp_size = 0, local_tp_rank = 0 % 0 => ZeroDivisionError
    with pytest.raises(ZeroDivisionError):
        compute_dp_attention_local_info(True, 0, 0, 0, 0) # 903ns -> 844ns (6.99% faster)

# --- Large Scale Test Cases ---

def test_large_scale_tp_size_1000():
    # Large tp_size, dp_size, and tp_rank
    # tp_size=1000, dp_size=500, tp_rank=999, moe_dense_tp_size=0
    # local_tp_size = 1000
    # local_tp_rank = 999
    # local_dp_size = max(1, 500 // (1000//1000)) = 500
    # local_attn_tp_size = 1000 // 500 = 2
    # local_attn_dp_rank = 999 // 2 = 499
    # local_attn_tp_rank = 999 % 2 = 1
    codeflash_output = compute_dp_attention_local_info(True, 999, 1000, 500, 0) # 1.26μs -> 929ns (35.2% faster)

def test_large_scale_moe_dense_tp_size_500():
    # Large moe_dense_tp_size
    # tp_size=1000, dp_size=250, tp_rank=750, moe_dense_tp_size=500
    # local_tp_size = 500
    # local_tp_rank = 750 % 500 = 250
    # local_dp_size = max(1, 250 // (1000//500)) = max(1, 250 // 2) = 125
    # local_attn_tp_size = 500 // 125 = 4
    # local_attn_dp_rank = 250 // 4 = 62
    # local_attn_tp_rank = 250 % 4 = 2
    codeflash_output = compute_dp_attention_local_info(True, 750, 1000, 250, 500) # 1.01μs -> 754ns (34.5% faster)

def test_large_scale_all_ranks():
    # Test all tp_rank in [0, 999] for large tp_size
    results = []
    for tp_rank in range(1000):
        codeflash_output = compute_dp_attention_local_info(True, tp_rank, 1000, 100, 0); result = codeflash_output # 282μs -> 230μs (22.3% faster)
        results.append(result)
    # Ensure all possible local_attn_tp_rank and local_attn_dp_rank values are present
    attn_tp_ranks = set(r[0] for r in results)
    attn_dp_ranks = set(r[2] for r in results)

def test_large_scale_dp_size_max():
    # dp_size at max, tp_size at max
    # tp_size=1000, dp_size=1000, tp_rank=999, moe_dense_tp_size=0
    # local_tp_size = 1000
    # local_tp_rank = 999
    # local_dp_size = max(1, 1000 // (1000//1000)) = 1000
    # local_attn_tp_size = 1000 // 1000 = 1
    # local_attn_dp_rank = 999 // 1 = 999
    # local_attn_tp_rank = 999 % 1 = 0
    codeflash_output = compute_dp_attention_local_info(True, 999, 1000, 1000, 0) # 1.25μs -> 900ns (39.2% faster)

def test_large_scale_moe_dense_tp_size_max():
    # moe_dense_tp_size at max, tp_size at max
    # tp_size=1000, moe_dense_tp_size=1000, dp_size=1000, tp_rank=500
    # local_tp_size = 1000
    # local_tp_rank = 500
    # local_dp_size = max(1, 1000 // (1000//1000)) = 1000
    # local_attn_tp_size = 1000 // 1000 = 1
    # local_attn_dp_rank = 500 // 1 = 500
    # local_attn_tp_rank = 500 % 1 = 0
    codeflash_output = compute_dp_attention_local_info(True, 500, 1000, 1000, 1000) # 1.11μs -> 812ns (36.8% 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-compute_dp_attention_local_info-mholhstj and push.