Repository navigation
Use NVFP4 draft vocabulary shards for Qwen3.8 tensor parallelism - #66
Merged
Merged
Conversation
Prepare an NVFP4 proposal-head copy for each tensor-parallel rank and pass it through the existing vocabulary projection and gather. This extends the TP=1 draft-head policy to sharded execution. The target decoder, prompt scoring, and speculative verification retain their BF16 weights. Package metadata is unchanged; packages need re-export to use the quantized draft shards. Validation: 12 focused Qwen MTP and CUDA projection tests pass on GB10 with base a773e68. Compiled replay covers two and four logical shards on one GPU, row counts 1 and 4, shape transitions, untouched output tails, immutable BF16 weights, and allocation stability. A TP=2 package exported with base bcf7177 and this change served on two 128 GB DGX Sparks with MTP and a server concurrency limit of 8. All 14 benchmark and warmup requests completed without errors. Adaptive decode measured 44.67 to 52.38 tokens/s; completion lengths differed. Full-model export and service were not repeated with base a773e68; multi-node TP=4 was not tested. Design-Signed-Off-By: Codex (GPT-6); design-review=6aeb
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.
Qwen3.8 MTP uses an NVFP4 vocabulary head at TP=1, but its tensor-parallel path projects draft logits through BF16 shards. Prepare an NVFP4 copy of each draft vocabulary shard and select it in the existing vocabulary projection and gather.
The target decoder, prompt scoring, and speculative verification retain their original BF16 weights. This change adds no kernel, transport, runtime-policy, or package-schema changes. Existing packages remain usable; re-export is required to use NVFP4 draft shards. Each shard has its own quantization scale, so draft logits need not be identical to TP=1.
Validation
a773e688plus this change:tests/recipes/test_qwen38_mtp.pyandtest_quantized_draft_shards_preserve_logits_under_compiled_replay.nvidia/Qwen3.8-Flash-Next-FP4-Localwas exported and served with MTP on nvspark2/nvspark1, two 128 GB DGX Sparks. Serving concurrency was capped at 8. Across four benchmark cases, all 14 measured and warmup requests completed without errors.Performance
The native TP=2 service comparison used base
bcf7177dwith and without this change, the same native binary, AIPerf concurrency 1, and thespeed_bench_codingdataset.Adaptive cases used three measured requests and one warmup each, with output length 8192. Fixed cases used two measured requests and one warmup each, with output length 1024. Generated completion lengths and adaptive verification widths differed. GPU-round counters include warmup; AIPerf throughput and latency include measured requests only.
Full-model export and service measurements were not repeated with base
a773e688; the focused tests above passed with that base.