From f5fd27433b6c7ec046b68b197edb21f4ed942fb8 Mon Sep 17 00:00:00 2001 From: Mergen Nachin Date: Tue, 15 Sep 2026 12:12:53 -0700 Subject: [PATCH 1/3] Add a narrow tile candidate for small ROCm INT4 matmuls Offer a 16-column tile for ROCm queries with M=1..4 to increase parallelism. Retain every existing candidate so export-time autotuning can choose for the target GPU rather than imposing the MI300X choice. CUDA and other query sizes keep their existing candidate set; kernel arithmetic and the runtime are unchanged. Validated with 46 focused tests, lintrunner, and full Voxtral W4/BF16 PTE exports executed by the native C++ runner on Linux MI300X. Inference on 30 seconds of audio improved from 12.270 to 7.332 seconds, and on 360 seconds from 156.649 to 97.771 seconds, with identical transcripts and matching token counts. The added candidate also compiled for representative RDNA2/3/4 targets and gfx950; those GPUs were not execution-tested. ROCm 7.1 remains unverified. Lint was run directly because the pre-commit bootstrap cannot build its untokenize dependency under Python 3.14. Authored with assistance from OpenAI Codex. --- backends/cuda/tests/test_int4_matmul.py | 52 +++++++++++++++++++++ backends/cuda/triton/kernels/int4_matmul.py | 19 +++++++- 2 files changed, 70 insertions(+), 1 deletion(-) 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 From 0e64460a171a7aba7f0ca9ed497616ad370666bd Mon Sep 17 00:00:00 2001 From: Mergen Nachin Date: Tue, 15 Sep 2026 13:06:55 -0700 Subject: [PATCH 2/3] Temporarily uncomment CI --- .github/workflows/rocm.yml | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/.github/workflows/rocm.yml b/.github/workflows/rocm.yml index 154454eb145..27b2bac5a2e 100644 --- a/.github/workflows/rocm.yml +++ b/.github/workflows/rocm.yml @@ -229,12 +229,11 @@ jobs: # This scarce RDNA runner is limited to manual runs and direct changes to the # Voxtral ROCm execution path; it does not participate in broad sampling. - # Temporarily disabled while the self-hosted runner teardown is unstable. test-voxtral-realtime-rocm-gfx1100: name: test-voxtral-realtime-rocm-gfx1100-rocm${{ matrix.rocm-version }} needs: [voxtral-run-decision] if: | - false && needs.voxtral-run-decision.outputs.run-gfx1100 == 'true' && + needs.voxtral-run-decision.outputs.run-gfx1100 == 'true' && (github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name == github.repository) concurrency: From ad9adc077d806de30fc9575280793afc613c67dd Mon Sep 17 00:00:00 2001 From: Mergen Nachin Date: Wed, 16 Sep 2026 08:42:04 -0700 Subject: [PATCH 3/3] Revert back CI config --- .github/workflows/rocm.yml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/workflows/rocm.yml b/.github/workflows/rocm.yml index 27b2bac5a2e..154454eb145 100644 --- a/.github/workflows/rocm.yml +++ b/.github/workflows/rocm.yml @@ -229,11 +229,12 @@ jobs: # This scarce RDNA runner is limited to manual runs and direct changes to the # Voxtral ROCm execution path; it does not participate in broad sampling. + # Temporarily disabled while the self-hosted runner teardown is unstable. test-voxtral-realtime-rocm-gfx1100: name: test-voxtral-realtime-rocm-gfx1100-rocm${{ matrix.rocm-version }} needs: [voxtral-run-decision] if: | - needs.voxtral-run-decision.outputs.run-gfx1100 == 'true' && + false && needs.voxtral-run-decision.outputs.run-gfx1100 == 'true' && (github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name == github.repository) concurrency: