@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 0d12618e98

[Spec Decode] Support hybrid (SWA + full attention) DFlash drafters (#47914)

Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Michael Goin committed Jul 8, 2026 at 11:12 UTC 0d12618e98ff2d21d36081e0e9b4eb23573b6d38

9 files changed +79 -30

vllm/config/vllm.py

+17

@@ -542,6 +542,11 @@ class VllmConfig:

542 ):

543 return True

544

545 + # Mixed sliding/full DFlash drafts need multiple KV groups (V2 only);

546 + # force V2 as for dspark, since a hybrid target otherwise defaults to V1.

547 + if self._dflash_needs_multi_kv_group():

548 + return True

549 +

550 if self.model_config is not None and self.model_config.is_diffusion:

551 return True

552

@@ -565,6 +570,18 @@ class VllmConfig:

570

571 return True

572

573 + def _dflash_needs_multi_kv_group(self) -> bool:

574 + """Whether a DFlash draft mixes sliding-window and full attention."""

575 + spec = self.speculative_config

576 + if spec is None or spec.method != "dflash":

577 + return False

578 + draft_config = getattr(spec, "draft_model_config", None)

579 + if draft_config is None:

580 + return False

581 + layer_types = getattr(draft_config.hf_config, "layer_types", None) or []

582 + num_sliding = sum(lt == "sliding_attention" for lt in layer_types)

583 + return 0 < num_sliding < len(layer_types)

584 +

585 def _is_default_v2_model_runner_model(self) -> bool:

586 model_config = self.model_config

587 if model_config is None:

vllm/model_executor/models/qwen3_dflash.py

+17 -9

@@ -86,14 +86,15 @@ def _resolve_layer_attention(

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.

89 + # Mixed sliding/full attention needs multiple KV groups (V2 runner only).

90 + if (

91 + 0 < num_sliding < len(layer_types)

92 + and not get_current_vllm_config().use_v2_model_runner

93 + ):

94 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."

95 + "DFlash drafters with mixed sliding/full attention require "

96 + "the V2 model runner; relaunch with "

97 + "VLLM_USE_V2_MODEL_RUNNER=1."

98 )

100 default_causal = False

@@ -207,8 +208,7 @@ class DFlashQwen3Attention(nn.Module):

208 attn_type=attn_type,

209 sinks=self.attention_sink_bias,

210 )

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

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

211 + self.causal = causal

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

@@ -709,6 +709,14 @@ class DFlashQwen3ForCausalLM(Qwen3ForCausalLM):

709 ) -> torch.Tensor:

710 return self.model(input_ids, positions, inputs_embeds)

711

712 + def get_draft_kv_cache_layer_names(self) -> list[str]:

713 + return [layer.self_attn.attn.layer_name for layer in self.model.layers]

714 +

715 + def get_draft_attn_causal(self) -> list[bool]:

716 + """Per-layer attention causality, aligned with

717 + get_draft_kv_cache_layer_names."""

718 + return [layer.self_attn.causal for layer in self.model.layers]

719 +

720 def compute_logits(

721 self,

722 hidden_states: torch.Tensor,

vllm/v1/worker/gpu/attn_utils.py

+5 -3

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from collections.abc import Iterable, Sequence

3 +from collections.abc import Iterable, Mapping, Sequence

4 from dataclasses import dataclass, replace

5 from math import prod

6 from typing import Any, cast

@@ -573,7 +573,7 @@ def build_attn_metadata(

573 mm_req_doc_ranges: dict[int, list[tuple[int, int]]] | None = None,

574 model_specific_attn_metadata: ModelSpecificAttnMetadata | None = None,

575 for_cudagraph_capture: bool = False,

576 - causal: bool = True,

576 + causal: bool | Mapping[int, bool] = True,

577 rswa_prefix_lens: torch.Tensor | None = None,

578 ) -> dict[str, Any]:

579 seq_lens = seq_lens[:num_reqs]

@@ -587,6 +587,8 @@ def build_attn_metadata(

587 for i in range(num_kv_cache_groups):

588 block_table = block_tables[i]

589 slot_mapping = slot_mappings[i]

590 + # Per-group causal for hybrid drafters (mixed SWA/full attention).

591 + group_causal = causal if isinstance(causal, bool) else causal.get(i, True)

592

593 common_attn_metadata_extra_kwargs = (

594 model_specific_attn_metadata.get_extra_common_attn_kwargs(i, num_reqs)

@@ -604,7 +606,7 @@ def build_attn_metadata(

606 max_query_len=max_query_len,

607 block_table_tensor=block_table,

608 slot_mapping=slot_mapping,

607 - causal=causal,

609 + causal=group_causal,

610 dcp_local_seq_lens=dcp_local_seq_lens,

611 positions=positions,

612 mm_req_doc_ranges=mm_req_doc_ranges,

vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py

+4 -7

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from collections.abc import Callable

3 +from collections.abc import Callable, Mapping

5 import torch

@@ -29,7 +29,7 @@ def _prepare_dflash_inputs_to_capture(

29 kv_cache_config: KVCacheConfig,

30 max_model_len: int,

31 skip_attn: bool,

32 - causal: bool,

32 + causal: bool | Mapping[int, bool],

33 ) -> AttentionState:

34 input_batch = InputBatch.make_dummy(num_reqs, num_tokens, input_buffers)

35 input_block_tables = block_tables.get_dummy_block_tables(num_reqs)

@@ -63,10 +63,6 @@ class DFlashCudaGraphManager(CudaGraphManager):

63 """DFlash CudaGraphManager for the parallel-drafting query forward,

64 building its own attention metadata from scratch."""

66 - def __init__(self, *args, causal: bool = False, **kwargs) -> None:

67 - super().__init__(*args, **kwargs)

68 - self.causal = causal

69 -

66 def capture(

67 self,

68 forward_fn: Callable,

@@ -75,6 +71,7 @@ class DFlashCudaGraphManager(CudaGraphManager):

71 attn_groups: list[list[AttentionGroup]],

72 kv_cache_config: KVCacheConfig,

73 max_model_len: int,

74 + causal: bool | Mapping[int, bool],

75 progress_bar_desc: str = "Capturing CUDA graphs",

76 ) -> None:

77 def create_forward_fn(

@@ -97,7 +94,7 @@ class DFlashCudaGraphManager(CudaGraphManager):

94 kv_cache_config,

95 max_model_len,

96 skip_attn=(desc.cg_mode == CUDAGraphMode.PIECEWISE),

100 - causal=self.causal,

97 + causal=causal,

98 )

99 attn_metadata, slot_mappings = attn_state

100

vllm/v1/worker/gpu/spec_decode/dflash/speculator.py

+15 -5

@@ -1,5 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3 +from collections.abc import Mapping

4 from typing import Any

6 import torch

@@ -94,7 +95,6 @@ class DFlashSpeculator(DraftModelSpeculator):

95 self.device,

96 cudagraph_mode,

97 decode_query_len=self.num_query_per_req,

97 - causal=self.dflash_causal,

98 )

100 def capture(self, attn_states: dict | None = None) -> None:

@@ -112,6 +112,7 @@ class DFlashSpeculator(DraftModelSpeculator):

112 self.attn_groups,

113 self.kv_cache_config,

114 self.max_model_len,

115 + causal=self._group_causal,

116 progress_bar_desc=f"Capturing {self._speculator_name.lower()} CUDA graphs",

117 )

118

@@ -148,7 +149,10 @@ class DFlashSpeculator(DraftModelSpeculator):

149 # of the kv-cache group its cache belongs to. Models that share a single group

150 # leave this as None and share one context slot mapping.

151 self._layer_group_idx: list[int] | None = None

152 + # Per-KV-group causal, falling back to the scalar dflash_causal.

153 + self._group_causal: dict[int, bool] | bool = self.dflash_causal

154 if hasattr(self.model, "get_draft_kv_cache_layer_names"):

155 + layer_names = self.model.get_draft_kv_cache_layer_names()

156 name_to_gid = {

157 ln: gid

158 for gid, group in enumerate(kv_cache_config.kv_cache_groups)

@@ -156,9 +160,15 @@ class DFlashSpeculator(DraftModelSpeculator):

160 }

161 gid_to_idx = {gid: i for i, gid in enumerate(self.draft_kv_cache_group_ids)}

162 self._layer_group_idx = [

159 - gid_to_idx[name_to_gid[name]]

160 - for name in self.model.get_draft_kv_cache_layer_names()

163 + gid_to_idx[name_to_gid[name]] for name in layer_names

164 ]

165 + if hasattr(self.model, "get_draft_attn_causal"):

166 + self._group_causal = {

167 + name_to_gid[name]: layer_causal

168 + for name, layer_causal in zip(

169 + layer_names, self.model.get_draft_attn_causal()

170 + )

171 + }

172

173 @torch.inference_mode()

174 def _run_model(

@@ -226,7 +236,7 @@ class DFlashSpeculator(DraftModelSpeculator):

236 num_reqs_padded: int,

237 num_tokens_padded: int,

238 num_query_per_req: int | None = None,

229 - causal: bool = False,

239 + causal: bool | Mapping[int, bool] = False,

240 ) -> dict[str, Any] | None:

241 if not self.draft_attn_layer_names:

242 return None

@@ -383,7 +393,7 @@ class DFlashSpeculator(DraftModelSpeculator):

393 num_reqs=num_reqs,

394 num_reqs_padded=num_reqs_padded,

395 num_tokens_padded=num_tokens_padded,

386 - causal=self.dflash_causal,

396 + causal=self._group_causal,

397 )

398 draft_slot_mappings_by_layer = build_slot_mappings_by_layer(

399 self.block_tables.slot_mappings[:, :num_tokens_padded],

vllm/v1/worker/gpu/spec_decode/dflash/utils.py

+5 -2

@@ -5,7 +5,10 @@ import torch.nn as nn

5 from vllm.config import ModelConfig, VllmConfig, replace

6 from vllm.distributed.parallel_state import get_pp_group

7 from vllm.model_executor.model_loader import get_model

8 -from vllm.v1.worker.gpu.spec_decode.eagle.utils import _should_share

8 +from vllm.v1.worker.gpu.spec_decode.eagle.utils import (

9 + _should_share,

10 + get_target_lm_head,

11 +)

14 def get_dflash_causal(draft_model_config: ModelConfig) -> bool:

@@ -57,7 +60,7 @@ def load_dflash_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo

60 del draft_inner.embed_tokens

61 draft_inner.embed_tokens = target_embed

60 - target_lm_head = getattr(target_model, "lm_head", None)

63 + target_lm_head = get_target_lm_head(target_model, target_language_model)

64 draft_lm_head = getattr(dflash_model, "lm_head", None)

65 if target_lm_head is not None and _should_share(

66 dflash_model, "has_own_lm_head", draft_lm_head, target_lm_head

vllm/v1/worker/gpu/spec_decode/dspark/utils.py

+5 -2

@@ -6,7 +6,10 @@ import torch.nn as nn

6 from vllm.config import VllmConfig, replace

7 from vllm.distributed.parallel_state import get_pp_group

8 from vllm.model_executor.model_loader import get_model

9 -from vllm.v1.worker.gpu.spec_decode.eagle.utils import _should_share

9 +from vllm.v1.worker.gpu.spec_decode.eagle.utils import (

10 + _should_share,

11 + get_target_lm_head,

12 +)

15 def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Module:

@@ -52,7 +55,7 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo

55 del draft_inner.embed_tokens

56 draft_inner.embed_tokens = target_embed

55 - target_lm_head = getattr(target_model, "lm_head", None)

58 + target_lm_head = get_target_lm_head(target_model, target_language_model)

59 draft_lm_head = getattr(draft_model, "lm_head", None)

60 if target_lm_head is not None and _should_share(

61 draft_model, "has_own_lm_head", draft_lm_head, target_lm_head

vllm/v1/worker/gpu/spec_decode/eagle/utils.py

+9 -1

@@ -25,6 +25,14 @@ def _should_share(eagle: nn.Module, flag: str, draft, target) -> bool:

25 return torch.equal(w, target.weight)

28 +def get_target_lm_head(target_model: nn.Module, target_language_model: nn.Module):

29 + """The target's lm_head — from get_language_model() for

30 + *ForConditionalGeneration targets, else the top-level module."""

31 + return getattr(target_language_model, "lm_head", None) or getattr(

32 + target_model, "lm_head", None

33 + )

34 +

35 +

36 def load_eagle_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Module:

37 from vllm.compilation.backends import set_model_tag

@@ -64,7 +72,7 @@ def load_eagle_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mod

72 del draft_inner.embed_tokens

73 draft_inner.embed_tokens = target_embed

67 - target_lm_head = getattr(target_model, "lm_head", None)

75 + target_lm_head = get_target_lm_head(target_model, target_language_model)

76 draft_lm_head = getattr(eagle_model, "lm_head", None)

77 if target_lm_head is not None and _should_share(

78 eagle_model, "has_own_lm_head", draft_lm_head, target_lm_head

vllm/v1/worker/gpu/spec_decode/speculator.py

+2 -1

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 from abc import ABC, abstractmethod

4 +from collections.abc import Mapping

5 from typing import Any

7 import torch

@@ -202,7 +203,7 @@ class DraftModelSpeculator(BaseSpeculator):

203 num_reqs_padded: int,

204 num_tokens_padded: int,

205 num_query_per_req: int = 1,

205 - causal: bool = True,

206 + causal: bool | Mapping[int, bool] = True,

207 ) -> dict[str, Any] | None:

208 # Uniform query: query_start_loc[i] = min(i, num_reqs) * num_query_per_req.

209 # Clamp keeps the series non-decreasing past num_reqs, which some