fix(ck): route scattered page_size=1 paged-KV batch-prefill to GLOBAL…

[mohbasit]

fix(ck): route scattered page_size=1 paged-KV batch-prefill to GLOBAL_LOAD_LDS

ISSUE ID : ROCm/aiter#3824

Target: ROCm/rocm-libraries (develop)
Files: projects/composablekernel/example/ck_tile/01_fmha/fmha_fwd.hpp,
projects/composablekernel/test/ck_tile/fmha/test_fmha_fwd.cpp
Additive follow-up to #9214.

Motivation

#9214 fixed a 32-bit SRD address overflow in the ck_tile FMHA mha_batch_prefill paged-KV gather by switching to GLOBAL_LOAD_LDS when base_lo32 + pool_bytes exceeds INT32_MAX. That fix removes the high-base-VA fault, but a second, distinct fault in the same page_block_size < kN0 BUFFER_LOAD arm remains for a scattered 1D paged KV layout (page_size=1 LINEAR + SGLANG_PAGE_TABLE_1D), as used by MiniMax-M3 under causal chunked prefill with a prefix cache (q_len != kv_len).

In that layout the per-page SRD voffset is physical_page * stride_page_block + within_page, where physical_page is read from a page table and indexes into the entire KV pool. #9214's overflow check estimates the maximum voffset as
num_total_pages * batch_stride * element_bytes, but for a scattered page table num_total_pages (this dispatch's page count) does not bound the physical page indices — they can point anywhere in the global pool. So the signed-int32
voffset can wrap on its own, independent of the pool's contiguous byte size and of the base address. #9214's check evaluates to "no overflow", keeps BUFFER_LOAD, the SRD reads a wrapped address, and the kernel takes a GPU
memory-access fault (Memory access fault by GPU node-N ... Reason: Unknown → coredump → Fatal Python error: Aborted).

Technical Details

Keep #9214's fast path and address-overflow check exactly as-is, and add a single guard for the scattered single-token-page case:

if(page_block_size >= kN0)
    return ck_tile::BlockAttentionKVCacheLoadModeEnum::BUFFER_LOAD;   // #9214 fast path

// Scattered 1D paged KV (page_size=1 LINEAR + SGLANG_PAGE_TABLE_1D): the page
// table holds arbitrary physical-page indices into the whole KV pool, so the
// per-page SRD voffset (physical_page * stride_page_block + within_page) is not
// bounded by this dispatch's num_total_pages and can wrap the signed int32 on
// its own, independent of base[31:0] + pool_bytes. The address-size check below
// cannot see this, so force the 64-bit-safe path for single-token pages.
if(page_block_size == 1)
    return ck_tile::BlockAttentionKVCacheLoadModeEnum::GLOBAL_LOAD_LDS;

// ... unchanged #9214 base_lo32 + pool_bytes > INT32_MAX check for K and V ...

• GLOBAL_LOAD_LDS forms the full 64-bit address via tile_scatter_gather and is always safe for the scattered case.
• The BUFFER_LOAD fast path for 1 < page_block_size < kN0 with in-bounds addresses is preserved — this change is strictly additive to fix: FMHA batch-prefill paged-KV 32-bit VA overflow at high GPU base addresses #9214 (+9 lines, no behavior change for any case fix: FMHA batch-prefill paged-KV 32-bit VA overflow at high GPU base addresses #9214 already handled).

Changed files:

• example/ck_tile/01_fmha/fmha_fwd.hpp — one extra page_block_size == 1 early return in fmha_batch_prefill_select_kv_load_mode.
• test/ck_tile/fmha/test_fmha_fwd.cpp — host-only unit test FmhaBatchPrefillKvLoadMode.SinglePageAlwaysGlobalLoad covering the new guard (extends the selector tests fix: FMHA batch-prefill paged-KV 32-bit VA overflow at high GPU base addresses #9214 added).

Test Plan

Full end-to-end serving repro on MI35x (gfx950), 4 GPUs, sglang with the aiter attention/MoE backends. Model: amd/MiniMax-M3-MXFP4, TP=4.

Server (image built with CK at develop HEAD incl. #9214, then this patch):

export SGLANG_USE_AITER=1
export SGLANG_OPT_USE_BF16_ROUTER_GEMM=0
python -m sglang.launch_server \
  --model-path amd/MiniMax-M3-MXFP4 --served-model-name MiniMaxAI/MiniMax-M3 \
  --host 0.0.0.0 --port 30040 \
  --tp-size 4 --trust-remote-code --attention-backend aiter --moe-runner-backend aiter \
  --reasoning-parser auto --tool-call-parser auto --page-size 128 \
  --max-running-requests 256 --max-queued-requests 128 --disable-custom-all-reduce --enable-metrics

Load (drives large mixed prefills with 2048-token shared prefix → prefix-cache hits at 128 concurrency, the exact conditions that manufacture the faulting batch-prefill dispatch):

python3 InferenceX/utils/bench_serving/benchmark_serving.py \
  --model amd/MiniMax-M3-MXFP4 --backend vllm --base-url http://127.0.0.1:30040 \
  --dataset-name random --random-input-len 3500 --random-output-len 1024 \
  --random-prefix-len 2048 --random-range-ratio 0.3 --num-prompts 3000 \
  --max-concurrency 128 --request-rate 12 --ignore-eos --trust-remote-code --num-warmups 0

Test Result

Build	Result
CK develop (with #9214), unpatched	Crashes — Memory access fault at ~2106/3000 requests
CK develop + this change	3000/3000 successful, 0 faults, server healthy throughout

Right before the fault (unpatched) the batch state is #running-req: 128,
#cached-token: 2048 (prefix-cache hits) with large mixed prefills — all four TP ranks (KFD nodes) fault. With this change the identical run completes cleanly at ~22.1k tok/s total throughput and the BUFFER_LOAD fast path is retained for the
page_block_size >= kN0 and 1 < page_block_size < kN0 in-bounds cases.

Submission Checklist

• Looked over the contributing guidelines at
  https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests
• Additive, no behavior change for cases already handled by fix: FMHA batch-prefill paged-KV 32-bit VA overflow at high GPU base addresses #9214
• Added host-only unit test for the new selector branch
• Validated end-to-end (MiniMax-M3-MXFP4, TP=4, gfx950)

[therock-pr-bot]

✅ All Checks Passed — Ready for Review

Check	Status	Details
📝 PR Description	✅ Pass	—
⛔ Forbidden Files	✅ Pass	—
🧪 Unit Test	✅ Pass	—
🔎 pre-commit	✅ Pass	—
🚫 Draft PR	🔜 To Be Enabled	—
🚩 Feature Flag	🔜 To Be Enabled	—
📊 Code Coverage	🔜 To Be Enabled	—
🤖 therock-pr-bot	✅ Pass	—

🎉 All checks passed! This PR is ready for review.

📖 Need help? See the Policy FAQ for details on every check and how to fix failures.

🙋 Wish to Override Policy?

[therock-pr-bot]

🎉 All checks passed! This PR is ready for review.

[Jeff-Huang]

Thanks for the fix! Quick question on num_total_pages semantics: our assumption is that it's the full KV-cache pool size so any valid page index should already be bounded by it regardless of scatter order, and #9214's INT32_MAX check should cover this case.

Could you help confirm whether, in the SGLang integration, k/v is always the full pool (num_total_pages reflects the total page count of the k/v pool) or whether CK can sometimes receive a smaller view than the index space kv_page_indices references? Thanks!

[mohbasit]

@Jeff-Huang upon probing more in depth, I found your assumption to be true

On num_total_pages semantics (your question): I traced the SGLang → aiter → CK path and instrumented mha_batch_prefill at runtime. Confirmed:

• SGLang passes the full per-layer KV pool (token_to_kv_pool.get_kv_buffer(layer_id)) as k, so aiter sets num_total_pages = k.size(0) (the whole pool) and batch_stride_k = k.stride(0).
• The page indices are bounded by it: e.g. num_total_pages = 3,613,312, observed max(kv_page_indices) = 256–313 ≪ 3.6M (max_idx < num_total_pages always true).

So CK does not receive a smaller view than the index space, and #9214's base_lo32 + pool_bytes > INT32_MAX is a valid (here, exact) bound. My PR's "num_total_pages doesn't bound the scattered indices" explanation is incorrect — apologies for that.

The real cause: even so, this MiniMax-M3 page_size=1 path still faults, and it is not the VA overflow #9214 models. I rebuilt CK at current develop (includes #9214, mirror 22ee914), force-cleared all stale mha_batch_prefill/mha_varlen .so + build dirs, verified a fresh build, and re-ran the load. It reproduces a GPU Memory access fault at ~3/3000 requests. At the faulting dispatch the selector had chosen BUFFER_LOAD on an in-bounds case that #9214 correctly deems safe:

page_size=1  kshape=(3613312,1,128)  stride0=128  elem_bytes=2
pool_bytes=0.861GiB  base_lo32=0x10a00000  lo32+pool=1.121GiB  over_INT32MAX=False
max_page_idx=256  ->  actual voffset ~ 256*128*2 = 64 KB  (no 32-bit wrap possible)

Forcing that arm to GLOBAL_LOAD_LDS (this PR) makes the identical run complete 3000/3000 with zero faults; leaving it on BUFFER_LOAD faults regardless of base address or pool size. So the defect is in the page_block_size == 1 BUFFER_LOAD gather arm itself on gfx950 (in the same spirit as the gfx950-incompatible batch-prefill tile #9214 had to drop), independent of the address-overflow check.

Net: the single-token-page BUFFER_LOAD gather faults on gfx950 even for in-bounds addresses #9214 considers safe, so we conservatively route page_block_size == 1 to the 64-bit-safe GLOBAL_LOAD_LDS path. #9214's fast path for 1 < page_block_size < kN0 is untouched.

[Jeff-Huang]

Thanks for the detailed investigation. One thing we still can't explain is why the BUFFER_LOAD arm faults while GLOBAL_LOAD_LDS does not, given the addresses are provably in range. With voffset ~64 KB and lo32+pool = 1.121 GiB, both paths should resolve to the same physical addresses, so the fault mechanism isn't clear to us yet. Without it, we're a little concerned this may be routing around a defect rather than fixing it, and that the same defect could still be reachable through other configurations.

Could you establish the root cause before this is merged? The natural starting points:

1. The faulting virtual address
2. The physical page id at the faulting access, and the page-id range for that dispatch

[geyyer]

LGTM! @kensclin, could you confirm this PR is safe to merge?

[mohbasit]

@Jeff-Huang
With HIP_LAUNCH_BLOCKING=1 (to pin the fault to its dispatch) and per-dispatch logging of the K/V base pointers + page-id range:

Faulting VA: 0x7f8937388000 (node-6 / dev0).
That dispatch's V base v_ptr = 0x7f8937400000; so the fault is v_ptr − 0x78000 = 491,520 bytes (480 KiB, exactly 1920 pages) below the V pool base — an OOB read before the allocation.

Page-id range for the dispatch: [256, 359939], all < num_total_pages = 3,613,312 → indices are valid; the bad address is from offset arithmetic, not an out-of-range page.

The dispatch is a causal chunked prefill with #cached-token = 2048 (prefix cache), q_len != kv_len.

So BUFFER_LOAD and GLOBAL_LOAD_LDS do not resolve to the same address here — the BUFFER_LOAD path computes a negative V offset. GLOBAL_LOAD_LDS forms the correct per-element 64-bit address from the page table and stays in bounds, which is why routing to it fixes the fault. This is fixing a real defect, not masking one.

Notably the underflow is exactly 2048 − 128 pages = prefix_len − page_size, which strongly implicates the V-gather deriving its KV base from the query offset (only valid when q_len == kv_len); with a prefix cache it underflows by the prefix length.

[Jeff-Huang]

@mohbasit
Thanks for tracking down the faulting VA — the prefix_len − page_size match is a very clean signal.

What's worth chasing is why the BUFFER_LOAD path lands below v_ptr at all. For page_size == 1, within_page is always 0 (kInPageOffsetMask == 0), the page ids you measured are all positive ([256, 359939]),
and the gather dim of coord is zeroed. All three terms should be non-negative, so base + offset shouldn't be able to go below base.

Would it be possible to dump the actual physical_page and coord_offset values at the faulting lane (rather than the range of kv_page_indices)? That should show directly which term goes negative.

One more question: was the patched run validated for numerical correctness against a reference, or only for absence of faults? Both paths share the same load_physical_pages lookup, so if the underlying issue
is a wrong page id or offset, GLOBAL_LOAD_LDS would read the same wrong location without faulting.

[mohbasit]

@Jeff-Huang
Answering your second question, the run was mainly tested for absence of faults as the prior one couldnt get to do the whole serving because of crashing.

As for the first query,
The mha_batch_prefill path uses the async pipeline (block_fmha_batch_prefill_pipeline_qr_ks_vs_async.hpp), which builds the V global address through its own tile-window / tile_scatter_gather coordinate math (the async direct-to-LDS buffer_load), not through get_block_ptr. I feel the negative/OOB term lives there — i.e., in the gather-dim coordinate or the tile-origin offset the pipeline feeds into the V buffer_load, (which evidently is not zeroed on this path under q_len != kv_len). That's consistent with the 1920 = prefix − kN0 signature pointing at the causal/prefix tile-origin computation.

physical_page — not negative/garbage. Neither the PageBlockNavigator path (block_index never < 0) nor load_physical_pages (global_token_idx never < 0) produces a bad page id.

within_page — is 0 for page_size=1 (kInPageOffsetMask == 0)

So the negative/OOB term is not in either page-index computation.

It's formed inside tile_scatter_gather's address assembly — the SRD base/voffset construction from physical_pages_ + page_stride_elements_ + coord, i.e., the coord/SRD-base handling in include/ck_tile/core/tensor/tile_scatter_gather.hpp, not the page lookup

I am not sure how to exactly pin point to the instruction which is causing this though. The OOB is sensitive to register allocation / instruction scheduling / SRD-register placement, that is why it seems latent in production.

[Jeff-Huang]

@mohbasit
Thanks for the follow-up.

Would it be possible to run a numerical correctness check on the case that used to crash? Our concern is just this: we still don't have a confirmed root cause for the negative offset in the BUFFER_LOAD path. Until we know what makes that address computation go wrong, we can't be sure the same underlying issue doesn't also affect GLOBAL_LOAD_LDS.

So at minimum it would be good to confirm that the GLOBAL_LOAD_LDS path produces correct values for this case.

[Jeff-Huang]

Approving as a workaround for the production crash. Root cause of the negative offset in the BUFFER_LOAD path is still open — tracking separately.

[AmosLewis]

Gardener triage: this failure is infra, not your change.

The only real failure is Test RPM Install - sles16
(job, 42s);
Multi-Arch CI Summary is red only because it aggregates that job. Everything else on the run is
green, including rhel8 / rhel10 / ubuntu2404 install, all build stages, PyTorch, sanity and
pre-commit.

Every RPM 404s while the repo metadata downloads fine, so it is a URL problem, not missing
artifacts. The dev version 10.1.0.dev0+<sha> puts a + in each RPM filename; on the same S3
prefix a literal + returns 404 and %2B returns 200. dnf escapes it and passes, zypper sends
the literal + and fails on all 367 packages. The sles16 lane has been red on every unrelated PR I
sampled, going back to at least June 20.

Your change only touches fmha_fwd.hpp and test_fmha_fwd.cpp, which cannot affect RPM download.

Tracking issue: ROCm/TheRock#7161

Happy to override-merge on that basis whenever you are ready — just ping me once the other CI you
are waiting on has landed.

[geyyer]

@AmosLewis, thank you, our CI passed, could you override-merge this PR?