Add Laguna XS.2.1 DFlash drafter support - #46853
Conversation
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in PRs do not trigger a full CI run by default. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can either: Add If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. Agent GuidelinesIMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban. 🚀 |
benchislett
left a comment
There was a problem hiding this comment.
Might be worth rebasing on top of #46104 which adds support for all-sliding-window DFlash.
Is there really enough complexity here to warrant a completely new model definition here? Ideally we'd be reusing a decent chunk of components and complexity from the existing dflash implementation.
benchislett
left a comment
There was a problem hiding this comment.
Also, needs some kind of E2E test coverage or else it will be broken for sure
Thanks @benchislett, I looked into #46104. Once it lands, I'm happy to rebase on top of it. I also tested cherry-picking #46104's non-merge commits locally and the Laguna-side adaptation looks fairly small: remove the local
The part that remains model-specific is our draft module definition. We don't really fit existing DFlash implementation assumptions. Laguna has custom attention gating, |
Happy to track adding one as a follow-up once the HF checkpoints are public. At the moment any test would need to be env-gated and would not run in CI by default. The changes here are scoped to the new Laguna DFlash draft model and should not affect existing DFlash users. I can add an optional smoke test now if you still prefer having it in this PR. |
|
@adamkbaranowski it still seems worthwhile to inherit from DFlash and reduce duplication in the parts that are shared. |
42f87d2 to
abd6f10
Compare
|
This pull request has merge conflicts that must be resolved before it can be |
d31efa1 to
e0ede93
Compare
Signed-off-by: Adam Baranowski <adam.baranowski@poolside.ai>
Signed-off-by: Adam Baranowski <adam.baranowski@poolside.ai>
Signed-off-by: Adam Baranowski <adam.baranowski@poolside.ai>
After vLLM vllm-project#46104, DFlashQwen3Model.embed_input_ids expects mask_token_id, mask_embedding, and has_separate_mask_embedding to exist. Laguna DFlash reuses that control flow without calling DFlashQwen3Model.__init__, so initialize the same mask state locally. Signed-off-by: Adam Baranowski <adam.baranowski@poolside.ai>
a775523 to
5ac9e27
Compare
Signed-off-by: Adam Baranowski <adam.baranowski@poolside.ai>
Signed-off-by: Adam Baranowski <adam.baranowski@poolside.ai>
Signed-off-by: Adam Baranowski <adam.baranowski@poolside.ai>
Signed-off-by: Adam Baranowski <adam.baranowski@poolside.ai>
Adds DSpark (semi-autoregressive block speculative decoding) support for Gemma4-12B, for the deepseek-ai/dspark_gemma4_12b_block7 checkpoint. Following the Laguna DFlash review (vllm-project#46853), this is a thin reuse of the existing stacks rather than a standalone model: - Gemma4DSparkModel subclasses DFlashQwen3Model and INHERITS the fused context-KV precompute (precompute_and_store_context_kv + _normalize_context_k + RoPE + cache-write). Gemma4's attention_k_eq_v is isolated behind the _build_context_kv_buffers / _project_context_kv hooks: the context is projected once through k_proj for all layers, K is returned raw (the inherited path applies k_norm + RoPE) and V = v_norm(that projection) with no RoPE. Only forward (sandwich norms + layer_scalar) and embed_input_ids (sqrt(hidden) scaling) are overridden. - Gemma4DSparkForCausalLM subclasses Qwen3DSparkForCausalLM: markov_embed/ markov_bias, embed_input_ids, forward, compute_logits and precompute delegation are inherited; only the self-contained checkpoint bits (own embed_tokens/lm_head, final-logit softcap, no draft->target id remap) and the flat weight loader are overridden. - The Gemma4 attention/decoder-layer/MLP reuse the native Gemma4 MTP draft (Gemma4MTPAttention/Gemma4MTPDecoderLayer/Gemma4MLP); the DSpark attention only adds the K/V projections + norms. registry.py registers Gemma4DSparkModel; config/speculative.py detects it as method=dspark and normalizes target_layer_ids -> dspark_target_layer_ids and block_size -> n_predict. Validated on 1x H100 (TP=1): draft loads (67 params), server healthy, greedy generation correct, and DSpark drafting/acceptance metrics recorded through the inherited fused precompute (block-7, ~3 mean acceptance length). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Adds DSpark (semi-autoregressive block speculative decoding) support for Gemma4-12B, for the deepseek-ai/dspark_gemma4_12b_block7 checkpoint. Following the Laguna DFlash review (vllm-project#46853), this is a thin reuse of the existing stacks rather than a standalone model: - Gemma4DSparkModel subclasses DFlashQwen3Model and INHERITS the fused context-KV precompute (precompute_and_store_context_kv + _normalize_context_k + RoPE + cache-write). Gemma4's attention_k_eq_v is isolated behind the _build_context_kv_buffers / _project_context_kv hooks: the context is projected once through k_proj for all layers, K is returned raw (the inherited path applies k_norm + RoPE) and V = v_norm(that projection) with no RoPE. Only forward (sandwich norms + layer_scalar) and embed_input_ids (sqrt(hidden) scaling) are overridden. - Gemma4DSparkForCausalLM subclasses Qwen3DSparkForCausalLM: markov_embed/ markov_bias, embed_input_ids, forward, compute_logits and precompute delegation are inherited; only the self-contained checkpoint bits (own embed_tokens/lm_head, final-logit softcap, no draft->target id remap) and the flat weight loader are overridden. - The Gemma4 attention/decoder-layer/MLP reuse the native Gemma4 MTP draft (Gemma4MTPAttention/Gemma4MTPDecoderLayer/Gemma4MLP); the DSpark attention only adds the K/V projections + norms. registry.py registers Gemma4DSparkModel; config/speculative.py detects it as method=dspark and normalizes target_layer_ids -> dspark_target_layer_ids and block_size -> n_predict. Validated on 1x H100 (TP=1): draft loads (67 params), server healthy, greedy generation correct, and DSpark drafting/acceptance metrics recorded through the inherited fused precompute (block-7, ~3 mean acceptance length). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Adds DSpark (semi-autoregressive block speculative decoding) support for Gemma4-12B, for the deepseek-ai/dspark_gemma4_12b_block7 checkpoint. Following the Laguna DFlash review (vllm-project#46853), this is a thin reuse of the existing stacks rather than a standalone model: - Gemma4DSparkModel subclasses DFlashQwen3Model and INHERITS the fused context-KV precompute (precompute_and_store_context_kv + _normalize_context_k + RoPE + cache-write). Gemma4's attention_k_eq_v is isolated behind the _build_context_kv_buffers / _project_context_kv hooks: the context is projected once through k_proj for all layers, K is returned raw (the inherited path applies k_norm + RoPE) and V = v_norm(that projection) with no RoPE. Only forward (sandwich norms + layer_scalar) and embed_input_ids (sqrt(hidden) scaling) are overridden. - Gemma4DSparkForCausalLM subclasses Qwen3DSparkForCausalLM: markov_embed/ markov_bias, embed_input_ids, forward, compute_logits and precompute delegation are inherited; only the self-contained checkpoint bits (own embed_tokens/lm_head, final-logit softcap, no draft->target id remap) and the flat weight loader are overridden. - The Gemma4 attention/decoder-layer/MLP reuse the native Gemma4 MTP draft (Gemma4MTPAttention/Gemma4MTPDecoderLayer/Gemma4MLP); the DSpark attention only adds the K/V projections + norms. registry.py registers Gemma4DSparkModel; config/speculative.py detects it as method=dspark and normalizes target_layer_ids -> dspark_target_layer_ids and block_size -> n_predict. Validated on 1x H100 (TP=1): draft loads (67 params), server healthy, greedy generation correct, and DSpark drafting/acceptance metrics recorded through the inherited fused precompute (block-7, ~3 mean acceptance length). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: DiegoCao <DiegoCao@users.noreply.github.com>
Adds DSpark (semi-autoregressive block speculative decoding) support for Gemma4-12B, for the deepseek-ai/dspark_gemma4_12b_block7 checkpoint. Following the Laguna DFlash review (vllm-project#46853), this is a thin reuse of the existing stacks rather than a standalone model: - Gemma4DSparkModel subclasses DFlashQwen3Model and INHERITS the fused context-KV precompute (precompute_and_store_context_kv + _normalize_context_k + RoPE + cache-write). Gemma4's attention_k_eq_v is isolated behind the _build_context_kv_buffers / _project_context_kv hooks: the context is projected once through k_proj for all layers, K is returned raw (the inherited path applies k_norm + RoPE) and V = v_norm(that projection) with no RoPE. Only forward (sandwich norms + layer_scalar) and embed_input_ids (sqrt(hidden) scaling) are overridden. - Gemma4DSparkForCausalLM subclasses Qwen3DSparkForCausalLM: markov_embed/ markov_bias, embed_input_ids, forward, compute_logits and precompute delegation are inherited; only the self-contained checkpoint bits (own embed_tokens/lm_head, final-logit softcap, no draft->target id remap) and the flat weight loader are overridden. - The Gemma4 attention/decoder-layer/MLP reuse the native Gemma4 MTP draft (Gemma4MTPAttention/Gemma4MTPDecoderLayer/Gemma4MLP); the DSpark attention only adds the K/V projections + norms. registry.py registers Gemma4DSparkModel; config/speculative.py detects it as method=dspark and normalizes target_layer_ids -> dspark_target_layer_ids and block_size -> n_predict. Validated on 1x H100 (TP=1): draft loads (67 params), server healthy, greedy generation correct, and DSpark drafting/acceptance metrics recorded through the inherited fused precompute (block-7, ~3 mean acceptance length). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: DiegoCao <DiegoCao@users.noreply.github.com>
Adds DSpark (semi-autoregressive block speculative decoding) support for Gemma4-12B, for the deepseek-ai/dspark_gemma4_12b_block7 checkpoint. Following the Laguna DFlash review (vllm-project#46853), this is a thin reuse of the existing stacks rather than a standalone model: - Gemma4DSparkModel subclasses DFlashQwen3Model and INHERITS the fused context-KV precompute (precompute_and_store_context_kv + _normalize_context_k + RoPE + cache-write). Gemma4's attention_k_eq_v is isolated behind the _build_context_kv_buffers / _project_context_kv hooks: the context is projected once through k_proj for all layers, K is returned raw (the inherited path applies k_norm + RoPE) and V = v_norm(that projection) with no RoPE. Only forward (sandwich norms + layer_scalar) and embed_input_ids (sqrt(hidden) scaling) are overridden. - Gemma4DSparkForCausalLM subclasses Qwen3DSparkForCausalLM: markov_embed/ markov_bias, embed_input_ids, forward, compute_logits and precompute delegation are inherited; only the self-contained checkpoint bits (own embed_tokens/lm_head, final-logit softcap, no draft->target id remap) and the flat weight loader are overridden. - The Gemma4 attention/decoder-layer/MLP reuse the native Gemma4 MTP draft (Gemma4MTPAttention/Gemma4MTPDecoderLayer/Gemma4MLP); the DSpark attention only adds the K/V projections + norms. registry.py registers Gemma4DSparkModel; config/speculative.py detects it as method=dspark and normalizes target_layer_ids -> dspark_target_layer_ids and block_size -> n_predict. Validated on 1x H100 (TP=1): draft loads (67 params), server healthy, greedy generation correct, and DSpark drafting/acceptance metrics recorded through the inherited fused precompute (block-7, ~3 mean acceptance length). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: DiegoCao <DiegoCao@users.noreply.github.com>
Adds DSpark (semi-autoregressive block speculative decoding) support for Gemma4-12B, for the deepseek-ai/dspark_gemma4_12b_block7 checkpoint. Following the Laguna DFlash review (vllm-project#46853), this is a thin reuse of the existing stacks rather than a standalone model: - Gemma4DSparkModel subclasses DFlashQwen3Model and INHERITS the fused context-KV precompute (precompute_and_store_context_kv + _normalize_context_k + RoPE + cache-write). Gemma4's attention_k_eq_v is isolated behind the _build_context_kv_buffers / _project_context_kv hooks: the context is projected once through k_proj for all layers, K is returned raw (the inherited path applies k_norm + RoPE) and V = v_norm(that projection) with no RoPE. Only forward (sandwich norms + layer_scalar) and embed_input_ids (sqrt(hidden) scaling) are overridden. - Gemma4DSparkForCausalLM subclasses Qwen3DSparkForCausalLM: markov_embed/ markov_bias, embed_input_ids, forward, compute_logits and precompute delegation are inherited; only the self-contained checkpoint bits (own embed_tokens/lm_head, final-logit softcap, no draft->target id remap) and the flat weight loader are overridden. - The Gemma4 attention/decoder-layer/MLP reuse the native Gemma4 MTP draft (Gemma4MTPAttention/Gemma4MTPDecoderLayer/Gemma4MLP); the DSpark attention only adds the K/V projections + norms. registry.py registers Gemma4DSparkModel; config/speculative.py detects it as method=dspark and normalizes target_layer_ids -> dspark_target_layer_ids and block_size -> n_predict. Validated on 1x H100 (TP=1): draft loads (67 params), server healthy, greedy generation correct, and DSpark drafting/acceptance metrics recorded through the inherited fused precompute (block-7, ~3 mean acceptance length). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: DiegoCao <DiegoCao@users.noreply.github.com>
Adds DSpark (semi-autoregressive block speculative decoding) support for Gemma4-12B, for the deepseek-ai/dspark_gemma4_12b_block7 checkpoint. Following the Laguna DFlash review (vllm-project#46853), this is a thin reuse of the existing stacks rather than a standalone model: - Gemma4DSparkModel subclasses DFlashQwen3Model and INHERITS the fused context-KV precompute (precompute_and_store_context_kv + _normalize_context_k + RoPE + cache-write). Gemma4's attention_k_eq_v is isolated behind the _build_context_kv_buffers / _project_context_kv hooks: the context is projected once through k_proj for all layers, K is returned raw (the inherited path applies k_norm + RoPE) and V = v_norm(that projection) with no RoPE. Only forward (sandwich norms + layer_scalar) and embed_input_ids (sqrt(hidden) scaling) are overridden. - Gemma4DSparkForCausalLM subclasses Qwen3DSparkForCausalLM: markov_embed/ markov_bias, embed_input_ids, forward, compute_logits and precompute delegation are inherited; only the self-contained checkpoint bits (own embed_tokens/lm_head, final-logit softcap, no draft->target id remap) and the flat weight loader are overridden. - The Gemma4 attention/decoder-layer/MLP reuse the native Gemma4 MTP draft (Gemma4MTPAttention/Gemma4MTPDecoderLayer/Gemma4MLP); the DSpark attention only adds the K/V projections + norms. registry.py registers Gemma4DSparkModel; config/speculative.py detects it as method=dspark and normalizes target_layer_ids -> dspark_target_layer_ids and block_size -> n_predict. Validated on 1x H100 (TP=1): draft loads (67 params), server healthy, greedy generation correct, and DSpark drafting/acceptance metrics recorded through the inherited fused precompute (block-7, ~3 mean acceptance length). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: DiegoCao <DiegoCao@users.noreply.github.com>
Signed-off-by: Adam Baranowski <adam.baranowski@poolside.ai>
Signed-off-by: Adam Baranowski <adam.baranowski@poolside.ai> Signed-off-by: mayuyuace <qiming1.zhang@intel.com>
Signed-off-by: Adam Baranowski <adam.baranowski@poolside.ai>
Signed-off-by: Adam Baranowski <adam.baranowski@poolside.ai>
Signed-off-by: Adam Baranowski <adam.baranowski@poolside.ai>
Signed-off-by: Adam Baranowski <adam.baranowski@poolside.ai> Signed-off-by: root <root@smci355-ccs-aus-m02-09.cs-aus.dcgpu>
Purpose
Add vLLM model support for Poolside Laguna XS.2.1 DFlash draft checkpoints.
This registers
DFlashLagunaForCausalLMand implements the Laguna-specific DFlash drafter used withmethod="dflash"speculative decoding. The drafter consumes verifier auxiliary hidden states from the Laguna target model, precomputes context K/V for DFlash attention, and shares the target model token embeddings and LM head.Changes:
vllm/model_executor/models/laguna_dflash.py.DFlashLagunaForCausalLMin the model registry.embed_tokensandlm_head.Test Plan
Test Result
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.