diff --git a/deepinfer/model_executor/models/v4_modules/index_attn.py b/deepinfer/model_executor/models/v4_modules/index_attn.py index 4122290..345b12c 100644 --- a/deepinfer/model_executor/models/v4_modules/index_attn.py +++ b/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 diff --git a/deepinfer/model_executor/models/v4_modules/mlp.py b/deepinfer/model_executor/models/v4_modules/mlp.py index af26b66..8f3fddc 100644 --- a/deepinfer/model_executor/models/v4_modules/mlp.py +++ b/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 diff --git a/deepinfer/neuron_ops/static_moe.py b/deepinfer/neuron_ops/static_moe.py index 8745ab0..a76580c 100644 --- a/deepinfer/neuron_ops/static_moe.py +++ b/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