Move the ROCm devcontainer to ROCm 7.14 / py3.14 / torch 2.12 with AITER 0.1.16 by demandal25 · Pull Request #294 · AMD-Ecosystem/flashinfer

demandal25 · 2026-08-24T22:00:01Z

Summary

Moves the ROCm development container to ROCm 7.14 / Ubuntu 26.04 / Python 3.14 / PyTorch 2.12, drops micromamba for a plain venv, and moves AITER from 0.1.10 to 0.1.16.post3. AITER changed its C++ ABI and API between those releases, so the shims, the vendored argument struct, and the JIT glue all had to move with it.

The upgrade also exposed a soft-cap defect in AITER's mha_varlen_fwd. That is not fixed here — it is fenced off so the default code path cannot return silently wrong numbers, and documented with the evidence.

What changed

Container

.devcontainer/rocm/Dockerfile — new base rocm/pytorch:rocm7.14_ubuntu26.04_py3.14_pytorch_release_2.12.0; micromamba replaced by /opt/venv (the system Python 3.14 is PEP 668 managed); AITER installed as a pinned wheel rather than built from source. The four version args select the base image tag, so they must name a tag that exists — they are not independent knobs.

AITER 0.1.16 API/ABI

include/flashinfer/attention/aiter/mha_fwd_args.h — 10 new fields. Two pointers were inserted before sink_ptr, so every later field shifts; the struct is passed by value through a dlsym'd pointer, so a wrong layout is silent corruption with no diagnostic. Layout confirmed against the installed headers at 424 bytes and 8-of-8 offsets.
flashinfer/csrc_rocm/aiter_tensor_compat.h — new at::Tensor → aiter_tensor_t adapter, built on AITER's real aiter_tensor.h rather than a vendored copy so a layout change is a compile error.
flashinfer/csrc_rocm/activation_aiter.cu — ported to the POD API. Includes AITER's real activation.h; the old hand-written at::Tensor& declaration kept compiling and then failed at load with undefined symbol.
flashinfer/jit/aiter_source.py — build_module arguments filtered against its signature, since third_party is required at 0.1.16 and absent at 0.1.10.
flashinfer/jit/rope.py — AITER renamed the module to module_rope_2c_cached_positions_fwd.
flashinfer/compilation_context_hip.py — -DUSE_ROCM=1. From torch 2.12 the c10::hip compat namespace is gated behind it, so every *_aiter.cu shim otherwise fails on c10::hip::getCurrentHIPStream().

Correctness fixes found while validating

flashinfer/csrc_rocm/activation_aiter.cu — propagate torch's stream explicitly. The POD API launches on aiter::getCurrentHIPStream(), a thread_local defaulting to nullptr that only AITER's Python layer sets; the kernel was running on the default stream while surrounding torch ops ran elsewhere. Also added a contiguity check — AITER's torch entry point used to validate arguments, the POD path has no wrapper left to do it, and to_aiter passes strides to a kernel that indexes linearly.
flashinfer/prefill_rocm.py — _AITER_LAST_VALIDATED corrected to the full pinned string. Under PEP 440 0.1.16.post3.dev0+g... sorts below 0.1.16.post3, so naming the release claimed validation for builds never exercised.

AITER soft-cap defect

mha_varlen_fwd applies logits_soft_cap incorrectly. A non-zero cap disables every AITER assembly path, leaving the CK _logits kernel, which is off by ~0.17 against an fp32 reference while cap=0 is exact to 8e-4.

The defect is AITER's, not ours. A probe calling aiter.ops.mha.mha_varlen_fwd directly, with no FlashInfer in the call path, reproduces it on all 10 failing shapes, while mha_batch_prefill — the entry point AITER's own op_tests cover at logits_soft_cap=[0.0, 30.0] — is tanh-exact on the same inputs. Scored against both CK soft-cap formulas (tanh, and the softsign variant selected by CK_TILE_ATTENTION_LOGITS_SOFT_CAP_DEFAULT) it matches neither, so it is not a build-flag mismatch. Our mha_fwd_args is field-equivalent to AITER's own builder, including mask_type=2 (bottom-right). No fix commits exist through v0.1.21.

Scope is measured, not assumed: causal only — test_logits_cap_hip.py exercises non-causal soft-cap at 30.0/50.0 across head_dim 128/256 and passes — plus head_dim=128 and kv_len >= 512.

It affects batch prefill too, not just the single-shot API. BatchPrefillWithRaggedKVCacheWrapper always dispatches through mha_varlen_fwd and is off by 0.0595–0.1915 at logits_soft_cap=8.0 — the same magnitudes, from the same kernel. That path is untested upstream and here (the suite parametrizes logits_soft_cap=[0.0] only), which is why it went unnoticed.

Skipping the affected tests alone would have left the default path silently wrong: backend="auto" is the default at five public entry points, and Gemma-2 / Grok shapes sit squarely in the defect region. _auto_select_prefill_backend now falls back to fa2 with the same one-time warning it uses for its other AITER constraints, and explicit backend="aiter" raises rather than returning plausible-looking wrong numbers.

Paged prefill is guarded only where the defect can occur: its native-paging route uses mha_batch_prefill, which is exact at cap=8 on all 10 shapes measured, so an unconditional guard would trade correct fast kernels for slower ones. A page size outside the native set forces flat-gather, which does go through mha_varlen_fwd, and that is ruled out at plan() time. Residual gap: when native paging is claimed but the runtime probe later falls back to flat-gather, the call escapes the guard — closing it means moving the probe ahead of backend selection.

The guard is deliberately not version-gated. Auto-expiring above a known-bad AITER version would silently re-enable a wrong-answer path on a nightly bump; _AITER_SOFTCAP_DEFECT_THROUGH records the bound for a human to widen after re-measuring.

Follow-ups, deliberately not in this PR

Routing soft-cap to mha_batch_prefill instead of fa2 — measured correct, but needs a page table synthesised from contiguous KV and inherits AITER's page-size probing. That would recover the performance the fallback gives up.
Closing the paged flat-gather residual gap described above.
docker/Dockerfile.rocm_ci stays on ROCm 7.1.1 / AITER 0.1.10. Its parameterised -f repo.radeon.com/rocm-rel-${ROCM_VERSION} install cannot reach 7.14 (that directory 404s), so bumping it means restructuring onto a rocm/pytorch base. Note the rope and activation changes make AITER 0.1.16 a hard floor for those two shims.
A compile-time sizeof/offsetof guard on mha_fwd_args against AITER's real header.
Reporting both defects to ROCm/aiter.

Test plan

tests/rocm_tests/test_single_prefill_kernels_hip.py, test_activation_aiter_hip.py, test_logits_cap_hip.py, test_rope_aiter_hip.py — green on gfx950 / MI350X / ROCm 7.14 / torch 2.12.0+rocm7.14.0 / AITER 0.1.16.post3.
Both new tests A/B'd against the guard. Routing: 7/7 pass, and exactly the two defect-region cases fail when _aiter_softcap_defect is neutered. Ragged numerics: 4/4 pass, 4/4 fail neutered.
.devcontainer/rocm/Dockerfile built, then pip install --no-build-isolation -ve . and import flashinfer verified inside it as the non-root devuser.
pre-commit run -a

Not run: the full tests/rocm_tests suite on this final tree — the affected files were run individually, and test_batch_prefill_kernels_hip.py got 31% of the way through cleanly before I cut it as redundant. Nothing on gfx942.

Moves the dev container to rocm/pytorch:rocm7.14_ubuntu26.04_py3.14_pytorch_release_2.12.0 and to amd-aiter 0.1.16.post3, which is the only AITER build published as a cp314 wheel. Four independent breakages had to be fixed to get there; the full "not slow" suite goes from 60 failures to 26 on gfx950. AITER's aiter::mha_fwd_args gained ten fields between 0.1.10 and 0.1.16, two of them inserted before sink_ptr so every later offset shifts. We pass that struct by value through a dlsym'd pointer, which does no type checking, so a stale mirror corrupted silently: 99.9% of elements wrong, no crash and no link error. The mirror is regenerated and checked against the installed aiter_meta headers (sizeof 424, offsets verified); mha_batch_prefill_args was checked too and had not drifted. torch 2.12 wraps the c10::hip compatibility namespace in `#ifdef USE_ROCM`, so COMMON_HIPCC_FLAGS now defines it. Without it every *_aiter.cu shim fails to compile on "no member named getCurrentHIPStream in namespace c10::hip", which reads like a missing include rather than a missing define. AITER moved its C++ API off torch types: 15 of 16 installed headers now take aiter_tensor_t, and silu_and_mul also gained a `limit` parameter. activation.h no longer pulls pybind11, so it is now included for real rather than forward-declared -- that is what turns the next signature change into a compile error instead of an undefined symbol at load. The at::Tensor -> aiter_tensor_t adapter deliberately includes AITER's own aiter_tensor.h for the same reason: layout drift should not be something we hand-maintain. AITER also split the rope modules by variant. "module_rope_pos_fwd" is simply not registered at 0.1.16; asking for it yields an empty source list and the JIT dies on `assert len(sources) > 0`. The entry point this shim calls now lives in module_rope_2c_cached_positions_fwd. rope.h itself still uses torch::Tensor, so that shim needed no port. Known remaining: 26 failures, all logits_soft_cap=8.0, on the mha_varlen_fwd group-mode path. cap=0 is exact (8e-4) while cap=8 is off by 0.166, so the cap is applied but wrongly. Not shipped as fixed. Co-Authored-By: Claude <noreply@anthropic.com>

Two fixes from reviewing the 0.1.16 migration. silu_and_mul launched on the wrong stream. AITER's old torch-typed entry point read torch's current stream itself, but the POD API launches on aiter::getCurrentHIPStream() -- a thread_local in aiter_stream.h that defaults to nullptr and is otherwise only ever set by AITER's Python layer. The shim's OptionalHIPGuardMasqueradingAsCUDA restores the device, not the stream, so after the migration the kernel ran on the default stream while the surrounding torch ops ran on whatever stream the caller had current. That is correct on the default stream and an ordering hazard anywhere else -- CUDA graphs, multi-stream serving -- and no test on the default stream can catch it. Set the stream explicitly. block_scale_size_q/kv now default to 128 rather than 0. Every one of AITER's own arg builders (mha_fwd_kernels.cu, mha_varlen_fwd_kernels.cu, asm_mha_fwd.cu, asm_mha_varlen_fwd.cu) passes 128 unconditionally, so 0 is a value the pipeline never sees from AITER itself. This did not change any measured result, but matching the reference construction is right regardless. Co-Authored-By: Claude <noreply@anthropic.com>

Review findings against the two preceding commits on this branch. The devcontainer could not build the package it ships. The micromamba layer that was removed installed cmake and scikit-build-core>=0.4.3; the replacement pip block did not, so the documented `pip install --no-build-isolation -ve .` fails on a fresh container with ModuleNotFoundError. Dropping --no-build-isolation is not a workaround -- isolation resolves `torch >= 2.7` from PyPI and pulls a non-ROCm wheel, which is what the build-system comments and the Dockerfile's HIP guard exist to prevent. These become dead once the setuptools convergence (#291) lands. silu_and_mul gained a contiguity check. AITER's old torch-typed entry point validated its arguments itself; the POD aiter_tensor_t API has no torch wrapper left, and to_aiter passes strides through faithfully to a kernel that indexes linearly. A sliced or transposed input previously aborted on TORCH_CHECK and now returned silently wrong values. Tests only ever pass freshly-allocated contiguous tensors, so nothing covered it. Also: the aiter_source docstring still described the forward-declaration strategy that activation_aiter.cu just abandoned -- it is the one place a contributor reads to learn the shim convention, so it taught the pattern that caused the undefined-symbol failure. Comment blocks in the touched files are trimmed to the conclusion per the repo comment cap; the evidence they carried (the four AITER arg builders that hard-code block_scale_size 128, the torch 2.12 USE_ROCM gating) is here instead. Verified: silu_and_mul(..., float limit = 0.0f) is AITER's declared default in activation.h, so passing 0.0f preserves prior behaviour rather than clamping.

AITER's mha_varlen_fwd applies logits_soft_cap incorrectly. A non-zero cap disables every AITER assembly path (aiter/ops/mha.py gates can_impl_fmha_v3_fwd and friends on logits_soft_cap == 0.0), leaving the CK _logits kernel, which is off by ~0.17 against an fp32 reference while cap=0 is exact to 8e-4. The defect is AITER's, not ours: a probe calling aiter.ops.mha.mha_varlen_fwd directly, with no FlashInfer in the call path, reproduces it on all 10 failing shapes, and mha_batch_prefill -- the entry point AITER's own op_tests cover at logits_soft_cap=[0.0, 30.0] -- is tanh-exact on the same inputs. Scored against both CK soft-cap formulas (tanh and the softsign variant selected by CK_TILE_ATTENTION_LOGITS_SOFT_CAP_DEFAULT) it matches neither, so it is not a build-flag mismatch. Our mha_fwd_args is field-equivalent to AITER's own builder, including mask_type=2 (bottom-right), so there is nothing to fix on the caller side. No fix commits exist through v0.1.21. Scope is narrow and measured, not assumed: causal only (test_logits_cap_hip.py exercises non-causal soft-cap at 30.0/50.0 across head_dim 128/256 and passes), head_dim=128, kv_len >= 512. Guarding only the test would have left the default path silently wrong -- backend="auto" is the default at five public entry points, and Gemma-2/Grok shapes (causal, hd128, cap 30/50, long context) sit squarely in the defect region. _auto_select_prefill_backend now falls back to fa2 with the same one-time warning it uses for its other AITER constraints, and explicit backend="aiter" warns rather than returning plausible-looking wrong numbers. The routing is tested rather than the numerics, since the failure is silent. A/B: 7/7 cases pass with the guard; neutering _aiter_softcap_defect fails exactly the two defect-region cases and nothing else. _AITER_LAST_VALIDATED also corrected to the full pinned string. Under PEP 440 "0.1.16.post3.dev0+g..." sorts BELOW "0.1.16.post3", so naming the release claimed validation for builds never exercised -- the exact silent widening of the support boundary the adjacent comment warns against.

The Dockerfile moved to ROCm 7.14 / Ubuntu 26.04 / Python 3.14 / torch 2.12 and AITER 0.1.16.post3; every doc still described 7.2 / py3.12 / torch 2.9.1 / AITER 0.1.10. The build-arg block was the most acutely wrong -- it documents `docker build` against .devcontainer/rocm/Dockerfile and every default in it had changed underneath, including two args (AITER_VERSION, AITER_INDEX) that were undocumented entirely. ROCm 7.14 needs its own wording rather than a version bump: repo.radeon.com publishes no rocm-rel-7.14/ directory, so the `pip install torch -f ...` recipe has no valid form for it and torch must come from the base image. CLAUDE.md also claimed torch installs via --index-url while README said the opposite; README was right (the radeon repo is a flat listing, not a PEP 503 index). The AITER section now carries both pins and says why they differ -- the CI image stays on 0.1.10 from pypi.amd.com, and the devcontainer takes the only cp314 wheel that exists, from the vllm-cdna nightlies, spelled out in full because pip will not select a local version from a loose specifier. The prior claim that both images install the same version is no longer true. Known Limitations gains the AITER soft-cap defect, scoped to causal since non-causal soft-capped prefill is exercised and correct. Left alone deliberately: the published rocm/flashinfer image table and the micromamba note under it (that image is still 7.2 and still micromamba-based), docker/Dockerfile.rocm_ci (still 7.1.1, and its parameterised -f install cannot reach 7.14), and the "since amd-aiter >= 0.1.10" feature floors, which document history rather than the current stack.

Found by building the devcontainer image and running it as the devuser: the README run command fails with "No CUDA GPUs are available". The old base image shipped a `render` group at the host's conventional GID, so `--group-add render` worked. The rewritten Dockerfile does `groupadd -r render` when the base has none, which takes an arbitrary free system GID -- 995 here against a host device owned by GID 109 -- and docker resolves the *name* against the image, so the container user never gets access to /dev/dri/renderD*. The Dockerfile comment already noted callers must pass the numeric GID; the docs were never updated to match. Missed until now because every test so far ran as root, which bypasses the group check entirely. Verified after the change: pip install --no-build-isolation -ve . succeeds in the image, torch reports MI350X, and import flashinfer loads.

The previous commit guarded only single prefill, which left the production path exposed: BatchPrefillWithRaggedKVCacheWrapper always dispatches through mha_varlen_fwd, so it carries the same defect, and it is what a serving stack actually calls. Measured on gfx950, ragged at logits_soft_cap=8.0 is off by 0.0595-0.1915 against an fp32 reference -- the same magnitudes single prefill produced, from the same kernel -- while cap=0 and fa2 are exact. That also corrects an earlier claim of mine that batch prefill was merely untested at cap>0. It is untested (the suite parametrizes logits_soft_cap=[0.0] only, which is why this went unnoticed) and it is broken. Paged is guarded only where the defect can actually occur. Its native-paging route calls mha_batch_prefill, which is exact at cap=8 on all 10 shapes measured, so guarding it unconditionally would trade correct fast kernels for slower ones. A page size outside the native set forces flat-gather, which does go through mha_varlen_fwd, and that is the case ruled out at plan() time. One residual gap: when native paging is claimed but the runtime probe later falls back to flat-gather, the call escapes the guard. Closing that means moving the probe ahead of backend selection, which is a larger change than this fix. The new test drives the ragged wrapper end to end rather than asserting on routing, because routing is not the property that matters. A/B: 4/4 pass with the guard, 4/4 fail without it. Two review suggestions declined. Version-gating the guard on _AITER_SOFTCAP_DEFECT_THROUGH would auto-expire it on a nightly AITER bump and silently re-enable a wrong-answer path; deoptimizing until someone re-measures is the safer failure. And the test skip is not over-broad: it drops 28 cases, 26 of which are the measured failures, not 1.

Copilot

Pull request overview

Updates the ROCm development environment and AITER integration to match a newer ROCm/PyTorch toolchain, while guarding against a known AITER soft-cap correctness defect and adding regression coverage for the new routing behavior.

Changes:

Move the ROCm devcontainer to a rocm/pytorch:rocm7.14_ubuntu26.04_py3.14_pytorch_release_2.12.0 base and switch to the base image venv for Python tooling.
Update FlashInfer’s AITER shims/adapters/JIT glue for AITER 0.1.16+ API/ABI changes, including stream propagation and argument struct updates.
Add soft-cap defect fencing (auto fallback + explicit-backend raise) and new ROCm tests/docs describing the defect and expected routing.

Reviewed changes

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

Show a summary per file

File	Description
tests/rocm_tests/test_single_prefill_kernels_hip.py	Skips known-bad AITER soft-cap region in numeric test; adds routing regression test for `backend="auto"`.
tests/rocm_tests/test_batch_prefill_kernels_hip.py	Adds ragged wrapper regression test to ensure `backend="auto"` remains numerically correct under soft-cap.
README.md	Documents ROCm 7.14/PyTorch 2.12 support and the AITER soft-cap defect; updates devcontainer build/run instructions and pinned AITER wheel guidance.
pyproject.toml	Updates torch install comment example and adds Python 3.14 classifier.
include/flashinfer/attention/aiter/mha_fwd_args.h	Updates vendored `mha_fwd_args` layout to match AITER 0.1.16+ fields/offsets.
flashinfer/prefill_rocm.py	Adds soft-cap defect detection, routes `auto` away from AITER for the defect region, and raises on explicit `backend="aiter"` in that region.
flashinfer/jit/rope.py	Updates AITER rope module name used for JIT flag extraction.
flashinfer/jit/aiter_source.py	Filters `build_module` kwargs against the installed AITER signature (for 0.1.10 vs 0.1.16+).
flashinfer/hip_utils.py	Extends ROCm-version-to-arch compatibility grouping to include 7.14.
flashinfer/csrc_rocm/aiter_tensor_compat.h	Introduces an `at::Tensor` → `aiter_tensor_t` adapter for the POD API (AITER 0.1.16+).
flashinfer/csrc_rocm/activation_aiter.cu	Ports SiLU+mul shim to AITER POD API, adds contiguity checks, and explicitly sets the HIP stream.
flashinfer/compilation_context_hip.py	Adds `-DUSE_ROCM=1` to satisfy torch 2.12 HIP compat gating.
CLAUDE.md	Updates Torch install guidance and AITER pinning guidance for the new environment.
amd-flashinfer-jit-cache/pyproject.toml	Updates torch install comment example and adds Python 3.14 classifier.
.devcontainer/rocm/Dockerfile	Switches to rocm/pytorch base + venv workflow, pins AITER wheel install, and updates user/group setup.
.claude/skills/debug-rocm-crash/SKILL.md	Updates AITER install troubleshooting guidance to reference the README pinned-wheel section.

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

+# Silence the warning about out-of-range UID/GID, then create the user.
 RUN sed -i 's/^\(UID_MAX\s*\).*$/\11000000000/' /etc/login.defs && \
    sed -i 's/^\(GID_MAX\s*\).*$/\11000000000/' /etc/login.defs && \
    groupadd --gid $USER_GID $USERNAME && \
    useradd --uid $USER_UID --gid $USER_GID -m $USERNAME && \


+pip install amd_aiter==0.1.16.post3.dev0+g620287969.d20260725 \
+  --extra-index-url https://rocm.frameworks-nightlies.amd.com/whl-multi-arch/vllm-cdna/


+        logits_soft_cap=soft_cap,
+        kv_len=kv_len,
+    )
+    assert chosen == ("aiter" if expect_aiter else "fa2")


demandal25 and others added 7 commits August 22, 2026 11:42

Copilot AI lite review requested due to automatic review settings August 24, 2026 22:00

Copilot started reviewing on behalf of demandal25 August 24, 2026 22:00 View session

Copilot AI reviewed Aug 24, 2026

View reviewed changes

Navigation Menu

Uh oh!

Move the ROCm devcontainer to ROCm 7.14 / py3.14 / torch 2.12 with AITER 0.1.16 - #294

Move the ROCm devcontainer to ROCm 7.14 / py3.14 / torch 2.12 with AITER 0.1.16#294
demandal25 wants to merge 7 commits into
amd-integrationfrom
devcontainer-rocm714

demandal25 commented Aug 24, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		pip install amd_aiter==0.1.16.post3.dev0+g620287969.d20260725 \
		--extra-index-url https://rocm.frameworks-nightlies.amd.com/whl-multi-arch/vllm-cdna/

Uh oh!

Conversation

demandal25 commented Aug 24, 2026

Summary

What changed

Container

AITER 0.1.16 API/ABI

Correctness fixes found while validating

AITER soft-cap defect

Follow-ups, deliberately not in this PR

Test plan

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants