Enable fast ROCm Voxtral decode#22005

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
digantdesai merged 2 commits into from
Aug 25, 2026
Merged

Enable fast ROCm Voxtral decode #22005

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
1282a09
Enable fast ROCm Voxtral decode
digantdesai Aug 21, 2026
f7a058c
Temporarily disable gfx1100 Voxtral CI
digantdesai Aug 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion .github/workflows/rocm.yml
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -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:
Expand Down
4 changes: 2 additions & 2 deletions backends/cuda/aoti_packed_int4_tensor.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -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"]
Expand Down
44 changes: 34 additions & 10 deletions backends/cuda/tests/test_sdpa_splitk_replacement.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Expand DownExpand Up@@ -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
)
Expand All@@ -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}",
)
Comment on lines +143 to +148
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
)
Expand All@@ -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."""
Expand Down
4 changes: 3 additions & 1 deletion backends/cuda/triton/replacement_pass.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Expand Down
37 changes: 6 additions & 31 deletions examples/models/voxtral_realtime/README.md
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -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

Expand DownExpand Up@@ -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.
Expand DownExpand Up@@ -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.

Expand Down
26 changes: 5 additions & 21 deletions examples/models/voxtral_realtime/export_voxtral_rt.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -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."""
Expand All@@ -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_(
Expand All@@ -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
)
Expand DownExpand Up@@ -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."""
Expand DownExpand Up@@ -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:
Expand DownExpand Up@@ -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."""
Expand DownExpand Up@@ -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:
Expand DownExpand Up@@ -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")
Expand DownExpand Up@@ -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",
Expand DownExpand Up@@ -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:
Expand Down
11 changes: 0 additions & 11 deletions examples/models/voxtral_realtime/run_rocm_e2e.sh
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -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).
Expand All@@ -33,19 +32,13 @@ 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}"

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) ;;
Expand DownExpand Up@@ -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)"
Expand Down
Loading
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Add copy buttons to all 
 blocks\n(function() {\n function addCopyButtons() {\n document.querySelectorAll('pre code').forEach(function(codeBlock) {\n if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;\n codeBlock.parentElement.setAttribute('data-copy-added', 'true');\n \n var btn = document.createElement('button');\n btn.textContent = 'Copy';\n btn.style.cssText = 'position:absolute;top:4px;right:4px;padding:2px 8px;font-size:11px;background:#4ecdc4;border:none;border-radius:4px;color:#1a1a2e;cursor:pointer;opacity:0.7;transition:opacity 0.2s;';\n btn.onmouseover = function() { this.style.opacity = '1'; };\n btn.onmouseout = function() { this.style.opacity = '0.7'; };\n btn.onclick = function() {\n navigator.clipboard.writeText(codeBlock.textContent).then(function() {\n btn.textContent = 'Copied!';\n setTimeout(function() { btn.textContent = 'Copy'; }, 1500);\n });\n };\n codeBlock.parentElement.style.position = 'relative';\n codeBlock.parentElement.appendChild(btn);\n });\n }\n \n addCopyButtons();\n \n // Re-run on dynamic content\n var observer = new MutationObserver(addCopyButtons);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Add Copy Buttons to Code Blocks");
}
} catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); }
})();
(function(){
try {
var __m = "github.com";
var __re = new RegExp('^' + "github\\.com" + '

Skip to content
Enable fast ROCm Voxtral decode#22005

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

digantdesai
merged 2 commits into

from

Aug 25, 2026
 Merged


Enable fast ROCm Voxtral decode
#22005
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


2 commits


Select commit
Hold shift + click to select a range
1282a09

Enable fast ROCm Voxtral decode


digantdesai Aug 21, 2026
f7a058c

Temporarily disable gfx1100 Voxtral CI


digantdesai Aug 24, 2026

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


3 changes: 2 additions & 1 deletion
.github/workflows/rocm.yml

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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:
Expand Down

4 changes: 2 additions & 2 deletions
backends/cuda/aoti_packed_int4_tensor.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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"]
Expand Down

44 changes: 34 additions & 10 deletions
backends/cuda/tests/test_sdpa_splitk_replacement.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Expand DownExpand Up@@ -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
 )
Expand All@@ -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}",
 )

Comment on lines

+143
 to 
+148

 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
 )
Expand All@@ -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."""
Expand Down

4 changes: 3 additions & 1 deletion
backends/cuda/triton/replacement_pass.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Expand Down

37 changes: 6 additions & 31 deletions
examples/models/voxtral_realtime/README.md

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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

Expand DownExpand Up@@ -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.
Expand DownExpand Up@@ -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.

Expand Down

26 changes: 5 additions & 21 deletions
examples/models/voxtral_realtime/export_voxtral_rt.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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."""
Expand All@@ -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_(
Expand All@@ -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
 )
Expand DownExpand Up@@ -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."""
Expand DownExpand Up@@ -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:
Expand DownExpand Up@@ -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."""
Expand DownExpand Up@@ -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:
Expand DownExpand Up@@ -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")
Expand DownExpand Up@@ -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",
Expand DownExpand Up@@ -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:
Expand Down

11 changes: 0 additions & 11 deletions
examples/models/voxtral_realtime/run_rocm_e2e.sh

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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).
Expand All@@ -33,19 +32,13 @@ 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}"

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) ;;
Expand DownExpand Up@@ -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)"
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Force GitHub README to respect dark mode\n(function() {\n var style = document.createElement('style');\n style.textContent = '\n .markdown-body {\n color-scheme: dark light;\n }\n .markdown-body pre { background: #161b22 !important; }\n .markdown-body code { background: rgba(110, 118, 129, 0.4) !important; }\n .markdown-body table th, .markdown-body table td { border-color: #30363d !important; }\n .markdown-body img { background: #0d1117; }\n .markdown-body blockquote { border-left-color: #8b949e; }\n .markdown-body hr { border-color: #30363d; }\n ';\n document.head.appendChild(style);\n})();", "GitHub Dark Mode README Fix");
}
} catch(__e) { console.warn('[Userscript:GitHub Dark Mode README Fix]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
Enable fast ROCm Voxtral decode#22005

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

digantdesai
merged 2 commits into

from

Aug 25, 2026
 Merged


Enable fast ROCm Voxtral decode
#22005
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


2 commits


Select commit
Hold shift + click to select a range
1282a09

Enable fast ROCm Voxtral decode


digantdesai Aug 21, 2026
f7a058c

Temporarily disable gfx1100 Voxtral CI


digantdesai Aug 24, 2026

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


3 changes: 2 additions & 1 deletion
.github/workflows/rocm.yml

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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:
Expand Down

4 changes: 2 additions & 2 deletions
backends/cuda/aoti_packed_int4_tensor.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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"]
Expand Down

44 changes: 34 additions & 10 deletions
backends/cuda/tests/test_sdpa_splitk_replacement.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Expand DownExpand Up@@ -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
 )
Expand All@@ -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}",
 )

Comment on lines

+143
 to 
+148

 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
 )
Expand All@@ -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."""
Expand Down

4 changes: 3 additions & 1 deletion
backends/cuda/triton/replacement_pass.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Expand Down

37 changes: 6 additions & 31 deletions
examples/models/voxtral_realtime/README.md

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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

Expand DownExpand Up@@ -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.
Expand DownExpand Up@@ -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.

Expand Down

26 changes: 5 additions & 21 deletions
examples/models/voxtral_realtime/export_voxtral_rt.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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."""
Expand All@@ -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_(
Expand All@@ -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
 )
Expand DownExpand Up@@ -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."""
Expand DownExpand Up@@ -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:
Expand DownExpand Up@@ -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."""
Expand DownExpand Up@@ -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:
Expand DownExpand Up@@ -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")
Expand DownExpand Up@@ -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",
Expand DownExpand Up@@ -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:
Expand Down

11 changes: 0 additions & 11 deletions
examples/models/voxtral_realtime/run_rocm_e2e.sh

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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).
Expand All@@ -33,19 +32,13 @@ 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}"

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) ;;
Expand DownExpand Up@@ -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)"
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Highlight search terms from Google/DuckDuckGo/Bing referrer\n(function() {\n var ref = document.referrer;\n var terms = [];\n \n if (ref.includes('google.com') || ref.includes('duckduckgo.com') || ref.includes('bing.com')) {\n var url = new URL(ref);\n var q = url.searchParams.get('q') || url.searchParams.get('p');\n if (q) {\n terms = q.split(/\\s+/).filter(function(t) { return t.length > 2; });\n }\n }\n \n if (terms.length === 0) return;\n \n var style = document.createElement('style');\n style.textContent = '.userscript-highlight { background: #fbbf24; color: #1a1a2e; padding: 1px 3px; border-radius: 2px; }';\n document.head.appendChild(style);\n \n function highlight(node) {\n if (node.nodeType === 3) { // text node\n var text = node.textContent;\n var found = false;\n terms.forEach(function(term) {\n var regex = new RegExp('(' + term.replace(/[.*+?^${}()|[\\]\\\\]/g, '\\\\') + ')', 'gi');\n if (regex.test(text)) {\n found = true;\n var frag = document.createDocumentFragment();\n var parts = text.split(regex);\n parts.forEach(function(part, i) {\n if (i % 2 === 0) {\n frag.appendChild(document.createTextNode(part));\n } else {\n var span = document.createElement('span');\n span.className = 'userscript-highlight';\n span.textContent = part;\n frag.appendChild(span);\n }\n });\n node.parentNode.replaceChild(frag, node);\n }\n });\n } else if (node.nodeType === 1 && node.childNodes) { // element\n var skipTags = ['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEXTAREA', 'INPUT', 'SELECT'];\n if (!skipTags.includes(node.tagName)) {\n Array.from(node.childNodes).forEach(highlight);\n }\n }\n }\n \n highlight(document.body);\n \n // Re-highlight on dynamic content\n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1 || node.nodeType === 3) highlight(node);\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Highlight Search Terms");
}
} catch(__e) { console.warn('[Userscript:Highlight Search Terms]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
Enable fast ROCm Voxtral decode#22005

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

digantdesai
merged 2 commits into

from

Aug 25, 2026
 Merged


Enable fast ROCm Voxtral decode
#22005
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


2 commits


Select commit
Hold shift + click to select a range
1282a09

Enable fast ROCm Voxtral decode


digantdesai Aug 21, 2026
f7a058c

Temporarily disable gfx1100 Voxtral CI


digantdesai Aug 24, 2026

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


3 changes: 2 additions & 1 deletion
.github/workflows/rocm.yml

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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:
Expand Down

4 changes: 2 additions & 2 deletions
backends/cuda/aoti_packed_int4_tensor.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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"]
Expand Down

44 changes: 34 additions & 10 deletions
backends/cuda/tests/test_sdpa_splitk_replacement.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Expand DownExpand Up@@ -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
 )
Expand All@@ -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}",
 )

Comment on lines

+143
 to 
+148

 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
 )
Expand All@@ -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."""
Expand Down

4 changes: 3 additions & 1 deletion
backends/cuda/triton/replacement_pass.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Expand Down

37 changes: 6 additions & 31 deletions
examples/models/voxtral_realtime/README.md

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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

Expand DownExpand Up@@ -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.
Expand DownExpand Up@@ -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.

Expand Down

26 changes: 5 additions & 21 deletions
examples/models/voxtral_realtime/export_voxtral_rt.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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."""
Expand All@@ -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_(
Expand All@@ -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
 )
Expand DownExpand Up@@ -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."""
Expand DownExpand Up@@ -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:
Expand DownExpand Up@@ -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."""
Expand DownExpand Up@@ -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:
Expand DownExpand Up@@ -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")
Expand DownExpand Up@@ -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",
Expand DownExpand Up@@ -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:
Expand Down

11 changes: 0 additions & 11 deletions
examples/models/voxtral_realtime/run_rocm_e2e.sh

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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).
Expand All@@ -33,19 +32,13 @@ 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}"

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) ;;
Expand DownExpand Up@@ -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)"
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Strip utm_, fbclid, gclid, etc. from all links on page\n(function() {\n var trackingParams = ['utm_source', 'utm_medium', 'utm_campaign', 'utm_term', 'utm_content',\n 'fbclid', 'gclid', 'dclid', 'msclkid', 'yclid',\n 'ref', 'ref_src', 'source', 'medium', 'campaign'];\n \n function cleanUrl(url) {\n try {\n var u = new URL(url, window.location.origin);\n var changed = false;\n trackingParams.forEach(function(p) {\n if (u.searchParams.has(p)) {\n u.searchParams.delete(p);\n changed = true;\n }\n });\n return changed ? u.toString() : url;\n } catch (e) {\n return url;\n }\n }\n \n function cleanLinks() {\n document.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n \n cleanLinks();\n \n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1) {\n if (node.tagName === 'A') cleanLinks();\n node.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Remove Tracking Parameters from Links");
}
} catch(__e) { console.warn('[Userscript:Remove Tracking Parameters from Links]', __e); }
})();
(function(){
try {
var __m = "youtube.com";
var __re = new RegExp('^' + "youtube\\.com" + '

Skip to content
Enable fast ROCm Voxtral decode#22005

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

digantdesai
merged 2 commits into

from

Aug 25, 2026
 Merged


Enable fast ROCm Voxtral decode
#22005
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


2 commits


Select commit
Hold shift + click to select a range
1282a09

Enable fast ROCm Voxtral decode


digantdesai Aug 21, 2026
f7a058c

Temporarily disable gfx1100 Voxtral CI


digantdesai Aug 24, 2026

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


3 changes: 2 additions & 1 deletion
.github/workflows/rocm.yml

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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:
Expand Down

4 changes: 2 additions & 2 deletions
backends/cuda/aoti_packed_int4_tensor.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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"]
Expand Down

44 changes: 34 additions & 10 deletions
backends/cuda/tests/test_sdpa_splitk_replacement.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Expand DownExpand Up@@ -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
 )
Expand All@@ -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}",
 )

Comment on lines

+143
 to 
+148

 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
 )
Expand All@@ -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."""
Expand Down

4 changes: 3 additions & 1 deletion
backends/cuda/triton/replacement_pass.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Expand Down

37 changes: 6 additions & 31 deletions
examples/models/voxtral_realtime/README.md

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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

Expand DownExpand Up@@ -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.
Expand DownExpand Up@@ -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.

Expand Down

26 changes: 5 additions & 21 deletions
examples/models/voxtral_realtime/export_voxtral_rt.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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."""
Expand All@@ -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_(
Expand All@@ -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
 )
Expand DownExpand Up@@ -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."""
Expand DownExpand Up@@ -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:
Expand DownExpand Up@@ -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."""
Expand DownExpand Up@@ -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:
Expand DownExpand Up@@ -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")
Expand DownExpand Up@@ -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",
Expand DownExpand Up@@ -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:
Expand Down

11 changes: 0 additions & 11 deletions
examples/models/voxtral_realtime/run_rocm_e2e.sh

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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).
Expand All@@ -33,19 +32,13 @@ 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}"

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) ;;
Expand DownExpand Up@@ -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)"
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Auto-enable theater mode on YouTube\n(function() {\n function tryTheater() {\n var btn = document.querySelector('button[aria-label=\"Theater mode\"], ytd-player #player button[title=\"Theater mode\"]');\n if (btn && !btn.classList.contains('activated')) {\n btn.click();\n }\n }\n \n // Try immediately\n tryTheater();\n \n // Try after navigation (SPA)\n var lastUrl = location.href;\n setInterval(function() {\n if (location.href !== lastUrl) {\n lastUrl = location.href;\n setTimeout(tryTheater, 500);\n }\n }, 1000);\n \n // Also try on player load\n var observer = new MutationObserver(tryTheater);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "YouTube Theater Mode Default");
}
} catch(__e) { console.warn('[Userscript:YouTube Theater Mode Default]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
Enable fast ROCm Voxtral decode#22005

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

digantdesai
merged 2 commits into

from

Aug 25, 2026
 Merged


Enable fast ROCm Voxtral decode
#22005
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


2 commits


Select commit
Hold shift + click to select a range
1282a09

Enable fast ROCm Voxtral decode


digantdesai Aug 21, 2026
f7a058c

Temporarily disable gfx1100 Voxtral CI


digantdesai Aug 24, 2026

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


3 changes: 2 additions & 1 deletion
.github/workflows/rocm.yml

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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:
Expand Down

4 changes: 2 additions & 2 deletions
backends/cuda/aoti_packed_int4_tensor.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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"]
Expand Down

44 changes: 34 additions & 10 deletions
backends/cuda/tests/test_sdpa_splitk_replacement.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Expand DownExpand Up@@ -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
 )
Expand All@@ -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}",
 )

Comment on lines

+143
 to 
+148

 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
 )
Expand All@@ -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."""
Expand Down

4 changes: 3 additions & 1 deletion
backends/cuda/triton/replacement_pass.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Expand Down

37 changes: 6 additions & 31 deletions
examples/models/voxtral_realtime/README.md

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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

Expand DownExpand Up@@ -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.
Expand DownExpand Up@@ -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.

Expand Down

26 changes: 5 additions & 21 deletions
examples/models/voxtral_realtime/export_voxtral_rt.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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."""
Expand All@@ -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_(
Expand All@@ -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
 )
Expand DownExpand Up@@ -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."""
Expand DownExpand Up@@ -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:
Expand DownExpand Up@@ -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."""
Expand DownExpand Up@@ -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:
Expand DownExpand Up@@ -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")
Expand DownExpand Up@@ -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",
Expand DownExpand Up@@ -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:
Expand Down

11 changes: 0 additions & 11 deletions
examples/models/voxtral_realtime/run_rocm_e2e.sh

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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).
Expand All@@ -33,19 +32,13 @@ 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}"

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) ;;
Expand DownExpand Up@@ -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)"
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Remove or un-stick sticky/fixed headers that block content\n(function() {\n function unstick() {\n document.querySelectorAll('header, nav, [role=\"banner\"], .header, .navbar, .sticky, .fixed-top, [style*=\"position: fixed\"], [style*=\"position:sticky\"]').forEach(function(el) {\n if (el.style.position === 'fixed' || el.style.position === 'sticky' || \n getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {\n el.style.position = 'static';\n el.style.top = 'auto';\n el.style.zIndex = 'auto';\n }\n });\n }\n \n unstick();\n \n var observer = new MutationObserver(unstick);\n observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });\n})();", "Kill Sticky Headers");
}
} catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
Enable fast ROCm Voxtral decode#22005

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

digantdesai
merged 2 commits into

from

Aug 25, 2026
 Merged


Enable fast ROCm Voxtral decode
#22005
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


2 commits


Select commit
Hold shift + click to select a range
1282a09

Enable fast ROCm Voxtral decode


digantdesai Aug 21, 2026
f7a058c

Temporarily disable gfx1100 Voxtral CI


digantdesai Aug 24, 2026

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


3 changes: 2 additions & 1 deletion
.github/workflows/rocm.yml

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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:
Expand Down

4 changes: 2 additions & 2 deletions
backends/cuda/aoti_packed_int4_tensor.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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"]
Expand Down

44 changes: 34 additions & 10 deletions
backends/cuda/tests/test_sdpa_splitk_replacement.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Expand DownExpand Up@@ -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
 )
Expand All@@ -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}",
 )

Comment on lines

+143
 to 
+148

 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
 )
Expand All@@ -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."""
Expand Down

4 changes: 3 additions & 1 deletion
backends/cuda/triton/replacement_pass.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Expand Down

37 changes: 6 additions & 31 deletions
examples/models/voxtral_realtime/README.md

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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

Expand DownExpand Up@@ -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.
Expand DownExpand Up@@ -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.

Expand Down

26 changes: 5 additions & 21 deletions
examples/models/voxtral_realtime/export_voxtral_rt.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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."""
Expand All@@ -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_(
Expand All@@ -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
 )
Expand DownExpand Up@@ -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."""
Expand DownExpand Up@@ -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:
Expand DownExpand Up@@ -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."""
Expand DownExpand Up@@ -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:
Expand DownExpand Up@@ -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")
Expand DownExpand Up@@ -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",
Expand DownExpand Up@@ -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:
Expand Down

11 changes: 0 additions & 11 deletions
examples/models/voxtral_realtime/run_rocm_e2e.sh

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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).
Expand All@@ -33,19 +32,13 @@ 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}"

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) ;;
Expand DownExpand Up@@ -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)"
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Universal Dark Mode - works on any site\n(function() {\n var enabled = true;\n \n function applyDarkMode() {\n if (!enabled) return;\n \n // Create style element if it doesn't exist\n var style = document.getElementById('universal-dark-mode-style');\n if (!style) {\n style = document.createElement('style');\n style.id = 'universal-dark-mode-style';\n document.head.appendChild(style);\n }\n \n // Dark mode CSS - inverts colors but preserves images/video\n style.textContent = '\n /* Invert everything except media */\n html {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #1a1a2e !important;\n }\n \n /* Restore images, videos, iframes, canvas */\n img, video, iframe, canvas, svg, picture, [style*=\"background-image\"] {\n filter: invert(1) hue-rotate(180deg) !important;\n }\n \n /* Preserve specific elements that should not be inverted */\n .no-dark-mode, .no-dark-mode *,\n [data-theme=\"light\"], [data-theme=\"light\"],\n .ace_editor, .ace_editor *,\n .CodeMirror, .CodeMirror *,\n .monaco-editor, .monaco-editor *,\n .markdown-body pre, .markdown-body pre *,\n .highlight, .highlight *,\n pre code, pre code * {\n filter: none !important;\n }\n \n /* Fix common UI elements */\n .modal, .popup, .dropdown-menu, .tooltip, .popover {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #2d2d44 !important;\n border-color: #444 !important;\n }\n \n /* Scrollbars */\n ::-webkit-scrollbar { background: #1a1a2e !important; }\n ::-webkit-scrollbar-thumb { background: #444 !important; }\n ::-webkit-scrollbar-thumb:hover { background: #555 !important; }\n \n /* Selection */\n ::selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ::-moz-selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ';\n }\n \n function removeDarkMode() {\n var style = document.getElementById('universal-dark-mode-style');\n if (style) style.remove();\n }\n \n // Toggle with Alt+Shift+D\n document.addEventListener('keydown', function(e) {\n if (e.altKey && e.shiftKey && e.key === 'D') {\n e.preventDefault();\n enabled = !enabled;\n if (enabled) {\n applyDarkMode();\n console.log('[Universal Dark Mode] Enabled');\n } else {\n removeDarkMode();\n console.log('[Universal Dark Mode] Disabled');\n }\n }\n });\n \n // Apply on load\n applyDarkMode();\n \n // Re-apply on dynamic content\n var observer = new MutationObserver(function(mutations) {\n if (enabled && !document.getElementById('universal-dark-mode-style')) {\n applyDarkMode();\n }\n });\n observer.observe(document.head, { childList: true });\n \n console.log('[Universal Dark Mode] Loaded - Press Alt+Shift+D to toggle');\n})();", "Universal Dark Mode");
}
} catch(__e) { console.warn('[Userscript:Universal Dark Mode]', __e); }
})();
})();

Skip to content
Enable fast ROCm Voxtral decode#22005

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

digantdesai
merged 2 commits into

from

Aug 25, 2026
 Merged


Enable fast ROCm Voxtral decode
#22005
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


2 commits


Select commit
Hold shift + click to select a range
1282a09

Enable fast ROCm Voxtral decode


digantdesai Aug 21, 2026
f7a058c

Temporarily disable gfx1100 Voxtral CI


digantdesai Aug 24, 2026

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


3 changes: 2 additions & 1 deletion
.github/workflows/rocm.yml

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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:
Expand Down

4 changes: 2 additions & 2 deletions
backends/cuda/aoti_packed_int4_tensor.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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"]
Expand Down

44 changes: 34 additions & 10 deletions
backends/cuda/tests/test_sdpa_splitk_replacement.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Expand DownExpand Up@@ -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
 )
Expand All@@ -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}",
 )

Comment on lines

+143
 to 
+148

 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
 )
Expand All@@ -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."""
Expand Down

4 changes: 3 additions & 1 deletion
backends/cuda/triton/replacement_pass.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Expand Down

37 changes: 6 additions & 31 deletions
examples/models/voxtral_realtime/README.md

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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

Expand DownExpand Up@@ -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.
Expand DownExpand Up@@ -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.

Expand Down

26 changes: 5 additions & 21 deletions
examples/models/voxtral_realtime/export_voxtral_rt.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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."""
Expand All@@ -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_(
Expand All@@ -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
 )
Expand DownExpand Up@@ -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."""
Expand DownExpand Up@@ -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:
Expand DownExpand Up@@ -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."""
Expand DownExpand Up@@ -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:
Expand DownExpand Up@@ -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")
Expand DownExpand Up@@ -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",
Expand DownExpand Up@@ -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:
Expand Down

11 changes: 0 additions & 11 deletions
examples/models/voxtral_realtime/run_rocm_e2e.sh

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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).
Expand All@@ -33,19 +32,13 @@ 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}"

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) ;;
Expand DownExpand Up@@ -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)"
Expand Down
Loading