From 1282a09fd6851a43a1afc8e473640a454612dce2 Mon Sep 17 00:00:00 2001 From: Digant Desai Date: Thu, 20 Aug 2026 21:27:56 -0700 Subject: [PATCH 1/2] Enable fast ROCm Voxtral decode Use the packed INT4 matvec automatically for ROCm W4 decoder exports. Keep ROCm fixed-shape decode on standard SDPA after split-K produced non-finite logits. CUDA and generic packed tensor defaults are unchanged. With assistance from Claude Code and Codex. --- backends/cuda/aoti_packed_int4_tensor.py | 4 +- .../tests/test_sdpa_splitk_replacement.py | 44 ++++++++++++++----- backends/cuda/triton/replacement_pass.py | 4 +- examples/models/voxtral_realtime/README.md | 37 +++------------- .../voxtral_realtime/export_voxtral_rt.py | 26 +++-------- .../models/voxtral_realtime/run_rocm_e2e.sh | 11 ----- 6 files changed, 50 insertions(+), 76 deletions(-) diff --git a/backends/cuda/aoti_packed_int4_tensor.py b/backends/cuda/aoti_packed_int4_tensor.py index 096de622c08..2c300808b79 100644 --- a/backends/cuda/aoti_packed_int4_tensor.py +++ b/backends/cuda/aoti_packed_int4_tensor.py @@ -18,8 +18,8 @@ class AotiPackedInt4Tensor(TorchAOBaseTensor): """Symmetric groupwise INT4 weight consumed by AOTI Triton kernels. - Linears use ``triton::int4_matmul`` by default; the opt-in fixed-shape path - uses ``triton::int4_matvec_bf16``. + Linears use ``triton::int4_matmul`` by default; a fixed-shape caller can + select ``triton::int4_matvec_bf16``. """ tensor_data_names = ["qdata", "scale"] diff --git a/backends/cuda/tests/test_sdpa_splitk_replacement.py b/backends/cuda/tests/test_sdpa_splitk_replacement.py index 465b0b7ecf4..28accad20bd 100644 --- a/backends/cuda/tests/test_sdpa_splitk_replacement.py +++ b/backends/cuda/tests/test_sdpa_splitk_replacement.py @@ -7,8 +7,8 @@ """Test ReplaceEdgeOpWithTritonOpPass split-K SDPA kernel selection. Exports a minimal model containing F.scaled_dot_product_attention through the -CUDA backend and verifies that the pass routes to split-K for decode -(L_q==1, L_kv >= 256) and standard SDPA otherwise. +CUDA backend and verifies that CUDA routes eligible decode shapes to split-K, +while ROCm and other shapes use standard SDPA. """ import logging @@ -127,8 +127,8 @@ def test_below_threshold_uses_standard(self): f"Expected 1 SDPA replaced with standard kernel. Log: {msgs}", ) - def test_at_threshold_uses_splitk(self): - """L_kv=256 == threshold -> split-K selected (boundary, inclusive).""" + def test_at_threshold_uses_backend_kernel(self): + """L_kv=256 selects split-K on CUDA and standard SDPA on ROCm.""" model = SDPAModule(n_heads=4, n_kv_heads=2, head_dim=64, kv_len=256).to( torch.bfloat16 ) @@ -140,11 +140,23 @@ def test_at_threshold_uses_splitk(self): _, msgs = _capture_pass_logs(lambda: _export_through_cuda_backend(model, args)) splitk = [m for m in msgs if "split-K" in m] - self.assertEqual(len(splitk), 1, f"Expected 1 split-K selection. Log: {msgs}") - self.assertIn("L_kv=256", splitk[0]) + expected = 0 if torch.version.hip is not None else 1 + self.assertEqual( + len(splitk), + expected, + f"Expected {expected} split-K selections. Log: {msgs}", + ) + if expected: + self.assertIn("L_kv=256", splitk[0]) + + replaced = [m for m in msgs if "Replaced" in m] + self.assertTrue( + any("1 nodes" in m for m in replaced), + f"Expected 1 SDPA replaced with a Triton kernel. Log: {msgs}", + ) - def test_large_kv_cache_uses_splitk(self): - """L_kv=4096 > threshold -> split-K selected for decode.""" + def test_large_kv_cache_uses_backend_kernel(self): + """L_kv=4096 selects split-K on CUDA and standard SDPA on ROCm.""" model = SDPAModule(n_heads=4, n_kv_heads=2, head_dim=64, kv_len=4096).to( torch.bfloat16 ) @@ -156,8 +168,20 @@ def test_large_kv_cache_uses_splitk(self): _, msgs = _capture_pass_logs(lambda: _export_through_cuda_backend(model, args)) splitk = [m for m in msgs if "split-K" in m] - self.assertEqual(len(splitk), 1, f"Expected 1 split-K selection. Log: {msgs}") - self.assertIn("L_kv=4096", splitk[0]) + expected = 0 if torch.version.hip is not None else 1 + self.assertEqual( + len(splitk), + expected, + f"Expected {expected} split-K selections. Log: {msgs}", + ) + if expected: + self.assertIn("L_kv=4096", splitk[0]) + + replaced = [m for m in msgs if "Replaced" in m] + self.assertTrue( + any("1 nodes" in m for m in replaced), + f"Expected 1 SDPA replaced with a Triton kernel. Log: {msgs}", + ) def test_non_pow2_head_dim_uses_standard(self): """Non-power-of-2 head_dim -> standard SDPA even with large L_kv.""" diff --git a/backends/cuda/triton/replacement_pass.py b/backends/cuda/triton/replacement_pass.py index d8284e50260..18d592aeb15 100644 --- a/backends/cuda/triton/replacement_pass.py +++ b/backends/cuda/triton/replacement_pass.py @@ -104,8 +104,10 @@ def _pick_sdpa_kernel(node: Node): L_q, D = q_shape[2], q_shape[3] L_kv = k_shape[2] + # TODO: Re-enable split-K after validating ROCm Voxtral decode numerics. if ( - isinstance(L_q, int) + torch.version.hip is None + and isinstance(L_q, int) and L_q == 1 and isinstance(L_kv, int) and L_kv >= _SPLITK_LKV_THRESHOLD diff --git a/examples/models/voxtral_realtime/README.md b/examples/models/voxtral_realtime/README.md index c150d42188b..9d9346c021e 100644 --- a/examples/models/voxtral_realtime/README.md +++ b/examples/models/voxtral_realtime/README.md @@ -206,32 +206,10 @@ before model loading. The packed path performs dequantization inside the GPU kernel and does not materialize a full BF16 weight for each invocation. -The default packed path retains the existing dynamic decoder export and uses -the packed INT4 matmul kernel. CUDA and other non-ROCm exports are unchanged. -Encoder linears also use packed INT4 matmul. - -An experimental ROCm-only matvec export is available for performance testing: - -```bash -python export_voxtral_rt.py \ - --model-path ~/models/Voxtral-Mini-4B-Realtime-2602 \ - --backend rocm \ - --dtype bf16 \ - --streaming \ - --sliding-window 2048 \ - --rocm-packed-matvec \ - --output-dir ./voxtral_rt_rocm_w4_bf16_matvec \ - --qlinear-encoder 4w \ - --qlinear 4w \ - --qembedding 8w -``` - -This specializes the decoder to its actual one-token runner input and uses a -BF16-rounded packed matvec. On MI300X it roughly doubled decode throughput for -the 30-second test clip, but greedy output differed from the dynamic matmul -baseline. It is off by default; verify transcript quality and performance on -the target GPU before enabling it. Kernel and export-graph tests cover this -option, but CI does not run a full-model transcript check with it. +The ROCm W4 decoder is specialized to the runner's one-token input and uses +the packed INT4 matvec kernel. Encoder linears use packed INT4 matmul. ROCm +uses the standard SDPA kernel because split-K decode produced non-finite logits +for this fixed-shape workload. CUDA and other non-ROCm exports are unchanged. #### Metal export examples @@ -377,8 +355,6 @@ python export_voxtral_rt.py \ | `--streaming` | off | Export streaming model with ring buffer KV caches (unlimited duration) | | `--max-enc-len` | `750` | Encoder sliding window size (streaming only) | | `--sliding-window` | from `params.json` | Decoder sliding window size (streaming only; ignored in offline mode). Smaller values reduce memory and improve decode speed but limit context | -| `--rocm-packed-matvec` | off | Experimental fixed-shape packed INT4 decoder matvec; requires ROCm and decoder `4w` | - **Notes:** - `fpa4w` quantization requires `--backend metal`. - The model was trained with `--delay-tokens 6`. Other values may degrade accuracy. @@ -435,9 +411,8 @@ examples/models/voxtral_realtime/run_rocm_e2e.sh \ ``` The third argument selects `bf16`, `w4-bf16`, or both precision modes. The -fourth selects `streaming`, `offline`, or both execution modes. Set -`ROCM_PACKED_MATVEC=1` to opt into the experimental fixed-shape decoder matvec. -Set `ROCM_PATH` if ROCm is installed outside `/opt/rocm`. +fourth selects `streaming`, `offline`, or both execution modes. Set `ROCM_PATH` +if ROCm is installed outside `/opt/rocm`. The script reports model export time, PTE/PTD sizes, and RTF computed as runner inference time divided by WAV duration. diff --git a/examples/models/voxtral_realtime/export_voxtral_rt.py b/examples/models/voxtral_realtime/export_voxtral_rt.py index 2aabe62d9b5..2a1c65cd98c 100644 --- a/examples/models/voxtral_realtime/export_voxtral_rt.py +++ b/examples/models/voxtral_realtime/export_voxtral_rt.py @@ -142,7 +142,6 @@ def _export_decoder_and_embedding( qembedding, qembedding_group_size, use_aoti_packed_int4=False, - use_aoti_matvec=False, device="cpu", ): """Export text_decoder and token_embedding into programs dict.""" @@ -155,6 +154,7 @@ def _export_decoder_and_embedding( text_decoder.eval() packed_linear_count = 0 + use_packed_matvec = use_aoti_packed_int4 and qlinear == "4w" if qlinear: print(f" Quantizing decoder ({qlinear})...") quantize_model_( @@ -163,16 +163,14 @@ def _export_decoder_and_embedding( qlinear_group_size=qlinear_group_size, qlinear_packing_format=qlinear_packing_format, ) - if use_aoti_packed_int4 and qlinear == "4w": + if use_packed_matvec: packed_linear_count = _pack_aoti_int4_weights( text_decoder, - use_matvec=use_aoti_matvec, + use_matvec=True, ) - if use_aoti_matvec: - # TODO: Resolve fixed-shape greedy-output drift before enabling this by - # default; the same drift reproduces with int4_matmul. - # Both native runner paths invoke the decoder one token at a time. + # Native runners decode one token per call; static M=1 enables matvec dispatch. + if use_packed_matvec: sample_embeds = torch.randn( 1, 1, model.config.dim, dtype=param_dtype, device=device ) @@ -232,7 +230,6 @@ def export_all( qembedding=None, qembedding_group_size=None, use_aoti_packed_int4=False, - use_aoti_matvec=False, backend="xnnpack", ): """Export all three model components with per-component quantization.""" @@ -297,7 +294,6 @@ def export_all( qembedding, qembedding_group_size, use_aoti_packed_int4=use_aoti_packed_int4, - use_aoti_matvec=use_aoti_matvec, device=device, ) if packed_linear_count: @@ -331,7 +327,6 @@ def export_streaming( qembedding=None, qembedding_group_size=None, use_aoti_packed_int4=False, - use_aoti_matvec=False, backend="xnnpack", ): """Export streaming model components with per-component quantization.""" @@ -392,7 +387,6 @@ def export_streaming( qembedding, qembedding_group_size, use_aoti_packed_int4=use_aoti_packed_int4, - use_aoti_matvec=use_aoti_matvec, device=device, ) if packed_linear_count: @@ -583,15 +577,11 @@ def _validate_rocm_args(parser, args): "tile_packed_to_4d requires a CUDA-only int4 fallback; " "omit the packing format for ROCm" ) - if args.rocm_packed_matvec and args.qlinear != "4w": - parser.error("--rocm-packed-matvec requires --qlinear=4w") def _validate_export_args(parser, args, backend_for_export): if args.backend == "rocm": _validate_rocm_args(parser, args) - elif args.rocm_packed_matvec: - parser.error("--rocm-packed-matvec requires --backend=rocm") if args.qlinear == "fpa4w" and backend_for_export != "metal": parser.error("--qlinear=fpa4w can only be used with --backend=metal") @@ -708,11 +698,6 @@ def main(): "typically 8192). Smaller values reduce memory and improve decode speed " "but limit how far back the decoder can attend. Only used with --streaming.", ) - parser.add_argument( - "--rocm-packed-matvec", - action="store_true", - help="Use the experimental fixed-shape packed INT4 decoder matvec on ROCm.", - ) parser.add_argument( "--dtype", default="fp32", @@ -771,7 +756,6 @@ def main(): "qembedding": args.qembedding, "qembedding_group_size": args.qembedding_group_size, "use_aoti_packed_int4": args.backend == "rocm", - "use_aoti_matvec": args.rocm_packed_matvec, "backend": backend_for_export, } if args.streaming: diff --git a/examples/models/voxtral_realtime/run_rocm_e2e.sh b/examples/models/voxtral_realtime/run_rocm_e2e.sh index 78ae3ee3a60..950ffaeecac 100755 --- a/examples/models/voxtral_realtime/run_rocm_e2e.sh +++ b/examples/models/voxtral_realtime/run_rocm_e2e.sh @@ -9,7 +9,6 @@ # SKIP_EXPORT=1 Use existing model.pte and aoti_cuda_blob.ptd files. # DEVICE_INDEX Visible GPU index (default: 0). # SLIDING_WINDOW Decoder window (default: 2048). -# ROCM_PACKED_MATVEC=1 Use the experimental fixed-shape decoder matvec. # OFFLINE_MAX_NEW_TOKENS Offline token limit (default: 500). # VOXTRAL_PYTHON Python executable (default: python). # ROCM_PATH ROCm installation (default: /opt/rocm). @@ -33,7 +32,6 @@ EXECUTION_MODE="${4:-streaming}" OUTPUT_ROOT="${5:-$PWD/voxtral_rt_rocm}" DEVICE_INDEX="${DEVICE_INDEX:-0}" SLIDING_WINDOW="${SLIDING_WINDOW:-2048}" -ROCM_PACKED_MATVEC="${ROCM_PACKED_MATVEC:-0}" OFFLINE_MAX_NEW_TOKENS="${OFFLINE_MAX_NEW_TOKENS:-500}" VOXTRAL_PYTHON="${VOXTRAL_PYTHON:-python}" ROCM_ROOT="${ROCM_PATH:-/opt/rocm}" @@ -41,11 +39,6 @@ ROCM_ROOT="${ROCM_PATH:-/opt/rocm}" export HIP_VISIBLE_DEVICES="$DEVICE_INDEX" export CUDA_VISIBLE_DEVICES="$DEVICE_INDEX" -if [[ "$ROCM_PACKED_MATVEC" != "0" && "$ROCM_PACKED_MATVEC" != "1" ]]; then - echo "ERROR: ROCM_PACKED_MATVEC must be 0 or 1" >&2 - exit 1 -fi - case "$PRECISION_MODE" in bf16) PRECISIONS=(bf16) ;; w4-bf16) PRECISIONS=(w4-bf16) ;; @@ -158,10 +151,6 @@ for precision in "${PRECISIONS[@]}"; do if [[ "$execution" == "streaming" ]]; then export_args+=(--streaming --sliding-window "$SLIDING_WINDOW") fi - if [[ "$precision" == "w4-bf16" && "$ROCM_PACKED_MATVEC" == "1" ]]; then - export_args+=(--rocm-packed-matvec) - fi - export_elapsed_ms=-1 if [[ "${SKIP_EXPORT:-0}" != "1" ]]; then export_start_ms="$(monotonic_ms)" From f7a058cb2bbb9fa46b69adcc327b5e13a1aad5ee Mon Sep 17 00:00:00 2001 From: Digant Desai Date: Mon, 24 Aug 2026 06:44:37 -0700 Subject: [PATCH 2/2] Temporarily disable gfx1100 Voxtral CI Keep the job definition and trigger conditions in place, but force it off while the self-hosted runner teardown remains unstable. With assistance from Claude Code and Codex. --- .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 3b7844208c3..ae047996122 100644 --- a/.github/workflows/rocm.yml +++ b/.github/workflows/rocm.yml @@ -222,11 +222,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: