Skip to content

Commit fe9354f

Browse files
dusterbloomdusterbloom
committed
feat(gemma4): add dflash draft runtime
1 parent 686581a commit fe9354f

5 files changed

Lines changed: 761 additions & 8 deletions

File tree

‎dflash/CMakeLists.txt‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -464,6 +464,11 @@ if(DFLASH27B_TESTS)
464464
target_include_directories(smoke_load_gemma4_draft PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS})
465465
target_link_libraries(smoke_load_gemma4_draft PRIVATE dflash27b ggml ggml-cuda)
466466
endif()
467+
if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/gemma4/smoke_gemma4_draft_forward.cpp")
468+
add_executable(smoke_gemma4_draft_forward test/gemma4/smoke_gemma4_draft_forward.cpp)
469+
target_include_directories(smoke_gemma4_draft_forward PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS})
470+
target_link_libraries(smoke_gemma4_draft_forward PRIVATE dflash27b ggml ggml-cuda)
471+
endif()
467472
if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/smoke_load_target_laguna.cpp")
468473
add_executable(smoke_load_target_laguna test/smoke_load_target_laguna.cpp)
469474
target_include_directories(smoke_load_target_laguna PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS})

‎dflash/src/gemma4_dflash_graph.cpp‎

Lines changed: 336 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,35 @@
1-
// Loaders for the Gemma4 DFlash draft model weights.
2-
// Graph builders (build_draft_kv_prefill_graph, build_gemma4_draft_graph)
3-
// land in the dflash-runtime PR.
1+
// Builds ggml compute graphs for the Gemma4 DFlash draft model
2+
// (5-layer block-diffusion model with KV cache and logit softcapping).
43
//
54
// Architecture:
6-
// - 5-layer block-diffusion draft (4 SWA + 1 full attention)
7-
// - 6 captured target layers
8-
// - FC input = 6 * target_hidden (target_hidden = 4096 for all Gemma4 variants),
9-
// giving FC width = 24576
5+
// - 6 captured target layers (Qwen3 used 5)
6+
// - FC input = 6 * target_hidden, where target_hidden = 4096 for all Gemma4
7+
// variants (31B dense and 26B-A4B MoE), giving FC width = 24576
108
// - Logit softcapping: tanh(logits / cap) * cap, cap = 30.0
119
// - Tied lm_head: uses tok_embd transposed (or a provided lm_head weight)
1210
// - Vocab = 262144
13-
// - Draft has its own lm_head + softcap
11+
// - Draft has its own lm_head + softcap — it does NOT rely on the target's
12+
// lm_head (unlike the Qwen3 draft which shares the target's projection)
13+
// - KV cache (prefix-direct): target features are projected into per-layer
14+
// K/V entries and stored in GemmaTargetCache::draft_k/draft_v.
15+
// build_draft_kv_prefill_graph materializes the context K/V;
16+
// build_gemma4_draft_graph writes block K/V and attends over the full cache.
17+
// - Layer types: 4 SWA (sliding_attention) + 1 full attention
18+
// The attention kernel itself is the same ggml_flash_attn_ext call in both
19+
// cases; the caller controls the mask to implement the sliding window.
20+
//
21+
// Two-step per-decode:
22+
// 1. build_draft_kv_prefill_graph: project new committed context tokens into
23+
// draft KV cache (side-effect only; nullptr returned).
24+
// 2. build_gemma4_draft_graph: attend over context+block K/V and return logits.
25+
//
26+
// build_gemma4_draft_graph takes:
27+
// - draft_embed [draft_hidden, n_tokens] f32 (MASK token embeddings)
28+
// - positions [n_tokens] i32 (absolute token positions)
29+
// - attn_mask [kv_pad, q_pad] f16 (causal over context+block)
30+
// - kv_start = cache.draft_kv_pos (context length before this block)
31+
// and returns:
32+
// - logits [n_vocab, n_tokens] f32 (after softcapping)
1433
//
1534
// Safetensors tensor naming (actual file, no model. prefix):
1635
// fc.weight → fc
@@ -57,6 +76,315 @@
5776

5877
namespace dflash27b {
5978

79+
// ─── Draft SWA truncation toggle ──────────────────────────────────────────
80+
// Set DFLASH_DRAFT_SWA_TRUNC=1 to enable per-layer K/V truncation in the
81+
// draft graph for SWA layers (last n-1 layers — the final layer is full).
82+
// Mirrors PR #129 for the qwen3 drafter, ported to gemma4's cached layout.
83+
static inline bool draft_swa_trunc_enabled() {
84+
static int e = -1;
85+
if (e < 0) {
86+
const char * v = std::getenv("DFLASH_DRAFT_SWA_TRUNC");
87+
e = (v && std::atoi(v) != 0) ? 1 : 0;
88+
if (e) {
89+
std::fprintf(stderr, "[draft-swa-trunc] enabled\n");
90+
}
91+
}
92+
return e == 1;
93+
}
94+
95+
// ─── Draft RoPE wrapper with optional YaRN extrapolation ──────────────────
96+
// Set DFLASH_DRAFT_YARN=1 to enable YaRN scaling for draft RoPE; assumes the
97+
// draft was effectively trained at DFLASH_DRAFT_YARN_NCTX_ORIG (default 32768)
98+
// despite config.json claiming a larger max_position_embeddings.
99+
static inline ggml_tensor * draft_rope(ggml_context * ctx, ggml_tensor * x,
100+
ggml_tensor * positions, int head_dim,
101+
float rope_base) {
102+
static struct {
103+
int nctx;
104+
float ext;
105+
float bf;
106+
float bs;
107+
bool init;
108+
} p = {0, 0.0f, 0.0f, 0.0f, false};
109+
if (!p.init) {
110+
const char * en = std::getenv("DFLASH_DRAFT_YARN");
111+
if (en && std::atoi(en) != 0) {
112+
const char * nc = std::getenv("DFLASH_DRAFT_YARN_NCTX_ORIG");
113+
p.nctx = nc ? std::atoi(nc) : 32768;
114+
p.ext = 1.0f;
115+
p.bf = 32.0f;
116+
p.bs = 1.0f;
117+
std::fprintf(stderr,
118+
"[draft-yarn] enabled: n_ctx_orig=%d ext_factor=%.2f beta_fast=%.1f beta_slow=%.1f\n",
119+
p.nctx, p.ext, p.bf, p.bs);
120+
}
121+
p.init = true;
122+
}
123+
return ggml_rope_ext(ctx, x, positions, /*freq_factors=*/nullptr,
124+
head_dim, GGML_ROPE_TYPE_NEOX, p.nctx,
125+
rope_base, /*freq_scale=*/1.0f,
126+
p.ext, /*attn_factor=*/1.0f, p.bf, p.bs);
127+
}
128+
129+
// ─── Graph builders ───────────────────────────────────────────────────────
130+
131+
// build_draft_kv_prefill_graph — prefix-direct KV materialisation (SGLang style).
132+
//
133+
// Projects n_tokens new context positions through the draft model's Wk / Wv
134+
// (after FC → ctx_hidden) and writes the resulting K, V tensors into
135+
// cache.draft_k[il] / cache.draft_v[il] starting at offset cache.draft_kv_pos.
136+
//
137+
// The function is side-effect only: it expands ggml_cpy ops into gf and
138+
// returns nullptr. The caller must ggml_graph_compute(gf) to materialise
139+
// the cache entries, then increment cache.draft_kv_pos by n_tokens.
140+
//
141+
// target_feat [6*target_hidden, n_tokens] f32
142+
// positions [n_tokens] i32 (absolute positions for RoPE)
143+
ggml_tensor * build_draft_kv_prefill_graph(
144+
ggml_context * ctx,
145+
ggml_cgraph * gf,
146+
const GemmaDraftWeights & w,
147+
GemmaTargetCache & cache,
148+
ggml_tensor * target_feat,
149+
ggml_tensor * positions,
150+
int n_tokens)
151+
{
152+
// Guard: writing cache.draft_kv_pos..cache.draft_kv_pos+n_tokens-1 must fit.
153+
if (cache.draft_k.empty() ||
154+
cache.draft_kv_pos < 0 ||
155+
cache.draft_kv_pos + n_tokens > (int)cache.draft_k[0]->ne[2]) {
156+
const int tensor_cap = cache.draft_k.empty() ? -1 : (int)cache.draft_k[0]->ne[2];
157+
GGML_ABORT("draft KV prefill out of bounds: draft_kv_pos=%d n_tokens=%d cap=%d tensor_cap=%d",
158+
cache.draft_kv_pos, n_tokens, cache.draft_kv_cap, tensor_cap);
159+
}
160+
161+
const int n_kv = w.n_head_kv;
162+
const int head_dim = w.head_dim;
163+
const float eps = GEMMA4_RMS_EPS;
164+
const float rope_base = w.rope_theta;
165+
166+
// ── 1. FC projection: ctx_hidden = fc @ target_feat → [n_embd, n_tokens]
167+
ggml_tensor * ctx_hidden = ggml_mul_mat(ctx, w.fc, target_feat);
168+
// hidden_norm: RMSNorm applied right after the fc projection
169+
// (matches qwen3_dflash_graph.cpp:57-59)
170+
ctx_hidden = ggml_rms_norm(ctx, ctx_hidden, eps);
171+
ctx_hidden = ggml_mul(ctx, ctx_hidden, w.hidden_norm);
172+
ggml_set_name(ctx_hidden, "draft_kv_prefill_ctx_hidden");
173+
174+
// ── 2. Per-layer K / V projection, normalisation, RoPE, cache write
175+
for (int il = 0; il < w.n_layer; il++) {
176+
const GemmaDraftLayer & L = w.layers[il];
177+
178+
// K = Wk @ ctx_hidden → [kv_dim, n_tokens] → [head_dim, n_kv, n_tokens]
179+
ggml_tensor * Kb = ggml_mul_mat(ctx, L.wk, ctx_hidden);
180+
Kb = ggml_reshape_3d(ctx, Kb, head_dim, n_kv, n_tokens);
181+
Kb = ggml_rms_norm(ctx, Kb, eps);
182+
Kb = ggml_mul(ctx, Kb, L.k_norm);
183+
Kb = draft_rope(ctx, Kb, positions, head_dim, rope_base);
184+
185+
// V = Wv @ ctx_hidden → [kv_dim, n_tokens] → [head_dim, n_kv, n_tokens]
186+
ggml_tensor * Vb = ggml_mul_mat(ctx, L.wv, ctx_hidden);
187+
Vb = ggml_reshape_3d(ctx, Vb, head_dim, n_kv, n_tokens);
188+
189+
// Write K into cache.draft_k[il] at offset cache.draft_kv_pos
190+
ggml_tensor * k_dst = ggml_view_3d(ctx, cache.draft_k[il],
191+
head_dim, n_kv, n_tokens,
192+
cache.draft_k[il]->nb[1], cache.draft_k[il]->nb[2],
193+
(size_t)cache.draft_kv_pos * cache.draft_k[il]->nb[2]);
194+
ggml_build_forward_expand(gf, ggml_cpy(ctx, Kb, k_dst));
195+
196+
// Write V into cache.draft_v[il] at offset cache.draft_kv_pos
197+
ggml_tensor * v_dst = ggml_view_3d(ctx, cache.draft_v[il],
198+
head_dim, n_kv, n_tokens,
199+
cache.draft_v[il]->nb[1], cache.draft_v[il]->nb[2],
200+
(size_t)cache.draft_kv_pos * cache.draft_v[il]->nb[2]);
201+
ggml_build_forward_expand(gf, ggml_cpy(ctx, Vb, v_dst));
202+
}
203+
204+
return nullptr;
205+
}
206+
207+
// build_gemma4_draft_graph — KV-cached draft forward.
208+
//
209+
// Attends over the full draft KV cache (context K/V already materialised by
210+
// build_draft_kv_prefill_graph, plus newly written block K/V) and returns
211+
// logits for the n_tokens block positions.
212+
//
213+
// draft_embed [n_embd, n_tokens] f32 (MASK token embeddings)
214+
// positions [n_tokens] i32 (absolute token positions)
215+
// attn_mask [kv_pad, q_pad] f16 (causal over context+block)
216+
// kv_start context length before this block (= cache.draft_kv_pos)
217+
//
218+
// Returns logits [n_vocab, n_tokens] f32 (softcapped).
219+
ggml_tensor * build_gemma4_draft_graph(
220+
ggml_context * ctx,
221+
ggml_cgraph * gf,
222+
const GemmaDraftWeights & w,
223+
GemmaTargetCache & cache,
224+
ggml_tensor * draft_embed,
225+
ggml_tensor * positions,
226+
ggml_tensor * attn_mask,
227+
int n_tokens,
228+
int kv_start)
229+
{
230+
// Validate KV cache write range before any graph nodes touch it.
231+
if (kv_start < 0 || kv_start + n_tokens > cache.draft_kv_cap) {
232+
GGML_ABORT("draft KV write out of bounds: kv_start=%d n_tokens=%d cap=%d",
233+
kv_start, n_tokens, cache.draft_kv_cap);
234+
}
235+
236+
const int n_head = w.n_head;
237+
const int n_kv = w.n_head_kv;
238+
const int head_dim = w.head_dim;
239+
const float eps = GEMMA4_RMS_EPS;
240+
const float rope_base = w.rope_theta;
241+
const int kv_len = kv_start + n_tokens;
242+
243+
// Gemma4 scales embeddings by sqrt(hidden_size) — the draft shares the
244+
// target's tok_embd, so it must apply the same scaling. Reference:
245+
// vLLM qwen3_dflash.py embed_normalizer = target_config.hidden_size**0.5
246+
ggml_tensor * hidden = ggml_scale(ctx, draft_embed, std::sqrt((float)w.n_embd));
247+
ggml_set_name(hidden, "gemma4_draft_scaled_embed");
248+
249+
// ── 2. Transformer layers ─────────────────────────────────────────
250+
for (int il = 0; il < w.n_layer; il++) {
251+
const GemmaDraftLayer & L = w.layers[il];
252+
253+
// ── 2a. Attention pre-norm
254+
ggml_tensor * cur = ggml_rms_norm(ctx, hidden, eps);
255+
cur = ggml_mul(ctx, cur, L.attn_norm);
256+
257+
// ── 2b. Q / K / V projections from block hidden state
258+
ggml_tensor * Q = ggml_mul_mat(ctx, L.wq, cur); // [q_dim, n_tokens]
259+
ggml_tensor * Kb = ggml_mul_mat(ctx, L.wk, cur); // [kv_dim, n_tokens]
260+
ggml_tensor * Vb = ggml_mul_mat(ctx, L.wv, cur); // [kv_dim, n_tokens]
261+
262+
// ── 2c. Reshape + per-head RMSNorm for Q and block K
263+
Q = ggml_reshape_3d(ctx, Q, head_dim, n_head, n_tokens);
264+
Q = ggml_rms_norm(ctx, Q, eps);
265+
Q = ggml_mul(ctx, Q, L.q_norm);
266+
267+
Kb = ggml_reshape_3d(ctx, Kb, head_dim, n_kv, n_tokens);
268+
Kb = ggml_rms_norm(ctx, Kb, eps);
269+
Kb = ggml_mul(ctx, Kb, L.k_norm);
270+
271+
Vb = ggml_reshape_3d(ctx, Vb, head_dim, n_kv, n_tokens);
272+
273+
// ── 2d. RoPE on Q and block K
274+
Q = draft_rope(ctx, Q, positions, head_dim, rope_base);
275+
Kb = draft_rope(ctx, Kb, positions, head_dim, rope_base);
276+
277+
// ── 2e. Write block K / V into draft KV cache at [kv_start..kv_start+n_tokens)
278+
ggml_tensor * k_dst = ggml_view_3d(ctx, cache.draft_k[il],
279+
head_dim, n_kv, n_tokens,
280+
cache.draft_k[il]->nb[1], cache.draft_k[il]->nb[2],
281+
(size_t)kv_start * cache.draft_k[il]->nb[2]);
282+
ggml_build_forward_expand(gf, ggml_cpy(ctx, Kb, k_dst));
283+
284+
ggml_tensor * v_dst = ggml_view_3d(ctx, cache.draft_v[il],
285+
head_dim, n_kv, n_tokens,
286+
cache.draft_v[il]->nb[1], cache.draft_v[il]->nb[2],
287+
(size_t)kv_start * cache.draft_v[il]->nb[2]);
288+
ggml_build_forward_expand(gf, ggml_cpy(ctx, Vb, v_dst));
289+
290+
// ── 2f. Full K / V view (context + block) from draft KV cache
291+
// Optional SWA truncation: when enabled and this is an SWA layer
292+
// with kv_len exceeding sliding_window, restrict K/V (and the mask)
293+
// to the last (sliding_window + n_tokens) slots. Matches the draft
294+
// model's training-time SWA pattern.
295+
const bool layer_is_swa = (il < (int)w.layer_is_swa.size())
296+
? w.layer_is_swa[il] : false;
297+
const bool use_swa_trunc = draft_swa_trunc_enabled()
298+
&& layer_is_swa
299+
&& w.sliding_window > 0
300+
&& kv_len > (w.sliding_window + n_tokens);
301+
const int eff_kv_len = use_swa_trunc
302+
? (w.sliding_window + n_tokens)
303+
: kv_len;
304+
const int kv_offset = kv_len - eff_kv_len; // 0 if no truncation
305+
306+
ggml_tensor * K_full = ggml_view_3d(ctx, cache.draft_k[il],
307+
head_dim, n_kv, eff_kv_len,
308+
cache.draft_k[il]->nb[1], cache.draft_k[il]->nb[2],
309+
(size_t)kv_offset * cache.draft_k[il]->nb[2]);
310+
ggml_tensor * V_full = ggml_view_3d(ctx, cache.draft_v[il],
311+
head_dim, n_kv, eff_kv_len,
312+
cache.draft_v[il]->nb[1], cache.draft_v[il]->nb[2],
313+
(size_t)kv_offset * cache.draft_v[il]->nb[2]);
314+
315+
// ── 2g. Permute into flash_attn_ext layout
316+
// Q: [head_dim, n_tokens, n_head, 1]
317+
// K_full: [head_dim, eff_kv_len, n_head_kv, 1]
318+
// V_full: [head_dim, eff_kv_len, n_head_kv, 1]
319+
Q = ggml_cont(ctx, ggml_permute(ctx, Q, 0, 2, 1, 3));
320+
K_full = ggml_cont(ctx, ggml_permute(ctx, K_full, 0, 2, 1, 3));
321+
V_full = ggml_cont(ctx, ggml_permute(ctx, V_full, 0, 2, 1, 3));
322+
323+
// SWA-truncated mask view: take the last eff_kv_len rows along the
324+
// kv axis (axis 0). Mask shape is [kv_pad, q_pad] with kv_pad >= kv_len,
325+
// so the slice [kv_offset .. kv_offset+eff_kv_len) gives the same
326+
// causal pattern for the surviving K positions.
327+
ggml_tensor * eff_mask = attn_mask;
328+
if (use_swa_trunc && kv_offset > 0) {
329+
// ggml_view_2d would produce a non-contiguous tensor (row stride is
330+
// unchanged at kv_pad * elt). FA requires contiguous mask, so we
331+
// copy the slice into a fresh tensor.
332+
ggml_tensor * mask_view = ggml_view_2d(ctx, attn_mask,
333+
eff_kv_len, attn_mask->ne[1],
334+
attn_mask->nb[1],
335+
(size_t)kv_offset * ggml_element_size(attn_mask));
336+
eff_mask = ggml_cont(ctx, mask_view);
337+
}
338+
339+
// ── 2h. Flash attention over full context+block KV
340+
// scale = 1 / sqrt(head_dim); no logit softcap at attention level
341+
const float scale = 1.0f / std::sqrt((float)head_dim);
342+
ggml_tensor * attn = ggml_flash_attn_ext(ctx, Q, K_full, V_full, eff_mask,
343+
scale, /*max_bias=*/0.0f,
344+
/*logit_softcap=*/0.0f);
345+
// attn: [head_dim, n_head, n_tokens, 1]
346+
attn = ggml_reshape_2d(ctx, attn, head_dim * n_head, n_tokens);
347+
348+
// ── 2i. Output projection + residual
349+
ggml_tensor * attn_out = ggml_mul_mat(ctx, L.wo, attn);
350+
hidden = ggml_add(ctx, hidden, attn_out);
351+
352+
// ── 2j. FFN pre-norm
353+
ggml_tensor * hf = ggml_rms_norm(ctx, hidden, eps);
354+
hf = ggml_mul(ctx, hf, L.ffn_norm);
355+
356+
// ── 2k. SwiGLU FFN: down(silu(gate(x)) * up(x))
357+
ggml_tensor * g = ggml_mul_mat(ctx, L.w_gate, hf);
358+
g = ggml_silu(ctx, g);
359+
ggml_tensor * u = ggml_mul_mat(ctx, L.w_up, hf);
360+
ggml_tensor * gu = ggml_mul(ctx, g, u);
361+
ggml_tensor * ffn_out = ggml_mul_mat(ctx, L.w_down, gu);
362+
363+
hidden = ggml_add(ctx, hidden, ffn_out);
364+
}
365+
366+
// ── 3. Final output norm
367+
ggml_tensor * out = ggml_rms_norm(ctx, hidden, eps);
368+
out = ggml_mul(ctx, out, w.out_norm);
369+
ggml_set_name(out, "gemma4_draft_hidden_out");
370+
371+
// ── 4. LM head (tied: transpose of tok_embd)
372+
// tok_embd: [draft_hidden, n_vocab] ggml ne[0]=draft_hidden, ne[1]=n_vocab
373+
// out: [draft_hidden, n_tokens]
374+
// logits: [n_vocab, n_tokens]
375+
ggml_tensor * logits = ggml_mul_mat(ctx, w.tok_embd, out);
376+
ggml_set_name(logits, "gemma4_draft_logits_pre_cap");
377+
378+
// ── 5. Logit softcapping: logits = cap * tanh(logits / cap)
379+
const float cap = w.logit_softcap;
380+
logits = ggml_scale(ctx, logits, 1.0f / cap);
381+
logits = ggml_tanh(ctx, logits);
382+
logits = ggml_scale(ctx, logits, cap);
383+
ggml_set_name(logits, "gemma4_draft_logits");
384+
385+
return logits;
386+
}
387+
60388
// ─── Safetensors loader ───────────────────────────────────────────────────
61389

62390
namespace {

0 commit comments

Comments
 (0)