Description
mha_batch_prefill_func returns all-zero outputs when the paged KV cache has more than ~8,388,608 pages (with page_block_size=1). The root cause is an int32 multiplication overflow in CK's tensor_descriptor.hpp that corrupts buffer_size_, causing AMD GPU buffer loads to silently return zeros for valid memory.
Minimal Reproduction
import torch
from aiter.ops.mha import mha_batch_prefill_func
NUM_PAGES = 8_401_239 # > 8,388,608 threshold
KV_START = 16_435 # > 12,630 (beyond overflowed SRD range)
k = torch.randn(NUM_PAGES, 8, 64, dtype=torch.bfloat16, device="cuda")
v = torch.randn(NUM_PAGES, 8, 64, dtype=torch.bfloat16, device="cuda")
q = torch.randn(512, 16, 64, dtype=torch.bfloat16, device="cuda")
cu = torch.arange(0, 513, 128, dtype=torch.int32, device="cuda")
idx = torch.arange(KV_START, KV_START + 512, dtype=torch.int32, device="cuda")
o = mha_batch_prefill_func(q, k, v, cu, cu, idx, 128, 128, causal=True)
print(f"max_abs={o.float().abs().max():.4e}") # 0.0000e+00 — BUG!
# Change NUM_PAGES to 5,000,000 → works fine (output ≈ 3.34e+00)Platform: AMD MI300X, ROCm, aiter with CK eb033ef20 (also verified on CK tag rocm-7.2.1)
Root Cause
In CK's tensor_descriptor.hpp (line ~244 in calculate_element_space_size_impl):
auto acc_new = acc_old + (lengths[i] - number<1>{}) * strides[i];
// ^^^^^^^^^ int32 × ^^^^^^^^^ int32 → overflow!With num_pages=8,401,239 and stride=512 (= num_kv_heads × head_dim = 8 × 64):
(8,401,238) × 512 = 4,301,433,856 — exceeds uint32_max (4,294,967,296)- Wraps to
6,466,560, making buffer_size_ = 6,466,624
- SRD range =
6,466,624 × 2 bytes = 12,933,248 bytes (12 MB instead of 8.6 GB)
- Any buffer load at byte offset > 12 MB → AMD GPU returns hardware zero (silent OOB)
Threshold: num_pages > 2^32 / stride = 2^32 / 512 = 8,388,608
Controlled Experiment
| num_pages |
kv_start |
output |
Why |
| 5,000,000 |
16,435 |
nonzero ✓ |
elem_space = 2.56B < 2^32, no overflow |
| 8,401,239 |
100 |
nonzero ✓ |
Overflow, but page 100 within 12MB SRD range |
| 8,401,239 |
16,435 |
ZERO ✗ |
Overflow + page 16K (16MB) exceeds 12MB SRD |
Both conditions needed: large pool (> 8.4M pages) AND page indices beyond ~12,630.
Impact
This affects any model using paged KV cache with page_block_size=1 when total KV pool exceeds ~8.4M tokens (e.g., ~17 GB of KV cache for num_kv_heads=8, head_dim=64, bf16). In our case, this manifests as garbage outputs during multi-node GRPO training with sglang rollout engine for a Mamba/Attention hybrid model (lfm2).
Workaround
Hardcode buffer_size_ to INT32_MAX in CK's buffer_view.hpp:
// Before:
buffer_size_{buffer_size / PackedSize}
// After:
buffer_size_{static_cast<BufferSizeType>(0x7FFFFFFF)}Suggested Proper Fix
Cast to long_index_t before multiplication in tensor_descriptor.hpp:
auto acc_new = acc_old + static_cast<long_index_t>(lengths[i] - number<1>{})
* static_cast<long_index_t>(strides[i]);Plus cap the SRD range field to 0xFFFFFFFF in buffer_view.hpp:init_raw() and amd_buffer_addressing.hpp non-raw load paths to handle uint32_t truncation for byte sizes > 4 GB.
Description
mha_batch_prefill_funcreturns all-zero outputs when the paged KV cache has more than ~8,388,608 pages (withpage_block_size=1). The root cause is anint32multiplication overflow in CK'stensor_descriptor.hppthat corruptsbuffer_size_, causing AMD GPU buffer loads to silently return zeros for valid memory.Minimal Reproduction
Platform: AMD MI300X, ROCm, aiter with CK
eb033ef20(also verified on CK tagrocm-7.2.1)Root Cause
In CK's
tensor_descriptor.hpp(line ~244 incalculate_element_space_size_impl):With
num_pages=8,401,239andstride=512(=num_kv_heads × head_dim = 8 × 64):(8,401,238) × 512 = 4,301,433,856— exceedsuint32_max(4,294,967,296)6,466,560, makingbuffer_size_ = 6,466,6246,466,624 × 2 bytes = 12,933,248 bytes(12 MB instead of 8.6 GB)Threshold:
num_pages > 2^32 / stride = 2^32 / 512 = 8,388,608Controlled Experiment
elem_space = 2.56B < 2^32, no overflowBoth conditions needed: large pool (> 8.4M pages) AND page indices beyond ~12,630.
Impact
This affects any model using paged KV cache with
page_block_size=1when total KV pool exceeds ~8.4M tokens (e.g., ~17 GB of KV cache fornum_kv_heads=8, head_dim=64, bf16). In our case, this manifests as garbage outputs during multi-node GRPO training with sglang rollout engine for a Mamba/Attention hybrid model (lfm2).Workaround
Hardcode
buffer_size_toINT32_MAXin CK'sbuffer_view.hpp:Suggested Proper Fix
Cast to
long_index_tbefore multiplication intensor_descriptor.hpp:Plus cap the SRD
rangefield to0xFFFFFFFFinbuffer_view.hpp:init_raw()andamd_buffer_addressing.hppnon-raw load paths to handleuint32_ttruncation for byte sizes > 4 GB.