Skip to content

Enable stride aware indexing on top of strided blocks for block_size and heads (assumes contiguous head dim) of fused_qk_norm_rope_cache_pts_quant_shuffle() #4312

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
valarLip merged 5 commits into from
Jul 29, 2026
Merged

Enable stride aware indexing on top of strided blocks for block_size and heads (assumes contiguous head dim) of fused_qk_norm_rope_cache_pts_quant_shuffle() #4312

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
0780adb
Enabled stride aware indexing for block, token, head, and head dim (a…
jhu960213 Jul 21, 2026
74cc059
Modified the unit test for this fusion to verify kv layout correctnes…
jhu960213 Jul 21, 2026
0dae8f3
Merge branch 'main' into jhu96/stride-aware-pt3
jhu960213 Jul 27, 2026
7dc2374
Merge branch 'main' into jhu96/stride-aware-pt3
jhu960213 Jul 27, 2026
028b87b
Merge branch 'main' into jhu96/stride-aware-pt3
jhu960213 Jul 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 26 additions & 10 deletions csrc/kernels/fused_qk_norm_rope_cache_quant.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3271,22 +3271,26 @@ 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();
auto qkv_dtype = qkv.dtype();
AITER_CHECK(positions.dim() == 1, "positions must be 1D");
float per_tensor_k_scale_ = *reinterpret_cast<float*>(per_tensor_k_scale.data_ptr());
float per_tensor_v_scale_ = *reinterpret_cast<float*>(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;
Expand Down Expand Up @@ -3327,7 +3331,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
{
Expand Down Expand Up @@ -3374,7 +3382,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
{
Expand Down Expand Up @@ -3417,7 +3429,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
Expand Down
60 changes: 42 additions & 18 deletions csrc/kernels/rope/rope_common.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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<int64_t>(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<int64_t>(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
{
Expand All @@ -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<int64_t>(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);
}
Expand Down Expand Up @@ -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<int64_t>(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<int64_t>(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
{
Expand All @@ -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<int64_t>(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);
}
Expand All @@ -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 <typename T, int M, typename KVT>
void fused_mrope_rms_set_kv(const T* qkv,
const T* q_w,
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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 \
{ \
Expand Down Expand Up @@ -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)
Expand Down
86 changes: 66 additions & 20 deletions op_tests/test_fused_qk_norm_rope_cache_quant.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2871,7 +2871,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()
Expand Down Expand Up @@ -2901,22 +2904,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"
)
Expand Down Expand Up @@ -2955,25 +2978,46 @@ 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}, "
f"Hq={num_heads_q}, Hkv={num_heads_kv}, D={head_size}, rotary_dim={rotary_dim}, "
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),
Expand Down Expand Up @@ -3372,7 +3416,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:
Expand Down
Loading