diff --git a/aiter/ops/triton/_gluon_kernels/gfx1250/attention/mla.py b/aiter/ops/triton/_gluon_kernels/gfx1250/attention/mla.py index 45fef0cfd61..80e14480569 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx1250/attention/mla.py +++ b/aiter/ops/triton/_gluon_kernels/gfx1250/attention/mla.py @@ -2878,13 +2878,19 @@ def _mla_decode_fwd_kernel_non_pipelined( acc, mask=query_mask_0_pv[:, None] & query_mask_1_pv[:, None], ) - segm_offset = ( - query_offset_0_qk.to(gl.int64) * (num_query_heads * NUM_SEGMENTS_PER_SEQ) - + query_offset_1_qk * NUM_SEGMENTS_PER_SEQ - + segm_idx - ) - gl.store(segm_max_ptr + segm_offset, M, mask=query_mask_0_qk & query_mask_1_qk) - gl.store(segm_expsum_ptr + segm_offset, L, mask=query_mask_0_qk & query_mask_1_qk) + # NUM_SEGMENTS_PER_SEQ == 1 skips the reduce kernel, and the host then aliases + # segm_max_ptr / segm_expsum_ptr onto the output buffer as dummy pointers, so + # writing M / L here would clobber the attention result. + if NUM_SEGMENTS_PER_SEQ > 1: + segm_offset = ( + query_offset_0_qk.to(gl.int64) * (num_query_heads * NUM_SEGMENTS_PER_SEQ) + + query_offset_1_qk * NUM_SEGMENTS_PER_SEQ + + segm_idx + ) + gl.store(segm_max_ptr + segm_offset, M, mask=query_mask_0_qk & query_mask_1_qk) + gl.store( + segm_expsum_ptr + segm_offset, L, mask=query_mask_0_qk & query_mask_1_qk + ) _mla_decode_fwd_reduce_kernel_repr = make_kernel_repr( diff --git a/aiter/ops/triton/_triton_kernels/attention/mla.py b/aiter/ops/triton/_triton_kernels/attention/mla.py index 1f535dcb263..50cc34ee708 100644 --- a/aiter/ops/triton/_triton_kernels/attention/mla.py +++ b/aiter/ops/triton/_triton_kernels/attention/mla.py @@ -639,13 +639,18 @@ def _mla_decode_fwd_kernel( acc, mask=query_mask_0[:, None] & query_mask_1[:, None], ) - segm_offset = ( - query_offset_0.to(tl.int64) * (num_query_heads * NUM_SEGMENTS_PER_SEQ) - + query_offset_1 * NUM_SEGMENTS_PER_SEQ - + segm_idx - ) - tl.store(segm_max_ptr + segm_offset, M, mask=query_mask_0 & query_mask_1) - tl.store(segm_expsum_ptr + segm_offset, L, mask=query_mask_0 & query_mask_1) + # Same guard as the gfx1250 Gluon kernel. select_3d_config floors this + # branch at MIN_SEGMENTS >= 8 so one segment is unreachable from here + # today, but the aliasing is shared: at one segment the host hands both + # pointers the output buffer. + if NUM_SEGMENTS_PER_SEQ > 1: + segm_offset = ( + query_offset_0.to(tl.int64) * (num_query_heads * NUM_SEGMENTS_PER_SEQ) + + query_offset_1 * NUM_SEGMENTS_PER_SEQ + + segm_idx + ) + tl.store(segm_max_ptr + segm_offset, M, mask=query_mask_0 & query_mask_1) + tl.store(segm_expsum_ptr + segm_offset, L, mask=query_mask_0 & query_mask_1) _mla_decode_fwd_reduce_kernel_repr = make_kernel_repr(