Conversation
5 tasks
Jiminator
force-pushed
the
feat/dflash2-selector-topk-guard
branch
from
August 19, 2026 23:00
d8603b3 to
26a94fa
Compare
DFlash2's candidate selector borrows the target lm_head and matmuls draft hiddens against .weight directly, so a quantized target (ModelOpt NVFP4 packed, compressed-tensors FP8 per-channel) fails the dense-head screen and DFLASH refuses to serve such checkpoints. Route an applicable quantized head through quant_method.apply — the same gate (should_apply_lm_head_quant_method) and kernel the target's own logits path uses. A packed weight cannot be row-sliced to the org vocab like the dense path, so the padded tail is masked to -inf out of the top-k instead: a crop view would either break the contiguity flashinfer's radix top-k CHECK_INPUTs or copy the whole valid vocab to restore it. Both TP paths project through one helper; their control flow is unchanged. The worker's pre-capture screen admits a gate-passing quantized head into the graph-folded selector sampler rather than leaving it on the eager per-round fallback, so quantized and dense heads take the same code path; the kernel it captures is the one the target's own logits path already runs under CUDA graphs. The dense-only screen still guards the non-selector static-matmul sampler. Folded and eager measure equal on the same cell (accept lengths bit-identical), so this is path consistency, not the speedup itself. Co-authored-by: LING ZHI <1747985437lz@gmail.com>
The selector masks a quantized head's padded vocab tail to -inf rather than cropping it, so a top_k wider than the org vocab would select those columns and emit out-of-vocab candidate ids instead of raising. The dense path errors on its own; assert so both paths fail the same way. Unreachable with current checkpoints (top_k is a handful of candidates); this guards a malformed draft config, which nothing else validates. Co-authored-by: LING ZHI <1747985437lz@gmail.com>
Jiminator
force-pushed
the
feat/dflash2-selector-topk-guard
branch
from
August 19, 2026 23:32
26a94fa to
10912a7
Compare
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.
Motivation
Depends on #35496 (its commit is included below; review only the second commit).
#35496 masks a quantized head's padded vocab tail to
-infrather than cropping it, which keeps the logits contiguous but leaves the padded columns in the tensor. Aselector_top_kwider than the org vocab would then select those columns and return out-of-vocab candidate ids silently, where the dense path raises fromtopk. Nothing validatesselector_top_kagainst the vocab today.Unreachable with current checkpoints —
selector_top_kis a handful of candidates against a vocab of tens of thousands — so this is a guard against a malformed draft config, not a fix for an observed failure.Modifications
models/dflash.py: assertnum_org >= top_kin_project_candidate_logits, so both the quantized and dense paths fail the same way. Host-side integer compare: free during graph capture, negligible on the eager path.Accuracy Tests
test_dflash_logits.py: 7/7 pass. The new test drives a padded head withtop_k > org_vocab_sizeand asserts theAssertionError; mutation-verified (removing the assert turns it red).Speed Tests and Profiling
None; the check is an integer comparison outside the kernel path.
Checklist
CI States
Latest PR Test (Base): ❌ Run #32313659722
Latest PR Test (Extra): ❌ Run #32313659600