Skip to content

perf(dflash): 16-byte loads in the draft GEMV + barrier-free draft attention - #678

Closed
fansilas wants to merge 1 commit into
gittensor-ai-lab:mainfrom
fansilas:perf/dflash-gemv-16b-draft-attn
Closed

perf(dflash): 16-byte loads in the draft GEMV + barrier-free draft attention#678
fansilas wants to merge 1 commit into
gittensor-ai-lab:mainfrom
fansilas:perf/dflash-gemv-16b-draft-attn

Conversation

@fansilas

Copy link
Copy Markdown
Contributor

Summary

Two fixes in the DFlash draft block, found by profiling the draft step and reading the SASS.
DFlash 362.49 → 398.91 tok/s (+10.0%), mean accept τ unchanged, one file.

  • The batched GEMV never actually vector-loaded. k_gemv_batched documents a uint4 load as "the kernel's dominant cost", but (const bf16*)&wrow4[k8] only takes an address into global memory — so it compiled to 8 + 8*BATCH = 136 scalar 2-byte loads per lane per chunk where 17 wide loads carry the same bytes. Assigning each chunk to a local uint4 first keeps it in registers: SASS on sm_120 goes 136 LDG.E.U16 → 34 LDG.E.64, kernel 33.46 → 15.29 ms. Same operands in the same per-(row, batch) order, so bit-exact. +5.4% on its own.
  • Draft attention barriered once per key. k_attn reduces every key across the whole block (~9 __syncthreads()) for one FMA of real work per thread. k_attn_warp gives each warp a strided slice of keys with its own online softmax — shuffle reduction, accumulator in registers, one barrier for the whole kernel at the merge; 16 warps/block measured best. Kernel 15.52 → 1.54 ms. Not bit-identical, which nothing needs: the draft only proposes, and the early-exit verify accepts only block[i+1] == p, so emitted tokens stay the target's greedy argmax.

Both behind env flags defaulting onSPARKINFER_DFLASH_GEMV_VEC=0 and SPARKINFER_DFLASH_ATTN_WARP=0 restore the old kernels, so before/after is one binary.

Proof of speedup

  • Tested on RTX 5090 (sm_120)

Decode tok/s — DFlash decode from qwen3_gguf_dflash_bench (the metric pr_dflash_bot.py scores), same box, same binary, interleaved, median of 5 passes at 128 tokens:

decode tok/s
before (main) 362.49
after (this PR) 398.91

Prefill pp tok/s (qwen3_gguf_bench, --ctx 32768, median of 3 interleaved passes). This PR does not touch prefill; measured only to show no regression — the +0.03% is run-to-run noise, not a claimed gain:

prefill pp tok/s
before prefill (main) 27599.63
after prefill (this PR) 27608.20
# DFlash decode — 101-token held-out prompt (seed=fixed), 128 new tokens, one binary
MAIN  362.6820 362.6804 362.3844 362.4890 362.3851   median 362.4890
NEW   399.0286 399.0362 398.8984 398.9084 398.7994   median 398.9084   -> +10.05%
mean accept t = 8.2500 on every run, both sides

# AR decode + prefill, ctx=32768 — unchanged, as expected for a draft-only change
MAIN  decode tg 478.67 tok/s   prefill pp 27599.63 tok/s
NEW   decode tg 478.61 tok/s   prefill pp 27608.20 tok/s

$ bench/scripts/dflash_accuracy.sh
METRIC SPEC_AGREE 32/32 = 1.0000
VERDICT PASS

ctest --test-dir build is 9/9. Draft-only, so the AR accuracy gate is untouched by construction — these kernels run only inside DFlashDraftModel::forward_block, which qwen3_gguf_score never instantiates; re-running bench/scripts/accuracy.sh with them disabled gives byte-identical top-1/KL.

@skyrocket2026 skyrocket2026 added area:runtime subsystem (emission weight 0.26) test-on-5090 Maintainer-approved to evaluate on RTX 5090 (greenlight) labels Jul 31, 2026
@skyrocket2026

Copy link
Copy Markdown
Member

Closed: no activity for 2+ days

This PR has had no updates (commits, comments, reviews, or label changes) for 2 days (threshold: 2 days).

Reopen this PR or open a fresh one when you're ready to continue.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:runtime subsystem (emission weight 0.26) test-on-5090 Maintainer-approved to evaluate on RTX 5090 (greenlight)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants