[peizhang56]

On the WITHIN_2GB=False path the KV offset kv_loc * stride_kv_c_bs is computed in int32 and overflows (that overflow is precisely why this branch exists), yielding a wild 64-bit pointer in the global_load fallback and a SIGSEGV in _mla_gluon.

Promote stride_kv_c_bs and stride_k_pe_bs to int64 when not WITHIN_2GB so the whole offset computes in 64-bit. The buffer_load path (WITHIN_2GB=True) keeps 32-bit offsets by hardware requirement; the guard is constexpr so it folds out there.

Motivation

_mla_gluon crashes with a SIGSEGV when the KV cache exceeds 2GB. In that case the kernel selects the WITHIN_2GB=False code path, but the KV offset kv_loc * stride_kv_c_bs is still computed in int32. For a >2GB cache this multiplication overflows, and the overflowed offset is added to the 64-bit base pointer in the global_load fallback, producing an out-of-bounds pointer and a segfault. The overflow is precisely the condition that selects this branch, so the offset must be computed in 64-bit.

Technical Details

Promote stride_kv_c_bs and stride_k_pe_bs to gl.int64 when not WITHIN_2GB, so every downstream offset computation (kv_loc * stride_kv_c_bs + ... for the KV cache, and the corresponding stride_k_pe_bs expressions for K_pe) is evaluated in 64-bit and no longer overflows. The reassignment happens before the first offset is built, so all load sites pick up the widened strides. The buffer_load path (WITHIN_2GB=True) intentionally keeps 32-bit offsets, which is a hardware requirement of buffer_load_to_shared; since WITHIN_2GB is a gl.constexpr, the guard folds out entirely on that path and there is no runtime or performance impact.

Test Plan

Ran the MLA gluon kernel with a KV cache sized above 2GB to force the WITHIN_2GB=False / global_load path, and confirmed the same configuration crashed before the change. Also exercised the WITHIN_2GB=True path to confirm no behavioral or performance change (the guard folds out at compile time).

Test Result

Before: SIGSEGV in _mla_gluon on the >2GB KV-cache configuration. After: the kernel runs correctly and produces expected output; the WITHIN_2GB=True path is unchanged. No new unit test is added because reproducing this reliably requires allocating a >2GB KV cache, which is impractical in the standard op_test harness.

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 4474 --add-label <label>

[dllehr-amd]

Thanks @peizhang56 ! Looks good on my side!

[peizhang56]

Hi @valarLip will you be able to help to review this code change.
There was a vllm k3 coredump due to interger overflow from K3 when prefix cache + triton_mla + dspark spec decode enabled.
I believe different ppl have tested with this, which fixed the coredump for vllm.