config: a parallel size has one spelling; a patched scope declares its own (#36621)

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Cheng Wan
2026-08-27 12:56:42 -07:00
committed by GitHub
co-authored by Claude Opus 5
parent ca1d7ed8e6
commit fd40a331bf
62 changed files with 439 additions and 1313 deletions
@@ -109,9 +109,12 @@ def mixer2_gated_norm_tensor_parallel(
import sglang.srt.layers.attention.mamba.mixer2_rms_norm_gated as m2
# Force attn-TP rank through the context (the weight loader reads it via
# get_parallel().attn_tp_rank); avoids calling initialize_dp_attention.
with get_parallel().override(attn_tp_rank=local_rank):
# Force the TP topology through the context (the weight loader reads
# get_parallel().attn_tp_rank, Mixer2RMSNormGated reads tp_size / tp_rank);
# avoids calling initialize_dp_attention.
with get_parallel().override(
attn_tp_rank=local_rank, tp_size=world_size, tp_rank=local_rank
):
# create gated-norm with TP
mixer = m2.Mixer2RMSNormGated(
full_hidden_size=hidden_size,