跳转至

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 
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
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 replace(self, name: str, entry: T) -> T:
        """Register ``entry`` under ``name``, overwriting any existing entry.

        Unlike :meth:`register`, this does not raise on a duplicate — it is
        the explicit-override path used by
        :func:`llm.runtime.plugins.load_entry_point_registry(overwrite=True)`
        so third-party plugins can override a built-in implementation (RIL
        ISS-061). For an unknown name it behaves exactly like ``register``.
        """
        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

replace 

replace(name, entry)

Register entry under name, overwriting any existing entry.

Unlike :meth:register, this does not raise on a duplicate — it is the explicit-override path used by :func:llm.runtime.plugins.load_entry_point_registry(overwrite=True) so third-party plugins can override a built-in implementation (RIL ISS-061). For an unknown name it behaves exactly like register.

源代码位于： src/llm/runtime/registry.py 
24
25
26
27
28
29
30
31
32
33
34
def replace(self, name: str, entry: T) -> T:
    """Register ``entry`` under ``name``, overwriting any existing entry.

    Unlike :meth:`register`, this does not raise on a duplicate — it is
    the explicit-override path used by
    :func:`llm.runtime.plugins.load_entry_point_registry(overwrite=True)`
    so third-party plugins can override a built-in implementation (RIL
    ISS-061). For an unknown name it behaves exactly like ``register``.
    """
    self._entries[name] = entry
    return entry

decorator_register 

decorator_register(registry)

Class decorator factory compatible with legacy register_model usage.

源代码位于： src/llm/runtime/registry.py 
49
50
51
52
53
54
55
56
57
58
59
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 
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
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 
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
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 
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
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 
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
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,
        "pos_encoding_learned": config.pos_encoding_learned,
        "mlp_activation": config.mlp_activation,
        "norm_first": config.norm_first,
        "qkv_bias": config.qkv_bias,
        "mlp_bias": config.mlp_bias,
        "lm_head_bias": config.lm_head_bias,
        "use_rope": config.use_rope,
        "rope_theta": config.rope_theta,
        "use_alibi": config.use_alibi,
        "attn_sparse": config.attn_sparse,
    }
    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. A plugin that fails to load is logged (not fatal) — the rest of the group still registers (RIL ISS-131).

源代码位于： src/llm/runtime/plugins.py 
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
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. A plugin that fails to load is
    logged (not fatal) — the rest of the group still registers (RIL ISS-131).
    """
    loaded: list[str] = []
    for ep in _iter_group_entry_points(group):
        preexisting = ep.name in registry
        if not overwrite and preexisting:
            # A third-party plugin claiming a built-in's name (e.g. an
            # exporter named ``onnx``) is silently kept out of the registry:
            # the built-in wins without a hint that the plugin was dropped.
            # Log it so the packaging conflict (which export/registry.py's
            # docstring claims "raises loudly") is at least observable
            # (RIL ISS-163).
            logger.warning(
                "Skipping entry point '%s' in group '%s': name already "
                "registered (built-in or earlier plugin); pass overwrite=True "
                "to replace it.",
                ep.name,
                group,
            )
            continue
        factory = _load_one(ep.name, group, ep)
        if factory is None:
            continue
        if preexisting:
            # ``overwrite=True`` and the name is already registered:
            # ``Registry.register`` would raise, so use the explicit
            # replace path (RIL ISS-061).
            registry.replace(ep.name, factory)
        else:
            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.

Hooks that fail to load or raise are logged, not fatal — the remaining hooks still run (RIL ISS-131).

源代码位于： src/llm/runtime/plugins.py 
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
def load_entry_point_hooks(group: str) -> list[str]:
    """Invoke zero-arg registration hooks from entry points.

    Hooks that fail to load or raise are logged, not fatal — the remaining
    hooks still run (RIL ISS-131).
    """
    invoked: list[str] = []
    for ep in _iter_group_entry_points(group):
        hook = _load_one(ep.name, group, ep)
        if hook is None:
            continue
        if not callable(hook):
            logger.error(
                "Entry point '%s' in group '%s' is not callable; skipped.",
                ep.name,
                group,
            )
            continue
        try:
            hook()
        except Exception as exc:  # noqa: BLE001 - one bad hook must not block others
            logger.error(
                "Hook '%s' in group '%s' raised %s: %s",
                ep.name,
                group,
                type(exc).__name__,
                exc,
            )
            continue
        invoked.append(ep.name)
    return invoked