[MoE] Generalize masked activation for padded layouts

csrc/libtorch_stable/activation_kernels.cu

@@ -496,24 +496,67 @@ __global__ void situ_and_mul_kernel(
   }
 }
 
-template <typename scalar_t>
+constexpr int kMaxMaskedTokenBlocks = 32;
+
+template <bool BATCHED_EXPERTS>
+__device__ __forceinline__ bool get_masked_row_range(
+    const int* __restrict__ valid_token_counts, const int max_num_tokens,
+    int64_t& first_row, int64_t& end_row, int64_t& row_stride) {
+  if constexpr (BATCHED_EXPERTS) {
+    // [E, T, *]: z lanes grid-stride over one expert's valid token prefix.
+    const int expert = blockIdx.y;
+    const int num_tokens =
+        max(0, min(valid_token_counts[expert], max_num_tokens));
+    const int token_block = blockIdx.z;
+    if (token_block >= num_tokens) {
+      return false;
+    }
+    const int64_t expert_first_row =
+        static_cast<int64_t>(expert) * max_num_tokens;
+    first_row = expert_first_row + token_block;
+    end_row = expert_first_row + num_tokens;
+    row_stride = gridDim.z;
+  } else {
+    // [T, *]: each block handles one row, masked by a single valid prefix.
+    const int row = blockIdx.x;
+    const int num_tokens = max(0, min(valid_token_counts[0], max_num_tokens));
+    if (row >= num_tokens) {
+      return false;
+    }
+    first_row = row;
+    end_row = row + 1;
+    row_stride = 1;
+  }
+  return first_row < end_row;
+}
+
+template <bool BATCHED_EXPERTS>
+__device__ __forceinline__ int get_masked_feature_index() {
+  // Feature tiles occupy x for batched layouts and y for flat layouts.
+  if constexpr (BATCHED_EXPERTS) {
+    return blockIdx.x * blockDim.x + threadIdx.x;
+  } else {
+    return blockIdx.y * blockDim.x + threadIdx.x;
+  }
+}
+
+template <typename scalar_t, bool BATCHED_EXPERTS>
 __global__ void masked_situ_and_mul_kernel(
     scalar_t* __restrict__ out, const scalar_t* __restrict__ input,
-    const int* __restrict__ expert_num_tokens, const int max_num_tokens,
+    const int* __restrict__ valid_token_counts, const int max_num_tokens,
     const int d, const float beta, const float linear_beta) {
-  const int expert = blockIdx.y;
-  const int num_tokens = expert_num_tokens[expert];
-  const int idx = blockIdx.x * blockDim.x + threadIdx.x;
-  if (idx >= d || num_tokens == 0) {
+  int64_t first_row, end_row, row_stride;
+  const int idx = get_masked_feature_index<BATCHED_EXPERTS>();
+  if (idx >= d ||
+      !get_masked_row_range<BATCHED_EXPERTS>(valid_token_counts, max_num_tokens,
+                                             first_row, end_row, row_stride)) {
     return;
   }
 
   const bool clamp_up = linear_beta > 0.0f;
   const float inv_beta = 1.0f / beta;
   const float inv_linear_beta = clamp_up ? 1.0f / linear_beta : 0.0f;
-  const int64_t expert_row = static_cast<int64_t>(expert) * max_num_tokens;
-  for (int token = 0; token < num_tokens; ++token) {
-    const int64_t row = expert_row + token;
+  for (int64_t row = first_row; row < end_row; row += row_stride) {
     const scalar_t* gate_ptr = input + row * 2 * d;
     const scalar_t* up_ptr = gate_ptr + d;
     scalar_t* out_ptr = out + row * d;
@@ -526,6 +569,119 @@ __global__ void masked_situ_and_mul_kernel(
   }
 }
 
+template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&, const float),
+          bool act_first, bool HAS_CLAMP, bool BATCHED_EXPERTS>
+__global__ void masked_act_and_mul_kernel(
+    scalar_t* __restrict__ out, const scalar_t* __restrict__ input,
+    const int* __restrict__ valid_token_counts, const int max_num_tokens,
+    const int d, const float limit, const float alpha, const float beta) {
+  int64_t first_row, end_row, row_stride;
+  const int idx = get_masked_feature_index<BATCHED_EXPERTS>();
+  if (idx >= d ||
+      !get_masked_row_range<BATCHED_EXPERTS>(valid_token_counts, max_num_tokens,
+                                             first_row, end_row, row_stride)) {
+    return;
+  }
+
+  for (int64_t row = first_row; row < end_row; row += row_stride) {
+    const scalar_t* x_ptr = input + row * 2 * d;
+    const scalar_t* y_ptr = x_ptr + d;
+    scalar_t* out_ptr = out + row * d;
+    const scalar_t x = VLLM_LDG(&x_ptr[idx]);
+    const scalar_t y = VLLM_LDG(&y_ptr[idx]);
+    out_ptr[idx] = compute<scalar_t, ACT_FN, act_first, HAS_CLAMP>(x, y, limit,
+                                                                   alpha, beta);
+  }
+}
+
+template <typename scalar_t, bool BATCHED_EXPERTS>
+__global__ void masked_swigluoai_and_mul_kernel(
+    scalar_t* __restrict__ out, const scalar_t* __restrict__ input,
+    const int* __restrict__ valid_token_counts, const int max_num_tokens,
+    const int d, const float alpha, const float limit) {
+  int64_t first_row, end_row, row_stride;
+  const int idx = get_masked_feature_index<BATCHED_EXPERTS>();
+  if (idx >= d ||
+      !get_masked_row_range<BATCHED_EXPERTS>(valid_token_counts, max_num_tokens,
+                                             first_row, end_row, row_stride)) {
+    return;
+  }
+
+  for (int64_t row = first_row; row < end_row; row += row_stride) {
+    const scalar_t* in_ptr = input + row * 2 * d;
+    scalar_t* out_ptr = out + row * d;
+    out_ptr[idx] =
+        swigluoai_and_mul(in_ptr[2 * idx], in_ptr[2 * idx + 1], alpha, limit);
+  }
+}
+
+template <typename scalar_t, bool BATCHED_EXPERTS>
+__global__ void masked_swiglustep_and_mul_kernel(
+    scalar_t* __restrict__ out, const scalar_t* __restrict__ input,
+    const int* __restrict__ valid_token_counts, const int max_num_tokens,
+    const int d, const float limit) {
+  int64_t first_row, end_row, row_stride;
+  const int idx = get_masked_feature_index<BATCHED_EXPERTS>();
+  if (idx >= d ||
+      !get_masked_row_range<BATCHED_EXPERTS>(valid_token_counts, max_num_tokens,
+                                             first_row, end_row, row_stride)) {
+    return;
+  }
+
+  for (int64_t row = first_row; row < end_row; row += row_stride) {
+    const scalar_t* gate_ptr = input + row * 2 * d;
+    const scalar_t* up_ptr = gate_ptr + d;
+    scalar_t* out_ptr = out + row * d;
+    const float gate = (float)VLLM_LDG(&gate_ptr[idx]);
+    const float up = (float)VLLM_LDG(&up_ptr[idx]);
+    const float gate_silu = gate / (1.0f + expf(-gate));
+    const float gate_clamped = fminf(gate_silu, limit);
+    const float up_clamped = fmaxf(fminf(up, limit), -limit);
+    out_ptr[idx] = (scalar_t)(gate_clamped * up_clamped);
+  }
+}
+
+template <typename T>
+__device__ __forceinline__ T silu_no_mul_kernel(const T& x) {
+  return silu_kernel(x, 1.0f);
+}
+
+template <typename T>
+__device__ __forceinline__ T gelu_no_mul_kernel(const T& x) {
+  return gelu_kernel(x, 1.0f);
+}
+
+template <typename T>
+__device__ __forceinline__ T gelu_tanh_no_mul_kernel(const T& x) {
+  return gelu_tanh_kernel(x, 1.0f);
+}
+
+template <typename T>
+__device__ __forceinline__ T relu2_no_mul_kernel(const T& x) {
+  const float value = fmaxf((float)x, 0.0f);
+  return (T)(value * value);
+}
+
+template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
+          bool BATCHED_EXPERTS>
+__global__ void masked_activation_kernel(
+    scalar_t* __restrict__ out, const scalar_t* __restrict__ input,
+    const int* __restrict__ valid_token_counts, const int max_num_tokens,
+    const int d) {
+  int64_t first_row, end_row, row_stride;
+  const int idx = get_masked_feature_index<BATCHED_EXPERTS>();
+  if (idx >= d ||
+      !get_masked_row_range<BATCHED_EXPERTS>(valid_token_counts, max_num_tokens,
+                                             first_row, end_row, row_stride)) {
+    return;
+  }
+
+  for (int64_t row = first_row; row < end_row; row += row_stride) {
+    const int64_t offset = row * d + idx;
+    out[offset] = ACT_FN(VLLM_LDG(&input[offset]));
+  }
+}
+
 }  // namespace vllm
 
 #define LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(KERNEL, PACKED_KERNEL, PARAM) \
@@ -646,23 +802,133 @@ void masked_situ_and_mul(torch::stable::Tensor& out,    // [E, T, d]
   int num_experts = input.size(0);
   int max_num_tokens = input.size(1);
   int d = input.size(2) / 2;
-  if (num_experts == 0 || max_num_tokens == 0) {
+  if (num_experts == 0 || max_num_tokens == 0 || d == 0) {
     return;
   }
   constexpr int block_size = 256;
-  dim3 grid((d + block_size - 1) / block_size, num_experts);
+  const int token_blocks =
+      std::min(max_num_tokens, vllm::kMaxMaskedTokenBlocks);
+  dim3 grid((d + block_size - 1) / block_size, num_experts, token_blocks);
   dim3 block(block_size);
   const torch::stable::accelerator::DeviceGuard device_guard(
       input.get_device_index());
   const cudaStream_t stream = get_current_cuda_stream();
   VLLM_STABLE_DISPATCH_FLOATING_TYPES(
       input.scalar_type(), "masked_situ_and_mul_kernel", [&] {
-        vllm::masked_situ_and_mul_kernel<scalar_t><<<grid, block, 0, stream>>>(
-            out.mutable_data_ptr<scalar_t>(), input.const_data_ptr<scalar_t>(),
-            expert_num_tokens.const_data_ptr<int>(), max_num_tokens, d,
-            (float)beta, (float)linear_beta);
+        vllm::masked_situ_and_mul_kernel<scalar_t, true>
+            <<<grid, block, 0, stream>>>(
+                out.mutable_data_ptr<scalar_t>(),
+                input.const_data_ptr<scalar_t>(),
+                expert_num_tokens.const_data_ptr<int>(), max_num_tokens, d,
+                (float)beta, (float)linear_beta);
       });
 }
+
+#define LAUNCH_MASKED_ACT_AND_MUL(KERNEL, ACT_FIRST, HAS_CLAMP, LIMIT, ALPHA, \
+                                  BETA)                                       \
+  vllm::masked_act_and_mul_kernel<scalar_t, KERNEL<scalar_t>, ACT_FIRST,      \
+                                  HAS_CLAMP, BATCHED_EXPERTS>                 \
+      <<<grid, block, 0, stream>>>(out.mutable_data_ptr<scalar_t>(),          \
+                                   input.const_data_ptr<scalar_t>(),          \
+                                   valid_token_counts.const_data_ptr<int>(),  \
+                                   max_num_tokens, d, LIMIT, ALPHA, BETA)
+
+#define LAUNCH_MASKED_ACTIVATION(KERNEL)                                      \
+  vllm::masked_activation_kernel<scalar_t, KERNEL<scalar_t>, BATCHED_EXPERTS> \
+      <<<grid, block, 0, stream>>>(                                           \
+          out.mutable_data_ptr<scalar_t>(), input.const_data_ptr<scalar_t>(), \
+          valid_token_counts.const_data_ptr<int>(), max_num_tokens, d)
+
+void masked_moe_activation(
+    torch::stable::Tensor& out,    // [T, d] or [E, T, d]
+    torch::stable::Tensor& input,  // [T, d] or [E, T, d], gated has 2 * d
+    const torch::stable::Tensor& valid_token_counts,
+    const std::string& activation, double clamp_limit, double alpha,
+    double beta, double situ_beta, double situ_linear_beta) {
+  const bool batched_experts = input.dim() == 3;
+  const int num_experts = batched_experts ? input.size(0) : 1;
+  const int max_num_tokens = batched_experts ? input.size(1) : input.size(0);
+  const int d = out.size(-1);
+  if (num_experts == 0 || max_num_tokens == 0 || d == 0) {
+    return;
+  }
+
+  constexpr int block_size = 256;
+  const int feature_blocks = (d + block_size - 1) / block_size;
+  const int token_blocks =
+      std::min(max_num_tokens, vllm::kMaxMaskedTokenBlocks);
+  // Batched grid: (feature tile, expert, token lane); flat grid: (row,
+  // feature tile). Token lanes grid-stride the valid prefix.
+  dim3 grid = batched_experts ? dim3(feature_blocks, num_experts, token_blocks)
+                              : dim3(max_num_tokens, feature_blocks);
+  dim3 block(block_size);
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      input.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream();
+  auto launch = [&]<bool BATCHED_EXPERTS>() {
+    VLLM_STABLE_DISPATCH_FLOATING_TYPES(
+        input.scalar_type(), "masked_moe_activation", [&] {
+          if (activation == "silu") {
+            LAUNCH_MASKED_ACT_AND_MUL(vllm::silu_kernel, true, false, 0.0f,
+                                      1.0f, 0.0f);
+          } else if (activation == "silu_with_clamp") {
+            LAUNCH_MASKED_ACT_AND_MUL(vllm::silu_kernel, true, true,
+                                      (float)clamp_limit, (float)alpha,
+                                      (float)beta);
+          } else if (activation == "gelu") {
+            LAUNCH_MASKED_ACT_AND_MUL(vllm::gelu_kernel, true, false, 0.0f,
+                                      1.0f, 0.0f);
+          } else if (activation == "gelu_tanh") {
+            LAUNCH_MASKED_ACT_AND_MUL(vllm::gelu_tanh_kernel, true, false, 0.0f,
+                                      1.0f, 0.0f);
+          } else if (activation == "situ") {
+            vllm::masked_situ_and_mul_kernel<scalar_t, BATCHED_EXPERTS>
+                <<<grid, block, 0, stream>>>(
+                    out.mutable_data_ptr<scalar_t>(),
+                    input.const_data_ptr<scalar_t>(),
+                    valid_token_counts.const_data_ptr<int>(), max_num_tokens, d,
+                    (float)situ_beta, (float)situ_linear_beta);
+          } else if (activation == "swigluoai") {
+            vllm::masked_swigluoai_and_mul_kernel<scalar_t, BATCHED_EXPERTS>
+                <<<grid, block, 0, stream>>>(
+                    out.mutable_data_ptr<scalar_t>(),
+                    input.const_data_ptr<scalar_t>(),
+                    valid_token_counts.const_data_ptr<int>(), max_num_tokens, d,
+                    (float)alpha, (float)clamp_limit);
+          } else if (activation == "swigluoai_uninterleave") {
+            LAUNCH_MASKED_ACT_AND_MUL(vllm::silu_kernel, true, true,
+                                      (float)clamp_limit, (float)alpha,
+                                      (float)beta);
+          } else if (activation == "swiglustep") {
+            vllm::masked_swiglustep_and_mul_kernel<scalar_t, BATCHED_EXPERTS>
+                <<<grid, block, 0, stream>>>(
+                    out.mutable_data_ptr<scalar_t>(),
+                    input.const_data_ptr<scalar_t>(),
+                    valid_token_counts.const_data_ptr<int>(), max_num_tokens, d,
+                    (float)clamp_limit);
+          } else if (activation == "silu_no_mul") {
+            LAUNCH_MASKED_ACTIVATION(vllm::silu_no_mul_kernel);
+          } else if (activation == "gelu_no_mul") {
+            LAUNCH_MASKED_ACTIVATION(vllm::gelu_no_mul_kernel);
+          } else if (activation == "gelu_tanh_no_mul") {
+            LAUNCH_MASKED_ACTIVATION(vllm::gelu_tanh_no_mul_kernel);
+          } else if (activation == "relu2_no_mul") {
+            LAUNCH_MASKED_ACTIVATION(vllm::relu2_no_mul_kernel);
+          } else {
+            STD_TORCH_CHECK(false,
+                            "Unsupported masked MoE activation: ", activation);
+          }
+        });
+  };
+  if (batched_experts) {
+    launch.template operator()<true>();
+  } else {
+    launch.template operator()<false>();
+  }
+}
+
+#undef LAUNCH_MASKED_ACT_AND_MUL
+#undef LAUNCH_MASKED_ACTIVATION
 namespace vllm {
 
 // Element-wise activation kernel template.

csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu

@@ -146,16 +146,17 @@ __global__ void preprocessTopkIdKernel(int* topk_id_ptr, int size,
   }
   __syncthreads();
 
-  // query global expert id in expert map.
-  // if global expert id = -1 in exert map, plus n_expert
-  // else set global expert id = exert map[global expert id]
   if (offset + tidx < bound) {
     auto topk_id = topk_id_ptr[offset + tidx];
-    auto local_expert_idx = smem_expert_map[topk_id];
-    if (local_expert_idx == -1) {
-      topk_id += num_experts;
+    if (topk_id < 0 || topk_id >= num_experts) {
+      topk_id = num_experts;
     } else {
-      topk_id = local_expert_idx;
+      auto local_expert_idx = smem_expert_map[topk_id];
+      if (local_expert_idx == -1) {
+        topk_id += num_experts;
+      } else {
+        topk_id = local_expert_idx;
+      }
     }
     __syncwarp();
     topk_id_ptr[offset + tidx] = topk_id;

csrc/libtorch_stable/ops.h

@@ -559,6 +559,13 @@ void masked_situ_and_mul(torch::stable::Tensor& out,
                          torch::stable::Tensor& input,
                          const torch::stable::Tensor& expert_num_tokens,
                          double beta = 1.0, double linear_beta = -1.0);
+void masked_moe_activation(torch::stable::Tensor& out,
+                           torch::stable::Tensor& input,
+                           const torch::stable::Tensor& valid_token_counts,
+                           const std::string& activation,
+                           double clamp_limit = 0.0, double alpha = 1.0,
+                           double beta = 0.0, double situ_beta = 1.0,
+                           double situ_linear_beta = -1.0);
 void gelu_new(torch::stable::Tensor& out, torch::stable::Tensor& input);
 void gelu_fast(torch::stable::Tensor& out, torch::stable::Tensor& input);
 void gelu_quick(torch::stable::Tensor& out, torch::stable::Tensor& input);

csrc/libtorch_stable/torch_bindings.cpp

@@ -632,6 +632,11 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
   ops.def(
       "masked_situ_and_mul(Tensor! out, Tensor input, Tensor "
       "expert_num_tokens, float beta=1.0, float linear_beta=-1.0) -> ()");
+  ops.def(
+      "masked_moe_activation(Tensor! out, Tensor input, Tensor "
+      "valid_token_counts, str activation, float clamp_limit=0.0, float "
+      "alpha=1.0, float beta=0.0, float situ_beta=1.0, float "
+      "situ_linear_beta=-1.0) -> ()");
 
   // GELU implementation used in GPT-2.
   ops.def("gelu_new(Tensor! out, Tensor input) -> ()");
@@ -853,6 +858,7 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
   ops.impl("swigluoai_and_mul", TORCH_BOX(&swigluoai_and_mul));
   ops.impl("situ_and_mul", TORCH_BOX(&situ_and_mul));
   ops.impl("masked_situ_and_mul", TORCH_BOX(&masked_situ_and_mul));
+  ops.impl("masked_moe_activation", TORCH_BOX(&masked_moe_activation));
   ops.impl("gelu_new", TORCH_BOX(&gelu_new));
   ops.impl("gelu_fast", TORCH_BOX(&gelu_fast));
   ops.impl("gelu_quick", TORCH_BOX(&gelu_quick));