Add caller-owned output ops gemm_a8w8_bpreshuffle_out and gemm_a4w4_out

aiter/ops/gemm_op_a4w4.py

@@ -189,7 +189,6 @@ def gemm_a4w4(
         )
     m = A.numel() // A.shape[-1]
     n = B.shape[0]
-    k = A.shape[-1] * 2
     gfx_arch = get_gfx()
     if gfx_arch in ["gfx1250"]:
         require_gfx1250_asm("gemm_a4w4")
@@ -208,6 +207,25 @@ def gemm_a4w4(
         )
         return out.view(*A.shape[:-1], out.shape[-1])
     out = torch.empty(((m + 31) // 32 * 32, n), dtype=dtype, device=A.device)
+    return _gemm_a4w4_into(A, B, A_scale, B_scale, out, bias, alpha, beta, bpreshuffle)
+
+
+def _gemm_a4w4_into(
+    A: Tensor,
+    B: Tensor,
+    A_scale: Tensor,
+    B_scale: Tensor,
+    out: Tensor,  # [ceil(M / 32) * 32, N]
+    bias: Tensor | None,
+    alpha: float | None,
+    beta: float | None,
+    bpreshuffle: bool | None,
+) -> Tensor:
+    """Run the gfx950 A4W4 kernels into ``out``; returns the first M rows."""
+    m = A.numel() // A.shape[-1]
+    n = B.shape[0]
+    k = A.shape[-1] * 2
+    gfx_arch = get_gfx()
     if gfx_arch in ["gfx942"]:
         raise RuntimeError(
             f"A4W4 GEMM kernel is not supported on gfx942, but got {gfx_arch}!"
@@ -253,6 +271,39 @@ def gemm_a4w4(
     return out[:m].view(*A.shape[:-1], n)
 
 
+# Only ``out`` is written, so it is the only argument declared mutable. Returning
+# nothing keeps the op free of outputs that alias an input under torch.compile.
+@torch_compile_guard(mutates_args=["out"], gen_fake=lambda *args, **kwargs: None)
+def gemm_a4w4_out(
+    out: Tensor,  # [ceil(M / 32) * 32, N] caller-owned output
+    A: Tensor,  # A:[M, K/2] f4x2
+    B: Tensor,  # B:[N, K/2] f4x2
+    A_scale: Tensor,  # A_scale:[M, K/32] e8m0 padded
+    B_scale: Tensor,  # B_scale:[N, K/32] e8m0 padded
+    bias: Tensor | None = None,  # bias:[1, N] f32
+    alpha: float | None = 1.0,
+    beta: float | None = 0.0,
+    bpreshuffle: bool | None = True,
+) -> None:
+    """gemm_a4w4 into a caller-owned buffer. The kernels write whole 32-row
+    tiles, so ``out`` is [ceil(M / 32) * 32, N] and the result is out[:M].
+    Output dtype is ``out.dtype``. gfx950 MXFP4 only."""
+    if get_gfx() in ["gfx1250"]:
+        raise NotImplementedError("gemm_a4w4_out is not supported on gfx1250")
+    m = A.numel() // A.shape[-1]
+    n = B.shape[0]
+    m_padded = (m + 31) // 32 * 32
+    if out.dtype not in (dtypes.bf16, dtypes.fp16):
+        raise ValueError(f"out dtype must be bf16 or fp16, got {out.dtype}")
+    if out.shape != (m_padded, n) or out.device != A.device or not out.is_contiguous():
+        raise ValueError(
+            f"out must be a contiguous tensor of shape {(m_padded, n)} "
+            f"(M padded to a multiple of 32) on {A.device}; got shape "
+            f"{tuple(out.shape)}, stride {tuple(out.stride())}, device {out.device}"
+        )
+    _gemm_a4w4_into(A, B, A_scale, B_scale, out, bias, alpha, beta, bpreshuffle)
+
+
 @compile_ops(
     "module_gemm_a4w4_asm",
     fc_name="gemm_a4w4_asm",

aiter/ops/gemm_op_a8w8.py

@@ -777,6 +777,18 @@ def gemm_a8w8_bpreshuffle(
     assert WQ.dtype == dtypes.fp8, "gemm_a8w8_bpreshuffle only support fp8 now"
     assert bias is None, "gemm_a8w8_bpreshuffle does not support bias now"
     Y = torch.empty(m, n, dtype=dtype, device=XQ.device)
+    return _gemm_a8w8_bpreshuffle_into(XQ, WQ, x_scale, w_scale, Y)
+
+
+def _gemm_a8w8_bpreshuffle_into(
+    XQ: Tensor, WQ: Tensor, x_scale: Tensor, w_scale: Tensor, Y: Tensor
+) -> Tensor:
+    """Run the tuned (or default) bpreshuffle kernel into the dense [M, N] ``Y``."""
+    m = XQ.shape[0]
+    n = WQ.shape[0]
+    k = XQ.shape[-1]
+    w_k = WQ.shape[-1]
+    dtype = Y.dtype
 
     # CKTile only supports bf16 dtype
     config = get_GEMM_config_with_quant_type(
@@ -838,6 +850,38 @@ def gemm_a8w8_bpreshuffle(
         ) from e
 
 
+# Only ``out`` is written, so it is the only argument declared mutable. Returning
+# nothing keeps the op free of outputs that alias an input under torch.compile.
+@torch_compile_guard(mutates_args=["out"], gen_fake=lambda *args, **kwargs: None)
+def gemm_a8w8_bpreshuffle_out(
+    out: Tensor,  # [M, N] caller-owned output, bf16 or fp16
+    XQ: Tensor,
+    WQ: Tensor,
+    x_scale: Tensor,
+    w_scale: Tensor,
+) -> None:
+    """gemm_a8w8_bpreshuffle into a caller-owned contiguous [M, N] buffer.
+    Output dtype is ``out.dtype``. CK and CKTile write a dense row-major [M, N],
+    so a strided ``out`` (e.g. a column slice of a wider buffer) is rejected."""
+    m = XQ.shape[0]
+    n = WQ.shape[0]
+    if out.dtype not in (dtypes.bf16, dtypes.fp16):
+        raise ValueError(f"out dtype must be bf16 or fp16, got {out.dtype}")
+    if out.shape != (m, n) or out.device != XQ.device or not out.is_contiguous():
+        raise ValueError(
+            f"out must be a contiguous tensor of shape {(m, n)} on {XQ.device}; "
+            f"got shape {tuple(out.shape)}, stride {tuple(out.stride())}, "
+            f"device {out.device}"
+        )
+    if WQ.shape[-1] < XQ.shape[-1]:
+        raise RuntimeError(
+            f"gemm_a8w8_bpreshuffle requires WQ K >= XQ K, got WQ K={WQ.shape[-1]}, "
+            f"XQ K={XQ.shape[-1]}"
+        )
+    assert WQ.dtype == dtypes.fp8, "gemm_a8w8_bpreshuffle only support fp8 now"
+    _gemm_a8w8_bpreshuffle_into(XQ, WQ, x_scale, w_scale, out)
+
+
 # M at or above which the triton kernel beats the untuned CK fallback, which
 # runs one fixed tile shape at every M. Measured per arch; do not extrapolate.
 _BLOCKSCALE_TRITON_FALLBACK_MIN_M = {"gfx942": 2048, "gfx950": 384}

op_tests/test_gemm_out.py

@@ -0,0 +1,202 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
+"""Contract of gemm_a8w8_bpreshuffle_out and gemm_a4w4_out: the result lands in
+the caller-owned ``out``, matches the allocating op bit for bit, only ``out`` is
+declared mutable, and a buffer the kernels cannot write densely is rejected."""
+
+import pytest
+import torch
+
+import aiter
+from aiter import dtypes
+from aiter.ops import gemm_op_a4w4
+from aiter.ops.shuffle import shuffle_weight
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="GEMM out tests require a CUDA/HIP device"
+)
+
+
+def _gfx():
+    from aiter.jit.utils.chip_info import get_gfx_runtime
+
+    return get_gfx_runtime()
+
+
+def _mutated_args(op):
+    schema = op.default._schema
+    return [
+        a.name
+        for a in schema.arguments
+        if a.alias_info is not None and a.alias_info.is_write
+    ]
+
+
+def _a8w8_inputs(m, n, k):
+    torch.manual_seed(0)
+    x = torch.randn((m, k), dtype=dtypes.bf16, device="cuda")
+    w = torch.randn((n, k), dtype=dtypes.bf16, device="cuda")
+    xq, x_scale = aiter.pertoken_quant(x, quant_dtype=dtypes.fp8)
+    wq, w_scale = aiter.pertoken_quant(w, quant_dtype=dtypes.fp8)
+    return xq, shuffle_weight(wq, layout=(16, 16)), x_scale, w_scale
+
+
+def test_a8w8_bpreshuffle_out_declares_only_out_mutable():
+    assert _mutated_args(torch.ops.aiter.gemm_a8w8_bpreshuffle_out) == ["out"]
+
+
+@pytest.mark.parametrize(
+    "m,n,k,dtype",
+    [
+        (1, 256, 512, dtypes.bf16),
+        (37, 1280, 512, dtypes.bf16),
+        (256, 7168, 3584, dtypes.bf16),
+        (37, 1280, 512, dtypes.fp16),
+    ],
+)
+def test_a8w8_bpreshuffle_out_matches(m, n, k, dtype):
+    xq, wq, x_scale, w_scale = _a8w8_inputs(m, n, k)
+    ref = aiter.gemm_a8w8_bpreshuffle(xq, wq, x_scale, w_scale, None, dtype)
+    out = torch.full((m, n), float("nan"), dtype=dtype, device="cuda")
+
+    res = aiter.gemm_a8w8_bpreshuffle_out(out, xq, wq, x_scale, w_scale)
+
+    assert res is None
+    torch.testing.assert_close(out, ref, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("backend", ["aot_eager", "inductor"])
+def test_a8w8_bpreshuffle_out_under_torch_compile(backend):
+    m, n, k = 64, 1280, 512
+    xq, wq, x_scale, w_scale = _a8w8_inputs(m, n, k)
+    ref = aiter.gemm_a8w8_bpreshuffle(xq, wq, x_scale, w_scale, None, dtypes.bf16)
+
+    @torch.compile(backend=backend, fullgraph=True)
+    def fn(out, xq, wq, x_scale, w_scale):
+        aiter.gemm_a8w8_bpreshuffle_out(out, xq, wq, x_scale, w_scale)
+        return out * 1
+
+    torch._dynamo.reset()
+    out = torch.full((m, n), float("nan"), dtype=dtypes.bf16, device="cuda")
+    res = fn(out, xq, wq, x_scale, w_scale)
+    torch.testing.assert_close(out, ref, rtol=0, atol=0)
+    torch.testing.assert_close(res, ref, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize(
+    "make_out",
+    [
+        pytest.param(
+            lambda m, n: torch.empty((m + 1, n), dtype=dtypes.bf16), id="shape"
+        ),
+        pytest.param(lambda m, n: torch.empty((m, n), dtype=dtypes.fp32), id="dtype"),
+        pytest.param(
+            lambda m, n: torch.empty((m, 2 * n), dtype=dtypes.bf16)[:, :n],
+            id="column-slice",
+        ),
+        pytest.param(
+            lambda m, n: torch.empty((n, m), dtype=dtypes.bf16).t(), id="transposed"
+        ),
+    ],
+)
+def test_a8w8_bpreshuffle_out_rejected(make_out):
+    m, n, k = 32, 256, 512
+    xq, wq, x_scale, w_scale = _a8w8_inputs(m, n, k)
+    with torch.device("cuda"):
+        out = make_out(m, n)
+    base = out._base if out._base is not None else out
+    before = base.clone()
+
+    with pytest.raises(ValueError, match="out "):
+        aiter.gemm_a8w8_bpreshuffle_out(out, xq, wq, x_scale, w_scale)
+    assert torch.equal(base.view(torch.uint8), before.view(torch.uint8))
+
+
+def _a4w4_inputs(m, n, k):
+    torch.manual_seed(0)
+    quant = aiter.get_triton_quant(aiter.QuantType.per_1x32)
+    x = torch.randn((m, k), dtype=dtypes.bf16, device="cuda")
+    w = torch.randn((n, k), dtype=dtypes.bf16, device="cuda")
+    xq, x_scale = quant(x, shuffle=True)
+    wq, w_scale = quant(w, shuffle=True)
+    return xq, shuffle_weight(wq, layout=(16, 16)), x_scale, w_scale
+
+
+a4w4_only = pytest.mark.skipif(
+    torch.cuda.is_available() and _gfx() != "gfx950",
+    reason="gemm_a4w4_out runs the gfx950 kernels",
+)
+
+
+def test_a4w4_out_declares_only_out_mutable():
+    assert _mutated_args(torch.ops.aiter.gemm_a4w4_out) == ["out"]
+
+
+@a4w4_only
+@pytest.mark.parametrize("m", [1, 32, 45])
+def test_a4w4_padded_out_matches(m):
+    n, k = 512, 512
+    xq, wq, x_scale, w_scale = _a4w4_inputs(m, n, k)
+    ref = aiter.gemm_a4w4(xq, wq, x_scale, w_scale, bpreshuffle=True)
+    m_padded = (m + 31) // 32 * 32
+    out = torch.full((m_padded, n), float("nan"), dtype=dtypes.bf16, device="cuda")
+
+    res = aiter.gemm_a4w4_out(out, xq, wq, x_scale, w_scale)
+
+    assert res is None
+    torch.testing.assert_close(out[:m], ref, rtol=0, atol=0)
+
+
+@a4w4_only
+def test_a4w4_padded_out_under_torch_compile():
+    # aot_eager only: gemm_a4w4 already fails to lower under inductor on main
+    # ("auto_functionalized_v2 was not removed"), with or without out.
+    m, n, k = 45, 512, 512
+    xq, wq, x_scale, w_scale = _a4w4_inputs(m, n, k)
+    ref = aiter.gemm_a4w4(xq, wq, x_scale, w_scale, bpreshuffle=True)
+
+    @torch.compile(backend="aot_eager", fullgraph=True)
+    def fn(out, xq, wq, x_scale, w_scale):
+        aiter.gemm_a4w4_out(out, xq, wq, x_scale, w_scale)
+        return out[:m] * 1
+
+    torch._dynamo.reset()
+    out = torch.full((64, n), float("nan"), dtype=dtypes.bf16, device="cuda")
+    res = fn(out, xq, wq, x_scale, w_scale)
+    torch.testing.assert_close(out[:m], ref, rtol=0, atol=0)
+    torch.testing.assert_close(res, ref, rtol=0, atol=0)
+
+
+@a4w4_only
+@pytest.mark.parametrize(
+    "make_out",
+    [
+        pytest.param(
+            lambda m, n: torch.empty((m, n), dtype=dtypes.bf16), id="unpadded"
+        ),
+        pytest.param(lambda m, n: torch.empty((64, n), dtype=dtypes.fp32), id="dtype"),
+        pytest.param(
+            lambda m, n: torch.empty((64, 2 * n), dtype=dtypes.bf16)[:, :n],
+            id="column-slice",
+        ),
+    ],
+)
+def test_a4w4_out_rejected(make_out):
+    m, n, k = 45, 512, 512
+    xq, wq, x_scale, w_scale = _a4w4_inputs(m, n, k)
+    with torch.device("cuda"):
+        out = make_out(m, n)
+
+    with pytest.raises(ValueError, match="out "):
+        aiter.gemm_a4w4_out(out, xq, wq, x_scale, w_scale)
+
+
+def test_a4w4_out_not_supported_on_gfx1250(monkeypatch):
+    monkeypatch.setattr(gemm_op_a4w4, "get_gfx", lambda: "gfx1250")
+    a = torch.zeros((32, 256), dtype=torch.uint8, device="cuda")
+    b = torch.zeros((256, 256), dtype=torch.uint8, device="cuda")
+    scale = torch.zeros((32, 16), dtype=torch.uint8, device="cuda")
+    out = torch.empty((32, 256), dtype=dtypes.bf16, device="cuda")
+
+    with pytest.raises(NotImplementedError, match="gfx1250"):
+        gemm_op_a4w4.gemm_a4w4_out(out, a, b, scale, scale)