Implement mxfp4 split-k gemm (#958)

The core things added are split-k gemm, and it is tested for (1)
generation of the `buffer_atomic_pk_add_bf16` instruction that we wanted
to use, and (2) for gemm correctness.

Overview of changes unrelated to wave_asm:

- `remove_global_indexing` in `general_utils.py`: Zeroes out tiling
constraint starts (e.g. `K_SPLIT_OFF`) alongside workgroup IDs before
dimension scaling, so that the subtraction of the start offset doesn't
mix scaled and unscaled units (K vs K/32 for MXFP4 scales).

- Fixing spurious bounds on split-K tiling that prevented scale vector
merging: TilingConstraint.get_index_bound was conservatively generating
bounds for the split-K case because sympy could not prove that
ceiling(Min(K, f(wg)) / tile)
* tile <= K. These bounds prevented merge_contiguous_reads from
combining scalar scale reads into vector<4xi8> loads (it skips reads
that already have bounds). Add _work_may_exceed_dim() to structurally
detect the aligned split-k pattern and prove no overshoot, avoiding the
spurious bound. (This was necessary to get scale_preshuffle to have 4x
vector loads when combined with split-k.)

---------

Signed-off-by: William G Hatch <william@hatch.uno>

Author: willghatch

examples/python/7.1_schedule.py

@@ -4,10 +4,14 @@
 Double-buffered MXFP4 GEMM with 4-wave and 8-wave configurations.
 Uses get_tagged_mxfp4_gemm (templates) + get_mxfp4_dbuf_schedule (schedules).
 
+The --splitk N flag enables split-K with N splits for supported tests.
+
 Usage:
     python 7.1_schedule.py --test test_dbuf_4wave_mxfp_gemm
-    python 7.1_schedule.py --test test_dbuf_8wave_mxfp_gemm
-    python 7.1_schedule.py --test test_dbuf_8wave_mxfp_gemm --debug
+    python 7.1_schedule.py --test test_dbuf_4wave_mxfp_gemm --splitk 2
+    python 7.1_schedule.py --test test_dbuf_8wave_pingpong_mxfp_gemm
+    python 7.1_schedule.py --test test_dbuf_8wave_pingpong_mxfp_gemm --splitk 2
+    python 7.1_schedule.py --test test_splitk_preshuffle_scales_gemm_cpp
     python 7.1_schedule.py --list_tests
 """
 
@@ -20,6 +24,7 @@
     SHARED_ADDRESS_SPACE,
 )
 from wave_lang.kernel.wave.compile import wave_compile
+from wave_lang.kernel.wave.constraints import ScaledMMAType
 from wave_lang.kernel.wave.schedules import (
     get_mxfp4_asymmetric_schedule,
     get_mxfp4_dbuf_mixed_pingpong_schedule,
@@ -35,6 +40,9 @@
     get_tagged_mxfp4_gemm_preshuffle_b_wide_store,
     get_tagged_mxfp4_gemm_preshuffle_scales,
     get_tagged_mxfp4_gemm_preshuffle_scales_and_B,
+    get_tagged_splitk_mxfp4_gemm,
+    get_tagged_splitk_mxfp4_gemm_preshuffle_b,
+    get_tagged_splitk_mxfp4_gemm_preshuffle_scales,
 )
 from wave_lang.kernel.wave.utils.mxfp_utils import (
     b_preshuffle,
@@ -61,7 +69,14 @@ def _run_mxfp_gemm(gemm, shape):
 
 
 def _run_mxfp_gemm_preshuffle(
-    gemm, shape, all=False, only_scale=False, only_b=False, output_dtype=torch.float32
+    gemm,
+    shape,
+    all=False,
+    only_scale=False,
+    only_b=False,
+    output_dtype=torch.float32,
+    atol=None,
+    rtol=None,
 ):
     """Run compiled GEMM kernel with preshuffled B and B_scale, verify against reference.
 
@@ -90,8 +105,13 @@ def _run_mxfp_gemm_preshuffle(
 
     gemm(x, x_scales_ps, w_t_ps, w_scales_ps, out)
 
+    tol_kwargs = {}
+    if atol is not None:
+        tol_kwargs["atol"] = atol
+    if rtol is not None:
+        tol_kwargs["rtol"] = rtol
     torch.testing.assert_close(
-        torch_out, out.cpu(), check_dtype=False, check_device=False
+        torch_out, out.cpu(), check_dtype=False, check_device=False, **tol_kwargs
     )
 
 
@@ -137,10 +157,17 @@ def _get_8wave_shape_from_block(block):
 
 
 def test_dbuf_4wave_mxfp_gemm(
-    is_debug=False, shape=(1024, 1024, 8192), block=(256, 256, 256)
+    is_debug=False, shape=(1024, 1024, 8192), block=(256, 256, 256), splitk=None
 ):
     """Double-buffered MXFP4 GEMM, 4 waves, no stagger."""
-    gemm, options = get_tagged_mxfp4_gemm(shape, block, wave_shape=(2, 2))
+    if splitk and block == (256, 256, 256):
+        block = (128, 128, 256)
+    if splitk:
+        gemm, options = get_tagged_splitk_mxfp4_gemm(
+            shape, num_splits=splitk, block_shape=block, wave_shape=(2, 2)
+        )
+    else:
+        gemm, options = get_tagged_mxfp4_gemm(shape, block, wave_shape=(2, 2))
     schedule = get_mxfp4_dbuf_schedule(use_stagger=False)
 
     options.print_ir_after = "all" if is_debug else []
@@ -150,11 +177,16 @@ def test_dbuf_4wave_mxfp_gemm(
     gemm = wave_compile(options, gemm, schedule)
 
     _run_mxfp_gemm(gemm, shape)
-    print("MXFP GEMM double-buffer 4-wave test passed!")
+    sk = f" split-K({splitk})" if splitk else ""
+    print(f"MXFP GEMM double-buffer 4-wave{sk} test passed!")
 
 
 def test_dbuf_8wave_pingpong_mxfp_gemm(
-    is_debug=False, shape=(1024, 1024, 8192), block=(256, 256, 256), dynamic=False
+    is_debug=False,
+    shape=(1024, 1024, 8192),
+    block=(256, 256, 256),
+    dynamic=False,
+    splitk=None,
 ):
     """Double-buffered MXFP4 GEMM, 8 waves, ping-pong with stagger.
     A&B scales are preshuffled and read from global memory directly to VGPRs.
@@ -163,10 +195,17 @@ def test_dbuf_8wave_pingpong_mxfp_gemm(
     Note: for dynamic mode, keep block MxN at or below 128x256 or 256x128
     to avoid exceeding shared-memory limits.
     """
+    if splitk and block == (256, 256, 256):
+        block = (128, 256, 256)
     wave_shape = _get_8wave_shape_from_block(block)
-    gemm, options = get_tagged_mxfp4_gemm_preshuffle_scales(
-        shape, block, wave_shape=wave_shape
-    )
+    if splitk:
+        gemm, options = get_tagged_splitk_mxfp4_gemm_preshuffle_scales(
+            shape, num_splits=splitk, block_shape=block, wave_shape=wave_shape
+        )
+    else:
+        gemm, options = get_tagged_mxfp4_gemm_preshuffle_scales(
+            shape, block, wave_shape=wave_shape
+        )
     options.specialize = True
     options.use_buffer_ops = True
     options.minimize_shared_allocs = True
@@ -185,19 +224,28 @@ def test_dbuf_8wave_pingpong_mxfp_gemm(
 
     _run_mxfp_gemm_preshuffle(gemm, shape, only_scale=True)
     mode = "dynamic" if dynamic else "static"
+    sk = f", split-K({splitk})" if splitk else ""
     print(
-        f"MXFP GEMM double-buffer 8-wave ping pong with scale shuffling ({mode}) test passed!"
+        f"MXFP GEMM double-buffer 8-wave ping pong with scale shuffling ({mode}{sk}) test passed!"
     )
 
 
 def test_dbuf_8wave_pingpong_mxfp_gemm_Bshuffle(
-    is_debug=False, shape=(1024, 1024, 8192), block=(256, 256, 256), dynamic=False
+    is_debug=False,
+    shape=(1024, 1024, 8192),
+    block=(256, 256, 256),
+    dynamic=False,
+    splitk=None,
 ):
     """Double-buffered MXFP4 GEMM, 8 waves, ping-pong with stagger.
     A&B scales are preshuffled and read from global memory directly to VGPRs.
     Same for B data. However, prefetching shuffled B directly to VGPR consumes too many VGPRs and causes spilling.
     A is read from global memory directly to LDS.
     """
+    if splitk:
+        raise NotImplementedError(
+            "split-K is not yet supported with B-shuffled ping-pong schedule"
+        )
     wave_shape = _get_8wave_shape_from_block(block)
     gemm, options = get_tagged_mxfp4_gemm_preshuffle_scales_and_B(
         shape, block, wave_shape=wave_shape
@@ -390,7 +438,7 @@ def test_dbuf_8wave_pingpong_mxfp_gemm_Bshuffle_lds_transposed(
 
 
 def test_dbuf_8wave_mixed_pingpong_mxfp_gemm(
-    is_debug=False, shape=(1024, 1024, 8192), block=(256, 256, 256)
+    is_debug=False, shape=(1024, 1024, 8192), block=(256, 256, 256), splitk=None
 ):
     """Double-buffered MXFP4 GEMM, 8 waves, with stagger.
 
@@ -411,6 +459,10 @@ def test_dbuf_8wave_mixed_pingpong_mxfp_gemm(
     barrier releases, effectively hiding the second barrier's latency behind
     the early loads and compute.
     """
+    if splitk:
+        raise NotImplementedError(
+            "split-K is not yet supported with the mixed ping-pong schedule"
+        )
     gemm, options = get_tagged_mxfp4_gemm(shape, block, wave_shape=(4, 2))
     options.specialize = True
     options.use_buffer_ops = True
@@ -426,15 +478,18 @@ def test_dbuf_8wave_mixed_pingpong_mxfp_gemm(
 
 
 def test_dbuf_8wave_mixed_pingpong_shuffle_mxfp_gemm(
-    is_debug=False, shape=(16384, 16384, 16384), block=(256, 256, 256)
+    is_debug=False, shape=(16384, 16384, 16384), block=(256, 256, 256), splitk=None
 ):
     """Like :func:`test_dbuf_8wave_mixed_pingpong_mxfp_gemm` but with A_scale & B_scale
     preshuffled and prefetched to VGPRs.
 
     Note: preshuffling B and loading it directly to VGPRs combined with prefetching
     consumes too many VGPRs and causes spilling.
     """
-
+    if splitk:
+        raise NotImplementedError(
+            "split-K is not yet supported with the mixed ping-pong shuffle schedule"
+        )
     gemm, options = get_tagged_mxfp4_gemm_preshuffle_scales(
         shape, block, wave_shape=(4, 2)
     )
@@ -453,9 +508,13 @@ def test_dbuf_8wave_mixed_pingpong_shuffle_mxfp_gemm(
 
 
 def test_dbuf_4wave_mxfp_asymmetric_gemm(
-    is_debug=False, shape=(1024, 1024, 8192), block=(256, 256, 256)
+    is_debug=False, shape=(1024, 1024, 8192), block=(256, 256, 256), splitk=None
 ):
     """Asymmetric-prefetch MXFP4 GEMM: A through LDS (2x prefetch), B direct from global."""
+    if splitk:
+        raise NotImplementedError(
+            "split-K is not yet supported with the asymmetric schedule"
+        )
     gemm, options = get_tagged_mxfp4_gemm(
         shape, block, wave_shape=(1, 4), b_address_space=GLOBAL_ADDRESS_SPACE
     )
@@ -481,9 +540,17 @@ def test_dbuf_4wave_mxfp_preshuffle_b_gemm(
     shape=(1024, 1024, 8192),
     block=(128, 256, 256),
     eliminate_epilogue=True,
+    splitk=None,
 ):
     """Asymmetric MXFP4 GEMM with preshuffled B data and B scales."""
-    gemm, options = get_tagged_mxfp4_gemm_preshuffle_b(shape, block, wave_shape=(1, 4))
+    if splitk:
+        gemm, options = get_tagged_splitk_mxfp4_gemm_preshuffle_b(
+            shape, num_splits=splitk, block_shape=block, wave_shape=(1, 4)
+        )
+    else:
+        gemm, options = get_tagged_mxfp4_gemm_preshuffle_b(
+            shape, block, wave_shape=(1, 4)
+        )
     options.minimize_shared_allocs = True
     options.linearize_shared_access = True
     options.use_buffer_ops = True
@@ -497,13 +564,18 @@ def test_dbuf_4wave_mxfp_preshuffle_b_gemm(
     gemm = wave_compile(options, gemm, schedule)
 
     _run_mxfp_gemm_preshuffle(gemm, shape, all=True)
-    print("MXFP GEMM preshuffle-B 4-wave test passed!")
+    sk = f" split-K({splitk})" if splitk else ""
+    print(f"MXFP GEMM preshuffle-B 4-wave{sk} test passed!")
 
 
 def test_dbuf_4wave_mxfp_asymmetric_gemm_cpp(
-    is_debug=False, shape=(1024, 1024, 8192), block=(128, 256, 256)
+    is_debug=False, shape=(1024, 1024, 8192), block=(128, 256, 256), splitk=None
 ):
     """Asymmetric MXFP4 GEMM using C++ WaveASM backend (no preshuffle)."""
+    if splitk:
+        raise NotImplementedError(
+            "split-K is not yet supported with the asymmetric C++ backend"
+        )
     gemm, options = get_tagged_mxfp4_gemm(
         shape, block, wave_shape=(1, 4), b_address_space=GLOBAL_ADDRESS_SPACE
     )
@@ -524,8 +596,14 @@ def test_dbuf_4wave_mxfp_preshuffle_b_gemm_cpp(
     shape=(512, 1024, 8192),  # 4*T0, 4*T1, 8192
     block=(128, 256, 256),
     eliminate_epilogue=True,
+    splitk=None,
 ):
     """Preshuffle-B MXFP4 GEMM using C++ WaveASM backend."""
+    if splitk:
+        raise NotImplementedError(
+            "split-K with WaveASM backend hits register alignment errors in "
+            "the assembler; use test_dbuf_4wave_mxfp_preshuffle_b_gemm instead"
+        )
     gemm, options = get_tagged_mxfp4_gemm_preshuffle_b(shape, block, wave_shape=(1, 4))
     options.backend = "asm"
     options.use_buffer_ops = True
@@ -551,8 +629,13 @@ def test_dbuf_4wave_mxfp_dynamic_preshuffle_b_gemm(
     shape=(1024, 1024, 8192),
     block=(128, 256, 256),
     eliminate_epilogue=True,
+    splitk=None,
 ):
     """Preshuffle-B MXFP4 GEMM with dynamic M, N, K."""
+    if splitk:
+        raise NotImplementedError(
+            "split-K is not yet supported with the dynamic preshuffle-B schedule"
+        )
     gemm, options = get_tagged_mxfp4_gemm_preshuffle_b(shape, block, wave_shape=(1, 4))
     # Make M, N, K dynamic so the compiler does not specialize on problem size.
     dynamic_symbols = [tkl.sym.M, tkl.sym.N, tkl.sym.K]
@@ -647,6 +730,33 @@ def test_dbuf_4wave_mxfp_dynamic_preshuffle_b_gemm_asm(
     )
 
 
+def test_splitk_preshuffle_scales_gemm_cpp(
+    is_debug=False,
+    shape=(1024, 1024, 8192),
+    block=(128, 128, 256),
+    splitk=None,
+):
+    """Split-K MXFP4 GEMM using C++ WaveASM backend (preshuffled scales)."""
+    num_splits = splitk if splitk else 2
+    splitk_fn, options = get_tagged_splitk_mxfp4_gemm_preshuffle_scales(
+        shape,
+        num_splits=num_splits,
+        mfma_variant=ScaledMMAType.F32_16x16x128_F8F6F4,
+        block_shape=block,
+        wave_shape=(2, 2),
+        output_type=tkl.f32,
+    )
+    options.backend = "asm"
+    options.wave_runtime = True
+    options.use_wave_asm_backend = True
+    options.print_ir_after = "all" if is_debug else []
+    options = set_default_run_config(options)
+    gemm = wave_compile(options, splitk_fn)
+
+    _run_mxfp_gemm_preshuffle(gemm, shape, only_scale=True)
+    print("Split-K MXFP4 GEMM (preshuffled scales, WaveASM backend) test passed!")
+
+
 if __name__ == "__main__":
     args = parse_args()
 
@@ -667,5 +777,6 @@ def test_dbuf_4wave_mxfp_dynamic_preshuffle_b_gemm_asm(
         args.shape,
         args.block,
         args.eliminate_epilogue,
+        args.splitk,
     )
     exit(0 if success else 1)

examples/python/utils.py

@@ -36,6 +36,13 @@ def parse_args():
         default=None,
         help="Enable epilogue elimination (true/false)",
     )
+    parser.add_argument(
+        "--splitk",
+        type=int,
+        default=None,
+        metavar="NUM_SPLITS",
+        help="Enable split-K with the given number of splits (e.g. --splitk 2)",
+    )
 
     args = parser.parse_args()
 
@@ -64,6 +71,7 @@ def run_test(
     shape=None,
     block=None,
     eliminate_epilogue=None,
+    splitk=None,
 ):
     """Run a test function multiple times."""
     if test_name not in module_globals:
@@ -78,6 +86,8 @@ def run_test(
         kwargs["block"] = block
     if eliminate_epilogue is not None:
         kwargs["eliminate_epilogue"] = eliminate_epilogue
+    if splitk is not None:
+        kwargs["splitk"] = splitk
 
     for i in range(repeat):
         try:

lit_tests/kernel/wave/attention/decode_attention.py

@@ -339,7 +339,7 @@ def test_gqa_flash_decoding():
     # CHECK-LABEL:          func.func @phase_0
     # CHECK:                   scf.for
     # CHECK:                      amdgpu.lds_barrier
-    # CHECK-COUNT-16:             vector.maskedload
+    # CHECK-COUNT-16:             vector.load
     # CHECK-COUNT-16:             vector.store
     # CHECK:                      amdgpu.lds_barrier
     # CHECK-COUNT-4:              amdgpu.mfma