[WS2][kernels] Deterministic Qwen3 SwiGLU forward (CUDA SM90 + Triton) - #258
Open
bitborne wants to merge 2 commits into
Open
[WS2][kernels] Deterministic Qwen3 SwiGLU forward (CUDA SM90 + Triton)#258bitborne wants to merge 2 commits into
bitborne wants to merge 2 commits into
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (5)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughAdds BF16 SwiGLU implementations for CUDA SM90 and Triton. It integrates hardware-gated dispatch, build configuration, operator specifications, validation tests, CI smoke checks, benchmarks, and documentation. ChangesSwiGLU backend integration
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant OpRegistry
participant SwiGLUSM90Op
participant swiglu_forward_sm90
participant CUDAKernel
OpRegistry->>SwiGLUSM90Op: select SM90 backend on Hopper
SwiGLUSM90Op->>swiglu_forward_sm90: pass validated BF16 inputs
swiglu_forward_sm90->>CUDAKernel: launch SiLU(gate) × up
CUDAKernel-->>SwiGLUSM90Op: return BF16 output
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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 |
bitborne
changed the title
This was referenced Aug 5, 2026
Add BF16 CUDA SM90 and Triton forward implementations for SiLU(gate) * up, including dispatch, validation, CI, docs, and the Qwen3-8B TP-local 6144-wide test contract. Refs RL-Align#239 Signed-off-by: Schatten <czhengt@qq.com>
bitborne
force-pushed
the
codex/qwen3-swiglu-forward-sm90
branch
from
August 7, 2026 06:19
cb732a9 to
50b3430
Compare
bitborne
marked this pull request as ready for review
August 7, 2026 06:19
bitborne
requested review from
EthanZero2Hero,
Flink-ddd,
KJLdefeated and
inaniloquentee
as code owners
August 7, 2026 06:19
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (2)
tests/test_swiglu.py (1)
168-169: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winKeep a deterministic SwiGLU registry fallback test.
This change removes the only supplied assertion for
kernel_registry.get_op("swiglu"). Add a CPU-targeted assertion forNativeSwiGLUOp. This validates the operator key and fallback without requiring CUDA or Triton.Proposed test
def test_registry_dispatches_native_silu(): assert isinstance(kernel_registry.get_op("silu"), NativeSiLUOp) + +def test_registry_dispatches_native_swiglu_on_cpu(): + assert isinstance( + kernel_registry.get_op("swiglu", device="cpu"), + NativeSwiGLUOp, + )🤖 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 `@tests/test_swiglu.py` around lines 168 - 169, Add a CPU-targeted test alongside test_registry_dispatches_native_silu that calls kernel_registry.get_op("swiglu") and asserts the result is a NativeSwiGLUOp, preserving deterministic coverage of the SwiGLU registry key and native fallback without CUDA or Triton.tests/test_swiglu_forward_backends.py (1)
31-33: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winTest the Triton fallback on supported non-Hopper CUDA devices.
The registry selects
TRITON_SWIGLUfor CUDA before the PyTorch backend. This marker skips every Triton test unless the device is SM90.Gate these tests on CUDA, Triton, and BF16 support instead of Hopper. If Triton is intentionally Hopper-only, add the same hardware gate to registry dispatch.
🤖 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 `@tests/test_swiglu_forward_backends.py` around lines 31 - 33, Update the requires_triton_sm90 marker to gate tests on CUDA availability, Triton availability, and BF16 support rather than requiring _IS_SM90, so Triton fallback tests run on supported non-Hopper devices. Preserve the existing skip behavior and reason text unless the registry dispatch is instead changed to restrict TRITON_SWIGLU to Hopper hardware.
🤖 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 `@ci/run_gpu_ci.sh`:
- Line 141: Update the KERNEL_ALIGN_ACTIVATION_SM90 export in ci/run_gpu_ci.sh
to preserve the variable’s raw value without adding literal single-quote
characters, so empty and numeric values remain valid for envs.env_flag and
setup.py.
In `@rl_engine/kernels/registry.py`:
- Around line 364-368: Make CUDA backend selection in get_op device-aware
instead of relying on the hardware adjustment performed by
_adjust_priority_for_hardware() at registry construction. Before caching or
returning a backend for device "cuda:N", validate SM90 eligibility against that
requested device and exclude CUDA_SWIGLU_SM90 on non-Hopper devices. Update the
CUDA backend cache key or structure so eligibility is not shared across
incompatible devices, preserving fallback to another SwiGLU backend.
In `@setup.py`:
- Around line 184-189: Update the activation SM90 setup logic around
enable_activation_sm90 so that enabling the backend validates
csrc/cuda/activation/swiglu_sm90.cu exists and raises a build error when it is
missing; only append the source and corresponding flags when validation
succeeds.
---
Nitpick comments:
In `@tests/test_swiglu_forward_backends.py`:
- Around line 31-33: Update the requires_triton_sm90 marker to gate tests on
CUDA availability, Triton availability, and BF16 support rather than requiring
_IS_SM90, so Triton fallback tests run on supported non-Hopper devices. Preserve
the existing skip behavior and reason text unless the registry dispatch is
instead changed to restrict TRITON_SWIGLU to Hopper hardware.
In `@tests/test_swiglu.py`:
- Around line 168-169: Add a CPU-targeted test alongside
test_registry_dispatches_native_silu that calls kernel_registry.get_op("swiglu")
and asserts the result is a NativeSwiGLUOp, preserving deterministic coverage of
the SwiGLU registry key and native fallback without CUDA or Triton.
🪄 Autofix
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 Plus
Run ID: b3da3fbe-8b78-42ae-858f-25c7a3197d22
📒 Files selected for processing (19)
benchmarks/benchmark_swiglu.pyci/run_gpu_ci.shcsrc/cuda/activation/swiglu_sm90.cucsrc/ops.cppdocs/operators/activation.mdenvs.pyrl_engine/_C.pyirl_engine/kernels/gtest/operator_inputs.pyrl_engine/kernels/gtest/operator_specs.pyrl_engine/kernels/ops/cuda/activation/__init__.pyrl_engine/kernels/ops/cuda/activation/swiglu.pyrl_engine/kernels/ops/triton/activation/__init__.pyrl_engine/kernels/ops/triton/activation/swiglu.pyrl_engine/kernels/registry.pyscripts/check_operator.pyscripts/ci_smoke.pysetup.pytests/test_swiglu.pytests/test_swiglu_forward_backends.py
| export FORCE_CUDA=1 | ||
| export MAX_JOBS=8 | ||
| export KERNEL_ALIGN_FORCE_SM90="'"${KERNEL_ALIGN_FORCE_SM90}"'" | ||
| export KERNEL_ALIGN_ACTIVATION_SM90="'"${KERNEL_ALIGN_ACTIVATION_SM90}"'" |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Remove literal quote characters from the export value.
This expression exports '' for the empty default and '1' when the value is 1. envs.env_flag rejects both values, so setup.py raises ValueError and blocks the GPU CI build.
Proposed fix
-export KERNEL_ALIGN_ACTIVATION_SM90="'"${KERNEL_ALIGN_ACTIVATION_SM90}"'"
+export KERNEL_ALIGN_ACTIVATION_SM90="${KERNEL_ALIGN_ACTIVATION_SM90}"📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| export KERNEL_ALIGN_ACTIVATION_SM90="'"${KERNEL_ALIGN_ACTIVATION_SM90}"'" | |
| export KERNEL_ALIGN_ACTIVATION_SM90="${KERNEL_ALIGN_ACTIVATION_SM90}" |
🧰 Tools
🪛 Shellcheck (0.11.0)
[info] 141-159: Expressions don't expand in single quotes, use double quotes for that.
(SC2016)
🤖 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 `@ci/run_gpu_ci.sh` at line 141, Update the KERNEL_ALIGN_ACTIVATION_SM90 export
in ci/run_gpu_ci.sh to preserve the variable’s raw value without adding literal
single-quote characters, so empty and numeric values remain valid for
envs.env_flag and setup.py.
Source: Linters/SAST tools
Signed-off-by: Schatten <czhengt@qq.com>
Collaborator
|
please resolve the code conflicts, Thanks. |
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 the missing Qwen3 SwiGLU forward activation boundary:
It includes:
Integration contract
This implementation follows the fixed layout in #239.
M_local = B * S / 2for the fixed CP=2 layout.The intended integration entry point is:
Scope
This draft only covers the forward activation operator.
The following remain outside this PR:
This PR references the missing activation TODO in #239 and does not close the issue.
Backend behavior
On SM90, registry dispatch prefers:
Each optimized backend is bitwise invariant to unrelated batch, chunk, and padding rows. CUDA and Triton are compared against the independent FP32 oracle using the shared #108 elementwise tolerance.
Cross-backend bitwise equality is not currently required because the CUDA and Triton exponential implementations may differ.
Validation
Passed locally:
pre-commit run --from-ref upstream/main --to-ref HEADcompileallclang-format --dry-run --Werror csrc/cuda/activation/swiglu_sm90.cubash -n ci/run_gpu_ci.shgit diff --check upstream/main...HEADSigned-off-byGPU runtime validation was not available locally. Before this PR is marked ready, the CUDA and Triton paths should be exercised on H100 through the repository's
needs-gpu-ciworkflow and in the composed Gate/Up → SwiGLU → Down integration.Integration follow-up
The downstream integration should confirm:
Refs #239
Summary by CodeRabbit