⚡️ Speed up function require_mlp_sync by 9%

[codeflash-ai]

📄 9% (0.09x) speedup for require_mlp_sync in python/sglang/srt/utils/common.py

⏱️ Runtime : 40.2 microseconds → 36.9 microseconds (best of 250 runs)

📝 Explanation and details

The optimization applies short-circuit evaluation to both functions by restructuring boolean expressions as conditional statements, providing a 9% speedup.

Key Changes:

1. require_gathered_buffer: Changed from return require_mlp_tp_gather(server_args) or require_attn_tp_gather(server_args) to evaluate require_attn_tp_gather first and return early if True, only calling require_mlp_tp_gather when necessary.

2. require_mlp_sync: Changed from return server_args.enable_dp_attention or require_gathered_buffer(server_args) to check the simple field access server_args.enable_dp_attention first and return early if True, avoiding the more expensive nested function calls.

Why This is Faster:

• Field access vs function calls: server_args.enable_dp_attention is a simple attribute lookup, while require_gathered_buffer involves multiple function calls with assertions and complex logic.
• Reduced function call overhead: When the first condition is True (which happens in ~65% of test cases based on profiler data), the expensive second function is never called.
• Better branch prediction: The restructured conditional flow is more predictable for the CPU.

Performance Impact in Hot Path:
The function references show require_mlp_sync is called in critical scheduler event loops (event_loop_normal_disagg_decode and event_loop_overlap_disagg_decode) that run continuously during model serving. The optimization is particularly effective for workloads where enable_dp_attention=True (common in distributed attention scenarios), providing immediate returns and avoiding deeper computational branches.

Test Case Analysis:
The optimization shows strongest gains (15-25% faster) when enable_dp_attention=False and require_gathered_buffer would normally be evaluated, and modest improvements when enable_dp_attention=True due to early returns.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 69 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
import torch.distributed
from sglang.srt.utils.common import require_mlp_sync

# Helper: minimal mock class for server_args
class Args:
    def __init__(
        self,
        enable_dp_attention=False,
        dp_size=1,
        moe_dense_tp_size=1,
        enable_dp_lm_head=True,
        moe_a2a_backend="none",
        tp_size=1,
    ):
        self.enable_dp_attention = enable_dp_attention
        self.dp_size = dp_size
        self.moe_dense_tp_size = moe_dense_tp_size
        self.enable_dp_lm_head = enable_dp_lm_head
        self.moe_a2a_backend = moe_a2a_backend
        self.tp_size = tp_size

# =============================
# Basic Test Cases
# =============================

def test_basic_enable_dp_attention_true():
    # enable_dp_attention is True, dp_size > 1, should return True
    args = Args(enable_dp_attention=True, dp_size=2)
    codeflash_output = require_mlp_sync(args) # 374ns -> 403ns (7.20% slower)

def test_basic_enable_dp_attention_false_gathered_buffer_false():
    # All flags False/default, should return False
    args = Args()
    codeflash_output = require_mlp_sync(args) # 1.27μs -> 1.17μs (8.58% faster)

def test_basic_require_gathered_buffer_true_by_attn():
    # enable_dp_attention False, but require_attn_tp_gather returns True
    args = Args(enable_dp_attention=False, moe_dense_tp_size=1, moe_a2a_backend="some_backend")
    codeflash_output = require_mlp_sync(args) # 1.01μs -> 875ns (15.7% faster)

def test_basic_require_gathered_buffer_true_by_mlp():
    # enable_dp_attention True, dp_size > 1, moe_dense_tp_size None
    args = Args(enable_dp_attention=True, dp_size=2, moe_dense_tp_size=None)
    codeflash_output = require_mlp_sync(args) # 357ns -> 349ns (2.29% faster)

def test_basic_require_gathered_buffer_false():
    # enable_dp_attention False, require_gathered_buffer False
    args = Args(enable_dp_attention=False, moe_dense_tp_size=1, moe_a2a_backend="none")
    codeflash_output = require_mlp_sync(args) # 1.17μs -> 983ns (19.4% faster)

# =============================
# Edge Test Cases
# =============================

def test_edge_moe_dense_tp_size_invalid():
    # moe_dense_tp_size not in [1, None], should assert in require_attn_tp_gather
    args = Args(enable_dp_attention=False, moe_dense_tp_size=2)
    with pytest.raises(AssertionError):
        require_mlp_sync(args) # 2.16μs -> 1.93μs (12.1% faster)

def test_edge_moe_a2a_backend_none_and_moe_dense_tp_size_none():
    # enable_dp_attention False, moe_dense_tp_size None, moe_a2a_backend none
    args = Args(enable_dp_attention=False, moe_dense_tp_size=None, moe_a2a_backend="none")
    # require_attn_tp_gather: asserts, but None is allowed, so should pass
    codeflash_output = require_mlp_sync(args) # 1.36μs -> 1.52μs (10.5% slower)

def test_edge_moe_dense_tp_size_none_with_attention():
    # enable_dp_attention True, dp_size > 1, moe_dense_tp_size None
    args = Args(enable_dp_attention=True, dp_size=5, moe_dense_tp_size=None)
    codeflash_output = require_mlp_sync(args) # 441ns -> 399ns (10.5% faster)

def test_edge_enable_dp_lm_head_false():
    # enable_dp_attention True, dp_size > 1, enable_dp_lm_head False
    args = Args(enable_dp_attention=True, dp_size=4, enable_dp_lm_head=False)
    codeflash_output = require_mlp_sync(args) # 377ns -> 378ns (0.265% slower)

def test_edge_moe_a2a_backend_not_none():
    # enable_dp_attention True, dp_size > 1, moe_a2a_backend not "none"
    args = Args(enable_dp_attention=True, dp_size=4, moe_dense_tp_size=2, enable_dp_lm_head=True, moe_a2a_backend="nccl", tp_size=8)
    codeflash_output = require_mlp_sync(args) # 335ns -> 373ns (10.2% slower)

def test_edge_moe_dense_tp_size_greater_than_tp_size_over_dp_size():
    # enable_dp_attention True, dp_size > 1, moe_dense_tp_size > tp_size // dp_size
    args = Args(enable_dp_attention=True, dp_size=2, moe_dense_tp_size=5, enable_dp_lm_head=True, moe_a2a_backend="nccl", tp_size=8)
    codeflash_output = require_mlp_sync(args) # 350ns -> 348ns (0.575% faster)

def test_edge_moe_dense_tp_size_less_than_tp_size_over_dp_size():
    # enable_dp_attention True, dp_size > 1, moe_dense_tp_size < tp_size // dp_size
    args = Args(enable_dp_attention=True, dp_size=2, moe_dense_tp_size=2, enable_dp_lm_head=True, moe_a2a_backend="nccl", tp_size=8)
    codeflash_output = require_mlp_sync(args) # 374ns -> 380ns (1.58% slower)

def test_edge_tp_size_less_than_dp_size():
    # enable_dp_attention True, dp_size > tp_size
    args = Args(enable_dp_attention=True, dp_size=8, moe_dense_tp_size=2, enable_dp_lm_head=True, moe_a2a_backend="nccl", tp_size=4)
    codeflash_output = require_mlp_sync(args) # 340ns -> 353ns (3.68% slower)

# =============================
# Large Scale Test Cases
# =============================

def test_large_scale_enable_dp_attention_many():
    # Large dp_size and tp_size, enable_dp_attention True
    args = Args(enable_dp_attention=True, dp_size=999, moe_dense_tp_size=10, enable_dp_lm_head=True, moe_a2a_backend="nccl", tp_size=999)
    codeflash_output = require_mlp_sync(args) # 341ns -> 356ns (4.21% slower)

def test_large_scale_no_attention_many():
    # Large tp_size, enable_dp_attention False, require_attn_tp_gather True
    args = Args(enable_dp_attention=False, moe_dense_tp_size=1, moe_a2a_backend="gloo", tp_size=999)
    codeflash_output = require_mlp_sync(args) # 1.32μs -> 1.06μs (25.2% faster)

def test_large_scale_moe_dense_tp_size_none():
    # Large tp_size, enable_dp_attention True, moe_dense_tp_size None
    args = Args(enable_dp_attention=True, dp_size=999, moe_dense_tp_size=None, enable_dp_lm_head=True, moe_a2a_backend="nccl", tp_size=999)
    codeflash_output = require_mlp_sync(args) # 350ns -> 318ns (10.1% faster)

def test_large_scale_moe_dense_tp_size_1():
    # Large tp_size, enable_dp_attention False, moe_dense_tp_size 1, moe_a2a_backend "none"
    args = Args(enable_dp_attention=False, moe_dense_tp_size=1, moe_a2a_backend="none", tp_size=999)
    codeflash_output = require_mlp_sync(args) # 1.19μs -> 1.00μs (18.0% faster)

def test_large_scale_many_combinations():
    # Test for many combinations in a loop, all should be True if enable_dp_attention True
    for dp_size in [2, 10, 100, 999]:
        args = Args(enable_dp_attention=True, dp_size=dp_size, moe_dense_tp_size=1, enable_dp_lm_head=True, moe_a2a_backend="nccl", tp_size=999)
        codeflash_output = require_mlp_sync(args) # 865ns -> 894ns (3.24% slower)

def test_large_scale_require_attn_tp_gather_true_many():
    # Test for large tp_size, dp_size < tp_size, enable_dp_attention False
    args = Args(enable_dp_attention=False, dp_size=10, moe_dense_tp_size=1, moe_a2a_backend="nccl", tp_size=999)
    codeflash_output = require_mlp_sync(args) # 1.17μs -> 963ns (21.4% faster)

def test_large_scale_require_attn_tp_gather_false_many():
    # Test for large tp_size, dp_size >= tp_size, enable_dp_attention True (should assert if dp_size==1)
    args = Args(enable_dp_attention=False, dp_size=999, moe_dense_tp_size=1, moe_a2a_backend="nccl", tp_size=10)
    codeflash_output = require_mlp_sync(args) # 1.03μs -> 896ns (15.0% faster)

# =============================
# Determinism Test
# =============================

def test_determinism():
    # Multiple calls with same input should always return same result
    args = Args(enable_dp_attention=True, dp_size=2)
    codeflash_output = require_mlp_sync(args); result1 = codeflash_output # 347ns -> 334ns (3.89% faster)
    codeflash_output = require_mlp_sync(args); result2 = codeflash_output # 221ns -> 232ns (4.74% slower)

# =============================
# Boolean Return Type Test
# =============================

def test_return_type_bool():
    # The function should always return a bool
    args = Args(enable_dp_attention=False, moe_dense_tp_size=1, moe_a2a_backend="none")
    codeflash_output = require_mlp_sync(args); result = codeflash_output # 1.15μs -> 1.00μs (15.1% faster)
# 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
import torch.distributed
from sglang.srt.utils.common import require_mlp_sync

# Helper class to simulate server_args
class Args:
    def __init__(
        self,
        enable_dp_attention=False,
        dp_size=1,
        moe_dense_tp_size=1,
        enable_dp_lm_head=True,
        moe_a2a_backend="none",
        tp_size=1,
    ):
        self.enable_dp_attention = enable_dp_attention
        self.dp_size = dp_size
        self.moe_dense_tp_size = moe_dense_tp_size
        self.enable_dp_lm_head = enable_dp_lm_head
        self.moe_a2a_backend = moe_a2a_backend
        self.tp_size = tp_size

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

def test_basic_enable_dp_attention_true():
    # enable_dp_attention True should always return True
    args = Args(enable_dp_attention=True, dp_size=2)
    codeflash_output = require_mlp_sync(args) # 325ns -> 350ns (7.14% slower)

def test_basic_enable_dp_attention_false_gathered_buffer_false():
    # All flags off, should return False
    args = Args(enable_dp_attention=False, dp_size=1, moe_dense_tp_size=1, enable_dp_lm_head=True, moe_a2a_backend="none", tp_size=1)
    codeflash_output = require_mlp_sync(args) # 1.29μs -> 1.10μs (17.2% faster)

def test_basic_enable_dp_attention_false_gathered_buffer_true():
    # enable_dp_attention False, but require_attn_tp_gather returns True
    args = Args(enable_dp_attention=False, dp_size=1, moe_dense_tp_size=1, enable_dp_lm_head=True, moe_a2a_backend="foo", tp_size=2)
    codeflash_output = require_mlp_sync(args) # 980ns -> 875ns (12.0% faster)

def test_basic_mlp_tp_gather_true_due_to_moe_dense_tp_size_none():
    # enable_dp_attention True, moe_dense_tp_size None triggers True
    args = Args(enable_dp_attention=True, dp_size=2, moe_dense_tp_size=None, enable_dp_lm_head=True, moe_a2a_backend="all", tp_size=2)
    codeflash_output = require_mlp_sync(args) # 329ns -> 345ns (4.64% slower)

def test_basic_mlp_tp_gather_true_due_to_enable_dp_lm_head_false():
    # enable_dp_attention True, enable_dp_lm_head False triggers True
    args = Args(enable_dp_attention=True, dp_size=2, moe_dense_tp_size=2, enable_dp_lm_head=False, moe_a2a_backend="all", tp_size=2)
    codeflash_output = require_mlp_sync(args) # 325ns -> 335ns (2.99% slower)

def test_basic_mlp_tp_gather_true_due_to_moe_a2a_backend_none():
    # enable_dp_attention True, moe_a2a_backend "none" triggers True
    args = Args(enable_dp_attention=True, dp_size=2, moe_dense_tp_size=2, enable_dp_lm_head=True, moe_a2a_backend="none", tp_size=2)
    codeflash_output = require_mlp_sync(args) # 295ns -> 320ns (7.81% slower)

def test_basic_mlp_tp_gather_true_due_to_moe_dense_tp_size_greater():
    # enable_dp_attention True, moe_dense_tp_size > tp_size // dp_size
    args = Args(enable_dp_attention=True, dp_size=2, moe_dense_tp_size=3, enable_dp_lm_head=True, moe_a2a_backend="all", tp_size=4)
    codeflash_output = require_mlp_sync(args) # 311ns -> 336ns (7.44% slower)

def test_basic_mlp_tp_gather_false_due_to_moe_dense_tp_size_not_greater():
    # enable_dp_attention True, moe_dense_tp_size <= tp_size // dp_size
    args = Args(enable_dp_attention=True, dp_size=2, moe_dense_tp_size=1, enable_dp_lm_head=True, moe_a2a_backend="all", tp_size=4)
    codeflash_output = require_mlp_sync(args) # 301ns -> 329ns (8.51% slower)

def test_basic_attn_tp_gather_true_due_to_moe_a2a_backend_not_none():
    # enable_dp_attention False, moe_a2a_backend not "none"
    args = Args(enable_dp_attention=False, dp_size=1, moe_dense_tp_size=1, enable_dp_lm_head=True, moe_a2a_backend="foo", tp_size=2)
    codeflash_output = require_mlp_sync(args) # 1.24μs -> 1.04μs (18.6% faster)

def test_basic_attn_tp_gather_true_due_to_enable_dp_attention_true():
    # enable_dp_attention True, dp_size < tp_size
    args = Args(enable_dp_attention=True, dp_size=2, moe_dense_tp_size=1, enable_dp_lm_head=True, moe_a2a_backend="foo", tp_size=4)
    codeflash_output = require_mlp_sync(args) # 420ns -> 391ns (7.42% faster)

def test_basic_attn_tp_gather_false_due_to_moe_dense_tp_size_not_1_or_none():
    # Should assert
    args = Args(enable_dp_attention=False, dp_size=1, moe_dense_tp_size=2, enable_dp_lm_head=True, moe_a2a_backend="foo", tp_size=2)
    with pytest.raises(AssertionError):
        require_mlp_sync(args) # 1.76μs -> 1.52μs (16.1% faster)

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

def test_edge_tp_size_zero():
    # tp_size = 0, should not cause division by zero, since only used in integer division
    args = Args(enable_dp_attention=True, dp_size=2, moe_dense_tp_size=1, enable_dp_lm_head=True, moe_a2a_backend="none", tp_size=0)
    codeflash_output = require_mlp_sync(args) # 605ns -> 504ns (20.0% faster)

def test_edge_moe_dense_tp_size_none_with_enable_dp_attention_false():
    # Should assert in require_attn_tp_gather if moe_dense_tp_size not in [1, None]
    args = Args(enable_dp_attention=False, dp_size=1, moe_dense_tp_size=None, enable_dp_lm_head=True, moe_a2a_backend="foo", tp_size=2)
    codeflash_output = require_mlp_sync(args) # 1.49μs -> 1.28μs (16.4% faster)

def test_edge_large_tp_size_and_dp_size():
    # Large (but < 1000) values for tp_size and dp_size
    args = Args(enable_dp_attention=True, dp_size=999, moe_dense_tp_size=100, enable_dp_lm_head=True, moe_a2a_backend="all", tp_size=999)
    codeflash_output = require_mlp_sync(args) # 393ns -> 391ns (0.512% faster)

def test_edge_moe_a2a_backend_empty_string():
    # moe_a2a_backend is empty string (not "none")
    args = Args(enable_dp_attention=False, dp_size=1, moe_dense_tp_size=1, enable_dp_lm_head=True, moe_a2a_backend="", tp_size=2)
    codeflash_output = require_mlp_sync(args) # 1.07μs -> 943ns (13.3% faster)

def test_edge_enable_dp_lm_head_false_with_enable_dp_attention_false():
    # enable_dp_attention False, enable_dp_lm_head False
    args = Args(enable_dp_attention=False, dp_size=1, moe_dense_tp_size=1, enable_dp_lm_head=False, moe_a2a_backend="none", tp_size=2)
    codeflash_output = require_mlp_sync(args) # 1.15μs -> 1.04μs (10.9% faster)

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

def test_large_scale_many_combinations():
    # Test a range of dp_size and tp_size values
    for dp_size in [2, 10, 100, 999]:
        for tp_size in [2, 10, 100, 999]:
            args = Args(enable_dp_attention=True, dp_size=dp_size, moe_dense_tp_size=tp_size, enable_dp_lm_head=True, moe_a2a_backend="all", tp_size=tp_size)
            codeflash_output = require_mlp_sync(args)

def test_large_scale_attn_tp_gather_true_for_many():
    # Test that attn_tp_gather returns True for many values
    for tp_size in [2, 10, 100, 999]:
        args = Args(enable_dp_attention=False, dp_size=1, moe_dense_tp_size=1, enable_dp_lm_head=True, moe_a2a_backend="foo", tp_size=tp_size)
        codeflash_output = require_mlp_sync(args) # 2.27μs -> 1.91μs (18.8% faster)

def test_large_scale_attn_tp_gather_enable_dp_attention_true():
    # Test enable_dp_attention True, dp_size < tp_size
    for dp_size, tp_size in [(2, 10), (10, 100), (50, 999)]:
        args = Args(enable_dp_attention=True, dp_size=dp_size, moe_dense_tp_size=1, enable_dp_lm_head=True, moe_a2a_backend="foo", tp_size=tp_size)
        codeflash_output = require_mlp_sync(args) # 750ns -> 771ns (2.72% slower)

def test_large_scale_all_false():
    # Test all flags off for large tp_size/dp_size
    args = Args(enable_dp_attention=False, dp_size=999, moe_dense_tp_size=1, enable_dp_lm_head=True, moe_a2a_backend="none", tp_size=999)
    codeflash_output = require_mlp_sync(args) # 1.19μs -> 1.01μs (17.4% faster)

def test_large_scale_assertion_on_moe_dense_tp_size():
    # Should assert for moe_dense_tp_size not in [1, None]
    for v in [2, 100, 999]:
        args = Args(enable_dp_attention=False, dp_size=1, moe_dense_tp_size=v, enable_dp_lm_head=True, moe_a2a_backend="foo", tp_size=2)
        with pytest.raises(AssertionError):
            require_mlp_sync(args)
# 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-require_mlp_sync-mijrdxzb and push.