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:

test/registered/4-gpu-models/test_nvidia_nemotron_3_super_nvfp4.py

@@ -12,7 +12,7 @@ from sglang.test.test_utils import (
     popen_launch_server,
 )
 
-register_cuda_ci(est_time=540, suite="nightly-4-gpu-b200", nightly=True)
+register_cuda_ci(est_time=810, suite="nightly-4-gpu-b200", nightly=True)
 
 NEMOTRON_3_SUPER_NVFP4_MODEL = "nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4"
 
@@ -29,6 +29,31 @@ NEMOTRON_3_SUPER_NVFP4_ARGS = [
     '{"enable_multithread_load": true, "num_threads": 17}',
 ]
 
+DP_ATTENTION_EP_ARGS = [
+    "--dp-size",
+    "4",
+    "--enable-dp-attention",
+    "--enable-dp-lm-head",
+    "--ep-size",
+    "4",
+    "--moe-a2a-backend",
+    "flashinfer",
+    "--moe-runner-backend",
+    "flashinfer_cutedsl",
+    "--mamba-full-memory-ratio",
+    "5.0",
+    "--mamba-radix-cache-strategy",
+    "extra_buffer",
+    "--attention-backend",
+    "trtllm_mha",
+    "--max-running-requests",
+    "1024",
+    "--mem-fraction-static",
+    "0.93",
+    "--max-prefill-tokens",
+    "8192",
+]
+
 MTP_ARGS = [
     "--speculative-algorithm",
     "EAGLE",
@@ -107,5 +132,32 @@ class TestNvidiaNemotron3SuperNVFP4MTP(CustomTestCase):
         _run_gsm8k(self)
 
 
+class TestNvidiaNemotron3SuperNVFP4DPAttentionEP(CustomTestCase):
+    """DP attention + EP with the FlashInfer one-sided A2A and CuteDSL MoE runner."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = NEMOTRON_3_SUPER_NVFP4_MODEL
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        with (
+            envs.SGLANG_ENABLE_ASYNC_ASSERT.override(0),
+            envs.SGLANG_FLASHINFER_NUM_MAX_DISPATCH_TOKENS_PER_RANK.override(4096),
+            envs.SGLANG_FLASHINFER_WORKSPACE_SIZE.override(1024 * 1024 * 1024),
+        ):
+            cls.process = popen_launch_server(
+                cls.model,
+                cls.base_url,
+                timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+                other_args=NEMOTRON_3_SUPER_NVFP4_ARGS + DP_ATTENTION_EP_ARGS,
+            )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_gsm8k(self):
+        _run_gsm8k(self)
+
+
 if __name__ == "__main__":
     unittest.main()