Conversation
ormandj
requested review from
ByronHsu,
Duyi-Wang,
HaiShaw,
ShangmingCai,
Ying1123,
alphabetc1,
hanming-lu,
hnyls2002,
huangtingwei9988,
hzh0425,
ispobock,
merrymercy,
sogalin,
xiezhq-hermann and
yizhang2077
as code owners
September 2, 2026 04:37
This was referenced Sep 2, 2026
3 tasks
ormandj
changed the title
ormandj
changed the title
This was referenced Sep 6, 2026
ormandj
force-pushed
the
pr/hicache-scale-bytes
branch
from
September 8, 2026 02:45
c9853eb to
d55fc5f
Compare
Contributor
Author
|
The adapted row-geometry dependency in #38212 now keeps this constructor guard and selects separate sidecars before allocation for a single MLA/MTP draft whose stored row width or dtype differs from the target. Matching layouts retain packed transfers. Mismatched multiple runners and FP4 MLA KV storage remain explicitly unsupported; sidecars add host memory beyond the target pool budget. The updated #38212 body includes author CPU tests on |
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.
ormandj
force-pushed
the
pr/hicache-scale-bytes
branch
from
September 8, 2026 21:43
d55fc5f to
a24b8c4
Compare
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
1 participant
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
HiCache can copy packed DSA cache rows using the wrong byte stride when the hybrid Mamba host pool is narrower than its device pool. Backup or restore then addresses different bytes on the two sides, risking incorrect cache contents or a GPU memory fault. A nominal 576-byte row does not hold the packed 656-byte DSA row, which also contains dequantization scales.
This PR makes the affected host-pool constructors use the device's actual packed width and reject incompatible target/draft geometry before allocation. The hybrid packed-DSA path is introduced by #36507; it was not reachable at this PR's original main base. The decode-offload constructor is corrected too.
This is separate from missing index-buffer transfers: #38212 (the main-based successor to #38161) includes an adapted version of this row-width fix and additionally restores the DSA indexes needed by hybrid HiCache loadback.
Modifications
Accuracy Tests
Author CPU validation at
dcb154c1f5: 36 tests passed, with CUDA hidden. Its Python runtime and test trees are unchanged at refreshed heada24b8c4a74on mainafe90a8bc9. Earlier GPU and serving results retain their stated source scope.The CPU reproducer builds a small packed device-pool fixture and checks that the host buffer width and transfer stride match it. Unpatched main produces
12 != 16. Additional cases require incompatible target/draft widths or dtypes to fail before allocation and check unchanged nominal/MHA behavior.Author-reported results: 17 assembler/geometry tests and 10 DCP host-pool tests passed. These tests construct pools; they do not execute GPU transfers. An author-observed downstream GLM-5.3-Flash TP2 setup stopped crashing under concurrent streams and forced restores after its width correction, but those logs are not attached to this PR.
Speed Tests and Profiling
No speed benchmark. For 576-byte nominal versus 656-byte packed rows, fixed-token transfers and host storage require about 13.9% more bytes; a fixed-size host pool holds about 12.2% fewer tokens.
Checklist
Developed with AI assistance.
CI States
Latest PR Test (Base): ❌ Run #34282197312
Latest PR Test (Extra): ❌ Run #34282197060
Latest PR Test (AMD ROCm 7.2): ❌ Run #34282197226