Conversation
build_hybrid_mamba_stack built the MLA host pool through build_kv_host_pool without override_kv_cache_dim, so MLATokenToKVPoolHost sized host rows as kv_lora_rank + qk_rope_head_dim. A DSA pool that stores packed fp8 rows has a wider kv_cache_dim, and every HiCache transfer path assumes host rows have the device pool's row geometry (the kernel paths use the host stride for both sides, the direct path copies row slices). Pass the device pool's kv_cache_dim when use_mla, as the DSA strategy and the HiRadixCache DSA entrypoint already do through build_anchor_sidecar_stack / build_kv_only_group. The override is gated on use_mla because MHA host pools take no such keyword. DecodeKVCacheOffloadManager had the same gap (its MLATokenToKVPool isinstance gate admits DSATokenToKVPool); pass the override there too. Make MLATokenToKVPoolHost check its row geometry at construction, before allocating host memory: the host row width must equal device_pool.kv_cache_dim (a narrower host row is exactly this defect, so a caller that omits the override for a packed pool now fails clearly), and every packed MTP draft pool must have the device pool's store_dtype and kv_cache_dim (reachable with --speculative-draft-kv-cache-dtype), instead of allocating and then transferring with mismatched rows. MLATokenToKVPool always sets kv_cache_dim; the lookup falls back to the host width only for duck-typed device pools that do not model row packing. On main a packed DSA pool does not reach build_hybrid_mamba_stack; HybridLinearKVPool with use_mla builds a plain MLATokenToKVPool. It does with the GLM-5.3-Flash series (sgl-project#36507), which makes a DSATokenToKVPool the hybrid pool's full_kv_pool. Add CPU unit tests: host row width follows a packed device pool, a non-MLA pool exposing kv_cache_dim gets no override, a draft pool with the target's geometry is packed, draft pools with a different dtype or (same dtype) a different width are rejected, a packed device pool without the override is rejected by the constructor (alone and with a same-dtype nominal-width draft), a packed target with the override and a matching draft allocates, and a nominal-width pool needs no override.
github-actions
Bot
added
hicache
ormandj
changed the title
|
Hi @ormandj — a heads-up and a big thank-you: your fix resolved a nasty production corruption for us (the exact #38031 symptom set: intermittent One problem: #36507 just landed in main (97c6978) without this fix, and this PR got closed unmerged — so right now main (and any image built from it) is exposed to the corruption again, and there is no successor PR. I verified the diff applies cleanly to current main @ 97c6978 — all 7 files (4 code + 3 tests), zero conflicts: Could you please re-file it against |
|
Refiled as #38212 against |
3 participants
Motivation
GLM-5.3-Flash can resume a cached prompt with incorrect attention state after HiCache moves that prompt from GPU memory to host memory and back. The KV values and recurrent state are restored, but the DSA index buffers used to select which tokens to attend to are not. Attention can therefore select the wrong parts of the prompt even though the cache reports a successful restore. This affects the target model even with speculative decoding disabled.
A second defect lets different prompt suffixes share a compressed index row when a prefix-cache split falls inside a compression group. The change keeps those groups intact so one suffix cannot overwrite another suffix's index state.
#38031 reports incorrect GLM-5.3-Flash output after HiCache loadback, including with MTP disabled. This PR repairs the missing-index and shared-row paths; the reporter's 8x H100 full-model configuration has not been validated with this patch. Stacked on #36507; includes the adapted packed-row dependency from #37534.
Modifications
Accuracy Tests
The GPU regression simulates eviction and loadback without model weights: fill target/draft index buffers with deterministic bytes, back up noncontiguous pages in two transfers, clear device storage, then restore into different pages and compare every byte. Moving the destination pages prevents an unchanged device buffer from hiding a missing transfer. It also checks that the index tier shares KV addresses without allocating or freeing those addresses independently. CPU tests check compressed-prefix alignment and host-pool construction. These are state-transfer regressions, not a full-model answer-quality reproducer.
Author run at
66f6da3a21on RTX PRO 6000 Blackwell (SM120), Transformers 5.12.1 / Tokenizers 0.22.2: 20/20 transfer cases across four layout/backend pairs, 100/100 mocked controller cycles and ratio-mode assembly passed; Compute Sanitizer reportedERROR SUMMARY: 0 errors. The two hybrid-pool CPU test files passed 30 tests and six subtests.Speed Tests and Profiling
No isolated speed benchmark.
Checklist
Developed with AI assistance.
CI States
Latest PR Test (Base): ❌ Run #34000673291
Latest PR Test (Extra): ❌ Run #34000673170
Latest PR Test (AMD ROCm 7.2): ❌ Run #34000673275