@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 9969466a59

[Spec Decode] Support SWA + DFlash for MiMo (#46104)

Benjamin Chislett committed Jun 30, 2026 at 23:34 UTC 9969466a597810db6e06b4942dd6cc2086885ee2

4 files changed +243 -25

tests/models/registry.py

+1 -1

@@ -1408,7 +1408,7 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {

1408 # [DFlash]

1409 "DFlashDraftModel": _HfExamplesInfo(

1410 "Qwen/Qwen3.5-4B",

1411 - speculative_model="z-lab/Qwen3.5-4B-DFlash",

1411 + speculative_model="z-lab/Qwen3-4B-DFlash-b16",

1412 use_original_num_layers=True, # Need all layers since DFlash has >1 layer,

1413 max_model_len=8192, # Reduce max len to ensure test runs in low-VRAM CI env

1414 max_num_seqs=32,

vllm/model_executor/models/mimo_v2.py

+16 -3

@@ -53,7 +53,12 @@ from vllm.sequence import IntermediateTensors

53 from vllm.v1.attention.backend import AttentionType

54 from vllm.v1.attention.backends.registry import AttentionBackendEnum

56 -from .interfaces import MixtureOfExperts, SupportsPP

56 +from .interfaces import (

57 + EagleModelMixin,

58 + MixtureOfExperts,

59 + SupportsEagle3,

60 + SupportsPP,

61 +)

62 from .utils import (

63 AutoWeightsLoader,

64 PPMissingLayer,

@@ -539,7 +544,7 @@ def _shard_fp8_qkv_proj(

544

545

546 @support_torch_compile

542 -class MiMoV2Model(nn.Module):

547 +class MiMoV2Model(nn.Module, EagleModelMixin):

548 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

549 super().__init__()

550

@@ -602,10 +607,16 @@ class MiMoV2Model(nn.Module):

607 hidden_states = intermediate_tensors["hidden_states"]

608 residual = intermediate_tensors["residual"]

609

610 + aux_hidden_states = self._maybe_add_hidden_state(

611 + [], self.start_layer, hidden_states, residual

612 + )

613 for idx, layer in enumerate(

614 islice(self.layers, self.start_layer, self.end_layer)

615 ):

616 hidden_states, residual = layer(positions, hidden_states, residual)

617 + self._maybe_add_hidden_state(

618 + aux_hidden_states, idx + 1, hidden_states, residual

619 + )

620

621 if not get_pp_group().is_last_rank:

622 return IntermediateTensors(

@@ -614,6 +625,8 @@ class MiMoV2Model(nn.Module):

625

626 hidden_states, _ = self.norm(hidden_states, residual)

627

628 + if len(aux_hidden_states) > 0:

629 + return hidden_states, aux_hidden_states

630 return hidden_states

631

632 def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:

@@ -822,7 +835,7 @@ class MiMoV2Model(nn.Module):

835 return True

836

837

825 -class MiMoV2FlashForCausalLM(nn.Module, SupportsPP, MixtureOfExperts):

838 +class MiMoV2FlashForCausalLM(nn.Module, SupportsPP, MixtureOfExperts, SupportsEagle3):

839 packed_modules_mapping = {

840 "qkv_proj": ["q_proj", "k_proj", "v_proj"],

841 "gate_up_proj": ["gate_proj", "up_proj"],

vllm/model_executor/models/qwen3_dflash.py

+193 -4

@@ -1,6 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 +import io

5 from collections.abc import Iterable

7 import torch

@@ -11,7 +12,10 @@ from transformers import Qwen3Config

12 from vllm import _custom_ops as ops

13 from vllm.compilation.decorators import support_torch_compile

14 from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config

14 -from vllm.distributed import get_tensor_model_parallel_world_size

15 +from vllm.distributed import (

16 + get_tensor_model_parallel_rank,

17 + get_tensor_model_parallel_world_size,

18 +)

19 from vllm.logger import init_logger

20 from vllm.model_executor.layers.attention import Attention

21 from vllm.model_executor.layers.layernorm import RMSNorm

@@ -33,6 +37,7 @@ from vllm.model_executor.model_loader.weight_utils import (

37 )

38 from vllm.multimodal.inputs import NestedTensors

39 from vllm.transformers_utils.config import set_default_rope_theta

40 +from vllm.transformers_utils.repo_utils import get_hf_file_bytes

41 from vllm.v1.attention.backend import AttentionType

43 from .qwen2 import Qwen2MLP as Qwen3MLP

@@ -47,6 +52,76 @@ from .utils import (

52 logger = init_logger(__name__)

55 +def _resolve_layer_attention(

56 + config: Qwen3Config, layer_idx: int

57 +) -> tuple[int | None, bool]:

58 + """Resolve ``(sliding_window, causal)`` for one DFlash draft layer.

59 +

60 + +----------------------+-------------------------+--------------------------------+

61 + | Config | ``layer_type`` | *``causal`` |

62 + +======================+=========================+================================+

63 + | ``layer_types`` | SWA if ``use_swa`` | True if ``layer_types[i]=SWA`` |

64 + | | else ``layer_types[i]`` | else False |

65 + +----------------------+-------------------------+--------------------------------+

66 + | ``layer_types=None`` | SWA | False |

67 + | + ``use_swa=True`` | | |

68 + +----------------------+-------------------------+--------------------------------+

69 + | ``layer_types=None`` | Full | False |

70 + | + ``use_swa=False`` | | |

71 + +----------------------+-------------------------+--------------------------------+

72 + * If ``dflash_config.causal`` is set, its value overrides ``causal`` for all layers.

73 +

74 + This is to support a varied ecosystem of checkpoints, including:

75 + - XiaomiMiMo/MiMo-V2.5-Pro-FP4-DFlash (sets "use_swa", assumes non-causal)

76 + - z-lab/gemma-4-31B-it-DFlash (has mixed layer types, assumes causal only for SWA)

77 + - z-lab/Qwen3.5-9B-DFlash ("standard" DFlash, all full attn, assumes non-causal)

78 + """

79 + dflash_config = getattr(config, "dflash_config", None) or {}

80 + layer_types = getattr(config, "layer_types", None)

81 + use_swa = dflash_config.get("use_swa", False)

82 + config_causal = dflash_config.get("causal", None)

83 +

84 + SLIDING_ATTENTION = "sliding_attention"

85 + any_sliding = False

86 + if layer_types is not None:

87 + num_sliding = sum(lt == SLIDING_ATTENTION for lt in layer_types)

88 + any_sliding = num_sliding > 0

89 + all_sliding = num_sliding == len(layer_types)

90 + if any_sliding and not all_sliding:

91 + # Mixed sliding/full attention needs per-layer causal metadata and

92 + # multiple KV-cache groups, which DFlash does not yet support.

93 + raise NotImplementedError(

94 + "DFlash does not yet support mixed sliding/full attention via "

95 + "layer_types; see "

96 + "https://github.com/vllm-project/vllm/issues/40898."

97 + )

98 +

99 + default_causal = False

100 + if layer_types is None or (use_swa and not any_sliding):

101 + # An absent ``layer_types`` (or the all-"full_attention" one that may

102 + # be synthesized when the checkpoint omits it) must not override

103 + # ``dflash_config.use_swa``, which forces SWA on every layer.

104 + is_sliding = use_swa

105 + else:

106 + is_sliding = layer_types[layer_idx] == SLIDING_ATTENTION

107 + # Full-attention layers default non-causal; SWA layers default causal.

108 + default_causal = is_sliding

109 +

110 + sliding_window = None

111 + if is_sliding:

112 + sliding_window = dflash_config.get(

113 + "swa_window_size", getattr(config, "sliding_window", None)

114 + )

115 + if sliding_window is None:

116 + raise ValueError(

117 + "DFlash sliding attention requires a window size configured in "

118 + "dflash_config.swa_window_size or the top-level sliding_window."

119 + )

120 +

121 + causal = config_causal if config_causal is not None else default_causal

122 + return sliding_window, causal

123 +

124 +

125 class DFlashQwen3Attention(nn.Module):

126 """Attention for DFlash speculative decoding.

127

@@ -64,6 +139,9 @@ class DFlashQwen3Attention(nn.Module):

139 head_dim: int | None = None,

140 rms_norm_eps: float = 1e-06,

141 attention_bias: bool = False,

142 + add_swa_attention_sink_bias: bool = False,

143 + sliding_window: int | None = None,

144 + causal: bool = False,

145 cache_config: CacheConfig | None = None,

146 quant_config: QuantizationConfig | None = None,

147 prefix: str = "",

@@ -109,6 +187,14 @@ class DFlashQwen3Attention(nn.Module):

187 max_position=max_position,

188 rope_parameters=rope_parameters,

189 )

190 +

191 + self.attention_sink_bias = (

192 + torch.nn.Parameter(torch.empty(self.num_heads), requires_grad=False)

193 + if add_swa_attention_sink_bias

194 + else None

195 + )

196 +

197 + self.sliding_window = sliding_window

198 self.attn = Attention(

199 self.num_heads,

200 self.head_dim,

@@ -116,9 +202,13 @@ class DFlashQwen3Attention(nn.Module):

202 num_kv_heads=self.num_kv_heads,

203 cache_config=cache_config,

204 quant_config=quant_config,

205 + per_layer_sliding_window=sliding_window,

206 prefix=f"{prefix}.attn",

207 attn_type=attn_type,

208 + sinks=self.attention_sink_bias,

209 )

210 + # NOTE: `causal` is currently unused here, but will be needed in the future

211 + # to support models with different causality per-layer.

212 self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)

213 self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)

214

@@ -156,6 +246,7 @@ class DFlashQwen3DecoderLayer(nn.Module):

246 vllm_config: VllmConfig,

247 *,

248 config: Qwen3Config,

249 + layer_idx: int,

250 cache_config: CacheConfig | None = None,

251 quant_config: QuantizationConfig | None = None,

252 prefix: str = "",

@@ -165,6 +256,18 @@ class DFlashQwen3DecoderLayer(nn.Module):

256 set_default_rope_theta(config, default_theta=1000000)

257 attn_type = AttentionType.DECODER

258

259 + # DFlash drafts store the sink-bias flag inside dflash_config; fall back

260 + # to the top-level attribute used by other (e.g. MiMo) configs.

261 + dflash_config = getattr(config, "dflash_config", None) or {}

262 + add_swa_attention_sink_bias = dflash_config.get(

263 + "attention_sink_bias",

264 + getattr(config, "add_swa_attention_sink_bias", False),

265 + )

266 +

267 + # Resolve this layer's attention mode (full vs sliding window, causal vs

268 + # non-causal) from the draft config.

269 + sliding_window, causal = _resolve_layer_attention(config, layer_idx)

270 +

271 self.self_attn = DFlashQwen3Attention(

272 hidden_size=self.hidden_size,

273 num_heads=config.num_attention_heads,

@@ -172,6 +275,9 @@ class DFlashQwen3DecoderLayer(nn.Module):

275 num_kv_heads=config.num_key_value_heads,

276 rms_norm_eps=config.rms_norm_eps,

277 attention_bias=getattr(config, "attention_bias", False),

278 + add_swa_attention_sink_bias=add_swa_attention_sink_bias,

279 + sliding_window=sliding_window,

280 + causal=causal,

281 head_dim=getattr(config, "head_dim", None),

282 cache_config=cache_config,

283 quant_config=quant_config,

@@ -243,11 +349,24 @@ class DFlashQwen3Model(nn.Module):

349 prefix=maybe_prefix(prefix, "embed_tokens"),

350 )

351

352 + # Masked query slots are fed to the draft as `mask_token_id`. Most DFlash

353 + # checkpoints will have the mask embedding in the vocabulary embedding table

354 + # at that slot id. Some checkpoints (XiaomiMiMo/MiMo-V2.5-Pro-FP4-DFlash) ship

355 + # with a separate mask embedding tensor to use instead. When present, we load it

356 + # and substitute it for embed_tokens[mask_token_id] when computing embeddings.

357 + self.mask_token_id = drafter_config.get("mask_token_id")

358 + self.mask_embedding = nn.Parameter(

359 + torch.zeros(self.config.hidden_size, dtype=vllm_config.model_config.dtype),

360 + requires_grad=False,

361 + )

362 + self.has_separate_mask_embedding = False

363 +

364 self.layers = nn.ModuleList(

365 [

366 DFlashQwen3DecoderLayer(

367 current_vllm_config,

368 config=self.config,

369 + layer_idx=layer_idx,

370 cache_config=current_vllm_config.cache_config,

371 quant_config=self.quant_config,

372 prefix=maybe_prefix(prefix, f"layers.{layer_idx + start_layer_id}"),

@@ -284,7 +403,12 @@ class DFlashQwen3Model(nn.Module):

403 )

404

405 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

287 - return self.embed_tokens(input_ids)

406 + embeds = self.embed_tokens(input_ids)

407 + if self.has_separate_mask_embedding and self.mask_token_id is not None:

408 + # Replace masked slots with the dedicated mask embedding.

409 + is_mask = (input_ids == self.mask_token_id).unsqueeze(-1)

410 + embeds = torch.where(is_mask, self.mask_embedding.to(embeds.dtype), embeds)

411 + return embeds

412

413 def _build_fused_kv_buffers(self) -> None:

414 """Build fused weight buffers for precompute_and_store_context_kv.

@@ -469,6 +593,8 @@ class DFlashQwen3Model(nn.Module):

593 ]

594 params_dict = dict(self.named_parameters())

595 loaded_params: set[str] = set()

596 + tp_rank = get_tensor_model_parallel_rank()

597 + tp_size = get_tensor_model_parallel_world_size()

598 for name, loaded_weight in weights:

599 if "midlayer." in name:

600 name = name.replace("midlayer.", "layers.0.")

@@ -476,6 +602,18 @@ class DFlashQwen3Model(nn.Module):

602 name = maybe_remap_kv_scale_name(name, params_dict)

603 if name is None:

604 continue

605 + if "attention_sink_bias" in name:

606 + if name not in params_dict:

607 + continue

608 + # Sink bias is per-head; shard it across TP ranks like the

609 + # attention heads themselves.

610 + param = params_dict[name]

611 + heads_per_rank = loaded_weight.shape[0] // tp_size

612 + head_start = tp_rank * heads_per_rank

613 + narrow_weight = loaded_weight.narrow(0, head_start, heads_per_rank)

614 + param.data.copy_(narrow_weight)

615 + loaded_params.add(name)

616 + continue

617 for param_name, weight_name, shard_id in stacked_params_mapping:

618 if weight_name not in name:

619 continue

@@ -495,7 +633,8 @@ class DFlashQwen3Model(nn.Module):

633 class DFlashQwen3ForCausalLM(Qwen3ForCausalLM):

634 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

635 nn.Module.__init__(self)

498 - self.config = vllm_config.speculative_config.draft_model_config.hf_config

636 + self.draft_model_config = vllm_config.speculative_config.draft_model_config

637 + self.config = self.draft_model_config.hf_config

638 if getattr(self.config, "draft_vocab_size", None) is None:

639 self.config.draft_vocab_size = getattr(self.config, "vocab_size", None)

640 target_layer_num = vllm_config.model_config.get_num_layers(

@@ -589,7 +728,9 @@ class DFlashQwen3ForCausalLM(Qwen3ForCausalLM):

728 includes_embed_tokens = False

729 for name, loaded_weight in weights:

730 assert "mask_hidden" not in name, (

592 - "DFlash should use mask_token_id to embed the padding hidden state"

731 + "DFlash embeds masked slots via mask_token_id (optionally "

732 + "overridden by a mask_embedding.pt file); it should not ship a "

733 + "mask_hidden weight."

734 )

735 if "t2d" in name:

736 continue

@@ -603,6 +744,13 @@ class DFlashQwen3ForCausalLM(Qwen3ForCausalLM):

744 model_weights[name] = loaded_weight

745 process_eagle_weight(self, name)

746

747 + # Route the separately-trained mask embedding (if shipped) through the

748 + # standard weight loader alongside the rest of the draft weights.

749 + mask_embedding = self._read_mask_embedding()

750 + if mask_embedding is not None:

751 + model_weights["model.mask_embedding"] = mask_embedding

752 + self.model.has_separate_mask_embedding = True

753 +

754 skip_substrs = []

755 if not includes_draft_id_mapping:

756 skip_substrs.append("draft_id_to_target_id")

@@ -610,6 +758,8 @@ class DFlashQwen3ForCausalLM(Qwen3ForCausalLM):

758 skip_substrs.append("embed_tokens")

759 if not self.model.use_aux_hidden_state:

760 skip_substrs.append("fc.")

761 + if not self.model.has_separate_mask_embedding:

762 + skip_substrs.append("mask_embedding")

763 loader = AutoWeightsLoader(

764 self,

765 skip_prefixes=None,

@@ -617,3 +767,42 @@ class DFlashQwen3ForCausalLM(Qwen3ForCausalLM):

767 )

768 loader.load_weights(model_weights.items())

769 self.model._build_fused_kv_buffers()

770 +

771 + def _read_mask_embedding(self) -> torch.Tensor | None:

772 + """Checks for an override mask embedding in `mask_embedding.pt` and returns it.

773 +

774 + Some checkpoints ship a separately-trained mask embedding for the mask token,

775 + which we use to overwrite the embedding for `mask_token_id`. This helper

776 + checks for the file, loads the pytorch tensor, and returns the embedding to use.

777 +

778 + Returns None if the override file is not present.

779 + """

780 + mask_token_id = self.model.mask_token_id

781 + if mask_token_id is None:

782 + return None

783 +

784 + MASK_EMBEDDING_FILENAME = "mask_embedding.pt"

785 + data = get_hf_file_bytes(

786 + MASK_EMBEDDING_FILENAME,

787 + self.draft_model_config.model,

788 + self.draft_model_config.revision,

789 + )

790 + if data is None:

791 + return None

792 +

793 + state = torch.load(io.BytesIO(data), weights_only=True)

794 + if isinstance(state, dict):

795 + if state.get("mask_token_id", mask_token_id) != mask_token_id:

796 + raise ValueError(

797 + f"{MASK_EMBEDDING_FILENAME} mask_token_id does not match "

798 + f"dflash_config.mask_token_id ({mask_token_id}). "

799 + f"Got {state.get('mask_token_id')}."

800 + )

801 + state = state["embedding"]

802 +

803 + logger.info(

804 + "Loaded DFlash mask embedding for mask_token_id %s from %s",

805 + mask_token_id,

806 + MASK_EMBEDDING_FILENAME,

807 + )

808 + return state.reshape(-1)

vllm/v1/attention/backends/flash_attn.py

+33 -17

@@ -252,6 +252,8 @@ class FlashAttentionMetadata:

252

253 causal: bool | torch.Tensor = True

254

255 + sliding_window: tuple[int, int] | None = None

256 +

257 # PrefixLM bidirectional ranges for multimodal tokens.

258 # Shape: (num_seqs, max_ranges, 2) int32, [start, end] per range.

259 mm_prefix_range_tensor: torch.Tensor | None = None

@@ -284,6 +286,20 @@ def _get_sliding_window_configs(

286 return sliding_window_configs

287

288

289 +def _maybe_symmetrize_window(

290 + window: tuple[int, int] | None,

291 + causal: bool | torch.Tensor,

292 +) -> tuple[int, int] | None:

293 + """Make a causal sliding window ``(w, 0)`` symmetric ``(w, w)`` when attention

294 + is non-causal, so bidirectional queries attend in both directions. Leaves

295 + full-attention ``(-1, -1)`` and already-symmetric windows untouched.

296 + """

297 + non_causal = isinstance(causal, torch.Tensor) or causal is False

298 + if window is not None and window[0] >= 0 and window[1] == 0 and non_causal:

299 + return (window[0], window[0])

300 + return window

301 +

302 +

303 class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetadata]):

304 # FA3:

305 # Supports full cudagraphs for all cases.

@@ -487,7 +503,9 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad

503 cu_seqlens_q=cu_query_lens,

504 page_size=self.block_size,

505 causal=causal,

490 - window_size=self.aot_sliding_window,

506 + window_size=_maybe_symmetrize_window(

507 + self.aot_sliding_window, causal

508 + ),

509 num_splits=max_num_splits,

510 )

511 return None

@@ -579,6 +597,13 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad

597 if isinstance(causal, torch.Tensor) and causal.dtype != torch.int32:

598 causal = causal.to(torch.int32)

599

600 + # Symmetrize the spec's sliding_window for non-causal attention

601 + group_sliding_window = getattr(self.kv_cache_spec, "sliding_window", None)

602 + base_window = (

603 + (-1, -1) if group_sliding_window is None else (group_sliding_window - 1, 0)

604 + )

605 + effective_sliding_window = _maybe_symmetrize_window(base_window, causal)

606 +

607 attn_metadata = FlashAttentionMetadata(

608 num_actual_tokens=num_actual_tokens,

609 max_query_len=max_query_len,

@@ -598,6 +623,7 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad

623 prefix_scheduler_metadata=prefix_scheduler_metadata,

624 max_num_splits=max_num_splits,

625 causal=causal,

626 + sliding_window=effective_sliding_window,

627 )

628

629 # Compute mm_prefix range tensor if the batch contains

@@ -844,28 +870,18 @@ class FlashAttentionImpl(AttentionImpl):

870 )

871 return output

872 else:

873 + window = (

874 + attn_metadata.sliding_window

875 + if attn_metadata.sliding_window is not None

876 + else self.sliding_window

877 + )

878 sliding_window_size: list[int] | None = (

848 - list(self.sliding_window)

849 - if self.sliding_window is not None

850 - else None

879 + list(window) if window is not None else None

880 )

881

882 causal = attn_metadata.causal

883 is_dynamic_causal = isinstance(causal, torch.Tensor)

884

856 - # For non-causal (bidirectional) attention, make the

857 - # sliding window symmetric so queries attend in both

858 - # directions.

859 - if (

860 - sliding_window_size is not None

861 - and sliding_window_size[1] == 0

862 - and (is_dynamic_causal or causal is False)

863 - ):

864 - sliding_window_size = [

865 - sliding_window_size[0],

866 - sliding_window_size[0],

867 - ]

868 -

885 mm_prefix_ranges = attn_metadata.mm_prefix_range_tensor

886 mm_mask_mod = None

887 mm_aux = None