Repository navigation
Conversation
The vLLM MultiSwitch heads store a -1e9 counting mask and integer adapter ids in the paged KV cache. float16 and fp8 e5m2 turn the mask into -inf, so the one-hot query computes 0 * -inf = NaN and every write address is wrong; fp8 e4m3 rounds adapter ids above 16 to a neighbour. In both cases routing is silently wrong while the base model keeps generating. --dtype float16 fails too: the mask is built in the model dtype, where it either overflows or fills -inf depending on the code path. MultiSwitch.__init__ now calls check_switch_kv_dtype, which accepts only bfloat16 or float32 storage (what the switch has been validated on) and raises a ValueError naming the flag to change. The check lives in a vLLM-free module, _kv_dtype.py, so tests/unit/test_switch_kv_dtype.py runs it on CPU along with the arithmetic it rests on. docs/MULTISWITCH_EXPLAINED.md section 8 listed fp8 as unguarded and gave the memory keys (code(n) * 28) as a cause; codeword entries are +/-1/8, so those keys peak at 3.5 and fit in fp8. The section now gives the measured causes and documents the check; CLAUDE.md gotcha 10 follows. Signed-off-by: samadarsh <samadarsh14@gmail.com>
Bring the fp8/float16 KV-cache guard up to date with main (vLLM 0.26-0.30 range, generative-computing#144 docs cleanup, SWITCH kernels write-up). No conflicts. Signed-off-by: samadarsh <samadarsh14@gmail.com>
samadarsh
requested review from
antonpibm,
aviv1ron1,
freunda and
yairallouche
as code owners
October 28, 2026 20:45
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
This branch has not been deployed
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.
Summary
The vLLM MultiSwitch stores a
-1e9counting mask and integer adapter ids in the paged KV cache. With a float16 or fp8 cache they don't survive, and adapter routing goes silently wrong while the base model keeps generating.docs/MULTISWITCH_EXPLAINED.md§8 listed--kv-cache-dtype fp8as unguarded; this adds the guard.Measured on CPU by rebuilding both heads and round-tripping K/V through each dtype (bf16 model):
-inf,0 * -inf = NaN, every write address wrong--dtype float16-inf, depending on the code pathChanges
vllm/switch/_kv_dtype.py(new, vLLM-free):check_switch_kv_dtypeaccepts only bf16/fp32 storage and raises aValueErrornaming the flag to change.vllm/switch/multi.py:MultiSwitch.__init__calls it before building the attention heads.tests/unit/test_switch_kv_dtype.py(new, CPU): accepted vs rejected configs, plus the arithmetic the guard rests on.docs/MULTISWITCH_EXPLAINED.md§8,CLAUDE.mdgotcha 10: measured causes. The old text blamed the memory keys (code(n) * 28), but codeword entries are ±1/8, so those keys peak at 3.5 and fit in fp8.Judgment call: fp8 e4m3 with ≤16 adapters passes the CPU simulation but is still rejected, because it hasn't been validated under vLLM's fp8 attention backends. Happy to relax it if you'd prefer.
Testing
cache_config.cache_dtypeat construction, with a safe fallback if the attribute is absent; the vLLM GPU suites use the bf16 default, so it shouldn't trip them. The publishedgranite-switch-4.1-*-previewcheckpoints are saved as bf16, so they pass the guard with default settings.