From a984c78330a8a5085a997818c396353262870429 Mon Sep 17 00:00:00 2001 From: iridiumine <42236072+iridiumine@users.noreply.github.com> Date: Sat, 12 Sep 2026 11:59:23 +0800 Subject: [PATCH] [NPU] Support DFlash speculative decoding for MiMo-V2.5-Pro (mxfp4) (#37565) --- .../sglang/srt/arg_groups/speculative_hook.py | 6 +- python/sglang/srt/configs/model_config.py | 9 + .../npu/attention/ascend_backend.py | 78 ++++- .../npu/graph_runner/npu_graph_runner.py | 116 ++++++- python/sglang/srt/layers/quantization/fp8.py | 97 +++++- python/sglang/srt/model_loader/loader.py | 3 + python/sglang/srt/models/dflash.py | 28 +- python/sglang/srt/models/mimo_v2.py | 67 +++- .../srt/speculative/dflash_disaggregation.py | 2 +- python/sglang/srt/speculative/dflash_utils.py | 27 ++ .../srt/speculative/dflash_worker_v2.py | 324 +++++++++++++++--- python/sglang/srt/speculative/spec_info.py | 8 + .../sglang/test/ascend/test_ascend_utils.py | 6 + .../test_npu_mimo_v2_5_pro_fp4_dflash.py | 47 +++ .../test_dspark_stacked_ctx_kv_parity.py | 3 + 15 files changed, 741 insertions(+), 80 deletions(-) create mode 100644 test/manual/ascend/llm_models/test_npu_mimo_v2_5_pro_fp4_dflash.py diff --git a/python/sglang/srt/arg_groups/speculative_hook.py b/python/sglang/srt/arg_groups/speculative_hook.py index 7db89b2c6..954e78d24 100644 --- a/python/sglang/srt/arg_groups/speculative_hook.py +++ b/python/sglang/srt/arg_groups/speculative_hook.py @@ -206,9 +206,11 @@ def _handle_dflash(server_args: ServerArgs) -> None: "DFLASH speculative decoding only supports CUDA, NPU and XPU devices." ) - if resolved_view(server_args).enable_dp_attention: + # DFLASH + dp attention is validated on NPU only. + if cfg.enable_dp_attention and not cfg.device == "npu": raise ValueError( - "Currently DFLASH speculative decoding does not support dp attention." + "Currently DFLASH speculative decoding does not support dp " + "attention on non-NPU devices." ) if cfg.pp_size != 1: diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 4de15464b..5c481158c 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -510,6 +510,15 @@ class ModelConfig: if self.is_fp4_experts: logger.info("Detected mixed checkpoint layout: routed experts are MXFP4.") + # MiMo-V2 mxfp4 ckpts declare the routed-expert layout via store_dtype. + if ( + not self.is_fp4_experts + and _hf_arch(self.hf_config) in MIMO_V2_MODEL_ARCHS + and str(quantization_config.get("store_dtype") or "").lower() == "mxfp4" + ): + self.is_fp4_experts = True + logger.info("Detected MiMo-V2 mxfp4 routed-expert layout.") + # DSV4 mxfp4 layout applies only when the ckpt does not opt in above. if is_deepseek_v4(self.hf_config) and routed_experts_quant_method is None: self.is_fp4_experts = envs.SGLANG_DSV4_FP4_EXPERTS.get() diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py index 8f131c352..6626f89aa 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py @@ -476,12 +476,23 @@ class AscendAttnBackend(AttentionBackend): self.forward_metadata = ForwardMetadata() seq_lens_max = forward_batch.seq_lens.max() if forward_batch.forward_mode.is_target_verify(): - spec_tokens_per_req = int(forward_batch.spec_info.draft_token_num) - # Overlap scheduling can publish the CPU sequence length one step - # ahead of the device tensor. FIA consumes seq_lens_cpu below, so - # derive the block-table width from the same source. Otherwise a - # page-aligned request can expose KV_S=N while asking FIA for N+1. - seq_lens_max = forward_batch.seq_lens_cpu.max().item() + spec_tokens_per_req + if ( + forward_batch.spec_algorithm is not None + and forward_batch.spec_algorithm.is_dflash() + ): + # dflash_worker_v2 already publishes seq_lens_cpu as prefix + + # one verify block, which already covers the draft block. + seq_lens_max = forward_batch.seq_lens_cpu.max().item() + else: + # Overlap scheduling can publish the CPU sequence length one + # step ahead of the device tensor. FIA consumes seq_lens_cpu + # below, so derive the block-table width from the same source. + # Otherwise a page-aligned request can expose KV_S=N while + # asking FIA for N+1. + spec_tokens_per_req = int(forward_batch.spec_info.draft_token_num) + seq_lens_max = ( + forward_batch.seq_lens_cpu.max().item() + spec_tokens_per_req + ) elif ( forward_batch.forward_mode.is_decode_or_idle() and forward_batch.spec_info is not None @@ -522,6 +533,10 @@ class AscendAttnBackend(AttentionBackend): ).int() self.forward_metadata.seq_lens_cpu_int = forward_batch.seq_lens_cpu.int() + # In graph mode (see _init_cuda_graph_metadata) seq_lens_cpu_int stays + # None so forward_mtp binds seq_lens_cpu_list instead: graph.update can + # only rebind the Host-side IntArray when captured as a Python list. + if ( not forward_batch.forward_mode.is_draft_extend_v2() and not forward_batch.forward_mode.is_target_verify() @@ -531,8 +546,10 @@ class AscendAttnBackend(AttentionBackend): if forward_batch.forward_mode.is_target_verify(): spec_algorithm = forward_batch.spec_algorithm - if spec_algorithm is None or not spec_algorithm.is_dspark(): - self.forward_metadata.seq_lens_cpu_int += spec_tokens_per_req + if spec_algorithm is None or not spec_algorithm.is_dflash_family(): + self.forward_metadata.seq_lens_cpu_int += int( + forward_batch.spec_info.draft_token_num + ) elif ( forward_batch.forward_mode.is_decode_or_idle() and forward_batch.spec_info is not None @@ -680,6 +697,8 @@ class AscendAttnBackend(AttentionBackend): metadata.swa_out_cache_loc = self.cuda_graph_swa_out_cache_loc[:num_tokens] metadata.seq_lens_cpu_list = seq_lens.cpu().int().tolist() metadata.seq_lens = seq_lens + # Deliberately no seq_lens_cpu_int here: forward_mtp must bind + # seq_lens_cpu_list so graph.update can rebind it. if forward_mode.is_target_verify() or forward_mode.is_draft_extend_v2(): metadata.actual_seq_lengths_q = torch.arange( self.speculative_num_draft_tokens, @@ -775,8 +794,17 @@ class AscendAttnBackend(AttentionBackend): metadata.block_tables_swa[:bs, max_seq_pages:].fill_(0) metadata.block_tables_swa[bs:, :].fill_(0) - # Update SWA mask: True = masked out (don't attend), False = attend - seq_lens_int = seq_lens[:bs].int() + # Update SWA mask: True = masked out (don't attend), False = attend. + # DFlash verify seq_lens is prefix-only, so use seq_lens_cpu + # (= prefix + block_size) to keep draft KV inside the mask window. + if ( + forward_mode.is_target_verify() + and _is_dflash_verify(spec_info) + and seq_lens_cpu is not None + ): + seq_lens_int = seq_lens_cpu[:bs].int() + else: + seq_lens_int = seq_lens[:bs].int() starts = torch.clamp(seq_lens_int - self.sliding_window_size, min=0) indices = self.graph_metadata["swa_indices"] start_exp = starts.unsqueeze(1) @@ -796,6 +824,13 @@ class AscendAttnBackend(AttentionBackend): if forward_mode.is_target_verify(): seq_lens = seq_lens + self.speculative_num_draft_tokens + # For DFlash, seq_lens_cpu (= prefix + block_size) is the true KV + # length; other spec algorithms already added the draft tokens above. + if _is_dflash_verify(spec_info) and seq_lens_cpu is not None: + kv_lens = seq_lens_cpu[:bs] + else: + kv_lens = seq_lens[:bs] + metadata.seq_lens_cpu_list = kv_lens.cpu().int().tolist() elif forward_mode.is_decode_or_idle() and spec_info is not None: seq_lens = seq_lens + self.speculative_step_offset_npu metadata.seq_lens[:bs].copy_(seq_lens[:bs]) @@ -2111,13 +2146,26 @@ class AscendAttnBackend(AttentionBackend): if not self.graph_mode: num_token_padding = query.shape[0] query = query[: forward_batch.global_num_token_non_padded_cpu] + # Trim DP padding rows so actualSeqLengthsKv matches the + # operator's batchSize (TND layout); only target_verify has a + # uniform per-request width. + if forward_batch.forward_mode.is_target_verify(): + real_bs = query.shape[0] // self.speculative_num_draft_tokens if self.forward_metadata.seq_lens_cpu_int is None: + # Graph mode: bind the Python list, which graph.update can + # rebind (a captured CPU tensor would be baked as constant). actual_seq_lengths_kv = self.forward_metadata.seq_lens_cpu_list else: actual_seq_lengths_kv = ( self.forward_metadata.seq_lens_cpu_int.cpu().int().tolist() ) + if ( + not self.graph_mode + and forward_batch.forward_mode.is_target_verify() + and len(actual_seq_lengths_kv) > real_bs + ): + actual_seq_lengths_kv = actual_seq_lengths_kv[:real_bs] if forward_batch.forward_mode.is_draft_extend_v2(): extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu @@ -2127,6 +2175,7 @@ class AscendAttnBackend(AttentionBackend): ] actual_seq_lengths = np.array(extend_seq_lens_cpu).cumsum().tolist() else: + # Static across replays ([spec_draft, 2*spec_draft, ...]). actual_seq_lengths = np.arange( self.speculative_num_draft_tokens, self.speculative_num_draft_tokens + query.shape[0], @@ -2146,6 +2195,13 @@ class AscendAttnBackend(AttentionBackend): block_table = self.forward_metadata.block_tables_swa else: block_table = self.forward_metadata.block_tables + if ( + not self.graph_mode + and forward_batch.forward_mode.is_target_verify() + and block_table.shape[0] > real_bs + ): + # Drop DP padding rows (see real_bs comment above). + block_table = block_table[:real_bs] if layer.attn_type == AttentionType.ENCODER_ONLY: mask = None @@ -2182,7 +2238,7 @@ class AscendAttnBackend(AttentionBackend): query, k_cache, v_cache, - block_table=self.forward_metadata.block_tables, + block_table=block_table, block_size=self.page_size, num_heads=layer.tp_q_head_num, num_key_value_heads=layer.tp_k_head_num, diff --git a/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py b/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py index b82c973bc..d30393fe7 100644 --- a/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py +++ b/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py @@ -40,9 +40,14 @@ from sglang.srt.configs.model_config import ( is_deepseek_dsa, is_deepseek_v4, ) -from sglang.srt.distributed.parallel_state import GroupCoordinator +from sglang.srt.distributed.parallel_state import ( + GroupCoordinator, +) from sglang.srt.environ import envs from sglang.srt.model_executor.runner import DecodeCudaGraphRunner +from sglang.srt.model_executor.runner.decode_cuda_graph_runner import ( + build_replay_fb_view, +) from sglang.srt.utils import ( empty_context, get_bool_env_var, @@ -62,7 +67,12 @@ if TYPE_CHECKING: from sglang.srt.model_executor.model_runner import ModelRunner from sglang.srt.layers.logits_processor import LogitsProcessorOutput -from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors +from sglang.srt.model_executor.forward_batch_info import ( + ForwardBatch, + PPProxyTensors, + compute_local_num_token_non_padded_cpu, + enable_num_token_non_padded, +) @contextmanager @@ -109,6 +119,11 @@ class NPUGraphRunner(DecodeCudaGraphRunner): self.update_attr_name = None self.update_attr_type = None self.model_runner = model_runner + # DFLASH verify under dp attention replays through the generic DP + # graph machinery: the scheduler-level DP vote keeps all DP ranks on + # the same graph/eager decision, and load_batch pads every rank to + # the same global max bucket so the captured dp-gather geometry stays + # valid when per-rank batch sizes diverge. self._init_arch_map() self.use_fia = get_bool_env_var("ASCEND_USE_FIA", "False") self.if_use_v2 = any( @@ -133,6 +148,8 @@ class NPUGraphRunner(DecodeCudaGraphRunner): self.attr_type: Dict[str, Union[list, torch.Tensor]] = { AttentionArch.MLA: [], AttentionArch.MHA: torch.Tensor(), + # TARGET_VERIFY must use a Python list: graph.update can only + # rebind the Host-side IntArray when captured as a list. "TARGET_VERIFY": [], } @@ -215,6 +232,46 @@ class NPUGraphRunner(DecodeCudaGraphRunner): self.load_batch(forward_batch, pp_proxy_tensors) else: # In speculative decoding, these two fields are still needed. + # NPU skips the DFLASH verify pre-planning, so load_batch may + # never have recorded the padded batch shapes; recompute them on + # every batch (the verify batch size varies with concurrency). + raw_bs = forward_batch.batch_size + if self.require_mlp_tp_gather: + bs = self._pad_to_bucket( + self._max_dp_batch_size(forward_batch), self.capture_bs + ) + else: + bs = self._pad_to_bucket(raw_bs, self.capture_bs) + self.raw_bs = raw_bs + self.raw_num_token = raw_bs * self.captured_req_width + self.bs = bs + # Restore the DeepEP dispatch mode recorded at capture time + # (mirrors load_batch); an interleaved eager extend may have + # switched it. + self.deepep_adapter.replay() + # Refresh the static DP token buffers bound by the captured + # graph (stale values misalign dp-gather segments across ranks); + # mirror the capture-side uniform [padded_num_tokens] * dp_size. + if self.require_mlp_tp_gather: + _padded_num_tokens = bs * self.captured_req_width + self.buffers.global_num_tokens_gpu.fill_(_padded_num_tokens) + self.buffers.global_num_tokens_for_logprob_gpu.fill_(_padded_num_tokens) + if ( + enable_num_token_non_padded() + and self.require_gathered_buffer + and not self.enable_prefill_cp + ): + self.buffers.num_token_non_padded.fill_( + compute_local_num_token_non_padded_cpu( + global_num_token_non_padded=( + forward_batch.global_num_token_non_padded_cpu + ), + num_tokens_per_dp=bs * self.captured_req_width, + sharded=self.model_runner.attn_tp_sequence_sharded( + bs * self.captured_req_width + ), + ) + ) self.buffers.input_ids[: self.raw_num_token].copy_(forward_batch.input_ids) self.buffers.positions[: self.raw_num_token].copy_(forward_batch.positions) if ( @@ -233,6 +290,42 @@ class NPUGraphRunner(DecodeCudaGraphRunner): forward_batch.mrope_positions ) + # The pre-planned path skipped init_forward_metadata_out_graph; + # refresh attention metadata so replay reads correct KV pages. + self.buffers.seq_lens[: self.raw_bs].copy_( + forward_batch.seq_lens_cpu[: self.raw_bs] + ) + self.buffers.seq_lens[self.raw_bs : self.bs].fill_(self.seq_len_fill_value) + self.buffers.seq_lens_cpu[: self.raw_bs].copy_( + forward_batch.seq_lens_cpu[: self.raw_bs] + ) + self.buffers.seq_lens_cpu[self.raw_bs : self.bs].fill_( + self.seq_len_fill_value + ) + self.buffers.req_pool_indices[: self.raw_bs].copy_( + forward_batch.req_pool_indices[: self.raw_bs] + ) + self.buffers.req_pool_indices[self.raw_bs : self.bs].fill_(0) + # Refresh the static out_cache_loc bound by the captured graph + # for full-pool KV writes in save_kv_cache (replay would + # otherwise write verify KV to stale capture-time slots). + if forward_batch.out_cache_loc is not None: + _padded_num_token = self.bs * self.captured_req_width + _n = min(self.raw_num_token, forward_batch.out_cache_loc.shape[0]) + self.buffers.out_cache_loc[:_n].copy_(forward_batch.out_cache_loc[:_n]) + self.buffers.out_cache_loc[_n:_padded_num_token].zero_() + fb_view = build_replay_fb_view( + forward_batch=forward_batch, + buffers=self.buffers, + bs=self.bs, + raw_bs=self.raw_bs, + num_tokens=self.bs * self.captured_req_width, + seq_len_fill_value=self.seq_len_fill_value, + capture_forward_mode=self.capture_forward_mode, + is_encoder_decoder=self.is_encoder_decoder, + ) + self._replay_attn_backend().init_forward_metadata_out_graph(fb_view) + graph_key = self._make_graph_key(self.bs) if not ( @@ -240,8 +333,23 @@ class NPUGraphRunner(DecodeCudaGraphRunner): or is_deepseek_v4(self.model_runner.model_config.hf_config) ): if forward_batch.forward_mode.is_target_verify(): - seq_lens_cpu = forward_batch.seq_lens.cpu() + self.captured_req_width - seq_lens = seq_lens_cpu.tolist() + [0] * (self.bs - self.raw_bs) + _attn = self._replay_attn_backend() + _meta = getattr(_attn, "forward_metadata", None) + _meta_list = getattr(_meta, "seq_lens_cpu_list", None) + if _meta_list is not None: + # graph.update must carry the exact KV length already + # computed in forward_metadata.seq_lens_cpu_list (it + # already includes the draft block for DFlash); do not + # recompute and double-add here. + seq_lens = list(_meta_list) + else: + # Wrapper backends (e.g. hybrid linear attention) keep + # forward_metadata only on their children, so it stays + # None here; fall back to the pre-DFlash computation. + seq_lens_cpu = ( + forward_batch.seq_lens.cpu() + self.captured_req_width + ) + seq_lens = seq_lens_cpu.tolist() + [0] * (self.bs - self.raw_bs) else: seq_lens = forward_batch.seq_lens.cpu().tolist() + [0] * ( self.bs - self.raw_bs diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 999500280..eaa25d6f3 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -261,6 +261,7 @@ class Fp8Config(QuantizationConfig): # model_loader from ModelConfig. Default False off the DSV4 path. self.is_fp4_experts = is_fp4_experts self.dequant_fp4_to_fp8 = False + self.is_dsv4_fp4_experts = False self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized if is_checkpoint_fp8_serialized: log_info_on_rank0(logger, "Detected fp8 checkpoint.") @@ -409,7 +410,14 @@ class Fp8Config(QuantizationConfig): ) return fp8_method - if self.is_fp4_experts and is_npu_arch35(): + if ( + self.is_fp4_experts + and is_npu_arch35() + # NPUW4A4Fp4MoEMethod is DSV4-specific (deepep dispatch, + # swiglu_limit); other FP4-expert checkpoints fall through + # to Fp8MoEMethod. + and self.is_dsv4_fp4_experts + ): from sglang.srt.hardware_backend.npu.quantization.fp4_moe_methods import ( NPUW4A4Fp4MoEMethod, ) @@ -505,6 +513,15 @@ class Fp8LinearMethod(LinearMethodBase): self.w8a8_mxfp8_linear = dispatch_w8a8_mxfp8_linear() else: self.w8a8_block_fp8_linear = dispatch_w8a8_block_fp8_linear() + if _is_npu and is_npu_arch35() and self.quant_config.scale_fmt != "ue8m0": + # The A5 backend expects the ue8m0 weight layout installed by + # the arch35 load path; keep plain block-FP8 checkpoints on + # the generic triton backend. + from sglang.srt.layers.quantization.fp8_utils import ( + triton_w8a8_block_fp8_linear, + ) + + self.w8a8_block_fp8_linear = triton_w8a8_block_fp8_linear self.is_checkpoint_fp8_serialized = ( self.quant_config.is_checkpoint_fp8_serialized ) @@ -708,7 +725,7 @@ class Fp8LinearMethod(LinearMethodBase): layer.weight_scale_inv.format_ue8m0 = True self._process_mxfp8_linear_weight_scale(layer) return - elif _is_npu and is_npu_arch35(): + elif _is_npu and is_npu_arch35() and self.quant_config.scale_fmt == "ue8m0": from sglang.srt.hardware_backend.npu.quantization.w8a8_mxfp8 import ( process_npu_arch35_mxfp8_linear_weights, ) @@ -1767,6 +1784,10 @@ class Fp8MoEMethod(FusedMoEMethodBase): logger.warning_once("Dequantized FP4 expert weights to FP8.") if self.is_fp4_expert: + if _is_npu: + self._process_npu_fp4_expert_weights(layer) + return + if get_moe_runner_backend().is_marlin(): layer.w13_weight.data = layer.w13_weight.data.view(torch.int8) layer.w2_weight.data = layer.w2_weight.data.view(torch.int8) @@ -2454,6 +2475,20 @@ class Fp8MoEMethod(FusedMoEMethodBase): ): self._owns_moe_runner = False self.moe_runner_config = moe_runner_config + + # MXFP4 experts on NPU run through the ASCEND runner, which + # quantizes activations internally. + if _is_npu and self.is_fp4_expert: + from sglang.srt.hardware_backend.npu.quantization.moe_methods import ( + NPUW4A8MXFP4MoEMethod, + ) + + layer.w13_kernel = NPUW4A8MXFP4MoEMethod() + layer.w2_kernel = NPUW4A8MXFP4MoEMethod() + moe_runner_config.layer = layer + self.runner = MoeRunner(MoeRunnerBackend.ASCEND, moe_runner_config) + return + moe_runner_backend = get_moe_runner_backend() if moe_runner_backend.is_auto(): @@ -2505,6 +2540,53 @@ class Fp8MoEMethod(FusedMoEMethodBase): block_shape=self.weight_block_size, ) + def _process_npu_fp4_expert_weights(self, layer: torch.nn.Module) -> None: + """Convert HF MXFP4 experts to the NPU W4A8 kernel layout. + + HF stores packed FP4 weights as int8 and block scales as float32; the + NPU kernel expects the weight in FRACTAL_NZ (transposed) and the scale + as e8m0 (uint8) reshaped to [E, K//64, N, 2]. + """ + from sglang.srt.hardware_backend.npu.quantization.linear_method_npu import ( + _get_float4_e2m1fn_x2_dtype, + ) + from sglang.srt.hardware_backend.npu.utils import npu_format_cast + + fp4_dtype = _get_float4_e2m1fn_x2_dtype() + if fp4_dtype is None: + raise RuntimeError("NPU W4A8 MXFP MoE requires float4 support.") + + for prefix in ("w13", "w2"): + weight = getattr(layer, f"{prefix}_weight") + weight.data = npu_format_cast( + weight.data.view(torch.uint8), + customize_dtype=torch.float8_e4m3fn, + input_dtype=fp4_dtype, + ).transpose(-1, -2) + + # Two e8m0 storage conventions in the fp32 scales: (a) the value + # IS the e8m0 byte (int in 0..255, MiMo-V2.5-Pro); (b) the value + # is 2^(e-127) -> recover e from the exponent field. Never + # re-encode with round(log2)+127: it re-biases the exponent. + scale_inv = getattr(layer, f"{prefix}_weight_scale_inv") + scale = scale_inv.data.to(torch.float32) + s_flat = scale.detach().float().flatten() + is_int_like = bool( + torch.all(s_flat >= 0) + and torch.all(s_flat <= 255) + and torch.allclose(s_flat, s_flat.round()) + ) + if is_int_like: + e8m0 = scale.to(torch.uint8) + else: + e8m0 = (scale.view(torch.int32) >> 23 & 0xFF).to(torch.uint8) + scale_inv.data = e8m0.reshape( + scale.shape[0], + scale.shape[1], + scale.shape[2] // 2, + 2, + ).transpose(1, 2) + def apply( self, layer: torch.nn.Module, @@ -2516,6 +2598,17 @@ class Fp8MoEMethod(FusedMoEMethodBase): x = dispatch_output.hidden_states moe_runner_config = self.moe_runner_config + if _is_npu and self.is_fp4_expert: + from sglang.srt.layers.moe.moe_runner.ascend import AscendQuantInfo + + quant_info = AscendQuantInfo( + w13_weight=layer.w13_weight, + w2_weight=layer.w2_weight, + w13_weight_scale=layer.w13_weight_scale_inv, + w2_weight_scale=layer.w2_weight_scale_inv, + ) + return self.runner.run(dispatch_output, quant_info) + if use_intel_amx_backend(layer): from sglang.srt.layers.moe.topk import apply_topk_weights_cpu diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index 13284e46d..6ac65b479 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -214,6 +214,9 @@ def _get_quantization_config( if isinstance(quant_config, Fp8Config): quant_config.is_fp4_experts = model_config.is_fp4_experts + from sglang.srt.configs.model_config import is_deepseek_v4 + + quant_config.is_dsv4_fp4_experts = is_deepseek_v4(model_config.hf_config) quant_config.dequant_fp4_to_fp8 = envs.SGLANG_DSV4_FP4_DEQUANT.get() # Handle hybrid NVFP4 moe (nvidia/DeepSeek-V4-Pro-NVFP4) nvfp4_meta = model_config.nvfp4_moe_meta diff --git a/python/sglang/srt/models/dflash.py b/python/sglang/srt/models/dflash.py index b2ed848da..7dc69d513 100644 --- a/python/sglang/srt/models/dflash.py +++ b/python/sglang/srt/models/dflash.py @@ -212,6 +212,7 @@ class DFlashAttention(nn.Module): base=rope_theta, rope_scaling=rope_scaling, is_neox_style=rope_is_neox_style, + partial_rotary_factor=float(getattr(config, "partial_rotary_factor", 1.0)), ) self.scaling = head_dim**-0.5 @@ -225,6 +226,8 @@ class DFlashAttention(nn.Module): self.sliding_window_size, self.attn_type = _get_dflash_layer_attention_params( config, layer_id ) + draft_cfg = parse_dflash_draft_config(draft_hf_config=config) + self.v_scale: Optional[float] = draft_cfg.attention_value_scale self.attention_sink_bias = None if is_nemotron_35_draft_config(config) and bool( getattr(config, "attention_sink_bias", False) @@ -243,6 +246,16 @@ class DFlashAttention(nn.Module): self.attention_sink_bias, {"weight_loader": sharded_weight_loader(0)}, ) + elif draft_cfg.attention_sink_bias: + # Per-head sink bias; each TP rank owns its slice of the + # all-heads checkpoint tensor. + self.attention_sink_bias = nn.Parameter( + torch.empty(self.num_heads, dtype=torch.float32), requires_grad=False + ) + set_weight_attrs( + self.attention_sink_bias, + {"weight_loader": sharded_weight_loader(0)}, + ) self.attn = RadixAttention( num_heads=self.num_heads, head_dim=head_dim, @@ -301,6 +314,8 @@ class DFlashAttention(nn.Module): q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) q, k = apply_qk_norm(q, k, self.q_norm, self.k_norm, self.head_dim) q, k = self.rotary_emb(positions, q, k) + if self.v_scale is not None: + v = v * self.v_scale if self.attention_sink_bias is None: attn_output = self.attn(q, k, v, forward_batch) else: @@ -334,11 +349,14 @@ class DFlashAttention(nn.Module): ) kv = F.linear(hidden_states, weight, bias) k, v = kv.split([self.kv_size, self.kv_size], dim=-1) - return k, v - - # Fallback: compute full QKV and discard Q (keeps compatibility with quantized weights). - qkv, _ = self.qkv_proj(hidden_states) - _, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + else: + # Fallback: compute full QKV and discard Q (keeps compatibility with quantized weights). + qkv, _ = self.qkv_proj(hidden_states) + _, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + # Keep V scaling consistent with forward() so ctx and self-generated V + # align in the draft cache. + if self.v_scale is not None: + v = v * self.v_scale return k, v def apply_k_norm(self, k: torch.Tensor) -> torch.Tensor: diff --git a/python/sglang/srt/models/mimo_v2.py b/python/sglang/srt/models/mimo_v2.py index eb02bca01..eb5e042ae 100644 --- a/python/sglang/srt/models/mimo_v2.py +++ b/python/sglang/srt/models/mimo_v2.py @@ -31,6 +31,10 @@ from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_r from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo from sglang.srt.layers.activation import SiluAndMul +from sglang.srt.layers.aux_hidden_states import ( + AuxHiddenStateAccumulator, + AuxHiddenStatePacker, +) from sglang.srt.layers.communicator import ( LayerCommunicator, LayerScatterModes, @@ -873,10 +877,16 @@ class MiMoV2DecoderLayer(nn.Module): hidden_states: torch.Tensor, forward_batch: ForwardBatch, residual: Optional[torch.Tensor], + captured_last_layer_outputs: Optional[AuxHiddenStateAccumulator] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: # Self Attention - hidden_states, residual = self.layer_communicator.prepare_attn( - hidden_states, residual, forward_batch + hidden_states, residual = ( + self.layer_communicator.prepare_attn_and_capture_last_layer_outputs( + hidden_states, + residual, + forward_batch, + captured_last_layer_outputs=captured_last_layer_outputs, + ) ) if hidden_states.shape[0] != 0: @@ -994,6 +1004,7 @@ class MiMoV2Model(nn.Module): self.padding_idx = getattr(config, "pad_token_id", None) self.vocab_size = config.vocab_size self.pp_group = get_pp_group() + self.layers_to_capture = [] if self.pp_group.is_first_rank: self.embed_tokens = VocabParallelEmbedding( @@ -1053,7 +1064,8 @@ class MiMoV2Model(nn.Module): hidden_states = pp_proxy_tensors["hidden_states"] residual = pp_proxy_tensors["residual"] - if forward_batch.can_run_tbo: + aux_hidden_states = AuxHiddenStatePacker(len(self.layers_to_capture)) + if forward_batch.can_run_tbo and not self.layers_to_capture: tbo_start_layer = self.start_layer tbo_end_layer = self.end_layer @@ -1088,8 +1100,22 @@ class MiMoV2Model(nn.Module): hidden_states, forward_batch, residual, + captured_last_layer_outputs=( + aux_hidden_states if i in self.layers_to_capture else None + ), ) + # A draft targeting the final layer ("after layer + # num_hidden_layers-1") maps to capture index num_hidden_layers, + # past the layer loop; capture the pre-norm output here instead. + if ( + self.pp_group.is_last_rank + and self.config.num_hidden_layers in self.layers_to_capture + ): + aux_hidden_states.append( + hidden_states if residual is None else hidden_states + residual + ) + hidden_states_before_norm = None if not self.pp_group.is_last_rank: return PPProxyTensors( @@ -1109,7 +1135,9 @@ class MiMoV2Model(nn.Module): else: hidden_states, _ = self.norm(hidden_states, residual) - return hidden_states, hidden_states_before_norm + if len(aux_hidden_states) == 0: + return hidden_states, hidden_states_before_norm + return hidden_states, hidden_states_before_norm, aux_hidden_states.finalize() # If this function is called, it should always initialize KV cache scale # factors (or else raise an exception). Thus, handled exceptions should @@ -1196,6 +1224,7 @@ class MiMoV2ForCausalLM(nn.Module, AudioEncoderMixin): self.logits_processor = ( LogitsProcessor(config) if not self.config.encoder_only else None ) + self.capture_aux_hidden_states = False vision_config = getattr(config, "vision_config", None) audio_config = getattr(config, "audio_config", None) @@ -1353,7 +1382,16 @@ class MiMoV2ForCausalLM(nn.Module, AudioEncoderMixin): "forward() should not be called in encoder_only mode" ) - if self._is_multimodal: + aux_hidden_states = None + if self.capture_aux_hidden_states: + hidden_states, hidden_states_before_norm, aux_hidden_states = self.model( + input_ids, + positions, + forward_batch, + input_embeds, + pp_proxy_tensors=pp_proxy_tensors, + ) + elif self._is_multimodal: hidden_states, hidden_states_before_norm = general_mm_embed_routine( input_ids=input_ids, forward_batch=forward_batch, @@ -1378,6 +1416,7 @@ class MiMoV2ForCausalLM(nn.Module, AudioEncoderMixin): self.lm_head, forward_batch, hidden_states_before_norm=hidden_states_before_norm, + aux_hidden_states=aux_hidden_states, ) else: return hidden_states @@ -1390,6 +1429,20 @@ class MiMoV2ForCausalLM(nn.Module, AudioEncoderMixin): def end_layer(self): return self.model.end_layer if self.model is not None else 0 + def set_dflash_layers_to_capture(self, layer_ids: List[int]): + if not self.pp_group.is_last_rank: + return + + if layer_ids is None: + raise ValueError( + "DFLASH requires explicit layer_ids for aux hidden capture." + ) + + self.capture_aux_hidden_states = True + # target_layer_ids are "after layer X" ids; capture before layer X+1, + # matching the draft's extract_context_feature (offset=1). + self.model.layers_to_capture = [val + 1 for val in layer_ids] + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): stacked_params_mapping = [ # (param_name, shard_name, shard_id) @@ -1582,6 +1635,10 @@ class MiMoV2ForCausalLM(nn.Module, AudioEncoderMixin): if weight_name not in name: continue name = name.replace(weight_name, param_name) + # mxfp4 ckpts store expert scales without the `_inv` suffix, + # while Fp8MoEMethod registers them as *_weight_scale_inv. + if name.endswith("weight_scale") and (name + "_inv" in params_dict): + name = name + "_inv" param = params_dict[name] weight_loader = param.weight_loader weight_loader( diff --git a/python/sglang/srt/speculative/dflash_disaggregation.py b/python/sglang/srt/speculative/dflash_disaggregation.py index 1cb2be7b2..1a672b810 100644 --- a/python/sglang/srt/speculative/dflash_disaggregation.py +++ b/python/sglang/srt/speculative/dflash_disaggregation.py @@ -13,7 +13,7 @@ if TYPE_CHECKING: from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2 -def build_dflash_family_disagg_draft_input( +def build_dflash_disagg_draft_input( batch: ScheduleBatch, last_tokens_tensor: torch.Tensor, future_map: FutureMap, diff --git a/python/sglang/srt/speculative/dflash_utils.py b/python/sglang/srt/speculative/dflash_utils.py index 49273b240..098c34c63 100644 --- a/python/sglang/srt/speculative/dflash_utils.py +++ b/python/sglang/srt/speculative/dflash_utils.py @@ -544,6 +544,8 @@ class DFlashDraftConfig: pure_draft_prefix_len: Optional[int] gru_hidden_dim: Optional[int] emb_dim: Optional[int] + attention_sink_bias: bool = False + attention_value_scale: Optional[float] = None @property def is_domino(self) -> bool: @@ -708,6 +710,29 @@ def parse_dflash_draft_config(*, draft_hf_config: Any) -> DFlashDraftConfig: f"got {mask_token_id}." ) + # MiMo DFlash draft extras: per-head attention sink bias and V value scale. + raw_attention_sink_bias = dflash_cfg.get("attention_sink_bias", False) + if not isinstance(raw_attention_sink_bias, bool): + raise ValueError( + "DFLASH dflash_config.attention_sink_bias must be a bool, " + f"got {raw_attention_sink_bias!r} (type={type(raw_attention_sink_bias).__name__})." + ) + attention_sink_bias = bool(raw_attention_sink_bias) + + raw_attention_value_scale = dflash_cfg.get("attention_value_scale", None) + if raw_attention_value_scale is None: + attention_value_scale: Optional[float] = None + else: + if isinstance(raw_attention_value_scale, bool) or not isinstance( + raw_attention_value_scale, (int, float) + ): + raise ValueError( + "DFLASH dflash_config.attention_value_scale must be int|float|None, " + f"got {raw_attention_value_scale!r} " + f"(type={type(raw_attention_value_scale).__name__})." + ) + attention_value_scale = float(raw_attention_value_scale) + projector_type = dflash_cfg.get( "projector_type", _cfg_get(draft_hf_config, "projector_type", None) ) @@ -792,6 +817,8 @@ def parse_dflash_draft_config(*, draft_hf_config: Any) -> DFlashDraftConfig: pure_draft_prefix_len=pure_draft_prefix_len, gru_hidden_dim=gru_hidden_dim, emb_dim=emb_dim, + attention_sink_bias=attention_sink_bias, + attention_value_scale=attention_value_scale, ) diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index 25f6186e7..28c5e42cd 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/python/sglang/srt/speculative/dflash_worker_v2.py @@ -1,9 +1,11 @@ import logging import math +import os from dataclasses import replace from typing import List, Optional, Tuple import torch +import torch.distributed as dist from sglang.kernels.ops.speculative.cache_locs import ( assign_extend_cache_locs_func, @@ -32,6 +34,7 @@ from sglang.srt.model_executor.forward_batch_info import ( ForwardBatch, ForwardMode, compute_position, + enable_num_token_non_padded, ) from sglang.srt.model_executor.model_runner import SamplingPrewarmResult from sglang.srt.model_executor.runner_utils.pool import ( @@ -40,6 +43,7 @@ from sglang.srt.model_executor.runner_utils.pool import ( ) from sglang.srt.runtime_context import ( get_exec, + get_parallel, get_schedule, get_spec, mamba_track_grid, @@ -81,8 +85,10 @@ from sglang.srt.speculative.spec_utils import ( GrammarTree, assign_req_to_token_pool_func, build_grammar_vocab_mask, + draft_tp_context, ) from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu +from sglang.srt.utils.common import empty_context _is_npu = is_npu() @@ -379,16 +385,33 @@ class DFlashWorkerV2(BaseSpecWorker): self._warned_sampling_fallback = False self._draft_probs_buf = None self._logged_first_verify = False - self._tp_sync = SpecTpSync(get_tp_group()) - - bundle = build_draft_tp_worker( - server_args=server_args, - gpu_id=gpu_id, - ps=replace(ps, pp_rank=0, pp_size=1), - nccl_port=nccl_port, - target_model_config=target_worker.model_runner.model_config, - algo_label="DFLASH", + self._full_embed_gpu: Optional[torch.Tensor] = None + # Under dp attention, peer DP ranks run different (idle) paths, so + # spec broadcasts must stay within the attn-TP group. + self._tp_sync = SpecTpSync( + get_parallel().attn_tp_group + if get_parallel().enable_dp_attention + else get_tp_group() ) + + # Under dp attention, the draft worker runs on the per-DP attn-TP + # group, independent of idle peer DP ranks. + self.draft_tp_context = ( + draft_tp_context if get_parallel().enable_dp_attention else empty_context + ) + if get_parallel().enable_dp_attention: + draft_init_ctx = draft_tp_context(get_parallel().attn_tp_group) + else: + draft_init_ctx = empty_context() + with draft_init_ctx: + bundle = build_draft_tp_worker( + server_args=server_args, + gpu_id=gpu_id, + ps=replace(ps, pp_rank=0, pp_size=1), + nccl_port=nccl_port, + target_model_config=target_worker.model_runner.model_config, + algo_label="DFLASH", + ) self._draft_worker = bundle.draft_worker self.draft_model_runner = bundle.draft_model_runner self._draft_sampler = None @@ -466,6 +489,8 @@ class DFlashWorkerV2(BaseSpecWorker): if hasattr(target_model, "get_dflash_noise_embedding_scale") else 1.0 ) + self._maybe_merge_trained_mask_embedding() + self._cache_full_embed_weight() if self.ps.tp_rank == 0: logger.info( "Initialized DFLASH draft runner. attention_backend=%s, model=%s, block_size=%s, draft_window_size=%s, compact_cache=%s", @@ -574,7 +599,8 @@ class DFlashWorkerV2(BaseSpecWorker): ) def init_attention_backends(self): - self._draft_worker.init_attention_backends() + with self.draft_tp_context(self.draft_model_runner.tp_group): + self._draft_worker.init_attention_backends() self._need_mamba_verify_commit = mambaish_config( self.model_runner.model_config ) is not None and hasattr( @@ -583,33 +609,44 @@ class DFlashWorkerV2(BaseSpecWorker): ) def init_cuda_graphs(self): - capture_decode_cuda_graph = ( - get_exec().graph.cuda_graph_config.decode.backend != Backend.DISABLED - ) - if is_cuda() and capture_decode_cuda_graph: - available_mem = self._tp_sync.available_memory_gb( - SpecTpSyncSite.DFLASH_MEM, - self.device, - self.gpu_id, - group=get_tp_group(), + with self.draft_tp_context(self.draft_model_runner.tp_group): + capture_decode_cuda_graph = ( + get_exec().graph.cuda_graph_config.decode.backend != Backend.DISABLED ) - if available_mem < 1.0: + if get_parallel().enable_dp_attention and capture_decode_cuda_graph: + # Idle DP ranks skip the draft step, so they cannot join a + # shared graph capture/replay; keep the draft eager under dp + # attention. capture_decode_cuda_graph = False - logger.warning( - "Disable DFLASH draft cuda graph because only %.2f GB GPU " - "memory is available after target backend initialization.", - available_mem, + if self.ps.tp_rank == 0: + logger.warning( + "Disable DFLASH draft cuda graph because dp attention " + "is enabled (draft runs eager)." + ) + if is_cuda() and capture_decode_cuda_graph: + available_mem = self._tp_sync.available_memory_gb( + SpecTpSyncSite.DFLASH_MEM, + self.device, + self.gpu_id, + group=get_tp_group(), ) - if capture_decode_cuda_graph: - # Must run before capture so the draft graph folds the head in. - self._draft_sampler = self._maybe_build_draft_sampler() - if self._draft_sampler is not None: - self.draft_model_runner.capture_tail_hooks.append( - make_draft_sampler_capture_hook(self._draft_sampler) - ) - self._draft_worker.init_cuda_graphs( - capture_decode_cuda_graph=capture_decode_cuda_graph - ) + if available_mem < 1.0: + capture_decode_cuda_graph = False + logger.warning( + "Disable DFLASH draft cuda graph because only %.2f GB GPU " + "memory is available after target backend initialization.", + available_mem, + ) + if capture_decode_cuda_graph: + # Must run before capture so the draft graph folds the head in. + self._draft_sampler = self._maybe_build_draft_sampler() + if self._draft_sampler is not None: + self.draft_model_runner.capture_tail_hooks.append( + make_draft_sampler_capture_hook(self._draft_sampler) + ) + self._draft_worker.init_cuda_graphs( + capture_decode_cuda_graph=capture_decode_cuda_graph + ) def _prewarm_batch_size(self, block_size: int) -> int: """Largest batch the non-greedy verify path can see in one step.""" @@ -1143,6 +1180,101 @@ class DFlashWorkerV2(BaseSpecWorker): bs, ) + def _maybe_merge_trained_mask_embedding(self) -> None: + """Merge a trained mask embedding (mask_embedding.pt) into the target table. + + The target's own mask-token row is typically an untrained added-vocab + slot; overwrite it with the learned vector. VocabParallelEmbedding-aware: + each rank only updates the row if the token falls within its shard. + """ + from sglang.srt.arg_groups.overrides import resolving_view + + draft_model_path = resolving_view(self.server_args).speculative_draft_model_path + if draft_model_path is None: + return + + mask_emb_path = os.path.join(draft_model_path, "mask_embedding.pt") + if not os.path.exists(mask_emb_path): + return + + saved = torch.load(mask_emb_path, map_location=self.device, weights_only=True) + embedding_tensor = saved["embedding"] + saved_token_id = int(saved["mask_token_id"]) + + if saved_token_id != self._mask_token_id: + raise ValueError( + f"DFLASH mask_embedding.pt was trained with mask_token_id={saved_token_id}, " + f"but the current resolved mask_token_id={self._mask_token_id}. " + "These must match." + ) + + target_model = self._target_worker.model_runner.model + embed_module = target_model.get_input_embeddings() + + if saved.get("per_position"): + raise NotImplementedError( + "DFLASH per-position mask embeddings are not supported in " + "inference yet (mask_embedding.pt per_position=True)." + ) + + token_id = self._mask_token_id + shard_indices = getattr(embed_module, "shard_indices", None) + with torch.no_grad(): + if shard_indices is not None: + # VocabParallelEmbedding: only update if this token is in our shard. + start = shard_indices.org_vocab_start_index + end = shard_indices.org_vocab_end_index + if start <= token_id < end: + local_idx = token_id - start + embed_module.weight[local_idx].copy_( + embedding_tensor.to(embed_module.weight.dtype) + ) + else: + embed_module.weight[token_id].copy_( + embedding_tensor.to(embed_module.weight.dtype) + ) + + if self.ps.tp_rank == 0: + logger.info( + "Merged trained mask embedding into target model " + "(mask_token_id=%s, source=%s)", + self._mask_token_id, + mask_emb_path, + ) + + def _cache_full_embed_weight(self) -> None: + """Cache the full target embedding replicated on GPU during init. + + Under dp attention, idle DP ranks skip the draft step, so the attn-TP + all_reduce inside VocabParallelEmbedding gets mismatched calls. Gather + the full embedding once during init and keep it replicated, making + the draft block-id lookup a collective-free on-device F.embedding. + """ + if not get_parallel().enable_dp_attention: + return + + tp_group = get_tp_group() + tp_size = int(tp_group.world_size) + if tp_size <= 1: + return + + target_model = self._target_worker.model_runner.model + embed_module = target_model.get_input_embeddings() + local_w = embed_module.weight.data + shard = getattr(embed_module, "shard_indices", None) + num_org = int(shard.num_org_elements) if shard else local_w.shape[0] + vocab_size = int(self._target_worker.model_runner.model_config.vocab_size) + + shard_t = local_w[:num_org].contiguous() + parts = [torch.empty_like(shard_t) for _ in range(tp_size)] + dist.all_gather(parts, shard_t, group=tp_group.device_group) + self._full_embed_gpu = torch.cat(parts, dim=0)[:vocab_size] + if self.ps.tp_rank == 0: + logger.info( + "DFLASH cached full embed on GPU for dp attention: shape=%s", + list(self._full_embed_gpu.shape), + ) + def _resolve_mask_token_id( self, *, mask_token: str, mask_token_id: Optional[int] = None ) -> int: @@ -1601,6 +1733,22 @@ class DFlashWorkerV2(BaseSpecWorker): f"DFLASH positions must be 1D, got shape={tuple(positions.shape)}." ) num_tokens = int(target_hidden.shape[0]) + # Drop trailing alignment-padding rows (from TP/EP-size rounding) + # before materializing into the draft KV; they would corrupt other + # requests' cache slots. + expected_tokens = int(cache_loc.numel()) + if num_tokens > expected_tokens: + if not getattr(self, "_logged_padding_trim", False): + logger.warning( + "DFLASH target_hidden has %d trailing padding row(s); trimming " + "to cache_loc length=%d (target_hidden=%d). Logged once per worker.", + num_tokens - expected_tokens, + expected_tokens, + num_tokens, + ) + self._logged_padding_trim = True + target_hidden = target_hidden[:expected_tokens] + num_tokens = expected_tokens if int(cache_loc.numel()) != num_tokens: raise ValueError( "DFLASH cache_loc length mismatch: " @@ -1650,7 +1798,10 @@ class DFlashWorkerV2(BaseSpecWorker): if commit_lens.dtype != torch.int32: commit_lens = commit_lens.to(torch.int32) - with torch.inference_mode(): + with ( + torch.inference_mode(), + self.draft_tp_context(self.draft_model_runner.tp_group), + ): ctx_hidden = self.draft_model.project_target_hidden(target_hidden) if cache_loc_2d is not None: @@ -1737,6 +1888,9 @@ class DFlashWorkerV2(BaseSpecWorker): ) if _is_npu: _, k, v = attn.forward_prepare_npu(ctx_positions, layer_ctx_hidden) + # Keep V scaling consistent with forward(). + if attn.v_scale is not None: + v = v * attn.v_scale else: k, v = attn.kv_proj_only(layer_ctx_hidden) k = attn.apply_k_norm(k) @@ -2054,6 +2208,15 @@ class DFlashWorkerV2(BaseSpecWorker): if on_publish is not None: on_publish(batch_output.new_seq_lens) + # An idle DP rank runs the empty target prefill above to stay in + # the DP collective, but must skip the draft KV materialization, + # which needs per-request extend info. + if batch.forward_mode.is_idle(): + batch_output.next_draft_input = DFlashDraftInputV2.create_idle_input( + device=self.device + ) + return batch_output + if logits_output.hidden_states is None: raise RuntimeError( "DFLASH requires target aux hidden capture for prefill, but got None. " @@ -2110,6 +2273,32 @@ class DFlashWorkerV2(BaseSpecWorker): ) if batch.forward_mode.is_idle(): + # Under dp attention an idle DP rank must still run the target + # verify forward (IDLE mode) so its cross-DP collectives stay in + # lockstep with the active DP group; the draft block's + # collectives are within-rank and skipped. + if get_parallel().enable_dp_attention: + idle_verify_input = DFlashVerifyInput( + draft_token=torch.empty((0,), dtype=torch.long, device=self.device), + positions=torch.empty((0,), dtype=torch.int64, device=self.device), + draft_token_num=int(self.block_size), + custom_mask=None, + capture_hidden_mode=CaptureHiddenMode.FULL, + ) + idle_verify_forward_batch, _ = idle_verify_input.prepare_for_verify( + batch, self._target_worker + ) + # Force eager: the active ranks run eager verify when any DP + # rank is idle (see the idle-guard below); a graph-replaying + # idle rank would disagree with them on DP-gather segment + # offsets (padded bucket vs raw counts). + idle_verify_forward_batch.can_run_decode_cuda_graph = False + self._target_worker.forward_batch_generation( + batch=None, + forward_batch=idle_verify_forward_batch, + is_verify=True, + skip_attn_backend_init=True, + ) empty_ids = torch.empty((0,), dtype=torch.int64, device=self.device) empty_lens = torch.empty((0,), dtype=torch.int32, device=self.device) next_draft_input = self._make_next_draft_input_decode( @@ -2221,7 +2410,14 @@ class DFlashWorkerV2(BaseSpecWorker): ) verify_out_cache_loc_2d.copy_(verify_out_cache_loc.view(bs, block_size)) - noise_embedding = embed_module(block_ids) + if self._full_embed_gpu is not None: + # Replicated lookup avoids the mismatched attn-TP all_reduce + # inside VocabParallelEmbedding under dp attention. + noise_embedding = torch.nn.functional.embedding( + block_ids, self._full_embed_gpu + ) + else: + noise_embedding = embed_module(block_ids) if self._noise_embed_scale != 1.0: noise_embedding = noise_embedding * self._noise_embed_scale input_embeds = noise_embedding.view(-1, noise_embedding.shape[-1]) @@ -2281,6 +2477,12 @@ class DFlashWorkerV2(BaseSpecWorker): spec_algorithm=SpeculativeAlgorithm.DFLASH, spec_info=self._draft_block_spec_info, capture_hidden_mode=CaptureHiddenMode.NULL, + num_token_non_padded=( + torch.tensor(bs * block_size, dtype=torch.int32, device=device) + if enable_num_token_non_padded() + else None + ), + global_num_token_non_padded_cpu=bs * block_size, ) if self.selector is not None: @@ -2291,7 +2493,10 @@ class DFlashWorkerV2(BaseSpecWorker): bs=bs, sampling_info=batch.sampling_info ) - with torch.inference_mode(): + with ( + torch.inference_mode(), + self.draft_tp_context(self.draft_model_runner.tp_group), + ): draft_out = self.draft_model_runner.forward(forward_batch) draft_logits_output = draft_out.logits_output @@ -2349,24 +2554,26 @@ class DFlashWorkerV2(BaseSpecWorker): self._draft_sampler.q_out[:bs], ) elif self.selector is not None: - draft_next = self._propose_selector_block( - draft_logits_output=draft_logits_output, - bs=bs, - lm_head=lm_head, - anchor_token_ids=block_ids[:, 0], - sampling_info=batch.sampling_info, - ) + with self.draft_tp_context(self.draft_model_runner.tp_group): + draft_next = self._propose_selector_block( + draft_logits_output=draft_logits_output, + bs=bs, + lm_head=lm_head, + anchor_token_ids=block_ids[:, 0], + sampling_info=batch.sampling_info, + ) else: draft_hidden = draft_logits_output.hidden_states if draft_hidden is None: raise RuntimeError("DFLASH draft model returned no hidden states.") draft_hidden = draft_hidden.view(bs, int(self.block_size), -1) - draft_next = self._greedy_sample_from_vocab_parallel_head( - hidden_states=draft_hidden[:, 1:, :].reshape( - -1, draft_hidden.shape[-1] - ), - lm_head=lm_head, - ).view(bs, int(self.block_size) - 1) + with self.draft_tp_context(self.draft_model_runner.tp_group): + draft_next = self._greedy_sample_from_vocab_parallel_head( + hidden_states=draft_hidden[:, 1:, :].reshape( + -1, draft_hidden.shape[-1] + ), + lm_head=lm_head, + ).view(bs, int(self.block_size) - 1) draft_tokens = self._draft_block_tokens_buf[:bs] draft_tokens[:, 0].copy_(block_ids[:, 0]) @@ -2413,6 +2620,23 @@ class DFlashWorkerV2(BaseSpecWorker): batch.seq_lens_cpu = seq_lens_cpu_backup batch.seq_lens_sum = seq_lens_sum_backup + # Idle-DP guard: an idle rank's eager verify contributes raw scheduler + # counts while a graph-replaying rank assumes the padded bucket, so + # their DP-gather segment offsets disagree. Fall back to eager verify + # whenever any DP rank is idle this round. + if ( + get_parallel().enable_dp_attention + and verify_forward_batch.original_global_num_tokens_cpu is not None + and min(verify_forward_batch.original_global_num_tokens_cpu) == 0 + ): + verify_forward_batch.can_run_decode_cuda_graph = False + + # Mixed-round guard: an extend rank's raw token counts disagree with + # the verify batch's spec-scaled counts on the DP-gather layout; run + # eager, symmetric with the idle guard above. + if get_parallel().enable_dp_attention and batch.is_extend_in_batch: + verify_forward_batch.can_run_decode_cuda_graph = False + target_out = self.target_worker.forward_batch_generation( batch=None, forward_batch=verify_forward_batch, diff --git a/python/sglang/srt/speculative/spec_info.py b/python/sglang/srt/speculative/spec_info.py index 11da6592a..0f3f598b2 100644 --- a/python/sglang/srt/speculative/spec_info.py +++ b/python/sglang/srt/speculative/spec_info.py @@ -211,6 +211,14 @@ class SpeculativeAlgorithm(Enum): return build_dspark_disagg_draft_input( batch, last_tokens_tensor, future_map ) + if self.is_dflash(): + from sglang.srt.speculative.dflash_disaggregation import ( + build_dflash_disagg_draft_input, + ) + + return build_dflash_disagg_draft_input( + batch, last_tokens_tensor, future_map + ) return None def need_topk(self) -> bool: diff --git a/python/sglang/test/ascend/test_ascend_utils.py b/python/sglang/test/ascend/test_ascend_utils.py index b7682a903..db19c7bca 100644 --- a/python/sglang/test/ascend/test_ascend_utils.py +++ b/python/sglang/test/ascend/test_ascend_utils.py @@ -149,6 +149,12 @@ META_LLAMA_3_1_8B_INSTRUCT = os.path.join( ) MIMO_7B_RL_WEIGHTS_PATH = os.path.join(MODEL_WEIGHTS_DIR, "XiaomiMiMo/MiMo-7B-RL") MIMO_V2_FLASH_WEIGHTS_PATH = os.path.join(MODEL_WEIGHTS_DIR, "XiaomiMiMo/MiMo-V2-Flash") +MIMO_V2_5_PRO_FP4_DFLASH_WEIGHTS_PATH = os.path.join( + MODEL_WEIGHTS_DIR, "XiaomiMiMo/MiMo-V2.5-Pro-FP4-DFlash" +) +MIMO_V2_5_PRO_FP4_DFLASH_DRAFT_WEIGHTS_PATH = os.path.join( + MIMO_V2_5_PRO_FP4_DFLASH_WEIGHTS_PATH, "dflash" +) MIMO_V2_5_W8A8_WEIGHTS_PATH = os.path.join(MODEL_WEIGHTS_DIR, "solinliu/MiMo-V2.5-W8A8") MINICPM3_4B_WEIGHTS_PATH = os.path.join(MODEL_WEIGHTS_DIR, "OpenBMB/MiniCPM3-4B") MISTRAL_7B_INSTRUCT_V0_2_WEIGHTS_PATH = os.path.join( diff --git a/test/manual/ascend/llm_models/test_npu_mimo_v2_5_pro_fp4_dflash.py b/test/manual/ascend/llm_models/test_npu_mimo_v2_5_pro_fp4_dflash.py new file mode 100644 index 000000000..0e1d96491 --- /dev/null +++ b/test/manual/ascend/llm_models/test_npu_mimo_v2_5_pro_fp4_dflash.py @@ -0,0 +1,47 @@ +import unittest + +from sglang.test.ascend.gsm8k_ascend_mixin import GSM8KAscendMixin +from sglang.test.ascend.test_ascend_utils import ( + MIMO_V2_5_PRO_FP4_DFLASH_DRAFT_WEIGHTS_PATH, + MIMO_V2_5_PRO_FP4_DFLASH_WEIGHTS_PATH, +) +from sglang.test.test_utils import CustomTestCase + + +class TestMiMoV25ProFP4GraphWithDFlash(GSM8KAscendMixin, CustomTestCase): + """Testcase: Verify the inference accuracy of MiMo-V2.5-Pro-FP4 on GSM8K with npu graph and DFlash speculative decoding. + + [Test Category] Model + [Test Target] MiMo-V2.5-Pro-FP4-DFlash + [Test Config] Prefill+Decode, npu graph enabled, DFLASH speculative decoding, mxfp4 quantization, dp attention + """ + + model = MIMO_V2_5_PRO_FP4_DFLASH_WEIGHTS_PATH + accuracy = 0.9 + other_args = [ + "--trust-remote-code", + "--mem-fraction-static", + "0.87", + "--attention-backend", + "ascend", + "--tp-size", + "8", + "--moe-a2a-backend", + "deepep", + "--deepep-mode", + "auto", + "--speculative-algorithm", + "DFLASH", + "--speculative-draft-model-path", + MIMO_V2_5_PRO_FP4_DFLASH_DRAFT_WEIGHTS_PATH, + "--speculative-num-draft-tokens", + "8", + "--dp-size", + "2", + "--enable-dp-attention", + "--enable-dp-lm-head", + ] + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/spec/dspark/test_dspark_stacked_ctx_kv_parity.py b/test/registered/spec/dspark/test_dspark_stacked_ctx_kv_parity.py index 6a5f6dda7..c3891d027 100644 --- a/test/registered/spec/dspark/test_dspark_stacked_ctx_kv_parity.py +++ b/test/registered/spec/dspark/test_dspark_stacked_ctx_kv_parity.py @@ -58,6 +58,9 @@ def _make_attn(rope, *, eps=EPS, has_bias=False, quantized=False, g=None): k_norm.weight.copy_(torch.randn(HEAD_DIM, device=DEVICE, generator=g)) attn.k_norm = k_norm attn.rotary_emb = rope + # Real DFlashAttention always sets v_scale in __init__; the mock must match + # so kv_proj_only's v_scale read does not AttributeError. + attn.v_scale = None for name in ("kv_proj_only", "apply_k_norm", "apply_k_rope"): setattr(attn, name, types.MethodType(getattr(DFlashAttention, name), attn)) return attn