perf(hip): adaptive 2D-grid row-split for act_and_mul kernels

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()

flashinfer/csrc_rocm/activation.cu

@@ -35,15 +35,17 @@ __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();
 
   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<c_type, silu><<<gridDim, blockDim, 0, stream>>>(
         static_cast<c_type*>(out.data_ptr()), static_cast<c_type*>(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<c_type, gelu_tanh><<<gridDim, blockDim, 0, stream>>>(
         static_cast<c_type*>(out.data_ptr()), static_cast<c_type*>(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<c_type, gelu><<<gridDim, blockDim, 0, stream>>>(
         static_cast<c_type*>(out.data_ptr()), static_cast<c_type*>(input.data_ptr()), d);

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<c_type, {{ act_func_name }}>;
 

include/flashinfer/attention/generic/activation.cuh

@@ -5,6 +5,9 @@
 #ifndef FLASHINFER_ACTIVATION_CUH_
 #define FLASHINFER_ACTIVATION_CUH_
 
+#include <algorithm>
+
+#include "gpu_iface/gpu_runtime_compat.hpp"
 #include "gpu_iface/math_ops.hpp"
 #include "gpu_iface/platform.hpp"
 #include "gpu_iface/utils.cuh"
@@ -14,37 +17,80 @@ 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<int64_t>(1, num_tokens);
+    blocks_per_row =
+        std::max(1u, std::min((uint32_t)ceil_div<uint64_t>(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 <typename T, float (*Activation)(const float&)>
 __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<float, vec_size> 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);
 #pragma unroll
     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))

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)

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)
 ]