From 3c42d1293c337d81e232a4090aa54e61f6abd877 Mon Sep 17 00:00:00 2001
From: bjf-frz <240110042+bjf-frz@users.noreply.github.com>
Date: Mon, 2 Nov 2026 12:48:46 +0000
Subject: [PATCH 1/6] perf(kernels): vectorize misaligned SiLU inputs with a
 small-input fallback

---
 csrc/activation/silu_and_mul.cu | 98 ++++++++++++++++++++++++++++++++-
 docs/kernels.md                 | 29 +++++++---
 tests/test_activation.py        | 78 ++++++++++++++++++++++----
 3 files changed, 183 insertions(+), 22 deletions(-)

diff --git a/csrc/activation/silu_and_mul.cu b/csrc/activation/silu_and_mul.cu
index faadcda1..75616429 100644
--- a/csrc/activation/silu_and_mul.cu
+++ b/csrc/activation/silu_and_mul.cu
@@ -77,12 +77,67 @@ __device__ __forceinline__ T compute_scalar(T gate, T up) {
   return Math<T>::mul(activated, up);
 }
 
-template <typename T, bool vectorized>
+template <typename T>
+__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 = 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;
+}
+
+template <typename T, bool aligned>
 __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 (aligned) {
     constexpr int elements = 16 / sizeof(T);
     constexpr int pairs = 16 / sizeof(typename Math<T>::Pair);
     const int64_t vectors = count / elements;
@@ -97,6 +152,27 @@ __global__ void silu_and_mul_kernel(const T* gate, const T* up, T* out, int64_t
       *reinterpret_cast<uint4*>(out + i * elements) = *reinterpret_cast<uint4*>(&result);
     }
     tail_start = vectors * elements;
+  } else {
+    constexpr int elements = 16 / sizeof(T);
+    constexpr int pairs = 16 / sizeof(typename Math<T>::Pair);
+    const int64_t 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) {
+      const int64_t index = prefix + i * elements;
+      Vector<T> g = load_vector(gate, index, count);
+      Vector<T> u = load_vector(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 + index) = *reinterpret_cast<const uint4*>(&result);
+    }
+    tail_start = prefix + vectors * elements;
   }
   // Only valid scalar addresses are touched, including for count < one vector.
   for (int64_t i = tail_start + thread; i < count; i += stride) {
@@ -104,6 +180,16 @@ __global__ void silu_and_mul_kernel(const T* gate, const T* up, T* out, int64_t
   }
 }
 
+template <typename T>
+__global__ void silu_and_mul_scalar_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;
+  for (int64_t i = thread; i < count; i += stride) {
+    out[i] = compute_scalar(gate[i], up[i]);
+  }
+}
+
 template <typename T>
 void launch(const at::Tensor& gate, const at::Tensor& up, at::Tensor& out) {
   const T* g = reinterpret_cast<const T*>(gate.const_data_ptr());
@@ -114,7 +200,11 @@ 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;
+  // B300 measurements including independently shifted inputs and all vector
+  // remainders favor scalar access below this total element count.
+  constexpr int64_t misaligned_vector_threshold = 352 * 1024;
+  const bool scalar = !vectorized && count < misaligned_vector_threshold;
+  const int64_t work = scalar ? count : std::max<int64_t>(count / (16 / sizeof(T)), 1);
   constexpr int threads = 256;
   int max_blocks = 4096;
   if constexpr (std::is_same_v<T, float>) {
@@ -129,6 +219,8 @@ void launch(const at::Tensor& gate, const at::Tensor& up, at::Tensor& out) {
   auto stream = c10::cuda::getCurrentCUDAStream(gate.get_device());
   if (vectorized) {
     silu_and_mul_kernel<T, true><<<blocks, threads, 0, stream>>>(g, u, result, count);
+  } else if (scalar) {
+    silu_and_mul_scalar_kernel<T><<<blocks, threads, 0, stream>>>(g, u, result, count);
   } else {
     silu_and_mul_kernel<T, false><<<blocks, threads, 0, stream>>>(g, u, result, count);
   }
diff --git a/docs/kernels.md b/docs/kernels.md
index a1fda4e1..ffe0788d 100644
--- a/docs/kernels.md
+++ b/docs/kernels.md
@@ -38,10 +38,22 @@ 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. Below
+`352 * 1024` total elements, these inputs use a scalar kernel with one work item
+per element. At or above this cutoff, 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. The output is processed in aligned vectors with a scalar
+tail; a scalar prefix also handles a misaligned output address. No input copy
+or extra kernel launch is needed. Fully aligned inputs retain their direct
+vector load/store path even for small inputs. Empty inputs return an empty output
+without launching. The cutoff is based on B300 measurements across all three
+dtypes, independent input offsets and vector remainders. It applies to total
+element count, independent of tensor rank or model width, and is not a universal
+hardware optimum.
 
 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 +67,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
diff --git a/tests/test_activation.py b/tests/test_activation.py
index dbedb464..f1e69079 100644
--- a/tests/test_activation.py
+++ b/tests/test_activation.py
@@ -160,7 +160,54 @@ def test_silu_and_mul_cuda_vectors_tails_and_storage_offsets(dtype, shape, offse
 
 @pytest.mark.gpu
 @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
-def test_silu_and_mul_cuda_special_values_and_underflow(dtype):
+def test_silu_and_mul_cuda_all_alignment_pairs_match_aligned_path(dtype):
+    # The two inputs can have different offsets. Cover every 16-byte alignment
+    # pair, including vectors adjacent to either view boundary and short inputs.
+    elements = 16 // torch.empty((), dtype=dtype).element_size()
+    counts = [*range(2 * elements + 2), 17 * elements - 1, 17 * elements + 1]
+    integer = torch.int32 if dtype == torch.float32 else torch.int16
+    generator = torch.Generator(device="cuda").manual_seed(42)
+    for count in counts:
+        for gate_offset in range(elements):
+            for up_offset in range(elements):
+                gate = torch.randn(
+                    count + gate_offset, device="cuda", dtype=dtype, generator=generator
+                )[gate_offset:]
+                up = torch.randn(
+                    count + up_offset, device="cuda", dtype=dtype, generator=generator
+                )[up_offset:]
+                expected = silu_and_mul(gate.clone(), up.clone())
+                actual = silu_and_mul(gate, up)
+                assert torch.equal(actual.view(integer), expected.view(integer)), (
+                    dtype,
+                    count,
+                    gate_offset,
+                    up_offset,
+                )
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("count", [352 * 1024 - 1, 352 * 1024, 352 * 1024 + 1, 352 * 1024 + 7])
+def test_silu_and_mul_cuda_threshold_and_vector_offsets_match_aligned_path(dtype, count):
+    # Short-offset tests now take the scalar fallback. Exercise all independent
+    # input alignments on both sides of the dispatch threshold and scalar tails.
+    elements = 16 // torch.empty((), dtype=dtype).element_size()
+    integer = torch.int32 if dtype == torch.float32 else torch.int16
+    generator = torch.Generator(device="cuda").manual_seed(43)
+    for go in range(elements):
+        for uo in range(elements):
+            gate = torch.randn(count + go, device="cuda", dtype=dtype, generator=generator)[go:]
+            up = torch.randn(count + uo, device="cuda", dtype=dtype, generator=generator)[uo:]
+            expected = silu_and_mul(gate.clone(), up.clone())
+            actual = silu_and_mul(gate, up)
+            assert torch.equal(actual.view(integer), expected.view(integer)), (dtype, count, go, uo)
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("offsets", [(0, 0), (1, 3)])
+def test_silu_and_mul_cuda_special_values_and_underflow(dtype, offsets):
     info = torch.finfo(dtype)
     values = [
         float("nan"),
@@ -184,6 +231,13 @@ def test_silu_and_mul_cuda_special_values_and_underflow(dtype):
         indexing="ij",
     )
     gate, up = gate.contiguous(), up.contiguous()
+    if offsets != (0, 0):
+        count = gate.numel()
+        shifted_gate = torch.empty(count + offsets[0], device="cuda", dtype=dtype)[offsets[0] :]
+        shifted_up = torch.empty(count + offsets[1], device="cuda", dtype=dtype)[offsets[1] :]
+        shifted_gate.copy_(gate.flatten())
+        shifted_up.copy_(up.flatten())
+        gate, up = shifted_gate.view(gate.shape), shifted_up.view(up.shape)
     torch.testing.assert_close(
         silu_and_mul(gate, up), silu_and_mul_reference(gate, up), equal_nan=True
     )
@@ -210,11 +264,13 @@ def test_silu_and_mul_cuda_rejects_autograd_but_allows_inference():
 
 @pytest.mark.gpu
 @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
-def test_silu_and_mul_cuda_current_stream_and_graph_replay(dtype):
+@pytest.mark.parametrize("offsets", [(0, 0), (1, 3)])
+@pytest.mark.parametrize("count", [4097, 352 * 1024 + 1])
+def test_silu_and_mul_cuda_current_stream_and_graph_replay(dtype, offsets, count):
     stream = torch.cuda.Stream()
     with torch.cuda.stream(stream):
-        gate = torch.randn(4097, device="cuda", dtype=dtype)
-        up = torch.randn_like(gate)
+        gate = torch.randn(count + offsets[0], device="cuda", dtype=dtype)[offsets[0] :]
+        up = torch.randn(count + offsets[1], device="cuda", dtype=dtype)[offsets[1] :]
         for _ in range(3):
             silu_and_mul(gate, up)
     stream.synchronize()
@@ -238,9 +294,10 @@ def test_silu_and_mul_cuda_fp32_streaming_grid_boundary(count):
 
 @pytest.mark.gpu
 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
-@pytest.mark.parametrize("offset", [0, 1])
-def test_silu_and_mul_cuda_low_precision_storage_patterns(dtype, offset):
-    # Exercise all gate bit patterns with both vector/tail and scalar-only access.
+@pytest.mark.parametrize("offsets", [(0, 0), (1, 0), (1, 3), (7, 5)])
+def test_silu_and_mul_cuda_low_precision_storage_patterns(dtype, offsets):
+    # Exercise all gate bit patterns through aligned and independently shifted
+    # vector loads, including subnormal products at the scalar tail.
     values = torch.arange(65536, device="cuda", dtype=torch.int32).to(torch.int16).view(dtype)
     info = torch.finfo(dtype)
     subnormal = info.tiny / (1024 if dtype == torch.float16 else 128)
@@ -265,10 +322,9 @@ def test_silu_and_mul_cuda_low_precision_storage_patterns(dtype, offset):
         device="cuda",
         dtype=dtype,
     )
-    gate = torch.empty(values.numel() * factors.numel() + 1 + offset, device="cuda", dtype=dtype)[
-        offset:
-    ]
-    up = torch.empty_like(gate)
+    count = values.numel() * factors.numel() + 1
+    gate = torch.empty(count + offsets[0], device="cuda", dtype=dtype)[offsets[0] :]
+    up = torch.empty(count + offsets[1], device="cuda", dtype=dtype)[offsets[1] :]
     gate[:-1].copy_(values.repeat(factors.numel()))
     up[:-1].copy_(factors.repeat_interleave(values.numel()))
     # The aligned path ends with one scalar multiplication yielding a subnormal.

From b5b44f694ecb38a887757411ff1fdc3d65812fa5 Mon Sep 17 00:00:00 2001
From: bjf-frz <frz123db@gmail.com>
Date: Tue, 3 Nov 2026 09:21:52 +0000
Subject: [PATCH 2/6] perf(kernels): scale SiLU vector cutoff with GPU size

---
 csrc/activation/silu_and_mul.cu | 73 +++++++++++++----------------
 docs/kernels.md                 | 33 +++++++------
 tests/test_activation.py        | 82 +++++++++++++++++++++++++++++++--
 3 files changed, 129 insertions(+), 59 deletions(-)

diff --git a/csrc/activation/silu_and_mul.cu b/csrc/activation/silu_and_mul.cu
index 75616429..b92e25e3 100644
--- a/csrc/activation/silu_and_mul.cu
+++ b/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,11 +78,11 @@ __device__ __forceinline__ T compute_scalar(T gate, T up) {
   return Math<T>::mul(activated, up);
 }
 
-template <typename T>
+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 = reinterpret_cast<uintptr_t>(input + index) & 15;
+  const int offset_bytes = aligned ? 0 : reinterpret_cast<uintptr_t>(input + index) & 15;
   Vector<T> result;
   if (offset_bytes == 0) {
     *reinterpret_cast<uint4*>(&result) =
@@ -132,39 +133,29 @@ __device__ __forceinline__ Vector<T> load_vector(const T* input, int64_t index,
   return result;
 }
 
-template <typename T, bool aligned>
+enum class AccessMode { Aligned, Reconstruct, 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 (aligned) {
+  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;
-    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);
-#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);
+    int64_t prefix = 0;
+    if constexpr (mode == AccessMode::Reconstruct) {
+      prefix = min(count, static_cast<int64_t>(
+          ((16 - (reinterpret_cast<uintptr_t>(out) & 15)) & 15) / sizeof(T)));
     }
-    tail_start = vectors * elements;
-  } else {
-    constexpr int elements = 16 / sizeof(T);
-    constexpr int pairs = 16 / sizeof(typename Math<T>::Pair);
-    const int64_t 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) {
       const int64_t index = prefix + i * elements;
-      Vector<T> g = load_vector(gate, index, count);
-      Vector<T> u = load_vector(up, index, count);
+      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) {
@@ -174,22 +165,12 @@ __global__ void silu_and_mul_kernel(const T* gate, const T* up, T* out, int64_t
     }
     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]);
   }
 }
 
-template <typename T>
-__global__ void silu_and_mul_scalar_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;
-  for (int64_t i = thread; i < count; i += stride) {
-    out[i] = compute_scalar(gate[i], up[i]);
-  }
-}
-
 template <typename T>
 void launch(const at::Tensor& gate, const at::Tensor& up, at::Tensor& out) {
   const T* g = reinterpret_cast<const T*>(gate.const_data_ptr());
@@ -200,12 +181,22 @@ 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();
-  // B300 measurements including independently shifted inputs and all vector
-  // remainders favor scalar access below this total element count.
-  constexpr int64_t misaligned_vector_threshold = 352 * 1024;
-  const bool scalar = !vectorized && count < misaligned_vector_threshold;
-  const int64_t work = scalar ? count : std::max<int64_t>(count / (16 / sizeof(T)), 1);
   constexpr int threads = 256;
+  bool scalar = false;
+  if (!vectorized) {
+    // Use 160 BF16-vector blocks on the measured 148-SM B300 as the reference.
+    // Scale that workload per SM and use the same element cutoff for all
+    // dtypes; FP32 still processes four elements per 16-byte vector.
+    constexpr int reference_sms = 148;
+    constexpr int reference_blocks = 160;
+    constexpr int reference_elements = 16 / sizeof(__nv_bfloat16);
+    const int sms = at::cuda::getDeviceProperties(gate.get_device())->multiProcessorCount;
+    const int64_t target_blocks =
+        (static_cast<int64_t>(sms) * reference_blocks + reference_sms - 1) / reference_sms;
+    const int64_t threshold = target_blocks * threads * reference_elements;
+    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
@@ -218,11 +209,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_scalar_kernel<T><<<blocks, threads, 0, stream>>>(g, u, result, count);
+    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::Reconstruct><<<blocks, threads, 0, stream>>>(g, u, result, count);
   }
   C10_CUDA_KERNEL_LAUNCH_CHECK();
 }
diff --git a/docs/kernels.md b/docs/kernels.md
index ffe0788d..032c22c1 100644
--- a/docs/kernels.md
+++ b/docs/kernels.md
@@ -40,20 +40,27 @@ 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 may have independently misaligned storage offsets. Below
-`352 * 1024` total elements, these inputs use a scalar kernel with one work item
-per element. At or above this cutoff, 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. The output is processed in aligned vectors with a scalar
-tail; a scalar prefix also handles a misaligned output address. No input copy
-or extra kernel launch is needed. Fully aligned inputs retain their direct
-vector load/store path even for small inputs. Empty inputs return an empty output
-without launching. The cutoff is based on B300 measurements across all three
+Contiguous views may have independently misaligned storage offsets. For these
+inputs, the scalar/vector cutoff scales with the input GPU's SM count:
+`ceil(SMs * 160 / 148) * 256 * 8` total elements. This uses a B300 benchmark
+reference of 160 BF16 vector blocks on 148 SMs, 256 threads per block and eight
+BF16 elements per 16-byte vector. The eight-element reference is also used for
+FP32 so that the B300 cutoff is `320 * 1024` elements for all three dtypes; an
+FP32 vector itself still contains four elements. Some FP16/BF16 inputs with a
+scalar tail near the cutoff were approximately 2.7% slower than the scalar path
+on B300. Scaling from B300 is a heuristic, not a measured optimum on other GPUs.
+
+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 reconstructed inputs share the
+same vector 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. The B300 cutoff was tested across all three
 dtypes, independent input offsets and vector remainders. It applies to total
-element count, independent of tensor rank or model width, and is not a universal
-hardware optimum.
+element count, independent of tensor rank or model width.
 
 SiLU uses FP32 arithmetic and rounds to the input dtype before multiplying by
 `up`, preserving the reference's intermediate rounding. FP16/BF16 vector paths
diff --git a/tests/test_activation.py b/tests/test_activation.py
index f1e69079..762fc00d 100644
--- a/tests/test_activation.py
+++ b/tests/test_activation.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 """Numerics, tails, layouts and stream contracts for SiLU-and-Mul."""
 
+import json
 import math
 import subprocess
 import sys
@@ -162,7 +163,7 @@ def test_silu_and_mul_cuda_vectors_tails_and_storage_offsets(dtype, shape, offse
 @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
 def test_silu_and_mul_cuda_all_alignment_pairs_match_aligned_path(dtype):
     # The two inputs can have different offsets. Cover every 16-byte alignment
-    # pair, including vectors adjacent to either view boundary and short inputs.
+    # pair in the scalar fallback, including zero and short input lengths.
     elements = 16 // torch.empty((), dtype=dtype).element_size()
     counts = [*range(2 * elements + 2), 17 * elements - 1, 17 * elements + 1]
     integer = torch.int32 if dtype == torch.float32 else torch.int16
@@ -188,10 +189,12 @@ def test_silu_and_mul_cuda_all_alignment_pairs_match_aligned_path(dtype):
 
 @pytest.mark.gpu
 @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
-@pytest.mark.parametrize("count", [352 * 1024 - 1, 352 * 1024, 352 * 1024 + 1, 352 * 1024 + 7])
-def test_silu_and_mul_cuda_threshold_and_vector_offsets_match_aligned_path(dtype, count):
+@pytest.mark.parametrize("delta", range(-1, 8))
+def test_silu_and_mul_cuda_threshold_and_vector_offsets_match_aligned_path(dtype, delta):
     # Short-offset tests now take the scalar fallback. Exercise all independent
     # input alignments on both sides of the dispatch threshold and scalar tails.
+    sms = torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count
+    count = math.ceil(sms * 160 / 148) * 256 * 8 + delta
     elements = 16 // torch.empty((), dtype=dtype).element_size()
     integer = torch.int32 if dtype == torch.float32 else torch.int16
     generator = torch.Generator(device="cuda").manual_seed(43)
@@ -204,6 +207,53 @@ def test_silu_and_mul_cuda_threshold_and_vector_offsets_match_aligned_path(dtype
             assert torch.equal(actual.view(integer), expected.view(integer)), (dtype, count, go, uo)
 
 
+@pytest.mark.gpu
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_silu_and_mul_cuda_uses_device_scaled_threshold(dtype, tmp_path):
+    # Check actual launch geometry on both sides, including FP32: B300's
+    # measured cutoff is in elements, rather than FP32 vector work items.
+    sms = torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count
+    threshold = math.ceil(sms * 160 / 148) * 256 * 8
+    elements = 16 // torch.empty((), dtype=dtype).element_size()
+    cases = [(threshold - 1, 1, 3), (threshold, 1, 3), (threshold, 0, 0)]
+    tensors = [
+        (
+            torch.randn(count + go, device="cuda", dtype=dtype)[go:],
+            torch.randn(count + uo, device="cuda", dtype=dtype)[uo:],
+        )
+        for count, go, uo in cases
+    ]
+    for gate, up in tensors:
+        silu_and_mul(gate, up)
+    torch.cuda.synchronize()
+    with torch.profiler.profile(
+        activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA]
+    ) as prof:
+        for gate, up in tensors:
+            silu_and_mul(gate, up)
+        torch.cuda.synchronize()
+    trace = tmp_path / "silu-dispatch.json"
+    prof.export_chrome_trace(str(trace))
+    events = json.loads(trace.read_text())["traceEvents"]
+    kernels = sorted(
+        (
+            event
+            for event in events
+            if event.get("cat") == "kernel" and "silu_and_mul_kernel" in event.get("name", "")
+        ),
+        key=lambda event: event["ts"],
+    )
+    assert len(kernels) == len(cases)
+    scalar_blocks = min(math.ceil((threshold - 1) / 256), 4096)
+    vector_blocks = min(math.ceil((threshold // elements) / 256), 4096)
+    assert [event["args"]["grid"][0] for event in kernels] == [
+        scalar_blocks,
+        vector_blocks,
+        vector_blocks,
+    ]
+    assert all(event["args"]["block"] == [256, 1, 1] for event in kernels)
+
+
 @pytest.mark.gpu
 @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
 @pytest.mark.parametrize("offsets", [(0, 0), (1, 3)])
@@ -265,8 +315,11 @@ def test_silu_and_mul_cuda_rejects_autograd_but_allows_inference():
 @pytest.mark.gpu
 @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
 @pytest.mark.parametrize("offsets", [(0, 0), (1, 3)])
-@pytest.mark.parametrize("count", [4097, 352 * 1024 + 1])
-def test_silu_and_mul_cuda_current_stream_and_graph_replay(dtype, offsets, count):
+@pytest.mark.parametrize("large", [False, True])
+def test_silu_and_mul_cuda_current_stream_and_graph_replay(dtype, offsets, large):
+    sms = torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count
+    threshold = math.ceil(sms * 160 / 148) * 256 * 8
+    count = threshold + 1 if large else min(4097, threshold - 1)
     stream = torch.cuda.Stream()
     with torch.cuda.stream(stream):
         gate = torch.randn(count + offsets[0], device="cuda", dtype=dtype)[offsets[0] :]
@@ -283,6 +336,25 @@ def test_silu_and_mul_cuda_current_stream_and_graph_replay(dtype, offsets, count
     torch.testing.assert_close(out, silu_and_mul_reference(gate, up))
 
 
+@pytest.mark.gpu
+def test_silu_and_mul_first_misaligned_call_during_graph_capture():
+    # The device-property lookup used by dispatch must also work on first use.
+    script = textwrap.dedent("""
+        import torch
+        from vllm_rlt.kernels.activation import silu_and_mul, silu_and_mul_reference
+        n = 4097
+        gate = torch.randn(n + 1, device='cuda', dtype=torch.bfloat16)[1:]
+        up = torch.randn(n + 3, device='cuda', dtype=torch.bfloat16)[3:]
+        torch.cuda.synchronize()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            output = silu_and_mul(gate, up)
+        graph.replay()
+        torch.testing.assert_close(output, silu_and_mul_reference(gate, up), rtol=0, atol=0)
+    """)
+    subprocess.run([sys.executable, "-c", script], check=True, capture_output=True, timeout=60)
+
+
 @pytest.mark.gpu
 @pytest.mark.parametrize("count", [(1 << 24) - 1, 1 << 24, (1 << 24) + 1])
 def test_silu_and_mul_cuda_fp32_streaming_grid_boundary(count):

From fe71d62de595df89d0d7ebf0dbccbbd837c3dfa6 Mon Sep 17 00:00:00 2001
From: bjf-frz <frz123db@gmail.com>
Date: Tue, 3 Nov 2026 09:41:49 +0000
Subject: [PATCH 3/6] perf(kernels): use per-SM vectorization threshold

---
 csrc/activation/silu_and_mul.cu | 15 +++-----
 docs/kernels.md                 | 24 ++++++------
 tests/test_activation.py        | 66 ++++++++++++++++++++++-----------
 3 files changed, 63 insertions(+), 42 deletions(-)

diff --git a/csrc/activation/silu_and_mul.cu b/csrc/activation/silu_and_mul.cu
index b92e25e3..3a5b9d13 100644
--- a/csrc/activation/silu_and_mul.cu
+++ b/csrc/activation/silu_and_mul.cu
@@ -184,16 +184,13 @@ void launch(const at::Tensor& gate, const at::Tensor& up, at::Tensor& out) {
   constexpr int threads = 256;
   bool scalar = false;
   if (!vectorized) {
-    // Use 160 BF16-vector blocks on the measured 148-SM B300 as the reference.
-    // Scale that workload per SM and use the same element cutoff for all
-    // dtypes; FP32 still processes four elements per 16-byte vector.
-    constexpr int reference_sms = 148;
-    constexpr int reference_blocks = 160;
-    constexpr int reference_elements = 16 / sizeof(__nv_bfloat16);
+    // 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 target_blocks =
-        (static_cast<int64_t>(sms) * reference_blocks + reference_sms - 1) / reference_sms;
-    const int64_t threshold = target_blocks * threads * reference_elements;
+    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);
diff --git a/docs/kernels.md b/docs/kernels.md
index 032c22c1..fc1f1d22 100644
--- a/docs/kernels.md
+++ b/docs/kernels.md
@@ -41,14 +41,13 @@ 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 may have independently misaligned storage offsets. For these
-inputs, the scalar/vector cutoff scales with the input GPU's SM count:
-`ceil(SMs * 160 / 148) * 256 * 8` total elements. This uses a B300 benchmark
-reference of 160 BF16 vector blocks on 148 SMs, 256 threads per block and eight
-BF16 elements per 16-byte vector. The eight-element reference is also used for
-FP32 so that the B300 cutoff is `320 * 1024` elements for all three dtypes; an
-FP32 vector itself still contains four elements. Some FP16/BF16 inputs with a
-scalar tail near the cutoff were approximately 2.7% slower than the scalar path
-on B300. Scaling from B300 is a heuristic, not a measured optimum on other GPUs.
+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
@@ -58,9 +57,12 @@ read would extend outside the view. Aligned and reconstructed inputs share the
 same vector 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. The B300 cutoff was tested across all three
-dtypes, independent input offsets and vector remainders. It applies to total
-element count, independent of tensor rank or model width.
+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
diff --git a/tests/test_activation.py b/tests/test_activation.py
index 762fc00d..0b877a0b 100644
--- a/tests/test_activation.py
+++ b/tests/test_activation.py
@@ -6,6 +6,7 @@
 import subprocess
 import sys
 import textwrap
+from itertools import product
 
 import pytest
 import torch
@@ -192,29 +193,32 @@ def test_silu_and_mul_cuda_all_alignment_pairs_match_aligned_path(dtype):
 @pytest.mark.parametrize("delta", range(-1, 8))
 def test_silu_and_mul_cuda_threshold_and_vector_offsets_match_aligned_path(dtype, delta):
     # Short-offset tests now take the scalar fallback. Exercise all independent
-    # input alignments on both sides of the dispatch threshold and scalar tails.
+    # input alignments above the threshold and each possible scalar tail.
     sms = torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count
-    count = math.ceil(sms * 160 / 148) * 256 * 8 + delta
     elements = 16 // torch.empty((), dtype=dtype).element_size()
+    count = sms * 160 * 256 * elements + delta
     integer = torch.int32 if dtype == torch.float32 else torch.int16
     generator = torch.Generator(device="cuda").manual_seed(43)
-    for go in range(elements):
-        for uo in range(elements):
-            gate = torch.randn(count + go, device="cuda", dtype=dtype, generator=generator)[go:]
-            up = torch.randn(count + uo, device="cuda", dtype=dtype, generator=generator)[uo:]
-            expected = silu_and_mul(gate.clone(), up.clone())
-            actual = silu_and_mul(gate, up)
-            assert torch.equal(actual.view(integer), expected.view(integer)), (dtype, count, go, uo)
+    gate_storage = torch.randn(
+        count + elements - 1, device="cuda", dtype=dtype, generator=generator
+    )
+    up_storage = torch.randn(count + elements - 1, device="cuda", dtype=dtype, generator=generator)
+    offsets = product(range(elements), repeat=2) if delta == 1 else [(1, 3), (elements - 1, 1)]
+    for go, uo in offsets:
+        gate = gate_storage[go : go + count]
+        up = up_storage[uo : uo + count]
+        expected = silu_and_mul(gate.clone(), up.clone())
+        actual = silu_and_mul(gate, up)
+        assert torch.equal(actual.view(integer), expected.view(integer)), (dtype, count, go, uo)
 
 
 @pytest.mark.gpu
 @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
 def test_silu_and_mul_cuda_uses_device_scaled_threshold(dtype, tmp_path):
-    # Check actual launch geometry on both sides, including FP32: B300's
-    # measured cutoff is in elements, rather than FP32 vector work items.
+    # Both paths can have the same capped grid, so also inspect the kernel mode.
     sms = torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count
-    threshold = math.ceil(sms * 160 / 148) * 256 * 8
     elements = 16 // torch.empty((), dtype=dtype).element_size()
+    threshold = sms * 160 * 256 * elements
     cases = [(threshold - 1, 1, 3), (threshold, 1, 3), (threshold, 0, 0)]
     tensors = [
         (
@@ -244,14 +248,23 @@ def test_silu_and_mul_cuda_uses_device_scaled_threshold(dtype, tmp_path):
         key=lambda event: event["ts"],
     )
     assert len(kernels) == len(cases)
-    scalar_blocks = min(math.ceil((threshold - 1) / 256), 4096)
-    vector_blocks = min(math.ceil((threshold // elements) / 256), 4096)
+
+    def blocks_for(work):
+        cap = 16384 if dtype == torch.float32 and work >= 16384 * 256 else 4096
+        return min(math.ceil(work / 256), cap)
+
+    scalar_blocks = blocks_for(threshold - 1)
+    vector_blocks = blocks_for(threshold // elements)
     assert [event["args"]["grid"][0] for event in kernels] == [
         scalar_blocks,
         vector_blocks,
         vector_blocks,
     ]
     assert all(event["args"]["block"] == [256, 1, 1] for event in kernels)
+    assert all(
+        f"AccessMode){mode}>" in event["name"]
+        for event, mode in zip(kernels, (2, 1, 0), strict=True)
+    )
 
 
 @pytest.mark.gpu
@@ -318,7 +331,8 @@ def test_silu_and_mul_cuda_rejects_autograd_but_allows_inference():
 @pytest.mark.parametrize("large", [False, True])
 def test_silu_and_mul_cuda_current_stream_and_graph_replay(dtype, offsets, large):
     sms = torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count
-    threshold = math.ceil(sms * 160 / 148) * 256 * 8
+    elements = 16 // torch.empty((), dtype=dtype).element_size()
+    threshold = sms * 160 * 256 * elements
     count = threshold + 1 if large else min(4097, threshold - 1)
     stream = torch.cuda.Stream()
     with torch.cuda.stream(stream):
@@ -394,15 +408,23 @@ def test_silu_and_mul_cuda_low_precision_storage_patterns(dtype, offsets):
         device="cuda",
         dtype=dtype,
     )
-    count = values.numel() * factors.numel() + 1
+    pattern_count = values.numel() * factors.numel()
+    sms = torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count
+    threshold = sms * 160 * 256 * (16 // values.element_size())
+    count = max(pattern_count + 1, threshold + 1) if offsets != (0, 0) else pattern_count + 1
     gate = torch.empty(count + offsets[0], device="cuda", dtype=dtype)[offsets[0] :]
     up = torch.empty(count + offsets[1], device="cuda", dtype=dtype)[offsets[1] :]
-    gate[:-1].copy_(values.repeat(factors.numel()))
-    up[:-1].copy_(factors.repeat_interleave(values.numel()))
-    # The aligned path ends with one scalar multiplication yielding a subnormal.
+    gate.zero_()
+    up.zero_()
+    gate[:pattern_count].copy_(values.repeat(factors.numel()))
+    up[:pattern_count].copy_(factors.repeat_interleave(values.numel()))
+    # A scalar tail multiplication produces a subnormal.
     gate[-1], up[-1] = 1.0, subnormal
     actual = silu_and_mul(gate, up)
-    expected = silu_and_mul_reference(gate, up)
-    torch.testing.assert_close(actual, expected, rtol=0, atol=0, equal_nan=True)
+    expected = silu_and_mul_reference(gate[:pattern_count], up[:pattern_count])
+    actual_patterns = actual[:pattern_count]
+    torch.testing.assert_close(actual_patterns, expected, rtol=0, atol=0, equal_nan=True)
     valid = ~torch.isnan(expected)
-    assert torch.equal(actual[valid].view(torch.int16), expected[valid].view(torch.int16))
+    assert torch.equal(actual_patterns[valid].view(torch.int16), expected[valid].view(torch.int16))
+    expected_tail = silu_and_mul_reference(gate[-1:], up[-1:])
+    torch.testing.assert_close(actual[-1:], expected_tail, rtol=0, atol=0)

From 4e3c361d94c40e2bb48012d1d466cd424f2e347e Mon Sep 17 00:00:00 2001
From: bjf-frz <frz123db@gmail.com>
Date: Tue, 3 Nov 2026 09:52:32 +0000
Subject: [PATCH 4/6] refactor(kernels): name misaligned SiLU access mode

---
 csrc/activation/silu_and_mul.cu | 6 +++---
 docs/kernels.md                 | 4 ++--
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/csrc/activation/silu_and_mul.cu b/csrc/activation/silu_and_mul.cu
index 3a5b9d13..08c05c7a 100644
--- a/csrc/activation/silu_and_mul.cu
+++ b/csrc/activation/silu_and_mul.cu
@@ -133,7 +133,7 @@ __device__ __forceinline__ Vector<T> load_vector(const T* input, int64_t index,
   return result;
 }
 
-enum class AccessMode { Aligned, Reconstruct, Scalar };
+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) {
@@ -144,7 +144,7 @@ __global__ void silu_and_mul_kernel(const T* gate, const T* up, T* out, int64_t
     constexpr int elements = 16 / sizeof(T);
     constexpr int pairs = 16 / sizeof(typename Math<T>::Pair);
     int64_t prefix = 0;
-    if constexpr (mode == AccessMode::Reconstruct) {
+    if constexpr (mode == AccessMode::Misaligned) {
       prefix = min(count, static_cast<int64_t>(
           ((16 - (reinterpret_cast<uintptr_t>(out) & 15)) & 15) / sizeof(T)));
     }
@@ -210,7 +210,7 @@ void launch(const at::Tensor& gate, const at::Tensor& up, at::Tensor& out) {
   } else if (scalar) {
     silu_and_mul_kernel<T, AccessMode::Scalar><<<blocks, threads, 0, stream>>>(g, u, result, count);
   } else {
-    silu_and_mul_kernel<T, AccessMode::Reconstruct><<<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();
 }
diff --git a/docs/kernels.md b/docs/kernels.md
index fc1f1d22..9f01fc48 100644
--- a/docs/kernels.md
+++ b/docs/kernels.md
@@ -53,8 +53,8 @@ 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 reconstructed inputs share the
-same vector calculation and output loop; output prefix and tail use scalar
+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

From 77e5ac7920bfbd6add81c468b8855923e2249088 Mon Sep 17 00:00:00 2001
From: bjf-frz <frz123db@gmail.com>
Date: Tue, 3 Nov 2026 10:52:36 +0000
Subject: [PATCH 5/6] perf(kernels): remove output alignment prefix

---
 csrc/activation/silu_and_mul.cu | 30 +++++++--------
 docs/kernels.md                 |  6 ++-
 tests/test_activation.py        | 67 +++++++++++++++++++++++++++++++++
 3 files changed, 86 insertions(+), 17 deletions(-)

diff --git a/csrc/activation/silu_and_mul.cu b/csrc/activation/silu_and_mul.cu
index 08c05c7a..506f8b69 100644
--- a/csrc/activation/silu_and_mul.cu
+++ b/csrc/activation/silu_and_mul.cu
@@ -143,17 +143,9 @@ __global__ void silu_and_mul_kernel(const T* gate, const T* up, T* out, int64_t
   if constexpr (mode != AccessMode::Scalar) {
     constexpr int elements = 16 / sizeof(T);
     constexpr int pairs = 16 / sizeof(typename Math<T>::Pair);
-    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;
+    const int64_t vectors = count / elements;
     for (int64_t i = thread; i < vectors; i += stride) {
-      const int64_t index = prefix + i * elements;
+      const int64_t index = 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;
@@ -161,9 +153,16 @@ __global__ void silu_and_mul_kernel(const T* gate, const T* up, T* out, int64_t
       for (int j = 0; j < pairs; ++j) {
         result.pairs[j] = Math<T>::compute(g.pairs[j], u.pairs[j]);
       }
-      *reinterpret_cast<uint4*>(out + index) = *reinterpret_cast<const uint4*>(&result);
+      if ((reinterpret_cast<uintptr_t>(out) & 15) == 0) {
+        *reinterpret_cast<uint4*>(out + index) = *reinterpret_cast<const uint4*>(&result);
+      } else {
+        uint2* dest = reinterpret_cast<uint2*>(out + index);
+        const uint2* src = reinterpret_cast<const uint2*>(&result);
+        dest[0] = src[0];
+        dest[1] = src[1];
+      }
     }
-    tail_start = prefix + vectors * elements;
+    tail_start = vectors * elements;
   }
   // Scalar mode starts at zero; vector modes handle only their remaining tail.
   for (int64_t i = tail_start + thread; i < count; i += stride) {
@@ -176,10 +175,11 @@ void launch(const at::Tensor& gate, const at::Tensor& up, at::Tensor& out) {
   const T* g = reinterpret_cast<const T*>(gate.const_data_ptr());
   const T* u = reinterpret_cast<const T*>(up.const_data_ptr());
   T* result = reinterpret_cast<T*>(out.mutable_data_ptr());
+  TORCH_CHECK((reinterpret_cast<uintptr_t>(result) & 7) == 0,
+              "silu_and_mul requires an output aligned to 8 bytes");
   // A contiguous storage-offset view is not necessarily 16-byte aligned.
-  const bool vectorized = ((reinterpret_cast<uintptr_t>(g) |
-                            reinterpret_cast<uintptr_t>(u) |
-                            reinterpret_cast<uintptr_t>(result)) & 15) == 0;
+  const bool vectorized =
+      ((reinterpret_cast<uintptr_t>(g) | reinterpret_cast<uintptr_t>(u)) & 15) == 0;
   const int64_t count = gate.numel();
   constexpr int threads = 256;
   bool scalar = false;
diff --git a/docs/kernels.md b/docs/kernels.md
index 9f01fc48..bb2fd864 100644
--- a/docs/kernels.md
+++ b/docs/kernels.md
@@ -54,8 +54,10 @@ 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
+same calculation and output loop. A 16-byte-aligned output uses 16-byte stores;
+an 8-byte-aligned output uses two 8-byte stores, and an incomplete tail uses
+scalar computation. Output addresses must be aligned to at least 8 bytes. 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
diff --git a/tests/test_activation.py b/tests/test_activation.py
index 0b877a0b..bbe1c549 100644
--- a/tests/test_activation.py
+++ b/tests/test_activation.py
@@ -3,6 +3,7 @@
 
 import json
 import math
+import shutil
 import subprocess
 import sys
 import textwrap
@@ -267,6 +268,72 @@ def blocks_for(work):
     )
 
 
+@pytest.mark.gpu
+def test_silu_and_mul_cuda_eight_byte_aligned_output(tmp_path):
+    nvcc = shutil.which("nvcc")
+    if nvcc is None:
+        pytest.skip("nvcc is required to build the custom CUDA allocator")
+    source = tmp_path / "shifted_allocator.cu"
+    library = tmp_path / "shifted_allocator.so"
+    source.write_text(
+        textwrap.dedent("""
+        #include <cuda_runtime_api.h>
+        #include <cstddef>
+        extern "C" void* alloc8(ssize_t size, int device, cudaStream_t stream) {
+            (void)device; (void)stream;
+            void* ptr = nullptr;
+            if (cudaMalloc(&ptr, static_cast<size_t>(size) + 16) != cudaSuccess) return nullptr;
+            return static_cast<char*>(ptr) + 8;
+        }
+        extern "C" void free8(void* ptr, size_t size, int device, cudaStream_t stream) {
+            (void)size; (void)device; (void)stream;
+            if (ptr) cudaFree(static_cast<char*>(ptr) - 8);
+        }
+    """)
+    )
+    subprocess.run(
+        [nvcc, "-shared", "-Xcompiler", "-fPIC", "-O2", str(source), "-o", str(library)],
+        check=True,
+        capture_output=True,
+        timeout=120,
+    )
+    script = textwrap.dedent(f"""
+        import json
+        import torch
+        from vllm_rlt.kernels import silu_and_mul
+        allocator = torch.cuda.memory.CUDAPluggableAllocator({str(library)!r}, 'alloc8', 'free8')
+        pool = torch.cuda.MemPool(allocator.allocator())
+        sms = torch.cuda.get_device_properties(0).multi_processor_count
+        modes = []
+        with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU,
+                                                torch.profiler.ProfilerActivity.CUDA]) as prof:
+            for dtype in (torch.bfloat16, torch.float16, torch.float32):
+                elements = 16 // torch.empty((), dtype=dtype).element_size()
+                large = sms * 160 * 256 * elements + 3
+                for count, gate_offset, up_offset in ((4097, 0, 0), (large, 1, 3)):
+                    gate_storage = torch.full((count + gate_offset,), 1, device='cuda', dtype=dtype)
+                    up_storage = torch.full((count + up_offset,), 2, device='cuda', dtype=dtype)
+                    gate, up = gate_storage[gate_offset:], up_storage[up_offset:]
+                    with torch.cuda.use_mem_pool(pool):
+                        out = silu_and_mul(gate, up)
+                    torch.cuda.synchronize()
+                    assert out.data_ptr() % 16 == 8
+                    expected = (torch.nn.functional.silu(gate[:1]) * up[:1]).item()
+                    assert torch.all(out == expected)
+                    modes.append(0 if gate_offset == 0 else 1)
+        trace = {str(tmp_path / "eight-byte-output.json")!r}
+        prof.export_chrome_trace(trace)
+        with open(trace) as f:
+            events = json.load(f)['traceEvents']
+        kernels = sorted((event for event in events if event.get('cat') == 'kernel' and
+                          'silu_and_mul_kernel' in event.get('name', '')), key=lambda e: e['ts'])
+        assert len(kernels) == len(modes), (len(kernels), len(modes))
+        assert all(f'AccessMode){{mode}}>' in event['name']
+                   for event, mode in zip(kernels, modes))
+    """)
+    subprocess.run([sys.executable, "-c", script], check=True, capture_output=True, timeout=120)
+
+
 @pytest.mark.gpu
 @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
 @pytest.mark.parametrize("offsets", [(0, 0), (1, 3)])

From 21c706e646ae350ab20844ec24bd6dc63ed6e27b Mon Sep 17 00:00:00 2001
From: bjf-frz <frz123db@gmail.com>
Date: Tue, 3 Nov 2026 11:04:18 +0000
Subject: [PATCH 6/6] Revert "perf(kernels): remove output alignment prefix"

This reverts commit 77e5ac7920bfbd6add81c468b8855923e2249088.
---
 csrc/activation/silu_and_mul.cu | 30 +++++++--------
 docs/kernels.md                 |  6 +--
 tests/test_activation.py        | 67 ---------------------------------
 3 files changed, 17 insertions(+), 86 deletions(-)

diff --git a/csrc/activation/silu_and_mul.cu b/csrc/activation/silu_and_mul.cu
index 506f8b69..08c05c7a 100644
--- a/csrc/activation/silu_and_mul.cu
+++ b/csrc/activation/silu_and_mul.cu
@@ -143,9 +143,17 @@ __global__ void silu_and_mul_kernel(const T* gate, const T* up, T* out, int64_t
   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) {
-      const int64_t index = 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;
@@ -153,16 +161,9 @@ __global__ void silu_and_mul_kernel(const T* gate, const T* up, T* out, int64_t
       for (int j = 0; j < pairs; ++j) {
         result.pairs[j] = Math<T>::compute(g.pairs[j], u.pairs[j]);
       }
-      if ((reinterpret_cast<uintptr_t>(out) & 15) == 0) {
-        *reinterpret_cast<uint4*>(out + index) = *reinterpret_cast<const uint4*>(&result);
-      } else {
-        uint2* dest = reinterpret_cast<uint2*>(out + index);
-        const uint2* src = reinterpret_cast<const uint2*>(&result);
-        dest[0] = src[0];
-        dest[1] = src[1];
-      }
+      *reinterpret_cast<uint4*>(out + index) = *reinterpret_cast<const uint4*>(&result);
     }
-    tail_start = vectors * elements;
+    tail_start = prefix + vectors * elements;
   }
   // Scalar mode starts at zero; vector modes handle only their remaining tail.
   for (int64_t i = tail_start + thread; i < count; i += stride) {
@@ -175,11 +176,10 @@ void launch(const at::Tensor& gate, const at::Tensor& up, at::Tensor& out) {
   const T* g = reinterpret_cast<const T*>(gate.const_data_ptr());
   const T* u = reinterpret_cast<const T*>(up.const_data_ptr());
   T* result = reinterpret_cast<T*>(out.mutable_data_ptr());
-  TORCH_CHECK((reinterpret_cast<uintptr_t>(result) & 7) == 0,
-              "silu_and_mul requires an output aligned to 8 bytes");
   // A contiguous storage-offset view is not necessarily 16-byte aligned.
-  const bool vectorized =
-      ((reinterpret_cast<uintptr_t>(g) | reinterpret_cast<uintptr_t>(u)) & 15) == 0;
+  const bool vectorized = ((reinterpret_cast<uintptr_t>(g) |
+                            reinterpret_cast<uintptr_t>(u) |
+                            reinterpret_cast<uintptr_t>(result)) & 15) == 0;
   const int64_t count = gate.numel();
   constexpr int threads = 256;
   bool scalar = false;
diff --git a/docs/kernels.md b/docs/kernels.md
index bb2fd864..9f01fc48 100644
--- a/docs/kernels.md
+++ b/docs/kernels.md
@@ -54,10 +54,8 @@ 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. A 16-byte-aligned output uses 16-byte stores;
-an 8-byte-aligned output uses two 8-byte stores, and an incomplete tail uses
-scalar computation. Output addresses must be aligned to at least 8 bytes. No
-input copy or extra kernel launch is needed. Fully aligned inputs
+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
diff --git a/tests/test_activation.py b/tests/test_activation.py
index bbe1c549..0b877a0b 100644
--- a/tests/test_activation.py
+++ b/tests/test_activation.py
@@ -3,7 +3,6 @@
 
 import json
 import math
-import shutil
 import subprocess
 import sys
 import textwrap
@@ -268,72 +267,6 @@ def blocks_for(work):
     )
 
 
-@pytest.mark.gpu
-def test_silu_and_mul_cuda_eight_byte_aligned_output(tmp_path):
-    nvcc = shutil.which("nvcc")
-    if nvcc is None:
-        pytest.skip("nvcc is required to build the custom CUDA allocator")
-    source = tmp_path / "shifted_allocator.cu"
-    library = tmp_path / "shifted_allocator.so"
-    source.write_text(
-        textwrap.dedent("""
-        #include <cuda_runtime_api.h>
-        #include <cstddef>
-        extern "C" void* alloc8(ssize_t size, int device, cudaStream_t stream) {
-            (void)device; (void)stream;
-            void* ptr = nullptr;
-            if (cudaMalloc(&ptr, static_cast<size_t>(size) + 16) != cudaSuccess) return nullptr;
-            return static_cast<char*>(ptr) + 8;
-        }
-        extern "C" void free8(void* ptr, size_t size, int device, cudaStream_t stream) {
-            (void)size; (void)device; (void)stream;
-            if (ptr) cudaFree(static_cast<char*>(ptr) - 8);
-        }
-    """)
-    )
-    subprocess.run(
-        [nvcc, "-shared", "-Xcompiler", "-fPIC", "-O2", str(source), "-o", str(library)],
-        check=True,
-        capture_output=True,
-        timeout=120,
-    )
-    script = textwrap.dedent(f"""
-        import json
-        import torch
-        from vllm_rlt.kernels import silu_and_mul
-        allocator = torch.cuda.memory.CUDAPluggableAllocator({str(library)!r}, 'alloc8', 'free8')
-        pool = torch.cuda.MemPool(allocator.allocator())
-        sms = torch.cuda.get_device_properties(0).multi_processor_count
-        modes = []
-        with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU,
-                                                torch.profiler.ProfilerActivity.CUDA]) as prof:
-            for dtype in (torch.bfloat16, torch.float16, torch.float32):
-                elements = 16 // torch.empty((), dtype=dtype).element_size()
-                large = sms * 160 * 256 * elements + 3
-                for count, gate_offset, up_offset in ((4097, 0, 0), (large, 1, 3)):
-                    gate_storage = torch.full((count + gate_offset,), 1, device='cuda', dtype=dtype)
-                    up_storage = torch.full((count + up_offset,), 2, device='cuda', dtype=dtype)
-                    gate, up = gate_storage[gate_offset:], up_storage[up_offset:]
-                    with torch.cuda.use_mem_pool(pool):
-                        out = silu_and_mul(gate, up)
-                    torch.cuda.synchronize()
-                    assert out.data_ptr() % 16 == 8
-                    expected = (torch.nn.functional.silu(gate[:1]) * up[:1]).item()
-                    assert torch.all(out == expected)
-                    modes.append(0 if gate_offset == 0 else 1)
-        trace = {str(tmp_path / "eight-byte-output.json")!r}
-        prof.export_chrome_trace(trace)
-        with open(trace) as f:
-            events = json.load(f)['traceEvents']
-        kernels = sorted((event for event in events if event.get('cat') == 'kernel' and
-                          'silu_and_mul_kernel' in event.get('name', '')), key=lambda e: e['ts'])
-        assert len(kernels) == len(modes), (len(kernels), len(modes))
-        assert all(f'AccessMode){{mode}}>' in event['name']
-                   for event, mode in zip(kernels, modes))
-    """)
-    subprocess.run([sys.executable, "-c", script], check=True, capture_output=True, timeout=120)
-
-
 @pytest.mark.gpu
 @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
 @pytest.mark.parametrize("offsets", [(0, 0), (1, 3)])