The MoE handoff: where the gate runs

        if compute_gate_on_attention:
            self.gate = native.GateLinear(
                config.hidden_size,
                config.n_routed_experts,
                out_dtype=router_dtype,
                prefix=f"{prefix}.gate",
            )
            if getattr(config, "topk_method", None) == "noaux_tc":
                self.gate.e_score_correction_bias = nn.Parameter(
                    torch.empty(config.n_routed_experts, dtype=torch.float32),
                )
            else:
                self.gate.e_score_correction_bias = None
        else:
            self.gate = None

        ep_size = native.get_ep_group().device_group.size()
        self.n_routed_experts = int(config.n_routed_experts)
        self.n_shared_experts = int(config.n_shared_experts)
        self.n_redundant_experts = parallel_config.eplb_config.num_redundant_experts
        self.n_logical_experts = self.n_routed_experts
        self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts
        self.n_local_physical_experts = self.n_physical_experts // ep_size
        # No local shared-expert MLP is constructed, so fused shared experts
        # can never be enabled for this shell.
        self.is_fused_shared_expert_enabled = False
        self.experts = AFDAttentionFusedMoE(
            layer_idx=layer_idx,
            external_routing=compute_gate_on_attention,
        )

        # ### PATCH START: compute the Attention-side gate before transport.
        num_tokens, hidden_dim = hidden_states.shape
        hidden_states = hidden_states.view(-1, hidden_dim)
        if self.is_sequence_parallel and not already_sequence_parallel:
            hidden_states = native.sequence_parallel_chunk(hidden_states)
        router_logits, _ = self.gate(hidden_states)
        final_hidden_states = self.experts(
            hidden_states=hidden_states,
            router_logits=router_logits,
        )

        num_layers = max(int(self.num_layers or 0), 1)
        experts_layer_indices = frozenset(
            model.get_experts_layer_indices(),
        )
        layer_indices = (
            tuple(sorted(experts_layer_indices))
            if self.afd_config.compute_gate_on_attention
            else tuple(range(num_layers))
        )

                    if (
                        uses_remote_experts
                        and self.afd_config.compute_gate_on_attention
                    ):
                        router_logits = payload.router_logits
                        assert router_logits is not None
                        rank_ffn_output = model.compute_experts_output(
                            hidden_states,
                            layer_idx,
                            router_logits,
                        )
                    else:
                        rank_ffn_output = self._execute_eager_mode(
                            hidden_states,
                            layer_idx,
                            input_ids=payload.input_ids,
                        )

        # v0.30 runner contract: gate is None means the runner consumes the
        # router logits transported from Attention (external routing).
        if self.mlp.experts.gate is not None:
            raise RuntimeError("FFN native runner must use external routing")
        return self.mlp.experts(
            hidden_states=hidden_states,
            router_logits=router_logits,
        )