diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 48b7bd163..e97ce2769 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -755,6 +755,7 @@ class ServerArgs: enable_deterministic_inference: bool = False rl_on_policy_target: Optional[str] = None enable_attn_tp_input_scattered: bool = False + disable_attn_tp_gather: bool = False gc_threshold: Optional[List[int]] = None # Context parallelism used in the long sequence prefill phase of DeepSeek v3.2 enable_dsa_prefill_context_parallel: bool = False @@ -6526,6 +6527,18 @@ class ServerArgs: action="store_true", help="Allow input of attention to be scattered when only using tensor parallelism, to reduce the computational load of operations such as qkv latent.", ) + parser.add_argument( + "--disable-attn-tp-gather", + action="store_true", + help="Disable scheduler-side attn_tp_gather (the upstream SP path " + "that pads num_tokens to attn_tp_size and pre-allocates a gathered " + "buffer). Use for models that manage SP scatter/gather at the " + "model level (e.g., perform their own all_gather/reduce_scatter " + "inside attention) and do not consume the upstream gathered_buffer. " + "Without this, the cuda graph runner pads num_tokens to attn_tp_size, " + "which can cause kernel autotuners to select wrong-sized variants " + "at small batches.", + ) parser.add_argument( "--enable-dsa-prefill-context-parallel", dest="enable_dsa_prefill_context_parallel", diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 31eb48e8f..c21ae98b1 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -3089,6 +3089,13 @@ def require_attn_tp_gather(server_args: ServerArgs): """ Check if the input of attention is scattered. """ + # Opt-out for models that manage SP scatter/gather at the model level + # and do not consume the upstream gathered_buffer. Without this, the + # cuda graph runner pads num_tokens to attn_tp_size, which can cause + # autotuners to pick suboptimal kernel variants at small batches. + if server_args.disable_attn_tp_gather: + return False + from sglang.srt.layers.moe.utils import get_moe_a2a_backend assert server_args.moe_dense_tp_size in [1, None]