dflash.py
18.6 KB · 471 lines · python Raw
1 import os
2 from functools import partial
3 from typing import Callable, Optional
4
5 import torch
6 from torch import nn
7 from torch.nn.attention.flex_attention import BlockMask, flex_attention
8 from transformers import DynamicCache
9 from transformers.cache_utils import Cache
10 from transformers.modeling_outputs import CausalLMOutputWithPast
11 from transformers.models.qwen3.modeling_qwen3 import (
12 ALL_ATTENTION_FUNCTIONS,
13 FlashAttentionKwargs,
14 GradientCheckpointingLayer,
15 Qwen3Config,
16 Qwen3MLP,
17 Qwen3PreTrainedModel,
18 Qwen3RMSNorm,
19 Qwen3RotaryEmbedding,
20 eager_attention_forward,
21 rotate_half,
22 )
23 from typing_extensions import Tuple, Unpack
24
25 # FlashAttention-4 flex backend, opt-in via env SPECFORGE_DRAFT_FLEX_BACKEND=fa4.
26 # flex_attention with kernel_options={"BACKEND": "FLASH"} runs the FA4 kernel instead
27 # of the Triton flex kernel (ref: meta-pytorch/attention-gym flex_flash_attention.py).
28 #
29 # STATUS on this stack (torch 2.11 / GB300 sm_10.3): the FA4 kernel works for a small
30 # head_dim (64/128) BUT ONLY for a mask_mod that captures no tensors (it needs the
31 # asymmetric block sparsity BLOCK_SIZE=(q=256, kv=128); see core/dflash.py). The DSpark
32 # dual-source mask (`create_dflash_block_mask`) MUST capture per-sample `anchor_positions`
33 # / `block_keep_mask` tensors, and the FA4 CuteDSL template fails on any captured-tensor
34 # mask_mod ("CuteDSL template failed"), both dynamic=True and False. => FA4 is currently
35 # NOT usable for the DSpark drafter; the default Triton flex backend (used when this env
36 # is unset) handles the captured-tensor mask correctly and is the supported path.
37 # (The DeepSeek-V4 DSpark draft was likewise FA4-ruled-out, there for head_dim 512.)
38 # The code path is kept, gated + off by default, for a future stack / a captured-tensor-
39 # free mask formulation. dynamic=True: draft Q/context lengths vary per batch.
40 _FLEX_FA4_COMPILED = None
41
42
43 def _flex_fa4():
44 global _FLEX_FA4_COMPILED
45 if _FLEX_FA4_COMPILED is None:
46 _FLEX_FA4_COMPILED = torch.compile(
47 partial(flex_attention, kernel_options={"BACKEND": "FLASH"}),
48 dynamic=True,
49 )
50 return _FLEX_FA4_COMPILED
51
52
53 def sample(logits: torch.Tensor, temperature: float = 0.0) -> torch.Tensor:
54 if temperature < 1e-5:
55 return torch.argmax(logits, dim=-1)
56 bsz, seq_len, vocab_size = logits.shape
57 logits = logits.view(-1, vocab_size)
58 logits = logits / temperature
59 probs = torch.softmax(logits, dim=-1)
60 return torch.multinomial(probs, num_samples=1).view(bsz, seq_len)
61
62
63 def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
64 cos = cos.unsqueeze(unsqueeze_dim)
65 sin = sin.unsqueeze(unsqueeze_dim)
66 q_len = q.size(-2)
67 q_embed = (q * cos[..., -q_len:, :]) + (rotate_half(q) * sin[..., -q_len:, :])
68 k_embed = (k * cos) + (rotate_half(k) * sin)
69 return q_embed, k_embed
70
71
72 class Qwen3DFlashAttention(nn.Module):
73 """Multi-headed attention from 'Attention Is All You Need' paper"""
74
75 def __init__(self, config: Qwen3Config, layer_idx: int):
76 super().__init__()
77 self.config = config
78 self.layer_idx = layer_idx
79 self.head_dim = getattr(
80 config, "head_dim", config.hidden_size // config.num_attention_heads
81 )
82 num_attention_heads = int(config.num_attention_heads)
83 num_key_value_heads = int(config.num_key_value_heads)
84 if (
85 num_attention_heads <= 0
86 or num_key_value_heads <= 0
87 or num_attention_heads % num_key_value_heads != 0
88 ):
89 raise ValueError(
90 "Qwen3DFlashAttention requires positive attention head counts and "
91 "num_attention_heads divisible by num_key_value_heads, got "
92 f"num_attention_heads={num_attention_heads}, "
93 f"num_key_value_heads={num_key_value_heads}."
94 )
95 self.num_key_value_groups = num_attention_heads // num_key_value_heads
96 self.scaling = self.head_dim**-0.5
97 self.attention_dropout = config.attention_dropout
98 self.is_causal = False
99 self.q_proj = nn.Linear(
100 config.hidden_size,
101 config.num_attention_heads * self.head_dim,
102 bias=config.attention_bias,
103 )
104 self.k_proj = nn.Linear(
105 config.hidden_size,
106 config.num_key_value_heads * self.head_dim,
107 bias=config.attention_bias,
108 )
109 self.v_proj = nn.Linear(
110 config.hidden_size,
111 config.num_key_value_heads * self.head_dim,
112 bias=config.attention_bias,
113 )
114 self.o_proj = nn.Linear(
115 config.num_attention_heads * self.head_dim,
116 config.hidden_size,
117 bias=config.attention_bias,
118 )
119 self.q_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps)
120 self.k_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps)
121 self.sliding_window = (
122 config.sliding_window
123 if config.layer_types[layer_idx] == "sliding_attention"
124 else None
125 )
126
127 # Keep the dual-source attention OUT of any outer torch.compile region: the
128 # flex_attention block-mask HOP fails inductor lowering when nested inside a
129 # larger dynamic-shape graph on this stack (CantSplit / "unsupported operand &"),
130 # even though it compiles fine on its own (HF's flex integration compiles it
131 # separately). Marking the attention compiler-disabled lets SPECFORGE_COMPILE_DRAFT
132 # fuse the rest of the block (RoPE, RMSNorm, MLP, residual) while the flex
133 # attention keeps its own (block-sparse, separately-compiled) fast path.
134 @torch.compiler.disable
135 def forward(
136 self,
137 hidden_states: torch.Tensor,
138 target_hidden: torch.Tensor,
139 position_embeddings: tuple[torch.Tensor, torch.Tensor],
140 attention_mask: Optional[torch.Tensor],
141 past_key_values: Optional[Cache] = None,
142 cache_position: Optional[torch.LongTensor] = None,
143 **kwargs: Unpack[FlashAttentionKwargs],
144 ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
145 bsz, q_len = hidden_states.shape[:-1]
146 ctx_len = target_hidden.shape[1]
147 q = self.q_proj(hidden_states)
148 q = q.view(bsz, q_len, -1, self.head_dim)
149 q = self.q_norm(q).transpose(1, 2)
150 k_ctx = self.k_proj(target_hidden)
151 k_noise = self.k_proj(hidden_states)
152 v_ctx = self.v_proj(target_hidden)
153 v_noise = self.v_proj(hidden_states)
154 k = torch.cat([k_ctx, k_noise], dim=1).view(
155 bsz, ctx_len + q_len, -1, self.head_dim
156 )
157 v = torch.cat([v_ctx, v_noise], dim=1).view(
158 bsz, ctx_len + q_len, -1, self.head_dim
159 )
160 k = self.k_norm(k).transpose(1, 2)
161 v = v.transpose(1, 2)
162 cos, sin = position_embeddings
163 q, k = apply_rotary_pos_emb(q, k, cos, sin)
164 if past_key_values is not None:
165 cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
166 k, v = past_key_values.update(k, v, self.layer_idx, cache_kwargs)
167 # FA4 flex path (opt-in): call flex_attention with the FLASH kernel directly on
168 # the prebuilt dual-source BlockMask, bypassing HF's Triton-flex wrapper. q/k/v
169 # are already [B, H, S, D]; flex returns [B, H, S, D] -> transpose to [B, S, H, D]
170 # to match the reshape below. GQA (num_kv_heads < num_heads) via enable_gqa.
171 if (
172 self.config._attn_implementation in ("flex_attention", "flex")
173 and os.environ.get("SPECFORGE_DRAFT_FLEX_BACKEND") == "fa4"
174 and isinstance(attention_mask, BlockMask)
175 ):
176 attn_output = _flex_fa4()(
177 q,
178 k,
179 v,
180 block_mask=attention_mask,
181 scale=self.scaling,
182 enable_gqa=(self.num_key_value_groups > 1),
183 )
184 attn_output = attn_output.transpose(1, 2).contiguous()
185 attn_weights = None
186 else:
187 attn_fn: Callable = eager_attention_forward
188 if self.config._attn_implementation != "eager":
189 attn_fn = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
190 attn_output, attn_weights = attn_fn(
191 self,
192 q,
193 k,
194 v,
195 attention_mask,
196 dropout=0.0 if not self.training else self.attention_dropout,
197 scaling=self.scaling,
198 sliding_window=self.sliding_window,
199 **kwargs,
200 )
201 attn_output = attn_output.reshape(bsz, q_len, -1)
202 attn_output = self.o_proj(attn_output)
203 return attn_output, attn_weights
204
205
206 class Qwen3DFlashDecoderLayer(GradientCheckpointingLayer):
207 def __init__(self, config: Qwen3Config, layer_idx: int):
208 super().__init__()
209 self.hidden_size = config.hidden_size
210 self.self_attn = Qwen3DFlashAttention(config=config, layer_idx=layer_idx)
211 self.mlp = Qwen3MLP(config)
212 self.input_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
213 self.post_attention_layernorm = Qwen3RMSNorm(
214 config.hidden_size, eps=config.rms_norm_eps
215 )
216
217 def forward(
218 self,
219 target_hidden: Optional[torch.Tensor] = None,
220 hidden_states: Optional[torch.Tensor] = None,
221 attention_mask: Optional[torch.Tensor] = None,
222 position_ids: Optional[torch.LongTensor] = None,
223 past_key_value: Optional[Cache] = None,
224 output_attentions: Optional[bool] = False,
225 use_cache: Optional[bool] = False,
226 cache_position: Optional[torch.LongTensor] = None,
227 position_embeddings: Optional[
228 Tuple[torch.Tensor, torch.Tensor]
229 ] = None, # necessary, but kept here for BC
230 **kwargs: Unpack[FlashAttentionKwargs],
231 ) -> Tuple[
232 torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]
233 ]:
234 residual = hidden_states
235 hidden_states = self.input_layernorm(hidden_states)
236 hidden_states = self.self_attn(
237 hidden_states=hidden_states,
238 target_hidden=target_hidden,
239 attention_mask=attention_mask,
240 position_ids=position_ids,
241 past_key_values=past_key_value,
242 output_attentions=output_attentions,
243 use_cache=use_cache,
244 cache_position=cache_position,
245 position_embeddings=position_embeddings,
246 **kwargs,
247 )[0]
248 hidden_states = residual + hidden_states
249 residual = hidden_states
250 hidden_states = self.post_attention_layernorm(hidden_states)
251 hidden_states = self.mlp(hidden_states)
252 hidden_states = residual + hidden_states
253 return hidden_states
254
255
256 def build_target_layer_ids(num_target_layers: int, num_draft_layers: int):
257 if num_draft_layers == 1:
258 return [(num_target_layers // 2)]
259 start = 1
260 end = num_target_layers - 3
261 span = end - start
262 target_layer_ids = [
263 int(round(start + (i * span) / (num_draft_layers - 1)))
264 for i in range(num_draft_layers)
265 ]
266 return target_layer_ids
267
268
269 def extract_context_feature(
270 hidden_states: list[torch.Tensor],
271 layer_ids: Optional[list[int]],
272 ) -> torch.Tensor:
273 offset = 1
274 selected_states = []
275 for layer_id in layer_ids:
276 selected_states.append(hidden_states[layer_id + offset])
277 target_hidden = torch.cat(selected_states, dim=-1)
278 return target_hidden
279
280
281 class DFlashDraftModel(Qwen3PreTrainedModel):
282 config_class = Qwen3Config
283 _no_split_modules = ["Qwen3DFlashDecoderLayer"]
284
285 def __init__(self, config) -> None:
286 super().__init__(config)
287 self.config = config
288 self.layers = nn.ModuleList(
289 [
290 Qwen3DFlashDecoderLayer(config, layer_idx)
291 for layer_idx in range(config.num_hidden_layers)
292 ]
293 )
294 dflash_config = getattr(config, "dflash_config", {}) or {}
295 self.target_layer_ids = dflash_config.get(
296 "target_layer_ids",
297 build_target_layer_ids(config.num_target_layers, config.num_hidden_layers),
298 )
299 self.norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
300 self.rotary_emb = Qwen3RotaryEmbedding(config)
301 self.fc = nn.Linear(
302 len(self.target_layer_ids) * config.hidden_size,
303 config.hidden_size,
304 bias=False,
305 )
306 self.hidden_norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
307 self.block_size = config.block_size
308 self.mask_token_id = dflash_config.get("mask_token_id", None)
309 self.projector_type = dflash_config.get("projector_type", None)
310 self.pure_draft_prefix_len = dflash_config.get("pure_draft_prefix_len", 0)
311 self.shift_label = dflash_config.get("shift_label", False)
312
313 if self.projector_type == "domino":
314 self.emb_dim = dflash_config["emb_dim"]
315 self.gru_hidden_dim = dflash_config["gru_hidden_dim"]
316 self.prefix_gru = nn.GRU(
317 input_size=config.hidden_size,
318 hidden_size=self.gru_hidden_dim,
319 num_layers=1,
320 batch_first=True,
321 bias=False,
322 )
323 in_dim = config.hidden_size + self.gru_hidden_dim
324 self.embed_proj = nn.Sequential(
325 nn.Linear(in_dim, self.emb_dim, bias=False),
326 nn.SiLU(),
327 nn.Linear(self.emb_dim, config.vocab_size, bias=False),
328 )
329 elif self.projector_type is not None:
330 raise ValueError(f"Unknown draft projector_type: {self.projector_type}")
331 self.post_init()
332
333 def forward(
334 self,
335 position_ids: torch.LongTensor,
336 attention_mask: Optional[torch.Tensor] = None,
337 noise_embedding: Optional[torch.Tensor] = None,
338 target_hidden: Optional[torch.Tensor] = None,
339 past_key_values: Optional[Cache] = None,
340 use_cache: bool = False,
341 **kwargs,
342 ) -> CausalLMOutputWithPast:
343 hidden_states = noise_embedding
344 target_hidden = self.hidden_norm(self.fc(target_hidden))
345 position_embeddings = self.rotary_emb(hidden_states, position_ids)
346 for layer in self.layers:
347 hidden_states = layer(
348 hidden_states=hidden_states,
349 target_hidden=target_hidden,
350 attention_mask=attention_mask,
351 position_ids=position_ids,
352 past_key_value=past_key_values,
353 use_cache=use_cache,
354 position_embeddings=position_embeddings,
355 **kwargs,
356 )
357 return self.norm(hidden_states)
358
359 @torch.inference_mode()
360 def spec_generate(
361 self,
362 target: nn.Module,
363 input_ids: torch.LongTensor,
364 max_new_tokens: int,
365 stop_token_ids: list[int],
366 temperature: float,
367 ):
368 self.eval()
369 num_input_tokens = input_ids.shape[1]
370 max_length = num_input_tokens + max_new_tokens
371
372 block_size = self.block_size
373 output_ids = torch.full(
374 (1, max_length + block_size),
375 self.mask_token_id,
376 dtype=torch.long,
377 device=target.device,
378 )
379 position_ids = torch.arange(
380 output_ids.shape[1], device=target.device
381 ).unsqueeze(0)
382
383 past_key_values_target = DynamicCache()
384 past_key_values_draft = DynamicCache()
385
386 # Prefill stage
387 output = target(
388 input_ids,
389 position_ids=position_ids[:, :num_input_tokens],
390 past_key_values=past_key_values_target,
391 use_cache=True,
392 logits_to_keep=1,
393 output_hidden_states=True,
394 )
395
396 output_ids[:, :num_input_tokens] = input_ids
397 output_ids[:, num_input_tokens : num_input_tokens + 1] = sample(
398 output.logits, temperature
399 )
400 target_hidden = extract_context_feature(
401 output.hidden_states, self.target_layer_ids
402 )
403
404 # Decode stage
405 acceptance_lengths = []
406 start = input_ids.shape[1]
407 while start < max_length:
408 block_output_ids = output_ids[:, start : start + block_size].clone()
409 block_position_ids = position_ids[:, start : start + block_size]
410 noise_embedding = target.model.embed_tokens(block_output_ids)
411 draft_logits = target.lm_head(
412 self(
413 target_hidden=target_hidden,
414 noise_embedding=noise_embedding,
415 position_ids=position_ids[
416 :, past_key_values_draft.get_seq_length() : start + block_size
417 ],
418 past_key_values=past_key_values_draft,
419 use_cache=True,
420 is_causal=False,
421 )[:, -block_size + 1 :, :]
422 )
423 past_key_values_draft.crop(start)
424 block_output_ids[:, 1:] = sample(draft_logits)
425
426 output = target(
427 block_output_ids,
428 position_ids=block_position_ids,
429 past_key_values=past_key_values_target,
430 use_cache=True,
431 output_hidden_states=True,
432 )
433
434 posterior = sample(output.logits, temperature)
435 acceptance_length = (
436 (block_output_ids[:, 1:] == posterior[:, :-1])
437 .cumprod(dim=1)
438 .sum(dim=1)[0]
439 .item()
440 )
441 output_ids[:, start : start + acceptance_length + 1] = block_output_ids[
442 :, : acceptance_length + 1
443 ]
444 output_ids[:, start + acceptance_length + 1] = posterior[
445 :, acceptance_length
446 ]
447 start += acceptance_length + 1
448 past_key_values_target.crop(start)
449 target_hidden = extract_context_feature(
450 output.hidden_states, self.target_layer_ids
451 )[:, : acceptance_length + 1, :]
452 acceptance_lengths.append(acceptance_length + 1)
453 if stop_token_ids is not None and any(
454 stop_token_id in output_ids[:, num_input_tokens:]
455 for stop_token_id in stop_token_ids
456 ):
457 break
458 output_ids = output_ids[:, :max_length]
459 output_ids = output_ids[:, output_ids[0] != self.mask_token_id]
460 if stop_token_ids is not None:
461 stop_token_ids = torch.tensor(stop_token_ids, device=output_ids.device)
462 stop_token_indices = torch.isin(
463 output_ids[0][num_input_tokens:], stop_token_ids
464 ).nonzero(as_tuple=True)[0]
465 if stop_token_indices.numel() > 0:
466 output_ids = output_ids[
467 :, : num_input_tokens + stop_token_indices[0] + 1
468 ]
469
470 return output_ids
471