Method forward_experts

tensorrt_llm/layers/moe.py:1290–1384

(self, hidden_states, token_selected_experts,
                        token_final_scales, lora_layer_params, side_stream_id)

        )

    def forward_experts(self, hidden_states, token_selected_experts,
                        token_final_scales, lora_layer_params, side_stream_id):
        assert side_stream_id == SideStreamIDType.disable, "MoeOOTB does not support using side stream"
        # TODO: https://nvbugspro.nvidia.com/bug/4781396 after this nvbug is fixed, we will remove this check.
        if lora_layer_params is not None:
            for module in ["mlp_h_to_4h", "mlp_4h_to_h", "mlp_gate"]:
                if lora_layer_params.get_runtime_params(0, module) is not None:
                    raise RuntimeError(
                        f"MoE  OOTB does not support {module} LoRA module, please enable MoE plugin"
                    )

        topk_indices = token_selected_experts
        topk_values = token_final_scales

        hidden_size = shape(hidden_states, -1)
        # [B*sq, hidden]
        inputs_merged = hidden_states.view(concat([-1, hidden_size]))
        flat_topk_indices = topk_indices.view(
            concat([-1, shape(topk_indices, -1)]))
        flat_topk_values = topk_values.view(concat([-1,
                                                    shape(topk_values, -1)]))

        # Create output space
        zero_buffer = inputs_merged * 0.0
        output = zero_buffer

        expert_indices_stack = []
        indices_stack = []
        # When topk indices are equal to expert index, the expert will inference the tokens.
        # Bundle all indices and experts index, then do mask once.
        for i, expert in enumerate(self.experts):
            if self.mapping.has_moe_ep():
                index = i + self.experts_per_node * self.mapping.moe_ep_rank
            else:
                index = i
            expert_indices_stack.append(
                flat_topk_indices.view(concat([1, shape(flat_topk_indices)])))

            indices_stack.append(constant(int32_array(index)))

        all_expert_indices = concat(expert_indices_stack, dim=0)
        indices = expand(
            concat(indices_stack).view(concat([len(self.experts), 1, 1])),
            shape(all_expert_indices))

        # Create all experts mask
        all_expert_mask = all_expert_indices == indices

        experts_weights = cast(
            sum(flat_topk_values *
                cast(all_expert_mask, flat_topk_values.dtype),
                dim=-1,
                keepdim=True), self.dtype)

        all_expert_mask = cast(
            sum(cast(all_expert_mask, flat_topk_values.dtype),
                dim=-1,
                keepdim=True), 'bool')