Enable fast ROCm Voxtral decode

[digantdesai]

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.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22005

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 2 Unclassified Failures

As of commit f7a058c with merge base 9a2d135 ():

UNCLASSIFIED FAILURES - DrCI could not classify the following jobs because the workflow did not run on the merge base. The failures may be pre-existing on trunk or introduced by this PR:

• Test CUDA Builds / test-model-cuda-e2e-Qwen3.5-35B-A3B-HQQ-INT4-quantized-int4-tile-packed / linux-job (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  RuntimeError: Command bash /__w/_temp/exec_script failed with exit code 1
• Test CUDA Builds / unittest-cuda / linux-job (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  test_device_planned_method_allocates_on_the_device

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[Copilot]

Pull request overview

Enables automatic packed INT4 matvec decoding for ROCm W4 Voxtral exports while keeping ROCm on standard SDPA after split-K.

Changes:

• Removes the experimental matvec opt-in.
• Updates ROCm export logic, documentation, tests, and backend behavior.
• Disables the gfx1100 ROCm workflow job.

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Review findings
examples/models/voxtral_realtime/run_rocm_e2e.sh	No final comments supplied.
examples/models/voxtral_realtime/README.md	No final comments supplied.
examples/models/voxtral_realtime/export_voxtral_rt.py	Critical (1 vote): Add a quality opt-out or validation gate before enabling matvec for all ROCm W4 exports.
backends/cuda/triton/replacement_pass.py	No final comments supplied.
backends/cuda/tests/test_sdpa_splitk_replacement.py	Moderate (3 votes): Verify Replaced 1 nodes for ROCm standard SDPA, including the assertion around line 165.
backends/cuda/aoti_packed_int4_tensor.py	No final comments supplied.
.github/workflows/rocm.yml	Moderate (3 votes): Retain equivalent gfx1100 coverage or document and track the missing architecture validation.

Suppressed comments (1)

backends/cuda/tests/test_sdpa_splitk_replacement.py:170

• As in the threshold case above, the ROCm branch only asserts absence of the split-K log and never confirms that this eligible shape uses standard triton.sdpa. The test would still pass if the replacement pass stopped replacing this attention node altogether. Add a replacement-count assertion for the expected == 0 case, matching test_below_threshold_uses_standard.

        expected = 0 if torch.version.hip is not None else 1
        self.assertEqual(
            len(splitk),
            expected,
            f"Expected {expected} split-K selections. Log: {msgs}",
        )

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

Pull request overview

Copilot reviewed 7 out of 7 changed files in this pull request and generated no new comments.

Suppressed comments (4)

.github/workflows/rocm.yml:230

• Because of false &&, this job can never run, even when the decision job sets run-gfx1100 for workflow_dispatch and direct Voxtral changes. That removes the gfx1100 canary while this PR changes the ROCm export path, and the documentation still says gfx1100 is covered; please re-enable the condition before relying on this job, or update the coverage documentation and track the disablement separately.

      false && needs.voxtral-run-decision.outputs.run-gfx1100 == 'true' &&

backends/cuda/tests/test_sdpa_splitk_replacement.py:147

• On ROCm, expected is 0, so this assertion only proves that no split-K log was emitted; it still passes if the replacement pass failed and left the original SDPA node untouched. The test docstring promises standard SDPA on ROCm, so also assert that one node was replaced (and that the resulting path is standard).

        expected = 0 if torch.version.hip is not None else 1
        self.assertEqual(
            len(splitk),
            expected,
            f"Expected {expected} split-K selections. Log: {msgs}",

backends/cuda/tests/test_sdpa_splitk_replacement.py:169

• The ROCm branch again sets expected to 0 and only checks for absence of split-K, so this large-cache case also passes when no SDPA replacement occurred at all. Add the same replacement assertion here so the test actually covers the promised standard-SDPA path.

        expected = 0 if torch.version.hip is not None else 1
        self.assertEqual(
            len(splitk),
            expected,
            f"Expected {expected} split-K selections. Log: {msgs}",

examples/models/voxtral_realtime/export_voxtral_rt.py:157

• This makes the previously opt-in BF16-rounded matvec path the default for every ROCm W4 decoder. The removed TODO/README documented greedy-output drift for this fixed-shape path, while the new end-to-end check only verifies that one token appears, so a transcript regression can still pass silently. Please keep this opt-in until the SDPA change is validated against a full expected transcript/quality check, or add that regression check before changing the default.

    use_packed_matvec = use_aoti_packed_int4 and qlinear == "4w"

[Copilot]

Pull request overview

Copilot reviewed 7 out of 7 changed files in this pull request and generated no new comments.

Suppressed comments (2)

.github/workflows/rocm.yml:230

• false && unconditionally prevents this job from running, including on workflow_dispatch, so the gfx1100 end-to-end path no longer receives any validation for this ROCm W4 change. If runner teardown instability requires a temporary pause, please use an explicitly temporary gate and restore the actual decision condition before landing rather than committing an unconditional skip.

      false && needs.voxtral-run-decision.outputs.run-gfx1100 == 'true' &&

examples/models/voxtral_realtime/export_voxtral_rt.py:170

• This makes the previously opt-in int4_matvec_bf16 path the default for every ROCm W4 decoder. The removed README/TODO documented greedy transcript drift, and the kernel still uses BF16-rounded weights and a different reduction order, so new ROCm W4 exports can silently change transcription quality. Please keep this opt-in until parity/quality is demonstrated, or add an explicit accuracy gate/fallback before enabling it globally.

        if use_packed_matvec:
            packed_linear_count = _pack_aoti_int4_weights(
                text_decoder,
                use_matvec=True,
            )

[shoumikhin]

The speedup here is real and the description undersells it. Comparing the gfx950
job log on this branch against the same job on main: decoder input goes from
[1, 4, 3072] to [1, 1, 3072] and RTF goes from 0.459 to 0.283, so about 1.6x
faster decode. That is the headline and it is not in the description.

My concern is that four separate changes are bundled here and three of them are
not mentioned in the description. Since the repo squash-merges on the PR body,
that context is lost at merge.

1. The split-K guard is not a ROCm issue.

replacement_pass.py:107 gates on torch.version.hip is None, and the README
attributes the non-finite logits to ROCm. The non-finite logits are fp32 exponent
range, not a vendor problem. I reproduced them on an H100 where
torch.version.hip is None, with the same shape, and #22133 fixes them in the
kernel for both vendors.

The guard is also narrower than it looks. sdpa() at sdpa.py:1012 dispatches
L_q == 1 straight into _launch_decode_splitk with no hip check, and these
model transforms reach the kernel without going through the replacement pass at
all:

• examples/models/qwen3_5_moe/model.py:298
• examples/models/gemma4_31b/cuda_source_transformations.py:176
• examples/models/muse-glimmer/source_transformations/cuda.py:234 and :540
• examples/models/voxtral_tts/model.py:383

So on ROCm it stops one of five paths. If it lands ahead of #22133, could the
guard go in sdpa_decode_splitk itself, say plainly that it works around the
fixed-phi overflow, and name #22133 as the removal condition? Otherwise it reads
as a permanent vendor exclusion for a bug that is about to not exist.

Related: _pick_sdpa_kernel's docstring at :91-101 still describes selection as
purely an L_kv tradeoff and does not mention that a vendor is excluded.

2. The greedy-output drift.

This removes the TODO in export_voxtral_rt.py that said to resolve the
fixed-shape greedy-output drift before enabling it by default, and the README
paragraph asking users to verify transcript quality on the target GPU, and then
enables it by default with no opt-out.

The gfx950 logs suggest the drift is still there. Same clip, same model, same
runner, main versus this branch. The transcripts agree word for word up to
"...occur most readily to the mind" and then diverge:

main:   the mind he has grave doubts whether sir frederick leighton's work is really greek after all and can
branch: the mind. He has grave doubts whether Sir Frederick Layton's work is really Greek after all, and can

Some of that is better (punctuation, casing) and some is worse ("Leighton" to
"Layton"). The CI check is grep -qi 'Quilter'
(.ci/scripts/test-rocm-voxtral.sh:133), so it passes either way and cannot see
this.

Was the drift resolved, or is it being accepted? Either answer is fine, it just
needs to be the stated one, ideally with a transcript comparison or a WER number
rather than a one-word grep. My guess is that the drift and the split-K
non-finite logits are the same defect: fixed-shape export is what makes L_q a
concrete 1, which is what selects split-K, which is what hits the overflow. If
that is right, #22133 resolves item 2 and item 1 together.

3. Disabling the gfx1100 job.

rocm.yml:230 prepends false && to the job condition. A few things about this:

• It is not in the PR description, and it arrived in a follow-up commit.
• There is no precedent for the idiom. false && and if: false appear zero
  times across the 71 workflow files on main; this adds the only one. The
  existing way to park CI work is a commented-out matrix entry with a TODO:
  naming the cause, as in cuda-windows.yml:78-82.
• The voxtral-run-decision selector fires RUN_GFX1100=true for exactly the
  files this PR touches, and line 230 overrides it. gfx1100 is wave32 and gfx950
  is wave64, so this is the PR that changes the ROCm decode default and also the
  PR that turns off the only wave32 end-to-end check of it.
• backends/cuda/rocm.md:10 and
  examples/models/voxtral_realtime/README.md:164 both still say CI exercises
  gfx950 and gfx1100.

Could this go in its own PR with a tracking issue and an owner? "Temporarily" plus
false && tends to be permanent, since nothing surfaces it again.

4. Flag removal.

--rocm-packed-matvec and the ROCM_PACKED_MATVEC env var are removed with no
compatibility path, along with the use_aoti_matvec keyword on export_all and
export_streaming. The flag fails loudly through argparse, which is fine, but the
env var fails silently: a script setting ROCM_PACKED_MATVEC=0 to stay on the old
path now gets the new one instead. They were experimental and recently added, so
this may well be acceptable, but a one-release no-op alias would be cheap.

Small one: the README diff removes the --rocm-packed-matvec table row and
the blank line after it, so **Notes:** no longer starts its own paragraph and
gets absorbed into the block above it when rendered.

On CI: the unittest-cuda failure on this branch looks unrelated to the diff.
It fails on PybindingsTest.test_device_planned_method_allocates_on_the_device,
which main added in #22058 after this branched, so the workflow asks for a test
the branch does not have. Main is green and at least one unrelated open PR shows
the same failure. A rebase should clear it. Worth doing before merge though,
because that step exits before pytest runs, so neither arm of your new test
branch is actually executing in CI right now.

[digantdesai]

The speedup here is real and the description undersells it. Comparing the gfx950 job log on this branch against the same job on main: decoder input goes from [1, 4, 3072] to [1, 1, 3072] and RTF goes from 0.459 to 0.283, so about 1.6x faster decode. That is the headline and it is not in the description.

My concern is that four separate changes are bundled here and three of them are not mentioned in the description. Since the repo squash-merges on the PR body, that context is lost at merge.

1. The split-K guard is not a ROCm issue.

replacement_pass.py:107 gates on torch.version.hip is None, and the README attributes the non-finite logits to ROCm. The non-finite logits are fp32 exponent range, not a vendor problem. I reproduced them on an H100 where torch.version.hip is None, with the same shape, and #22133 fixes them in the kernel for both vendors.

The guard is also narrower than it looks. sdpa() at sdpa.py:1012 dispatches L_q == 1 straight into _launch_decode_splitk with no hip check, and these model transforms reach the kernel without going through the replacement pass at all:

• examples/models/qwen3_5_moe/model.py:298
• examples/models/gemma4_31b/cuda_source_transformations.py:176
• examples/models/muse-glimmer/source_transformations/cuda.py:234 and :540
• examples/models/voxtral_tts/model.py:383

So on ROCm it stops one of five paths. If it lands ahead of #22133, could the guard go in sdpa_decode_splitk itself, say plainly that it works around the fixed-phi overflow, and name #22133 as the removal condition? Otherwise it reads as a permanent vendor exclusion for a bug that is about to not exist.

Related: _pick_sdpa_kernel's docstring at :91-101 still describes selection as purely an L_kv tradeoff and does not mention that a vendor is excluded.

2. The greedy-output drift.

This removes the TODO in export_voxtral_rt.py that said to resolve the fixed-shape greedy-output drift before enabling it by default, and the README paragraph asking users to verify transcript quality on the target GPU, and then enables it by default with no opt-out.

The gfx950 logs suggest the drift is still there. Same clip, same model, same runner, main versus this branch. The transcripts agree word for word up to "...occur most readily to the mind" and then diverge:

main:   the mind he has grave doubts whether sir frederick leighton's work is really greek after all and can
branch: the mind. He has grave doubts whether Sir Frederick Layton's work is really Greek after all, and can

Some of that is better (punctuation, casing) and some is worse ("Leighton" to "Layton"). The CI check is grep -qi 'Quilter' (.ci/scripts/test-rocm-voxtral.sh:133), so it passes either way and cannot see this.

Was the drift resolved, or is it being accepted? Either answer is fine, it just needs to be the stated one, ideally with a transcript comparison or a WER number rather than a one-word grep. My guess is that the drift and the split-K non-finite logits are the same defect: fixed-shape export is what makes L_q a concrete 1, which is what selects split-K, which is what hits the overflow. If that is right, #22133 resolves item 2 and item 1 together.

3. Disabling the gfx1100 job.

rocm.yml:230 prepends false && to the job condition. A few things about this:

• It is not in the PR description, and it arrived in a follow-up commit.
• There is no precedent for the idiom. false && and if: false appear zero
  times across the 71 workflow files on main; this adds the only one. The
  existing way to park CI work is a commented-out matrix entry with a TODO:
  naming the cause, as in cuda-windows.yml:78-82.
• The voxtral-run-decision selector fires RUN_GFX1100=true for exactly the
  files this PR touches, and line 230 overrides it. gfx1100 is wave32 and gfx950
  is wave64, so this is the PR that changes the ROCm decode default and also the
  PR that turns off the only wave32 end-to-end check of it.
• backends/cuda/rocm.md:10 and
  examples/models/voxtral_realtime/README.md:164 both still say CI exercises
  gfx950 and gfx1100.

Could this go in its own PR with a tracking issue and an owner? "Temporarily" plus false && tends to be permanent, since nothing surfaces it again.

4. Flag removal.

--rocm-packed-matvec and the ROCM_PACKED_MATVEC env var are removed with no compatibility path, along with the use_aoti_matvec keyword on export_all and export_streaming. The flag fails loudly through argparse, which is fine, but the env var fails silently: a script setting ROCM_PACKED_MATVEC=0 to stay on the old path now gets the new one instead. They were experimental and recently added, so this may well be acceptable, but a one-release no-op alias would be cheap.

Small one: the README diff removes the --rocm-packed-matvec table row and the blank line after it, so **Notes:** no longer starts its own paragraph and gets absorbed into the block above it when rendered.

On CI: the unittest-cuda failure on this branch looks unrelated to the diff. It fails on PybindingsTest.test_device_planned_method_allocates_on_the_device, which main added in #22058 after this branched, so the workflow asks for a test the branch does not have. Main is green and at least one unrelated open PR shows the same failure. A rebase should clear it. Worth doing before merge though, because that step exits before pytest runs, so neither arm of your new test branch is actually executing in CI right now.

1 - yes, this is something needs fixing on cuda side first see #22133
2 - yeah, no need to have other path, the previous one was buggy
3 - gfx1100 will be enabled in the future, ci infra issue
4 - for the flag, we do not want any alternate path

Thanks for the feedback :)