From 811ccee2b24f86a5e282311976180dbe92acb4be Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Matsumoto=20Takaya=20=28=E6=9D=BE=E6=9C=AC=20=E8=B2=B4?= =?UTF-8?q?=E4=B9=9F=29?= Date: Fri, 4 Sep 2026 11:02:14 +0900 Subject: [PATCH] fix(kernels): give the V tensor its own row pitch in the fp8 KV store `quantize_kv_to_cache` passed `k.stride(0)` as the only source pitch and `_kv_quant_scatter_kernel` used it for both tensors: src = t * stride_xs + h * D + d xk = tl.load(k_src + src, ...) xv = tl.load(v_src + src, ...) The guard above it checks only the inner stride (`k.stride(1) == 1 and v.stride(1) == 1`), never `k.stride(0) == v.stride(0)`, so the kernel carries an undocumented contract: K and V must share one row pitch. When they do not, V is read at K's pitch. In the failing test K is a view of the qkv slice (pitch 1152) while V is materialised by `.clamp()` (pitch 384), so with 8 tokens of 3072 elements: token 0 reads 0 correct by coincidence token 1-2 reads 1152, 2304 in range, WRONG rows token 3-7 reads 3456 .. 8064 past the initialised data 2684 of 3072 codes wrong, all in V, K byte-perfect. Deterministic addressing; only the contents of the uninitialised tail vary with allocator history, which is why the mismatch count drifts (2684 / 2663 / 2676 across runs) while the mismatching positions do not -- the in-range half is exactly 764 every time. Found with `compute-sanitizer --tool initcheck` (TRITON_DISABLE_LINE_INFO=0), which named `kv_quant.py:110`. `memcheck` reports 0 errors because PyTorch's caching allocator rounds allocations up and the bad read stays inside the pooled segment; `racecheck` reports 0 hazards because it is not a race. Fix: pass `v.stride(0)` as its own kernel argument and load each tensor with its own pitch. Verified on RTX 4090 (sm_89): tests/kernels/test_kv_fp8.py 2 failed -> 1 failed, the flip being test_codes_match_the_reference_quantizer_and_reconstruction_is_close; five consecutive standalone runs give K 0/3072 and V 0/3072 with got.sort() == exp.sort(); compute-sanitizer initcheck reports 0 errors on the patched build. Independently confirmed on RTX 5090 D (sm_120) by @Kaempferia: same single flip, same multiset property, 5 runs clean. Note for reviewers: the test's SECOND assertion (dequantised error <= 0.08) passes at 0.035 while V is wrong, so a reconstruction-level check does not catch this class. Only the exact-code assertion does. Co-Authored-By: Claude Opus 5 --- python/freetoken/kernel/triton/kv_quant.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/python/freetoken/kernel/triton/kv_quant.py b/python/freetoken/kernel/triton/kv_quant.py index 6866316a2..e041e0466 100644 --- a/python/freetoken/kernel/triton/kv_quant.py +++ b/python/freetoken/kernel/triton/kv_quant.py @@ -89,7 +89,9 @@ def _kv_quant_scatter_kernel( k_scale, v_scale, idx_ptr, - stride_xs, # source row pitch, in elements (the qkv slice is wider than one row) + stride_xs, # K source row pitch, in elements (the qkv slice is wider than one row) + stride_vx, # V source row pitch. K and V need not share one: a .clamp() on one side + # leaves it densely packed, so reusing K's pitch reads V off its rows. stride_kd, # K cache row pitch, in elements (== HEADS * D) stride_vd, stride_ks, # scale row pitch, in elements (== HEADS) @@ -105,9 +107,9 @@ def _kv_quant_scatter_kernel( d = tl.arange(0, BLOCK_D) mask = d < D - src = t * stride_xs + h * D + d - xk = tl.load(k_src + src, mask=mask, other=0.0).to(tl.float32) - xv = tl.load(v_src + src, mask=mask, other=0.0).to(tl.float32) + off = h * D + d + xk = tl.load(k_src + t * stride_xs + off, mask=mask, other=0.0).to(tl.float32) + xv = tl.load(v_src + t * stride_vx + off, mask=mask, other=0.0).to(tl.float32) # 448 == e4m3 finite max; 1e-10 is the amax floor of the activation quant in # kernel/triton/fp8_block_linear.py (literals keep the kernel self-contained). @@ -172,6 +174,7 @@ def quantize_kv_to_cache( v_scale, out_loc, k.stride(0), + v.stride(0), k_cache.stride(0), v_cache.stride(0), k_scale.stride(0),