perf(hip): adaptive 2D-grid row-split for act_and_mul kernels#265
Merged
Conversation
silu_and_mul / gelu / gelu_tanh launched one block per token, which underfills the GPU for small num_tokens (decode / small-batch serving) with large hidden dim, leaving HBM bandwidth idle. Output elements are independent, so split each row across blocks on gridDim.y (no atomics): blockIdx.x = token, blockIdx.y = column-tile. blocks_per_row is chosen host-side to oversubscribe the CU count by 2x only when num_tokens is below that threshold, so prefill (gridDim.y == 1) is byte-for-byte unchanged and the CUDA 1D-launch path is unaffected. Measured 1.16-1.39x on MI300X for large-d, small/medium-token shapes; flat elsewhere. The oversubscription factor past ~1x CU coverage is bandwidth-neutral (swept 1-32), so 2x is the minimal safe choice. The launch heuristic lives in one shared inline helper in activation.cuh called by both the AOT launcher and the JIT template (no drift), and CU count is queried via a new cached getMultiProcessorCount. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
demandal25
force-pushed
the
perf-silu-and-mul-adaptive-grid
branch
from
20e244a to
12cf818
Compare
There was a problem hiding this comment.
Pull request overview
This PR improves ROCm performance for silu_and_mul / gelu / gelu_tanh activation-and-multiply kernels in small-num_tokens, large-d (decode/small-batch) regimes by switching the kernel to a 2D grid and adding a shared launch heuristic that splits each row across multiple blocks to better fill the CU array.
Changes:
- Generalize
act_and_mul_kernelto support a 2D grid (blockIdx.ycolumn-tiling) and add a sharedact_and_mul_launch_dimsheuristic. - Add a cached SM/CU count query (
getMultiProcessorCount) used by the heuristic, and route ROCm AOT + HIP JIT launchers through the shared helper. - Extend ROCm activation tests with shapes that exercise both row-split and prefill-scale paths; add a new ROCm benchmark script.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 6 comments.
Show a summary per file
| File | Description |
|---|---|
include/flashinfer/attention/generic/activation.cuh |
Implements 2D-grid kernel + adaptive launch-dims helper for ROCm activation kernels. |
include/gpu_iface/gpu_runtime_compat.hpp |
Adds cached multiprocessor-count query used by the new launch heuristic. |
flashinfer/csrc_rocm/activation.cu |
Updates ROCm AOT launchers to use the shared adaptive launch helper. |
flashinfer/jit/activation.py |
Updates HIP JIT template launcher to call the shared adaptive launch helper. |
tests/rocm_tests/test_activation_hip.py |
Adds test shapes to cover row-split and no-split (prefill-scale) behavior. |
benchmarks/rocm_benchmarks/bench_silu_and_mul.py |
Adds a benchmark sweep to measure decode vs. prefill regimes for silu_and_mul. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
- getMultiProcessorCount: make the per-device cache thread_local so concurrent callers (multi-threaded Python) cannot race on it. - silu/gelu/gelu_tanh AOT launchers + HIP JIT template: early-return on num_tokens == 0 so empty tensors are a no-op instead of a 0-sized (invalid) grid launch. - Reword the kernel's "byte-for-byte identical" comment to "same memory-access pattern" (codegen differs; behavior does not). Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Comment on lines
35
to
+38
| void silu_and_mul(at::Tensor& out, at::Tensor& input, bool enable_pdl) { | ||
| int d = input.size(-1) / 2; | ||
| int64_t num_tokens = input.numel() / input.size(-1); | ||
| if (num_tokens == 0) return; // empty input → no-op (a 0-sized grid is an invalid launch) |
Comment on lines
60
to
+63
| void gelu_tanh_and_mul(at::Tensor& out, at::Tensor& input, bool enable_pdl) { | ||
| int d = input.size(-1) / 2; | ||
| int64_t num_tokens = input.numel() / input.size(-1); | ||
| if (num_tokens == 0) return; // empty input → no-op (a 0-sized grid is an invalid launch) |
Comment on lines
84
to
88
| void gelu_and_mul(at::Tensor& out, at::Tensor& input, bool enable_pdl) { | ||
| int d = input.size(-1) / 2; | ||
| int64_t num_tokens = input.numel() / input.size(-1); | ||
| if (num_tokens == 0) return; // empty input → no-op (a 0-sized grid is an invalid launch) | ||
| const c10::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard(out.device()); |
Comment on lines
28
to
32
| void {{ func_name }}(at::Tensor& out, at::Tensor& input, bool enable_pdl) { | ||
| int d = input.size(-1) / 2; | ||
| int64_t num_tokens = input.numel() / input.size(-1); | ||
| if (num_tokens == 0) return; // empty input → no-op (a 0-sized grid is an invalid launch) | ||
|
|
Comment on lines
+84
to
+87
| // Scalar remainder over [num_vec*vec_size, d), column-tiled the same way. | ||
| // Always empty for the fp16/bf16 dispatch (16-byte alignment forces d % vec_size | ||
| // == 0); kept defensive. Do NOT key this off d % (blockDim.x * vec_size) — that | ||
| // assumes blockDim.x is the global stride, which is false under column-tiling. |
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
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
The
silu_and_mul/gelu/gelu_tanhactivation kernels launched one block per token (gridDim = num_tokens). For decode and small-batch serving — few tokens, large hidden dim — onlynum_tokensblocks launch on a 304-CU MI300X, so the GPU sits mostly idle and most of HBM bandwidth (the bottleneck for this memory-bound kernel) goes unused. This PR makes the launch adaptive: it splits each row across multiple blocks to fill the CU array whennum_tokensis small, while leaving the already-saturated prefill path unchanged.What changed
Kernel
include/flashinfer/attention/generic/activation.cuh— generalizedact_and_mul_kernelto a 2D grid:blockIdx.xselects the token (row),blockIdx.yselects a column-tile of that row. Output elements are independent (no cross-element reduction), so a row can be split acrossgridDim.yblocks with no atomics. WhengridDim.y == 1the kernel collapses to one block per token, byte-for-byte identical to the original — so prefill and the unchanged CUDA 1D-launch path are unaffected. The intra-row column index stays 32-bit (row-base addresses are 64-bit) to keep the inner-loop address math identical to the original hot loop.Launch heuristic
include/flashinfer/attention/generic/activation.cuh— added a single sharedact_and_mul_launch_dimshelper that picksblocks_per_rowto oversubscribe the CU count by 2x, but only whennum_tokensis below that threshold (so prefill stays atblocks_per_row == 1). It is called from both the AOT launcher and the JIT template, so the two paths cannot drift. Also fixes a latentblockDim == 0launch crash ford < vec_sizethat the AOT path had (the JIT path already guarded it).include/gpu_iface/gpu_runtime_compat.hpp— addedgetMultiProcessorCount, a cached cross-platform CU/SM-count query (the heuristic needs it on every launch; the cache avoids a per-launch driver call).flashinfer/csrc_rocm/activation.cu— all three AOT launchers now call the shared helper.flashinfer/jit/activation.py— the HIP JIT template (the runtime path, sinceautoresolves to native) calls the shared helper.Tests / benchmark
tests/rocm_tests/test_activation_hip.py— added two shapes:(2, 14336)exercisesblocks_per_row > 1(row split across blocks) and(4096, 4096)locks the prefillblocks_per_row == 1path.benchmarks/rocm_benchmarks/bench_silu_and_mul.py— new benchmark sweeping num_tokens x d x dtype through the existingRocmProfilerharness.Design notes
The oversubscription factor only needs to cross ~1x CU coverage; this kernel is bandwidth-bound, so once enough blocks are resident to saturate the HBM controllers, adding more concurrent blocks cannot move data faster. A sweep of the factor over {1, 2, 4, 8, 16, 32} at d=14336 was flat within noise, so
2xis chosen as the minimal value that reliably crosses the fill threshold while avoiding pointless extra blocks (and tail overhead) oncenum_tokensalone fills the GPU. CU count is a runtime device property (varies by SKU), so it is queried at runtime rather than templated.Benchmark results
MI300X (gfx942), median kernel time, optimized vs. baseline (silu, both dtypes), measured back-to-back through
bench_silu_and_mul.py:Largest wins land at large d with small/medium num_tokens (the decode / small-batch regime). Shapes that are launch-overhead-bound (small d) or already GPU-saturated (prefill) are flat, as expected.
Test plan
pytest tests/rocm_tests/test_activation_hip.py— 73 passed (includes the two new shapes x 3 activations x 2 dtypes)pytest tests/rocm_tests/test_activation_aiter_hip.py— passed (AITER backend unaffected)pytest tests/utils/test_activation.py— 1890 passed (broader dim/batch sweep)python benchmarks/rocm_benchmarks/bench_silu_and_mul.py --timing-only— before/after comparison abovepre-commit run -a(clang-format / ruff / mypy on the changed files)