Pramodith/partial topk#18
Draft
pramodith wants to merge 16 commits into
Draft
Conversation
Instrument test_dflash (DFLASH_TOPK_CALIB=1, chain verify path) to measure how
large a draft-derived candidate shortlist M must be to cover the target's greedy
argmax / top-k / top-p nucleus, for a candidate-restricted LM head
(ggml_get_rows over the 248k-vocab head). Reports set-coverage@M and top-p
nucleus MASS-coverage@M (the quality-relevant metric for sampling) on an M grid
up to 65536.
Findings (5 prompts, 2550 positions, Qwen3.6-27B Q4 target + Q4 3.6 draft):
greedy viable (M~1k -> ~97%); sampling set-coverage poor at small M but nucleus
MASS-coverage ~99.6% at M=4096 (large M is cheap: head ~6% of step traffic).
Adds calib/{prep_prompts.py,run_calib.sh,aggregate.py} + docs writeup for
scale-up (incl. B200 sm_100 build).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Full top-k/top-p LM-head shortlist calibration on Qwen3.6-27B Q4_K_M target + Q4_K_M 3.6 draft, run on a single B200 via calib/calibrate.py over 198 balanced prompts (HumanEval/GSM8K/MATH-500), 102,045 chain-verify positions. Verdict: greedy viable (target argmax in draft top-1024 at 99.46%); sampling exact-set-coverage impractical but mass-coverage faithful (top_p=0.95 nucleus mass 99.8% at M=8192). Also documents the batched-prefill NaN bug and the DFLASH27B_PREFILL_UBATCH=1 workaround used to capture all prompts. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The draft top-K + logsumexp kernel launched only n_positions (~15) blocks,
leaving most of the GPU's SMs idle, and kept its per-thread top-K in a
data-dependent insertion index that the compiler spilled to local memory and
re-read on every vocab element. Both made the ~9 MB vocab scan run at a small
fraction of peak DRAM bandwidth.
Rework into a split-K two-pass design:
- pass 1 (draft_topk_partial) splits each position's vocab scan across many
blocks (2D grid n_positions x split) so all SMs stay busy;
- pass 2 (draft_topk_combine) merges the per-split partials per position.
Template both kernels on K (compile-time) so the top-K stays register-resident
via a branchless unrolled bubble instead of spilling, and read logits as float4
(one coalesced 16-byte transaction per 4 logits) with a scalar fallback when a
row base is not 16-byte aligned (vocab % 4 != 0), preserving any-vocab
correctness. split is auto-tuned (env override DFLASH_TOPK_SPLIT).
Measured on an RTX 3090 (n=15, vocab=151936, K=8):
- GPU kernel time: 392 us -> 36.3 us (30.6 partial + 5.75 combine), 10.8x
- full call (kernel+sync+D2H): 0.407 ms -> 0.053 ms, 7.7x
Full-call speedup is 5.9-8.4x across n in {7,15,31,63}. Output is bit-for-bit
equivalent to the CPU reference (id_mismatches=0) across K in {1,2,4,8} and
both aligned and odd vocab; compute-sanitizer memcheck clean on both paths.
Adds bench_topk.cu, a standalone microbenchmark + CPU-reference correctness
harness (not wired into the build) used to profile and A/B this change.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01T55hNb5cgyCwNYnNAE1hun
# Conflicts: # server/CMakeLists.txt # server/src/qwen35/qwen35_dflash_target.cpp # server/test/test_dflash.cpp
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
1 participant
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.
No description provided.