[merceod]

On AMD GPUs the norm, RoPE and sampler resources now run AMD's FlashInfer fork (amd-flashinfer 0.6.18+amd.1). The kernel is picked per platform through PlatformDispatch from #603, and NVIDIA keeps calling FlashInfer as it does today.

• RMSNorm. mstar::rms_norm's CUDA kernel dispatches to FlashInfer or the fork. Fused add + norm gets its own op with a fake, mstar::fused_add_rms_norm, so dynamo keeps it in the graph on ROCm.
• RoPE. On ROCm the position manager builds an fp32 cos/sin cache (angles computed in fp64) and runs the fork's cached kernel. A layer that passes RoPE arguments the cache wasn't built from gets the on-the-fly kernel, and a rope_table_dtype table ([Model] Add Fish Audio S2-Pro: text-to-speech with voice cloning and exact streaming #492) takes precedence over the cache on ROCm as it does on CUDA.
• Sampling. The fork's top-p and top-k/top-p samplers take a device seed and offset per row, so captured sampling works like it does on NVIDIA. Main's HF-order filter and min-tokens floor ([MiniCPM Prereq] Block-causal ragged attention + Sampler windowed penalty, min-tokens floor, top-p-first order #381) draw and renormalise through the same dispatch, since the fork's ops raise when dynamo traces into them. The Triton prep kernel's autotune drops configs whose block can't launch with 64-wide wavefronts (32 warps x 64 > 1024 threads).

Why not AITER's norm kernels: they're 20-40% faster per call, but they round bf16 stores toward zero. In a drift test that shrank the residual by 9.1% over 64 layers, against 1.05% (and no bias) for the fork.

NVIDIA: no behaviour change. Every platform but ROCm gets today's FlashInfer calls (the sampler dispatch makes FlashInfer its default), forward_residual keeps its direct FlashInfer call and its graph break (moving NVIDIA onto the op needs a worker_phases A/B first), no RoPE cache gets built, and the autotune prune keeps every config at warp size 32.

scripts/ci/select_cpu_tests.py lists test/rocm with the modules CPU Core doesn't run (they need an AMD GPU), since main's build check wants every test module accounted for.

PlatformDispatch.reset_all() lets the modular tests' platform fixture drop every dispatch's cached kernel between tests, the same way it clears the cached platform, so a test that patches a kernel module gets its own.

Tests

What	Result
ruff check .	clean
test/modular/test_rocm_op_dispatch.py (new), dispatch per platform on CPU	passes
norm, rope, sampler and dispatch tests (test_rocm_op_dispatch.py, test_rms_norm_op.py, test_fused_rmsnorm.py, test_position_autocast.py, test_platform_dispatch.py, test_sampler_*, test_sampling_*, test_qwen3_5_config.py, test_resource_runner.py)	pass in the run of the whole integration branch on MI300X below
test/rocm/test_norm_kernels.py, test_rope.py, test_sampler.py, test_triton_kernels.py on MI300X	226 passed. Parity, chi-square for top-k/top-p, greedy ties, same-seed repeats, capture and replay, a compiled caller
test/modular, test/distributed and test/sampling_test on an MI300X node with GPUs hidden, on the integration branch that carries this whole stack, rebased on main	only the 7 environment failures main has there too (rsync, ftfy)
test/rust there	359 passed
test/rocm on MI300X, same branch	674 passed, plus the fused all-reduce harness at 2, 4 and 8 GPUs
CI's CPU Core jobs (since #581 all of test/modular, test/distributed and the model test dirs, plus the TP lockstep sim; torch 2.9.1 CPU, merged with main)	green on every cut so far (last run 2026-10-21: 5877 passed, 749 skipped), TP sim passes. One earlier run timed out once in test_eval_framework.py::test_sigterm_to_the_runner_stops_the_managed_server, a subprocess timing test that passed on the rerun and on the other PRs
Rebase onto main 54f29da (2026-10-26, the whole stack)	CPU on a devel node, GPUs hidden: test/modular only fails the two test_dump_stacks environment cases main fails too (7174 passed before the last three test fixes, targeted rerun on the tip 398 passed), other CPU suites 248 passed, TP lockstep sim 7 passed, test/rust 368 passed. One MI300X: test/rocm 675 passed, fused-MoE and sampler GPU tests 228 passed, served smoke rows for Qwen3.5 0.8B, Kokoro, Whisper large-v3 and Cosmos3 DROID match the earlier runs

Numbers

End to end against the same tree without this PR, one MI300X (VM slice, indicative), benchmark.runner closed loop at c=16, 96 requests, two runs per arm.

	before	after
Orpheus (configs/orpheus_colocated.yaml with flashinfer_backend: fa2), audio s/s	21.4-21.7	22.5-22.9 (+4-6%)
Orpheus, RTF mean	0.69-0.70	0.65-0.67
Qwen3-TTS (configs/qwen3tts.yaml), audio s/s	60.8	61.5-62.2 (+1-2%)

Whisper on the 5 clips per model per run: WER 0 on nearly every clip in both arms, at most one wrong word on one clip per run. Exclusive-node tables follow.

Not in this PR

• Most models leave rope_theta out of PositionConfig, so they get the on-the-fly kernel instead of the cache.
• The fork's fused add + norm is 2.5x slower than AITER's Triton one at 512+ rows (Qwen3.5 prefill). A faster round-to-nearest kernel for those sizes would help.

Checks on this rebase

The stack now sits on main f7cb5ca (2026-10-29), every branch cut from one integration tree.

Check	Result
CPU on a devel node with the GPUs hidden: ruff, test/ci, the TP lockstep sim, test/rust, CPU Core	ruff clean, sim 14 passed, test/rust 436 passed, CPU Core fails nothing main doesn't (main's own 34: test_dump_stacks x2, loop_output_ownership x32)
test/rocm on one MI325X (a VM slice, torch 2.12.1+rocm7.2)	715 passed, 9 skipped, 0 failed
MoE and sampler tests there	MoE EP 53 passed, sampler 145 passed. The 75 test_*moe* failures are main's GLM-5.2 kernel tests, the same set on the previous head
Served rows there through the correctness runner	all PASS: Qwen3.5 0.8B and 9B text (no divergence away from near ties, 4/4 needles), Whisper large-v3 (WER 1.64%, 99/100 transcripts equal to the reference), Kokoro (WER delta 0), Qwen3-Omni colocated (50/50, spoken WER 0), pi0.5 and Cosmos3 DROID (two instances bit-same)
Qwen3.5 9B TP2 on two GPUs there	PASS (21/32 exact, 11 near ties, 0 unexplained), the fused all-reduce on both ranks

Nothing fails on this tree that the previous head doesn't: the failing-set diff against it is empty for test/rocm, MoE, the sampler and the rows.