Add MultipleBuffer, MBSK, LocalSplitU, and TreeStreamK split-K MXFP4 …

tests/kernel/wave_gemm_test.py

@@ -62,7 +62,11 @@
     get_persistent_reordering_kernel,
 )
 from wave_lang.kernel.wave.templates.tagged_mxfp4_gemm import (
+    get_tagged_lsu_mxfp4_gemm,
+    get_tagged_mbsk_splitk_mxfp4_gemm,
+    get_tagged_multibuffer_splitk_mxfp4_gemm,
     get_tagged_splitk_mxfp4_gemm,
+    get_tagged_tree_streamk_mxfp4_gemm,
 )
 from wave_lang.kernel.wave.templates.test_kernels import (
     get_gemm_prefetch_kernel_and_schedule,
@@ -83,6 +87,7 @@
 )
 
 from wave_lang.kernel.lang import DataType
+import math
 import os
 import json
 from torch.testing import assert_close
@@ -3647,6 +3652,9 @@ def testSplitKGemm(
         ((256, 256, 512), 4),
         ((256, 256, 1024), 4),
         ((512, 512, 1024), 4),
+        ((256, 256, 1024), 2),
+        ((256, 256, 2048), 2),
+        ((256, 256, 8192), 2),
     ],
 )
 @pytest.mark.parametrize("output_type", [torch.float32, torch.bfloat16])
@@ -3688,3 +3696,229 @@ def testSplitKMxfp4Gemm(
         # partial-sum magnitudes of ~250-430 the ULP is 2-4, so with 2-4
         # splits the worst-case rounding error is O(num_splits * ULP).
         assert_close(c_gpu.cpu().to(torch.float32), torch_ref, rtol=1e-1, atol=4.0)
+
+
+@require_e2e
+@require_cdna4
+@pytest.mark.parametrize(
+    "shape, num_splits",
+    [
+        ((256, 256, 256), 2),
+        ((256, 256, 512), 2),
+        ((256, 256, 512), 4),
+        ((256, 256, 1024), 4),
+        ((512, 512, 1024), 4),
+    ],
+)
+def testMultiBufferSplitKMxfp4Gemm(
+    shape: tuple[int, int, int],
+    num_splits: int,
+):
+    (
+        main_fn,
+        main_options,
+        reduction_fn,
+        reduction_options,
+    ) = get_tagged_multibuffer_splitk_mxfp4_gemm(
+        shape,
+        num_splits=num_splits,
+        block_shape=(128, 128, 128),
+        mfma_variant=ScaledMMAType.F32_16x16x128_F8F6F4,
+        output_type=TORCH_DTYPE_TO_WAVE[torch.float32],
+    )
+    main_options = set_default_run_config(main_options)
+    reduction_options = set_default_run_config(reduction_options)
+    schedule = get_mxfp4_dbuf_schedule(use_stagger=True, k_partitions=1)
+    compiled_main = wave_compile(main_options, main_fn, schedule)
+    compiled_reduction = wave_compile(reduction_options, reduction_fn)
+
+    m, n, _ = shape
+    x, w, x_scales, w_scales = generate_gemm_afp4wfp4_inputs(
+        shape, device=torch.device("cpu")
+    )
+    torch_ref = torchScaledGemmMXFP4(x, w, x_scales, w_scales)
+
+    x_gpu = x.cuda()
+    w_t_gpu = w.T.contiguous().cuda()
+    x_scales_gpu = x_scales.cuda()
+    w_scales_gpu = w_scales.cuda()
+    workspace = device_zeros(num_splits, m, n, dtype=torch.float32)
+    c_gpu = device_zeros(m, n, dtype=torch.float32)
+
+    compiled_main(x_gpu, x_scales_gpu, w_t_gpu, w_scales_gpu, workspace)
+    compiled_reduction(workspace, c_gpu)
+    assert_close(c_gpu.cpu(), torch_ref, rtol=1e-3, atol=1e-2)
+
+
+@require_e2e
+@require_cdna4
+@pytest.mark.parametrize(
+    "shape, num_splits",
+    [
+        ((256, 256, 256), 2),
+        ((256, 256, 512), 2),
+        ((256, 256, 512), 4),
+        ((256, 256, 1024), 4),
+    ],
+)
+def testMBSKSplitKMxfp4Gemm(
+    shape: tuple[int, int, int],
+    num_splits: int,
+):
+    block_shape = (128, 128, 128)
+    mbsk_gemm, options = get_tagged_mbsk_splitk_mxfp4_gemm(
+        shape,
+        num_splits=num_splits,
+        block_shape=block_shape,
+        mfma_variant=ScaledMMAType.F32_16x16x128_F8F6F4,
+        output_type=TORCH_DTYPE_TO_WAVE[torch.float32],
+    )
+    options = set_default_run_config(options)
+    schedule = get_mxfp4_dbuf_schedule(use_stagger=True, k_partitions=1)
+    compiled = wave_compile(options, mbsk_gemm, schedule)
+
+    m, n, _ = shape
+    x, w, x_scales, w_scales = generate_gemm_afp4wfp4_inputs(
+        shape, device=torch.device("cpu")
+    )
+    torch_ref = torchScaledGemmMXFP4(x, w, x_scales, w_scales)
+
+    x_gpu = x.cuda()
+    w_t_gpu = w.T.contiguous().cuda()
+    x_scales_gpu = x_scales.cuda()
+    w_scales_gpu = w_scales.cuda()
+    block_m, block_n, _ = block_shape
+    num_tiles = math.ceil(m / block_m) * math.ceil(n / block_n)
+    workspace = device_zeros(num_splits, m, n, dtype=torch.float32)
+    sync_buffer = device_zeros(num_tiles, dtype=torch.int32)
+    c_gpu = device_zeros(m, n, dtype=torch.float32)
+
+    compiled(
+        x_gpu,
+        x_scales_gpu,
+        w_t_gpu,
+        w_scales_gpu,
+        workspace,
+        sync_buffer,
+        c_gpu,
+    )
+    assert_close(c_gpu.cpu(), torch_ref, rtol=1e-3, atol=1e-2)
+
+
+@require_e2e
+@require_cdna4
+@pytest.mark.parametrize(
+    "shape, lsu_factor",
+    [
+        ((256, 256, 256), 2),
+        ((256, 256, 512), 2),
+        ((256, 256, 1024), 2),
+        ((512, 512, 1024), 2),
+        ((256, 256, 512), 4),
+        ((256, 256, 1024), 4),
+        ((256, 256, 2048), 4),
+    ],
+)
+@pytest.mark.parametrize("output_type", [torch.float32])
+def testLocalSplitUMxfp4Gemm(
+    shape: tuple[int, int, int],
+    lsu_factor: int,
+    output_type: torch.dtype,
+):
+    """LocalSplitU: 2-kernel split-K with workspace reduction, no atomics on C.
+
+    Main kernel writes f32 partials to workspace.
+    Reduction kernel sums partials and writes the final result to C.
+    """
+    block_shape = (128, 128, 128)
+    (
+        main_fn,
+        main_options,
+        reduction_fn,
+        reduction_options,
+    ) = get_tagged_lsu_mxfp4_gemm(
+        shape,
+        lsu_factor=lsu_factor,
+        block_shape=block_shape,
+        mfma_variant=ScaledMMAType.F32_16x16x128_F8F6F4,
+        output_type=TORCH_DTYPE_TO_WAVE[output_type],
+    )
+
+    main_options = set_default_run_config(main_options)
+    reduction_options = set_default_run_config(reduction_options)
+    schedule = get_mxfp4_dbuf_schedule(use_stagger=True, k_partitions=1)
+    compiled_main = wave_compile(main_options, main_fn, schedule)
+    compiled_reduction = wave_compile(reduction_options, reduction_fn)
+
+    m, n, k = shape
+    x, w, x_scales, w_scales = generate_gemm_afp4wfp4_inputs(
+        shape, device=torch.device("cpu")
+    )
+    torch_ref = torchScaledGemmMXFP4(x, w, x_scales, w_scales)
+
+    x_gpu = x.cuda()
+    w_t_gpu = w.T.contiguous().cuda()
+    x_scales_gpu = x_scales.cuda()
+    w_scales_gpu = w_scales.cuda()
+    block_m, block_n, _ = block_shape
+    num_tiles = math.ceil(m / block_m) * math.ceil(n / block_n)
+    workspace = device_zeros(lsu_factor, m, n, dtype=torch.float32)
+    sync_counter = device_zeros(num_tiles, dtype=torch.int32)
+    c_gpu = device_zeros(m, n, dtype=output_type)
+
+    compiled_main(x_gpu, x_scales_gpu, w_t_gpu, w_scales_gpu, workspace, sync_counter)
+    compiled_reduction(workspace, c_gpu)
+    assert_close(c_gpu.cpu(), torch_ref, rtol=1e-3, atol=1e-2)
+
+
+@require_e2e
+@require_cdna4
+@pytest.mark.parametrize(
+    "shape, num_ctas",
+    [
+        ((128, 128, 128), 1),
+        ((128, 128, 256), 2),
+        ((256, 256, 256), 8),
+        ((256, 256, 512), 16),
+        ((512, 512, 1024), 16),
+    ],
+)
+@pytest.mark.parametrize("output_type", [torch.float32])
+def testTreeStreamKMxfp4Gemm(
+    shape: tuple[int, int, int],
+    num_ctas: int,
+    output_type: torch.dtype,
+):
+    """StreamK with tree reduction: single-kernel, O(log n) fixup depth.
+
+    Uses workspace + flag buffer for binary tree reduction instead of
+    linear spinlock or atomic_add.
+    """
+    m, n, k = shape
+    block_m, block_n, block_k = 128, 128, 128
+
+    streamk_gemm, options = get_tagged_tree_streamk_mxfp4_gemm(
+        shape,
+        block_shape=(block_m, block_n, block_k),
+        mfma_variant=ScaledMMAType.F32_16x16x128_F8F6F4,
+        output_type=TORCH_DTYPE_TO_WAVE[output_type],
+        num_ctas=num_ctas,
+    )
+
+    options = set_default_run_config(options)
+    compiled = wave_compile(options, streamk_gemm)
+
+    x, w, x_scales, w_scales = generate_gemm_afp4wfp4_inputs(
+        shape, device=torch.device("cpu")
+    )
+    torch_ref = torchScaledGemmMXFP4(x, w, x_scales, w_scales)
+
+    x_gpu = x.cuda()
+    w_t_gpu = w.T.contiguous().cuda()
+    x_scales_gpu = x_scales.cuda()
+    w_scales_gpu = w_scales.cuda()
+
+    c_gpu = device_zeros(m, n, dtype=output_type)
+
+    compiled(x_gpu, x_scales_gpu, w_t_gpu, w_scales_gpu, c_gpu)
+    assert_close(c_gpu.cpu(), torch_ref, rtol=1e-3, atol=1e-2)

wave_lang/kernel/wave/schedules/gemm_mxfp4_double_buffer.py

@@ -45,9 +45,13 @@ def get_mxfp4_dbuf_schedule(use_stagger: bool = True, k_partitions: int = 2):
     Args:
         use_stagger: Enable wave staggering + WorkgroupBarrier in cluster 0.
             Recommended for 8-wave configs; disable for 4-wave.
+            Forced off when ``k_partitions=1`` because the stagger interacts
+            incorrectly with single-partition pipeline iterations > 1.
         k_partitions: ``2`` for standard GEMM (partition along K).  Use ``1`` for
             split-K MXFP4 kernels where the K loop has only one expansion id.
     """
+    if k_partitions == 1:
+        use_stagger = False
     K = tkl.sym.K
 
     @wave_schedule.wave_schedule()
@@ -289,11 +293,15 @@ def get_mxfp4_dbuf_pingpong_schedule(
     Args:
         use_stagger: Enable wave staggering + WorkgroupBarrier in cluster 0.
             Recommended for 8-wave configs; disable for 4-wave.
+            Forced off when ``k_partitions=1`` because the stagger interacts
+            incorrectly with single-partition pipeline iterations > 1.
         shape: Tuple of (M, N, K) dimensions. If provided and bigger than
             (1024, 1024, 1024), an extra WorkgroupBarrier will be added
             after the first SchedulingBarrier in cluster 0.
         k_partitions: ``2`` for standard GEMM.  Use ``1`` for split-K MXFP4 kernels.
     """
+    if k_partitions == 1:
+        use_stagger = False
     K = tkl.sym.K
 
     @wave_schedule.wave_schedule()

wave_lang/kernel/wave/templates/__init__.py

@@ -7,6 +7,7 @@
 from .attention_common import AttentionShape
 from .tagged_attention import get_tagged_bshd_attention_kernel
 from .tagged_mxfp4_gemm import (
+    get_tagged_lsu_mxfp4_gemm,
     get_tagged_mxfp4_gemm,
     get_tagged_mxfp4_gemm_preshuffle_b,
     get_tagged_mxfp4_gemm_preshuffle_b_wide_store,
@@ -15,11 +16,13 @@
     get_tagged_splitk_mxfp4_gemm,
     get_tagged_splitk_mxfp4_gemm_preshuffle_b,
     get_tagged_splitk_mxfp4_gemm_preshuffle_scales,
+    get_tagged_tree_streamk_mxfp4_gemm,
 )
 
 __all__ = [
     "AttentionShape",
     "get_tagged_bshd_attention_kernel",
+    "get_tagged_lsu_mxfp4_gemm",
     "get_tagged_mxfp4_gemm",
     "get_tagged_mxfp4_gemm_preshuffle_b",
     "get_tagged_mxfp4_gemm_preshuffle_b_wide_store",
@@ -28,4 +31,5 @@
     "get_tagged_splitk_mxfp4_gemm",
     "get_tagged_splitk_mxfp4_gemm_preshuffle_b",
     "get_tagged_splitk_mxfp4_gemm_preshuffle_scales",
+    "get_tagged_tree_streamk_mxfp4_gemm",
 ]