[Bug] ck_tile FMHA `mha_batch_prefill` paged-KV gather: 32-bit *address* overflow at high GPU base addresses

[zx3xyy]

ck_tile FMHA mha_batch_prefill paged-KV gather: 32-bit address overflow at high GPU base addresses (separate from #2517)

Description

mha_batch_prefill_func reads out of bounds (GPU memory access fault, or NaN
output when the OOB read hits mapped memory) when the paged KV cache lands at a
high GPU virtual address — even though the KV pool is small and its
element_space is far below INT32_MAX.

This is not the tensor_descriptor.hpp calculate_element_space_size_impl
overflow fixed by ROCm/rocm-libraries#6653. That fix addresses the descriptor size (it now casts
to long_index_t before the multiply and clamps the size to INT32_MAX). The
bug here is in the per-access address computation in the buffer-load paged-KV
gather path, which is still 32-bit. With #6653 applied, the original
all-zero repro no longer reproduces, but this one still faults.

Observed in production as NaN attention output on the last (global-attention)
layer of a large model whose KV cache sits at the top of GPU VA.

Minimal reproduction

element_space = num_pages * num_kv_heads * head_dim = 1,903,888 * 4 * 64 = 487,395,328 << INT32_MAX, so #6653's element_space path does not engage — this
isolates the address overflow.

import sys, torch
from aiter.ops.mha import mha_batch_prefill_func
dev = "cuda:0"
NUM_PAGES, NKH, HD, NQH, L = 1_903_888, 4, 64, 32, 287
pad_gib = float(sys.argv[1]) if len(sys.argv) > 1 else 0.0
# push the KV cache to a high VA (mimics a model with weights already resident)
pad = torch.empty(int(pad_gib*1024**3)//2, dtype=torch.bfloat16, device=dev) if pad_gib else None
k = torch.randn(NUM_PAGES, NKH, HD, dtype=torch.bfloat16, device=dev)
v = torch.randn(NUM_PAGES, NKH, HD, dtype=torch.bfloat16, device=dev)
q = torch.randn(L+1, NQH, HD, dtype=torch.bfloat16, device=dev)
cu  = torch.tensor([0, L], dtype=torch.int32, device=dev)
idx = torch.arange(L, dtype=torch.int32, device=dev)
o = mha_batch_prefill_func(q, k, v, cu, cu, idx, L, L, causal=True)
torch.cuda.synchronize()
print(f"pad={pad_gib}GiB k_addr=0x{k.data_ptr():x} max_abs={o.float().abs().max():.3e} nan={int(torch.isnan(o).sum())}")

The bug depends on the absolute VA the KV cache lands on, so sweep pad (each in
its own process, since a fault kills the process):

for p in 0 2 4 6 8 10 12 14 16; do python repro.py $p; done

Sample output (MI350X, ROCm 7.2, aiter @ CK with #6653):

pad=0.0GiB  k_addr=0x7f109e800000  max_abs=3.406e+00 nan=0
pad=2GiB    -> Memory access fault by GPU node-6 on address 0x7ee42cc10000
pad=4.0GiB  k_addr=0x7f969d000000  max_abs=2.844e+00 nan=0
pad=10GiB   -> Memory access fault by GPU node-6 on address 0x7ef12e610000
pad=16GiB   -> Memory access fault by GPU node-6 on address 0x7f712e610000

Root cause (analysis)

The fault address is deterministic relative to the KV base:

fault_addr == k_cache_addr - 0xffff0000   (== base - 2^32 + 0x10000)

i.e. a KV-tile address advance drops a carry out of bit 31 → a 32-bit address
wrap to base - 2^32 + tile_offset. When that wrapped address is unmapped it
faults; when it is mapped (the common case in a real engine) it returns garbage
→ NaN.

Controlled experiments narrowing it down:

• element_space < INT32_MAX (so #6653's clamp is inactive) — still faults ⇒
  not the descriptor-size path.
• A ~4 GiB pool whose accesses cross an in-pool 2³² byte boundary does not
  fault (HW 48-bit base + 32-bit voffset handles that correctly) ⇒ the overflow
  is driven by the absolute base VA, not by the in-pool offset magnitude.
• Reproduces only when the KV cache lands at certain high VAs (≈1/3 of the
  pad sweep) ⇒ address-dependent, consistent with a 32-bit truncation of an
  address that should be 64-bit.

This points at the buffer-load paged-KV gather in
projects/composablekernel/include/ck_tile/core/tensor/tile_scatter_gather.hpp:
the GLOBAL_LOAD_LDS path forms the gather address in 64-bit
(static_cast<long_index_t>(physical_page) * page_stride_elements_ + ...),
while the BUFFER_LOAD / async path merges the page offset into an
array<index_t> (int32) coordinate, so the address is formed in 32 bits. The
load-mode selector fmha_batch_prefill_select_kv_load_mode
(example/ck_tile/01_fmha/fmha_fwd.hpp) only routes to the 64-bit path when
kv_pool_bytes > INT32_MAX, which does not catch an address-driven overflow.

Suggested fix

1. Root fix: compute the gather address in long_index_t in the
   BUFFER_LOAD/async path of tile_scatter_gather.hpp::load(), mirroring the
   existing 64-bit GLOBAL_LOAD_LDS path (advance the bottom-tensor base by
   static_cast<long_index_t>(page_offset) * stride). Then BUFFER_LOAD is
   address-safe and the selector can choose it purely for performance.
2. Or, conservative selector fix: since the overflow is address-driven (not
   size-driven), the kv_pool_bytes > INT32_MAX predicate cannot guard it;
   route the paged small-page case to GLOBAL_LOAD_LDS unconditionally.

Environment

• AMD MI350X (gfx950), ROCm 7.2
• aiter with CK including #6653 (verified #6653's all-zero repro no longer
  reproduces here)

[zx3xyy]

cc @HaiShaw @asleepzzz @poyenc @Jeff-Huang

[kensclin]

Check K and V independently for overflow. Observed so far: high GPU base addresses only affect V; K has not faulted under the same conditions. Will run more experiments to confirm K's behavior / whether other cases can trigger it before deciding whether to keep the K check.