diff --git a/docs/advanced_features/customization.md b/docs/advanced_features/customization.md index fa71be886..ebc6affd7 100644 --- a/docs/advanced_features/customization.md +++ b/docs/advanced_features/customization.md @@ -85,6 +85,47 @@ 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. +## DFlash-family attention modes + +DFlash-family draft models select their attention parameterization with +`dflash_config.attention_mode`: `gqa` (the default), `mha`, or `mla`. The mode +swaps only the attention projections inside the shared decoder layer; the +`DFlashDraftModel`, `DominoDraftModel`, and `DSparkDraftModel` architectures, +target-context injection, per-layer full/sliding masks, objectives, capture +contract, and `eager`/`sdpa`/`flex_attention` backend selection are identical +across modes. Multi-head Latent Attention is therefore a draft JSON change, +not a new architecture: + +```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 dimensions and behavior use the standard top-level Hugging Face fields: +`q_lora_rank` (`null` selects a direct query projection), `kv_lora_rank`, the +head dims (`qk_nope_head_dim` may be zero; `qk_rope_head_dim` must be even), +and `rope_interleave` for the rotation convention (omitted means interleaved, +the DeepSeek default; `false` selects NeoX-style half rotation). Omitting +`attention_mode` preserves the existing GQA path; `"mha"` requires equal +query/KV head counts. + +MLA is a training-side mode: checkpoints train, evaluate through +`spec_generate`, and export through `--to hf`. SGLang serving of DFlash-family +drafts currently implements the GQA/MHA layout only, so plan benchmarks +accordingly. + ## Draft architectures Draft classes register through `@register_draft`. The key defaults to the diff --git a/scripts/gates/normalize_dflash_export.py b/scripts/gates/normalize_dflash_export.py index 82622e1df..feeb6a1a7 100644 --- a/scripts/gates/normalize_dflash_export.py +++ b/scripts/gates/normalize_dflash_export.py @@ -81,6 +81,15 @@ def normalize_export(config_path: str, expected_block_size: int) -> Dict[str, An "export is not DFlash-family: " f"dflash_config.projector_type={projector_type!r}" ) + attention_mode = method_config.get("attention_mode", "gqa") + if not isinstance(attention_mode, str) or attention_mode.lower() not in { + "gqa", + "mha", + }: + raise ValueError( + "SGLang DFlash-family serving supports only GQA/MHA exports, " + f"got dflash_config.attention_mode={attention_mode!r}" + ) if projector_type == "dspark": _normalize_dspark(config, method_config) diff --git a/specforge/algorithms/model_providers.py b/specforge/algorithms/model_providers.py index 285433fcd..4ee7ceb38 100644 --- a/specforge/algorithms/model_providers.py +++ b/specforge/algorithms/model_providers.py @@ -479,6 +479,7 @@ def apply_dflash_overrides(cfg: Config, draft_config: Any) -> None: from specforge.modeling.draft.dflash import ( build_target_layer_ids, resolve_dflash_attention_layout, + validate_dflash_attention_config, ) requested_layers = cfg.model.draft_num_hidden_layers @@ -502,6 +503,7 @@ def apply_dflash_overrides(cfg: Config, draft_config: Any) -> None: } resolve_dflash_attention_layout(draft_config) + validate_dflash_attention_config(draft_config) __all__ = [ diff --git a/specforge/modeling/draft/dflash.py b/specforge/modeling/draft/dflash.py index 4acd30e11..53b0d61f3 100644 --- a/specforge/modeling/draft/dflash.py +++ b/specforge/modeling/draft/dflash.py @@ -1,3 +1,5 @@ +import copy +import math from typing import Callable, Optional import torch @@ -70,6 +72,99 @@ def resolve_dflash_attention_layout( return layer_types, sliding_window +def resolve_dflash_attention_mode(config: Qwen3Config) -> str: + """Validate and return the configured draft attention mode. + + ``gqa`` and ``mha`` share :class:`Qwen3DFlashAttention`; ``mla`` swaps in + the latent parameterization while retaining the family decoder, + target-context injection, masks, and objectives. + """ + + dflash_config = getattr(config, "dflash_config", None) or {} + attention_mode = str(dflash_config.get("attention_mode", "gqa")).lower() + if attention_mode not in _DFLASH_ATTENTION_CLASSES: + raise ValueError( + "DFlash dflash_config.attention_mode must be one of " + f"{sorted(_DFLASH_ATTENTION_CLASSES)}, got {attention_mode!r}" + ) + return attention_mode + + +def _require_bool_config(value: object, field: str) -> bool: + if not isinstance(value, bool): + raise ValueError(f"DFlash {field} must be a boolean, got {value!r}") + return value + + +def _resolve_mla_rope_interleaved(config: Qwen3Config) -> bool: + """Rotation convention from the standard MLA ``rope_interleave`` field.""" + + return _require_bool_config( + getattr(config, "rope_interleave", True), + "config.rope_interleave", + ) + + +def validate_dflash_mla_config(config: Qwen3Config) -> None: + """Validate the standard MLA dimension fields 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}") + + for name in ("kv_lora_rank", "qk_rope_head_dim", "v_head_dim"): + 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}") + _resolve_mla_rope_interleaved(config) + + +def validate_dflash_attention_config(config: Qwen3Config) -> str: + """Validate the selected attention parameterization and return its mode.""" + + attention_mode = resolve_dflash_attention_mode(config) + if attention_mode == "mha" and int(config.num_key_value_heads) != int( + config.num_attention_heads + ): + raise ValueError( + "attention_mode 'mha' requires num_key_value_heads == " + f"num_attention_heads, got {config.num_key_value_heads} and " + f"{config.num_attention_heads}" + ) + if attention_mode == "mla": + validate_dflash_mla_config(config) + return attention_mode + + +def _rope_config(config: Qwen3Config, attention_mode: str) -> Qwen3Config: + """Rotary config for the mode: MLA rotates only the partial-RoPE slice.""" + + if attention_mode != "mla": + return config + rope_config = copy.deepcopy(config) + rope_config.head_dim = config.qk_rope_head_dim + return rope_config + + 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) @@ -79,6 +174,31 @@ def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): return q_embed, k_embed +def _rotate_half_interleaved(x: torch.Tensor) -> torch.Tensor: + """Rotate consecutive pairs, the DeepSeek-style MLA RoPE convention.""" + + paired = x.reshape(*x.shape[:-1], -1, 2) + first, second = paired.unbind(dim=-1) + return torch.stack((-second, first), dim=-1).flatten(-2) + + +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) + + def _prepare_dflash_eager_mask( attention_mask: Optional[torch.Tensor], dtype: torch.dtype, @@ -96,8 +216,17 @@ def _prepare_dflash_eager_mask( return additive_mask, valid_queries -class Qwen3DFlashAttention(nn.Module): - """Multi-headed attention from 'Attention Is All You Need' paper""" +class Qwen3DFlashAttentionBase(nn.Module): + """Shared scaffold for the DFlash-family attention modes. + + Subclasses own only the projection parameterization: ``_init_projections`` + builds the weights and must define ``scaling``, ``num_key_value_groups``, + and ``o_proj`` (the shared forward relies on them); ``_compute_qkv`` + returns rotated ``(q, k, v)`` in ``(batch, heads, seq, dim)`` layout with + keys ordered context-then-draft. Everything the modes must agree on — + KV-cache updates, backend dispatch, fully-masked-query zeroing, and the + output projection — lives here so it is maintained in exactly one place. + """ def __init__( self, @@ -108,46 +237,33 @@ def __init__( super().__init__() self.config = config self.layer_idx = layer_idx - self.head_dim = getattr( - config, "head_dim", config.hidden_size // config.num_attention_heads - ) - self.num_key_value_groups = ( - config.num_attention_heads // config.num_key_value_heads - ) - self.scaling = self.head_dim**-0.5 self.attention_dropout = config.attention_dropout if config._attn_implementation == "flex_attention": assert ( config.attention_dropout == 0.0 ), "DFlash FlexAttention requires attention_dropout=0.0" self.is_causal = False - self.q_proj = nn.Linear( - config.hidden_size, - config.num_attention_heads * self.head_dim, - bias=config.attention_bias, - ) - self.k_proj = nn.Linear( - config.hidden_size, - config.num_key_value_heads * self.head_dim, - bias=config.attention_bias, - ) - self.v_proj = nn.Linear( - config.hidden_size, - config.num_key_value_heads * self.head_dim, - bias=config.attention_bias, - ) - self.o_proj = nn.Linear( - config.num_attention_heads * self.head_dim, - config.hidden_size, - bias=config.attention_bias, - ) - self.q_norm = kernels.make_rms_norm(self.head_dim, config.rms_norm_eps) - self.k_norm = kernels.make_rms_norm(self.head_dim, config.rms_norm_eps) self.sliding_window = ( config.sliding_window if config.layer_types[layer_idx] == SLIDING_ATTENTION else None ) + self._init_projections(config, kernels) + for attribute in ("scaling", "num_key_value_groups", "o_proj"): + assert hasattr( + self, attribute + ), f"_init_projections must define {attribute}" + + def _init_projections(self, config: Qwen3Config, kernels: DFlashKernels) -> None: + raise NotImplementedError + + def _compute_qkv( + self, + hidden_states: torch.Tensor, + target_hidden: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + raise NotImplementedError def forward( self, @@ -160,25 +276,9 @@ def forward( **kwargs: Unpack[FlashAttentionKwargs], ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: bsz, q_len = hidden_states.shape[:-1] - ctx_len = target_hidden.shape[1] - q = self.q_proj(hidden_states) - q = q.view(bsz, q_len, -1, self.head_dim) - q = self.q_norm(q).transpose(1, 2) - k_ctx = self.k_proj(target_hidden) - k_noise = self.k_proj(hidden_states) - v_ctx = self.v_proj(target_hidden) - v_noise = self.v_proj(hidden_states) - k = torch.cat([k_ctx, k_noise], dim=1).view( - bsz, ctx_len + q_len, -1, self.head_dim - ) - v = torch.cat([v_ctx, v_noise], dim=1).view( - bsz, ctx_len + q_len, -1, self.head_dim - ) - k = self.k_norm(k).transpose(1, 2) - v = v.transpose(1, 2) - cos, sin = position_embeddings - q, k = apply_rotary_pos_emb(q, k, cos, sin) + q, k, v = self._compute_qkv(hidden_states, target_hidden, position_embeddings) if past_key_values is not None: + cos, sin = position_embeddings cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} k, v = past_key_values.update(k, v, self.layer_idx, cache_kwargs) valid_queries = None @@ -230,6 +330,219 @@ def forward( return attn_output, attn_weights +class Qwen3DFlashAttention(Qwen3DFlashAttentionBase): + """GQA/MHA projections over the family's context-then-draft KV layout.""" + + def _init_projections(self, config: Qwen3Config, kernels: DFlashKernels) -> None: + self.head_dim = getattr( + config, "head_dim", config.hidden_size // config.num_attention_heads + ) + self.num_key_value_groups = ( + config.num_attention_heads // config.num_key_value_heads + ) + self.scaling = self.head_dim**-0.5 + self.q_proj = nn.Linear( + config.hidden_size, + config.num_attention_heads * self.head_dim, + bias=config.attention_bias, + ) + self.k_proj = nn.Linear( + config.hidden_size, + config.num_key_value_heads * self.head_dim, + bias=config.attention_bias, + ) + self.v_proj = nn.Linear( + config.hidden_size, + config.num_key_value_heads * self.head_dim, + bias=config.attention_bias, + ) + self.o_proj = nn.Linear( + config.num_attention_heads * self.head_dim, + config.hidden_size, + bias=config.attention_bias, + ) + self.q_norm = kernels.make_rms_norm(self.head_dim, config.rms_norm_eps) + self.k_norm = kernels.make_rms_norm(self.head_dim, config.rms_norm_eps) + + def _compute_qkv( + self, + hidden_states: torch.Tensor, + target_hidden: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + bsz, q_len = hidden_states.shape[:-1] + ctx_len = target_hidden.shape[1] + q = self.q_proj(hidden_states) + q = q.view(bsz, q_len, -1, self.head_dim) + q = self.q_norm(q).transpose(1, 2) + k_ctx = self.k_proj(target_hidden) + k_noise = self.k_proj(hidden_states) + v_ctx = self.v_proj(target_hidden) + v_noise = self.v_proj(hidden_states) + k = torch.cat([k_ctx, k_noise], dim=1).view( + bsz, ctx_len + q_len, -1, self.head_dim + ) + v = torch.cat([v_ctx, v_noise], dim=1).view( + bsz, ctx_len + q_len, -1, self.head_dim + ) + k = self.k_norm(k).transpose(1, 2) + v = v.transpose(1, 2) + cos, sin = position_embeddings + q, k = apply_rotary_pos_emb(q, k, cos, sin) + return q, k, v + + +class Qwen3DFlashMLAAttention(Qwen3DFlashAttentionBase): + """Multi-head Latent Attention projections for DFlash-family drafts. + + Standard MLA parameterization: an optional low-rank Q path, a shared + compressed KV latent, and partial RoPE (interleaved or NeoX, from the + standard ``rope_interleave`` field). K/V are expanded per head for + training so the mode runs through the same masks and attention backends + as :class:`Qwen3DFlashAttention`. + """ + + def _init_projections(self, config: Qwen3Config, kernels: DFlashKernels) -> None: + 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 + # DeepSeek YaRN applies mscale_all_dim to the full QK logits; the + # rotary attention factor separately scales only the partial-RoPE slice. + rope_parameters = config.rope_parameters + if rope_parameters.get("rope_type", "default") != "default": + mscale_all_dim = rope_parameters.get("mscale_all_dim", 0) + if mscale_all_dim: + factor = rope_parameters["factor"] + mscale = ( + 1.0 + if factor <= 1 + else 0.1 * mscale_all_dim * math.log(factor) + 1.0 + ) + self.scaling *= mscale * mscale + self.rope_interleaved = _resolve_mla_rope_interleaved(config) + + hidden_size = int(config.hidden_size) + bias = bool(config.attention_bias) + if self.q_lora_rank is None: + self.q_proj = nn.Linear( + hidden_size, + self.num_heads * self.qk_head_dim, + bias=False, + ) + else: + self.q_a_proj = nn.Linear(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=False, + ) + self.kv_a_proj_with_mqa = nn.Linear( + 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=False, + ) + self.o_proj = nn.Linear( + self.num_heads * self.v_head_dim, + hidden_size, + bias=bias, + ) + + 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]: + bsz, seq_len = hidden_states.shape[:2] + kv_compressed, k_rope = self.kv_a_proj_with_mqa(hidden_states).split( + [self.kv_lora_rank, self.qk_rope_head_dim], + dim=-1, + ) + kv = self.kv_b_proj(self.kv_a_layernorm(kv_compressed)).view( + bsz, + 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 _compute_qkv( + self, + hidden_states: torch.Tensor, + target_hidden: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + bsz, q_len = hidden_states.shape[:2] + query = self._project_q(hidden_states).view( + bsz, + q_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, + ) + ctx_k_nope, ctx_k_rope, ctx_value = self._project_kv(target_hidden) + noise_k_nope, noise_k_rope, noise_value = self._project_kv(hidden_states) + k_nope = torch.cat((ctx_k_nope, noise_k_nope), dim=1) + k_rope = torch.cat((ctx_k_rope, noise_k_rope), dim=1) + v = torch.cat((ctx_value, noise_value), dim=1) + # The model-level rotary carries qk_rope_head_dim; queries take the + # trailing positions exactly like apply_rotary_pos_emb. + cos, sin = position_embeddings + q_rope = apply_mla_rope( + q_rope.transpose(1, 2), + cos[:, -q_len:], + sin[:, -q_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) + q = torch.cat((q_nope.transpose(1, 2), q_rope), dim=-1) + k = torch.cat((k_nope.transpose(1, 2), k_rope), dim=-1) + return q, k, v.transpose(1, 2) + + +_DFLASH_ATTENTION_CLASSES = { + "gqa": Qwen3DFlashAttention, + "mha": Qwen3DFlashAttention, + "mla": Qwen3DFlashMLAAttention, +} + + class Qwen3DFlashDecoderLayer(GradientCheckpointingLayer): def __init__( self, @@ -239,7 +552,8 @@ def __init__( ): super().__init__() self.hidden_size = config.hidden_size - self.self_attn = Qwen3DFlashAttention( + attention_cls = _DFLASH_ATTENTION_CLASSES[resolve_dflash_attention_mode(config)] + self.self_attn = attention_cls( config=config, layer_idx=layer_idx, kernels=kernels, @@ -368,6 +682,7 @@ def __init__( super().__init__(config) self.config = config self.layer_types, self.sliding_window = resolve_dflash_attention_layout(config) + self.attention_mode = validate_dflash_attention_config(config) kernels = dflash_kernels or DEFAULT_DFLASH_KERNELS self.layers = nn.ModuleList( [ @@ -381,7 +696,9 @@ def __init__( build_target_layer_ids(config.num_target_layers, config.num_hidden_layers), ) self.norm = kernels.make_rms_norm(config.hidden_size, config.rms_norm_eps) - self.rotary_emb = Qwen3RotaryEmbedding(config) + self.rotary_emb = Qwen3RotaryEmbedding( + _rope_config(config, self.attention_mode) + ) self.fc = nn.Linear( len(self.target_layer_ids) * config.hidden_size, config.hidden_size, diff --git a/specforge/modeling/draft/dspark.py b/specforge/modeling/draft/dspark.py index 1c5853d8b..aebdb43bb 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_mode from .registry import register_draft @@ -287,30 +287,24 @@ class DSparkDraftModel(DFlashDraftModel): def __init__(self, config) -> None: dflash_config = dict(getattr(config, "dflash_config", None) or {}) + attention_mode = resolve_dflash_attention_mode(config) num_heads = int(config.num_attention_heads) num_kv_heads = int(config.num_key_value_heads) - if num_heads % num_kv_heads: + # MLA carries its own head geometry; the query/KV head-count policy + # below only constrains the GQA/MHA projections. + if attention_mode != "mla" 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}" ) - attention_mode = str(dflash_config.get("attention_mode", "gqa")).lower() - if attention_mode not in {"gqa", "mha"}: - raise ValueError( - "DSpark dflash_config.attention_mode must be 'gqa' or 'mha', " - f"got {attention_mode!r}" - ) if attention_mode == "gqa" 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: - raise ValueError( - "DSpark MHA opt-in requires num_key_value_heads == " - f"num_attention_heads, got {num_kv_heads} and {num_heads}" - ) + # 'mha' head-count consistency is enforced by the shared + # validate_dflash_attention_config at model init. dflash_config["attention_mode"] = attention_mode projector_type = dflash_config.get("projector_type") if projector_type is None: diff --git a/tests/test_modeling/test_dflash_mla.py b/tests/test_modeling/test_dflash_mla.py new file mode 100644 index 000000000..0468e115a --- /dev/null +++ b/tests/test_modeling/test_dflash_mla.py @@ -0,0 +1,528 @@ +import copy +import math +import unittest + +import torch +from torch.testing import assert_close +from transformers import DynamicCache, Qwen3Config +from transformers.models.qwen3.modeling_qwen3 import ( + Qwen3ForCausalLM, + Qwen3RotaryEmbedding, +) + +from specforge.algorithms.common.dflash_family_model import create_dflash_block_mask +from specforge.modeling.draft.dflash import ( + DFlashDraftModel, + Qwen3DFlashAttention, + Qwen3DFlashMLAAttention, + apply_mla_rope, +) +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], + } + 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 _position_embeddings( + config: Qwen3Config, + key_len: int, + *, + position_offset: int = 0, + device=None, + dtype: torch.dtype = torch.float32, +) -> tuple[torch.Tensor, torch.Tensor]: + """Model-level rotary embeddings over the MLA partial-RoPE slice.""" + + rope_config = copy.deepcopy(config) + rope_config.head_dim = config.qk_rope_head_dim + rotary = Qwen3RotaryEmbedding(rope_config).to(device) + position_ids = torch.arange( + position_offset, + position_offset + key_len, + device=device, + ).unsqueeze(0) + reference = torch.zeros((1, key_len, 1), device=device, dtype=dtype) + return rotary(reference, position_ids) + + +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_embeddings = _position_embeddings( + attention.config, + key_len, + position_offset=position_offset, + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + return attention( + hidden_states=hidden_states, + target_hidden=target_hidden, + position_embeddings=position_embeddings, + 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) + # Omitted rope_interleave defaults to the DeepSeek interleaved layout. + self.assertTrue(model.layers[0].self_attn.rope_interleaved) + # The model-level rotary rotates only the partial-RoPE slice. + cos, _ = model.rotary_emb(torch.zeros(1, 2, 1), torch.arange(2).unsqueeze(0)) + self.assertEqual(cos.shape[-1], model.config.qk_rope_head_dim) + + def test_omitting_attention_mode_preserves_gqa(self): + config = _mla_config() + del config.dflash_config["attention_mode"] + model = DFlashDraftModel(config) + + self.assertEqual(model.attention_mode, "gqa") + self.assertIsInstance(model.layers[0].self_attn, Qwen3DFlashAttention) + cos, _ = model.rotary_emb(torch.zeros(1, 2, 1), torch.arange(2).unsqueeze(0)) + self.assertEqual(cos.shape[-1], config.head_dim) + + def test_dspark_uses_the_same_generic_mla_backbone(self): + config = _mla_config(architecture="DSparkDraftModel") + # MLA carries its own head geometry, so DSpark's GQA/MHA head-count + # policy must not apply: 3 does not divide 4. + config.num_key_value_heads = 3 + model = DSparkDraftModel(config) + + 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_projection_biases_match_deepseek_mla(self): + config = _mla_config() + config.attention_bias = True + attention = DFlashDraftModel(config).layers[0].self_attn + + self.assertIsNotNone(attention.q_a_proj.bias) + self.assertIsNone(attention.q_b_proj.bias) + self.assertIsNotNone(attention.kv_a_proj_with_mqa.bias) + self.assertIsNone(attention.kv_b_proj.bias) + self.assertIsNotNone(attention.o_proj.bias) + + direct_config = _mla_config(q_lora_rank=None) + direct_config.attention_bias = True + direct_attention = DFlashDraftModel(direct_config).layers[0].self_attn + self.assertIsNone(direct_attention.q_proj.bias) + + def test_deepseek_yarn_scales_attention_logits(self): + config = _mla_config() + config.rope_parameters = { + "rope_type": "yarn", + "rope_theta": 10_000.0, + "factor": 40.0, + "original_max_position_embeddings": 4_096, + "mscale": 0.707, + "mscale_all_dim": 0.707, + } + attention = DFlashDraftModel(config).layers[0].self_attn + + mscale = 0.1 * 0.707 * math.log(40.0) + 1.0 + qk_head_dim = config.qk_nope_head_dim + config.qk_rope_head_dim + expected = mscale**2 / math.sqrt(qk_head_dim) + self.assertAlmostEqual(attention.scaling, expected) + + 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", + "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_rope_conventions_match_independent_reference_values(self): + hidden = torch.tensor([[[[1.0, 2.0, 3.0, 4.0]]]]) + cos = torch.tensor([[[0.8, 0.6, 0.8, 0.6]]]) + sin = torch.tensor([[[0.6, 0.8, 0.6, 0.8]]]) + expected = { + True: torch.tensor([[[[-0.4, 2.2, -1.4, 4.8]]]]), + False: torch.tensor([[[[-1.0, -2.0, 3.0, 4.0]]]]), + } + + for interleaved, reference in expected.items(): + with self.subTest(interleaved=interleaved): + actual = apply_mla_rope( + hidden, + cos, + sin, + interleaved=interleaved, + ) + assert_close(actual, reference) + + def test_query_rope_uses_the_suffix_positions(self): + torch.manual_seed(17) + config = _mla_config(q_lora_rank=None) + attention = Qwen3DFlashMLAAttention( + config, + layer_idx=0, + kernels=DEFAULT_DFLASH_KERNELS, + ).eval() + hidden_states = torch.randn(1, 2, config.hidden_size) + target_hidden = torch.randn(1, 3, config.hidden_size) + cos, sin = _position_embeddings(config, key_len=5) + + query, _, _ = attention._compute_qkv( + hidden_states, + target_hidden, + (cos, sin), + ) + raw_query = attention.q_proj(hidden_states).view( + 1, + 2, + attention.num_heads, + attention.qk_head_dim, + ) + raw_query = raw_query.transpose(1, 2) + raw_query_rope = raw_query[..., attention.qk_nope_head_dim :] + expected_rope = apply_mla_rope( + raw_query_rope, + cos[:, -hidden_states.shape[1] :], + sin[:, -hidden_states.shape[1] :], + interleaved=True, + ) + + assert_close(query[..., attention.qk_nope_head_dim :], expected_rope) + + 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): + torch.manual_seed(13) + attention = Qwen3DFlashMLAAttention( + _mla_config(), + layer_idx=0, + kernels=DEFAULT_DFLASH_KERNELS, + ).eval() + cache = DynamicCache() + first_hidden = torch.randn(1, 2, 24) + first_target = torch.randn(1, 3, 24) + second_hidden = torch.randn(1, 2, 24) + second_target = torch.randn(1, 1, 24) + + first_output, _ = _attention_forward( + attention, + first_hidden, + first_target, + 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, + second_hidden, + second_target, + None, + past_key_values=cache, + position_offset=5, + ) + self.assertEqual(second_output.shape, (1, 2, 24)) + self.assertEqual(cache.get_seq_length(), 8) + + full_target = torch.cat( + [first_target, first_hidden, second_target], + dim=1, + ) + uncached_output, _ = _attention_forward( + attention, + second_hidden, + full_target, + None, + ) + assert_close(second_output, uncached_output, rtol=1e-5, atol=1e-6) + + @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 = 24 + 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 TestDFlashMLASpecGenerate(unittest.TestCase): + def test_spec_generate_decode_smoke(self): + torch.manual_seed(3) + target_config = Qwen3Config( + hidden_size=24, + intermediate_size=48, + num_attention_heads=4, + num_key_value_heads=2, + num_hidden_layers=2, + head_dim=6, + max_position_embeddings=128, + vocab_size=64, + tie_word_embeddings=False, + ) + target_config._attn_implementation = "sdpa" + target = Qwen3ForCausalLM(target_config).eval() + + config = _mla_config() + config.dflash_config["mask_token_id"] = 0 + model = DFlashDraftModel(config).eval() + + input_ids = torch.randint(1, 64, (1, 6)) + output_ids = model.spec_generate( + target, + input_ids, + max_new_tokens=8, + stop_token_ids=None, + temperature=0.0, + ) + + self.assertEqual(output_ids.shape[0], 1) + self.assertLessEqual(output_ids.shape[1], input_ids.shape[1] + 8) + self.assertTrue(torch.equal(output_ids[:, :6], input_ids)) + + +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) + + def test_rejects_non_boolean_rope_interleave(self): + # A JSON null or string must fail loudly instead of silently flipping + # the rotation convention. + for value in (None, "false"): + with self.subTest(value=value): + config = _mla_config() + config.rope_interleave = value + with self.assertRaisesRegex(ValueError, "rope_interleave"): + DFlashDraftModel(config) + + def test_rope_interleave_false_selects_neox_rotation(self): + config = _mla_config() + config.rope_interleave = False + + attention = DFlashDraftModel(config).layers[0].self_attn + + self.assertFalse(attention.rope_interleaved) + + def test_mha_requires_equal_head_counts(self): + config = _mla_config() + config.dflash_config["attention_mode"] = "mha" + self.assertNotEqual(config.num_key_value_heads, config.num_attention_heads) + + with self.assertRaisesRegex(ValueError, "num_key_value_heads"): + DFlashDraftModel(config) + + config.num_key_value_heads = config.num_attention_heads + model = DFlashDraftModel(config) + self.assertIsInstance(model.layers[0].self_attn, Qwen3DFlashAttention) + self.assertEqual(model.layers[0].self_attn.num_key_value_groups, 1) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/tests/test_runtime/test_model_loading.py b/tests/test_runtime/test_model_loading.py index 419b28ffe..b6e1c8ae9 100644 --- a/tests/test_runtime/test_model_loading.py +++ b/tests/test_runtime/test_model_loading.py @@ -225,6 +225,29 @@ def test_dflash_layer_override_rejects_ambiguous_hybrid_resize(self): provider=_draft_config_provider("dflash"), ) + def test_dflash_resolution_validates_mla_before_model_construction(self): + with tempfile.TemporaryDirectory() as directory: + path = os.path.join(directory, "draft.json") + payload = _draft_payload("DFlashDraftModel", block_size=16) + payload.update( + q_lora_rank=16, + kv_lora_rank=8, + qk_nope_head_dim=4, + qk_rope_head_dim=4, + v_head_dim=8, + rope_interleave="false", + ) + payload["dflash_config"] = {"attention_mode": "mla"} + with open(path, "w", encoding="utf-8") as stream: + json.dump(payload, stream) + + cfg = _run_config("dflash", draft_model_config=path) + with self.assertRaisesRegex(ValueError, "rope_interleave"): + resolve_draft_config( + cfg, + provider=_draft_config_provider("dflash"), + ) + def test_local_json_and_directory_are_equivalent_sources(self): with tempfile.TemporaryDirectory() as directory: path = os.path.join(directory, "config.json") diff --git a/tests/test_scripts/test_gate_orchestration.py b/tests/test_scripts/test_gate_orchestration.py index 30ced360a..e277e6405 100644 --- a/tests/test_scripts/test_gate_orchestration.py +++ b/tests/test_scripts/test_gate_orchestration.py @@ -178,6 +178,25 @@ def test_normalizes_dspark_for_sglang(self): ) self.assertEqual(json.loads(path.read_text()), normalized) + def test_rejects_mla_before_rewriting_the_export(self): + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "config.json" + original = { + "architectures": ["DFlashDraftModel"], + "auto_map": {"AutoModel": "dflash.DFlashDraftModel"}, + "block_size": 16, + "dflash_config": { + "projector_type": "dflash", + "attention_mode": "mla", + }, + } + path.write_text(json.dumps(original), encoding="utf-8") + + with self.assertRaisesRegex(ValueError, "only GQA/MHA"): + self.module.normalize_export(str(path), 16) + + self.assertEqual(json.loads(path.read_text()), original) + def test_rejects_dspark_without_a_positive_markov_rank(self): with tempfile.TemporaryDirectory() as tmp: path = Path(tmp) / "config.json"