[Kernels] Vectorize misaligned SiLU inputs with a per-SM cutoff

csrc/activation/silu_and_mul.cu

@@ -2,6 +2,7 @@
 #include "../ops.h"
 
 #include <ATen/core/grad_mode.h>
+#include <ATen/cuda/CUDAContext.h>
 #include <c10/cuda/CUDAGuard.h>
 #include <c10/cuda/CUDAException.h>
 #include <c10/cuda/CUDAStream.h>
@@ -77,28 +78,94 @@ __device__ __forceinline__ T compute_scalar(T gate, T up) {
   return Math<T>::mul(activated, up);
 }
 
-template <typename T, bool vectorized>
+template <typename T, bool aligned>
+__device__ __forceinline__ Vector<T> load_vector(const T* input, int64_t index,
+                                                int64_t count) {
+  constexpr int elements = 16 / sizeof(T);
+  const int offset_bytes = aligned ? 0 : reinterpret_cast<uintptr_t>(input + index) & 15;
+  Vector<T> result;
+  if (offset_bytes == 0) {
+    *reinterpret_cast<uint4*>(&result) =
+        *reinterpret_cast<const uint4*>(input + index);
+    return result;
+  }
+
+  const int offset = offset_bytes / sizeof(T);
+  const int64_t lower = index - offset;
+  Vector<T> head{}, tail{};
+  // Never read before or after the input view, even if its storage is larger.
+  if (lower >= 0) {
+    *reinterpret_cast<uint4*>(&head) = *reinterpret_cast<const uint4*>(input + lower);
+  } else {
+#pragma unroll
+    for (int j = 0; j < elements; ++j) {
+      if (j >= offset) reinterpret_cast<T*>(&head)[j] = input[lower + j];
+    }
+  }
+  if (lower + 2 * elements <= count) {
+    *reinterpret_cast<uint4*>(&tail) =
+        *reinterpret_cast<const uint4*>(input + lower + elements);
+  } else {
+#pragma unroll
+    for (int j = 0; j < elements; ++j) {
+      if (j < offset) reinterpret_cast<T*>(&tail)[j] = input[lower + elements + j];
+    }
+  }
+
+  // Reassemble the requested logical vector with register moves and bit shifts.
+  uint4 lo = *reinterpret_cast<const uint4*>(&head);
+  uint4 hi = *reinterpret_cast<const uint4*>(&tail);
+  if (offset_bytes & 8) {
+    lo = make_uint4(lo.z, lo.w, hi.x, hi.y);
+    hi = make_uint4(hi.z, hi.w, 0, 0);
+  }
+  if (offset_bytes & 4) {
+    lo = make_uint4(lo.y, lo.z, lo.w, hi.x);
+    hi.x = hi.y;
+  }
+  if constexpr (sizeof(T) == 2) {
+    if (offset_bytes & 2) {
+      lo = make_uint4(__funnelshift_r(lo.x, lo.y, 16), __funnelshift_r(lo.y, lo.z, 16),
+                      __funnelshift_r(lo.z, lo.w, 16), __funnelshift_r(lo.w, hi.x, 16));
+    }
+  }
+  *reinterpret_cast<uint4*>(&result) = lo;
+  return result;
+}
+
+enum class AccessMode { Aligned, Misaligned, Scalar };
+
+template <typename T, AccessMode mode>
 __global__ void silu_and_mul_kernel(const T* gate, const T* up, T* out, int64_t count) {
   const int64_t thread = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
   const int64_t stride = static_cast<int64_t>(gridDim.x) * blockDim.x;
   int64_t tail_start = 0;
-  if constexpr (vectorized) {
+  if constexpr (mode != AccessMode::Scalar) {
     constexpr int elements = 16 / sizeof(T);
     constexpr int pairs = 16 / sizeof(typename Math<T>::Pair);
-    const int64_t vectors = count / elements;
+    int64_t prefix = 0;
+    if constexpr (mode == AccessMode::Misaligned) {
+      prefix = min(count, static_cast<int64_t>(
+          ((16 - (reinterpret_cast<uintptr_t>(out) & 15)) & 15) / sizeof(T)));
+    }
+    for (int64_t i = thread; i < prefix; i += stride) {
+      out[i] = compute_scalar(gate[i], up[i]);
+    }
+    const int64_t vectors = (count - prefix) / elements;
     for (int64_t i = thread; i < vectors; i += stride) {
-      Vector<T> g, u, result;
-      *reinterpret_cast<uint4*>(&g) = *reinterpret_cast<const uint4*>(gate + i * elements);
-      *reinterpret_cast<uint4*>(&u) = *reinterpret_cast<const uint4*>(up + i * elements);
+      const int64_t index = prefix + i * elements;
+      Vector<T> g = load_vector<T, mode == AccessMode::Aligned>(gate, index, count);
+      Vector<T> u = load_vector<T, mode == AccessMode::Aligned>(up, index, count);
+      Vector<T> result;
 #pragma unroll
       for (int j = 0; j < pairs; ++j) {
         result.pairs[j] = Math<T>::compute(g.pairs[j], u.pairs[j]);
       }
-      *reinterpret_cast<uint4*>(out + i * elements) = *reinterpret_cast<uint4*>(&result);
+      *reinterpret_cast<uint4*>(out + index) = *reinterpret_cast<const uint4*>(&result);
     }
-    tail_start = vectors * elements;
+    tail_start = prefix + vectors * elements;
   }
-  // Only valid scalar addresses are touched, including for count < one vector.
+  // Scalar mode starts at zero; vector modes handle only their remaining tail.
   for (int64_t i = tail_start + thread; i < count; i += stride) {
     out[i] = compute_scalar(gate[i], up[i]);
   }
@@ -114,8 +181,19 @@ void launch(const at::Tensor& gate, const at::Tensor& up, at::Tensor& out) {
                             reinterpret_cast<uintptr_t>(u) |
                             reinterpret_cast<uintptr_t>(result)) & 15) == 0;
   const int64_t count = gate.numel();
-  const int64_t work = vectorized ? std::max<int64_t>(count / (16 / sizeof(T)), 1) : count;
   constexpr int threads = 256;
+  bool scalar = false;
+  if (!vectorized) {
+    // Use 160 reference vector blocks per SM to determine the element cutoff.
+    // The actual launch grid is calculated separately and may be capped.
+    constexpr int reference_blocks_per_sm = 160;
+    constexpr int elements_per_vector = 16 / sizeof(T);
+    const int sms = at::cuda::getDeviceProperties(gate.get_device())->multiProcessorCount;
+    const int64_t threshold =
+        static_cast<int64_t>(sms) * reference_blocks_per_sm * threads * elements_per_vector;
+    scalar = count < threshold;
+  }
+  const int64_t work = scalar ? count : std::max<int64_t>(count / (16 / sizeof(T)), 1);
   int max_blocks = 4096;
   if constexpr (std::is_same_v<T, float>) {
     // More blocks improve large FP32 streaming workloads. Require enough work
@@ -128,9 +206,11 @@ void launch(const at::Tensor& gate, const at::Tensor& up, at::Tensor& out) {
   const int blocks = static_cast<int>(std::min<int64_t>((work + threads - 1) / threads, max_blocks));
   auto stream = c10::cuda::getCurrentCUDAStream(gate.get_device());
   if (vectorized) {
-    silu_and_mul_kernel<T, true><<<blocks, threads, 0, stream>>>(g, u, result, count);
+    silu_and_mul_kernel<T, AccessMode::Aligned><<<blocks, threads, 0, stream>>>(g, u, result, count);
+  } else if (scalar) {
+    silu_and_mul_kernel<T, AccessMode::Scalar><<<blocks, threads, 0, stream>>>(g, u, result, count);
   } else {
-    silu_and_mul_kernel<T, false><<<blocks, threads, 0, stream>>>(g, u, result, count);
+    silu_and_mul_kernel<T, AccessMode::Misaligned><<<blocks, threads, 0, stream>>>(g, u, result, count);
   }
   C10_CUDA_KERNEL_LAUNCH_CHECK();
 }

docs/kernels.md

@@ -38,10 +38,31 @@ nor copies incompatible inputs.
 when the extension is installed; CPU inputs and installations without the
 extension use the Torch reference. Import it from `vllm_rlt.kernels`. The model
 call site does not check devices or extension availability. The native kernel
-flattens contiguous inputs, using 16-byte
-loads/stores for the aligned vector body and scalar accesses for the remaining
-elements in the same kernel. Contiguous views with misaligned storage offsets
-use scalar accesses. Empty inputs return an empty output without launching.
+flattens contiguous inputs, using 16-byte loads/stores for the aligned vector
+body and scalar accesses for the remaining elements in the same kernel.
+Contiguous views may have independently misaligned storage offsets. For these
+inputs, the scalar/vector cutoff scales with the input GPU's actual SM count:
+`SMs * 160 * 256 * (16 / sizeof(dtype))` total elements. The reference is 160
+vector blocks per SM, with 256 threads and 16 bytes per vector: eight FP16/BF16
+elements or four FP32 elements. On a 148-SM B300, the corresponding cutoffs are
+48,496,640 FP16/BF16 elements and 24,248,320 FP32 elements. The reference block
+count specifies the workload cutoff; the actual launch grid is computed below
+and may be capped. This threshold is a heuristic, not a measured optimum.
+
+Below the cutoff, misaligned inputs use scalar computation with one work item
+per element. At or above it, the kernel loads each input's adjacent aligned
+vectors and reassembles the required logical vector with register moves and bit
+shifts. Scalar loads handle input prefixes and suffixes where a full aligned
+read would extend outside the view. Aligned and misaligned vector paths share the
+same calculation and output loop; output prefix and tail use scalar
+computation. No input copy or extra kernel launch is needed. Fully aligned inputs
+retain their direct vector loads even for small inputs. Empty inputs return an
+empty output without launching. On B300, correctness was tested for all
+independent 16-byte input offsets just above the cutoff and for every vector
+remainder with selected offsets. The cutoff applies to total element count,
+independent of tensor rank or model width. Raising it to 160 reference blocks
+per SM also sends many previously vectorized inputs through the slower scalar
+path; the performance impact depends on the workload.
 
 SiLU uses FP32 arithmetic and rounds to the input dtype before multiplying by
 `up`, preserving the reference's intermediate rounding. FP16/BF16 vector paths
@@ -55,10 +76,11 @@ stream.
 
 Blocks have 256 threads. The grid is capped at 4096 blocks, except for FP32
 workloads with at least `16384 * 256` work items: these use up to 16384 blocks.
-Each aligned work item processes four FP32 elements; an unaligned work item
-processes one scalar. The threshold requires enough work to fill the larger
-grid. This keeps small and medium shapes' launch geometry unchanged. Grid sizes,
-16/32-byte transfers and scalar/packed FP32 products were compared on B300;
+Each vector work item processes four FP32 elements or eight FP16/BF16 elements;
+the scalar path for small misaligned inputs processes one element per work item.
+The streaming threshold requires
+enough work to fill the larger grid. Aligned inputs retain their existing launch
+geometry. Grid sizes, 16/32-byte transfers and scalar/packed FP32 products were compared on B300;
 the selected configuration is not a universal hardware optimum.
 
 Ouro's inference MLP calls this registered operator directly. Availability is