Guard Humming DeepEP padding

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

Author: mgoin

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;

tests/kernels/core/test_activation.py

@@ -75,6 +75,27 @@ def test_masked_moe_activation_rejects_unsupported_activation() -> None:
         )
 
 
+def test_masked_moe_activation_validates_tensor_contract() -> None:
+    input = torch.empty(1, 1, 2)
+    valid_token_counts = torch.ones(1, dtype=torch.int32)
+
+    with pytest.raises(AssertionError, match="dtypes must match"):
+        apply_moe_activation(
+            MoEActivation.SILU,
+            torch.empty(1, 1, 1, dtype=torch.float16),
+            input,
+            valid_token_counts=valid_token_counts,
+        )
+
+    with pytest.raises(AssertionError, match="Input must be contiguous"):
+        apply_moe_activation(
+            MoEActivation.SILU,
+            torch.empty(1, 1, 1),
+            torch.empty(1, 1, 4)[..., ::2],
+            valid_token_counts=valid_token_counts,
+        )
+
+
 @pytest.mark.parametrize(
     "activation",
     [

tests/kernels/moe/test_deepep_v2_moe.py

@@ -487,14 +487,10 @@ def _make_mxfp4_humming_experts(
     return experts, quant_config, layer.w13_weight, layer.w2_weight
 
 
-def _deep_ep_v2_moe_cudagraph(
+def _run_deep_ep_v2_backend_case(
     pgi: ProcessGroupInfo,
     dp_size: int,
     config: TestConfig,
-    w1: torch.Tensor,
-    w2: torch.Tensor,
-    w1_scale: torch.Tensor | None,
-    w2_scale: torch.Tensor | None,
     moe_backend: str,
     activation: MoEActivation,
     use_cudagraph: bool,
@@ -909,13 +905,9 @@ def _launch_deep_ep_v2_case(
 
     parallel_launch(
         world_size,
-        _deep_ep_v2_moe_cudagraph,
+        _run_deep_ep_v2_backend_case,
         dp_size,
         config,
-        None,  # weights created inside worker
-        None,
-        None,
-        None,
         moe_backend,
         activation,
         use_cudagraph,

tests/kernels/moe/test_moe_permute_unpermute.py

@@ -286,3 +286,39 @@ def test_moe_permute_reuses_scratch_buffers(dtype: torch.dtype):
         permuted_idx_1.untyped_storage().data_ptr()
         == scratch.permuted_idx.untyped_storage().data_ptr()
     )
+
+
+def test_moe_permute_ignores_invalid_expert_ids_with_scratch() -> None:
+    if not moe_permute_unpermute_supported():
+        pytest.skip("moe_permute_unpermute is not supported on this platform.")
+
+    hidden_states = torch.arange(5 * 16, dtype=torch.bfloat16, device="cuda").view(
+        5, 16
+    )
+    topk_ids = torch.tensor([[0], [-1], [1], [4], [2]], device="cuda")
+    expert_map = torch.tensor([0, 1, -1, -1], dtype=torch.int32, device="cuda")
+    scratch = MoEPermuteScratch(
+        max_num_tokens=5,
+        topk=1,
+        num_experts=4,
+        num_local_experts=2,
+        device=hidden_states.device,
+        hidden_size=16,
+        hidden_dtype=hidden_states.dtype,
+    )
+
+    permuted, _, expert_offsets, _, _ = moe_permute(
+        hidden_states=hidden_states,
+        a1q_scale=None,
+        topk_ids=topk_ids,
+        n_expert=4,
+        n_local_expert=2,
+        expert_map=expert_map,
+        scratch=scratch,
+    )
+
+    torch.testing.assert_close(
+        expert_offsets,
+        torch.tensor([0, 1, 2], dtype=torch.int64, device="cuda"),
+    )
+    torch.testing.assert_close(permuted[:2], hidden_states[[0, 2]])

vllm/model_executor/layers/fused_moe/activation.py

@@ -235,10 +235,18 @@ def _apply_moe_activation_masked(
 
     assert input.dim() in (2, 3), "Masked input must be 2D or 3D"
     _validate_moe_activation_shapes(activation, output, input, expected_dim=input.dim())
+    assert input.dtype == output.dtype, "Input and output dtypes must match"
+    assert input.device == output.device, "Input and output devices must match"
+    assert input.is_contiguous(), "Input must be contiguous"
+    assert output.is_contiguous(), "Output must be contiguous"
     assert valid_token_counts.dtype == torch.int32, (
         "valid_token_counts must use torch.int32"
     )
     assert valid_token_counts.dim() == 1, "valid_token_counts must be 1D"
+    assert valid_token_counts.device == input.device, (
+        "valid_token_counts must be on the input device"
+    )
+    assert valid_token_counts.is_contiguous(), "valid_token_counts must be contiguous"
     expected_counts = input.size(0) if input.dim() == 3 else 1
     assert valid_token_counts.size(0) == expected_counts, (
         f"valid_token_counts must have {expected_counts} element(s) for "