A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 12:15 UTC
[Spec Decode] Allow causal DFlash (#43445)
Benjamin Chislett committed
May 28, 2026 at 17:18 UTC
9202ea6fda056702ceaca1f996c470a602748041
1 file changed
+17
-15
vllm/v1/spec_decode/dflash.py
+17
-15
@@ -68,6 +68,8 @@ class DFlashProposer(SpecDecodeBaseProposer):
68
# For DFlash we use the input embeddings to embed the mask token
69
self.parallel_drafting_hidden_state_tensor = None
70
71
+ self.dflash_causal = self.dflash_config.get("causal", False)
72
+
73
@override
74
def _create_draft_vllm_config(self) -> VllmConfig:
75
base = super()._create_draft_vllm_config()
@@ -75,7 +77,7 @@ class DFlashProposer(SpecDecodeBaseProposer):
77
base,
78
attention_config=replace(
79
base.attention_config,
78
- use_non_causal=True,
80
+ use_non_causal=not self.dflash_causal,
81
),
82
)
83
@@ -184,7 +186,7 @@ class DFlashProposer(SpecDecodeBaseProposer):
186
max_seq_len=cad.max_seq_len + num_query_per_req,
187
block_table_tensor=cad.block_table_tensor,
188
slot_mapping=query_slot_mapping,
187
- causal=False, # Non-causal attention is required for DFlash
189
+ causal=self.dflash_causal,
190
)
191
192
return num_query_total, token_indices_to_sample, new_cad
@@ -281,20 +283,20 @@ class DFlashProposer(SpecDecodeBaseProposer):
283
per_group, per_layer = super().build_per_group_and_layer_attn_metadata(
284
cad, draft_index
285
)
284
- for layer_name, attn_metadata in per_layer.items():
285
- assert getattr(attn_metadata, "causal", None) is False, (
286
- f"Attention metadata for layer {layer_name} does not have"
287
- " non-causal support, which is required for DFlash."
288
- " Consider using a different attention backend, such as FlashAttention."
289
- )
286
+ if not self.dflash_causal:
287
+ # Require all layers to support non-causal attention when required by DFlash
288
+ for layer_name, attn_metadata in per_layer.items():
289
+ assert getattr(attn_metadata, "causal", None) is False, (
290
+ f"Attention metadata for layer {layer_name} does not have"
291
+ " non-causal support, which is required for DFlash."
292
+ " Consider using a different attention backend, e.g FlashAttention."
293
+ )
294
return per_group, per_layer
295
296
@override
297
def _get_eagle3_use_aux_hidden_state_from_config(self):
294
- use_aux_hidden_state = True
295
- dflash_config = getattr(
296
- self.draft_model_config.hf_config, "dflash_config", None
297
- )
298
- if dflash_config is not None:
299
- use_aux_hidden_state = dflash_config.get("use_aux_hidden_state", True)
300
- return use_aux_hidden_state
298
+ return self.dflash_config.get("use_aux_hidden_state", True)
299
+
300
+ @property
301
+ def dflash_config(self):
302
+ return getattr(self.draft_model_config.hf_config, "dflash_config", None) or {}