Repository navigation
Escha-W2 dense + OSCAR INT2 KV + HiCache on stock sglang 0.5.19 - #1
Open
Nakanokensetsu wants to merge 8 commits into
Open
Nakanokensetsu wants to merge 8 commits into
Nakanokensetsu wants to merge 8 commits into
Conversation
Port of the OSCAR 2-bit KV cache (FutureMLS-Lab/OSCAR, upstream PR sgl-project#32129) onto stock v0.5.19, plus the Escha-W2 quantization layers it needs. Packed codes are [slots, head_num, head_dim//4] uint8; the per-group (scale, zero) pairs live in a separate k_scales_zeros/v_scales_zeros tensor, which matters for anything that moves a KV slot. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Cherry-picked behaviour from upstream sgl-project#38191 (zero-length insert in cache_finished_req), sgl-project#37943 and sgl-project#39526. sgl-project#39342 (--enable-mixed-chunk with extra_buffer) is covered by the sgl-project#39526 backport. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
UnifiedInt2HPKVPool keeps a high-precision prefix/recent window next to the quant arena. It is NOT reached by the hybrid GDN model this fork targets -- that path always builds MHATokenToKVPool(dtype="int2") -- so the code is carried but never constructed. Kept for the hybrid wiring work. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The stock MHATokenToKVPoolHost treats one token as head_num * head_dim * itemsize contiguous bytes. INT2 breaks that twice: packed codes are only head_dim//4 bytes per head, and the scales/zeros live in a separate tensor that must travel with them, or a restored slot is read with the previous occupant's quantization parameters and corrupts silently. The host body now mirrors the packed codes only and the scales/zeros get their own host buffer, so the host token stride, element_dim and the device row are all the same width. _assert_transfer_geometry() checks the three at startup: a mismatch otherwise surfaces as "CUDA illegal memory access" on the first long request, which is expensive to diagnose. element_dim moved behind an overridable _compute_element_dim(), because the parent derives it from the device head_dim, which is not a byte count for a quantized pool. L3 storage backends and MTP draft pools raise NotImplementedError: the flat data page carries the packed codes only. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The mamba backup kernel stores to the host pool through its host pointer, which requires that memory to be mapped into the device address space. On WSL2 that does not hold for malloc'd memory registered with cudaHostRegister -- with any flag combination -- and the kernel faults. Measured with the production shapes (tools/mamba_transfer_repro.py): dst on device -> ok dst torch pin_memory -> ok (cudaHostAlloc) dst cudaHostRegister -> illegal memory access (flags 0/1/2/3) dst pageable -> illegal memory access SGLANG_HICACHE_HOST_ALLOC=pin|register overrides the detection. Also fixes the error path: cudaGetErrorString() only accepts torch._C._cudart.cudaError, so passing the raw int raised TypeError and swallowed the real reason a registration failed. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…rable Backport of upstream PR sgl-project#38157 (issue sgl-project#38156): ranks reach the sizing guard at different times, so a rank arriving after its peers have pinned their pools sees the reduced "available" while still dividing by the full rank count. The peers are charged twice and the guard rejects pools that fit. A MIN all-reduce over the TP cpu_group is also a barrier, so every rank reads before any of them allocates. SGLANG_HICACHE_HOST_RESERVE_GB overrides the fixed 10 GiB reserve, which is sized for machines much larger than a 32 GB host: the pools are built while the weights are still loading, i.e. at the low point of MemAvailable. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The kernel only sees cached layer base pointers and two index tensors, so a stale pointer, an out-of-range index or a tensor that is not tightly packed all surface as an asynchronous "illegal memory access" at whatever synchronises next -- which is usually unrelated code. Check them on the first call of each transfer shape and log the real geometry. SGLANG_DEBUG_MAMBA_TRANSFER=0 disables it. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
What runs in production, why each number is what it is, and how it fails: the KV pool (not the mamba cache) is the long-context ceiling, HiCache needs cudaHostAlloc under WSL2, --hicache-ratio 2 pins 10.3 GB and makes a 32 GB host swap, layer_first is rejected by MambaPoolHost, L3 cannot carry INT2. Includes the verification tools, two of which need no GPU. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
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.
Stock v0.5.19 → the configuration this fork actually runs in production:
dense Qwen3.8-27B-Escha-W2 with an OSCAR INT2 KV cache and HiCache (L2,
host RAM) on 2× 12 GB consumer GPUs, 128K context, 210,000-token KV pool.
The base branch
escha-0519-baseis the unmodifiedv0.5.19tag, so thisdiff is exactly what the fork adds. The seven commits are separable on
purpose — the four
hicache:ones are independent of the INT2 work and twoof them are not specific to this fork at all.
Commits
escha: OSCAR INT2 KV cacheescha: backport upstream cache fixesescha: mixed HP+INT2 KV windowhicache: host KV pool for the INT2 device poolhicache: use cudaHostAlloc where cudaHostRegister is not device-mappablecudaHostRegistered memory, with any flaghicache: read host memory once per TP grouphicache: verify the mamba transfer geometryillegal memory accessinto a named assertionWhy the two HiCache fixes exist
Enabling
--enable-hierarchical-cacheon this stack died withCUDA error: an illegal memory access was encounteredon the first longrequest. Two independent causes:
The INT2 host pool disagreed with itself about how wide a token is. The
host stride,
element_dimand the device row were 192 B, 512 B and128 B.
_assert_transfer_geometry()now refuses to start on a mismatch.The mamba (GDN state) backup kernel stores to the host pool through its
host pointer. On WSL2 that memory is not device-mappable when it comes
from
cudaHostRegister. Isolated withtools/mamba_transfer_repro.py:The second one is the reason HiCache looked like it "silently breaks hybrid
GDN models" here; it is unrelated to sgl-project#39830.
Measured (2× RTX 5060 Ti 12 GB, TP2, 32 GB host, WSL2)
HiCache on (
--hicache-ratio 1) vs off, same build:--hicache-ratio 2pins 10.3 GB of host memory here (KV 2.58 GB + mamba2.55 GB, per rank, ×2 ranks). Pinned pages cannot be paged out, so
everything else is pushed to swap: page cache collapses to 589 MB and the
machine swaps continuously while idle. Ratio 1 pins 5.2 GB and is fine.
--hicache-mem-layout layer_firstis rejected byMambaPoolHost; usepage_first. L3 storage backends raiseNotImplementedErrorfor INT2because the flat data page carries the packed codes without their
scales/zeros.
Not included
Code comments are in Japanese in the parts written for this fork.
CI States
Latest PR Test (Base): ❌ Run #35548496035
Latest PR Test (Extra): ❌ Run #35548495844
Latest PR Test (AMD ROCm 7.2): ❌ Run #35548496059