Model wrappers: the role-aware decoder layer

        # ### PATCH START: require an explicit AFD role before allocation.
        afd_config = parse_afd_config(vllm_config, validate=False)
        afd_role = afd_config.role

        torch.nn.Module.__init__(self)
        # ### PATCH END

            if self.uses_remote_experts:
                self.mlp = AFDDeepseekV2RemoteExpertsMoE(
                    config=config,
                    parallel_config=parallel_config,
                    layer_idx=layer_idx,
                    prefix=f"{prefix}.mlp",
                    compute_gate_on_attention=self.compute_gate_on_attention,
                )
            elif self.compute_gate_on_attention and self.is_moe_layer:
                self.mlp = GateOnlyRemoteMoE(
                    config=config,
                    layer_idx=layer_idx,
                    prefix=f"{prefix}.mlp",
                    vllm_config=vllm_config,
                )
            elif self.compute_gate_on_attention:
                self.mlp = native.DeepseekV2MLP(
                    hidden_size=config.hidden_size,
                    intermediate_size=config.intermediate_size,
                    hidden_act=config.hidden_act,
                    quant_config=quant_config,
                    prefix=f"{prefix}.mlp",
                )
            else:
                self.mlp = RemoteFFNProxy(layer_idx=layer_idx)

        elif afd_role == "ffn":
            self.self_attn = native.PPMissingLayer()
            if self.compute_gate_on_attention and not self.is_moe_layer:
                self.mlp = native.PPMissingLayer()
            elif self.is_moe_layer:
                # vLLM models bind FusedMoEFactory at module import time. AFD
                # can import native DeepSeek before vLLM-Ascend patches the
                # package factory, so refresh that binding after platform
                # initialization and before constructing the NPU FFN MoE.
                if device_type == "npu":
                    native.FusedMoEFactory = fused_moe.FusedMoEFactory
                moe_config = config
                if afd_config.connector == AFD_ASYNC_CONNECTOR:
                    moe_config = copy(config)
                    # HF validates this public field as int, but pinned native
                    # MoE uses only None to omit shared construction (0 builds
                    # a zero-width MLP). Apply its sentinel to this private
                    # construction copy; preserve the real model config.
                    object.__setattr__(moe_config, "n_shared_experts", None)
                self.mlp = native.DeepseekV2MoE(
                    config=moe_config,
                    parallel_config=parallel_config,
                    quant_config=quant_config,
                    prefix=f"{prefix}.mlp",
                    # aiter applies routed_scaling_factor internally
                    apply_routed_scale_to_output=(
                        not native.rocm_aiter_ops.is_fused_moe_enabled()
                    ),
                )
                if self.compute_gate_on_attention and device_type == "cuda":
                    # Keep the gate weights loadable, but with the runner's
                    # gate cleared it consumes the Attention-side router
                    # logits instead of computing its own (gate is None means
                    # external routing in the v0.30 runner contract).
                    self.mlp.experts.gate = None
            else:
                self.mlp = native.DeepseekV2MLP(
                    hidden_size=config.hidden_size,
                    intermediate_size=config.intermediate_size,
                    hidden_act=config.hidden_act,
                    quant_config=quant_config,
                    prefix=f"{prefix}.mlp",
                )

    def _send_and_receive(
        self,
        hidden_states: torch.Tensor,
        **send_kwargs: torch.Tensor,
    ) -> torch.Tensor:
        afd_metadata = get_afd_metadata_from_forward_context()
        if afd_metadata is None:
            raise RuntimeError("RemoteFFNProxy requires AFD forward metadata")
        forward_context = get_forward_context()
        stage_idx = int(
            getattr(forward_context, "ubatch_idx", afd_metadata.stage_idx),
        )
        afd_metadata.stage_idx = stage_idx
        metadata = AFDTransferMetadata.create_attention_metadata(
            layer_idx=self.layer_idx,
            stage_idx=stage_idx,
            seq_len=int(hidden_states.shape[0]),
        )
        context = AFDTransferContext(metadata=metadata)
        afd_metadata.connector.send_attn_output(
            hidden_states,
            context,
            **send_kwargs,
        )
        hidden_states = maybe_apply_dbo_yield(
            hidden_states,
            role="attention",
        )
        return afd_metadata.connector.recv_ffn_output(
            ref_tensor=hidden_states,
            ubatch_idx=stage_idx,
        )

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        return super().load_weights(
            _iter_role_weights(
                weights,
                role=self.afd_role,
                config=self.config,
                compute_gate_on_attention=self.afd_config.compute_gate_on_attention,
                attention_shared_experts=self.afd_config.connector
                == AFD_ASYNC_CONNECTOR,
            )
        )

def _checkpoint_weight_roles(
    name: str,
    config: _DeepseekAdapterConfig,
    *,
    compute_gate_on_attention: bool,
    attention_shared_experts: bool = False,
) -> frozenset[str]:
    """Classify one native checkpoint path by its AFD execution owner."""
    layer_path = _weight_layer_path(name)
    if layer_path is None:
        return _BOTH_ROLES

    layer_idx, stage, remainder = layer_path
    if stage == "self_attn":
        return _ATTENTION_ROLE
    if stage != "mlp":
        return _BOTH_ROLES

    if not _is_moe_layer(config, layer_idx):
        return _ATTENTION_ROLE if compute_gate_on_attention else _FFN_ROLE

    if attention_shared_experts and remainder and remainder[0] == "shared_experts":
        return _ATTENTION_ROLE
    is_moe_gate = bool(remainder) and remainder[0] == "gate"
    if is_moe_gate and compute_gate_on_attention:
        return _BOTH_ROLES
    return _FFN_ROLE