Fix broken Nemotron DP attention (#33123)

Co-authored-by: Brayden Zhong <brayden@radixark.ai>

python/sglang/srt/environ.py

@@ -686,6 +686,9 @@ class Envs:
     SGLANG_FLASHINFER_USE_PAGED = EnvBool(False)
     # Default to the pick from flashinfer
     SGLANG_FLASHINFER_WORKSPACE_SIZE = EnvInt(384 * 1024 * 1024)
+    # Per-rank dispatch capacity of the FlashInfer MoE A2A dispatcher. Unset
+    # means each call site keeps its own default.
+    SGLANG_FLASHINFER_NUM_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(None)
     # Enable NVFP4 per-token activation scaling path for FlashInfer TRT-LLM MoE.
     SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION = EnvBool(False)
     # Launch the TRT-LLM MoE grouped GEMMs with PDL only at or below this

python/sglang/srt/layers/attention/mamba/mamba2_metadata.py

@@ -240,7 +240,7 @@ class Mamba2Metadata(ForwardMetadata):
         batch_size = getattr(forward_batch, "_original_batch_size", None)
         if batch_size is None:
             batch_size = len(forward_batch.seq_lens)
-        num_decodes = batch_size - num_prefills
+        num_decodes = max(0, batch_size - num_prefills)
         context_lens_tensor = forward_batch.extend_prefix_lens
         assert context_lens_tensor is not None
         has_initial_states = context_lens_tensor > 0

python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py

@@ -29,7 +29,6 @@ from sglang.srt.layers.moe.topk import (
 from sglang.srt.layers.moe.utils import get_moe_runner_backend
 from sglang.srt.runtime_context import get_schedule, get_spec
 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
-from sglang.srt.utils import get_int_env_var
 
 try:
     from flashinfer import nvfp4_block_scale_interleave
@@ -125,9 +124,13 @@ class FlashinferDispatcher(BaseDispatcher):
         # (which warms up at batch_size = req_to_token_pool.size).
         cps = get_schedule().chunked_prefill_size
         default_max_tokens = max(cps if cps and cps > 0 else 4096, 4096)
-        self.max_num_tokens = get_int_env_var(
-            "SGLANG_FLASHINFER_NUM_MAX_DISPATCH_TOKENS_PER_RANK",
-            default_max_tokens,
+        configured_max_tokens = (
+            envs.SGLANG_FLASHINFER_NUM_MAX_DISPATCH_TOKENS_PER_RANK.get()
+        )
+        self.max_num_tokens = (
+            configured_max_tokens
+            if configured_max_tokens is not None
+            else default_max_tokens
         )
 
         # Calculate workspace size. For eagle mode, use the larger workspace size since nextn layer will be unquantized.

python/sglang/srt/model_executor/model_runner_components/layer_setup.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 from typing import TYPE_CHECKING, Any, NamedTuple
 
 import msgspec
+from torch import nn
 
 if TYPE_CHECKING:
     from sglang.srt.configs.model_config import ModelConfig
@@ -23,7 +24,10 @@ def compute_attention_and_moe_layers(layer_model: Any) -> AttentionAndMoeLayers:
     moe_fusions: list[Any] = []
     dsa_indexers: list[Any] = []
     mha_companion_layers: list[Any] = []
-    for layer in layer_model.layers:
+    layers = layer_model.layers
+    if isinstance(layers, nn.ModuleDict):
+        layers = layers.values()
+    for layer in layers:
         attn_layer = None
         mha_companion_layer = None
         if hasattr(layer, "self_attn"):

python/sglang/srt/models/nemotron_h_mtp.py

@@ -288,13 +288,14 @@ class NemotronHMultiTokenPredictor(nn.Module):
     def forward(
         self,
         input_ids: torch.Tensor,
-        hidden_states: torch.Tensor,
+        positions: torch.Tensor,
         forward_batch: ForwardBatch,
         inputs_embeds: torch.Tensor | None = None,
     ) -> torch.Tensor:
         if inputs_embeds is None:
             inputs_embeds = self.get_input_embeddings(input_ids)
 
+        hidden_states = forward_batch.spec_info.hidden_states
         residual = None
 
         for i in range(self.pattern_len):
@@ -352,11 +353,9 @@ class NemotronHForCausalLMMTP(NemotronHForCausalLM):
         input_embeds: torch.Tensor | None = None,
         **kwargs,
     ) -> torch.Tensor:
-        hidden_states = forward_batch.spec_info.hidden_states
-
         hidden_states = self.model(
             input_ids,
-            hidden_states,
+            positions,
             forward_batch,
             input_embeds,
         )

python/sglang/srt/server_args.py

@@ -67,7 +67,6 @@ from sglang.srt.utils.common import (
     get_device,
     get_device_memory_capacity,
     get_device_sm,
-    get_int_env_var,
     get_quantization_config,
     human_readable_int,
     is_blackwell_supported,
@@ -6625,8 +6624,8 @@ class ServerArgs:
         ):
             return
         required_tokens = self.cutedsl_moe_max_num_tokens()
-        max_dispatch_tokens_per_rank = get_int_env_var(
-            "SGLANG_FLASHINFER_NUM_MAX_DISPATCH_TOKENS_PER_RANK", 1024
+        max_dispatch_tokens_per_rank = (
+            envs.SGLANG_FLASHINFER_NUM_MAX_DISPATCH_TOKENS_PER_RANK.get() or 1024
         )
         max_cutedsl_tokens = max_dispatch_tokens_per_rank * view.ep_size
         if max_cutedsl_tokens < required_tokens: