Skip to content

Commit b55d8fc

Browse files
authored
Enable fast ROCm Voxtral decode (pytorch#22005)
Use the packed INT4 matvec automatically for ROCm W4 decoder exports. Keep ROCm fixed-shape decode on standard SDPA after split-K. CUDA and generic packed tensor defaults are unchanged.
1 parent 0db103b commit b55d8fc

7 files changed

Lines changed: 52 additions & 77 deletions

File tree

‎.github/workflows/rocm.yml‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -222,11 +222,12 @@ jobs:
222222
223223
# This scarce RDNA runner is limited to manual runs and direct changes to the
224224
# Voxtral ROCm execution path; it does not participate in broad sampling.
225+
# Temporarily disabled while the self-hosted runner teardown is unstable.
225226
test-voxtral-realtime-rocm-gfx1100:
226227
name: test-voxtral-realtime-rocm-gfx1100-rocm${{ matrix.rocm-version }}
227228
needs: [voxtral-run-decision]
228229
if: |
229-
needs.voxtral-run-decision.outputs.run-gfx1100 == 'true' &&
230+
false && needs.voxtral-run-decision.outputs.run-gfx1100 == 'true' &&
230231
(github.event_name != 'pull_request' ||
231232
github.event.pull_request.head.repo.full_name == github.repository)
232233
concurrency:

‎backends/cuda/aoti_packed_int4_tensor.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,8 @@
1818
class AotiPackedInt4Tensor(TorchAOBaseTensor):
1919
"""Symmetric groupwise INT4 weight consumed by AOTI Triton kernels.
2020
21-
Linears use ``triton::int4_matmul`` by default; the opt-in fixed-shape path
22-
uses ``triton::int4_matvec_bf16``.
21+
Linears use ``triton::int4_matmul`` by default; a fixed-shape caller can
22+
select ``triton::int4_matvec_bf16``.
2323
"""
2424

2525
tensor_data_names = ["qdata", "scale"]

‎backends/cuda/tests/test_sdpa_splitk_replacement.py‎

Lines changed: 34 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,8 @@
77
"""Test ReplaceEdgeOpWithTritonOpPass split-K SDPA kernel selection.
88
99
Exports a minimal model containing F.scaled_dot_product_attention through the
10-
CUDA backend and verifies that the pass routes to split-K for decode
11-
(L_q==1, L_kv >= 256) and standard SDPA otherwise.
10+
CUDA backend and verifies that CUDA routes eligible decode shapes to split-K,
11+
while ROCm and other shapes use standard SDPA.
1212
"""
1313

1414
import logging
@@ -127,8 +127,8 @@ def test_below_threshold_uses_standard(self):
127127
f"Expected 1 SDPA replaced with standard kernel. Log: {msgs}",
128128
)
129129

130-
def test_at_threshold_uses_splitk(self):
131-
"""L_kv=256 == threshold -> split-K selected (boundary, inclusive)."""
130+
def test_at_threshold_uses_backend_kernel(self):
131+
"""L_kv=256 selects split-K on CUDA and standard SDPA on ROCm."""
132132
model = SDPAModule(n_heads=4, n_kv_heads=2, head_dim=64, kv_len=256).to(
133133
torch.bfloat16
134134
)
@@ -140,11 +140,23 @@ def test_at_threshold_uses_splitk(self):
140140
_, msgs = _capture_pass_logs(lambda: _export_through_cuda_backend(model, args))
141141

142142
splitk = [m for m in msgs if "split-K" in m]
143-
self.assertEqual(len(splitk), 1, f"Expected 1 split-K selection. Log: {msgs}")
144-
self.assertIn("L_kv=256", splitk[0])
143+
expected = 0 if torch.version.hip is not None else 1
144+
self.assertEqual(
145+
len(splitk),
146+
expected,
147+
f"Expected {expected} split-K selections. Log: {msgs}",
148+
)
149+
if expected:
150+
self.assertIn("L_kv=256", splitk[0])
151+
152+
replaced = [m for m in msgs if "Replaced" in m]
153+
self.assertTrue(
154+
any("1 nodes" in m for m in replaced),
155+
f"Expected 1 SDPA replaced with a Triton kernel. Log: {msgs}",
156+
)
145157

146-
def test_large_kv_cache_uses_splitk(self):
147-
"""L_kv=4096 > threshold -> split-K selected for decode."""
158+
def test_large_kv_cache_uses_backend_kernel(self):
159+
"""L_kv=4096 selects split-K on CUDA and standard SDPA on ROCm."""
148160
model = SDPAModule(n_heads=4, n_kv_heads=2, head_dim=64, kv_len=4096).to(
149161
torch.bfloat16
150162
)
@@ -156,8 +168,20 @@ def test_large_kv_cache_uses_splitk(self):
156168
_, msgs = _capture_pass_logs(lambda: _export_through_cuda_backend(model, args))
157169

158170
splitk = [m for m in msgs if "split-K" in m]
159-
self.assertEqual(len(splitk), 1, f"Expected 1 split-K selection. Log: {msgs}")
160-
self.assertIn("L_kv=4096", splitk[0])
171+
expected = 0 if torch.version.hip is not None else 1
172+
self.assertEqual(
173+
len(splitk),
174+
expected,
175+
f"Expected {expected} split-K selections. Log: {msgs}",
176+
)
177+
if expected:
178+
self.assertIn("L_kv=4096", splitk[0])
179+
180+
replaced = [m for m in msgs if "Replaced" in m]
181+
self.assertTrue(
182+
any("1 nodes" in m for m in replaced),
183+
f"Expected 1 SDPA replaced with a Triton kernel. Log: {msgs}",
184+
)
161185

162186
def test_non_pow2_head_dim_uses_standard(self):
163187
"""Non-power-of-2 head_dim -> standard SDPA even with large L_kv."""

‎backends/cuda/triton/replacement_pass.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -104,8 +104,10 @@ def _pick_sdpa_kernel(node: Node):
104104
L_q, D = q_shape[2], q_shape[3]
105105
L_kv = k_shape[2]
106106

107+
# TODO: Re-enable split-K after validating ROCm Voxtral decode numerics.
107108
if (
108-
isinstance(L_q, int)
109+
torch.version.hip is None
110+
and isinstance(L_q, int)
109111
and L_q == 1
110112
and isinstance(L_kv, int)
111113
and L_kv >= _SPLITK_LKV_THRESHOLD

‎examples/models/voxtral_realtime/README.md‎

Lines changed: 6 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -206,32 +206,10 @@ before model loading.
206206
The packed path performs dequantization inside the GPU kernel and does not
207207
materialize a full BF16 weight for each invocation.
208208

209-
The default packed path retains the existing dynamic decoder export and uses
210-
the packed INT4 matmul kernel. CUDA and other non-ROCm exports are unchanged.
211-
Encoder linears also use packed INT4 matmul.
212-
213-
An experimental ROCm-only matvec export is available for performance testing:
214-
215-
```bash
216-
python export_voxtral_rt.py \
217-
--model-path ~/models/Voxtral-Mini-4B-Realtime-2602 \
218-
--backend rocm \
219-
--dtype bf16 \
220-
--streaming \
221-
--sliding-window 2048 \
222-
--rocm-packed-matvec \
223-
--output-dir ./voxtral_rt_rocm_w4_bf16_matvec \
224-
--qlinear-encoder 4w \
225-
--qlinear 4w \
226-
--qembedding 8w
227-
```
228-
229-
This specializes the decoder to its actual one-token runner input and uses a
230-
BF16-rounded packed matvec. On MI300X it roughly doubled decode throughput for
231-
the 30-second test clip, but greedy output differed from the dynamic matmul
232-
baseline. It is off by default; verify transcript quality and performance on
233-
the target GPU before enabling it. Kernel and export-graph tests cover this
234-
option, but CI does not run a full-model transcript check with it.
209+
The ROCm W4 decoder is specialized to the runner's one-token input and uses
210+
the packed INT4 matvec kernel. Encoder linears use packed INT4 matmul. ROCm
211+
uses the standard SDPA kernel because split-K decode produced non-finite logits
212+
for this fixed-shape workload. CUDA and other non-ROCm exports are unchanged.
235213

236214
#### Metal export examples
237215

@@ -377,8 +355,6 @@ python export_voxtral_rt.py \
377355
| `--streaming` | off | Export streaming model with ring buffer KV caches (unlimited duration) |
378356
| `--max-enc-len` | `750` | Encoder sliding window size (streaming only) |
379357
| `--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 |
380-
| `--rocm-packed-matvec` | off | Experimental fixed-shape packed INT4 decoder matvec; requires ROCm and decoder `4w` |
381-
382358
**Notes:**
383359
- `fpa4w` quantization requires `--backend metal`.
384360
- The model was trained with `--delay-tokens 6`. Other values may degrade accuracy.
@@ -435,9 +411,8 @@ examples/models/voxtral_realtime/run_rocm_e2e.sh \
435411
```
436412

437413
The third argument selects `bf16`, `w4-bf16`, or both precision modes. The
438-
fourth selects `streaming`, `offline`, or both execution modes. Set
439-
`ROCM_PACKED_MATVEC=1` to opt into the experimental fixed-shape decoder matvec.
440-
Set `ROCM_PATH` if ROCm is installed outside `/opt/rocm`.
414+
fourth selects `streaming`, `offline`, or both execution modes. Set `ROCM_PATH`
415+
if ROCm is installed outside `/opt/rocm`.
441416
The script reports model export time, PTE/PTD sizes, and RTF computed as runner
442417
inference time divided by WAV duration.
443418

‎examples/models/voxtral_realtime/export_voxtral_rt.py‎

Lines changed: 5 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -142,7 +142,6 @@ def _export_decoder_and_embedding(
142142
qembedding,
143143
qembedding_group_size,
144144
use_aoti_packed_int4=False,
145-
use_aoti_matvec=False,
146145
device="cpu",
147146
):
148147
"""Export text_decoder and token_embedding into programs dict."""
@@ -155,6 +154,7 @@ def _export_decoder_and_embedding(
155154
text_decoder.eval()
156155

157156
packed_linear_count = 0
157+
use_packed_matvec = use_aoti_packed_int4 and qlinear == "4w"
158158
if qlinear:
159159
print(f" Quantizing decoder ({qlinear})...")
160160
quantize_model_(
@@ -163,16 +163,14 @@ def _export_decoder_and_embedding(
163163
qlinear_group_size=qlinear_group_size,
164164
qlinear_packing_format=qlinear_packing_format,
165165
)
166-
if use_aoti_packed_int4 and qlinear == "4w":
166+
if use_packed_matvec:
167167
packed_linear_count = _pack_aoti_int4_weights(
168168
text_decoder,
169-
use_matvec=use_aoti_matvec,
169+
use_matvec=True,
170170
)
171171

172-
if use_aoti_matvec:
173-
# TODO: Resolve fixed-shape greedy-output drift before enabling this by
174-
# default; the same drift reproduces with int4_matmul.
175-
# Both native runner paths invoke the decoder one token at a time.
172+
# Native runners decode one token per call; static M=1 enables matvec dispatch.
173+
if use_packed_matvec:
176174
sample_embeds = torch.randn(
177175
1, 1, model.config.dim, dtype=param_dtype, device=device
178176
)
@@ -232,7 +230,6 @@ def export_all(
232230
qembedding=None,
233231
qembedding_group_size=None,
234232
use_aoti_packed_int4=False,
235-
use_aoti_matvec=False,
236233
backend="xnnpack",
237234
):
238235
"""Export all three model components with per-component quantization."""
@@ -297,7 +294,6 @@ def export_all(
297294
qembedding,
298295
qembedding_group_size,
299296
use_aoti_packed_int4=use_aoti_packed_int4,
300-
use_aoti_matvec=use_aoti_matvec,
301297
device=device,
302298
)
303299
if packed_linear_count:
@@ -331,7 +327,6 @@ def export_streaming(
331327
qembedding=None,
332328
qembedding_group_size=None,
333329
use_aoti_packed_int4=False,
334-
use_aoti_matvec=False,
335330
backend="xnnpack",
336331
):
337332
"""Export streaming model components with per-component quantization."""
@@ -392,7 +387,6 @@ def export_streaming(
392387
qembedding,
393388
qembedding_group_size,
394389
use_aoti_packed_int4=use_aoti_packed_int4,
395-
use_aoti_matvec=use_aoti_matvec,
396390
device=device,
397391
)
398392
if packed_linear_count:
@@ -583,15 +577,11 @@ def _validate_rocm_args(parser, args):
583577
"tile_packed_to_4d requires a CUDA-only int4 fallback; "
584578
"omit the packing format for ROCm"
585579
)
586-
if args.rocm_packed_matvec and args.qlinear != "4w":
587-
parser.error("--rocm-packed-matvec requires --qlinear=4w")
588580

589581

590582
def _validate_export_args(parser, args, backend_for_export):
591583
if args.backend == "rocm":
592584
_validate_rocm_args(parser, args)
593-
elif args.rocm_packed_matvec:
594-
parser.error("--rocm-packed-matvec requires --backend=rocm")
595585

596586
if args.qlinear == "fpa4w" and backend_for_export != "metal":
597587
parser.error("--qlinear=fpa4w can only be used with --backend=metal")
@@ -708,11 +698,6 @@ def main():
708698
"typically 8192). Smaller values reduce memory and improve decode speed "
709699
"but limit how far back the decoder can attend. Only used with --streaming.",
710700
)
711-
parser.add_argument(
712-
"--rocm-packed-matvec",
713-
action="store_true",
714-
help="Use the experimental fixed-shape packed INT4 decoder matvec on ROCm.",
715-
)
716701
parser.add_argument(
717702
"--dtype",
718703
default="fp32",
@@ -771,7 +756,6 @@ def main():
771756
"qembedding": args.qembedding,
772757
"qembedding_group_size": args.qembedding_group_size,
773758
"use_aoti_packed_int4": args.backend == "rocm",
774-
"use_aoti_matvec": args.rocm_packed_matvec,
775759
"backend": backend_for_export,
776760
}
777761
if args.streaming:

‎examples/models/voxtral_realtime/run_rocm_e2e.sh‎

Lines changed: 0 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@
99
# SKIP_EXPORT=1 Use existing model.pte and aoti_cuda_blob.ptd files.
1010
# DEVICE_INDEX Visible GPU index (default: 0).
1111
# SLIDING_WINDOW Decoder window (default: 2048).
12-
# ROCM_PACKED_MATVEC=1 Use the experimental fixed-shape decoder matvec.
1312
# OFFLINE_MAX_NEW_TOKENS Offline token limit (default: 500).
1413
# VOXTRAL_PYTHON Python executable (default: python).
1514
# ROCM_PATH ROCm installation (default: /opt/rocm).
@@ -33,19 +32,13 @@ EXECUTION_MODE="${4:-streaming}"
3332
OUTPUT_ROOT="${5:-$PWD/voxtral_rt_rocm}"
3433
DEVICE_INDEX="${DEVICE_INDEX:-0}"
3534
SLIDING_WINDOW="${SLIDING_WINDOW:-2048}"
36-
ROCM_PACKED_MATVEC="${ROCM_PACKED_MATVEC:-0}"
3735
OFFLINE_MAX_NEW_TOKENS="${OFFLINE_MAX_NEW_TOKENS:-500}"
3836
VOXTRAL_PYTHON="${VOXTRAL_PYTHON:-python}"
3937
ROCM_ROOT="${ROCM_PATH:-/opt/rocm}"
4038

4139
export HIP_VISIBLE_DEVICES="$DEVICE_INDEX"
4240
export CUDA_VISIBLE_DEVICES="$DEVICE_INDEX"
4341

44-
if [[ "$ROCM_PACKED_MATVEC" != "0" && "$ROCM_PACKED_MATVEC" != "1" ]]; then
45-
echo "ERROR: ROCM_PACKED_MATVEC must be 0 or 1" >&2
46-
exit 1
47-
fi
48-
4942
case "$PRECISION_MODE" in
5043
bf16) PRECISIONS=(bf16) ;;
5144
w4-bf16) PRECISIONS=(w4-bf16) ;;
@@ -158,10 +151,6 @@ for precision in "${PRECISIONS[@]}"; do
158151
if [[ "$execution" == "streaming" ]]; then
159152
export_args+=(--streaming --sliding-window "$SLIDING_WINDOW")
160153
fi
161-
if [[ "$precision" == "w4-bf16" && "$ROCM_PACKED_MATVEC" == "1" ]]; then
162-
export_args+=(--rocm-packed-matvec)
163-
fi
164-
165154
export_elapsed_ms=-1
166155
if [[ "${SKIP_EXPORT:-0}" != "1" ]]; then
167156
export_start_ms="$(monotonic_ms)"

0 commit comments

Comments
 (0)