Conversation
5 tasks
Jiminator
force-pushed
the
refactor/dflash2-selector-tp-shape
branch
from
August 19, 2026 23:00
3bccc89 to
3ef310f
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>
compute_candidates duplicates the projection and the radix top-k across its TP=1 and sharded branches, which differ only by the crop width and the global id offset. Resolve those two up front so the projection and the top-k each have a single call site. Behavior is unchanged on both paths; the existing TP unit test pins the crop, the offset, and the fp32 cast before the all-gather.
Jiminator
force-pushed
the
refactor/dflash2-selector-tp-shape
branch
from
August 19, 2026 23:32
3ef310f to
4b1d1eb
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). Optional cleanup — no behavior change, no user-visible effect; opened separately so it does not ride along with a feature PR.
compute_candidatesduplicates the projection call and the radix top-k across its TP=1 and sharded branches, although the branches differ only by the crop width and the global id offset. Every change to the projection so far has had to be applied to both call sites.Modifications
models/dflash.py: resolvenum_org/org_vocab_startonce for both TP modes, then project and top-k once. The TP=1 early return moves after the top-k.Accuracy Tests
test_dflash_logits.py: 6/6 pass unchanged — the TP test from #35496 pins the per-shard crop, the global id offset, and the fp32 cast before the all-gather, which is what this restructure moves.Speed Tests and Profiling
None; identical kernels in identical order on both paths.
Checklist
CI States
Latest PR Test (Base): ❌ Run #32313662431
Latest PR Test (Extra): ❌ Run #32313662219