minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Feature] Megatron LayerNorm sequence parallelism (--enable-layernorm-sp) (#30915)

Browse Source
This commit is contained in:
Martin Hua
2026-08-31 19:27:28 -07:00
committed by GitHub
parent 783af667fb
commit 562b661e0e
11 changed files with 623 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/arg_groups/layernorm_sp_hook.py
+63
View File
@@ -0,0 +1,63 @@
from __future__ import annotations
import logging
from typing import TYPE_CHECKING, Optional
if TYPE_CHECKING:
from sglang.srt.server_args import ServerArgs
from sglang.srt.arg_groups.overrides import model_config_of, resolving_view
logger = logging.getLogger(__name__)
def handle_layernorm_sp(server_args: ServerArgs) -> None:
"""Validate --enable-layernorm-sp against the resolved parallelism config.
Runs in the resolution pipeline rather than in the layers so a model that
never builds a LayerCommunicator rejects the flag instead of ignoring it.
"""
cfg = resolving_view(server_args)
if not cfg.enable_layernorm_sp:
return
architectures = model_config_of(server_args).hf_config.architectures
validate_layernorm_sp(
architecture=architectures[0] if architectures else None,
tp_size=cfg.tp_size,
enable_dp_attention=cfg.enable_dp_attention,
speculative_algorithm=cfg.speculative_algorithm,
)
def validate_layernorm_sp(
*,
architecture: Optional[str],
tp_size: int,
enable_dp_attention: bool,
speculative_algorithm: Optional[str],
) -> None:
"""Fail loud for unsupported / incompatible configs. Callers gate on the flag."""
from sglang.srt.layers.layernorm_sp import SP_SUPPORTED_ARCHITECTURES
if architecture not in SP_SUPPORTED_ARCHITECTURES:
raise ValueError(
"--enable-layernorm-sp is only supported for "
f"{sorted(SP_SUPPORTED_ARCHITECTURES)}; got {architecture}."
)
if tp_size <= 1:
raise ValueError(
"--enable-layernorm-sp requires tp_size > 1: there is no sequence to "
"shard across a single TP rank."
)
if enable_dp_attention:
raise ValueError(
"--enable-layernorm-sp is not compatible with --enable-dp-attention: "
"SP shards the sequence across the full TP group, which under DP "
"attention spans data-parallel groups holding different sequences."
)
if speculative_algorithm is not None:
raise ValueError(
"--enable-layernorm-sp is not compatible with speculative decoding "
"(EAGLE/EAGLE3): the captured aux hidden states would be "
"sequence-sharded."
)
python/sglang/srt/arg_groups/pipeline.py
+5
View File
@@ -318,6 +318,11 @@ def run_resolution_pipeline(server_args: Any) -> None:
handle_speculative_decoding(server_args)
# After the speculative hook so speculative_algorithm is final.
from sglang.srt.arg_groups.layernorm_sp_hook import handle_layernorm_sp
handle_layernorm_sp(server_args)
# Validate the CuteDSL A2A token budget now that num_tokens_per_req is final.
validate_cutedsl_a2a_token_budget(server_args)
python/sglang/srt/distributed/bootstrap.py
+5
View File
@@ -27,6 +27,7 @@ from sglang.srt.distributed.parallel_state import (
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.environ import envs
from sglang.srt.layers.dp_attention import initialize_dp_attention
from sglang.srt.layers.layernorm_sp import initialize_layernorm_sp
from sglang.srt.platforms import current_platform
from sglang.srt.runtime_context import (
get_disagg,
@@ -280,6 +281,10 @@ def _init_parallel_groups(
server_args=server_args,
model_config=model_config,
)
initialize_layernorm_sp(
server_args=server_args,
model_config=model_config,
)
if is_npu():
register_sgl_tp_rank(gpu_id)
python/sglang/srt/layers/communicator.py
+69 -1
View File
@@ -31,6 +31,7 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory,
)
from sglang.srt.environ import envs
from sglang.srt.layers import layernorm_sp
from sglang.srt.layers.attention.dsa.utils import (
dsa_use_prefill_cp,
is_dsa_enable_prefill_cp,
@@ -72,7 +73,7 @@ from sglang.srt.model_executor.cuda_graph_config import (
Phase,
check_cuda_graph_backend,
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.runtime_context import (
get_exec,
get_forward,
@@ -483,6 +484,7 @@ class LayerCommunicator:
force_layernorm_before_dp_gather: bool = False,
enable_fused_ar_quant: bool = False,
fused_ar_quant_keep_bf16: bool = False,
_is_sp_variant: bool = False,
):
self.layer_scatter_modes = layer_scatter_modes
self.input_layernorm = input_layernorm
@@ -503,6 +505,30 @@ class LayerCommunicator:
get_spec().speculative_algorithm
)
# Under LayerNorm SP the norm/residual run on the sequence shard with no
# collectives, so delegate to an all-SCATTERED sibling while the region is
# active. _is_sp_variant stops the sibling from building its own.
self._sp_variant: Optional[LayerCommunicator] = None
if not _is_sp_variant and layernorm_sp.layernorm_sp_enabled():
self._sp_variant = LayerCommunicator(
layer_scatter_modes=LayerScatterModes(
layer_input_mode=ScatterMode.SCATTERED,
attn_mode=ScatterMode.SCATTERED,
mlp_mode=ScatterMode.SCATTERED,
middle_residual_mode=ScatterMode.SCATTERED,
layer_output_mode=ScatterMode.SCATTERED,
),
input_layernorm=input_layernorm,
post_attention_layernorm=post_attention_layernorm,
allow_reduce_scatter=allow_reduce_scatter,
is_last_layer=is_last_layer,
qkv_latent_func=qkv_latent_func,
force_layernorm_before_dp_gather=force_layernorm_before_dp_gather,
enable_fused_ar_quant=enable_fused_ar_quant,
fused_ar_quant_keep_bf16=fused_ar_quant_keep_bf16,
_is_sp_variant=True,
)
def _post_init_communicate(self):
self._communicate_simple_fn = CommunicateSimpleFn.get_fn(
input_mode=self.layer_scatter_modes.layer_input_mode,
@@ -592,6 +618,24 @@ class LayerCommunicator:
quant_format: str = "",
post_residual_addition: Optional[torch.Tensor] = None,
):
# residual is None marks the first decoder layer, where the SP region
# opens: re-evaluated per forward so a crash mid-loop cannot leak into
# the next one.
if self._sp_variant is not None:
if residual is None:
get_forward().set(
"sp_active", forward_batch.forward_mode == ForwardMode.EXTEND
)
if get_forward().sp_active:
hidden_states = layernorm_sp.sp_entry_scatter(hidden_states)
if get_forward().sp_active:
return self._sp_variant.prepare_attn(
hidden_states,
residual,
forward_batch,
quant_format,
post_residual_addition,
)
if get_attn_tp_context().input_scattered:
hidden_states, residual = self._tp_reduce_scatter(
hidden_states,
@@ -788,6 +832,10 @@ class LayerCommunicator:
forward_batch: ForwardBatch,
cache=None,
):
if self._sp_variant is not None and get_forward().sp_active:
return self._sp_variant.prepare_mlp(
hidden_states, residual, forward_batch, cache
)
if cache is not None:
self._context.cache = cache
@@ -805,6 +853,10 @@ class LayerCommunicator:
residual: torch.Tensor,
forward_batch: ForwardBatch,
):
if self._sp_variant is not None and get_forward().sp_active:
return self._sp_variant.postprocess_layer(
hidden_states, residual, forward_batch
)
return self._communicate_summable_tensor_pair_fn(
hidden_states=hidden_states,
residual=residual,
@@ -1028,6 +1080,22 @@ class CommunicateWithAllReduceAndLayerNormFn:
):
return CommunicateWithAllReduceAndLayerNormFn._simple
if (
hidden_states_input_mode == ScatterMode.SCATTERED
and residual_input_mode == ScatterMode.SCATTERED
and hidden_states_output_mode == ScatterMode.SCATTERED
and residual_output_mode == ScatterMode.SCATTERED
):
# Megatron LayerNorm sequence parallelism (layers/layernorm_sp.py):
# activations stay sequence-sharded across the attn->mlp boundary, so
# there is nothing to gather or scatter here -- just the residual add
# plus LayerNorm on the local shard. The row-parallel o_proj already
# issued the reduce-scatter (g-bar) that the all-reduce would have
# done, and the g all-gather is fused into the next column-parallel
# linear. Distinct from the branch above because under pure TP
# attn_tp_size == tp_size > 1, so that gate does not fire.
return CommunicateWithAllReduceAndLayerNormFn._simple
if (
(hidden_states_input_mode == ScatterMode.TP_ATTN_FULL)
and (
python/sglang/srt/layers/layernorm_sp.py
+245
View File
@@ -0,0 +1,245 @@
# Copyright 2023-2024 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Megatron-style LayerNorm sequence parallelism (SP, arXiv:2205.05198).
Under pure tensor parallelism the row-parallel ``all_reduce`` is algebraically a
``reduce_scatter`` (g-bar) followed by an ``all_gather`` (g). Splitting it that
way lets the LayerNorm / residual regions run on sequence-sharded activations --
each rank holds 1/tp of the tokens -- which cuts the transient activation memory
of long-context prefill with no extra communication volume (all_reduce and
reduce_scatter+all_gather move the same bytes).
Everything SP lives here so the feature stays decoupled from model code and from
``dp_attention.py``:
- which models opt in (the Qwen3-dense allowlist) and config validation,
- the per-forward ``sp_active`` flag (a ForwardFlags bool) read at depth by the
participant linears and the ``LayerCommunicator``,
- the entry-scatter / exit-gather collectives, and
- the fused matmul + collective fast-paths for the participant linears.
SP runs for prefill (EXTEND) only and is off by default; with the flag off,
nothing in this module executes.
"""
from __future__ import annotations
from typing import Optional
import torch
from sglang.srt.distributed import get_tp_group
from sglang.srt.runtime_context import get_flags, get_forward
from sglang.srt.utils.common import ceil_align
# Architectures whose decoder layers route attention/MLP through
# ``LayerCommunicator`` with the standard participant linears, and for which SP
# has been validated. Other models reject --enable-layernorm-sp at construction.
# The mechanism is generic; extend the allowlist as families are validated.
SP_SUPPORTED_ARCHITECTURES = frozenset({"Qwen3ForCausalLM"})
def initialize_layernorm_sp(*, server_args, model_config) -> None:
"""Materialize ``flags.sp.enabled``; runs once per worker after distributed
setup, alongside ``initialize_dp_attention``."""
architectures = model_config.hf_config.architectures
get_flags().sp.enabled = bool(
server_args.enable_layernorm_sp
and architectures
and architectures[0] in SP_SUPPORTED_ARCHITECTURES
)
def layernorm_sp_enabled() -> bool:
return get_flags().sp.enabled
def runs_sp(forward_mode) -> bool:
"""Whether this forward runs SP: an enabled model, prefill only.
Code outside the CUDA-graph-captured region must use this and not
``get_forward().sp_active``: Python writes made inside that region do not
re-execute on graph replay, so the flag is stale there.
"""
from sglang.srt.model_executor.forward_batch_info import ForwardMode
return layernorm_sp_enabled() and forward_mode == ForwardMode.EXTEND
class _SPForwardState:
"""Real (unpadded) token count of the current SP forward.
An instance attribute, not a ForwardFlags int slot: this is read inside
torch.compile-traced linear code, where dynamo gives an attribute-source int
automatic-dynamic, while a dict-slot int recompiles per sequence length.
"""
num_tokens: int = 0
_sp_state = _SPForwardState()
def set_sp_num_tokens(num_tokens: int) -> None:
_sp_state.num_tokens = num_tokens
def sp_num_tokens() -> int:
return _sp_state.num_tokens
# --- entry scatter / exit gather (once per forward, at the boundary) ----------
def sp_entry_scatter(hidden_states: torch.Tensor) -> torch.Tensor:
"""Shard the replicated ``[M, h]`` hidden states along the token dim.
Pads M up to a multiple of tp_size; the padding rows are dropped by the exit
gather. The input is replicated across the TP group, so this is a local slice.
"""
num_tokens = hidden_states.shape[0]
set_sp_num_tokens(num_tokens)
tp_group = get_tp_group()
tp_size = tp_group.world_size
if tp_size == 1:
return hidden_states
padded = ceil_align(num_tokens, tp_size)
if padded != num_tokens:
hidden_states = torch.nn.functional.pad(
hidden_states, (0, 0, 0, padded - num_tokens)
)
return hidden_states.tensor_split(tp_size)[tp_group.rank_in_group].contiguous()
def sp_exit_gather(hidden_states: torch.Tensor, num_tokens: int) -> torch.Tensor:
"""g: all-gather the per-rank shards back to the full sequence along dim 0,
then narrow to ``num_tokens`` (dropping the entry-scatter padding)."""
tp_group = get_tp_group()
tp_size = tp_group.world_size
if tp_size == 1:
return hidden_states[:num_tokens]
hidden_states = hidden_states.contiguous()
output = hidden_states.new_empty(
(hidden_states.shape[0] * tp_size, *hidden_states.shape[1:])
)
tp_group.all_gather_into_tensor(output, hidden_states)
return output[:num_tokens]
def maybe_exit_gather(
*,
hidden_states: torch.Tensor,
hidden_states_before_norm: Optional[torch.Tensor],
input_ids: Optional[torch.Tensor],
forward_mode,
) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
"""Undo the sequence sharding before the LM head, and leave the region.
No-op unless this forward runs SP. The token count comes from ``input_ids``
and the predicate from ``runs_sp``, so this stays correct on CUDA graph
replay, where the writes made inside the captured region do not re-execute.
"""
if not runs_sp(forward_mode) or input_ids is None:
return hidden_states, hidden_states_before_norm
num_tokens = input_ids.shape[0]
hidden_states = sp_exit_gather(hidden_states, num_tokens=num_tokens)
if hidden_states_before_norm is not None:
hidden_states_before_norm = sp_exit_gather(
hidden_states_before_norm, num_tokens=num_tokens
)
get_forward().set("sp_active", False)
return hidden_states, hidden_states_before_norm
# --- fused matmul + collective fast-paths for the participant linears ---------
# Fused matmul+reduce-scatter (g-bar) and all-gather+matmul (g) overlap the
# collective with the GEMM. Availability is probed once at import; TP groups are
# registered for symmetric memory lazily by the fused ops on first use (the old
# enable_symm_mem_for_group is a deprecated no-op), so we only import the module
# to register the torch.ops.symm_mem namespace the probe checks. NVLink/NVSwitch.
try:
import torch.distributed._symmetric_memory # noqa: F401
_HAS_TORCH_SYMM_MEM_FUSED = hasattr(
torch.ops.symm_mem, "fused_matmul_reduce_scatter"
) and hasattr(torch.ops.symm_mem, "fused_all_gather_matmul")
except Exception:
_HAS_TORCH_SYMM_MEM_FUSED = False
def sp_fused_matmul_eligible(linear) -> bool:
"""Whether the torch symm_mem fused matmul+collective fast-path applies: the
ops are available and ``linear`` is unquantized, bias-free, bf16/fp16 (the
case the fused ops support). Depends only on static layer properties, so the
decision is identical across TP ranks.
"""
if not _HAS_TORCH_SYMM_MEM_FUSED or linear.bias is not None:
return False
from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod
if not isinstance(linear.quant_method, UnquantizedLinearMethod):
return False
return linear.weight.dtype in (torch.bfloat16, torch.float16)
def column_parallel_g_matmul(
linear, input_parallel: torch.Tensor, bias
) -> torch.Tensor:
"""Megatron SP g for a ColumnParallelLinear participant (qkv / gate_up).
The input is this rank's sequence shard ``[M_pad/tp, K]``; all-gather it back
to the full sequence, matmul, and narrow to the real token count (recorded at
the entry scatter). Uses the fused symm-mem kernel when eligible (all-gather +
GEMM in one shot), else a plain all-gather + matmul.
"""
num_tokens = sp_num_tokens()
if sp_fused_matmul_eligible(linear):
group_name = get_tp_group().device_group.group_name
_, mm_outputs = torch.ops.symm_mem.fused_all_gather_matmul(
input_parallel.contiguous(),
[linear.weight.t()],
gather_dim=0,
group_name=group_name,
)
return mm_outputs[0][:num_tokens]
gathered = sp_exit_gather(input_parallel, num_tokens=num_tokens)
return linear.quant_method.apply(linear, gathered, bias)
def row_parallel_gbar_matmul(linear, input_: torch.Tensor, bias) -> torch.Tensor:
"""Megatron SP g-bar for a RowParallelLinear participant (o_proj / down).
Computes ``input_ @ weight.T`` and reduce-scatters (sum) the result across the
TP group along dim 0, leaving this rank's ``[M_pad/tp, h]`` shard. The token
dim is padded to a multiple of tp_size (padding rows are zeros, dropped by the
exit gather). Uses the fused symm-mem kernel when eligible, else matmul + a
plain reduce-scatter.
"""
tp_size = linear.tp_size
x = input_.contiguous()
num_tokens = x.shape[0]
padded = ceil_align(num_tokens, tp_size)
if padded != num_tokens:
x = torch.nn.functional.pad(x, (0, 0, 0, padded - num_tokens))
if sp_fused_matmul_eligible(linear):
group_name = get_tp_group().device_group.group_name
return torch.ops.symm_mem.fused_matmul_reduce_scatter(
x,
linear.weight.t(),
"sum",
scatter_dim=0,
group_name=group_name,
)
full = linear.quant_method.apply(linear, x, bias)
output = full.new_empty((padded // tp_size, *full.shape[1:]))
get_tp_group().reduce_scatter_tensor(output, full)
return output
python/sglang/srt/layers/linear.py
+25 -1
View File
@@ -25,6 +25,7 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory,
)
from sglang.srt.environ import envs
from sglang.srt.layers import layernorm_sp
from sglang.srt.layers.dp_attention import (
is_allocation_symmetric,
)
@@ -39,7 +40,7 @@ from sglang.srt.layers.parameter import (
_ColumnvLLMParameter,
)
from sglang.srt.layers.utils import pad_or_narrow_weight
from sglang.srt.runtime_context import get_exec, get_parallel
from sglang.srt.runtime_context import get_exec, get_forward, get_parallel
from sglang.srt.utils import get_bool_env_var, is_cpu, is_hip, is_npu, set_weight_attrs
if TYPE_CHECKING:
@@ -492,6 +493,14 @@ class ColumnParallelLinear(LinearBase):
def forward(self, input_):
bias = self.bias if not self.skip_bias_add else None
# Megatron SP "g": the input is this rank's [M_pad/tp, K] sequence shard;
# all-gather to the full sequence and matmul. Participants (qkv/gate_up)
# have gather_output=False, so there is no output all-gather to reconcile.
if get_forward().sp_active and self.tp_size > 1:
output = layernorm_sp.column_parallel_g_matmul(self, input_, bias)
output_bias = self.bias if self.skip_bias_add else None
return output, output_bias
# Matrix multiply.
assert self.quant_method is not None
output_parallel = self.quant_method.apply(self, input_, bias)
@@ -1621,6 +1630,21 @@ class RowParallelLinear(LinearBase):
# Only fuse bias add into GEMM for rank 0 (this ensures that
# bias will not get added more than once in TP>1 case)
bias_ = None if (self.tp_rank > 0 or self.skip_bias_add) else self.bias
# Megatron SP "g-bar": reduce-scatter along the token dim instead of
# all-reduce, leaving the output sharded for the next SP LayerNorm region.
# Fires regardless of reduce_results: o_proj / down are built
# reduce_results=False, so under SP the linear owns the reduction.
if (
get_forward().sp_active
and self.tp_size > 1
and not skip_all_reduce
and output_tensor is None
):
output = layernorm_sp.row_parallel_gbar_matmul(self, input_parallel, bias_)
output_bias = self.bias if self.skip_bias_add else None
return output, output_bias
if self.use_dp_attention_reduce:
symm_ctx = use_symmetric_memory(get_parallel().attn_tp_group)
else:
python/sglang/srt/layers/logits_processor.py
+10
View File
@@ -28,6 +28,7 @@ from sglang.srt.beam_search.logits_capture import BeamLogitsCapture
from sglang.srt.distributed import get_tp_group
from sglang.srt.distributed.device_communicators import triton_symm_mem_ag
from sglang.srt.environ import envs
from sglang.srt.layers import layernorm_sp
from sglang.srt.layers.aux_hidden_states import (
AuxHiddenStates,
pack_aux_hidden_states,
@@ -439,6 +440,15 @@ class LogitsProcessor(nn.Module):
if _autotune_run_lm_head is False:
return LogitsProcessorOutput(next_token_logits=None)
# Under LayerNorm SP the decoder loop leaves these sequence-sharded; undo
# that before the LM head, which must not participate.
hidden_states, hidden_states_before_norm = layernorm_sp.maybe_exit_gather(
hidden_states=hidden_states,
hidden_states_before_norm=hidden_states_before_norm,
input_ids=input_ids,
forward_mode=logits_metadata.forward_mode,
)
# Multi-item scoring only for prefill-only requests with pre-computed indices.
if multi_item_delimiter_indices is not None and logits_metadata.is_prefill_only:
return self.compute_logprobs_for_multi_item_scoring(
python/sglang/srt/runtime_context.py
+14 -1
View File
@@ -522,6 +522,15 @@ class DpFlags(_FlagGroupBase):
buffer_device: Any = None
@dataclasses.dataclass
class SpFlags(_FlagGroupBase):
"""LayerNorm sequence-parallelism flags, materialized by
``initialize_layernorm_sp`` (after distributed setup; reads the model
config). See ``layers.layernorm_sp``."""
enabled: bool = False
@dataclasses.dataclass
class Flags(_FlagGroupBase):
"""Root of the runtime-flags tier.
@@ -529,12 +538,13 @@ class Flags(_FlagGroupBase):
Resolved configuration lives in the config bags below (projected from the
declarations at publish) — this tier only carries genuine runtime
state whose value is not a function of the configuration alone, grouped
by lifecycle (``capture``) or subsystem (``moe`` / ``dp``).
by lifecycle (``capture``) or subsystem (``moe`` / ``dp`` / ``sp``).
"""
capture: CaptureFlags = dataclasses.field(default_factory=CaptureFlags)
moe: MoeFlags = dataclasses.field(default_factory=MoeFlags)
dp: DpFlags = dataclasses.field(default_factory=DpFlags)
sp: SpFlags = dataclasses.field(default_factory=SpFlags)
@dataclasses.dataclass
@@ -606,6 +616,8 @@ class ForwardFlags:
"fuse_mlp_allreduce": False,
"mlp_reduce_scatter": False,
"flashinfer_trtllm_bypass": False,
# LayerNorm sequence parallelism region; see layers/layernorm_sp.py.
"sp_active": False,
}
# Read/written inside compiled graphs (vocab embedding, communicator,
@@ -620,6 +632,7 @@ class ForwardFlags:
"fuse_mlp_allreduce",
"mlp_reduce_scatter",
"flashinfer_trtllm_bypass",
"sp_active",
}
)
python/sglang/srt/server_args.py
+9
View File
@@ -1124,6 +1124,15 @@ class ServerArgs:
"Shard dense MLP weights across the attention TP group under DP attention.",
NS("parallel"),
] = False
enable_layernorm_sp: A[
bool,
"Enable Megatron-style sequence parallelism (arXiv:2205.05198) for the "
"LayerNorm/residual regions under pure tensor parallelism: the row-parallel "
"all-reduce becomes reduce-scatter + all-gather, so LayerNorm runs on "
"sequence-sharded activations with no extra communication volume. "
"Prefill only; Qwen3 dense; requires tp_size > 1 and NVLink/NVSwitch.",
NS("parallel"),
] = False
disable_attn_tp_gather: A[
bool,
"Disable scheduler-side attn_tp_gather (the upstream SP path "