gemm_a16w16_asm silently returns wrong results past 2**31 elements (32-bit buffer addressing)

[JohnQinAMD]

Problem

gemm_a16w16_asm silently returns wrong results once the activation exceeds
2**31 elements. There is no error, no warning, and no fallback — the call
succeeds and the output is garbage. On gfx950 with bf16 and K=12288 the
boundary is at M = 174762, and it is sharp:

M	M*K / 2**31	max abs err	ASM mean|.|	torch mean|.|	correct
160,000	0.916	2.0	88.44	88.44	✅
174,000	0.996	0.0	88.44	88.44	✅
175,000	1.001	776.0	0.24	88.45	❌
180,000	1.030	890.0	5.15	88.45	❌
240,000	1.373	1006.0	48.09	88.45	❌
360,000	2.060	948.0	7.72	88.44	❌

This matters because the wrong path is also faster — it measures 1.11–1.45×
against hipBLASLt at these shapes — so it looks like a win in any benchmark that
does not check numerics. It is not a win: the kernel wraps its addresses and
re-reads a cached window instead of streaming the whole activation, so it is
fast because it is skipping work. Correcting it (see below) puts it back at
1.000–1.005× of hipBLASLt.

Any caller doing long-sequence work — large batches, long-context transformers,
video models — reaches these row counts routinely and gets no signal that the
result is wrong.

Reproduction

Needs only torch and aiter. Tested on gfx950, ROCm 7.2.4, torch
2.10.0+rocm7.2.4, aiter built with GPU_ARCHS=gfx950.

import torch
from aiter import gemm_a16w16_asm

K, N = 12288, 8192
LIMIT = 2**31

print(f"K={K} N={N}  bf16, so the 2**31-element boundary is at M={LIMIT // K}\n")
print(f"{'M':>9} {'M*K/2**31':>11} {'max abs err':>13} {'ASM mean|.|':>12} "
      f"{'torch mean|.|':>14} {'correct':>8}")

for m in (160_000, 174_000, 175_000, 180_000, 240_000, 360_000):
    a = torch.randn(m, K, device="cuda", dtype=torch.bfloat16)
    w = torch.randn(N, K, device="cuda", dtype=torch.bfloat16)
    out = torch.empty(m, N, device="cuda", dtype=torch.bfloat16)

    gemm_a16w16_asm(a, w, out)
    torch.cuda.synchronize()
    ref = torch.nn.functional.linear(a, w)

    err = (out.float() - ref.float()).abs().max().item()
    ok = torch.allclose(out.float(), ref.float(), rtol=3e-2, atol=3e-2)
    print(f"{m:>9,} {m * K / LIMIT:>11.3f} {err:>13.1f} "
          f"{out.float().abs().mean():>12.2f} {ref.float().abs().mean():>14.2f} {str(ok):>8}")

    del a, w, out, ref
    torch.cuda.empty_cache()

Note for anyone writing a regression test: compare the whole output. An
earlier check here compared only the first 4096 rows and reported the broken
kernel as correct — the corruption lives at high offsets, which is exactly where
a truncated check does not look.

Root cause

The kernel addresses memory exclusively through buffer instructions, and buffer
addressing on CDNA is 32-bit. Disassembling
aiter_meta/hsa/gfx950/bf16gemm/bf16gemm_bf16_tn_256x256.co:

$ llvm-objdump -d --mcpu=gfx950 bf16gemm_bf16_tn_256x256.co

buffer_load    96      buffer_load_dwordx4 v210, s[16:19], 0 offen lds
buffer_store   96
global_load     0      <- none
v_mul_lo_u32    8      v_mul_lo_u32 v200, s41, v5    <- 32-bit address multiply
v_mul_hi_u32    4
v_mad_u64_u32   0      <- no 64-bit address arithmetic anywhere
v_mad_i64_i32   0
v_lshl_add_u64  0

With offen, the per-lane offset is a 32-bit unsigned byte offset into the
V# in s[16:19]. bf16 is 2 bytes, so 2**31 elements is exactly 2**32 bytes —
the wrap point. That is where the boundary in the table above comes from, and it
is why the threshold scales with element size rather than being a fixed row
count.

This is not a missing v_mad_u64_u32. Widening the multiply does not help,
because the buffer voffset register is 32-bit by ISA definition and num_records
in the descriptor is 32-bit as well. A single descriptor cannot address more
than 4 GB, so the addressing scheme itself has to change.

(The s41/s42 operands are read as strides and v5 as a row index from
context; the disassembly is stripped. The conclusion — all buffer addressing,
zero 64-bit address arithmetic — does not depend on that reading.)

How to fix

Recommended: re-base the buffer descriptor per workgroup. Instead of one
descriptor spanning the whole matrix with a large voffset, fold the workgroup's
row-block offset into the descriptor's 48-bit base at kernel entry:

s_add_u32   s16, s16, s_tile_byte_off_lo
s_addc_u32  s17, s17, s_tile_byte_off_hi   ; carry into the high half

Each lane's voffset is then relative to its own tile and stays small. Cost is two
scalar instructions per workgroup. This must be done for all three descriptors
— A, B and D. The output overflows independently of the input: at K=6144,
N=8192, M=360360, M*K is under the limit while M*N = 2.86e9 is over, so
fixing only the activation still returns wrong results.

Alternative: global_load/global_store, which take a 64-bit vaddr. Costs a
VGPR pair per address and loses the hardware bounds check num_records provides.
It is also a larger change than it looks, because this kernel uses direct-to-LDS
(buffer_load_dwordx4 ... lds).

Caller-side mitigation, until the kernel is fixed: either refuse the shape,
or tile along M. A slice of a contiguous [M, K] tensor is itself contiguous and
its data_ptr already points at the first row of the slice, so a[i:j] restarts
the kernel's offsets at zero:

def chunked_asm(a, w, out, limit=2**31):
    m, k = a.shape
    n = w.shape[0]
    rows = limit // max(k, n)      # the output overflows too -- bound by both
    rows -= rows % 256             # keep chunks on the kernel's tile boundary
    for i in range(0, m, rows):
        j = min(i + rows, m)
        gemm_a16w16_asm(a[i:j], w, out[i:j])
    return out

Verified to restore correctness. It does not recover the apparent speed, for the
reason in the first section — that speed was never real.

Whatever the fix, a guard in the host wrapper would be worth having on its
own. csrc/py_itfs_cu/asm_gemm_a16w16.cu already validates K % 64 == 0 and the
N/tileN relationship via AITER_CHECK; an equivalent check on
M * max(K, N) * sizeof(dtype) against 2**32 would turn a silent wrong answer
into an error at the call site.