From a14c870886bc273706c25fe9cc06b2cbd9a80c20 Mon Sep 17 00:00:00 2001 From: Brayden Zhong Date: Wed, 5 Aug 2026 14:27:05 -0700 Subject: [PATCH] Fix broken Nemotron DP attention (#33123) Co-authored-by: Brayden Zhong --- python/sglang/srt/environ.py | 3 ++ .../layers/attention/mamba/mamba2_metadata.py | 2 +- .../layers/moe/token_dispatcher/flashinfer.py | 11 ++-- .../model_runner_components/layer_setup.py | 6 ++- python/sglang/srt/models/nemotron_h_mtp.py | 7 ++- python/sglang/srt/server_args.py | 5 +- .../test_nvidia_nemotron_3_super_nvfp4.py | 54 ++++++++++++++++++- 7 files changed, 74 insertions(+), 14 deletions(-) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index b9f3cc605..3ba3e809a 100644 --- a/python/sglang/srt/environ.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py b/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py index 42d8a0c97..2cb3dd4f4 100644 --- a/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py +++ b/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 diff --git a/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py b/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py index 670fb428c..35759b919 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py +++ b/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. diff --git a/python/sglang/srt/model_executor/model_runner_components/layer_setup.py b/python/sglang/srt/model_executor/model_runner_components/layer_setup.py index c6485abca..83a94a63b 100644 --- a/python/sglang/srt/model_executor/model_runner_components/layer_setup.py +++ b/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"): diff --git a/python/sglang/srt/models/nemotron_h_mtp.py b/python/sglang/srt/models/nemotron_h_mtp.py index b3cd79fad..3c375e90c 100644 --- a/python/sglang/srt/models/nemotron_h_mtp.py +++ b/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, ) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index fdd0d8599..a40d9607b 100644 --- a/python/sglang/srt/server_args.py +++ b/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: diff --git a/test/registered/4-gpu-models/test_nvidia_nemotron_3_super_nvfp4.py b/test/registered/4-gpu-models/test_nvidia_nemotron_3_super_nvfp4.py index 4737d0461..ba37a7ae9 100644 --- a/test/registered/4-gpu-models/test_nvidia_nemotron_3_super_nvfp4.py +++ b/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()