Skip to content

fused_qk_rope_cat_and_cache_mla: Fix Triton compilation error and batch size constraint and output tensor sizing #1407

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 2 commits into from
Nov 15, 2025
Merged

fused_qk_rope_cat_and_cache_mla: Fix Triton compilation error and batch size constraint and output tensor sizing #1407

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
97c507b
Fix Triton compilation error by nesting OUTPUT_Q_NOPE_ZEROS condition
xudonlyu Nov 13, 2025
b980a66
Correct batch size constraint and output tensor sizing
xudonlyu Nov 14, 2025
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
21 changes: 12 additions & 9 deletions aiter/ops/triton/_triton_kernels/fused_kv_cache.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -229,15 +229,18 @@ def _fused_qk_rope_cat_and_cache_mla_kernel(
q_pe.to(decode_q_pe_out_ptr.dtype.element_ty),
)

if OUTPUT_Q_NOPE_ZEROS and pid < num_decode_toks_for_zeros * QH:
z = tl.zeros((BLOCK_DK_nope,), dtype=q_nope_zeros_out_ptr.dtype.element_ty)
tl.store(
q_nope_zeros_out_ptr
+ pid_b * q_nope_zeros_out_stride_b
+ pid_hq * q_nope_zeros_out_stride_h
+ dk_nope_offs * q_nope_zeros_out_stride_d,
z,
)
if OUTPUT_Q_NOPE_ZEROS:
if pid < num_decode_toks_for_zeros * QH:
z = tl.zeros(
(BLOCK_DK_nope,), dtype=q_nope_zeros_out_ptr.dtype.element_ty
)
tl.store(
q_nope_zeros_out_ptr
+ pid_b * q_nope_zeros_out_stride_b
+ pid_hq * q_nope_zeros_out_stride_h
+ dk_nope_offs * q_nope_zeros_out_stride_d,
z,
)

if pid_hq % QH_PER_KH == 0:
pid_slot = tl.load(slot_mapping_ptr + pid_b).to(tl.int64)
Expand Down
11 changes: 6 additions & 5 deletions aiter/ops/triton/fused_kv_cache.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,9 +60,10 @@ def fused_qk_rope_cat_and_cache_mla(
b_cache, h_cache, d_cache = kv_cache.shape
(b_slot,) = slot_mapping.shape

# allow bk >= b to support prefill + decode mixed scenario
assert (
b_slot <= b and b == b2 == bk == bk2
), "batch dimension should be identical for q_nope, q_pe, k_nope, and k_pe, and the batch dimeion of slot_mapping should be no more than that of q_nope, q_pe, k_nope, and k_pe"
b == b2 and bk == bk2 and b_slot <= bk and b <= bk
), "Q batch dimensions should be identical (b == b2), K batch dimensions should be identical (bk == bk2), slot_mapping should not exceed K batch size (b_slot <= bk), and Q batch should not exceed K batch (b <= bk)"
assert qh == qh2, "Q head should be identical"
assert kh == kh2 == h_cache, "K head should be identical"
assert d_pe == dk2, "D dimension of q_pe and k_pe should be identical"
Expand Down Expand Up @@ -105,12 +106,12 @@ def fused_qk_rope_cat_and_cache_mla(
), "decode_q_pe_out shape mismatch"

if k_pe_out is None:
k_pe_out = torch.empty((b, kh, d_pe), dtype=k_pe.dtype, device=k_pe.device)
k_pe_out = torch.empty((bk, kh, d_pe), dtype=k_pe.dtype, device=k_pe.device)
else:
b_k_pe_out, hk_k_pe_out, d_k_pe_out = k_pe_out.shape
assert (
b == b_k_pe_out and kh == hk_k_pe_out and d_pe == d_k_pe_out
), "k_pe_out shape mismatch"
bk == b_k_pe_out and kh == hk_k_pe_out and d_pe == d_k_pe_out
), "k_pe_out shape mismatch, expected (bk, kh, d_pe)"

q_nope_zeros_out = None
if num_decode_toks_for_zeros > 0:
Expand Down