[2/3] kv: cross-request prefix caching

docs/adding_models.rst

@@ -233,6 +233,77 @@ 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.
 
+Cross-request prefix reuse
+~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Two requests that begin with the same tokens fill the same KV pages, and the second can
+attend the pages the first filled instead of recomputing them. A model opts in with two
+methods. Both have defaults, and a model that overrides neither is never cached. The
+deployment switch and the per-request opt-out are described under **Resource overrides**
+in :doc:`serving`.
+
+``prefix_key_streams(self) -> dict[str, dict[str, PrefixStream]]``
+   Return ``{resource_key: {label: PrefixStream}}``, one entry for each KV cache stream
+   whose pages are keyed. A :class:`mstar.model.base.PrefixStream` has four fields:
+
+   - ``tensor`` names the request tensor that holds the stream's token ids, as
+     ``process_prompt`` returns it.
+   - ``keyed_by`` is ``"ids"``: each page is keyed by its token ids, chained from the key
+     of the page before it.
+   - ``walk`` names the graph walk that writes the keyed span. Only this walk is probed
+     for a match. If any walk other than ``walk`` and ``decode_walk`` writes into the
+     stream, the stream stops being keyed.
+   - ``decode_walk`` names the walk whose input is the token just sampled, so that the
+     generated tokens are chained onto the prompt's keys and a later request that
+     re-sends the reply matches it too. ``None`` keys the prompt only.
+
+   BAGEL keys the text walk of its LLM and the decode that follows it:
+
+   .. code-block:: python
+
+      return {"kv": {"main": PrefixStream(
+          "text_inputs", "ids", "prefill_text", "decode",
+      )}}
+
+   Orpheus keys its prompt walk and its decode walk:
+
+   .. code-block:: python
+
+      return {
+          KV_CACHE: {"main": PrefixStream("text_inputs", "ids", "prefill", "decode")},
+      }
+
+``checkpoint_path(self) -> str | None``
+   Return the local directory that holds the weights and ``config.json``. The cache
+   hashes the checkpoint's safetensors index and config into the root of every key, so
+   naming the checkpoint is what keeps two builds off each other's pages. The default
+   returns ``None``, which leaves the weights out of the root. A model without a
+   checkpoint therefore needs ``prefix_cache_salt`` set to use the cache at all. Without
+   one, the engine keeps the cache shut and logs why.
+
+A model that declares a stream is checked at load time. Loading fails in three cases:
+
+- A position resource over a keyed cache uses a scheme other than sequential. A page is
+  reusable only where its tokens land at the positions they were written at. The message
+  is ``resource '<key>' places positions by <scheme>, but it sits over ['<kv key>'],
+  which the model keyed for prefix reuse; either drop the stream or give the resource a
+  sequential scheme``.
+- A resource on a keyed node does not set ``prefix_skip_safe``. A hit takes the front of
+  the request away from every resource on the node, including resources that never saw
+  those tokens. The message is ``resource '<key>' cannot serve a request whose prefix
+  came from the cache, but it sits on ['<node>'], which the model keyed for prefix
+  reuse; either drop the stream or give the resource the prefix hooks``.
+- A stream names a walk that the model does not run. The message is ``<Model> keys
+  '<resource key>'/'<label>' on ['<walk>'], which it never runs; its walks are [...]``.
+
+Some requests stay uncached even when the model declares a stream. The engine checks each
+request's inputs on the keyed walk and cuts a prefix only from inputs that the default
+``split_inputs`` can slice. A request whose inputs carry ``tensor_inputs``, ``kwargs`` or
+``resource_step_info`` is never probed; BAGEL's guided requests carry ``requires_cfg``
+there. The keys stop at a prompt's first input that is not a token, so a BAGEL prompt is
+keyed over its leading text only. A stream under a sliding window files pages only up to
+its protected prefix, which the window never releases.
+
 .. _Step 2a — Declare your resources:
 
 Step 2a — Declare your resources
@@ -267,9 +338,10 @@ The spec types are:
      - What it builds
    * - ``KVSpec(config=KVConfig(...))``
      - A paged KV cache. ``KVConfig`` holds ``num_layers``, ``num_kv_heads``,
-       ``head_dim``, ``max_seq_len`` and ``num_qo_heads``. It also holds three fields that
-       a deployment can tune: ``max_num_pages``, ``page_size`` and ``cpu_offload_pages``
-       (the number of pinned host pages used for offload; 0 disables offload).
+       ``head_dim``, ``max_seq_len`` and ``num_qo_heads``. It also holds five fields that
+       a deployment can tune: ``max_num_pages``, ``page_size``, ``cpu_offload_pages``
+       (the number of pinned host pages used for offload; 0 disables offload),
+       ``prefix_cache`` and ``prefix_cache_salt``.
    * - ``AttentionSpec(config=AttentionConfig(kv_cache=...))``
      - Self-attention planned over the named cache. ``backend`` selects
        ``AttnBackend.FLASHINFER`` (the default) or ``AttnBackend.DENSE``.
@@ -1210,7 +1282,8 @@ that a misspelled setting is never silently ignored:
    * - Spec
      - Accepts
    * - ``KVSpec``
-     - ``max_num_pages``, ``page_size``, ``max_seq_len``, ``cpu_offload_pages``
+     - ``max_num_pages``, ``page_size``, ``max_seq_len``, ``cpu_offload_pages``,
+       ``prefix_cache``, ``prefix_cache_salt``
    * - ``AttentionSpec`` / ``CrossAttentionSpec``
      - ``backend`` (``flashinfer`` / ``dense``), ``flashinfer_backend``
        (``auto`` / ``fa2`` / ``fa3``)
@@ -1589,6 +1662,7 @@ Checklist
          [ ] postprocess
          [ ] get_submodule
          [ ] (optional) get_request_resource_configs
+         [ ] (optional) prefix_key_streams + checkpoint_path   — cross-request prefix reuse
    [ ] mstar/model/registry.py                    — add to MODEL_REGISTRY (+ HF_MODELS)
    [ ] configs/<your_model>.yaml                  — node_groups → ranks (+ resources: overrides)
    [ ] (optional) async partitions if pipelined

docs/serving.rst

@@ -179,6 +179,9 @@ A config maps the model's computation-graph nodes to physical GPU ranks. The key
    * - ``resources``
      - *(optional)* Per-resource overrides, keyed by the model's resource names (see
        below).
+   * - ``max_concurrent_requests``
+     - *(optional)* How many requests run at once. Requests past the cap wait in the
+       conductor's queue.
    * - ``model_kwargs``
      - *(optional)* Server-init model parameters (see below).
 
@@ -194,7 +197,8 @@ them by that name:
 
    resources:
      kv_cache:
-       max_num_pages: 1024      # also: page_size, max_seq_len, cpu_offload_pages
+       max_num_pages: 1024      # also: page_size, max_seq_len, cpu_offload_pages,
+                                # prefix_cache, prefix_cache_salt
      talker_attn:
        flashinfer_backend: fa2  # also: backend (flashinfer / dense)
      vit_attn:                  # a cacheless (ragged) attention resource
@@ -206,6 +210,18 @@ that the model does not declare, or sets a key that the resource does not accept
 fails. A misspelled setting is never silently ignored. A top-level ``kv_cache:`` block is
 no longer read. It raises an error with a message describing the migration.
 
+``prefix_cache: false`` turns cross-request prefix reuse off for that cache, and
+``prefix_cache_salt`` keeps two otherwise identical deployments off each other's
+cached pages. A single request opts out by sending ``prefix_cache=False``, which
+travels with its other ``model_kwargs``.
+
+A cached prompt no longer reserves its pages, so the pool no longer limits how many
+requests are admitted. It has to hold the decode of every request allowed to run at once,
+and a decode step that finds no free page holds its requests until they time out. Two
+settings prevent that. ``max_concurrent_requests`` caps how many requests run at once, and
+the pool's pages divided by the pages one request needs, prompt and output together, is a
+safe value. ``cpu_offload_pages`` gives a full pool a request to move to the host.
+
 **Single GPU.** Everything on rank 0:
 
 .. code-block:: yaml

mstar/api_server/data_worker.py

@@ -24,7 +24,10 @@
 )
 from mstar.communication.communicator import BaseCommunicator, CommProtocol, make_communicator
 from mstar.communication.tensors import NameToTensorList, create_tensor_communication_manager
-from mstar.model.base import Model
+from mstar.engine.resources.kv.config import KVSpec
+from mstar.engine.resources.kv.keys import chain
+from mstar.engine.resources.spec import apply_yaml_overrides
+from mstar.model.base import Model, ProcessPromptOutput
 from mstar.profile.format import InputInfo, RxInfo, TxInfo
 from mstar.utils.ipc_format import (
     AbortRequest,
@@ -47,6 +50,16 @@ def _preprocess_loop(**kwargs):
     worker.run()
 
 
+def _kv_page_sizes(model: Model, model_config: dict) -> dict[str, int]:
+    """Page size per KV resource, after this deployment's YAML has been applied."""
+    specs = model.get_node_resources()
+    apply_yaml_overrides(specs, model_config)
+    return {
+        spec.resource_key: spec.config.page_size
+        for spec in specs if isinstance(spec, KVSpec)
+    }
+
+
 NameToLoopIndices = dict[str, NestedLoopIndices]
 
 
@@ -58,7 +71,8 @@ def __init__(
         socket_path_prefix: str = "/tmp/mstar",
         tensor_comm_protocol: CommProtocol = CommProtocol.RDMA,
         tcp_transfer_device="",
-        enable_prof: bool=False
+        enable_prof: bool=False,
+        model_config: dict | None = None,
     ):
         self.request_input_queue = queue.Queue()
         self.result_tensor_input_queue = queue.Queue()
@@ -108,7 +122,8 @@ def __init__(
                 communicator=self.communicator,
                 tensor_manager=self.tensor_manager,
                 model=model,
-                enable_prof=enable_prof
+                enable_prof=enable_prof,
+                model_config=model_config,
             )
         )
         self.thread.start()
@@ -260,8 +275,19 @@ def __init__(
         tensor_manager,
         device: str = "cpu",
         model: Model | None = None,
-        enable_prof: bool=False
+        enable_prof: bool=False,
+        model_config: dict | None = None,
     ):
+        # keying a stream needs the deployment's page size as well as the
+        # model's declaration, so a worker built without a config keys no stream
+        self._prefix_streams = (
+            model.prefix_key_streams()
+            if model is not None and model_config else {}
+        )
+        # resolved as the worker does at load, so both split a prompt into the same pages
+        self._prefix_page_sizes = (
+            _kv_page_sizes(model, model_config) if self._prefix_streams else {}
+        )
         self.in_queue = in_queue
         self.result_tensor_queue = result_tensor_queue
         self.cleanup_request_queue = cleanup_request_queue
@@ -332,6 +358,11 @@ def _process_input(
         # image_grid_thw, audio_features, audio_seqlens from the raw tensors
         # loaded above).  process_prompt receives the raw multimodal tensors
         # and returns any additional tensors to merge into the final dict.
+        model_kwargs = dict(input.model_kwargs or {})
+        # only this worker keys a prompt: a client that sent its own could name
+        # another request's pages and be served that request's KV
+        for name in ("prefix_keys", "prefix_tail", "prefix_decode"):
+            model_kwargs.pop(name, None)
         if self.model is not None:
             prompt_tensors = self.model.process_prompt(
                 input.text,
@@ -340,10 +371,21 @@ def _process_input(
                 tensors=tensors,
                 input_metadata=input_metadata,
                 prompt_parts=input.prompt_parts,
-                **(input.model_kwargs or {}),
+                **model_kwargs,
             )
+            if isinstance(prompt_tensors, ProcessPromptOutput):
+                model_kwargs.update(prompt_tensors.metadata)
+                prompt_tensors = prompt_tensors.new_input_tensors
             if prompt_tensors:
                 tensors.update(prompt_tensors)
+            # after the update: the chain keys the tensors the request will
+            # actually be prefilled with
+            prefix_keys, prefix_tail, prefix_decode = self._prefix_keys(tensors)
+            if prefix_keys:
+                model_kwargs["prefix_keys"] = prefix_keys
+                model_kwargs["prefix_tail"] = prefix_tail
+                if prefix_decode:
+                    model_kwargs["prefix_decode"] = prefix_decode
         elif input.text is not None:
             # Fallback: encode as UTF-8 bytes -> uint8 tensor
             byte_data = input.text.encode("utf-8")
@@ -368,7 +410,7 @@ def _process_input(
                 input.request_id, info.uuid, persist=True
             )
 
-        self.request_model_kwargs[input.request_id] = input.model_kwargs or {}
+        self.request_model_kwargs[input.request_id] = model_kwargs
         msg = ConductorMessage(
             message_type=ConductorMessageType.NEW_REQUEST,
             body=NewRequestConductor(
@@ -377,7 +419,7 @@ def _process_input(
                 initial_input_modalities=input.input_modalities,
                 initial_output_modalities=input.output_modalities,
                 input_metadata=input_metadata,
-                model_kwargs=input.model_kwargs
+                model_kwargs=model_kwargs
             ),
         )
         self.communicator.send("conductor", msg)
@@ -394,6 +436,33 @@ def _process_input(
                 inputs=self._summarize_inputs(input),
             ))
 
+    def _prefix_keys(self, tensors: dict) -> tuple[dict, dict, dict]:
+        """Key each page of every declared stream, by resource and label.
+
+        Returns the keys, the prompt tail past the last whole page, and the output
+        tensor each stream's sampled ids arrive in. The keys are unrooted.
+        """
+        keys: dict[str, dict[str, list[bytes]]] = {}
+        tails: dict[str, dict[str, list[int]]] = {}
+        decode: dict[str, dict[str, str]] = {}
+        for resource_key, by_label in self._prefix_streams.items():
+            page_size = self._prefix_page_sizes[resource_key]
+            for label, stream in by_label.items():
+                ids = tensors.get(stream.tensor)
+                if stream.keyed_by != "ids" or not ids:
+                    continue
+                flat = ids[0].flatten().tolist()
+                pages = [
+                    flat[at:at + page_size]
+                    for at in range(0, len(flat), page_size)
+                ]
+                keys.setdefault(resource_key, {})[label] = chain(pages)
+                whole = len(flat) // page_size
+                tails.setdefault(resource_key, {})[label] = flat[whole * page_size:]
+                if stream.decode_walk is not None:
+                    decode.setdefault(resource_key, {})[label] = stream.tensor
+        return keys, tails, decode
+
     @staticmethod
     def _summarize_inputs(input: PreprocessInput) -> list[InputInfo]:
         """Aggregate the *raw* (pre-decoding) inputs into per-modality sizes.

mstar/api_server/entrypoint.py

@@ -184,6 +184,7 @@ def __init__(
         model_name: str = "dummy",
         log_stats: bool = False,
         log_stats_file: str | None = None,
+        model_config: dict | None = None,
     ):
         self.upload_dir = Path(upload_dir)
         self.upload_dir.mkdir(parents=True, exist_ok=True)
@@ -204,6 +205,7 @@ def __init__(
 
         self.preprocess_worker = PreprocessWorker(
             model=model,
+            model_config=model_config,
             hostname=hostname,
             socket_path_prefix=socket_path_prefix,
             tensor_comm_protocol=tensor_comm_protocol,
@@ -1055,6 +1057,7 @@ def main(argv: list[str] | None = None):
         tensor_comm_protocol=CommProtocol(args.tensor_comm_protocol),
         model=model,
         model_name=model_name,
+        model_config=config,
         tcp_transfer_device=args.tcp_transfer_device,
         log_stats=log_stats,
         log_stats_file=args.log_stats_file,

mstar/conductor/conductor.py

@@ -26,7 +26,7 @@
 )
 from mstar.distributed.base import ShardingConfig
 from mstar.distributed.communication import GlobalParallelConfig, WorkerParallelGroups
-from mstar.engine.resources import ResourceReqConfig
+from mstar.engine.resources import KVReqConfig, ResourceReqConfig
 from mstar.graph.base import GraphEdge, NodeAndGraphWalk, TensorPointerInfo
 from mstar.graph.loop_indices import NestedLoopIndices
 from mstar.model.base import ForwardPassArgs, Model, WorkerGraph
@@ -413,10 +413,16 @@ def _get_resource_configs(
         Resolved once, here, and carried on the request: the worker hands each
         config to its resource at ingest. KV shape is not part of this — that
         is a deployment-wide property the model declares in its resource specs.
+
+        A declared stream gets a config even when the model returned none, or its
+        keys would have nowhere to go.
         """
-        return self.model.get_request_resource_configs(
+        configs = self.model.get_request_resource_configs(
             partition_fwd_args=partition_fwd_args, model_kwargs=model_kwargs
         )
+        for key, streams in self.model.prefix_key_streams().items():
+            configs.setdefault(key, KVReqConfig(needed_labels=list(streams)))
+        return configs
 
     def _derive_worker_info(self):
         """Derive per-rank worker info from the worker graphs."""
@@ -886,8 +892,20 @@ def _do_ingest_request(
         request_data.resource_configs = self._get_resource_configs(
             model_kwargs, partition_fwd_args
         )
-        for cfg in request_data.resource_configs.values():
-            cfg.apply_conductor_config(seed=seed)
+        # keyed by resource, so each config is handed only its own chain
+        kwargs = model_kwargs or {}
+        prefix_keys = kwargs.get("prefix_keys") or {}
+        prefix_tail = kwargs.get("prefix_tail") or {}
+        prefix_decode = kwargs.get("prefix_decode") or {}
+        prefix_cache = kwargs.get("prefix_cache")
+        for key, cfg in request_data.resource_configs.items():
+            cfg.apply_conductor_config(
+                seed=seed,
+                prefix_keys=prefix_keys.get(key),
+                prefix_tail=prefix_tail.get(key),
+                prefix_decode=prefix_decode.get(key),
+                prefix_cache=prefix_cache,
+            )
 
         # Send NewRequest to each worker with the appropriate partition's inputs
         for worker_id, worker_graph_ids in worker_to_worker_graph_ids.items():