From 3d1820296c79e6dcf66cd7a7c6253c5ef3c554f3 Mon Sep 17 00:00:00 2001 From: maocheng23 <35615230+maocheng23@users.noreply.github.com> Date: Fri, 7 Aug 2026 21:34:43 -0700 Subject: [PATCH 1/2] Add config-selectable MLA draft attention --- docs/advanced_features/customization.md | 39 +++ specforge/modeling/draft/dflash.py | 336 +++++++++++++++++++++++- specforge/modeling/draft/dspark.py | 16 +- tests/test_modeling/test_dflash_mla.py | 312 ++++++++++++++++++++++ 4 files changed, 693 insertions(+), 10 deletions(-) create mode 100644 tests/test_modeling/test_dflash_mla.py diff --git a/docs/advanced_features/customization.md b/docs/advanced_features/customization.md index dd4776680..98059bf57 100644 --- a/docs/advanced_features/customization.md +++ b/docs/advanced_features/customization.md @@ -85,6 +85,45 @@ EAGLE3 offline sequence parallelism is selected with tracking are also config features rather than custom launchers; see the [training guide](../basic_usage/training.md) for their validated combinations. +## MLA draft attention + +DFlash-family draft models can select Multi-head Latent Attention (MLA) in the +draft JSON without introducing a model-specific architecture. Keep the existing +`DFlashDraftModel`, `DominoDraftModel`, or `DSparkDraftModel` architecture and +set `dflash_config.attention_mode` to `"mla"`: + +```json +{ + "architectures": ["DSparkDraftModel"], + "hidden_size": 4096, + "num_attention_heads": 32, + "q_lora_rank": 1536, + "kv_lora_rank": 512, + "qk_nope_head_dim": 128, + "qk_rope_head_dim": 64, + "v_head_dim": 128, + "dflash_config": { + "projector_type": "dspark", + "attention_mode": "mla", + "mla_rope_interleaved": true, + "mla_use_output_gate": false + } +} +``` + +`q_lora_rank` may be `null` to use a direct query projection. The KV rank and +all head dimensions must be positive except `qk_nope_head_dim`, which may be +zero; the rotary dimension must be even. Interleaved partial RoPE is the +default MLA convention and can be disabled with +`dflash_config.mla_rope_interleaved: false`. The output gate is optional and +disabled by default. + +MLA changes only the draft attention parameterization. The DFlash-family +target-context injection, per-layer full/sliding masks, objectives, capture +contract, and `eager`, `sdpa`, or `flex_attention` training backend selection +remain unchanged. Omitting `attention_mode` preserves the existing GQA path; +explicit `"mha"` continues to select equal query/KV head counts for DSpark. + ## Draft architectures Draft classes register through `@register_draft`. The key defaults to the diff --git a/specforge/modeling/draft/dflash.py b/specforge/modeling/draft/dflash.py index 4acd30e11..35b069c43 100644 --- a/specforge/modeling/draft/dflash.py +++ b/specforge/modeling/draft/dflash.py @@ -1,3 +1,4 @@ +import copy from typing import Callable, Optional import torch @@ -25,6 +26,7 @@ FULL_ATTENTION = "full_attention" SLIDING_ATTENTION = "sliding_attention" _VALID_DFLASH_LAYER_TYPES = {FULL_ATTENTION, SLIDING_ATTENTION} +_VALID_DFLASH_ATTENTION_MODES = {"gqa", "mha", "mla"} def sample(logits: torch.Tensor, temperature: float = 0.0) -> torch.Tensor: @@ -70,6 +72,57 @@ def resolve_dflash_attention_layout( return layer_types, sliding_window +def resolve_dflash_attention_mode(config: Qwen3Config) -> str: + """Return the configured draft attention architecture. + + GQA and MHA share the existing Qwen attention implementation. MLA swaps + only the attention projections while retaining the DFlash-family decoder, + target-context injection, masks, and objective. + """ + + dflash_config = dict(getattr(config, "dflash_config", None) or {}) + attention_mode = str(dflash_config.get("attention_mode", "gqa")).lower() + if attention_mode not in _VALID_DFLASH_ATTENTION_MODES: + raise ValueError( + "DFlash dflash_config.attention_mode must be one of " + f"{sorted(_VALID_DFLASH_ATTENTION_MODES)}, got {attention_mode!r}" + ) + return attention_mode + + +def validate_dflash_mla_config(config: Qwen3Config) -> None: + """Validate the standard MLA dimensions carried by a draft config.""" + + required = ( + "kv_lora_rank", + "qk_nope_head_dim", + "qk_rope_head_dim", + "v_head_dim", + ) + missing = [name for name in required if getattr(config, name, None) is None] + if missing: + raise ValueError(f"MLA draft config is missing required fields: {missing}") + + q_lora_rank = getattr(config, "q_lora_rank", None) + if q_lora_rank is not None and int(q_lora_rank) <= 0: + raise ValueError(f"q_lora_rank must be positive or null, got {q_lora_rank}") + + positive = ("kv_lora_rank", "qk_rope_head_dim", "v_head_dim") + for name in positive: + value = int(getattr(config, name)) + if value <= 0: + raise ValueError(f"{name} must be positive, got {value}") + + qk_nope_head_dim = int(config.qk_nope_head_dim) + if qk_nope_head_dim < 0: + raise ValueError( + f"qk_nope_head_dim must be non-negative, got {qk_nope_head_dim}" + ) + qk_rope_head_dim = int(config.qk_rope_head_dim) + if qk_rope_head_dim % 2: + raise ValueError(f"qk_rope_head_dim must be even, got {qk_rope_head_dim}") + + def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) @@ -230,6 +283,275 @@ def forward( return attn_output, attn_weights +def _rotate_half_interleaved(x: torch.Tensor) -> torch.Tensor: + """Rotate consecutive pairs, as used by DeepSeek-style MLA RoPE.""" + + paired = x.float().reshape(*x.shape[:-1], -1, 2) + first, second = paired.unbind(dim=-1) + return torch.stack((-second, first), dim=-1).flatten(-2).to(x.dtype) + + +def _apply_mla_rope( + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + *, + interleaved: bool, +) -> torch.Tensor: + if interleaved: + half = cos.shape[-1] // 2 + cos = cos[..., :half].repeat_interleave(2, dim=-1) + sin = sin[..., :half].repeat_interleave(2, dim=-1) + rotated = _rotate_half_interleaved(x) + else: + rotated = rotate_half(x) + return x * cos.unsqueeze(1) + rotated * sin.unsqueeze(1) + + +class Qwen3DFlashMLAAttention(nn.Module): + """Config-selected Multi-head Latent Attention for DFlash-family drafts. + + The module uses the standard MLA parameterization: an optional low-rank Q + path, a shared compressed KV latent, partial RoPE, and expanded per-head + K/V projections for training. Expanding K/V makes the implementation work + with the same eager, SDPA, and FlexAttention interfaces and masks as the + existing GQA/MHA path while preserving MLA's parameterization. + """ + + def __init__( + self, + config: Qwen3Config, + layer_idx: int, + kernels: DFlashKernels, + ): + super().__init__() + validate_dflash_mla_config(config) + self.config = config + self.layer_idx = layer_idx + self.hidden_size = int(config.hidden_size) + self.num_heads = int(config.num_attention_heads) + self.num_key_value_groups = 1 + self.q_lora_rank = ( + None + if getattr(config, "q_lora_rank", None) is None + else int(config.q_lora_rank) + ) + self.kv_lora_rank = int(config.kv_lora_rank) + self.qk_nope_head_dim = int(config.qk_nope_head_dim) + self.qk_rope_head_dim = int(config.qk_rope_head_dim) + self.v_head_dim = int(config.v_head_dim) + self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim + self.head_dim = self.qk_head_dim + self.scaling = self.qk_head_dim**-0.5 + self.attention_dropout = float(config.attention_dropout) + self.is_causal = False + + dflash_config = dict(getattr(config, "dflash_config", None) or {}) + self.rope_interleaved = bool( + dflash_config.get( + "mla_rope_interleaved", + getattr(config, "rope_interleave", True), + ) + ) + self.use_output_gate = bool( + dflash_config.get( + "mla_use_output_gate", + getattr(config, "mla_use_output_gate", False), + ) + ) + + bias = bool(config.attention_bias) + if self.q_lora_rank is None: + self.q_proj = nn.Linear( + self.hidden_size, + self.num_heads * self.qk_head_dim, + bias=bias, + ) + else: + self.q_a_proj = nn.Linear( + self.hidden_size, + self.q_lora_rank, + bias=bias, + ) + self.q_a_layernorm = kernels.make_rms_norm( + self.q_lora_rank, + config.rms_norm_eps, + ) + self.q_b_proj = nn.Linear( + self.q_lora_rank, + self.num_heads * self.qk_head_dim, + bias=bias, + ) + + self.kv_a_proj_with_mqa = nn.Linear( + self.hidden_size, + self.kv_lora_rank + self.qk_rope_head_dim, + bias=bias, + ) + self.kv_a_layernorm = kernels.make_rms_norm( + self.kv_lora_rank, + config.rms_norm_eps, + ) + self.kv_b_proj = nn.Linear( + self.kv_lora_rank, + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), + bias=bias, + ) + if self.use_output_gate: + self.g_proj = nn.Linear( + self.hidden_size, + self.num_heads * self.v_head_dim, + bias=False, + ) + self.o_proj = nn.Linear( + self.num_heads * self.v_head_dim, + self.hidden_size, + bias=bias, + ) + + rope_config = copy.deepcopy(config) + rope_config.head_dim = self.qk_rope_head_dim + self.rotary_emb = Qwen3RotaryEmbedding(rope_config) + self.sliding_window = ( + config.sliding_window + if config.layer_types[layer_idx] == SLIDING_ATTENTION + else None + ) + + def _project_q(self, hidden_states: torch.Tensor) -> torch.Tensor: + if self.q_lora_rank is None: + return self.q_proj(hidden_states) + return self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) + + def _project_kv( + self, + hidden_states: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + batch_size, seq_len = hidden_states.shape[:2] + kv_combined = self.kv_a_proj_with_mqa(hidden_states) + kv_compressed, k_rope = kv_combined.split( + [self.kv_lora_rank, self.qk_rope_head_dim], + dim=-1, + ) + kv = self.kv_b_proj(self.kv_a_layernorm(kv_compressed)).view( + batch_size, + seq_len, + self.num_heads, + self.qk_nope_head_dim + self.v_head_dim, + ) + k_nope, value = kv.split( + [self.qk_nope_head_dim, self.v_head_dim], + dim=-1, + ) + return k_nope, k_rope, value + + def forward( + self, + hidden_states: torch.Tensor, + target_hidden: torch.Tensor, + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]], + attention_mask: Optional[torch.Tensor], + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + cache_position: Optional[torch.LongTensor] = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + del position_embeddings + if position_ids is None: + raise ValueError("MLA attention requires position_ids") + + batch_size, query_len = hidden_states.shape[:2] + context_len = target_hidden.shape[1] + key_len = context_len + query_len + if position_ids.shape[-1] != key_len: + raise ValueError( + "MLA position_ids must cover target context and draft queries; " + f"expected {key_len}, got {position_ids.shape[-1]}" + ) + + query = self._project_q(hidden_states).view( + batch_size, + query_len, + self.num_heads, + self.qk_head_dim, + ) + q_nope, q_rope = query.split( + [self.qk_nope_head_dim, self.qk_rope_head_dim], + dim=-1, + ) + + context_k_nope, context_k_rope, context_value = self._project_kv(target_hidden) + noise_k_nope, noise_k_rope, noise_value = self._project_kv(hidden_states) + k_nope = torch.cat((context_k_nope, noise_k_nope), dim=1) + k_rope = torch.cat((context_k_rope, noise_k_rope), dim=1) + value = torch.cat((context_value, noise_value), dim=1) + + cos, sin = self.rotary_emb(k_rope, position_ids) + q_rope = _apply_mla_rope( + q_rope.transpose(1, 2), + cos[:, -query_len:], + sin[:, -query_len:], + interleaved=self.rope_interleaved, + ) + k_rope = _apply_mla_rope( + k_rope.unsqueeze(1), + cos, + sin, + interleaved=self.rope_interleaved, + ).expand(-1, self.num_heads, -1, -1) + query = torch.cat((q_nope.transpose(1, 2), q_rope), dim=-1) + key = torch.cat((k_nope.transpose(1, 2), k_rope), dim=-1) + value = value.transpose(1, 2) + + if past_key_values is not None: + cache_kwargs = {"cache_position": cache_position} + key, value = past_key_values.update( + key, + value, + self.layer_idx, + cache_kwargs, + ) + + valid_queries = None + attn_fn: Callable = eager_attention_forward + if self.config._attn_implementation == "eager": + attention_mask, valid_queries = _prepare_dflash_eager_mask( + attention_mask, + query.dtype, + ) + else: + attn_fn = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + attn_output, attn_weights = attn_fn( + self, + query, + key, + value, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + sliding_window=self.sliding_window, + **kwargs, + ) + if valid_queries is not None and attn_weights is not None: + attn_weights = attn_weights.masked_fill(~valid_queries, 0) + + if self.use_output_gate: + gate = torch.sigmoid(self.g_proj(hidden_states)).view( + batch_size, + query_len, + self.num_heads, + self.v_head_dim, + ) + attn_output = attn_output * gate + attn_output = self.o_proj(attn_output.reshape(batch_size, query_len, -1)) + if valid_queries is not None: + attn_output = attn_output.masked_fill( + ~valid_queries.any(dim=1), + 0, + ) + return attn_output, attn_weights + + class Qwen3DFlashDecoderLayer(GradientCheckpointingLayer): def __init__( self, @@ -239,7 +561,12 @@ def __init__( ): super().__init__() self.hidden_size = config.hidden_size - self.self_attn = Qwen3DFlashAttention( + attention_cls = ( + Qwen3DFlashMLAAttention + if resolve_dflash_attention_mode(config) == "mla" + else Qwen3DFlashAttention + ) + self.self_attn = attention_cls( config=config, layer_idx=layer_idx, kernels=kernels, @@ -368,6 +695,7 @@ def __init__( super().__init__(config) self.config = config self.layer_types, self.sliding_window = resolve_dflash_attention_layout(config) + self.attention_mode = resolve_dflash_attention_mode(config) kernels = dflash_kernels or DEFAULT_DFLASH_KERNELS self.layers = nn.ModuleList( [ @@ -463,7 +791,11 @@ def forward( ) -> CausalLMOutputWithPast: hidden_states = noise_embedding target_hidden = self.hidden_norm(self.fc(target_hidden)) - position_embeddings = self.rotary_emb(hidden_states, position_ids) + position_embeddings = ( + None + if self.attention_mode == "mla" + else self.rotary_emb(hidden_states, position_ids) + ) for layer_type, layer in zip(self.layer_types, self.layers): layer_attention_mask = ( attention_mask[layer_type] diff --git a/specforge/modeling/draft/dspark.py b/specforge/modeling/draft/dspark.py index 1c5853d8b..54c87736b 100644 --- a/specforge/modeling/draft/dspark.py +++ b/specforge/modeling/draft/dspark.py @@ -286,19 +286,19 @@ class DSparkDraftModel(DFlashDraftModel): expected_projector_type = "dspark" def __init__(self, config) -> None: - dflash_config = dict(getattr(config, "dflash_config", None) or {}) num_heads = int(config.num_attention_heads) num_kv_heads = int(config.num_key_value_heads) - if num_heads % num_kv_heads: + dflash_config = dict(getattr(config, "dflash_config", None) or {}) + attention_mode = str(dflash_config.get("attention_mode", "gqa")).lower() + if attention_mode not in {"gqa", "mha", "mla"}: raise ValueError( - "DSpark requires num_key_value_heads to divide " - f"num_attention_heads, got {num_kv_heads} and {num_heads}" + "DSpark dflash_config.attention_mode must be 'gqa', 'mha', " + f"or 'mla', got {attention_mode!r}" ) - attention_mode = str(dflash_config.get("attention_mode", "gqa")).lower() - if attention_mode not in {"gqa", "mha"}: + if attention_mode != "mla" and num_heads % num_kv_heads: raise ValueError( - "DSpark dflash_config.attention_mode must be 'gqa' or 'mha', " - f"got {attention_mode!r}" + "DSpark requires num_key_value_heads to divide " + f"num_attention_heads, got {num_kv_heads} and {num_heads}" ) if attention_mode == "gqa" and num_kv_heads >= num_heads: raise ValueError( diff --git a/tests/test_modeling/test_dflash_mla.py b/tests/test_modeling/test_dflash_mla.py new file mode 100644 index 000000000..0b2ee2948 --- /dev/null +++ b/tests/test_modeling/test_dflash_mla.py @@ -0,0 +1,312 @@ +import copy +import unittest + +import torch +from torch.testing import assert_close +from transformers import DynamicCache, Qwen3Config + +from specforge.algorithms.common.dflash_family_model import create_dflash_block_mask +from specforge.modeling.draft.dflash import DFlashDraftModel, Qwen3DFlashMLAAttention +from specforge.modeling.draft.dflash_kernels import DEFAULT_DFLASH_KERNELS +from specforge.modeling.draft.dspark import DSparkDraftModel + + +def _mla_config( + *, + architecture: str = "DFlashDraftModel", + implementation: str = "sdpa", + q_lora_rank: int | None = 12, +) -> Qwen3Config: + dflash_config = { + "attention_mode": "mla", + "target_layer_ids": [1], + "mla_rope_interleaved": True, + "mla_use_output_gate": True, + } + if architecture == "DSparkDraftModel": + dflash_config.update( + { + "projector_type": "dspark", + "markov_rank": 4, + "enable_confidence_head": True, + } + ) + config = Qwen3Config( + architectures=[architecture], + block_size=3, + hidden_size=24, + intermediate_size=48, + num_attention_heads=4, + num_key_value_heads=1, + num_hidden_layers=1, + num_target_layers=4, + head_dim=6, + q_lora_rank=q_lora_rank, + kv_lora_rank=8, + qk_nope_head_dim=4, + qk_rope_head_dim=4, + v_head_dim=6, + max_position_embeddings=64, + vocab_size=64, + layer_types=["full_attention"], + attention_bias=False, + attention_dropout=0.0, + dflash_config=dflash_config, + ) + config._attn_implementation = implementation + return config + + +def _attention_forward( + attention: Qwen3DFlashMLAAttention, + hidden_states: torch.Tensor, + target_hidden: torch.Tensor, + attention_mask: torch.Tensor | None, + *, + past_key_values: DynamicCache | None = None, + position_offset: int = 0, +): + key_len = target_hidden.shape[1] + hidden_states.shape[1] + position_ids = torch.arange( + position_offset, + position_offset + key_len, + device=hidden_states.device, + ).unsqueeze(0) + return attention( + hidden_states=hidden_states, + target_hidden=target_hidden, + position_embeddings=None, + position_ids=position_ids, + attention_mask=attention_mask, + past_key_values=past_key_values, + ) + + +class TestDFlashMLASelection(unittest.TestCase): + def test_dflash_selects_mla_from_config(self): + model = DFlashDraftModel(_mla_config()) + + self.assertEqual(model.attention_mode, "mla") + self.assertIsInstance(model.layers[0].self_attn, Qwen3DFlashMLAAttention) + self.assertTrue(model.layers[0].self_attn.rope_interleaved) + self.assertTrue(model.layers[0].self_attn.use_output_gate) + + def test_dspark_uses_the_same_generic_mla_backbone(self): + model = DSparkDraftModel(_mla_config(architecture="DSparkDraftModel")) + + self.assertIsInstance(model.layers[0].self_attn, Qwen3DFlashMLAAttention) + self.assertEqual(model.config.dflash_config["attention_mode"], "mla") + self.assertIsNotNone(model.markov_head) + self.assertIsNotNone(model.confidence_head) + + def test_direct_query_projection_is_supported(self): + model = DFlashDraftModel(_mla_config(q_lora_rank=None)) + attention = model.layers[0].self_attn + + self.assertTrue(hasattr(attention, "q_proj")) + self.assertFalse(hasattr(attention, "q_a_proj")) + + def test_forward_backward_is_finite(self): + torch.manual_seed(7) + model = DFlashDraftModel(_mla_config()).train() + hidden_states = torch.randn(2, 3, 24, requires_grad=True) + target_hidden = torch.randn(2, 5, 24) + position_ids = torch.arange(8).expand(2, -1) + output = model( + position_ids=position_ids, + noise_embedding=hidden_states, + target_hidden=target_hidden, + attention_mask=torch.ones(2, 1, 3, 8, dtype=torch.bool), + ) + + self.assertEqual(output.shape, (2, 3, 24)) + self.assertTrue(torch.isfinite(output).all()) + output.square().mean().backward() + self.assertIsNotNone(hidden_states.grad) + self.assertTrue(torch.isfinite(hidden_states.grad).all()) + attention = model.layers[0].self_attn + for name in ( + "q_a_proj", + "q_b_proj", + "kv_a_proj_with_mqa", + "kv_b_proj", + "g_proj", + "o_proj", + ): + parameter = getattr(attention, name).weight + self.assertIsNotNone(parameter.grad, name) + self.assertTrue(torch.isfinite(parameter.grad).all(), name) + + +class TestDFlashMLAAttention(unittest.TestCase): + def test_eager_matches_sdpa_with_boolean_mask(self): + torch.manual_seed(11) + eager_config = _mla_config(implementation="eager") + sdpa_config = _mla_config(implementation="sdpa") + eager = Qwen3DFlashMLAAttention( + eager_config, + layer_idx=0, + kernels=DEFAULT_DFLASH_KERNELS, + ).eval() + sdpa = Qwen3DFlashMLAAttention( + sdpa_config, + layer_idx=0, + kernels=DEFAULT_DFLASH_KERNELS, + ).eval() + sdpa.load_state_dict(eager.state_dict()) + + eager_hidden = torch.randn(1, 3, 24, requires_grad=True) + eager_target = torch.randn(1, 4, 24, requires_grad=True) + sdpa_hidden = eager_hidden.detach().clone().requires_grad_(True) + sdpa_target = eager_target.detach().clone().requires_grad_(True) + mask = torch.tensor( + [ + [ + [ + [True, True, True, True, True, False, False], + [True, True, False, True, True, True, False], + [False, False, False, False, False, False, False], + ] + ] + ] + ) + + eager_output, eager_weights = _attention_forward( + eager, + eager_hidden, + eager_target, + mask, + ) + sdpa_output, _ = _attention_forward( + sdpa, + sdpa_hidden, + sdpa_target, + mask, + ) + + assert_close(eager_output, sdpa_output, rtol=1e-5, atol=1e-6) + self.assertIsNotNone(eager_weights) + self.assertEqual(eager_output[:, -1].abs().sum().item(), 0.0) + forbidden = eager_weights.masked_select(~mask.expand_as(eager_weights)) + assert_close(forbidden, torch.zeros_like(forbidden), rtol=0, atol=0) + + output_grad = torch.randn_like(eager_output) + eager_grads = torch.autograd.grad( + (eager_output * output_grad).sum(), + (eager_hidden, eager_target), + ) + sdpa_grads = torch.autograd.grad( + (sdpa_output * output_grad).sum(), + (sdpa_hidden, sdpa_target), + ) + for eager_grad, sdpa_grad in zip(eager_grads, sdpa_grads): + assert_close(eager_grad, sdpa_grad, rtol=1e-5, atol=1e-6) + + def test_dynamic_cache_accepts_different_key_and_value_dims(self): + attention = Qwen3DFlashMLAAttention( + _mla_config(), + layer_idx=0, + kernels=DEFAULT_DFLASH_KERNELS, + ).eval() + cache = DynamicCache() + + first_output, _ = _attention_forward( + attention, + torch.randn(1, 2, 24), + torch.randn(1, 3, 24), + None, + past_key_values=cache, + ) + self.assertEqual(first_output.shape, (1, 2, 24)) + self.assertEqual(cache.get_seq_length(), 5) + + second_output, _ = _attention_forward( + attention, + torch.randn(1, 2, 24), + torch.randn(1, 1, 24), + None, + past_key_values=cache, + position_offset=5, + ) + self.assertEqual(second_output.shape, (1, 2, 24)) + self.assertEqual(cache.get_seq_length(), 8) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required") + def test_flex_attention_forward_backward_cuda(self): + config = _mla_config(implementation="flex_attention") + # CUDA FlexAttention requires both QK and V head dimensions >= 16. + config.qk_nope_head_dim = 12 + config.v_head_dim = 16 + attention = ( + Qwen3DFlashMLAAttention( + config, + layer_idx=0, + kernels=DEFAULT_DFLASH_KERNELS, + ) + .to(device="cuda", dtype=torch.bfloat16) + .train() + ) + hidden_states = torch.randn( + 1, + 6, + 24, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + target_hidden = torch.randn( + 1, + 6, + 24, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + block_mask = create_dflash_block_mask( + anchor_positions=torch.tensor([[2, 5]], device="cuda"), + block_keep_mask=torch.tensor([[True, True]], device="cuda"), + S=6, + block_size=3, + device=torch.device("cuda"), + ) + + output, _ = _attention_forward( + attention, + hidden_states, + target_hidden, + block_mask, + ) + self.assertEqual(output.shape, (1, 6, 24)) + self.assertTrue(torch.isfinite(output).all()) + output.float().square().mean().backward() + for tensor in (hidden_states, target_hidden): + self.assertIsNotNone(tensor.grad) + self.assertTrue(torch.isfinite(tensor.grad).all()) + + +class TestDFlashMLAConfigValidation(unittest.TestCase): + def test_rejects_missing_or_invalid_dimensions(self): + cases = { + "kv_lora_rank": None, + "q_lora_rank": 0, + "qk_nope_head_dim": -1, + "qk_rope_head_dim": 3, + "v_head_dim": 0, + } + for field, value in cases.items(): + with self.subTest(field=field, value=value): + config = copy.deepcopy(_mla_config()) + setattr(config, field, value) + with self.assertRaisesRegex(ValueError, field): + DFlashDraftModel(config) + + def test_rejects_unknown_attention_mode(self): + config = _mla_config() + config.dflash_config["attention_mode"] = "latent-ish" + + with self.assertRaisesRegex(ValueError, "attention_mode"): + DFlashDraftModel(config) + + +if __name__ == "__main__": + unittest.main(verbosity=2) From 7afa170bb8a7da4314669f6ec366f77bf337b2ec Mon Sep 17 00:00:00 2001 From: maocheng23 <35615230+maocheng23@users.noreply.github.com> Date: Sat, 15 Aug 2026 00:53:41 -0700 Subject: [PATCH 2/2] Add configurable GQA and MLA SWA layers --- docs/advanced_features/customization.md | 5 + docs/basic_usage/training.md | 35 +++ specforge/modeling/draft/dflash.py | 64 ++++-- specforge/modeling/draft/dspark.py | 21 +- tests/test_modeling/test_dflash_swa_modes.py | 219 +++++++++++++++++++ 5 files changed, 318 insertions(+), 26 deletions(-) create mode 100644 tests/test_modeling/test_dflash_swa_modes.py diff --git a/docs/advanced_features/customization.md b/docs/advanced_features/customization.md index 98059bf57..9b5ac8813 100644 --- a/docs/advanced_features/customization.md +++ b/docs/advanced_features/customization.md @@ -123,6 +123,11 @@ target-context injection, per-layer full/sliding masks, objectives, capture contract, and `eager`, `sdpa`, or `flex_attention` training backend selection remain unchanged. Omitting `attention_mode` preserves the existing GQA path; explicit `"mha"` continues to select equal query/KV head counts for DSpark. +For a mixed projection stack, replace the scalar with an `attention_modes` +list containing one `"gqa"`, `"mha"`, or `"mla"` entry per draft layer. This +list composes with `layer_types`, so both SWA-GQA and SWA-MLA layers can appear +in the same architecture; see the [training guide](../basic_usage/training.md) +for the paired config contract. ## Draft architectures diff --git a/docs/basic_usage/training.md b/docs/basic_usage/training.md index 39efc3e5a..8f1167ad3 100644 --- a/docs/basic_usage/training.md +++ b/docs/basic_usage/training.md @@ -179,6 +179,41 @@ mixed layout must be edited explicitly in the draft JSON. The `eager`, `sdpa`, and `flex_attention` backends support both layouts. +The context layout and projection layout are independent. Use +`dflash_config.attention_mode` as a uniform `"gqa"`, `"mha"`, or `"mla"` +shorthand, or use `dflash_config.attention_modes` to select one mode per draft +layer. The per-layer list must match `num_hidden_layers`; do not set the scalar +and list forms together. + +```json +{ + "num_hidden_layers": 4, + "layer_types": [ + "sliding_attention", + "sliding_attention", + "full_attention", + "full_attention" + ], + "use_sliding_window": true, + "sliding_window": 2048, + "dflash_config": { + "attention_modes": ["gqa", "mla", "gqa", "mla"] + } +} +``` + +The two lists are paired by layer index: + +| Context layout | Projection mode | Result | +|---|---|---| +| `sliding_attention` | `gqa` | SWA-GQA | +| `sliding_attention` | `mla` | SWA-MLA | +| `full_attention` | `gqa` or `mha` | full-context standard attention | +| `full_attention` | `mla` | full-context MLA | + +Any configuration containing an MLA layer must also provide the standard MLA +dimension fields described in [MLA draft attention](../advanced_features/customization.md#mla-draft-attention). + Domino and DSpark need their projector/head metadata, so they require an explicit draft config (or a pretrained warm-start source that contains `config.json`). The old Domino parser exposed an optional config flag, but its diff --git a/specforge/modeling/draft/dflash.py b/specforge/modeling/draft/dflash.py index 35b069c43..9d180ae51 100644 --- a/specforge/modeling/draft/dflash.py +++ b/specforge/modeling/draft/dflash.py @@ -72,22 +72,54 @@ def resolve_dflash_attention_layout( return layer_types, sliding_window -def resolve_dflash_attention_mode(config: Qwen3Config) -> str: - """Return the configured draft attention architecture. +def resolve_dflash_attention_modes(config: Qwen3Config) -> tuple[str, ...]: + """Return the validated per-layer draft attention architectures. - GQA and MHA share the existing Qwen attention implementation. MLA swaps - only the attention projections while retaining the DFlash-family decoder, - target-context injection, masks, and objective. + ``attention_mode`` remains the uniform shorthand. ``attention_modes`` is + the per-layer form and is intentionally orthogonal to ``layer_types``: + each layer independently chooses its projection layout (GQA/MHA/MLA) and + its context layout (full/sliding). """ dflash_config = dict(getattr(config, "dflash_config", None) or {}) - attention_mode = str(dflash_config.get("attention_mode", "gqa")).lower() - if attention_mode not in _VALID_DFLASH_ATTENTION_MODES: + if "attention_mode" in dflash_config and "attention_modes" in dflash_config: + raise ValueError( + "DFlash dflash_config must set only one of attention_mode or " + "attention_modes" + ) + + configured_modes = dflash_config.get("attention_modes") + if configured_modes is None: + mode = str(dflash_config.get("attention_mode", "gqa")).lower() + attention_modes = (mode,) * int(config.num_hidden_layers) + else: + if not isinstance(configured_modes, (list, tuple)): + raise ValueError( + "DFlash dflash_config.attention_modes must be a per-layer list" + ) + attention_modes = tuple(str(mode).lower() for mode in configured_modes) + if len(attention_modes) != int(config.num_hidden_layers): + raise ValueError( + "DFlash dflash_config.attention_modes must contain exactly " + f"num_hidden_layers={config.num_hidden_layers} entries, got " + f"{len(attention_modes)}" + ) + + invalid = set(attention_modes) - _VALID_DFLASH_ATTENTION_MODES + if invalid: raise ValueError( - "DFlash dflash_config.attention_mode must be one of " - f"{sorted(_VALID_DFLASH_ATTENTION_MODES)}, got {attention_mode!r}" + "DFlash dflash_config.attention_mode/attention_modes values must be " + "selected from " + f"{sorted(_VALID_DFLASH_ATTENTION_MODES)}, got {sorted(invalid)}" ) - return attention_mode + return attention_modes + + +def resolve_dflash_attention_mode(config: Qwen3Config) -> str: + """Return the uniform mode, or ``mixed`` for a per-layer configuration.""" + + attention_modes = resolve_dflash_attention_modes(config) + return attention_modes[0] if len(set(attention_modes)) == 1 else "mixed" def validate_dflash_mla_config(config: Qwen3Config) -> None: @@ -561,10 +593,9 @@ def __init__( ): super().__init__() self.hidden_size = config.hidden_size + attention_mode = resolve_dflash_attention_modes(config)[layer_idx] attention_cls = ( - Qwen3DFlashMLAAttention - if resolve_dflash_attention_mode(config) == "mla" - else Qwen3DFlashAttention + Qwen3DFlashMLAAttention if attention_mode == "mla" else Qwen3DFlashAttention ) self.self_attn = attention_cls( config=config, @@ -695,7 +726,10 @@ def __init__( super().__init__(config) self.config = config self.layer_types, self.sliding_window = resolve_dflash_attention_layout(config) - self.attention_mode = resolve_dflash_attention_mode(config) + self.attention_modes = resolve_dflash_attention_modes(config) + self.attention_mode = ( + self.attention_modes[0] if len(set(self.attention_modes)) == 1 else "mixed" + ) kernels = dflash_kernels or DEFAULT_DFLASH_KERNELS self.layers = nn.ModuleList( [ @@ -793,7 +827,7 @@ def forward( target_hidden = self.hidden_norm(self.fc(target_hidden)) position_embeddings = ( None - if self.attention_mode == "mla" + if all(mode == "mla" for mode in self.attention_modes) else self.rotary_emb(hidden_states, position_ids) ) for layer_type, layer in zip(self.layer_types, self.layers): diff --git a/specforge/modeling/draft/dspark.py b/specforge/modeling/draft/dspark.py index 54c87736b..549879bb3 100644 --- a/specforge/modeling/draft/dspark.py +++ b/specforge/modeling/draft/dspark.py @@ -8,7 +8,7 @@ import torch from torch import nn -from .dflash import DFlashDraftModel +from .dflash import DFlashDraftModel, resolve_dflash_attention_modes from .registry import register_draft @@ -289,29 +289,28 @@ def __init__(self, config) -> None: num_heads = int(config.num_attention_heads) num_kv_heads = int(config.num_key_value_heads) dflash_config = dict(getattr(config, "dflash_config", None) or {}) - attention_mode = str(dflash_config.get("attention_mode", "gqa")).lower() - if attention_mode not in {"gqa", "mha", "mla"}: - raise ValueError( - "DSpark dflash_config.attention_mode must be 'gqa', 'mha', " - f"or 'mla', got {attention_mode!r}" - ) - if attention_mode != "mla" and num_heads % num_kv_heads: + attention_modes = resolve_dflash_attention_modes(config) + standard_modes = set(attention_modes) - {"mla"} + if standard_modes and num_heads % num_kv_heads: raise ValueError( "DSpark requires num_key_value_heads to divide " f"num_attention_heads, got {num_kv_heads} and {num_heads}" ) - if attention_mode == "gqa" and num_kv_heads >= num_heads: + if "gqa" in standard_modes and num_kv_heads >= num_heads: raise ValueError( "DSpark defaults to GQA and requires num_key_value_heads < " "num_attention_heads; set dflash_config.attention_mode='mha' " "to opt into equal query/KV head counts" ) - if attention_mode == "mha" and num_kv_heads != num_heads: + if "mha" in standard_modes and num_kv_heads != num_heads: raise ValueError( "DSpark MHA opt-in requires num_key_value_heads == " f"num_attention_heads, got {num_kv_heads} and {num_heads}" ) - dflash_config["attention_mode"] = attention_mode + if "attention_modes" in dflash_config: + dflash_config["attention_modes"] = list(attention_modes) + else: + dflash_config["attention_mode"] = attention_modes[0] projector_type = dflash_config.get("projector_type") if projector_type is None: dflash_config["projector_type"] = self.expected_projector_type diff --git a/tests/test_modeling/test_dflash_swa_modes.py b/tests/test_modeling/test_dflash_swa_modes.py new file mode 100644 index 000000000..e67c289db --- /dev/null +++ b/tests/test_modeling/test_dflash_swa_modes.py @@ -0,0 +1,219 @@ +import copy +import unittest + +import torch +from transformers import Qwen3Config + +from specforge.algorithms.common.dflash_family_model import ( + create_dflash_block_mask, + create_dflash_sdpa_mask, +) +from specforge.modeling.draft.dflash import ( + DFlashDraftModel, + Qwen3DFlashAttention, + Qwen3DFlashMLAAttention, + resolve_dflash_attention_mode, + resolve_dflash_attention_modes, +) +from specforge.modeling.draft.dspark import DSparkDraftModel + + +def _swa_config( + *, + attention_modes=("gqa", "mla", "gqa", "mla"), + layer_types=( + "sliding_attention", + "sliding_attention", + "full_attention", + "full_attention", + ), + architecture="DFlashDraftModel", + implementation="sdpa", +): + dflash_config = { + "attention_modes": list(attention_modes), + "target_layer_ids": [1], + "mla_rope_interleaved": True, + "mla_use_output_gate": False, + } + if architecture == "DSparkDraftModel": + dflash_config.update( + { + "projector_type": "dspark", + "markov_rank": 4, + "enable_confidence_head": False, + } + ) + config = Qwen3Config( + architectures=[architecture], + block_size=2, + hidden_size=32, + intermediate_size=64, + num_attention_heads=2, + num_key_value_heads=1, + num_hidden_layers=len(layer_types), + num_target_layers=4, + head_dim=16, + q_lora_rank=16, + kv_lora_rank=8, + qk_nope_head_dim=12, + qk_rope_head_dim=4, + v_head_dim=16, + max_position_embeddings=64, + vocab_size=64, + layer_types=list(layer_types), + sliding_window=4, + use_sliding_window=True, + attention_bias=False, + attention_dropout=0.0, + dflash_config=dflash_config, + ) + config._attn_implementation = implementation + return config + + +def _masks(*, device, flex): + factory = create_dflash_block_mask if flex else create_dflash_sdpa_mask + kwargs = { + "anchor_positions": torch.tensor([[4]], device=device), + "block_keep_mask": torch.tensor([[True]], device=device), + "S": 6, + "block_size": 2, + "device": device, + } + return { + "full_attention": factory(**kwargs), + "sliding_attention": factory(**kwargs, sliding_window=4), + } + + +class TestDFlashSWAProjectionModes(unittest.TestCase): + def test_config_selects_swa_gqa_and_swa_mla_per_layer(self): + model = DFlashDraftModel(_swa_config()) + + self.assertEqual(model.attention_mode, "mixed") + self.assertEqual(model.attention_modes, ("gqa", "mla", "gqa", "mla")) + expected_classes = ( + Qwen3DFlashAttention, + Qwen3DFlashMLAAttention, + Qwen3DFlashAttention, + Qwen3DFlashMLAAttention, + ) + for index, expected_class in enumerate(expected_classes): + with self.subTest(index=index): + attention = model.layers[index].self_attn + self.assertIsInstance(attention, expected_class) + expected_window = 4 if index < 2 else None + self.assertEqual(attention.sliding_window, expected_window) + + def test_mixed_swa_gqa_mla_forward_backward(self): + torch.manual_seed(17) + model = DFlashDraftModel(_swa_config()).train() + noise_embedding = torch.randn(1, 2, 32, requires_grad=True) + target_hidden = torch.randn(1, 6, 32, requires_grad=True) + output = model( + position_ids=torch.arange(8).unsqueeze(0), + noise_embedding=noise_embedding, + target_hidden=target_hidden, + attention_mask=_masks(device=torch.device("cpu"), flex=False), + ) + + self.assertEqual(output.shape, (1, 2, 32)) + self.assertTrue(torch.isfinite(output).all()) + output.square().mean().backward() + for tensor in (noise_embedding, target_hidden): + self.assertIsNotNone(tensor.grad) + self.assertTrue(torch.isfinite(tensor.grad).all()) + for index in range(4): + attention = model.layers[index].self_attn + projection = ( + attention.q_a_proj + if isinstance(attention, Qwen3DFlashMLAAttention) + else attention.q_proj + ) + self.assertIsNotNone(projection.weight.grad, index) + self.assertTrue(torch.isfinite(projection.weight.grad).all(), index) + + def test_dspark_uses_the_same_mixed_swa_backbone(self): + config = _swa_config(architecture="DSparkDraftModel") + model = DSparkDraftModel(config) + + self.assertEqual(model.attention_modes, ("gqa", "mla", "gqa", "mla")) + self.assertIsInstance(model.layers[0].self_attn, Qwen3DFlashAttention) + self.assertIsInstance(model.layers[1].self_attn, Qwen3DFlashMLAAttention) + self.assertEqual( + model.config.dflash_config["attention_modes"], + ["gqa", "mla", "gqa", "mla"], + ) + + def test_uniform_attention_mode_shorthand_remains_compatible(self): + config = _swa_config(attention_modes=("mla",) * 4) + config.dflash_config.pop("attention_modes") + config.dflash_config["attention_mode"] = "MLA" + + self.assertEqual(resolve_dflash_attention_modes(config), ("mla",) * 4) + self.assertEqual(resolve_dflash_attention_mode(config), "mla") + + config.dflash_config.pop("attention_mode") + self.assertEqual(resolve_dflash_attention_modes(config), ("gqa",) * 4) + self.assertEqual(resolve_dflash_attention_mode(config), "gqa") + + def test_rejects_invalid_per_layer_attention_modes(self): + cases = ( + ("not-a-list", "per-layer list"), + (["gqa"], "num_hidden_layers"), + (["gqa", "mla", "latent-ish", "gqa"], "latent-ish"), + ) + for modes, message in cases: + with self.subTest(modes=modes): + config = copy.deepcopy(_swa_config()) + config.dflash_config["attention_modes"] = modes + with self.assertRaisesRegex(ValueError, message): + resolve_dflash_attention_modes(config) + + config = _swa_config() + config.dflash_config["attention_mode"] = "gqa" + with self.assertRaisesRegex(ValueError, "only one"): + resolve_dflash_attention_modes(config) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required") + def test_mixed_swa_gqa_mla_flex_forward_backward_cuda(self): + config = _swa_config( + attention_modes=("gqa", "mla"), + layer_types=("sliding_attention", "sliding_attention"), + implementation="flex_attention", + ) + model = DFlashDraftModel(config).to(device="cuda", dtype=torch.bfloat16).train() + noise_embedding = torch.randn( + 1, + 2, + 32, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + target_hidden = torch.randn( + 1, + 6, + 32, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + output = model( + position_ids=torch.arange(8, device="cuda").unsqueeze(0), + noise_embedding=noise_embedding, + target_hidden=target_hidden, + attention_mask=_masks(device=torch.device("cuda"), flex=True), + ) + + self.assertEqual(output.shape, (1, 2, 32)) + self.assertTrue(torch.isfinite(output).all()) + output.float().square().mean().backward() + for tensor in (noise_embedding, target_hidden): + self.assertIsNotNone(tensor.grad) + self.assertTrue(torch.isfinite(tensor.grad).all()) + + +if __name__ == "__main__": + unittest.main(verbosity=2)