feat(ci): cross-platform Docker matrix and hardware CI for CUDA/ROCm (closes #173)#206
Open
haoruilee wants to merge 6 commits into
Open
feat(ci): cross-platform Docker matrix and hardware CI for CUDA/ROCm (closes #173)#206haoruilee wants to merge 6 commits into
haoruilee wants to merge 6 commits into
Conversation
Closes RL-Align#173. Add reproducible CUDA and ROCm Docker images, portable hardware CI runners, GitHub/GitLab workflow integration, and source-build contributor documentation. Harden CUDA extension builds and ROCm HIP compatibility based on hardware validation across NVIDIA and AMD devices.
# Conflicts: # csrc/cuda/fused_linear_logp_sm90.cu # csrc/utils/tma_utils.cuh # rl_engine/kernels/ops/cuda/loss/logp.py
…d naming
runpodctl 2.6.x reports the pod SSH endpoint in a top-level `.ssh {ip, port}`
block; the previous grep heuristic never matched it, so the orchestrator
waited until retry exhaustion on a healthy pod (reproduced on a live H100).
Resolve the endpoint with jq (`.ssh` first, then the older privatePort==22
port-mapping layout, regex as last resort) and resolve the pod id the same
way. Also restore the explicit POD_* / RunPod naming used on main instead of
provider-neutral wording that no longer matched the implementation.
…disk The two RunPod jobs duplicated the runpodctl install/auth/SSH-key steps; move them to .github/actions/setup-runpod (checked out from the base SHA, so it stays trusted under pull_request_target). Image builds now free ~25GB of preinstalled toolchains first: ubuntu-latest has ~14GB free while the CUDA/ROCm devel bases are tens of GB, which made PR builds flaky.
The 1e-2 atol was masking an environment bug, not Triton drift: cloud GPU images that set NVIDIA_TF32_OVERRIDE=1 force cuBLAS into TF32 regardless of torch.backends settings, degrading the native fp32 oracle itself to ~1.4e-2 error (Triton with input_precision="ieee" stays at ~2e-5 vs fp64). Pin the override off in conftest.py before the first cuBLAS call and restore the original 1e-3 tolerance; measured worst-case drift is ~2e-5 over 200 seeds. Also probe CPU BLAS fp32 batch invariance in test_lm_head.py and skip the bitwise Axis-A cases where the BLAS genuinely lacks the property (MKL on AVX2 picks M-dependent K-blocking even single-threaded), so hardware CI does not fail machine-dependently.
haoruilee
requested review from
Flink-ddd,
KJLdefeated,
bitborne and
inaniloquentee
as code owners
📝 WalkthroughWalkthroughThis PR adds CUDA/ROCm Docker images, RunPod and self-hosted GPU CI orchestration, SM90 kernel fixes, HIP compatibility updates, dispatch validation, packaging changes, and expanded tests/docs. ChangesCI/CD Docker images and GPU CI orchestration
SM90 kernel fixes and HIP portability
Estimated code review effort: 4 (Complex) | ~80 minutes Sequence Diagram(s)sequenceDiagram
participant Workflow as gpu-ci.yml job
participant SetupAction as setup-runpod action
participant RunGpuCi as ci/run_gpu_ci.sh
participant RunPod as RunPod pod
participant RunCudaTests as ci/run_cuda_tests.sh
Workflow->>SetupAction: configure runpodctl and SSH key
Workflow->>RunGpuCi: start RunPod orchestration
RunGpuCi->>RunPod: create pod and resolve SSH endpoint
RunGpuCi->>RunPod: scp run_cuda_tests.sh
RunGpuCi->>RunCudaTests: ssh bash /tmp/run_cuda_tests.sh
RunCudaTests->>RunCudaTests: install deps and run smoke tests
RunCudaTests->>RunCudaTests: run pytest
RunCudaTests-->>RunGpuCi: exit code
RunGpuCi->>RunPod: release pod
RunGpuCi-->>Workflow: final status
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Collaborator
|
Hi @haoruilee , Thank you for your excellent work. Could you please add some test results in Markdown format? We currently don't have multi-GPU tests for ROCm, only multi-GPU tests for CUDA. So this is our next step work. Thanks again. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
.github/workflows/build-ci-image.yml (1)
86-107: 🔒 Security & Privacy | 🟡 Minor | ⚡ Quick winSame missing
persist-credentials: falseon the push-job checkout.Lower risk here since this job only runs on trusted
pushtomain, but still worth hardening per the same zizmor finding, especially since this job also holdspackages: write.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In @.github/workflows/build-ci-image.yml around lines 86 - 107, The checkout step in the build-ci-image workflow is missing the same hardening flag as the other job. Update the actions/checkout@v4 step in this push-job to set persist-credentials to false, matching the existing security pattern and keeping GitHub credentials out of the local checkout before the docker/login-action and image build/push steps run.Source: Linters/SAST tools
🧹 Nitpick comments (6)
docker/Dockerfile.cuda (1)
26-37: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winPin dev-tooling versions for reproducibility.
black,isort,mypy,packaging,pre-commit,psutil,pytest,ruffare installed unpinned, so the image drifts each rebuild even though the PR's goal is a "reproducible" CI image (base image is digest-pinned, but tooling isn't).♻️ Example fix
- && python -m pip install \ - black \ - isort \ - mypy \ - packaging \ - pre-commit \ - psutil \ - pytest \ - ruff + && python -m pip install \ + black==24.10.0 \ + isort==5.13.2 \ + mypy==1.13.0 \ + packaging==24.2 \ + pre-commit==4.0.1 \ + psutil==6.1.0 \ + pytest==8.3.3 \ + ruff==0.7.4🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docker/Dockerfile.cuda` around lines 26 - 37, The dev-tooling packages installed in the Dockerfile.cuda build are unpinned, so the image can drift between rebuilds. Update the pip install step that installs black, isort, mypy, packaging, pre-commit, psutil, pytest, and ruff to use fixed versions (or another explicit pinning strategy) so the Docker image remains reproducible. Keep the change localized to the existing RUN block in Dockerfile.cuda.docker/Dockerfile.rocm (2)
27-39: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winPin dev-tooling versions for reproducibility.
Same as the CUDA Dockerfile:
black,einops,isort,mypy,ninja,packaging,pre-commit,psutil,pytest,ruffare unpinned, undermining the stated reproducibility goal of this image.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docker/Dockerfile.rocm` around lines 27 - 39, The ROCm Dockerfile’s dev-tooling install block leaves several packages unpinned, which breaks reproducible builds. Update the `RUN python -m pip install ...` step in `docker/Dockerfile.rocm` to use fixed versions for the same tooling packages installed there (`black`, `einops`, `isort`, `mypy`, `ninja`, `packaging`, `pre-commit`, `psutil`, `pytest`, `ruff`), matching the version-pinning approach used in the CUDA Dockerfile. Keep the change limited to the dev-tooling install list so the image remains reproducible.
1-42: 🔒 Security & Privacy | 🔵 Trivial | 💤 Low valueNo non-root
USERin the image.Trivy flags this image as running as root with no
USERdirective (rule DS-0002). Consider adding a non-root user for the tooling/lint stage, keeping root only where GPU device access truly requires it (e.g., insideci/run_rocm_container.sh'sdocker run, which already grants device group access via--group-add).🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docker/Dockerfile.rocm` around lines 1 - 42, The Dockerfile.rocm image currently has no non-root USER, so the tooling/lint container runs as root by default. Update the Dockerfile to create and switch to a non-root user before the final WORKDIR/command context, while keeping root only where necessary for GPU access in ci/run_rocm_container.sh’s docker run flow. Use the existing build setup in Dockerfile.rocm and preserve the current package installs, but ensure the default runtime user is non-root.Source: Linters/SAST tools
docs/contributing/testing.md (1)
106-117: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUpdate docs once the label re-trigger issue in
gpu-ci.ymlis fixed.If the
needs-gpu-ci/needs-gpu-ci-sm90trigger types are restricted tolabeledonly (see companion comment ongpu-ci.yml), this section should clarify that maintainers must re-apply the label after each new push to re-run hardware CI, since pushing new commits will no longer auto-retrigger these jobs.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docs/contributing/testing.md` around lines 106 - 117, The RunPod CUDA runner docs need to reflect the new label-trigger behavior for the GPU CI flow. Update the “Path B — RunPod CUDA runner” section to state that, when using needs-gpu-ci or needs-gpu-ci-sm90, maintainers must re-apply the label after each new push to re-run hardware CI because labeled-only triggers in gpu-ci.yml will not auto-restart on new commits. Keep the guidance aligned with the existing run_gpu_ci.sh and run_cuda_tests.sh workflow and mention both label names explicitly.csrc/cuda/fused_linear_logp_sm90.cu (1)
378-387: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winConsider a launch check after the first kernel too.
C10_CUDA_KERNEL_LAUNCH_CHECK()is only invoked after the combine kernel. A launch/config failure onfused_linear_logp_sm90_kernel(e.g., the large dynamic-smem path) would not be surfaced before the combine kernel is enqueued against partially-initializedpart_*buffers, and the resulting error is attributed to the wrong launch. Adding a check right after the first launch improves error attribution and avoids running combine on bad state.🛡️ Proposed change
fused_linear_logp_sm90_kernel<<<grid, WG_THREADS, smem, stream>>>( h_tmap, w_tmap, target_i.data_ptr<int>(), bias_ptr, part_max.data_ptr<float>(), part_sum.data_ptr<float>(), part_zt.data_ptr<float>(), N, D, V, n_split); + C10_CUDA_KERNEL_LAUNCH_CHECK(); const int combine_threads = 256;🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@csrc/cuda/fused_linear_logp_sm90.cu` around lines 378 - 387, The first CUDA launch in fused_linear_logp_sm90 currently has no immediate launch/config check, so failures can be misattributed to the later combine step. Add a C10_CUDA_KERNEL_LAUNCH_CHECK() right after the fused_linear_logp_sm90_kernel launch, before enqueueing fused_linear_logp_sm90_combine_kernel, so errors from the first kernel are surfaced at the correct point and the combine kernel is not run on invalid part_* buffers.setup.py (1)
34-40: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low valueOverride path may drop the arch-native
asuffix.When
KERNEL_ALIGN_SM90_ARCHis set, the returned string is normalized but never guaranteed to end ina(e.g.,sm_90→90), whereas the device-derived branches always appenda. Per the comment at Line 140-141, WGMMA/TMA require an arch-nativeatarget, so an override without it could silently produce a non-abuild. Consider normalizing/appendingaunless the user is expected to always pass the full value.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@setup.py` around lines 34 - 40, The override handling in _sm90_arch_from_env_or_device can return a non-arch-native target when KERNEL_ALIGN_SM90_ARCH is set, unlike the cc_major/cc_minor branches that always produce an `a` suffix. Update the override normalization in _sm90_arch_from_env_or_device so the returned arch is guaranteed to include the trailing `a` (or otherwise validate that the env var already provides it), keeping behavior consistent with the WGMMA/TMA arch-native requirement referenced near the build logic.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In @.github/actions/setup-runpod/action.yml:
- Around line 17-23: The Runpod installation step in the setup-runpod action
currently downloads from a moving latest release, which can make CI
non-deterministic. Update the runpodctl install block to fetch a specific
released version instead of latest, and add checksum verification before moving
the binary into place. Use the existing install step in the action.yml workflow
to pin both the version and the expected hash so the downloaded binary is
validated.
In @.github/workflows/build-ci-image.yml:
- Around line 40-51: The checkout step in the PR build workflow still persists
the GitHub token into the workspace, which is unsafe for a Docker build driven
by PR-controlled files. Update the actions/checkout usage in this job to set
persist-credentials to false so the token is not written into .git/config before
building the CUDA/ROCm images, and keep the change scoped to the checkout step
used by the build-ci-image workflow.
In @.github/workflows/gpu-ci.yml:
- Around line 57-60: The base-branch checkout steps in the GPU CI workflow are
missing credential hardening. Update each `actions/checkout@v4` step that uses
`github.event.pull_request.base.sha` to also set `persist-credentials: false`,
matching the other secure checkout pattern and covering all four occurrences in
the hardware jobs.
- Around line 13-20: The GPU CI workflow currently triggers on synchronize as
well as labeled, so once the hardware label is added it can keep rerunning with
secrets on new pushes. Update the workflow in gpu-ci.yml to restrict execution
to fresh label events only, either by removing synchronize from the
pull_request_target trigger and keeping labeled, or by adding an explicit
github.event.action == 'labeled' check in each hardware job that uses the
setup-runpod and GPU test steps.
In `@ci/run_cuda_tests.sh`:
- Around line 137-165: The NCCL smoke test in run_cuda_tests.sh writes the
temporary script to a predictable /tmp path, which is symlink/TOCTOU-attackable.
Replace the hardcoded /tmp/rl_kernel_nccl_smoke.py target in the NCCL smoke test
block with a securely created unique temp file (for example via mktemp) and use
that variable consistently for both writing and executing the script. Keep the
existing GPU_COUNT, NCCL_P2P_DISABLE, and timeout flow intact, but ensure the
temp script path is not guessable or pre-creatable by another process.
---
Outside diff comments:
In @.github/workflows/build-ci-image.yml:
- Around line 86-107: The checkout step in the build-ci-image workflow is
missing the same hardening flag as the other job. Update the actions/checkout@v4
step in this push-job to set persist-credentials to false, matching the existing
security pattern and keeping GitHub credentials out of the local checkout before
the docker/login-action and image build/push steps run.
---
Nitpick comments:
In `@csrc/cuda/fused_linear_logp_sm90.cu`:
- Around line 378-387: The first CUDA launch in fused_linear_logp_sm90 currently
has no immediate launch/config check, so failures can be misattributed to the
later combine step. Add a C10_CUDA_KERNEL_LAUNCH_CHECK() right after the
fused_linear_logp_sm90_kernel launch, before enqueueing
fused_linear_logp_sm90_combine_kernel, so errors from the first kernel are
surfaced at the correct point and the combine kernel is not run on invalid
part_* buffers.
In `@docker/Dockerfile.cuda`:
- Around line 26-37: The dev-tooling packages installed in the Dockerfile.cuda
build are unpinned, so the image can drift between rebuilds. Update the pip
install step that installs black, isort, mypy, packaging, pre-commit, psutil,
pytest, and ruff to use fixed versions (or another explicit pinning strategy) so
the Docker image remains reproducible. Keep the change localized to the existing
RUN block in Dockerfile.cuda.
In `@docker/Dockerfile.rocm`:
- Around line 27-39: The ROCm Dockerfile’s dev-tooling install block leaves
several packages unpinned, which breaks reproducible builds. Update the `RUN
python -m pip install ...` step in `docker/Dockerfile.rocm` to use fixed
versions for the same tooling packages installed there (`black`, `einops`,
`isort`, `mypy`, `ninja`, `packaging`, `pre-commit`, `psutil`, `pytest`,
`ruff`), matching the version-pinning approach used in the CUDA Dockerfile. Keep
the change limited to the dev-tooling install list so the image remains
reproducible.
- Around line 1-42: The Dockerfile.rocm image currently has no non-root USER, so
the tooling/lint container runs as root by default. Update the Dockerfile to
create and switch to a non-root user before the final WORKDIR/command context,
while keeping root only where necessary for GPU access in
ci/run_rocm_container.sh’s docker run flow. Use the existing build setup in
Dockerfile.rocm and preserve the current package installs, but ensure the
default runtime user is non-root.
In `@docs/contributing/testing.md`:
- Around line 106-117: The RunPod CUDA runner docs need to reflect the new
label-trigger behavior for the GPU CI flow. Update the “Path B — RunPod CUDA
runner” section to state that, when using needs-gpu-ci or needs-gpu-ci-sm90,
maintainers must re-apply the label after each new push to re-run hardware CI
because labeled-only triggers in gpu-ci.yml will not auto-restart on new
commits. Keep the guidance aligned with the existing run_gpu_ci.sh and
run_cuda_tests.sh workflow and mention both label names explicitly.
In `@setup.py`:
- Around line 34-40: The override handling in _sm90_arch_from_env_or_device can
return a non-arch-native target when KERNEL_ALIGN_SM90_ARCH is set, unlike the
cc_major/cc_minor branches that always produce an `a` suffix. Update the
override normalization in _sm90_arch_from_env_or_device so the returned arch is
guaranteed to include the trailing `a` (or otherwise validate that the env var
already provides it), keeping behavior consistent with the WGMMA/TMA arch-native
requirement referenced near the build logic.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: e187c86a-5ba3-4617-8b9e-54fb7dddde6a
📒 Files selected for processing (28)
.github/actionlint.yaml.github/actions/setup-runpod/action.yml.github/workflows/build-ci-image.yml.github/workflows/ci.yml.github/workflows/gpu-ci.ymlbenchmarks/profiler.pyci/run_cuda_tests.shci/run_gpu_ci.shci/run_rocm_ci.shci/run_rocm_container.shcsrc/cuda/fused_linear_logp_sm90.cucsrc/cuda/fused_logp_sm90.cucsrc/fused_logp_kernel.cucsrc/ops.cppcsrc/utils/tma_utils.cuhdocker/Dockerfile.cudadocker/Dockerfile.rocmdocs/contributing/testing.mddocs/getting_started/installation.mdpyproject.tomlrequirements.txtsetup.pytests/conftest.pytests/test_kernel_registry.pytests/test_linear_logp.pytests/test_lm_head.pytests/test_op_accuracy.pytests/test_profiler.py
There was a problem hiding this comment.
🧹 Nitpick comments (1)
.github/workflows/gpu-ci.yml (1)
85-103: 🔒 Security & Privacy | 🔵 TrivialSelf-hosted runners execute PR head code directly — prefer ephemeral/isolated runners.
Unlike the cloud jobs (which only run the trusted base-branch orchestrator on the GitHub runner and push PR code to isolated pods),
cuda-self-hostedandrocm-self-hostedclone and runPR_SHAcode directly on persistent hardware viaci/run_cuda_tests.sh/ci/run_rocm_ci.sh. The maintainer label gate is a good first line of defense, but a labeled PR can still compromise the runner host, persist between jobs, and reach anything on that machine.Consider running these on ephemeral, single-use runners (fresh VM/container per job) and scoping/rotating any long-lived credentials on the host. The label guard mitigates casual abuse but does not isolate execution.
Also applies to: 140-159
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In @.github/workflows/gpu-ci.yml around lines 85 - 103, The self-hosted GPU jobs currently run untrusted PR head code on persistent hardware, so update the `cuda-self-hosted` and `rocm-self-hosted` workflows to use ephemeral, isolated runners instead of long-lived hosts. Keep the existing trusted checkout/orchestrator pattern, but ensure `Run CUDA tests` and `Run ROCm CI` execute on fresh single-use VM/container runners with scoped or rotated credentials so PR code cannot persist or affect the host between jobs.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In @.github/workflows/gpu-ci.yml:
- Around line 85-103: The self-hosted GPU jobs currently run untrusted PR head
code on persistent hardware, so update the `cuda-self-hosted` and
`rocm-self-hosted` workflows to use ephemeral, isolated runners instead of
long-lived hosts. Keep the existing trusted checkout/orchestrator pattern, but
ensure `Run CUDA tests` and `Run ROCm CI` execute on fresh single-use
VM/container runners with scoped or rotated credentials so PR code cannot
persist or affect the host between jobs.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: fb8f0fb3-50f8-4935-b5c3-193932bc5117
📒 Files selected for processing (4)
.github/actions/setup-runpod/action.yml.github/workflows/build-ci-image.yml.github/workflows/gpu-ci.ymlci/run_cuda_tests.sh
🚧 Files skipped from review as they are similar to previous changes (3)
- .github/actions/setup-runpod/action.yml
- .github/workflows/build-ci-image.yml
- ci/run_cuda_tests.sh
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR adds reproducible CUDA and ROCm Docker environments, plus label-gated hardware CI for CUDA, SM90/Hopper, and ROCm paths. Closes #173 .
The infrastructure work uncovered several real kernel/build/test issues, so this PR includes both the CI setup and the fixes needed for those hardware paths to pass reliably.
Changes
Docker and reproducible build environments
pytorch/pytorchdevel image, and ROCm uses the officialrocm/pytorchbase.build-ci-image.ymlvalidates both images in a PR matrix and publishesrl-kernel-ci:{cuda,rocm}frommain. The workflow also frees runner disk before building because CUDA/ROCm base images are large.Hardware CI matrix
gpu-ci.ymlso expensive GPU jobs only run when requested.KERNEL_ALIGN_FORCE_SM90=1for the SM90/TMA build.ci/run_cuda_tests.sh, whileci/run_gpu_ci.shhandles RunPod provisioning, cleanup, and endpoint discovery. RunPod setup is shared through.github/actions/setup-runpod.Fallback coverage
Documentation
pull_request_targetsecurity model.Fixes Needed To Make The Matrix Pass
CUDA SM90/TMA fixes
csrc/cuda/fused_logp_sm90.cu,csrc/cuda/fused_linear_logp_sm90.cu, andcsrc/utils/tma_utils.cuhnow launch kernels on the current CUDA stream instead of the default stream.C10_CUDA_KERNEL_LAUNCH_CHECKso launch failures are reported immediately.mbarrier_arrive_expect_txhappens before the TMA issue.fused_logp_sm90now keeps its TMA/MMA mbarriers out of dynamic shared memory.ROCm/HIP build fixes
csrc/fused_logp_kernel.cuandcsrc/ops.cppadd HIP portability guards for warp mask width and bf16 handling.setup.pynow points ROCm builds at the realcsrc/fused_logp_kernel.cusource file. The previous referencedcsrc/fused_logp_kernel.cppfile did not exist, so ROCm extension builds were broken on main.FORCE_CUDA=1can force the CUDA extension build path when CUDA is available.Dependency and test stability fixes
pyproject.tomlandsetup.pysplit optional dependencies intotest,hf, andbench, so CI can install only the dependency groups each path actually needs.accelerateandtransformersare no longer unconditional core dependencies; they are only needed for HF integration coverage.tabulatemoved out of core dependencies because it is only used by benchmark/profiler output and test tooling.rocmextra no longer tries to installaiterdirectly; ROCm dependencies that require source builds are handled separately in the ROCm setup docs.requirements.txtnow matches the package metadata and drops the bogusloggingPyPI package.tests/conftest.pypinsNVIDIA_TF32_OVERRIDE=0so cloud CUDA images cannot silently weaken fp32 oracle accuracy through forced TF32.tests/test_lm_head.pyprobes CPU BLAS fp32 batch invariance and skips bitwise Axis-A cases on machines where the BLAS genuinely lacks that property.Validation
KERNEL_ALIGN_FORCE_SM90=1; device capability detected as 9.0 and the SM90/TMA path was exercised. 356 passed, 82 skipped, 0 failed.ci/run_rocm_container.shusing a verified ROCm image. 93 passed, 91 skipped, 0 failed.FORCE_CUDA=1set. 343 passed, 95 skipped, 0 failed.Notes
Hardware jobs are label-gated to control cost. Default PR CI still runs unconditionally for lint, docs, CPU tests, and mocked fallback coverage.
Summary by CodeRabbit
New Features
Bug Fixes
Documentation