Model registration: the AFD<arch> wrappers

Code version: afd-plugin @ 7452d52 · 2026-10-27 · vLLM @ ced6857

The idea#

AFD needs a role-aware version of each model, so the attention side can replace the FFN stage with a remote proxy. It supplies one AFD<arch> class per architecture. It registers these classes with vLLM's ModelRegistry, as lazy "module:Class" strings, so vLLM imports a class only when it is needed.

The model config names a plain architecture, for example DeepseekV2ForCausalLM. At build time, the AFD worker swaps that name to AFDDeepseekV2ForCausalLM, so vLLM builds the AFD class instead of the native one. The swap copies the config, so the native config is left unchanged.

An everyday comparison. Think of a catalogue with an "AFD" edition of each item. You order by the plain name; the clerk swaps in the AFD edition. The comparison stops here: the swap keeps one hidden link intact, so vLLM Ascend can still tell a text model from a multimodal one.

The picture#

Read panel 1 as load time: register the names. Read panel 2 as build time: detect and swap the identity. Read panel 3 as the exact mappings. The Model wrappers: the role-aware decoder layer page shows what the wrappers do inside.

In the code#

Step 1: Register each AFD name#

At the end of register_afd, a loop registers each architecture under the name "AFD" + arch (afd_plugin/__init__.py:211-214). The source map holds eight architectures (afd_plugin/__init__.py:96-139).

    from vllm.model_executor.models import ModelRegistry

    for model_arch, model_cls in _MODEL_REGISTRATIONS.items():
        ModelRegistry.register_model(f"AFD{model_arch}", model_cls)

Step 2: The lazy "module:Class" form#

The registered value is a string, not a class. vLLM stores it and imports the class only when the model is built (vllm:vllm/model_executor/models/registry.py:1126-1141). So loading the plugin does not import every model file.

Step 3: Detect an AFD model#

has_afd_model_registration returns true when the model's architecture, with any AFD prefix removed, is in the map (afd_plugin/model_executor/models/model_utils.py:16-22). The runtime validator uses it to require a registered AFD model (afd_plugin/validation.py:136-141).

def has_afd_model_registration(model_config: ModelConfig) -> bool:
    """Return whether the model resolves to a registered AFD implementation."""

    return any(
        model_arch.removeprefix("AFD") in _MODEL_REGISTRATIONS
        for model_arch in model_config.hf_config.architectures
    )

Step 4: Swap the model identity#

get_afd_model_config finds the architecture in the map, deep-copies the model config, and sets its architecture to AFD<arch> (afd_plugin/model_executor/models/model_utils.py:25-46). The worker then uses the returned config (afd_plugin/v1/worker/attention_worker.py:75-80).

    for model_arch in model_config.hf_config.architectures:
        if model_arch in _MODEL_REGISTRATIONS:
            if model_arch in _QWEN3_5_MODEL_REGISTRATIONS and device_type != "cuda":
                raise ValueError(
                    "AFD Qwen3.5/3.6 supports CUDA execution only; "
                    f"got device_type={device_type!r}",
                )
            # deepcopy preserves aliasing within the copied object graph, so
            # the pure-text identity hf_text_config is hf_config is retained
            # automatically. vLLM Ascend uses that identity to distinguish
            # text models from multimodal models.
            afd_model_config = deepcopy(model_config)
            afd_model_config.hf_config.architectures = [f"AFD{model_arch}"]
            return afd_model_config
    return model_config

Step 5: Which architectures map to which wrapper#

The map registers these exact strings (afd_plugin/__init__.py:96-131).

The map also registers DeepseekForCausalLM and GlmMoeDsaForCausalLM (afd_plugin/__init__.py:96-116).

Deep dive#

What can go wrong#

• Model not registered. A model whose architecture is not in the map fails the runtime check: "AFD ModelRunnerV2 requires a registered AFD model" (afd_plugin/validation.py:136-141).
• Qwen3.5 on NPU. get_afd_model_config raises for a non-CUDA device (afd_plugin/model_executor/models/model_utils.py:34-38).
• Model file import error. Because the registration is a lazy string, an import error in the model file surfaces when vLLM builds the model, not when the plugin loads (vllm:vllm/model_executor/models/registry.py:1126-1141).

Why it matters. The worker runs an AFD model only after the identity swap. If the architecture is not registered, the swap does nothing and the runtime check stops start-up.