[rocm-libraries] ROCm/rocm-libraries#9214 (commit 04892a3)

fix: FMHA batch-prefill paged-KV 32-bit VA overflow at high
 GPU base addresses (#9214)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## Motivation

Fixes ROCm/aiter#3824 — a 32-bit SRD address overflow in the ck_tile
FMHA mha_batch_prefill paged-KV gather path that causes wrong-page reads
(garbage output / memory faults) when K or V is allocated at a high GPU
virtual address.

When page_block_size < kN0, the batch-prefill pipeline gathers pages
with buffer_load through a single SRD whose base spans the whole K (or
V) pool. The hardware computes the effective address as:
```c++
effective_addr = (base[63:32] << 32) | ((base[31:0] + voffset) & 0xFFFFFFFF)
```

The carry out of bit 31 is silently dropped, so the kernel reads the
wrong physical page whenever base[31:0] + max_voffset overflows the low
32 bits. This can trigger even with a sub-2GB pool if the allocator
places the tensor high in the address space (e.g. base_lo32 ≈ 3.7GB +
pool ≈ 0.9GB → 4.6GB). The fault is latent/probabilistic: it depends on
the random base VA from the HIP allocator and whether the wrapped
address lands on unmapped memory.

## Technical Details

The runtime selector fmha_batch_prefill_select_kv_load_mode now receives
both K and V base pointers and their separate page strides, and decides
BUFFER_LOAD vs GLOBAL_LOAD_LDS per tensor:

- Fast path unchanged: if page_block_size >= kN0, the SRD is rebased per
page, so no wide voffset is needed → always BUFFER_LOAD.
- Otherwise, compute base_lo32 + pool_bytes independently for K and V
(they are allocated separately and may span differently), and if either
would overflow, return GLOBAL_LOAD_LDS.

Threshold is INT32_MAX (2GB), not 0xFFFFFFFF (4GB). Although the
hardware masks at 0xFFFFFFFF, the voffset value flows through
ck_tile::index_t (= std::int32_t, signed) before reaching the SRD. Once
base_lo32 + pool_bytes exceeds INT32_MAX, the signed intermediate has
already wrapped negative, so the safe bound is 2GB. Using the 4GB bound
leaves a (2GB, 4GB] "danger band" that still faults — this PR closes
that band.

  Changed files (composable_kernel):
- example/ck_tile/01_fmha/fmha_fwd.hpp — selector takes K+V pointers and
separate strides; per-tensor INT32_MAX overflow check; added <cstdint>
for INT32_MAX. Optional env-gated (CK_KV_LOAD_DEBUG) diagnostic that
prints the chosen mode and per-tensor band classification.
- example/ck_tile/01_fmha/codegen/ops/fmha_batch_prefill.py — dispatch
predicate passes a.k_ptr, a.v_ptr, a.batch_stride_k, a.batch_stride_v,
and per-dtype element bytes; removed a gfx950-incompatible d256
batch-prefill tile (bn0=32) whose async K-load descriptor computes
NumIssues = 0 (empty tensor → space_filling_curve static_assert
failure); the remaining bn0=128 d256 tile covers all CU counts.

## Test Plan

### 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.
```python
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):
```shell
for p in 0 2 4 6 8 10 12 14 16; do python repro.py $p; done
```

## Test Result

- Fixed build (2GB / INT32_MAX): 0 faults across the full page-size
sweep and repeated stress runs; all high-base cases correctly route to
GLOBAL_LOAD_LDS, low-base cases stay on BUFFER_LOAD.

## Submission Checklist

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

  ---