From 0780adb57f7d1bea014fffb54b3bf92475c85e88 Mon Sep 17 00:00:00 2001 From: Jack Hu Date: Tue, 21 Jul 2026 01:53:37 +0000 Subject: [PATCH 1/2] Enabled stride aware indexing for block, token, head, and head dim (assumed always contiguous) for pts quant shuffle fused qk norm rope and cache kernel --- .../kernels/fused_qk_norm_rope_cache_quant.cu | 36 +++++++---- csrc/kernels/rope/rope_common.h | 60 +++++++++++++------ 2 files changed, 68 insertions(+), 28 deletions(-) diff --git a/csrc/kernels/fused_qk_norm_rope_cache_quant.cu b/csrc/kernels/fused_qk_norm_rope_cache_quant.cu index 2be37a5a264..79cf0e959d4 100644 --- a/csrc/kernels/fused_qk_norm_rope_cache_quant.cu +++ b/csrc/kernels/fused_qk_norm_rope_cache_quant.cu @@ -3272,12 +3272,11 @@ void fused_qk_norm_rope_cache_pts_quant_shuffle(aiter_tensor_t& qkv, AITER_CHECK(qkv.is_contiguous() && qw.is_contiguous() && kw.is_contiguous() && cos_sin.is_contiguous()); AITER_CHECK(slot_mapping.is_contiguous()); - if(!(k_cache.is_contiguous() && v_cache.is_contiguous())) - { - // Non-contiguous block dim (e.g. vLLM [num_blocks, 2, ...] after unbind(1)) is OK as long as each block is internally contiguous. - AITER_CHECK(is_contiguous_from_dim1(k_cache) && is_contiguous_from_dim1(v_cache), - "k_cache/v_cache must be contiguous within a block (dims >= 1)"); - } + // block/token/head are read as strides below; only the innermost head_size + // dim must stay contiguous for the vectorized store. + AITER_CHECK(k_cache.stride(k_cache.dim() - 1) == 1 && + v_cache.stride(v_cache.dim() - 1) == 1, + "k_cache/v_cache innermost (head_size) dim must be contiguous"); HipDeviceGuard device_guard(qkv.device_id); const hipStream_t stream = aiter::getCurrentHIPStream(); auto kv_cache_dtype = k_cache.dtype(); @@ -3285,9 +3284,14 @@ void fused_qk_norm_rope_cache_pts_quant_shuffle(aiter_tensor_t& qkv, AITER_CHECK(positions.dim() == 1, "positions must be 1D"); float per_tensor_k_scale_ = *reinterpret_cast(per_tensor_k_scale.data_ptr()); float per_tensor_v_scale_ = *reinterpret_cast(per_tensor_v_scale.data_ptr()); - // Per-block (dim-0) stride: == num_heads_k*HEAD_SIZE*block_size when contiguous (old formula); larger for an interleaved [num_blocks, 2, ...] cache. + // K/V cache indexing is stride-aware in block/token/head (innermost head_size + // assumed contiguous), so most future KV layout changes need no change here. int64_t k_cache_block_stride = k_cache.stride(0); int64_t v_cache_block_stride = v_cache.stride(0); + int64_t k_cache_token_stride = k_cache.stride(1); + int64_t k_cache_head_stride = k_cache.stride(2); + int64_t v_cache_token_stride = v_cache.stride(1); + int64_t v_cache_head_stride = v_cache.stride(2); VLLM_DISPATCH_FLOATING_TYPES_rmTorch( qkv_dtype, "fused_qk_norm_rope_cache_pts_quant_shuffle", [&] { using T = scalar_t; @@ -3328,7 +3332,11 @@ void fused_qk_norm_rope_cache_pts_quant_shuffle(aiter_tensor_t& qkv, x, rotary_dim, k_cache_block_stride, - v_cache_block_stride); + v_cache_block_stride, + k_cache_token_stride, + k_cache_head_stride, + v_cache_token_stride, + v_cache_head_stride); } else { @@ -3375,7 +3383,11 @@ void fused_qk_norm_rope_cache_pts_quant_shuffle(aiter_tensor_t& qkv, x, rotary_dim, k_cache_block_stride, - v_cache_block_stride); + v_cache_block_stride, + k_cache_token_stride, + k_cache_head_stride, + v_cache_token_stride, + v_cache_head_stride); } else { @@ -3418,7 +3430,11 @@ void fused_qk_norm_rope_cache_pts_quant_shuffle(aiter_tensor_t& qkv, x, rotary_dim, k_cache_block_stride, - v_cache_block_stride); + v_cache_block_stride, + k_cache_token_stride, + k_cache_head_stride, + v_cache_token_stride, + v_cache_head_stride); } } else diff --git a/csrc/kernels/rope/rope_common.h b/csrc/kernels/rope/rope_common.h index 62ead401e05..4e70c73ca2e 100644 --- a/csrc/kernels/rope/rope_common.h +++ b/csrc/kernels/rope/rope_common.h @@ -7789,7 +7789,11 @@ __global__ void fused_mrope_rms_kv_kernel(const T* qkv, int rotary_dim = 0, int64_t k_block_stride = 0, int64_t v_block_stride = 0, - bool gemma_norm = false) + bool gemma_norm = false, + int64_t k_token_stride = 0, + int64_t k_head_stride = 0, + int64_t v_token_stride = 0, + int64_t v_head_stride = 0) { constexpr int VEC_SIZE = HEAD_SIZE / WARP_SIZE; constexpr int HALF_HEAD_SIZE = HEAD_SIZE / 2; @@ -8007,15 +8011,18 @@ __global__ void fused_mrope_rms_kv_kernel(const T* qkv, } else { - // block_size == 0 => non-paged cache (flat [num_slots, num_heads_k, HEAD_SIZE]): - // index directly by slot. Otherwise the cache is paged and K/V are interleaved - // per block, so index with the cache's real per-block stride (k_block_stride): - // offset = block_id*block_stride + block_offset*slot_size + head*HEAD_SIZE + elem + // Non-paged (block_size==0) indexes by slot; paged offset = + // block*block_stride + slot*tok + head*head + elem. const int64_t slot_size = static_cast(num_heads_k) * HEAD_SIZE; + // Each stride 0 => contiguous fallback (num_heads*HEAD_SIZE, + // HEAD_SIZE); non-zero writes a strided view (e.g. #44455) in place. + const int64_t tok_stride = (k_token_stride != 0) ? k_token_stride : slot_size; + const int64_t head_stride = (k_head_stride != 0) ? k_head_stride + : static_cast(HEAD_SIZE); int64_t offset; if(block_size == 0) { - offset = slot_id * slot_size + head_id_k * HEAD_SIZE + access_id_in_head; + offset = slot_id * tok_stride + head_id_k * head_stride + access_id_in_head; } else { @@ -8024,8 +8031,8 @@ __global__ void fused_mrope_rms_kv_kernel(const T* qkv, const int64_t block_stride = (k_block_stride != 0) ? k_block_stride : static_cast(block_size) * slot_size; - offset = block_id * block_stride + block_offset * slot_size + - head_id_k * HEAD_SIZE + access_id_in_head; + offset = block_id * block_stride + block_offset * tok_stride + + head_id_k * head_stride + access_id_in_head; } out_kv_vec.store(k_cache + offset); } @@ -8061,13 +8068,16 @@ __global__ void fused_mrope_rms_kv_kernel(const T* qkv, } else { - // Same scheme as the K path above, for the V cache. - // block_size == 0 => non-paged cache, index directly by slot. + // Same stride-aware scheme as K (block/token/head; 0 => contiguous + // fallback, non-zero writes a strided view such as the #44455 packed). const int64_t slot_size = static_cast(num_heads_v) * HEAD_SIZE; + const int64_t tok_stride = (v_token_stride != 0) ? v_token_stride : slot_size; + const int64_t head_stride = (v_head_stride != 0) ? v_head_stride + : static_cast(HEAD_SIZE); int64_t offset; if(block_size == 0) { - offset = slot_id * slot_size + head_id_v * HEAD_SIZE + access_id_in_head; + offset = slot_id * tok_stride + head_id_v * head_stride + access_id_in_head; } else { @@ -8076,8 +8086,8 @@ __global__ void fused_mrope_rms_kv_kernel(const T* qkv, const int64_t block_stride = (v_block_stride != 0) ? v_block_stride : static_cast(block_size) * slot_size; - offset = block_id * block_stride + block_offset * slot_size + - head_id_v * HEAD_SIZE + access_id_in_head; + offset = block_id * block_stride + block_offset * tok_stride + + head_id_v * head_stride + access_id_in_head; } out_kv_vec.store(v_cache + offset); } @@ -8091,8 +8101,8 @@ __global__ void fused_mrope_rms_kv_kernel(const T* qkv, } } -// mrope-3D launcher: intentionally relies on the default-0 (contiguous) block stride -// in fused_mrope_rms_kv_kernel; the stride-aware path is the pts launcher below. +// mrope-3D launcher: relies on default-0 (contiguous) strides; the fully +// stride-aware path is the pts launcher (fused_rope_rms_set_kv) below. template void fused_mrope_rms_set_kv(const T* qkv, const T* q_w, @@ -8259,7 +8269,11 @@ void fused_rope_rms_set_kv(const T* qkv, int64_t x = 0, int64_t rotary_dim = 0, int64_t k_block_stride = 0, - int64_t v_block_stride = 0) + int64_t v_block_stride = 0, + int64_t k_token_stride = 0, + int64_t k_head_stride = 0, + int64_t v_token_stride = 0, + int64_t v_head_stride = 0) { TORCH_CHECK(head_size == 64 || head_size == 128 || head_size == 256); constexpr int THREAD_BLOCK_SIZE = 256; @@ -8300,7 +8314,12 @@ void fused_rope_rms_set_kv(const T* qkv, x, \ (int)rotary_dim, \ k_block_stride, \ - v_block_stride); \ + v_block_stride, \ + /*gemma_norm=*/false,\ + k_token_stride, \ + k_head_stride, \ + v_token_stride, \ + v_head_stride); \ } \ else \ { \ @@ -8332,7 +8351,12 @@ void fused_rope_rms_set_kv(const T* qkv, x, \ (int)rotary_dim, \ k_block_stride, \ - v_block_stride); \ + v_block_stride, \ + /*gemma_norm=*/false,\ + k_token_stride, \ + k_head_stride, \ + v_token_stride, \ + v_head_stride); \ } switch(head_size) From 74cc0599d6e4f34cb679de66c2c272002e3fafb5 Mon Sep 17 00:00:00 2001 From: Jack Hu Date: Tue, 21 Jul 2026 02:29:06 +0000 Subject: [PATCH 2/2] Modified the unit test for this fusion to verify kv layout correctness parity between blocks first, kv first, and head dim packed first --- .../test_fused_qk_norm_rope_cache_quant.py | 86 ++++++++++++++----- 1 file changed, 66 insertions(+), 20 deletions(-) diff --git a/op_tests/test_fused_qk_norm_rope_cache_quant.py b/op_tests/test_fused_qk_norm_rope_cache_quant.py index 5cb642a7c47..0f59aa895ea 100644 --- a/op_tests/test_fused_qk_norm_rope_cache_quant.py +++ b/op_tests/test_fused_qk_norm_rope_cache_quant.py @@ -2874,7 +2874,10 @@ def test_pts_quant_shuffle_block_layout_parity( eps=1e-6, use_shuffle_layout=True, ): - """Parity: the pts write must give identical KV cache for the original [2, num_blocks, ...] and new [num_blocks, 2, ...] (unbind(1)) paged layouts, for both use_shuffle_layout=True and False (both honor the cache's per-block stride).""" + """Parity: the pts write must give an identical KV cache across paged layouts + that differ only in block/token/head strides -- [2, num_blocks, ...], + [num_blocks, 2, ...] (unbind(1)), and the packed + [num_blocks, num_heads_kv, block_size, 2*head_size] (non-shuffle only).""" cache_dtype = cache_dtype or dtype # None => auto (cache dtype == qkv dtype) x = ( 16 // torch.empty(0, dtype=cache_dtype).element_size() @@ -2904,22 +2907,42 @@ def test_pts_quant_shuffle_block_layout_parity( per_tensor_k_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda") per_tensor_v_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda") - def run(blocks_first: bool): - # blocks_first=True: new [num_blocks,2,...] (K=cache[:,0], stride 2x); False: original [2,num_blocks,...] (K=cache[0], contiguous). - if blocks_first: + def run(layout: str): + # Three physically distinct paged KV layouts that must all write the same + # logical [num_blocks, block_size, num_heads_kv, head_size] K/V cache. The + # kernel is stride-aware in block/token/head; only the innermost head_size + # dim must stay contiguous (stride 1). + if layout == "kv_first": + # original [2, num_blocks, ...]: K=cache[0], per-block contiguous. + kv = torch.zeros( + (2, num_blocks, block_size, num_heads_kv, head_size), + dtype=cache_dtype, + device="cuda", + ) + k_cache, v_cache = kv[0], kv[1] + elif layout == "blocks_first": + # [num_blocks, 2, ...]: K=cache[:,0], K/V interleaved per block (block stride 2x). kv = torch.zeros( (num_blocks, 2, block_size, num_heads_kv, head_size), dtype=cache_dtype, device="cuda", ) k_cache, v_cache = kv[:, 0], kv[:, 1] - else: + elif layout == "packed_headdim": + # [num_blocks, num_heads_kv, block_size, 2*head_size]: K/V packed along + # the innermost dim (K=[...,:D], V=[...,D:]) with head/token order + # transposed vs the others. permute(0,2,1,3) -> + # [num_blocks, block_size, num_heads_kv, head_size] so the kernel's + # (block, token, head) stride reads line up; head_size stays contiguous. kv = torch.zeros( - (2, num_blocks, block_size, num_heads_kv, head_size), + (num_blocks, num_heads_kv, block_size, 2 * head_size), dtype=cache_dtype, device="cuda", ) - k_cache, v_cache = kv[0], kv[1] + k_cache = kv[..., :head_size].permute(0, 2, 1, 3) + v_cache = kv[..., head_size:].permute(0, 2, 1, 3) + else: + raise ValueError(f"unknown layout {layout!r}") q_out = torch.empty( (num_tokens, num_heads_q, head_size), dtype=dtype, device="cuda" ) @@ -2958,8 +2981,20 @@ def run(blocks_first: bool): ) return q_out, k_cache, v_cache - q_orig, k_orig, v_orig = run(blocks_first=False) # [2, num_blocks, ...] - q_new, k_new, v_new = run(blocks_first=True) # [num_blocks, 2, ...] + # Reference layout + the variants under test; every layout must produce a + # bit-identical logical K/V cache (only block/token/head strides differ). + q_ref, k_ref, v_ref = run("kv_first") # [2, num_blocks, ...] + variants = { + "blocks_first": run("blocks_first"), # [num_blocks, 2, ...] + } + # Packed [num_blocks, Hkv, block_size, 2*head_size]: head_size contiguous in the + # innermost dim (K=[...,:D], V=[...,D:]). Non-shuffle only -- the x-shuffle repacks + # head_size non-contiguously and honors only the block stride, so a contiguous-head + # packed layout can't be represented under it. + if not use_shuffle_layout: + variants["packed_headdim"] = run( + "packed_headdim" + ) # [num_blocks, Hkv, block_size, 2*D] tag = ( f"block_layout_parity qkv={dtype}, cache={cache_dtype}, tokens={num_tokens}, " @@ -2967,16 +3002,25 @@ def run(blocks_first: bool): f"block_size={block_size}, blocks={num_blocks}, neox={is_neox_style}, " f"shuffle={use_shuffle_layout}" ) - # Only the block stride differs -> must match exactly; checkAllclose logs but doesn't raise, so assert on its ratio. - for name, a, b in ( - ("q_out", q_orig, q_new), - ("k_cache", k_orig, k_new), - ("v_cache", v_orig, v_new), - ): - err = checkAllclose( - a.float(), b.float(), rtol=0, atol=0, printLog=False, msg=f"{name} {tag}" - ) - assert err == 0, f"{name} [2,B] vs [B,2] parity MISMATCH (err={err}): {tag}" + # Only block/token/head strides differ -> must match exactly; checkAllclose + # logs but doesn't raise, so assert on its returned ratio. + for layout, (q_v, k_v, v_v) in variants.items(): + for name, a, b in ( + ("q_out", q_ref, q_v), + ("k_cache", k_ref, k_v), + ("v_cache", v_ref, v_v), + ): + err = checkAllclose( + a.float(), + b.float(), + rtol=0, + atol=0, + printLog=False, + msg=f"{name} kv_first vs {layout} {tag}", + ) + assert ( + err == 0 + ), f"{name} kv_first vs {layout} parity MISMATCH (err={err}): {tag}" print(f"[PASS] {tag}", flush=True) return { "qkv_dtype": str(dtype), @@ -3376,7 +3420,9 @@ def run(blocks_first: bool): df_md = df.to_markdown(index=False) aiter.logger.info("partial_rotary_pts_quant summary (markdown):\n%s", df_md) - # Stride-aware shuffle parity: original [2, num_blocks, ...] vs new [num_blocks, 2, ...] paged layout must give identical KV cache. + # Stride-aware shuffle parity: [2, num_blocks, ...], [num_blocks, 2, ...] and + # packed [num_blocks, Hkv, block_size, 2*head_size] paged layouts must all + # give an identical KV cache. fp8 = get_dtype_fp8() df = [] for is_neox_style in args.is_neox_styles: