fix: MSE_BITS==8 OOB kernel patches + FORCE_SYNC removal for spec-TQ84 - #1
Open
amd-dlimpus wants to merge 4 commits into
Open
fix: MSE_BITS==8 OOB kernel patches + FORCE_SYNC removal for spec-TQ84#1amd-dlimpus wants to merge 4 commits into
amd-dlimpus wants to merge 4 commits into
Conversation
… read _tq_decode_stage1 was missing the 1-byte-per-element fast path for 8-bit MSE keys (spec-TQ84). The generic raw16 path read mse_addrs0+1, which is one byte past the MSE region for d=D-1. Under cudagraph, this OOB read lands on an unmapped page causing a HIP memory fault. Same fix already applied to _tq_full_dequant_kv and _tq_load_k_tile (v3) in commit b876718. This closes the last missed site. Co-Authored-By: Claude Sonnet 4 <noreply@anthropic.com>
The per-dispatch cuda sync was masking an OOB read in _tq_decode_stage1 for MSE_BITS==8. Commit 5692d0f fixed the root cause. The workaround is now dead code and adds ~200us serialization overhead per decode step when enabled. Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
triton_turboquant_decode_v2.py had the same missing 1-byte-per-element path as _tq_decode_stage1 in the v1 kernel. The raw16 path reads mse_addrs0+1, causing OOB for MSE_BITS==8 (spec-TQ84). v2 is opt-in (VLLM_TQ_DECODE_V2=1) but must be safe before upstreaming. Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
Add turboquant_k8v4_nc and turboquant_k8v4_spec_nc to PRESET_EXPECTED. Update key_quant_bits exclusivity assertion to include 8 (MSE 8-bit keys used by spec-TQ84). Co-Authored-By: Claude Sonnet 4 <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
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.
feat: spec-TQ84 — 8-bit Lloyd-Max KV codec with cudagraph fix and OOB kernel patches
Base branch:
feat/spec-tq84Target branch:
feat/spec-tq84-kernel-fixRemote:
origin(aditi-amd/vllm)Summary
Spec-TQ84 is an 8-bit Lloyd-Max quantized KV cache codec for vLLM, implementing the approach from arXiv 2504.19874. It uses 256-centroid Lloyd-Max codebooks with online Hadamard rotation and per-token L2 norm preservation to quantize KV cache keys to 8 bits while maintaining distributional fidelity.
The upstream TQ84 path used a bare FP8 cast with no rotation or codebook optimization, causing distributional mismatch and accuracy degradation. Spec-TQ84 closes this gap: at cap=16384, LCB multi-turn accuracy reaches 0.654 vs TQ44 at 0.372 (delta +28pp, McNemar p=1.17e-4).
This PR (
feat/spec-tq84-kernel-fix) contains three targeted fixes on top offeat/spec-tq84:VLLM_TQ_FORCE_SYNCworkaround that was masking the OOB bugChanges in this branch
5692d0fbdMSE_BITS==8branch to_tq_decode_stage1to fix cudagraph OOB read542d61ad8VLLM_TQ_FORCE_SYNCworkaround (superseded by5692d0fbd)7d3f8f495MSE_BITS==8branch to v2 decode kernel (same OOB as5692d0fbd)Files changed: 3 files, +44 / -61 lines (net -17)
Bug fixes (critical)
MSE_BITS==8 OOB kernel read (
5692d0fbd,7d3f8f495)_tq_decode_stage1(v1) and_tq_decode_stage1_v2(v2) were both missing a 1-byte-per-element fast path for 8-bit MSE keys (spec-TQ84 uses 256 centroids = 8 bits per index).The generic
raw16bit-extraction path loads two consecutive bytes (mse_addrs0andmse_addrs0 + 1) and combines them into a 16-bit word before shifting. For MSE_BITS==8, the second byte load atd = D-1reads one byte past the end of the MSE region. Under eager dispatch this OOB read happened to land on valid (but wrong) memory and was masked. Under cudagraph replay, the same address lands on an unmapped page, causing a HIP memory fault:Fix: When
MSE_BITS == 8, use a direct single-byte load per element (no bit-packing, noraw16combine). This is both correct and faster for the 8-bit case.The same fix was already applied to
_tq_full_dequant_kvand_tq_load_k_tilein commitb87671809on the parent branch;_tq_decode_stage1(v1 and v2) were the last missed sites.VLLM_TQ_FORCE_SYNC removal (
542d61ad8)The
VLLM_TQ_FORCE_SYNCenvironment variable enabled a per-dispatchcuda.synchronize()workaround that was introduced to mask the OOB fault above. With the root cause fixed, this dead code path adds ~200us serialization overhead per decode step when enabled and serves no purpose. Removed from bothtriton_turboquant_decode.pyandtriton_turboquant_unified_attention.py.Preset
The parent branch (
feat/spec-tq84) introduces theturboquant_k8v4_spec_ncpreset:This PR does not modify the preset; it only fixes the kernel to correctly decode 8-bit MSE keys.
Test results
--enforce-eager)Branch relationship
feat/spec-tq84-kernel-fixis 3 commits ahead offeat/spec-tq84feat/spec-tq84-kernel-fixis 0 commits behindfeat/spec-tq84feat/spec-tq84(722397c0b)feat/spec-tq84must be merged to the main integration branch first; thenfeat/spec-tq84-kernel-fixcan fast-forward merge on topKnown limitations
Review checklist
MSE_BITS==8direct-load path produces identical output to Python reference decodeturboquant_k8v4_spec_ncpreset (no--enforce-eager)VLLM_TQ_FORCE_SYNCremoval does not regress non-TQ84 configs (TQ44, TQ34)