Neuron: skip the vacuous candidate-block filter; mask padding rows in static MoE (#26)

- When runtime_max_model_len <= candidate_topk_blocks * candidate_block_size,
  the candidate blocks hold every position, so the filter keeps them all and
  the dense indexer logits give the same top-k. On Neuron the candidate path
  gathers all candidate keys per query token: [128, 16384] rows a layer for a
  prompt chunk, the NEFF that still ran HBM out in rg16.
- static_moe_forward takes the decode metadata's valid_token_mask: padding rows
  select no expert.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

deepinfer/model_executor/models/v4_modules/index_attn.py

@@ -98,6 +98,15 @@ class MQA(nn.Module):
         candidate_cfg = shared_topology.candidate_source_cfg[candidate_source] if candidate_source is not None else None
         self.is_candidate_topk_source = shared_topology is not None and shared_topology.is_candidate_source_layer(self.layer_idx)
         self.use_candidate_topk = candidate_source is not None and not self.is_candidate_topk_source
+        if (
+            PLATFORM.is_neuron
+            and self.use_candidate_topk
+            and cfg.runtime_max_model_len <= candidate_cfg.candidate_topk_blocks * candidate_cfg.candidate_block_size
+        ):
+            # Neuron: the candidate blocks hold every position the model admits, so the filter keeps all of them and
+            # the dense indexer logits give the same top-k. That skips a per-query-token gather of all candidate keys,
+            # which for a prompt chunk run as a decode step (issue #26) is [tokens, 16384] rows a layer.
+            self.use_candidate_topk = False
         if self.use_candidate_topk:
             assert candidate_cfg is not None
             assert self.n_head <= 32 and self.head_dim == 128

deepinfer/model_executor/models/v4_modules/mlp.py

@@ -590,7 +590,9 @@ class DeepseekMoE(nn.Module):
         topk_idx, topk_weights = self.gate_forward(hidden_states_bf16, input_metadata)
         if PLATFORM.is_neuron and not input_metadata.is_prompt:
             # Fixed shapes for the decode graphs: all-gather + dense local experts + reduce-scatter.
-            return static_moe_forward(self, hidden_states, topk_idx, topk_weights, hidden_states_bf16)
+            return static_moe_forward(
+                self, hidden_states, topk_idx, topk_weights, hidden_states_bf16, getattr(input_metadata, "valid_token_mask", None)
+            )
 
         try:
             # Dispatch

deepinfer/neuron_ops/static_moe.py

@@ -81,7 +81,13 @@ def static_moe_forward(
     topk_idx: torch.Tensor,
     topk_weights: torch.Tensor,
     hidden_states_bf16: torch.Tensor | None = None,
+    valid: torch.Tensor | None = None,
 ) -> torch.Tensor:
+    """``valid`` ([tokens] bool, the decode metadata's valid_token_mask): padding rows of a graph batch select no
+    expert, so they cost nothing past the dense compute and never reach a real row's output."""
+    if valid is not None:
+        topk_idx = torch.where(valid.view(-1, 1), topk_idx, torch.full_like(topk_idx, -1))
+        topk_weights = torch.where(valid.view(-1, 1), topk_weights, torch.zeros_like(topk_weights))
     if hidden_states_bf16 is not None and nki_gemm.has_kernel(moe.gate_up_proj_weight):
         return _nki_moe_forward(moe, hidden_states_bf16, topk_idx, topk_weights)
     group, size, rank = moe.ep_group, moe.ep_size, moe.ep_rank