[HIP] [Feature] a8w8 blockscale asm: out_zeroed skips the split-K output memset

aiter/ops/gemm_op_a8w8.py

@@ -404,6 +404,7 @@ def _gemm_a8w8_blockscale_bpreshuffle_asm(
     kernelName: str | None = None,
     bpreshuffle: int = 1,
     zero_bias_buf: Tensor | None = None,
+    out_zeroed: bool = False,
 ) -> None: ...
 
 
@@ -432,7 +433,10 @@ def gemm_a8w8_blockscale_bpreshuffle_asm(
     kernelName: str | None = None,
     bpreshuffle: bool | None = True,
     zero_bias_buf: Tensor | None = None,
+    out_zeroed: bool = False,
 ) -> Tensor:
+    """out_zeroed: `out` already holds zeros, so a split-K kernel accumulates into it
+    without the hipMemsetAsync launched before it otherwise."""
     if bias is None and zero_bias_buf is None:
         zero_bias_buf = get_zero_bias_buf(B)
     _gemm_a8w8_blockscale_bpreshuffle_asm(
@@ -446,6 +450,7 @@ def gemm_a8w8_blockscale_bpreshuffle_asm(
         kernelName,
         int(bpreshuffle) if bpreshuffle is not None else 1,
         zero_bias_buf,
+        out_zeroed,
     )
     return out
 
@@ -1200,6 +1205,7 @@ def gemm_a8w8_blockscale_bpreshuffle_fake(
     w_scale: Tensor,
     dtype: torch.dtype = dtypes.bf16,
     out: Tensor | None = None,
+    out_zeroed: bool = False,
 ) -> Tensor:
     # Must mirror the real signature (incl. out=): the abstract_impl forwards the
     # out= kwarg here at trace time. When out is given the real fn returns it, so
@@ -1225,6 +1231,7 @@ def gemm_a8w8_blockscale_bpreshuffle(
     w_scale: Tensor,
     dtype: torch.dtype = dtypes.bf16,
     out: Tensor | None = None,
+    out_zeroed: bool = False,
 ) -> Tensor:
     assert dtype in [
         dtypes.bf16,
@@ -1239,6 +1246,10 @@ def gemm_a8w8_blockscale_bpreshuffle(
     # positionally and write into it, so honoring `out` there is free. The
     # triton-fallback branches self-allocate their output and CANNOT honor `out`
     # -> assert loud rather than silently returning a different address.
+    # `out_zeroed`: `out` already holds zeros (written by its producer), so an asm
+    # split-K row accumulates into it without its memset; every other row
+    # overwrites `out` and ignores the flag.
+    assert out is not None or not out_zeroed, "out_zeroed needs the caller's out"
     if out is not None:
         assert out.shape == (m, n) and out.dtype == dtype and out.device == XQ.device, (
             f"gemm_a8w8_blockscale_bpreshuffle: out buffer {tuple(out.shape)}/"
@@ -1428,7 +1439,14 @@ def gemm_a8w8_blockscale_bpreshuffle(
             if dtype == dtypes.bf16:
                 splitK = config["splitK"]
                 return gemm_a8w8_blockscale_bpreshuffle_asm(
-                    XQ, WQ, Y, x_scale, w_scale, splitK=splitK, kernelName=kernelName
+                    XQ,
+                    WQ,
+                    Y,
+                    x_scale,
+                    w_scale,
+                    splitK=splitK,
+                    kernelName=kernelName,
+                    out_zeroed=out_zeroed,
                 )
             _warn_blockscale_asm_row_not_bf16(n, k, dtype)
         elif libtype == "opus":

csrc/py_itfs_cu/asm_a8w8_blockscale_bpreshuffle.cu

@@ -265,8 +265,9 @@ AITER_CTYPES_DEFINE_ENTRYPOINT_VOID(
     const char*  kernelName,
     int          bpreshuffle,
     aiter_tensor_t* zero_bias_buf,
+    int          out_zeroed,
     hipStream_t  stream),
-    (A, B, out, A_scale, B_scale, bias, splitK, kernelName, bpreshuffle, zero_bias_buf, stream))
+    (A, B, out, A_scale, B_scale, bias, splitK, kernelName, bpreshuffle, zero_bias_buf, out_zeroed, stream))
 {
     validate_inputs(A, B, out, A_scale, B_scale);
     int Mdim = A->size(0);
@@ -333,19 +334,19 @@ AITER_CTYPES_DEFINE_ENTRYPOINT_VOID(
                    (Kdim / k_per_split_aligned) == (selectedsplitK - 1),
                    __func__, " Kdim alignment check failed for splitK!");
 
-        // Step 3: Zero output buffer for atomic accumulation across splits.
         if (selectedsplitK > 1) {
-            // Split-K faults on output rows that are not 4-byte aligned.
+            // Split-K requires aligned output even when the producer clears it.
             AITER_CHECK(reinterpret_cast<uintptr_t>(out->ptr) % 4 == 0 && stride_c % 4 == 0,
                         __func__, " split-K needs 4-byte aligned output rows, got ptr offset ",
                         reinterpret_cast<uintptr_t>(out->ptr) % 4, " and row stride ",
                         stride_c, " bytes");
-            size_t row_bytes = static_cast<size_t>(Ndim) * out->element_size();
-            if (static_cast<size_t>(stride_c) == row_bytes) {
-                HIP_CALL(hipMemsetAsync(out->ptr, 0, out->numel() * out->element_size(), stream));
-            } else {
-                // The kernel writes rows stride_c bytes apart, so clear only the M x N view.
-                HIP_CALL(hipMemset2DAsync(out->ptr, stride_c, 0, row_bytes, Mdim, stream));
+            if (!out_zeroed) {
+                size_t row_bytes = static_cast<size_t>(Ndim) * out->element_size();
+                if (static_cast<size_t>(stride_c) == row_bytes) {
+                    HIP_CALL(hipMemsetAsync(out->ptr, 0, out->numel() * out->element_size(), stream));
+                } else {
+                    HIP_CALL(hipMemset2DAsync(out->ptr, stride_c, 0, row_bytes, Mdim, stream));
+                }
             }
         }
     }

op_tests/test_gemm_a8w8_blockscale_bpreshuffle_out.py

@@ -92,6 +92,62 @@ def test_split_k_out_needs_four_byte_alignment():
         run_backend("asm_sk4", x, w, x_scale, w_scale, wide[:, 1 : N + 1])
 
 
+@pytest.mark.parametrize("row_padded", [False, True])
+@pytest.mark.parametrize("dispatch", [False, True])
+def test_pre_zeroed_split_k_out(monkeypatch, row_padded, dispatch):
+    m = 32
+    x, w, x_scale, w_scale = make_inputs(m)
+    ref = torch.empty(m, N, dtype=dtypes.bf16, device="cuda")
+    run_backend("asm_sk4", x, w, x_scale, w_scale, ref)
+    width = N + PAD if row_padded else N
+    storage = torch.full((m, width), SENTINEL, dtype=dtypes.bf16, device="cuda")
+    out = storage[:, :N]
+    if dispatch:
+        asm_row = {"libtype": "asm", "splitK": 4, "kernelName": ASM_KERNEL}
+        monkeypatch.setattr(gemm_mod, "get_CKGEMM_config", lambda *args, **kw: asm_row)
+
+    def launch():
+        if dispatch:
+            return gemm_mod.gemm_a8w8_blockscale_bpreshuffle(
+                x, w, x_scale, w_scale, dtypes.bf16, out=out, out_zeroed=True
+            )
+        return aiter.gemm_a8w8_blockscale_bpreshuffle_asm(
+            x,
+            w,
+            out,
+            x_scale,
+            w_scale,
+            splitK=4,
+            kernelName=ASM_KERNEL,
+            out_zeroed=True,
+        )
+
+    out.zero_()
+    launch()
+    torch.testing.assert_close(out.float(), ref.float(), atol=0.5, rtol=0.02)
+    assert bool((storage[:, N:] == SENTINEL).all()), "cleared output padding"
+    # Deliberately violate the contract to prove the explicit flag skips clearing.
+    out.fill_(SENTINEL)
+    launch()
+    assert float((out.float() - ref.float()).mean()) > 6.0
+
+
+def test_pre_zeroed_split_k_still_checks_alignment():
+    x, w, x_scale, w_scale = make_inputs(32)
+    wide = torch.zeros(32, N + 2, dtype=dtypes.bf16, device="cuda")
+    with pytest.raises(RuntimeError, match="4-byte aligned"):
+        aiter.gemm_a8w8_blockscale_bpreshuffle_asm(
+            x,
+            w,
+            wide[:, 1 : N + 1],
+            x_scale,
+            w_scale,
+            splitK=4,
+            kernelName=ASM_KERNEL,
+            out_zeroed=True,
+        )
+
+
 def test_fp16_fallback_warns_once_per_shape(monkeypatch):
     asm_row = {"libtype": "asm", "splitK": 4, "kernelName": ASM_KERNEL}
     monkeypatch.setattr(gemm_mod, "get_CKGEMM_config", lambda *args, **kwargs: asm_row)