@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 4c3c64fcf7

Add Laguna XS.2.1 DFlash drafter support (#46853)

Signed-off-by: Adam Baranowski <adam.baranowski@poolside.ai>

adamkbaranowski committed Jul 3, 2026 at 03:09 UTC 4c3c64fcf76450d3d0bbfd3d1725eade0f214710
7 files changed +494 -47
tests/models/registry.py
+7
@@ -1410,6 +1410,13 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
1410 max_model_len=8192, # Reduce max len to ensure test runs in low-VRAM CI env
1411 max_num_seqs=32,
1412 ),
1413 + "DFlashLagunaForCausalLM": _HfExamplesInfo(
1414 + "poolside/Laguna-XS-2.1",
1415 + speculative_model="poolside/Laguna-XS-2.1-DFlash",
1416 + use_original_num_layers=True,
1417 + max_model_len=8192, # Reduce max len to ensure test runs in low-VRAM CI env
1418 + max_num_seqs=32,
1419 + ),
1420 "DFlashQwen3NextDraftModel": _HfExamplesInfo(
1421 "Qwen/Qwen3-Coder-Next",
1422 speculative_model="z-lab/Qwen3-Coder-Next-DFlash",
tests/v1/e2e/spec_decode/test_laguna_dflash.py new
+60
@@ -0,0 +1,60 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +import torch
6 +
7 +from tests.utils import large_gpu_mark
8 +from vllm import LLM, SamplingParams
9 +from vllm.distributed import cleanup_dist_env_and_memory
10 +
11 +
12 +def _get_counter(metrics, name: str) -> float:
13 + metric = next((m for m in metrics if m.name == name), None)
14 + assert metric is not None, f"Missing metric: {name}"
15 + return metric.value
16 +
17 +
18 +@pytest.mark.slow_test
19 +@large_gpu_mark(min_gb=80)
20 +def test_laguna_dflash_hf_pair_smoke(monkeypatch):
21 + """Smoke-test the public Laguna XS-2.1 base/DFlash checkpoint pair."""
22 + monkeypatch.setenv("VLLM_USE_FLASHINFER_SAMPLER", "0")
23 +
24 + llm = LLM(
25 + model="poolside/Laguna-XS-2.1",
26 + trust_remote_code=True,
27 + speculative_config={
28 + "method": "dflash",
29 + "model": "poolside/Laguna-XS-2.1-DFlash",
30 + "num_speculative_tokens": 15,
31 + },
32 + max_model_len=8192,
33 + max_num_batched_tokens=65536,
34 + max_num_seqs=32,
35 + enforce_eager=True,
36 + disable_log_stats=False,
37 + )
38 +
39 + try:
40 + outputs = llm.generate(
41 + [
42 + "What is the capital of the United Kingdom?",
43 + "Write a Python function that returns the square of a number.",
44 + ],
45 + SamplingParams(temperature=0.0, max_tokens=32, ignore_eos=True),
46 + )
47 + assert len(outputs) == 2
48 + assert all(output.outputs[0].text for output in outputs)
49 +
50 + metrics = llm.get_metrics()
51 + num_drafts = _get_counter(metrics, "vllm:spec_decode_num_drafts")
52 + num_accepted = _get_counter(metrics, "vllm:spec_decode_num_accepted_tokens")
53 +
54 + assert num_drafts > 0
55 + acceptance_len = 1 + (num_accepted / num_drafts)
56 + assert acceptance_len > 1.0
57 + finally:
58 + del llm
59 + torch.accelerator.empty_cache()
60 + cleanup_dist_env_and_memory()
vllm/model_executor/models/laguna.py
+18 -6
@@ -269,8 +269,14 @@ class LagunaAttention(nn.Module):
269 quant_config: QuantizationConfig | None = None,
270 prefix: str = "",
271 attention_sink: bool = False,
272 + layer_idx: int | None = None,
273 + attention_prefix: str | None = None,
274 ) -> None:
275 super().__init__()
276 + if layer_idx is None:
277 + layer_idx = extract_layer_index(prefix)
278 + if attention_prefix is None:
279 + attention_prefix = prefix
280 self.hidden_size = hidden_size
281 tp_size = get_tensor_model_parallel_world_size()
282 self.total_num_heads = num_heads
@@ -294,7 +300,6 @@ class LagunaAttention(nn.Module):
300 # Per-layer sliding window (follows Gemma2/Cohere2 convention)
301 layer_types = getattr(config, "layer_types", None)
302 if layer_types is not None:
297 - layer_idx = extract_layer_index(prefix)
303 is_sliding = layer_types[layer_idx] == "sliding_attention"
304 self.sliding_window = config.sliding_window if is_sliding else None
305 else:
@@ -361,9 +366,7 @@ class LagunaAttention(nn.Module):
366 # crash `get_rope`'s cache lookup, so always pull out the layer's
367 # sub-dict before forwarding.
368 layer_type = (
364 - layer_types[extract_layer_index(prefix)]
365 - if layer_types is not None
366 - else "full_attention"
369 + layer_types[layer_idx] if layer_types is not None else "full_attention"
370 )
371 is_sliding = layer_type == "sliding_attention"
372
@@ -415,7 +418,7 @@ class LagunaAttention(nn.Module):
418 cache_config=cache_config,
419 quant_config=quant_config,
420 per_layer_sliding_window=self.sliding_window,
418 - prefix=f"{prefix}.attn",
421 + prefix=f"{attention_prefix}.attn",
422 sinks=sinks,
423 )
424
@@ -468,10 +471,13 @@ class LagunaDecoderLayer(nn.Module):
471 quant_config: QuantizationConfig | None = None,
472 prefix: str = "",
473 enable_eplb: bool = False,
474 + layer_idx: int | None = None,
475 + attention_prefix: str | None = None,
476 ) -> None:
477 super().__init__()
478 self.hidden_size = config.hidden_size
474 - layer_idx = extract_layer_index(prefix)
479 + if layer_idx is None:
480 + layer_idx = extract_layer_index(prefix)
481
482 # Determine if this layer uses sliding window attention
483 layer_types = getattr(config, "layer_types", None)
@@ -503,6 +509,12 @@ class LagunaDecoderLayer(nn.Module):
509 quant_config=quant_config,
510 prefix=f"{prefix}.self_attn",
511 attention_sink=attention_sink,
512 + layer_idx=layer_idx,
513 + attention_prefix=(
514 + f"{attention_prefix}.self_attn"
515 + if attention_prefix is not None
516 + else None
517 + ),
518 )
519
520 # Check if this layer uses MoE or dense MLP (matches Qwen2/Qwen3 convention)
vllm/model_executor/models/laguna_dflash.py new
+342
@@ -0,0 +1,342 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""DFlash speculator for Laguna target models.
4 +
5 +Laguna DFlash uses a uniform drafter layer flavor (`layer_types` all full
6 +or all sliding). The draft checkpoint shares token embedding and lm_head
7 +weights with the target model through the generic spec-decode proposer.
8 +"""
9 +
10 +from collections.abc import Iterable
11 +
12 +import torch
13 +from torch import nn
14 +
15 +from vllm import _custom_ops as ops
16 +from vllm.compilation.decorators import support_torch_compile
17 +from vllm.config import VllmConfig
18 +from vllm.logger import init_logger
19 +from vllm.model_executor.layers.layernorm import RMSNorm
20 +from vllm.model_executor.layers.linear import ReplicatedLinear
21 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
22 +from vllm.model_executor.layers.vocab_parallel_embedding import (
23 + ParallelLMHead,
24 + VocabParallelEmbedding,
25 +)
26 +from vllm.model_executor.model_loader.weight_utils import (
27 + default_weight_loader,
28 + maybe_remap_kv_scale_name,
29 +)
30 +from vllm.model_executor.models.interfaces import EagleModelMixin, SupportsEagle3
31 +from vllm.multimodal.inputs import NestedTensors
32 +
33 +from .laguna import LagunaDecoderLayer
34 +from .qwen3_dflash import DFlashQwen3Model
35 +from .utils import (
36 + AutoWeightsLoader,
37 + get_draft_quant_config,
38 + maybe_prefix,
39 + process_eagle_weight,
40 +)
41 +
42 +logger = init_logger(__name__)
43 +
44 +
45 +def _get_dflash_layer_types(config) -> tuple[str, ...]:
46 + layer_types = getattr(config, "layer_types", None)
47 + if layer_types is None:
48 + raise ValueError("Laguna DFlash config requires `layer_types`.")
49 + if len(layer_types) != config.num_hidden_layers:
50 + raise ValueError(
51 + f"DFlash layer_types length {len(layer_types)} does not match "
52 + f"num_hidden_layers {config.num_hidden_layers}."
53 + )
54 + # Laguna DFlash checkpoints use a uniform drafter attention flavor.
55 + if len(set(layer_types)) > 1:
56 + raise NotImplementedError(
57 + "Laguna DFlash drafter requires a uniform `layer_types` "
58 + f"(got {sorted(set(layer_types))})."
59 + )
60 + return tuple(layer_types)
61 +
62 +
63 +@support_torch_compile
64 +class DFlashLagunaModel(DFlashQwen3Model, EagleModelMixin):
65 + def __init__(
66 + self,
67 + *,
68 + vllm_config: VllmConfig,
69 + prefix: str = "",
70 + ) -> None:
71 + nn.Module.__init__(self)
72 + self.config = vllm_config.speculative_config.draft_model_config.hf_config
73 + self.vocab_size = self.config.vocab_size
74 + self.quant_config = get_draft_quant_config(vllm_config)
75 +
76 + target_layer_ids = self.config.dflash_config["target_layer_ids"]
77 + if not target_layer_ids:
78 + raise ValueError(
79 + "Laguna DFlash config requires non-empty "
80 + "`dflash_config.target_layer_ids`."
81 + )
82 +
83 + self.embed_tokens = VocabParallelEmbedding(
84 + self.config.vocab_size,
85 + self.config.hidden_size,
86 + prefix=maybe_prefix(prefix, "embed_tokens"),
87 + )
88 +
89 + self.mask_token_id = self.config.dflash_config.get("mask_token_id")
90 + self.register_buffer(
91 + "mask_embedding",
92 + torch.zeros(
93 + self.config.hidden_size,
94 + dtype=vllm_config.model_config.dtype,
95 + ),
96 + persistent=False,
97 + )
98 + self.has_separate_mask_embedding = False
99 +
100 + self.layer_types = _get_dflash_layer_types(self.config)
101 + target_layer_count = self.config.target_layer_count
102 + self.layers = nn.ModuleList(
103 + [
104 + LagunaDecoderLayer(
105 + prefix=maybe_prefix(prefix, f"layers.{layer_idx}"),
106 + config=self.config,
107 + cache_config=vllm_config.cache_config,
108 + quant_config=self.quant_config,
109 + layer_idx=layer_idx,
110 + attention_prefix=maybe_prefix(
111 + prefix, f"layers.{layer_idx + target_layer_count}"
112 + ),
113 + )
114 + for layer_idx in range(self.config.num_hidden_layers)
115 + ]
116 + )
117 + for layer in self.layers:
118 + if getattr(layer.self_attn, "sliding_window", None) is not None:
119 + # DFlash inserts verifier-context K/V at absolute cache slots.
120 + # Keep full KV allocation; SWA remains a compute-time limit.
121 + layer.self_attn.attn.sliding_window = None
122 + num_features_to_use = len(target_layer_ids)
123 + target_hidden_size = vllm_config.model_config.get_hidden_size()
124 + fc_input_size = target_hidden_size * num_features_to_use
125 + self.num_aux_slices = num_features_to_use
126 + self.aux_hidden_norms = nn.ModuleList(
127 + [
128 + RMSNorm(
129 + fc_input_size // num_features_to_use,
130 + eps=self.config.rms_norm_eps,
131 + )
132 + for _ in range(num_features_to_use)
133 + ]
134 + )
135 + self.fc = ReplicatedLinear(
136 + input_size=fc_input_size,
137 + output_size=self.config.hidden_size,
138 + bias=False,
139 + params_dtype=vllm_config.model_config.dtype,
140 + quant_config=self.quant_config,
141 + prefix=maybe_prefix(prefix, "fc"),
142 + return_bias=False,
143 + )
144 + self.hidden_norm = RMSNorm(
145 + self.config.hidden_size,
146 + eps=self.config.rms_norm_eps,
147 + )
148 + self.norm = RMSNorm(
149 + self.config.hidden_size,
150 + eps=self.config.rms_norm_eps,
151 + )
152 +
153 + def _build_context_kv_buffers(
154 + self,
155 + layers_attn: list[nn.Module],
156 + has_bias: bool,
157 + ) -> None:
158 + self._kv_weights = torch.stack(
159 + [a.qkv_proj.weight[a.q_size :] for a in layers_attn], dim=0
160 + ).contiguous()
161 + if has_bias:
162 + self._kv_biases: torch.Tensor | None = torch.stack(
163 + [a.qkv_proj.bias[a.q_size :] for a in layers_attn], dim=0
164 + ).contiguous()
165 + else:
166 + self._kv_biases = None
167 + self._input_layernorm_weights = torch.stack(
168 + [layer.input_layernorm.weight.data for layer in self.layers], dim=0
169 + ).contiguous()
170 + self._k_norm_weights = torch.stack(
171 + [a.k_norm.weight.data for a in layers_attn], dim=0
172 + ).contiguous()
173 +
174 + def _project_context_kv(
175 + self,
176 + context_states: torch.Tensor,
177 + num_ctx: int,
178 + num_layers: int,
179 + num_kv_heads: int,
180 + head_dim: int,
181 + ) -> tuple[torch.Tensor, torch.Tensor]:
182 + normed_context_states = torch.empty(
183 + (num_layers, num_ctx, context_states.shape[-1]),
184 + dtype=context_states.dtype,
185 + device=context_states.device,
186 + )
187 + ops.rms_norm(
188 + normed_context_states,
189 + context_states.unsqueeze(0).expand(num_layers, -1, -1),
190 + self._input_layernorm_weights,
191 + self._rms_norm_eps,
192 + )
193 + all_kv_flat = torch.bmm(
194 + normed_context_states,
195 + self._kv_weights.transpose(1, 2),
196 + )
197 + if self._kv_biases is not None:
198 + all_kv_flat += self._kv_biases[:, None, :]
199 + all_kv = (
200 + all_kv_flat.view(num_layers, num_ctx, 2, num_kv_heads, head_dim)
201 + .permute(2, 0, 1, 3, 4)
202 + .contiguous()
203 + )
204 + all_k = all_kv[0]
205 + all_v = all_kv[1]
206 + return all_k, all_v
207 +
208 + def _normalize_context_k(self, all_k: torch.Tensor) -> torch.Tensor:
209 + all_k_normed = torch.empty_like(all_k)
210 + ops.rms_norm(
211 + all_k_normed,
212 + all_k,
213 + self._k_norm_weights,
214 + self._rms_norm_eps,
215 + )
216 + return all_k_normed
217 +
218 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
219 + params_dict = dict(self.named_parameters())
220 + loaded_params: set[str] = set()
221 + for name, loaded_weight in weights:
222 + if self.quant_config is not None and (
223 + scale_name := self.quant_config.get_cache_scale(name)
224 + ):
225 + param = params_dict[scale_name]
226 + weight_loader = getattr(param, "weight_loader", default_weight_loader)
227 + loaded_weight = (
228 + loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
229 + )
230 + weight_loader(param, loaded_weight)
231 + loaded_params.add(scale_name)
232 + continue
233 + if "scale" in name:
234 + name = maybe_remap_kv_scale_name(name, params_dict)
235 + if name is None:
236 + continue
237 + param = params_dict[name]
238 + weight_loader = getattr(param, "weight_loader", default_weight_loader)
239 + weight_loader(param, loaded_weight)
240 + loaded_params.add(name)
241 + return loaded_params
242 +
243 +
244 +class DFlashLagunaForCausalLM(nn.Module, SupportsEagle3):
245 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
246 + nn.Module.__init__(self)
247 + self.config = vllm_config.speculative_config.draft_model_config.hf_config
248 + if getattr(self.config, "draft_vocab_size", None) is None:
249 + raise ValueError("Laguna DFlash config requires `draft_vocab_size`.")
250 + self.has_own_embed_tokens = False
251 + self.has_own_lm_head = False
252 + target_layer_num = vllm_config.model_config.get_num_layers(
253 + vllm_config.parallel_config
254 + )
255 + self.config.target_layer_count = target_layer_num
256 + target_vocab_size = vllm_config.model_config.get_vocab_size()
257 + if self.config.draft_vocab_size != target_vocab_size:
258 + raise ValueError(
259 + "Laguna DFlash shares the target lm_head and requires "
260 + "`draft_vocab_size` to match the target vocabulary size "
261 + f"({self.config.draft_vocab_size} != {target_vocab_size})."
262 + )
263 + self.model = DFlashLagunaModel(
264 + vllm_config=vllm_config,
265 + prefix="model",
266 + )
267 +
268 + self.lm_head = ParallelLMHead(
269 + self.config.draft_vocab_size,
270 + self.config.hidden_size,
271 + prefix=maybe_prefix(prefix, "lm_head"),
272 + )
273 + self.logits_processor = LogitsProcessor(self.config.draft_vocab_size)
274 +
275 + def embed_input_ids(
276 + self,
277 + input_ids: torch.Tensor,
278 + multimodal_embeddings: NestedTensors | None = None,
279 + is_multimodal: torch.Tensor | None = None,
280 + ) -> torch.Tensor:
281 + return self.model.embed_input_ids(input_ids)
282 +
283 + def forward(
284 + self,
285 + input_ids: torch.Tensor,
286 + positions: torch.Tensor,
287 + inputs_embeds: torch.Tensor | None = None,
288 + ) -> torch.Tensor:
289 + return self.model(input_ids, positions, inputs_embeds)
290 +
291 + def compute_logits(
292 + self,
293 + hidden_states: torch.Tensor,
294 + ) -> torch.Tensor | None:
295 + return self.logits_processor(self.lm_head, hidden_states)
296 +
297 + def precompute_and_store_context_kv(
298 + self,
299 + context_states: torch.Tensor,
300 + context_positions: torch.Tensor,
301 + context_slot_mapping: torch.Tensor | None = None,
302 + ) -> None:
303 + self.model.precompute_and_store_context_kv(
304 + context_states, context_positions, context_slot_mapping
305 + )
306 +
307 + def combine_hidden_states(
308 + self,
309 + hidden_states: torch.Tensor,
310 + ) -> torch.Tensor:
311 + # Normalize each verifier hidden-state slice, concatenate them, then
312 + # project into the drafter hidden size used as DFlash context.
313 + needs_squeeze = hidden_states.dim() == 1
314 + if needs_squeeze:
315 + hidden_states = hidden_states.unsqueeze(0)
316 + num_slices = self.model.num_aux_slices
317 + slice_size = hidden_states.shape[-1] // num_slices
318 + slices = hidden_states.view(hidden_states.shape[0], num_slices, slice_size)
319 + normed = torch.empty_like(slices)
320 + for i, norm in enumerate(self.model.aux_hidden_norms):
321 + normed[:, i, :] = norm(slices[:, i, :])
322 + hidden_states = normed.reshape(hidden_states.shape[0], -1)
323 + result = self.model.fc(hidden_states)
324 + result = self.model.hidden_norm(result)
325 + if needs_squeeze:
326 + result = result.squeeze(0)
327 + return result
328 +
329 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
330 + model_weights = {}
331 + for name, loaded_weight in weights:
332 + if "lm_head" not in name:
333 + name = "model." + name
334 + model_weights[name] = loaded_weight
335 + process_eagle_weight(self, name)
336 +
337 + loader = AutoWeightsLoader(self)
338 + loaded_weight_names = loader.load_weights(model_weights.items())
339 + loaded_weight_names.add("lm_head.weight")
340 + loaded_weight_names.add("model.embed_tokens.weight")
341 + self.model._build_fused_kv_buffers()
342 + return loaded_weight_names
vllm/model_executor/models/qwen3_dflash.py
+64 -41
@@ -410,18 +410,11 @@ class DFlashQwen3Model(nn.Module):
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.
415 -
416 - Must be called after weights are loaded. Stacks the KV-projection
417 - weights, K-norm weights, and RoPE parameters from every attention
418 - layer so that precompute_and_store_context_kv can run one fused
419 - GEMM for all layers at once. Also aliases the weight of the hidden_norm.
420 - """
421 - layers_attn = [layer.self_attn for layer in self.layers]
422 - attn0 = layers_attn[0]
423 - has_bias = attn0.qkv_proj.bias is not None
424 -
413 + def _build_context_kv_buffers(
414 + self,
415 + layers_attn: list[nn.Module],
416 + has_bias: bool,
417 + ) -> None:
418 self._hidden_norm_weight = self.hidden_norm.weight.data
419
420 # KV projection weights: [num_layers * 2 * kv_size, hidden_size]
@@ -439,6 +432,20 @@ class DFlashQwen3Model(nn.Module):
432 [a.k_norm.weight.data for a in layers_attn], dim=0
433 ).contiguous()
434
435 + def _build_fused_kv_buffers(self) -> None:
436 + """Build fused weight buffers for precompute_and_store_context_kv.
437 +
438 + Must be called after weights are loaded. Stacks the KV-projection
439 + weights, K-norm weights, and RoPE parameters from every attention
440 + layer so that precompute_and_store_context_kv can run one fused
441 + GEMM for all layers at once. Also aliases the weight of the hidden_norm.
442 + """
443 + layers_attn = [layer.self_attn for layer in self.layers]
444 + attn0 = layers_attn[0]
445 + has_bias = attn0.qkv_proj.bias is not None
446 +
447 + self._build_context_kv_buffers(layers_attn, has_bias)
448 +
449 # RoPE parameters
450 self._rope_head_size = attn0.rotary_emb.head_size
451 self._rope_cos_sin_cache = attn0.rotary_emb.cos_sin_cache
@@ -468,6 +475,49 @@ class DFlashQwen3Model(nn.Module):
475 # References to inner Attention layers for direct cache writes
476 self._attn_layers = [layer.self_attn.attn for layer in self.layers]
477
478 + def _project_context_kv(
479 + self,
480 + context_states: torch.Tensor,
481 + num_ctx: int,
482 + num_layers: int,
483 + num_kv_heads: int,
484 + head_dim: int,
485 + ) -> tuple[torch.Tensor, torch.Tensor]:
486 + # --- Fused KV projection (one GEMM for all layers) ---
487 + normed_context_states = torch.empty_like(context_states)
488 + ops.rms_norm(
489 + normed_context_states,
490 + context_states,
491 + self._hidden_norm_weight,
492 + self._rms_norm_eps,
493 + )
494 + all_kv_flat = F.linear(
495 + normed_context_states, self._fused_kv_weight, self._fused_kv_bias
496 + )
497 + # Single contiguous copy that separates K/V and transposes to
498 + # layer-major layout. Result: [2, L, num_ctx, nkv, hd] contiguous.
499 + # Indexing dim-0 gives contiguous [L, num_ctx, nkv, hd] for K and V.
500 + all_kv = (
501 + all_kv_flat.view(num_ctx, num_layers, 2, num_kv_heads, head_dim)
502 + .permute(2, 1, 0, 3, 4)
503 + .contiguous()
504 + )
505 + all_k = all_kv[0] # [L, num_ctx, nkv, hd], contiguous
506 + all_v = all_kv[1] # [L, num_ctx, nkv, hd], contiguous
507 + return all_k, all_v
508 +
509 + def _normalize_context_k(self, all_k: torch.Tensor) -> torch.Tensor:
510 + # --- Grouped RMSNorm K across all layers ([L, num_ctx, nkv, hd]) ---
511 + # The weight is selected per layer by the outermost (layer) index.
512 + all_k_normed = torch.empty_like(all_k)
513 + ops.rms_norm(
514 + all_k_normed,
515 + all_k,
516 + self._k_norm_weights,
517 + self._rms_norm_eps,
518 + )
519 + return all_k_normed
520 +
521 def precompute_and_store_context_kv(
522 self,
523 context_states: torch.Tensor,
@@ -499,35 +549,8 @@ class DFlashQwen3Model(nn.Module):
549 hd = self._head_dim
550 nkv = self._num_kv_heads
551
502 - # --- Fused KV projection (one GEMM for all layers) ---
503 - normed_context_states = torch.empty_like(context_states)
504 - ops.rms_norm(
505 - normed_context_states,
506 - context_states,
507 - self._hidden_norm_weight,
508 - self._rms_norm_eps,
509 - )
510 - all_kv_flat = F.linear(
511 - normed_context_states, self._fused_kv_weight, self._fused_kv_bias
512 - )
513 - # Single contiguous copy that separates K/V and transposes to
514 - # layer-major layout. Result: [2, L, num_ctx, nkv, hd] contiguous.
515 - # Indexing dim-0 gives contiguous [L, num_ctx, nkv, hd] for K and V.
516 - all_kv = (
517 - all_kv_flat.view(num_ctx, L, 2, nkv, hd).permute(2, 1, 0, 3, 4).contiguous()
518 - )
519 - all_k = all_kv[0] # [L, num_ctx, nkv, hd], contiguous
520 - all_v = all_kv[1] # [L, num_ctx, nkv, hd], contiguous
521 -
522 - # --- Grouped RMSNorm K across all layers ([L, num_ctx, nkv, hd]) ---
523 - # The weight is selected per layer by the outermost (layer) index.
524 - all_k_normed = torch.empty_like(all_k)
525 - ops.rms_norm(
526 - all_k_normed,
527 - all_k,
528 - self._k_norm_weights,
529 - self._rms_norm_eps,
530 - )
552 + all_k, all_v = self._project_context_kv(context_states, num_ctx, L, nkv, hd)
553 + all_k_normed = self._normalize_context_k(all_k)
554
555 # --- Fused RoPE across all layers ---
556 # View as [L * num_ctx, kv] so RoPE sees one big batch (no copy).
vllm/model_executor/models/registry.py
+1
@@ -591,6 +591,7 @@ _SPECULATIVE_DECODING_MODELS = {
591 "DFlashDraftModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"),
592 "DSparkDraftModel": ("vllm.models.deepseek_v4", "DSparkDeepseekV4ForCausalLM"),
593 "Qwen3DSparkModel": ("qwen3_dspark", "Qwen3DSparkForCausalLM"),
594 + "DFlashLagunaForCausalLM": ("laguna_dflash", "DFlashLagunaForCausalLM"),
595 "PEagleDraftModel": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
596 "PeagleLlamaForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
597 "Eagle3LlamaForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
vllm/v1/spec_decode/llm_base_proposer.py
+2
@@ -28,6 +28,7 @@ from vllm.model_executor.model_loader import get_model
28 from vllm.model_executor.models import supports_multimodal
29 from vllm.model_executor.models.deepseek_eagle3 import Eagle3DeepseekV2ForCausalLM
30 from vllm.model_executor.models.interfaces import SupportsMultiModal
31 +from vllm.model_executor.models.laguna_dflash import DFlashLagunaForCausalLM
32 from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM
33 from vllm.model_executor.models.qwen3_dflash import DFlashQwen3ForCausalLM
34 from vllm.model_executor.models.qwen3_eagle3 import Eagle3Qwen3ForCausalLM
@@ -533,6 +534,7 @@ class SpecDecodeBaseProposer:
534 Eagle3DeepseekV2ForCausalLM,
535 DFlashQwen3ForCausalLM,
536 Eagle3Qwen3ForCausalLM,
537 + DFlashLagunaForCausalLM,
538 ),
539 )
540 target_hidden_states = self.model.combine_hidden_states(