Add tirx silu_mul kernel - #78
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a TIRx-based implementation of the fused SwiGLU activation function (silu(gate) * up) with custom autograd forward and backward passes. The review feedback highlights critical safety improvements, specifically recommending early returns for empty tensors (n == 0) in both the forward and backward passes to prevent CUDA driver crashes, and adding assertions to verify that the input tensors reside on the same CUDA device.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| out = torch.empty_like(gate) | ||
| dtype, elem_bytes = _DTYPE_MAP[gate.dtype] | ||
| n = gate.numel() | ||
| assert n % (_VEC_BYTES // elem_bytes) == 0, "numel must be vector-divisible" | ||
| block, tile = _select_config(n) | ||
| _fwd_exec(dtype, block, tile, elem_bytes)( | ||
| gate.reshape(-1), | ||
| up.reshape(-1), | ||
| out.reshape(-1), | ||
| ) | ||
| return out |
There was a problem hiding this comment.
If the input tensors are empty (i.e., n == 0), launching the CUDA kernel with 0 blocks will cause a CUDA driver error or crash. Add a check to return early when n == 0. Additionally, improve the vector-divisibility assertion message to be more descriptive.
| out = torch.empty_like(gate) | |
| dtype, elem_bytes = _DTYPE_MAP[gate.dtype] | |
| n = gate.numel() | |
| assert n % (_VEC_BYTES // elem_bytes) == 0, "numel must be vector-divisible" | |
| block, tile = _select_config(n) | |
| _fwd_exec(dtype, block, tile, elem_bytes)( | |
| gate.reshape(-1), | |
| up.reshape(-1), | |
| out.reshape(-1), | |
| ) | |
| return out | |
| out = torch.empty_like(gate) | |
| n = gate.numel() | |
| if n == 0: | |
| return out | |
| dtype, elem_bytes = _DTYPE_MAP[gate.dtype] | |
| vec_size = _VEC_BYTES // elem_bytes | |
| assert n % vec_size == 0, f"numel ({n}) must be divisible by vector size ({vec_size})" | |
| block, tile = _select_config(n) | |
| _fwd_exec(dtype, block, tile, elem_bytes)( | |
| gate.reshape(-1), | |
| up.reshape(-1), | |
| out.reshape(-1), | |
| ) | |
| return out |
| grad_gate = torch.empty_like(gate) | ||
| grad_up = torch.empty_like(up) | ||
| dtype, elem_bytes = _DTYPE_MAP[gate.dtype] | ||
| n = gate.numel() | ||
| block, tile = _select_config(n) | ||
| _bwd_exec(dtype, block, tile, elem_bytes)( | ||
| grad_output.reshape(-1), | ||
| gate.reshape(-1), | ||
| up.reshape(-1), | ||
| grad_gate.reshape(-1), | ||
| grad_up.reshape(-1), | ||
| ) | ||
| return grad_gate, grad_up |
There was a problem hiding this comment.
Similarly to the forward pass, if the input tensors are empty (n == 0), the backward pass will attempt to launch a CUDA kernel with 0 blocks and crash. Add an early return check for n == 0 in the backward pass.
| grad_gate = torch.empty_like(gate) | |
| grad_up = torch.empty_like(up) | |
| dtype, elem_bytes = _DTYPE_MAP[gate.dtype] | |
| n = gate.numel() | |
| block, tile = _select_config(n) | |
| _bwd_exec(dtype, block, tile, elem_bytes)( | |
| grad_output.reshape(-1), | |
| gate.reshape(-1), | |
| up.reshape(-1), | |
| grad_gate.reshape(-1), | |
| grad_up.reshape(-1), | |
| ) | |
| return grad_gate, grad_up | |
| grad_gate = torch.empty_like(gate) | |
| grad_up = torch.empty_like(up) | |
| n = gate.numel() | |
| if n == 0: | |
| return grad_gate, grad_up | |
| dtype, elem_bytes = _DTYPE_MAP[gate.dtype] | |
| block, tile = _select_config(n) | |
| _bwd_exec(dtype, block, tile, elem_bytes)( | |
| grad_output.reshape(-1), | |
| gate.reshape(-1), | |
| up.reshape(-1), | |
| grad_gate.reshape(-1), | |
| grad_up.reshape(-1), | |
| ) | |
| return grad_gate, grad_up |
| assert gate.shape == up.shape, f"shape mismatch: {gate.shape} vs {up.shape}" | ||
| assert gate.dtype == up.dtype, f"dtype mismatch: {gate.dtype} vs {up.dtype}" | ||
| assert gate.is_contiguous(), "gate must be contiguous" | ||
| assert up.is_contiguous(), "up must be contiguous" | ||
| assert gate.dtype in _DTYPE_MAP, f"unsupported dtype: {gate.dtype}" | ||
| ctx.save_for_backward(gate, up) |
There was a problem hiding this comment.
Add assertions to verify that both input tensors are on the same CUDA device and that they are indeed CUDA tensors. Since the TVM target is hardcoded to "cuda", passing CPU tensors or tensors on mismatched devices will cause runtime crashes or undefined behavior.
| assert gate.shape == up.shape, f"shape mismatch: {gate.shape} vs {up.shape}" | |
| assert gate.dtype == up.dtype, f"dtype mismatch: {gate.dtype} vs {up.dtype}" | |
| assert gate.is_contiguous(), "gate must be contiguous" | |
| assert up.is_contiguous(), "up must be contiguous" | |
| assert gate.dtype in _DTYPE_MAP, f"unsupported dtype: {gate.dtype}" | |
| ctx.save_for_backward(gate, up) | |
| assert gate.shape == up.shape, f"shape mismatch: {gate.shape} vs {up.shape}" | |
| assert gate.dtype == up.dtype, f"dtype mismatch: {gate.dtype} vs {up.dtype}" | |
| assert gate.device == up.device, f"device mismatch: {gate.device} vs {up.device}" | |
| assert gate.is_cuda, "gate must be on a CUDA device" | |
| assert gate.is_contiguous(), "gate must be contiguous" | |
| assert up.is_contiguous(), "up must be contiguous" | |
| assert gate.dtype in _DTYPE_MAP, f"unsupported dtype: {gate.dtype}" | |
| ctx.save_for_backward(gate, up) |
|
@claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| # Flat-grid configs as (BLOCK threads, ELEMS_PER_CTA), selected by element count. | ||
| _CONFIG_SMALL = (256, 2048) # small problems | ||
| _CONFIG_MID = (128, 2048) # medium problems ("current") | ||
| _CONFIG_LARGE = (128, 1024) # large problems | ||
|
|
||
|
|
||
| def _select_config(n: int) -> tuple[int, int]: | ||
| """Pick (BLOCK, ELEMS_PER_CTA) for a flat array of ``n`` elements.""" | ||
| if n < 4194304: | ||
| return _CONFIG_SMALL | ||
| if n < 16777216: | ||
| return _CONFIG_MID | ||
| return _CONFIG_LARGE |
There was a problem hiding this comment.
Compactness: The sibling Triton silu_mul.py runs every shape through one fixed _BLOCK_SIZE = 1024; this port adds three configs plus _select_config with two magic thresholds. The PR's own benchmark table is ~1.00–1.05x across all sizes, so the extra branches appear to buy near-nothing over a single config. Unless the SMALL/LARGE variants measurably beat MID in isolation, collapse to the one _CONFIG_MID constant (matching the sibling) and drop _select_config — it removes the thresholds, two config tuples, and a param threaded through every builder/cache function.
Compactness ReviewClean, self-contained TIRx port of the sibling Triton |
| @@ -0,0 +1,258 @@ | |||
| """ | |||
| TIRx implementation of the fused SwiGLU-style ``silu(gate) * up`` autograd | |||
| function (a port of the Triton kernels in ``silu_mul.py``). | |||
There was a problem hiding this comment.
Consistency: silu_mul.py is this file's own basename, so a reader can't tell it points at the sibling Triton module (pithtrain/operators/silu_mul.py) rather than self-referencing. Qualify the path.
| function (a port of the Triton kernels in ``silu_mul.py``). | |
| function (a port of the Triton kernels in ``operators/silu_mul.py``). |
|
|
||
| # Flat-grid configs as (BLOCK threads, ELEMS_PER_CTA), selected by element count. | ||
| _CONFIG_SMALL = (256, 2048) # small problems | ||
| _CONFIG_MID = (128, 2048) # medium problems ("current") |
There was a problem hiding this comment.
Consistency: ("current") is a leftover tuning annotation that describes nothing to a future reader — and it misleads, since MID is not an always-selected default (config is chosen by n in _select_config). Drop it.
| _CONFIG_MID = (128, 2048) # medium problems ("current") | |
| _CONFIG_MID = (128, 2048) # medium problems |
Consistency ReviewAdd-only PR (one new file, |
Performance ReviewNo performance regression identified. The new TIRx kernel is a standalone addition ( Two notes for when/if this replaces the Triton kernel: (1) the sole call site ( |
| return ex | ||
|
|
||
|
|
||
| class _SiLUMulTirx(torch.autograd.Function): |
There was a problem hiding this comment.
Correctness: This is a custom-autograd kernel (defines the backward for silu(gate)*up), but the PR's only evidence is a speedup table — no numerical/gradient parity is shown, and nothing tests it: tests/test_silu_mul.py imports the Triton silu_mul, not silu_mul_tirx, and the op is not wired into any model. A wrong grad_gate/grad_up here would leave forward and loss identical while silently corrupting gradients, and there is no test to catch it.
Add a test analogous to test_silu_mul.py that compares forward output and both grads against F.silu(gate)*up / the Triton kernel. Note the shapes must be vector-divisible (numel % 8 for bf16/fp16) — the forward asserts this (line 210), so this port does not accept the arbitrary shapes the Triton version handles (e.g. the (1023, 777) case in the existing test would raise).
| n = gate.numel() | ||
| assert n % (_VEC_BYTES // elem_bytes) == 0, "numel must be vector-divisible" | ||
| block, tile = _select_config(n) |
There was a problem hiding this comment.
Correctness: Empty input (n == 0) launches a zero-block grid — (0 + TILE - 1)//TILE == 0 at line 83 — which raises a CUDA invalid-configuration error at launch. The Triton reference tolerates grid=(0,) as a no-op, so this port diverges on the empty-group edge case (an MoE expert receiving 0 tokens). Guard it with an early return in both forward and backward:
| n = gate.numel() | |
| assert n % (_VEC_BYTES // elem_bytes) == 0, "numel must be vector-divisible" | |
| block, tile = _select_config(n) | |
| n = gate.numel() | |
| if n == 0: | |
| return out | |
| assert n % (_VEC_BYTES // elem_bytes) == 0, "numel must be vector-divisible" | |
| block, tile = _select_config(n) |
Backward needs the same guard (if n == 0: return grad_gate, grad_up after line 232).
Correctness ReviewThe kernel math is a faithful port of the Triton reference: the forward The gaps are around evidence and edge cases:
Both are inline. One caveat worth noting: the forward asserts numel is vector-divisible, so this is not a drop-in for the arbitrary shapes the Triton version accepts. |
|
We'll need to test whether all kernels work on H100/200. |
|
Also, would TVM be a new dependency? I believe the current environment won't install it. |
How are we testing on Hopper? Are we using modal? |
3 participants
silu_mulTIRx vs Triton — speedup (bf16, B200)Pure-kernel time, min-of-N; speedup = triton_time / tirx_time
(>1.00 = TIRx faster).