Description
Every tensor-parallel prefill and decode fails on Qwen3.5-4B with:
select_batch: logits vocab 248077 != scratch vocab 248320
248077 is the tokenizer's decodable vocabulary; 248320 is the checkpoint's padded vocab_size.
#676 restricted token selection to the decodable vocabulary: Config35::selection_vocab is narrowed at load time when the tokenizer stops short of the checkpoint width, and every logits arena is sized from it — prefill, the decode buffers, and the unified forward path, along with the SampleScratch those paths build. The tensor-parallel executor was not part of that change; it still builds its SampleScratch from vocab_size. That line predates #676 (it arrived with the TP phase 1 work in #627), and tp_executor.rs does not appear in #676's changed files.
select_batch requires the scratch width to equal the logits width. The check is deliberate — a mismatch would otherwise be a silent kernel out-of-bounds rather than a clean error — so it fails closed and takes every TP request with it.
Scope of the breakage: any checkpoint whose tokenizer defines fewer ids than the checkpoint pads to, on the tensor-parallel path only. Single-GPU serving is unaffected because it sizes both sides from selection_vocab. Qwen3 and Kimi have no equivalent narrowing and are unaffected.
Why CI is green: the four tests covering this path — the two hf_golden_gate TP2 logits gates, e2e_scheduler TP2, and serving_tp2 — are all #[ignore], since they need two CUDA devices, NCCL, and weights.
The fix is to size the TP scratch from selection_vocab like every other call site. Verification is a red/green differential on two GPUs: the current tree reproduces the error, the fixed tree passes all four tests.
Description
Every tensor-parallel prefill and decode fails on Qwen3.5-4B with:
248077 is the tokenizer's decodable vocabulary; 248320 is the checkpoint's padded
vocab_size.#676 restricted token selection to the decodable vocabulary:
Config35::selection_vocabis narrowed at load time when the tokenizer stops short of the checkpoint width, and every logits arena is sized from it — prefill, the decode buffers, and the unified forward path, along with theSampleScratchthose paths build. The tensor-parallel executor was not part of that change; it still builds itsSampleScratchfromvocab_size. That line predates #676 (it arrived with the TP phase 1 work in #627), andtp_executor.rsdoes not appear in #676's changed files.select_batchrequires the scratch width to equal the logits width. The check is deliberate — a mismatch would otherwise be a silent kernel out-of-bounds rather than a clean error — so it fails closed and takes every TP request with it.Scope of the breakage: any checkpoint whose tokenizer defines fewer ids than the checkpoint pads to, on the tensor-parallel path only. Single-GPU serving is unaffected because it sizes both sides from
selection_vocab. Qwen3 and Kimi have no equivalent narrowing and are unaffected.Why CI is green: the four tests covering this path — the two
hf_golden_gateTP2 logits gates,e2e_schedulerTP2, andserving_tp2— are all#[ignore], since they need two CUDA devices, NCCL, and weights.The fix is to size the TP scratch from
selection_vocablike every other call site. Verification is a red/green differential on two GPUs: the current tree reproduces the error, the fixed tree passes all four tests.