diff --git a/benchmarks/rocm_benchmarks/bench_silu_and_mul.py b/benchmarks/rocm_benchmarks/bench_silu_and_mul.py new file mode 100644 index 0000000000..f21963fafb --- /dev/null +++ b/benchmarks/rocm_benchmarks/bench_silu_and_mul.py @@ -0,0 +1,97 @@ +""" +Copyright (c) 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. + +silu_and_mul benchmark: sweeps num_tokens x d x dtype to expose the small-batch +(decode) regime where one-block-per-token underfills the GPU, alongside the +large-batch (prefill) regime that is already saturated. + +This kernel is memory-bandwidth bound (reads 2*d, writes d per token), so the +roofline sits on the HBM ceiling and tokens/sec is the headline metric. + +Run: + python benchmarks/rocm_benchmarks/bench_silu_and_mul.py # full pipeline + python benchmarks/rocm_benchmarks/bench_silu_and_mul.py --timing-only # no profiling + python benchmarks/rocm_benchmarks/bench_silu_and_mul.py --replot # regenerate plot +""" + +import logging +import sys +from pathlib import Path + +import torch + +import flashinfer +from flashinfer.jit.core import logger as _jit_logger + +_jit_logger.setLevel(logging.WARNING) + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent / "rocm_profiler")) +from rocm_profiler import KernelConfig, RocmProfiler + +_OUTPUT_DIR = str(Path(__file__).parent) + +# num_tokens sweep crosses the CU-fill boundary: tiny (decode), small batch, +# up to prefill-scale where blocks_per_row resolves to 1. +_NUM_TOKENS = [1, 8, 32, 64, 128, 256, 1024, 4096] +# d = hidden_size // 2 for representative MLP intermediate sizes. +_DIMS = [4096, 14336] +_DTYPES = [(torch.float16, "f16"), (torch.bfloat16, "bf16")] + + +@torch.inference_mode() +def _make_configs() -> list[KernelConfig]: + configs = [] + for dtype, dt_name in _DTYPES: + itemsize = torch.tensor([], dtype=dtype).element_size() + for d in _DIMS: + for nt in _NUM_TOKENS: + x = torch.randn(nt, 2 * d, device="cuda", dtype=dtype) + out = torch.empty(nt, d, device="cuda", dtype=dtype) + # Bandwidth-bound: read gate+up (2*d) and write (d) per token. + theo_bytes = nt * 3 * d * itemsize + # One mul per output element; FLOPs are not the bottleneck but the + # profiler needs a nonzero value for arithmetic intensity. + theo_flops = nt * d + configs.append( + KernelConfig( + name=f"silu_{dt_name}_nt{nt}_d{d}", + run_fn=torch.inference_mode()( + lambda x=x, out=out: flashinfer.activation.silu_and_mul( + x, out=out + ) + ), + theoretical_flops=theo_flops, + theoretical_bytes=theo_bytes, + num_tokens=nt, + label=f"{dt_name} nt={nt:>5d} d={d:>5d}", + ) + ) + return configs + + +if __name__ == "__main__": + _skip_gpu = "--replot" in sys.argv or "--list-presets" in sys.argv + profiler = RocmProfiler( + configs=[] if _skip_gpu else _make_configs(), + num_warmup=3, + dry_run_ms=100, + repeat_ms=1000, + counters="roofline", + kernel_name_regex="act_and_mul_kernel", + output_dir=_OUTPUT_DIR, + label="silu_and_mul", + roofline=True, + ) + profiler.run() diff --git a/flashinfer/csrc_rocm/activation.cu b/flashinfer/csrc_rocm/activation.cu index 88900a9d9d..e669670f1f 100644 --- a/flashinfer/csrc_rocm/activation.cu +++ b/flashinfer/csrc_rocm/activation.cu @@ -35,6 +35,7 @@ __device__ __forceinline__ float gelu_tanh(const float& val) { 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) const c10::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard(out.device()); const hipStream_t stream = at::hip::getCurrentHIPStream(); @@ -42,8 +43,9 @@ void silu_and_mul(at::Tensor& out, at::Tensor& input, bool enable_pdl) { DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(input.scalar_type(), c_type, [&] { uint32_t vec_size = 16 / sizeof(c_type); - uint64_t gridDim = num_tokens; - uint64_t blockDim = std::min(d / vec_size, 1024U); + dim3 gridDim, blockDim; + activation::act_and_mul_launch_dims(d, num_tokens, vec_size, out.get_device(), gridDim, + blockDim); activation::act_and_mul_kernel<<>>( static_cast(out.data_ptr()), static_cast(input.data_ptr()), d); @@ -58,14 +60,16 @@ void silu_and_mul(at::Tensor& out, at::Tensor& input, bool enable_pdl) { 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) const c10::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard(out.device()); const hipStream_t stream = at::hip::getCurrentHIPStream(); DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(input.scalar_type(), c_type, [&] { uint32_t vec_size = 16 / sizeof(c_type); - uint64_t gridDim = num_tokens; - uint64_t blockDim = std::min(d / vec_size, 1024U); + dim3 gridDim, blockDim; + activation::act_and_mul_launch_dims(d, num_tokens, vec_size, out.get_device(), gridDim, + blockDim); activation::act_and_mul_kernel<<>>( static_cast(out.data_ptr()), static_cast(input.data_ptr()), d); @@ -80,14 +84,16 @@ void gelu_tanh_and_mul(at::Tensor& out, at::Tensor& input, bool enable_pdl) { 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()); const hipStream_t stream = at::hip::getCurrentHIPStream(); DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(input.scalar_type(), c_type, [&] { uint32_t vec_size = 16 / sizeof(c_type); - uint64_t gridDim = num_tokens; - uint64_t blockDim = std::min(d / vec_size, 1024U); + dim3 gridDim, blockDim; + activation::act_and_mul_launch_dims(d, num_tokens, vec_size, out.get_device(), gridDim, + blockDim); activation::act_and_mul_kernel<<>>( static_cast(out.data_ptr()), static_cast(input.data_ptr()), d); diff --git a/flashinfer/jit/activation.py b/flashinfer/jit/activation.py index 47f4b0276f..91ffe8784d 100644 --- a/flashinfer/jit/activation.py +++ b/flashinfer/jit/activation.py @@ -28,14 +28,16 @@ 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) const c10::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard(out.device()); auto stream = at::hip::getCurrentHIPStream(); DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(input.scalar_type(), c_type, [&] { uint32_t vec_size = 16 / sizeof(c_type); - uint32_t block_size = std::max(1U, std::min(d / vec_size, 1024U)); - dim3 gridDim(num_tokens); - dim3 blockDim(block_size); + + dim3 gridDim, blockDim; + flashinfer::activation::act_and_mul_launch_dims(d, num_tokens, vec_size, + out.get_device(), gridDim, blockDim); auto kernel = flashinfer::activation::act_and_mul_kernel; diff --git a/include/flashinfer/attention/generic/activation.cuh b/include/flashinfer/attention/generic/activation.cuh index 68b6b8ec0c..cbcc1ea718 100644 --- a/include/flashinfer/attention/generic/activation.cuh +++ b/include/flashinfer/attention/generic/activation.cuh @@ -5,6 +5,9 @@ #ifndef FLASHINFER_ACTIVATION_CUH_ #define FLASHINFER_ACTIVATION_CUH_ +#include + +#include "gpu_iface/gpu_runtime_compat.hpp" #include "gpu_iface/math_ops.hpp" #include "gpu_iface/platform.hpp" #include "gpu_iface/utils.cuh" @@ -14,20 +17,60 @@ namespace flashinfer { using namespace gpu_iface::vec_dtypes; namespace activation { +// Adaptive launch config for act_and_mul_kernel. One block per token underfills +// the GPU when num_tokens is small (decode / small batch), so split each row +// across blocks_per_row blocks on gridDim.y until the total block count covers +// the CU array. For large num_tokens this resolves to blocks_per_row == 1, i.e. +// the original one-block-per-token launch. Single definition shared by the AOT +// launcher (flashinfer/csrc_rocm/activation.cu) and the JIT template +// (flashinfer/jit/activation.py) so the two paths cannot drift. +inline void act_and_mul_launch_dims(int d, int64_t num_tokens, uint32_t vec_size, int dev_id, + dim3& grid_dim, dim3& block_dim) { + uint32_t vecs = std::max(1U, (uint32_t)(d / vec_size)); + uint32_t block_size = std::max(1U, std::min(vecs, 1024U)); + // Oversubscribe CUs by 2x: enough to fill the GPU when num_tokens is small, + // without splitting rows once num_tokens already covers the CU array (extra + // splitting only adds launch/tail overhead — empirically bandwidth-neutral). + const uint32_t target_blocks = (uint32_t)getMultiProcessorCount(dev_id) * 2u; + const uint32_t max_bpr = ceil_div(vecs, block_size); + uint32_t blocks_per_row = 1u; + if ((uint64_t)num_tokens < target_blocks) { + const uint64_t nt = (uint64_t)std::max(1, num_tokens); + blocks_per_row = + std::max(1u, std::min((uint32_t)ceil_div(target_blocks, nt), max_bpr)); + } + grid_dim = dim3((unsigned)num_tokens, blocks_per_row, 1); + block_dim = dim3(block_size, 1, 1); +} + +// 2D grid: blockIdx.x selects the token (row), blockIdx.y selects a column-tile +// of that row. Output elements are independent (no cross-element reduction), so a +// row can be split across gridDim.y blocks with no atomics. When gridDim.y == 1 +// (e.g. any 1D launch, including the CUDA path) this collapses to one block per +// token with the same memory-access pattern as the original kernel. template __global__ void act_and_mul_kernel(T* __restrict__ out, const T* __restrict__ input, const int d) { constexpr uint32_t vec_size = 16 / sizeof(T); + // Row-base addresses are 64-bit (token_idx * 2 * d can exceed 2^31); the + // intra-row column index stays 32-bit to keep the inner-loop address math + // identical to the original one-block-per-token kernel (no 64-bit multiplies + // in the hot loop). col_block <= 65535, blockDim.x <= 1024 → products fit u32. const int64_t token_idx = blockIdx.x; - const int64_t thread_idx = threadIdx.x; - const int64_t stride = blockDim.x; - const int64_t offset = token_idx * 2 * d; + const int64_t offset = token_idx * 2 * d; // input row base (gate || up) + const int64_t out_base = token_idx * d; // output row base + const uint32_t col_block = blockIdx.y; // 0 when 1D-equivalent + const uint32_t num_col_blocks = gridDim.y; // 1 when 1D-equivalent + const uint32_t thread_idx = threadIdx.x; + const uint32_t col_stride = blockDim.x * num_col_blocks; // == blockDim.x when 1D + const uint32_t num_vec = d / vec_size; + const uint32_t vec_start = col_block * blockDim.x + thread_idx; #if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) asm volatile("griddepcontrol.wait;"); #endif #pragma unroll 1 - for (uint32_t idx = thread_idx; idx < d / vec_size; idx += stride) { + for (uint32_t idx = vec_start; idx < num_vec; idx += col_stride) { vec_t x_vec, y_vec, out_vec; x_vec.cast_load(input + offset + idx * vec_size); y_vec.cast_load(input + offset + d + idx * vec_size); @@ -35,16 +78,19 @@ __global__ void act_and_mul_kernel(T* __restrict__ out, const T* __restrict__ in for (uint32_t i = 0; i < vec_size; ++i) { out_vec[i] = Activation(x_vec[i]) * y_vec[i]; } - out_vec.cast_store(out + token_idx * d + idx * vec_size); + out_vec.cast_store(out + out_base + idx * vec_size); } - const int64_t remaining_offset = d - d % (stride * vec_size); - // process the remaining elements + // 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. + const uint32_t scalar_base = num_vec * vec_size; + const uint32_t scalar_count = (uint32_t)d - scalar_base; #pragma unroll 1 - for (int64_t idx = thread_idx; idx < d % (stride * vec_size); idx += stride) { - float x = input[offset + remaining_offset + idx], - y = input[offset + remaining_offset + d + idx]; - out[token_idx * d + remaining_offset + idx] = Activation(x) * y; + for (uint32_t s = vec_start; s < scalar_count; s += col_stride) { + const uint32_t e = scalar_base + s; + out[out_base + e] = Activation((float)input[offset + e]) * (float)input[offset + d + e]; } #if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) diff --git a/include/gpu_iface/gpu_runtime_compat.hpp b/include/gpu_iface/gpu_runtime_compat.hpp index 1ea64dbe83..ae71dd3930 100644 --- a/include/gpu_iface/gpu_runtime_compat.hpp +++ b/include/gpu_iface/gpu_runtime_compat.hpp @@ -135,6 +135,25 @@ } \ } while (0) +/// Returns the number of multiprocessors (CUs on CDNA / SMs on CUDA). +/// +/// Cached per device id to avoid repeating the attribute query on every kernel +/// launch. The cache is thread_local so concurrent callers (e.g. multi-threaded +/// Python) never race on it. 0 is treated as "not cached" — a valid count is +/// always > 0, so a device reporting 0 simply isn't memoized rather than +/// poisoning the cache. +/// +/// @param dev_id Device ID +/// @return Multiprocessor (CU/SM) count +inline int getMultiProcessorCount(int dev_id) { + static thread_local int cache[64] = {0}; + if (dev_id >= 0 && dev_id < 64 && cache[dev_id] > 0) return cache[dev_id]; + int count = 0; + FI_GPU_CALL(gpuDeviceGetAttribute(&count, gpuDevAttrMultiProcessorCount, dev_id)); + if (dev_id >= 0 && dev_id < 64 && count > 0) cache[dev_id] = count; + return count; +} + inline int getMaxSharedMemPerMultiprocessor(int dev_id) { int max_smem_per_sm = 0; #if defined(PLATFORM_CUDA_DEVICE) diff --git a/tests/rocm_tests/test_activation_hip.py b/tests/rocm_tests/test_activation_hip.py index 926c087350..ecfa667f5d 100644 --- a/tests/rocm_tests/test_activation_hip.py +++ b/tests/rocm_tests/test_activation_hip.py @@ -45,6 +45,8 @@ def _gelu_and_mul_ref(x: torch.Tensor) -> torch.Tensor: (16, 7168), # Mistral-7B ffn_dim // 2 (1, 8192), # hits 1024-thread cap (4, 14336), # Llama-3-70B ffn_dim // 2, multiple stride iterations + (2, 14336), # few tokens, large d → blocks_per_row > 1 (row split across blocks) + (4096, 4096), # prefill-scale token count → blocks_per_row == 1 (no row split) ]