fix(tune): restore late AS and split issue/wait for M-major W4A8 grouped GEMM

humming/tune/sm90_policies.py

@@ -397,6 +397,8 @@ def _set_w4a8_config(config: dict, block_m: int, shape_k: int = 0) -> None:
         use_warp_spec=True,
         use_stream_k=False,
         use_packed_k_layout=True,
+        wgmma_use_late_as=True,
+        wgmma_split_issue_wait=block_m <= 160,
         raster_group_m=16,
         multi_cast_size_a=1,
         multi_cast_size_b=1,

tests/test_sm90_heuristics.py

@@ -3,7 +3,8 @@
 import pytest
 
 from humming import dtypes
-from humming.config import GemmType, LayerConfig
+from humming.config import GemmType, LayerConfig, TuningConfig
+from humming.config.mma import get_register_budget_error
 from humming.device import DeviceInfo
 from humming.tune import _get_heuristics_config
 from humming.tune.sm90 import Sm90Heuristics
@@ -385,6 +386,14 @@ def test_grouped_w4a8_ranges_match_direct_selection(monkeypatch, device_name, nu
         assert lower == previous_upper and lower < upper
         assert config["use_packed_k_layout"]
         assert config["warp_shape"][1:] == (16, 128)
+        assert config["wgmma_use_late_as"]
+        assert config["wgmma_split_issue_wait"] == (config["warp_shape"][0] <= 160)
+        tuning_keys = ("mma_type", "block_shape", "warp_shape", "num_stages", "use_warp_spec")
+        tuning = TuningConfig(
+            **{key: config[key] for key in tuning_keys},
+            wgmma_split_issue_wait=config["wgmma_split_issue_wait"],
+        )
+        assert get_register_budget_error(layer, tuning, registers_per_sm=65536) is None
         for shape_m in (lower + 1, min(upper, lower + num_experts * 2048)):
             assert select(layer, shape_m=shape_m, **kwargs) == config
         previous_upper = upper