diff --git a/deepinfer/model_executor/models/v4_modules/index_attn.py b/deepinfer/model_executor/models/v4_modules/index_attn.py index 345b12c..9d4f955 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.is_candidate_topk_source + and candidate_cfg is not None + and cfg.runtime_max_model_len <= candidate_cfg.candidate_topk_blocks * candidate_cfg.candidate_block_size + ): + # Neuron: every consumer skips the vacuous filter (below), so nothing reads the candidate blocks; and with + # fewer blocks than candidate_topk_blocks the block top-k would not be defined. + self.is_candidate_topk_source = False if ( PLATFORM.is_neuron and self.use_candidate_topk