From 392fd78363f7f3e4abd2dba0c941b0c3c4db8182 Mon Sep 17 00:00:00 2001 From: alextmagro Date: Mon, 31 Aug 2026 11:28:01 -0500 Subject: [PATCH 1/3] Reduce scatter overlap kernels for gfx950 + bf16 --- ci/pytorch.sh | 2 +- .../distributed/test_comm_gemm_overlap.py | 193 +----- .../distributed/test_rocm_fused_overlap.py | 348 ++++++++++ .../comm_gemm_overlap/comm_gemm_overlap.cpp | 21 +- .../rocm_comm_gemm_overlap.cpp | 133 +++- .../common/gemm/kittens/CMakeLists.txt | 2 +- .../kittens/cdna4/overlap/bulk_rs_gemm_nt.cuh | 444 +++++++++++++ .../comm_gemm.cpp} | 299 ++++++++- .../cdna4/{ => overlap}/fused_ag_gemm_nn.cuh | 84 +-- .../cdna4/{ => overlap}/fused_ag_gemm_tn.cuh | 80 +-- .../cdna4/overlap/fused_rs_gemm_tn.cuh | 617 ++++++++++++++++++ .../kittens/cdna4/overlap/overlap_common.cuh | 70 ++ .../common/gemm/kittens/comm_gemm.h | 74 +++ .../common/gemm/kittens/fused_ag_gemm.h | 41 -- .../common/gemm/kittens/kittens_common.cpp | 78 ++- .../transformer_engine/comm_gemm_overlap.h | 40 ++ .../pytorch/csrc/extensions/gemm.cpp | 24 +- .../pytorch/csrc/extensions/pybind.cpp | 4 +- transformer_engine/pytorch/module/base.py | 42 +- .../pytorch/module/layernorm_linear.py | 10 +- .../pytorch/module/layernorm_mlp.py | 5 +- transformer_engine/pytorch/module/linear.py | 11 +- transformer_engine/pytorch/transformer.py | 7 +- 23 files changed, 2225 insertions(+), 404 deletions(-) create mode 100644 tests/pytorch/distributed/test_rocm_fused_overlap.py create mode 100644 transformer_engine/common/gemm/kittens/cdna4/overlap/bulk_rs_gemm_nt.cuh rename transformer_engine/common/gemm/kittens/cdna4/{fused_ag_gemm.cpp => overlap/comm_gemm.cpp} (55%) rename transformer_engine/common/gemm/kittens/cdna4/{ => overlap}/fused_ag_gemm_nn.cuh (94%) rename transformer_engine/common/gemm/kittens/cdna4/{ => overlap}/fused_ag_gemm_tn.cuh (93%) create mode 100644 transformer_engine/common/gemm/kittens/cdna4/overlap/fused_rs_gemm_tn.cuh create mode 100644 transformer_engine/common/gemm/kittens/cdna4/overlap/overlap_common.cuh create mode 100644 transformer_engine/common/gemm/kittens/comm_gemm.h delete mode 100644 transformer_engine/common/gemm/kittens/fused_ag_gemm.h diff --git a/ci/pytorch.sh b/ci/pytorch.sh index b5241f3d75..efdd8eeeae 100755 --- a/ci/pytorch.sh +++ b/ci/pytorch.sh @@ -132,7 +132,7 @@ run_test_config_mgpu(){ run_default_fa 1 test_gemm_sm_count.py run_default_fa 3 test_sanity_import.py run_default_fa 3 distributed/test_cast_master_weights_to_fp8.py - run_default_fa 3 distributed/test_comm_gemm_overlap.py + run_default_fa 3 distributed/test_rocm_fused_overlap.py run_default_fa 2 distributed/test_fusible_ops.py run_default_fa 2 distributed/test_numerics.py run_default_fa 1 distributed/test_torch_fsdp2.py diff --git a/tests/pytorch/distributed/test_comm_gemm_overlap.py b/tests/pytorch/distributed/test_comm_gemm_overlap.py index 4a796deea4..6e87061f8b 100644 --- a/tests/pytorch/distributed/test_comm_gemm_overlap.py +++ b/tests/pytorch/distributed/test_comm_gemm_overlap.py @@ -21,6 +21,7 @@ fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) +bulk_available = not IS_HIP_EXTENSION RNG_SEED: int = 42 SEQ_LENGTH: int = 1024 @@ -212,7 +213,7 @@ def test_split_reduce_scatter_overlaps(quantization, p2p, use_cublasmp): _run_gemm_with_overlap("RS", False, p2p, False, False, quantization, use_cublasmp) -@pytest.mark.skipif(IS_HIP_EXTENSION, reason="Bulk overlap is not yet supported on ROCm.") +@pytest.mark.skipif(not bulk_available, reason="Bulk overlap is not supported on ROCm.") @pytest.mark.parametrize( "comm_type, quantization, connections", [ @@ -237,6 +238,8 @@ def test_bulk_overlaps(comm_type, quantization, connections): Test bulk overlaps with direct calls to te.cpp_extensions.gemm or te.cpp_extensions.fp8_gemm. """ if connections == 8: + if IS_HIP_EXTENSION: + pytest.skip("CUDA_DEVICE_MAX_CONNECTIONS has no HIP equivalent.") if torch.cuda.get_device_properties(0).major != 9: pytest.skip( "CUDA_DEVICE_MAX_CONNECTIONS=8 test only applies to devices with compute capability" @@ -445,191 +448,3 @@ 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(nprocs, bulk=False, quantization="none"): - """Run the AG overlap harness with the fused backend, returning the completed process.""" - test_cmd = _fused_launch_cmd(nprocs) + [ - 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", - "--fused", - ] - test_cmd += ["--bulk-overlap"] if bulk else ["--p2p", f"--quantization={quantization}"] - return subprocess.run(test_cmd, env=os.environ, capture_output=True, check=False) - -ELIGIBLE_OUT_FEATURES_PER_RANK = 1536 -INELIGIBLE_OUT_FEATURES_PER_RANK = 1568 -UNALIGNED_SEQ_LENGTH = 1152 - -def _run_fused_layer(nprocs, extra_args, seq_length=SEQ_LENGTH): - """Run the layer harness on a column-parallel LayerNormLinear with the fused backend live.""" - test_cmd = ( - _fused_launch_cmd(nprocs) - + [ - str(TEST_ROOT / "run_layer_with_overlap.py"), - f"--seed={RNG_SEED}", - f"--seq-length={seq_length}", - f"--batch-size={BATCH_SIZE}", - f"--num-heads={NUM_HEADS}", - f"--head-dim={HEAD_DIM}", - f"--layer-type={te.LayerNormLinear.__name__}", - "--linear-parallel-mode=column", - "--num-layers=1", - "--use-bf16-params", - ] - + extra_args - ) - env = os.environ.copy() - env["PYTORCH_JIT"] = "0" - env["NVTE_TORCH_COMPILE"] = "0" - env["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] = "0" - return subprocess.run(test_cmd, env=env, capture_output=True, check=False) - - -def _reported_names(stdout, prefix): - """The layer name sets the harness printed under `prefix`.""" - for line in stdout.decode().splitlines(): - if prefix in line: - return set(line.split(prefix, 1)[1].split()) - return None - - -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)) - - -@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(nprocs, quantization=quantization) - 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) - _assert_numerics_passed(first) - second = _run_fused_ag(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" - - -@pytest.mark.skipif(not fused_available, reason=reason_for_no_fused) -@pytest.mark.parametrize("nprocs", FUSED_PROC_COUNTS) -def test_fused_bulk_ag_overlap_bf16(nprocs): - """The bulk all-gather that rides in an unrelated GEMM's grid.""" - _assert_numerics_passed(_run_fused_ag(nprocs, bulk=True)) - - -@pytest.mark.skipif(not fused_available, reason=reason_for_no_fused) -@pytest.mark.parametrize("nprocs", FUSED_PROC_COUNTS) -def test_fused_layer_bulk_dgrad_bf16(nprocs): - """A column-parallel layer whose dgrad dimensions clear the fused contract.""" - result = _run_fused_layer(nprocs, [f"--out-features={ELIGIBLE_OUT_FEATURES_PER_RANK * nprocs}"]) - _assert_numerics_passed(result) - fused = _reported_names(result.stdout, "UB FUSED NAMES: ") - assert fused is not None, f"harness printed no fused name set\n{result.stdout.decode()}" - assert "qkv_dgrad" in fused, fused - eligible = _reported_names(result.stdout, "UB BULK ELIGIBLE: ") - assert eligible is not None, f"harness printed no eligibility set\n{result.stdout.decode()}" - assert "qkv_dgrad" in eligible, eligible - - -@pytest.mark.skipif(not fused_available, reason=reason_for_no_fused) -@pytest.mark.parametrize("nprocs", FUSED_PROC_COUNTS) -def test_fused_layer_declines_ineligible_k(nprocs): - """A K the fused kernels cannot serve has to fall back to no overlap.""" - result = _run_fused_layer( - nprocs, [f"--out-features={INELIGIBLE_OUT_FEATURES_PER_RANK * nprocs}"] - ) - _assert_numerics_passed(result) - stderr = result.stderr.decode() - assert "ineligible shape" not in stderr, stderr - assert "failed to launch" not in stderr, stderr - fused = _reported_names(result.stdout, "UB FUSED NAMES: ") - disabled = _reported_names(result.stdout, "UB DISABLED NAMES: ") - assert fused is not None, f"harness printed no fused name set\n{result.stdout.decode()}" - assert "qkv_dgrad" in fused, fused - assert disabled is not None and "qkv_dgrad" not in disabled, disabled - eligible = _reported_names(result.stdout, "UB BULK ELIGIBLE: ") - assert eligible is not None, f"harness printed no eligibility set\n{result.stdout.decode()}" - assert "qkv_dgrad" not in eligible, eligible - - -@pytest.mark.skipif(not fused_available, reason=reason_for_no_fused) -@pytest.mark.parametrize("nprocs", FUSED_PROC_COUNTS) -def test_fused_layer_declines_unaligned_region(nprocs): - """A Userbuffers region the fused backend cannot serve declines at setup.""" - result = _run_fused_layer( - nprocs, - [f"--out-features={ELIGIBLE_OUT_FEATURES_PER_RANK * nprocs}"], - seq_length=UNALIGNED_SEQ_LENGTH, - ) - _assert_numerics_passed(result) - fused = _reported_names(result.stdout, "UB FUSED NAMES: ") - disabled = _reported_names(result.stdout, "UB DISABLED NAMES: ") - assert fused == set(), f"expected no fused communicators, got {fused}" - assert disabled is not None, f"harness printed no disabled name set\n{result.stdout.decode()}" - assert {"qkv_fprop", "qkv_dgrad", "qkv_wgrad"} <= disabled, disabled diff --git a/tests/pytorch/distributed/test_rocm_fused_overlap.py b/tests/pytorch/distributed/test_rocm_fused_overlap.py new file mode 100644 index 0000000000..0aed15059a --- /dev/null +++ b/tests/pytorch/distributed/test_rocm_fused_overlap.py @@ -0,0 +1,348 @@ +# Copyright (c) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +"""ROCm fused comm+GEMM overlap tests""" +import os +import subprocess +from pathlib import Path + +import pytest +import torch +import transformer_engine.pytorch as te + +from torch.utils.cpp_extension import IS_HIP_EXTENSION +from transformer_engine.pytorch.utils import get_device_compute_capability + + +if torch.cuda.device_count() < 2: + pytest.skip("Comm+GEMM overlap requires at least 2 GPUs.", allow_module_level=True) + +fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) +mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) + +RNG_SEED: int = 42 +SEQ_LENGTH: int = 1024 +BATCH_SIZE: int = 2 +NUM_HEADS: int = 32 +HEAD_DIM: int = 48 + +TEST_ROOT = Path(__file__).parent.resolve() + + +def _fused_shape_ok(nprocs: int) -> bool: + """The minimum both backends need: a 256-aligned per-rank chunk and a 256-aligned N. + + The reduce-scatter backend is stricter still -- its combine is specialised on 8 ranks -- but + that is asserted per test rather than hidden in collection, so a tp=4 decline stays visible. + """ + 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 <= torch.cuda.device_count() 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 comm+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(nprocs, bulk=False, quantization="none"): + """Run the AG overlap harness with the fused backend, returning the completed process.""" + test_cmd = _fused_launch_cmd(nprocs) + [ + 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", + "--fused", + ] + test_cmd += ["--bulk-overlap"] if bulk else ["--p2p", f"--quantization={quantization}"] + return subprocess.run(test_cmd, env=os.environ, capture_output=True, check=False) + + +ELIGIBLE_OUT_FEATURES_PER_RANK = 1536 +INELIGIBLE_OUT_FEATURES_PER_RANK = 1568 +UNALIGNED_SEQ_LENGTH = 1152 + + +def _run_fused_layer(nprocs, extra_args, seq_length=SEQ_LENGTH): + """Run the layer harness on a column-parallel LayerNormLinear with the fused backend live.""" + test_cmd = ( + _fused_launch_cmd(nprocs) + + [ + str(TEST_ROOT / "run_layer_with_overlap.py"), + f"--seed={RNG_SEED}", + f"--seq-length={seq_length}", + f"--batch-size={BATCH_SIZE}", + f"--num-heads={NUM_HEADS}", + f"--head-dim={HEAD_DIM}", + f"--layer-type={te.LayerNormLinear.__name__}", + "--linear-parallel-mode=column", + "--num-layers=1", + "--use-bf16-params", + ] + + extra_args + ) + env = os.environ.copy() + env["PYTORCH_JIT"] = "0" + env["NVTE_TORCH_COMPILE"] = "0" + env["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] = "0" + return subprocess.run(test_cmd, env=env, capture_output=True, check=False) + + +def _reported_names(stdout, prefix): + """The layer name sets the harness printed under `prefix`.""" + for line in stdout.decode().splitlines(): + if prefix in line: + return set(line.split(prefix, 1)[1].split()) + return None + + +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)) + + +@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(nprocs, quantization=quantization) + 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) + _assert_numerics_passed(first) + second = _run_fused_ag(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" + + +@pytest.mark.skipif(not fused_available, reason=reason_for_no_fused) +@pytest.mark.parametrize("nprocs", FUSED_PROC_COUNTS) +def test_fused_bulk_ag_overlap_bf16(nprocs): + """The bulk all-gather that rides in an unrelated GEMM's grid.""" + _assert_numerics_passed(_run_fused_ag(nprocs, bulk=True)) + + +@pytest.mark.skipif(not fused_available, reason=reason_for_no_fused) +@pytest.mark.parametrize("nprocs", FUSED_PROC_COUNTS) +def test_fused_layer_bulk_dgrad_bf16(nprocs): + """A column-parallel layer whose dgrad dimensions clear the fused contract.""" + result = _run_fused_layer(nprocs, [f"--out-features={ELIGIBLE_OUT_FEATURES_PER_RANK * nprocs}"]) + _assert_numerics_passed(result) + fused = _reported_names(result.stdout, "UB FUSED NAMES: ") + assert fused is not None, f"harness printed no fused name set\n{result.stdout.decode()}" + assert "qkv_dgrad" in fused, fused + eligible = _reported_names(result.stdout, "UB BULK ELIGIBLE: ") + assert eligible is not None, f"harness printed no eligibility set\n{result.stdout.decode()}" + assert "qkv_dgrad" in eligible, eligible + + +@pytest.mark.skipif(not fused_available, reason=reason_for_no_fused) +@pytest.mark.parametrize("nprocs", FUSED_PROC_COUNTS) +def test_fused_layer_bulk_wgrad_bf16(nprocs): + """The backward reduce-scatter hidden behind the wgrad GEMM, on the fused method. + + wgrad is in the default fused layer set on gfx950, so this needs no extra configuration. + """ + result = _run_fused_layer(nprocs, [f"--out-features={ELIGIBLE_OUT_FEATURES_PER_RANK * nprocs}"]) + _assert_numerics_passed(result) + fused = _reported_names(result.stdout, "UB FUSED NAMES: ") + assert fused is not None, f"harness printed no fused name set\n{result.stdout.decode()}" + assert "qkv_wgrad" in fused, fused + eligible = _reported_names(result.stdout, "UB BULK ELIGIBLE: ") + assert eligible is not None, f"harness printed no eligibility set\n{result.stdout.decode()}" + assert "qkv_wgrad" in eligible, eligible + + +def _run_fused_rs_layer(nprocs, extra_args, seq_length=SEQ_LENGTH): + """Run the layer harness on a ROW-parallel Linear, which is the fused reduce-scatter's path.""" + test_cmd = ( + _fused_launch_cmd(nprocs) + + [ + str(TEST_ROOT / "run_layer_with_overlap.py"), + f"--seed={RNG_SEED}", + f"--seq-length={seq_length}", + f"--batch-size={BATCH_SIZE}", + f"--num-heads={NUM_HEADS}", + f"--head-dim={HEAD_DIM}", + f"--layer-type={te.Linear.__name__}", + "--linear-parallel-mode=row", + "--num-layers=1", + "--use-bf16-params", + ] + + extra_args + ) + env = os.environ.copy() + env["PYTORCH_JIT"] = "0" + env["NVTE_TORCH_COMPILE"] = "0" + env["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] = "0" + return subprocess.run(test_cmd, env=env, capture_output=True, check=False) + + +@pytest.mark.skipif(not fused_available, reason=reason_for_no_fused) +@pytest.mark.parametrize("nprocs", FUSED_PROC_COUNTS) +def test_fused_rs_overlap_bf16(nprocs): + """bf16 at an aligned shape: the fused reduce-scatter runs and the result is correct. + + The combine is specialised on 8 ranks, so below that the backend must decline and fall back + rather than fail -- both outcomes are correct, and which one applies is asserted on nprocs. + """ + result = _run_fused_rs_layer( + nprocs, [f"--in-features={ELIGIBLE_OUT_FEATURES_PER_RANK * nprocs}"] + ) + _assert_numerics_passed(result) + stderr = result.stderr.decode() + assert "failed to launch" not in stderr, stderr + disabled = _reported_names(result.stdout, "UB DISABLED NAMES: ") + assert disabled is not None, f"harness printed no disabled name set\n{result.stdout.decode()}" + if nprocs == 8: + assert "proj_fprop" not in disabled, disabled + else: + assert "proj_fprop" in disabled, f"tp={nprocs} must decline the fused RS, got {disabled}" + + +@pytest.mark.skipif(not fused_available, reason=reason_for_no_fused) +@pytest.mark.parametrize("nprocs", FUSED_PROC_COUNTS) +@pytest.mark.parametrize("quantization", ("fp8_delayed_scaling", "mxfp8")) +def test_fused_rs_overlap_rejects_non_bf16(quantization, nprocs): + """Non-bf16 is outside the backend, and setup must decline rather than reach the kernel. + + The AG sibling asserts a non-zero exit because its region carries the operand dtype straight + into the kernel guard. The RS path is gated earlier -- _fused_rs_ub_supported rejects any + non-bf16 region at setup -- so the layer is disabled and the run stays correct. Either way the + kernel must never see a non-bf16 operand, which is what the stderr assertion pins. + """ + if quantization.startswith("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_rs_layer( + nprocs, + [ + f"--in-features={ELIGIBLE_OUT_FEATURES_PER_RANK * nprocs}", + f"--quantization={quantization}", + ], + ) + stderr = result.stderr.decode() + assert "non-bf16 operand" not in stderr, f"a non-bf16 operand reached the kernel\n{stderr}" + if result.returncode == 0: + disabled = _reported_names(result.stdout, "UB DISABLED NAMES: ") + assert disabled is not None, f"harness printed no disabled set\n{result.stdout.decode()}" + assert "proj_fprop" in disabled, f"non-bf16 was not declined at setup: {disabled}" + + +@pytest.mark.skipif(not fused_available, reason=reason_for_no_fused) +@pytest.mark.parametrize("nprocs", FUSED_PROC_COUNTS) +def test_fused_rs_declines_unaligned_region(nprocs): + """tokens %% (tp * BLOCK_ROW) != 0 has no whole band, so setup must decline.""" + result = _run_fused_rs_layer( + nprocs, + [f"--in-features={ELIGIBLE_OUT_FEATURES_PER_RANK * nprocs}"], + seq_length=UNALIGNED_SEQ_LENGTH, + ) + _assert_numerics_passed(result) + disabled = _reported_names(result.stdout, "UB DISABLED NAMES: ") + assert disabled is not None, f"harness printed no disabled name set\n{result.stdout.decode()}" + assert "proj_fprop" in disabled, disabled + + +@pytest.mark.skipif(not fused_available, reason=reason_for_no_fused) +@pytest.mark.parametrize("nprocs", FUSED_PROC_COUNTS) +def test_fused_rs_declines_ineligible_k(nprocs): + """A K the fused kernels cannot serve falls back without erroring.""" + result = _run_fused_rs_layer( + nprocs, [f"--in-features={INELIGIBLE_OUT_FEATURES_PER_RANK * nprocs}"] + ) + _assert_numerics_passed(result) + stderr = result.stderr.decode() + assert "ineligible shape" not in stderr, stderr + assert "failed to launch" not in stderr, stderr + + +@pytest.mark.skipif(not fused_available, reason=reason_for_no_fused) +@pytest.mark.parametrize("nprocs", FUSED_PROC_COUNTS) +def test_fused_rs_overlap_is_deterministic(nprocs): + """Two runs at the same seed agree: the collective's arrival order must not reach the output.""" + args = [f"--in-features={ELIGIBLE_OUT_FEATURES_PER_RANK * nprocs}"] + first, second = (_run_fused_rs_layer(nprocs, args) for _ in range(2)) + _assert_numerics_passed(first) + _assert_numerics_passed(second) + + +@pytest.mark.skipif(not fused_available, reason=reason_for_no_fused) +@pytest.mark.parametrize("nprocs", FUSED_PROC_COUNTS) +def test_fused_layer_declines_ineligible_k(nprocs): + """A K the fused kernels cannot serve has to fall back to no overlap.""" + result = _run_fused_layer( + nprocs, [f"--out-features={INELIGIBLE_OUT_FEATURES_PER_RANK * nprocs}"] + ) + _assert_numerics_passed(result) + stderr = result.stderr.decode() + assert "ineligible shape" not in stderr, stderr + assert "failed to launch" not in stderr, stderr + fused = _reported_names(result.stdout, "UB FUSED NAMES: ") + disabled = _reported_names(result.stdout, "UB DISABLED NAMES: ") + assert fused is not None, f"harness printed no fused name set\n{result.stdout.decode()}" + assert "qkv_dgrad" in fused, fused + assert disabled is not None and "qkv_dgrad" not in disabled, disabled + eligible = _reported_names(result.stdout, "UB BULK ELIGIBLE: ") + assert eligible is not None, f"harness printed no eligibility set\n{result.stdout.decode()}" + assert "qkv_dgrad" not in eligible, eligible + + +@pytest.mark.skipif(not fused_available, reason=reason_for_no_fused) +@pytest.mark.parametrize("nprocs", FUSED_PROC_COUNTS) +def test_fused_layer_declines_unaligned_region(nprocs): + """A Userbuffers region the fused backend cannot serve declines at setup.""" + result = _run_fused_layer( + nprocs, + [f"--out-features={ELIGIBLE_OUT_FEATURES_PER_RANK * nprocs}"], + seq_length=UNALIGNED_SEQ_LENGTH, + ) + _assert_numerics_passed(result) + fused = _reported_names(result.stdout, "UB FUSED NAMES: ") + disabled = _reported_names(result.stdout, "UB DISABLED NAMES: ") + assert fused == set(), f"expected no fused communicators, got {fused}" + assert disabled is not None, f"harness printed no disabled name set\n{result.stdout.decode()}" + assert {"qkv_fprop", "qkv_dgrad", "qkv_wgrad"} <= disabled, disabled diff --git a/transformer_engine/common/comm_gemm_overlap/comm_gemm_overlap.cpp b/transformer_engine/common/comm_gemm_overlap/comm_gemm_overlap.cpp index 0a235ea80b..349965f480 100644 --- a/transformer_engine/common/comm_gemm_overlap/comm_gemm_overlap.cpp +++ b/transformer_engine/common/comm_gemm_overlap/comm_gemm_overlap.cpp @@ -21,6 +21,10 @@ #include "common/util/system.h" #include "userbuffers/userbuffers.h" +#ifdef USE_HIPKITTENS_GEMM +#include "../gemm/kittens/comm_gemm.h" +#endif + #define HALF_BYTES 2 #define UB_MAX_SM 32 @@ -825,10 +829,19 @@ void CommOverlapP2PBase::initialize(const std::vector &buffer_shape, DTy int buffer_chunk_bytes = buffer_bytes / _tp_size; _num_ubuf_chunks = _tp_size; if (_is_reduce_scatter) { - // GEMM + RS overlap: Allocate `2 x tp_size - 1` buffers to hold recieved GEMM chunk - // outputs for reduction at the end of the pipelining. - buffer_bytes = buffer_bytes / _tp_size * (_tp_size * 2 - 1); - _num_ubuf_chunks = _tp_size * 2 - 1; +#ifdef USE_HIPKITTENS_GEMM + if (_fused) { + // Fused RS adds an additional 4K-aligned control block over pipeline. + _num_ubuf_chunks = _tp_size * 2; + buffer_bytes = kittens_fused_rs_region_bytes(buffer_chunk_bytes, _tp_size); + } else +#endif + { + // GEMM + RS overlap: Allocate `2 x tp_size - 1` buffers to hold recieved GEMM chunk + // outputs for reduction at the end of the pipelining. + buffer_bytes = buffer_bytes / _tp_size * (_tp_size * 2 - 1); + _num_ubuf_chunks = _tp_size * 2 - 1; + } } void *buffer_ptr; diff --git a/transformer_engine/common/comm_gemm_overlap/rocm_comm_gemm_overlap.cpp b/transformer_engine/common/comm_gemm_overlap/rocm_comm_gemm_overlap.cpp index e15541ddb8..71faa19ebd 100644 --- a/transformer_engine/common/comm_gemm_overlap/rocm_comm_gemm_overlap.cpp +++ b/transformer_engine/common/comm_gemm_overlap/rocm_comm_gemm_overlap.cpp @@ -15,7 +15,7 @@ #include "common/util/system.h" #include "userbuffers/userbuffers.h" #ifdef USE_HIPKITTENS_GEMM -#include "../gemm/kittens/fused_ag_gemm.h" +#include "../gemm/kittens/comm_gemm.h" #endif namespace transformer_engine { @@ -270,7 +270,7 @@ static bool hk_fused_ag_gemm(const TensorWrapper &A, bool transa, bool transb, T " tp_size=", tp_size, ")"); const int rank_round_tp = comm->myrank - tp_id; - KittensFusedAgGemmArgs args{ + KittensAgGemmArgs args{ A.dptr(), ubuf.dptr(), D.dptr(), reinterpret_cast(comm->gpu_ptrs) + reg * comm->nvsize * sizeof(void *), rank_round_tp % comm->nvsize, comm->nvsize, @@ -282,6 +282,39 @@ static bool hk_fused_ag_gemm(const TensorWrapper &A, bool transa, bool transb, T return kittens_fused_ag_gemm_bf16(args); } +// Fused GEMM + reduce-scatter, launched by fused_overlap_rs below. +static bool hk_fused_rs_gemm(const TensorWrapper &A, bool transa, const TensorWrapper &B, + bool transb, const TensorWrapper &bias, + const TensorWrapper &pre_gelu_out, TensorWrapper &rs_output, + 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, bool *eligible) { + // TODO: Add bias support + NVTE_CHECK(!accumulate && bias.numel() == 0 && pre_gelu_out.numel() == 0, + "fused GEMM+RS reached with an unsupported epilogue"); + NVTE_CHECK(A.dtype() == DType::kBFloat16 && B.dtype() == DType::kBFloat16 && rs_output.dtype() == DType::kBFloat16, + "fused GEMM+RS reached with a non-bf16 operand"); + NVTE_CHECK(transa, "fused GEMM+RS is TN only"); + + 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); + const size_t tokens = n_chunk * tp_size; + + const int rank_round_tp = comm->myrank - tp_id; + KittensRsGemmArgs args{ + B.dptr(), A.dptr(), rs_output.dptr(), ubuf.dptr(), + reinterpret_cast(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(GET_SEND_PTR_BY_INDEX(0, comm, reg, 0) - reinterpret_cast(comm->peer_ptr[0][0])), + static_cast(GET_RECV_PTR_BY_INDEX(1, comm, reg, 0) - GET_RECV_PTR_BY_INDEX(0, comm, reg, 0)), + signal, static_cast(m), static_cast(tokens), static_cast(k), + tp_id, tp_size, chunk.bytes(), workspace.dptr(), workspace.bytes(), stream}; + *eligible = kittens_fused_rs_gemm_eligible(args); + return *eligible && kittens_fused_rs_gemm_bf16(args); +} + // Bulk sibling of hk_fused_ag_gemm. AG is not associated with the GEMM. static bool hk_bulk_ag_gemm(const TensorWrapper &A, bool transa, const TensorWrapper &B, bool transb, TensorWrapper &D, const TensorWrapper &bias, @@ -308,7 +341,7 @@ static bool hk_bulk_ag_gemm(const TensorWrapper &A, bool transa, const TensorWra n_chunk * tp_size, " from the Userbuffers region."); const int rank_round_tp = comm->myrank - tp_id; - KittensFusedAgGemmArgs args{ + KittensAgGemmArgs args{ A.dptr(), B.dptr(), D.dptr(), reinterpret_cast(comm->gpu_ptrs) + reg * comm->nvsize * sizeof(void *), rank_round_tp % comm->nvsize, comm->nvsize, @@ -319,7 +352,76 @@ static bool hk_bulk_ag_gemm(const TensorWrapper &A, bool transa, const TensorWra tp_id, tp_size, chunk.bytes(), workspace.dptr(), workspace.bytes(), stream, ubuf.dptr()}; return kittens_bulk_ag_gemm_bf16(args); } + +// Bulk sibling of hk_fused_rs_gemm. RS is not associated with the GEMM. +static bool hk_bulk_rs_gemm(const TensorWrapper &A, bool transa, const TensorWrapper &B, bool transb, + TensorWrapper &D, const TensorWrapper &bias, + const TensorWrapper &pre_gelu_out, TensorWrapper &workspace, + bool accumulate, const TensorWrapper &rs_output, + const TensorWrapper &ubuf, communicator *comm, int reg, int tp_id, + int tp_size, uint64_t signal, cudaStream_t stream) { + if (transa || !transb || accumulate || bias.numel() != 0 || pre_gelu_out.numel() != 0) { + return false; + } + if (A.dtype() != DType::kBFloat16 || B.dtype() != DType::kBFloat16 || + D.dtype() != DType::kBFloat16 || ubuf.dtype() != DType::kBFloat16) { + return false; + } + if (rs_output.numel() != 0 || ubuf.numel() == 0) { + return false; + } + if (tp_size != 4 && tp_size != 8) { + return false; + } + + const size_t tokens = ubuf.size(0); + const size_t hidden = ubuf.size(1); + if (tokens % (static_cast(tp_size) * 256) != 0 || hidden % 256 != 0) { + return false; + } + + const size_t m = A.size(1); + const size_t k = A.size(0); + const size_t n = B.size(1); + if (m % 256 != 0 || n % 256 != 0 || k % 128 != 0 || k < 256) { + return false; + } + if (B.size(0) != k || k != tokens || m != hidden || D.size(0) != n || D.size(1) != m) { + return false; + } + + const int rank_round_tp = comm->myrank - tp_id; + KittensRsGemmArgs args{ + A.dptr(), B.dptr(), D.dptr(), ubuf.dptr(), + reinterpret_cast(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(GET_SEND_PTR_BY_INDEX(0, comm, reg, 0) - reinterpret_cast(comm->peer_ptr[0][0])), + static_cast(GET_RECV_PTR_BY_INDEX(1, comm, reg, 0) - GET_RECV_PTR_BY_INDEX(0, comm, reg, 0)), + signal, static_cast(m), static_cast(n), static_cast(k), tp_id, tp_size, + tokens / tp_size * hidden * ubuf.element_size(), workspace.dptr(), workspace.bytes(), stream}; + return kittens_bulk_rs_gemm_bf16(args); +} +#endif + +void CommOverlapP2PBase::fused_overlap_bulk_rs(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 &rs_output, + cudaStream_t stream_main) { +#ifdef USE_HIPKITTENS_GEMM + if (kittens_bulk_rs_gemm_supported(cuda::sm_arch())) { + const bool launched = hk_bulk_rs_gemm(A, transa, B, transb, D, bias, pre_gelu_out, workspace, + accumulate, rs_output, _ubuf, _ub_comm, _ub_reg, _tp_id, + _tp_size, _rs_signal_base + _tp_size, stream_main); + NVTE_CHECK(launched, "fused bulk RS failed to launch"); + _rs_signal_base += _tp_size; + return; + } #endif + NVTE_ERROR("fused bulk RS was selected but is not built into this library"); +} void CommOverlapP2PBase::fused_overlap_ag(const TensorWrapper &A, bool transa, const TensorWrapper &B, bool transb, TensorWrapper &D, TensorWrapper &bias, @@ -358,6 +460,31 @@ void CommOverlapP2PBase::fused_overlap_bulk_ag(const TensorWrapper &A, bool tran NVTE_ERROR("fused bulk AG was selected but is not built into this library"); } +void CommOverlapP2PBase::fused_overlap_rs(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 &rs_output, cudaStream_t stream_main) { +#ifdef USE_HIPKITTENS_GEMM + if (kittens_fused_rs_gemm_supported(cuda::sm_arch())) { + bool eligible = false; + const bool launched = hk_fused_rs_gemm(A, transa, B, transb, bias, pre_gelu_out, rs_output, + workspace, accumulate, _ubuf, _ubufs[0], _ub_comm, + _ub_reg, _tp_id, _tp_size, _ag_signal_base + _tp_size, + stream_main, &eligible); + if (!eligible) { + rocm_split_overlap_rs(A, transa, B, transb, D, bias, pre_gelu_out, workspace, grad, + accumulate, use_split_accumulator, rs_output, stream_main); + return; + } + NVTE_CHECK(launched, "fused GEMM+RS failed to launch on an eligible shape"); + _ag_signal_base += _tp_size; + return; + } +#endif + NVTE_ERROR("fused GEMM+RS 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, diff --git a/transformer_engine/common/gemm/kittens/CMakeLists.txt b/transformer_engine/common/gemm/kittens/CMakeLists.txt index 18db7c4c55..f362ac5709 100644 --- a/transformer_engine/common/gemm/kittens/CMakeLists.txt +++ b/transformer_engine/common/gemm/kittens/CMakeLists.txt @@ -84,7 +84,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 + cdna4/overlap/comm_gemm.cpp FLAGS -gline-tables-only SOURCE_FLAGS cdna4/blockwise_fp8_gemm.cpp "-ffast-math") diff --git a/transformer_engine/common/gemm/kittens/cdna4/overlap/bulk_rs_gemm_nt.cuh b/transformer_engine/common/gemm/kittens/cdna4/overlap/bulk_rs_gemm_nt.cuh new file mode 100644 index 0000000000..8cb0c0fc64 --- /dev/null +++ b/transformer_engine/common/gemm/kittens/cdna4/overlap/bulk_rs_gemm_nt.cuh @@ -0,0 +1,444 @@ +/************************************************************************* + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + * License for AMD contributions = MIT. See LICENSE for more information +*************************************************************************/ +#pragma once + +#include "hip/hip_runtime.h" +#include "kittens.cuh" +#include "overlap_common.cuh" +#include +#include +#include + +namespace hk_rs_nt { + +using namespace kittens; +using namespace hk_overlap; + +constexpr int BLOCK_SIZE = 256; +constexpr int HALF_BLOCK_SIZE = BLOCK_SIZE / 2; +constexpr int WARPS_M = 2; +constexpr int WARPS_N = 4; +constexpr int REG_BLOCK_M = BLOCK_SIZE / WARPS_M; +constexpr int REG_BLOCK_N = BLOCK_SIZE / WARPS_N; +constexpr int HALF_REG_BLOCK_M = REG_BLOCK_M / 2; +constexpr int HALF_REG_BLOCK_N = REG_BLOCK_N / 2; + +static_assert(WARPS_M * WARPS_N == NUM_WARPS, "bulk warp grid disagrees with hk_overlap"); + + +constexpr int RS_COMM_WG_DEFAULT = 8; +constexpr int RS_BAND_ROWS = 256; + +using _gl_A = gl; +using _gl_B = gl; +using _gl_C = gl; + +using G = kittens::group; + +// Per-PE pointer to each peer's Userbuffers region +struct PeerPtrs { + bf16 *base[8]; +}; + + +__device__ __forceinline__ +void rs_pull_fold(int w, int nred, int bands, int gband, int tp_size, int my_pe, + const PeerPtrs &peers, size_t shard_elems, size_t band_elems) +{ + typedef int v4i __attribute__((ext_vector_type(4))); + const size_t lines = band_elems / 8; // one int4 = 8 bf16 + const size_t soff = (size_t)my_pe * shard_elems; // our shard, same offset on every rank + bf16 *const out = peers.base[my_pe] + soff; + + for (int b0 = 0; b0 < bands; b0 += gband) { + const int g = (b0 + gband <= bands) ? gband : (bands - b0); + const size_t lines_g = lines * (size_t)g; + const size_t per_g = (lines_g + (size_t)nred - 1) / (size_t)nred; + const size_t l0 = (size_t)w * per_g; + if (l0 >= lines_g) continue; + const size_t l1 = (l0 + per_g < lines_g) ? (l0 + per_g) : lines_g; + const size_t boff = (size_t)b0 * band_elems; + v4i *const dst = (v4i *)(out + boff); + for (size_t l = l0 + threadIdx.x; l < l1; l += blockDim.x) { + v4i acc; +#pragma unroll 1 + for (int s = 0; s < tp_size; s++) { + const bf16 *base = peers.base[s] + soff + boff; + const v4i v = ((const v4i *)base)[l]; + if (s == 0) { + acc = v; + } else { + __hip_bfloat16 *a = reinterpret_cast<__hip_bfloat16 *>(&acc); + const __hip_bfloat16 *x = reinterpret_cast(&v); +#pragma unroll + for (int j = 0; j < 8; j++) a[j] = bf16_add(a[j], x[j]); + } + } + __builtin_nontemporal_store(acc, &dst[l]); + } + } + // Rendezvous the workgroup and retire its stores before the early return. + __syncthreads(); + asm volatile("s_waitcnt vmcnt(0)" ::: "memory"); + __syncthreads(); +} + +struct rs_globals { + _gl_A a; // dY, stored [K=tokens, M=out_local] + _gl_B b; // X, stored [K=tokens, N=hidden] + _gl_C c; // dW, [M=out_local, N=hidden] + gl w; // fp32 split-K partials [splits*M, N]; unused at splits==1 + int splits; + int nred; + int bands; + int gband; // bands folded per cycle; gband == bands is one pass + int tp_size; + int my_pe; + PeerPtrs peers; + size_t shard_elems; // (tokens/tp) * hidden + size_t band_elems; // RS_BAND_ROWS * hidden + hipStream_t stream; + int M = c.rows(); + int N = c.cols(); + int K = a.rows(); + dim3 grid() { return dim3((N / BLOCK_SIZE) * (M / BLOCK_SIZE) * splits + nred); } + dim3 block() { return dim3(NUM_THREADS); } + size_t dynamic_shared_memory() { return MAX_SHARED_MEMORY; } +}; + +__global__ __launch_bounds__(NUM_THREADS, 2) +void wgrad_rs_tk(const rs_globals g, int M, int N, int K) { + extern __shared__ alignment_dummy __shm[]; + shared_allocator al((int*)&__shm[0]); + using ST_AB = st_bf; // natural [K, X] + ST_AB (&As)[2][2] = al.allocate(); + ST_AB (&Bs)[2][2] = al.allocate(); + + rt A_tile; + rt B_tile_0; + rt B_tile_1; + rt_fl C_accum[2][2]; + zero(C_accum[0][0]); + zero(C_accum[0][1]); + zero(C_accum[1][0]); + zero(C_accum[1][1]); + + if ((int)blockIdx.x < g.nred) { + rs_pull_fold((int)blockIdx.x, g.nred, g.bands, g.gband, g.tp_size, g.my_pe, g.peers, + g.shard_elems, g.band_elems); + return; + } + + int wgid = ((blockIdx.y * gridDim.x) + blockIdx.x) - g.nred; + const int NUM_WGS = (int)(gridDim.x * gridDim.y) - g.nred; + + // Split-K + const int splits = g.splits; + const int split_id = wgid % splits; + wgid /= splits; + const int WGM = 8; + // Swizzle chiplet so that wgids are in the same XCD. + wgid = chiplet_transform_chunked(wgid, NUM_WGS / splits, NUM_XCDS, 64); + // Swizzle for better L2 within the same XCD. + const int num_pid_m = ceil_div(M, BLOCK_SIZE); + const int num_pid_n = ceil_div(N, BLOCK_SIZE); + const int num_wgid_in_group = WGM * num_pid_n; + int group_id = wgid / num_wgid_in_group; + int first_pid_m = group_id * WGM; + int group_size_m = min(num_pid_m - first_pid_m, WGM); + int pid_m = first_pid_m + ((wgid % num_wgid_in_group) % group_size_m); + int pid_n = (wgid % num_wgid_in_group) / group_size_m; + int row = pid_m; + int col = pid_n; + + const int warp_id = kittens::warpid(); + const int warp_row = warp_id / 4; + const int warp_col = warp_id % 4; + const int num_tiles = (K / K_STEP) / splits; + const int kt_base = split_id * num_tiles; + + const bf16* a_base = (bf16*)&g.a[{0, 0, 0, 0}]; + const bf16* b_base = (bf16*)&g.b[{0, 0, 0, 0}]; + const int a_row_stride = M * (int)sizeof(bf16); + const int b_row_stride = N * (int)sizeof(bf16); + i32x4 a_srsrc_base = make_srsrc(a_base, K * a_row_stride, a_row_stride); + i32x4 b_srsrc_base = make_srsrc(b_base, K * b_row_stride, b_row_stride); + + const int wid = warpid() % NUM_WARPS; + + constexpr int B_INSTS = (HALF_BLOCK_SIZE * K_STEP * (int)sizeof(bf16)) / (16 * NUM_THREADS); + uint32_t b_lane_off[B_INSTS]; + { + const int lane = kittens::laneid(); +#pragma unroll + for (int i = 0; i < B_INSTS; i++) { + const int q = (wid + i * NUM_WARPS) * 64 + lane; + const int W = q / 32, within = q % 32; + const int nb = W / 4, w = W % 4; + const int r = within / 2, c0 = (within % 2) * 8; + const int n = nb * 16 + c0; + const int k = 32 * (w / 2) + 8 * (r / 4) + (r % 4) + 4 * (w % 2); + b_lane_off[i] = (uint32_t)(((size_t)k * N + n) * sizeof(bf16)); + } + } + constexpr int epw_ab = (16 / sizeof(bf16)) * kittens::WARP_THREADS; + +#define LDSB(t, i, j) static_cast(__builtin_amdgcn_readfirstlane(static_cast( \ + reinterpret_cast(&t[i][j].data[0]) + wid * epw_ab * sizeof(bf16)))) + uint32_t a_lds[2][2] = {{LDSB(As,0,0), LDSB(As,0,1)}, {LDSB(As,1,0), LDSB(As,1,1)}}; + uint32_t b_lds[2][2] = {{LDSB(Bs,0,0), LDSB(Bs,0,1)}, {LDSB(Bs,1,0), LDSB(Bs,1,1)}}; +#undef LDSB + + using T = typename ST_AB::dtype; + constexpr int bpt = ST_AB::underlying_subtile_bytes_per_thread; + constexpr int bpm = bpt * NUM_THREADS; + constexpr int cps = K_STEP * BLOCK_SIZE * sizeof(T) / bpm; + uint32_t sw_A[cps/2], sw_B[cps/2]; + G::prefill_swizzled_offsets(As[0][0], g.a, sw_A); + G::prefill_swizzled_offsets(Bs[0][0], g.b, sw_B); + + + int tic = 0; + int toc = 1; + + + G::load(Bs[tic][0], g.b, {0, 0, kt_base + (0), col * 2 + 0}, sw_B, b_srsrc_base, b_base, b_lds[tic][0]); + G::load(As[tic][0], g.a, {0, 0, kt_base + (0), row * 2 + 0}, sw_A, a_srsrc_base, a_base, a_lds[tic][0]); + G::load(Bs[tic][1], g.b, {0, 0, kt_base + (0), col * 2 + 1}, sw_B, b_srsrc_base, b_base, b_lds[tic][1]); + G::load(As[tic][1], g.a, {0, 0, kt_base + (0), row * 2 + 1}, sw_A, a_srsrc_base, a_base, a_lds[tic][1]); + + if (warp_row == 1) { + __builtin_amdgcn_s_barrier(); + } + + asm volatile("s_waitcnt vmcnt(4)"); + __builtin_amdgcn_s_barrier(); + + G::load(Bs[toc][0], g.b, {0, 0, kt_base + (1), col * 2 + 0}, sw_B, b_srsrc_base, b_base, b_lds[toc][0]); + G::load(As[toc][0], g.a, {0, 0, kt_base + (1), row * 2 + 0}, sw_A, a_srsrc_base, a_base, a_lds[toc][0]); + G::load(Bs[toc][1], g.b, {0, 0, kt_base + (1), col * 2 + 1}, sw_B, b_srsrc_base, b_base, b_lds[toc][1]); + + asm volatile("s_waitcnt vmcnt(6)"); + __builtin_amdgcn_s_barrier(); + +#pragma unroll + for (int tile = 0; tile < num_tiles - 2; tile += 2) { + load(B_tile_0, subtile_inplace(Bs[0][0], {0, warp_col})); + load(A_tile, subtile_inplace(As[0][0], {0, warp_row})); + G::load(As[1][1], g.a, {0, 0, kt_base + (tile + 1), row * 2 + 1}, sw_A, a_srsrc_base, a_base, a_lds[1][1]); + asm volatile("s_waitcnt lgkmcnt(8)"); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + mma_AtB(C_accum[0][0], A_tile, B_tile_0, C_accum[0][0]); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + load(B_tile_1, subtile_inplace(Bs[0][1], {0, warp_col})); + G::load(Bs[0][0], g.b, {0, 0, kt_base + (tile + 2), col * 2 + 0}, sw_B, b_srsrc_base, b_base, b_lds[0][0]); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + mma_AtB(C_accum[0][1], A_tile, B_tile_1, C_accum[0][1]); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + + load(A_tile, subtile_inplace(As[0][1], {0, warp_row})); + G::load(As[0][0], g.a, {0, 0, kt_base + (tile + 2), row * 2 + 0}, sw_A, a_srsrc_base, a_base, a_lds[0][0]); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + mma_AtB(C_accum[1][0], A_tile, B_tile_0, C_accum[1][0]); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + asm volatile("s_waitcnt vmcnt(10)" ::: "memory"); + load(B_tile_0, subtile_inplace(Bs[1][0], {0, warp_col})); + G::load(Bs[0][1], g.b, {0, 0, kt_base + (tile + 2), col * 2 + 1}, sw_B, b_srsrc_base, b_base, b_lds[0][1]); + asm volatile("s_waitcnt vmcnt(6)"); + __builtin_amdgcn_s_barrier(); + + __builtin_amdgcn_s_setprio(1); + mma_AtB(C_accum[1][1], A_tile, B_tile_1, C_accum[1][1]); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + + + load(A_tile, subtile_inplace(As[1][0], {0, warp_row})); + G::load(As[0][1], g.a, {0, 0, kt_base + (tile + 2), row * 2 + 1}, sw_A, a_srsrc_base, a_base, a_lds[0][1]); + asm volatile("s_waitcnt lgkmcnt(8)"); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + mma_AtB(C_accum[0][0], A_tile, B_tile_0, C_accum[0][0]); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + load(B_tile_1, subtile_inplace(Bs[1][1], {0, warp_col})); + G::load(Bs[1][0], g.b, {0, 0, kt_base + (tile + 3), col * 2 + 0}, sw_B, b_srsrc_base, b_base, b_lds[1][0]); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + mma_AtB(C_accum[0][1], A_tile, B_tile_1, C_accum[0][1]); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + + load(A_tile, subtile_inplace(As[1][1], {0, warp_row})); + G::load(As[1][0], g.a, {0, 0, kt_base + (tile + 3), row * 2 + 0}, sw_A, a_srsrc_base, a_base, a_lds[1][0]); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + mma_AtB(C_accum[1][0], A_tile, B_tile_0, C_accum[1][0]); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + G::load(Bs[1][1], g.b, {0, 0, kt_base + (tile + 3), col * 2 + 1}, sw_B, b_srsrc_base, b_base, b_lds[1][1]); + asm volatile("s_waitcnt vmcnt(6)"); + __builtin_amdgcn_s_barrier(); + + __builtin_amdgcn_s_setprio(1); + mma_AtB(C_accum[1][1], A_tile, B_tile_1, C_accum[1][1]); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + } + + { + int tile = num_tiles - 2; + + load(B_tile_0, subtile_inplace(Bs[tic][0], {0, warp_col})); + load(A_tile, subtile_inplace(As[tic][0], {0, warp_row})); + G::load(As[toc][1], g.a, {0, 0, kt_base + (tile + 1), row * 2 + 1}, sw_A, a_srsrc_base, a_base, a_lds[toc][1]); + __builtin_amdgcn_s_barrier(); + asm volatile("s_waitcnt lgkmcnt(0)"); + + __builtin_amdgcn_s_setprio(1); + mma_AtB(C_accum[0][0], A_tile, B_tile_0, C_accum[0][0]); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + + load(B_tile_1, subtile_inplace(Bs[tic][1], {0, warp_col})); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + mma_AtB(C_accum[0][1], A_tile, B_tile_1, C_accum[0][1]); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + + load(A_tile, subtile_inplace(As[tic][1], {0, warp_row})); + asm volatile("s_waitcnt vmcnt(4)"); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + mma_AtB(C_accum[1][0], A_tile, B_tile_0, C_accum[1][0]); + mma_AtB(C_accum[1][1], A_tile, B_tile_1, C_accum[1][1]); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + tic^=1, toc^=1; + } + + { + load(B_tile_0, subtile_inplace(Bs[tic][0], {0, warp_col})); + load(A_tile, subtile_inplace(As[tic][0], {0, warp_row})); + asm volatile("s_waitcnt vmcnt(2)"); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + mma_AtB(C_accum[0][0], A_tile, B_tile_0, C_accum[0][0]); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + + load(B_tile_1, subtile_inplace(Bs[tic][1], {0, warp_col})); + asm volatile("s_waitcnt vmcnt(0)"); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + mma_AtB(C_accum[0][1], A_tile, B_tile_1, C_accum[0][1]); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + + load(A_tile, subtile_inplace(As[tic][1], {0, warp_row})); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + mma_AtB(C_accum[1][0], A_tile, B_tile_0, C_accum[1][0]); + mma_AtB(C_accum[1][1], A_tile, B_tile_1, C_accum[1][1]); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + } + + if (warp_row == 0) { + __builtin_amdgcn_s_barrier(); + } + + const int cr0 = (row * 2) * WARPS_M + warp_row; + const int cr1 = cr0 + WARPS_M; + const int cc0 = col * 2 * WARPS_N + warp_col; + const int cc1 = cc0 + WARPS_N; + if (splits == 1) { + store(g.c, C_accum[0][0], {0, 0, cr0, cc0}); + store(g.c, C_accum[0][1], {0, 0, cr0, cc1}); + store(g.c, C_accum[1][0], {0, 0, cr1, cc0}); + store(g.c, C_accum[1][1], {0, 0, cr1, cc1}); + } else { + const int so = split_id * (M / HALF_REG_BLOCK_M); + store(g.w, C_accum[0][0], {0, 0, so + cr0, cc0}); + store(g.w, C_accum[0][1], {0, 0, so + cr0, cc1}); + store(g.w, C_accum[1][0], {0, 0, so + cr1, cc0}); + store(g.w, C_accum[1][1], {0, 0, so + cr1, cc1}); + } +} + +// Deterministic split-K sum +__global__ void rs_reduce_splits(const float *__restrict__ w, uint16_t *__restrict__ c, + size_t n, int splits) { + size_t i = blockIdx.x * (size_t)blockDim.x + threadIdx.x; + if (i >= n) return; + float acc = 0.f; + for (int s = 0; s < splits; s++) acc += w[(size_t)s * n + i]; + union { float f; unsigned u; } cv; cv.f = acc; + c[i] = (uint16_t)(cv.u >> 16); +} + +static inline int select_split_k_shape(int M, int N, int K) { + const int nt_tiles = (M / BLOCK_SIZE) * (N / BLOCK_SIZE); + int splits = (nt_tiles <= 64) ? 4 : 1; + if ((K / K_STEP) % splits != 0) splits = 1; + return splits; +} + + +static inline bool bands_tile_shard(size_t shard_elems, int bands, size_t band_elems) { + if (bands <= 0 || band_elems == 0) return false; + if (band_elems % 8) return false; + return (size_t)bands * band_elems == shard_elems; +} + +static inline void dispatch(rs_globals g) { + unsigned long mem_size = g.dynamic_shared_memory(); + hipFuncSetAttribute((void *)wgrad_rs_tk, hipFuncAttributeMaxDynamicSharedMemorySize, mem_size); + wgrad_rs_tk<<>>(g, g.M, g.N, g.K); + if (g.splits > 1) { + const size_t n = (size_t)g.M * g.N; + rs_reduce_splits<<<(unsigned)((n + 255) / 256), 256, 0, g.stream>>>( + (const float *)g.w.raw_ptr, (uint16_t *)g.c.raw_ptr, n, g.splits); + } +} + +} // namespace hk_rs_nt diff --git a/transformer_engine/common/gemm/kittens/cdna4/fused_ag_gemm.cpp b/transformer_engine/common/gemm/kittens/cdna4/overlap/comm_gemm.cpp similarity index 55% rename from transformer_engine/common/gemm/kittens/cdna4/fused_ag_gemm.cpp rename to transformer_engine/common/gemm/kittens/cdna4/overlap/comm_gemm.cpp index 61be812500..15ade3a3aa 100644 --- a/transformer_engine/common/gemm/kittens/cdna4/fused_ag_gemm.cpp +++ b/transformer_engine/common/gemm/kittens/cdna4/overlap/comm_gemm.cpp @@ -4,9 +4,11 @@ *************************************************************************/ #include "hip/hip_runtime.h" -#include "../fused_ag_gemm.h" +#include "../../comm_gemm.h" #include "fused_ag_gemm_tn.cuh" #include "fused_ag_gemm_nn.cuh" +#include "bulk_rs_gemm_nt.cuh" +#include "fused_rs_gemm_tn.cuh" #include #include @@ -16,6 +18,28 @@ namespace { +constexpr size_t kFusedRsCtrlHalf = 32u << 10; +constexpr size_t kFusedRsCtrlBytes = 2 * kFusedRsCtrlHalf; + +struct FusedRsLayout { + size_t stage_bytes; + size_t recv_off; + size_t arrive_off; + size_t ready_off; + size_t total_bytes; +}; + +inline FusedRsLayout fused_rs_layout(size_t shard_bytes, int tp_size) { + FusedRsLayout l{}; + l.stage_bytes = shard_bytes * static_cast(tp_size); + l.recv_off = l.stage_bytes; + const size_t data = l.stage_bytes * 2; + l.arrive_off = (data + 4095) & ~static_cast(4095); + l.ready_off = l.arrive_off + kFusedRsCtrlHalf; + l.total_bytes = l.arrive_off + kFusedRsCtrlBytes; + return l; +} + uint64_t compute_warn_ticks() { int dev = 0, khz = 0; @@ -88,8 +112,7 @@ const std::vector *peer_bases(const void *peer_ub, int count) { // Peer-dedicated queues win when the queue is walked at most ~2 times and there are enough M-tiles // to fill the buckets. The 1024 is fixed, not derived from the grid cap. int auto_xcd_bucket(int num_tiles, int tiles_m) { - const int grid_cap = getenv("HK_GRID_CAP") ? atoi(getenv("HK_GRID_CAP")) : 256; - return (grid_cap > 0 && num_tiles <= 1024 && tiles_m >= 5) ? 1 : 0; + return (num_tiles <= 1024 && tiles_m >= 5) ? 1 : 0; } // Gatherer width. Tuned based off AG BW / GEMM FLOPS ratio. @@ -112,8 +135,8 @@ int gath_wg_bulk(int k, int tp_size) { return GATH_WG; } -template -bool upload_plan(AgPlan &plan, const std::vector &queue) { +template +bool upload_plan(Plan &plan, const std::vector &queue) { const size_t bytes = queue.size() * sizeof(TD); if (hipMalloc(&plan.queue, bytes) != hipSuccess) return false; if (hipMemcpy(plan.queue, queue.data(), bytes, hipMemcpyHostToDevice) != hipSuccess) { @@ -136,7 +159,7 @@ struct Carve { bool fits() const { return used <= cap; } }; -bool run_tn(const KittensFusedAgGemmArgs &args) { +bool run_tn(const KittensAgGemmArgs &args) { using namespace hk_ag_tn; const int M = args.n; @@ -218,7 +241,7 @@ struct NnSetup { float *cw; }; -bool prepare_nn(const KittensFusedAgGemmArgs &args, int S, void *peer_local, NnSetup &out) { +bool prepare_nn(const KittensAgGemmArgs &args, int S, void *peer_local, NnSetup &out) { using namespace hk_ag_nn; const int M = args.n; @@ -283,7 +306,7 @@ bool prepare_nn(const KittensFusedAgGemmArgs &args, int S, void *peer_local, NnS return true; } -int split_k_nn(const KittensFusedAgGemmArgs &args) { +int split_k_nn(const KittensAgGemmArgs &args) { using namespace hk_ag_nn; const int M = args.n; const int N_TOTAL = args.m; @@ -293,7 +316,7 @@ int split_k_nn(const KittensFusedAgGemmArgs &args) { return S; } -bool run_nn(const KittensFusedAgGemmArgs &args) { +bool run_nn(const KittensAgGemmArgs &args) { using namespace hk_ag_nn; const int M = args.n; @@ -321,7 +344,7 @@ bool run_nn(const KittensFusedAgGemmArgs &args) { return hipGetLastError() == hipSuccess; } -bool run_bulk_nn(const KittensFusedAgGemmArgs &args) { +bool run_bulk_nn(const KittensAgGemmArgs &args) { using namespace hk_ag_nn; const int M = args.n; @@ -349,8 +372,220 @@ bool run_bulk_nn(const KittensFusedAgGemmArgs &args) { return hipGetLastError() == hipSuccess; } -// Shape and pointer requirements shared by all entry points -bool guards_ok(const KittensFusedAgGemmArgs &args) { +int rs_comm_wg_bulk(int out_local, int tp_size) { + if (tp_size != 8) return hk_rs_nt::RS_COMM_WG_DEFAULT; + if (out_local >= 8192) return 2; + if (out_local >= 3584) return 4; + return hk_rs_nt::RS_COMM_WG_DEFAULT; +} + +static_assert(hk_rs_nt::BLOCK_SIZE == 256 && hk_rs_nt::K_STEP == 64 && hk_rs_nt::NUM_XCDS == 8, + "bulk_rs_shape_ok literals are stale against hk_rs_nt geometry"); + +bool bulk_rs_shape_ok(int M, int N, int K, int splits, int nred) { + const int nt = (splits >= 1 && (K / 64) % splits == 0) ? (K / 64) / splits : 0; + + // Order matters here + return M % 256 == 0 && N % 256 == 0 && K % 64 == 0 && + splits >= 1 && (K / 64) % splits == 0 && + nt >= 4 && (nt & 1) == 0 && + nred >= 0 && nred % 8 == 0; +} + +bool run_bulk_rs(const KittensRsGemmArgs &args) { + using namespace hk_rs_nt; + + const int M = args.n; // out_local + const int N = args.m; // hidden + const int K = args.k; // tokens + const int tp_size = args.nranks; + + int splits = select_split_k_shape(M, N, K); + const size_t partial_bytes = static_cast(splits) * M * N * sizeof(float); + if (splits > 1 && partial_bytes > args.workspace_size) splits = 1; + + const int wgs = (tp_size - 1) * rs_comm_wg_bulk(M, tp_size); + const int nred = (wgs + hk_rs_nt::NUM_XCDS - 1) / hk_rs_nt::NUM_XCDS * hk_rs_nt::NUM_XCDS; + + if (!bulk_rs_shape_ok(M, N, K, splits, nred)) return false; + + const size_t shard_elems = args.shard_bytes / sizeof(bf16); + const size_t band_elems = static_cast(RS_BAND_ROWS) * N; + const int bands = static_cast(shard_elems / band_elems); + if (!bands_tile_shard(shard_elems, bands, band_elems)) return false; + + std::lock_guard lock(g_mu); + + const std::vector *bases = peer_bases(args.peer_ub, args.peer_count); + if (!bases) return false; + PeerPtrs peers{}; + for (int c = 0; c < tp_size; c++) { + peers.base[c] = static_cast((*bases)[(args.peer_first + c) % args.peer_count]); + } + peers.base[args.rank] = static_cast(args.ub); + + if (args.arrive_peers && args.arrive_local) { + ag_ready_kernel<<<1, 64, 0, args.stream>>>( + static_cast(const_cast(args.arrive_peers)), args.arrive_offset, + static_cast(args.arrive_local), args.arrive_stride, args.arrive_value, + args.peer_first, args.peer_count, tp_size, ag_ready_warn_ticks()); + } + + rs_globals g{ + _gl_A(static_cast(const_cast(args.B)), 1, 1, K, M), + _gl_B(static_cast(const_cast(args.A)), 1, 1, K, N), + _gl_C(static_cast(args.D), 1, 1, M, N), + gl(static_cast(args.workspace), 1, 1, splits * M, N), + splits, nred, bands, bands, tp_size, args.rank, peers, shard_elems, band_elems, + args.stream}; + dispatch(g); + return hipGetLastError() == hipSuccess; +} + +// Fused TN GEMM + reduce-scatter +struct RsPlan { + void *queue = nullptr; + int num_tiles = 0; + int xcd_bucket = 0; + int off[hk_rs_tn::NUM_XCDS_AFF] = {}; + int cnt[hk_rs_tn::NUM_XCDS_AFF] = {}; + int sched[2] = {}; +}; + +std::map g_rs_plans; + +int rs_comm_wg_tn(int tokens, int hidden, int k_local, int tp_size) { + if (tp_size != 8) return COMM_WG; + if (hidden == 16384 && k_local == 6656 && + (tokens == 16384 || tokens == 32768 || tokens == 65536)) { + return 4; + } + if (hidden == 8192 && k_local == 1024 && (tokens == 32768 || tokens == 65536)) { + return 12; + } + if (hidden == 4096 && k_local == 512) { + if (tokens == 65536) return 16; + if (tokens == 32768) return 12; + } + return COMM_WG; +} + +bool run_fused_rs(const KittensRsGemmArgs &args) { + using namespace hk_rs_tn; + + const int M = args.n; + const int N_TOTAL = args.m; + const int K = args.k; + const int tp_size = args.nranks; + const int tiles_m = M / BLOCK_ROW; + const int bands = (M / tp_size) / BLOCK_ROW; + + std::lock_guard lock(g_mu); + + const PlanKey key{2, M, N_TOTAL, K, tp_size, args.rank, 1}; + auto it = g_rs_plans.find(key); + if (it == g_rs_plans.end()) { + RsPlan plan; + auto queue = build_rs_work_queue(M, N_TOTAL, K, tp_size, args.rank); + plan.num_tiles = static_cast(queue.size()); + plan.xcd_bucket = auto_xcd_bucket(plan.num_tiles, tiles_m); + if (plan.xcd_bucket) { + XcdBuckets bk{}; + queue = bucketize_by_xcd(queue, bk); + for (int b = 0; b < NUM_XCDS_AFF; b++) { + plan.off[b] = bk.off[b]; + plan.cnt[b] = bk.cnt[b]; + } + } + plan.sched[0] = 0; + plan.sched[1] = bands; + if (!upload_plan(plan, queue)) return false; + it = g_rs_plans.emplace(key, plan).first; + } + const RsPlan &plan = it->second; + + const size_t done_bytes = static_cast(tp_size) * bands * sizeof(unsigned int); + Carve ws{static_cast(args.workspace), 0, args.workspace_size}; + int *tile_counter = static_cast(ws.take(sizeof(int))); + int *bucket_ctr = static_cast(ws.take(NUM_XCDS_AFF * sizeof(int))); + unsigned int *done = static_cast(ws.take(done_bytes)); + const size_t counter_bytes = ws.used; + if (!ws.fits()) return false; + + const FusedRsLayout lay = fused_rs_layout(args.shard_bytes, tp_size); + + const std::vector *bases = peer_bases(args.peer_ub, args.peer_count); + if (!bases) return false; + RsPeers peers{}; + for (int c = 0; c < tp_size; c++) { + char *pb = static_cast((*bases)[(args.peer_first + c) % args.peer_count]); + peers.stage[c] = reinterpret_cast(pb); + peers.recv[c] = reinterpret_cast(pb + lay.recv_off); + peers.arrive[c] = reinterpret_cast(pb + lay.arrive_off); + peers.ready[c] = reinterpret_cast(pb + lay.ready_off); + } + char *lb = static_cast(args.ub); + peers.stage[args.rank] = reinterpret_cast(lb); + peers.recv[args.rank] = reinterpret_cast(lb + lay.recv_off); + peers.arrive[args.rank] = reinterpret_cast(lb + lay.arrive_off); + peers.ready[args.rank] = reinterpret_cast(lb + lay.ready_off); + + bf16 *local_stage = peers.stage[args.rank]; + + XcdBuckets buckets{}; + for (int b = 0; b < NUM_XCDS_AFF; b++) { + buckets.off[b] = plan.off[b]; + buckets.cnt[b] = plan.cnt[b]; + } + + if (hipMemsetAsync(args.workspace, 0, counter_bytes, args.stream) != hipSuccess) return false; + if (hipMemsetAsync(lb + lay.arrive_off, 0, kFusedRsCtrlBytes, args.stream) != hipSuccess) { + return false; + } + + if (args.arrive_peers && args.arrive_local) { + ag_ready_kernel<<<1, 64, 0, args.stream>>>( + static_cast(const_cast(args.arrive_peers)), args.arrive_offset, + static_cast(args.arrive_local), args.arrive_stride, args.arrive_value, + args.peer_first, args.peer_count, tp_size, ag_ready_warn_ticks()); + } + + RsLaunchCfg cfg; + cfg.comm_wg = rs_comm_wg_tn(M, N_TOTAL, K, tp_size); + cfg.xcd_bucket = plan.xcd_bucket; + + launch_persistent_rs(M, N_TOTAL, K, static_cast(const_cast(args.A)), + static_cast(const_cast(args.B)), local_stage, + static_cast(args.D), static_cast(plan.queue), + plan.num_tiles, tile_counter, peers, done, args.rank, tp_size, cfg, + buckets, bucket_ctr, args.stream); + return hipGetLastError() == hipSuccess; +} + +// Shape and pointer requirements shared by rs entry points +static_assert(hk_rs_tn::BLOCK_ROW == 256 && hk_rs_tn::BLOCK_COL == 256 && hk_rs_tn::K_STEP == 64, + "rs_guards_ok literals are stale against hk_rs_tn geometry"); + +bool rs_guards_ok(const KittensRsGemmArgs &args) { + const int M = args.n; + const int N_TOTAL = args.m; + const int K = args.k; + const int tp_size = args.nranks; + const int bands = (M / tp_size) / 256; + + // Order matters here + return (tp_size == 4 || tp_size == 8) && tp_size <= args.peer_count && + args.rank >= 0 && args.rank < tp_size && + M % (tp_size * 256) == 0 && N_TOTAL % 256 == 0 && + K % 128 == 0 && K >= 256 && bands >= 1 && + static_cast(tp_size) * bands * sizeof(unsigned int) <= kFusedRsCtrlHalf && + args.shard_bytes != 0 && + args.shard_bytes == static_cast(M) / tp_size * N_TOTAL * sizeof(uint16_t) && + args.workspace && args.ub && args.A && args.B && args.D && args.peer_ub; +} + +// Shape and pointer requirements shared by ag entry points +bool ag_guards_ok(const KittensAgGemmArgs &args) { const int M = args.n; const int N_TOTAL = args.m; const int K = args.k; @@ -366,27 +601,31 @@ bool guards_ok(const KittensFusedAgGemmArgs &args) { } // namespace -void kittens_fused_ag_gemm_reset_cdna4() { +void kittens_persistent_plans_reset_cdna4() { std::lock_guard lock(g_mu); for (auto &kv : g_plans) { if (kv.second.queue) static_cast(hipFree(kv.second.queue)); } + for (auto &kv : g_rs_plans) { + if (kv.second.queue) static_cast(hipFree(kv.second.queue)); + } g_plans.clear(); + g_rs_plans.clear(); g_peers.clear(); } -bool kittens_fused_ag_gemm_bf16_cdna4(const KittensFusedAgGemmArgs &args) { +bool kittens_fused_ag_gemm_bf16_cdna4(const KittensAgGemmArgs &args) { const int M = args.n; const int K = args.k; const int tp_size = args.nranks; - if (!guards_ok(args)) return false; + if (!ag_guards_ok(args)) return false; if (args.chunk_bytes != static_cast(M / tp_size) * K * sizeof(uint16_t)) return false; return args.transa ? run_tn(args) : run_nn(args); } -bool kittens_bulk_ag_gemm_bf16_cdna4(const KittensFusedAgGemmArgs &args) { +bool kittens_bulk_ag_gemm_bf16_cdna4(const KittensAgGemmArgs &args) { const int M = args.n; const int N_TOTAL = args.m; const int tp_size = args.nranks; @@ -394,10 +633,36 @@ bool kittens_bulk_ag_gemm_bf16_cdna4(const KittensFusedAgGemmArgs &args) { // NN only if (args.transa) return false; if (!args.gather_dst) return false; - if (!guards_ok(args)) return false; + if (!ag_guards_ok(args)) return false; if (args.chunk_bytes != static_cast(M / tp_size) * N_TOTAL * sizeof(uint16_t)) { return false; } return run_bulk_nn(args); } + +bool kittens_bulk_rs_gemm_bf16_cdna4(const KittensRsGemmArgs &args) { + const int tp_size = args.nranks; + if ((tp_size != 4 && tp_size != 8) || tp_size > args.peer_count) return false; + if (args.rank < 0 || args.rank >= tp_size) return false; + if (!args.workspace || !args.ub || !args.A || !args.B || !args.D || !args.peer_ub) return false; + if (args.shard_bytes == 0 || args.shard_bytes % sizeof(uint16_t) != 0) return false; + if (args.shard_bytes != static_cast(args.k) / tp_size * args.m * sizeof(uint16_t)) { + return false; + } + + return run_bulk_rs(args); +} + +size_t kittens_fused_rs_region_bytes_cdna4(size_t chunk_bytes, int tp_size) { + return fused_rs_layout(chunk_bytes, tp_size).total_bytes; +} + +bool kittens_fused_rs_gemm_eligible_cdna4(const KittensRsGemmArgs &args) { + return rs_guards_ok(args); +} + +bool kittens_fused_rs_gemm_bf16_cdna4(const KittensRsGemmArgs &args) { + if (!kittens_fused_rs_gemm_eligible_cdna4(args)) return false; + return run_fused_rs(args); +} diff --git a/transformer_engine/common/gemm/kittens/cdna4/fused_ag_gemm_nn.cuh b/transformer_engine/common/gemm/kittens/cdna4/overlap/fused_ag_gemm_nn.cuh similarity index 94% rename from transformer_engine/common/gemm/kittens/cdna4/fused_ag_gemm_nn.cuh rename to transformer_engine/common/gemm/kittens/cdna4/overlap/fused_ag_gemm_nn.cuh index 8015ef67a7..fdf9f0a13a 100644 --- a/transformer_engine/common/gemm/kittens/cdna4/fused_ag_gemm_nn.cuh +++ b/transformer_engine/common/gemm/kittens/cdna4/overlap/fused_ag_gemm_nn.cuh @@ -6,6 +6,7 @@ #include "hip/hip_runtime.h" #include "kittens.cuh" +#include "overlap_common.cuh" #include #include #include @@ -13,22 +14,7 @@ namespace hk_ag_nn { -using namespace kittens; - -constexpr int SCHED_ROUNDS = 2; -constexpr int NUM_WARPS = 8; -constexpr int WARPS_ROW = 2; -constexpr int WARPS_COL = 4; -constexpr int BLOCK_ROW = 256; -constexpr int BLOCK_COL = 256; -constexpr int K_STEP = 64; -constexpr int HALF_ROW = BLOCK_ROW / 2; -constexpr int HALF_COL = BLOCK_COL / 2; -constexpr int REG_M = BLOCK_ROW / WARPS_ROW / 2; -constexpr int REG_N = BLOCK_COL / WARPS_COL / 2; -constexpr int NUM_THREADS = NUM_WARPS * WARP_THREADS; - -using G_group = kittens::group; +using namespace hk_overlap; struct TileDesc { int chunk_id; @@ -37,27 +23,35 @@ struct TileDesc { int ks; }; +using namespace kittens; + +constexpr int SCHED_ROUNDS = 2; + +using G_group = kittens::group; + + // Per-PE pointer to each peer's [M,K] A buffer. struct PeerPtrs { bf16 *base[8]; }; -#ifndef GATH_WG -#define GATH_WG 8 -#endif -constexpr int NUM_XCDS_AFF = 8; -struct XcdBuckets { - int off[NUM_XCDS_AFF]; - int cnt[NUM_XCDS_AFF]; -}; -#define AG_PUBLISH(p) __hip_atomic_fetch_add((p), 1u, __ATOMIC_RELEASE, __HIP_MEMORY_SCOPE_AGENT) -#define AG_SPIN(p) __hip_atomic_load((p), __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT) -#define AG_ACQUIRE(p) ((void)__hip_atomic_load((p), __ATOMIC_ACQUIRE, __HIP_MEMORY_SCOPE_AGENT)) // NT selects nontemporal stores for the gathered shard for performance. +template +__device__ __forceinline__ +void gather_all(int my_pe, int gath_wg, int tiles_per_chunk, char *gb, const PeerPtrs &peers, + size_t chunk_bytes, unsigned int *arrive) { + const int pi = (int)blockIdx.x / gath_wg; + const int sub = (int)blockIdx.x % gath_wg; + const int peer = pi + (pi >= my_pe ? 1 : 0); + for (int tn = 0; tn < tiles_per_chunk; tn++) { + gather_peer_tile(peer, tn, sub, gath_wg, tiles_per_chunk, gb, peers, chunk_bytes, arrive); + } +} + template __device__ __forceinline__ void gather_copy_wg(void *__restrict__ dst, const void *__restrict__ src, size_t nbytes) { @@ -117,21 +111,6 @@ void gather_peer_tile(int peer, int tn, int sub, int gath_wg, int tiles_per_chun __syncthreads(); } -__host__ __device__ __forceinline__ -int tile_xcd(int chunk_id) { return chunk_id & (NUM_XCDS_AFF - 1); } - -template -__device__ __forceinline__ -void gather_all(int my_pe, int gath_wg, int tiles_per_chunk, char *gb, const PeerPtrs &peers, - size_t chunk_bytes, unsigned int *arrive) { - const int pi = (int)blockIdx.x / gath_wg; - const int sub = (int)blockIdx.x % gath_wg; - const int peer = pi + (pi >= my_pe ? 1 : 0); - for (int tn = 0; tn < tiles_per_chunk; tn++) { - gather_peer_tile(peer, tn, sub, gath_wg, tiles_per_chunk, gb, peers, chunk_bytes, arrive); - } -} - template __device__ __forceinline__ void store_c_tile(U *base, const RT &src, int row_unit, int col_unit, int row_stride, int lane) { @@ -243,14 +222,14 @@ void persistent_ag_bf16_gemm_body( }; auto read_B = [&](auto &bt, int p, int j) { const uint32_t addr = (uint32_t)(uintptr_t)(&Bs[p][j][0]) + (uint32_t)(warp_n * 4096 + kittens::laneid() * 8); - #define RDB(I, J, O0, O1) \ +#define RDB(I, J, O0, O1) \ asm volatile("ds_read_b64_tr_b16 %0, %2 offset:" #O0 "\n" \ "ds_read_b64_tr_b16 %1, %2 offset:" #O1 "\n" \ : "=v"(*reinterpret_cast(&bt.tiles[I][J].data[0])), \ "=v"(*reinterpret_cast(&bt.tiles[I][J].data[2])) \ : "v"(addr) : "memory") RDB(0,0, 0, 512); RDB(0,1,1024,1536); RDB(1,0,2048,2560); RDB(1,1,3072,3584); - #undef RDB +#undef RDB }; const int NGATH = (tp_size - 1) * gath_wg; @@ -620,25 +599,10 @@ static std::vector build_work_queue(int M, int N_total, int K, int tp_ return queue; } -// Stable-partition the queue into one contiguous segment per XCD -static std::vector bucketize_by_xcd(const std::vector &q, XcdBuckets &bk) { - std::vector out; - out.reserve(q.size()); - for (int b = 0; b < NUM_XCDS_AFF; b++) { - bk.off[b] = (int)out.size(); - for (size_t i = 0; i < q.size(); i++) { - if (tile_xcd(q[i].chunk_id) == b) out.push_back(q[i]); - } - bk.cnt[b] = (int)out.size() - bk.off[b]; - } - return out; -} - static int ag_grid(int tiles_M, int tiles_N, int ksplit, int tp_size, int gath_wg) { const int NGATH = (tp_size - 1) * gath_wg; - static const int grid_cap = getenv("HK_GRID_CAP") ? atoi(getenv("HK_GRID_CAP")) : 256; int grid = tiles_M * tiles_N * ksplit + NGATH; - if (grid_cap > 0 && grid > grid_cap) grid = grid_cap; + if (grid > GRID_CAP) grid = GRID_CAP; if (grid < NGATH) grid = NGATH; return grid; } diff --git a/transformer_engine/common/gemm/kittens/cdna4/fused_ag_gemm_tn.cuh b/transformer_engine/common/gemm/kittens/cdna4/overlap/fused_ag_gemm_tn.cuh similarity index 93% rename from transformer_engine/common/gemm/kittens/cdna4/fused_ag_gemm_tn.cuh rename to transformer_engine/common/gemm/kittens/cdna4/overlap/fused_ag_gemm_tn.cuh index 80d135a331..b5b2c7d66a 100644 --- a/transformer_engine/common/gemm/kittens/cdna4/fused_ag_gemm_tn.cuh +++ b/transformer_engine/common/gemm/kittens/cdna4/overlap/fused_ag_gemm_tn.cuh @@ -6,6 +6,7 @@ #include "hip/hip_runtime.h" #include "kittens.cuh" +#include "overlap_common.cuh" #include #include #include @@ -13,22 +14,7 @@ namespace hk_ag_tn { -using namespace kittens; - -constexpr int SCHED_ROUNDS = 2; -constexpr int NUM_WARPS = 8; -constexpr int WARPS_ROW = 2; -constexpr int WARPS_COL = 4; -constexpr int BLOCK_ROW = 256; -constexpr int BLOCK_COL = 256; -constexpr int K_STEP = 64; -constexpr int HALF_ROW = BLOCK_ROW / 2; -constexpr int HALF_COL = BLOCK_COL / 2; -constexpr int REG_M = BLOCK_ROW / WARPS_ROW / 2; -constexpr int REG_N = BLOCK_COL / WARPS_COL / 2; -constexpr int NUM_THREADS = NUM_WARPS * WARP_THREADS; - -using G_group = kittens::group; +using namespace hk_overlap; struct TileDesc { int chunk_id; @@ -36,27 +22,35 @@ struct TileDesc { int tile_n; }; +using namespace kittens; + +constexpr int SCHED_ROUNDS = 2; + +using G_group = kittens::group; + + // Per-PE pointer to each peer's [M,K] A buffer. struct PeerPtrs { bf16 *base[8]; }; -#ifndef GATH_WG -#define GATH_WG 8 -#endif -constexpr int NUM_XCDS_AFF = 8; -struct XcdBuckets { - int off[NUM_XCDS_AFF]; - int cnt[NUM_XCDS_AFF]; -}; -#define AG_PUBLISH(p) __hip_atomic_fetch_add((p), 1u, __ATOMIC_RELEASE, __HIP_MEMORY_SCOPE_AGENT) -#define AG_SPIN(p) __hip_atomic_load((p), __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT) -#define AG_ACQUIRE(p) ((void)__hip_atomic_load((p), __ATOMIC_ACQUIRE, __HIP_MEMORY_SCOPE_AGENT)) // NT selects nontemporal stores for the gathered shard for performance. +template +__device__ __forceinline__ +void gather_all(int my_pe, int gath_wg, int tiles_per_chunk, char *gb, const PeerPtrs &peers, + size_t chunk_bytes, unsigned int *arrive) { + const int pi = (int)blockIdx.x / gath_wg; + const int sub = (int)blockIdx.x % gath_wg; + const int peer = pi + (pi >= my_pe ? 1 : 0); + for (int tn = 0; tn < tiles_per_chunk; tn++) { + gather_peer_tile(peer, tn, sub, gath_wg, tiles_per_chunk, gb, peers, chunk_bytes, arrive); + } +} + template __device__ __forceinline__ void gather_copy_wg(void *__restrict__ dst, const void *__restrict__ src, size_t nbytes) { @@ -115,21 +109,6 @@ void gather_peer_tile(int peer, int tn, int sub, int gath_wg, int tiles_per_chun __syncthreads(); } -__host__ __device__ __forceinline__ -int tile_xcd(int chunk_id) { return chunk_id & (NUM_XCDS_AFF - 1); } - -template -__device__ __forceinline__ -void gather_all(int my_pe, int gath_wg, int tiles_per_chunk, char *gb, const PeerPtrs &peers, - size_t chunk_bytes, unsigned int *arrive) { - const int pi = (int)blockIdx.x / gath_wg; - const int sub = (int)blockIdx.x % gath_wg; - const int peer = pi + (pi >= my_pe ? 1 : 0); - for (int tn = 0; tn < tiles_per_chunk; tn++) { - gather_peer_tile(peer, tn, sub, gath_wg, tiles_per_chunk, gb, peers, chunk_bytes, arrive); - } -} - template __device__ __forceinline__ void store_c_tile(U *base, const RT &src, int row_unit, int col_unit, int row_stride, int lane) { @@ -533,20 +512,6 @@ static std::vector build_work_queue(int M, int N_total, int K, int tp_ return queue; } -// Stable-partition the queue into one contiguous segment per XCD. -static std::vector bucketize_by_xcd(const std::vector &q, XcdBuckets &bk) { - std::vector out; - out.reserve(q.size()); - for (int b = 0; b < NUM_XCDS_AFF; b++) { - bk.off[b] = (int)out.size(); - for (size_t i = 0; i < q.size(); i++) { - if (tile_xcd(q[i].chunk_id) == b) out.push_back(q[i]); - } - bk.cnt[b] = (int)out.size() - bk.off[b]; - } - return out; -} - static void launch_persistent(int M, int N_TOTAL, int K, bf16 *d_a, bf16 *d_b, bf16 *d_c, TileDesc *d_queue, int num_tiles, int *d_tile_counter, PeerPtrs peers, unsigned int *d_arrive, int my_pe, int tp_size, int gath_wg, int m_local, size_t chunk_bytes, int xcd_bucket, @@ -561,9 +526,8 @@ static void launch_persistent(int M, int N_TOTAL, int K, bf16 *d_a, bf16 *d_b, b // Gatherers are dedicated (blockIdx < NGATH) and only join the compute queue once their peer's chunk has landed. const int NGATH = (tp_size - 1) * gath_wg; - static const int grid_cap = getenv("HK_GRID_CAP") ? atoi(getenv("HK_GRID_CAP")) : 256; int grid = tiles_M * tiles_N + NGATH; - if (grid_cap > 0 && grid > grid_cap) grid = grid_cap; + if (grid > GRID_CAP) grid = GRID_CAP; if (grid < NGATH) grid = NGATH; persistent_ag_bf16_gemm<<>>( diff --git a/transformer_engine/common/gemm/kittens/cdna4/overlap/fused_rs_gemm_tn.cuh b/transformer_engine/common/gemm/kittens/cdna4/overlap/fused_rs_gemm_tn.cuh new file mode 100644 index 0000000000..f5ba2d8e58 --- /dev/null +++ b/transformer_engine/common/gemm/kittens/cdna4/overlap/fused_rs_gemm_tn.cuh @@ -0,0 +1,617 @@ +/************************************************************************* + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + * License for AMD contributions = MIT. See LICENSE for more information +*************************************************************************/ +#pragma once + +#include "hip/hip_runtime.h" +#include "kittens.cuh" +#include "overlap_common.cuh" +#include +#include +#include +#include +#include +#include + +namespace hk_rs_tn { + +using namespace hk_overlap; + +struct TileDesc { + int chunk_id; + int tile_m; + int tile_n; +}; + +using namespace kittens; +using G_group = kittens::group; + +constexpr int RS_MAX_TP = 8; + +#ifndef COMM_WG +#define COMM_WG 8 +#endif + +// Sentinel selectors, declared here so the epilogue clamp below can see them. +// Poison and clamp are quiet negative NaNs so neither can collide with a hardware NaN. +#define RS_SENT_BF16 0xFFAAu +#define RS_SENT_DW ((unsigned int)RS_SENT_BF16 * 0x00010001u) + +struct RsPeers { + unsigned int *arrive[RS_MAX_TP]; + bf16 *recv[RS_MAX_TP]; + unsigned int *ready[RS_MAX_TP]; + bf16 *stage[RS_MAX_TP]; +}; + + +#define RS_PUBLISH(p) __hip_atomic_fetch_add((p), 1u, __ATOMIC_RELEASE, __HIP_MEMORY_SCOPE_SYSTEM) +#define RS_SPIN(p) __hip_atomic_load((p), __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_SYSTEM) +#define RS_ACQUIRE(p) ((void)__hip_atomic_load((p), __ATOMIC_ACQUIRE, __HIP_MEMORY_SCOPE_SYSTEM)) +#define LC_PUBLISH_TILE(p) __hip_atomic_fetch_add((p), 1u, __ATOMIC_RELEASE, __HIP_MEMORY_SCOPE_SYSTEM) +#define RS_TRACE_ARG + +struct CdWalk { + int q, tn, dq, dr; +}; + +__host__ __device__ __forceinline__ +CdWalk cd_walk_init(int pos, int stride, int tiles_N) { + CdWalk w; + w.q = pos / tiles_N; + w.tn = pos - w.q * tiles_N; + w.dq = stride / tiles_N; + w.dr = stride - w.dq * tiles_N; + return w; +} + +template +__host__ __device__ __forceinline__ +TileDesc cd_walk_desc(CdWalk w, int my_pe, int bands) { + const int j = w.q % TP; + const int tm = w.q / TP; + TileDesc td; + td.chunk_id = (j == 0) ? my_pe : ((j - 1 < my_pe) ? j - 1 : j); + td.tile_m = td.chunk_id * bands + tm; + td.tile_n = w.tn; + return td; +} + +__host__ __device__ __forceinline__ +CdWalk cd_walk_next(CdWalk w, int tiles_N) { + w.q += w.dq; + w.tn += w.dr; + if (w.tn >= tiles_N) { + w.tn -= tiles_N; + w.q += 1; + } + return w; +} + +template +__device__ __forceinline__ +void store_c_tile(U *base, const RT &src, int row_unit, int col_unit, int row_stride, int lane) { + using T = float; + constexpr int packing = 2; + U *dst_ptr = base + (size_t)(row_unit * RT::rows) * row_stride + col_unit * RT::cols; + const int row_offset = RT::base_tile_stride * (lane / RT::base_tile_cols); + const int col_offset = lane % RT::base_tile_cols; + +#pragma unroll + for (int i = 0; i < RT::height; i++) { +#pragma unroll + for (int j = 0; j < RT::width; j++) { + const int col = j * RT::base_tile_cols + col_offset; +#pragma unroll + for (int k = 0; k < RT::base_tile_num_strides; k++) { + const int row = i * RT::base_tile_rows + row_offset + + k * RT::base_tile_elements_per_stride_group; +#pragma unroll + for (int l = 0; l < RT::base_tile_stride / packing; l++) { + const int idx = l + k * RT::base_tile_stride / packing; + dst_ptr[(row + l * 2) * row_stride + col] = + base_types::convertor::convert(src.tiles[i][j].data[idx].x); + dst_ptr[(row + l * 2 + 1) * row_stride + col] = + base_types::convertor::convert(src.tiles[i][j].data[idx].y); + } + } + } + } +} + +typedef int rs_v4i __attribute__((ext_vector_type(4))); +typedef const volatile __attribute__((address_space(1))) rs_v4i *rs_gvol4; + +__device__ __forceinline__ +rs_v4i sent_load16(const rs_v4i *p) { + return *(rs_gvol4)(size_t)p; +} + +// NOTE: a dword compare can match a half of a real value that happens to equal the poison half. +__device__ __forceinline__ +unsigned int sent_slot_pending(const int *__restrict__ v) { + unsigned int hit = 0u; +#pragma unroll + for (int j = 0; j < 4; j++) { + const unsigned int y = (unsigned int)v[j] ^ RS_SENT_DW; + hit |= (y - 0x00010001u) & ~y & 0x80008000u; + } + return hit; +} + +template +__device__ __forceinline__ +void pull_reduce_all_sent(int my_pe, int ncomm, int bands, int tiles_N, + const bf16 *__restrict__ local_stage, const RsPeers &peers, + unsigned int *__restrict__ done, bf16 *__restrict__ out, + size_t band_elems) { + const int w = (int)blockIdx.x; + const size_t lines = band_elems / 8; + const int peer = (w < TP - 1) ? (w + (w >= my_pe ? 1 : 0)) : -1; + typedef int v4i __attribute__((ext_vector_type(4))); + + { + const int b0 = 0; + const size_t lines_g = lines * (size_t)bands; + const size_t per_g = (lines_g + (size_t)ncomm - 1) / (size_t)ncomm; + const size_t g0 = (size_t)w * per_g; + if (g0 >= lines_g) return; + const size_t g1 = (g0 + per_g < lines_g) ? (g0 + per_g) : lines_g; + const size_t soff = ((size_t)my_pe * bands + b0) * band_elems; + v4i *dst = (v4i *)(out + (size_t)b0 * band_elems); + + for (size_t l = g0 + threadIdx.x; l < g1; ) { + v4i acc; + unsigned int pend = 0u; +#pragma unroll 1 + for (int s = 0; s < TP; s++) { + const bf16 *base = (s == my_pe) ? local_stage : peers.stage[s]; + const v4i v = sent_load16((const rs_v4i *)(base + soff) + l); + pend |= sent_slot_pending((const int *)&v); + if (s == 0) { + acc = v; + } else { + __hip_bfloat16 *a = reinterpret_cast<__hip_bfloat16 *>(&acc); + const __hip_bfloat16 *x = reinterpret_cast(&v); +#pragma unroll + for (int j = 0; j < 8; j++) a[j] = bf16_add(a[j], x[j]); + } + } + if (pend != 0u) { + continue; + } + if (NT) { + __builtin_nontemporal_store(acc, &dst[l]); + } else { + dst[l] = acc; + } + l += blockDim.x; + } + } +} + +// Grid layout: [0, ncomm) communication workgroups, the rest draining the tile queue. +template +__global__ __launch_bounds__(NUM_THREADS, 2) +void persistent_rs_bf16_gemm(const gl A, const gl B, + bf16 *__restrict__ local_stage, bf16 *__restrict__ out, + const TileDesc *__restrict__ work_queue, int num_tiles, + int *__restrict__ tile_counter, const RsPeers peers, + unsigned int *__restrict__ done, int my_pe, int comm_wg, int ncomm, + int bands, int tiles_N, + int xcd_bucket, const XcdBuckets buckets, + int *__restrict__ bucket_ctr RS_TRACE_ARG) { + const int M = A.rows(); + const int K = A.cols(); + const int N_TOTAL = B.rows(); + const int k_tiles = K / K_STEP; + (void)M; + + const size_t band_elems = (size_t)BLOCK_ROW * N_TOTAL; + const size_t band_bytes = band_elems * sizeof(bf16); + + using ST_A = st_bf; + using ST_B = st_bf; + using RT_A = rt_bf; + using RT_B = rt_bf; + using RT_C = rt_fl; + + __shared__ ST_A As[2][2]; + __shared__ ST_B Bs[2][2]; + + const int warp_m = warpid() / WARPS_COL; + const int warp_n = warpid() % WARPS_COL; + const int wid = warpid() % NUM_WARPS; + + using T = bf16; + constexpr int bpt = ST_A::underlying_subtile_bytes_per_thread; + constexpr int bpm = bpt * NUM_THREADS; + constexpr int copies_A = HALF_ROW * K_STEP * sizeof(T) / bpm; + constexpr int copies_B = HALF_COL * K_STEP * sizeof(T) / bpm; + + constexpr int elem_per_warp = (16 / sizeof(T)) * kittens::WARP_THREADS; + const uintptr_t lds_warp_off = wid * elem_per_warp * sizeof(T); + + uint32_t a_lds_00 = __builtin_amdgcn_readfirstlane(static_cast(reinterpret_cast(&As[0][0].data[0]) + lds_warp_off)); + uint32_t a_lds_01 = __builtin_amdgcn_readfirstlane(static_cast(reinterpret_cast(&As[0][1].data[0]) + lds_warp_off)); + uint32_t a_lds_10 = __builtin_amdgcn_readfirstlane(static_cast(reinterpret_cast(&As[1][0].data[0]) + lds_warp_off)); + uint32_t a_lds_11 = __builtin_amdgcn_readfirstlane(static_cast(reinterpret_cast(&As[1][1].data[0]) + lds_warp_off)); + uint32_t b_lds_00 = __builtin_amdgcn_readfirstlane(static_cast(reinterpret_cast(&Bs[0][0].data[0]) + lds_warp_off)); + uint32_t b_lds_01 = __builtin_amdgcn_readfirstlane(static_cast(reinterpret_cast(&Bs[0][1].data[0]) + lds_warp_off)); + uint32_t b_lds_10 = __builtin_amdgcn_readfirstlane(static_cast(reinterpret_cast(&Bs[1][0].data[0]) + lds_warp_off)); + uint32_t b_lds_11 = __builtin_amdgcn_readfirstlane(static_cast(reinterpret_cast(&Bs[1][1].data[0]) + lds_warp_off)); + + uint32_t sw_A[copies_A], sw_B[copies_B]; + G_group::prefill_swizzled_offsets(As[0][0], A, sw_A); + G_group::prefill_swizzled_offsets(Bs[0][0], B, sw_B); + + const T *a_base = (const T *)&A[{0, 0, 0, 0}]; + const int a_row_stride = A.template stride<2>() * sizeof(T); + const int b_row_stride = B.template stride<2>() * sizeof(T); + i32x4 a_srd = make_srsrc(a_base, (size_t)M * a_row_stride, a_row_stride); + const T *b_base = (const T *)&B[{0, 0, 0, 0}]; + i32x4 b_srd = make_srsrc(b_base, (size_t)N_TOTAL * b_row_stride, b_row_stride); + + if ((int)blockIdx.x < ncomm) { + // The comm workgroups are the reduce-scatter: they read all eight sources and write `out` once. + (void)comm_wg; (void)band_bytes; + pull_reduce_all_sent(my_pe, ncomm, bands, tiles_N, local_stage, peers, done, out, + band_elems); + } + + const int hy_ncw = (int)gridDim.x - ncomm; + const int hy_vid = (int)blockIdx.x - ncomm; + const int hy_R = (hy_ncw > 0) ? (num_tiles / hy_ncw) : 0; + const int hy_S = hy_R * hy_ncw; + int hy_left = (hy_vid >= 0 && !xcd_bucket) ? hy_R : 0; + int hy_tick = hy_vid; + CdWalk hy_cd = cd_walk_init((hy_vid >= 0) ? hy_vid : 0, (hy_ncw > 0) ? hy_ncw : 1, tiles_N); + while (true) { + __shared__ int s_tile_idx; + const bool hy_static = (hy_left > 0); + if (threadIdx.x == 0) { + if (xcd_bucket) { + // Steal order: own bucket, then the local chunk, then the other XCDs. + int found = -1; + const int b0 = (int)blockIdx.x % NUM_XCDS_AFF; + for (int s = 0; s <= NUM_XCDS_AFF; s++) { + int bb; + if (s == 0) { + bb = b0; + } else if (s == 1) { + bb = my_pe; + } else { + bb = (b0 + s - 1) & (NUM_XCDS_AFF - 1); + } + if (buckets.cnt[bb] == 0) continue; + if (__hip_atomic_load(&bucket_ctr[bb], __ATOMIC_RELAXED, + __HIP_MEMORY_SCOPE_AGENT) >= buckets.cnt[bb]) continue; + const int idx = atomicAdd(&bucket_ctr[bb], 1); + if (idx < buckets.cnt[bb]) { + found = buckets.off[bb] + idx; + break; + } + } + s_tile_idx = (found < 0) ? num_tiles : found; + } else { + if (hy_static) { + s_tile_idx = hy_tick; + } else { + s_tile_idx = hy_S + atomicAdd(tile_counter, 1); + } + } + } + __syncthreads(); + int tile_idx = s_tile_idx; + if (tile_idx >= num_tiles) break; + + TileDesc desc; + if (hy_static) { + desc = cd_walk_desc(hy_cd, my_pe, bands); + hy_cd = cd_walk_next(hy_cd, tiles_N); + hy_tick += hy_ncw; + hy_left -= 1; + } else { + desc = work_queue[tile_idx]; + } + + const int owner = desc.chunk_id; + const int lrow = desc.tile_m - owner * bands; + + int block_row = desc.tile_m; + int block_col = desc.tile_n; + + RT_A a_tile; + RT_B b_tile_0, b_tile_1; + RT_C c00, c01, c10, c11; + zero(c00); zero(c01); zero(c10); zero(c11); + + int tic = 0, toc = 1; + + G_group::load(Bs[tic][0], B, {0, 0, block_col * 2, 0}, sw_B, b_srd, b_base, b_lds_00); + G_group::load(As[tic][0], A, {0, 0, block_row * 2, 0}, sw_A, a_srd, a_base, a_lds_00); + G_group::load(Bs[tic][1], B, {0, 0, block_col * 2 + 1, 0}, sw_B, b_srd, b_base, b_lds_01); + G_group::load(As[tic][1], A, {0, 0, block_row * 2 + 1, 0}, sw_A, a_srd, a_base, a_lds_01); + + if (warp_m == 1) { + __builtin_amdgcn_s_barrier(); + } + + asm volatile("s_waitcnt vmcnt(4)"); + __builtin_amdgcn_s_barrier(); + + G_group::load(Bs[toc][0], B, {0, 0, block_col * 2, 1}, sw_B, b_srd, b_base, b_lds_10); + G_group::load(As[toc][0], A, {0, 0, block_row * 2, 1}, sw_A, a_srd, a_base, a_lds_10); + G_group::load(Bs[toc][1], B, {0, 0, block_col * 2 + 1, 1}, sw_B, b_srd, b_base, b_lds_11); + + asm volatile("s_waitcnt vmcnt(4)"); + __builtin_amdgcn_s_barrier(); + +#pragma unroll 1 + for (int tile = 0; tile < k_tiles - 2; tile += 2) { + + auto st_b = subtile_inplace(Bs[0][0], {warp_n, 0}); + load(b_tile_0, st_b); + auto st_a = subtile_inplace(As[0][0], {warp_m, 0}); + load(a_tile, st_a); + G_group::load(As[1][1], A, {0, 0, block_row * 2 + 1, tile + 1}, sw_A, a_srd, a_base, a_lds_11); + asm volatile("s_waitcnt lgkmcnt(8)"); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + mma_ABt(c00, a_tile, b_tile_0, c00); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + st_b = subtile_inplace(Bs[0][1], {warp_n, 0}); + load(b_tile_1, st_b); + G_group::load(Bs[0][0], B, {0, 0, block_col * 2, tile + 2}, sw_B, b_srd, b_base, b_lds_00); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + mma_ABt(c01, a_tile, b_tile_1, c01); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + + st_a = subtile_inplace(As[0][1], {warp_m, 0}); + load(a_tile, st_a); + G_group::load(As[0][0], A, {0, 0, block_row * 2, tile + 2}, sw_A, a_srd, a_base, a_lds_00); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + mma_ABt(c10, a_tile, b_tile_0, c10); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + st_b = subtile_inplace(Bs[1][0], {warp_n, 0}); + load(b_tile_0, st_b); + G_group::load(Bs[0][1], B, {0, 0, block_col * 2 + 1, tile + 2}, sw_B, b_srd, b_base, b_lds_01); + asm volatile("s_waitcnt vmcnt(6)"); + __builtin_amdgcn_s_barrier(); + + __builtin_amdgcn_s_setprio(1); + mma_ABt(c11, a_tile, b_tile_1, c11); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + + st_a = subtile_inplace(As[1][0], {warp_m, 0}); + load(a_tile, st_a); + G_group::load(As[0][1], A, {0, 0, block_row * 2 + 1, tile + 2}, sw_A, a_srd, a_base, a_lds_01); + asm volatile("s_waitcnt lgkmcnt(8)"); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + mma_ABt(c00, a_tile, b_tile_0, c00); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + st_b = subtile_inplace(Bs[1][1], {warp_n, 0}); + load(b_tile_1, st_b); + G_group::load(Bs[1][0], B, {0, 0, block_col * 2, tile + 3}, sw_B, b_srd, b_base, b_lds_10); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + mma_ABt(c01, a_tile, b_tile_1, c01); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + + st_a = subtile_inplace(As[1][1], {warp_m, 0}); + load(a_tile, st_a); + G_group::load(As[1][0], A, {0, 0, block_row * 2, tile + 3}, sw_A, a_srd, a_base, a_lds_10); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + mma_ABt(c10, a_tile, b_tile_0, c10); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + G_group::load(Bs[1][1], B, {0, 0, block_col * 2 + 1, tile + 3}, sw_B, b_srd, b_base, b_lds_11); + asm volatile("s_waitcnt vmcnt(4)"); + __builtin_amdgcn_s_barrier(); + + __builtin_amdgcn_s_setprio(1); + mma_ABt(c11, a_tile, b_tile_1, c11); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + } + + { + int tile = k_tiles - 2; + + auto st_b = subtile_inplace(Bs[tic][0], {warp_n, 0}); + load(b_tile_0, st_b); + auto st_a = subtile_inplace(As[tic][0], {warp_m, 0}); + load(a_tile, st_a); + G_group::load(As[toc][1], A, {0, 0, block_row * 2 + 1, tile + 1}, sw_A, a_srd, a_base, a_lds_11); + __builtin_amdgcn_s_barrier(); + asm volatile("s_waitcnt lgkmcnt(0)"); + + __builtin_amdgcn_s_setprio(1); + mma_ABt(c00, a_tile, b_tile_0, c00); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + + st_b = subtile_inplace(Bs[tic][1], {warp_n, 0}); + load(b_tile_1, st_b); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + mma_ABt(c01, a_tile, b_tile_1, c01); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + + st_a = subtile_inplace(As[tic][1], {warp_m, 0}); + load(a_tile, st_a); + asm volatile("s_waitcnt vmcnt(4)"); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + mma_ABt(c10, a_tile, b_tile_0, c10); + mma_ABt(c11, a_tile, b_tile_1, c11); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + tic ^= 1; toc ^= 1; + } + + { + auto st_b = subtile_inplace(Bs[tic][0], {warp_n, 0}); + load(b_tile_0, st_b); + auto st_a = subtile_inplace(As[tic][0], {warp_m, 0}); + load(a_tile, st_a); + asm volatile("s_waitcnt vmcnt(2)"); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + mma_ABt(c00, a_tile, b_tile_0, c00); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + + st_b = subtile_inplace(Bs[tic][1], {warp_n, 0}); + load(b_tile_1, st_b); + asm volatile("s_waitcnt vmcnt(0)"); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + mma_ABt(c01, a_tile, b_tile_1, c01); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + + st_a = subtile_inplace(As[tic][1], {warp_m, 0}); + load(a_tile, st_a); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + mma_ABt(c10, a_tile, b_tile_0, c10); + mma_ABt(c11, a_tile, b_tile_1, c11); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + } + + if (warp_m == 0) { + __builtin_amdgcn_s_barrier(); + } + + bf16 *sbase = local_stage + ((size_t)owner * bands + lrow) * band_elems; + + const int rf0 = __builtin_amdgcn_readfirstlane(warp_m); + const int rf1 = __builtin_amdgcn_readfirstlane(WARPS_ROW + warp_m); + const int cf0 = __builtin_amdgcn_readfirstlane(block_col * WARPS_COL * 2 + warp_n); + const int cf1 = __builtin_amdgcn_readfirstlane(block_col * WARPS_COL * 2 + WARPS_COL + warp_n); + int lane_epi = kittens::laneid(); + asm volatile("" : "+v"(lane_epi)); + + store_c_tile(sbase, c00, rf0, cf0, N_TOTAL, lane_epi); + store_c_tile(sbase, c01, rf0, cf1, N_TOTAL, lane_epi); + store_c_tile(sbase, c10, rf1, cf0, N_TOTAL, lane_epi); + store_c_tile(sbase, c11, rf1, cf1, N_TOTAL, lane_epi); + + __syncthreads(); + if (threadIdx.x == 0) LC_PUBLISH_TILE(&done[(size_t)owner * bands + lrow]); + __syncthreads(); + } // end persistent loop + +} + +static std::vector build_rs_work_queue(int M, int N_total, int K, int tp_size, int my_pe) { + (void)K; + const int tiles_N = N_total / BLOCK_COL; + const int bands = (M / tp_size) / BLOCK_ROW; + + std::vector queue; + queue.reserve((size_t)tiles_N * bands * tp_size); + + auto emit = [&](int chunk, int tm) { + for (int tn = 0; tn < tiles_N; tn++) { + TileDesc td{}; + td.chunk_id = chunk; + td.tile_m = chunk * bands + tm; + td.tile_n = tn; + queue.push_back(td); + } + }; + + for (int tm = 0; tm < bands; tm++) { + emit(my_pe, tm); + for (int c = 0; c < tp_size; c++) { + if (c != my_pe) { + emit(c, tm); + } + } + } + return queue; +} + +struct RsLaunchCfg { + int comm_wg = COMM_WG; + int xcd_bucket = 0; +}; + +static void launch_persistent_rs(int M, int N_TOTAL, int K, bf16 *d_a, bf16 *d_b, + bf16 *d_local_stage, bf16 *d_out, TileDesc *d_queue, int num_tiles, + int *d_tile_counter, RsPeers peers, unsigned int *d_done, + int my_pe, int tp_size, const RsLaunchCfg &cfg, + XcdBuckets buckets, int *d_bucket_ctr, hipStream_t stream, + unsigned long long *d_trace = nullptr) { + (void)d_trace; + const int tiles_M = M / BLOCK_ROW; + const int tiles_N = N_TOTAL / BLOCK_COL; + const int bands = (M / tp_size) / BLOCK_ROW; + + gl A_gl(d_a, nullptr, nullptr, (size_t)M, (size_t)K); + gl B_gl(d_b, nullptr, nullptr, (size_t)N_TOTAL, (size_t)K); + + const int ncomm = (tp_size - 1) * cfg.comm_wg; + + int grid = tiles_M * tiles_N + ncomm; + if (grid > GRID_CAP) grid = GRID_CAP; + if (grid < ncomm + 1) grid = ncomm + 1; +#define RS_LAUNCH(TPV) \ + persistent_rs_bf16_gemm<<>>( \ + A_gl, B_gl, d_local_stage, d_out, d_queue, num_tiles, d_tile_counter, peers, d_done, \ + my_pe, cfg.comm_wg, ncomm, bands, tiles_N, cfg.xcd_bucket, buckets, d_bucket_ctr) + + switch (tp_size) { + case 8: RS_LAUNCH(8); break; + case 4: RS_LAUNCH(4); break; + default: break; + } +#undef RS_LAUNCH +} + +} // namespace hk_rs_tn diff --git a/transformer_engine/common/gemm/kittens/cdna4/overlap/overlap_common.cuh b/transformer_engine/common/gemm/kittens/cdna4/overlap/overlap_common.cuh new file mode 100644 index 0000000000..1baa7333f0 --- /dev/null +++ b/transformer_engine/common/gemm/kittens/cdna4/overlap/overlap_common.cuh @@ -0,0 +1,70 @@ +/************************************************************************* + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + * License for AMD contributions = MIT. See LICENSE for more information +*************************************************************************/ + +#pragma once + +#include "hip/hip_runtime.h" +#include "kittens.cuh" +#include + +namespace hk_overlap { + +constexpr int BLOCK_ROW = 256; +constexpr int BLOCK_COL = 256; +constexpr int K_STEP = 64; + +constexpr int WARPS_ROW = 2; +constexpr int WARPS_COL = 4; +constexpr int NUM_WARPS = WARPS_ROW * WARPS_COL; + +constexpr int NUM_THREADS = NUM_WARPS * kittens::WARP_THREADS; + +constexpr int REG_M = BLOCK_ROW / WARPS_ROW / 2; +constexpr int REG_N = BLOCK_COL / WARPS_COL / 2; + +constexpr int HALF_ROW = BLOCK_ROW / 2; +constexpr int HALF_COL = BLOCK_COL / 2; + +constexpr int NUM_XCDS_AFF = 8; + +constexpr int GRID_CAP = 256; + +struct XcdBuckets { + int off[NUM_XCDS_AFF]; + int cnt[NUM_XCDS_AFF]; +}; + +__host__ __device__ __forceinline__ +int tile_xcd(int chunk_id) { return chunk_id & (NUM_XCDS_AFF - 1); } + +template +static std::vector bucketize_by_xcd(const std::vector &q, XcdBuckets &bk) { + std::vector out; + out.reserve(q.size()); + for (int b = 0; b < NUM_XCDS_AFF; b++) { + bk.off[b] = (int)out.size(); + for (size_t i = 0; i < q.size(); i++) { + if (tile_xcd(q[i].chunk_id) == b) { + out.push_back(q[i]); + } + } + bk.cnt[b] = (int)out.size() - bk.off[b]; + } + return out; +} + +__device__ __forceinline__ +__hip_bfloat16 bf16_add(__hip_bfloat16 x, __hip_bfloat16 y) { + return __float2bfloat16(__bfloat162float(x) + __bfloat162float(y)); +} + +#ifndef GATH_WG +#define GATH_WG 8 +#endif +#define AG_PUBLISH(p) __hip_atomic_fetch_add((p), 1u, __ATOMIC_RELEASE, __HIP_MEMORY_SCOPE_AGENT) +#define AG_SPIN(p) __hip_atomic_load((p), __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT) +#define AG_ACQUIRE(p) ((void)__hip_atomic_load((p), __ATOMIC_ACQUIRE, __HIP_MEMORY_SCOPE_AGENT)) + +} // namespace hk_overlap diff --git a/transformer_engine/common/gemm/kittens/comm_gemm.h b/transformer_engine/common/gemm/kittens/comm_gemm.h new file mode 100644 index 0000000000..4106b92d3e --- /dev/null +++ b/transformer_engine/common/gemm/kittens/comm_gemm.h @@ -0,0 +1,74 @@ +/************************************************************************* + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + * License for AMD contributions = MIT. See LICENSE for more information +*************************************************************************/ + +#pragma once + +#include +#include +#include + +struct KittensAgGemmArgs { + const void *A; + void *ub; + void *D; + const void *peer_ub; + int peer_first; + int peer_count; + void *arrive_local; + const void *arrive_peers; + size_t arrive_offset; + size_t arrive_stride; + uint64_t arrive_value; + int m, n, k; + bool transa; + int rank, nranks; + size_t chunk_bytes; + void *workspace; + size_t workspace_size; + hipStream_t stream; + void *gather_dst; // Bulk all-gather only +}; + +struct KittensRsGemmArgs { + const void *A; + const void *B; + void *D; + void *ub; + const void *peer_ub; + int peer_first; + int peer_count; + void *arrive_local; + const void *arrive_peers; + size_t arrive_offset; + size_t arrive_stride; + uint64_t arrive_value; + int m, n, k; + int rank, nranks; + size_t shard_bytes; + void *workspace; + size_t workspace_size; + hipStream_t stream; +}; + +bool kittens_fused_ag_gemm_supported(int sm_arch); + +// Drops the cached work-queue plans and peer base pointers +void kittens_comm_gemm_reset(); + +bool kittens_fused_ag_gemm_bf16(const KittensAgGemmArgs &args); + +bool kittens_bulk_ag_gemm_bf16(const KittensAgGemmArgs &args); + +bool kittens_bulk_rs_gemm_supported(int sm_arch); + +bool kittens_bulk_rs_gemm_bf16(const KittensRsGemmArgs &args); + +bool kittens_fused_rs_gemm_supported(int sm_arch); + +bool kittens_fused_rs_gemm_eligible(const KittensRsGemmArgs &args); + +bool kittens_fused_rs_gemm_bf16(const KittensRsGemmArgs &args); + +size_t kittens_fused_rs_region_bytes(size_t chunk_bytes, int tp_size); diff --git a/transformer_engine/common/gemm/kittens/fused_ag_gemm.h b/transformer_engine/common/gemm/kittens/fused_ag_gemm.h deleted file mode 100644 index e85879388d..0000000000 --- a/transformer_engine/common/gemm/kittens/fused_ag_gemm.h +++ /dev/null @@ -1,41 +0,0 @@ -/************************************************************************* - * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. - * License for AMD contributions = MIT. See LICENSE for more information -*************************************************************************/ - -#pragma once - -#include -#include -#include - -struct KittensFusedAgGemmArgs { - const void *A; - void *ub; - void *D; - const void *peer_ub; - int peer_first; - int peer_count; - void *arrive_local; - const void *arrive_peers; - size_t arrive_offset; - size_t arrive_stride; - uint64_t arrive_value; - int m, n, k; - bool transa; - int rank, nranks; - size_t chunk_bytes; - void *workspace; - size_t workspace_size; - hipStream_t stream; - void *gather_dst; // Bulk all-gather only -}; - -bool kittens_fused_ag_gemm_supported(int sm_arch); - -// Drops the cached work-queue plans and peer base pointers -void kittens_fused_ag_gemm_reset(); - -bool kittens_fused_ag_gemm_bf16(const KittensFusedAgGemmArgs &args); - -bool kittens_bulk_ag_gemm_bf16(const KittensFusedAgGemmArgs &args); diff --git a/transformer_engine/common/gemm/kittens/kittens_common.cpp b/transformer_engine/common/gemm/kittens/kittens_common.cpp index 4dd7af2ffb..4df8e72120 100644 --- a/transformer_engine/common/gemm/kittens/kittens_common.cpp +++ b/transformer_engine/common/gemm/kittens/kittens_common.cpp @@ -4,13 +4,17 @@ *************************************************************************/ #include "hip/hip_runtime.h" -#include "fused_ag_gemm.h" +#include "comm_gemm.h" #include "kittens_common.h" #ifdef KITTENS_HAVE_CDNA4 -bool kittens_fused_ag_gemm_bf16_cdna4(const KittensFusedAgGemmArgs &args); -bool kittens_bulk_ag_gemm_bf16_cdna4(const KittensFusedAgGemmArgs &args); -void kittens_fused_ag_gemm_reset_cdna4(); +bool kittens_fused_ag_gemm_bf16_cdna4(const KittensAgGemmArgs &args); +bool kittens_bulk_ag_gemm_bf16_cdna4(const KittensAgGemmArgs &args); +bool kittens_bulk_rs_gemm_bf16_cdna4(const KittensRsGemmArgs &args); +bool kittens_fused_rs_gemm_bf16_cdna4(const KittensRsGemmArgs &args); +bool kittens_fused_rs_gemm_eligible_cdna4(const KittensRsGemmArgs &args); +size_t kittens_fused_rs_region_bytes_cdna4(size_t chunk_bytes, int tp_size); +void kittens_persistent_plans_reset_cdna4(); #endif bool kittens_fused_ag_gemm_supported(int sm_arch) { @@ -23,13 +27,13 @@ bool kittens_fused_ag_gemm_supported(int sm_arch) { return false; } -void kittens_fused_ag_gemm_reset() { +void kittens_comm_gemm_reset() { #ifdef KITTENS_HAVE_CDNA4 - kittens_fused_ag_gemm_reset_cdna4(); + kittens_persistent_plans_reset_cdna4(); #endif } -bool kittens_fused_ag_gemm_bf16(const KittensFusedAgGemmArgs &args) { +bool kittens_fused_ag_gemm_bf16(const KittensAgGemmArgs &args) { #ifdef KITTENS_HAVE_CDNA4 return kittens_fused_ag_gemm_bf16_cdna4(args); #else @@ -38,7 +42,7 @@ bool kittens_fused_ag_gemm_bf16(const KittensFusedAgGemmArgs &args) { #endif } -bool kittens_bulk_ag_gemm_bf16(const KittensFusedAgGemmArgs &args) { +bool kittens_bulk_ag_gemm_bf16(const KittensAgGemmArgs &args) { #ifdef KITTENS_HAVE_CDNA4 return kittens_bulk_ag_gemm_bf16_cdna4(args); #else @@ -47,6 +51,64 @@ bool kittens_bulk_ag_gemm_bf16(const KittensFusedAgGemmArgs &args) { #endif } +bool kittens_bulk_rs_gemm_supported(int sm_arch) { +#ifdef KITTENS_HAVE_CDNA4 + if (sm_arch == 95) { + return true; + } +#endif + static_cast(sm_arch); + return false; +} + +bool kittens_bulk_rs_gemm_bf16(const KittensRsGemmArgs &args) { +#ifdef KITTENS_HAVE_CDNA4 + return kittens_bulk_rs_gemm_bf16_cdna4(args); +#else + static_cast(args); + return false; +#endif +} + +bool kittens_fused_rs_gemm_supported(int sm_arch) { +#ifdef KITTENS_HAVE_CDNA4 + if (sm_arch == 95) { + return true; + } +#endif + static_cast(sm_arch); + return false; +} + + +size_t kittens_fused_rs_region_bytes(size_t chunk_bytes, int tp_size) { +#ifdef KITTENS_HAVE_CDNA4 + return kittens_fused_rs_region_bytes_cdna4(chunk_bytes, tp_size); +#else + static_cast(chunk_bytes); + static_cast(tp_size); + return 0; +#endif +} + +bool kittens_fused_rs_gemm_eligible(const KittensRsGemmArgs &args) { +#ifdef KITTENS_HAVE_CDNA4 + return kittens_fused_rs_gemm_eligible_cdna4(args); +#else + static_cast(args); + return false; +#endif +} + +bool kittens_fused_rs_gemm_bf16(const KittensRsGemmArgs &args) { +#ifdef KITTENS_HAVE_CDNA4 + return kittens_fused_rs_gemm_bf16_cdna4(args); +#else + static_cast(args); + return false; +#endif +} + BlockwiseGemmBackend *BlockwiseGemmBackend::get_for_arch(int sm_arch) { #ifdef KITTENS_HAVE_CDNA4 if (sm_arch == 95) { diff --git a/transformer_engine/common/include/transformer_engine/comm_gemm_overlap.h b/transformer_engine/common/include/transformer_engine/comm_gemm_overlap.h index 5ebe50c462..f38278b4a3 100644 --- a/transformer_engine/common/include/transformer_engine/comm_gemm_overlap.h +++ b/transformer_engine/common/include/transformer_engine/comm_gemm_overlap.h @@ -214,6 +214,22 @@ class CommOverlapCore { cudaStream_t stream_main) { NVTE_ERROR("Operation is not implemented."); } + + virtual void fused_overlap_bulk_rs(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 &rs_output, cudaStream_t stream_main) { + NVTE_ERROR("Operation is not implemented."); + } + + virtual void fused_overlap_rs(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 &rs_output, cudaStream_t stream_main) { + NVTE_ERROR("Operation is not implemented."); + } }; // CommOverlapCore class CommOverlapBase : public CommOverlapCore { @@ -222,6 +238,9 @@ class CommOverlapBase : public CommOverlapCore { bool _rs_overlap_first_gemm; cudaStream_t _stream_comm; cudaEvent_t _start_d2dcopy; +#ifdef __HIP_PLATFORM_AMD__ + uint64_t _rs_signal_base = 0; +#endif private: void initialize(const std::vector &buffer_shape, DType buffer_dtype, @@ -256,6 +275,9 @@ class CommOverlapBase : public CommOverlapCore { bool use_split_accumulator, CommOverlapType comm_type, TensorWrapper &rs_output, cudaStream_t stream_main) override; +#ifdef __HIP_PLATFORM_AMD__ +#endif + void atomic_gemm_overlap_ag(const TensorWrapper &A, bool transa, const TensorWrapper &B, bool transb, TensorWrapper &D, TensorWrapper &bias, TensorWrapper &pre_gelu_out, TensorWrapper &workspace, bool grad, @@ -443,6 +465,24 @@ class CommOverlapP2PBase : public CommOverlapCore { bool accumulate, bool use_split_accumulator, cudaStream_t stream_main) override; + /* + ** ROCm fused bulk ReduceScatter implemented with hipKittens + */ + void fused_overlap_bulk_rs(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 &rs_output, + cudaStream_t stream_main) override; + + /* + ** ROCm fused ReduceScatter + GEMM implemented with hipKittens + */ + void fused_overlap_rs(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 &rs_output, + cudaStream_t stream_main) override; + bool is_fused() override { return _fused; } /* diff --git a/transformer_engine/pytorch/csrc/extensions/gemm.cpp b/transformer_engine/pytorch/csrc/extensions/gemm.cpp index 12689b10ce..c528296746 100644 --- a/transformer_engine/pytorch/csrc/extensions/gemm.cpp +++ b/transformer_engine/pytorch/csrc/extensions/gemm.cpp @@ -351,6 +351,13 @@ std::vector gemm(py::handle A, bool transa, py::handle B, bool trans bias_tensor, te_pre_gelu_out, te_workspace, grad, accumulate, use_split_accumulator, main_stream); }); + } else if (comm_overlap->is_fused() && comm_type.value() == CommOverlapType::RS) { + NVTE_SCOPED_GIL_RELEASE({ + comm_overlap->fused_overlap_bulk_rs(A_tensor, transa, B_tensor, transb, out_tensor, + bias_tensor, te_pre_gelu_out, te_workspace, grad, + accumulate, use_split_accumulator, + extra_output_tensor, main_stream); + }); } else #endif { @@ -407,12 +414,19 @@ std::vector gemm(py::handle A, bool transa, py::handle B, bool trans #ifdef __HIP_PLATFORM_AMD__ if (comm_overlap->is_p2p_overlap()) { NVTE_SCOPED_GIL_RELEASE({ - comm_overlap->rocm_split_overlap_rs(A_tensor, transa, B_tensor, transb, out_tensor, - bias_tensor, te_pre_gelu_out, te_workspace, grad, - accumulate, use_split_accumulator, extra_output_tensor, - main_stream); + if (comm_overlap->is_fused()) { + comm_overlap->fused_overlap_rs(A_tensor, transa, B_tensor, transb, out_tensor, + bias_tensor, te_pre_gelu_out, te_workspace, grad, + accumulate, use_split_accumulator, + extra_output_tensor, main_stream); + } else { + comm_overlap->rocm_split_overlap_rs(A_tensor, transa, B_tensor, transb, out_tensor, + bias_tensor, te_pre_gelu_out, te_workspace, grad, + accumulate, use_split_accumulator, extra_output_tensor, + main_stream); + } }); - } else + } else #endif { NVTE_SCOPED_GIL_RELEASE({ diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index e1dcb8c68e..ff23b97e27 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -9,7 +9,7 @@ #include "pybind.h" #ifdef USE_HIPKITTENS_GEMM -#include "common/gemm/kittens/fused_ag_gemm.h" +#include "common/gemm/kittens/comm_gemm.h" #endif #include @@ -735,7 +735,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "reset_fused_ag_gemm_cache", []() { #ifdef USE_HIPKITTENS_GEMM - kittens_fused_ag_gemm_reset(); + kittens_comm_gemm_reset(); #endif }, "Drop cached fused AG+GEMM peer base pointers"); diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 913c08c0fc..8a2f7c29b8 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -349,12 +349,12 @@ def initialize_ub( "qkv_wgrad", "fc1_wgrad", ] - # gfx950 runs the fused backend or nothing by default. + # gfx950 runs the fused backend by default. _fused_default = get_device_compute_capability() == (9, 5) methods = { "ring_exchange": [] if _fused_default else list(_rocm_layers), "pipeline": [], - # TODO: Investigate issues with qkv_dgrad and fc1_dgrad overlap on ROCm + # Bulk regions are rejected on ROCm. "bulk": [], "fused": list(_rocm_layers) if _fused_default else [], } @@ -432,10 +432,10 @@ def add_ub( if method == "fused" and get_device_compute_capability() != (9, 5): raise ValueError(f"At {name}, `fused` overlap method requires a gfx950 device.") if method == "fused" and is_reduce_scatter: - # TODO: Add RS support. - _ub_disabled_names.add(name) - return - if method == "fused" and not _fused_ub_supported(shape, tp_size, dtype): + if not _fused_rs_ub_supported(shape, tp_size, dtype): + _ub_disabled_names.add(name) + return + elif method == "fused" and not _fused_ub_supported(shape, tp_size, dtype): _ub_disabled_names.add(name) return if with_cublasmp and method in ("bulk", "external", "fused"): @@ -661,6 +661,17 @@ def _fused_ub_supported(shape: Union[list, tuple], tp_size: int, dtype: torch.dt return (shape[0] // tp_size) % 256 == 0 +def _fused_rs_ub_supported(shape: Union[list, tuple], tp_size: int, dtype: torch.dtype) -> bool: + """Whether the fused reduce-scatter backend can serve a region of this shape.""" + if tp_size not in (4, 8): + return False + if dtype != torch.bfloat16: + return False + if shape[0] % (tp_size * 256) != 0: + return False + return shape[1] % 256 == 0 + + def fused_ag_gemm_eligible( name: str, inp: torch.Tensor, @@ -703,6 +714,25 @@ def fused_bulk_ag_eligible( return eligible +def fused_bulk_rs_eligible( + name: str, + inp: torch.Tensor, + weight: torch.Tensor, + dtype: torch.dtype, + tp_size: int, + fp8: bool, +) -> bool: + """Whether this call may use the bulk reduce-scatter overlap.""" + if not IS_HIP_EXTENSION: + return True + eligible = _ub_is_fused(name) and not fp8 and dtype == torch.bfloat16 + if eligible: + m, k, n_chunk = _fused_gemm_dims(inp, weight, is_dgrad=False) + eligible = _fused_gemm_shape_ok(m, k, n_chunk, tp_size) + _ub_fused_bulk_decisions[name] = eligible + return eligible + + def _ub_is_fused(name: str) -> bool: """Whether `name` was configured with the fused overlap method.""" return name in _ub_fused_names diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 868aff6bd8..fbe5469412 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -26,9 +26,11 @@ fill_userbuffers_buffer_for_all_gather, fused_ag_gemm_eligible, fused_bulk_ag_eligible, + fused_bulk_rs_eligible, get_ub, get_ub_is_fp8, is_ub_initialized, + _ub_is_fused, ub_overlap_disabled, using_cublasmp_backend, quantize_weight, @@ -229,6 +231,10 @@ def forward( ub_name + "_dgrad", inp, weight, activation_dtype, tp_size, fp8, ): ub_bulk_dgrad = False + if ub_bulk_wgrad and not fused_bulk_rs_eligible( + ub_name + "_wgrad", inp, weight, activation_dtype, tp_size, fp8, + ): + ub_bulk_wgrad = False if ub_overlap_rs_fprop: ub_obj = get_ub(ub_name + "_fprop", fp8) ub_type = tex.CommOverlapType.RS @@ -1475,7 +1481,9 @@ def __init__( self.ub_overlap_ag_dgrad = False self.ub_overlap_rs_dgrad = False self.ub_bulk_dgrad = False - if ub_overlap_disabled(ub_name + "_wgrad"): + if ub_overlap_disabled(ub_name + "_wgrad") or ( + IS_HIP_EXTENSION and not _ub_is_fused(ub_name + "_wgrad") + ): self.ub_bulk_wgrad = False if any( diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index cd3cd6bf6b..e75ce12ed6 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -30,6 +30,7 @@ get_ub, get_ub_is_fp8, is_ub_initialized, + _ub_is_fused, ub_overlap_disabled, using_cublasmp_backend, quantize_weight, @@ -2139,7 +2140,9 @@ def __init__( if ub_overlap_disabled("fc1_dgrad"): self.ub_overlap_rs_dgrad = False self.ub_bulk_dgrad = False - if ub_overlap_disabled("fc1_wgrad"): + if ub_overlap_disabled("fc1_wgrad") or ( + IS_HIP_EXTENSION and not _ub_is_fused("fc1_wgrad") + ): self.ub_bulk_wgrad = False if any( diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index c630d2938d..0ac951b89f 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -23,6 +23,8 @@ fill_userbuffers_buffer_for_all_gather, fused_ag_gemm_eligible, fused_bulk_ag_eligible, + fused_bulk_rs_eligible, + _ub_is_fused, ub_overlap_disabled, get_dummy_wgrad, get_ub, @@ -1665,7 +1667,9 @@ def __init__( self.ub_overlap_ag_dgrad = False self.ub_overlap_rs_dgrad = False self.ub_bulk_dgrad = False - if ub_overlap_disabled(ub_name + "_wgrad"): + if ub_overlap_disabled(ub_name + "_wgrad") or ( + IS_HIP_EXTENSION and not _ub_is_fused(ub_name + "_wgrad") + ): self.ub_bulk_wgrad = False if any( @@ -2020,6 +2024,11 @@ def forward( self.activation_dtype, self.tp_size, self.fp8, ): ub_bulk_dgrad = False + if ub_bulk_wgrad and not fused_bulk_rs_eligible( + self.ub_name + "_wgrad", inp, weight_tensor, + self.activation_dtype, self.tp_size, self.fp8, + ): + ub_bulk_wgrad = False wgrad_store = self.wgrad_store if self.wgrad_store.delay_wgrad_compute() else None fwd_args = LinearFwdArgs( # tensors diff --git a/transformer_engine/pytorch/transformer.py b/transformer_engine/pytorch/transformer.py index 7829b10b5f..d377e5f3b3 100644 --- a/transformer_engine/pytorch/transformer.py +++ b/transformer_engine/pytorch/transformer.py @@ -1,5 +1,3 @@ -# 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. @@ -11,7 +9,6 @@ from typing import Any, Callable, List, Optional, Tuple, Union import torch -from torch.utils.cpp_extension import IS_HIP_EXTENSION from transformer_engine.pytorch.torch_version import torch_version from transformer_engine.pytorch.module import LayerNormMLP, LayerNorm, RMSNorm @@ -341,10 +338,8 @@ def __init__( ub_overlap_ag: bool = True, ub_overlap_rs: bool = True, ub_overlap_rs_dgrad: bool = False, - # True on every platform: ROCm narrows this per-call in fused_bulk_ag_eligible(), which - # declines unless the layer was registered on the "fused" method -- gfx950 only. ub_bulk_dgrad: bool = True, - ub_bulk_wgrad: bool = not IS_HIP_EXTENSION, + ub_bulk_wgrad: bool = True, bias: bool = True, activation: str = "gelu", activation_params: Optional[dict] = None, From d83daa7bed9fbe0ca89008552386dadabe651495 Mon Sep 17 00:00:00 2001 From: alextmagro Date: Wed, 2 Sep 2026 20:06:37 -0500 Subject: [PATCH 2/3] Address claude comments --- .../distributed/run_layer_with_overlap.py | 17 +++++++ .../distributed/test_rocm_fused_overlap.py | 46 ++++++------------- .../rocm_comm_gemm_overlap.cpp | 4 +- .../gemm/kittens/cdna4/overlap/comm_gemm.cpp | 35 ++++++++++++-- .../cdna4/overlap/fused_rs_gemm_tn.cuh | 5 +- .../transformer_engine/comm_gemm_overlap.h | 5 -- 6 files changed, 66 insertions(+), 46 deletions(-) diff --git a/tests/pytorch/distributed/run_layer_with_overlap.py b/tests/pytorch/distributed/run_layer_with_overlap.py index 0ba1819cef..e631804ce5 100644 --- a/tests/pytorch/distributed/run_layer_with_overlap.py +++ b/tests/pytorch/distributed/run_layer_with_overlap.py @@ -11,6 +11,7 @@ import socket import subprocess import argparse +import hashlib import warnings import pprint import yaml @@ -319,6 +320,18 @@ def _parse_args(argv=None, namespace=None): return args +def _digest(tensors): + """A bit-exact, order-sensitive fingerprint of a tensor list.""" + digest = hashlib.sha256() + for tensor in tensors: + if tensor is None: + digest.update(b"") + continue + raw = tensor.detach().contiguous().cpu().flatten().view(torch.uint8) + digest.update(raw.numpy().tobytes()) + return digest.hexdigest() + + def _compare_tensors(name, test, ref, rtol, atol): # Make sure tensors aren't zero and we don't pass trivially if test.count_nonzero() == 0: @@ -575,6 +588,10 @@ def run_fwd_bwd(model, x): test_grads.append(test_param.grad) names.append(test_name + ".grad") + all_digests = [None] * opts.tp + dist.all_gather_object(all_digests, _digest(test_grads), group=nccl_world) + dist_print("OUTPUT HASH: " + " ".join(all_digests)) + torch.set_rng_state(torch_rng_state) torch.cuda.set_rng_state(cuda_rng_state, torch.device(f"cuda:{LOCAL_RANK}")) if opts.use_cuda_graphs: diff --git a/tests/pytorch/distributed/test_rocm_fused_overlap.py b/tests/pytorch/distributed/test_rocm_fused_overlap.py index 0aed15059a..4129820a7f 100644 --- a/tests/pytorch/distributed/test_rocm_fused_overlap.py +++ b/tests/pytorch/distributed/test_rocm_fused_overlap.py @@ -9,6 +9,7 @@ import pytest import torch import transformer_engine.pytorch as te +import transformer_engine.pytorch.cpp_extensions as tex from torch.utils.cpp_extension import IS_HIP_EXTENSION from transformer_engine.pytorch.utils import get_device_compute_capability @@ -30,11 +31,7 @@ def _fused_shape_ok(nprocs: int) -> bool: - """The minimum both backends need: a 256-aligned per-rank chunk and a 256-aligned N. - - The reduce-scatter backend is stricter still -- its combine is specialised on 8 ranks -- but - that is asserted per test rather than hidden in collection, so a tp=4 decline stays visible. - """ + """The minimum both backends need: 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 @@ -111,6 +108,12 @@ def _reported_names(stdout, prefix): return None +def _output_hashes(stdout): + """The per-rank digests the harness printed, in rank order.""" + prefix = "OUTPUT HASH: " + return [ln.split(prefix, 1)[1].strip() for ln in stdout.decode().splitlines() if prefix in ln] + + def _assert_numerics_passed(result): stdout, stderr = result.stdout.decode(), result.stderr.decode() assert result.returncode == 0, f"non-zero exit\n{stderr}" @@ -148,11 +151,7 @@ def test_fused_ag_overlap_is_deterministic(nprocs): second = _run_fused_ag(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) + first_hashes, second_hashes = _output_hashes(first.stdout), _output_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" @@ -181,10 +180,6 @@ def test_fused_layer_bulk_dgrad_bf16(nprocs): @pytest.mark.skipif(not fused_available, reason=reason_for_no_fused) @pytest.mark.parametrize("nprocs", FUSED_PROC_COUNTS) def test_fused_layer_bulk_wgrad_bf16(nprocs): - """The backward reduce-scatter hidden behind the wgrad GEMM, on the fused method. - - wgrad is in the default fused layer set on gfx950, so this needs no extra configuration. - """ result = _run_fused_layer(nprocs, [f"--out-features={ELIGIBLE_OUT_FEATURES_PER_RANK * nprocs}"]) _assert_numerics_passed(result) fused = _reported_names(result.stdout, "UB FUSED NAMES: ") @@ -196,7 +191,7 @@ def test_fused_layer_bulk_wgrad_bf16(nprocs): def _run_fused_rs_layer(nprocs, extra_args, seq_length=SEQ_LENGTH): - """Run the layer harness on a ROW-parallel Linear, which is the fused reduce-scatter's path.""" + """Run the layer harness on a row-parallel Linear, which is the fused reduce-scatter's path.""" test_cmd = ( _fused_launch_cmd(nprocs) + [ @@ -223,11 +218,7 @@ def _run_fused_rs_layer(nprocs, extra_args, seq_length=SEQ_LENGTH): @pytest.mark.skipif(not fused_available, reason=reason_for_no_fused) @pytest.mark.parametrize("nprocs", FUSED_PROC_COUNTS) def test_fused_rs_overlap_bf16(nprocs): - """bf16 at an aligned shape: the fused reduce-scatter runs and the result is correct. - - The combine is specialised on 8 ranks, so below that the backend must decline and fall back - rather than fail -- both outcomes are correct, and which one applies is asserted on nprocs. - """ + """bf16 at an aligned shape: the fused reduce-scatter runs and the result is correct.""" result = _run_fused_rs_layer( nprocs, [f"--in-features={ELIGIBLE_OUT_FEATURES_PER_RANK * nprocs}"] ) @@ -236,23 +227,13 @@ def test_fused_rs_overlap_bf16(nprocs): assert "failed to launch" not in stderr, stderr disabled = _reported_names(result.stdout, "UB DISABLED NAMES: ") assert disabled is not None, f"harness printed no disabled name set\n{result.stdout.decode()}" - if nprocs == 8: - assert "proj_fprop" not in disabled, disabled - else: - assert "proj_fprop" in disabled, f"tp={nprocs} must decline the fused RS, got {disabled}" + assert "proj_fprop" not in disabled, disabled @pytest.mark.skipif(not fused_available, reason=reason_for_no_fused) @pytest.mark.parametrize("nprocs", FUSED_PROC_COUNTS) @pytest.mark.parametrize("quantization", ("fp8_delayed_scaling", "mxfp8")) def test_fused_rs_overlap_rejects_non_bf16(quantization, nprocs): - """Non-bf16 is outside the backend, and setup must decline rather than reach the kernel. - - The AG sibling asserts a non-zero exit because its region carries the operand dtype straight - into the kernel guard. The RS path is gated earlier -- _fused_rs_ub_supported rejects any - non-bf16 region at setup -- so the layer is disabled and the run stays correct. Either way the - kernel must never see a non-bf16 operand, which is what the stderr assertion pins. - """ if quantization.startswith("fp8") and not fp8_available: pytest.skip(reason_for_no_fp8) if quantization == "mxfp8" and not mxfp8_available: @@ -308,6 +289,9 @@ def test_fused_rs_overlap_is_deterministic(nprocs): first, second = (_run_fused_rs_layer(nprocs, args) for _ in range(2)) _assert_numerics_passed(first) _assert_numerics_passed(second) + first_hashes, second_hashes = _output_hashes(first.stdout), _output_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" @pytest.mark.skipif(not fused_available, reason=reason_for_no_fused) diff --git a/transformer_engine/common/comm_gemm_overlap/rocm_comm_gemm_overlap.cpp b/transformer_engine/common/comm_gemm_overlap/rocm_comm_gemm_overlap.cpp index 71faa19ebd..6fc27f1da8 100644 --- a/transformer_engine/common/comm_gemm_overlap/rocm_comm_gemm_overlap.cpp +++ b/transformer_engine/common/comm_gemm_overlap/rocm_comm_gemm_overlap.cpp @@ -470,7 +470,7 @@ void CommOverlapP2PBase::fused_overlap_rs(const TensorWrapper &A, bool transa, c bool eligible = false; const bool launched = hk_fused_rs_gemm(A, transa, B, transb, bias, pre_gelu_out, rs_output, workspace, accumulate, _ubuf, _ubufs[0], _ub_comm, - _ub_reg, _tp_id, _tp_size, _ag_signal_base + _tp_size, + _ub_reg, _tp_id, _tp_size, _rs_signal_base + _tp_size, stream_main, &eligible); if (!eligible) { rocm_split_overlap_rs(A, transa, B, transb, D, bias, pre_gelu_out, workspace, grad, @@ -478,7 +478,7 @@ void CommOverlapP2PBase::fused_overlap_rs(const TensorWrapper &A, bool transa, c return; } NVTE_CHECK(launched, "fused GEMM+RS failed to launch on an eligible shape"); - _ag_signal_base += _tp_size; + _rs_signal_base += _tp_size; return; } #endif diff --git a/transformer_engine/common/gemm/kittens/cdna4/overlap/comm_gemm.cpp b/transformer_engine/common/gemm/kittens/cdna4/overlap/comm_gemm.cpp index 15ade3a3aa..84293e0a1a 100644 --- a/transformer_engine/common/gemm/kittens/cdna4/overlap/comm_gemm.cpp +++ b/transformer_engine/common/gemm/kittens/cdna4/overlap/comm_gemm.cpp @@ -41,6 +41,22 @@ inline FusedRsLayout fused_rs_layout(size_t shard_bytes, int tp_size) { } +// 90 D3. The fill pattern here and the compare constant in the device object must not drift apart; +// a stale build fails silently, with every slot reading as already-arrived. +bool sentinel_pattern_agrees() { + static_assert(static_cast(RS_SENT_BF16) * 0x00010001u == RS_SENT_DW, + "90 D3: the 16-bit fill pattern and the 32-bit compare pattern disagree."); + static bool checked = false; + static bool agrees = false; + if (!checked) { + unsigned int device_dw = 0; + agrees = hipMemcpyFromSymbol(&device_dw, HIP_SYMBOL(hk_rs_tn::rs_sent_dw_device), + sizeof(device_dw)) == hipSuccess && device_dw == RS_SENT_DW; + checked = true; + } + return agrees; +} + uint64_t compute_warn_ticks() { int dev = 0, khz = 0; if (hipGetDevice(&dev) != hipSuccess) return 0; @@ -543,13 +559,22 @@ bool run_fused_rs(const KittensRsGemmArgs &args) { return false; } - if (args.arrive_peers && args.arrive_local) { - ag_ready_kernel<<<1, 64, 0, args.stream>>>( - static_cast(const_cast(args.arrive_peers)), args.arrive_offset, - static_cast(args.arrive_local), args.arrive_stride, args.arrive_value, - args.peer_first, args.peer_count, tp_size, ag_ready_warn_ticks()); + // Arms the arrival sentinel: the only edge ordering a comm workgroup's read of a peer's stage + // against that peer's epilogue store. Must precede the rendezvous below. + if (!sentinel_pattern_agrees()) return false; + if (hipMemsetD16Async(reinterpret_cast(local_stage), RS_SENT_BF16, + lay.stage_bytes / sizeof(bf16), args.stream) != hipSuccess) { + return false; } + // Mandatory here, unlike the ag paths: it is what quiesces every rank's fill before any rank + // reads, so without it the fill above orders nothing. + if (!args.arrive_peers || !args.arrive_local) return false; + ag_ready_kernel<<<1, 64, 0, args.stream>>>( + static_cast(const_cast(args.arrive_peers)), args.arrive_offset, + static_cast(args.arrive_local), args.arrive_stride, args.arrive_value, + args.peer_first, args.peer_count, tp_size, ag_ready_warn_ticks()); + RsLaunchCfg cfg; cfg.comm_wg = rs_comm_wg_tn(M, N_TOTAL, K, tp_size); cfg.xcd_bucket = plan.xcd_bucket; diff --git a/transformer_engine/common/gemm/kittens/cdna4/overlap/fused_rs_gemm_tn.cuh b/transformer_engine/common/gemm/kittens/cdna4/overlap/fused_rs_gemm_tn.cuh index f5ba2d8e58..512bb31c97 100644 --- a/transformer_engine/common/gemm/kittens/cdna4/overlap/fused_rs_gemm_tn.cuh +++ b/transformer_engine/common/gemm/kittens/cdna4/overlap/fused_rs_gemm_tn.cuh @@ -33,10 +33,9 @@ constexpr int RS_MAX_TP = 8; #define COMM_WG 8 #endif -// Sentinel selectors, declared here so the epilogue clamp below can see them. -// Poison and clamp are quiet negative NaNs so neither can collide with a hardware NaN. -#define RS_SENT_BF16 0xFFAAu +// Arrival poison, a comm workgroup still reading it knows the producing GEMM has not stored there yet #define RS_SENT_DW ((unsigned int)RS_SENT_BF16 * 0x00010001u) +__device__ __constant__ unsigned int rs_sent_dw_device = RS_SENT_DW; struct RsPeers { unsigned int *arrive[RS_MAX_TP]; diff --git a/transformer_engine/common/include/transformer_engine/comm_gemm_overlap.h b/transformer_engine/common/include/transformer_engine/comm_gemm_overlap.h index f38278b4a3..130374d49b 100644 --- a/transformer_engine/common/include/transformer_engine/comm_gemm_overlap.h +++ b/transformer_engine/common/include/transformer_engine/comm_gemm_overlap.h @@ -238,9 +238,6 @@ class CommOverlapBase : public CommOverlapCore { bool _rs_overlap_first_gemm; cudaStream_t _stream_comm; cudaEvent_t _start_d2dcopy; -#ifdef __HIP_PLATFORM_AMD__ - uint64_t _rs_signal_base = 0; -#endif private: void initialize(const std::vector &buffer_shape, DType buffer_dtype, @@ -275,8 +272,6 @@ class CommOverlapBase : public CommOverlapCore { bool use_split_accumulator, CommOverlapType comm_type, TensorWrapper &rs_output, cudaStream_t stream_main) override; -#ifdef __HIP_PLATFORM_AMD__ -#endif void atomic_gemm_overlap_ag(const TensorWrapper &A, bool transa, const TensorWrapper &B, bool transb, TensorWrapper &D, TensorWrapper &bias, From 1548db73e03499f053be56b57d7a195394a42d05 Mon Sep 17 00:00:00 2001 From: alextmagro Date: Wed, 2 Sep 2026 23:49:38 -0500 Subject: [PATCH 3/3] Fix build error, sentinel correctness bug and address PR comments --- .../distributed/run_gemm_with_overlap.py | 5 +- .../gemm/kittens/cdna4/overlap/comm_gemm.cpp | 42 ++- .../cdna4/overlap/fused_ag_gemm_nn.cuh | 21 +- .../cdna4/overlap/fused_ag_gemm_tn.cuh | 243 +--------------- .../cdna4/overlap/fused_rs_gemm_tn.cuh | 262 +----------------- .../kittens/cdna4/overlap/overlap_common.cuh | 25 ++ .../kittens/cdna4/overlap/tn_mainloop.inc | 179 ++++++++++++ .../kittens/cdna4/overlap/tn_prologue.inc | 41 +++ .../gemm/kittens/cdna4/overlap/xcd_steal.inc | 22 ++ 9 files changed, 302 insertions(+), 538 deletions(-) create mode 100644 transformer_engine/common/gemm/kittens/cdna4/overlap/tn_mainloop.inc create mode 100644 transformer_engine/common/gemm/kittens/cdna4/overlap/tn_prologue.inc create mode 100644 transformer_engine/common/gemm/kittens/cdna4/overlap/xcd_steal.inc diff --git a/tests/pytorch/distributed/run_gemm_with_overlap.py b/tests/pytorch/distributed/run_gemm_with_overlap.py index 99c31ca600..db46cbbc1d 100644 --- a/tests/pytorch/distributed/run_gemm_with_overlap.py +++ b/tests/pytorch/distributed/run_gemm_with_overlap.py @@ -63,6 +63,9 @@ def _parse_args(argv=None, namespace=None): parser.add_argument( "-d", "--head-dim", type=int, default=48, help="Dimension of each attention head." ) + parser.add_argument( + "--ffn-hidden-size", type=int, default=0, help="MLP intermediate size; 0 means 4x hidden." + ) parser.add_argument("--seed", type=int, default=42, help="RNG seed.") parser.add_argument( "--quantization", @@ -425,7 +428,7 @@ def dist_print(msg, src=None, info=False, error=False, section=False, group=None # K = MLP intermediate size (usually 4x hidden size) # P = number of devices for sequence/tensor parallelism # NOTE: TE-GEMM is set up to work with a transposed kernels and non-transposed inputs. - ffn_hidden_size = 4 * hidden_size + ffn_hidden_size = opts.ffn_hidden_size if opts.ffn_hidden_size > 0 else 4 * hidden_size if opts.bulk_overlap: # Bulk overlap weight and input tensors are not relevant so they're globally sized local_kernel_t_shape = (ffn_hidden_size, hidden_size) diff --git a/transformer_engine/common/gemm/kittens/cdna4/overlap/comm_gemm.cpp b/transformer_engine/common/gemm/kittens/cdna4/overlap/comm_gemm.cpp index 84293e0a1a..2185881b91 100644 --- a/transformer_engine/common/gemm/kittens/cdna4/overlap/comm_gemm.cpp +++ b/transformer_engine/common/gemm/kittens/cdna4/overlap/comm_gemm.cpp @@ -41,8 +41,6 @@ inline FusedRsLayout fused_rs_layout(size_t shard_bytes, int tp_size) { } -// 90 D3. The fill pattern here and the compare constant in the device object must not drift apart; -// a stale build fails silently, with every slot reading as already-arrived. bool sentinel_pattern_agrees() { static_assert(static_cast(RS_SENT_BF16) * 0x00010001u == RS_SENT_DW, "90 D3: the 16-bit fill pattern and the 32-bit compare pattern disagree."); @@ -470,19 +468,15 @@ struct RsPlan { std::map g_rs_plans; -int rs_comm_wg_tn(int tokens, int hidden, int k_local, int tp_size) { - if (tp_size != 8) return COMM_WG; - if (hidden == 16384 && k_local == 6656 && - (tokens == 16384 || tokens == 32768 || tokens == 65536)) { - return 4; - } - if (hidden == 8192 && k_local == 1024 && (tokens == 32768 || tokens == 65536)) { - return 12; - } - if (hidden == 4096 && k_local == 512) { - if (tokens == 65536) return 16; - if (tokens == 32768) return 12; - } +std::map g_rs_epoch; + +int rs_comm_wg_tn(int tokens, int hidden, int k_local) { + const int tiles = (tokens / hk_rs_tn::BLOCK_ROW) * (hidden / hk_rs_tn::BLOCK_COL); + const int waves = tiles / hk_rs_tn::GRID_CAP; + if (waves < 8) return COMM_WG; + if (k_local >= 4096) return waves >= 16 ? 4 : COMM_WG; + if (k_local <= 768) return waves >= 16 ? 16 : 12; + if (k_local <= 1536) return waves >= 16 ? 12 : COMM_WG; return COMM_WG; } @@ -530,18 +524,20 @@ bool run_fused_rs(const KittensRsGemmArgs &args) { const FusedRsLayout lay = fused_rs_layout(args.shard_bytes, tp_size); + const size_t stage_off = (g_rs_epoch[args.ub]++ & 1ull) ? lay.recv_off : 0; + const std::vector *bases = peer_bases(args.peer_ub, args.peer_count); if (!bases) return false; RsPeers peers{}; for (int c = 0; c < tp_size; c++) { char *pb = static_cast((*bases)[(args.peer_first + c) % args.peer_count]); - peers.stage[c] = reinterpret_cast(pb); + peers.stage[c] = reinterpret_cast(pb + stage_off); peers.recv[c] = reinterpret_cast(pb + lay.recv_off); peers.arrive[c] = reinterpret_cast(pb + lay.arrive_off); peers.ready[c] = reinterpret_cast(pb + lay.ready_off); } char *lb = static_cast(args.ub); - peers.stage[args.rank] = reinterpret_cast(lb); + peers.stage[args.rank] = reinterpret_cast(lb + stage_off); peers.recv[args.rank] = reinterpret_cast(lb + lay.recv_off); peers.arrive[args.rank] = reinterpret_cast(lb + lay.arrive_off); peers.ready[args.rank] = reinterpret_cast(lb + lay.ready_off); @@ -559,16 +555,13 @@ bool run_fused_rs(const KittensRsGemmArgs &args) { return false; } - // Arms the arrival sentinel: the only edge ordering a comm workgroup's read of a peer's stage - // against that peer's epilogue store. Must precede the rendezvous below. + // Arms the sentinel: the only edge ordering a peer's stage read against its epilogue store. if (!sentinel_pattern_agrees()) return false; - if (hipMemsetD16Async(reinterpret_cast(local_stage), RS_SENT_BF16, - lay.stage_bytes / sizeof(bf16), args.stream) != hipSuccess) { + if (hipMemsetD32Async(reinterpret_cast(local_stage), RS_SENT_DW, + lay.stage_bytes / sizeof(unsigned int), args.stream) != hipSuccess) { return false; } - // Mandatory here, unlike the ag paths: it is what quiesces every rank's fill before any rank - // reads, so without it the fill above orders nothing. if (!args.arrive_peers || !args.arrive_local) return false; ag_ready_kernel<<<1, 64, 0, args.stream>>>( static_cast(const_cast(args.arrive_peers)), args.arrive_offset, @@ -576,7 +569,7 @@ bool run_fused_rs(const KittensRsGemmArgs &args) { args.peer_first, args.peer_count, tp_size, ag_ready_warn_ticks()); RsLaunchCfg cfg; - cfg.comm_wg = rs_comm_wg_tn(M, N_TOTAL, K, tp_size); + cfg.comm_wg = rs_comm_wg_tn(M, N_TOTAL, K); cfg.xcd_bucket = plan.xcd_bucket; launch_persistent_rs(M, N_TOTAL, K, static_cast(const_cast(args.A)), @@ -636,6 +629,7 @@ void kittens_persistent_plans_reset_cdna4() { } g_plans.clear(); g_rs_plans.clear(); + g_rs_epoch.clear(); g_peers.clear(); } diff --git a/transformer_engine/common/gemm/kittens/cdna4/overlap/fused_ag_gemm_nn.cuh b/transformer_engine/common/gemm/kittens/cdna4/overlap/fused_ag_gemm_nn.cuh index fdf9f0a13a..5032513ddb 100644 --- a/transformer_engine/common/gemm/kittens/cdna4/overlap/fused_ag_gemm_nn.cuh +++ b/transformer_engine/common/gemm/kittens/cdna4/overlap/fused_ag_gemm_nn.cuh @@ -27,8 +27,6 @@ using namespace kittens; constexpr int SCHED_ROUNDS = 2; -using G_group = kittens::group; - // Per-PE pointer to each peer's [M,K] A buffer. struct PeerPtrs { @@ -246,24 +244,7 @@ void persistent_ag_bf16_gemm_body( __shared__ int s_tile_idx; if (threadIdx.x == 0) { if (xcd_bucket) { - // Steal order: Own bucket, then the local chunk, then the other XCDs. - int found = -1; - const int b0 = (int)blockIdx.x % NUM_XCDS_AFF; - for (int s = 0; s <= NUM_XCDS_AFF; s++) { - int bb; - if (s == 0) bb = b0; - else if (s == 1) bb = my_pe; - else bb = (b0 + s - 1) & (NUM_XCDS_AFF - 1); - if (buckets.cnt[bb] == 0) continue; - if (__hip_atomic_load(&bucket_ctr[bb], __ATOMIC_RELAXED, - __HIP_MEMORY_SCOPE_AGENT) >= buckets.cnt[bb]) continue; - const int idx = atomicAdd(&bucket_ctr[bb], 1); - if (idx < buckets.cnt[bb]) { - found = buckets.off[bb] + idx; - break; - } - } - s_tile_idx = (found < 0) ? num_tiles : found; +#include "xcd_steal.inc" } else { s_tile_idx = static_sched ? (int)(blockIdx.x + (long)sched_iter * gridDim.x) : atomicAdd(tile_counter, 1); diff --git a/transformer_engine/common/gemm/kittens/cdna4/overlap/fused_ag_gemm_tn.cuh b/transformer_engine/common/gemm/kittens/cdna4/overlap/fused_ag_gemm_tn.cuh index b5b2c7d66a..47e660f707 100644 --- a/transformer_engine/common/gemm/kittens/cdna4/overlap/fused_ag_gemm_tn.cuh +++ b/transformer_engine/common/gemm/kittens/cdna4/overlap/fused_ag_gemm_tn.cuh @@ -26,8 +26,6 @@ using namespace kittens; constexpr int SCHED_ROUNDS = 2; -using G_group = kittens::group; - // Per-PE pointer to each peer's [M,K] A buffer. struct PeerPtrs { @@ -150,47 +148,7 @@ void persistent_ag_bf16_gemm(const gl A, const gl; - using ST_B = st_bf; - using RT_A = rt_bf; - using RT_B = rt_bf; - using RT_C = rt_fl; - - __shared__ ST_A As[2][2]; - __shared__ ST_B Bs[2][2]; - - const int warp_m = warpid() / WARPS_COL; - const int warp_n = warpid() % WARPS_COL; - const int wid = warpid() % NUM_WARPS; - - using T = bf16; - constexpr int bpt = ST_A::underlying_subtile_bytes_per_thread; - constexpr int bpm = bpt * NUM_THREADS; - constexpr int copies_A = HALF_ROW * K_STEP * sizeof(T) / bpm; - constexpr int copies_B = HALF_COL * K_STEP * sizeof(T) / bpm; - - constexpr int elem_per_warp = (16 / sizeof(T)) * kittens::WARP_THREADS; - const uintptr_t lds_warp_off = wid * elem_per_warp * sizeof(T); - - uint32_t a_lds_00 = __builtin_amdgcn_readfirstlane(static_cast(reinterpret_cast(&As[0][0].data[0]) + lds_warp_off)); - uint32_t a_lds_01 = __builtin_amdgcn_readfirstlane(static_cast(reinterpret_cast(&As[0][1].data[0]) + lds_warp_off)); - uint32_t a_lds_10 = __builtin_amdgcn_readfirstlane(static_cast(reinterpret_cast(&As[1][0].data[0]) + lds_warp_off)); - uint32_t a_lds_11 = __builtin_amdgcn_readfirstlane(static_cast(reinterpret_cast(&As[1][1].data[0]) + lds_warp_off)); - uint32_t b_lds_00 = __builtin_amdgcn_readfirstlane(static_cast(reinterpret_cast(&Bs[0][0].data[0]) + lds_warp_off)); - uint32_t b_lds_01 = __builtin_amdgcn_readfirstlane(static_cast(reinterpret_cast(&Bs[0][1].data[0]) + lds_warp_off)); - uint32_t b_lds_10 = __builtin_amdgcn_readfirstlane(static_cast(reinterpret_cast(&Bs[1][0].data[0]) + lds_warp_off)); - uint32_t b_lds_11 = __builtin_amdgcn_readfirstlane(static_cast(reinterpret_cast(&Bs[1][1].data[0]) + lds_warp_off)); - - uint32_t sw_A[copies_A], sw_B[copies_B]; - G_group::prefill_swizzled_offsets(As[0][0], A, sw_A); - G_group::prefill_swizzled_offsets(Bs[0][0], B, sw_B); - - const T *a_base = (const T *)&A[{0, 0, 0, 0}]; - const int a_row_stride = A.template stride<2>() * sizeof(T); - const int b_row_stride = B.template stride<2>() * sizeof(T); - i32x4 a_srd = make_srsrc(a_base, (size_t)M * a_row_stride, a_row_stride); - const T *b_base = (const T *)&B[{0, 0, 0, 0}]; - i32x4 b_srd = make_srsrc(b_base, (size_t)N_TOTAL * b_row_stride, b_row_stride); +#include "tn_prologue.inc" bf16 *c_base = (bf16 *)&C[{0, 0, 0, 0}]; const int NGATH = (tp_size - 1) * gath_wg; @@ -205,24 +163,7 @@ void persistent_ag_bf16_gemm(const gl A, const gl= buckets.cnt[bb]) continue; - const int idx = atomicAdd(&bucket_ctr[bb], 1); - if (idx < buckets.cnt[bb]) { - found = buckets.off[bb] + idx; - break; - } - } - s_tile_idx = (found < 0) ? num_tiles : found; +#include "xcd_steal.inc" } else { s_tile_idx = static_sched ? (int)(blockIdx.x + (long)sched_iter * gridDim.x) : atomicAdd(tile_counter, 1); @@ -277,185 +218,7 @@ void persistent_ag_bf16_gemm(const gl A, const gl(Bs[0][0], {warp_n, 0}); - load(b_tile_0, st_b); - auto st_a = subtile_inplace(As[0][0], {warp_m, 0}); - load(a_tile, st_a); - G_group::load(As[1][1], A, {0, 0, block_row * 2 + 1, tile + 1}, sw_A, a_srd, a_base, a_lds_11); - asm volatile("s_waitcnt lgkmcnt(8)"); - __builtin_amdgcn_s_barrier(); - - asm volatile("s_waitcnt lgkmcnt(0)"); - __builtin_amdgcn_s_setprio(1); - mma_ABt(c00, a_tile, b_tile_0, c00); - __builtin_amdgcn_s_setprio(0); - __builtin_amdgcn_s_barrier(); - __builtin_amdgcn_sched_barrier(0); - - st_b = subtile_inplace(Bs[0][1], {warp_n, 0}); - load(b_tile_1, st_b); - G_group::load(Bs[0][0], B, {0, 0, block_col * 2, tile + 2}, sw_B, b_srd, b_base, b_lds_00); - __builtin_amdgcn_s_barrier(); - - asm volatile("s_waitcnt lgkmcnt(0)"); - __builtin_amdgcn_s_setprio(1); - mma_ABt(c01, a_tile, b_tile_1, c01); - __builtin_amdgcn_s_setprio(0); - __builtin_amdgcn_s_barrier(); - - st_a = subtile_inplace(As[0][1], {warp_m, 0}); - load(a_tile, st_a); - G_group::load(As[0][0], A, {0, 0, block_row * 2, tile + 2}, sw_A, a_srd, a_base, a_lds_00); - __builtin_amdgcn_s_barrier(); - - asm volatile("s_waitcnt lgkmcnt(0)"); - __builtin_amdgcn_s_setprio(1); - mma_ABt(c10, a_tile, b_tile_0, c10); - __builtin_amdgcn_s_setprio(0); - __builtin_amdgcn_s_barrier(); - __builtin_amdgcn_sched_barrier(0); - - st_b = subtile_inplace(Bs[1][0], {warp_n, 0}); - load(b_tile_0, st_b); - G_group::load(Bs[0][1], B, {0, 0, block_col * 2 + 1, tile + 2}, sw_B, b_srd, b_base, b_lds_01); - asm volatile("s_waitcnt vmcnt(6)"); - __builtin_amdgcn_s_barrier(); - - __builtin_amdgcn_s_setprio(1); - mma_ABt(c11, a_tile, b_tile_1, c11); - __builtin_amdgcn_s_setprio(0); - __builtin_amdgcn_s_barrier(); - - st_a = subtile_inplace(As[1][0], {warp_m, 0}); - load(a_tile, st_a); - G_group::load(As[0][1], A, {0, 0, block_row * 2 + 1, tile + 2}, sw_A, a_srd, a_base, a_lds_01); - asm volatile("s_waitcnt lgkmcnt(8)"); - __builtin_amdgcn_s_barrier(); - - asm volatile("s_waitcnt lgkmcnt(0)"); - __builtin_amdgcn_s_setprio(1); - mma_ABt(c00, a_tile, b_tile_0, c00); - __builtin_amdgcn_s_setprio(0); - __builtin_amdgcn_s_barrier(); - __builtin_amdgcn_sched_barrier(0); - - st_b = subtile_inplace(Bs[1][1], {warp_n, 0}); - load(b_tile_1, st_b); - G_group::load(Bs[1][0], B, {0, 0, block_col * 2, tile + 3}, sw_B, b_srd, b_base, b_lds_10); - __builtin_amdgcn_s_barrier(); - - asm volatile("s_waitcnt lgkmcnt(0)"); - __builtin_amdgcn_s_setprio(1); - mma_ABt(c01, a_tile, b_tile_1, c01); - __builtin_amdgcn_s_setprio(0); - __builtin_amdgcn_s_barrier(); - - st_a = subtile_inplace(As[1][1], {warp_m, 0}); - load(a_tile, st_a); - G_group::load(As[1][0], A, {0, 0, block_row * 2, tile + 3}, sw_A, a_srd, a_base, a_lds_10); - __builtin_amdgcn_s_barrier(); - - asm volatile("s_waitcnt lgkmcnt(0)"); - __builtin_amdgcn_s_setprio(1); - mma_ABt(c10, a_tile, b_tile_0, c10); - __builtin_amdgcn_s_setprio(0); - __builtin_amdgcn_s_barrier(); - __builtin_amdgcn_sched_barrier(0); - - G_group::load(Bs[1][1], B, {0, 0, block_col * 2 + 1, tile + 3}, sw_B, b_srd, b_base, b_lds_11); - asm volatile("s_waitcnt vmcnt(4)"); - __builtin_amdgcn_s_barrier(); - - __builtin_amdgcn_s_setprio(1); - mma_ABt(c11, a_tile, b_tile_1, c11); - __builtin_amdgcn_s_setprio(0); - __builtin_amdgcn_s_barrier(); - } - - // Epilogue: K-tile k_tiles-2 (stage tic) - { - int tile = k_tiles - 2; - - auto st_b = subtile_inplace(Bs[tic][0], {warp_n, 0}); - load(b_tile_0, st_b); - auto st_a = subtile_inplace(As[tic][0], {warp_m, 0}); - load(a_tile, st_a); - G_group::load(As[toc][1], A, {0, 0, block_row * 2 + 1, tile + 1}, sw_A, a_srd, a_base, a_lds_11); - __builtin_amdgcn_s_barrier(); - asm volatile("s_waitcnt lgkmcnt(0)"); - - __builtin_amdgcn_s_setprio(1); - mma_ABt(c00, a_tile, b_tile_0, c00); - __builtin_amdgcn_s_setprio(0); - __builtin_amdgcn_s_barrier(); - - st_b = subtile_inplace(Bs[tic][1], {warp_n, 0}); - load(b_tile_1, st_b); - __builtin_amdgcn_s_barrier(); - - asm volatile("s_waitcnt lgkmcnt(0)"); - __builtin_amdgcn_s_setprio(1); - mma_ABt(c01, a_tile, b_tile_1, c01); - __builtin_amdgcn_s_setprio(0); - __builtin_amdgcn_s_barrier(); - - st_a = subtile_inplace(As[tic][1], {warp_m, 0}); - load(a_tile, st_a); - asm volatile("s_waitcnt vmcnt(4)"); - __builtin_amdgcn_s_barrier(); - - asm volatile("s_waitcnt lgkmcnt(0)"); - __builtin_amdgcn_s_setprio(1); - mma_ABt(c10, a_tile, b_tile_0, c10); - mma_ABt(c11, a_tile, b_tile_1, c11); - __builtin_amdgcn_s_setprio(0); - __builtin_amdgcn_s_barrier(); - tic ^= 1; toc ^= 1; - } - - // Epilogue: K-tile k_tiles-1 (stage tic after flip) - { - auto st_b = subtile_inplace(Bs[tic][0], {warp_n, 0}); - load(b_tile_0, st_b); - auto st_a = subtile_inplace(As[tic][0], {warp_m, 0}); - load(a_tile, st_a); - asm volatile("s_waitcnt vmcnt(2)"); - __builtin_amdgcn_s_barrier(); - - asm volatile("s_waitcnt lgkmcnt(0)"); - __builtin_amdgcn_s_setprio(1); - mma_ABt(c00, a_tile, b_tile_0, c00); - __builtin_amdgcn_s_setprio(0); - __builtin_amdgcn_s_barrier(); - - st_b = subtile_inplace(Bs[tic][1], {warp_n, 0}); - load(b_tile_1, st_b); - asm volatile("s_waitcnt vmcnt(0)"); - __builtin_amdgcn_s_barrier(); - - asm volatile("s_waitcnt lgkmcnt(0)"); - __builtin_amdgcn_s_setprio(1); - mma_ABt(c01, a_tile, b_tile_1, c01); - __builtin_amdgcn_s_setprio(0); - __builtin_amdgcn_s_barrier(); - - st_a = subtile_inplace(As[tic][1], {warp_m, 0}); - load(a_tile, st_a); - __builtin_amdgcn_s_barrier(); - - asm volatile("s_waitcnt lgkmcnt(0)"); - __builtin_amdgcn_s_setprio(1); - mma_ABt(c10, a_tile, b_tile_0, c10); - mma_ABt(c11, a_tile, b_tile_1, c11); - __builtin_amdgcn_s_setprio(0); - __builtin_amdgcn_s_barrier(); - } - - if (warp_m == 0) { - __builtin_amdgcn_s_barrier(); - } +#include "tn_mainloop.inc" const int rf0 = __builtin_amdgcn_readfirstlane(block_row * WARPS_ROW * 2 + warp_m); const int rf1 = __builtin_amdgcn_readfirstlane(block_row * WARPS_ROW * 2 + WARPS_ROW + warp_m); diff --git a/transformer_engine/common/gemm/kittens/cdna4/overlap/fused_rs_gemm_tn.cuh b/transformer_engine/common/gemm/kittens/cdna4/overlap/fused_rs_gemm_tn.cuh index 512bb31c97..aac7d15ed5 100644 --- a/transformer_engine/common/gemm/kittens/cdna4/overlap/fused_rs_gemm_tn.cuh +++ b/transformer_engine/common/gemm/kittens/cdna4/overlap/fused_rs_gemm_tn.cuh @@ -25,7 +25,6 @@ struct TileDesc { }; using namespace kittens; -using G_group = kittens::group; constexpr int RS_MAX_TP = 8; @@ -34,6 +33,7 @@ constexpr int RS_MAX_TP = 8; #endif // Arrival poison, a comm workgroup still reading it knows the producing GEMM has not stored there yet +#define RS_SENT_BF16 0xFFAAu #define RS_SENT_DW ((unsigned int)RS_SENT_BF16 * 0x00010001u) __device__ __constant__ unsigned int rs_sent_dw_device = RS_SENT_DW; @@ -45,9 +45,6 @@ struct RsPeers { }; -#define RS_PUBLISH(p) __hip_atomic_fetch_add((p), 1u, __ATOMIC_RELEASE, __HIP_MEMORY_SCOPE_SYSTEM) -#define RS_SPIN(p) __hip_atomic_load((p), __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_SYSTEM) -#define RS_ACQUIRE(p) ((void)__hip_atomic_load((p), __ATOMIC_ACQUIRE, __HIP_MEMORY_SCOPE_SYSTEM)) #define LC_PUBLISH_TILE(p) __hip_atomic_fetch_add((p), 1u, __ATOMIC_RELEASE, __HIP_MEMORY_SCOPE_SYSTEM) #define RS_TRACE_ARG @@ -141,24 +138,21 @@ unsigned int sent_slot_pending(const int *__restrict__ v) { template __device__ __forceinline__ -void pull_reduce_all_sent(int my_pe, int ncomm, int bands, int tiles_N, +void pull_reduce_all_sent(int my_pe, int ncomm, int bands, const bf16 *__restrict__ local_stage, const RsPeers &peers, - unsigned int *__restrict__ done, bf16 *__restrict__ out, - size_t band_elems) { + bf16 *__restrict__ out, size_t band_elems) { const int w = (int)blockIdx.x; const size_t lines = band_elems / 8; - const int peer = (w < TP - 1) ? (w + (w >= my_pe ? 1 : 0)) : -1; typedef int v4i __attribute__((ext_vector_type(4))); { - const int b0 = 0; const size_t lines_g = lines * (size_t)bands; const size_t per_g = (lines_g + (size_t)ncomm - 1) / (size_t)ncomm; const size_t g0 = (size_t)w * per_g; if (g0 >= lines_g) return; const size_t g1 = (g0 + per_g < lines_g) ? (g0 + per_g) : lines_g; - const size_t soff = ((size_t)my_pe * bands + b0) * band_elems; - v4i *dst = (v4i *)(out + (size_t)b0 * band_elems); + const size_t soff = (size_t)my_pe * bands * band_elems; + v4i *dst = (v4i *)out; for (size_t l = g0 + threadIdx.x; l < g1; ) { v4i acc; @@ -210,53 +204,12 @@ void persistent_rs_bf16_gemm(const gl A, const gl; - using ST_B = st_bf; - using RT_A = rt_bf; - using RT_B = rt_bf; - using RT_C = rt_fl; - - __shared__ ST_A As[2][2]; - __shared__ ST_B Bs[2][2]; - - const int warp_m = warpid() / WARPS_COL; - const int warp_n = warpid() % WARPS_COL; - const int wid = warpid() % NUM_WARPS; - - using T = bf16; - constexpr int bpt = ST_A::underlying_subtile_bytes_per_thread; - constexpr int bpm = bpt * NUM_THREADS; - constexpr int copies_A = HALF_ROW * K_STEP * sizeof(T) / bpm; - constexpr int copies_B = HALF_COL * K_STEP * sizeof(T) / bpm; - - constexpr int elem_per_warp = (16 / sizeof(T)) * kittens::WARP_THREADS; - const uintptr_t lds_warp_off = wid * elem_per_warp * sizeof(T); - - uint32_t a_lds_00 = __builtin_amdgcn_readfirstlane(static_cast(reinterpret_cast(&As[0][0].data[0]) + lds_warp_off)); - uint32_t a_lds_01 = __builtin_amdgcn_readfirstlane(static_cast(reinterpret_cast(&As[0][1].data[0]) + lds_warp_off)); - uint32_t a_lds_10 = __builtin_amdgcn_readfirstlane(static_cast(reinterpret_cast(&As[1][0].data[0]) + lds_warp_off)); - uint32_t a_lds_11 = __builtin_amdgcn_readfirstlane(static_cast(reinterpret_cast(&As[1][1].data[0]) + lds_warp_off)); - uint32_t b_lds_00 = __builtin_amdgcn_readfirstlane(static_cast(reinterpret_cast(&Bs[0][0].data[0]) + lds_warp_off)); - uint32_t b_lds_01 = __builtin_amdgcn_readfirstlane(static_cast(reinterpret_cast(&Bs[0][1].data[0]) + lds_warp_off)); - uint32_t b_lds_10 = __builtin_amdgcn_readfirstlane(static_cast(reinterpret_cast(&Bs[1][0].data[0]) + lds_warp_off)); - uint32_t b_lds_11 = __builtin_amdgcn_readfirstlane(static_cast(reinterpret_cast(&Bs[1][1].data[0]) + lds_warp_off)); - - uint32_t sw_A[copies_A], sw_B[copies_B]; - G_group::prefill_swizzled_offsets(As[0][0], A, sw_A); - G_group::prefill_swizzled_offsets(Bs[0][0], B, sw_B); - - const T *a_base = (const T *)&A[{0, 0, 0, 0}]; - const int a_row_stride = A.template stride<2>() * sizeof(T); - const int b_row_stride = B.template stride<2>() * sizeof(T); - i32x4 a_srd = make_srsrc(a_base, (size_t)M * a_row_stride, a_row_stride); - const T *b_base = (const T *)&B[{0, 0, 0, 0}]; - i32x4 b_srd = make_srsrc(b_base, (size_t)N_TOTAL * b_row_stride, b_row_stride); +#include "tn_prologue.inc" if ((int)blockIdx.x < ncomm) { // The comm workgroups are the reduce-scatter: they read all eight sources and write `out` once. (void)comm_wg; (void)band_bytes; - pull_reduce_all_sent(my_pe, ncomm, bands, tiles_N, local_stage, peers, done, out, - band_elems); + pull_reduce_all_sent(my_pe, ncomm, bands, local_stage, peers, out, band_elems); } const int hy_ncw = (int)gridDim.x - ncomm; @@ -271,28 +224,7 @@ void persistent_rs_bf16_gemm(const gl A, const gl 0); if (threadIdx.x == 0) { if (xcd_bucket) { - // Steal order: own bucket, then the local chunk, then the other XCDs. - int found = -1; - const int b0 = (int)blockIdx.x % NUM_XCDS_AFF; - for (int s = 0; s <= NUM_XCDS_AFF; s++) { - int bb; - if (s == 0) { - bb = b0; - } else if (s == 1) { - bb = my_pe; - } else { - bb = (b0 + s - 1) & (NUM_XCDS_AFF - 1); - } - if (buckets.cnt[bb] == 0) continue; - if (__hip_atomic_load(&bucket_ctr[bb], __ATOMIC_RELAXED, - __HIP_MEMORY_SCOPE_AGENT) >= buckets.cnt[bb]) continue; - const int idx = atomicAdd(&bucket_ctr[bb], 1); - if (idx < buckets.cnt[bb]) { - found = buckets.off[bb] + idx; - break; - } - } - s_tile_idx = (found < 0) ? num_tiles : found; +#include "xcd_steal.inc" } else { if (hy_static) { s_tile_idx = hy_tick; @@ -348,183 +280,7 @@ void persistent_rs_bf16_gemm(const gl A, const gl(Bs[0][0], {warp_n, 0}); - load(b_tile_0, st_b); - auto st_a = subtile_inplace(As[0][0], {warp_m, 0}); - load(a_tile, st_a); - G_group::load(As[1][1], A, {0, 0, block_row * 2 + 1, tile + 1}, sw_A, a_srd, a_base, a_lds_11); - asm volatile("s_waitcnt lgkmcnt(8)"); - __builtin_amdgcn_s_barrier(); - - asm volatile("s_waitcnt lgkmcnt(0)"); - __builtin_amdgcn_s_setprio(1); - mma_ABt(c00, a_tile, b_tile_0, c00); - __builtin_amdgcn_s_setprio(0); - __builtin_amdgcn_s_barrier(); - __builtin_amdgcn_sched_barrier(0); - - st_b = subtile_inplace(Bs[0][1], {warp_n, 0}); - load(b_tile_1, st_b); - G_group::load(Bs[0][0], B, {0, 0, block_col * 2, tile + 2}, sw_B, b_srd, b_base, b_lds_00); - __builtin_amdgcn_s_barrier(); - - asm volatile("s_waitcnt lgkmcnt(0)"); - __builtin_amdgcn_s_setprio(1); - mma_ABt(c01, a_tile, b_tile_1, c01); - __builtin_amdgcn_s_setprio(0); - __builtin_amdgcn_s_barrier(); - - st_a = subtile_inplace(As[0][1], {warp_m, 0}); - load(a_tile, st_a); - G_group::load(As[0][0], A, {0, 0, block_row * 2, tile + 2}, sw_A, a_srd, a_base, a_lds_00); - __builtin_amdgcn_s_barrier(); - - asm volatile("s_waitcnt lgkmcnt(0)"); - __builtin_amdgcn_s_setprio(1); - mma_ABt(c10, a_tile, b_tile_0, c10); - __builtin_amdgcn_s_setprio(0); - __builtin_amdgcn_s_barrier(); - __builtin_amdgcn_sched_barrier(0); - - st_b = subtile_inplace(Bs[1][0], {warp_n, 0}); - load(b_tile_0, st_b); - G_group::load(Bs[0][1], B, {0, 0, block_col * 2 + 1, tile + 2}, sw_B, b_srd, b_base, b_lds_01); - asm volatile("s_waitcnt vmcnt(6)"); - __builtin_amdgcn_s_barrier(); - - __builtin_amdgcn_s_setprio(1); - mma_ABt(c11, a_tile, b_tile_1, c11); - __builtin_amdgcn_s_setprio(0); - __builtin_amdgcn_s_barrier(); - - st_a = subtile_inplace(As[1][0], {warp_m, 0}); - load(a_tile, st_a); - G_group::load(As[0][1], A, {0, 0, block_row * 2 + 1, tile + 2}, sw_A, a_srd, a_base, a_lds_01); - asm volatile("s_waitcnt lgkmcnt(8)"); - __builtin_amdgcn_s_barrier(); - - asm volatile("s_waitcnt lgkmcnt(0)"); - __builtin_amdgcn_s_setprio(1); - mma_ABt(c00, a_tile, b_tile_0, c00); - __builtin_amdgcn_s_setprio(0); - __builtin_amdgcn_s_barrier(); - __builtin_amdgcn_sched_barrier(0); - - st_b = subtile_inplace(Bs[1][1], {warp_n, 0}); - load(b_tile_1, st_b); - G_group::load(Bs[1][0], B, {0, 0, block_col * 2, tile + 3}, sw_B, b_srd, b_base, b_lds_10); - __builtin_amdgcn_s_barrier(); - - asm volatile("s_waitcnt lgkmcnt(0)"); - __builtin_amdgcn_s_setprio(1); - mma_ABt(c01, a_tile, b_tile_1, c01); - __builtin_amdgcn_s_setprio(0); - __builtin_amdgcn_s_barrier(); - - st_a = subtile_inplace(As[1][1], {warp_m, 0}); - load(a_tile, st_a); - G_group::load(As[1][0], A, {0, 0, block_row * 2, tile + 3}, sw_A, a_srd, a_base, a_lds_10); - __builtin_amdgcn_s_barrier(); - - asm volatile("s_waitcnt lgkmcnt(0)"); - __builtin_amdgcn_s_setprio(1); - mma_ABt(c10, a_tile, b_tile_0, c10); - __builtin_amdgcn_s_setprio(0); - __builtin_amdgcn_s_barrier(); - __builtin_amdgcn_sched_barrier(0); - - G_group::load(Bs[1][1], B, {0, 0, block_col * 2 + 1, tile + 3}, sw_B, b_srd, b_base, b_lds_11); - asm volatile("s_waitcnt vmcnt(4)"); - __builtin_amdgcn_s_barrier(); - - __builtin_amdgcn_s_setprio(1); - mma_ABt(c11, a_tile, b_tile_1, c11); - __builtin_amdgcn_s_setprio(0); - __builtin_amdgcn_s_barrier(); - } - - { - int tile = k_tiles - 2; - - auto st_b = subtile_inplace(Bs[tic][0], {warp_n, 0}); - load(b_tile_0, st_b); - auto st_a = subtile_inplace(As[tic][0], {warp_m, 0}); - load(a_tile, st_a); - G_group::load(As[toc][1], A, {0, 0, block_row * 2 + 1, tile + 1}, sw_A, a_srd, a_base, a_lds_11); - __builtin_amdgcn_s_barrier(); - asm volatile("s_waitcnt lgkmcnt(0)"); - - __builtin_amdgcn_s_setprio(1); - mma_ABt(c00, a_tile, b_tile_0, c00); - __builtin_amdgcn_s_setprio(0); - __builtin_amdgcn_s_barrier(); - - st_b = subtile_inplace(Bs[tic][1], {warp_n, 0}); - load(b_tile_1, st_b); - __builtin_amdgcn_s_barrier(); - - asm volatile("s_waitcnt lgkmcnt(0)"); - __builtin_amdgcn_s_setprio(1); - mma_ABt(c01, a_tile, b_tile_1, c01); - __builtin_amdgcn_s_setprio(0); - __builtin_amdgcn_s_barrier(); - - st_a = subtile_inplace(As[tic][1], {warp_m, 0}); - load(a_tile, st_a); - asm volatile("s_waitcnt vmcnt(4)"); - __builtin_amdgcn_s_barrier(); - - asm volatile("s_waitcnt lgkmcnt(0)"); - __builtin_amdgcn_s_setprio(1); - mma_ABt(c10, a_tile, b_tile_0, c10); - mma_ABt(c11, a_tile, b_tile_1, c11); - __builtin_amdgcn_s_setprio(0); - __builtin_amdgcn_s_barrier(); - tic ^= 1; toc ^= 1; - } - - { - auto st_b = subtile_inplace(Bs[tic][0], {warp_n, 0}); - load(b_tile_0, st_b); - auto st_a = subtile_inplace(As[tic][0], {warp_m, 0}); - load(a_tile, st_a); - asm volatile("s_waitcnt vmcnt(2)"); - __builtin_amdgcn_s_barrier(); - - asm volatile("s_waitcnt lgkmcnt(0)"); - __builtin_amdgcn_s_setprio(1); - mma_ABt(c00, a_tile, b_tile_0, c00); - __builtin_amdgcn_s_setprio(0); - __builtin_amdgcn_s_barrier(); - - st_b = subtile_inplace(Bs[tic][1], {warp_n, 0}); - load(b_tile_1, st_b); - asm volatile("s_waitcnt vmcnt(0)"); - __builtin_amdgcn_s_barrier(); - - asm volatile("s_waitcnt lgkmcnt(0)"); - __builtin_amdgcn_s_setprio(1); - mma_ABt(c01, a_tile, b_tile_1, c01); - __builtin_amdgcn_s_setprio(0); - __builtin_amdgcn_s_barrier(); - - st_a = subtile_inplace(As[tic][1], {warp_m, 0}); - load(a_tile, st_a); - __builtin_amdgcn_s_barrier(); - - asm volatile("s_waitcnt lgkmcnt(0)"); - __builtin_amdgcn_s_setprio(1); - mma_ABt(c10, a_tile, b_tile_0, c10); - mma_ABt(c11, a_tile, b_tile_1, c11); - __builtin_amdgcn_s_setprio(0); - __builtin_amdgcn_s_barrier(); - } - - if (warp_m == 0) { - __builtin_amdgcn_s_barrier(); - } +#include "tn_mainloop.inc" bf16 *sbase = local_stage + ((size_t)owner * bands + lrow) * band_elems; diff --git a/transformer_engine/common/gemm/kittens/cdna4/overlap/overlap_common.cuh b/transformer_engine/common/gemm/kittens/cdna4/overlap/overlap_common.cuh index 1baa7333f0..1a94c7811a 100644 --- a/transformer_engine/common/gemm/kittens/cdna4/overlap/overlap_common.cuh +++ b/transformer_engine/common/gemm/kittens/cdna4/overlap/overlap_common.cuh @@ -21,6 +21,8 @@ constexpr int NUM_WARPS = WARPS_ROW * WARPS_COL; constexpr int NUM_THREADS = NUM_WARPS * kittens::WARP_THREADS; +using G_group = kittens::group; + constexpr int REG_M = BLOCK_ROW / WARPS_ROW / 2; constexpr int REG_N = BLOCK_COL / WARPS_COL / 2; @@ -39,6 +41,29 @@ struct XcdBuckets { __host__ __device__ __forceinline__ int tile_xcd(int chunk_id) { return chunk_id & (NUM_XCDS_AFF - 1); } +// Steal order: own bucket, then the local chunk, then the other XCDs. +__device__ __forceinline__ +int steal_bucket_tile(const XcdBuckets &buckets, int *__restrict__ bucket_ctr, int my_pe, + int num_tiles) { + const int b0 = (int)blockIdx.x % NUM_XCDS_AFF; + for (int s = 0; s <= NUM_XCDS_AFF; s++) { + int bb; + if (s == 0) { + bb = b0; + } else if (s == 1) { + bb = my_pe; + } else { + bb = (b0 + s - 1) & (NUM_XCDS_AFF - 1); + } + if (buckets.cnt[bb] == 0) continue; + if (__hip_atomic_load(&bucket_ctr[bb], __ATOMIC_RELAXED, + __HIP_MEMORY_SCOPE_AGENT) >= buckets.cnt[bb]) continue; + const int idx = atomicAdd(&bucket_ctr[bb], 1); + if (idx < buckets.cnt[bb]) return buckets.off[bb] + idx; + } + return num_tiles; +} + template static std::vector bucketize_by_xcd(const std::vector &q, XcdBuckets &bk) { std::vector out; diff --git a/transformer_engine/common/gemm/kittens/cdna4/overlap/tn_mainloop.inc b/transformer_engine/common/gemm/kittens/cdna4/overlap/tn_mainloop.inc new file mode 100644 index 0000000000..e0ad18ccc3 --- /dev/null +++ b/transformer_engine/common/gemm/kittens/cdna4/overlap/tn_mainloop.inc @@ -0,0 +1,179 @@ + for (int tile = 0; tile < k_tiles - 2; tile += 2) { + + auto st_b = subtile_inplace(Bs[0][0], {warp_n, 0}); + load(b_tile_0, st_b); + auto st_a = subtile_inplace(As[0][0], {warp_m, 0}); + load(a_tile, st_a); + G_group::load(As[1][1], A, {0, 0, block_row * 2 + 1, tile + 1}, sw_A, a_srd, a_base, a_lds_11); + asm volatile("s_waitcnt lgkmcnt(8)"); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + mma_ABt(c00, a_tile, b_tile_0, c00); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + st_b = subtile_inplace(Bs[0][1], {warp_n, 0}); + load(b_tile_1, st_b); + G_group::load(Bs[0][0], B, {0, 0, block_col * 2, tile + 2}, sw_B, b_srd, b_base, b_lds_00); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + mma_ABt(c01, a_tile, b_tile_1, c01); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + + st_a = subtile_inplace(As[0][1], {warp_m, 0}); + load(a_tile, st_a); + G_group::load(As[0][0], A, {0, 0, block_row * 2, tile + 2}, sw_A, a_srd, a_base, a_lds_00); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + mma_ABt(c10, a_tile, b_tile_0, c10); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + st_b = subtile_inplace(Bs[1][0], {warp_n, 0}); + load(b_tile_0, st_b); + G_group::load(Bs[0][1], B, {0, 0, block_col * 2 + 1, tile + 2}, sw_B, b_srd, b_base, b_lds_01); + asm volatile("s_waitcnt vmcnt(6)"); + __builtin_amdgcn_s_barrier(); + + __builtin_amdgcn_s_setprio(1); + mma_ABt(c11, a_tile, b_tile_1, c11); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + + st_a = subtile_inplace(As[1][0], {warp_m, 0}); + load(a_tile, st_a); + G_group::load(As[0][1], A, {0, 0, block_row * 2 + 1, tile + 2}, sw_A, a_srd, a_base, a_lds_01); + asm volatile("s_waitcnt lgkmcnt(8)"); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + mma_ABt(c00, a_tile, b_tile_0, c00); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + st_b = subtile_inplace(Bs[1][1], {warp_n, 0}); + load(b_tile_1, st_b); + G_group::load(Bs[1][0], B, {0, 0, block_col * 2, tile + 3}, sw_B, b_srd, b_base, b_lds_10); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + mma_ABt(c01, a_tile, b_tile_1, c01); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + + st_a = subtile_inplace(As[1][1], {warp_m, 0}); + load(a_tile, st_a); + G_group::load(As[1][0], A, {0, 0, block_row * 2, tile + 3}, sw_A, a_srd, a_base, a_lds_10); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + mma_ABt(c10, a_tile, b_tile_0, c10); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + G_group::load(Bs[1][1], B, {0, 0, block_col * 2 + 1, tile + 3}, sw_B, b_srd, b_base, b_lds_11); + asm volatile("s_waitcnt vmcnt(4)"); + __builtin_amdgcn_s_barrier(); + + __builtin_amdgcn_s_setprio(1); + mma_ABt(c11, a_tile, b_tile_1, c11); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + } + + // Epilogue: K-tile k_tiles-2 (stage tic) + { + int tile = k_tiles - 2; + + auto st_b = subtile_inplace(Bs[tic][0], {warp_n, 0}); + load(b_tile_0, st_b); + auto st_a = subtile_inplace(As[tic][0], {warp_m, 0}); + load(a_tile, st_a); + G_group::load(As[toc][1], A, {0, 0, block_row * 2 + 1, tile + 1}, sw_A, a_srd, a_base, a_lds_11); + __builtin_amdgcn_s_barrier(); + asm volatile("s_waitcnt lgkmcnt(0)"); + + __builtin_amdgcn_s_setprio(1); + mma_ABt(c00, a_tile, b_tile_0, c00); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + + st_b = subtile_inplace(Bs[tic][1], {warp_n, 0}); + load(b_tile_1, st_b); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + mma_ABt(c01, a_tile, b_tile_1, c01); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + + st_a = subtile_inplace(As[tic][1], {warp_m, 0}); + load(a_tile, st_a); + asm volatile("s_waitcnt vmcnt(4)"); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + mma_ABt(c10, a_tile, b_tile_0, c10); + mma_ABt(c11, a_tile, b_tile_1, c11); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + tic ^= 1; toc ^= 1; + } + + // Epilogue: K-tile k_tiles-1 (stage tic after flip) + { + auto st_b = subtile_inplace(Bs[tic][0], {warp_n, 0}); + load(b_tile_0, st_b); + auto st_a = subtile_inplace(As[tic][0], {warp_m, 0}); + load(a_tile, st_a); + asm volatile("s_waitcnt vmcnt(2)"); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + mma_ABt(c00, a_tile, b_tile_0, c00); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + + st_b = subtile_inplace(Bs[tic][1], {warp_n, 0}); + load(b_tile_1, st_b); + asm volatile("s_waitcnt vmcnt(0)"); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + mma_ABt(c01, a_tile, b_tile_1, c01); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + + st_a = subtile_inplace(As[tic][1], {warp_m, 0}); + load(a_tile, st_a); + __builtin_amdgcn_s_barrier(); + + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_s_setprio(1); + mma_ABt(c10, a_tile, b_tile_0, c10); + mma_ABt(c11, a_tile, b_tile_1, c11); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_s_barrier(); + } + + if (warp_m == 0) { + __builtin_amdgcn_s_barrier(); + } diff --git a/transformer_engine/common/gemm/kittens/cdna4/overlap/tn_prologue.inc b/transformer_engine/common/gemm/kittens/cdna4/overlap/tn_prologue.inc new file mode 100644 index 0000000000..037f93ae37 --- /dev/null +++ b/transformer_engine/common/gemm/kittens/cdna4/overlap/tn_prologue.inc @@ -0,0 +1,41 @@ + using ST_A = st_bf; + using ST_B = st_bf; + using RT_A = rt_bf; + using RT_B = rt_bf; + using RT_C = rt_fl; + + __shared__ ST_A As[2][2]; + __shared__ ST_B Bs[2][2]; + + const int warp_m = warpid() / WARPS_COL; + const int warp_n = warpid() % WARPS_COL; + const int wid = warpid() % NUM_WARPS; + + using T = bf16; + constexpr int bpt = ST_A::underlying_subtile_bytes_per_thread; + constexpr int bpm = bpt * NUM_THREADS; + constexpr int copies_A = HALF_ROW * K_STEP * sizeof(T) / bpm; + constexpr int copies_B = HALF_COL * K_STEP * sizeof(T) / bpm; + + constexpr int elem_per_warp = (16 / sizeof(T)) * kittens::WARP_THREADS; + const uintptr_t lds_warp_off = wid * elem_per_warp * sizeof(T); + + uint32_t a_lds_00 = __builtin_amdgcn_readfirstlane(static_cast(reinterpret_cast(&As[0][0].data[0]) + lds_warp_off)); + uint32_t a_lds_01 = __builtin_amdgcn_readfirstlane(static_cast(reinterpret_cast(&As[0][1].data[0]) + lds_warp_off)); + uint32_t a_lds_10 = __builtin_amdgcn_readfirstlane(static_cast(reinterpret_cast(&As[1][0].data[0]) + lds_warp_off)); + uint32_t a_lds_11 = __builtin_amdgcn_readfirstlane(static_cast(reinterpret_cast(&As[1][1].data[0]) + lds_warp_off)); + uint32_t b_lds_00 = __builtin_amdgcn_readfirstlane(static_cast(reinterpret_cast(&Bs[0][0].data[0]) + lds_warp_off)); + uint32_t b_lds_01 = __builtin_amdgcn_readfirstlane(static_cast(reinterpret_cast(&Bs[0][1].data[0]) + lds_warp_off)); + uint32_t b_lds_10 = __builtin_amdgcn_readfirstlane(static_cast(reinterpret_cast(&Bs[1][0].data[0]) + lds_warp_off)); + uint32_t b_lds_11 = __builtin_amdgcn_readfirstlane(static_cast(reinterpret_cast(&Bs[1][1].data[0]) + lds_warp_off)); + + uint32_t sw_A[copies_A], sw_B[copies_B]; + G_group::prefill_swizzled_offsets(As[0][0], A, sw_A); + G_group::prefill_swizzled_offsets(Bs[0][0], B, sw_B); + + const T *a_base = (const T *)&A[{0, 0, 0, 0}]; + const int a_row_stride = A.template stride<2>() * sizeof(T); + const int b_row_stride = B.template stride<2>() * sizeof(T); + i32x4 a_srd = make_srsrc(a_base, (size_t)M * a_row_stride, a_row_stride); + const T *b_base = (const T *)&B[{0, 0, 0, 0}]; + i32x4 b_srd = make_srsrc(b_base, (size_t)N_TOTAL * b_row_stride, b_row_stride); diff --git a/transformer_engine/common/gemm/kittens/cdna4/overlap/xcd_steal.inc b/transformer_engine/common/gemm/kittens/cdna4/overlap/xcd_steal.inc new file mode 100644 index 0000000000..ef7b4f2c62 --- /dev/null +++ b/transformer_engine/common/gemm/kittens/cdna4/overlap/xcd_steal.inc @@ -0,0 +1,22 @@ +// Steal order: own bucket, then the local chunk, then the other XCDs. +int found = -1; +const int b0 = (int)blockIdx.x % NUM_XCDS_AFF; +for (int s = 0; s <= NUM_XCDS_AFF; s++) { + int bb; + if (s == 0) { + bb = b0; + } else if (s == 1) { + bb = my_pe; + } else { + bb = (b0 + s - 1) & (NUM_XCDS_AFF - 1); + } + if (buckets.cnt[bb] == 0) continue; + if (__hip_atomic_load(&bucket_ctr[bb], __ATOMIC_RELAXED, + __HIP_MEMORY_SCOPE_AGENT) >= buckets.cnt[bb]) continue; + const int idx = atomicAdd(&bucket_ctr[bb], 1); + if (idx < buckets.cnt[bb]) { + found = buckets.off[bb] + idx; + break; + } +} +s_tile_idx = (found < 0) ? num_tiles : found;