diff --git a/backends/cuda/tests/test_int4_matmul.py b/backends/cuda/tests/test_int4_matmul.py index 1c038d82a19..e99d9e26426 100644 --- a/backends/cuda/tests/test_int4_matmul.py +++ b/backends/cuda/tests/test_int4_matmul.py @@ -17,9 +17,12 @@ """ import unittest +from unittest.mock import patch import torch from executorch.backends.cuda.triton.kernels.int4_matmul import ( + _INT4_MATMUL_CONFIGS, + _int4_matmul_prune, dequant_w4_to_bf16, int4_matmul, int4_matvec, @@ -94,6 +97,35 @@ def _eager_int4_matmul(x, w_ref): return (x.float() @ w_ref.float().T).to(torch.bfloat16) +class TestInt4MatmulConfig(unittest.TestCase): + def test_small_rocm_queries_include_narrow_tile(self): + with patch.object(torch.version, "hip", "7.2"): + for m in (1, 2, 3, 4): + with self.subTest(m=m): + configs = _int4_matmul_prune(_INT4_MATMUL_CONFIGS, {"M": m}) + self.assertEqual(configs, _INT4_MATMUL_CONFIGS) + self.assertIn( + {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 16, "BLOCK_SIZE_K": 128}, + [config.kwargs for config in configs], + ) + + def test_other_queries_keep_existing_configs(self): + expected = [c for c in _INT4_MATMUL_CONFIGS if c.kwargs["BLOCK_SIZE_N"] != 16] + for hip, m in ( + (None, 1), + (None, 4), + (None, 16), + ("7.2", 0), + ("7.2", 5), + ("7.2", 16), + ("7.2", None), + ): + with self.subTest(hip=hip, m=m), patch.object(torch.version, "hip", hip): + self.assertEqual( + _int4_matmul_prune(_INT4_MATMUL_CONFIGS, {}, M=m), expected + ) + + class TestDequantW4ToBf16(unittest.TestCase): """Tests for dequant_w4_to_bf16 Triton kernel.""" @@ -204,6 +236,26 @@ def test_non_power_of_two_N(self): def test_small(self): self._run_matmul(1, 16, 64, 32) + def test_streaming_encoder_shapes(self): + for m, n, k in [ + (4, 1280, 5120), + (4, 2048, 1280), + (4, 5120, 1280), + (1, 3072, 5120), + ]: + with self.subTest(m=m, n=n, k=k): + self._run_matmul(m, n, k, 32) + + def test_small_query_tail_and_strides(self): + m, n, k, group_size = 3, 37, 96, 32 + w = torch.randn(n, k, dtype=torch.bfloat16, device=DEVICE) + packed, scale, w_ref = _quantize_simple(w, group_size) + x = torch.randn(k * 2, m * 2, dtype=torch.bfloat16, device=DEVICE)[::2, ::2].T + packed = packed.T.contiguous().T + scale = scale.T.contiguous().T + actual = int4_matmul(x, packed, scale, group_size) + _assert_snr(self, actual, _eager_int4_matmul(x, w_ref), "strided small query") + class TestInt4Matvec(unittest.TestCase): """Tests for int4_matvec Triton kernel (M=1 decode).""" diff --git a/backends/cuda/triton/kernels/int4_matmul.py b/backends/cuda/triton/kernels/int4_matmul.py index e5c91735143..0d282e78238 100644 --- a/backends/cuda/triton/kernels/int4_matmul.py +++ b/backends/cuda/triton/kernels/int4_matmul.py @@ -26,6 +26,11 @@ # -- Autotune configs --------------------------------------------------------- _INT4_MATMUL_CONFIGS = [ + triton.Config( + {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 16, "BLOCK_SIZE_K": 128}, + num_warps=4, + num_stages=5, + ), # Large-M prefill configs (tensor core saturated) triton.Config( {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 128}, @@ -91,10 +96,22 @@ ] +def _int4_matmul_prune(configs, nargs, **kwargs): + m = kwargs.get("M", nargs.get("M")) + # Let small ROCm queries trade tile width for more CTAs on the target GPU. + if torch.version.hip is not None and isinstance(m, int) and 1 <= m <= 4: + return configs + return [config for config in configs if config.kwargs["BLOCK_SIZE_N"] != 16] + + # -- Triton kernel ------------------------------------------------------------ -@triton.autotune(configs=_INT4_MATMUL_CONFIGS, key=["M", "N", "K"]) +@triton.autotune( + configs=_INT4_MATMUL_CONFIGS, + key=["M", "N", "K"], + prune_configs_by={"early_config_prune": _int4_matmul_prune}, +) @triton.jit def _int4_matmul_kernel( # Pointers