llm.runtime — Plugin Registries and Model Factory

The runtime layer wires extension points (model architectures, attention, MLP, norms, generation backends) via the generic :class:~llm.runtime.registry.Registry kernel. The model factory turns a ModelConfig into an instantiated DecoderModel.

Registries

registry

Generic plugin registry for runtime extensibility.

Registry

Name-to-object registry with explicit registration and lookup.

源代码位于： src/llm/runtime/registry.py

class Registry[T]:
    """Name-to-object registry with explicit registration and lookup."""

    def __init__(self, name: str) -> None:
        self._name = name
        self._entries: dict[str, T] = {}

    def register(self, name: str, entry: T) -> T:
        if name in self._entries:
            raise ValueError(f"'{name}' is already registered in {self._name} registry.")
        self._entries[name] = entry
        return entry

    def get(self, name: str) -> T:
        if name not in self._entries:
            available = ", ".join(sorted(self._entries))
            raise ValueError(f"'{name}' not found in {self._name} registry. Available: {available}")
        return self._entries[name]

    def names(self) -> list[str]:
        return sorted(self._entries)

    def __contains__(self, name: str) -> bool:
        return name in self._entries

decorator_register

decorator_register(registry)

Class decorator factory compatible with legacy register_model usage.

源代码位于： src/llm/runtime/registry.py

def decorator_register(registry: Registry[type]) -> Callable[[str], Callable[[type], type]]:
    """Class decorator factory compatible with legacy register_model usage."""

    def register(name: str) -> Callable[[type], type]:
        def wrapper(cls: type) -> type:
            registry.register(name, cls)
            return cls

        return wrapper

    return register

Model Factory

model_factory

Central model construction for training, serving, and compat loaders.

ModelFactory

Resolve registered model builders from typed or raw configuration.

源代码位于： src/llm/runtime/model_factory.py

class ModelFactory:
    """Resolve registered model builders from typed or raw configuration."""

    @staticmethod
    def from_config(config: ModelConfig, *, model_type: str = "decoder", **overrides: Any) -> nn.Module:
        if model_type == "decoder":
            kwargs = decoder_kwargs_from_config(config, **overrides)
        else:
            kwargs = {
                "hidden_size": config.hidden_size,
                "intermediate_size": config.intermediate_size,
                "dropout_p": config.dropout,
                "use_glu": config.use_glu,
            }
            kwargs.update(overrides)
        return ModelFactory.build(model_type, **kwargs)

    @staticmethod
    def build(model_type: str = "decoder", **kwargs: Any) -> nn.Module:
        builder = MODEL_REGISTRY.get(model_type)
        return builder(**kwargs)

build_decoder

build_decoder(*, vocab_size, hidden_size, num_layers, num_heads, max_seq_len=512, intermediate_size=None, embedding_dropout_p=0.1, attn_dropout_p=0.1, mlp_dropout_p=0.1, num_experts=0, top_k=0, num_kv_heads=None, use_glu=False, attn_impl='mha', mlp_impl='mlp', norm_eps=1e-05, norm_impl='layer_norm', device=None, dtype=None, **kwargs)

Construct a DecoderModel from explicit architecture kwargs.

源代码位于： src/llm/runtime/model_factory.py

def build_decoder(
    *,
    vocab_size: int,
    hidden_size: int,
    num_layers: int,
    num_heads: int,
    max_seq_len: int = 512,
    intermediate_size: int | None = None,
    embedding_dropout_p: float = 0.1,
    attn_dropout_p: float = 0.1,
    mlp_dropout_p: float = 0.1,
    num_experts: int = 0,
    top_k: int = 0,
    num_kv_heads: int | None = None,
    use_glu: bool = False,
    attn_impl: str = "mha",
    mlp_impl: str = "mlp",
    norm_eps: float = 1e-5,
    norm_impl: str = "layer_norm",
    device: torch.device | str | None = None,
    dtype: torch.dtype | None = None,
    **kwargs: Any,
) -> DecoderModel:
    """Construct a DecoderModel from explicit architecture kwargs."""
    return DecoderModel(
        vocab_size=vocab_size,
        hidden_size=hidden_size,
        num_layers=num_layers,
        num_heads=num_heads,
        max_seq_len=max_seq_len,
        intermediate_size=intermediate_size,
        embedding_dropout_p=embedding_dropout_p,
        attn_dropout_p=attn_dropout_p,
        mlp_dropout_p=mlp_dropout_p,
        num_experts=num_experts,
        top_k=top_k,
        num_kv_heads=num_kv_heads,
        use_glu=use_glu,
        attn_impl=attn_impl,
        mlp_impl=mlp_impl,
        norm_eps=norm_eps,
        norm_impl=norm_impl,
        device=device,
        dtype=dtype,
        **kwargs,
    )

build_regression_mlp

build_regression_mlp(*, hidden_size, intermediate_size=None, dropout_p=0.1, use_glu=False, **_)

Construct a standalone MLP for the synthetic regression demo task.

源代码位于： src/llm/runtime/model_factory.py

def build_regression_mlp(
    *,
    hidden_size: int,
    intermediate_size: int | None = None,
    dropout_p: float = 0.1,
    use_glu: bool = False,
    **_: Any,
) -> nn.Module:
    """Construct a standalone MLP for the synthetic regression demo task."""
    from llm.core.mlp import MLP

    return MLP(
        hidden_size=hidden_size,
        intermediate_size=intermediate_size,
        dropout_p=dropout_p,
        use_glu=use_glu,
    )

decoder_kwargs_from_config

decoder_kwargs_from_config(config, **overrides)

Map a ModelConfig into DecoderModel constructor kwargs.

源代码位于： src/llm/runtime/model_factory.py

def decoder_kwargs_from_config(config: ModelConfig, **overrides: Any) -> dict[str, Any]:
    """Map a ModelConfig into DecoderModel constructor kwargs."""
    kwargs: dict[str, Any] = {
        "vocab_size": config.vocab_size,
        "hidden_size": config.hidden_size,
        "num_layers": config.num_layers,
        "num_heads": config.num_heads,
        "max_seq_len": config.max_seq_len,
        "intermediate_size": config.intermediate_size,
        "embedding_dropout_p": config.dropout,
        "attn_dropout_p": config.dropout,
        "mlp_dropout_p": config.dropout,
        "num_experts": config.num_experts,
        "top_k": config.top_k,
        "num_kv_heads": config.num_kv_heads,
        "use_glu": config.use_glu,
        "attn_impl": config.attn_impl,
        "mlp_impl": config.mlp_impl,
        "norm_impl": config.norm_impl,
    }
    kwargs.update(overrides)
    return kwargs

Plugin Loader

plugins

Discover third-party plugins via setuptools entry points.

load_entry_point_registry

load_entry_point_registry(group, registry, *, overwrite=False)

Load callables from entry points into a registry.

Returns names that were newly registered.

源代码位于： src/llm/runtime/plugins.py

def load_entry_point_registry[T](
    group: str,
    registry: Registry[T],
    *,
    overwrite: bool = False,
) -> list[str]:
    """Load callables from entry points into a registry.

    Returns names that were newly registered.
    """
    loaded: list[str] = []
    for ep in _iter_group_entry_points(group):
        if not overwrite and ep.name in registry:
            continue
        factory = ep.load()
        registry.register(ep.name, factory)
        loaded.append(ep.name)
    return loaded

load_entry_point_hooks

load_entry_point_hooks(group)

Invoke zero-arg registration hooks from entry points.

源代码位于： src/llm/runtime/plugins.py

def load_entry_point_hooks(group: str) -> list[str]:
    """Invoke zero-arg registration hooks from entry points."""
    invoked: list[str] = []
    for ep in _iter_group_entry_points(group):
        hook = ep.load()
        if not callable(hook):
            raise TypeError(f"Entry point '{ep.name}' in group '{group}' must be callable.")
        hook()
        invoked.append(ep.name)
    return invoked