[JohnQin-llm]

Summary

Add deterministic regression coverage for the greater-than-2-GiB Gluon MLA
KV-cache offset overflow fixed by #4474. This PR is test-only and depends on
#4474; it no longer carries a competing runtime fix.

The test allocates a 4.72-GiB BF16 cache and indexes row 4,255,488, beyond the
signed 32-bit wrap point for a 576-element row. It covers both a full
BLOCK_N=128 tile and a partial 100-token trailing tile.

Why existing coverage missed it

The existing harness allocates 2,097,152 rows, so it selects the
WITHIN_2GB=False branch but cannot index past the wrap point:

Case	Maximum row	Compared with 2**31 // 576
existing random indices	2,097,151	1.78x short
sequential, batch 256 x context 8,192	about 2,107,152	1.77x short
this test	4,255,488	crosses the boundary

The child process is intentional. An HSA memory fault aborts the interpreter;
subprocess isolation turns that into one actionable test failure instead of
terminating the remaining CI shard. The parent also checks that sufficient
memory is available before launching the 4.72-GiB case.

Public-image reproduction

Hardware: AMD Instinct MI355X (gfx950). Public base image:
vllm/vllm-openai-rocm:kimi-k3@sha256:5aa7e626ff73672f5ca7aae46754570488c23d33ca1ac90756a1d2d1a3fe099b.
Test PR head: 81e524ccb559731d5bf9ba98e62ea49ddd80b816.
Runtime dependency #4474 head:
1e4135b1070153135de288af87258f94351fdaba.

Reconstruct only the two explicit PR deltas; do not merge a moving main:

git clone https://github.com/ROCm/aiter.git aiter-source
git -C aiter-source fetch origin \
  refs/pull/4474/head:pr-4474 \
  refs/pull/4488/head:pr-4488
test "$(git -C aiter-source rev-parse pr-4474)" = \
  1e4135b1070153135de288af87258f94351fdaba
test "$(git -C aiter-source rev-parse pr-4474^)" = \
  3a79aec3b71393a36351e2d0512414d9a1ba7813
test "$(git -C aiter-source rev-parse pr-4488)" = \
  81e524ccb559731d5bf9ba98e62ea49ddd80b816
test "$(git -C aiter-source rev-parse pr-4488^)" = \
  702aacd62c8a6e2fbbc260da338c510ab76f1b1c
git -C aiter-source diff \
  3a79aec3b71393a36351e2d0512414d9a1ba7813..pr-4474 \
  --binary > aiter-4474.patch
git -C aiter-source diff \
  702aacd62c8a6e2fbbc260da338c510ab76f1b1c..pr-4488 \
  --binary > aiter-4488.patch

git clone aiter-source aiter-fixed-tested
git -C aiter-fixed-tested checkout --detach \
  702aacd62c8a6e2fbbc260da338c510ab76f1b1c
git -C aiter-fixed-tested apply --index ../aiter-4474.patch
git -C aiter-fixed-tested apply --index ../aiter-4488.patch
git -C aiter-fixed-tested diff --cached --check

Save this as Dockerfile.repro:

# syntax=docker/dockerfile:1
FROM vllm/vllm-openai-rocm:kimi-k3@sha256:5aa7e626ff73672f5ca7aae46754570488c23d33ca1ac90756a1d2d1a3fe099b
COPY --from=aiter_source . /workspace/aiter
COPY --from=aiter_source aiter/ /usr/local/lib/python3.12/dist-packages/aiter/
COPY --from=aiter_source csrc/ /usr/local/lib/python3.12/dist-packages/aiter_meta/csrc/
COPY --from=aiter_source hsa/ /usr/local/lib/python3.12/dist-packages/aiter_meta/hsa/
COPY --from=source_manifests / /opt/repro-manifests/
WORKDIR /workspace/aiter
RUN cd /usr/local/lib/python3.12/dist-packages && \
    sha256sum --quiet --check /opt/repro-manifests/aiter.sha256 && \
    cd aiter_meta && \
    sha256sum --quiet --check /opt/repro-manifests/aiter-meta.sha256 && \
    cd /workspace/aiter && \
    sha256sum --quiet --check /opt/repro-manifests/op-tests.sha256 && \
    uv run --no-project --offline python -m compileall -q \
    aiter op_tests/test_mla.py

Build and run only the isolated child case:

mkdir -p manifests-4488
(cd aiter-fixed-tested && \
  find aiter -type f -print0 | sort -z | xargs -0 sha256sum) \
  >manifests-4488/aiter.sha256
(cd aiter-fixed-tested && \
  find csrc hsa -type f -print0 | sort -z | xargs -0 sha256sum) \
  >manifests-4488/aiter-meta.sha256
(cd aiter-fixed-tested && \
  find op_tests -type f -print0 | sort -z | xargs -0 sha256sum) \
  >manifests-4488/op-tests.sha256
docker build --file Dockerfile.repro \
  --build-context aiter_source=./aiter-fixed-tested \
  --build-context source_manifests=./manifests-4488 \
  --tag aiter-pr4488-large-kv-regression .
docker run --rm --ipc host --shm-size 16g \
  --device /dev/kfd --device /dev/dri \
  --group-add video --group-add render \
  --security-opt label=disable -e HIP_VISIBLE_DEVICES=0 \
  --entrypoint bash aiter-pr4488-large-kv-regression \
  -lc 'uv run --no-project --offline python op_tests/test_mla.py --_large-kv-child'

Observed on MI355X with #4474: both context lengths passed and the child exited
zero. On the unfixed source, the child exits via SIGABRT after the HSA memory
fault. The parent disables core dumps before starting that intentionally unsafe
control, and the large allocation runs only for the default direct invocation,
not for each of aiter_test.sh's five argument sweeps. Black, Ruff, and
git diff --check passed.

Dependency and duplication

#4474 owns the runtime fix. This PR contributes only the missing deterministic
boundary regression and should merge after, or together with, #4474. OpenAI
Codex assisted with testing and review; the author reviewed every changed line
and the final evidence.

[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 4488 --add-label <label>

[Copilot]

Pull request overview

This PR fixes a GPU memory fault in the Gluon MLA kernel by preventing int32 overflow in KV-cache row offset arithmetic when the KV cache exceeds ~2GB, ensuring correct address generation for large KV caches.

Changes:

• Moves KV offset computation inside the WITHIN_2GB branches to preserve the existing fast-path behavior.
• Widens the KV row index to int64 in the WITHIN_2GB == False path so row * stride does not wrap in int32 for large caches.
• Applies the same widening pattern across K_nope (slice 0/1) and optional K_pe loads in both prologue and loop/epilogue sections.

Suppressed comments (3)

aiter/ops/triton/gluon/mla_gluon.py:467

• global_load_to_shared in the >2GB branch is missing the same split_kv_end mask used on the buffer_load_to_shared path. For partial tiles this can read K_pe from invalid rows (or from garbage kv_loc_pe if the page-table load was masked).

            offs_k_pe = kv_loc_pe.to(gl.int64)[None, :] * stride_k_pe_bs + offs_d_kpe_1[:, None] + KV_PE_OFFSET
            gl.amd.cdna4.async_copy.global_load_to_shared(bufs_kpe.index(0), K_pe_cache + offs_k_pe)

aiter/ops/triton/gluon/mla_gluon.py:488

• global_load_to_shared in the >2GB branch omits the split_kv_end mask that exists on the buffer_load_to_shared path. When the slice-1 half-tile straddles split_kv_end, this can issue out-of-bounds reads.

        offs_k_c1 = kv_loc1.to(gl.int64)[None, :] * stride_kv_c_bs + offs_d_ckv_10[:, None]
        gl.amd.cdna4.async_copy.global_load_to_shared(bufs_kv1, Kv_c_cache + offs_k_c1)

aiter/ops/triton/gluon/mla_gluon.py:658

• In the epilogue K_pe load, the >2GB global_load_to_shared path is missing the split_kv_end mask that the buffer_load_to_shared path uses. This can read past the valid KV range for partial tiles.

                # legalizes). buffer_load needs int32, so only this branch.
                offs_k_pe = kv_loc_pe.to(gl.int64)[None, :] * stride_k_pe_bs + offs_d_kpe_1[:, None] + KV_PE_OFFSET
                gl.amd.cdna4.async_copy.global_load_to_shared(bufs_kpe.index(async_idx), K_pe_cache + offs_k_pe)

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[JohnQin-llm]

Closed as test coverage is not perf blocker and no review for a while