跳转至

llm.serving — FastAPI Application

REST + OpenAI-compatible chat-completions API. Boot the server with uv run llm-serve or uv run uvicorn llm.serving.api:app.

Application Factory

api

FastAPI application factory for the serving API.

This module wires the pieces together:

  • :class:RequestIDMiddleware (request correlation + access log)
  • exception handlers that map every failure to the structured envelope in :mod:llm.serving.errors
  • the three routers (health, generate, chat)
  • :class:prometheus_fastapi_instrumentator.Instrumentator

All endpoint logic lives in the routers; this file stays focused on process wiring (logging setup, lifespan, app construction) so adding a new endpoint doesn't grow it.

get_api_key async 

get_api_key(api_key_header_value=Security(api_key_header), auth_header=Security(authorization_header))

Verify the API key from X-API-Key or Authorization: Bearer.

Comparison uses hmac.compare_digest to avoid leaking key bytes via timing. If the module-level config.api_key is unset, auth is disabled and the function returns None (the public-host guard in :mod:llm.serving.cli blocks starting the server on a non-loopback interface without auth).

源代码位于： src/llm/serving/auth.py 
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
async def get_api_key(
    api_key_header_value: str | None = Security(api_key_header),
    auth_header: str | None = Security(authorization_header),
) -> str | None:
    """Verify the API key from ``X-API-Key`` or ``Authorization: Bearer``.

    Comparison uses ``hmac.compare_digest`` to avoid leaking key bytes via
    timing. If the module-level ``config.api_key`` is unset, auth is
    disabled and the function returns ``None`` (the public-host guard in
    :mod:`llm.serving.cli` blocks starting the server on a non-loopback
    interface without auth).
    """
    from llm.serving.api import config as _config

    if not _config.api_key:
        return None

    expected = _config.api_key
    # Check X-API-Key header first.
    if api_key_header_value is not None and hmac.compare_digest(api_key_header_value, expected):
        return api_key_header_value

    # Check Bearer token.
    bearer = _extract_bearer_token(auth_header)
    if bearer is not None and hmac.compare_digest(bearer, expected):
        return bearer

    raise APIError(ErrorCode.UNAUTHORIZED, "Could not validate credentials")

is_loopback 

is_loopback(host)

Return True if host is a loopback address.

Covers 127.0.0.0/8 and ::1. Anything else (0.0.0.0, *, LAN IPs, public hostnames) is treated as non-loopback.

源代码位于： src/llm/serving/auth.py 
63
64
65
66
67
68
69
70
71
def is_loopback(host: str) -> bool:
    """Return True if ``host`` is a loopback address.

    Covers ``127.0.0.0/8`` and ``::1``. Anything else (``0.0.0.0``, ``*``,
    LAN IPs, public hostnames) is treated as non-loopback.
    """
    if host in ("127.0.0.1", "localhost", "::1"):
        return True
    return bool(host.startswith("127."))

lifespan async 

lifespan(_app)

FastAPI lifespan manager — load model, wire routers, log config.

源代码位于： src/llm/serving/api.py 
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
@contextlib.asynccontextmanager
async def lifespan(_app: FastAPI) -> AsyncGenerator[None, Any]:
    """FastAPI lifespan manager — load model, wire routers, log config."""
    logger.info("Starting up...")

    # Refuse to serve anonymous inference on a non-loopback interface, with
    # no api_key configured.  This guard must live in the *lifespan* (which
    # every entrypoint runs — uvicorn, the Docker CMD that launches
    # ``llm.serving.api:app`` directly, and the ``llm-serve`` CLI) rather
    # than only in ``cli.main``: a Docker/uvicorn launch that skipped the CLI
    # used to bind 0.0.0.0 with ``api_key=None`` and serve /generate fully
    # anonymously by default.
    #
    # The bind and the guard now share ONE address: the Docker CMD binds
    # ``$LLM_SERVING_HOST`` and this lifespan validates the same
    # ``ServingConfig.host``, so the fail-closed check fires for direct
    # uvicorn launches too (RIL ISS-164).
    from llm.serving.auth import assert_safe_bind

    assert_safe_bind(config.host, config.api_key)

    generation_service = ServingGenerationService.from_config(config)
    model = generation_service.model
    tokenizer = generation_service.tokenizer
    # The ``batched`` backend owns its own ContinuousBatchingEngine (built
    # inside ``from_config`` so ``generation_backend=batched`` can start the
    # server at all — RIL ISS-150); reuse it for the step observer and
    # shutdown instead of building a second, unwired engine the service never
    # uses. Every other backend has no engine, so build one for the observer/
    # unload wiring as before.
    engine = getattr(generation_service.backend, "engine", None) or ContinuousBatchingEngine.from_serving_config(
        config,
        model=model,
        tokenizer=tokenizer,
    )

    # Publish ``llm_batch_fill_ratio`` after every ``engine.step()``.
    # Called under the engine's step lock so observers see a consistent
    # post-step snapshot of the slot allocator.
    # ``record_batch_fill_ratio`` takes keyword-only args while the engine
    # invokes observers with a positional ``StepStats``; adapt at the wiring
    # point so every step publishes the metric.
    engine.set_step_observer(_step_observer)

    # Inject dependencies into the routers. Done here (not at import time)
    # so unit tests can replace them via ``monkeypatch.setattr`` on the
    # module-level attributes without instantiating a real model.
    generate_router.configure(config, generation_service, inference_semaphore, METRICS)
    chat_router.configure(config, inference_semaphore, METRICS)

    _log_server_config(generation_service, config)

    try:
        yield
    finally:
        engine.unload_model()
        generation_service.shutdown()
        logger.info("Shutting down...")

Configuration

config

ServingConfig

Bases: BaseSettings

Serving Configuration using environment variables.

源代码位于： src/llm/serving/config.py 
  9
 10
 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
 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
 85
 86
 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
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
class ServingConfig(BaseSettings):
    """
    Serving Configuration using environment variables.
    """

    # Model configuration
    model_path: str | None = None  # Path to training checkpoint (.pt)
    tokenizer_path: str | None = None  # Path to tokenizer pickle or HF repo
    tokenizer_type: str = Field("simple", pattern="^(simple|hf)$")
    device: str = "auto"

    # Security & Observability
    api_key: str | None = None  # If set, requires this key for access
    log_level: str = "INFO"
    host: str = "127.0.0.1"  # Set LLM_SERVING_HOST=0.0.0.0 for container bind-all

    # Generation
    generation_backend: str = "eager"  # eager | batched

    # Performance. Kept > 0 deliberately: 0 concurrency builds
    # ``asyncio.Semaphore(0)`` (every request blocks then 504s), a
    # non-positive timeout makes every request time out instantly, and
    # ``max_blocks``/``block_size``/``max_prefixes`` feed directly into
    # block arithmetic that divides by ``block_size``.
    compile_model: bool = False  # Enable torch.compile for acceleration
    max_concurrent_requests: int = Field(4, ge=1, description="Max concurrent inference requests")
    max_request_bytes: int = Field(
        16 * 1024 * 1024,
        ge=1,
        description=(
            "Maximum accepted request-body size in bytes; larger bodies get "
            "413 before the model does any work (RIL ISS-171)."
        ),
    )
    request_timeout: float = Field(60.0, gt=0.0, description="Request timeout in seconds")

    # Chat template (OpenAI /v1/chat/completions).
    # Each message is rendered using ``chat_message_template.format(role=...,
    # content=...)``. The rendered messages are joined with newlines and the
    # ``chat_generation_prefix`` is appended at the end so the model knows
    # where to start producing assistant tokens. Override either field to
    # match your fine-tuned model's expected format (e.g., ChatML,
    # Llama-2-chat, Vicuna). None means: use the built-in defaults.
    chat_message_template: str | None = Field(
        default=None,
        description=(
            "Python format string applied to each chat message. "
            "Available placeholders: {role}, {content}. "
            "If None, falls back to '{role}: {content}'."
        ),
    )
    chat_generation_prefix: str | None = Field(
        default=None,
        description=(
            "String appended after all messages, signaling the model to "
            "start generating the assistant response. "
            "If None, falls back to 'Assistant: '."
        ),
    )

    # Paged Attention (block allocator sidecar; model forward still uses KVCache)
    use_paged_attention: bool = False
    max_blocks: int = Field(256, ge=1)
    block_size: int = Field(16, ge=1)

    # Prefix Cache
    enable_prefix_cache: bool = False
    max_prefixes: int = Field(10, ge=1)

    # Model Params (for dummy init if no ckpt)
    hidden_size: int = 64
    num_layers: int = 2
    num_heads: int = 4
    max_seq_len: int = 128

    # Advanced Arch Params
    num_kv_heads: int | None = None
    num_experts: int = 0
    top_k: int = 0
    attn_impl: str = "mha"
    mlp_impl: str = "mlp"

    # PEFT adapter loading (T2 PEFT #49).
    # Bridges the training-side PEFT save/load surface into the serving
    # loader: trained adapters (LoRA / IA³ / BitFit / Adapter / Pfeiffer
    # / AdaLoRA / QLoRA / Prefix Tuning) are applied to the base model
    # at startup and (optionally) folded into the base weights.
    #
    # Typical workflow:
    #   1. Train with ``TrainingConfig.peft_method="lora"`` +
    #      ``peft_save_path=...`` — the trainer's
    #      ``PEFTAdapterCheckpointCallback`` writes the sidecar on
    #      ``on_train_end``.
    #   2. Serve with ``LLM_SERVING_PEFT_METHOD=lora`` +
    #      ``LLM_SERVING_PEFT_ADAPTER_PATH=...`` — the loader applies
    #      the method and copies the saved adapter values into the
    #      model before the first request.
    peft_method: str | None = Field(
        default=None,
        description=(
            "Registered PEFT method name (e.g. 'lora', 'ia3', 'bitfit'). "
            "When set, the loader applies the method to the base model "
            "after loading the checkpoint. Must be a key in PEFT_REGISTRY."
        ),
    )
    peft_kwargs: dict[str, Any] = Field(
        default_factory=dict,
        description=(
            "Kwargs forwarded to ``apply_peft`` (e.g. ``{'rank': 8, "
            "'alpha': 16.0}`` for LoRA). Overridden by ``peft_kwargs`` "
            "stored inside the sidecar when ``peft_adapter_path`` is set."
        ),
    )
    peft_adapter_path: str | None = Field(
        default=None,
        description=(
            "Path to a PEFT sidecar file written by ``save_peft`` (typically "
            "``{checkpoint_dir}/peft_adapter_{method}.bin`` from "
            "``PEFTAdapterCheckpointCallback``). When set, the loader "
            "loads the saved adapter values into the model after applying "
            "the method."
        ),
    )
    peft_merge: bool = Field(
        default=False,
        description=(
            "If True, fold the adapter into the base weights at serve "
            "time (saves per-token routing overhead). Loses the ability "
            "to ``disable_peft`` / ``unmerge_peft`` at runtime. Refused "
            "for methods that don't expose a merge helper (bitfit / "
            "qlora / prefix_tuning)."
        ),
    )

    model_config = SettingsConfigDict(env_prefix="LLM_SERVING_")

    @classmethod
    def from_yaml(cls, path: str | Path) -> ServingConfig:
        """Load configuration from a YAML file.

        Shaped after :meth:`llm.training.core.config.Config.from_yaml`, with
        one deliberate divergence (RIL DEC-106): a missing file returns the
        default :class:`ServingConfig` (``llm-serve`` treats an absent config
        as "all defaults"), whereas the training side **raises
        FileNotFoundError** on a missing ``--config-path`` (an explicit path
        is always a user error there). Both behaviours are pinned by tests.

        The YAML keys are the unprefixed field names (``api_key``,
        ``model_path``, ``peft_method``, ...). The same config can be set
        via env vars (``LLM_SERVING_API_KEY``, ``LLM_SERVING_MODEL_PATH``,
        ``LLM_SERVING_PEFT_METHOD``, ...).

        Returns:
            A :class:`ServingConfig` instance built from the YAML content —
            or the default config when the path does not exist.
            Validators (``_validate_peft_method``, cross-field consistency)
            run on construction, so unknown PEFT methods and
            inconsistent ``peft_*`` combinations fail loud at load time.
        """
        path = Path(path)
        if not path.exists():
            return cls()
        with path.open(encoding="utf-8") as f:
            config_dict = yaml.safe_load(f) or {}
        return cls.model_validate(config_dict)

    # Methods that don't expose a merge helper. ``peft_merge=True`` is
    # rejected for these in the model_validator below — failing loud at
    # startup is better than silently no-op'ing at first-request time.
    _NON_MERGEABLE_METHODS: frozenset[str] = frozenset({"bitfit", "qlora", "prefix_tuning"})

    @field_validator("peft_method")
    @classmethod
    def _validate_peft_method(cls, value: str | None) -> str | None:
        """Reject unknown PEFT method names at config-load time.

        Built-ins are registered lazily — calling
        :func:`ensure_methods_registered` here is idempotent so the
        validator triggers the bootstrap without surprising the caller
        (subsequent calls return immediately).
        """
        if value is None:
            return None
        # Lazy import to avoid pulling PEFT into the serving import
        # graph just to instantiate the config (the auth / CLI guards
        # build a config at startup with no PEFT).
        from llm.core.peft import PEFT_REGISTRY
        from llm.core.peft.registry import ensure_methods_registered

        ensure_methods_registered()
        try:
            PEFT_REGISTRY.get(value)
        except ValueError as exc:
            raise ValueError(f"Unknown PEFT method {value!r}. Registered methods: {PEFT_REGISTRY.names()}.") from exc
        return value

    @model_validator(mode="after")
    def _validate_peft_field_consistency(self) -> ServingConfig:
        """Cross-field PEFT validation.

        - ``peft_adapter_path`` and ``peft_kwargs`` imply a method is set.
        - ``peft_merge=True`` requires a method that exposes a merge
          helper (raises ``ValueError`` for bitfit / qlora /
          prefix_tuning).
        """
        if self.peft_method is None:
            if self.peft_adapter_path is not None:
                raise ValueError(
                    "peft_adapter_path is set but peft_method is None. "
                    "Set peft_method (e.g. 'lora') to identify which "
                    "method the sidecar was saved with."
                )
            if self.peft_kwargs:
                raise ValueError(
                    "peft_kwargs is set but peft_method is None. Set peft_method (e.g. 'lora') to use the kwargs."
                )
            return self

        if self.peft_merge and self.peft_method in self._NON_MERGEABLE_METHODS:
            raise ValueError(
                f"peft_merge=True is not supported for method "
                f"{self.peft_method!r}: it does not expose a merge "
                f"helper. Either set peft_merge=False or pick a "
                f"different method."
            )
        return self

from_yaml classmethod 

from_yaml(path)

Load configuration from a YAML file.

Shaped after :meth:llm.training.core.config.Config.from_yaml, with one deliberate divergence (RIL DEC-106): a missing file returns the default :class:ServingConfig (llm-serve treats an absent config as "all defaults"), whereas the training side raises FileNotFoundError on a missing --config-path (an explicit path is always a user error there). Both behaviours are pinned by tests.

The YAML keys are the unprefixed field names (api_key, model_path, peft_method, ...). The same config can be set via env vars (LLM_SERVING_API_KEY, LLM_SERVING_MODEL_PATH, LLM_SERVING_PEFT_METHOD, ...).

返回：

名称 类型 描述
A ServingConfig

class:ServingConfig instance built from the YAML content —
ServingConfig

or the default config when the path does not exist.
ServingConfig

Validators (_validate_peft_method, cross-field consistency)
ServingConfig

run on construction, so unknown PEFT methods and
ServingConfig

inconsistent peft_* combinations fail loud at load time.
源代码位于： src/llm/serving/config.py 
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
@classmethod
def from_yaml(cls, path: str | Path) -> ServingConfig:
    """Load configuration from a YAML file.

    Shaped after :meth:`llm.training.core.config.Config.from_yaml`, with
    one deliberate divergence (RIL DEC-106): a missing file returns the
    default :class:`ServingConfig` (``llm-serve`` treats an absent config
    as "all defaults"), whereas the training side **raises
    FileNotFoundError** on a missing ``--config-path`` (an explicit path
    is always a user error there). Both behaviours are pinned by tests.

    The YAML keys are the unprefixed field names (``api_key``,
    ``model_path``, ``peft_method``, ...). The same config can be set
    via env vars (``LLM_SERVING_API_KEY``, ``LLM_SERVING_MODEL_PATH``,
    ``LLM_SERVING_PEFT_METHOD``, ...).

    Returns:
        A :class:`ServingConfig` instance built from the YAML content —
        or the default config when the path does not exist.
        Validators (``_validate_peft_method``, cross-field consistency)
        run on construction, so unknown PEFT methods and
        inconsistent ``peft_*`` combinations fail loud at load time.
    """
    path = Path(path)
    if not path.exists():
        return cls()
    with path.open(encoding="utf-8") as f:
        config_dict = yaml.safe_load(f) or {}
    return cls.model_validate(config_dict)

Continuous Batching Engine

batch_engine

StepStats dataclass

Per-step stats returned by :meth:ContinuousBatchingEngine.step.

scheduled is the number of sequences that ran a forward pass in the step (a.k.a. effective batch size). total_active_slots is the engine's full slot pool — used as the denominator for the llm_batch_fill_ratio Prometheus gauge.

源代码位于： src/llm/serving/batch_engine.py 
43
44
45
46
47
48
49
50
51
52
53
54
@dataclass(frozen=True)
class StepStats:
    """Per-step stats returned by :meth:`ContinuousBatchingEngine.step`.

    ``scheduled`` is the number of sequences that ran a forward pass in
    the step (a.k.a. effective batch size). ``total_active_slots`` is the
    engine's full slot pool — used as the denominator for the
    ``llm_batch_fill_ratio`` Prometheus gauge.
    """

    scheduled: int
    total_active_slots: int

SlotPrefixCache

Maps token prefixes to KV cache slots for reuse across requests.

源代码位于： src/llm/serving/batch_engine.py 
 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
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
class SlotPrefixCache:
    """Maps token prefixes to KV cache slots for reuse across requests."""

    def __init__(self, max_prefixes: int = 10, min_prefix_len: int = 4) -> None:
        self.max_prefixes = max_prefixes
        self.min_prefix_len = min_prefix_len
        self._entries: OrderedDict[str, tuple[int, int]] = OrderedDict()

    @staticmethod
    def hash_tokens(tokens: list[int]) -> str:
        # Token ids from real tokenizers (BPE, SentencePiece) routinely
        # exceed 255 (vocabularies of 50 k-128 k are the norm).  ``bytes()``
        # rejects values outside ``[0, 256)`` and would crash the prefix
        # cache on the first request with a large vocabulary.  Use
        # ``array.array('i', ...)`` to pack each id as a 4-byte C int —
        # deterministic, collision-free for distinct lists, and works
        # across the full ``int`` range.
        return hashlib.sha256(array.array("i", tokens).tobytes()).hexdigest()

    def get(self, tokens: list[int]) -> tuple[int, int] | None:
        if len(tokens) < self.min_prefix_len:
            return None
        key = self.hash_tokens(tokens)
        if key in self._entries:
            # Promote to most-recently-used so LRU eviction evicts the
            # correct entry.
            self._entries.move_to_end(key)
            return self._entries[key]
        return None

    def put(self, tokens: list[int], slot: int, prefix_len: int) -> None:
        if len(tokens) < self.min_prefix_len:
            return
        key = self.hash_tokens(tokens)
        if len(self._entries) >= self.max_prefixes and key not in self._entries:
            self._entries.popitem(last=False)
        self._entries[key] = (slot, prefix_len)
        self._entries.move_to_end(key)

    def invalidate_for_slot(self, slot: int) -> None:
        """Drop every prefix entry that points at ``slot``.

        When a sequence finishes, its KV slot returns to the free pool and a
        later request may be allocated the same slot, overwriting the cached
        K/V.  If the stale prefix entry were left in place, a later request
        with the *same* prompt would hit the cache and replay another
        request's (now-overwritten or in-flight) K/V as its own prefix —
        a use-after-free of the cached KV.  Entries are removed rather than
        re-pointed so the most-recent-usage ordering is unaffected for the
        remaining prefixes.
        """
        if not self._entries:
            return
        stale = [key for key, (cached_slot, _len) in self._entries.items() if cached_slot == slot]
        for key in stale:
            del self._entries[key]

invalidate_for_slot 

invalidate_for_slot(slot)

Drop every prefix entry that points at slot.

When a sequence finishes, its KV slot returns to the free pool and a later request may be allocated the same slot, overwriting the cached K/V. If the stale prefix entry were left in place, a later request with the same prompt would hit the cache and replay another request's (now-overwritten or in-flight) K/V as its own prefix — a use-after-free of the cached KV. Entries are removed rather than re-pointed so the most-recent-usage ordering is unaffected for the remaining prefixes.

源代码位于： src/llm/serving/batch_engine.py 
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
def invalidate_for_slot(self, slot: int) -> None:
    """Drop every prefix entry that points at ``slot``.

    When a sequence finishes, its KV slot returns to the free pool and a
    later request may be allocated the same slot, overwriting the cached
    K/V.  If the stale prefix entry were left in place, a later request
    with the *same* prompt would hit the cache and replay another
    request's (now-overwritten or in-flight) K/V as its own prefix —
    a use-after-free of the cached KV.  Entries are removed rather than
    re-pointed so the most-recent-usage ordering is unaffected for the
    remaining prefixes.
    """
    if not self._entries:
        return
    stale = [key for key, (cached_slot, _len) in self._entries.items() if cached_slot == slot]
    for key in stale:
        del self._entries[key]

SlotAllocator

Manages allocation of KV cache slots.

源代码位于： src/llm/serving/batch_engine.py 
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
class SlotAllocator:
    """Manages allocation of KV cache slots."""

    def __init__(self, total_slots: int):
        self.total_slots = total_slots
        self.free_slots: set[int] = set(range(total_slots))
        self.seq_to_slot: dict[str, int] = {}  # request_id -> slot_id

    @property
    def num_free(self) -> int:
        """Number of available slots."""
        return len(self.free_slots)

    def allocate(self, request_id: str) -> int:
        if request_id in self.seq_to_slot:
            return self.seq_to_slot[request_id]
        if not self.free_slots:
            raise RuntimeError("No free slots available in KV cache.")
        slot = self.free_slots.pop()
        self.seq_to_slot[request_id] = slot
        return slot

    def free(self, request_id: str):
        if request_id in self.seq_to_slot:
            slot = self.seq_to_slot.pop(request_id)
            self.free_slots.add(slot)

    def get_slot(self, request_id: str) -> int:
        return self.seq_to_slot.get(request_id, -1)

num_free property 

num_free

Number of available slots.

ContinuousBatchingEngine

Inference engine supporting continuous batching (iteration-level scheduling).

This is the primary serving engine. It manages request states, schedules sequences at an iteration level, and orchestrates the forward pass.

源代码位于： src/llm/serving/batch_engine.py 
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
class ContinuousBatchingEngine:
    """
    Inference engine supporting continuous batching (iteration-level scheduling).

    This is the primary serving engine. It manages request states, schedules
    sequences at an iteration level, and orchestrates the forward pass.
    """

    def __init__(
        self,
        model: DecoderModel,
        tokenizer: BaseTokenizer,
        device: str | torch.device = "cuda",
        max_batch_size: int = 16,
        max_seq_len: int = 512,
        dtype: torch.dtype = torch.float16,
        *,
        enable_prefix_cache: bool = False,
        max_prefixes: int = 10,
        use_paged_attention: bool = False,
        max_blocks: int = 256,
        block_size: int = 16,
    ):
        """
        Initialize the engine with an already-loaded model and tokenizer.

        Args:
            model: The loaded DecoderModel instance.
            tokenizer: The tokenizer instance (must have encode/decode methods).
            device: Target device ("cuda", "cpu", or torch.device).
            max_batch_size: Maximum number of concurrent sequences.
            max_seq_len: Maximum sequence length for KV cache.
            dtype: Data type for model and cache.
        """
        if isinstance(device, str):
            if device == "auto":
                device = "cuda" if torch.cuda.is_available() else "cpu"
            self.device = torch.device(device)
        else:
            self.device = device

        self.dtype = dtype
        self.max_batch_size = max_batch_size
        self.max_seq_len = max_seq_len
        self.tokenizer = tokenizer

        # Model setup
        self.model = model
        self.model.to(self.device, dtype=self.dtype)
        self.model.eval()

        # Scheduler and Slot Allocator
        self.scheduler = Scheduler(max_batch_size=max_batch_size)
        self.slot_allocator = SlotAllocator(total_slots=max_batch_size)

        # Initialize KV Cache Pool. The dense ``KVCache`` pool is only
        # built when paged attention is disabled — when enabled, the
        # block-allocator pool below replaces it for the model forward
        # path (and building both would waste memory).
        self.kv_caches: list[KVCache] = []
        if not use_paged_attention:
            first_attn = self.model.transformer_blocks[0].self_attn
            if not isinstance(first_attn, _CacheSizedAttention):
                raise TypeError(f"attention backend {type(first_attn).__name__} must expose num_kv_heads/head_dim")
            self.kv_caches = KVCache.from_model_config(
                max_batch_size=self.max_batch_size,
                max_seq_len=self.max_seq_len,
                num_layers=len(self.model.transformer_blocks),
                num_kv_heads=first_attn.num_kv_heads,
                head_dim=first_attn.head_dim,
                device=self.device,
                dtype=self.dtype,
            )

        self.enable_prefix_cache = enable_prefix_cache
        self.prefix_cache = SlotPrefixCache(max_prefixes=max_prefixes) if enable_prefix_cache else None
        self.paged_kv_cache = None
        if use_paged_attention:
            from llm.core.paged_attention.paged_kv_cache import PagedKVCache

            paged_attn = self.model.transformer_blocks[0].self_attn
            if not isinstance(paged_attn, _CacheSizedAttention):
                raise TypeError(f"attention backend {type(paged_attn).__name__} must expose num_kv_heads/head_dim")
            self.paged_kv_cache = PagedKVCache(
                num_layers=len(self.model.transformer_blocks),
                num_kv_heads=paged_attn.num_kv_heads,
                head_dim=paged_attn.head_dim,
                num_blocks=max_blocks,
                block_size=block_size,
                device=str(self.device),
                dtype=self.dtype,
                enable_prefix_cache=enable_prefix_cache,
                max_prefixes=max_prefixes,
            )

        # Concurrency control. ``step()`` mutates Python bookkeeping
        # (``self._seq_len``, ``self.free_slots``, ``self.kv_caches``, prefix
        # cache) that is not thread-safe. FastAPI's ``run_in_threadpool`` calls
        # ``service.generate`` from multiple worker threads, so we serialize.
        # PyTorch CUDA ops have their own internal serialization; this lock
        # only guards the Python-side state machine. Future async refactors
        # should release this lock during the inner model forward.
        self._step_lock = threading.Lock()
        # Cache for the served model's sparse/streaming pattern mask (RIL
        # TASK-247): built once on first use and reused across decode steps
        # instead of recomputing a ``[k_len, k_len]`` boolean mask every step.
        # Immutable for the engine's lifetime, keyed on (scheme, k_len) so a
        # scheme change or config resize invalidates it.
        self._sparse_mask_cache: tuple | None = None
        # Optional callback invoked once per ``step()`` with the resulting
        # :class:`StepStats`. The serving tier uses it to publish
        # ``llm_batch_fill_ratio``. Called under ``self._step_lock`` so the
        # callback sees consistent post-step state.
        self._on_step: Callable[[StepStats], None] | None = None

    def _sparse_pattern_mask(self, k_len: int) -> torch.Tensor | None:
        """Return the served model's sparse pattern mask-out tensor over ``k_len``
        keys (``True`` = mask out), cached across steps; ``None`` when the model
        has no ``attn_sparse`` scheme.

        The pattern is folded into the per-step causal ``run_attn_mask`` (RIL
        TASK-246). It is absolute-position based and constant for the engine's
        lifetime, so it is built once and cached keyed on (scheme, k_len) — a
        scheme change or a config resize invalidates the cache (RIL TASK-247).
        The cached tensor lives on ``self.device`` so the OR with the causal
        mask is device-correct (the builders return CPU tensors).
        """
        spar = getattr(self.model, "attn_sparse", None) if self.model is not None else None
        if not spar:
            self._sparse_mask_cache = None
            return None
        key = (tuple(sorted(spar.items())), k_len)
        cached = self._sparse_mask_cache
        if cached is not None and cached[0] == key:
            return cached[1]
        from llm.core.attn.sparse import build_sparse_attention_mask

        _params = dict(spar)
        _kind = _params.pop("kind")
        _params.pop("causal", None)  # causality comes from the causal mask
        _allow = build_sparse_attention_mask(_kind, k_len, causal=False, **_params)
        mask_out = (~_allow).bool().to(self.device)
        self._sparse_mask_cache = (key, mask_out)
        return mask_out

    @classmethod
    def from_serving_config(cls, config, model: DecoderModel, tokenizer: BaseTokenizer) -> ContinuousBatchingEngine:
        """Build an engine from ServingConfig flags.

        Paged Attention is fully wired through the continuous batching
        forward path (``docs/adr/004-paged-attention-serving.md`` was
        flipped to "Accepted" with this slice). When
        ``config.use_paged_attention=True`` the engine builds a
        :class:`PagedKVCache`, passes it to the model forward, and
        frees per-request blocks on sequence completion.
        """
        return cls(
            model=model,
            tokenizer=tokenizer,
            device=config.device,
            max_batch_size=config.max_concurrent_requests,
            max_seq_len=config.max_seq_len,
            enable_prefix_cache=config.enable_prefix_cache,
            max_prefixes=config.max_prefixes,
            use_paged_attention=config.use_paged_attention,
            max_blocks=config.max_blocks,
            block_size=config.block_size,
        )

    def _copy_kv_between_slots(self, src_slot: int, dst_slot: int, length: int) -> None:
        if self.paged_kv_cache is not None:
            # Prefix replay on the paged path does not copy K/V between dense
            # slots — ``_lock_step_pre`` instead stages the cached blocks into
            # the new sequence via ``PagedKVCache.stage_prefix`` (shared, no
            # data copy). Nothing to do here.
            return
        for cache in self.kv_caches:
            cache.k_cache[dst_slot, :, :length, :] = cache.k_cache[src_slot, :, :length, :].clone()
            cache.v_cache[dst_slot, :, :length, :] = cache.v_cache[src_slot, :, :length, :].clone()

    def add_request(self, request: GenerationRequest) -> str:
        """Add a request to the engine."""
        encoded = self.tokenizer.encode(request.prompt)
        if isinstance(encoded, list):
            input_ids = encoded
        elif isinstance(encoded, torch.Tensor):
            input_ids = encoded.tolist()
            if isinstance(input_ids[0], list):
                input_ids = input_ids[0]
        else:
            input_ids = list(encoded)

        # Reject requests that would outgrow the context window BEFORE they
        # are scheduled. The model computes a position id for every prompt and
        # generated token; past the capacity the positional-encoding table is
        # indexed out of bounds, which surfaces on CUDA as a device-side
        # assert that corrupts the CUDA context and can crash the whole
        # serving process, not just this request.
        #
        # The hard ceiling is the MODEL's positional capacity, not the engine's
        # KV-window budget — the engine window may be configured (or resized)
        # larger than the checkpoint's ``max_seq_len``, and a request that fits
        # the engine budget but exceeds the model capacity used to crash in the
        # positional-encoding range check (RIL round-73 serving deep-dive:
        # ``max_seq_len=128`` engine serving a ``max_seq_len=8`` model).
        budget = min(getattr(self.model, "max_seq_len", self.max_seq_len), self.max_seq_len)
        if len(input_ids) + request.max_new_tokens > budget:
            raise ValueError(
                f"Prompt has {len(input_ids)} tokens and max_new_tokens="
                f"{request.max_new_tokens}, but the model's context window is "
                f"{budget}; the request would exceed the context window and "
                "crash the forward pass"
            )

        req_id = request.request_id or uuid.uuid4().hex

        seq = Sequence(
            request_id=req_id,
            prompt=request.prompt,
            input_ids=input_ids,
            status=RequestState.WAITING,
            max_new_tokens=request.max_new_tokens,
            temperature=request.temperature,
            top_k=request.top_k,
            top_p=request.top_p,
            repetition_penalty=request.repetition_penalty,
            frequency_penalty=request.frequency_penalty,
            presence_penalty=request.presence_penalty,
            logit_bias=request.logit_bias,
            stop=request.stop,
        )

        # Reject a *different* request that reuses an active request_id (RIL
        # ISS-123/F3). The slot allocator keys KV slots by request_id, so two
        # distinct requests claiming the same id would both write into the
        # same slot — cross-request contamination plus a freed-slot reuse when
        # the first finishes. We must NOT reject unconditionally: the engine's
        # own ``generate_request`` -> ``stream_request`` re-adds the SAME
        # logical request (double-add contract), and the streaming reap loop
        # explicitly removes every copy. Only a request whose content differs
        # from the active holder's is a genuine collision.
        if request.request_id is not None:
            added = self.scheduler.add_sequence_if_not_conflicting(
                seq,
                matches=lambda existing: (
                    existing.prompt == seq.prompt
                    and existing.max_new_tokens == seq.max_new_tokens
                    and existing.temperature == seq.temperature
                    and existing.top_k == seq.top_k
                    and existing.top_p == seq.top_p
                    and existing.repetition_penalty == seq.repetition_penalty
                    and existing.frequency_penalty == seq.frequency_penalty
                    and existing.presence_penalty == seq.presence_penalty
                    and existing.logit_bias == seq.logit_bias
                    and existing.stop == seq.stop
                ),
            )
            if not added:
                raise ValueError(
                    f"request_id '{req_id}' is already in use by a different active "
                    "request; duplicate request ids would share one KV slot"
                )
        else:
            self.scheduler.add_sequence(seq)
        return req_id

    def _emit_tokens(
        self,
        sequence: Sequence,
        new_token_ids: list[int],
        stops: list[str] | None,
        max_stop_len: int,
        buffer: str,
    ) -> tuple[list[str], bool, str]:
        """Decode ``new_token_ids`` into text chunks, honouring stop sequences.

        Shared by both the finished-sequence drain and the post-step token
        emission paths in :meth:`stream_request`. ``buffer`` carries over
        any un-yielded tail from the previous call (the last
        ``<max_stop_len>`` characters, held for cross-token-boundary suffix
        checking). Returns a triple:

        - ``chunks``: list of text chunks to yield to the caller.
        - ``stop_hit``: True when a stop sequence matched as a suffix — the
          caller should terminate the streaming generator.
        - ``buffer``: remaining un-yielded buffer text for the next call.

        When ``stops`` is None / empty, each decoded chunk is returned
        directly and the buffer is returned unchanged (it stays empty).
        """
        chunks: list[str] = []
        stop_hit = False

        eos_id = getattr(self.tokenizer, "eos_token_id", None)
        for token_id in new_token_ids:
            # The EOS token ends generation; never emit its decoded text
            # (parity with the eager/speculative backends — RIL ISS-96/98).
            # ``_lock_step_post`` already appended it and marked the sequence
            # FINISHED, so the drain here simply stops short of the EOS.
            if eos_id is not None and token_id == eos_id:
                break
            text_chunk = self.tokenizer.decode([token_id])
            if stops and text_chunk:
                buffer += text_chunk
                for s in stops:
                    if buffer.endswith(s):
                        prefix = buffer[: len(buffer) - len(s)]
                        if prefix:
                            chunks.append(prefix)
                        sequence.status = RequestState.FINISHED
                        stop_hit = True
                        return chunks, stop_hit, ""
                # No stop match — yield the safe prefix (everything beyond
                # the last max_stop_len characters) and keep the tail.
                if len(buffer) > max_stop_len:
                    safe_len = len(buffer) - max_stop_len
                    chunks.append(buffer[:safe_len])
                    buffer = buffer[safe_len:]
            else:
                chunks.append(text_chunk)

        return chunks, stop_hit, buffer

    def stream_request(
        self,
        request: GenerationRequest,
    ):
        """Run a request to completion, yielding decoded text chunks.

        Honours ``stop`` sequences: when the accumulated generated text
        (post-prompt) ends with any stop string, generation halts and the
        stop string itself is excluded from the yielded output (OpenAI
        semantics).
        """
        req_id = self.add_request(request)
        from llm.generation.eager import _normalize_stop

        stops = _normalize_stop(request.stop)
        max_stop_len = max((len(s) for s in stops), default=0) if stops else 0
        buffer = ""
        emitted = 0
        try:
            while True:
                seq = self.scheduler.get_sequence(req_id)
                if seq is None:
                    break
                if seq.is_finished():
                    chunks, stop_hit, buffer = self._emit_tokens(
                        seq, seq.generated_ids[emitted:], stops, max_stop_len, buffer
                    )
                    yield from chunks
                    if stop_hit:
                        # A stop-sequence match finished the request/sequence
                        # *outside* a step, so ``_lock_step_post`` never ran its
                        # free path — release the slot here or it leaks (ISS-044).
                        with self._step_lock:
                            self._release_request_slot_by_id(req_id)
                        return
                    if stops and buffer:
                        # The sequence is already finished, so after draining the
                        # tail buffer there is nothing left to stream. Return
                        # (not ``break``): the post-loop ``yield buffer`` would
                        # emit the SAME tail a second time (RIL ISS-054).
                        yield buffer
                    return
                self.step()
                # Reuse the ``seq`` reference captured BEFORE the step (from
                # the ``get_sequence`` at the top of this iteration). The
                # post-step re-fetch was a race: a concurrent generator's
                # ``step()`` -> ``Scheduler.schedule()`` can evict a just-
                # finished sequence from ``running`` between our ``step()``
                # releasing ``_step_lock`` and this ``get_sequence``
                # re-acquiring the scheduler lock, so the re-fetch returned
                # ``None`` and the final step's token(s) were silently
                # dropped (RIL F1 — concurrent stream truncation). The
                # ``Sequence`` object itself is never destroyed until the
                # ``finally`` reap below, and ``step()`` appended this
                # iteration's tokens to ``seq.generated_ids`` in place, so
                # the already-held reference drains them correctly.
                chunks, stop_hit, buffer = self._emit_tokens(
                    seq, seq.generated_ids[emitted:], stops, max_stop_len, buffer
                )
                yield from chunks
                if stop_hit:
                    # Same leak vector as above.
                    with self._step_lock:
                        self._release_request_slot_by_id(req_id)
                    return
                emitted = len(seq.generated_ids)
                if seq.is_finished():
                    break
            if stops and buffer:
                yield buffer
        finally:
            # Cleanup for ANY exit path — including an *abandoned* generator.
            # ``stream_request`` is the sequence's only stepper; when the
            # consumer disconnects (gen.close() / GC) it can never advance the
            # sequence again, so a RUNNING sequence with its slot allocated
            # would otherwise leak one ``max_batch_size`` slot per disconnect
            # (RIL ISS-105). Reap the sequence and release its slot here;
            # ``_release_request_slot_by_id`` is a no-op if it was already
            # freed (normal completion, stop-halt, forward failure). Loop
            # because a caller may have added the request twice (a stray
            # ``add_request`` before streaming) — every copy must be reaped.
            #
            # ``scheduler.remove`` mutates the LIVE ``running``/``waiting``
            # lists, so it MUST run under ``_step_lock`` like every other
            # engine mutation — otherwise it races with ``_lock_step_pre``
            # iterating the ``running`` list a concurrent ``step()`` returned,
            # and the mid-iteration pop silently skips/duplicates a sequence
            # for that step (RIL ISS-117, a regression the ISS-105 fix
            # originally left un-serialized).
            with self._step_lock:
                while self.scheduler.remove(req_id) is not None:
                    self._release_request_slot_by_id(req_id)

    def generate_request(self, request: GenerationRequest) -> str:
        """Run a request to completion and return prompt + generated text."""
        chunks = list(self.stream_request(request))
        return request.prompt + "".join(chunks)

    def batch_generate_requests(self, requests: list[GenerationRequest]) -> list[str]:
        """Run multiple requests sequentially through the batching engine."""
        return [self.generate_request(request) for request in requests]

    @torch.no_grad()
    def step(self) -> StepStats:
        """Run one inference step (sync wrapper).

        The whole step — pre-compute, model forward, post-compute — holds
        ``self._step_lock`` so two concurrent ``step()`` calls can never both
        run the forward against the same slots and append two tokens to the
        same sequence from one logical step (a real corruption when the
        ``batched`` backend serves concurrent HTTP requests from FastAPI's
        threadpool). Holding the lock across the forward does NOT serialize
        request enqueueing: :meth:`add_request` never takes the step lock, so
        new requests still arrive in parallel. The forward is the only code
        that mutates the KV caches / slot bookkeeping the post step reads, so
        this is exactly what the lock needs to guard.

        Returns:
            :class:`StepStats` describing the step. ``scheduled`` is the
            effective batch size; ``total_active_slots`` is the engine's
            full slot pool (denominator for ``llm_batch_fill_ratio``).
        """
        with self._step_lock:
            inputs = self._lock_step_pre()
            if inputs is None:
                stats = StepStats(scheduled=0, total_active_slots=self.slot_allocator.total_slots)
            else:
                result = self._forward_and_sample(inputs)
                stats = self._lock_step_post(result)
            if self._on_step is not None:
                self._on_step(stats)
        return stats

    async def step_async(self) -> StepStats:
        """Run one inference step, yielding to the event loop during the forward.

        Identical contract to :meth:`step`, but the model forward runs
        in a worker thread via :func:`asyncio.to_thread`. The lock is
        only held for the bookkeeping portions (pre + post). This lets
        the FastAPI event loop keep processing I/O (other requests,
        health checks, /metrics scrapes) while a forward pass runs.

        .. warning::
           On the PAGED path (``paged_kv_cache`` set) the forward mutates the
           block manager (allocation / extension / copy-on-write of shared
           prefix blocks — RIL TASK-065) which is NOT thread-safe. Two
           overlapping ``step_async`` calls would interleave those mutations
           and corrupt the block table. Production serving uses the
           synchronous :meth:`step` (which holds the lock across the whole
           forward) via ``run_in_threadpool``; callers that enable paged
           attention + prefix caching must use :meth:`step`, not
           ``step_async``, until the cache is externally synchronized.

        Returns:
            :class:`StepStats` (same fields as :meth:`step`).
        """
        with self._step_lock:
            inputs = self._lock_step_pre()
        if inputs is None:
            stats = StepStats(scheduled=0, total_active_slots=self.slot_allocator.total_slots)
        else:
            result = await asyncio.to_thread(self._forward_and_sample, inputs)
            with self._step_lock:
                stats = self._lock_step_post(result)
        if self._on_step is not None:
            with self._step_lock:
                self._on_step(stats)
        return stats

    def _lock_step_pre(self) -> _StepInputs | None:
        """Acquire work from the scheduler and build the dense batch.

        Caller MUST hold ``self._step_lock``. Returns ``None`` when
        there is no work to do (idle engine).
        """
        running_sequences = self.scheduler.schedule()
        if not running_sequences:
            return None

        batch_size = len(running_sequences)

        batch_input_ids_list: list[list[int]] = []
        batch_position_ids_list: list[list[int]] = []
        batch_slots_list: list[int] = []
        seq_input_lengths: list[int] = []
        prefix_full_hits: list[bool] = []

        for seq in running_sequences:
            slot = self.slot_allocator.allocate(seq.request_id)
            batch_slots_list.append(slot)
            prefix_full_hit = False

            if len(seq.generated_ids) == 0:
                # Dense path prefix replay: ``_copy_kv_between_slots`` copies
                # the cached K/V into a fresh slot and only the final prompt
                # token runs through the model.
                use_prefix_shortcut = self.prefix_cache is not None and self.paged_kv_cache is None
                cached = self.prefix_cache.get(seq.input_ids) if use_prefix_shortcut else None
                if cached is not None and cached[1] == len(seq.input_ids):
                    src_slot, prefix_len = cached
                    if src_slot != slot:
                        self._copy_kv_between_slots(src_slot, slot, prefix_len)
                    ids = [seq.input_ids[-1]]
                    pos_ids = [prefix_len - 1]
                    prefix_full_hit = True
                elif (
                    self.paged_kv_cache is not None
                    and self.paged_kv_cache.enable_prefix_cache
                    and len(seq.input_ids) > 0
                ):
                    # Paged path prefix replay (RIL TASK-065). An exact
                    # full-prompt match stages the cached blocks into the new
                    # sequence SHARED (refcounted — ``stage_prefix`` forks
                    # them) and runs only the final prompt token, mirroring
                    # the dense shortcut: the model re-writes that token's K/V
                    # at position N-1, which is idempotent with the cached
                    # value, and ``PagedKVCache.update`` copy-on-writes if the
                    # boundary block is still shared so the cache owner's K/V
                    # is never corrupted.
                    prefix_blocks = self.paged_kv_cache.try_get_prefix_blocks(seq.input_ids)
                    if prefix_blocks is not None:
                        self.paged_kv_cache.stage_prefix(slot, prefix_blocks, len(seq.input_ids) - 1)
                        ids = [seq.input_ids[-1]]
                        pos_ids = [len(seq.input_ids) - 1]
                        prefix_full_hit = True
                    else:
                        ids = seq.input_ids
                        pos_ids = list(range(len(ids)))
                else:
                    ids = seq.input_ids
                    pos_ids = list(range(len(ids)))
            else:
                ids = [seq.generated_ids[-1]]
                pos_val = seq.total_len - 1
                pos_ids = [pos_val]

            batch_input_ids_list.append(ids)
            batch_position_ids_list.append(pos_ids)
            seq_input_lengths.append(len(ids))
            prefix_full_hits.append(prefix_full_hit)

        max_len = max(seq_input_lengths)

        padded_input_ids = torch.zeros((batch_size, max_len), dtype=torch.long, device=self.device)
        padded_position_ids = torch.zeros((batch_size, max_len), dtype=torch.long, device=self.device)
        batch_indices = torch.tensor(batch_slots_list, dtype=torch.long, device=self.device)

        pad_id = 0
        if hasattr(self.tokenizer, "pad_token_id") and self.tokenizer.pad_token_id is not None:
            pad_id = self.tokenizer.pad_token_id

        padded_input_ids.fill_(pad_id)

        q_len = max_len
        k_len = self.max_seq_len

        col_indices = torch.arange(k_len, device=self.device).reshape(1, 1, 1, -1)
        # The causal attention mask must be built from the *real* position
        # ids, which are only known after the per-row fill below. Building
        # it from the freshly-allocated (all-zero) ``padded_position_ids``
        # collapsed every query row to q_pos=0, so decode attention could
        # only ever see the first prompt token's KV and outputs diverged
        # from the eager backend after a few decode steps. All-True
        # (masked) by default; the visible causal region for each query
        # row is set after the fill.
        run_attn_mask = torch.ones((batch_size, 1, q_len, k_len), dtype=torch.bool, device=self.device)

        # Sparse/streaming scheme (RIL TASK-246): when the served model carries an
        # ``attn_sparse`` scheme, fold its sink+window pattern into the mask so the
        # batched/paged serving path actually constrains keys instead of silently
        # running dense (the decoder's own auto-mask is bypassed here because the
        # engine always supplies an explicit ``attn_mask``). The pattern is
        # absolute-position based over the full key window; per query row it is
        # indexed by the real position id below and OR-ed with the causal mask.
        sparse_mask_out = self._sparse_pattern_mask(k_len)

        for i, length in enumerate(seq_input_lengths):
            input_row = torch.tensor(batch_input_ids_list[i], dtype=torch.long, device=self.device)
            pos_row = torch.tensor(batch_position_ids_list[i], dtype=torch.long, device=self.device)

            padded_input_ids[i, :length] = input_row
            padded_position_ids[i, :length] = pos_row

            # True = mask out columns > position[s] (the sdpa wrapper's
            # convention), so each query position s sees keys 0..position[s].
            q_pos_row = padded_position_ids[i, :length].reshape(1, 1, length, 1)
            run_attn_mask[i, :, :length, :] = col_indices > q_pos_row
            if sparse_mask_out is not None:
                # Fold the sparse pattern for each query row's real absolute
                # position (sink + window), on top of the causal mask.
                pos = padded_position_ids[i, :length].long()
                run_attn_mask[i, :, :length, :] = run_attn_mask[i, :, :length, :] | sparse_mask_out[pos]

        return _StepInputs(
            running_sequences=running_sequences,
            batch_slots_list=batch_slots_list,
            seq_input_lengths=seq_input_lengths,
            prefix_full_hits=prefix_full_hits,
            padded_input_ids=padded_input_ids,
            padded_position_ids=padded_position_ids,
            batch_indices=batch_indices,
            run_attn_mask=run_attn_mask,
            batch_size=batch_size,
        )

    def _forward_and_sample(self, inputs: _StepInputs) -> _StepResult:
        """Run the model forward and sampling WITHOUT holding the lock.

        This is the expensive path: ~ ms of GPU/CPU work depending on
        batch size and model size. The lock is released for the entire
        duration so other threads can pre-/post-compute in parallel.

        On forward failure we record the exception in the result so
        the caller can free slots + clean up state under the lock
        (so the engine stays consistent even when a forward raises).
        """
        try:
            if self.model is None:
                raise RuntimeError("engine model was unloaded")
            logits, _ = self.model(
                input_ids=inputs.padded_input_ids,
                position_ids=inputs.padded_position_ids,
                kv_caches=self.kv_caches if self.paged_kv_cache is None else None,
                paged_kv_cache=self.paged_kv_cache,
                use_cache=True,
                batch_indices=inputs.batch_indices,
                attn_mask=inputs.run_attn_mask,
            )

            next_token_ids: list[int] = []
            for i, length in enumerate(inputs.seq_input_lengths):
                seq = inputs.running_sequences[i]
                seq_logits = logits[i, length - 1, :]
                # The pad token must never be emitted (the eager backend
                # masks it the same way); without this the engine can
                # sample pad and diverge from eager.
                pad_id = getattr(self.tokenizer, "pad_token_id", None)
                if pad_id is not None and 0 <= pad_id < seq_logits.size(-1):
                    seq_logits = seq_logits.clone()
                    seq_logits[pad_id] = float("-inf")
                # Sample only tokenizer-decodable ids: a padded-vocab or
                # BPE/HF model served with a smaller-vocab tokenizer would
                # otherwise sample a tail id and crash in ``_emit_tokens`` at
                # ``tokenizer.decode([id])`` (KeyError), while eager and
                # speculative mask it on every step (RIL ISS-125).
                mask_undecodable_logits(seq_logits, getattr(self.tokenizer, "vocab_size", None))
                context_ids = seq.input_ids + seq.generated_ids
                if seq.repetition_penalty != 1.0:
                    seq_logits = apply_repetition_penalty(seq_logits, context_ids, seq.repetition_penalty)
                if seq.frequency_penalty != 0.0:
                    seq_logits = apply_frequency_penalty(seq_logits, context_ids, seq.frequency_penalty)
                if seq.presence_penalty != 0.0:
                    seq_logits = apply_presence_penalty(seq_logits, context_ids, seq.presence_penalty)
                if seq.logit_bias:
                    seq_logits = apply_logit_bias(seq_logits, seq.logit_bias)
                next_token_ids.append(
                    sample_next_token(
                        seq_logits,
                        temperature=seq.temperature,
                        top_k=seq.top_k,
                        top_p=seq.top_p,
                    )
                )
        except BaseException as exc:  # noqa: BLE001 - propagate via result
            return _StepResult(inputs=inputs, forward_failed=exc)

        return _StepResult(inputs=inputs, next_token_ids=next_token_ids)

    def _release_request_slots(self, request_id: str, slot: int) -> None:
        """Return a request's KV slot, prefix-cache entry and paged blocks.

        Shared by the finished-in-step path (:meth:`_lock_step_post`), the
        stop-sequence termination path (:meth:`stream_request`) and the
        forward-failure path. Omitting this anywhere leaves the slot
        permanently leaked (RIL ISS-044): the scheduler filters FINISHED
        sequences out of ``running`` without ever freeing their slot, so a
        slot leaked per stop-terminated request eventually exhausts the pool
        and every ``allocate()`` raises ``No free slots available in KV
        cache``.
        """
        self.slot_allocator.free(request_id)
        if self.prefix_cache is not None:
            # The slot is back in the free pool; a later request may reuse it
            # and overwrite the cached K/V. A leftover prefix entry would
            # replay stale/in-flight K/V on a prompt match.
            self.prefix_cache.invalidate_for_slot(slot)
        if self.paged_kv_cache is not None:
            # ``seq_id`` == slot id in the paged path.
            self.paged_kv_cache.free(slot)

    def _release_request_slot_by_id(self, request_id: str) -> None:
        """Release a request's slot, resolving the slot id from the allocator.

        Caller MUST hold ``self._step_lock`` (the same contract as
        :meth:`_lock_step_post`) — this mutates the shared allocator /
        prefix cache / paged pool. No-op when the request no longer holds a
        slot (e.g. it already finished inside a step and was released there).
        """
        slot = self.slot_allocator.get_slot(request_id)
        if slot < 0:
            return
        self._release_request_slots(request_id, slot)

    def _lock_step_post(self, result: _StepResult) -> StepStats:
        """Append sampled tokens, free slots, mark finished sequences.

        Caller MUST hold ``self._step_lock``. If the forward failed,
        we free the slots we allocated in pre but don't append any
        token — the sequences are left in their previous state.
        """
        inputs = result.inputs
        if result.forward_failed is not None:
            # Free the slots we allocated in pre, and mark the sequences
            # FINISHED so the next ``schedule()`` drops them. Without the
            # status flip a persistently-failing sequence (OOM, bad token
            # id, shape mismatch) stays RUNNING and is re-scheduled every
            # step — re-allocated, re-forwarded, re-failed — livelocking
            # the whole engine past the first error (RIL ISS-051).
            for i, seq in enumerate(inputs.running_sequences):
                seq.status = RequestState.FINISHED
                self._release_request_slots(seq.request_id, inputs.batch_slots_list[i])
            raise result.forward_failed

        for i, seq in enumerate(inputs.running_sequences):
            token_id = result.next_token_ids[i]
            seq.append_token_id(token_id)

            # Dense path maintains the dense SlotPrefixCache; the paged path
            # maintains its own block cache via PagedKVCache.add_prefix.
            if (
                self.prefix_cache
                and self.paged_kv_cache is None
                and len(seq.generated_ids) == 1
                and not inputs.prefix_full_hits[i]
            ):
                self.prefix_cache.put(seq.input_ids, inputs.batch_slots_list[i], len(seq.input_ids))

            if (
                self.paged_kv_cache is not None
                and self.paged_kv_cache.enable_prefix_cache
                and len(seq.generated_ids) == 1
                and not inputs.prefix_full_hits[i]
            ):
                # Register the just-prefilled prompt's blocks for future
                # replay. At this instant the block table holds exactly the
                # prompt's blocks (decode has not extended it yet), so the
                # cached prefix is the full prompt. Skipped for a prefix-hit
                # sequence (it shared the owner's blocks — re-registering the
                # same hash to this seq would make its free() wrongly evict
                # the owner's entry).
                slot = inputs.batch_slots_list[i]
                self.paged_kv_cache.add_prefix(slot, seq.input_ids, self.paged_kv_cache.get_block_table(slot))

            if (
                (hasattr(self.tokenizer, "eos_token_id") and token_id == self.tokenizer.eos_token_id)
                or len(seq.generated_ids) >= seq.max_new_tokens
                or seq.total_len >= self.max_seq_len
            ):
                seq.status = RequestState.FINISHED
                self._release_request_slots(seq.request_id, inputs.batch_slots_list[i])

        return StepStats(
            scheduled=inputs.batch_size,
            total_active_slots=self.slot_allocator.total_slots,
        )

    def set_step_observer(self, callback: Callable[[StepStats], None] | None) -> None:
        """Install or clear a per-step observer (used for metric publishing).

        The callback runs at the end of every :meth:`step` call, under
        ``self._step_lock``, with the :class:`StepStats` for that step.
        Pass ``None`` to remove a previously installed observer.
        """
        self._on_step = callback

    def unload_model(self):
        """Release model, KV caches and all scheduler state.

        Clears model weights, both legacy and paged KV caches, prefix
        cache, slot allocator mappings, and scheduler queues so that
        GPU memory is freed and the engine is reusable after this call.
        """
        self.model = None
        self.kv_caches = []
        self.paged_kv_cache = None
        self._on_step = None

        self.scheduler.clear()
        self.slot_allocator.seq_to_slot.clear()
        self.slot_allocator.free_slots = set(range(self.slot_allocator.total_slots))

        if self.prefix_cache is not None:
            self.prefix_cache._entries.clear()

        if self.device.type == "cuda":
            torch.cuda.empty_cache()

from_serving_config classmethod 

from_serving_config(config, model, tokenizer)

Build an engine from ServingConfig flags.

Paged Attention is fully wired through the continuous batching forward path (docs/adr/004-paged-attention-serving.md was flipped to "Accepted" with this slice). When config.use_paged_attention=True the engine builds a :class:PagedKVCache, passes it to the model forward, and frees per-request blocks on sequence completion.

源代码位于： src/llm/serving/batch_engine.py 
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
@classmethod
def from_serving_config(cls, config, model: DecoderModel, tokenizer: BaseTokenizer) -> ContinuousBatchingEngine:
    """Build an engine from ServingConfig flags.

    Paged Attention is fully wired through the continuous batching
    forward path (``docs/adr/004-paged-attention-serving.md`` was
    flipped to "Accepted" with this slice). When
    ``config.use_paged_attention=True`` the engine builds a
    :class:`PagedKVCache`, passes it to the model forward, and
    frees per-request blocks on sequence completion.
    """
    return cls(
        model=model,
        tokenizer=tokenizer,
        device=config.device,
        max_batch_size=config.max_concurrent_requests,
        max_seq_len=config.max_seq_len,
        enable_prefix_cache=config.enable_prefix_cache,
        max_prefixes=config.max_prefixes,
        use_paged_attention=config.use_paged_attention,
        max_blocks=config.max_blocks,
        block_size=config.block_size,
    )

add_request 

add_request(request)

Add a request to the engine.

源代码位于： src/llm/serving/batch_engine.py 
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
def add_request(self, request: GenerationRequest) -> str:
    """Add a request to the engine."""
    encoded = self.tokenizer.encode(request.prompt)
    if isinstance(encoded, list):
        input_ids = encoded
    elif isinstance(encoded, torch.Tensor):
        input_ids = encoded.tolist()
        if isinstance(input_ids[0], list):
            input_ids = input_ids[0]
    else:
        input_ids = list(encoded)

    # Reject requests that would outgrow the context window BEFORE they
    # are scheduled. The model computes a position id for every prompt and
    # generated token; past the capacity the positional-encoding table is
    # indexed out of bounds, which surfaces on CUDA as a device-side
    # assert that corrupts the CUDA context and can crash the whole
    # serving process, not just this request.
    #
    # The hard ceiling is the MODEL's positional capacity, not the engine's
    # KV-window budget — the engine window may be configured (or resized)
    # larger than the checkpoint's ``max_seq_len``, and a request that fits
    # the engine budget but exceeds the model capacity used to crash in the
    # positional-encoding range check (RIL round-73 serving deep-dive:
    # ``max_seq_len=128`` engine serving a ``max_seq_len=8`` model).
    budget = min(getattr(self.model, "max_seq_len", self.max_seq_len), self.max_seq_len)
    if len(input_ids) + request.max_new_tokens > budget:
        raise ValueError(
            f"Prompt has {len(input_ids)} tokens and max_new_tokens="
            f"{request.max_new_tokens}, but the model's context window is "
            f"{budget}; the request would exceed the context window and "
            "crash the forward pass"
        )

    req_id = request.request_id or uuid.uuid4().hex

    seq = Sequence(
        request_id=req_id,
        prompt=request.prompt,
        input_ids=input_ids,
        status=RequestState.WAITING,
        max_new_tokens=request.max_new_tokens,
        temperature=request.temperature,
        top_k=request.top_k,
        top_p=request.top_p,
        repetition_penalty=request.repetition_penalty,
        frequency_penalty=request.frequency_penalty,
        presence_penalty=request.presence_penalty,
        logit_bias=request.logit_bias,
        stop=request.stop,
    )

    # Reject a *different* request that reuses an active request_id (RIL
    # ISS-123/F3). The slot allocator keys KV slots by request_id, so two
    # distinct requests claiming the same id would both write into the
    # same slot — cross-request contamination plus a freed-slot reuse when
    # the first finishes. We must NOT reject unconditionally: the engine's
    # own ``generate_request`` -> ``stream_request`` re-adds the SAME
    # logical request (double-add contract), and the streaming reap loop
    # explicitly removes every copy. Only a request whose content differs
    # from the active holder's is a genuine collision.
    if request.request_id is not None:
        added = self.scheduler.add_sequence_if_not_conflicting(
            seq,
            matches=lambda existing: (
                existing.prompt == seq.prompt
                and existing.max_new_tokens == seq.max_new_tokens
                and existing.temperature == seq.temperature
                and existing.top_k == seq.top_k
                and existing.top_p == seq.top_p
                and existing.repetition_penalty == seq.repetition_penalty
                and existing.frequency_penalty == seq.frequency_penalty
                and existing.presence_penalty == seq.presence_penalty
                and existing.logit_bias == seq.logit_bias
                and existing.stop == seq.stop
            ),
        )
        if not added:
            raise ValueError(
                f"request_id '{req_id}' is already in use by a different active "
                "request; duplicate request ids would share one KV slot"
            )
    else:
        self.scheduler.add_sequence(seq)
    return req_id

stream_request 

stream_request(request)

Run a request to completion, yielding decoded text chunks.

Honours stop sequences: when the accumulated generated text (post-prompt) ends with any stop string, generation halts and the stop string itself is excluded from the yielded output (OpenAI semantics).

源代码位于： src/llm/serving/batch_engine.py 
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
def stream_request(
    self,
    request: GenerationRequest,
):
    """Run a request to completion, yielding decoded text chunks.

    Honours ``stop`` sequences: when the accumulated generated text
    (post-prompt) ends with any stop string, generation halts and the
    stop string itself is excluded from the yielded output (OpenAI
    semantics).
    """
    req_id = self.add_request(request)
    from llm.generation.eager import _normalize_stop

    stops = _normalize_stop(request.stop)
    max_stop_len = max((len(s) for s in stops), default=0) if stops else 0
    buffer = ""
    emitted = 0
    try:
        while True:
            seq = self.scheduler.get_sequence(req_id)
            if seq is None:
                break
            if seq.is_finished():
                chunks, stop_hit, buffer = self._emit_tokens(
                    seq, seq.generated_ids[emitted:], stops, max_stop_len, buffer
                )
                yield from chunks
                if stop_hit:
                    # A stop-sequence match finished the request/sequence
                    # *outside* a step, so ``_lock_step_post`` never ran its
                    # free path — release the slot here or it leaks (ISS-044).
                    with self._step_lock:
                        self._release_request_slot_by_id(req_id)
                    return
                if stops and buffer:
                    # The sequence is already finished, so after draining the
                    # tail buffer there is nothing left to stream. Return
                    # (not ``break``): the post-loop ``yield buffer`` would
                    # emit the SAME tail a second time (RIL ISS-054).
                    yield buffer
                return
            self.step()
            # Reuse the ``seq`` reference captured BEFORE the step (from
            # the ``get_sequence`` at the top of this iteration). The
            # post-step re-fetch was a race: a concurrent generator's
            # ``step()`` -> ``Scheduler.schedule()`` can evict a just-
            # finished sequence from ``running`` between our ``step()``
            # releasing ``_step_lock`` and this ``get_sequence``
            # re-acquiring the scheduler lock, so the re-fetch returned
            # ``None`` and the final step's token(s) were silently
            # dropped (RIL F1 — concurrent stream truncation). The
            # ``Sequence`` object itself is never destroyed until the
            # ``finally`` reap below, and ``step()`` appended this
            # iteration's tokens to ``seq.generated_ids`` in place, so
            # the already-held reference drains them correctly.
            chunks, stop_hit, buffer = self._emit_tokens(
                seq, seq.generated_ids[emitted:], stops, max_stop_len, buffer
            )
            yield from chunks
            if stop_hit:
                # Same leak vector as above.
                with self._step_lock:
                    self._release_request_slot_by_id(req_id)
                return
            emitted = len(seq.generated_ids)
            if seq.is_finished():
                break
        if stops and buffer:
            yield buffer
    finally:
        # Cleanup for ANY exit path — including an *abandoned* generator.
        # ``stream_request`` is the sequence's only stepper; when the
        # consumer disconnects (gen.close() / GC) it can never advance the
        # sequence again, so a RUNNING sequence with its slot allocated
        # would otherwise leak one ``max_batch_size`` slot per disconnect
        # (RIL ISS-105). Reap the sequence and release its slot here;
        # ``_release_request_slot_by_id`` is a no-op if it was already
        # freed (normal completion, stop-halt, forward failure). Loop
        # because a caller may have added the request twice (a stray
        # ``add_request`` before streaming) — every copy must be reaped.
        #
        # ``scheduler.remove`` mutates the LIVE ``running``/``waiting``
        # lists, so it MUST run under ``_step_lock`` like every other
        # engine mutation — otherwise it races with ``_lock_step_pre``
        # iterating the ``running`` list a concurrent ``step()`` returned,
        # and the mid-iteration pop silently skips/duplicates a sequence
        # for that step (RIL ISS-117, a regression the ISS-105 fix
        # originally left un-serialized).
        with self._step_lock:
            while self.scheduler.remove(req_id) is not None:
                self._release_request_slot_by_id(req_id)

generate_request 

generate_request(request)

Run a request to completion and return prompt + generated text.

源代码位于： src/llm/serving/batch_engine.py 
596
597
598
599
def generate_request(self, request: GenerationRequest) -> str:
    """Run a request to completion and return prompt + generated text."""
    chunks = list(self.stream_request(request))
    return request.prompt + "".join(chunks)

batch_generate_requests 

batch_generate_requests(requests)

Run multiple requests sequentially through the batching engine.

源代码位于： src/llm/serving/batch_engine.py 
601
602
603
def batch_generate_requests(self, requests: list[GenerationRequest]) -> list[str]:
    """Run multiple requests sequentially through the batching engine."""
    return [self.generate_request(request) for request in requests]

step 

step()

Run one inference step (sync wrapper).

The whole step — pre-compute, model forward, post-compute — holds self._step_lock so two concurrent step() calls can never both run the forward against the same slots and append two tokens to the same sequence from one logical step (a real corruption when the batched backend serves concurrent HTTP requests from FastAPI's threadpool). Holding the lock across the forward does NOT serialize request enqueueing: :meth:add_request never takes the step lock, so new requests still arrive in parallel. The forward is the only code that mutates the KV caches / slot bookkeeping the post step reads, so this is exactly what the lock needs to guard.

返回：

类型 描述
StepStats

class:StepStats describing the step. scheduled is the
StepStats

effective batch size; total_active_slots is the engine's
StepStats

full slot pool (denominator for llm_batch_fill_ratio).
源代码位于： src/llm/serving/batch_engine.py 
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
@torch.no_grad()
def step(self) -> StepStats:
    """Run one inference step (sync wrapper).

    The whole step — pre-compute, model forward, post-compute — holds
    ``self._step_lock`` so two concurrent ``step()`` calls can never both
    run the forward against the same slots and append two tokens to the
    same sequence from one logical step (a real corruption when the
    ``batched`` backend serves concurrent HTTP requests from FastAPI's
    threadpool). Holding the lock across the forward does NOT serialize
    request enqueueing: :meth:`add_request` never takes the step lock, so
    new requests still arrive in parallel. The forward is the only code
    that mutates the KV caches / slot bookkeeping the post step reads, so
    this is exactly what the lock needs to guard.

    Returns:
        :class:`StepStats` describing the step. ``scheduled`` is the
        effective batch size; ``total_active_slots`` is the engine's
        full slot pool (denominator for ``llm_batch_fill_ratio``).
    """
    with self._step_lock:
        inputs = self._lock_step_pre()
        if inputs is None:
            stats = StepStats(scheduled=0, total_active_slots=self.slot_allocator.total_slots)
        else:
            result = self._forward_and_sample(inputs)
            stats = self._lock_step_post(result)
        if self._on_step is not None:
            self._on_step(stats)
    return stats

step_async async 

step_async()

Run one inference step, yielding to the event loop during the forward.

Identical contract to :meth:step, but the model forward runs in a worker thread via :func:asyncio.to_thread. The lock is only held for the bookkeeping portions (pre + post). This lets the FastAPI event loop keep processing I/O (other requests, health checks, /metrics scrapes) while a forward pass runs.

.. warning:: On the PAGED path (paged_kv_cache set) the forward mutates the block manager (allocation / extension / copy-on-write of shared prefix blocks — RIL TASK-065) which is NOT thread-safe. Two overlapping step_async calls would interleave those mutations and corrupt the block table. Production serving uses the synchronous :meth:step (which holds the lock across the whole forward) via run_in_threadpool; callers that enable paged attention + prefix caching must use :meth:step, not step_async, until the cache is externally synchronized.

返回：

类型 描述
StepStats

class:StepStats (same fields as :meth:step).
源代码位于： src/llm/serving/batch_engine.py 
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
async def step_async(self) -> StepStats:
    """Run one inference step, yielding to the event loop during the forward.

    Identical contract to :meth:`step`, but the model forward runs
    in a worker thread via :func:`asyncio.to_thread`. The lock is
    only held for the bookkeeping portions (pre + post). This lets
    the FastAPI event loop keep processing I/O (other requests,
    health checks, /metrics scrapes) while a forward pass runs.

    .. warning::
       On the PAGED path (``paged_kv_cache`` set) the forward mutates the
       block manager (allocation / extension / copy-on-write of shared
       prefix blocks — RIL TASK-065) which is NOT thread-safe. Two
       overlapping ``step_async`` calls would interleave those mutations
       and corrupt the block table. Production serving uses the
       synchronous :meth:`step` (which holds the lock across the whole
       forward) via ``run_in_threadpool``; callers that enable paged
       attention + prefix caching must use :meth:`step`, not
       ``step_async``, until the cache is externally synchronized.

    Returns:
        :class:`StepStats` (same fields as :meth:`step`).
    """
    with self._step_lock:
        inputs = self._lock_step_pre()
    if inputs is None:
        stats = StepStats(scheduled=0, total_active_slots=self.slot_allocator.total_slots)
    else:
        result = await asyncio.to_thread(self._forward_and_sample, inputs)
        with self._step_lock:
            stats = self._lock_step_post(result)
    if self._on_step is not None:
        with self._step_lock:
            self._on_step(stats)
    return stats

set_step_observer 

set_step_observer(callback)

Install or clear a per-step observer (used for metric publishing).

The callback runs at the end of every :meth:step call, under self._step_lock, with the :class:StepStats for that step. Pass None to remove a previously installed observer.

源代码位于： src/llm/serving/batch_engine.py 
969
970
971
972
973
974
975
976
def set_step_observer(self, callback: Callable[[StepStats], None] | None) -> None:
    """Install or clear a per-step observer (used for metric publishing).

    The callback runs at the end of every :meth:`step` call, under
    ``self._step_lock``, with the :class:`StepStats` for that step.
    Pass ``None`` to remove a previously installed observer.
    """
    self._on_step = callback

unload_model 

unload_model()

Release model, KV caches and all scheduler state.

Clears model weights, both legacy and paged KV caches, prefix cache, slot allocator mappings, and scheduler queues so that GPU memory is freed and the engine is reusable after this call.

源代码位于： src/llm/serving/batch_engine.py 
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
def unload_model(self):
    """Release model, KV caches and all scheduler state.

    Clears model weights, both legacy and paged KV caches, prefix
    cache, slot allocator mappings, and scheduler queues so that
    GPU memory is freed and the engine is reusable after this call.
    """
    self.model = None
    self.kv_caches = []
    self.paged_kv_cache = None
    self._on_step = None

    self.scheduler.clear()
    self.slot_allocator.seq_to_slot.clear()
    self.slot_allocator.free_slots = set(range(self.slot_allocator.total_slots))

    if self.prefix_cache is not None:
        self.prefix_cache._entries.clear()

    if self.device.type == "cuda":
        torch.cuda.empty_cache()

Scheduler

scheduler

Scheduler

Simple FCFS scheduler for continuous batching with backpressure.

Limits the waiting queue to max_waiting to prevent unbounded memory growth under load. When the queue is full, :meth:add_sequence raises RuntimeError so callers can apply backpressure (HTTP 503).

Thread safety: every method that reads or mutates waiting / running takes the scheduler's own :attr:_lock. The engine's step() path calls :meth:schedule under the engine's _step_lock, while streaming/backpressure paths call :meth:add_sequence, :meth:get_sequence and :attr:has_pending_work WITHOUT the engine lock — without a scheduler-internal lock those concurrent popleft / iteration calls crashed with RuntimeError: deque mutated during iteration under concurrent streaming (RIL ISS-069). Lock ordering stays acyclic: schedule (engine lock → scheduler lock) never nests the engine lock inside the scheduler lock.

源代码位于： src/llm/serving/scheduler.py 
  8
  9
 10
 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
 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
 85
 86
 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
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
class Scheduler:
    """Simple FCFS scheduler for continuous batching with backpressure.

    Limits the ``waiting`` queue to ``max_waiting`` to prevent unbounded
    memory growth under load. When the queue is full, :meth:`add_sequence`
    raises ``RuntimeError`` so callers can apply backpressure (HTTP 503).

    Thread safety: every method that reads or mutates ``waiting`` /
    ``running`` takes the scheduler's own :attr:`_lock`. The engine's
    ``step()`` path calls :meth:`schedule` under the engine's
    ``_step_lock``, while streaming/backpressure paths call
    :meth:`add_sequence`, :meth:`get_sequence` and :attr:`has_pending_work`
    WITHOUT the engine lock — without a scheduler-internal lock those
    concurrent ``popleft`` / iteration calls crashed with ``RuntimeError:
    deque mutated during iteration`` under concurrent streaming (RIL
    ISS-069). Lock ordering stays acyclic: ``schedule`` (engine lock →
    scheduler lock) never nests the engine lock inside the scheduler lock.
    """

    DEFAULT_MAX_WAITING = 1024

    def __init__(self, max_batch_size: int = 16, *, max_waiting: int | None = None):
        self.waiting: deque[Sequence] = deque()
        self.running: list[Sequence] = []
        self.max_batch_size = max_batch_size
        self.max_waiting = max_waiting if max_waiting is not None else self.DEFAULT_MAX_WAITING
        self._lock = Lock()

    @property
    def has_pending_work(self) -> bool:
        with self._lock:
            return len(self.waiting) > 0 or len(self.running) > 0

    def add_sequence(self, seq: Sequence):
        """Add a new sequence to the waiting queue.

        Raises:
            RuntimeError: If the waiting queue is at capacity.
        """
        with self._lock:
            if len(self.waiting) >= self.max_waiting:
                raise RuntimeError(
                    f"Waiting queue full ({len(self.waiting)}/{self.max_waiting}); retry later or increase max_waiting."
                )
            self.waiting.append(seq)

    def add_sequence_if_not_conflicting(self, seq: Sequence, matches: Callable[[Sequence], bool] | None = None) -> bool:
        """Add ``seq`` unless an ACTIVE sequence already occupies its request_id.

        Guard for RIL ISS-123: two distinct requests sharing one client
        request_id hash to the SAME KV slot (``SlotAllocator`` keys slots by
        request_id), so their K/V and generated tokens contaminate each other
        and a ``free()`` by either returns the other's live slot to the pool.

        The engine's internal *double-add* contract (``generate_request`` ->
        ``stream_request`` re-adding the SAME logical request so the reap loop
        can remove every copy) stays intact: when an ``matches`` callback is
        supplied, a duplicate whose content equals the active holder is
        allowed (it is the same request re-added). By default (no callback)
        ANY active duplicate is rejected.

        Returns ``True`` when the sequence was enqueued, ``False`` when it was
        rejected as a conflicting duplicate.

        Raises:
            RuntimeError: If the waiting queue is at capacity.
        """
        with self._lock:
            for s in list(self.running) + list(self.waiting):
                if s.request_id == seq.request_id and not s.is_finished():
                    if matches is not None and matches(s):
                        continue  # same logical request re-added — permit
                    return False
            if len(self.waiting) >= self.max_waiting:
                raise RuntimeError(
                    f"Waiting queue full ({len(self.waiting)}/{self.max_waiting}); retry later or increase max_waiting."
                )
            self.waiting.append(seq)
            return True

    def schedule(self) -> list[Sequence]:
        """
        Schedule sequences for the next inference step.
        Promotes waiting sequences to running if there is capacity.
        """
        with self._lock:
            # Clean up finished sequences (engine should ideally handle this, or we handle it here pre-schedule)
            # But if engine updates state to FINISHED, we can filter them out.
            # However, we usually want to return 'Finished' status to user once before removing.
            # Let's assume engine calls `free_completed` explicitly or we filter here.
            # Better: Filter out finished ones from `running` at the start.
            self.running = [s for s in self.running if not s.is_finished()]

            # Fill available slots
            while self.waiting and len(self.running) < self.max_batch_size:
                seq = self.waiting.popleft()
                seq.status = RequestState.RUNNING
                self.running.append(seq)

            return self.running

    def get_sequence(self, request_id: str) -> Sequence | None:
        """Find a sequence by its request_id."""
        with self._lock:
            for s in self.running:
                if s.request_id == request_id:
                    return s
            for s in self.waiting:
                if s.request_id == request_id:
                    return s
            return None

    def remove(self, request_id: str) -> Sequence | None:
        """Drop a sequence from the waiting queue / running list.

        Used by the streaming generator's cleanup path when the consumer
        abandons mid-generation (RIL ISS-105): the abandoned generator is the
        sequence's only stepper and can never advance it again, so leaving it
        RUNNING permanently consumes a KV slot (``schedule`` only filters
        FINISHED). Idempotent: returns the removed sequence or ``None``.
        """
        with self._lock:
            for i, s in enumerate(self.running):
                if s.request_id == request_id:
                    return self.running.pop(i)
            for i, s in enumerate(self.waiting):
                if s.request_id == request_id:
                    del self.waiting[i]
                    return s
            return None

    def clear(self) -> None:
        """Empty both queues (engine teardown)."""
        with self._lock:
            self.waiting.clear()
            self.running.clear()

add_sequence 

add_sequence(seq)

Add a new sequence to the waiting queue.

引发：

类型 描述
RuntimeError

If the waiting queue is at capacity.
源代码位于： src/llm/serving/scheduler.py 
41
42
43
44
45
46
47
48
49
50
51
52
def add_sequence(self, seq: Sequence):
    """Add a new sequence to the waiting queue.

    Raises:
        RuntimeError: If the waiting queue is at capacity.
    """
    with self._lock:
        if len(self.waiting) >= self.max_waiting:
            raise RuntimeError(
                f"Waiting queue full ({len(self.waiting)}/{self.max_waiting}); retry later or increase max_waiting."
            )
        self.waiting.append(seq)

add_sequence_if_not_conflicting 

add_sequence_if_not_conflicting(seq, matches=None)

Add seq unless an ACTIVE sequence already occupies its request_id.

Guard for RIL ISS-123: two distinct requests sharing one client request_id hash to the SAME KV slot (SlotAllocator keys slots by request_id), so their K/V and generated tokens contaminate each other and a free() by either returns the other's live slot to the pool.

The engine's internal double-add contract (generate_request -> stream_request re-adding the SAME logical request so the reap loop can remove every copy) stays intact: when an matches callback is supplied, a duplicate whose content equals the active holder is allowed (it is the same request re-added). By default (no callback) ANY active duplicate is rejected.

Returns True when the sequence was enqueued, False when it was rejected as a conflicting duplicate.

引发：

类型 描述
RuntimeError

If the waiting queue is at capacity.
源代码位于： src/llm/serving/scheduler.py 
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
85
86
def add_sequence_if_not_conflicting(self, seq: Sequence, matches: Callable[[Sequence], bool] | None = None) -> bool:
    """Add ``seq`` unless an ACTIVE sequence already occupies its request_id.

    Guard for RIL ISS-123: two distinct requests sharing one client
    request_id hash to the SAME KV slot (``SlotAllocator`` keys slots by
    request_id), so their K/V and generated tokens contaminate each other
    and a ``free()`` by either returns the other's live slot to the pool.

    The engine's internal *double-add* contract (``generate_request`` ->
    ``stream_request`` re-adding the SAME logical request so the reap loop
    can remove every copy) stays intact: when an ``matches`` callback is
    supplied, a duplicate whose content equals the active holder is
    allowed (it is the same request re-added). By default (no callback)
    ANY active duplicate is rejected.

    Returns ``True`` when the sequence was enqueued, ``False`` when it was
    rejected as a conflicting duplicate.

    Raises:
        RuntimeError: If the waiting queue is at capacity.
    """
    with self._lock:
        for s in list(self.running) + list(self.waiting):
            if s.request_id == seq.request_id and not s.is_finished():
                if matches is not None and matches(s):
                    continue  # same logical request re-added — permit
                return False
        if len(self.waiting) >= self.max_waiting:
            raise RuntimeError(
                f"Waiting queue full ({len(self.waiting)}/{self.max_waiting}); retry later or increase max_waiting."
            )
        self.waiting.append(seq)
        return True

schedule 

schedule()

Schedule sequences for the next inference step. Promotes waiting sequences to running if there is capacity.

源代码位于： src/llm/serving/scheduler.py 
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
def schedule(self) -> list[Sequence]:
    """
    Schedule sequences for the next inference step.
    Promotes waiting sequences to running if there is capacity.
    """
    with self._lock:
        # Clean up finished sequences (engine should ideally handle this, or we handle it here pre-schedule)
        # But if engine updates state to FINISHED, we can filter them out.
        # However, we usually want to return 'Finished' status to user once before removing.
        # Let's assume engine calls `free_completed` explicitly or we filter here.
        # Better: Filter out finished ones from `running` at the start.
        self.running = [s for s in self.running if not s.is_finished()]

        # Fill available slots
        while self.waiting and len(self.running) < self.max_batch_size:
            seq = self.waiting.popleft()
            seq.status = RequestState.RUNNING
            self.running.append(seq)

        return self.running

get_sequence 

get_sequence(request_id)

Find a sequence by its request_id.

源代码位于： src/llm/serving/scheduler.py 
109
110
111
112
113
114
115
116
117
118
def get_sequence(self, request_id: str) -> Sequence | None:
    """Find a sequence by its request_id."""
    with self._lock:
        for s in self.running:
            if s.request_id == request_id:
                return s
        for s in self.waiting:
            if s.request_id == request_id:
                return s
        return None

remove 

remove(request_id)

Drop a sequence from the waiting queue / running list.

Used by the streaming generator's cleanup path when the consumer abandons mid-generation (RIL ISS-105): the abandoned generator is the sequence's only stepper and can never advance it again, so leaving it RUNNING permanently consumes a KV slot (schedule only filters FINISHED). Idempotent: returns the removed sequence or None.

源代码位于： src/llm/serving/scheduler.py 
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
def remove(self, request_id: str) -> Sequence | None:
    """Drop a sequence from the waiting queue / running list.

    Used by the streaming generator's cleanup path when the consumer
    abandons mid-generation (RIL ISS-105): the abandoned generator is the
    sequence's only stepper and can never advance it again, so leaving it
    RUNNING permanently consumes a KV slot (``schedule`` only filters
    FINISHED). Idempotent: returns the removed sequence or ``None``.
    """
    with self._lock:
        for i, s in enumerate(self.running):
            if s.request_id == request_id:
                return self.running.pop(i)
        for i, s in enumerate(self.waiting):
            if s.request_id == request_id:
                del self.waiting[i]
                return s
        return None

clear 

clear()

Empty both queues (engine teardown).

源代码位于： src/llm/serving/scheduler.py 
139
140
141
142
143
def clear(self) -> None:
    """Empty both queues (engine teardown)."""
    with self._lock:
        self.waiting.clear()
        self.running.clear()

Generation Service

generation_service

Serving-side generation service backed by GenerationBackend.

ServingGenerationService dataclass

Shared generation entry point for REST and chat APIs.

源代码位于： src/llm/serving/generation_service.py 
 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
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 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
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
@dataclass
class ServingGenerationService:
    """Shared generation entry point for REST and chat APIs."""

    model: DecoderModel
    tokenizer: Any
    backend: GenerationBackend
    device: torch.device

    def shutdown(self) -> None:
        """Release references to model so GPU memory can be freed."""
        self.model = None  # type: ignore

    @classmethod
    def from_config(
        cls,
        config: ServingConfig,
        *,
        engine: ContinuousBatchingEngine | None = None,
    ) -> ServingGenerationService:
        device = _resolve_device(config.device)
        model, tokenizer = load_model_and_tokenizer(config)
        if engine is None:
            model.to(device)
            model.eval()
            # The "batched" backend requires a ContinuousBatchingEngine, but
            # ``api.lifespan`` constructs the service before the engine —
            # so build one from the just-loaded model here instead of
            # letting ``get_generation_backend`` raise ValueError and abort
            # server startup (RIL ISS-150).
            if config.generation_backend == "batched":
                engine = ContinuousBatchingEngine.from_serving_config(
                    config,
                    model=model,
                    tokenizer=tokenizer,
                )
        else:
            if engine.model is None:
                raise RuntimeError("engine model was unloaded")
            model = engine.model
            device = engine.device
        backend = get_generation_backend(config.generation_backend, engine=engine)
        return cls(model=model, tokenizer=tokenizer, backend=backend, device=device)

    def _generation_config(
        self,
        *,
        max_new_tokens: int,
        temperature: float = 1.0,
        top_k: int | None = None,
        top_p: float | None = None,
        repetition_penalty: float = 1.0,
        frequency_penalty: float = 0.0,
        presence_penalty: float = 0.0,
        logit_bias: dict[int, float] | None = None,
        stop: str | list[str] | None = None,
    ) -> GenerationConfig:
        # Reject requests that cannot fit in the model's context window up
        # front, instead of letting the eager backend truncate the prompt to an
        # *empty* input tensor (a 0-length slice when max_new_tokens >=
        # max_seq_len) and 500 on the forward. max_new_tokens min
        # max_seq_len is the absolute floor; realistic sampling always needs
        # at least 1 token of prompt, but this is the backend-agnostic guard.
        model_max_seq = getattr(self.model, "max_seq_len", None)
        if model_max_seq is not None and max_new_tokens >= model_max_seq:
            raise ValueError(
                f"max_new_tokens ({max_new_tokens}) must be less than the model's "
                f"max_seq_len ({model_max_seq}); the prompt would have no room "
                "to fit in the context window."
            )
        return GenerationConfig(
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
            repetition_penalty=repetition_penalty,
            frequency_penalty=frequency_penalty,
            presence_penalty=presence_penalty,
            logit_bias=logit_bias,
            use_cache=True,
            stop=stop,
        )

    def generate(
        self,
        prompt: str,
        *,
        request_id: str | None = None,
        max_new_tokens: int,
        temperature: float = 1.0,
        top_k: int | None = None,
        top_p: float | None = None,
        repetition_penalty: float = 1.0,
        frequency_penalty: float = 0.0,
        presence_penalty: float = 0.0,
        logit_bias: dict[int, float] | None = None,
        stop: str | list[str] | None = None,
    ) -> str:
        gen_config = self._generation_config(
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
            repetition_penalty=repetition_penalty,
            frequency_penalty=frequency_penalty,
            presence_penalty=presence_penalty,
            logit_bias=logit_bias,
            stop=stop,
        )
        return self.backend.generate(self.model, self.tokenizer, prompt, gen_config, request_id=request_id)

    def stream(
        self,
        prompt: str,
        *,
        request_id: str | None = None,
        max_new_tokens: int,
        temperature: float = 1.0,
        top_k: int | None = None,
        top_p: float | None = None,
        repetition_penalty: float = 1.0,
        frequency_penalty: float = 0.0,
        presence_penalty: float = 0.0,
        logit_bias: dict[int, float] | None = None,
        stop: str | list[str] | None = None,
    ):
        gen_config = self._generation_config(
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
            repetition_penalty=repetition_penalty,
            frequency_penalty=frequency_penalty,
            presence_penalty=presence_penalty,
            logit_bias=logit_bias,
            stop=stop,
        )
        yield from self.backend.stream(self.model, self.tokenizer, prompt, gen_config, request_id=request_id)

    def batch_generate(
        self,
        prompts: list[str],
        *,
        max_new_tokens: int,
        temperature: float = 1.0,
        top_k: int | None = None,
        top_p: float | None = None,
        repetition_penalty: float = 1.0,
        frequency_penalty: float = 0.0,
        presence_penalty: float = 0.0,
        logit_bias: dict[int, float] | None = None,
        stop: str | list[str] | None = None,
    ) -> list[str]:
        gen_config = self._generation_config(
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
            repetition_penalty=repetition_penalty,
            frequency_penalty=frequency_penalty,
            presence_penalty=presence_penalty,
            logit_bias=logit_bias,
            stop=stop,
        )
        return self.backend.batch_generate(self.model, self.tokenizer, prompts, gen_config)

shutdown 

shutdown()

Release references to model so GPU memory can be freed.

源代码位于： src/llm/serving/generation_service.py 
33
34
35
def shutdown(self) -> None:
    """Release references to model so GPU memory can be freed."""
    self.model = None  # type: ignore

Errors

errors

Structured error envelope for the serving API (Finding K).

All endpoints in src/llm/serving raise APIError (or an HTTPException that the global handler maps to an envelope). The envelope shape is::

{
  "error": {
    "code": "<stable_id>",   # machine-readable, e.g. "invalid_request"
    "message": "<human>",    # one-line description
    "details": {...},        # structured context (field-level errors, etc.)
    "request_id": "<uuid>"   # X-Request-ID echo
  }
}

The HTTP status is set by the status_code field on APIError (or by mapping ErrorCode to its default status). The X-Request-ID middleware sets request.state.request_id and echoes it on the response header.

ErrorCode

Bases: StrEnum

Stable machine-readable error identifiers.

These strings are part of the public API. Adding new codes is fine; renaming or removing them is a breaking change.

源代码位于： src/llm/serving/errors.py 
30
31
32
33
34
35
36
37
38
39
40
41
class ErrorCode(StrEnum):
    """Stable machine-readable error identifiers.

    These strings are part of the public API. Adding new codes is fine;
    renaming or removing them is a breaking change.
    """

    INVALID_REQUEST = "invalid_request"
    UNAUTHORIZED = "unauthorized"
    TIMEOUT = "timeout"
    MODEL_UNAVAILABLE = "model_unavailable"
    INTERNAL = "internal"

APIError

Bases: Exception

A typed error with a stable code, human message, and structured details.

Endpoints raise APIError instead of HTTPException. The global exception handler installed in :mod:llm.serving.api converts it to the standard envelope.

参数：

名称 类型 描述 默认
code ErrorCode | str

One of :class:ErrorCode (or a custom string for new codes).

 必需
message str

One-line human description.

 必需
status_code int | None

HTTP status to return. Defaults to the canonical status for code; override only if you have a strong reason.

 None
details dict[str, Any] | None

Optional structured context (e.g. field-level validation errors, retry-after seconds, etc.).

 None
源代码位于： src/llm/serving/errors.py 
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
85
86
87
88
89
90
91
92
93
94
class APIError(Exception):
    """A typed error with a stable code, human message, and structured details.

    Endpoints raise ``APIError`` instead of ``HTTPException``. The global
    exception handler installed in :mod:`llm.serving.api` converts it to
    the standard envelope.

    Args:
        code: One of :class:`ErrorCode` (or a custom string for new codes).
        message: One-line human description.
        status_code: HTTP status to return. Defaults to the canonical
            status for ``code``; override only if you have a strong reason.
        details: Optional structured context (e.g. field-level validation
            errors, retry-after seconds, etc.).
    """

    def __init__(
        self,
        code: ErrorCode | str,
        message: str,
        *,
        status_code: int | None = None,
        details: dict[str, Any] | None = None,
    ) -> None:
        super().__init__(message)
        if isinstance(code, ErrorCode):
            resolved_code = code.value
            resolved_status = status_code if status_code is not None else _CODE_TO_STATUS[code]
        else:
            resolved_code = str(code)
            resolved_status = status_code if status_code is not None else 500
        self.code = resolved_code
        self.message = message
        self.status_code = resolved_status
        self.details: dict[str, Any] = details or {}

default_status_for 

default_status_for(code)

HTTP status code for a given :class:ErrorCode.

源代码位于： src/llm/serving/errors.py 
55
56
57
def default_status_for(code: ErrorCode) -> int:
    """HTTP status code for a given :class:`ErrorCode`."""
    return _CODE_TO_STATUS[code]

to_envelope 

to_envelope(*, code, message, request_id, details=None)

Build the canonical error envelope dict.

参数：

名称 类型 描述 默认
code str

Machine-readable error code (e.g. "invalid_request").

 必需
message str

One-line human description.

 必需
request_id str

The X-Request-ID for this request.

 必需
details dict[str, Any] | None

Optional structured context.

 None

返回：

类型 描述
dict[str, Any]

{"error": {"code", "message", "details", "request_id"}}
源代码位于： src/llm/serving/errors.py 
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
def to_envelope(
    *,
    code: str,
    message: str,
    request_id: str,
    details: dict[str, Any] | None = None,
) -> dict[str, Any]:
    """Build the canonical error envelope dict.

    Args:
        code: Machine-readable error code (e.g. ``"invalid_request"``).
        message: One-line human description.
        request_id: The X-Request-ID for this request.
        details: Optional structured context.

    Returns:
        ``{"error": {"code", "message", "details", "request_id"}}``
    """
    return {
        "error": {
            "code": code,
            "message": message,
            "details": details or {},
            "request_id": request_id,
        }
    }

envelope_from_http_exception 

envelope_from_http_exception(exc, request_id)

Convert a fastapi.HTTPException to the envelope shape.

The detail field may be a string or a dict; we pass it through as details so callers don't lose context.

源代码位于： src/llm/serving/errors.py 
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
def envelope_from_http_exception(exc: HTTPException, request_id: str) -> dict[str, Any]:
    """Convert a ``fastapi.HTTPException`` to the envelope shape.

    The detail field may be a string or a dict; we pass it through as
    ``details`` so callers don't lose context.
    """
    details: dict[str, Any] = {}
    if isinstance(exc.detail, dict):
        details = exc.detail
    elif exc.detail is not None:
        details = {"detail": exc.detail}
    # Map well-known statuses to a stable code. Unknown statuses map to
    # "internal" (so operators see something they can search for).
    status_to_code = {
        400: "invalid_request",
        401: "unauthorized",
        403: "unauthorized",
        404: "not_found",
        408: "timeout",
        422: "invalid_request",
        429: "rate_limited",
        500: "internal",
        502: "model_unavailable",
        503: "model_unavailable",
        504: "timeout",
    }
    code = status_to_code.get(exc.status_code, "internal")
    return to_envelope(
        code=code,
        message=str(exc.detail) if exc.detail is not None else code,
        request_id=request_id,
        details=details,
    )

envelope_from_validation_error 

envelope_from_validation_error(exc, request_id)

Convert a pydantic / FastAPI RequestValidationError to the envelope.

源代码位于： src/llm/serving/errors.py 
160
161
162
163
164
165
166
167
def envelope_from_validation_error(exc: RequestValidationError, request_id: str) -> dict[str, Any]:
    """Convert a pydantic / FastAPI ``RequestValidationError`` to the envelope."""
    return to_envelope(
        code="invalid_request",
        message="Request validation failed.",
        request_id=request_id,
        details={"errors": exc.errors()},
    )

envelope_from_api_error 

envelope_from_api_error(exc, request_id)

Convert an :class:APIError to the envelope shape.

源代码位于： src/llm/serving/errors.py 
170
171
172
173
174
175
176
177
def envelope_from_api_error(exc: APIError, request_id: str) -> dict[str, Any]:
    """Convert an :class:`APIError` to the envelope shape."""
    return to_envelope(
        code=exc.code,
        message=exc.message,
        request_id=request_id,
        details=exc.details,
    )

envelope_from_unexpected 

envelope_from_unexpected(exc, request_id, *, logger=None)

Convert an unexpected Exception to the envelope shape.

Logs the exception (class + traceback) with the request_id so operators can correlate. The envelope does NOT leak internal details to clients — not even the exception class name (RIL ISS-168): type(exc).__name__ previously let a client fingerprint the framework/backend from any 500 ("RuntimeError", "CUDA error" classes, filesystem paths inside backend str(exc) messages).

源代码位于： src/llm/serving/errors.py 
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
def envelope_from_unexpected(exc: Exception, request_id: str, *, logger=None) -> dict[str, Any]:
    """Convert an unexpected ``Exception`` to the envelope shape.

    Logs the exception (class + traceback) with the request_id so operators
    can correlate. The envelope does NOT leak internal details to clients —
    not even the exception class name (RIL ISS-168): ``type(exc).__name__``
    previously let a client fingerprint the framework/backend from any 500
    (``"RuntimeError"``, ``"CUDA error"`` classes, filesystem paths inside
    backend ``str(exc)`` messages).
    """
    if logger is not None:
        logger.exception("unexpected error during request %s", request_id, exc_info=exc)
    return to_envelope(
        code="internal",
        message="Internal server error.",
        request_id=request_id,
    )

get_request_id 

get_request_id(request)

Return the request_id stored on the request state.

Falls back to "unknown" if the middleware hasn't run (which would be a programming error in production — see :class:llm.serving.middleware.RequestIDMiddleware).

源代码位于： src/llm/serving/errors.py 
199
200
201
202
203
204
205
206
def get_request_id(request: Request) -> str:
    """Return the request_id stored on the request state.

    Falls back to ``"unknown"`` if the middleware hasn't run (which would
    be a programming error in production — see
    :class:`llm.serving.middleware.RequestIDMiddleware`).
    """
    return getattr(request.state, "request_id", "unknown")

envelope_response 

envelope_response(payload, status_code, request_id)

Build the canonical JSON response: envelope body + X-Request-ID header.

源代码位于： src/llm/serving/errors.py 
209
210
211
212
213
214
215
def envelope_response(payload: dict[str, Any], status_code: int, request_id: str) -> JSONResponse:
    """Build the canonical JSON response: envelope body + ``X-Request-ID`` header."""
    return JSONResponse(
        status_code=status_code,
        content=payload,
        headers={"X-Request-ID": request_id},
    )

register_exception_handlers 

register_exception_handlers(app, *, logger=None)

Wire the standard FastAPI exception handlers onto app.

All errors raised inside the serving API (typed :class:APIError, FastAPI :class:HTTPException, pydantic validation errors, and unexpected exceptions) flow through this single registration point and come out shaped as :func:to_envelope.

参数：

名称 类型 描述 默认
app FastAPI

The :class:fastapi.FastAPI instance to register on.

 必需
logger Any | None

Optional logger to receive unexpected exceptions. If None, unexpected errors are still converted to a 500 envelope but not logged here (the :class:RequestIDMiddleware access log line still records the failure).

 None
源代码位于： src/llm/serving/errors.py 
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
def register_exception_handlers(app: FastAPI, *, logger: Any | None = None) -> None:
    """Wire the standard FastAPI exception handlers onto ``app``.

    All errors raised inside the serving API (typed :class:`APIError`,
    FastAPI :class:`HTTPException`, pydantic validation errors, and
    unexpected exceptions) flow through this single registration point
    and come out shaped as :func:`to_envelope`.

    Args:
        app: The :class:`fastapi.FastAPI` instance to register on.
        logger: Optional logger to receive unexpected exceptions. If
            ``None``, unexpected errors are still converted to a 500
            envelope but not logged here (the :class:`RequestIDMiddleware`
            access log line still records the failure).
    """
    # Functions below are defined in module scope earlier in this file —
    # no import needed.

    @app.exception_handler(RequestValidationError)
    async def _validation_handler(request: Request, exc: RequestValidationError) -> Response:
        return envelope_response(
            envelope_from_validation_error(exc, get_request_id(request)),
            status_code=422,
            request_id=get_request_id(request),
        )

    @app.exception_handler(APIError)
    async def _api_error_handler(request: Request, exc: APIError) -> Response:
        """Typed errors raised from endpoints / dependencies.

        Registered explicitly (not via the generic ``Exception`` handler)
        because FastAPI dispatches dependency exceptions through the
        same handler chain but the generic handler loses the typed
        status code when an exception is raised from a dependency.
        """
        return envelope_response(
            envelope_from_api_error(exc, get_request_id(request)),
            status_code=exc.status_code,
            request_id=get_request_id(request),
        )

    @app.exception_handler(Exception)
    async def _unhandled_handler(request: Request, exc: Exception) -> Response:
        request_id = get_request_id(request)
        if isinstance(exc, HTTPException):
            return envelope_response(
                envelope_from_http_exception(exc, request_id),
                status_code=exc.status_code,
                request_id=request_id,
            )
        return envelope_response(
            envelope_from_unexpected(exc, request_id, logger=logger),
            status_code=500,
            request_id=request_id,
        )

Auth

auth

Authentication for the serving API.

Currently a single shared API key compared in constant time (hmac.compare_digest) to avoid timing leaks. Supports both X-API-Key: <key> and Authorization: Bearer <key> headers.

A future multi-tenant extension can replace the body of :func:get_api_key without changing call sites, as long as the return contract (the key on success, raising :class:APIError with code unauthorized on failure) is preserved.

get_api_key async 

get_api_key(api_key_header_value=Security(api_key_header), auth_header=Security(authorization_header))

Verify the API key from X-API-Key or Authorization: Bearer.

Comparison uses hmac.compare_digest to avoid leaking key bytes via timing. If the module-level config.api_key is unset, auth is disabled and the function returns None (the public-host guard in :mod:llm.serving.cli blocks starting the server on a non-loopback interface without auth).

源代码位于： src/llm/serving/auth.py 
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
async def get_api_key(
    api_key_header_value: str | None = Security(api_key_header),
    auth_header: str | None = Security(authorization_header),
) -> str | None:
    """Verify the API key from ``X-API-Key`` or ``Authorization: Bearer``.

    Comparison uses ``hmac.compare_digest`` to avoid leaking key bytes via
    timing. If the module-level ``config.api_key`` is unset, auth is
    disabled and the function returns ``None`` (the public-host guard in
    :mod:`llm.serving.cli` blocks starting the server on a non-loopback
    interface without auth).
    """
    from llm.serving.api import config as _config

    if not _config.api_key:
        return None

    expected = _config.api_key
    # Check X-API-Key header first.
    if api_key_header_value is not None and hmac.compare_digest(api_key_header_value, expected):
        return api_key_header_value

    # Check Bearer token.
    bearer = _extract_bearer_token(auth_header)
    if bearer is not None and hmac.compare_digest(bearer, expected):
        return bearer

    raise APIError(ErrorCode.UNAUTHORIZED, "Could not validate credentials")

is_loopback 

is_loopback(host)

Return True if host is a loopback address.

Covers 127.0.0.0/8 and ::1. Anything else (0.0.0.0, *, LAN IPs, public hostnames) is treated as non-loopback.

源代码位于： src/llm/serving/auth.py 
63
64
65
66
67
68
69
70
71
def is_loopback(host: str) -> bool:
    """Return True if ``host`` is a loopback address.

    Covers ``127.0.0.0/8`` and ``::1``. Anything else (``0.0.0.0``, ``*``,
    LAN IPs, public hostnames) is treated as non-loopback.
    """
    if host in ("127.0.0.1", "localhost", "::1"):
        return True
    return bool(host.startswith("127."))

assert_safe_bind 

assert_safe_bind(host, api_key, *, source='ServingConfig.host')

Fail-closed guard: refuse to serve anonymously on a non-loopback bind.

The single home for the check every entry point validates — the llm-serve CLI (:func:llm.serving.cli.main), the FastAPI lifespan that the Docker image's direct uvicorn llm.serving.api:app launch runs, etc. Because all of them now read the SAME host the server actually binds (LLM_SERVING_HOST / ServingConfig.host), a Docker/uvicorn launch that skipped the CLI can no longer bind 0.0.0.0 with api_key=None and serve instruction-generating endpoints fully anonymously (RIL ISS-164). host=0.0.0.0 without auth fails at startup rather than silently at runtime.

参数：

名称 类型 描述 默认
host str

The bind address the server will use.

 必需
api_key str | None

The configured ServingConfig.api_key (None when auth is disabled).

 必需
source str

Name used in the error message for context.

 'ServingConfig.host'
源代码位于： src/llm/serving/auth.py 
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
def assert_safe_bind(
    host: str,
    api_key: str | None,
    *,
    source: str = "ServingConfig.host",
) -> None:
    """Fail-closed guard: refuse to serve anonymously on a non-loopback bind.

    The single home for the check every entry point validates — the
    ``llm-serve`` CLI (:func:`llm.serving.cli.main`), the FastAPI lifespan
    that the Docker image's direct ``uvicorn llm.serving.api:app`` launch
    runs, etc. Because all of them now read the SAME ``host`` the server
    actually binds (``LLM_SERVING_HOST`` / ``ServingConfig.host``), a
    Docker/uvicorn launch that skipped the CLI can no longer bind
    ``0.0.0.0`` with ``api_key=None`` and serve instruction-generating
    endpoints fully anonymously (RIL ISS-164). ``host=0.0.0.0`` without
    auth fails at startup rather than silently at runtime.

    Args:
        host: The bind address the server will use.
        api_key: The configured ``ServingConfig.api_key`` (``None`` when
            auth is disabled).
        source: Name used in the error message for context.
    """
    if not is_loopback(host) and not api_key:
        raise RuntimeError(
            f"Refusing to start: {source}={host!r} binds to a "
            f"non-loopback address but api_key is not set. Anonymous access on a "
            f"public interface is unsafe. Either set LLM_SERVING_HOST to a loopback "
            f"address (127.0.0.1) or set LLM_SERVING_API_KEY."
        )

Middleware

middleware

ASGI middleware for the serving API.

Right now this holds the :class:RequestIDMiddleware, which assigns a stable X-Request-ID to every request (honoring an inbound header), echoes it on the response, and logs a structured access line per request so operators can correlate uvicorn access logs with application logs.

RequestBodySizeLimit

Pure-ASGI middleware rejecting request bodies over max_bytes.

Defense-in-depth for client memory-exhaustion DoS (RIL ISS-171): without a limit, a single multi-hundred-MB JSON body (one giant prompt, many chat messages) is fully buffered and tokenized. Two layers here:

  • Content-Length fast reject — when the client declares a length over the cap, respond 413 Payload Too Large before reading any body.
  • Incremental cap for chunked / unknown-length bodies — counts bytes as the ASGI http.request messages stream in and rejects as soon as the cap is crossed, discarding (never buffering) the remainder.

Implemented as raw ASGI (not :class:BaseHTTPMiddleware) so it wraps the transport layer and does NOT interfere with SSE response streaming.

源代码位于： src/llm/serving/middleware.py 
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
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
class RequestBodySizeLimit:
    """Pure-ASGI middleware rejecting request bodies over ``max_bytes``.

    Defense-in-depth for client memory-exhaustion DoS (RIL ISS-171): without
    a limit, a single multi-hundred-MB JSON body (one giant ``prompt``, many
    chat messages) is fully buffered and tokenized. Two layers here:

    - **Content-Length fast reject** — when the client declares a length
      over the cap, respond ``413 Payload Too Large`` before reading any body.
    - **Incremental cap for chunked / unknown-length bodies** — counts bytes
      as the ASGI ``http.request`` messages stream in and rejects as soon as
      the cap is crossed, discarding (never buffering) the remainder.

    Implemented as raw ASGI (not :class:`BaseHTTPMiddleware`) so it wraps the
    transport layer and does NOT interfere with SSE response streaming.
    """

    def __init__(self, app, max_bytes: int):
        self.app = app
        self.max_bytes = max_bytes

    async def __call__(self, scope, receive, send):
        if scope["type"] != "http":
            return await self.app(scope, receive, send)

        headers = dict(scope.get("headers") or [])
        content_length = headers.get(b"content-length")
        if content_length is not None:
            try:
                declared = int(content_length)
            except TypeError, ValueError:
                declared = None
            if declared is not None and declared > self.max_bytes:
                return await send(self._too_large_response())

        received = 0

        async def limited_receive():
            nonlocal received
            message = await receive()
            if message["type"] == "http.request":
                received += len(message.get("body") or b"")
                if received > self.max_bytes:
                    # Drain the remainder so the connection is reusable, then
                    # abort the request; the bytes are counted, not buffered.
                    while message.get("more_body", False):
                        message = await receive()
                    raise _RequestBodyTooLargeError()
            return message

        try:
            return await self.app(scope, limited_receive, send)
        except _RequestBodyTooLargeError:
            await send(self._too_large_response())
            return

    @staticmethod
    def _too_large_response() -> dict:
        """An ASGI ``send`` single-call payload (413 with a JSON envelope)."""
        body = b'{"error":{"message":"Request body too large","type":"invalid_request_error"}}'
        return {
            "type": "http.response.start",
            "status": 413,
            "headers": [
                (b"content-type", b"application/json"),
                (b"content-length", str(len(body)).encode()),
            ],
        }

RequestIDMiddleware

Bases: BaseHTTPMiddleware

Assign, propagate, and log X-Request-ID for every request.

Behavior: - If the client sent X-Request-ID, reuse it (so callers can stitch retries to a single trace). - Otherwise, generate a new UUID4 hex. - Store on request.state.request_id so handlers and exception handlers can include it in error envelopes. - Echo on the response X-Request-ID header. - Log a structured INFO line on response (method, path, status, duration_ms, request_id).

源代码位于： src/llm/serving/middleware.py 
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
class RequestIDMiddleware(BaseHTTPMiddleware):
    """Assign, propagate, and log ``X-Request-ID`` for every request.

    Behavior:
    - If the client sent ``X-Request-ID``, reuse it (so callers can stitch
      retries to a single trace).
    - Otherwise, generate a new UUID4 hex.
    - Store on ``request.state.request_id`` so handlers and exception
      handlers can include it in error envelopes.
    - Echo on the response ``X-Request-ID`` header.
    - Log a structured INFO line on response (method, path, status,
      duration_ms, request_id).
    """

    async def dispatch(self, request: Request, call_next) -> Response:
        inbound = request.headers.get(REQUEST_ID_HEADER)
        request_id = inbound if inbound else uuid.uuid4().hex
        request.state.request_id = request_id

        start = time.perf_counter()
        status_code = 500  # default if call_next raises before returning
        try:
            response = await call_next(request)
            status_code = response.status_code
            response.headers[REQUEST_ID_HEADER] = request_id
            return response
        finally:
            duration_ms = (time.perf_counter() - start) * 1000
            logger.info(
                "request",
                extra={
                    "event": "request",
                    "request_id": request_id,
                    "method": request.method,
                    "path": request.url.path,
                    "status": status_code,
                    "duration_ms": round(duration_ms, 2),
                },
            )

Custom Prometheus Metrics

metrics

Custom Prometheus metrics for the serving API (Finding AX, T2 #22).

prometheus-fastapi-instrumentator already gives generic HTTP RED metrics (rate, errors, duration) per route. Domain-specific signals — tokens generated per request, batch fill ratio, KV-cache hit rate, queue depth — are not visible to operators without these counters.

Layout

:class:ServingMetrics is a thin container around a CollectorRegistry that holds the six metrics required by the audit ticket. The module also exposes a module-level :data:METRICS singleton wired into the default Prometheus registry so :func:prometheus_fastapi_instrumentator.Instrumentator.expose makes them visible at /metrics.

Usage

::

from llm.serving.metrics import METRICS

with METRICS.request_timer(endpoint="generate") as timer:
    async with METRICS.track_inflight():
        text = await run_in_threadpool(_sync_generate, ...)
    METRICS.observe_tokens(endpoint="generate", token_count=len(text))
    timer.set_status(200)

The engine can publish per-step stats via :func:ServingMetrics.record_batch_fill_ratio, called from the on_step hook passed to :class:~llm.serving.batch_engine.ContinuousBatchingEngine.

ServingMetrics

Container for the serving tier's domain Prometheus metrics.

参数：

名称 类型 描述 默认
registry CollectorRegistry | None

The Prometheus registry to register metrics against. Defaults to the module-level default (:data:prometheus_client.REGISTRY) so /metrics exposes them. Tests should pass a private :class:CollectorRegistry for isolation.

 None
源代码位于： src/llm/serving/metrics.py 
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
class ServingMetrics:
    """Container for the serving tier's domain Prometheus metrics.

    Args:
        registry: The Prometheus registry to register metrics against.
            Defaults to the module-level default (:data:`prometheus_client.REGISTRY`)
            so ``/metrics`` exposes them. Tests should pass a private
            :class:`CollectorRegistry` for isolation.
    """

    def __init__(self, registry: CollectorRegistry | None = None) -> None:
        self.registry: CollectorRegistry = registry if registry is not None else REGISTRY

        self.tokens_generated_total: Counter = Counter(
            "llm_tokens_generated_total",
            "Total completion tokens generated, by endpoint.",
            labelnames=("endpoint",),
            registry=self.registry,
        )
        self.tokens_per_request: Histogram = Histogram(
            "llm_tokens_per_request",
            "Distribution of completion tokens per request, by endpoint.",
            labelnames=("endpoint",),
            buckets=_TOKEN_BUCKETS,
            registry=self.registry,
        )
        self.request_duration_seconds: Histogram = Histogram(
            "llm_request_duration_seconds",
            "End-to-end request duration, by endpoint and HTTP status.",
            labelnames=("endpoint", "status"),
            buckets=_DURATION_BUCKETS,
            registry=self.registry,
        )
        self.batch_fill_ratio: Gauge = Gauge(
            "llm_batch_fill_ratio",
            "Fraction of continuous-batching slots currently in use (0-1).",
            registry=self.registry,
        )
        self.kv_cache_hit_ratio: Gauge = Gauge(
            "llm_kv_cache_hit_ratio",
            "Fraction of recent prefix lookups served from the KV prefix cache (0-1).",
            registry=self.registry,
        )
        self.inflight_requests: Gauge = Gauge(
            "llm_inflight_requests",
            "Number of in-flight inference requests currently holding a slot.",
            registry=self.registry,
        )

    # --- Observation helpers ---------------------------------------------

    def observe_tokens(self, *, endpoint: str, token_count: int) -> None:
        """Record one request's completion token count.

        Bumps both the cumulative counter and the per-request histogram
        so dashboards can show throughput and p95 in one set of queries.
        """
        if token_count < 0:
            raise ValueError(f"token_count must be non-negative, got {token_count}")
        self.tokens_generated_total.labels(endpoint=endpoint).inc(token_count)
        self.tokens_per_request.labels(endpoint=endpoint).observe(token_count)

    def record_batch_fill_ratio(self, *, scheduled: int, total_active_slots: int) -> None:
        """Update ``llm_batch_fill_ratio`` from a ``step()`` result.

        ``total_active_slots`` is the engine's full slot pool; ``scheduled``
        is how many of those were used in the most recent step. Callers
        passing ``total_active_slots=0`` get a ``ValueError`` because the
        ratio is undefined (engine not yet initialized).
        """
        if total_active_slots <= 0:
            raise ValueError(f"Cannot compute batch fill ratio: total_active_slots={total_active_slots}")
        ratio = scheduled / total_active_slots
        if ratio < 0.0 or ratio > 1.0:
            raise ValueError(
                f"Batch fill ratio {ratio} is outside [0, 1] "
                f"(scheduled={scheduled}, total_active_slots={total_active_slots})"
            )
        self.batch_fill_ratio.set(ratio)

    def record_kv_cache_hit_ratio(self, ratio: float) -> None:
        """Update ``llm_kv_cache_hit_ratio``. ``ratio`` must be in [0, 1]."""
        if ratio < 0.0 or ratio > 1.0:
            raise ValueError(f"KV cache hit ratio {ratio} is outside [0, 1]")
        self.kv_cache_hit_ratio.set(ratio)

    @contextmanager
    def track_inflight(self) -> Iterator[None]:
        """Context manager that inc/dec ``llm_inflight_requests``.

        Use to wrap the section of a request that holds an inference
        slot (typically ``async with inference_semaphore``). Decrement
        runs even on exception so the gauge never sticks above zero.
        """
        self.inflight_requests.inc()
        try:
            yield
        finally:
            self.inflight_requests.dec()

    def request_timer(self, *, endpoint: str) -> _RequestTimer:
        """Return a context manager that times a request and labels by status.

        Use :meth:`_RequestTimer.set_status` before exiting to record the
        HTTP status; the default is ``"error"`` so a missing call still
        emits a labelled observation.
        """
        return _RequestTimer(self.request_duration_seconds, endpoint)

observe_tokens 

observe_tokens(*, endpoint, token_count)

Record one request's completion token count.

Bumps both the cumulative counter and the per-request histogram so dashboards can show throughput and p95 in one set of queries.

源代码位于： src/llm/serving/metrics.py 
148
149
150
151
152
153
154
155
156
157
def observe_tokens(self, *, endpoint: str, token_count: int) -> None:
    """Record one request's completion token count.

    Bumps both the cumulative counter and the per-request histogram
    so dashboards can show throughput and p95 in one set of queries.
    """
    if token_count < 0:
        raise ValueError(f"token_count must be non-negative, got {token_count}")
    self.tokens_generated_total.labels(endpoint=endpoint).inc(token_count)
    self.tokens_per_request.labels(endpoint=endpoint).observe(token_count)

record_batch_fill_ratio 

record_batch_fill_ratio(*, scheduled, total_active_slots)

Update llm_batch_fill_ratio from a step() result.

total_active_slots is the engine's full slot pool; scheduled is how many of those were used in the most recent step. Callers passing total_active_slots=0 get a ValueError because the ratio is undefined (engine not yet initialized).

源代码位于： src/llm/serving/metrics.py 
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
def record_batch_fill_ratio(self, *, scheduled: int, total_active_slots: int) -> None:
    """Update ``llm_batch_fill_ratio`` from a ``step()`` result.

    ``total_active_slots`` is the engine's full slot pool; ``scheduled``
    is how many of those were used in the most recent step. Callers
    passing ``total_active_slots=0`` get a ``ValueError`` because the
    ratio is undefined (engine not yet initialized).
    """
    if total_active_slots <= 0:
        raise ValueError(f"Cannot compute batch fill ratio: total_active_slots={total_active_slots}")
    ratio = scheduled / total_active_slots
    if ratio < 0.0 or ratio > 1.0:
        raise ValueError(
            f"Batch fill ratio {ratio} is outside [0, 1] "
            f"(scheduled={scheduled}, total_active_slots={total_active_slots})"
        )
    self.batch_fill_ratio.set(ratio)

record_kv_cache_hit_ratio 

record_kv_cache_hit_ratio(ratio)

Update llm_kv_cache_hit_ratio. ratio must be in [0, 1].

源代码位于： src/llm/serving/metrics.py 
177
178
179
180
181
def record_kv_cache_hit_ratio(self, ratio: float) -> None:
    """Update ``llm_kv_cache_hit_ratio``. ``ratio`` must be in [0, 1]."""
    if ratio < 0.0 or ratio > 1.0:
        raise ValueError(f"KV cache hit ratio {ratio} is outside [0, 1]")
    self.kv_cache_hit_ratio.set(ratio)

track_inflight 

track_inflight()

Context manager that inc/dec llm_inflight_requests.

Use to wrap the section of a request that holds an inference slot (typically async with inference_semaphore). Decrement runs even on exception so the gauge never sticks above zero.

源代码位于： src/llm/serving/metrics.py 
183
184
185
186
187
188
189
190
191
192
193
194
195
@contextmanager
def track_inflight(self) -> Iterator[None]:
    """Context manager that inc/dec ``llm_inflight_requests``.

    Use to wrap the section of a request that holds an inference
    slot (typically ``async with inference_semaphore``). Decrement
    runs even on exception so the gauge never sticks above zero.
    """
    self.inflight_requests.inc()
    try:
        yield
    finally:
        self.inflight_requests.dec()

request_timer 

request_timer(*, endpoint)

Return a context manager that times a request and labels by status.

Use :meth:_RequestTimer.set_status before exiting to record the HTTP status; the default is "error" so a missing call still emits a labelled observation.

源代码位于： src/llm/serving/metrics.py 
197
198
199
200
201
202
203
204
def request_timer(self, *, endpoint: str) -> _RequestTimer:
    """Return a context manager that times a request and labels by status.

    Use :meth:`_RequestTimer.set_status` before exiting to record the
    HTTP status; the default is ``"error"`` so a missing call still
    emits a labelled observation.
    """
    return _RequestTimer(self.request_duration_seconds, endpoint)

Chat Template

chat_template

Render OpenAI-style chat messages into a single prompt string.

This is intentionally tiny and dependency-free so the serving tier stays out of the tokenizer's way. The model is expected to have been trained on the rendered format (or one configured via ServingConfig).

messages_to_prompt 

messages_to_prompt(messages, *, message_template=None, generation_prefix=None)

Convert chat messages to a single prompt string.

Each message is rendered with message_template.format(role=..., content=...) (default: "{role}: {content}"). The rendered messages are joined with newlines and generation_prefix (default: "Assistant: ") is appended so the model knows where to start producing assistant tokens.

Override message_template and generation_prefix to match a fine-tuned model's expected format (ChatML, Llama-2-chat, Vicuna, …).

源代码位于： src/llm/serving/chat_template.py 
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
def messages_to_prompt(
    messages: Iterable,
    *,
    message_template: str | None = None,
    generation_prefix: str | None = None,
) -> str:
    """Convert chat messages to a single prompt string.

    Each message is rendered with ``message_template.format(role=..., content=...)``
    (default: ``"{role}: {content}"``). The rendered messages are joined
    with newlines and ``generation_prefix`` (default: ``"Assistant: "``) is
    appended so the model knows where to start producing assistant tokens.

    Override ``message_template`` and ``generation_prefix`` to match a
    fine-tuned model's expected format (ChatML, Llama-2-chat, Vicuna, …).
    """
    mt = message_template or DEFAULT_CHAT_MESSAGE_TEMPLATE
    gp = generation_prefix or DEFAULT_CHAT_GENERATION_PREFIX
    parts = [mt.format(role=msg.role, content=msg.content) for msg in messages]
    parts.append(gp)
    return "\n".join(parts)

CLI

cli

CLI entry point for the serving API.

reload=True is intentionally disabled because uvicorn's file-watcher conflicts with from llm.serving.api import app (the watch import path is incompatible with production use). For local development with auto-reload, run uvicorn directly::

uvicorn llm.serving.api:app --reload --host 127.0.0.1 --port 8000

The function refuses to start when the server would bind to a non-loopback address without an api_key configured. host=0.0.0.0 without auth exposes the inference endpoint to the network; this guard makes that mistake fail loudly at startup rather than silently at runtime.

main 

main(config=None)

Entry point for the llm-serve CLI.

参数：

名称 类型 描述 默认
config ServingConfig | None

Optional pre-built :class:ServingConfig. When None (the common CLI case), a fresh instance is constructed from env vars. The :mod:llm.serving.api wrapper passes its module-level config so test code that monkey-patches api.config.host / api.config.api_key exercises the guard without rebuilding a fresh config.

 None
源代码位于： src/llm/serving/cli.py 
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
def main(config: ServingConfig | None = None) -> None:
    """Entry point for the ``llm-serve`` CLI.

    Args:
        config: Optional pre-built :class:`ServingConfig`. When ``None``
            (the common CLI case), a fresh instance is constructed from
            env vars. The :mod:`llm.serving.api` wrapper passes its
            module-level config so test code that monkey-patches
            ``api.config.host`` / ``api.config.api_key`` exercises the
            guard without rebuilding a fresh config.
    """
    if config is None:
        config = ServingConfig()
    # Shared fail-closed guard — see llm.serving.auth.assert_safe_bind.
    assert_safe_bind(config.host, config.api_key)

    reload = os.environ.get("LLM_SERVING_RELOAD", "").lower() in ("1", "true", "yes")
    # Imported lazily so ``main`` is cheap to import (e.g. for `--help`).
    import uvicorn

    uvicorn.run("llm.serving.api:app", host=config.host, port=8000, reload=reload)

Schemas

schemas

RequestState

Bases: Enum

Internal state of a request in the scheduler.

源代码位于： src/llm/serving/schemas.py 
48
49
50
51
52
53
54
class RequestState(Enum):
    """Internal state of a request in the scheduler."""

    WAITING = auto()
    RUNNING = auto()
    PENDING = auto()  # Waiting for preemption to complete
    FINISHED = auto()

Sequence dataclass

Internal representation of a sequence state for the engine.

源代码位于： src/llm/serving/schemas.py 
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
85
86
87
88
89
90
91
@dataclass
class Sequence:
    """
    Internal representation of a sequence state for the engine.
    """

    request_id: str
    prompt: str
    input_ids: list[int]

    status: RequestState = RequestState.WAITING
    generated_ids: list[int] = field(default_factory=list)
    output_text: str = ""
    max_new_tokens: int = 50
    temperature: float = 1.0
    top_k: int | None = None
    top_p: float | None = None
    repetition_penalty: float = 1.0
    frequency_penalty: float = 0.0
    presence_penalty: float = 0.0
    logit_bias: dict[str, float] | None = None
    stop: str | list[str] | None = None

    def __post_init__(self):
        self._prompt_len = len(self.input_ids)

    @property
    def total_len(self) -> int:
        return self._prompt_len + len(self.generated_ids)

    def is_finished(self) -> bool:
        return self.status == RequestState.FINISHED

    def append_token_id(self, token_id: int):
        self.generated_ids.append(token_id)

GenerationRequest

Bases: BaseModel

Generation request model.

源代码位于： src/llm/serving/schemas.py 
 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
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
class GenerationRequest(BaseModel):
    """Generation request model."""

    request_id: str | None = Field(None, description="Client-provided request ID.")
    prompt: str = Field(
        ...,
        min_length=1,
        max_length=MAX_PROMPT_CHARS,
        description=(
            f"Input prompt text (1..{MAX_PROMPT_CHARS} characters). An empty prompt is "
            "rejected: the eager backend cannot decode a zero-length context."
        ),
    )
    max_new_tokens: int = Field(50, ge=1, le=4096, description="Maximum number of tokens to generate.")
    temperature: float = Field(1.0, ge=0.0, description="Controls randomness. 0 for Greedy Search.")
    top_k: int | None = Field(None, ge=1, description="Top-k sampling parameter. None to disable.")
    top_p: float | None = Field(
        None,
        gt=0.0,
        le=1.0,
        description="Nucleus sampling (top-p) parameter. 1.0 means no truncation (OpenAI default).",
    )
    repetition_penalty: float = Field(1.0, ge=1.0, description="Repetition penalty. 1.0 means no penalty.")
    frequency_penalty: float = Field(0.0, ge=-2.0, le=2.0, description="OpenAI-compatible per-frequency penalty.")
    presence_penalty: float = Field(
        0.0,
        ge=-2.0,
        le=2.0,
        description="OpenAI-compatible per-presence penalty.",
    )
    logit_bias: dict[str, float] | None = Field(
        None,
        description="OpenAI-compatible per-token additive bias. JSON keys are token ids (as strings); values are added to the affected logits before sampling.",
    )
    stream: bool = Field(False, description="Whether to use streaming output (SSE).")
    stop: str | list[str] | None = Field(
        None,
        description="Stop sequence(s). Generation halts the moment the generated text contains any of these as a suffix; the stop string itself is excluded from the response. Accepts a single string or a list of up to 4 strings. None means no stop.",
    )

    @field_validator("stop")
    @classmethod
    def _check_stop(cls, v):
        return _validate_stop(v)

GenerationResponse

Bases: BaseModel

Generation response model.

源代码位于： src/llm/serving/schemas.py 
140
141
142
143
144
class GenerationResponse(BaseModel):
    """Generation response model."""

    generated_text: str = Field(..., description="Generated text.")
    token_count: int | None = Field(None, description="Number of generated tokens.")

BatchGenerationRequest

Bases: BaseModel

Batch generation request model.

源代码位于： src/llm/serving/schemas.py 
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
class BatchGenerationRequest(BaseModel):
    """Batch generation request model."""

    prompts: list[str] = Field(
        ...,
        min_length=1,
        max_length=MAX_BATCH_PROMPTS,
        description=f"List of input prompts (1..{MAX_BATCH_PROMPTS} entries, each at most {MAX_PROMPT_CHARS} chars).",
    )

    @field_validator("prompts")
    @classmethod
    def _check_prompts(cls, v):
        for s in v:
            if not s:
                raise ValueError("each prompt must be a non-empty string")
            if len(s) > MAX_PROMPT_CHARS:
                raise ValueError(f"each prompt may be at most {MAX_PROMPT_CHARS} characters")
        return v

    max_new_tokens: int = Field(50, ge=1, le=4096, description="Maximum tokens to generate per prompt.")
    temperature: float = Field(1.0, ge=0.0, description="Sampling temperature. 0 for greedy.")
    top_k: int | None = Field(None, ge=1, description="Top-k sampling parameter.")
    top_p: float | None = Field(
        None, gt=0.0, le=1.0, description="Nucleus sampling parameter. 1.0 means no truncation."
    )
    repetition_penalty: float = Field(1.0, ge=1.0, description="Repetition penalty.")
    frequency_penalty: float = Field(0.0, ge=-2.0, le=2.0, description="OpenAI-compatible per-frequency penalty.")
    presence_penalty: float = Field(
        0.0,
        ge=-2.0,
        le=2.0,
        description="OpenAI-compatible per-presence penalty.",
    )
    logit_bias: dict[str, float] | None = Field(
        None,
        description="Per-token additive logit bias. Keys are token ids (as strings); values are added to the affected logits before sampling.",
    )
    stop: str | list[str] | None = Field(
        None,
        description="Stop sequence(s). Generation halts the moment the generated text contains any of these as a suffix; the stop string itself is excluded from the response. Accepts a single string or a list of up to 4 strings. None means no stop.",
    )

    @field_validator("stop")
    @classmethod
    def _check_stop(cls, v):
        return _validate_stop(v)

BatchGenerationResponse

Bases: BaseModel

Batch generation response model.

源代码位于： src/llm/serving/schemas.py 
196
197
198
199
class BatchGenerationResponse(BaseModel):
    """Batch generation response model."""

    results: list[GenerationResponse] = Field(..., description="List of generation results.")

ChatMessage

Bases: BaseModel

OpenAI-compatible chat message.

源代码位于： src/llm/serving/schemas.py 
205
206
207
208
209
210
211
212
213
class ChatMessage(BaseModel):
    """OpenAI-compatible chat message."""

    role: Literal["system", "user", "assistant"] = Field(..., description="Role of the message author.")
    content: str = Field(
        ...,
        max_length=MAX_PROMPT_CHARS,
        description=f"Content of the message (max {MAX_PROMPT_CHARS} characters).",
    )

ChatCompletionRequest

Bases: BaseModel

OpenAI-compatible chat completion request.

源代码位于： src/llm/serving/schemas.py 
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
class ChatCompletionRequest(BaseModel):
    """OpenAI-compatible chat completion request."""

    model: str = Field("llm", description="Model ID (ignored, for compatibility).")
    messages: list[ChatMessage] = Field(
        ...,
        min_length=1,
        max_length=MAX_MESSAGES,
        description=f"List of messages (1..{MAX_MESSAGES}).",
    )
    max_tokens: int = Field(50, ge=1, le=4096, description="Maximum tokens to generate.")
    temperature: float = Field(1.0, ge=0.0, le=2.0, description="Sampling temperature.")
    top_p: float | None = Field(
        None, gt=0.0, le=1.0, description="Nucleus sampling parameter. 1.0 means no truncation."
    )
    stream: bool = Field(False, description="Whether to stream responses.")
    stop: list[str] | str | None = Field(
        None,
        description="Stop sequence(s). Generation halts the moment the generated text contains any of these as a suffix; the stop string itself is excluded from the response. Accepts a single string or a list of up to 4 strings. None means no stop.",
    )

    @field_validator("stop")
    @classmethod
    def _check_stop(cls, v):
        return _validate_stop(v)

    presence_penalty: float = Field(
        0.0,
        ge=-2.0,
        le=2.0,
        description="OpenAI-compatible per-presence penalty.",
    )
    frequency_penalty: float = Field(0.0, ge=-2.0, le=2.0, description="OpenAI-compatible per-frequency penalty.")
    logit_bias: dict[str, float] | None = Field(
        None,
        description="OpenAI-compatible per-token additive bias. JSON keys are token ids (as strings); values are added to the affected logits before sampling.",
    )

ChatCompletionUsage

Bases: BaseModel

Token usage statistics.

源代码位于： src/llm/serving/schemas.py 
255
256
257
258
259
260
class ChatCompletionUsage(BaseModel):
    """Token usage statistics."""

    prompt_tokens: int = Field(..., description="Number of tokens in the prompt.")
    completion_tokens: int = Field(..., description="Number of tokens in the completion.")
    total_tokens: int = Field(..., description="Total tokens used.")

ChatCompletionChoiceMessage

Bases: BaseModel

Message in a chat completion choice.

源代码位于： src/llm/serving/schemas.py 
263
264
265
266
267
class ChatCompletionChoiceMessage(BaseModel):
    """Message in a chat completion choice."""

    role: Literal["assistant"] = "assistant"
    content: str = Field(..., description="Generated content.")

ChatCompletionChoice

Bases: BaseModel

A single completion choice.

源代码位于： src/llm/serving/schemas.py 
270
271
272
273
274
275
class ChatCompletionChoice(BaseModel):
    """A single completion choice."""

    index: int = Field(..., description="Index of this choice.")
    message: ChatCompletionChoiceMessage = Field(..., description="Generated message.")
    finish_reason: Literal["stop", "length"] | None = Field(None, description="Reason for stopping.")

ChatCompletionResponse

Bases: BaseModel

OpenAI-compatible chat completion response.

源代码位于： src/llm/serving/schemas.py 
278
279
280
281
282
283
284
285
286
class ChatCompletionResponse(BaseModel):
    """OpenAI-compatible chat completion response."""

    id: str = Field(default_factory=lambda: f"chatcmpl-{uuid.uuid4().hex[:12]}")
    object: Literal["chat.completion"] = "chat.completion"
    created: int = Field(default_factory=lambda: int(time.time()))
    model: str = Field("llm", description="Model used.")
    choices: list[ChatCompletionChoice] = Field(..., description="List of choices.")
    usage: ChatCompletionUsage = Field(..., description="Token usage.")

ChatCompletionChunkDelta

Bases: BaseModel

Delta content in a streaming chunk.

源代码位于： src/llm/serving/schemas.py 
289
290
291
292
293
class ChatCompletionChunkDelta(BaseModel):
    """Delta content in a streaming chunk."""

    role: Literal["assistant"] | None = None
    content: str | None = None

ChatCompletionChunkChoice

Bases: BaseModel

A choice in a streaming chunk.

源代码位于： src/llm/serving/schemas.py 
296
297
298
299
300
301
class ChatCompletionChunkChoice(BaseModel):
    """A choice in a streaming chunk."""

    index: int = 0
    delta: ChatCompletionChunkDelta = Field(..., description="Delta content.")
    finish_reason: Literal["stop", "length"] | None = None

ChatCompletionChunk

Bases: BaseModel

OpenAI-compatible streaming chunk.

源代码位于： src/llm/serving/schemas.py 
304
305
306
307
308
309
310
311
class ChatCompletionChunk(BaseModel):
    """OpenAI-compatible streaming chunk."""

    id: str = Field(default_factory=lambda: f"chatcmpl-{uuid.uuid4().hex[:12]}")
    object: Literal["chat.completion.chunk"] = "chat.completion.chunk"
    created: int = Field(default_factory=lambda: int(time.time()))
    model: str = "llm"
    choices: list[ChatCompletionChunkChoice] = Field(..., description="List of chunk choices.")

Routers

chat

OpenAI-compatible chat completions endpoint.

configure 

configure(config_, semaphore_, metrics_=None)

Bind the module-level config reference, semaphore, and metrics.

Called during lifespan startup.

源代码位于： src/llm/serving/routers/chat.py 
64
65
66
67
68
69
70
71
72
73
74
75
76
77
def configure(
    config_: ServingConfig,
    semaphore_: asyncio.Semaphore,
    metrics_: ServingMetrics | None = None,
) -> None:
    """Bind the module-level config reference, semaphore, and metrics.

    Called during lifespan startup.
    """
    global config, inference_semaphore, metrics
    config = config_
    inference_semaphore = semaphore_
    if metrics_ is not None:
        metrics = metrics_

chat_completions async 

chat_completions(request, config_, _api_key)

OpenAI-compatible chat completions endpoint.

源代码位于： src/llm/serving/routers/chat.py 
 80
 81
 82
 83
 84
 85
 86
 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
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
@router.post("/v1/chat/completions", response_model=ChatCompletionResponse)
async def chat_completions(
    request: ChatCompletionRequest,
    config_: Annotated[ServingConfig, Depends(lambda: config)],
    _api_key: Annotated[str, Depends(get_api_key)],
) -> ChatCompletionResponse | StreamingResponse:
    """OpenAI-compatible chat completions endpoint."""
    prompt = messages_to_prompt(
        request.messages,
        message_template=config_.chat_message_template,
        generation_prefix=config_.chat_generation_prefix,
    )
    # ``presence_penalty`` flows through its own kwarg to the service —
    # the OpenAI chat endpoint does not expose ``repetition_penalty``,
    # so the legacy ``1.0 + presence_penalty`` alias was removed in
    # Tier 3 #37 (presence and frequency now use flat-per-token math).
    repetition_penalty = 1.0

    if request.stream:
        # Validate before the SSE starts so an un-encodable message or a
        # context-overflowing ``max_tokens`` is a real 4xx, not an
        # "Error: ..." chunk inside a 200 stream — the sibling ``/generate``
        # stream pre-validates the same way (RIL ISS-113, ISS-149).
        _validate_stream_request(prompt, request.max_tokens)
        return StreamingResponse(
            _chat_stream_generator(
                request,
                prompt,
                repetition_penalty,
                request.frequency_penalty,
                request.presence_penalty,
                request.logit_bias,
            ),
            media_type="text/event-stream",
        )

    timer = metrics.request_timer(endpoint="chat_completions")
    with timer as t:
        try:
            async with asyncio.timeout(config_.request_timeout):
                async with inference_semaphore or _null_cm():
                    with metrics.track_inflight():
                        generated_text = await run_in_threadpool(
                            _sync_generate,
                            prompt=prompt,
                            max_new_tokens=request.max_tokens,
                            temperature=request.temperature,
                            top_p=request.top_p,
                            repetition_penalty=repetition_penalty,
                            frequency_penalty=request.frequency_penalty,
                            presence_penalty=request.presence_penalty,
                            logit_bias=request.logit_bias,
                            stop=request.stop,
                        )
        except TimeoutError as exc:
            t.set_status(504)
            raise APIError(ErrorCode.TIMEOUT, "Request timeout") from exc
        except RuntimeError as exc:
            # Do NOT echo the backend exception text back to the client (RIL
            # ISS-168): it can contain filesystem paths / framework internals.
            t.set_status(503)
            logger.error("backend generation failed: %s", exc)
            raise APIError(ErrorCode.MODEL_UNAVAILABLE, "Model unavailable during generation") from exc
        # Character-level tokenizers raise ``KeyError`` for characters outside
        # their corpus (the same exception the ``/generate`` sibling maps to a
        # client 400 — RIL ISS-147). Without it this client-caused error would
        # fall through to the generic handler and surface as a spurious 500.
        except (ValueError, KeyError) as exc:
            t.set_status(400)
            raise APIError(ErrorCode.INVALID_REQUEST, f"Invalid request: {exc}", details={"field": str(exc)}) from exc
        except APIError as exc:
            t.set_status(exc.status_code)
            raise
        except Exception as exc:
            logger.exception("Unexpected error in chat_completions")
            t.set_status(500)
            raise APIError(ErrorCode.INTERNAL, "Internal server error") from exc
        else:
            t.set_status(200)

    # Strip only the prompt prefix if the model echoed it back. The
    # completion itself is returned verbatim — `.strip()` removed a model's
    # own leading/trailing whitespace (e.g. a leading newline char-tokenizers
    # commonly emit), diverging from the streaming path which emits chunks
    # verbatim (RIL ISS-114).
    completion = generated_text[len(prompt) :] if generated_text.startswith(prompt) else generated_text

    # Per the OpenAI spec, a completion truncated by the ``max_tokens`` budget
    # is ``finish_reason="length"``, not ``"stop"``. Token counts come from
    # re-encoding the text with the serving tokenizer (exact for char
    # tokenizers, and correct for HF/BPE tokenizers where one token spans
    # several chars) — matching the ``token_count`` the streaming path
    # accumulates per decoded token (RIL ISS-151 / round-73 FINDING 4).
    completion_tokens = _token_count(completion)
    metrics.observe_tokens(endpoint="chat_completions", token_count=completion_tokens)
    finish_reason = "length" if completion_tokens >= request.max_tokens else "stop"

    return ChatCompletionResponse(
        model=request.model,
        choices=[
            ChatCompletionChoice(
                index=0,
                message=ChatCompletionChoiceMessage(content=completion),
                finish_reason=finish_reason,
            )
        ],
        usage=ChatCompletionUsage(
            prompt_tokens=_token_count(prompt),
            completion_tokens=completion_tokens,
            total_tokens=_token_count(prompt) + completion_tokens,
        ),
    )

generate

Text generation endpoints (/generate, /batch_generate).

configure 

configure(config_, generation_service_, semaphore_, metrics_=None)

Bind the module-level references.

Called once during FastAPI lifespan startup. Importing this module without configuring first will yield None values and the endpoints will refuse to serve (the RuntimeError below is a programming error, not a runtime condition).

源代码位于： src/llm/serving/routers/generate.py 
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
def configure(
    config_: ServingConfig,
    generation_service_,
    semaphore_: asyncio.Semaphore,
    metrics_: ServingMetrics | None = None,
) -> None:
    """Bind the module-level references.

    Called once during FastAPI lifespan startup. Importing this module
    without configuring first will yield None values and the endpoints
    will refuse to serve (the ``RuntimeError`` below is a programming
    error, not a runtime condition).
    """
    global config, generation_service, inference_semaphore, metrics
    config = config_
    generation_service = generation_service_
    inference_semaphore = semaphore_
    if metrics_ is not None:
        metrics = metrics_

generate_text async 

generate_text(request, config_, _api_key)

Generate text from a single prompt. Supports streaming and non-streaming.

源代码位于： src/llm/serving/routers/generate.py 
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
@router.post("/generate", response_model=GenerationResponse)
async def generate_text(
    request: GenerationRequest,
    config_: Annotated[ServingConfig, Depends(lambda: config)],
    _api_key: Annotated[str, Depends(get_api_key)],
) -> GenerationResponse | StreamingResponse:
    """Generate text from a single prompt. Supports streaming and non-streaming."""
    if request.stream:
        # Validate before the SSE starts so an un-encodable prompt or a
        # context-overflowing ``max_new_tokens`` is a real 4xx, not an
        # "Error: ..." chunk inside a 200 stream (RIL ISS-113, ISS-149).
        _validate_stream_request(request.prompt, request.max_new_tokens)
        return StreamingResponse(_stream_generator(request), media_type="text/event-stream")

    timer = metrics.request_timer(endpoint="generate")
    with timer as t:
        try:
            async with asyncio.timeout(config_.request_timeout):
                async with inference_semaphore or _null_cm():
                    with metrics.track_inflight():
                        generated_text = await run_in_threadpool(
                            _sync_generate,
                            prompt=request.prompt,
                            # Thread the client's request_id to the engine so
                            # batched/continuous-batching duplicate-request-id
                            # protection (RIL ISS-123/F3) works from the HTTP
                            # path — it was previously dropped at every layer
                            # (round-73 FINDING 3 / ISS-224).
                            request_id=request.request_id,
                            max_new_tokens=request.max_new_tokens,
                            temperature=request.temperature,
                            top_k=request.top_k,
                            top_p=request.top_p,
                            repetition_penalty=request.repetition_penalty,
                            frequency_penalty=request.frequency_penalty,
                            presence_penalty=request.presence_penalty,
                            logit_bias=request.logit_bias,
                            stop=request.stop,
                        )
        except TimeoutError as exc:
            t.set_status(504)
            raise APIError(ErrorCode.TIMEOUT, "Request timeout") from exc
        except RuntimeError as exc:
            # Do NOT echo the backend exception text back to the client (RIL
            # ISS-168): it can contain filesystem paths / framework internals
            # that aid fingerprinting. Log it server-side, send a fixed message.
            t.set_status(503)
            logger.error("backend generation failed: %s", exc)
            raise APIError(ErrorCode.MODEL_UNAVAILABLE, "Model unavailable during generation") from exc
        except (ValueError, KeyError) as exc:
            t.set_status(400)
            raise APIError(ErrorCode.INVALID_REQUEST, f"Invalid request: {exc}", details={"field": str(exc)}) from exc
        except APIError as exc:
            t.set_status(exc.status_code)
            raise
        except Exception as exc:
            logger.exception("Unexpected error in generate_text")
            t.set_status(500)
            raise APIError(ErrorCode.INTERNAL, "Internal server error") from exc
        else:
            t.set_status(200)
    # The eager backend's ``generate()`` returns ``prompt + completion`` (its
    # ``stream()`` yields chunks only). Reported as-is, the prompt would be
    # echoed into ``generated_text`` and counted as generated tokens,
    # diverging from both this route's streaming mode and ``/v1/chat``
    # (chat.py strips the same prior-echo — RIL ISS-148). Strip the prompt
    # if the model echoed it back, mirroring chat.py's completion extraction.
    completion = generated_text[len(request.prompt) :] if generated_text.startswith(request.prompt) else generated_text

    metrics.observe_tokens(endpoint="generate", token_count=_token_count(completion))
    return GenerationResponse(generated_text=completion, token_count=_token_count(completion))

batch_generate_text async 

batch_generate_text(request, config_, _api_key)

Generate text for a batch of prompts in one call.

源代码位于： src/llm/serving/routers/generate.py 
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
@router.post("/batch_generate", response_model=BatchGenerationResponse)
async def batch_generate_text(
    request: BatchGenerationRequest,
    config_: Annotated[ServingConfig, Depends(lambda: config)],
    _api_key: Annotated[str, Depends(get_api_key)],
) -> BatchGenerationResponse:
    """Generate text for a batch of prompts in one call."""
    timer = metrics.request_timer(endpoint="batch_generate")
    with timer as t:
        try:
            async with asyncio.timeout(config_.request_timeout):
                async with inference_semaphore or _null_cm():
                    with metrics.track_inflight():
                        results = await run_in_threadpool(
                            _sync_batch_generate,
                            prompts=request.prompts,
                            max_new_tokens=request.max_new_tokens,
                            temperature=request.temperature,
                            top_k=request.top_k,
                            top_p=request.top_p,
                            repetition_penalty=request.repetition_penalty,
                            frequency_penalty=request.frequency_penalty,
                            presence_penalty=request.presence_penalty,
                            logit_bias=request.logit_bias,
                            stop=request.stop,
                        )
        except TimeoutError as exc:
            t.set_status(504)
            raise APIError(ErrorCode.TIMEOUT, "Request timeout") from exc
        except RuntimeError as exc:
            # Do NOT echo the backend exception text back to the client (RIL
            # ISS-168): it can contain filesystem paths / framework internals
            # that aid fingerprinting. Log it server-side, send a fixed message.
            t.set_status(503)
            logger.error("backend generation failed: %s", exc)
            raise APIError(ErrorCode.MODEL_UNAVAILABLE, "Model unavailable during generation") from exc
        except (ValueError, KeyError) as exc:
            t.set_status(400)
            raise APIError(ErrorCode.INVALID_REQUEST, f"Invalid request: {exc}", details={"field": str(exc)}) from exc
        except APIError as exc:
            t.set_status(exc.status_code)
            raise
        except Exception as exc:
            logger.exception("Unexpected error in batch_generate_text")
            t.set_status(500)
            raise APIError(ErrorCode.INTERNAL, "Internal server error") from exc
        else:
            t.set_status(200)
    # Backend ``batch_generate`` returns ``prompt + completion`` per row (it
    # delegates to ``generate()``); strip the echoed prompt so the response is
    # the completion only, consistent with ``/generate`` and streaming
    # (RIL ISS-148). ``strict=False`` preserves the lenient results-driven
    # behavior this endpoint always had (production backends return one row
    # per prompt, so no truncation happens there).
    stripped: list[str] = []
    for prompt, text in zip(request.prompts, results, strict=False):
        stripped.append(text[len(prompt) :] if text.startswith(prompt) else text)
    # Record per-prompt token count; the counter is cumulative across
    # the whole batch, the histogram is per-prompt.
    for text in stripped:
        metrics.observe_tokens(endpoint="batch_generate", token_count=_token_count(text))
    return BatchGenerationResponse(
        results=[GenerationResponse(generated_text=text, token_count=len(text)) for text in stripped]
    )

health

Health and readiness endpoints.

health_check async 

health_check()

Liveness probe. Returns {"status": "ok"} if the process is up.

源代码位于： src/llm/serving/routers/health.py 
10
11
12
13
@router.get("/health")
async def health_check() -> dict[str, str]:
    """Liveness probe. Returns ``{"status": "ok"}`` if the process is up."""
    return {"status": "ok"}

Model Loader

loader

Load models and tokenizers for inference serving.

TrainingCheckpoint dataclass

Minimal view of a training checkpoint file.

源代码位于： src/llm/serving/loader.py 
38
39
40
41
42
43
44
45
46
47
@dataclass(frozen=True)
class TrainingCheckpoint:
    """Minimal view of a training checkpoint file."""

    path: Path
    model_state: dict[str, torch.Tensor]
    model_config: dict[str, Any] | None = None
    epoch: int | None = None
    loss: float | None = None
    model_obj: torch.nn.Module | None = None

load_training_checkpoint 

load_training_checkpoint(path)

Load a training checkpoint produced by CheckpointManager.

Accepts both the v2 split layout (<stem>.safetensors + <stem>.meta.json + <stem>.extra_state.pt, the modern CheckpointManager format) and the legacy v0.0.5 single-file .pt blob, plus the bare nn.Module pickle emitted by llm-quantize (a quantized model with GPTQQuantizedLinear / AWQQuantizedLinear / SmoothQuantLinear submodules whose per-layer quantization parameters live on the module instances). path may be a stem, any of the three v2 sidecar paths, or a legacy .pt path.

源代码位于： src/llm/serving/loader.py 
 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
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
def load_training_checkpoint(path: str | Path) -> TrainingCheckpoint:
    """Load a training checkpoint produced by CheckpointManager.

    Accepts both the v2 split layout (``<stem>.safetensors`` +
    ``<stem>.meta.json`` + ``<stem>.extra_state.pt``, the modern
    ``CheckpointManager`` format) and the legacy v0.0.5 single-file
    ``.pt`` blob, plus the bare ``nn.Module`` pickle emitted by
    ``llm-quantize`` (a quantized model with ``GPTQQuantizedLinear`` /
    ``AWQQuantizedLinear`` / ``SmoothQuantLinear`` submodules whose
    per-layer quantization parameters live on the module instances).
    ``path`` may be a stem, any of the three v2 sidecar paths, or a
    legacy ``.pt`` path.
    """
    ckpt_path = Path(path)
    # A bare-module blob (e.g. ``llm-quantize --output model.pt``) is not a
    # checkpoint dict: detect it before the layout resolution below so it is
    # not routed through the legacy path (which would emit a misleading
    # "run llm-migrate-ckpt" DeprecationWarning for a quantized artifact).
    #
    # Security (RIL ISS-170): the probe loads with ``weights_only=True`` so an
    # untrusted ``.pt`` cannot execute arbitrary ``__reduce__`` code at serve
    # time. The framework's own module classes are allowlisted (see
    # ``_register_framework_safe_globals``), so legitimate quantized blobs load
    # normally; a pickle referencing anything outside the allowlist raises
    # UnpicklingError and is REFUSED rather than silently re-loaded with
    # ``weights_only=False`` (matching the already-hardened PEFT sidecar path).
    if ckpt_path.exists() and ckpt_path.suffix == ".pt":
        _register_framework_safe_globals()
        try:
            obj = torch.load(ckpt_path, map_location="cpu", weights_only=True)
        except Exception:  # noqa: BLE001 - probe only; any failure falls through
            # to the checkpoint-layout resolution below, which raises the
            # proper format-specific error for this path.
            obj = None
        if isinstance(obj, torch.nn.Module):
            logger.info("Loading bare quantized model blob from %s", ckpt_path)
            return TrainingCheckpoint(
                path=ckpt_path,
                model_state=obj.state_dict(),
                model_obj=obj,
            )
    payload = load_checkpoint_payload(ckpt_path)
    if payload is None:
        raise FileNotFoundError(f"Checkpoint not found: {ckpt_path} (checked the legacy .pt and the v2 split layout)")
    if "model_state" not in payload:
        raise ValueError(f"Checkpoint missing 'model_state': {ckpt_path}")

    return TrainingCheckpoint(
        path=ckpt_path,
        model_state=_strip_compile_prefix(payload["model_state"]),
        model_config=payload.get("model_config"),
        epoch=payload.get("epoch"),
        loss=payload.get("loss"),
    )

infer_vocab_size 

infer_vocab_size(state_dict)

Infer vocabulary size from an LM head or embedding weight tensor.

The embedding fallback key was a misspelling (embedding.token_embedding vs the real embedding_layer.token_embeddings), making the branch unreachable — a config-less tied-embedding checkpoint (no lm_head.weight, per llm.compat.weight_mapping) wrongly raised Cannot infer vocab_size even though the embedding tensor was present and loadable (RIL ISS-167).

源代码位于： src/llm/serving/loader.py 
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
def infer_vocab_size(state_dict: dict[str, torch.Tensor]) -> int:
    """Infer vocabulary size from an LM head or embedding weight tensor.

    The embedding fallback key was a misspelling (``embedding.token_embedding``
    vs the real ``embedding_layer.token_embeddings``), making the branch
    unreachable — a config-less **tied-embedding** checkpoint (no
    ``lm_head.weight``, per ``llm.compat.weight_mapping``) wrongly raised
    ``Cannot infer vocab_size`` even though the embedding tensor was present
    and loadable (RIL ISS-167).
    """
    if "lm_head.weight" in state_dict:
        return int(state_dict["lm_head.weight"].shape[0])
    if "embedding_layer.token_embeddings.weight" in state_dict:
        return int(state_dict["embedding_layer.token_embeddings.weight"].shape[0])
    raise ValueError("Cannot infer vocab_size from checkpoint state dict")

infer_num_layers 

infer_num_layers(state_dict)

Count transformer blocks present in a state dict.

源代码位于： src/llm/serving/loader.py 
123
124
125
126
127
128
129
130
131
132
133
134
def infer_num_layers(state_dict: dict[str, torch.Tensor]) -> int | None:
    """Count transformer blocks present in a state dict."""
    indices: set[int] = set()
    prefix = "transformer_blocks."
    for key in state_dict:
        if not key.startswith(prefix):
            continue
        rest = key[len(prefix) :]
        index_str, _, _ = rest.partition(".")
        if index_str.isdigit():
            indices.add(int(index_str))
    return max(indices) + 1 if indices else None

load_tokenizer 

load_tokenizer(config)

Load tokenizer from config or fall back to a printable character tokenizer.

源代码位于： src/llm/serving/loader.py 
183
184
185
def load_tokenizer(config: ServingConfig) -> Any:
    """Load tokenizer from config or fall back to a printable character tokenizer."""
    return TokenizerFactory.from_serving_config(config)

load_model_and_tokenizer 

load_model_and_tokenizer(config)

Build model/tokenizer for serving, loading weights when model_path is set.

PEFT integration (T2 PEFT #49): if config.peft_method is set, the loader applies the method to the freshly loaded base model and (optionally) loads the sidecar from config.peft_adapter_path. Without any PEFT fields, the loader behavior is unchanged.

The apply+load step is fail-loud: a missing or corrupt sidecar raises FileNotFoundError / ValueError / RuntimeError so the serving process refuses to start with a partial config — better than silently serving the un-adapted base model.

源代码位于： src/llm/serving/loader.py 
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
def load_model_and_tokenizer(config: ServingConfig) -> tuple[DecoderModel, Any]:
    """Build model/tokenizer for serving, loading weights when model_path is set.

    PEFT integration (T2 PEFT #49): if ``config.peft_method`` is set, the
    loader applies the method to the freshly loaded base model and
    (optionally) loads the sidecar from ``config.peft_adapter_path``.
    Without any PEFT fields, the loader behavior is unchanged.

    The apply+load step is fail-loud: a missing or corrupt sidecar
    raises ``FileNotFoundError`` / ``ValueError`` /
    ``RuntimeError`` so the serving process refuses to start with a
    partial config — better than silently serving the un-adapted base
    model.
    """
    if not config.model_path:
        return _create_dummy_model_and_tokenizer(config)

    checkpoint = load_training_checkpoint(config.model_path)
    if checkpoint.model_obj is not None:
        # Quantized-model blob produced by ``llm-quantize``: the module is
        # self-contained (packed weights, scales, per-layer quantization
        # params), so skip ModelFactory reconstruction entirely.
        model = checkpoint.model_obj
        if not isinstance(model, DecoderModel):
            raise TypeError(f"quantized model blob must be a DecoderModel, got {type(model).__name__}")
    else:
        model = _build_decoder(
            serving_config=config,
            model_config=checkpoint.model_config,
            state_dict=checkpoint.model_state,
        )
    tokenizer = load_tokenizer(config)

    _apply_peft_if_configured(model, config)

    return model, tokenizer