diff --git a/tests/test_dflash2.py b/tests/test_dflash2.py index efc9c8a0..be4fd10c 100644 --- a/tests/test_dflash2.py +++ b/tests/test_dflash2.py @@ -157,6 +157,56 @@ def _batch(seed=0, all_masked=False): class TestDFlash2Config(unittest.TestCase): + def test_mla_mode_builds_k3_mla_attention(self): + from torchspec.models.draft.dspark import K3DSparkMLAAttention + + config = DFlash2Config( + **_tiny_config_kwargs( + num_hidden_layers=3, + attention_mode="mla", + layer_types=["sliding_attention", "sliding_attention", "full_attention"], + sliding_window=2048, + is_causal=False, + q_lora_rank=4, + kv_lora_rank=4, + qk_nope_head_dim=2, + qk_rope_head_dim=2, + v_head_dim=2, + rope_parameters={"rope_type": "default", "rope_theta": 10000.0}, + ) + ) + + model = DFlash2DraftModel(config) + + self.assertEqual(config.attention_mode, "mla") + self.assertTrue( + all(isinstance(layer.self_attn, K3DSparkMLAAttention) for layer in model.layers) + ) + + def test_mixed_layer_types_build_per_layer_mask_policies(self): + config = DFlash2Config( + **_tiny_config_kwargs( + num_hidden_layers=3, + layer_types=["sliding_attention", "sliding_attention", "full_attention"], + sliding_window=2048, + is_causal=False, + ) + ) + model = DFlash2Model( + DFlash2DraftModel(config), + block_size=config.block_size, + num_anchors=1, + ) + + self.assertEqual( + [model._block_mask_options_for_layer(i) for i in range(3)], + [ + {"is_causal": False, "sliding_window": 2048}, + {"is_causal": False, "sliding_window": 2048}, + {"is_causal": False, "sliding_window": None}, + ], + ) + def test_repository_config_dispatches_to_dflash2(self): config_path = ROOT / "torchspec" / "config" / "dflash2_draft_config.json" self.assertTrue(config_path.is_file()) @@ -372,7 +422,6 @@ def test_invalid_config_is_rejected(self): invalid_layer_types = ( (["full_attention", "full_attention"], 1, "one entry"), (["linear_attention"], 1, "Unsupported"), - (["full_attention", "sliding_attention"], 2, "mixed"), ) for layer_types, num_hidden_layers, error in invalid_layer_types: with self.subTest(layer_types=layer_types): diff --git a/torchspec/models/dflash.py b/torchspec/models/dflash.py index c4c245da..da9061ca 100644 --- a/torchspec/models/dflash.py +++ b/torchspec/models/dflash.py @@ -267,6 +267,17 @@ def _create_noise_embed( def _block_mask_options(self) -> dict: return {} + def _block_mask_options_for_layer(self, layer_id: int) -> dict: + """Return the attention-mask options for one draft layer. + + DFlash historically used one mask for every layer. Keeping this + method separate from ``_block_mask_options`` lets model variants use + per-layer attention schedules without changing the public training + wrapper API. + """ + del layer_id + return self._block_mask_options() + def _compute_logits( self, draft_hidden: torch.Tensor, @@ -331,21 +342,32 @@ def _draft_backbone( block_mask = None if device.type == "cuda": - mask_mod = _create_dflash_mask_mod( - anchor_positions=anchor_positions, - block_keep_mask=block_keep_mask, - ctx_len=seq_len, - block_size=self.block_size, - **self._block_mask_options(), - ) - block_mask = compile_friendly_create_block_mask( - mask_mod=mask_mod, - B=bsz, - H=None, - Q_LEN=draft_len, - KV_LEN=kv_len, - device=device, - ) + # A mixed full/SWA schedule needs a distinct BlockMask for each + # attention policy. Cache equal policies so the common uniform + # schedule still creates exactly one mask. + mask_cache = {} + block_masks = [] + for layer_id in range(self.draft_model.num_layers): + options = self._block_mask_options_for_layer(layer_id) + cache_key = tuple(sorted(options.items())) + if cache_key not in mask_cache: + mask_mod = _create_dflash_mask_mod( + anchor_positions=anchor_positions, + block_keep_mask=block_keep_mask, + ctx_len=seq_len, + block_size=self.block_size, + **options, + ) + mask_cache[cache_key] = compile_friendly_create_block_mask( + mask_mod=mask_mod, + B=bsz, + H=None, + Q_LEN=draft_len, + KV_LEN=kv_len, + device=device, + ) + block_masks.append(mask_cache[cache_key]) + block_mask = block_masks[0] if len(mask_cache) == 1 else block_masks # 6. Draft model forward — pass embeddings directly draft_hidden = self.draft_model( diff --git a/torchspec/models/dflash2.py b/torchspec/models/dflash2.py index 99d43664..d968f889 100644 --- a/torchspec/models/dflash2.py +++ b/torchspec/models/dflash2.py @@ -48,28 +48,41 @@ def __init__(self, *args, selector_loss_alpha: float = 1.0, **kwargs): attention_types = set(layer_types or ["full_attention"]) if not attention_types <= {"full_attention", "sliding_attention"}: raise ValueError(f"Unsupported DFlash2 layer types: {sorted(attention_types)}") - if len(attention_types) > 1: - raise ValueError("DFlash2 training does not support mixed full and sliding layers") - - uses_sliding_window = "sliding_attention" in attention_types explicit_causality = getattr(config, "is_causal", None) - self.attention_is_causal = ( - uses_sliding_window if explicit_causality is None else bool(explicit_causality) - ) configured_window = getattr(config, "sliding_window", None) - self.sliding_window = ( - int(configured_window) - if uses_sliding_window and configured_window is not None - else None - ) - if self.sliding_window is not None and self.sliding_window < 1: - raise ValueError(f"sliding_window must be positive, got {self.sliding_window}") + if "sliding_attention" in attention_types: + if configured_window is None: + raise ValueError( + "DFlash2 sliding_attention layers require an explicit positive sliding_window" + ) + configured_window = int(configured_window) + if configured_window < 1: + raise ValueError(f"sliding_window must be positive, got {configured_window}") + + normalized_layer_types = layer_types or ["full_attention"] * config.num_hidden_layers + self.layer_block_mask_options = [] + for layer_type in normalized_layer_types: + uses_sliding_window = layer_type == "sliding_attention" + self.layer_block_mask_options.append( + { + "is_causal": ( + uses_sliding_window + if explicit_causality is None + else bool(explicit_causality) + ), + "sliding_window": configured_window if uses_sliding_window else None, + } + ) + + # Retain the uniform-policy attributes for callers that inspect them. + self.attention_is_causal = self.layer_block_mask_options[0]["is_causal"] + self.sliding_window = self.layer_block_mask_options[0]["sliding_window"] def _block_mask_options(self) -> dict: - return { - "is_causal": self.attention_is_causal, - "sliding_window": self.sliding_window, - } + return self.layer_block_mask_options[0] + + def _block_mask_options_for_layer(self, layer_id: int) -> dict: + return self.layer_block_mask_options[layer_id] def _compute_logits( self, diff --git a/torchspec/models/draft/dflash.py b/torchspec/models/draft/dflash.py index aa2cec0a..af36c1b9 100644 --- a/torchspec/models/draft/dflash.py +++ b/torchspec/models/draft/dflash.py @@ -446,7 +446,8 @@ def forward( context_feature: [B, ctx_len, D] — projected context from target draft_position_ids: [B, draft_len] context_position_ids: [B, ctx_len] - block_mask: FlexAttention BlockMask + block_mask: FlexAttention BlockMask, or a per-layer sequence of + BlockMasks for mixed attention schedules noise_embedding: [B, draft_len, D] — pre-computed embeddings (from training wrapper) Returns: @@ -457,13 +458,16 @@ def forward( else: draft_hidden = self.embed_tokens(draft_input_ids).to(context_feature.dtype) - for layer in self.layers: + for layer_id, layer in enumerate(self.layers): + layer_block_mask = block_mask + if isinstance(block_mask, (list, tuple)): + layer_block_mask = block_mask[layer_id] draft_hidden = layer( draft_hidden=draft_hidden, context_hidden=context_feature, draft_position_ids=draft_position_ids, context_position_ids=context_position_ids, - block_mask=block_mask, + block_mask=layer_block_mask, ) return self.final_norm(draft_hidden) diff --git a/torchspec/models/draft/dflash2.py b/torchspec/models/draft/dflash2.py index c48c210a..1488d982 100644 --- a/torchspec/models/draft/dflash2.py +++ b/torchspec/models/draft/dflash2.py @@ -44,6 +44,7 @@ ) _NESTED_CONFIG_FIELDS = ( + "attention_mode", "block_size", "conv_group_size", "conv_kernel_size", @@ -101,6 +102,15 @@ def __init__( conv_group_size: int = 16, selector_rank: int = 256, selector_top_k: int = 16, + attention_mode: str = "gqa", + q_lora_rank: int | None = 1536, + kv_lora_rank: int = 512, + qk_nope_head_dim: int = 128, + qk_rope_head_dim: int = 64, + v_head_dim: int = 128, + mla_use_output_gate: bool = False, + rope_scaling: dict | None = None, + rope_parameters: dict | None = None, input_embedding_scale: float = 1.0, output_multiplier: float = 1.0, final_logit_softcapping: float | None = None, @@ -119,6 +129,7 @@ def __init__( conv_group_size = int(nested.get("conv_group_size", conv_group_size)) selector_rank = int(nested.get("selector_rank", selector_rank)) selector_top_k = int(nested.get("selector_top_k", selector_top_k)) + attention_mode = str(nested.get("attention_mode", attention_mode)).lower() input_embedding_scale = float(nested.get("input_embedding_scale", input_embedding_scale)) output_multiplier = float(nested.get("output_multiplier", output_multiplier)) final_logit_softcapping = nested.get("final_logit_softcapping", final_logit_softcapping) @@ -181,8 +192,10 @@ def __init__( kwargs["mask_token_id"] = int( nested.get("mask_token_id", kwargs.get("mask_token_id", 151669)) ) - rope_parameters = kwargs.get("rope_parameters") or {} - kwargs.setdefault("rope_theta", rope_parameters.get("rope_theta", 10000.0)) + rope_parameters = rope_scaling if rope_scaling is not None else rope_parameters + rope_parameters = dict(rope_parameters) if rope_parameters is not None else None + rope_parameters_for_validation = rope_parameters or {} + kwargs.setdefault("rope_theta", rope_parameters_for_validation.get("rope_theta", 10000.0)) kwargs.pop("model_type", None) hidden_act = kwargs.get("hidden_act", "silu") @@ -194,11 +207,22 @@ def __init__( raise ValueError("DFlash2 training requires attention_dropout=0") if kwargs.get("fc_norm", False): raise ValueError("DFlash2 training does not support fc_norm=True") - if kwargs.get("rope_scaling") is not None: - raise ValueError("DFlash2 training does not support rope_scaling") - unsupported_rope_keys = set(rope_parameters) - {"rope_theta", "rope_type"} - if rope_parameters.get("rope_type", "default") != "default" or unsupported_rope_keys: - raise ValueError("DFlash2 training supports only default rope_parameters") + if attention_mode not in {"gqa", "mla"}: + raise ValueError( + f"DFlash2 attention_mode must be 'gqa' or 'mla', got {attention_mode!r}" + ) + if attention_mode == "gqa": + if rope_scaling is not None: + raise ValueError("GQA DFlash2 training does not support rope_scaling") + unsupported_rope_keys = set(rope_parameters_for_validation) - { + "rope_theta", + "rope_type", + } + if ( + rope_parameters_for_validation.get("rope_type", "default") != "default" + or unsupported_rope_keys + ): + raise ValueError("GQA DFlash2 training supports only default rope_parameters") num_hidden_layers = int(kwargs.get("num_hidden_layers", 5)) use_sliding_window = bool(kwargs.get("use_sliding_window", False)) @@ -226,8 +250,6 @@ def __init__( attention_types = set(layer_types) if not attention_types <= {"full_attention", "sliding_attention"}: raise ValueError(f"Unsupported DFlash2 layer types: {sorted(attention_types)}") - if len(attention_types) > 1: - raise ValueError("DFlash2 training does not support mixed full and sliding layers") if "sliding_attention" in attention_types: if sliding_window is None: raise ValueError( @@ -275,6 +297,31 @@ def __init__( ) super().__init__(**kwargs) + self.attention_mode = attention_mode + self.q_lora_rank = q_lora_rank + self.kv_lora_rank = int(kv_lora_rank) + self.qk_nope_head_dim = int(qk_nope_head_dim) + self.qk_rope_head_dim = int(qk_rope_head_dim) + self.v_head_dim = int(v_head_dim) + self.mla_use_output_gate = bool(mla_use_output_gate) + self.rope_parameters = rope_parameters + if self.attention_mode == "mla": + # DeepSeekMLAAttention's shared rotary builder reads the legacy + # ``rope_scaling`` name. Keep it populated for Transformers + # versions where it is not an alias of ``rope_parameters``. + self.rope_scaling = rope_parameters + if rope_parameters is not None: + nested_rope_theta = rope_parameters.get("rope_theta") + if nested_rope_theta is not None: + self.rope_theta = float(nested_rope_theta) + if rope_parameters.get("rope_type", rope_parameters.get("type")) == "yarn": + for key, default in { + "beta_fast": 32.0, + "beta_slow": 1.0, + "mscale": 1.0, + "mscale_all_dim": 0.0, + }.items(): + rope_parameters.setdefault(key, default) self.block_size = block_size self.conv_kernel_size = conv_kernel_size self.conv_group_size = conv_group_size @@ -286,6 +333,7 @@ def __init__( nested.update( { + "attention_mode": self.attention_mode, "block_size": self.block_size, "conv_kernel_size": self.conv_kernel_size, "conv_group_size": self.conv_group_size, @@ -418,6 +466,10 @@ def score_candidates( class DFlash2DecoderLayer(DFlashDecoderLayer): def __init__(self, config: DFlash2Config): + if config.attention_mode == "mla": + from torchspec.models.draft.dspark import K3DSparkMLAAttention + + self.attention_class = K3DSparkMLAAttention super().__init__(config) conv_args = { "hidden_size": config.hidden_size,