跳转至

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
@contextlib.asynccontextmanager
async def lifespan(_app: FastAPI) -> AsyncGenerator[None, Any]:
    """FastAPI lifespan manager — load model, wire routers, log config."""
    logger.info("Starting up...")

    generation_service = ServingGenerationService.from_config(config)
    engine = ContinuousBatchingEngine.from_serving_config(
        config,
        model=generation_service.model,
        tokenizer=generation_service.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.
    engine.set_step_observer(METRICS.record_batch_fill_ratio)

    # 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()
        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
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
    compile_model: bool = False  # Enable torch.compile for acceleration
    max_concurrent_requests: int = 4  # Max concurrent inference requests
    request_timeout: float = 60.0  # 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 = 256
    block_size: int = 16

    # Prefix Cache
    enable_prefix_cache: bool = False
    max_prefixes: int = 10

    # 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.

        Mirrors :meth:`llm.training.core.config.Config.from_yaml` — keeps the
        training-side and serving-side config loading surfaces symmetric so
        users can pick whichever input style fits their workflow (YAML files
        vs. environment variables).

        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.
            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() 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.

Mirrors :meth:llm.training.core.config.Config.from_yaml — keeps the training-side and serving-side config loading surfaces symmetric so users can pick whichever input style fits their workflow (YAML files vs. environment variables).

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

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

    Mirrors :meth:`llm.training.core.config.Config.from_yaml` — keeps the
    training-side and serving-side config loading surfaces symmetric so
    users can pick whichever input style fits their workflow (YAML files
    vs. environment variables).

    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.
        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() 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 
24
25
26
27
28
29
30
31
32
33
34
35
@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 
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
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:
        return hashlib.sha256(bytes(tokens)).hexdigest()

    def get(self, tokens: list[int]) -> tuple[int, int] | None:
        if len(tokens) < self.min_prefix_len:
            return None
        return self._entries.get(self.hash_tokens(tokens))

    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)

SlotAllocator

Manages allocation of KV cache slots.

源代码位于： src/llm/serving/batch_engine.py 
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
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

    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)

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 
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
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
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: object,
        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:
            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=self.model.transformer_blocks[0].self_attn.num_kv_heads,
                head_dim=self.model.transformer_blocks[0].self_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

            self.paged_kv_cache = PagedKVCache(
                num_layers=len(self.model.transformer_blocks),
                num_kv_heads=self.model.transformer_blocks[0].self_attn.num_kv_heads,
                head_dim=self.model.transformer_blocks[0].self_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()
        # 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

    @classmethod
    def from_serving_config(cls, config, model: DecoderModel, tokenizer: object) -> 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 cache replay across slots is only supported on the
            # dense KV-cache path. The paged-cache path reuses blocks via
            # ``PagedKVCache.add_prefix`` + ``try_get_prefix_blocks``;
            # wiring those into ``_lock_step_pre`` is a follow-up.
            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)

        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,
        )
        self.scheduler.add_sequence(seq)
        return req_id

    def stream_request(
        self,
        request: GenerationRequest,
    ):
        """Run a request to completion, yielding decoded text chunks."""
        req_id = self.add_request(request)
        emitted = 0
        while True:
            seq = self.scheduler.get_sequence(req_id)
            if seq is None:
                break
            if seq.is_finished():
                for token_id in seq.generated_ids[emitted:]:
                    yield self.tokenizer.decode([token_id])
                break
            self.step()
            seq = self.scheduler.get_sequence(req_id)
            if seq is None:
                break
            for token_id in seq.generated_ids[emitted:]:
                yield self.tokenizer.decode([token_id])
            emitted = len(seq.generated_ids)
            if seq.is_finished():
                break

    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).

        Bookend the model forward with lock acquire/release: lock for
        pre-compute (slot allocation, prefix-cache lookup, batch tensor
        construction), release for the forward, re-acquire for
        post-compute (append tokens, free slots, mark finished). The
        forward is the expensive part; freeing the lock around it lets
        other worker threads enqueue / dequeue requests in parallel.

        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)
            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

    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.

        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:
                cached = self.prefix_cache.get(seq.input_ids) if self.prefix_cache 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
                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)
        q_pos = padded_position_ids.unsqueeze(1).unsqueeze(-1)
        run_attn_mask = col_indices > q_pos

        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

            if length < q_len:
                run_attn_mask[i, :, length:, :] = True

        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:
            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, :]
                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)
                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 _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 so the engine stays
            # leak-free even when the model raises mid-forward. The
            # sequences themselves remain in their last-known state;
            # callers are expected to clean them up via the timeout
            # path.
            for i, seq in enumerate(inputs.running_sequences):
                self.slot_allocator.free(seq.request_id)
                if self.paged_kv_cache is not None:
                    # ``seq_id`` == slot id in the paged path.
                    self.paged_kv_cache.free(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)

            if self.prefix_cache 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 (
                (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.slot_allocator.free(seq.request_id)
                if self.paged_kv_cache is not None:
                    # Return the per-sequence blocks to the allocator.
                    self.paged_kv_cache.free(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):
        """Unload model."""
        self.model = None
        self.kv_caches = []

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 
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
@classmethod
def from_serving_config(cls, config, model: DecoderModel, tokenizer: object) -> 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 
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
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)

    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,
    )
    self.scheduler.add_sequence(seq)
    return req_id

stream_request 

stream_request(request)

Run a request to completion, yielding decoded text chunks.

源代码位于： src/llm/serving/batch_engine.py 
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
def stream_request(
    self,
    request: GenerationRequest,
):
    """Run a request to completion, yielding decoded text chunks."""
    req_id = self.add_request(request)
    emitted = 0
    while True:
        seq = self.scheduler.get_sequence(req_id)
        if seq is None:
            break
        if seq.is_finished():
            for token_id in seq.generated_ids[emitted:]:
                yield self.tokenizer.decode([token_id])
            break
        self.step()
        seq = self.scheduler.get_sequence(req_id)
        if seq is None:
            break
        for token_id in seq.generated_ids[emitted:]:
            yield self.tokenizer.decode([token_id])
        emitted = len(seq.generated_ids)
        if seq.is_finished():
            break

generate_request 

generate_request(request)

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

源代码位于： src/llm/serving/batch_engine.py 
316
317
318
319
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 
321
322
323
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).

Bookend the model forward with lock acquire/release: lock for pre-compute (slot allocation, prefix-cache lookup, batch tensor construction), release for the forward, re-acquire for post-compute (append tokens, free slots, mark finished). The forward is the expensive part; freeing the lock around it lets other worker threads enqueue / dequeue requests in parallel.

返回：

类型 描述
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 
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
@torch.no_grad()
def step(self) -> StepStats:
    """Run one inference step (sync wrapper).

    Bookend the model forward with lock acquire/release: lock for
    pre-compute (slot allocation, prefix-cache lookup, batch tensor
    construction), release for the forward, re-acquire for
    post-compute (append tokens, free slots, mark finished). The
    forward is the expensive part; freeing the lock around it lets
    other worker threads enqueue / dequeue requests in parallel.

    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)
        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

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.

返回：

类型 描述
StepStats

class:StepStats (same fields as :meth:step).
源代码位于： src/llm/serving/batch_engine.py 
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
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.

    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 
551
552
553
554
555
556
557
558
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()

Unload model.

源代码位于： src/llm/serving/batch_engine.py 
560
561
562
563
def unload_model(self):
    """Unload model."""
    self.model = None
    self.kv_caches = []

Scheduler

scheduler

Scheduler

A simple First-Come-First-Serve (FCFS) scheduler for continuous batching.

源代码位于： src/llm/serving/scheduler.py 
 6
 7
 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
class Scheduler:
    """
    A simple First-Come-First-Serve (FCFS) scheduler for continuous batching.
    """

    def __init__(self, max_batch_size: int = 16):
        self.waiting: deque[Sequence] = deque()
        self.running: list[Sequence] = []
        self.max_batch_size = max_batch_size

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

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

    def schedule(self) -> list[Sequence]:
        """
        Schedule sequences for the next inference step.
        Promotes waiting sequences to running if there is capacity.
        """
        # 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."""
        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

add_sequence 

add_sequence(seq)

Add a new sequence to the waiting queue.

源代码位于： src/llm/serving/scheduler.py 
20
21
22
def add_sequence(self, seq: Sequence):
    """Add a new sequence to the waiting queue."""
    self.waiting.append(seq)

schedule 

schedule()

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

源代码位于： src/llm/serving/scheduler.py 
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
def schedule(self) -> list[Sequence]:
    """
    Schedule sequences for the next inference step.
    Promotes waiting sequences to running if there is capacity.
    """
    # 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 
44
45
46
47
48
49
50
51
52
def get_sequence(self, request_id: str) -> Sequence | None:
    """Find a sequence by its request_id."""
    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

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
@dataclass
class ServingGenerationService:
    """Shared generation entry point for REST and chat APIs."""

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

    @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()
        else:
            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,
    ) -> GenerationConfig:
        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,
        )

    def generate(
        self,
        prompt: 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,
    ) -> 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,
        )
        return self.backend.generate(self.model, self.tokenizer, prompt, gen_config)

    def stream(
        self,
        prompt: 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,
    ):
        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,
        )
        yield from self.backend.stream(self.model, self.tokenizer, prompt, gen_config)

    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,
    ) -> 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,
        )
        return self.backend.batch_generate(self.model, self.tokenizer, prompts, gen_config)

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 with the request_id so operators can correlate. The envelope does NOT leak internal details to clients.

源代码位于： src/llm/serving/errors.py 
180
181
182
183
184
185
186
187
188
189
190
191
192
193
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 with the request_id so operators can correlate.
    The envelope does NOT leak internal details to clients.
    """
    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,
        details={"type": type(exc).__name__},
    )

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 
196
197
198
199
200
201
202
203
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 
206
207
208
209
210
211
212
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 
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
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).
    """
    from llm.serving.errors import (
        envelope_from_api_error,
        envelope_from_http_exception,
        envelope_from_unexpected,
        envelope_from_validation_error,
    )

    @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."))

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.

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 
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
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
48
49
50
51
52
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()
    if not is_loopback(config.host) and not config.api_key:
        raise RuntimeError(
            f"Refusing to start: ServingConfig.host='{config.host}' 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."
        )

    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 
10
11
12
13
14
15
16
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 
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
@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

    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 
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
class GenerationRequest(BaseModel):
    """Generation request model."""

    request_id: str | None = Field(None, description="Client-provided request ID.")
    prompt: str = Field(..., description="Input prompt text.")
    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, lt=1.0, description="Nucleus sampling (top-p) parameter.")
    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.")
    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).")

GenerationResponse

Bases: BaseModel

Generation response model.

源代码位于： src/llm/serving/schemas.py 
70
71
72
73
74
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 
77
78
79
80
81
82
83
84
85
86
87
88
89
90
class BatchGenerationRequest(BaseModel):
    """Batch generation request model."""

    prompts: list[str] = Field(..., min_length=1, max_length=32, description="List of input prompts.")
    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, lt=1.0, description="Nucleus sampling parameter.")
    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.")
    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.",
    )

BatchGenerationResponse

Bases: BaseModel

Batch generation response model.

源代码位于： src/llm/serving/schemas.py 
93
94
95
96
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 
102
103
104
105
106
class ChatMessage(BaseModel):
    """OpenAI-compatible chat message."""

    role: Literal["system", "user", "assistant"] = Field(..., description="Role of the message author.")
    content: str = Field(..., description="Content of the message.")

ChatCompletionRequest

Bases: BaseModel

OpenAI-compatible chat completion request.

源代码位于： src/llm/serving/schemas.py 
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
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, description="List of 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, lt=1.0, description="Nucleus sampling parameter.")
    stream: bool = Field(False, description="Whether to stream responses.")
    stop: list[str] | str | None = Field(None, description="Stop sequences (not implemented).")
    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 
132
133
134
135
136
137
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 
140
141
142
143
144
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 
147
148
149
150
151
152
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 
155
156
157
158
159
160
161
162
163
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 
166
167
168
169
170
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 
173
174
175
176
177
178
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 
181
182
183
184
185
186
187
188
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 
53
54
55
56
57
58
59
60
61
62
63
64
65
66
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 
 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
@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:
        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):
                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,
                    )
        except TimeoutError as exc:
            t.set_status(504)
            raise APIError(ErrorCode.TIMEOUT, "Request timeout") from exc
        except RuntimeError as exc:
            t.set_status(503)
            raise APIError(ErrorCode.MODEL_UNAVAILABLE, str(exc)) from exc
        except ValueError 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 the prompt prefix if the model echoed it back.
    completion = generated_text[len(prompt) :].strip() if generated_text.startswith(prompt) else generated_text.strip()

    metrics.observe_tokens(endpoint="chat_completions", token_count=len(completion))

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

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 
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
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 
 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
@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:
        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):
                with metrics.track_inflight():
                    generated_text = await run_in_threadpool(
                        _sync_generate,
                        prompt=request.prompt,
                        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,
                        logit_bias=request.logit_bias,
                    )
        except TimeoutError as exc:
            t.set_status(504)
            raise APIError(ErrorCode.TIMEOUT, "Request timeout") from exc
        except RuntimeError as exc:
            t.set_status(503)
            raise APIError(ErrorCode.MODEL_UNAVAILABLE, str(exc)) from exc
        except ValueError 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)
    metrics.observe_tokens(endpoint="generate", token_count=len(generated_text))
    return GenerationResponse(generated_text=generated_text, token_count=len(generated_text))

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 
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
@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):
                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,
                        logit_bias=request.logit_bias,
                    )
        except TimeoutError as exc:
            t.set_status(504)
            raise APIError(ErrorCode.TIMEOUT, "Request timeout") from exc
        except RuntimeError as exc:
            t.set_status(503)
            raise APIError(ErrorCode.MODEL_UNAVAILABLE, str(exc)) from exc
        except ValueError 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)
    # Record per-prompt token count; the counter is cumulative across
    # the whole batch, the histogram is per-prompt.
    for text in results:
        metrics.observe_tokens(endpoint="batch_generate", token_count=len(text))
    return BatchGenerationResponse(
        results=[GenerationResponse(generated_text=text, token_count=len(text)) for text in results]
    )

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 
21
22
23
24
25
26
27
28
29
@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

load_training_checkpoint 

load_training_checkpoint(path, *, map_location='cpu')

Load a training checkpoint produced by CheckpointManager.

源代码位于： src/llm/serving/loader.py 
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
def load_training_checkpoint(path: str | Path, *, map_location: str | torch.device = "cpu") -> TrainingCheckpoint:
    """Load a training checkpoint produced by CheckpointManager."""
    ckpt_path = Path(path)
    if not ckpt_path.exists():
        raise FileNotFoundError(f"Checkpoint not found: {ckpt_path}")

    payload = torch.load(ckpt_path, map_location=map_location, weights_only=False)
    if "model_state" not in payload:
        raise ValueError(f"Checkpoint missing 'model_state': {ckpt_path}")

    return TrainingCheckpoint(
        path=ckpt_path,
        model_state=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.

源代码位于： src/llm/serving/loader.py 
51
52
53
54
55
56
57
def infer_vocab_size(state_dict: dict[str, torch.Tensor]) -> int:
    """Infer vocabulary size from an LM head or embedding weight tensor."""
    if "lm_head.weight" in state_dict:
        return int(state_dict["lm_head.weight"].shape[0])
    if "embedding.token_embedding.weight" in state_dict:
        return int(state_dict["embedding.token_embedding.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 
60
61
62
63
64
65
66
67
68
69
70
71
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 
117
118
119
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 
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
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)
    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