[MoE] Harden masked activation backend support

Co-authored-by: OpenAI Codex <codex@openai.com>
Signed-off-by: mgoin <mgoin64@gmail.com>

Author: mgoin

csrc/libtorch_stable/activation_kernels.cu

@@ -496,17 +496,26 @@ __global__ void situ_and_mul_kernel(
   }
 }
 
+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& first_row, int64_t& end_row, int64_t& row_stride) {
   if constexpr (BATCHED_EXPERTS) {
-    // [E, T, *]: each block handles one expert's valid token prefix.
+    // [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));
-    first_row = static_cast<int64_t>(expert) * max_num_tokens;
-    end_row = first_row + 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;
@@ -516,6 +525,7 @@ __device__ __forceinline__ bool get_masked_row_range(
     }
     first_row = row;
     end_row = row + 1;
+    row_stride = 1;
   }
   return first_row < end_row;
 }
@@ -535,17 +545,18 @@ __global__ void masked_situ_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 beta, const float linear_beta) {
-  int64_t first_row, end_row;
+  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)) {
+  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;
-  for (int64_t row = first_row; row < end_row; ++row) {
+  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;
@@ -564,14 +575,15 @@ __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;
+  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)) {
+  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) {
+  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;
@@ -587,14 +599,15 @@ __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;
+  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)) {
+  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) {
+  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] =
@@ -607,14 +620,15 @@ __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;
+  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)) {
+  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) {
+  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;
@@ -654,14 +668,15 @@ __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;
+  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)) {
+  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) {
+  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]));
   }
@@ -787,11 +802,13 @@ 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());
@@ -832,14 +849,17 @@ void masked_moe_activation(
   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) {
+  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;
-  // Batched grid: (feature tile, expert); flat grid: (row, feature tile).
-  dim3 grid = batched_experts ? dim3(feature_blocks, num_experts)
+  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(

tests/kernels/core/test_activation.py

@@ -28,6 +28,7 @@
     apply_moe_activation,
     apply_moe_activation_masked_supported,
 )
+from vllm.platforms import current_platform
 from vllm.utils.torch_utils import set_random_seed
 
 DTYPES = [torch.half, torch.bfloat16, torch.float]
@@ -96,6 +97,27 @@ def test_masked_moe_activation_validates_tensor_contract() -> None:
         )
 
 
+def test_moe_silu_clamp_uses_native_xpu_fallback(
+    default_vllm_config, monkeypatch
+) -> None:
+    monkeypatch.setattr(current_platform, "is_xpu", lambda: True)
+    clamp_limit = 3.0
+    input = torch.tensor([[12.0, -12.0, 8.0, -8.0], [-2.0, 2.0, -4.0, 4.0]])
+    output = torch.empty(2, 2)
+
+    apply_moe_activation(
+        MoEActivation.SILU,
+        output,
+        input,
+        activation_config=ApplyMoEActivationConfig(clamp_limit=clamp_limit),
+    )
+
+    expected = SiluAndMulWithClamp(clamp_limit, compile_native=False).forward_native(
+        input
+    )
+    torch.testing.assert_close(output, expected)
+
+
 @pytest.mark.parametrize(
     "activation",
     [
@@ -345,17 +367,19 @@ def test_masked_situ_and_mul(
 )
 @pytest.mark.parametrize("dtype", [torch.half, torch.bfloat16])
 @pytest.mark.parametrize("mask_layout", ["flat", "batched_experts"])
+@pytest.mark.parametrize("d", [512, 513])
 @torch.inference_mode()
 def test_masked_moe_activation(
     default_vllm_config,
     activation: MoEActivation,
     activation_config: ApplyMoEActivationConfig,
     dtype: torch.dtype,
     mask_layout: str,
+    d: int,
 ) -> None:
     """Count-aware MoE activation computes only valid row prefixes."""
     device = CUDA_DEVICES[0]
-    num_experts, max_num_tokens, d = 4, 7, 128
+    num_experts, max_num_tokens = 4, 7
     input_dim = 2 * d if activation.is_gated else d
     if mask_layout == "flat":
         input = torch.randn(max_num_tokens, input_dim, dtype=dtype, device=device)
@@ -400,6 +424,98 @@ def test_masked_moe_activation(
         assert torch.all(batched_output[expert, num_tokens:] == 42.0)
 
 
+@pytest.mark.parametrize(
+    ("mask_layout", "counts"),
+    [
+        pytest.param("flat", [-2], id="flat_negative"),
+        pytest.param("flat", [70], id="flat_over_capacity"),
+        pytest.param("batched_experts", [-2, 70], id="batched"),
+    ],
+)
+@torch.inference_mode()
+def test_masked_moe_activation_clamps_valid_token_counts(
+    default_vllm_config,
+    mask_layout: str,
+    counts: list[int],
+) -> None:
+    """Counts outside [0, T] clamp without touching padded output rows."""
+    device = CUDA_DEVICES[0]
+    max_num_tokens, d = 67, 513
+    num_experts = len(counts) if mask_layout == "batched_experts" else 1
+    leading_shape = (
+        (max_num_tokens,) if mask_layout == "flat" else (num_experts, max_num_tokens)
+    )
+    input = torch.randn(*leading_shape, 2 * d, dtype=torch.half, device=device)
+    output = torch.full((*leading_shape, d), 42.0, dtype=torch.half, device=device)
+    valid_token_counts = torch.tensor(counts, dtype=torch.int32, device=device)
+
+    apply_moe_activation(
+        MoEActivation.SILU,
+        output,
+        input,
+        valid_token_counts=valid_token_counts,
+    )
+
+    batched_input = input.view(-1, max_num_tokens, 2 * d)
+    batched_output = output.view(-1, max_num_tokens, d)
+    for expert, raw_count in enumerate(counts):
+        num_tokens = max(0, min(raw_count, max_num_tokens))
+        if num_tokens:
+            expected = torch.empty((num_tokens, d), dtype=torch.half, device=device)
+            apply_moe_activation(
+                MoEActivation.SILU,
+                expected,
+                batched_input[expert, :num_tokens].clone(),
+            )
+            torch.testing.assert_close(
+                batched_output[expert, :num_tokens],
+                expected,
+                atol=get_default_atol(output),
+                rtol=get_default_rtol(output),
+            )
+        assert torch.all(batched_output[expert, num_tokens:] == 42.0)
+
+
+@pytest.mark.parametrize("mask_layout", ["flat", "batched_experts"])
+@torch.inference_mode()
+def test_masked_moe_activation_zero_features(
+    default_vllm_config,
+    mask_layout: str,
+) -> None:
+    """A zero-width activation is a no-op instead of an invalid launch."""
+    device = CUDA_DEVICES[0]
+    if mask_layout == "flat":
+        input = torch.empty(3, 0, dtype=torch.half, device=device)
+        output = torch.empty(3, 0, dtype=torch.half, device=device)
+        valid_token_counts = torch.tensor([2], dtype=torch.int32, device=device)
+    else:
+        input = torch.empty(2, 3, 0, dtype=torch.half, device=device)
+        output = torch.empty(2, 3, 0, dtype=torch.half, device=device)
+        valid_token_counts = torch.tensor([1, 3], dtype=torch.int32, device=device)
+
+    apply_moe_activation(
+        MoEActivation.SILU,
+        output,
+        input,
+        valid_token_counts=valid_token_counts,
+    )
+
+    assert output.numel() == 0
+
+
+@torch.inference_mode()
+def test_masked_situ_and_mul_zero_features(default_vllm_config) -> None:
+    """The legacy masked SITU entry point also accepts zero-width tensors."""
+    device = CUDA_DEVICES[0]
+    input = torch.empty(2, 3, 0, dtype=torch.half, device=device)
+    output = torch.empty(2, 3, 0, dtype=torch.half, device=device)
+    expert_num_tokens = torch.tensor([1, 3], dtype=torch.int32, device=device)
+
+    torch.ops._C.masked_situ_and_mul(output, input, expert_num_tokens, 1.5, -1.0)
+
+    assert output.numel() == 0
+
+
 @pytest.mark.parametrize("mask_layout", ["flat", "batched_experts"])
 @torch.inference_mode()
 def test_masked_moe_activation_opcheck(default_vllm_config, mask_layout: str) -> None: