Persistent AG + GEMM kernels w/ HipKittens & UserBuffers

tests/pytorch/distributed/run_gemm_with_overlap.py

@@ -1,12 +1,15 @@
 #!/usr/bin/python3
 
+# This file was modified for portability to AMDGPU
+# Copyright (c) 2025-2026, Advanced Micro Devices, Inc. All rights reserved.
 # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # See LICENSE for license information.
 
 import os
 import sys
 import socket
+import hashlib
 import warnings
 import subprocess
 import argparse
@@ -76,6 +79,12 @@ def _parse_args(argv=None, namespace=None):
     parser.add_argument(
         "--atomic", action="store_true", default=False, help="Test overlap with atomic GEMM."
     )
+    parser.add_argument(
+        "--fused",
+        action="store_true",
+        default=False,
+        help="Select the fused AG+GEMM backend (gfx950 only).",
+    )
     parser.add_argument(
         "--aggregate",
         action="store_true",
@@ -357,6 +366,7 @@ def dist_print(msg, src=None, info=False, error=False, section=False, group=None
             atomic_gemm=opts.atomic,
             aggregate=opts.aggregate,
             use_ce=not (opts.atomic and bool(int(os.getenv("NVTE_AG_P2P_MULTI_ATOMIC", "0")))),
+            fused=opts.fused,
         )
     else:
         ub_obj = tex.CommOverlap(
@@ -874,6 +884,11 @@ def _gemm():
 
         torch.cuda.synchronize()
         dist.barrier(tp_group)
+        # Bit-exact fingerprint of the output.
+        out_bytes = test_out.detach().contiguous().cpu().flatten().view(torch.uint8)
+        out_hash = hashlib.sha256(out_bytes.numpy().tobytes()).hexdigest()
+        dist_print(f"OUTPUT HASH: {out_hash}", section=True, group=tp_group)
+
         diff = torch.abs(test_out - ref_out).flatten()
         m = torch.argmax(diff)
         abs_err = diff[m].item()

tests/pytorch/distributed/test_comm_gemm_overlap.py

@@ -445,3 +445,97 @@ def test_multi_layer_with_overlap_fp8(
         num_layers,
         use_cublasmp=use_cublasmp,
     )
+
+
+# The fused AG+GEMM backend is deliberately exempt from MAX_GPUS_TO_USE.
+FUSED_MAX_GPUS_TO_USE = 8
+
+
+def _fused_shape_ok(nprocs: int) -> bool:
+    """The backend needs a 256-aligned per-rank chunk and a 256-aligned N."""
+    return (SEQ_LENGTH * BATCH_SIZE) % (256 * nprocs) == 0 and (NUM_HEADS * HEAD_DIM) % 256 == 0
+
+
+FUSED_PROC_COUNTS = [
+    n
+    for n in (4, 8)
+    if n <= min(torch.cuda.device_count(), FUSED_MAX_GPUS_TO_USE) and _fused_shape_ok(n)
+]
+
+fused_available = (
+    IS_HIP_EXTENSION and get_device_compute_capability() == (9, 5) and len(FUSED_PROC_COUNTS) > 0
+)
+reason_for_no_fused = (
+    "Fused AG+GEMM overlap requires a gfx950 device, tp_size in (4, 8) and a 256-aligned per-rank chunk."
+)
+
+
+def _fused_launch_cmd(nprocs: int):
+    """Same form as LAUNCH_CMD, but at a rank count the fused tests choose."""
+    if tex.ubuf_built_with_mpi():
+        return ["mpirun", "-np", str(nprocs), "--oversubscribe", "--quiet", "python3"]
+    return ["torchrun", f"--nproc_per_node={nprocs}"]
+
+
+def _run_fused_ag(quantization="none", nprocs=None):
+    """Run the AG overlap harness with the fused backend, returning the completed process."""
+    test_cmd = _fused_launch_cmd(nprocs if nprocs is not None else FUSED_PROC_COUNTS[0]) + [
+        str(TEST_ROOT / "run_gemm_with_overlap.py"),
+        "--check-numerics",
+        f"--seed={RNG_SEED}",
+        f"--seq-length={SEQ_LENGTH}",
+        f"--batch-size={BATCH_SIZE}",
+        f"--num-heads={NUM_HEADS}",
+        f"--head-dim={HEAD_DIM}",
+        "--comm-type=AG",
+        "--p2p",
+        "--fused",
+        f"--quantization={quantization}",
+    ]
+    return subprocess.run(test_cmd, env=os.environ, capture_output=True, check=False)
+
+
+def _assert_numerics_passed(result):
+    stdout, stderr = result.stdout.decode(), result.stderr.decode()
+    assert result.returncode == 0, f"non-zero exit\n{stderr}"
+    assert "NUMERICAL CHECK FAILED" not in stderr, stderr
+    assert "NUMERICAL CHECK PASSED" in stdout, stdout
+
+
+@pytest.mark.skipif(not fused_available, reason=reason_for_no_fused)
+@pytest.mark.parametrize("nprocs", FUSED_PROC_COUNTS)
+def test_fused_ag_overlap_bf16(nprocs):
+    """bf16 at an aligned shape: the fused backend runs and the result is correct."""
+    _assert_numerics_passed(_run_fused_ag(nprocs=nprocs))
+
+
+@pytest.mark.skipif(not fused_available, reason=reason_for_no_fused)
+@pytest.mark.parametrize("nprocs", FUSED_PROC_COUNTS)
+@pytest.mark.parametrize("quantization", ("fp8", "mxfp8"))
+def test_fused_ag_overlap_rejects_non_bf16(quantization, nprocs):
+    """Non-bf16 is currently outside the backend."""
+    if quantization == "fp8" and not fp8_available:
+        pytest.skip(reason_for_no_fp8)
+    if quantization == "mxfp8" and not mxfp8_available:
+        pytest.skip(reason_for_no_mxfp8)
+    result = _run_fused_ag(quantization=quantization, nprocs=nprocs)
+    assert result.returncode != 0, "fused AG+GEMM accepted a non-bf16 operand"
+    assert "non-bf16 operand" in result.stderr.decode(), result.stderr.decode()
+
+
+@pytest.mark.skipif(not fused_available, reason=reason_for_no_fused)
+@pytest.mark.parametrize("nprocs", FUSED_PROC_COUNTS)
+def test_fused_ag_overlap_is_deterministic(nprocs):
+    """Bitwise reproducibility across runs"""
+    first = _run_fused_ag(nprocs=nprocs)
+    _assert_numerics_passed(first)
+    second = _run_fused_ag(nprocs=nprocs)
+    _assert_numerics_passed(second)
+
+    def _hashes(out):
+        prefix = "OUTPUT HASH: "
+        return [ln.split(prefix, 1)[1].strip() for ln in out.decode().splitlines() if prefix in ln]
+
+    first_hashes, second_hashes = _hashes(first.stdout), _hashes(second.stdout)
+    assert first_hashes, f"harness printed no output hash\n{first.stdout.decode()}"
+    assert first_hashes == second_hashes, "two identical runs produced different outputs"

transformer_engine/common/comm_gemm_overlap/comm_gemm_overlap.cpp

@@ -799,11 +799,12 @@ CommOverlapP2PBase::CommOverlapP2PBase(const std::vector<size_t> &buffer_shape,
                                        CommOverlapType comm_type, int num_max_streams,
                                        int comm_cga_size, int gemm_priority, int comm_priority,
                                        int num_comm_sm, bool set_sm_margin, bool use_ce,
-                                       bool atomic_gemm, bool aggregate)
+                                       bool atomic_gemm, bool aggregate, bool fused)
     : CommOverlapCore(myrank, numranks, mylocal, numlocal, mynode, numnodes, tp_size,
                       allgather_handle, barrier_handle, tp_size, num_max_streams, comm_cga_size,
                       gemm_priority, comm_priority, num_comm_sm, set_sm_margin, use_ce,
-                      atomic_gemm) {
+                      atomic_gemm),
+      _fused(fused) {
   initialize(buffer_shape, buffer_dtype, comm_type, aggregate);
 }
 

transformer_engine/common/comm_gemm_overlap/rocm_comm_gemm_overlap.cpp

@@ -14,6 +14,9 @@
 #include "common/util/logging.h"
 #include "common/util/system.h"
 #include "userbuffers/userbuffers.h"
+#ifdef USE_HIPKITTENS_GEMM
+#include "../gemm/kittens/fused_ag_gemm.h"
+#endif
 
 namespace transformer_engine {
 #if 0
@@ -246,6 +249,59 @@ static_assert(rings_are_unique<7,8>(tp_next_8), "Rings overlap");
 static_assert(prev_is_inverse_of_next<2,4>(tp_next_4, tp_prev_4), "tp_prev_4 is not inverse of tp_next_4");
 static_assert(prev_is_inverse_of_next<7,8>(tp_next_8, tp_prev_8), "tp_prev_8 is not inverse of tp_next_8");
 
+#ifdef USE_HIPKITTENS_GEMM
+// Fused all-gather + GEMM, launched by fused_overlap_ag below.
+static bool hk_fused_ag_gemm(const TensorWrapper &A, bool transa, bool transb, TensorWrapper &D,
+                             const TensorWrapper &bias, const TensorWrapper &pre_gelu_out,
+                             const TensorWrapper &B_copy, TensorWrapper &workspace, bool accumulate,
+                             const TensorWrapper &ubuf, const TensorWrapper &chunk, communicator *comm,
+                             int reg, int tp_id, int tp_size, uint64_t signal, cudaStream_t stream) {
+  // TODO: Add bias support
+  NVTE_CHECK(!transb && !accumulate && bias.numel() == 0 && pre_gelu_out.numel() == 0 && B_copy.numel() == 0,
+             "fused AG+GEMM reached with an unsupported epilogue");
+  NVTE_CHECK(A.dtype() == DType::kBFloat16 && ubuf.dtype() == DType::kBFloat16 && D.dtype() == DType::kBFloat16,
+             "fused AG+GEMM reached with a non-bf16 operand");
+
+  const size_t m       = (transa) ? A.size(0) : A.size(1);
+  const size_t k       = (transa) ? A.size(1) : A.size(0);
+  const size_t n_chunk = chunk.size(0);
+  NVTE_CHECK((tp_size == 4 || tp_size == 8) && m % 256 == 0 && k % 128 == 0 && k >= 256 && n_chunk % 256 == 0,
+             "fused AG+GEMM reached with an ineligible shape (m=", m, " k=", k, " n_chunk=", n_chunk,
+             " tp_size=", tp_size, ")");
+
+  const int rank_round_tp = comm->myrank - tp_id;
+  KittensFusedAgGemmArgs args{
+      A.dptr(), ubuf.dptr(), D.dptr(),
+      reinterpret_cast<char *>(comm->gpu_ptrs) + reg * comm->nvsize * sizeof(void *),
+      rank_round_tp % comm->nvsize, comm->nvsize,
+      GET_RECV_PTR_BY_INDEX(rank_round_tp, comm, reg, 0), comm->gpu_ptrs,
+      static_cast<size_t>(GET_SEND_PTR_BY_INDEX(0, comm, reg, 0) - reinterpret_cast<char *>(comm->peer_ptr[0][0])),
+      static_cast<size_t>(GET_RECV_PTR_BY_INDEX(1, comm, reg, 0) - GET_RECV_PTR_BY_INDEX(0, comm, reg, 0)),
+      signal, static_cast<int>(m), static_cast<int>(n_chunk * tp_size), static_cast<int>(k), transa,
+      tp_id, tp_size, chunk.bytes(), workspace.dptr(), workspace.bytes(), stream};
+  return kittens_fused_ag_gemm_bf16(args);
+}
+#endif
+
+void CommOverlapP2PBase::fused_overlap_ag(const TensorWrapper &A, bool transa, const TensorWrapper &B,
+                                bool transb, TensorWrapper &D, TensorWrapper &bias,
+                                TensorWrapper &pre_gelu_out, TensorWrapper &workspace, bool grad,
+                                bool accumulate, bool use_split_accumulator, TensorWrapper &B_copy,
+                                cudaStream_t stream_main) {
+#ifdef USE_HIPKITTENS_GEMM
+  if (kittens_fused_ag_gemm_supported(cuda::sm_arch())) {
+    const bool launched = hk_fused_ag_gemm(A, transa, transb, D, bias, pre_gelu_out, B_copy,
+                                           workspace, accumulate, _ubuf, _ubufs[0], _ub_comm,
+                                           _ub_reg, _tp_id, _tp_size, _ag_signal_base + _tp_size,
+                                           stream_main);
+    NVTE_CHECK(launched, "fused AG+GEMM failed to launch");
+    _ag_signal_base += _tp_size;
+    return;
+  }
+#endif
+  NVTE_ERROR("fused AG+GEMM was selected but is not built into this library");
+}
+
 // TODO: Introduce HIPGraphs for dependency management.
 void CommOverlapP2PBase::rocm_split_overlap_ag(const TensorWrapper &A, bool transa, const TensorWrapper &B,
                                 bool transb, TensorWrapper &D, TensorWrapper &bias,

transformer_engine/common/gemm/kittens/CMakeLists.txt

@@ -87,6 +87,7 @@ else()
     # round_to_out_dtype() in cdna4/blockwise_fp8_gemm_helper.cuh.
     kittens_add_arch(TAG cdna4 GFX gfx950
                      SOURCES cdna4/blockwise_fp8_gemm.cpp cdna4/mxfp8_gemm.cpp
+                             cdna4/fused_ag_gemm.cpp
                      FLAGS -gline-tables-only
                      SOURCE_FLAGS cdna4/blockwise_fp8_gemm.cpp "-ffast-math")