fix(triton): improve RDNA INT8 autotune selection

Keep ten candidates per device while replacing an unused gfx12-oriented tile with a small-M configuration on gfx11. Retain the original gfx12 pool, preserve useful BK128 candidates, and prune only the unstable 128x128x128 tile for small M.

Keep device_index in the autotune key but mark it non-specializing in the generated kernels, so heterogeneous devices tune independently without duplicate code variants.

Signed-off-by: liminfei-amd <91481003+liminfei-amd@users.noreply.github.com>

Author: liminfei-amd

comfy_kitchen/backends/triton/quantization.py

@@ -857,7 +857,7 @@ def triton_quantize_and_rotate_rowwise(x: torch.Tensor, h: torch.Tensor, group_s
 
 
 def _int8_autotune_configs():
-    """Return the default and RDNA-tuned INT8 matmul config pools."""
+    """Return the default and generation-specific RDNA config pools."""
     default_configs = [
         triton.Config({'block_m': 128, 'block_n': 256, 'block_k': 64, 'group_size_m': 8}, num_stages=3, num_warps=8),
         triton.Config({'block_m': 64,  'block_n': 256, 'block_k': 32, 'group_size_m': 8}, num_stages=4, num_warps=4),
@@ -866,7 +866,7 @@ def _int8_autotune_configs():
         triton.Config({'block_m': 64,  'block_n': 128, 'block_k': 32, 'group_size_m': 8}, num_stages=4, num_warps=4),
         triton.Config({'block_m': 128, 'block_n': 32,  'block_k': 32, 'group_size_m': 8}, num_stages=4, num_warps=4),
     ]
-    rdna_configs = [
+    gfx12_configs = [
         triton.Config({'block_m': 128, 'block_n': 256, 'block_k': 64,  'group_size_m': 8}, num_stages=2, num_warps=8),
         triton.Config({'block_m': 128, 'block_n': 128, 'block_k': 128, 'group_size_m': 8}, num_stages=2, num_warps=8),
         triton.Config({'block_m': 128, 'block_n': 128, 'block_k': 64,  'group_size_m': 8}, num_stages=2, num_warps=4),
@@ -878,11 +878,26 @@ def _int8_autotune_configs():
         triton.Config({'block_m': 64,  'block_n': 128, 'block_k': 64,  'group_size_m': 8, 'waves_per_eu': 1}, num_stages=2, num_warps=4),
         triton.Config({'block_m': 64,  'block_n': 64,  'block_k': 64,  'group_size_m': 4, 'waves_per_eu': 1}, num_stages=2, num_warps=4),
     ]
-    return default_configs, rdna_configs
-
-
-_INT8_DEFAULT_CONFIGS, _INT8_RDNA_CONFIGS = _int8_autotune_configs()
-_INT8_MATMUL_CONFIGS = _INT8_DEFAULT_CONFIGS + _INT8_RDNA_CONFIGS
+    gfx11_configs = list(gfx12_configs)
+    gfx11_configs[6] = triton.Config(
+        {'block_m': 64, 'block_n': 64, 'block_k': 128,
+         'group_size_m': 4, 'waves_per_eu': 1},
+        num_stages=2,
+        num_warps=8,
+    )
+    return default_configs, gfx11_configs, gfx12_configs
+
+
+(
+    _INT8_DEFAULT_CONFIGS,
+    _INT8_GFX11_CONFIGS,
+    _INT8_GFX12_CONFIGS,
+) = _int8_autotune_configs()
+# Per-device pruning returns ten RDNA candidates. The decorator sees their
+# eleven-config union so no device pays to benchmark the other generation's tile.
+_INT8_MATMUL_CONFIGS = (
+    _INT8_DEFAULT_CONFIGS + _INT8_GFX12_CONFIGS + [_INT8_GFX11_CONFIGS[6]]
+)
 
 
 def _prune_int8_autotune_configs(configs, named_args, **kwargs):
@@ -894,13 +909,23 @@ def _prune_int8_autotune_configs(configs, named_args, **kwargs):
         arch = torch.cuda.get_device_properties(args['device_index']).gcnArchName.split(":")[0]
     except Exception:
         return _INT8_DEFAULT_CONFIGS
-    if not arch.startswith(("gfx11", "gfx12")):
+    if arch.startswith("gfx11"):
+        rdna_configs = _INT8_GFX11_CONFIGS
+    elif arch.startswith("gfx12"):
+        rdna_configs = _INT8_GFX12_CONFIGS
+    else:
         return _INT8_DEFAULT_CONFIGS
 
-    rdna_configs = _INT8_RDNA_CONFIGS
-    # BK128 was consistently slower for small-M calls on gfx1151.
+    # This large tile was unstable for small-M calls on gfx1151.
     if args['m'] <= 128:
-        rdna_configs = [config for config in rdna_configs if config.kwargs['block_k'] <= 64]
+        rdna_configs = [
+            config for config in rdna_configs
+            if not (
+                config.kwargs['block_m'] == 128
+                and config.kwargs['block_n'] == 128
+                and config.kwargs['block_k'] == 128
+            )
+        ]
     return rdna_configs
 
 
@@ -913,13 +938,13 @@ def _prune_int8_autotune_configs(configs, named_args, **kwargs):
 }
 
 @triton.autotune(**_INT8_AUTOTUNE_KWARGS)
-@triton.jit
+@triton.jit(do_not_specialize=["device_index"])
 def _int8_matmul_dequant_kernel(
     # Pointers
     a_ptr, b_ptr, c_ptr,
     a_scale_ptr, b_scale_ptr, bias_ptr,
     # Matrix Dimensions
-    m, n, k, device_index: tl.constexpr,
+    m, n, k, device_index,
     # Strides
     stride_am, stride_ak,
     stride_bk, stride_bn,
@@ -985,13 +1010,13 @@ def _int8_matmul_dequant_kernel(
     tl.store(c_ptrs, c, mask=c_mask)
 
 @triton.autotune(**_INT8_AUTOTUNE_KWARGS)
-@triton.jit
+@triton.jit(do_not_specialize=["device_index"])
 def _int8_matmul_dequant_per_row_kernel(
     # Pointers
     a_ptr, b_ptr, c_ptr,
     a_scale_ptr, b_scale_ptr, bias_ptr,
     # Matrix Dimensions
-    m, n, k, device_index: tl.constexpr,
+    m, n, k, device_index,
     # Strides
     stride_am, stride_ak,
     stride_bk, stride_bn,

tests/test_int8.py

@@ -230,16 +230,17 @@ def record_call(*args, **kwargs):
     "hip_version,arches,device_index,m,expected_pool",
     [
         (None, ["gfx1100"], 0, 1024, "default"),
-        ("7.2", ["gfx1100"], 0, 1024, "rdna"),
-        ("7.2", ["gfx1151"], 0, 1024, "rdna"),
-        ("7.2", ["gfx1201"], 0, 1024, "rdna"),
-        ("7.2", ["gfx1100"], 0, 128, "rdna-small-m"),
-        ("7.2", ["gfx1100"], 0, 129, "rdna"),
+        ("7.2", ["gfx1100"], 0, 1024, "gfx11"),
+        ("7.2", ["gfx1151"], 0, 1024, "gfx11"),
+        ("7.2", ["gfx1201"], 0, 1024, "gfx12"),
+        ("7.2", ["gfx1100"], 0, 128, "gfx11-small-m"),
+        ("7.2", ["gfx1201"], 0, 128, "gfx12-small-m"),
+        ("7.2", ["gfx1100"], 0, 129, "gfx11"),
         ("7.2", ["gfx908"], 0, 1024, "default"),
         ("7.2", ["gfx90a"], 0, 1024, "default"),
         ("7.2", ["gfx942"], 0, 1024, "default"),
         ("7.2", ["gfx950"], 0, 1024, "default"),
-        ("7.2", ["gfx1100", "gfx90a"], 0, 1024, "rdna"),
+        ("7.2", ["gfx1100", "gfx90a"], 0, 1024, "gfx11"),
         ("7.2", ["gfx1100", "gfx90a"], 1, 1024, "default"),
         ("7.2", ["gfx1300"], 0, 1024, "default"),
     ],
@@ -263,12 +264,24 @@ def test_int8_autotune_configs_are_device_scoped(
         m=m,
     )
 
-    if expected_pool == "rdna":
-        assert configs == quantization._INT8_RDNA_CONFIGS
-    elif expected_pool == "rdna-small-m":
+    expected_configs = {
+        "gfx11": quantization._INT8_GFX11_CONFIGS,
+        "gfx12": quantization._INT8_GFX12_CONFIGS,
+    }
+    if expected_pool in expected_configs:
+        assert configs == expected_configs[expected_pool]
+    elif expected_pool.endswith("-small-m"):
+        full_pool = expected_configs[expected_pool.removesuffix("-small-m")]
         assert len(configs) == 9
-        assert all(config in quantization._INT8_RDNA_CONFIGS for config in configs)
-        assert all(config.kwargs["block_k"] <= 64 for config in configs)
+        assert all(config in full_pool for config in configs)
+        if expected_pool.startswith("gfx11"):
+            assert any(config.kwargs["block_k"] == 128 for config in configs)
+        assert not any(
+            config.kwargs["block_m"] == 128
+            and config.kwargs["block_n"] == 128
+            and config.kwargs["block_k"] == 128
+            for config in configs
+        )
     else:
         signatures = [
             (config.kwargs, config.num_stages, config.num_warps)
@@ -302,7 +315,7 @@ def test_int8_autotune_configs_isolate_heterogeneous_devices(monkeypatch):
         quantization._INT8_MATMUL_CONFIGS, {}, device_index=0, m=1024
     )
 
-    assert rdna_configs == quantization._INT8_RDNA_CONFIGS
+    assert rdna_configs == quantization._INT8_GFX11_CONFIGS
     assert cdna_configs == quantization._INT8_DEFAULT_CONFIGS
     assert rdna_configs_again == rdna_configs
 
@@ -336,10 +349,13 @@ def test_int8_autotune_pool_creation_does_not_probe_devices(monkeypatch):
         lambda _: (_ for _ in ()).throw(AssertionError("unexpected device probe")),
     )
 
-    default_configs, rdna_configs = quantization._int8_autotune_configs()
+    default_configs, gfx11_configs, gfx12_configs = quantization._int8_autotune_configs()
 
     assert len(default_configs) == 6
-    assert len(rdna_configs) == 10
+    assert len(gfx11_configs) == 10
+    assert len(gfx12_configs) == 10
+    assert len(quantization._INT8_MATMUL_CONFIGS) == 17
+    assert gfx11_configs[6] != gfx12_configs[6]
 
 
 def test_int8_autotune_cache_key_includes_device():
@@ -351,8 +367,14 @@ def test_int8_autotune_cache_key_includes_device():
         "prune_configs_by": quantization._INT8_PRUNE_CONFIGS_BY,
     }
 
-    assert "device_index" in quantization._int8_matmul_dequant_kernel.keys
-    assert "device_index" in quantization._int8_matmul_dequant_per_row_kernel.keys
+    for kernel in (
+        quantization._int8_matmul_dequant_kernel,
+        quantization._int8_matmul_dequant_per_row_kernel,
+    ):
+        assert "device_index" in kernel.keys
+        device_index = kernel.fn.params[kernel.fn.arg_names.index("device_index")]
+        assert not device_index.is_constexpr
+        assert device_index.do_not_specialize
 
 
 # =============================================================================