feat(dflash): remap fused QKV weights for Laguna warm-start

src/speculators/convert/dflash/converter.py

@@ -117,7 +117,13 @@ def _build_config(
             # z-lab reads hidden_states[layer_id + 1] (index 0 is the embedding
             # output) while speculators uses the layer id directly.
             # Source: z-lab utils.extract_context_feature.
-            aux_hidden_state_layer_ids = [i + 1 for i in target_layer_ids]
+            # Exclude the last verifier layer — it is always included
+            # implicitly by launch_vllm.py (--include-last-layer) and split
+            # off as verifier_last_hidden_states during training.
+            num_verifier_layers = verifier_config_dict["num_hidden_layers"]
+            aux_hidden_state_layer_ids = [
+                i + 1 for i in target_layer_ids if i + 1 != num_verifier_layers
+            ]
 
         speculators_config = SpeculatorsConfig(
             algorithm="dflash",
@@ -142,6 +148,67 @@ def _build_config(
             speculators_config=speculators_config,
         )
 
+    def _remap_weights(
+        self,
+        weights: dict[str, torch.Tensor],
+        config: DFlashSpeculatorConfig,
+        model: DFlashDraftModel,
+    ) -> dict[str, torch.Tensor]:
+        """Remap checkpoint weights to match DFlashDraftModel's state dict.
+
+        Handles Laguna-style fused ``qkv_proj`` → separate ``q/k/v_proj``,
+        drops ``g_proj`` and ``aux_hidden_norms`` (no DFlash equivalent), and
+        slices ``fc.weight`` when the source has more target layers than needed.
+        """
+        has_fused_qkv = any("qkv_proj" in k for k in weights)
+        if not has_fused_qkv:
+            return weights
+
+        tl = config.transformer_layer_config
+        q_dim = tl.num_attention_heads * tl.head_dim
+        kv_dim = tl.num_key_value_heads * tl.head_dim
+
+        remapped: dict[str, torch.Tensor] = {}
+        dropped: list[str] = []
+
+        for key, tensor in weights.items():
+            if "qkv_proj" in key:
+                projection_keys = [
+                    key.replace("qkv_proj", proj)
+                    for proj in ("q_proj", "k_proj", "v_proj")
+                ]
+                conflicts = [pk for pk in projection_keys if pk in weights]
+                if conflicts:
+                    raise ValueError(
+                        f"Checkpoint contains both fused qkv_proj and separate "
+                        f"projection keys: {conflicts}"
+                    )
+                q, k, v = tensor.split([q_dim, kv_dim, kv_dim], dim=0)
+                remapped[projection_keys[0]] = q
+                remapped[projection_keys[1]] = k
+                remapped[projection_keys[2]] = v
+            elif ".g_proj." in key or key.startswith("aux_hidden_norms."):
+                dropped.append(key)
+            elif key == "fc.weight":
+                model_fc_dim = model.fc.in_features
+                if tensor.shape[1] > model_fc_dim:
+                    logger.info(
+                        f"Slicing fc.weight from {tensor.shape[1]} to {model_fc_dim}"
+                    )
+                    remapped[key] = tensor[:, :model_fc_dim]
+                else:
+                    remapped[key] = tensor
+            else:
+                remapped[key] = tensor
+
+        if dropped:
+            logger.info(f"Dropped {len(dropped)} incompatible keys: {dropped}")
+        logger.info(
+            f"Remapped {sum(1 for k in weights if 'qkv_proj' in k)} fused qkv_proj "
+            f"→ separate q/k/v_proj"
+        )
+        return remapped
+
     def _save(
         self,
         config: DFlashSpeculatorConfig,
@@ -151,6 +218,7 @@ def _save(
         model = DFlashDraftModel(config=config)
 
         body = {k: v for k, v in weights.items() if k not in ("t2d", "d2t")}
+        body = self._remap_weights(body, config, model)
         missing, unexpected = model.load_state_dict(body, strict=False)
         if unexpected:
             raise ValueError(

src/speculators/models/dflash/core.py

@@ -1,4 +1,5 @@
 import logging
+from copy import deepcopy
 from typing import ClassVar
 
 import torch
@@ -95,7 +96,19 @@ def __init__(
             config.transformer_layer_config.hidden_size,
             eps=config.transformer_layer_config.rms_norm_eps,  # type: ignore[arg-type]
         )
-        self.rotary_emb = Qwen3RotaryEmbedding(config.transformer_layer_config)  # type: ignore[arg-type]
+        rotary_config = config.transformer_layer_config
+        rope_params = getattr(rotary_config, "rope_parameters", None)
+        if rope_params and "sliding_attention" in rope_params:
+            if self.uses_full_attn:
+                logger.warning(
+                    "Flattening nested rope_parameters to the sliding_attention "
+                    "variant, but this model has %d full-attention layer(s). "
+                    "Full-attention layers may use incorrect rope scaling.",
+                    num_draft_layers - len(self.sliding_window_indices),
+                )
+            rotary_config = deepcopy(rotary_config)
+            rotary_config.rope_parameters = rope_params["sliding_attention"]
+        self.rotary_emb = Qwen3RotaryEmbedding(rotary_config)  # type: ignore[arg-type]
 
         self.fc = nn.Linear(
             len(self.target_layer_ids) * config.transformer_layer_config.hidden_size,

tests/unit/convert/test_dflash_converter.py

@@ -1,31 +1,40 @@
-"""Unit tests for DFlashConverter config building."""
+"""Unit tests for DFlashConverter config building and weight remapping."""
 
+import logging
 from unittest.mock import patch
 
 import pytest
+import torch
 from transformers import Qwen3Config
 
 from speculators.config import SpeculatorsConfig, VerifierConfig
 from speculators.convert.dflash.converter import DFlashConverter
-from speculators.models.dflash import DFlashSpeculatorConfig
+from speculators.models.dflash import DFlashDraftModel, DFlashSpeculatorConfig
 from speculators.proposals.greedy import GreedyTokenProposalConfig
 
+_HIDDEN = 16
+_NUM_HEADS = 2
+_NUM_KV_HEADS = 1
+_HEAD_DIM = 8
+_Q_DIM = _NUM_HEADS * _HEAD_DIM  # 16
+_KV_DIM = _NUM_KV_HEADS * _HEAD_DIM  # 8
 
-def _tiny_dflash_config():
+
+def _tiny_dflash_config(num_aux_layers=1):
     return DFlashSpeculatorConfig(
         transformer_layer_config=Qwen3Config(
             vocab_size=32,
-            hidden_size=16,
+            hidden_size=_HIDDEN,
             intermediate_size=32,
             num_hidden_layers=1,
-            num_attention_heads=2,
-            num_key_value_heads=1,
-            head_dim=8,
+            num_attention_heads=_NUM_HEADS,
+            num_key_value_heads=_NUM_KV_HEADS,
+            head_dim=_HEAD_DIM,
             max_position_embeddings=32,
         ),
         draft_vocab_size=32,
         block_size=4,
-        aux_hidden_state_layer_ids=[0],
+        aux_hidden_state_layer_ids=list(range(num_aux_layers)),
         mask_token_id=1,
         speculators_config=SpeculatorsConfig(
             algorithm="dflash",
@@ -63,7 +72,11 @@ class TestBuildConfig:
     @patch("speculators.convert.dflash.converter.PretrainedConfig.get_config_dict")
     def test_happy_path(self, mock_get_config):
         mock_get_config.return_value = (
-            {"hidden_size": 4096, "architectures": ["Qwen3ForCausalLM"]},
+            {
+                "hidden_size": 4096,
+                "num_hidden_layers": 36,
+                "architectures": ["Qwen3ForCausalLM"],
+            },
             None,
         )
         config = DFlashConverter()._build_config(
@@ -82,6 +95,23 @@ def test_happy_path(self, mock_get_config):
         assert config.transformer_layer_config.num_hidden_layers == 5
         assert not hasattr(config.transformer_layer_config, "dflash_config")
 
+    @patch("speculators.convert.dflash.converter.PretrainedConfig.get_config_dict")
+    def test_excludes_last_verifier_layer(self, mock_get_config):
+        mock_get_config.return_value = (
+            {
+                "hidden_size": 4096,
+                "num_hidden_layers": 34,
+                "architectures": ["Qwen3ForCausalLM"],
+            },
+            None,
+        )
+        # target_layer_ids [1, 9, 17, 25, 33] → +1 → [2, 10, 18, 26, 34]
+        # but 34 == num_hidden_layers, so it should be excluded
+        config = DFlashConverter()._build_config(
+            _source_config(), "Qwen/Qwen3-8B", None
+        )
+        assert config.aux_hidden_state_layer_ids == [2, 10, 18, 26]
+
     @patch("speculators.convert.dflash.converter.PretrainedConfig.get_config_dict")
     def test_explicit_aux_layer_ids_override(self, mock_get_config):
         mock_get_config.return_value = ({"hidden_size": 4096}, None)
@@ -98,12 +128,122 @@ def test_hidden_size_mismatch_raises(self, mock_get_config):
 
     @patch("speculators.convert.dflash.converter.PretrainedConfig.get_config_dict")
     def test_missing_target_layer_ids_raises(self, mock_get_config):
-        mock_get_config.return_value = ({"hidden_size": 4096}, None)
+        mock_get_config.return_value = (
+            {"hidden_size": 4096, "num_hidden_layers": 36},
+            None,
+        )
         source = _source_config(dflash_config={"mask_token_id": 151669})
         with pytest.raises(ValueError, match="target_layer_ids"):
             DFlashConverter()._build_config(source, "Qwen/Qwen3-8B", None)
 
 
+class TestRemapWeights:
+    def _make_fused_weights(self):
+        qkv = torch.randn(_Q_DIM + 2 * _KV_DIM, _HIDDEN)
+        return {
+            "layers.0.self_attn.qkv_proj.weight": qkv,
+            "layers.0.self_attn.g_proj.weight": torch.randn(2, _HIDDEN),
+            "layers.0.self_attn.o_proj.weight": torch.randn(_HIDDEN, _Q_DIM),
+            "aux_hidden_norms.0.weight": torch.randn(_HIDDEN),
+            "fc.weight": torch.randn(_HIDDEN, _HIDDEN * 2),
+            "norm.weight": torch.randn(_HIDDEN),
+        }
+
+    def test_splits_fused_qkv(self):
+        config = _tiny_dflash_config()
+        model = DFlashDraftModel(config=config)
+        weights = self._make_fused_weights()
+        qkv = weights["layers.0.self_attn.qkv_proj.weight"]
+
+        remapped = DFlashConverter()._remap_weights(weights, config, model)
+
+        assert "layers.0.self_attn.qkv_proj.weight" not in remapped
+        assert torch.equal(remapped["layers.0.self_attn.q_proj.weight"], qkv[:_Q_DIM])
+        assert torch.equal(
+            remapped["layers.0.self_attn.k_proj.weight"],
+            qkv[_Q_DIM : _Q_DIM + _KV_DIM],
+        )
+        assert torch.equal(
+            remapped["layers.0.self_attn.v_proj.weight"],
+            qkv[_Q_DIM + _KV_DIM :],
+        )
+
+    def test_drops_g_proj_and_aux_hidden_norms(self):
+        config = _tiny_dflash_config()
+        model = DFlashDraftModel(config=config)
+        remapped = DFlashConverter()._remap_weights(
+            self._make_fused_weights(), config, model
+        )
+        assert not any("g_proj" in k for k in remapped)
+        assert not any("aux_hidden_norms" in k for k in remapped)
+
+    def test_slices_fc_weight(self):
+        config = _tiny_dflash_config(num_aux_layers=1)
+        model = DFlashDraftModel(config=config)
+        weights = self._make_fused_weights()
+        # fc from checkpoint is wider than model expects
+        wide_fc = torch.randn(_HIDDEN, _HIDDEN * 3)
+        weights["fc.weight"] = wide_fc
+
+        remapped = DFlashConverter()._remap_weights(weights, config, model)
+        assert remapped["fc.weight"].shape[1] == model.fc.in_features
+        assert torch.equal(remapped["fc.weight"], wide_fc[:, : model.fc.in_features])
+
+    def test_passthrough_when_no_fused_qkv(self):
+        config = _tiny_dflash_config()
+        model = DFlashDraftModel(config=config)
+        weights = {"norm.weight": torch.randn(_HIDDEN)}
+        remapped = DFlashConverter()._remap_weights(weights, config, model)
+        assert remapped is weights
+
+    def test_preserves_other_keys(self):
+        config = _tiny_dflash_config()
+        model = DFlashDraftModel(config=config)
+        weights = self._make_fused_weights()
+        remapped = DFlashConverter()._remap_weights(weights, config, model)
+        assert "layers.0.self_attn.o_proj.weight" in remapped
+        assert "norm.weight" in remapped
+
+    def test_mixed_fused_and_separate_qkv_raises(self):
+        config = _tiny_dflash_config()
+        model = DFlashDraftModel(config=config)
+        weights = self._make_fused_weights()
+        weights["layers.0.self_attn.q_proj.weight"] = torch.randn(_Q_DIM, _HIDDEN)
+        with pytest.raises(ValueError, match="both fused qkv_proj and separate"):
+            DFlashConverter()._remap_weights(weights, config, model)
+
+
+class TestRopeParameters:
+    _ROPE_SLIDING = {
+        "rope_type": "default",
+        "rope_theta": 10000.0,
+    }
+
+    def test_nested_rope_params_warns_with_full_attn(self, caplog):
+        config = _tiny_dflash_config()
+        tl = config.transformer_layer_config
+        tl.layer_types = ["full_attention"]
+        tl.rope_parameters = {
+            "sliding_attention": dict(self._ROPE_SLIDING),
+            "full_attention": {"rope_type": "default", "rope_theta": 1000000.0},
+        }
+        with caplog.at_level(logging.WARNING, logger="speculators.models.dflash.core"):
+            DFlashDraftModel(config=config)
+        assert any("full-attention layer" in r.message for r in caplog.records)
+
+    def test_nested_rope_params_no_warn_all_sliding(self, caplog):
+        config = _tiny_dflash_config()
+        tl = config.transformer_layer_config
+        tl.layer_types = ["sliding_attention"]
+        tl.sliding_window = 512
+        tl.rope_parameters = {
+            "sliding_attention": dict(self._ROPE_SLIDING),
+        }
+        with caplog.at_level(logging.WARNING, logger="speculators.models.dflash.core"):
+            DFlashDraftModel(config=config)
+        assert not any("full-attention layer" in r.message for r in caplog.records)
+
+
 class TestSave:
     def test_missing_draft_weights_raise(self, tmp_path):
         # No source weights: every draft-body weight (fc, norm, hidden_norm,