[Bugfix] Fail fast when hierarchical cache is enabled with FP4 KV cache dtype - #36121
Open
jerry-he wants to merge 1 commit into
Open
[Bugfix] Fail fast when hierarchical cache is enabled with FP4 KV cache dtype#36121jerry-he wants to merge 1 commit into
jerry-he wants to merge 1 commit into
Conversation
…he dtype The HiCache L2 write-back/load-back path only transfers the packed KV data buffers and never the FP4 block scale buffers, so reloaded pages silently lose their scales; the host pool also sizes tokens with a full-width dtype formula. Reject the combination at server args validation until full FP4 HiCache support lands.
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.
Motivation
Enabling the hierarchical cache (HiCache) together with an FP4 KV cache dtype (
--kv-cache-dtype nvfp4/fp4_mx_block16) currently passes server args validation but silently produces corrupted outputs after write-back/load-back. Two concrete gaps on current main:python/sglang/srt/mem_cache/pool_host/mha.pyoperate exclusively on the packed data buffers (k_data_ptrs/v_data_ptrs);k_scale_buffer/v_scale_bufferdo not appear in any host transfer set. After a page is written back and later reloaded, its FP4 block scales are lost, so attention consumes packed FP4 data with missing/garbage scales — silent corruption, no error anywhere.MHATokenToKVPoolHost.get_size_per_token()(pool_host/mha.py, thehead_dim * head_num * layer_num * dtype.itemsize * 2formula) does not account for FP4's half-width packed data plus the separate scale buffers, so the host pool size is also wrong for FP4.There is no existing guard rejecting this combination (checked on v0.5.17, v0.5.18, and main).
Modifications
Add a minimal fail-fast guard in
ServerArgs._handle_hicache()(python/sglang/srt/server_args.py): when any of the paths that build a hierarchical-cache host pool is active (--enable-hierarchical-cache, decode KV-cache offload, or decodehost_poolretraction backup) and--kv-cache-dtypeisnvfp4orfp4_mx_block16, raise aValueErrorexplaining that the L2 write-back/load-back path does not transfer FP4 scale buffers and that the host pool sizes tokens with a full-width formula, directing users to pick either HiCache or an FP4 KV dtype.Also add a unit test (
TestHiCacheArgs.test_hicache_rejects_fp4_kv_cache_dtype) covering both FP4 dtypes.Why fail-fast instead of a fix
Full FP4 HiCache support requires host-side scale-buffer transfer in every write-back/load-back path, an FP4-aware host-pool capacity formula, and dedicated correctness validation — each of which is meaningful follow-up work on top of the FP4 KV cache roadmap (#29913). Until then, rejecting the combination at startup is the safe behavior; this mirrors the fail-fast precedent for unsupported NVFP4 KV cache platforms in vLLM (vLLM PR #43669). Related area: #36010.
Verification
python -m py_compileon both changed files.pytest test/registered/unit/server_args/test_server_args.py: 150 passed; the 2 failures present (TestFa4PageSizeAutoForce::test_explicit_prefill_fa4_forces_page_size_128,TestGrpcServerArgs::test_start_server_call_site_matches_native_signature) also fail on unmodified upstreammainin this CPU-only environment, i.e. pre-existing and unrelated.TestHiCacheArgstests: all pass.F401,F821,UP037), black, and isort clean on the changed files.CI States
Latest PR Test (Base): ❌ Run #32680970506
Latest PR Test (Extra): ❌ Run #32680970320
Latest PR Test (AMD ROCm 7.2): ❌ Run #32680970426