Add schedule for 256x224x256 macro tile

It is a no_unroll schedule to get under the register budget.
This gets the macro tile functional with the waveasm backend.

For the 7.1 example, it adds
- `--wave_shape` flag -- Previously (1,4) was hard-coded, but the 256x224x256 tile needed (2, 2) because the N dimension was not divisible by 4 after pipelining... I think was the reason we chose that.
- `--no_unroll` flag to access the new no_unroll schedule.

The particular 7.1 example target for this work was
`python examples/python/7.1_schedule.py --block 256,224,256 --shape 1024,896,8192 --wave_shape 2,2 --no-unroll --test test_dbuf_4wave_mxfp_preshuffle_b_gemm_cpp`

This also adds an e2e waveasm test.

At this stage no real effort has been made to make the schedule performant, just to get it working.

Author: willghatch

examples/python/7.1_schedule.py

@@ -27,6 +27,7 @@
     get_mxfp4_dbuf_pingpong_schedule,
     get_mxfp4_dbuf_mixed_pingpong_schedule,
     get_mxfp4_asymmetric_schedule,
+    get_mxfp4_asymmetric_nounroll_schedule,
     get_mxfp4_dbuf_mixed_pingpong_shuffle_schedule,
     get_mxfp4_dbuf_pingpong_schedule_Bshuffled,
     get_mxfp4_dbuf_pingpong_schedule_Bshuffled_lds,
@@ -372,19 +373,26 @@ def test_dbuf_4wave_mxfp_preshuffle_b_gemm_cpp(
     is_debug=False,
     shape=(512, 1024, 8192),  # 4*T0, 4*T1, 8192
     block=(128, 256, 256),
+    wave_shape=(1, 4),
     eliminate_epilogue=True,
+    no_unroll=False,
 ):
     """Preshuffle-B MXFP4 GEMM using C++ WaveASM backend."""
-    gemm, options = get_tagged_mxfp4_gemm_preshuffle_b(shape, block, wave_shape=(1, 4))
+    gemm, options = get_tagged_mxfp4_gemm_preshuffle_b(shape, block, wave_shape=wave_shape)
     options.backend = "asm"
     options.use_buffer_ops = True
-    options.wave_runtime = True
     options.use_wave_asm_backend = True
+    options.wave_runtime = True
     options.dump_intermediates = "build/intermediates"
     options.eliminate_epilogue = eliminate_epilogue
-    schedule = get_mxfp4_asymmetric_schedule(
-        eliminate_epilogue=eliminate_epilogue, is_bscale_shuffled=True
-    )
+    if no_unroll:
+        schedule = get_mxfp4_asymmetric_nounroll_schedule(
+            eliminate_epilogue=eliminate_epilogue, is_bscale_shuffled=True
+        )
+    else:
+        schedule = get_mxfp4_asymmetric_schedule(
+            eliminate_epilogue=eliminate_epilogue, is_bscale_shuffled=True
+        )
     options.print_ir_after = "all" if is_debug else []
     options = set_default_run_config(options)
     gemm = wave_compile(options, gemm, schedule)
@@ -444,5 +452,7 @@ def test_dbuf_4wave_mxfp_dynamic_preshuffle_b_gemm(
         args.shape,
         args.block,
         args.eliminate_epilogue,
+        args.wave_shape,
+        no_unroll=args.no_unroll,
     )
     exit(0 if success else 1)

examples/python/utils.py

@@ -36,6 +36,17 @@ def parse_args():
         default=None,
         help="Enable epilogue elimination (true/false)",
     )
+    parser.add_argument(
+        "--wave_shape",
+        type=str,
+        default=None,
+        help="Wave shape, e.g. 2,2",
+    )
+    parser.add_argument(
+        "--no-unroll",
+        action="store_true",
+        help="Use nounroll (unroll_factor=1) schedule variant",
+    )
 
     args = parser.parse_args()
 
@@ -44,6 +55,8 @@ def parse_args():
         args.shape = tuple(map(int, args.shape.split(",")))
     if isinstance(args.block, str):
         args.block = tuple(map(int, args.block.split(",")))
+    if isinstance(args.wave_shape, str):
+        args.wave_shape = tuple(map(int, args.wave_shape.split(",")))
 
     return args
 
@@ -64,6 +77,8 @@ def run_test(
     shape=None,
     block=None,
     eliminate_epilogue=None,
+    wave_shape=None,
+    no_unroll=False,
 ):
     """Run a test function multiple times."""
     if test_name not in module_globals:
@@ -78,6 +93,10 @@ def run_test(
         kwargs["block"] = block
     if eliminate_epilogue is not None:
         kwargs["eliminate_epilogue"] = eliminate_epilogue
+    if wave_shape is not None:
+        kwargs["wave_shape"] = wave_shape
+    if no_unroll:
+        kwargs["no_unroll"] = True
 
     for i in range(repeat):
         try:

tests/kernel/wave/asm/test_waveasm_e2e.py

@@ -1142,6 +1142,7 @@ def _dbuf_mxfp4_helper(
     wave_shape=None,
     reorder_workgroups=None,
     eliminate_epilogue=False,
+    no_unroll=False,
 ):
     """Shared helper for double-buffered MXFP4 scheduled GEMM tests.
 
@@ -1168,6 +1169,7 @@ def _dbuf_mxfp4_helper(
     from wave_lang.kernel.wave.schedules import (
         get_mxfp4_dbuf_schedule,
         get_mxfp4_asymmetric_schedule,
+        get_mxfp4_asymmetric_nounroll_schedule,
     )
     from wave_lang.kernel.wave.scheduling.schedule_enums import SchedulingType
     from wave_lang.kernel.wave.utils.run_utils import set_default_run_config
@@ -1200,9 +1202,14 @@ def _dbuf_mxfp4_helper(
         )
         options.eliminate_epilogue = eliminate_epilogue
         if use_schedule:
-            schedule = get_mxfp4_asymmetric_schedule(
-                eliminate_epilogue=eliminate_epilogue, is_bscale_shuffled=True
-            )
+            if no_unroll:
+                schedule = get_mxfp4_asymmetric_nounroll_schedule(
+                    eliminate_epilogue=eliminate_epilogue, is_bscale_shuffled=True
+                )
+            else:
+                schedule = get_mxfp4_asymmetric_schedule(
+                    eliminate_epilogue=eliminate_epilogue, is_bscale_shuffled=True
+                )
         else:
             schedule = None
             options.schedule = SchedulingType.NONE
@@ -1442,6 +1449,44 @@ def test_dbuf_4wave_mxfp4_gemm_cpp_backend(
     )
 
 
+@pytest.mark.parametrize("eliminate_epilogue", [True], ids=["ee"])
+@pytest.mark.parametrize(
+    "shape,block,wave_shape",
+    [
+        pytest.param((1024, 896, 8192), (256, 224, 256), (2, 2), id="256x224x256"),
+    ],
+)
+def test_dbuf_4wave_mxfp4_nounroll_gemm_cpp_backend(
+    shape,
+    block,
+    wave_shape,
+    eliminate_epilogue,
+    compiler,
+    dump_asm,
+):
+    """End-to-end test for asymmetric MXFP4 GEMM with no-unroll schedule.
+
+    The no-unroll schedule (unroll_factor=1) reduces register pressure by
+    not unrolling the K-loop body, allowing larger block sizes like
+    256x224x256 that would otherwise exceed the 256-VGPR hardware limit
+    with the standard asymmetric schedule.
+    """
+    _dbuf_mxfp4_helper(
+        shape=shape,
+        block=block,
+        num_waves=4,
+        use_stagger=False,
+        compiler=compiler,
+        dump_asm=dump_asm,
+        use_buffer_ops=True,
+        use_schedule=True,
+        output_dtype="f32",
+        wave_shape=wave_shape,
+        eliminate_epilogue=eliminate_epilogue,
+        no_unroll=True,
+    )
+
+
 @pytest.mark.parametrize(
     "shape,block,wave_shape",
     [

wave_lang/kernel/wave/schedules/__init__.py

@@ -18,6 +18,7 @@
     get_mxfp4_dbuf_mixed_pingpong_schedule,
     get_mxfp4_dbuf_mixed_pingpong_shuffle_schedule,
     get_mxfp4_asymmetric_schedule,
+    get_mxfp4_asymmetric_nounroll_schedule,
     get_mxfp4_dbuf_pingpong_schedule_Bshuffled,
     get_mxfp4_dbuf_pingpong_schedule_Bshuffled_lds,
 )
@@ -34,6 +35,7 @@
     "get_mxfp4_dbuf_pingpong_schedule_Bshuffled",
     "get_mxfp4_dbuf_pingpong_schedule_Bshuffled_lds",
     "get_mxfp4_asymmetric_schedule",
+    "get_mxfp4_asymmetric_nounroll_schedule",
     "get_mxfp4_dbuf_mixed_pingpong_shuffle_schedule",
     "get_attention_prefetch_schedule",
 ]