[Perf][MoE] Triton FP8 MoE: per-layer activation scale format and always-fused SiLU+quant

csrc/ops.h

@@ -239,7 +239,8 @@ void silu_and_mul_per_block_quant(torch::stable::Tensor& out,
                                   torch::stable::Tensor& scales,
                                   int64_t group_size,
                                   std::optional<torch::stable::Tensor> scale_ub,
-                                  bool is_scale_transposed);
+                                  bool is_scale_transposed, bool use_ue8m0,
+                                  std::optional<double> clamp_limit);
 
 // Positional encoding kernels (shared CUDA/ROCm)
 void rotary_embedding(torch::stable::Tensor& positions,

csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu

@@ -4,14 +4,15 @@
 #include "../../torch_utils.h"
 
 #include "../../dispatch_utils.h"
+#include "../../core/cuda_compat.h"
 #include "quant_conversions.cuh"
 
 namespace vllm {
 
 // Logic: one thread block per (token, group) pair
 
 template <typename scalar_t, typename scalar_out_t, bool is_scale_transposed,
-          int32_t group_size>
+          int32_t group_size, bool use_ue8m0, bool has_clamp>
 __global__ void silu_and_mul_per_block_quant_kernel(
     scalar_out_t* __restrict__ out,  // Output: [num_tokens, hidden_size] in
                                      // FP8/INT8
@@ -20,8 +21,8 @@ __global__ void silu_and_mul_per_block_quant_kernel(
                                  // num_tokens]
     scalar_t const* __restrict__ input,  // Input: [num_tokens, hidden_size * 2]
     float const* scale_ub,               // Optional scale upper bound
-    int32_t const hidden_size  // Output hidden size (input is 2x this)
-) {
+    int32_t const hidden_size,  // Output hidden size (input is 2x this)
+    float const clamp_limit) {
   static_assert((group_size & (group_size - 1)) == 0,
                 "group_size must be a power of 2 for correct reduction");
 
@@ -53,6 +54,10 @@ __global__ void silu_and_mul_per_block_quant_kernel(
   // Step 1: Each thread loads one element, computes SiLU, stores in register
   float gate = static_cast<float>(token_input_gate[tid]);
   float up = static_cast<float>(token_input_up[tid]);
+  if constexpr (has_clamp) {
+    gate = fminf(gate, clamp_limit);
+    up = fminf(fmaxf(up, -clamp_limit), clamp_limit);
+  }
 
   // Compute SiLU(gate) * up
   float sigmoid_gate = 1.0f / (1.0f + expf(-gate));
@@ -87,6 +92,10 @@ __global__ void silu_and_mul_per_block_quant_kernel(
     // Use minimum safe scaling factor
     group_scale = fmaxf(group_scale, min_scaling_factor<scalar_out_t>::val());
 
+    if constexpr (use_ue8m0) {
+      group_scale = exp2f(ceilf(log2f(group_scale)));
+    }
+
     // Store scale to global memory
     *group_scale_ptr = group_scale;
 
@@ -102,14 +111,101 @@ __global__ void silu_and_mul_per_block_quant_kernel(
       vllm::ScaledQuant<scalar_out_t, false>::quant_fn(result, group_scale);
 }
 
+// Vectorized variant: each group is handled by group_size / 8 lanes of a
+// warp, each loading 8 contiguous 16-bit elements of gate and up (16B loads)
+// and storing 8 quantized bytes. The group max is reduced with warp shuffles,
+// so there is no shared memory or __syncthreads. Requires 16B-aligned input
+// and 8B-aligned output rows.
+template <typename scalar_t, typename scalar_out_t, bool is_scale_transposed,
+          int32_t group_size, bool use_ue8m0, bool has_clamp>
+__global__ void silu_and_mul_per_block_quant_vec_kernel(
+    scalar_out_t* __restrict__ out, float* __restrict__ scales,
+    scalar_t const* __restrict__ input, float const* scale_ub,
+    int32_t const hidden_size, int32_t const num_groups,
+    int64_t const num_tokens, float const clamp_limit) {
+  constexpr int kVec = 16 / sizeof(scalar_t);
+  constexpr int kLanes = group_size / kVec;
+  static_assert(kVec == 8, "expects 16-bit input");
+  static_assert(kLanes > 0 && kLanes <= 32 && (32 % kLanes) == 0);
+
+  int64_t const gid =
+      (static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x) / kLanes;
+  int const lane = threadIdx.x % kLanes;
+  bool const valid = gid < num_tokens * num_groups;
+  int64_t const token_idx = valid ? gid / num_groups : 0;
+  int const group_idx = valid ? static_cast<int>(gid % num_groups) : 0;
+
+  int64_t const in_off = token_idx * (2 * static_cast<int64_t>(hidden_size)) +
+                         group_idx * group_size + lane * kVec;
+  alignas(16) scalar_t g[kVec];
+  alignas(16) scalar_t u[kVec];
+  if (valid) {
+    *reinterpret_cast<uint4*>(g) =
+        *reinterpret_cast<uint4 const*>(input + in_off);
+    *reinterpret_cast<uint4*>(u) =
+        *reinterpret_cast<uint4 const*>(input + in_off + hidden_size);
+  } else {
+    *reinterpret_cast<uint4*>(g) = make_uint4(0, 0, 0, 0);
+    *reinterpret_cast<uint4*>(u) = make_uint4(0, 0, 0, 0);
+  }
+
+  float r[kVec];
+  float local_max = 0.f;
+#pragma unroll
+  for (int i = 0; i < kVec; ++i) {
+    float gate = static_cast<float>(g[i]);
+    float up = static_cast<float>(u[i]);
+    if constexpr (has_clamp) {
+      gate = fminf(gate, clamp_limit);
+      up = fminf(fmaxf(up, -clamp_limit), clamp_limit);
+    }
+    float const sigmoid_gate = 1.0f / (1.0f + expf(-gate));
+    r[i] = gate * sigmoid_gate * up;
+    local_max = fmaxf(local_max, fabsf(r[i]));
+  }
+#pragma unroll
+  for (int off = kLanes / 2; off > 0; off >>= 1) {
+    local_max =
+        fmaxf(local_max, VLLM_SHFL_XOR_SYNC_WIDTH(local_max, off, kLanes));
+  }
+
+  float group_scale = local_max / quant_type_max_v<scalar_out_t>;
+  if (scale_ub != nullptr) {
+    group_scale = fminf(group_scale, *scale_ub);
+  }
+  group_scale = fmaxf(group_scale, min_scaling_factor<scalar_out_t>::val());
+  if constexpr (use_ue8m0) {
+    group_scale = exp2f(ceilf(log2f(group_scale)));
+  }
+  if (!valid) {
+    return;
+  }
+
+  alignas(8) scalar_out_t q[kVec];
+#pragma unroll
+  for (int i = 0; i < kVec; ++i) {
+    q[i] = vllm::ScaledQuant<scalar_out_t, false>::quant_fn(r[i], group_scale);
+  }
+  *reinterpret_cast<uint2*>(out + token_idx * hidden_size +
+                            group_idx * group_size + lane * kVec) =
+      *reinterpret_cast<uint2 const*>(q);
+  if (lane == 0) {
+    float* scale_ptr = is_scale_transposed
+                           ? scales + group_idx * num_tokens + token_idx
+                           : scales + token_idx * num_groups + group_idx;
+    *scale_ptr = group_scale;
+  }
+}
+
 }  // namespace vllm
 
 void silu_and_mul_per_block_quant(torch::stable::Tensor& out,
                                   torch::stable::Tensor const& input,
                                   torch::stable::Tensor& scales,
                                   int64_t group_size,
                                   std::optional<torch::stable::Tensor> scale_ub,
-                                  bool is_scale_transposed) {
+                                  bool is_scale_transposed, bool use_ue8m0,
+                                  std::optional<double> clamp_limit) {
   static torch::headeronly::ScalarType kFp8Type =
       is_fp8_ocp() ? torch::headeronly::ScalarType::Float8_e4m3fn
                    : torch::headeronly::ScalarType::Float8_e4m3fnuz;
@@ -144,6 +240,13 @@ void silu_and_mul_per_block_quant(torch::stable::Tensor& out,
 
   dim3 grid(num_tokens, num_groups);
   dim3 block(group_size);
+  // Below ~3k groups the one-block-per-group kernel has lower latency.
+  bool const use_vec =
+      static_cast<int64_t>(num_tokens) * num_groups >= 3072 &&
+      input.element_size() == 2 &&
+      reinterpret_cast<uintptr_t>(input.const_data_ptr()) % 16 == 0 &&
+      reinterpret_cast<uintptr_t>(out.const_data_ptr()) % 8 == 0;
+  constexpr int kVecBlock = 256;
 
   VLLM_STABLE_DISPATCH_FLOATING_TYPES(
       input.scalar_type(), "silu_and_mul_per_block_quant", [&] {
@@ -156,16 +259,48 @@ void silu_and_mul_per_block_quant(torch::stable::Tensor& out,
               VLLM_STABLE_DISPATCH_GROUP_SIZE(group_size, gs, [&] {
                 VLLM_STABLE_DISPATCH_BOOL(
                     is_scale_transposed, transpose_scale, [&] {
-                      vllm::silu_and_mul_per_block_quant_kernel<
-                          scalar_in_t, scalar_out_t, transpose_scale, gs>
-                          <<<grid, block, 0, stream>>>(
-                              out.mutable_data_ptr<scalar_out_t>(),
-                              scales.mutable_data_ptr<float>(),
-                              input.const_data_ptr<scalar_in_t>(),
-                              scale_ub.has_value()
-                                  ? scale_ub->const_data_ptr<float>()
-                                  : nullptr,
-                              hidden_size);
+                      VLLM_STABLE_DISPATCH_BOOL(use_ue8m0, ue8m0, [&] {
+                        VLLM_STABLE_DISPATCH_BOOL(
+                            clamp_limit.has_value(), clamp, [&] {
+                              if constexpr (sizeof(scalar_in_t) == 2) {
+                                if (use_vec) {
+                                  constexpr int kLanes = gs / 8;
+                                  int64_t const total =
+                                      static_cast<int64_t>(num_tokens) *
+                                      num_groups * kLanes;
+                                  dim3 vgrid(static_cast<unsigned>(
+                                      (total + kVecBlock - 1) / kVecBlock));
+                                  vllm::silu_and_mul_per_block_quant_vec_kernel<
+                                      scalar_in_t, scalar_out_t,
+                                      transpose_scale, gs, ue8m0, clamp>
+                                      <<<vgrid, kVecBlock, 0, stream>>>(
+                                          out.mutable_data_ptr<scalar_out_t>(),
+                                          scales.mutable_data_ptr<float>(),
+                                          input.const_data_ptr<scalar_in_t>(),
+                                          scale_ub.has_value()
+                                              ? scale_ub
+                                                    ->const_data_ptr<float>()
+                                              : nullptr,
+                                          hidden_size, num_groups, num_tokens,
+                                          static_cast<float>(
+                                              clamp_limit.value_or(0.0)));
+                                  return;
+                                }
+                              }
+                              vllm::silu_and_mul_per_block_quant_kernel<
+                                  scalar_in_t, scalar_out_t, transpose_scale,
+                                  gs, ue8m0, clamp><<<grid, block, 0, stream>>>(
+                                  out.mutable_data_ptr<scalar_out_t>(),
+                                  scales.mutable_data_ptr<float>(),
+                                  input.const_data_ptr<scalar_in_t>(),
+                                  scale_ub.has_value()
+                                      ? scale_ub->const_data_ptr<float>()
+                                      : nullptr,
+                                  hidden_size,
+                                  static_cast<float>(
+                                      clamp_limit.value_or(0.0)));
+                            });
+                      });
                     });
               });
             });

csrc/torch_bindings.cpp

@@ -400,7 +400,9 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       "Tensor! scales, "
       "int group_size, "
       "Tensor? scale_ub=None, "
-      "bool is_scale_transposed=False) -> ()");
+      "bool is_scale_transposed=False, "
+      "bool use_ue8m0=False, "
+      "float? clamp_limit=None) -> ()");
 
   // Rotary embedding
   // Apply GPT-NeoX or GPT-J style rotary embedding to query and key.

tests/kernels/core/test_fused_silu_mul_block_quant.py

@@ -212,3 +212,101 @@ def test_silu_block_quant_int32_offset_overflow(default_vllm_config):
     )
     assert torch.equal(scales[-1:], expected_scales)
     assert torch.equal(out[-1:], expected_out)
+
+
+@pytest.mark.parametrize("use_ue8m0", [False, True])
+@pytest.mark.parametrize("clamp_limit", [None, 7.0])
+@pytest.mark.parametrize("group_size", GROUP_SIZES)
+@pytest.mark.parametrize("is_scale_transposed", IS_SCALE_TRANSPOSED)
+@torch.inference_mode()
+def test_silu_and_mul_per_block_quant_ue8m0_clamp(
+    default_vllm_config,
+    use_ue8m0: bool,
+    clamp_limit: float | None,
+    group_size: int,
+    is_scale_transposed: bool,
+) -> None:
+    """UE8M0 scales and the SwiGLU clamp must match the unfused fp32 reference,
+    as consumed by MoE backends running with DeepGEMM E8M0 or swiglu_limit."""
+    torch.set_default_device("cuda")
+    torch.random.manual_seed(0)
+    num_tokens, hidden_size = 37, 1024
+    quant_dtype = current_platform.fp8_dtype()
+    x = torch.randn(num_tokens, hidden_size * 2, dtype=torch.bfloat16) * 4
+
+    gate, up = x.float().split(hidden_size, dim=-1)
+    if clamp_limit is not None:
+        gate = gate.clamp(max=clamp_limit)
+        up = up.clamp(-clamp_limit, clamp_limit)
+    ref_out, ref_scales = per_token_group_quant_fp8(
+        F.silu(gate) * up, group_size=group_size, use_ue8m0=use_ue8m0
+    )
+
+    out, scales = ops.silu_and_mul_per_block_quant(
+        x,
+        group_size,
+        quant_dtype,
+        is_scale_transposed=is_scale_transposed,
+        use_ue8m0=use_ue8m0,
+        clamp_limit=clamp_limit,
+    )
+
+    if use_ue8m0:
+        assert torch.equal(scales, torch.exp2(torch.log2(scales).round()))
+    torch.testing.assert_close(scales, ref_scales, rtol=1e-5, atol=0)
+    deq = out.float() * scales.repeat_interleave(group_size, dim=1)
+    ref_deq = ref_out.float() * ref_scales.repeat_interleave(group_size, dim=1)
+    # Allow rare one-ulp fp8 rounding flips from expf vs torch.exp.
+    mismatch = ~torch.isclose(deq, ref_deq, atol=5e-2, rtol=5e-2)
+    assert mismatch.float().mean() < 1e-3
+    torch.testing.assert_close(deq, ref_deq, atol=5e-2, rtol=0.13)
+
+    opcheck(
+        torch.ops._C.silu_and_mul_per_block_quant,
+        (
+            torch.empty_like(out),
+            x,
+            torch.empty_like(scales.t() if is_scale_transposed else scales),
+            group_size,
+            None,
+            is_scale_transposed,
+            use_ue8m0,
+            clamp_limit,
+        ),
+    )
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("quant_dtype", QUANT_DTYPES)
+@pytest.mark.parametrize("group_size", GROUP_SIZES)
+@pytest.mark.parametrize("use_ue8m0", [False, True])
+@pytest.mark.parametrize("is_scale_transposed", IS_SCALE_TRANSPOSED)
+@torch.inference_mode()
+def test_silu_and_mul_per_block_quant_vec_matches_scalar(
+    default_vllm_config,
+    dtype: torch.dtype,
+    quant_dtype: torch.dtype,
+    group_size: int,
+    use_ue8m0: bool,
+    is_scale_transposed: bool,
+) -> None:
+    """The vectorized kernel (large, 16B-aligned inputs) must match the
+    one-block-per-group kernel bitwise; a 2-byte-misaligned view of the same
+    values forces the latter."""
+    if quant_dtype == torch.int8 and use_ue8m0:
+        pytest.skip("UE8M0 scales are FP8-only")
+    num_tokens, hidden_size = 1024, 1536
+    x = torch.randn(num_tokens, 2 * hidden_size, dtype=dtype, device="cuda") * 4
+    buf = torch.empty(x.numel() + 1, dtype=dtype, device="cuda")
+    x_misaligned = buf[1:].view_as(x)
+    x_misaligned.copy_(x)
+    kwargs = dict(
+        group_size=group_size,
+        quant_dtype=quant_dtype,
+        is_scale_transposed=is_scale_transposed,
+        use_ue8m0=use_ue8m0,
+    )
+    out, scales = ops.silu_and_mul_per_block_quant(x, **kwargs)
+    ref_out, ref_scales = ops.silu_and_mul_per_block_quant(x_misaligned, **kwargs)
+    assert torch.equal(out.view(torch.uint8), ref_out.view(torch.uint8))
+    assert torch.equal(scales, ref_scales)