[WS1] Add CUDA and Triton SiLU/SwiGLU activation kernels - #280
[WS1] Add CUDA and Triton SiLU/SwiGLU activation kernels#280maxiaosong1124 wants to merge 5 commits into
Conversation
Implement batch-invariant SiLU and SwiGLU CUDA/Triton backends matching the PyTorch fp32 gold path, register them in the kernel registry and issue RL-Align#108 OP_SPECS harness, and extend tests for correctness, Axis-A invariance, and candidate-vs-gold consistency on Qwen3-8B intermediate shapes.
📝 WalkthroughWalkthroughThis change adds CUDA and Triton implementations for SiLU and SwiGLU, including forward and backward paths. It exposes CUDA extension APIs, updates backend dispatch and validation, documents contracts, and adds operator and edge-case coverage. ChangesSiLU and SwiGLU activation backends
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant Registry
participant CUDAOrTriton
participant Autograd
Caller->>Registry: request SiLU or SwiGLU
Registry->>CUDAOrTriton: select available backend
CUDAOrTriton->>Autograd: run forward and save inputs
Autograd-->>CUDAOrTriton: request backward gradients
CUDAOrTriton-->>Caller: return output or gradients
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 |
![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 (2)
585-585: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse a raw string for the regex pattern.
Ruff flags RUF043 here. The pattern contains regex metacharacters. Use a raw string to make the intent explicit and to satisfy lint.
- with pytest.raises(RuntimeError, match="same .*device"): + with pytest.raises(RuntimeError, match=r"same .*device"):🤖 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` at line 585, Update the pytest.raises call’s match argument in the relevant test to use a raw string literal for the regex pattern, preserving the existing “same .*device” matching behavior and satisfying Ruff RUF043.Source: Linters/SAST tools
528-531: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd
strict=tozip().Ruff flags B905 on line 531.
expectedandactualare both 3-tuples from_run, sostrict=Trueis safe and makes a future tuple-shape change fail loudly.- assert all(torch.equal(lhs, rhs) for lhs, rhs in zip(expected, actual)) + assert all(torch.equal(lhs, rhs) for lhs, rhs in zip(expected, actual, strict=True))🤖 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 528 - 531, Update the zip() call in the test loop around _run() to pass strict=True, preserving the existing equality assertion while ensuring mismatched expected and actual tuple lengths fail explicitly.Source: Linters/SAST tools
🤖 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 `@csrc/cuda/activation.cu`:
- Around line 104-108: Update check_cuda_contig() to accept only torch.float16,
torch.bfloat16, and torch.float32, rejecting float64 at the native boundary.
Ensure the relevant AT_DISPATCH_FLOATING_TYPES_AND2[...] usage also excludes
double so direct _C calls cannot dispatch float64.
In `@rl_engine/kernels/ops/triton/activation/swiglu.py`:
- Around line 137-139: Update the activation Function.forward implementations
around _launch_silu_fwd and the corresponding SwiGLU forward path to materialize
contiguous input tensors once, pass those tensors to the launch helpers, and
save the same tensors with ctx.save_for_backward. Remove duplicate contiguity
conversions inside the helpers or forward code while preserving the existing
launch and autograd behavior.
In `@tests/test_swiglu.py`:
- Around line 26-40: Guard the module-level TritonSiLUOp and TritonSwiGLUOp
import with an ImportError handler, setting unavailable symbols and a capability
flag consistently with the CUDA block. Update the test parametrization to
include the Triton case only when Triton is available, while preserving native
and CUDA test coverage in environments without Triton.
---
Nitpick comments:
In `@tests/test_swiglu.py`:
- Line 585: Update the pytest.raises call’s match argument in the relevant test
to use a raw string literal for the regex pattern, preserving the existing “same
.*device” matching behavior and satisfying Ruff RUF043.
- Around line 528-531: Update the zip() call in the test loop around _run() to
pass strict=True, preserving the existing equality assertion while ensuring
mismatched expected and actual tuple lengths fail explicitly.
🪄 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: dc56ccd2-4372-4e97-a7be-e8eb6a895c00
📒 Files selected for processing (14)
csrc/cuda/activation.cucsrc/ops.cppdocs/operators/activation.mdrl_engine/_C.pyirl_engine/kernels/gtest/operator_specs.pyrl_engine/kernels/ops/cuda/__init__.pyrl_engine/kernels/ops/cuda/activation/__init__.pyrl_engine/kernels/ops/cuda/activation/swiglu.pyrl_engine/kernels/ops/pytorch/activation/swiglu.pyrl_engine/kernels/ops/triton/activation/__init__.pyrl_engine/kernels/ops/triton/activation/swiglu.pyrl_engine/kernels/registry.pysetup.pytests/test_swiglu.py
| operator_names, | ||
| ) | ||
| from rl_engine.kernels.ops.pytorch.activation.swiglu import NativeSiLUOp, NativeSwiGLUOp | ||
| from rl_engine.kernels.ops.triton.activation.swiglu import TritonSiLUOp, TritonSwiGLUOp |
There was a problem hiding this comment.
This top-level Triton import makes the whole file fail during collection when Triton isn’t installed. Could we guard it with a try/except and skip, like the other Triton tests?
| ] | ||
|
|
||
|
|
||
| @requires_nvidia_cuda |
There was a problem hiding this comment.
This assert feels a bit too hard for a CUDA test. If the extension isn’t built yet, this fails instead of skipping even on a valid CUDA setup. Could we make it skip unless those symbols are actually present?
Guard Triton imports and soft-skip missing CUDA symbols so CPU-only and partial-extension environments can still collect tests. Reuse a single contiguous materialization in CUDA/Triton autograd forwards, reject float64 at the native boundary, and resolve the Qwen3 integration script's hard import of SwiGLUSM90Op so pytest collection no longer aborts GPU CI.
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
tests/test_qwen3_fwd_integration.py (2)
96-109: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winBuild the oracle independently of
_SwiGLUOp.
hidden_fp32is converted fromhidden_local. A defective SwiGLU implementation can therefore pass this validation because both the tested path and the oracle use the same activation output.Compute the reference activation from
gate_local.float()andup_local.float()with PyTorch SiLU. Apply the expected BF16 boundary cast before the down projection. Asserthidden_localagainst that reference before validatingfinal_output.🤖 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_qwen3_fwd_integration.py` around lines 96 - 109, The FP32 oracle currently reuses hidden_local, so it cannot detect defects in _SwiGLUOp. Build the reference independently from gate_local.float() and up_local.float() using PyTorch SiLU, apply the expected BF16 boundary cast before the down projection, and assert hidden_local matches this reference before the existing final_output oracle validation.
34-41: 🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy liftModel the TP and CP data ownership.
x_localuses42 + rank. Ranks 0 and 1, and ranks 2 and 3, are in the same TP group but receive different activations. The TP all-reduce therefore sums independent forward passes instead of TP shards of one forward pass.
cp_groupsis created but never used. Derive TP and CP coordinates. Seed or shard activations by CP coordinate. Seed or shard weights by TP coordinate. Use the CP group in a CP invariant or communication check.Also applies to: 63-73
🤖 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_qwen3_fwd_integration.py` around lines 34 - 41, Update the distributed setup and test flow around the TP/CP group construction and forward inputs so data ownership matches the parallel dimensions: derive each rank’s TP and CP coordinates, shard or seed activations by CP coordinate, and shard or seed weights by TP coordinate. Replace the unused cp_groups path with a CP invariant or communication check, while preserving TP all-reduce validation for shards belonging to the same forward pass.
🤖 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.
Outside diff comments:
In `@tests/test_qwen3_fwd_integration.py`:
- Around line 96-109: The FP32 oracle currently reuses hidden_local, so it
cannot detect defects in _SwiGLUOp. Build the reference independently from
gate_local.float() and up_local.float() using PyTorch SiLU, apply the expected
BF16 boundary cast before the down projection, and assert hidden_local matches
this reference before the existing final_output oracle validation.
- Around line 34-41: Update the distributed setup and test flow around the TP/CP
group construction and forward inputs so data ownership matches the parallel
dimensions: derive each rank’s TP and CP coordinates, shard or seed activations
by CP coordinate, and shard or seed weights by TP coordinate. Replace the unused
cp_groups path with a CP invariant or communication check, while preserving TP
all-reduce validation for shards belonging to the same forward pass.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: fd2eda07-820a-4182-927b-69bf73247a25
📒 Files selected for processing (5)
csrc/cuda/activation.curl_engine/kernels/ops/cuda/activation/swiglu.pyrl_engine/kernels/ops/triton/activation/swiglu.pytests/test_qwen3_fwd_integration.pytests/test_swiglu.py
🚧 Files skipped from review as they are similar to previous changes (4)
- rl_engine/kernels/ops/triton/activation/swiglu.py
- rl_engine/kernels/ops/cuda/activation/swiglu.py
- csrc/cuda/activation.cu
- tests/test_swiglu.py
3 participants
Summary
Adds CUDA and Triton implementations of SiLU and SwiGLU for the WS1 batch-invariant forward chain.
The new kernels provide deterministic, element-wise forward and backward paths and match the existing PyTorch FP32 ground-truth implementation under the shared #108 numerical contract.
This PR also wires both operators into the kernel registry and the shared operator validation harness.
Part of #271.
Related to #83 and #108.
Implementation
silu(x) = x * sigmoid(x)swiglu(gate, up) = silu(gate) * upforwardpath.forward_fp32as the FP32 ground-truth path.OP_SPECS.Numerical contract and invariance
SiLU and SwiGLU are element-wise and row-independent, so no cross-row reductions or atomics are required.
The tests verify:
For SwiGLU,
gateandupmust have the same shape, dtype, and device. Broadcasting is intentionally unsupported.Validation
Validation environment:
Commands and results: