Skip to content

[Gluon][GFX950][MLA] Fix mla decode accuracy issue with empty kv split #3641

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
leonling-ll merged 3 commits into from
Jun 12, 2026
Merged

[Gluon][GFX950][MLA] Fix mla decode accuracy issue with empty kv split #3641

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
22cc754
Fix mla decode empty kv split
leonling-ll Jun 9, 2026
34722d2
make early return when empty split
leonling-ll Jun 10, 2026
8930d24
Reduce only loop valid kv seq
leonling-ll Jun 11, 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
76 changes: 48 additions & 28 deletions aiter/ops/triton/gluon/mla_decode_gluon.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -117,6 +117,33 @@ def _mla_decode_gluon(
batch_page_start = gl.load(B_seq_len + cur_batch)
cur_batch_seq_len = gl.load(B_seq_len + cur_batch + 1) - batch_page_start

# split-KV: each program covers [split_kv_start, split_kv_end).
# OLD: ceil-based per_split. The LAST split could be empty (num_iter=0),
# which breaks the unconditional epilogue-2 consume. Kept here as commented
# reference; remove in cleanup.
# kv_len_per_split = gl.cdiv(cur_batch_seq_len, NUM_KV_SPLITS)
# split_kv_start = kv_len_per_split * split_kv_id
# split_kv_end = gl.minimum(split_kv_start + kv_len_per_split, cur_batch_seq_len)
#
# NEW: floor per_split with the last split absorbing the remainder
# (remainder = seq mod NUM_KV_SPLITS, in [0, NUM_KV_SPLITS)). Combined with
# the wrapper bound min_kv_seq_len >= NUM_KV_SPLITS this guarantees every
# split is non-empty (split_len >= floor >= 1, hence num_iter >= 1); bh64
# additionally bounds min_kv_seq_len so num_iter >= 3 for its gl.assume.
# Trade-off: at seqs just above the wrapper minimum the last CU does up to
# ~(floor + NUM_KV_SPLITS - 1)/floor more work than the others.
kv_len_per_split = cur_batch_seq_len // NUM_KV_SPLITS
split_kv_start = kv_len_per_split * split_kv_id
split_kv_end = split_kv_start + kv_len_per_split
if split_kv_id == NUM_KV_SPLITS - 1:
split_kv_end = cur_batch_seq_len
num_iter = gl.cdiv(split_kv_end - split_kv_start, BLOCK_N)
start_n = split_kv_start

# early return with empty kv slice to save compute
if split_kv_start >= split_kv_end:
return

######### layout setting begin #########
# Q-side layouts + mfma_layout: switch by BLOCK_H.
# bh64 has BLOCK_H=64; bh16bn128 and bh16bn64 share BLOCK_H=16 (identical Q layouts + mfma orientation).
Expand Down Expand Up @@ -325,29 +352,6 @@ def _mla_decode_gluon(
e_sum = gl.zeros([BLOCK_H], dtype=gl.float32, layout=gl.SliceLayout(1, mfma_layout))
acc = gl.zeros([BLOCK_H, HEAD_DIM_CKV], dtype=gl.float32, layout=mfma_layout)

# split-KV: each program covers [split_kv_start, split_kv_end).
# OLD: ceil-based per_split. The LAST split could be empty (num_iter=0),
# which breaks the unconditional epilogue-2 consume. Kept here as commented
# reference; remove in cleanup.
# kv_len_per_split = gl.cdiv(cur_batch_seq_len, NUM_KV_SPLITS)
# split_kv_start = kv_len_per_split * split_kv_id
# split_kv_end = gl.minimum(split_kv_start + kv_len_per_split, cur_batch_seq_len)
#
# NEW: floor per_split with the last split absorbing the remainder
# (remainder = seq mod NUM_KV_SPLITS, in [0, NUM_KV_SPLITS)). Combined with
# the wrapper bound min_kv_seq_len >= NUM_KV_SPLITS this guarantees every
# split is non-empty (split_len >= floor >= 1, hence num_iter >= 1); bh64
# additionally bounds min_kv_seq_len so num_iter >= 3 for its gl.assume.
# Trade-off: at seqs just above the wrapper minimum the last CU does up to
# ~(floor + NUM_KV_SPLITS - 1)/floor more work than the others.
kv_len_per_split = cur_batch_seq_len // NUM_KV_SPLITS
split_kv_start = kv_len_per_split * split_kv_id
split_kv_end = split_kv_start + kv_len_per_split
if split_kv_id == NUM_KV_SPLITS - 1:
split_kv_end = cur_batch_seq_len
num_iter = gl.cdiv(split_kv_end - split_kv_start, BLOCK_N)
start_n = split_kv_start

# Fold KV dequant scale into the QK temperature. For fp8 KV the real
# logits are (Q @ K_fp8^T) * kv_scale * sm_scale; softmax is shift- but
# not scale-invariant, so kv_scale must affect qk (not just acc).
Expand Down Expand Up @@ -658,6 +662,7 @@ def _mla_softmax_reducev_kernel(
Mid_lse,
O, # noqa: E741
Final_lse,
B_seq_len, # same seq_info as the decode kernel to derive empty kv splits
stride_l_b,
stride_l_h,
stride_l_s,
Expand All @@ -671,10 +676,22 @@ def _mla_softmax_reducev_kernel(
NUM_KV_SPLITS: tl.constexpr,
HEAD_DIM_CKV: tl.constexpr,
HAS_FINAL_LSE: tl.constexpr,
USE_2D_VIEW: tl.constexpr,
):
cur_batch = tl.program_id(0)
cur_head = tl.program_id(1)

# Recompute this batch's seq len exactly as the decode kernel did, so we can
# rederive which splits are empty. Stage-1 early-returns on empty splits
# (num_iter == 0) and writes nothing, so their logits_buf / mid_lse slots hold
# raw, uninitialised memory - they cannot be loaded or reduced.
if USE_2D_VIEW:
cur_batch_seq_len = tl.load(B_seq_len + cur_batch)
else:
batch_page_start = tl.load(B_seq_len + cur_batch)
cur_batch_seq_len = tl.load(B_seq_len + cur_batch + 1) - batch_page_start
kv_len_per_split = cur_batch_seq_len // NUM_KV_SPLITS

offs_d_ckv = tl.arange(0, HEAD_DIM_CKV)
offs_l = cur_batch * stride_l_b + cur_head * stride_l_h + offs_d_ckv
offs_ml = cur_batch * stride_ml_b + cur_head * stride_ml_h
Expand All @@ -683,23 +700,24 @@ def _mla_softmax_reducev_kernel(
e_max = -float("inf")
acc = tl.zeros([HEAD_DIM_CKV], dtype=tl.float32)

# all splits non-empty (floor-with-remainder-on-last), so merge unconditionally.
for split_kv_id in range(0, NUM_KV_SPLITS):
LOOP_START = NUM_KV_SPLITS - 1 if kv_len_per_split == 0 else 0
for split_kv_id in range(LOOP_START, NUM_KV_SPLITS):
logits = tl.load(Logits + offs_l + split_kv_id * stride_l_s)
logits_1 = tl.load(Mid_lse + offs_ml + split_kv_id * stride_ml_s)

n_e_max = tl.maximum(logits_1, e_max)
old_scale = tl.exp(e_max - n_e_max)
old_scale = tl.where(e_max == -float("inf"), 0.0, tl.exp(e_max - n_e_max))
acc *= old_scale
exp_logic = tl.exp(logits_1 - n_e_max)
exp_logic = tl.where(logits_1 == -float("inf"), 0.0, tl.exp(logits_1 - n_e_max))
acc += exp_logic * logits

e_sum = e_sum * old_scale + exp_logic
e_max = n_e_max

out = acc / e_sum if e_sum > 0.0 else tl.zeros([HEAD_DIM_CKV], dtype=tl.float32)
tl.store(
O + cur_batch * stride_o_b + cur_head * stride_o_h + offs_d_ckv,
acc / e_sum,
out,
)
if HAS_FINAL_LSE:
tl.store(
Expand Down Expand Up @@ -926,6 +944,7 @@ def mla_decode_gluon(
mid_lse,
o,
final_lse,
seq_info,
logits_buf.stride(0),
logits_buf.stride(1),
logits_buf.stride(2),
Expand All @@ -939,6 +958,7 @@ def mla_decode_gluon(
NUM_KV_SPLITS=NUM_KV_SPLITS,
HEAD_DIM_CKV=head_dim_ckv,
HAS_FINAL_LSE=return_lse,
USE_2D_VIEW=use_2d_view,
num_warps=8,
)

Expand Down
Loading