diff --git a/docs/adding_models.rst b/docs/adding_models.rst index ba6e83198a..548336f995 100644 --- a/docs/adding_models.rst +++ b/docs/adding_models.rst @@ -287,45 +287,131 @@ sampling parameters. The engine no longer reads it directly. Pass its result int through ``SamplingReqConfig`` are stored in buffers whose addresses do not change across replays, so they stay per-request and remain safe under CUDA graphs. +Paged KV layouts +~~~~~~~~~~~~~~~~ + +A model with a paged KV pool states, for every request, what its walks write into each +label of the pool. The KV manager turns that layout into the request's footprint, the most +pages each label can hold, before the request runs. A model declares its labels once and +returns each request's layout from ``process_prompt``. + +``kv_labels(self) -> dict[str, dict[str, KVLabel]]`` + Return ``{resource_key: {label: KVLabel}}`` for every label of every paged pool. A + :class:`mstar.model.base.KVLabel` has two fields: + + - ``growth``: ``(walk, loop)`` pairs, each a walk that keeps writing after admission, + bounded by the named loop. Only a named loop can bound growth, since an unnamed one + is renamed in every process. + - ``keys``: a :class:`mstar.model.base.PrefixKeys`, or ``None``. See **Cross-request + prefix reuse** below. + + A label written only by spans and scratch, such as a guidance branch, is ``KVLabel()``. + +``process_prompt`` returns, under ``kv_layout`` in its metadata, a +:class:`mstar.model.base.LabelLayout` for every label the request opens: the labels its +``KVReqConfig`` can name on any node and walk, ``main`` when it names none. A layout may be +empty where an opened label is never written. It may also name labels the request opens +only on another deployment, as an unguided BAGEL generation does for a CFG-parallel one. +A layout has five fields: + +- ``spans``: the committed writes known at intake, in walk order. A + :class:`mstar.model.base.Span` has a ``kind``: ``"ids"`` for token ids from the tensor + ``source`` names, ``"digest"`` for KV made from an uploaded file (``source`` is + ``(modality, index)``), and ``"rows"`` for KV made from anything else, such as another + node's outputs. ``length`` is the slots the write fills, ``advance`` how far it moves + the position counter. A span's ``scratch`` is the most one non-committing step writes + before the span commits into those slots, as a window's denoise steps do. +- ``growth``: ``(loop, slots)`` pairs, ``slots`` per iteration of ``loop``. +- ``scratch``: the most KV slots one non-committing step writes past the committed length. +- ``retention``: the ``RetentionPolicy`` the label's committing steps declare. +- ``keyed_spans``: how many leading spans the label's keys cover, by default every span + before the first ``rows`` span. + +The footprint adds the spans, the growth at each loop's iterations and the scratch, and +rounds up to whole pages per label. A span's own scratch counts only with the spans before +it, so a window's denoise steps and its commit count once. Under a ``RetentionPolicy`` it +takes the window's peak instead. + +``prompt_spans(plan, texts, text_walk, item_spans)`` builds the common case: a prompt's +text and attachments in ``prefill_plan`` order. Qwen3.5 lays out its one label with it: + +.. code-block:: python + + layout = LabelLayout( + spans=tuple(prompt_spans(prefill_plan(parts), segments, walk, image)), + growth=(("decode_loop", 1),), + ) + return ProcessPromptOutput(result, {"kv_layout": {KV_CACHE: {"main": layout}}}) + +``get_loop_caps(self, **model_kwargs) -> dict[str, int]`` + The iterations each named loop may run for this request. A loop runs the smaller of its + graph ``max_iters`` and its cap, and speculation never runs past it. A model whose stop + check ends a loop at a count returns that count, so the cap and the stop check end the + loop on the same iteration; ``test/modular/test_model_kv_layouts.py`` checks this for + every model. Growth on an uncapped loop is counted to its graph ``max_iters``. + +Where ``process_prompt`` computes a length that a walk computes another way, such as a +media item's rows, the walk asserts the two agree: BAGEL's ViT walk against +``vit_patches``, Qwen3-Omni's vision encoder against the processor's grid. + +Loading fails in three cases, and a request fails in two: + +- A paged pool has no labels. The message is `` declares no labels for its paged + pool ''; add them to `kv_labels`, and return each request's layout for them from + `process_prompt```. A ring pool needs none. +- A label's growth or keys name a walk the model never runs. The message is `` + declares ''/'