qwen35: DSpark speculative decoding (Qwen3.8-27B drafters) - #625
Draft
davide221 wants to merge 11 commits into
Draft
qwen35: DSpark speculative decoding (Qwen3.8-27B drafters)#625davide221 wants to merge 11 commits into
davide221 wants to merge 11 commits into
Conversation
Wire the DSpark drafter heads (low-rank Markov bigram correction + confidence head) into the qwen35 spec-decode loop, so Qwen3.8-27B DSpark drafters (e.g. RadixArk/Qwen3.8-27B-DSpark) run with full head support: - spec loop: markov-corrected greedy chain (fused single-graph variant with non-fused fallback) replaces plain argmax projection when the drafter ships DSpark heads; DDTree candidate top-k gets the markov bias too. Env-gated: DFLASH_QWEN35_DSPARK, DFLASH_QWEN35_FUSED_DSPARK, DFLASH_QWEN35_DSPARK_TREE (all default on). - target capture layers now follow the drafter GGUF's dflash.target_layer_ids instead of the evenly-spaced derivation; the Qwen3.8 drafter is trained on layers 4/16/28/40/52, not 1/16/31/46/61. - draft loader: dflash.mask_token_id from the drafter GGUF wins over the family default (Qwen3.8 drafter uses 248077, default was 248070), and optional YaRN rope scaling keys are parsed into DraftWeights. - draft graph: rope calls honor the drafter's YaRN config (previously hardcoded plain NEOX rope). - Qwen35DFlashTarget exposes lm_head for the fused head path. - convert_dflash_to_gguf.py: handle single-file DSpark releases (markov/ confidence heads inline in model.safetensors), transformers>=5 nested rope_parameters and dflash_config.mask_token_id, and emit YaRN scaling metadata. The confidence-gate adaptive block length is not wired yet (q_len sizes the per-request step buffers); the chain runs with the gate off.
Verify/accept now run over v_len (the drafted chain's actual length) instead of the buffer-sizing q_len, so the DSpark confidence gate's adaptive block truncation is structurally supported. The gate itself stays off by default (DFLASH_QWEN35_DSPARK_CONFIDENCE_THRESHOLD=0): with the RadixArk Qwen3.8 drafter, any threshold in 0.1-0.5 truncates to the same short chain regardless of value, so the confidence scores coming out of the shared head path look mis-scaled and need a separate investigation before the gate can help. threshold=0 is bench-verified regression-free.
- ggml_ssm_conv_step: one kernel for the causal-conv decode/verify step (history window + silu(conv) + in-place history write-back + optional rollback window copy) replacing transpose/concat/ssm_conv/silu/cpy. - ggml_gated_delta_net_set_raw_gates: the GDN kernel applies sigmoid(beta) and softplus(alpha + dt_bias) * A itself. - ADD + RMS_NORM + MUL fusion (residual add materialized alongside the normalized output) in the CUDA/HIP graph evaluator. - legacy pool MAX_BUFFERS 256 -> 1024: LUCE_Q8_MEMO holds ~300 pooled buffers per evaluation; a full pool freed in-flight buffers with cudaFree and produced illegal memory accesses on long prefills.
Rename the RDNA small-tile macro to GGML_CUDA_MMQ_SMALL_TILE and apply it to IQ4_XS/Q4_K/Q5_K/Q6_K/Q8_0 in addition to the ROCmFPX formats. At spec-decode verify widths (N<=16) the 128-row tile leaves a 5120-row projection with only 40 blocks on a 64-CU gfx1201; 64x64/4-warp tiles measured +12-23% on those shapes (verify step 43.8 -> 39.7 ms on Qwen3.8-27B) at ~8% prefill cost.
- loader places attn_gate|attn_qkv and ssm_beta|ssm_alpha back to back and exposes zero-copy stacked aliases (L.wqkv_z, L.ssm_ba): one GEMV each instead of two (DFLASH_QWEN35_NO_STACK=1 disables). - FFN uses ggml_swiglu_split so the backend fuses gate/up/GLU into one vector kernel at decode. - DeltaNet block: single l2_norm over the q|k slab, ggml_ssm_conv_step, raw-gate gated_delta_net (in place, no state copy), no q/k head repeat (the kernel broadcasts). DFLASH_QWEN35_NO_FUSED_KERNELS=1 keeps the op-by-op graph for A/B. - DFLASH_KV_ROTATE=0 skips the FWHT K/Q rotation (precision-neutral with q8_0/f16 caches, two fewer launches per attention layer). Qwen3.8-27B IQ4_XS on R9700: plain decode 30.4 -> 33.8 tok/s with identical greedy output.
- Qwen35AdaptiveSpecPolicy: EMA of accepted draft tokens per step; below 0.8*(spec_step_ratio-1) the loop runs a burst of plain-decode steps (seed-only verify, no drafter/heads/snapshot/rollback, features still captured) and probes again afterwards. Env DFLASH_QWEN35_SPEC_STEP_RATIO (default 1.7, 0 disables) and DFLASH_QWEN35_AR_BURST (default 40). Low-acceptance prose 28.1 -> 32.4 tok/s, code/mixed unchanged. - Confidence gate now uses the fused Markov graph and truncates on the host; DFLASH_QWEN35_DSPARK_CONF_DEBUG=1 prints per-position scores. - spec-profile hooks for the chain path (project/snapshot/verify/ rollback/feature).
launch_fattn was told the vec kernel consumes D keys per step; it walks nthreads (128) per step, so a 256-key window at head_dim 256 ran as one block per head. Passing nthreads lets it use two blocks per head plus the combine pass: Qwen3.8-27B plain decode 34.3 -> 34.6 tok/s on R9700, identical output.
The first spec step after a plain-decode burst updates the acceptance EMA with alpha 0.5 so a stream that became predictable leaves plain decode immediately; step ratio and start value keep the measured best balance (45.7 / 31.8 / 40.4 tok/s code / prose / mixed).
The break-even acceptance now follows live EMAs of the spec-step and plain-step wall times (default 1.9 until both are measured), so it is right for any drafter block size (width-8 DSpark and width-16 DFlash measure ~1.8 on gfx1201).
DFlash 2 (z-lab/inco, e.g. z-lab/Qwen3.8-27B-DFlash2) is the DFlash backbone plus a grouped dynamic causal conv around attention and MLP in every layer and a candidate selector head (top-k lm_head candidates per block position, one path scored by a low-rank bigram form). - converter: maps attention_conv/mlp_conv (base kernels F32, kernel projections) and candidate_selector tensors, emits dflash2.* metadata, reads block_size from dflash_config, emits SWA pattern for drafters with causal sliding layers. - loader: DraftConvWeights per layer, DraftSelectorWeights, shape checks. - draft graph: conv prepare/finish (two taps over the block, per-element base + per-group dynamic coefficient) in both the stateless and the cached-KV builders. - selector chain: top-k via the target's GPU top-k (kMaxK 8 -> 16), one cached graph for hproj + codebook row gathers, host path search. - spec loop uses the selector before the DSpark/argmax paths. Qwen3.8-27B IQ4_XS on R9700, q8_0 drafter, greedy: 109.9 code / 50.7 prose / 111.8 mixed tok/s (DSpark drafter: 45.6 / 32.4 / 38.6); avg 5.9-6.0 accepted tokens per 8-token block on code, ~2.7 on prose.
With the 64-row/4-warp tile the mmq_x=32 instantiation runs at 180 GB/s on gfx1201 (17408x5120 IQ4_XS) against 443 GB/s at mmq_x=16 and 315 at 48, so N=17..32 batches (DDTree budgets, prefill remainders) took 2.4x longer than N=16 or N=40. Choose the next tile instead.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds DSpark and DFlash 2 drafter support to the qwen35 family spec-decode loop, targeting the published Qwen3.8-27B checkpoints (
RadixArk/Qwen3.8-27B-DSpark: DFlash backbone + low-rank Markov bigram head + confidence head;z-lab/Qwen3.8-27B-DFlash2: DFlash backbone + grouped dynamic convs + candidate selector), plus the decode/prefill optimizations listed below.What's in here
Spec loop (
qwen35_backend.cpp)DFLASH_QWEN35_DSPARK,DFLASH_QWEN35_FUSED_DSPARK,DFLASH_QWEN35_DSPARK_TREE.dflash.target_layer_idsinstead of the evenly-spaced derivation. The Qwen3.8 drafter is trained on layers 4/16/28/40/52; capturing the derived 1/16/31/46/61 silently destroys acceptance.Draft loader (
draft_gguf_loader.cpp)dflash.mask_token_idfrom the drafter GGUF wins over the family default (this drafter trains with MASK 248077; the default was 248070 — another silent acceptance killer).rope.scaling.*) parsed intoDraftWeights.Draft graph (
draft_graph.cpp)Converter (
convert_dflash_to_gguf.py)model.safetensors, nodflash_aux_heads.ptsidecar) now convert directly.rope_parameters(theta + YaRN) anddflash_config.mask_token_id; emits YaRN scaling metadata.Fused head path:
Qwen35DFlashTargetexposeslm_head_tensor()(the fused chain needs the target lm_head; layer-split target falls back gracefully).Decode/prefill optimization commits (same branch)
Follow-up commits on top of the DSpark integration, all measured on the same box (R9700 / gfx1201, ROCm 7.2), greedy, 300-token generations, identical greedy output before/after unless noted:
ggml: fused DeltaNet decode kernels for HIP:ggml_ssm_conv_step(conv history + silu + write-back + rollback window in one kernel), raw-gategated_delta_net(sigmoid/softplus/A inside the kernel), residual ADD + RMS_NORM + MUL fusion, legacy poolMAX_BUFFERS256 -> 1024 (a full pool freed in-flightLUCE_Q8_MEMObuffers and crashed 6k-token prefills).ggml: 64x64 MMQ tiles for dense verify widths on RDNA:GGML_CUDA_MMQ_SMALL_TILEfor IQ4_XS/Q4_K/Q5_K/Q6_K/Q8_0 (verify step 43.8 -> 39.7 ms; ~8% prefill cost).qwen35: stacked projections and fused DeltaNet decode graph: zero-copy stackedattn_gate|attn_qkvandssm_beta|ssm_alphaGEMVs,ggml_swiglu_split(backend fuses gate/up/GLU), one l2_norm over q|k, in-place recurrent state, no q/k head repeat,DFLASH_KV_ROTATE=0. Launches per decode token 1929 -> ~1140.qwen35: adaptive speculation policy: EMA of accepted drafts/step, plain-decode bursts (seed-only verify, features still captured) when speculation is below break-even;DFLASH_QWEN35_SPEC_STEP_RATIO(0 disables),DFLASH_QWEN35_AR_BURST.Requantized target used for the numbers below: pure IQ4_XS body, Q5_K output, Q6_K attn_v/ssm_out (
llama-quantize --allow-requantize --pure --output-tensor-type q5_k --tensor-type ssm_out=q6_k --tensor-type attn_v=q6_k ... iq4_xs), 14.84 GiB; perplexity on a 24x1024 sample 3.5634 -> 3.5670 vs bartowski IQ4_XS.Numbers (R9700 / gfx1201, Qwen3.8-27B, greedy, 300-token gens)
Decode, tok/s:
Prefill, tok/s at 123 / 423 / 1623 / 6023 prompt tokens: 884 / 946 / 978 / 920 (hipfire published pp128/512/2048/8192: 772 / 759 / 737 / 663; upstream llama.cpp 944 / 1318 / 1277 / 1123).
Long context (128 new tokens; DFlash 2 q8_0 drafter): summarize-documentation task 102.0 / 75.4 / 60.6 / 55.8 tok/s at ~128 / 512 / 2048 / 8192 prompt tokens, code-continuation task 71.6 / 75.1 / 85.4 at ~128 / 2048 / 8192; plain decode 35.0 / 34.8 / 34.2 / 32.3 (hipfire 36.2 / 36.2 / 36.0 / 34.7, llama.cpp 31.7 / 31.7 / 31.5 / 29.8). Thinking mode (
reasoning_effort: high, max_tokens 6000): 0.77-0.78 acceptance per position, 113.9 tok/s on a GSM8K-style prompt; note that with max_tokens belowhard_limit_reply_budget(4096) the reasoning budget hook force-closes thinking immediately and the reply runs as plain decode.DFlash 2 (
z-lab/Qwen3.8-27B-DFlash2, converted withconvert_dflash_to_gguf.pyand quantized withquantize_dflash_draft.py --scheme q8_0) is the recommended drafter for Qwen3.8 on this branch: 5.9-6.2 accepted tokens per 8-token block on code/math, ~2.7 on prose. The DSpark and 3.6-DFlash drafters keep working; the z-lab Qwen3.6-27B DFlash drafter transfers to 3.8 unchanged (same hidden size, layer count and vocab). Greedy verification keeps the output exact whichever drafter is used.Serving config for the numbers:
--draft <drafter gguf> --fa-window 2048 --cache-type-k q8_0 --cache-type-v q8_0, envDFLASH_SINGLE_CHAIN_CHECKPOINT_F32=1 DFLASH_FAST_ROLLBACK_THRESHOLD=1 LUCE_Q8_MEMO=1 DFLASH_KV_ROTATE=0,GGML_HIP_GRAPHS=ONbuild. Prose acceptance with either drafter is below break-even, so the adaptive policy runs it mostly as plain decode (bounded by the ~34 tok/s plain-decode rate).Known limitations / follow-ups
LUCE_MMVQ_MAX_NCOLS=8) is 2x slower than MMQ on gfx1201.