From 441e94cc4daa6065b079b3bfab1d2551d39d7aa8 Mon Sep 17 00:00:00 2001 From: Gaurav Shah Date: Wed, 21 Oct 2026 05:50:50 +0000 Subject: [PATCH 01/26] kv: a page footprint from a label's layout --- mstar/engine/resources/kv/config.py | 38 +++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/mstar/engine/resources/kv/config.py b/mstar/engine/resources/kv/config.py index dcbcb6cb63..83ec9d9433 100644 --- a/mstar/engine/resources/kv/config.py +++ b/mstar/engine/resources/kv/config.py @@ -352,6 +352,44 @@ def __post_init__(self): raise ValueError(f"protected_prefix must be >= 0, got {self.protected_prefix}") +@dataclass(frozen=True) +class StreamLayout: + """Everything a request writes into one stream, known when its prompt is preprocessed. + + ``spans`` are the committed writes, in walk order. Each ``(loop, slots)`` in + ``growth`` writes ``slots`` per iteration of ``loop``. ``scratch`` is the most one + non-committing step writes past the committed length. + """ + spans: tuple[PrefixSpan, ...] = () + growth: tuple[tuple[str, int], ...] = () + scratch: int = 0 + retention: RetentionPolicy | None = None + + def pages(self, caps: Mapping[str, int], page_size: int) -> int: + """The most pages the stream can hold, where ``caps`` gives each growth loop's iterations. + + Under ``retention`` the stream holds a window. `KVManager._apply_retention` leaves + less than a page over the budget, so the window peaks at the protected prefix plus + the budget, or at the prefix rounded up to a page if that is larger. The bound adds + a page of overrun and the largest step past the prefix to that peak. + """ + written = sum(span.length for span in self.spans) + sum( + slots * caps[loop] for loop, slots in self.growth + ) + pages = -(-(written + self.scratch) // page_size) + if self.retention is None: + return pages + prefix = self.retention.protected_prefix + step, end = self.scratch, 0 + for span in self.spans: + end += span.length + if end > prefix: + step = max(step, span.length) + step = max([step, *(slots for _, slots in self.growth)]) + window = max(prefix + self.retention.context_budget, -(-prefix // page_size) * page_size) + return min(pages, -(-(window + page_size - 1 + step) // page_size)) + + @dataclass(frozen=True) class KVStep(ResourceStep): # write: bool # @nsagan: opting to remove this for now bc it's dead code From 8484a1a7a335b62c3d24ecb539cce005fe882649 Mon Sep 17 00:00:00 2001 From: Gaurav Shah Date: Wed, 21 Oct 2026 05:50:51 +0000 Subject: [PATCH 02/26] model, runtime: a loop's per-request cap becomes its max_iters --- mstar/conductor/conductor.py | 11 ++++++++++ mstar/graph/runtime/base.py | 6 ++++-- mstar/graph/runtime/python.py | 5 +++-- mstar/graph/runtime/rust.py | 2 ++ mstar/model/base.py | 9 +++++++++ mstar/utils/ipc_format.py | 2 ++ mstar/worker/node_manager_utils.py | 6 +++++- mstar/worker/worker.py | 1 + rust/src/graph/runtime.rs | 18 ++++++++++++++--- rust/src/graph/state.rs | 20 ++++++++++++++++--- test/modular/test_delivery_chunk_counts.py | 1 + test/modular/test_graph_runtime.py | 3 +++ test/modular/test_ingest_rollback.py | 2 +- test/modular/test_partition_bookkeeping.py | 1 + test/modular/test_rid_handles.py | 4 ++-- test/modular/test_startup_graph_validation.py | 2 ++ test/modular/test_stream_completion_order.py | 1 + test/modular/test_wire_codec.py | 2 +- test/modular/test_worker_admit_lifecycle.py | 1 + test/modular/test_worker_drain.py | 2 +- test/rust/test_runtime_lockstep.py | 2 +- test/rust/test_runtime_parity.py | 13 ++++++++++++ test/rust/test_rust_graph_runtime.py | 1 + test/rust/test_rust_wire_frames.py | 17 +++++++++------- 24 files changed, 108 insertions(+), 24 deletions(-) diff --git a/mstar/conductor/conductor.py b/mstar/conductor/conductor.py index 83584c424f..403d39d52d 100644 --- a/mstar/conductor/conductor.py +++ b/mstar/conductor/conductor.py @@ -641,6 +641,11 @@ def _derive_worker_info(self): worker_graph_id: set(worker_graph.section.get_loops()) for worker_graph_id, worker_graph in self.worker_graphs.items() } + # the largest max_iters of each loop name, so a cap clamped to it never cuts a longer loop of that name + self._loop_max_iters: dict[str, int] = {} + for worker_graph in self.worker_graphs.values(): + for name, loop in worker_graph.section.get_loops().items(): + self._loop_max_iters[name] = max(loop.max_iters, self._loop_max_iters.get(name, 0)) # set each group's _tp_rank (the worker's instance rank) per worker self.per_worker_sharding_config: dict[str, ShardingConfig] = {} @@ -1107,6 +1112,11 @@ def _do_ingest_request( model_kwargs = body.model_kwargs or {} max_output_tokens = self._max_output_tokens(model_kwargs) + caps = self.model.get_loop_caps(**model_kwargs) + # Clamp every cap to its loop's max_iters. Rust's graph runtime holds a cap + # as a u32, so a cap from a large max_output_tokens would fail the request + # there but not in Python. + loop_caps = {name: min(n, caps.get(name, n)) for name, n in self._loop_max_iters.items()} # Honor an explicit per-request seed (e.g. OpenAI ``seed``) when given; # otherwise derive a stable seed from the request id. ``_ingest_request`` # already rejected a bad one. @@ -1237,6 +1247,7 @@ def _do_ingest_request( partition_worker_graph_ids=partition_wg_ids, worker_graph_to_workers=worker_graph_to_workers, initial_inputs=inputs_per_worker.get(worker_id, []), + loop_caps=loop_caps, request_info=CurrentForwardPassInfo( request_id=body.request_id, graph_walk=fwd_args.full_metadata.graph_walk, diff --git a/mstar/graph/runtime/base.py b/mstar/graph/runtime/base.py index 525c0c4346..7e6099096e 100644 --- a/mstar/graph/runtime/base.py +++ b/mstar/graph/runtime/base.py @@ -542,10 +542,12 @@ def add_request( partition: str, graph_walk: str, partition_worker_graph_ids: list[int], - worker_graph_to_workers: ParallelList[int, list[str]] + worker_graph_to_workers: ParallelList[int, list[str]], + loop_caps: Mapping[str, int], ) -> int: """ - Returns the integer handle for this request. + Returns the integer handle for this request. Each loop named in + ``loop_caps`` runs at most that many iterations for it. """ pass diff --git a/mstar/graph/runtime/python.py b/mstar/graph/runtime/python.py index f4ca544c73..f2bb698c8e 100644 --- a/mstar/graph/runtime/python.py +++ b/mstar/graph/runtime/python.py @@ -213,7 +213,8 @@ def add_request( partition: str, graph_walk: str, partition_worker_graph_ids: list[int], - worker_graph_to_workers: ParallelList[int, list[str]] + worker_graph_to_workers: ParallelList[int, list[str]], + loop_caps: dict[str, int], ) -> int: # The conductor sends one NewRequest PER PARTITION, so this runs # several times for one request; they all share the first's handle. @@ -255,7 +256,7 @@ def add_request( for graph_id in partition_worker_graph_ids: if graph_id in self._queues: - self._queues[graph_id].add_request(handle) + self._queues[graph_id].add_request(handle, loop_caps) self._request_info[handle].partition_info[partition] = GraphRuntimePartitionInfo( graph_walk=graph_walk, diff --git a/mstar/graph/runtime/rust.py b/mstar/graph/runtime/rust.py index cfb1202ef2..1d9addbb9c 100644 --- a/mstar/graph/runtime/rust.py +++ b/mstar/graph/runtime/rust.py @@ -292,6 +292,7 @@ def add_request( self, request_id: str, partition: str, graph_walk: str, partition_worker_graph_ids: list[int], worker_graph_to_workers: ParallelList[int, list[str]], + loop_caps: dict[str, int], ) -> int: flat: list[str] = [] counts: list[int] = [] @@ -306,6 +307,7 @@ def add_request( # is derived from. list(partition_worker_graph_ids), list(worker_graph_to_workers.keys), flat, counts, + list(loop_caps), list(loop_caps.values()), ) # One NewRequest arrives PER PARTITION, all carrying the same map, so # the first one settles it. diff --git a/mstar/model/base.py b/mstar/model/base.py index 08584c7686..e306a467de 100644 --- a/mstar/model/base.py +++ b/mstar/model/base.py @@ -849,6 +849,15 @@ def get_submodule( def get_max_output_tokens(self, **model_kwargs): return model_kwargs.get("max_output_tokens", MAX_OUTPUT_TOKENS) + def get_loop_caps(self, **model_kwargs) -> dict[str, int]: + """The iterations each named loop may run for this request, each at least 1. + + A loop runs the smaller of its graph ``max_iters`` and its cap. A model + whose stop check ends a loop at a count sets the cap to that count, so the + loop ends on the same iteration and never speculates past it. + """ + return {} + def get_max_output_tokens_limit(self) -> int | None: """The largest ``max_output_tokens`` a request may ask for; above it the conductor refuses the request (a 400) rather than stopping short. None diff --git a/mstar/utils/ipc_format.py b/mstar/utils/ipc_format.py index 7162ca5ad7..9dae596d8e 100644 --- a/mstar/utils/ipc_format.py +++ b/mstar/utils/ipc_format.py @@ -50,6 +50,8 @@ class NewRequest(MessageBody): worker_graph_to_workers: dict[int, list[str]] initial_inputs: list[GraphEdge] request_info: CurrentForwardPassInfo + # loop name -> the most iterations it runs for this request (`Model.get_loop_caps`) + loop_caps: dict[str, int] class MessageSource(IntEnum): diff --git a/mstar/worker/node_manager_utils.py b/mstar/worker/node_manager_utils.py index e632ba7ce7..0dd45d2ccd 100644 --- a/mstar/worker/node_manager_utils.py +++ b/mstar/worker/node_manager_utils.py @@ -102,12 +102,16 @@ def is_done(self, request_id) -> bool: queue = self.per_request_queues[request_id] return queue.wg_state_registry.is_done - def add_request(self, request_id: int): + def add_request(self, request_id: int, loop_caps: dict[str, int]): """ Initialize queues for a new request """ section_copy = pickle.loads(self._section_blob) queue = WorkerGraphIO(section_copy, wg_id=self.worker_graph_id) + for name, cap in loop_caps.items(): + loop = queue.loops.get(name) + if loop is not None: + loop.max_iters = min(loop.max_iters, cap) queue.register_communication_info( self.tensor_manager, request_id ) diff --git a/mstar/worker/worker.py b/mstar/worker/worker.py index ec3256d3a8..8b86724ea7 100644 --- a/mstar/worker/worker.py +++ b/mstar/worker/worker.py @@ -795,6 +795,7 @@ def _add_new_request(self, body: NewRequest) -> None: graph_walk=body.request_info.graph_walk, partition_worker_graph_ids=body.partition_worker_graph_ids, worker_graph_to_workers=ParallelList.from_dict(body.worker_graph_to_workers), + loop_caps=body.loop_caps, ) body.request_info.rid_handle = request_id diff --git a/rust/src/graph/runtime.rs b/rust/src/graph/runtime.rs index c5d07e25ce..54df9016a2 100644 --- a/rust/src/graph/runtime.rs +++ b/rust/src/graph/runtime.rs @@ -760,7 +760,7 @@ pub struct GraphRuntime { let Some(state) = self.state(wg, rid) else { return false; }; - state.loop_iter(lid) + 1 < g.lp(lid).max_iters && !state.loop_finished(lid) + state.loop_iter(lid) + 1 < state.loop_max_iters(lid) && !state.loop_finished(lid) } /// One rid: ingest its chunks, check readiness, gather inputs. Returns @@ -1607,6 +1607,9 @@ impl GraphRuntime { /// /// `worker_graph_workers` is flattened; `workers_per_worker_graph` gives /// each worker graph's slice, parallel to `worker_graph_ids`. + /// + /// Each loop in `loop_names` runs at most the matching `loop_caps` + /// iterations for this request. // Struct-of-arrays across the boundary, as everywhere else here: a dict // per worker graph would be a Python allocation per admission. #[allow(clippy::too_many_arguments)] @@ -1619,6 +1622,8 @@ impl GraphRuntime { worker_graph_ids: Vec, worker_graph_workers: Vec, workers_per_worker_graph: Vec, + loop_names: Vec, + loop_caps: Vec, ) -> PyResult { let (handle, grew) = self.rids.intern(&request_id); if grew { @@ -1694,8 +1699,15 @@ impl GraphRuntime { // In range by construction: the vectors were grown for this // handle above. if self.states[wg as usize][handle as usize].is_none() { - self.states[wg as usize][handle as usize] = - Some(RequestState::new(self.graphs[wg as usize].clone())); + let g = &self.graphs[wg as usize]; + let mut state = RequestState::new(g.clone()); + for (name, &cap) in loop_names.iter().zip(&loop_caps) { + let lid = self.interner.get(name).and_then(|s| g.loop_by_name.get(&s)); + if let Some(&lid) = lid { + state.cap_loop(lid, cap); + } + } + self.states[wg as usize][handle as usize] = Some(state); } } diff --git a/rust/src/graph/state.rs b/rust/src/graph/state.rs index d8e27dd853..87965a4c0c 100644 --- a/rust/src/graph/state.rs +++ b/rust/src/graph/state.rs @@ -95,6 +95,7 @@ impl NodeState { #[derive(Clone, Default)] struct LoopState { + max_iters: u32, // Loop.max_iters, this request's copy curr_iter: u32, // Loop.curr_iter finish_signal: bool, // Loop._finish_signal done: bool, // Loop.is_done @@ -193,7 +194,11 @@ pub struct RequestState { impl RequestState { pub fn new(graph: GraphRef) -> Self { let nodes = graph.nodes.iter().map(|n| NodeState::new(n.inputs.len())).collect(); - let loops = vec![LoopState::default(); graph.loops.len()]; + let loops = graph + .loops + .iter() + .map(|l| LoopState { max_iters: l.max_iters, ..LoopState::default() }) + .collect(); let w = graph.n_ready_words(); let mut s = Self { nodes, loops, ready: vec![0; w], ready_streaming: vec![0; w], @@ -563,7 +568,7 @@ impl RequestState { let finishing = { let st = &self.loops[lid as usize]; - st.finish_signal || lspec.max_iters == st.curr_iter + 1 + st.finish_signal || st.max_iters == st.curr_iter + 1 }; if finishing { self.finish_loop(g, lid, out, filtered); @@ -764,6 +769,15 @@ impl RequestState { st.done || st.finish_signal } + pub fn loop_max_iters(&self, lid: LoopId) -> u32 { + self.loops[lid as usize].max_iters + } + + pub fn cap_loop(&mut self, lid: LoopId, cap: u32) { + let st = &mut self.loops[lid as usize]; + st.max_iters = st.max_iters.min(cap); + } + pub fn loop_indices(&self) -> Vec { self.loops.iter().map(|l| l.curr_iter).collect() } @@ -840,7 +854,7 @@ impl RequestState { n.clear(); } for l in self.loops.iter_mut() { - *l = LoopState::default(); + *l = LoopState { max_iters: l.max_iters, ..LoopState::default() }; } for w in self.ready.iter_mut() { *w = 0 } for w in self.ready_streaming.iter_mut() { *w = 0 } diff --git a/test/modular/test_delivery_chunk_counts.py b/test/modular/test_delivery_chunk_counts.py index ce1ad9f6f8..b41f182264 100644 --- a/test/modular/test_delivery_chunk_counts.py +++ b/test/modular/test_delivery_chunk_counts.py @@ -91,6 +91,7 @@ def drive(runtime, put, prefills: int, decode_steps: int, inline: bool = False) request_id="r1", partition="default", graph_walk=PREFILL, partition_worker_graph_ids=[0, 1], worker_graph_to_workers=ParallelList([0, 1], [[WORKER], [WORKER]]), + loop_caps={}, ) uuid = 100 diff --git a/test/modular/test_graph_runtime.py b/test/modular/test_graph_runtime.py index 5365652df8..617097764b 100644 --- a/test/modular/test_graph_runtime.py +++ b/test/modular/test_graph_runtime.py @@ -167,6 +167,7 @@ def _build(section, wg_id, walk, nodes, loops=frozenset(), worker_id="worker0"): graph_walk=walk, partition_worker_graph_ids=[wg_id], worker_graph_to_workers=ParallelList([wg_id], [[worker_id]]), + loop_caps={}, ) mgr.add_request(rid, fwd_info) return mgr, runtime, rid @@ -523,6 +524,7 @@ def test_route_batch_decodes_the_flat_rid_major_layout(): request_id="rid_b", partition="default", graph_walk="decode", partition_worker_graph_ids=[0], worker_graph_to_workers=ParallelList([0], [["worker0"]]), + loop_caps={}, ) mgr.add_request(rid_b, _fwd_info("decode")) for rid in (rid_a, rid_b): @@ -1011,6 +1013,7 @@ def _build_on_real_manager(section, nodes, tmp_path, loops=frozenset()): request_id="r", partition="default", graph_walk="w", partition_worker_graph_ids=[0], worker_graph_to_workers=ParallelList([0], [["worker_0"]]), + loop_caps={}, ) tm.register_request(rid, runtime.get_sharding_config(rid)) return tm, runtime, rid diff --git a/test/modular/test_ingest_rollback.py b/test/modular/test_ingest_rollback.py index 9fc68919b9..29899bc6e6 100644 --- a/test/modular/test_ingest_rollback.py +++ b/test/modular/test_ingest_rollback.py @@ -63,7 +63,7 @@ def test_the_worker_fails_only_the_request_whose_ingest_raised(): w._fail_requests = failed.append w._process_message_list = processed.extend body = SimpleNamespace( - request_id="wire", partition_worker_graph_ids=None, worker_graph_to_workers={}, + request_id="wire", partition_worker_graph_ids=None, worker_graph_to_workers={}, loop_caps={}, request_info=SimpleNamespace( partition_name="p", graph_walk="g", resource_configs=None, rid_handle=None, ), diff --git a/test/modular/test_partition_bookkeeping.py b/test/modular/test_partition_bookkeeping.py index 12415c416d..12f41336be 100644 --- a/test/modular/test_partition_bookkeeping.py +++ b/test/modular/test_partition_bookkeeping.py @@ -66,6 +66,7 @@ def add_both_partitions(rt): rid = rt.add_request( request_id="rid", partition=part, graph_walk=walk, partition_worker_graph_ids=[wg], worker_graph_to_workers=ParallelList([0, 1], [["worker_0"], ["worker_0"]]), + loop_caps={}, ) return rid diff --git a/test/modular/test_rid_handles.py b/test/modular/test_rid_handles.py index e5efb1a4e5..637593238c 100644 --- a/test/modular/test_rid_handles.py +++ b/test/modular/test_rid_handles.py @@ -257,7 +257,7 @@ def test_add_new_request_hands_the_handle_to_every_subsystem(): body = SimpleNamespace( request_id="wire-1", request_info=info, partition_worker_graph_ids={}, worker_graph_to_workers={}, - initial_inputs=[], + initial_inputs=[], loop_caps={}, ) w.request_state.per_request_info = {0: cfg} @@ -302,7 +302,7 @@ def new(partition, lead): ) body = SimpleNamespace( request_id="wire-1", request_info=info, partition_worker_graph_ids={}, - worker_graph_to_workers={}, initial_inputs=[], + worker_graph_to_workers={}, initial_inputs=[], loop_caps={}, ) Worker._add_new_request(w, body) diff --git a/test/modular/test_startup_graph_validation.py b/test/modular/test_startup_graph_validation.py index 5ce7934d2d..808d930a79 100644 --- a/test/modular/test_startup_graph_validation.py +++ b/test/modular/test_startup_graph_validation.py @@ -285,6 +285,7 @@ def __init__(self, node_groups, tmp_path): self.rid = self.rt.add_request( request_id="r", partition="C", graph_walk="chunk", partition_worker_graph_ids=ids, worker_graph_to_workers=ParallelList(ids, [["worker_0"]] * len(ids)), + loop_caps={}, ) self.minter = TensorUuidMinter("worker_0") @@ -502,6 +503,7 @@ def __init__(self, runtime, tmp_path): self.rid = self.rt.add_request( request_id="r", partition="default", graph_walk="w", partition_worker_graph_ids=[0], worker_graph_to_workers=ParallelList([0], [["worker_0"]]), + loop_caps={}, ) self.minter = TensorUuidMinter("worker_0") _ingest(self.rt, self.rid, [GraphEdge(next_node="step", name="x")]) diff --git a/test/modular/test_stream_completion_order.py b/test/modular/test_stream_completion_order.py index 7f46572300..0ee6ae00db 100644 --- a/test/modular/test_stream_completion_order.py +++ b/test/modular/test_stream_completion_order.py @@ -141,6 +141,7 @@ def __init__(self, kind, tmp_path): request_id=RID, partition="LLM", graph_walk="decode", partition_worker_graph_ids=[0], worker_graph_to_workers=ParallelList([0, 1], [[PRODUCER], [CONSUMER]]), + loop_caps={}, ) def step(self, uuid): diff --git a/test/modular/test_wire_codec.py b/test/modular/test_wire_codec.py index 984ec5ae34..add32d8793 100644 --- a/test/modular/test_wire_codec.py +++ b/test/modular/test_wire_codec.py @@ -158,7 +158,7 @@ def test_polymorphic_body_dispatch(): WorkerMessageType.SCHEDULE_TP), (NewRequest(request_id="r", partition_worker_graph_ids=["wg0"], worker_graph_to_workers={"wg0": ["w0", "w1"]}, - initial_inputs=[_edge()], request_info=_fwd_info()), + initial_inputs=[_edge()], request_info=_fwd_info(), loop_caps={}), WorkerMessageType.NEW_REQUEST), ]: got = roundtrip(WorkerMessage(message_type=mtype, body=body)) diff --git a/test/modular/test_worker_admit_lifecycle.py b/test/modular/test_worker_admit_lifecycle.py index 752eab4df1..5b86a082ba 100644 --- a/test/modular/test_worker_admit_lifecycle.py +++ b/test/modular/test_worker_admit_lifecycle.py @@ -113,6 +113,7 @@ def _new_request(request_id="r1"): partition_worker_graph_ids=[WG_ID], worker_graph_to_workers={WG_ID: [WORKER]}, initial_inputs=[], + loop_caps={}, ) diff --git a/test/modular/test_worker_drain.py b/test/modular/test_worker_drain.py index 774b34f08b..27286b0c59 100644 --- a/test/modular/test_worker_drain.py +++ b/test/modular/test_worker_drain.py @@ -426,7 +426,7 @@ def _state_add(request_id, request_info): info = SimpleNamespace(resource_configs={}, graph_walk="g", partition_name="p") Worker._add_new_request(w, SimpleNamespace( request_id="X", request_info=info, initial_inputs=[], - partition_worker_graph_ids=[], worker_graph_to_workers={}, + partition_worker_graph_ids=[], worker_graph_to_workers={}, loop_caps={}, )) handle = w._graph_runtime.get_rid_handle("X") diff --git a/test/rust/test_runtime_lockstep.py b/test/rust/test_runtime_lockstep.py index 1f800ef606..d8c6d1e692 100644 --- a/test/rust/test_runtime_lockstep.py +++ b/test/rust/test_runtime_lockstep.py @@ -190,7 +190,7 @@ def admit(self, rid_str="r1"): return self._both(f"add_request({rid_str})", lambda rt, b, s: rt.add_request( request_id=rid_str, partition="default", graph_walk=WALK, partition_worker_graph_ids=[WG_ID], - worker_graph_to_workers=ParallelList([WG_ID], [[WORKER]]))) + worker_graph_to_workers=ParallelList([WG_ID], [[WORKER]]), loop_caps={})) def put(self, uuid): """Store a tensor on both sides so routing has something to refer to.""" diff --git a/test/rust/test_runtime_parity.py b/test/rust/test_runtime_parity.py index de7910df1c..e925f05af9 100644 --- a/test/rust/test_runtime_parity.py +++ b/test/rust/test_runtime_parity.py @@ -176,6 +176,7 @@ def _admit(rt, rid="r1"): request_id=rid, partition="default", graph_walk=WALK, partition_worker_graph_ids=[WG_ID], worker_graph_to_workers=ParallelList([WG_ID], [[WORKER]]), + loop_caps={}, ) @@ -296,6 +297,7 @@ def test_one_handle_per_request_across_partitions(pair): request_id="r", partition=p, graph_walk=WALK, partition_worker_graph_ids=[WG_ID], worker_graph_to_workers=ParallelList([WG_ID], [[WORKER]]), + loop_caps={}, ) for p in ("Thinker", "Talker", "Code2Wav") ] @@ -310,6 +312,7 @@ def test_one_handle_per_request_across_partitions(pair): request_id="r2", partition="Thinker", graph_walk=WALK, partition_worker_graph_ids=[WG_ID], worker_graph_to_workers=ParallelList([WG_ID], [[WORKER]]), + loop_caps={}, ) == handles[0] @@ -928,6 +931,7 @@ def test_a_tensor_read_by_two_workers_holds_two_references(two_consumers): worker_graph_to_workers=ParallelList( [WG_ID, 1], [[WORKER], ["worker1", "worker2"]], ), + loop_caps={}, ) rt.ingest_inputs_batch(_ingest_block([rid], [_spec("prompt", "only")])) rt.pop_rids("only", WALK, [rid]) @@ -1109,6 +1113,7 @@ def test_the_node_becomes_ready_again_in_the_next_walk(two_walks): request_id="r1", partition="default", graph_walk="prefill", partition_worker_graph_ids=[0, 1], worker_graph_to_workers=ParallelList([0, 1], [[WORKER], [WORKER]]), + loop_caps={}, ) assert _run_walk(rt, book, store, rid, "prefill", 0, uuid=1) == [ ("LLM", "prefill", [rid]) @@ -1127,6 +1132,7 @@ def _one_rid_per_walk(rt): request_id=f"r{walk}", partition="default", graph_walk=walk, partition_worker_graph_ids=[0, 1], worker_graph_to_workers=ParallelList([0, 1], [[WORKER], [WORKER]]), + loop_caps={}, ) for walk in ("prefill", "decode") ] @@ -1205,6 +1211,7 @@ def test_a_batch_transitions_together(two_walks): request_id=f"r{i}", partition="default", graph_walk="prefill", partition_worker_graph_ids=[0, 1], worker_graph_to_workers=ParallelList([0, 1], [[WORKER], [WORKER]]), + loop_caps={}, ) for i in range(3) ] @@ -1291,6 +1298,7 @@ def test_an_edge_into_a_sibling_local_graph_is_ingested_not_only_sent( request_id="r1", partition="default", graph_walk=WALK, partition_worker_graph_ids=[0, 1], worker_graph_to_workers=ParallelList([0, 1], [[WORKER], [WORKER]]), + loop_caps={}, ) rt.ingest_inputs_batch(_ingest_block([rid], [_spec("image_inputs", "encoder")])) assert _ready(rt) == [("encoder", WALK, [rid])] @@ -1367,6 +1375,7 @@ def test_a_stream_to_a_sibling_graph_is_reported_local(streams_to_sibling): request_id="r1", partition="default", graph_walk=WALK, partition_worker_graph_ids=[0, 1], worker_graph_to_workers=ParallelList([0, 1], [[WORKER], [WORKER]]), + loop_caps={}, ) rt.ingest_inputs_batch(_ingest_block([rid], [_spec("text_inputs", "LLM")])) rt.pop_rids("LLM", WALK, [rid]) @@ -1403,6 +1412,7 @@ def test_a_sibling_graph_chunk_is_held_once_not_twice(streams_to_sibling): request_id="r1", partition="default", graph_walk=WALK, partition_worker_graph_ids=[0, 1], worker_graph_to_workers=ParallelList([0, 1], [[WORKER], [WORKER]]), + loop_caps={}, ) rt.ingest_inputs_batch(_ingest_block([rid], [_spec("text_inputs", "LLM")])) rt.pop_rids("LLM", WALK, [rid]) @@ -1527,6 +1537,7 @@ def test_a_finished_worker_graph_asks_for_request_info(streams_to_sibling): request_id="r1", partition="default", graph_walk=WALK, partition_worker_graph_ids=[0, 1], worker_graph_to_workers=ParallelList([0, 1], [[WORKER], [WORKER]]), + loop_caps={}, ) rt.ingest_inputs_batch(_ingest_block([rid], [_spec("text_inputs", "LLM")])) rt.pop_rids("LLM", WALK, [rid]) @@ -1865,6 +1876,7 @@ def _drive_tp(rt, book, store, steps=8): request_id="r1", partition="default", graph_walk=TP_WALK, partition_worker_graph_ids=[0, 1], worker_graph_to_workers=ParallelList([0, 1], [[WORKER], [WORKER, PEER]]), + loop_caps={}, ) rt.ingest_inputs_batch(_ingest_block([rid], [_spec("prompt", "src")])) @@ -2017,6 +2029,7 @@ def _admit_with_remote(rt, remote): request_id="r1", partition="default", graph_walk=WALK, partition_worker_graph_ids=[WG_ID, 1], worker_graph_to_workers=ParallelList([WG_ID, 1], [[WORKER], ["worker1"]]), + loop_caps={}, ) diff --git a/test/rust/test_rust_graph_runtime.py b/test/rust/test_rust_graph_runtime.py index e4ce37d9cf..f2e35fd657 100644 --- a/test/rust/test_rust_graph_runtime.py +++ b/test/rust/test_rust_graph_runtime.py @@ -121,6 +121,7 @@ def _admit(runtime, request_id="r1", walk=WALK): request_id=request_id, partition="default", graph_walk=walk, partition_worker_graph_ids=[WG_ID], worker_graph_to_workers=ParallelList([WG_ID], [[WORKER]]), + loop_caps={}, ) diff --git a/test/rust/test_rust_wire_frames.py b/test/rust/test_rust_wire_frames.py index d6e4be5bc9..9c66d68501 100644 --- a/test/rust/test_rust_wire_frames.py +++ b/test/rust/test_rust_wire_frames.py @@ -96,6 +96,7 @@ def admit(self, wgs=(WG_ID,), workers=(ME,)): # request's are the same list. "r1", "default", WALK, list(wgs), list(wgs), list(workers), [1] * len(wgs), + [], [], ) def run(self, rid, uuids, **send_kwargs): @@ -468,7 +469,7 @@ def _loop_mesh(tmp_path): bookkeeping=book._rust, me=ME, communicator=comm, ) rid = rt.add_request("r1", "default", WALK, [WG_ID, 1], [WG_ID, 1], - [ME, PEER], [1, 1]) + [ME, PEER], [1, 1], [], []) return rt, rid, inboxes @@ -517,7 +518,7 @@ def _emit_mesh(tmp_path, tp_rank, shard_dim=None): }, bookkeeping=book._rust, me=ME, communicator=comm, ) - rid = rt.add_request("r1", "default", WALK, [WG_ID], [WG_ID], [ME, PEER], [2]) + rid = rt.add_request("r1", "default", WALK, [WG_ID], [WG_ID], [ME, PEER], [2], [], []) info = TensorPointerInfo( dims=[2, 3], dtype=torch.bfloat16, nbytes=12, address=1, stride=[3, 1], uuid=1, source_session_id="h:1", source_entity=ME, @@ -562,7 +563,7 @@ def _persist_mesh(tmp_path, tp_rank): }, bookkeeping=book._rust, me=ME, communicator=comm, ) - rid = rt.add_request("r1", "default", WALK, [WG_ID], [WG_ID], [ME, PEER], [2]) + rid = rt.add_request("r1", "default", WALK, [WG_ID], [WG_ID], [ME, PEER], [2], [], []) info = TensorPointerInfo( dims=[2, 3], dtype=torch.bfloat16, nbytes=12, address=1, stride=[3, 1], uuid=1, source_session_id="h:1", source_entity=ME, @@ -662,7 +663,7 @@ def _fanout_mesh(tmp_path, shard_dim=None, dims=(4, 3), nbytes=24, stride=None): bookkeeping=book._rust, me=ME, communicator=comm, ) rid = rt.add_request("r1", "default", WALK, [WG_ID, 1], [WG_ID, 1], - [ME, PEER, P2], [1, 2]) + [ME, PEER, P2], [1, 2], [], []) # By default 4 rows of 6 bytes. book.put_tensor(1, _info( 1, dims=list(dims), nbytes=nbytes, @@ -796,7 +797,7 @@ def _gather_mesh(tmp_path, src_tp, dest_tp, my_rank=0, shard_dim=0, srcs = [ME, PEER][:src_tp] dests = [PEER, P2][:dest_tp] rid = rt.add_request("r1", "default", WALK, [WG_ID, 1], [WG_ID, 1], - srcs + dests, [src_tp, dest_tp]) + srcs + dests, [src_tp, dest_tp], [], []) book.put_tensor(1, _info(1, dims=[4, 3], nbytes=24)) book.increment_ref(1, 1) rt.ingest_inputs_batch(_one_signal(rid, "prompt", "src")) @@ -893,7 +894,7 @@ def _looping_emit_mesh(tmp_path): "tp_enabled_nodes": [], "sp_enabled_nodes": []}, bookkeeping=book._rust, me=ME, communicator=comm, ) - rid = rt.add_request("r1", "default", WALK, [WG_ID], [WG_ID], [ME], [1]) + rid = rt.add_request("r1", "default", WALK, [WG_ID], [WG_ID], [ME], [1], [], []) return rt, book, rid, inboxes @@ -997,6 +998,7 @@ def test_a_removed_rid_leaves_no_persist_signal_for_the_next_handle(tmp_path): keep = mesh.admit() doomed = mesh.rt.add_request( "doomed", "default", WALK, [WG_ID], [WG_ID], [ME], [1], + [], [], ) _put(mesh, 1) _put(mesh, 2) @@ -1014,6 +1016,7 @@ def test_a_removed_rid_leaves_no_persist_signal_for_the_next_handle(tmp_path): mesh.rt.remove_request(doomed) reused = mesh.rt.add_request( "fresh", "default", WALK, [WG_ID], [WG_ID], [ME], [1], + [], [], ) assert reused == doomed, "the handle must be recycled for this to bite" @@ -1052,7 +1055,7 @@ def test_one_route_reports_each_row_in_its_own_walks_graph(tmp_path): bookkeeping=book._rust, me=ME, communicator=comm, ) rids = [ - rt.add_request(name, "default", walk, [0, 1], [0, 1], [ME, ME], [1, 1]) + rt.add_request(name, "default", walk, [0, 1], [0, 1], [ME, ME], [1, 1], [], []) for name, walk in (("r_prefill", "prefill"), ("r_decode", "decode")) ] for rid in rids: From c195820781568b9f03aabe18f0614c9e5998bc9f Mon Sep 17 00:00:00 2001 From: Gaurav Shah Date: Wed, 21 Oct 2026 05:50:51 +0000 Subject: [PATCH 03/26] model, kv: process_prompt's layout reaches the KV manager --- mstar/api_server/data_worker.py | 32 +++++++++++++++++++-- mstar/conductor/conductor.py | 9 ++++++ mstar/engine/resources/kv/config.py | 30 ++++++++++++++++++- mstar/engine/resources/kv/manager.py | 12 ++++++-- mstar/model/base.py | 20 +++++++++++-- mstar/worker/worker.py | 3 ++ test/modular/test_kv_prefix_request_keys.py | 2 +- 7 files changed, 100 insertions(+), 8 deletions(-) diff --git a/mstar/api_server/data_worker.py b/mstar/api_server/data_worker.py index 389f7a6835..46205eb931 100644 --- a/mstar/api_server/data_worker.py +++ b/mstar/api_server/data_worker.py @@ -31,7 +31,7 @@ from mstar.engine.resources.kv.config import KVSpec, PagedKVConfig from mstar.engine.resources.kv.keys import PageItem, chain, end_key, fingerprint, page_items from mstar.engine.resources.spec import apply_yaml_overrides -from mstar.model.base import Model, PrefixStream, ProcessPromptOutput, Span +from mstar.model.base import LabelLayout, Model, PrefixStream, ProcessPromptOutput, Span from mstar.model.media_limits import ( RequestMediaBudget, charge_decoded_audio, @@ -140,6 +140,32 @@ def _kv_page_sizes(model: Model, model_config: dict) -> dict[str, int]: } +def _layout_rows(kv_layout: dict[str, dict[str, LabelLayout]]) -> dict: + """``kv_layout`` as plain rows, which keep the request's kwargs on the wire's plain msgpack path. + + Each label becomes ``[spans, growth, scratch, retention]``, a span ``[length, advance, walk]``. + """ + rows: dict[str, dict[str, list]] = {} + for resource_key, by_label in kv_layout.items(): + for label, layout in by_label.items(): + for span in layout.spans: + if span.kind not in ("ids", "digest", "rows"): + raise ValueError(f"a span of kind {span.kind!r}") + if not (isinstance(span.length, int) and isinstance(span.advance, int)): + raise ValueError( + f"a {span.kind} span's length and advance are " + f"{type(span.length).__name__} and {type(span.advance).__name__}" + ) + retention = layout.retention + rows.setdefault(resource_key, {})[label] = [ + [[span.length, span.advance, span.walk] for span in layout.spans], + [[loop, slots] for loop, slots in layout.growth], + layout.scratch, + None if retention is None else [retention.context_budget, retention.protected_prefix], + ] + return rows + + class PreprocessWorker: @@ -622,7 +648,7 @@ def _process_input( # another request's pages and be served that request's KV. Likewise # only process_prompt sets the model's INTERNAL_KWARGS. internal = getattr(self.model, "INTERNAL_KWARGS", ()) - for name in ("prefix_keys", "prefix_tail", "prefix_decode", "prefix_layout", *internal): + for name in ("prefix_keys", "prefix_tail", "prefix_decode", "prefix_layout", "kv_layout", *internal): model_kwargs.pop(name, None) if self.model is not None: prompt_tensors = self.model.process_prompt( @@ -640,6 +666,8 @@ def _process_input( if prompt_tensors: tensors.update(prompt_tensors) layouts = model_kwargs.pop("prefix_layout", {}) + if "kv_layout" in model_kwargs: + model_kwargs["kv_layout"] = _layout_rows(model_kwargs["kv_layout"]) # after the update: the chain keys the tensors the request will # actually be prefilled with if model_kwargs.get("prefix_cache") is not False: diff --git a/mstar/conductor/conductor.py b/mstar/conductor/conductor.py index 403d39d52d..cf7b1fabbe 100644 --- a/mstar/conductor/conductor.py +++ b/mstar/conductor/conductor.py @@ -32,6 +32,7 @@ from mstar.distributed.base import ShardingConfig from mstar.distributed.communication import GlobalParallelConfig, WorkerParallelGroups from mstar.engine.resources import KVReqConfig, ResourceReqConfig +from mstar.engine.resources.kv.config import without_layouts from mstar.graph.base import GraphEdge, NodeAndGraphWalk, TensorPointerInfo from mstar.model.base import ForwardPassArgs, Model, WorkerGraph from mstar.observability import catalog @@ -604,6 +605,8 @@ def _get_resource_configs( ) for key, streams in self.model.prefix_key_streams().items(): configs.setdefault(key, KVReqConfig(needed_labels=list(streams))) + for key in model_kwargs.get("kv_layout") or {}: + configs.setdefault(key, KVReqConfig()) return configs def _derive_worker_info(self): @@ -1211,6 +1214,7 @@ def _do_ingest_request( prefix_decode = kwargs.get("prefix_decode") or {} prefix_layout = kwargs.get("prefix_layout") or {} prefix_cache = kwargs.get("prefix_cache") + kv_layout = kwargs.get("kv_layout") or {} for key, cfg in request_data.resource_configs.items(): cfg.apply_conductor_config( seed=seed, @@ -1220,6 +1224,8 @@ def _do_ingest_request( prefix_decode=prefix_decode.get(key), prefix_layout=prefix_layout.get(key), prefix_cache=prefix_cache, + kv_layout=kv_layout.get(key), + loop_caps=loop_caps, ) # Send NewRequest to each worker with the appropriate partition's inputs @@ -1267,6 +1273,9 @@ def _do_ingest_request( ), ) + # each worker's KV resource opens with the layouts on its NewRequest, so later messages go without them + request_data.resource_configs = without_layouts(request_data.resource_configs) + # After the loop: a tensor's reader count must cover every worker and # partition it went to before the producer may free it. self._un_persist_tensors(body.request_id, [ diff --git a/mstar/engine/resources/kv/config.py b/mstar/engine/resources/kv/config.py index 83ec9d9433..52c523d693 100644 --- a/mstar/engine/resources/kv/config.py +++ b/mstar/engine/resources/kv/config.py @@ -6,7 +6,7 @@ from abc import ABC, abstractmethod from collections.abc import Mapping -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from enum import Enum from typing import TYPE_CHECKING @@ -248,6 +248,10 @@ class KVReqConfig(ResourceReqConfig): # label -> the layout its keys were chained over, one span per write prefix_layout: dict[str, list[PrefixSpan]] | None = None prefix_cache: bool = True + # label -> everything the request writes there + kv_layout: dict[str, "StreamLayout"] | None = None + # loop -> the most iterations it runs, for each loop a layout grows on + loop_caps: dict[str, int] | None = None def apply_conductor_config( self, @@ -256,6 +260,8 @@ def apply_conductor_config( prefix_decode: dict[str, str] | None=None, prefix_layout: dict[str, list[list]] | None=None, prefix_cache: bool | None=None, + kv_layout: dict[str, list] | None=None, + loop_caps: Mapping[str, int] | None=None, **kwargs, ): if prefix_keys is not None: @@ -272,6 +278,20 @@ def apply_conductor_config( } if prefix_cache is not None: self.prefix_cache = prefix_cache + if kv_layout is not None: + self.kv_layout = { + label: StreamLayout( + spans=tuple(PrefixSpan(*row) for row in spans), + growth=tuple((loop, slots) for loop, slots in growth), + scratch=scratch, + retention=None if retention is None else RetentionPolicy(*retention), + ) + for label, (spans, growth, scratch, retention) in kv_layout.items() + } + self.loop_caps = { + loop: loop_caps[loop] + for layout in self.kv_layout.values() for loop, _ in layout.growth + } def get_labels(self, node: str, walk: str): if (node, walk) in self.needed_labels_per_node_walk: @@ -305,6 +325,14 @@ def get_publish_labels( return mapping.get((node, walk), []) +def without_layouts(configs: Mapping[str, ResourceReqConfig]) -> dict[str, ResourceReqConfig]: + """``configs`` with each KV config's layouts and caps left out. Only the KV manager's ingest + reads them, and the request's info that carries the configs rides every frame between workers.""" + return { + key: replace(config, kv_layout=None, loop_caps=None) if isinstance(config, KVReqConfig) else config + for key, config in configs.items() + } + @dataclass class KVSpec(NodeResourceSpec): config: KVConfig diff --git a/mstar/engine/resources/kv/manager.py b/mstar/engine/resources/kv/manager.py index 2587164f0b..a1ad697bf9 100644 --- a/mstar/engine/resources/kv/manager.py +++ b/mstar/engine/resources/kv/manager.py @@ -234,6 +234,8 @@ class CacheStream: # its prefix lookup is in `KVManager.prefix_tokens`, so a re-probe after a # miss (a refused admit prepares the step again) is not counted twice prefix_counted: bool = False + # the most pages the request's layout lets this stream hold; None without a layout + footprint: int | None = None def reset(self, freed: bool=False, *, content_reset: bool=True): self.stored_len = 0 @@ -1266,9 +1268,15 @@ def ingest_request(self, rid, overrides: KVReqConfig | None=None): self._streams.setdefault(rid, {"main": CacheStream()}) self._overrides.setdefault(rid, overrides) for label, stream in self._streams[rid].items(): + stream.footprint = self._footprint(rid, label) if stream.chain is None: self._seed_keys(rid, label, stream) + def _footprint(self, rid: str, label: str) -> int | None: + overrides = self._overrides.get(rid) + layout = (overrides.kv_layout or {}).get(label) if overrides is not None else None + return None if layout is None else layout.pages(overrides.loop_caps, self.config.page_size) + def admit_retrieve( self, rid: str, node_name: str, @@ -2972,8 +2980,8 @@ def cleanup(self): def _ensure_label(self, rid: str, label: str) -> CacheStream: if label not in self._streams[rid]: - self._streams[rid][label] = CacheStream() - self._seed_keys(rid, label, self._streams[rid][label]) + stream = self._streams[rid][label] = CacheStream(footprint=self._footprint(rid, label)) + self._seed_keys(rid, label, stream) return self._streams[rid][label] def _check_ready( diff --git a/mstar/model/base.py b/mstar/model/base.py index e306a467de..dbbabf37b1 100644 --- a/mstar/model/base.py +++ b/mstar/model/base.py @@ -16,6 +16,7 @@ ) from mstar.distributed.base import ShardingConfig, ShardingGroup from mstar.engine.resources import NodeResourceSpec, ResourceReqConfig +from mstar.engine.resources.kv.config import RetentionPolicy from mstar.engine.resources.sampler.utils import SamplingConfig from mstar.graph.base import ( GraphEdge, @@ -63,13 +64,14 @@ class PrefixStream(NamedTuple): class Span(NamedTuple): - """One write to a keyed stream. A layout lists them in the order the walks write. + """One write to a KV stream. A layout lists them in the order the walks write. ``length`` is how many KV slots the write fills, and ``advance`` is how far it moves the position counter, which is the length for text. ``source`` is the tensor whose next entry an ids span keys, or the file of a digest span as ``(modality, index)``. ``params`` are the preprocessing arguments that - change the tensors made from the file. + change the tensors made from the file. A ``rows`` span writes KV made from + anything else, such as another node's outputs, and is never keyed. """ kind: str walk: str @@ -79,6 +81,20 @@ class Span(NamedTuple): params: tuple = () +class LabelLayout(NamedTuple): + """Everything a request writes into one paged KV label, from `process_prompt`'s ``kv_layout``. + + ``spans`` are the committed writes known at intake, in walk order. Each + ``(loop, slots)`` in ``growth`` writes ``slots`` per iteration of the named + loop. ``scratch`` is the most one non-committing step writes past the + committed length, and ``retention`` the policy the label's committing steps declare. + """ + spans: tuple[Span, ...] = () + growth: tuple[tuple[str, int], ...] = () + scratch: int = 0 + retention: RetentionPolicy | None = None + + def video_metadata_dict(metadata_obj) -> dict: # torchcodec>=0.9 reports pixel_aspect_ratio as a Fraction, which the wire codec rejects. return {k: float(v) if isinstance(v, Fraction) else v for k, v in asdict(metadata_obj).items()} diff --git a/mstar/worker/worker.py b/mstar/worker/worker.py index 8b86724ea7..86b6b8d60c 100644 --- a/mstar/worker/worker.py +++ b/mstar/worker/worker.py @@ -33,6 +33,7 @@ from mstar.engine import apply_torch_config from mstar.engine.engine import ExecutingBatch from mstar.engine.resources import AdmitFailedReason, AllocationFailed, StepContext +from mstar.engine.resources.kv.config import without_layouts from mstar.engine.resources.kv.transfer import ( TransferEngineInfo, make_deployment_kv_shm_dir, @@ -817,6 +818,8 @@ def _add_new_request(self, body: NewRequest) -> None: self._fail_requests({request_id: f"Error in worker: {type(err).__name__}: {err}"}) self._process_unprocessed(body.request_id) return + # the KV managers keep the layouts they ingested, and this info rides every frame the worker sends + body.request_info.resource_configs = without_layouts(body.request_info.resource_configs) # Create StreamBuffers for consumer connections on this worker. A request # with several partitions here arrives once per partition: keep the diff --git a/test/modular/test_kv_prefix_request_keys.py b/test/modular/test_kv_prefix_request_keys.py index 58660d6c38..dc2a32f09b 100644 --- a/test/modular/test_kv_prefix_request_keys.py +++ b/test/modular/test_kv_prefix_request_keys.py @@ -353,7 +353,7 @@ def _conductor_configs(model, model_kwargs: dict | None = None) -> dict: the model's own, a default for anything it declared a stream for, and then each handed its own chain.""" configs = Conductor._get_resource_configs( - SimpleNamespace(model=model), model_kwargs, {}, + SimpleNamespace(model=model), model_kwargs or {}, {}, ) for key, cfg in configs.items(): _apply_chain(cfg, key, model_kwargs) From fe10c6bef50f0dd8f739deb7af5b1dacae09140e Mon Sep 17 00:00:00 2001 From: Gaurav Shah Date: Wed, 21 Oct 2026 05:50:51 +0000 Subject: [PATCH 04/26] kv: PrefixStream becomes a mark with one span cursor --- mstar/api_server/data_worker.py | 131 +++++---------- mstar/conductor/conductor.py | 12 +- mstar/engine/engine.py | 20 +-- mstar/engine/resources/kv/config.py | 31 ++-- mstar/engine/resources/kv/manager.py | 171 ++++++++++++-------- mstar/model/bagel/bagel_model.py | 101 ++++++++---- mstar/model/bagel/submodules.py | 20 +++ mstar/model/base.py | 89 +++++++--- mstar/model/orpheus/orpheus_model.py | 17 +- mstar/worker/engine_manager.py | 37 ++--- test/modular/test_kv_prefix_bagel_parity.py | 6 +- test/modular/test_kv_prefix_batch.py | 2 +- test/modular/test_kv_prefix_chain_parent.py | 4 +- test/modular/test_kv_prefix_correctness.py | 26 +-- test/modular/test_kv_prefix_decode.py | 25 ++- test/modular/test_kv_prefix_eviction.py | 4 +- test/modular/test_kv_prefix_fill.py | 35 ++-- test/modular/test_kv_prefix_layout_lease.py | 32 ++-- test/modular/test_kv_prefix_lease.py | 41 ++--- test/modular/test_kv_prefix_logs.py | 23 +-- test/modular/test_kv_prefix_offload.py | 4 +- test/modular/test_kv_prefix_positions.py | 13 +- test/modular/test_kv_prefix_request_keys.py | 45 ++++-- test/modular/test_kv_prefix_root.py | 28 ++-- test/modular/test_kv_prefix_uncovered.py | 30 ++-- test/modular/test_media_decode_limits.py | 4 +- test/modular/test_multimodal_ordering.py | 30 ++-- test/modular/test_orpheus_empty_text.py | 2 +- test/modular/test_request_media_totals.py | 2 +- test/modular/test_video_frame_protocol.py | 2 +- test/modular/test_worker_metrics.py | 10 +- 31 files changed, 573 insertions(+), 424 deletions(-) diff --git a/mstar/api_server/data_worker.py b/mstar/api_server/data_worker.py index 46205eb931..e4fc718a19 100644 --- a/mstar/api_server/data_worker.py +++ b/mstar/api_server/data_worker.py @@ -31,7 +31,7 @@ from mstar.engine.resources.kv.config import KVSpec, PagedKVConfig from mstar.engine.resources.kv.keys import PageItem, chain, end_key, fingerprint, page_items from mstar.engine.resources.spec import apply_yaml_overrides -from mstar.model.base import LabelLayout, Model, PrefixStream, ProcessPromptOutput, Span +from mstar.model.base import Model, PrefixKeys, ProcessPromptOutput, Span, keyed_labels from mstar.model.media_limits import ( RequestMediaBudget, charge_decoded_audio, @@ -140,34 +140,6 @@ def _kv_page_sizes(model: Model, model_config: dict) -> dict[str, int]: } -def _layout_rows(kv_layout: dict[str, dict[str, LabelLayout]]) -> dict: - """``kv_layout`` as plain rows, which keep the request's kwargs on the wire's plain msgpack path. - - Each label becomes ``[spans, growth, scratch, retention]``, a span ``[length, advance, walk]``. - """ - rows: dict[str, dict[str, list]] = {} - for resource_key, by_label in kv_layout.items(): - for label, layout in by_label.items(): - for span in layout.spans: - if span.kind not in ("ids", "digest", "rows"): - raise ValueError(f"a span of kind {span.kind!r}") - if not (isinstance(span.length, int) and isinstance(span.advance, int)): - raise ValueError( - f"a {span.kind} span's length and advance are " - f"{type(span.length).__name__} and {type(span.advance).__name__}" - ) - retention = layout.retention - rows.setdefault(resource_key, {})[label] = [ - [[span.length, span.advance, span.walk] for span in layout.spans], - [[loop, slots] for loop, slots in layout.growth], - layout.scratch, - None if retention is None else [retention.context_budget, retention.protected_prefix], - ] - return rows - - - - class PreprocessWorker: def __init__( self, @@ -500,12 +472,12 @@ def __init__( timing_requested: threading.Event | None = None, worker_metrics_requested: threading.Event | 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 - declared = model.prefix_key_streams() if model is not None and model_config else {} + # keying a label needs the deployment's page size as well as the + # model's declaration, so a worker built without a config keys no label + declared = keyed_labels(model) 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 declared else {} - self._prefix_streams = { + self._keyed_labels = { key: by_label for key, by_label in declared.items() if key in self._prefix_page_sizes } self.in_queue = in_queue @@ -612,7 +584,7 @@ def _process_input( input_metadata[key] = [] for filepath in input.file_paths[modality]: - if self._prefix_streams: + if self._keyed_labels: # Record the file's state before loading it. If the file changes before it # is hashed, its KV would be keyed by bytes we never decoded. file_states.setdefault(modality, []).append(_file_state(filepath)) @@ -648,7 +620,7 @@ def _process_input( # another request's pages and be served that request's KV. Likewise # only process_prompt sets the model's INTERNAL_KWARGS. internal = getattr(self.model, "INTERNAL_KWARGS", ()) - for name in ("prefix_keys", "prefix_tail", "prefix_decode", "prefix_layout", "kv_layout", *internal): + for name in ("prefix_keys", "prefix_tail", "prefix_decode", "kv_layout", *internal): model_kwargs.pop(name, None) if self.model is not None: prompt_tensors = self.model.process_prompt( @@ -665,14 +637,18 @@ def _process_input( prompt_tensors = prompt_tensors.new_input_tensors if prompt_tensors: tensors.update(prompt_tensors) - layouts = model_kwargs.pop("prefix_layout", {}) - if "kv_layout" in model_kwargs: - model_kwargs["kv_layout"] = _layout_rows(model_kwargs["kv_layout"]) + layouts = model_kwargs.pop("kv_layout", {}) + rows = { + key: {label: layout.row() for label, layout in by_label.items()} + for key, by_label in layouts.items() + } + if rows: + model_kwargs["kv_layout"] = rows # after the update: the chain keys the tensors the request will # actually be prefilled with if model_kwargs.get("prefix_cache") is not False: model_kwargs.update(self._prefix_keys( - tensors, layouts, input.file_paths, file_states, + tensors, layouts, rows, input.file_paths, file_states, )) elif input.text is not None: # Fallback: encode as UTF-8 bytes -> uint8 tensor @@ -753,52 +729,32 @@ def _start_profile(input: PreprocessInput) -> DataWorkerProfile: ) def _prefix_keys( - self, tensors: dict, layouts: dict, file_paths: dict | None, + self, tensors: dict, layouts: dict, rows: dict, file_paths: dict | None, file_states: dict[str, list[tuple[int, int, int]]], ) -> dict: - """Key each page of every declared stream, by resource and label. + """Key each page of every label's keyed spans, by resource and label. Returns the kwargs that carry the keys, the prompt tail past the last whole - page, each stream's layout as ``[length, advance, walk, digest, end_key]`` - rows, and the output tensor each stream's sampled ids arrive in. The keys - are unrooted. + page, and the output tensor each label's sampled ids arrive in. The keys + are unrooted. Each keyed span's row in ``rows`` gains its digest and end key. """ - if not self._prefix_streams: + if not self._keyed_labels: return {} - for resource_key, by_label in layouts.items(): - for label in by_label: - if label not in self._prefix_streams.get(resource_key, {}): - raise ValueError( - f"{type(self.model).__name__} returned a prefix_layout for " - f"{resource_key!r}/{label!r}, which it does not declare" - ) keys: dict[str, dict[str, list[bytes]]] = {} tails: dict[str, dict[str, list[int]]] = {} - layout_rows: dict[str, dict[str, list[list]]] = {} decode: dict[str, dict[str, str]] = {} - for resource_key, by_label in self._prefix_streams.items(): + for resource_key, by_label in self._keyed_labels.items(): page_size = self._prefix_page_sizes[resource_key] - for label, stream in by_label.items(): + for label, kv_label in by_label.items(): layout = layouts.get(resource_key, {}).get(label) - if layout is None: - ids = tensors.get(stream.tensor) - if stream.keyed_by != "ids" or not ids: - continue - layout = [Span( - "ids", stream.walk, ids[0].numel(), ids[0].numel(), - stream.tensor, - )] - elif not stream.layout_walks: - raise ValueError( - f"{type(self.model).__name__} returned a prefix_layout " - f"for {resource_key!r}/{label!r}, which declares no " - "layout_walks" - ) + keyed = () if layout is None else layout.spans[:layout.keyed_count()] + if not keyed: + continue try: digests = [ self._digest(span, file_paths, file_states) if span.kind == "digest" else None - for span in layout + for span in keyed ] except _FileChangedError as changed: logger.warning( @@ -806,39 +762,37 @@ def _prefix_keys( ) return {"prefix_cache": False} pages, items = self._page_layout( - stream, layout, digests, tensors, page_size, + kv_label.keys, keyed, digests, tensors, page_size, ) page_keys = chain(pages, items) keys.setdefault(resource_key, {})[label] = page_keys - total = sum(span.length for span in layout) + total = sum(span.length for span in keyed) tails.setdefault(resource_key, {})[label] = ( pages[-1] if total % page_size else [] ) - rows, at = [], 0 - for span, digest in zip(layout, digests, strict=True): + at = 0 + span_rows = rows[resource_key][label][0][:len(keyed)] + for row, span, digest in zip(span_rows, keyed, digests, strict=True): at += span.length page, end = divmod(at, page_size) - rows.append([span.length, span.advance, span.walk, digest, end_key( + row += [digest, end_key( page_keys[page - 1] if page else b"", pages[page], items.get(page, ()), end, - ) if digest is not None and end else None]) - layout_rows.setdefault(resource_key, {})[label] = rows - if stream.decode_walk is not None: - decode.setdefault(resource_key, {})[label] = stream.tensor + ) if digest is not None and end else None] + if kv_label.keys.growth: + decode.setdefault(resource_key, {})[label] = kv_label.keys.tensor if not keys: return {} - kwargs = { - "prefix_keys": keys, "prefix_tail": tails, "prefix_layout": layout_rows, - } + kwargs = {"prefix_keys": keys, "prefix_tail": tails} if decode: kwargs["prefix_decode"] = decode return kwargs def _page_layout( - self, stream: PrefixStream, layout: list[Span], + self, prefix_keys: PrefixKeys, layout: tuple[Span, ...], digests: list[bytes | None], tensors: dict, page_size: int, ) -> tuple[list[list[int]], dict[int, list[PageItem]]]: - """Check the layout's spans and return each page's ids and items.""" - walks = {stream.walk, *stream.layout_walks} + """Check the keyed spans and return each page's ids and items.""" + walks = set(prefix_keys.walks) entries: dict[str, int] = {} total = sum(span.length for span in layout) pages: list[list[int]] = [[] for _ in range(-(-total // page_size))] @@ -849,12 +803,7 @@ def _page_layout( if span.walk not in walks: raise ValueError( f"a {span.kind} span names walk {span.walk!r}, which the " - f"stream does not declare; it declares {sorted(walks)}" - ) - if not (isinstance(span.length, int) and isinstance(span.advance, int)): - raise ValueError( - f"a {span.kind} span's length and advance are " - f"{type(span.length).__name__} and {type(span.advance).__name__}" + f"label's keys do not name; they name {sorted(walks)}" ) touched = range(at // page_size, -(-(at + span.length) // page_size)) if span.kind == "ids": diff --git a/mstar/conductor/conductor.py b/mstar/conductor/conductor.py index cf7b1fabbe..c17e835a19 100644 --- a/mstar/conductor/conductor.py +++ b/mstar/conductor/conductor.py @@ -597,16 +597,14 @@ def _get_resource_configs( 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. + A KV resource with declared labels gets a config even when the model + returned none, or its keys and layout would have nowhere to go. """ 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))) - for key in model_kwargs.get("kv_layout") or {}: - configs.setdefault(key, KVReqConfig()) + for key, labels in self.model.kv_labels().items(): + configs.setdefault(key, KVReqConfig(needed_labels=list(labels))) return configs def _derive_worker_info(self): @@ -1212,7 +1210,6 @@ def _do_ingest_request( prefix_keys = kwargs.get("prefix_keys") or {} prefix_tail = kwargs.get("prefix_tail") or {} prefix_decode = kwargs.get("prefix_decode") or {} - prefix_layout = kwargs.get("prefix_layout") or {} prefix_cache = kwargs.get("prefix_cache") kv_layout = kwargs.get("kv_layout") or {} for key, cfg in request_data.resource_configs.items(): @@ -1222,7 +1219,6 @@ def _do_ingest_request( prefix_keys=prefix_keys.get(key), prefix_tail=prefix_tail.get(key), prefix_decode=prefix_decode.get(key), - prefix_layout=prefix_layout.get(key), prefix_cache=prefix_cache, kv_layout=kv_layout.get(key), loop_caps=loop_caps, diff --git a/mstar/engine/engine.py b/mstar/engine/engine.py index 0ef8b9284e..f41ff60753 100644 --- a/mstar/engine/engine.py +++ b/mstar/engine/engine.py @@ -45,6 +45,7 @@ AdmitOutcome, ) from mstar.graph.runtime.base import GraphRuntime +from mstar.model.base import keyed_labels from mstar.model.submodule_base import ( EMPTY_INPUT_METADATA, ARNodeInputs, @@ -856,16 +857,15 @@ def _open_prefix_caches(self, specs_by_key, model) -> None: checkpoint_identity(checkpoint) if checkpoint is not None else b"", model.preprocess_fingerprint() if model is not None else "", ] - declared = model.prefix_key_streams() if model is not None else {} + declared = keyed_labels(model) if model is not None else {} self._prefix_model = type(model).__name__ - # node -> the walks a probe may run on: the keyed walk and the layout walks - # of every stream declared on a cache the node uses + # node -> the walks a probe may run on. Each keyed label names them in its + # keys, for every cache the node uses. self._keyed_walks: dict[str, set[str]] = {} for key, by_label in declared.items(): for node in specs_by_key[key].nodes: self._keyed_walks.setdefault(node, set()).update( - walk for stream in by_label.values() - for walk in (stream.walk, *stream.layout_walks) + walk for label in by_label.values() for walk in label.keys.walks ) # a cache another engine holds is taken as open when its config says so opened_keys = { @@ -893,8 +893,8 @@ def _open_prefix_caches(self, specs_by_key, model) -> None: if mark is not None: parts.append(mark) opened = resource.enable_prefix_cache(fingerprint(*shared, *parts), { - label: (stream.walk, stream.decode_walk, *stream.layout_walks) - for label, stream in declared.get(key, {}).items() + label: (*kv_label.keys.walks, *(walk for walk, _ in kv_label.growth if kv_label.keys.growth)) + for label, kv_label in declared.get(key, {}).items() }) if not opened or not declared.get(key): continue @@ -914,12 +914,12 @@ def _open_prefix_caches(self, specs_by_key, model) -> None: "prompt frees little", key, ) - # the nodes of a layout walk that don't use the stream's cache, such as an image + # the nodes of a keyed walk that don't use the label's cache, such as an image # encoder. The KV computed from their output is cached under their input's digest. self._item_encoders: set[str] = set() for key in opened_keys: - for stream in declared[key].values(): - for walk in stream.layout_walks: + for label in declared[key].values(): + for walk in label.keys.walks: nodes = model.get_graph_walk_graphs()[walk].get_nodes() self._item_encoders.update(set(nodes) - specs_by_key[key].nodes) diff --git a/mstar/engine/resources/kv/config.py b/mstar/engine/resources/kv/config.py index 52c523d693..ea530f70c1 100644 --- a/mstar/engine/resources/kv/config.py +++ b/mstar/engine/resources/kv/config.py @@ -245,8 +245,6 @@ class KVReqConfig(ResourceReqConfig): prefix_tail: dict[str, list[int]] | None = None # label -> the output tensor its sampled ids arrive in, if the stream keys generation prefix_decode: dict[str, str] | None = None - # label -> the layout its keys were chained over, one span per write - prefix_layout: dict[str, list[PrefixSpan]] | None = None prefix_cache: bool = True # label -> everything the request writes there kv_layout: dict[str, "StreamLayout"] | None = None @@ -258,7 +256,6 @@ def apply_conductor_config( prefix_keys: dict[str, list[bytes]] | None=None, prefix_tail: dict[str, list[int]] | None=None, prefix_decode: dict[str, str] | None=None, - prefix_layout: dict[str, list[list]] | None=None, prefix_cache: bool | None=None, kv_layout: dict[str, list] | None=None, loop_caps: Mapping[str, int] | None=None, @@ -270,23 +267,22 @@ def apply_conductor_config( self.prefix_tail = prefix_tail if prefix_decode is not None: self.prefix_decode = prefix_decode - if prefix_layout is not None: - # sent as plain lists, which keeps the request's kwargs on the wire's plain msgpack path - self.prefix_layout = { - label: [PrefixSpan(*row) for row in rows] - for label, rows in prefix_layout.items() - } if prefix_cache is not None: self.prefix_cache = prefix_cache if kv_layout is not None: + # a model lays out the labels a request may open on any deployment, and the + # request is held only to the ones it opens + opened = self.opened_labels() self.kv_layout = { label: StreamLayout( spans=tuple(PrefixSpan(*row) for row in spans), growth=tuple((loop, slots) for loop, slots in growth), scratch=scratch, retention=None if retention is None else RetentionPolicy(*retention), + keyed_spans=keyed_spans, ) - for label, (spans, growth, scratch, retention) in kv_layout.items() + for label, (spans, growth, scratch, retention, keyed_spans) in kv_layout.items() + if label in opened } self.loop_caps = { loop: loop_caps[loop] @@ -302,6 +298,13 @@ def get_labels(self, node: str, walk: str): return self.needed_labels return ["main"] + def opened_labels(self) -> set[str]: + """Every label `get_labels` can give, on any node and walk.""" + labels = set(self.needed_labels or ["main"]) + for named in (*self.needed_labels_per_node.values(), *self.needed_labels_per_node_walk.values()): + labels.update(named) + return labels + def get_publish_labels( self, node: str | None, @@ -382,16 +385,12 @@ def __post_init__(self): @dataclass(frozen=True) class StreamLayout: - """Everything a request writes into one stream, known when its prompt is preprocessed. - - ``spans`` are the committed writes, in walk order. Each ``(loop, slots)`` in - ``growth`` writes ``slots`` per iteration of ``loop``. ``scratch`` is the most one - non-committing step writes past the committed length. - """ + """`LabelLayout` as the KV manager holds it. The stream's keys cover its first ``keyed_spans`` spans.""" spans: tuple[PrefixSpan, ...] = () growth: tuple[tuple[str, int], ...] = () scratch: int = 0 retention: RetentionPolicy | None = None + keyed_spans: int = 0 def pages(self, caps: Mapping[str, int], page_size: int) -> int: """The most pages the stream can hold, where ``caps`` gives each growth loop's iterations. diff --git a/mstar/engine/resources/kv/manager.py b/mstar/engine/resources/kv/manager.py index a1ad697bf9..bedc3a0002 100644 --- a/mstar/engine/resources/kv/manager.py +++ b/mstar/engine/resources/kv/manager.py @@ -27,6 +27,7 @@ PagedKVConfig, PrefixSpan, RetentionPolicy, + StreamLayout, ) from mstar.engine.resources.kv.cpu_page_pool import CPUPagePool from mstar.engine.resources.kv.keys import PageItem, fingerprint, page_items, page_key @@ -128,17 +129,10 @@ class PrefixChain: keyed_pages: int # tokens the chain accounts for: its prompt, and every token sampled since covered_len: int - # The prompt's layout, and how far the stream is through it. - # `spans` has one span per write. + # `spans` are the layout's keyed spans, one per write. # `unkeyed_items` are the items on the page `unkeyed` fills. - # `done` counts the spans the stream holds, whether written, read in or served from the cache. - # `consumed` counts the tokens of those spans. - # `probed` is the span the latest probe answered for. spans: list[PrefixSpan] unkeyed_items: list[PageItem] - done: int = 0 - consumed: int = 0 - probed: int | None = None # The prompt's keys hashed under the root, kept from the first probe so a repeated probe hashes nothing. # `rooted` has one key per page. # `rooted_ends` has one key per image that ends inside a page, with the token the image ends at. @@ -187,6 +181,35 @@ def pages_filled(self, stored_len: int, page_size: int) -> int: return min(stored_len // page_size, self.keyed_pages) +@dataclass +class SpanCursor: + """How far a stream is through its request's layout.""" + spans: tuple[PrefixSpan, ...] + # `done` counts the spans the stream holds, whether written, read in or served from the cache. + # `consumed` counts the tokens of those spans. + # `probed` is the span the latest probe answered for. + done: int = 0 + consumed: int = 0 + probed: int | None = None + + def advance(self, start: int, end: int) -> list[tuple[int, int]]: + """Move past the spans a write over ``[start, end)`` completes. + + Returns ``(index, end)`` for each span the write touched. A span before + ``start`` that a probe served or a retrieve read in is passed, not touched. + """ + touched = [] + while self.done < len(self.spans) and self.consumed < end: + span_end = self.consumed + self.spans[self.done].length + if span_end > start: + touched.append((self.done, span_end)) + if end < span_end: + break + self.consumed = span_end + self.done += 1 + return touched + + @dataclass class CacheStream: """(request, label) cache stream metadata""" @@ -234,8 +257,10 @@ class CacheStream: # its prefix lookup is in `KVManager.prefix_tokens`, so a re-probe after a # miss (a refused admit prepares the step again) is not counted twice prefix_counted: bool = False - # the most pages the request's layout lets this stream hold; None without a layout + # the most pages the request's layout lets this stream hold, and how far through + # that layout it is; both None without a layout footprint: int | None = None + span_cursor: SpanCursor | None = None def reset(self, freed: bool=False, *, content_reset: bool=True): self.stored_len = 0 @@ -253,12 +278,14 @@ def reset(self, freed: bool=False, *, content_reset: bool=True): self.page_indices.clear() def forget_chain(self): - """Drop what the chain knows of this stream, for a run that starts over. + """Drop what the chain knows of this stream, and its place in the layout, for a run that starts over. Not part of `reset`, which an offload calls too: a reload brings the same tokens back, so the chain that describes them has to survive it. """ self.chain = None + if self.span_cursor is not None: + self.span_cursor = SpanCursor(self.span_cursor.spans) class KVPageCounts(NamedTuple): @@ -626,21 +653,21 @@ def resolve_cached_prefix( self._warn_unkeyed(rid, node_name, graph_walk) return None stream = self._ensure_label(rid, label) - chain = stream.chain - if chain is None or chain.done == len(chain.spans): + chain, cursor = stream.chain, stream.span_cursor + if chain is None or cursor.done >= len(chain.spans): return None - span = chain.spans[chain.done] + span = chain.spans[cursor.done] self._check_positions(span, graph_walk, inputs) if span.walk != graph_walk or span.length != inputs.input_seq_len: self._drop_layout(rid, label, stream, ( f"{graph_walk} writes {inputs.input_seq_len} slots where its " - f"layout's span {chain.done} has {span.walk} write {span.length}" + f"layout's span {cursor.done} has {span.walk} write {span.length}" )) return None if stream.converted: - return self._answer(label, chain, stream.stored_len) + return self._answer(label, cursor, stream.stored_len) if stream.lease is not None: - return self._answer(label, chain, stream.lease_len) + return self._answer(label, cursor, stream.lease_len) if stream.stored_len or stream.offloaded or stream.read_pending: return None keys, rooted = list(chain.keys), chain.rooted @@ -688,7 +715,7 @@ def resolve_cached_prefix( if not matched: return None stream.lease, stream.lease_len = matched, held - return self._answer(label, chain, stream.lease_len) + return self._answer(label, cursor, stream.lease_len) def _check_positions(self, span: PrefixSpan, graph_walk: str, inputs) -> None: assert span.digest is not None or inputs.custom_pos_ids is None, ( @@ -697,18 +724,18 @@ def _check_positions(self, span: PrefixSpan, graph_walk: str, inputs) -> None: ) @staticmethod - def _answer(label: str, chain: PrefixChain, present: int) -> CachedPrefix: + def _answer(label: str, cursor: SpanCursor, present: int) -> CachedPrefix: """This walk's share of the ``present`` tokens the stream holds, with the position past that share. Records the span as the one the probe answered for.""" - span = chain.spans[chain.done] - tokens = min(max(present - chain.consumed, 0), span.length) - position = sum(earlier.advance for earlier in chain.spans[:chain.done]) + span = cursor.spans[cursor.done] + tokens = min(max(present - cursor.consumed, 0), span.length) + position = sum(earlier.advance for earlier in cursor.spans[:cursor.done]) if span.digest is None: position += tokens elif tokens == span.length: # A partly held item doesn't move the position. Its walk still runs and advances it. position += span.advance - chain.probed = chain.done + cursor.probed = cursor.done return CachedPrefix(label, tokens, position) def _drop_layout( @@ -725,13 +752,13 @@ def _drop_layout( logger.warning( "KV %s: %s/%s %s; its chain ends here", self.name, rid, label, disagreement, ) - chain = stream.chain - if chain.done and stream.lease is not None: - keep, partial = divmod(chain.consumed, self.config.page_size) + cursor = stream.span_cursor + if cursor.done and stream.lease is not None: + keep, partial = divmod(cursor.consumed, self.config.page_size) if partial: raise RuntimeError( f"KV {self.name}: {rid}/{label}: the walks the cache served end " - f"inside a page ({chain.consumed} slots), so their KV cannot be kept" + f"inside a page ({cursor.consumed} slots), so their KV cannot be kept" ) self._adopt(stream, stream.lease[:keep]) self._arena.release(stream.lease[keep:]) @@ -756,19 +783,20 @@ def apply_cached_prefix( return if prefix is None: # an unprobed walk never reached the position check in `resolve_cached_prefix` - chain = self._ensure_label(rid, label).chain - if chain is not None and chain.done < len(chain.spans): - self._check_positions(chain.spans[chain.done], graph_walk, inputs) + stream = self._ensure_label(rid, label) + chain, cursor = stream.chain, stream.span_cursor + if chain is not None and cursor.done < len(chain.spans): + self._check_positions(chain.spans[cursor.done], graph_walk, inputs) stream = self._streams.get(rid, {}).get(label) if stream is None or stream.lease is None or stream.stored_len: return if prefix is None: self._drop_layout(rid, label, stream, f"{graph_walk} writes unprobed") return - chain = stream.chain - if prefix.tokens == chain.spans[chain.done].length: + cursor = stream.span_cursor + if prefix.tokens == cursor.spans[cursor.done].length: return - held = chain.consumed + prefix.tokens + held = cursor.consumed + prefix.tokens if held < stream.lease_len: keep = held // self.config.page_size self._arena.release(stream.lease[keep:]) @@ -782,14 +810,16 @@ def complete_cached_walk(self, rid: str, node_name: str, graph_walk: str) -> Non stream = self._streams.get(rid, {}).get(label) if stream is None or stream.chain is None: return - chain = stream.chain - chain.consumed += chain.spans[chain.done].length - chain.done += 1 + cursor = stream.span_cursor + cursor.consumed += cursor.spans[cursor.done].length + cursor.done += 1 def _index_filled_pages( self, segment: Segment, stream: CacheStream, ctx: StepContext, + touched: list[tuple[int, int]], ) -> None: - """Offer every page this commit filled to the index. + """Offer every page this commit filled to the index. ``touched`` holds + ``(index, end)`` for each span of the layout this commit wrote into. Only whole pages, because the tail is still being written into and a second owner would be reading bytes that are still moving. A dummy or @@ -819,22 +849,16 @@ def _index_filled_pages( if stream.released or stream.stored_len - segment.span > chain.covered_len + lag: stream.chain = None return - start = stream.stored_len - segment.span ended = None - while chain.done < len(chain.spans): - end = chain.consumed + chain.spans[chain.done].length - # Checking every span would end the chain at the first one the cache filled rather than this walk. - if chain.consumed < stream.stored_len and end > start and ( - chain.spans[chain.done].walk != walk - ): + for index, end in touched: + if index >= len(chain.spans): + break + span = chain.spans[index] + if span.walk != walk: stream.chain = None return - if stream.stored_len < end: - break - if end > start and chain.spans[chain.done].end_key is not None: - ended = end, chain.spans[chain.done].end_key - chain.consumed = end - chain.done += 1 + if end <= stream.stored_len and span.end_key is not None: + ended = end, span.end_key filled = chain.pages_filled(stream.stored_len, self.config.page_size) if stream.retention is not None: filled = min( @@ -942,14 +966,18 @@ def _warn_unkeyed(self, rid: str, node_name: str, graph_walk: str) -> None: ): return keys = overrides.prefix_keys or {} - layout = overrides.prefix_layout or {} + layouts = overrides.kv_layout or {} for label in overrides.get_labels(node_name, graph_walk): - if graph_walk in self._keyed_walks.get(label, ()) and not (keys.get(label) and layout.get(label)): + layout = layouts.get(label) + # a layout that keys no span leaves its request unkeyed on purpose + if graph_walk in self._keyed_walks.get(label, ()) and not keys.get(label) and ( + layout is None or layout.keyed_spans + ): self._warned_unkeyed.add(node_name) logger.warning( "KV %s: requests reach %s with no %s for %s, which the " "model declared for prefix reuse, so nothing will be cached", - self.name, node_name, "layout" if keys.get(label) else "keys", label, + self.name, node_name, "layout" if layout is None else "keys", label, ) return @@ -967,9 +995,9 @@ def _keyed_label( if overrides is None or not overrides.prefix_cache: return None keys = overrides.prefix_keys or {} - layout = overrides.prefix_layout or {} + layouts = overrides.kv_layout or {} for label in overrides.get_labels(node_name, graph_walk): - if keys.get(label) and layout.get(label): + if keys.get(label) and label in layouts: return label return None @@ -979,11 +1007,14 @@ def _seed_keys(self, rid: str, label: str, stream: CacheStream) -> None: if overrides is None or not overrides.prefix_cache: return keys = (overrides.prefix_keys or {}).get(label) - spans = (overrides.prefix_layout or {}).get(label) - if not keys or not spans: + # the chain reads the span cursor, which only a label the request opens has + layout = self._held_layout(rid, label) + if not keys or layout is None: return tail = list((overrides.prefix_tail or {}).get(label) or ()) - stream.chain = PrefixChain.seed(keys, tail, spans, self.config.page_size) + stream.chain = PrefixChain.seed( + keys, tail, layout.spans[:layout.keyed_spans], self.config.page_size, + ) @property def keys_prefix_chains(self) -> bool: @@ -1268,14 +1299,21 @@ def ingest_request(self, rid, overrides: KVReqConfig | None=None): self._streams.setdefault(rid, {"main": CacheStream()}) self._overrides.setdefault(rid, overrides) for label, stream in self._streams[rid].items(): - stream.footprint = self._footprint(rid, label) + if stream.span_cursor is None: + self._lay_out(rid, label, stream) if stream.chain is None: self._seed_keys(rid, label, stream) - def _footprint(self, rid: str, label: str) -> int | None: + def _lay_out(self, rid: str, label: str, stream: CacheStream) -> None: + """Set ``stream``'s footprint and span cursor from its request's layout for ``label``, if there is one.""" + layout = self._held_layout(rid, label) + if layout is not None: + stream.footprint = layout.pages(self._overrides[rid].loop_caps, self.config.page_size) + stream.span_cursor = SpanCursor(layout.spans) + + def _held_layout(self, rid: str, label: str) -> StreamLayout | None: overrides = self._overrides.get(rid) - layout = (overrides.kv_layout or {}).get(label) if overrides is not None else None - return None if layout is None else layout.pages(overrides.loop_caps, self.config.page_size) + return (overrides.kv_layout or {}).get(label) if overrides is not None else None def admit_retrieve( self, rid: str, @@ -1435,7 +1473,7 @@ def admit(self, step: KVStep, ctx: StepContext) -> AdmitOutcome: and not stream.offloaded and not stream.read_pending ): - assert stream.chain.probed == stream.chain.done, ( + assert stream.span_cursor.probed == stream.span_cursor.done, ( f"KV {self.name}: {segment.request_id}/{segment.label} admits a " "walk its lease was never probed for, which would write over " "the pages the lease holds for it" @@ -2253,8 +2291,12 @@ def commit(self, step: KVStep, ctx: StepContext): ) if policy is not None: self._adopt_retention(stream, policy, segment) + touched = [] + if stream.chain is not None: + written = stream.stored_len + stream.released + touched = stream.span_cursor.advance(written - segment.span, written) if index_on and stream.chain is not None: - self._index_filled_pages(segment, stream, ctx) + self._index_filled_pages(segment, stream, ctx, touched) # so a claim taken in a window the mark misses still fails # `_commit_offload`'s generation guard stream.generation += 1 @@ -2980,7 +3022,8 @@ def cleanup(self): def _ensure_label(self, rid: str, label: str) -> CacheStream: if label not in self._streams[rid]: - stream = self._streams[rid][label] = CacheStream(footprint=self._footprint(rid, label)) + stream = self._streams[rid][label] = CacheStream() + self._lay_out(rid, label, stream) self._seed_keys(rid, label, stream) return self._streams[rid][label] diff --git a/mstar/model/bagel/bagel_model.py b/mstar/model/bagel/bagel_model.py index c533582451..1d8312d4fb 100644 --- a/mstar/model/bagel/bagel_model.py +++ b/mstar/model/bagel/bagel_model.py @@ -81,10 +81,20 @@ VAEDecoderSubmodule, VAEEncoderSubmodule, ViTEncoderSubmodule, + vae_latents, vit_patches, vit_preprocess, ) -from mstar.model.base import DECODE, ForwardPassArgs, Model, PrefixStream, ProcessPromptOutput, Span +from mstar.model.base import ( + DECODE, + ForwardPassArgs, + KVLabel, + LabelLayout, + Model, + PrefixKeys, + ProcessPromptOutput, + Span, +) from mstar.model.loader import iter_safetensors_file, load_hf_weights from mstar.model.loader.base import LLAMA_STACKED_PARAMS, StackedParamRule from mstar.model.multimodal import ( @@ -615,11 +625,9 @@ def process_prompt( attachments gets its own span and no seam is re-tokenized. The schedule builder walks the same :func:`prefill_plan`. - An understanding prompt with images, given their tensors, also returns - the main stream's ``prefix_layout``, so it is keyed past its images. + Given the tensors, it also returns the main stream's ``kv_layout``. """ result: NameToTensorList = {} - layout = None if prompt is not None: target_output = output_modalities[0] if output_modalities else "text" @@ -670,13 +678,6 @@ def process_prompt( spans, len(segments), ) - # An edit gets no layout. Its image is resized below before the ViT - # sees it, and the file's digest doesn't cover that resize. - if is_understanding and "image" not in output_modalities and tensors is not None: - layout = self._prefix_layout( - parts, segments, tensors["image_inputs"], - vit_preprocess(kwargs.get("image_preprocess")), - ) if think_mode and not is_understanding: # Not part of the request layout: tokenized on its own and # prefilled ahead of the plan. @@ -704,31 +705,62 @@ def process_prompt( for img in tensors["image_inputs"] ] - if layout is None: + if prompt is None or tensors is None: return result - return ProcessPromptOutput(result, {"prefix_layout": {"kv": {"main": layout}}}) + images = result.get("image_inputs") or tensors.get("image_inputs", []) + layout = self._main_layout( + parts, segments, images, output_modalities, think_mode, + vit_preprocess(kwargs.get("image_preprocess")), + kwargs.get("height", 1024), kwargs.get("width", 1024), + ) + return ProcessPromptOutput(result, {"kv_layout": {"kv": {"main": layout}}}) + + def _main_layout( + self, parts: list[PromptPart], segments: list[torch.Tensor], images: list[torch.Tensor], + output_modalities: list[str], think_mode: bool, image_preprocess: str, height: int, width: int, + ) -> LabelLayout: + """The main stream's writes, in the order `_build_prefill_schedule` runs them. - def _prefix_layout( - self, parts: list[PromptPart], segments: list[torch.Tensor], - images: list[torch.Tensor], image_preprocess: str, - ) -> list[Span]: - """The main stream's writes in walk order: text by its ids, an image by its file.""" + A generation keys only its first span, if that is text: its images are resized + before their walks see them, and a file's digest doesn't cover the resize. + """ + cfg = self.config + understanding = (output_modalities or ["text"])[0] == "text" spans = [] - for step in prefill_plan(parts, leading_text=True): + if think_mode and not understanding: + spans.append(Span("ids", "prefill_text", segments[0].numel(), segments[0].numel(), "text_inputs")) + texts = segments[len(spans):] + for step in prefill_plan(parts, leading_text=understanding): if step.modality == TEXT: - length = segments[step.index].numel() + length = texts[step.index].numel() spans.append(Span("ids", "prefill_text", length, length, "text_inputs")) continue - _, height, width = images[step.index].shape + _, image_height, image_width = images[step.index].shape + item = (("image", step.index), (image_preprocess,)) + if not understanding: + latents = vae_latents( + image_height, image_width, image_preprocess, cfg.latent_downsample, cfg.max_latent_size, + ) + spans.append(Span("digest", "prefill_vae", latents + IMAGE_SENTINELS, IMAGE_ADVANCE, *item)) patches = vit_patches( - height, width, image_preprocess, - self.config.vit_config.patch_size, self.config.vit_max_num_patch_per_side, + image_height, image_width, image_preprocess, + cfg.vit_config.patch_size, cfg.vit_max_num_patch_per_side, ) - spans.append(Span( - "digest", "prefill_vit", patches + IMAGE_SENTINELS, IMAGE_ADVANCE, - ("image", step.index), (image_preprocess,), - )) - return spans + spans.append(Span("digest", "prefill_vit", patches + IMAGE_SENTINELS, IMAGE_ADVANCE, *item)) + keyed = int(spans[0].kind == "ids") + if understanding: + return LabelLayout( + spans=tuple(spans), growth=(("decode_loop", 1),), + keyed_spans=None if "image" not in output_modalities else keyed, + ) + # an edit generates at its input's size (see `get_initial_forward_pass_args`) + if images: + _, height, width = images[0].shape + latents = (height // cfg.latent_downsample) * (width // cfg.latent_downsample) + return LabelLayout( + spans=tuple(spans), growth=(("decode_loop", 1),) if think_mode else (), + scratch=latents + IMAGE_SENTINELS, keyed_spans=keyed, + ) @staticmethod def _resize_to_fit( @@ -797,16 +829,17 @@ def _kv_config(self) -> PagedKVConfig: num_qo_heads=self.config.num_attention_heads, ) - def prefix_key_streams(self) -> dict[str, dict[str, PrefixStream]]: + def kv_labels(self) -> dict[str, dict[str, KVLabel]]: """The text walk's prompt is its token ids, and so is every step it takes. An image in an understanding prompt is keyed by its file. An edit's VAE - walk writes this stream too, over positions no layout describes, so it - is not named. A chat sends a reply back as its own assistant turn, so - its pages match. + walk writes ``main`` too, but the keys of ``main`` leave that walk out (see + `_kv_layout`). A chat sends a reply back as + its own assistant turn, so its pages match. """ - return {"kv": {"main": PrefixStream( - "text_inputs", "ids", "prefill_text", "decode", ("prefill_vit",), + return {"kv": {"main": KVLabel( + growth=((DECODE, "decode_loop"),), + keys=PrefixKeys("text_inputs", ("prefill_text", "prefill_vit"), growth=True), )}} def checkpoint_path(self) -> str: diff --git a/mstar/model/bagel/submodules.py b/mstar/model/bagel/submodules.py index d1cafd2095..eaa753f885 100644 --- a/mstar/model/bagel/submodules.py +++ b/mstar/model/bagel/submodules.py @@ -125,6 +125,22 @@ def vit_patches( return (height // patch_size) * (width // patch_size) +def vae_latents( + height: int, width: int, image_preprocess: str, + latent_downsample: int, max_latent_size: int, +) -> int: + """How many latents the VAE encodes an image of this size into. Its ``prepare_inputs`` asserts the same count.""" + if image_preprocess == "vllm": + # the size `vllm_vae_resize` resizes to + stride, max_size = latent_downsample, max_latent_size * latent_downsample + scale = max(min(max_size / max(width, height), 1.0), min(256, max_size) / min(width, height)) + height = min(max(stride, round(height * scale / stride) * stride), max_size) + width = min(max(stride, round(width * scale / stride) * stride), max_size) + else: + height, width = _VAE_TRANSFORM.resize_transform.target_size(height, width) + return (height // latent_downsample) * (width // latent_downsample) + + class ViTEncoderSubmodule(NodeSubmodule): """SigLIP2 ViT + connector + vit_pos_embed: pixel patches -> ViT features. @@ -543,6 +559,10 @@ def prepare_inputs( _, img_h, img_w = image_tensor.shape h = (img_h // ds) w = (img_w // ds) + assert h * w == vae_latents(*img.shape[-2:], image_preprocess, ds, self.max_latent_size), ( + f"the VAE encoded a {tuple(img.shape[-2:])} image into {h * w} latents, " + "not the count `vae_latents` gives a layout to size it by" + ) packed_vae_position_ids = get_flattened_position_ids_extrapolate( img_h, img_w, diff --git a/mstar/model/base.py b/mstar/model/base.py index dbbabf37b1..c95de6d258 100644 --- a/mstar/model/base.py +++ b/mstar/model/base.py @@ -51,18 +51,6 @@ class ProcessPromptOutput(NamedTuple): metadata: dict -class PrefixStream(NamedTuple): - """How one cache stream's pages are keyed: by one input tensor's token ids, or by a layout over several walks.""" - tensor: str - keyed_by: str - # the walk that writes the keyed span; a write from a walk not named here ends the chain - walk: str - # the walk whose input is the last sampled token, to key generated pages; None keys the prompt only - decode_walk: str | None = None - # the other walks a layout may name - layout_walks: tuple[str, ...] = () - - class Span(NamedTuple): """One write to a KV stream. A layout lists them in the order the walks write. @@ -93,6 +81,67 @@ class LabelLayout(NamedTuple): growth: tuple[tuple[str, int], ...] = () scratch: int = 0 retention: RetentionPolicy | None = None + # the label's keys cover its first `keyed_spans` spans; None covers every span before the first rows span + keyed_spans: int | None = None + + def keyed_count(self) -> int: + """How many leading spans the label's keys cover.""" + if self.keyed_spans is not None: + return self.keyed_spans + return next((i for i, span in enumerate(self.spans) if span.kind == "rows"), len(self.spans)) + + def row(self) -> list: + """The layout as ``[spans, growth, scratch, retention, keyed spans]``, a span as + ``[length, advance, walk]``: plain values, which keep a request's kwargs on the + wire's msgpack path.""" + for span in self.spans: + if span.kind not in ("ids", "digest", "rows"): + raise ValueError(f"a span of kind {span.kind!r}") + if not (isinstance(span.length, int) and isinstance(span.advance, int)): + raise ValueError( + f"a {span.kind} span's length and advance are " + f"{type(span.length).__name__} and {type(span.advance).__name__}" + ) + retention = self.retention + return [ + [[span.length, span.advance, span.walk] for span in self.spans], + [[loop, slots] for loop, slots in self.growth], + self.scratch, + None if retention is None else [retention.context_budget, retention.protected_prefix], + self.keyed_count(), + ] + + +class PrefixKeys(NamedTuple): + """How a label's pages are keyed for cross-request prefix reuse. + + ``tensor`` holds the ids an ids span is keyed by. Only a write from one of + ``walks`` keeps the keys going. With ``growth``, the ids the label's growth + walk samples extend the keys. + """ + tensor: str + walks: tuple[str, ...] + growth: bool = False + + +class KVLabel(NamedTuple): + """A paged KV label, as the model's walks write it. + + Each ``(walk, loop)`` in ``growth`` writes after admission, bounded by the + named loop. A label with ``keys`` is keyed for prefix reuse. + """ + growth: tuple[tuple[str, str], ...] = () + keys: PrefixKeys | None = None + + +def keyed_labels(model: "Model") -> dict[str, dict[str, KVLabel]]: + """The labels of ``model``'s `kv_labels` that have keys, by resource, leaving out a resource with none.""" + keyed = {} + for key, labels in model.kv_labels().items(): + marked = {label: kv_label for label, kv_label in labels.items() if kv_label.keys is not None} + if marked: + keyed[key] = marked + return keyed def video_metadata_dict(metadata_obj) -> dict: @@ -674,16 +723,14 @@ def combined_walk_of(self) -> dict[tuple[str, str], str]: out[(node, walk)] = name return out - def prefix_key_streams(self) -> dict[str, dict[str, PrefixStream]]: - """Which input tensor keys which ``(resource, label)`` cache stream. - - Declaring one opts a node into cross-request prefix reuse: the - preprocess worker keys that tensor and the KV resource matches what it - already holds. Empty leaves the model uncached whatever the deployment - asks for. + def kv_labels(self) -> dict[str, dict[str, KVLabel]]: + """The paged KV labels the model's walks write, by resource and label. - A stream with ``layout_walks`` takes its ``prefix_layout`` from - `process_prompt`'s metadata, as ``{resource: {label: [Span, ...]}}``. + `process_prompt` returns a `LabelLayout` for each under ``kv_layout``, as + ``{resource: {label: LabelLayout}}``. A label with keys opts its node into + cross-request prefix reuse: the preprocess worker keys its layout's keyed + spans and the KV resource matches what it already holds. With no keyed + label the model is uncached whatever the deployment asks for. """ return {} diff --git a/mstar/model/orpheus/orpheus_model.py b/mstar/model/orpheus/orpheus_model.py index 2bf050293d..f3b32da74f 100644 --- a/mstar/model/orpheus/orpheus_model.py +++ b/mstar/model/orpheus/orpheus_model.py @@ -42,7 +42,7 @@ ) from mstar.graph.base import GraphEdge, GraphNode, GraphSection, Loop, TensorPointerInfo from mstar.graph.special_destinations import EMIT_TO_CLIENT, EMPTY_DESTINATION -from mstar.model.base import ForwardPassArgs, Model, PrefixStream +from mstar.model.base import ForwardPassArgs, KVLabel, LabelLayout, Model, PrefixKeys, ProcessPromptOutput, Span from mstar.model.orpheus.config import ATTN, KV_CACHE, ROPE, SAMPLER, OrpheusModelConfig from mstar.model.submodule_base import NodeSubmodule from mstar.streaming.topology import Connection, PartitionTopology, StreamingGraphEdge @@ -289,7 +289,7 @@ def process_prompt( output_modalities: list[str], tensors: NameToTensorList | None = None, **kwargs, - ) -> NameToTensorList: + ) -> ProcessPromptOutput: # Orpheus is text-only; raw multimodal tensors are unused. # Refused here, where it is a 400 before admission: the LLM's prefill # has no input without text, and blank text would generate from the voice @@ -327,7 +327,11 @@ def process_prompt( f"{self.get_max_output_tokens()} audio tokens it can generate)" ) - return {"text_inputs": [all_input_ids]} + length = all_input_ids.numel() + layout = LabelLayout( + spans=(Span("ids", "prefill", length, length, "text_inputs"),), growth=(("decode_loop", 1),), + ) + return ProcessPromptOutput({"text_inputs": [all_input_ids]}, {"kv_layout": {KV_CACHE: {"main": layout}}}) # ------------------------------------------------------------------- # Model ABC: forward pass args @@ -383,10 +387,13 @@ def get_initial_forward_pass_args( # Model ABC: resources # ------------------------------------------------------------------- - def prefix_key_streams(self) -> dict[str, dict[str, PrefixStream]]: + def kv_labels(self) -> dict[str, dict[str, KVLabel]]: """The LLM's prompt is its token ids. Its generated audio codes never come back in a prompt, so they are not keyed.""" - return {KV_CACHE: {"main": PrefixStream("text_inputs", "ids", "prefill")}} + return {KV_CACHE: {"main": KVLabel( + growth=(("decode", "decode_loop"),), + keys=PrefixKeys("text_inputs", ("prefill",)), + )}} def checkpoint_path(self) -> str: return _resolve_local_hf_snapshot( diff --git a/mstar/worker/engine_manager.py b/mstar/worker/engine_manager.py index c4b7676ba4..6132d62916 100644 --- a/mstar/worker/engine_manager.py +++ b/mstar/worker/engine_manager.py @@ -15,7 +15,7 @@ from mstar.engine.resources.position.config import PositionSpec, PosScheme from mstar.graph.runtime.base import GraphRuntime from mstar.graph.special_destinations import SPECIAL_DESTINATIONS -from mstar.model.base import Model +from mstar.model.base import Model, keyed_labels from mstar.utils.device import DeviceCapability from mstar.utils.streams import reset_device_scheduling @@ -30,7 +30,7 @@ def _refuse_uncacheable_positions( A cached page is reusable only at the positions it was written at, since its K is stored already rotated. """ - declared = set(model.prefix_key_streams()) + declared = set(keyed_labels(model)) if not declared: return for spec in specs: @@ -56,7 +56,7 @@ def _refuse_unskippable_resources( not just the ones that answered the probe, so one that cannot account for tokens it never saw would plan its step against a sequence that never ran. """ - declared = set(model.prefix_key_streams()) + declared = set(keyed_labels(model)) if not declared: return cached_nodes = { @@ -120,42 +120,43 @@ def refuse_combined_walks_across_tp_groups( def _refuse_unknown_walks(model: Model) -> None: - """Refuse a declared stream that names a walk the model never runs. + """Refuse a declared label that names a walk the model never runs. Only the named walks may file pages or be probed, so a name that matches nothing leaves that cache empty for good, and nothing says why. """ walks = set(model.get_graph_walk_graphs()) - for key, by_label in model.prefix_key_streams().items(): - for label, stream in by_label.items(): - unknown = sorted( - {stream.walk, stream.decode_walk, *stream.layout_walks} - walks - {None} - ) + for key, by_label in model.kv_labels().items(): + for label, kv_label in by_label.items(): + named = {walk for walk, _ in kv_label.growth} + if kv_label.keys is not None: + named.update(kv_label.keys.walks) + unknown = sorted(named - walks) if unknown: raise ValueError( - f"{type(model).__name__} keys {key!r}/{label!r} on " + f"{type(model).__name__} declares {key!r}/{label!r} over " f"{unknown}, which it never runs; its walks are " f"{sorted(walks)}" ) def _refuse_unservable_walks(specs: list[NodeResourceSpec], model: Model) -> None: - """Refuse a stream with layout walks if the cache could not serve one of its - walks whole: a walk where a keyed node feeds another node, or where no + """Refuse a keyed label over several walks if the cache could not serve one + of them whole: a walk where a keyed node feeds another node, or where no keyed node runs. A walk the cache serves whole completes with no outputs, so a node that - reads them would get nothing. Only a stream with layout walks is checked, - because a stream with one span always runs its last token. + reads them would get nothing. A label keyed over one walk is not checked, + because its one span always runs its last token. """ graphs = model.get_graph_walk_graphs() by_key = {spec.resource_key: spec for spec in specs} - for key, by_label in model.prefix_key_streams().items(): - for label, stream in by_label.items(): - if not stream.layout_walks: + for key, by_label in keyed_labels(model).items(): + for label, kv_label in by_label.items(): + if len(kv_label.keys.walks) < 2: continue keyed_nodes = by_key[key].nodes - for walk in (stream.walk, *stream.layout_walks): + for walk in kv_label.keys.walks: nodes = graphs[walk].get_nodes() ran = sorted(keyed_nodes & set(nodes)) if not ran: diff --git a/test/modular/test_kv_prefix_bagel_parity.py b/test/modular/test_kv_prefix_bagel_parity.py index da42a567a6..ea08d611f3 100644 --- a/test/modular/test_kv_prefix_bagel_parity.py +++ b/test/modular/test_kv_prefix_bagel_parity.py @@ -40,7 +40,7 @@ ) from mstar.engine.resources.base import EngineResourceInfo from mstar.engine.resources.kv import manager as manager_mod -from mstar.engine.resources.kv.config import KVReqConfig, KVSpec, KVStep, PrefixSpan +from mstar.engine.resources.kv.config import KVReqConfig, KVSpec, KVStep, PrefixSpan, StreamLayout from mstar.engine.resources.kv.keys import chain from mstar.engine.resources.kv.manager import KVManager from mstar.engine.resources.position.config import PositionStep @@ -187,11 +187,11 @@ def ingest(self, rid: str, tokens: list[int]) -> None: for at in range(0, len(tokens), PAGE_SIZE) ])}, prefix_tail={"main": tokens[whole * PAGE_SIZE:]}, - prefix_layout={"main": [PrefixSpan(len(tokens), len(tokens), WALK)]}, + kv_layout={"main": StreamLayout(spans=(PrefixSpan(len(tokens), len(tokens), WALK),), keyed_spans=1)}, )}) def resolve(self, rid: str) -> int: - inputs = ARNodeInputs(input_seq_len=self.kv._overrides[rid].prefix_layout["main"][0].length) + inputs = ARNodeInputs(input_seq_len=self.kv._overrides[rid].kv_layout["main"].spans[0].length) prefix = self.runner.resolve_cached_prefix(rid, NODE, WALK, inputs) self.runner.apply_cached_prefix(rid, NODE, WALK, inputs, prefix) return prefix.tokens if prefix is not None else 0 diff --git a/test/modular/test_kv_prefix_batch.py b/test/modular/test_kv_prefix_batch.py index bff7bce16d..56ce756ff2 100644 --- a/test/modular/test_kv_prefix_batch.py +++ b/test/modular/test_kv_prefix_batch.py @@ -48,7 +48,7 @@ def _run(batched: bool, lazy_steps: int | None = None, flush: bool = False): rids = ["r0", "r1"] prompts = {"r0": list(range(100)), "r1": list(range(50, 50 + H.PAGE_SIZE + 10))} for rid in rids: - H._ingest(kv, rid, prompts[rid]) + H._ingest(kv, rid, prompts[rid], decode=2 * H.PAGE_SIZE + 2) # the prefill of each, and the token it sampled for rid in rids: H._step(kv, rid, len(prompts[rid])) diff --git a/test/modular/test_kv_prefix_chain_parent.py b/test/modular/test_kv_prefix_chain_parent.py index 9608f13e04..181f42bfce 100644 --- a/test/modular/test_kv_prefix_chain_parent.py +++ b/test/modular/test_kv_prefix_chain_parent.py @@ -21,7 +21,7 @@ import torch from mstar.engine.resources.kv import manager as manager_mod -from mstar.engine.resources.kv.config import KVReqConfig, KVStep, PagedKVConfig, PrefixSpan +from mstar.engine.resources.kv.config import KVReqConfig, KVStep, PagedKVConfig, PrefixSpan, StreamLayout from mstar.engine.resources.kv.keys import chain, fingerprint from mstar.engine.resources.kv.manager import KVManager from mstar.engine.resources.step import Segment, StepContext @@ -85,7 +85,7 @@ def _ingest(kv: KVManager, rid: str, tokens: list[int]) -> None: tokens[at:at + PAGE_SIZE] for at in range(0, len(tokens), PAGE_SIZE) ])}, prefix_tail={"main": tokens[whole * PAGE_SIZE:]}, - prefix_layout={"main": [PrefixSpan(len(tokens), len(tokens), WALK)]}, + kv_layout={"main": StreamLayout(spans=(PrefixSpan(len(tokens), len(tokens), WALK),), keyed_spans=1)}, )) diff --git a/test/modular/test_kv_prefix_correctness.py b/test/modular/test_kv_prefix_correctness.py index e6ac113cbe..97b3e3d319 100644 --- a/test/modular/test_kv_prefix_correctness.py +++ b/test/modular/test_kv_prefix_correctness.py @@ -24,7 +24,7 @@ from mstar.engine.resources import PagedKVConfig, PositionConfig, StepContext, StepRunner from mstar.engine.resources.kv import manager as manager_mod -from mstar.engine.resources.kv.config import KVReqConfig, KVStep, PrefixSpan +from mstar.engine.resources.kv.config import KVReqConfig, KVStep, PrefixSpan, StreamLayout from mstar.engine.resources.kv.keys import chain from mstar.engine.resources.kv.manager import KVManager from mstar.engine.resources.position.config import PositionStep @@ -100,13 +100,13 @@ def ingest(self, rid: str, tokens: list[int]) -> None: for at in range(0, len(tokens), PAGE_SIZE) ])}, prefix_tail={"main": tokens[whole * PAGE_SIZE:]}, - prefix_layout={"main": [PrefixSpan(len(tokens), len(tokens), WALK)]}, + kv_layout={"main": StreamLayout(spans=(PrefixSpan(len(tokens), len(tokens), WALK),), keyed_spans=1)}, )}) def resolve(self, rid: str) -> int: """Both halves, in the order the engine runs them: every resource has to be told the length before any of them is asked to act on it.""" - inputs = ARNodeInputs(input_seq_len=self.kv._overrides[rid].prefix_layout["main"][0].length) + inputs = ARNodeInputs(input_seq_len=self.kv._overrides[rid].kv_layout["main"].spans[0].length) prefix = self.runner.resolve_cached_prefix(rid, NODE, WALK, inputs) self.runner.apply_cached_prefix(rid, NODE, WALK, inputs, prefix) return prefix.tokens if prefix is not None else 0 @@ -252,24 +252,24 @@ def test_two_schedules_over_the_same_tokens_agree_without_any_cache(): def test_orpheus_declares_the_tensor_its_llm_prefills_from(): - from mstar.model.base import PrefixStream + from mstar.model.base import PrefixKeys from mstar.model.orpheus.config import KV_CACHE from mstar.model.orpheus.orpheus_model import OrpheusModel # the method reads nothing off the model, so a stand-in is the whole of it - declared = OrpheusModel.prefix_key_streams(object()) + declared = OrpheusModel.kv_labels(object()) - assert declared == { - KV_CACHE: {"main": PrefixStream("text_inputs", "ids", "prefill")} - }, "the LLM node's prompt does not arrive under the name it keys" + assert declared[KV_CACHE]["main"].keys == PrefixKeys("text_inputs", ("prefill",)), ( + "the LLM node's prompt does not arrive under the name it keys" + ) def test_bagel_declares_the_tensor_its_text_walk_prefills_from(): from mstar.model.bagel.bagel_model import BagelModel - from mstar.model.base import PrefixStream + from mstar.model.base import PrefixKeys - declared = BagelModel.prefix_key_streams(object()) + declared = BagelModel.kv_labels(object()) - assert declared == { - "kv": {"main": PrefixStream("text_inputs", "ids", "prefill_text", "decode", ("prefill_vit",))} - }, "the text walk's prompt does not arrive under the name it keys" + assert declared["kv"]["main"].keys == PrefixKeys("text_inputs", ("prefill_text", "prefill_vit"), growth=True), ( + "the text walk's prompt does not arrive under the name it keys" + ) diff --git a/test/modular/test_kv_prefix_decode.py b/test/modular/test_kv_prefix_decode.py index ca94ec0305..d6e7ac8e93 100644 --- a/test/modular/test_kv_prefix_decode.py +++ b/test/modular/test_kv_prefix_decode.py @@ -23,7 +23,7 @@ import torch from mstar.engine.resources.kv import manager as manager_mod -from mstar.engine.resources.kv.config import KVReqConfig, KVStep, PagedKVConfig, PrefixSpan +from mstar.engine.resources.kv.config import KVReqConfig, KVStep, PagedKVConfig, PrefixSpan, StreamLayout from mstar.engine.resources.kv.keys import chain from mstar.engine.resources.kv.manager import KVManager from mstar.engine.resources.step import Segment, StepContext @@ -84,13 +84,21 @@ def _pages(tokens: list[int]) -> list[list[int]]: ] -def _ingest(kv: KVManager, rid: str, prompt: list[int], chains: bool = True): +def _layout(prompt: list[int]) -> StreamLayout: + """The prompt, then a slot per decode step.""" + return StreamLayout( + spans=(PrefixSpan(len(prompt), len(prompt), WALK),), growth=(("decode_loop", 1),), keyed_spans=1, + ) + + +def _ingest(kv: KVManager, rid: str, prompt: list[int], chains: bool = True, decode: int = 0): whole = len(prompt) // PAGE_SIZE kv.ingest_request(rid, KVReqConfig( prefix_keys={"main": chain(_pages(prompt))}, prefix_tail={"main": prompt[whole * PAGE_SIZE:]}, prefix_decode={"main": TENSOR} if chains else None, - prefix_layout={"main": [PrefixSpan(len(prompt), len(prompt), WALK)]}, + kv_layout={"main": _layout(prompt)}, + loop_caps={"decode_loop": decode}, )) @@ -133,7 +141,7 @@ def _indexed(kv: KVManager) -> int: def test_a_generated_page_is_indexed_at_the_commit_that_fills_it(): kv = _manager() prompt = list(range(PAGE_SIZE)) - _ingest(kv, "r0", prompt) + _ingest(kv, "r0", prompt, decode=PAGE_SIZE + 1) _prefill(kv, "r0", prompt) assert _indexed(kv) == 1, "the prompt's own page was not indexed" @@ -160,7 +168,7 @@ def test_a_generated_page_is_matched_by_the_request_that_asks_for_it_next(): kv = _manager() prompt = list(range(100)) generated = list(range(9000, 9000 + 28)) - _ingest(kv, "r0", prompt) + _ingest(kv, "r0", prompt, decode=len(generated)) _prefill(kv, "r0", prompt, sampled=generated[0]) for token in generated[1:]: _decode(kv, "r0", token) @@ -184,7 +192,7 @@ def test_a_generated_page_is_matched_by_the_request_that_asks_for_it_next(): def test_a_stop_on_the_boundary_leaves_the_page_one_slot_short(): kv = _manager() prompt = list(range(100)) - _ingest(kv, "r0", prompt) + _ingest(kv, "r0", prompt, decode=27) _prefill(kv, "r0", prompt, sampled=9000) # the request stops here: its last two sampled ids, the one before the @@ -207,7 +215,7 @@ def test_a_stop_on_the_boundary_leaves_the_page_one_slot_short(): def test_a_node_whose_decode_ids_are_not_the_sampled_token_stops_at_its_prompt(): kv = _manager() prompt = list(range(PAGE_SIZE)) - _ingest(kv, "r0", prompt, chains=False) + _ingest(kv, "r0", prompt, chains=False, decode=PAGE_SIZE + 1) _prefill(kv, "r0", prompt) after_prompt = _indexed(kv) @@ -232,7 +240,8 @@ def test_a_request_that_opted_out_chains_nothing(): prefix_keys={"main": chain(_pages(prompt))}, prefix_tail={"main": []}, prefix_decode={"main": TENSOR}, - prefix_layout={"main": [PrefixSpan(len(prompt), len(prompt), WALK)]}, + kv_layout={"main": _layout(prompt)}, + loop_caps={"decode_loop": PAGE_SIZE + 1}, prefix_cache=False, )) diff --git a/test/modular/test_kv_prefix_eviction.py b/test/modular/test_kv_prefix_eviction.py index 8c641db175..21383e60e1 100644 --- a/test/modular/test_kv_prefix_eviction.py +++ b/test/modular/test_kv_prefix_eviction.py @@ -22,7 +22,7 @@ import torch from mstar.engine.resources.kv import manager as manager_mod -from mstar.engine.resources.kv.config import KVReqConfig, KVStep, PagedKVConfig, PrefixSpan +from mstar.engine.resources.kv.config import KVReqConfig, KVStep, PagedKVConfig, PrefixSpan, StreamLayout from mstar.engine.resources.kv.keys import chain from mstar.engine.resources.kv.manager import KVManager from mstar.engine.resources.step import Segment, StepContext @@ -87,7 +87,7 @@ def _ingest(kv: KVManager, rid: str, tokens: list[int]) -> None: ] kv.ingest_request(rid, KVReqConfig( prefix_keys={"main": chain(pages)}, prefix_tail={"main": []}, - prefix_layout={"main": [PrefixSpan(len(tokens), len(tokens), WALK)]}, + kv_layout={"main": StreamLayout(spans=(PrefixSpan(len(tokens), len(tokens), WALK),), keyed_spans=1)}, )) diff --git a/test/modular/test_kv_prefix_fill.py b/test/modular/test_kv_prefix_fill.py index cba8c172d8..ba9b079594 100644 --- a/test/modular/test_kv_prefix_fill.py +++ b/test/modular/test_kv_prefix_fill.py @@ -21,7 +21,7 @@ import torch from mstar.engine.resources.kv import manager as manager_mod -from mstar.engine.resources.kv.config import KVReqConfig, KVStep, PagedKVConfig, PrefixSpan +from mstar.engine.resources.kv.config import KVReqConfig, KVStep, PagedKVConfig, PrefixSpan, StreamLayout from mstar.engine.resources.kv.keys import chain from mstar.engine.resources.kv.manager import KVManager from mstar.engine.resources.step import Segment, StepContext @@ -82,8 +82,8 @@ def _keys(n_tokens: int, first: int = 0) -> list[bytes]: ]) -def _layout(n_tokens: int) -> dict[str, list[PrefixSpan]]: - return {"main": [PrefixSpan(n_tokens, n_tokens, WALK)]} +def _layout(n_tokens: int) -> dict[str, StreamLayout]: + return {"main": StreamLayout(spans=(PrefixSpan(n_tokens, n_tokens, WALK),), keyed_spans=1)} def _step( @@ -110,7 +110,7 @@ def _indexed(kv: KVManager) -> int: def test_a_chunked_prefill_indexes_the_page_the_second_chunk_completes(): kv = _manager() - kv.ingest_request("r0", KVReqConfig(prefix_keys={"main": _keys(128)}, prefix_layout=_layout(128))) + kv.ingest_request("r0", KVReqConfig(prefix_keys={"main": _keys(128)}, kv_layout=_layout(128))) _step(kv, "r0", 100) after_first = _indexed(kv) @@ -125,7 +125,7 @@ def test_only_whole_pages_are_offered(): kv = _manager() # The prompt is three pages and this step writes part of it, so the chain # names page 2 while that page is still partial. - kv.ingest_request("r0", KVReqConfig(prefix_keys={"main": _keys(384)}, prefix_layout=_layout(384))) + kv.ingest_request("r0", KVReqConfig(prefix_keys={"main": _keys(384)}, kv_layout=_layout(384))) _step(kv, "r0", 300) @@ -136,10 +136,10 @@ def test_only_whole_pages_are_offered(): def test_a_second_request_with_the_same_prompt_matches_what_was_indexed(): kv = _manager() keys = _keys(512) - kv.ingest_request("r0", KVReqConfig(prefix_keys={"main": keys}, prefix_layout=_layout(512))) + kv.ingest_request("r0", KVReqConfig(prefix_keys={"main": keys}, kv_layout=_layout(512))) _step(kv, "r0", 512) - kv.ingest_request("r1", KVReqConfig(prefix_keys={"main": keys}, prefix_layout=_layout(512))) + kv.ingest_request("r1", KVReqConfig(prefix_keys={"main": keys}, kv_layout=_layout(512))) assert kv.resolve_cached_prefix("r1", NODE, WALK, ARNodeInputs(input_seq_len=512)).tokens == 3 * PAGE_SIZE, ( "the four indexed pages did not turn into a match of three" @@ -150,11 +150,11 @@ def test_a_second_request_with_the_same_prompt_matches_what_was_indexed(): def test_a_request_that_was_served_from_the_cache_does_not_reindex_it(): kv = _manager() keys = _keys(512) - kv.ingest_request("r0", KVReqConfig(prefix_keys={"main": keys}, prefix_layout=_layout(512))) + kv.ingest_request("r0", KVReqConfig(prefix_keys={"main": keys}, kv_layout=_layout(512))) _step(kv, "r0", 512) indexed = _indexed(kv) - kv.ingest_request("r1", KVReqConfig(prefix_keys={"main": keys}, prefix_layout=_layout(512))) + kv.ingest_request("r1", KVReqConfig(prefix_keys={"main": keys}, kv_layout=_layout(512))) matched = kv.resolve_cached_prefix("r1", NODE, WALK, ARNodeInputs(input_seq_len=512)).tokens _step(kv, "r1", 512 - matched) @@ -168,7 +168,7 @@ def test_the_first_request_to_fill_a_page_is_the_one_the_index_names(): kv = _manager() keys = _keys(256) for rid in ("r0", "r1"): - kv.ingest_request(rid, KVReqConfig(prefix_keys={"main": keys}, prefix_layout=_layout(256))) + kv.ingest_request(rid, KVReqConfig(prefix_keys={"main": keys}, kv_layout=_layout(256))) _step(kv, "r0", 256) winner = list(kv._index.pages()) @@ -185,7 +185,7 @@ def test_the_first_request_to_fill_a_page_is_the_one_the_index_names(): def test_a_rewound_stream_indexes_its_rerun_from_page_zero(): kv = _manager() - kv.ingest_request("r0", KVReqConfig(prefix_keys={"main": _keys(512)}, prefix_layout=_layout(512))) + kv.ingest_request("r0", KVReqConfig(prefix_keys={"main": _keys(512)}, kv_layout=_layout(512))) _step(kv, "r0", 512) kv.reset_request("r0") # nothing left from the first run, so the rerun is the only thing to index @@ -205,7 +205,7 @@ def test_a_rewound_stream_indexes_its_rerun_from_page_zero(): def test_a_capture_step_indexes_nothing(): kv = _manager() - kv.ingest_request("d0", KVReqConfig(prefix_keys={"main": _keys(256)}, prefix_layout=_layout(256))) + kv.ingest_request("d0", KVReqConfig(prefix_keys={"main": _keys(256)}, kv_layout=_layout(256))) # the rid is in the batch, so only the capture itself can keep it out _step(kv, "d0", 256, capture=True) @@ -219,7 +219,7 @@ def test_a_capture_step_indexes_nothing(): def test_a_padded_row_indexes_nothing(): kv = _manager() - kv.ingest_request("pad", KVReqConfig(prefix_keys={"main": _keys(256)}, prefix_layout=_layout(256))) + kv.ingest_request("pad", KVReqConfig(prefix_keys={"main": _keys(256)}, kv_layout=_layout(256))) _step(kv, "pad", 256, padded=True) padded = _indexed(kv) @@ -232,7 +232,10 @@ def test_a_padded_row_indexes_nothing(): def test_a_label_the_request_never_keyed_indexes_nothing(): kv = _manager() - kv.ingest_request("r0", KVReqConfig(prefix_keys={"main": _keys(256)}, prefix_layout=_layout(256))) + kv.ingest_request("r0", KVReqConfig( + prefix_keys={"main": _keys(256)}, + kv_layout={**_layout(256), "cfg_text": StreamLayout(spans=(PrefixSpan(256, 256, WALK),))}, + )) _step(kv, "r0", 256, label="cfg_text") unkeyed = _indexed(kv) @@ -246,9 +249,9 @@ def test_a_label_the_request_never_keyed_indexes_nothing(): def test_a_request_that_opted_out_indexes_nothing(): kv = _manager() kv.ingest_request("r0", KVReqConfig( - prefix_keys={"main": _keys(256)}, prefix_layout=_layout(256), prefix_cache=False, + prefix_keys={"main": _keys(256)}, kv_layout=_layout(256), prefix_cache=False, )) - kv.ingest_request("r1", KVReqConfig(prefix_keys={"main": _keys(256)}, prefix_layout=_layout(256))) + kv.ingest_request("r1", KVReqConfig(prefix_keys={"main": _keys(256)}, kv_layout=_layout(256))) _step(kv, "r0", 256) opted_out = _indexed(kv) diff --git a/test/modular/test_kv_prefix_layout_lease.py b/test/modular/test_kv_prefix_layout_lease.py index aa0726b9a3..5f6e05f17f 100644 --- a/test/modular/test_kv_prefix_layout_lease.py +++ b/test/modular/test_kv_prefix_layout_lease.py @@ -29,7 +29,7 @@ from mstar.engine.resources.kv.config import KVReqConfig, KVStep from mstar.engine.resources.position.config import PositionStep from mstar.engine.resources.step import Segment, SubmoduleStep -from mstar.model.base import PrefixStream, Span +from mstar.model.base import KVLabel, LabelLayout, PrefixKeys, Span from mstar.model.submodule_base import ARNodeInputs RID = "r1" @@ -59,35 +59,41 @@ def _walks(parts: list) -> list[tuple[str, int]]: ] -def _config(parts: list, text_walk: str = TEXT_WALK, image_walk: str = IMAGE_WALK) -> KVReqConfig: - """What the preprocess worker sends for ``parts``, keyed by its own code.""" - layout = [ +def _config(parts: list, text_walk: str = TEXT_WALK, image_walk: str = IMAGE_WALK, decode: int = 0) -> KVReqConfig: + """What the preprocess worker sends for ``parts`` and ``decode`` decode steps, keyed by its own code.""" + layouts = {"kv": {"main": LabelLayout(spans=tuple( Span("ids", text_walk, slots, slots, "text_inputs") if isinstance(part, list) # The span's source holds the digest itself, which the stand-in `_digest` returns. else Span("digest", image_walk, slots, 1, part[1] if isinstance(part, tuple) else DIGEST) for part, (_, slots) in zip(parts, _walks(parts), strict=True) - ] + ), growth=(("decode_loop", 1),))}} + walks = (text_walk, image_walk) worker = SimpleNamespace( - _prefix_streams={"kv": {"main": PrefixStream("text_inputs", "ids", text_walk, "decode", (image_walk,))}}, + _keyed_labels={"kv": {"main": KVLabel(keys=PrefixKeys("text_inputs", walks, growth=True))}}, _prefix_page_sizes={"kv": PAGE_SIZE}, model=None, _digest=lambda span, file_paths, file_states: span.source, _page_layout=lambda *args: PreprocessWorkerThread._page_layout(None, *args), ) tensors = {"text_inputs": [torch.tensor(part) for part in parts if isinstance(part, list)]} - sent = PreprocessWorkerThread._prefix_keys(worker, tensors, {"kv": {"main": layout}}, None, {}) + rows = {key: {label: layout.row() for label, layout in by_label.items()} for key, by_label in layouts.items()} + sent = PreprocessWorkerThread._prefix_keys(worker, tensors, layouts, rows, None, {}) config = KVReqConfig() - config.apply_conductor_config(**{name: by_resource["kv"] for name, by_resource in sent.items()}) + config.apply_conductor_config( + kv_layout=rows["kv"], loop_caps={"decode_loop": decode}, + **{name: by_resource["kv"] for name, by_resource in sent.items()}, + ) return config class _Node(_TextNode): """One request's cache and positions under a real runner, after an earlier - request ran every walk of ``seeded`` and indexed what it filled.""" + request ran every walk of ``seeded`` and indexed what it filled. The earlier + request is laid out for ``decode`` decode steps after them.""" - def __init__(self, seeded: list = PARTS, parts: list | None = None): + def __init__(self, seeded: list = PARTS, parts: list | None = None, decode: int = 0): super().__init__() self.kv.enable_prefix_cache(ROOT, {"main": (TEXT_WALK, "decode", IMAGE_WALK)}) - self.runner.ingest_request("seed", {"kv": _config(seeded)}) + self.runner.ingest_request("seed", {"kv": _config(seeded, decode=decode)}) for walk, length in _walks(seeded): self.step("seed", length, walk) self.runner.ingest_request(RID, {"kv": _config(parts or seeded)}) @@ -242,7 +248,7 @@ def test_a_missed_probe_asked_again_hashes_nothing_and_is_served_what_was_cached def test_a_page_a_reply_fills_after_an_image_is_matched_by_the_next_turn(): reply = list(range(400, 409)) # the image's last two slots, the question and the reply's nine ids fill page 3 - node = _Node([TEXT, IMAGE, ASK], [TEXT, IMAGE, ASK + reply + OTHER_ASK]) + node = _Node([TEXT, IMAGE, ASK], [TEXT, IMAGE, ASK + reply + OTHER_ASK], decode=len(reply)) for token in reply: node.kv.extend_prefix_chain("seed", NODE, "decode", {"text_inputs": [torch.tensor([token])]}) node.step("seed", 1, "decode") @@ -268,7 +274,7 @@ def test_a_different_image_is_never_served_another_images_end(): def test_a_prompt_prefilled_beside_a_decode_is_matched_by_its_repeat(): prompt = list(range(500, 500 + 3 * PAGE_SIZE)) - node = _Node(parts=[prompt]) + node = _Node(parts=[prompt], decode=1) node.runner.ingest_request("first", {"kv": _config([prompt])}) node.mixed_step({"first": len(prompt), "seed": 1}, {"first": TEXT_WALK, "seed": "decode"}) diff --git a/test/modular/test_kv_prefix_lease.py b/test/modular/test_kv_prefix_lease.py index ce75c24d62..195edea230 100644 --- a/test/modular/test_kv_prefix_lease.py +++ b/test/modular/test_kv_prefix_lease.py @@ -21,7 +21,7 @@ from mstar.engine.resources import Resource, StepRunner from mstar.engine.resources.base import CachedPrefix from mstar.engine.resources.kv import manager as manager_mod -from mstar.engine.resources.kv.config import KVReqConfig, KVStep, PagedKVConfig, PrefixSpan +from mstar.engine.resources.kv.config import KVReqConfig, KVStep, PagedKVConfig, PrefixSpan, StreamLayout from mstar.engine.resources.kv.keys import chain, fingerprint from mstar.engine.resources.kv.manager import KVManager from mstar.engine.resources.step import Segment, StepContext @@ -93,8 +93,8 @@ def _keys(tokens: list[int]) -> list[bytes]: ]) -def _layout(n_tokens: int) -> dict[str, list[PrefixSpan]]: - return {"main": [PrefixSpan(n_tokens, n_tokens, WALK)]} +def _layout(n_tokens: int) -> dict[str, StreamLayout]: + return {"main": StreamLayout(spans=(PrefixSpan(n_tokens, n_tokens, WALK),), keyed_spans=1)} def _grow(kv: KVManager, rid: str, span: int, label: str = "main", **step) -> None: @@ -114,7 +114,7 @@ def _seed(kv: KVManager, tokens: list[int]) -> list[bytes]: a later request can then match. """ keys = _keys(tokens) - kv.ingest_request("seed", KVReqConfig(prefix_keys={"main": keys}, prefix_layout=_layout(len(tokens)))) + kv.ingest_request("seed", KVReqConfig(prefix_keys={"main": keys}, kv_layout=_layout(len(tokens)))) _grow(kv, "seed", len(tokens)) pages = kv._streams["seed"]["main"].page_indices parent = None @@ -130,7 +130,7 @@ def _seed(kv: KVManager, tokens: list[int]) -> list[bytes]: def test_a_second_request_with_the_same_prompt_matches_what_the_first_wrote(): kv = _manager() keys = _seed(kv, list(range(100))) - kv.ingest_request("r1", KVReqConfig(prefix_keys={"main": keys}, prefix_layout=_layout(100))) + kv.ingest_request("r1", KVReqConfig(prefix_keys={"main": keys}, kv_layout=_layout(100))) matched = kv.resolve_cached_prefix("r1", NODE, WALK, ARNodeInputs(input_seq_len=100)).tokens @@ -143,7 +143,7 @@ def test_a_prompt_that_diverges_matches_only_what_they_share(): _seed(kv, list(range(100))) diverged = list(range(48)) + [999] * 52 - kv.ingest_request("r1", KVReqConfig(prefix_keys={"main": _keys(diverged)}, prefix_layout=_layout(100))) + kv.ingest_request("r1", KVReqConfig(prefix_keys={"main": _keys(diverged)}, kv_layout=_layout(100))) assert kv.resolve_cached_prefix("r1", NODE, WALK, ARNodeInputs(input_seq_len=100)).tokens == 48, ( "the match ran past the page where the two prompts stop agreeing" @@ -154,7 +154,7 @@ def test_a_prompt_that_diverges_matches_only_what_they_share(): def test_a_page_aligned_prompt_loses_its_last_page(): kv = _manager() keys = _seed(kv, list(range(96))) - kv.ingest_request("r1", KVReqConfig(prefix_keys={"main": keys}, prefix_layout=_layout(96))) + kv.ingest_request("r1", KVReqConfig(prefix_keys={"main": keys}, kv_layout=_layout(96))) assert kv.resolve_cached_prefix("r1", NODE, WALK, ARNodeInputs(input_seq_len=96)).tokens == 80, ( "a full match would leave the request with no token to sample" @@ -165,7 +165,7 @@ def test_a_page_aligned_prompt_loses_its_last_page(): def test_a_request_that_matches_nothing_gets_no_lease(): kv = _manager() _seed(kv, list(range(100))) - kv.ingest_request("r1", KVReqConfig(prefix_keys={"main": _keys([7] * 100)}, prefix_layout=_layout(100))) + kv.ingest_request("r1", KVReqConfig(prefix_keys={"main": _keys([7] * 100)}, kv_layout=_layout(100))) assert kv.resolve_cached_prefix("r1", NODE, WALK, ARNodeInputs(input_seq_len=100)) is None, ( "a prompt sharing no page with anything indexed was answered a length" @@ -182,7 +182,7 @@ def test_a_request_that_matches_nothing_gets_no_lease(): def test_a_refused_admit_answers_the_same_and_holds_one_lease(): kv = _manager(max_num_pages=16) keys = _seed(kv, list(range(100))) - kv.ingest_request("r1", KVReqConfig(prefix_keys={"main": keys}, prefix_layout=_layout(100))) + kv.ingest_request("r1", KVReqConfig(prefix_keys={"main": keys}, kv_layout=_layout(100))) first = kv.resolve_cached_prefix("r1", NODE, WALK, ARNodeInputs(input_seq_len=100)).tokens held = list(kv._streams["r1"]["main"].lease) owners = [kv._arena.num_owners[page] for page in held] @@ -209,7 +209,7 @@ def test_a_refused_admit_answers_the_same_and_holds_one_lease(): def test_remove_gives_back_a_lease_admit_never_took(): kv = _manager() keys = _seed(kv, list(range(100))) - kv.ingest_request("r1", KVReqConfig(prefix_keys={"main": keys}, prefix_layout=_layout(100))) + kv.ingest_request("r1", KVReqConfig(prefix_keys={"main": keys}, kv_layout=_layout(100))) kv.resolve_cached_prefix("r1", NODE, WALK, ARNodeInputs(input_seq_len=100)) leased = list(kv._streams["r1"]["main"].lease) owners = [kv._arena.num_owners[page] for page in leased] @@ -225,7 +225,7 @@ def test_remove_gives_back_a_lease_admit_never_took(): def test_a_reset_gives_back_a_lease_admit_never_took(): kv = _manager() keys = _seed(kv, list(range(100))) - kv.ingest_request("r1", KVReqConfig(prefix_keys={"main": keys}, prefix_layout=_layout(100))) + kv.ingest_request("r1", KVReqConfig(prefix_keys={"main": keys}, kv_layout=_layout(100))) kv.resolve_cached_prefix("r1", NODE, WALK, ARNodeInputs(input_seq_len=100)) kv.reset_request("r1") @@ -242,7 +242,7 @@ def test_a_reset_gives_back_a_lease_admit_never_took(): def test_admit_takes_the_lease_over_and_commit_clears_it(): kv = _manager() keys = _seed(kv, list(range(100))) - kv.ingest_request("r1", KVReqConfig(prefix_keys={"main": keys}, prefix_layout=_layout(100))) + kv.ingest_request("r1", KVReqConfig(prefix_keys={"main": keys}, kv_layout=_layout(100))) matched = kv.resolve_cached_prefix("r1", NODE, WALK, ARNodeInputs(input_seq_len=100)).tokens leased = list(kv._streams["r1"]["main"].lease) @@ -267,7 +267,7 @@ def _converted_by_a_refused_admit(kv: KVManager) -> list[int]: """ keys = _seed(kv, list(range(100))) kv.remove_request("seed") - kv.ingest_request("r1", KVReqConfig(prefix_keys={"main": keys}, prefix_layout=_layout(100))) + kv.ingest_request("r1", KVReqConfig(prefix_keys={"main": keys}, kv_layout=_layout(100))) matched = kv.resolve_cached_prefix("r1", NODE, WALK, ARNodeInputs(input_seq_len=100)).tokens leased = list(kv._streams["r1"]["main"].lease) kv.ingest_request("hog", KVReqConfig()) @@ -334,7 +334,7 @@ def _refused_then_cached(kv: KVManager, prompt: list[int]) -> list[int]: and refuses r2's; before the batch comes round again, another request fills and indexes the prompt r1 is waiting on. """ - kv.ingest_request("r1", KVReqConfig(prefix_keys={"main": _keys(prompt)}, prefix_layout=_layout(len(prompt)))) + kv.ingest_request("r1", KVReqConfig(prefix_keys={"main": _keys(prompt)}, kv_layout=_layout(len(prompt)))) kv.ingest_request("r2", KVReqConfig()) step = KVStep(segments=( Segment("r1", "main", len(prompt)), Segment("r2", "main", 2 * len(prompt)), @@ -394,7 +394,10 @@ def test_a_leased_stream_an_offload_has_claimed_keeps_its_pages(): def test_a_pre_fork_off_a_leased_stream_covers_the_whole_prefix(): kv = _manager() keys = _seed(kv, list(range(100))) - kv.ingest_request("r1", KVReqConfig(prefix_keys={"main": keys}, prefix_layout=_layout(100))) + kv.ingest_request("r1", KVReqConfig( + prefix_keys={"main": keys}, + kv_layout={**_layout(100), "cfg_text": StreamLayout(spans=(PrefixSpan(100, 100, WALK),))}, + )) matched = kv.resolve_cached_prefix("r1", NODE, WALK, ARNodeInputs(input_seq_len=100)).tokens _grow(kv, "r1", 100 - matched, pre_forks=(("main", "cfg_text"),)) @@ -433,7 +436,7 @@ def test_a_lease_is_cut_to_the_smallest_answer_and_the_rest_given_back(): kv = _manager() keys = _seed(kv, list(range(100))) kv.remove_request("seed") - kv.ingest_request("r1", KVReqConfig(prefix_keys={"main": keys}, prefix_layout=_layout(100))) + kv.ingest_request("r1", KVReqConfig(prefix_keys={"main": keys}, kv_layout=_layout(100))) runner = StepRunner( {"kv": kv, "other": _Holding(PAGE_SIZE)}, node_resources={NODE: ["kv", "other"]}, @@ -484,7 +487,7 @@ def test_a_probe_hashes_the_prompt_with_the_lock_down(monkeypatch): kv = _manager() keys = _seed(kv, list(range(100))) kv.remove_request("seed") - kv.ingest_request("r1", KVReqConfig(prefix_keys={"main": keys}, prefix_layout=_layout(100))) + kv.ingest_request("r1", KVReqConfig(prefix_keys={"main": keys}, kv_layout=_layout(100))) _assert_hashed_with_the_lock_down(kv, monkeypatch) assert kv.resolve_cached_prefix("r1", NODE, WALK, ARNodeInputs(input_seq_len=100)), "the probe matched nothing" @@ -521,7 +524,7 @@ def test_a_probe_raced_by_a_remove_raises_nothing(monkeypatch): kv = _manager() keys = _seed(kv, list(range(64))) kv.remove_request("seed") - kv.ingest_request("r0", KVReqConfig(prefix_keys={"main": keys}, prefix_layout=_layout(64))) + kv.ingest_request("r0", KVReqConfig(prefix_keys={"main": keys}, kv_layout=_layout(64))) _removed_in_the_window(kv, monkeypatch, "r0") try: @@ -545,7 +548,7 @@ def test_a_chain_extension_raced_by_a_remove_raises_nothing(monkeypatch): prefix_keys={"main": _keys(tokens)}, prefix_tail={"main": tokens[32:]}, prefix_decode={"main": "text_inputs"}, - prefix_layout=_layout(len(tokens)), + kv_layout=_layout(len(tokens)), )) _grow(kv, "r0", len(tokens)) _removed_in_the_window(kv, monkeypatch, "r0") diff --git a/test/modular/test_kv_prefix_logs.py b/test/modular/test_kv_prefix_logs.py index bc2e2abafd..2958982526 100644 --- a/test/modular/test_kv_prefix_logs.py +++ b/test/modular/test_kv_prefix_logs.py @@ -23,7 +23,7 @@ import torch from mstar.engine.resources.kv import manager as manager_mod -from mstar.engine.resources.kv.config import KVReqConfig, KVStep, PagedKVConfig, PrefixSpan +from mstar.engine.resources.kv.config import KVReqConfig, KVStep, PagedKVConfig, PrefixSpan, StreamLayout from mstar.engine.resources.kv.keys import chain from mstar.engine.resources.kv.manager import KVManager from mstar.engine.resources.kv.transfer import TransferEngineInfo @@ -81,7 +81,7 @@ def _manager(prefix_cache: bool = True, walks=None) -> KVManager: return kv -def _ingest(kv: KVManager, rid: str, tokens: list[int], keyed: bool = True) -> None: +def _ingest(kv: KVManager, rid: str, tokens: list[int], keyed: bool = True, decode: int = 0) -> None: if not keyed: kv.ingest_request(rid, KVReqConfig()) return @@ -91,7 +91,10 @@ def _ingest(kv: KVManager, rid: str, tokens: list[int], keyed: bool = True) -> N tokens[at:at + PAGE_SIZE] for at in range(0, len(tokens), PAGE_SIZE) ])}, prefix_tail={"main": tokens[whole * PAGE_SIZE:]}, - prefix_layout={"main": [PrefixSpan(len(tokens), len(tokens), WALK)]}, + kv_layout={"main": StreamLayout( + spans=(PrefixSpan(len(tokens), len(tokens), WALK),), growth=(("decode_loop", 1),), keyed_spans=1, + )}, + loop_caps={"decode_loop": decode}, )) @@ -141,10 +144,10 @@ def test_a_hit_reports_how_much_of_the_prompt_was_already_here(caplog): _run(kv, "a", len(tokens)) kv.remove_request("a") _ingest(kv, "b", tokens) - kv.resolve_cached_prefix("b", NODE, WALK, ARNodeInputs(input_seq_len=len(tokens))) + hit = kv.resolve_cached_prefix("b", NODE, WALK, ARNodeInputs(input_seq_len=len(tokens))).tokens with caplog.at_level(logging.INFO, logger=manager_mod.__name__): - _run(kv, "b", len(tokens)) + _run(kv, "b", len(tokens) - hit) line = _lines(caplog)[0] assert "matched 3 of its 4 pages" in line, line @@ -163,10 +166,10 @@ def test_a_whole_prompt_that_came_back_says_so(caplog): # every whole page this request has is one the first request left behind longer = tokens + list(range(9000, 9004)) _ingest(kv, "b", longer) - kv.resolve_cached_prefix("b", NODE, WALK, ARNodeInputs(input_seq_len=len(longer))) + hit = kv.resolve_cached_prefix("b", NODE, WALK, ARNodeInputs(input_seq_len=len(longer))).tokens with caplog.at_level(logging.INFO, logger=manager_mod.__name__): - _run(kv, "b", len(longer)) + _run(kv, "b", len(longer) - hit) assert "whole prompt already cached True" in _lines(caplog)[0], _lines(caplog)[0] @@ -208,7 +211,7 @@ def test_a_deployment_that_turned_the_cache_off_reports_nothing(caplog): def test_a_request_admitted_twice_is_still_one_line(caplog): kv = _manager() tokens = list(range(64)) - _ingest(kv, "r0", tokens) + _ingest(kv, "r0", tokens, decode=PAGE_SIZE) with caplog.at_level(logging.INFO, logger=manager_mod.__name__): _run(kv, "r0", len(tokens)) @@ -228,8 +231,8 @@ def _serve(kv: KVManager, count: int, keyed: bool) -> None: for n in range(count): rid = f"r{n}" _ingest(kv, rid, list(range(64)), keyed=keyed) - kv.resolve_cached_prefix(rid, NODE, WALK, ARNodeInputs(input_seq_len=64)) - _run(kv, rid, 64) + hit = kv.resolve_cached_prefix(rid, NODE, WALK, ARNodeInputs(input_seq_len=64)) + _run(kv, rid, 64 - (hit.tokens if hit is not None else 0)) kv.remove_request(rid) diff --git a/test/modular/test_kv_prefix_offload.py b/test/modular/test_kv_prefix_offload.py index 64e661e83c..7a860c0dfa 100644 --- a/test/modular/test_kv_prefix_offload.py +++ b/test/modular/test_kv_prefix_offload.py @@ -22,7 +22,7 @@ import torch from mstar.engine.resources.kv import manager as manager_mod -from mstar.engine.resources.kv.config import KVReqConfig, KVStep, PagedKVConfig, PrefixSpan +from mstar.engine.resources.kv.config import KVReqConfig, KVStep, PagedKVConfig, PrefixSpan, StreamLayout from mstar.engine.resources.kv.keys import chain from mstar.engine.resources.kv.manager import ( KVManager, @@ -100,7 +100,7 @@ def _ingest(kv: KVManager, rid: str, tokens: list[int]) -> None: kv.ingest_request(rid, KVReqConfig( prefix_keys={"main": _keys(tokens)}, prefix_tail={"main": tokens[whole * PAGE_SIZE:]}, - prefix_layout={"main": [PrefixSpan(len(tokens), len(tokens), WALK)]}, + kv_layout={"main": StreamLayout(spans=(PrefixSpan(len(tokens), len(tokens), WALK),), keyed_spans=1)}, )) diff --git a/test/modular/test_kv_prefix_positions.py b/test/modular/test_kv_prefix_positions.py index ff61aa552a..2dc3f7bcdb 100644 --- a/test/modular/test_kv_prefix_positions.py +++ b/test/modular/test_kv_prefix_positions.py @@ -34,7 +34,7 @@ from mstar.engine.resources.position.manager import RopeManager from mstar.engine.resources.spec import NodeResourceSpec from mstar.engine.resources.step import Segment, SubmoduleStep -from mstar.model.base import PrefixStream +from mstar.model.base import KVLabel, PrefixKeys from mstar.worker.engine_manager import ( _refuse_uncacheable_positions, _refuse_unknown_walks, @@ -213,8 +213,9 @@ class _Model: def __init__(self, *keyed: str): self._keyed = keyed - def prefix_key_streams(self): - return {key: {} for key in self._keyed} + def kv_labels(self): + keyed = KVLabel(keys=PrefixKeys("text_inputs", ("prefill",))) + return {key: {"main": keyed} for key in self._keyed} def test_a_declared_node_over_a_block_scheme_is_refused_at_load(): @@ -311,8 +312,10 @@ def __init__(self, walk: str, runs: tuple[str, ...] = ("prefill", "decode")): self._walk = walk self._runs = runs - def prefix_key_streams(self): - return {KV: {"main": PrefixStream("text_inputs", "ids", self._walk, "decode")}} + def kv_labels(self): + return {KV: {"main": KVLabel(growth=(("decode", "decode_loop"),), keys=PrefixKeys( + "text_inputs", (self._walk,), growth=True, + ))}} def get_graph_walk_graphs(self): return {walk: None for walk in self._runs} diff --git a/test/modular/test_kv_prefix_request_keys.py b/test/modular/test_kv_prefix_request_keys.py index dc2a32f09b..62dde9156b 100644 --- a/test/modular/test_kv_prefix_request_keys.py +++ b/test/modular/test_kv_prefix_request_keys.py @@ -25,7 +25,15 @@ from mstar.engine.resources import SamplingReqConfig from mstar.engine.resources.kv.config import KVReqConfig, KVSpec, PagedKVConfig from mstar.engine.resources.kv.keys import chain -from mstar.model.base import PrefixStream, ProcessPromptOutput, Span, TensorAndMetadata +from mstar.model.base import ( + KVLabel, + LabelLayout, + PrefixKeys, + ProcessPromptOutput, + Span, + TensorAndMetadata, + keyed_labels, +) from mstar.model.orpheus.config import OrpheusModelConfig from mstar.model.orpheus.orpheus_model import OrpheusModel @@ -60,22 +68,25 @@ def set_persist(self, uuid, persist): class _Model: - """A model that declares one id-keyed stream, or none when asked not to.""" + """A model that declares one id-keyed label unless ``declares`` is False, and lays it out + unless ``lays_out`` is False.""" def __init__( self, declares: bool = True, metadata: dict | None = None, - decode_walk: str | None = None, + decode_walk: str | None = None, lays_out: bool = True, ): self._declares = declares self._metadata = metadata self._decode_walk = decode_walk + self._lays_out = lays_out - def prefix_key_streams(self): + def kv_labels(self): if not self._declares: return {} - return {"kv": {"main": PrefixStream( - "text_inputs", "ids", "prefill", self._decode_walk, - )}} + growth = () if self._decode_walk is None else ((self._decode_walk, "decode_loop"),) + return {"kv": {"main": KVLabel(growth=growth, keys=PrefixKeys( + "text_inputs", ("prefill",), growth=self._decode_walk is not None, + ))}} def max_input_files(self): return {} @@ -93,9 +104,10 @@ def get_node_resources(self): def process_prompt(self, *args, **kwargs): tensors = {"text_inputs": [torch.tensor(PROMPT, dtype=torch.long)]} - if self._metadata is None: + if not self._lays_out: return tensors - return ProcessPromptOutput(tensors, self._metadata) + layout = LabelLayout(spans=(Span("ids", "prefill", len(PROMPT), len(PROMPT), "text_inputs"),)) + return ProcessPromptOutput(tensors, {"kv_layout": {"kv": {"main": layout}}, **(self._metadata or {})}) def _worker(model, model_config: dict | None) -> PreprocessWorkerThread: @@ -283,13 +295,13 @@ def test_the_metadata_fold_leaves_the_requests_own_kwargs_alone(): assert kwargs["from_process_prompt"] == 1, "the metadata never arrived" -def test_a_model_returning_only_tensors_still_carries_its_keys(): - worker = _worker(_Model(metadata=None), _deployment()) +def test_a_declared_label_with_no_layout_sends_no_keys(): + worker = _worker(_Model(lays_out=False), _deployment()) kwargs = _run(worker, {"voice": "tara"}) assert kwargs["voice"] == "tara", "the request's own kwargs were dropped" - assert "prefix_keys" in kwargs, "the plain-dict return lost the keys" + assert "prefix_keys" not in kwargs, "keys were guessed for a label its request never laid out" # ── the handover to the resource config ───────────────────────────────── @@ -378,7 +390,7 @@ def test_orpheus_is_handed_a_config_for_the_stream_it_declares(): configs = _conductor_configs(orpheus) - assert set(orpheus.prefix_key_streams()) <= set(configs), ( + assert set(keyed_labels(orpheus)) <= set(configs), ( "orpheus returns a sampler config alone, so the resource it declared " "a stream for had nothing to carry its chain" ) @@ -395,11 +407,12 @@ class _LaidOut(_Model): def __init__(self, image: Span, rewrites: bytes | None = None): text = Span("ids", "prefill", len(PROMPT), len(PROMPT), "text_inputs") - super().__init__(metadata={"prefix_layout": {"kv": {"main": [text, image]}}}) + super().__init__(metadata={"kv_layout": {"kv": {"main": LabelLayout(spans=(text, image))}}}) self._rewrites = rewrites - def prefix_key_streams(self): - return {"kv": {"main": PrefixStream("text_inputs", "ids", "prefill", None, ("prefill_image",))}} + def kv_labels(self): + walks = ("prefill", "prefill_image") + return {"kv": {"main": KVLabel(keys=PrefixKeys("text_inputs", walks))}} def preprocess_fingerprint(self): return "stub" diff --git a/test/modular/test_kv_prefix_root.py b/test/modular/test_kv_prefix_root.py index 958574cb08..296b165969 100644 --- a/test/modular/test_kv_prefix_root.py +++ b/test/modular/test_kv_prefix_root.py @@ -30,13 +30,13 @@ from mstar.engine.engine import Engine from mstar.engine.resources import Resource from mstar.engine.resources.kv import manager as manager_mod -from mstar.engine.resources.kv.config import KVReqConfig, KVSpec, PagedKVConfig, PrefixSpan +from mstar.engine.resources.kv.config import KVReqConfig, KVSpec, PagedKVConfig, PrefixSpan, StreamLayout from mstar.engine.resources.kv.manager import KVManager from mstar.engine.resources.position.config import PositionConfig, PositionSpec from mstar.engine.resources.position.manager import RopeManager from mstar.graph.base import GraphEdge, GraphNode, Sequential from mstar.graph.special_destinations import EMIT_TO_CLIENT -from mstar.model.base import PrefixStream +from mstar.model.base import KVLabel, PrefixKeys from mstar.worker.engine_manager import _refuse_unservable_walks KV = "kv" @@ -88,15 +88,20 @@ def checkpoint_path(self): def preprocess_fingerprint(self): return self._name - def prefix_key_streams(self): + def kv_labels(self): return {} class _Declaring(_Model): """Keys ``main`` on the KV resource from its text walk.""" - def prefix_key_streams(self): - return {KV: {"main": PrefixStream("text_inputs", "ids", "prefill_text", "decode")}} + def kv_labels(self): + return {KV: {"main": KVLabel(growth=(("decode", "decode_loop"),), keys=PrefixKeys( + "text_inputs", ("prefill_text",), growth=True, + ))}} + + def get_graph_walk_graphs(self): + return {"prefill_text": GraphNode(name="LLM", input_names=["text_inputs"], outputs=[])} def _checkpoint(root, config: bytes) -> str: @@ -312,13 +317,13 @@ def test_an_uncached_node_says_nothing_at_load(caplog, overrides, model): # ── the keys the request brought ──────────────────────────────────────── -_LAYOUT = {"main": [PrefixSpan(32, 32, "prefill_text")]} +_LAYOUT = {"main": StreamLayout(spans=(PrefixSpan(32, 32, "prefill_text"),), keyed_spans=1)} def test_ingest_puts_the_requests_keys_on_the_stream(): kv = _kv() - kv.ingest_request("r0", KVReqConfig(prefix_keys={"main": [b"k0", b"k1"]}, prefix_layout=_LAYOUT)) + kv.ingest_request("r0", KVReqConfig(prefix_keys={"main": [b"k0", b"k1"]}, kv_layout=_LAYOUT)) assert kv._streams["r0"]["main"].chain.keys == [b"k0", b"k1"], ( "the chain the request brought never reached its stream" @@ -327,7 +332,7 @@ def test_ingest_puts_the_requests_keys_on_the_stream(): def test_a_label_the_request_did_not_key_carries_none(): kv = _kv() - kv.ingest_request("r0", KVReqConfig(prefix_keys={"main": [b"k0"]}, prefix_layout=_LAYOUT)) + kv.ingest_request("r0", KVReqConfig(prefix_keys={"main": [b"k0"]}, kv_layout=_LAYOUT)) stream = kv._ensure_label("r0", "cfg_text") @@ -338,7 +343,7 @@ def test_a_request_can_opt_out_of_the_cache(): kv = _kv() kv.ingest_request("r0", KVReqConfig( - prefix_keys={"main": [b"k0"]}, prefix_layout=_LAYOUT, prefix_cache=False, + prefix_keys={"main": [b"k0"]}, kv_layout=_LAYOUT, prefix_cache=False, )) assert kv._streams["r0"]["main"].chain is None, ( @@ -428,9 +433,10 @@ def node(name, feeds): "prefill_image": Sequential([node("vit", "LLM"), node("LLM", image_feeds)]) if image_walk_runs_llm else node("vit", EMIT_TO_CLIENT), } - stream = PrefixStream("text_inputs", "ids", "prefill", None, ("prefill_image",)) + walks = ("prefill", "prefill_image") + label = KVLabel(keys=PrefixKeys("text_inputs", walks)) return SimpleNamespace( - prefix_key_streams=lambda: {KV: {"main": stream}}, get_graph_walk_graphs=lambda: graphs, + kv_labels=lambda: {KV: {"main": label}}, get_graph_walk_graphs=lambda: graphs, checkpoint_path=lambda: None, preprocess_fingerprint=lambda: "laid out", ) diff --git a/test/modular/test_kv_prefix_uncovered.py b/test/modular/test_kv_prefix_uncovered.py index c512ed86d6..df3d22003c 100644 --- a/test/modular/test_kv_prefix_uncovered.py +++ b/test/modular/test_kv_prefix_uncovered.py @@ -24,7 +24,7 @@ import torch from mstar.engine.resources.kv import manager as manager_mod -from mstar.engine.resources.kv.config import KVReqConfig, KVStep, PagedKVConfig, PrefixSpan +from mstar.engine.resources.kv.config import KVReqConfig, KVStep, PagedKVConfig, PrefixSpan, StreamLayout from mstar.engine.resources.kv.keys import chain from mstar.engine.resources.kv.manager import ( KVManager, @@ -84,15 +84,17 @@ def _manager() -> KVManager: return kv -def _ingest(kv: KVManager, rid: str, prompt: list[int] = PROMPT) -> None: +def _ingest(kv: KVManager, rid: str, prompt: list[int] = PROMPT, decode: int = 0, image: int = 0) -> None: whole = len(prompt) // PAGE_SIZE + spans = (PrefixSpan(len(prompt), len(prompt), WALK),) + ((PrefixSpan(image, 1, WALK),) if image else ()) kv.ingest_request(rid, KVReqConfig( prefix_keys={"main": chain([ prompt[at:at + PAGE_SIZE] for at in range(0, len(prompt), PAGE_SIZE) ])}, prefix_tail={"main": prompt[whole * PAGE_SIZE:]}, prefix_decode={"main": TENSOR}, - prefix_layout={"main": [PrefixSpan(len(prompt), len(prompt), WALK)]}, + kv_layout={"main": StreamLayout(spans=spans, growth=(("decode_loop", 1),), keyed_spans=1)}, + loop_caps={"decode_loop": decode}, )) @@ -119,7 +121,7 @@ def _indexed(kv: KVManager) -> int: def test_an_image_after_the_text_leaves_the_chain_at_the_text(): kv = _manager() - _ingest(kv, "r0") + _ingest(kv, "r0", decode=30, image=IMAGE) _step(kv, "r0", len(PROMPT)) # the image is the last thing prefilled, so its step samples the first token _step(kv, "r0", IMAGE) @@ -141,7 +143,7 @@ def test_an_image_after_the_text_leaves_the_chain_at_the_text(): def test_a_read_in_stream_keys_nothing_it_generates(): kv = _manager() - _ingest(kv, "r0") + _ingest(kv, "r0", decode=30) published = PublishedKVInfo.build_for_rank(0, 1, {"main": KVSequenceInfo( seq_len=len(PROMPT), latest_kv_transfer_info="peer", page_indices=list(range(3)), @@ -194,7 +196,7 @@ def test_a_peers_publish_drops_the_partial_chain_even_when_nothing_is_read(): def test_a_text_stream_keeps_its_chain_through_its_decode(): kv = _manager() - _ingest(kv, "r0") + _ingest(kv, "r0", decode=30) _step(kv, "r0", len(PROMPT)) _sampled(kv, "r0", 9000) @@ -210,7 +212,7 @@ def test_a_text_stream_keeps_its_chain_through_its_decode(): def test_a_stream_reading_its_own_publish_keeps_keying_its_samples(): kv = _manager() - _ingest(kv, "r0") + _ingest(kv, "r0", decode=30) _step(kv, "r0", len(PROMPT)) _sampled(kv, "r0", 9000) @@ -228,7 +230,7 @@ def test_a_stream_reading_its_own_publish_keeps_keying_its_samples(): def test_a_decode_step_that_commits_before_its_token_is_read_back_keeps_the_chain(): kv = _manager() - _ingest(kv, "r0") + _ingest(kv, "r0", decode=30) _step(kv, "r0", len(PROMPT)) # every step commits the token its predecessor sampled before that @@ -251,13 +253,13 @@ def test_a_decode_step_that_commits_before_its_token_is_read_back_keeps_the_chai def test_a_stream_that_released_its_front_keys_nothing_further(): kv = _manager() - _ingest(kv, "r0") + _ingest(kv, "r0", decode=30) _step(kv, "r0", len(PROMPT)) indexed = _indexed(kv) _sampled(kv, "r0", 9000) - # what a sliding window leaves behind: the pages moved, the lengths did not - kv._streams["r0"]["main"].released = PAGE_SIZE + # what a sliding window leaves behind: the front page goes, the positions after it stay + assert kv.release_oldest("r0", PAGE_SIZE) == PAGE_SIZE for token in range(9001, 9031): _step(kv, "r0", 1) _sampled(kv, "r0", token) @@ -304,7 +306,7 @@ def test_a_stream_with_no_window_indexes_every_page_it_fills(): def _first_write(walk: str) -> KVManager: - """Key 40 tokens, then commit 50 under ``walk`` before anything else.""" + """Key 40 tokens, then commit them under ``walk`` before anything else.""" kv = KVManager( cfg=PagedKVConfig( num_layers=1, num_kv_heads=1, head_dim=8, max_seq_len=4096, @@ -318,9 +320,9 @@ def _first_write(walk: str) -> KVManager: kv.ingest_request("r0", KVReqConfig( prefix_keys={"main": chain([prompt[:16], prompt[16:32], prompt[32:]])}, prefix_tail={"main": prompt[32:]}, - prefix_layout={"main": [PrefixSpan(len(prompt), len(prompt), "prefill_text")]}, + kv_layout={"main": StreamLayout(spans=(PrefixSpan(len(prompt), len(prompt), walk),), keyed_spans=1)}, )) - step = KVStep(segments=(Segment("r0", "main", 50),)) + step = KVStep(segments=(Segment("r0", "main", len(prompt)),)) ctx = StepContext(request_ids=("r0",), graph_walk=walk, slot=0, capture=False) assert kv.admit(step, ctx).ok kv.plan(step, ctx) diff --git a/test/modular/test_media_decode_limits.py b/test/modular/test_media_decode_limits.py index 9314a83d74..7e013410cb 100644 --- a/test/modular/test_media_decode_limits.py +++ b/test/modular/test_media_decode_limits.py @@ -162,7 +162,7 @@ def test_refused_image_reaches_the_client_as_a_400(tmp_path, no_decode): wt.in_flight_requests = set() wt.request_output_state = {} wt.device = "cpu" - wt._prefix_streams = {} + wt._keyed_labels = {} wt._live_requests = {} wt.model = types.SimpleNamespace( load_image=lambda p, device: Model.load_image(None, p, device), max_input_files=dict, @@ -514,7 +514,7 @@ def _preprocess_thread(model): wt.in_flight_requests = set() wt.request_output_state = {} wt.device = "cpu" - wt._prefix_streams = {} + wt._keyed_labels = {} wt._live_requests = {} wt.model = model return wt diff --git a/test/modular/test_multimodal_ordering.py b/test/modular/test_multimodal_ordering.py index 7835e01c5a..7b8a32a276 100644 --- a/test/modular/test_multimodal_ordering.py +++ b/test/modular/test_multimodal_ordering.py @@ -656,6 +656,7 @@ def _laid_out(bagel, modalities: list[str], output_modalities=("text",), tensors """Return what Bagel's `process_prompt` gives for the text and images in ``modalities``, with the images.""" bagel.config.vit_config = SimpleNamespace(patch_size=PATCH_SIZE) bagel.config.vit_max_num_patch_per_side = MAX_PATCHES_PER_SIDE + bagel.config.latent_downsample, bagel.config.max_latent_size = 16, 64 texts = iter(["what is", "in this picture", "and in this one"]) parts, images = [], [] for modality in modalities: @@ -695,7 +696,7 @@ def test_bagel_lays_each_image_out_in_schedule_order_by_its_own_file_over_its_vi bagel, modalities, image_preprocess, ): output, images = _laid_out(bagel, modalities, image_preprocess=image_preprocess) - layout = output.metadata["prefix_layout"]["kv"]["main"] + layout = output.metadata["kv_layout"]["kv"]["main"].spans schedule = bagel._build_prefill_schedule( input_modalities=modalities, is_understanding=True, input_signals={"text_inputs": output.new_input_tensors["text_inputs"], "image_inputs": images}, @@ -714,21 +715,26 @@ def test_bagel_lays_each_image_out_in_schedule_order_by_its_own_file_over_its_vi ) -def test_bagel_lays_out_only_the_understood_prompts_it_has_images_for(bagel): +def test_bagel_keys_an_understood_image_and_otherwise_only_a_leading_text_span(bagel): interleaved = ["text", "image", "text"] cases = { - "understood with its image": (interleaved, ("text",), True), - "text only": (["text"], ("text",), True), - "generation": (interleaved, ("image",), True), - "understood without tensors": (interleaved, ("text",), False), + "understood with its image": (interleaved, ("text",)), + "text only": (["text"], ("text",)), + "generation": (interleaved, ("image",)), + "an edit led by its image": (["image", "text"], ("image",)), # Text comes first, so the prompt is understood, but it also outputs an image, so its input is resized first. - "understood and edited": (interleaved, ("text", "image"), True), + "understood and edited": (interleaved, ("text", "image")), } - # A prompt with no layout comes back as a plain dict of its tensors. - laid_out = {name: not isinstance(_laid_out(bagel, *case)[0], dict) for name, case in cases.items()} + keyed = { + name: _laid_out(bagel, *case)[0].metadata["kv_layout"]["kv"]["main"].keyed_spans + for name, case in cases.items() + } - assert laid_out == {name: name == "understood with its image" for name in cases}, ( - "a layout went to a prompt with no understood image or none of its " - "tensors, or was missing from the one that has both" + assert keyed == { + "understood with its image": None, "text only": None, "generation": 1, + "an edit led by its image": 0, "understood and edited": 1, + }, "a prompt whose images are resized before their walks see them is keyed past its first text span" + assert isinstance(_laid_out(bagel, interleaved, tensors=False)[0], dict), ( + "a prompt was laid out without the images that size it" ) diff --git a/test/modular/test_orpheus_empty_text.py b/test/modular/test_orpheus_empty_text.py index aac3016437..6ee0359bb2 100644 --- a/test/modular/test_orpheus_empty_text.py +++ b/test/modular/test_orpheus_empty_text.py @@ -190,6 +190,6 @@ async def run(): def test_text_is_still_tokenized_with_its_voice(): - (ids,) = _orpheus().process_prompt("hi", ["text"], ["audio"], voice="tara")["text_inputs"] + (ids,) = _orpheus().process_prompt("hi", ["text"], ["audio"], voice="tara").new_input_tensors["text_inputs"] config = OrpheusModelConfig() assert ids.tolist() == [config.start_token_id, *map(ord, "tara: hi"), *config.end_token_ids] diff --git a/test/modular/test_request_media_totals.py b/test/modular/test_request_media_totals.py index 3e42779b9b..8d3336388c 100644 --- a/test/modular/test_request_media_totals.py +++ b/test/modular/test_request_media_totals.py @@ -58,7 +58,7 @@ def _preprocess_thread(model): wt.in_flight_requests = set() wt.request_output_state = {} wt.device = "cpu" - wt._prefix_streams = {} + wt._keyed_labels = {} wt._live_requests = {} wt.model = model return wt diff --git a/test/modular/test_video_frame_protocol.py b/test/modular/test_video_frame_protocol.py index 44822c2bdb..3d65fe9df2 100644 --- a/test/modular/test_video_frame_protocol.py +++ b/test/modular/test_video_frame_protocol.py @@ -397,7 +397,7 @@ def test_data_worker_cleanup_drops_all_frame_protocol_state(): worker.device = "cpu" worker.enable_prof = False worker.timing_requested = threading.Event() - worker._prefix_streams = {} + worker._keyed_labels = {} worker._live_requests = {} worker.communicator = SimpleNamespace(send=lambda *args: None) worker._process_input( diff --git a/test/modular/test_worker_metrics.py b/test/modular/test_worker_metrics.py index 175176279c..f3b091d75b 100644 --- a/test/modular/test_worker_metrics.py +++ b/test/modular/test_worker_metrics.py @@ -404,9 +404,9 @@ def test_prefix_lookups_count_each_request_once_through_the_engines_probe(): m = _metrics(eng) prompt = list(range(100)) keys = _seed(kv, prompt) - kv.ingest_request("hit", KVReqConfig(prefix_keys={"main": keys}, prefix_layout=_layout(100))) + kv.ingest_request("hit", KVReqConfig(prefix_keys={"main": keys}, kv_layout=_layout(100))) other = _keys(list(range(1000, 1100))) - kv.ingest_request("miss", KVReqConfig(prefix_keys={"main": other}, prefix_layout=_layout(100))) + kv.ingest_request("miss", KVReqConfig(prefix_keys={"main": other}, kv_layout=_layout(100))) assert _prefix(m) == ({}, {}) batch = _prepare(eng, PREFILL, hit=100, miss=100) @@ -467,15 +467,15 @@ def test_prefix_lookups_before_start_counting_are_not_counted(): eng = _engine([kv], keyed=True) m = _metrics(eng, counting=False) keys = _seed(kv, list(range(100))) - kv.ingest_request("warmup-0", KVReqConfig(prefix_keys={"main": keys}, prefix_layout=_layout(100))) + kv.ingest_request("warmup-0", KVReqConfig(prefix_keys={"main": keys}, kv_layout=_layout(100))) _prepare(eng, PREFILL, **{"warmup-0": 100}) assert _prefix(m) == ({}, {}) m.start_counting() assert _prefix(m) == ({(NODE,): 0}, {(NODE,): 0}) other = _keys(list(range(1000, 1100))) - kv.ingest_request("miss", KVReqConfig(prefix_keys={"main": other}, prefix_layout=_layout(100))) - kv.ingest_request("hit", KVReqConfig(prefix_keys={"main": keys}, prefix_layout=_layout(100))) + kv.ingest_request("miss", KVReqConfig(prefix_keys={"main": other}, kv_layout=_layout(100))) + kv.ingest_request("hit", KVReqConfig(prefix_keys={"main": keys}, kv_layout=_layout(100))) _prepare(eng, PREFILL, miss=100, hit=100) assert _prefix(m) == ({(NODE,): 200}, {(NODE,): 96}) From 79ca1bfa5503cc5acf2ae5ddd1870b96d14268a9 Mon Sep 17 00:00:00 2001 From: Gaurav Shah Date: Wed, 21 Oct 2026 05:50:51 +0000 Subject: [PATCH 05/26] kv: always-on page check and debug span checks --- mstar/engine/resources/kv/config.py | 11 ++-- mstar/engine/resources/kv/manager.py | 92 +++++++++++++++++++++++++++- 2 files changed, 96 insertions(+), 7 deletions(-) diff --git a/mstar/engine/resources/kv/config.py b/mstar/engine/resources/kv/config.py index ea530f70c1..eca5d1366c 100644 --- a/mstar/engine/resources/kv/config.py +++ b/mstar/engine/resources/kv/config.py @@ -392,6 +392,12 @@ class StreamLayout: retention: RetentionPolicy | None = None keyed_spans: int = 0 + def slots(self, caps: Mapping[str, int]) -> int: + """Every slot the stream commits, its spans then its growth.""" + return sum(span.length for span in self.spans) + sum( + slots * caps[loop] for loop, slots in self.growth + ) + def pages(self, caps: Mapping[str, int], page_size: int) -> int: """The most pages the stream can hold, where ``caps`` gives each growth loop's iterations. @@ -400,10 +406,7 @@ def pages(self, caps: Mapping[str, int], page_size: int) -> int: the budget, or at the prefix rounded up to a page if that is larger. The bound adds a page of overrun and the largest step past the prefix to that peak. """ - written = sum(span.length for span in self.spans) + sum( - slots * caps[loop] for loop, slots in self.growth - ) - pages = -(-(written + self.scratch) // page_size) + pages = -(-(self.slots(caps) + self.scratch) // page_size) if self.retention is None: return pages prefix = self.retention.protected_prefix diff --git a/mstar/engine/resources/kv/manager.py b/mstar/engine/resources/kv/manager.py index bedc3a0002..7b6ac5fef8 100644 --- a/mstar/engine/resources/kv/manager.py +++ b/mstar/engine/resources/kv/manager.py @@ -258,9 +258,12 @@ class CacheStream: # miss (a refused admit prepares the step again) is not counted twice prefix_counted: bool = False # the most pages the request's layout lets this stream hold, and how far through - # that layout it is; both None without a layout + # that layout it is. Both are None for a request with no layout. A label the + # request doesn't lay out gets a footprint of 0 and no cursor. footprint: int | None = None span_cursor: SpanCursor | None = None + # the most pages the stream has taken, kept only under `_DEBUG_ASSERTS` + peak: int = 0 def reset(self, freed: bool=False, *, content_reset: bool=True): self.stored_len = 0 @@ -518,6 +521,8 @@ def __init__( self._keyed_walks: dict[str, frozenset[str]] = {} # nodes already warned that a declared stream reached them unkeyed self._warned_unkeyed: set[str] = set() + # (label, whether the request opens it) for each overrun already reported + self._overran: set[tuple[str, bool]] = set() # streams whose lease still holds a cached image end for `plan` to copy self._image_ends: list[CacheStream] = [] self._rank = joint_comm_group.rank if joint_comm_group is not None else 0 @@ -1305,11 +1310,17 @@ def ingest_request(self, rid, overrides: KVReqConfig | None=None): self._seed_keys(rid, label, stream) def _lay_out(self, rid: str, label: str, stream: CacheStream) -> None: - """Set ``stream``'s footprint and span cursor from its request's layout for ``label``, if there is one.""" + """Set ``stream``'s footprint and span cursor from its request's layout for ``label``, if there is one. + + A laid-out request is held to no pages on a label it doesn't open, so a batch that + writes one for it anyway takes them past its footprint. + """ layout = self._held_layout(rid, label) if layout is not None: stream.footprint = layout.pages(self._overrides[rid].loop_caps, self.config.page_size) stream.span_cursor = SpanCursor(layout.spans) + elif self._overrides.get(rid) is not None and self._overrides[rid].kv_layout: + stream.footprint = 0 def _held_layout(self, rid: str, label: str) -> StreamLayout | None: overrides = self._overrides.get(rid) @@ -1444,6 +1455,8 @@ def admit(self, step: KVStep, ctx: StepContext) -> AdmitOutcome: with self._lock: if self._admit_reserved_pages: self._forget_step_reservations(ctx) + if _DEBUG_ASSERTS and not step.commit: + self._assert_scratch(step, ctx) # forks reserve here and copy later (plan for pre-, commit for post-), # so a step that never runs leaves pages resident but no page contents # moved — re-admitting it allocates nothing and re-copies nothing. @@ -2292,9 +2305,11 @@ def commit(self, step: KVStep, ctx: StepContext): if policy is not None: self._adopt_retention(stream, policy, segment) touched = [] - if stream.chain is not None: + if stream.span_cursor is not None and (stream.chain is not None or _DEBUG_ASSERTS): written = stream.stored_len + stream.released touched = stream.span_cursor.advance(written - segment.span, written) + if _DEBUG_ASSERTS: + self._assert_layout(segment, stream, touched, ctx, policy) if index_on and stream.chain is not None: self._index_filled_pages(segment, stream, ctx, touched) # so a claim taken in a window the mark misses still fails @@ -2318,6 +2333,26 @@ def commit(self, step: KVStep, ctx: StepContext): if _DEBUG_ASSERTS: self.assert_pages_conserved() + def _assert_layout( + self, segment: Segment, stream: CacheStream, touched: list[tuple[int, int]], + ctx: StepContext, policy: RetentionPolicy | None, + ) -> None: + """Check a commit against its request's layout. Debug only, since it looks the layout up per segment.""" + overrides = self._overrides[segment.request_id] + layout = overrides.kv_layout[segment.label] + walk = ctx.walk_of(segment.request_id) if ctx.request_walks else ctx.graph_walk + where = f"KV {self.name}: {segment.request_id}/{segment.label}" + for index, _ in touched: + assert stream.span_cursor.spans[index].walk == walk, ( + f"{where}: {walk} wrote into span {index}, which its layout gives to " + f"{stream.span_cursor.spans[index].walk}" + ) + written, bound = stream.stored_len + stream.released, layout.slots(overrides.loop_caps) + assert written <= bound, f"{where}: committed {written} slots, past its layout's spans and growth ({bound})" + assert policy is None or policy == layout.retention, ( + f"{where}: a step declares {policy}, where its layout declares {layout.retention}" + ) + # Partial release behind a protected prefix (windowed generation): a # request that generates in windows commits each window's K/V and, once # its context horizon fills, drops the oldest generated pages while the @@ -2908,6 +2943,8 @@ def remove_request(self, rid: str): if stream.read_future is not None: wait([stream.read_future]) with self._lock: + if _DEBUG_ASSERTS and rid in self._streams: + self._report_underuse(rid) if rid in self._streams: for stream in self._streams[rid].values(): self._release_lease(stream) @@ -2924,6 +2961,19 @@ def remove_request(self, rid: str): if _DEBUG_ASSERTS: self.assert_pages_conserved() + def _report_underuse(self, rid: str) -> None: + """Log each label whose stream stayed a page or more under its footprint, so an over-count shows. Debug only.""" + overrides = self._overrides[rid] + for label, layout in (overrides.kv_layout or {}).items(): + stream = self._streams[rid].get(label) + footprint = layout.pages(overrides.loop_caps, self.config.page_size) + peak = 0 if stream is None else stream.peak + if footprint > peak: + logger.warning( + "KV %s: request %s held at most %d pages on label %s against a footprint of %d", + self.name, rid, peak, label, footprint, + ) + def assert_pages_conserved(self) -> None: """Check the owner counts against the streams holding the pages. @@ -3170,8 +3220,44 @@ def _alloc( ) stream.page_indices.extend(new_pages) stream.generation += 1 + if stream.footprint is not None and len(stream.page_indices) > stream.footprint: + self._report_overrun(request_id, label, stream) + if _DEBUG_ASSERTS: + stream.peak = max(stream.peak, len(stream.page_indices)) return AllocResult(new_pages=new_pages) + def _assert_scratch(self, step: KVStep, ctx: StepContext) -> None: + """Check that a step that commits nothing writes no more than its layout's scratch. Debug only.""" + for segment in step.segments: + layout = self._held_layout(segment.request_id, segment.label) + if layout is None or self._on_sink(ctx, segment.request_id) or ( + (segment.request_id, segment.label) in step.sink + ): + continue + assert segment.span <= layout.scratch, ( + f"KV {self.name}: {segment.request_id}/{segment.label}: a step that commits nothing " + f"writes {segment.span} slots, past its layout's scratch of {layout.scratch}" + ) + + def _report_overrun(self, rid: str, label: str, stream: CacheStream) -> None: + """Log once per label that a request took more pages than its layout allows. + + The pool would then promise room it doesn't have, and the request's + layout is the likely cause, so the line prints it. A label the request + doesn't open is reported on its own, so the pages a batch writes there + for every row don't hide a real overrun. + """ + overrides = self._overrides[rid] + reported = (label, label in overrides.kv_layout) + if reported in self._overran: + return + self._overran.add(reported) + logger.error( + "KV %s: request %s took %d pages on label %s against a footprint of %d: layout %s, loop caps %s", + self.name, rid, len(stream.page_indices), label, stream.footprint, + overrides.kv_layout.get(label), overrides.loop_caps, + ) + ### Submodule-level functionality # Label / layer cursors come from `AttentionResource`; the readers resolve # them. TODO: rename the `set_layer_idx` call sites and drop this alias. From e5d052ed2e3afa09df569645435519f05508068f Mon Sep 17 00:00:00 2001 From: Gaurav Shah Date: Wed, 21 Oct 2026 05:50:51 +0000 Subject: [PATCH 06/26] model: orpheus caps decode where its stop check ends it --- mstar/model/orpheus/orpheus_model.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/mstar/model/orpheus/orpheus_model.py b/mstar/model/orpheus/orpheus_model.py index f3b32da74f..ea30b56d5b 100644 --- a/mstar/model/orpheus/orpheus_model.py +++ b/mstar/model/orpheus/orpheus_model.py @@ -387,6 +387,10 @@ def get_initial_forward_pass_args( # Model ABC: resources # ------------------------------------------------------------------- + def get_loop_caps(self, **model_kwargs) -> dict[str, int]: + # the decode stop check ends the loop after max_tokens iterations + return {"decode_loop": self.get_max_output_tokens(**model_kwargs)} + def kv_labels(self) -> dict[str, dict[str, KVLabel]]: """The LLM's prompt is its token ids. Its generated audio codes never come back in a prompt, so they are not keyed.""" From 760eba71021ac726c5facc62074bb2ccb208ac7b Mon Sep 17 00:00:00 2001 From: Gaurav Shah Date: Wed, 21 Oct 2026 05:50:51 +0000 Subject: [PATCH 07/26] model: bagel lays out its guidance labels and caps decode --- mstar/model/bagel/bagel_model.py | 143 +++++++++++------- .../model/bagel/components/modeling_utils.py | 20 ++- mstar/model/bagel/submodules.py | 7 +- test/modular/test_bagel_cfg_parallel.py | 2 +- test/modular/test_multimodal_ordering.py | 2 + 5 files changed, 107 insertions(+), 67 deletions(-) diff --git a/mstar/model/bagel/bagel_model.py b/mstar/model/bagel/bagel_model.py index 1d8312d4fb..cc7f2e1f31 100644 --- a/mstar/model/bagel/bagel_model.py +++ b/mstar/model/bagel/bagel_model.py @@ -625,7 +625,7 @@ def process_prompt( attachments gets its own span and no seam is re-tokenized. The schedule builder walks the same :func:`prefill_plan`. - Given the tensors, it also returns the main stream's ``kv_layout``. + Given the tensors, it also returns a ``kv_layout`` for each label the request may open. """ result: NameToTensorList = {} @@ -705,33 +705,43 @@ def process_prompt( for img in tensors["image_inputs"] ] - if prompt is None or tensors is None: + if tensors is None: return result + if prompt is None: + # the walks then prefill the attachments alone (see `_build_prefill_schedule`) + parts, segments = parts_from_modalities(input_modalities), [] images = result.get("image_inputs") or tensors.get("image_inputs", []) - layout = self._main_layout( - parts, segments, images, output_modalities, think_mode, - vit_preprocess(kwargs.get("image_preprocess")), - kwargs.get("height", 1024), kwargs.get("width", 1024), - ) - return ProcessPromptOutput(result, {"kv_layout": {"kv": {"main": layout}}}) + params = self._request_params(input_modalities, output_modalities, kwargs) + layout = self._kv_layout(parts, segments, images, output_modalities, params) + return ProcessPromptOutput(result, {"kv_layout": {"kv": layout}}) - def _main_layout( + def _kv_layout( self, parts: list[PromptPart], segments: list[torch.Tensor], images: list[torch.Tensor], - output_modalities: list[str], think_mode: bool, image_preprocess: str, height: int, width: int, - ) -> LabelLayout: - """The main stream's writes, in the order `_build_prefill_schedule` runs them. - - A generation keys only its first span, if that is text: its images are resized - before their walks see them, and a file's digest doesn't cover the resize. + output_modalities: list[str], params: dict, + ) -> dict[str, LabelLayout]: + """Each label's writes, in the order `_build_prefill_schedule` runs them. + + ``main`` holds every span. A generation keys only its first span, if that is + text: its images are resized before their walks see them, and a file's digest + doesn't cover the resize. Under guidance, ``cfg_img`` holds the text alone. + ``cfg_text`` forks before a text walk and after an image walk, so it holds all + but the last span when the last span is text, and every span when the prompt + ends on an image. Without guidance a generation still denoises on all three + labels on a CFG-parallel deployment, which isn't known here. """ cfg = self.config - understanding = (output_modalities or ["text"])[0] == "text" + target_output = (output_modalities or ["text"])[0] + understanding = target_output == "text" + think_mode, image_preprocess = params["think_mode"], vit_preprocess(params["image_preprocess"]) spans = [] - if think_mode and not understanding: + if think_mode and not understanding and segments: spans.append(Span("ids", "prefill_text", segments[0].numel(), segments[0].numel(), "text_inputs")) texts = segments[len(spans):] for step in prefill_plan(parts, leading_text=understanding): if step.modality == TEXT: + if step.index >= len(texts): + # no tensor for it: `_build_prefill_schedule` skips the step too + continue length = texts[step.index].numel() spans.append(Span("ids", "prefill_text", length, length, "text_inputs")) continue @@ -747,20 +757,29 @@ def _main_layout( cfg.vit_config.patch_size, cfg.vit_max_num_patch_per_side, ) spans.append(Span("digest", "prefill_vit", patches + IMAGE_SENTINELS, IMAGE_ADVANCE, *item)) - keyed = int(spans[0].kind == "ids") + keyed = int(bool(spans) and spans[0].kind == "ids") if understanding: - return LabelLayout( + main = LabelLayout( spans=tuple(spans), growth=(("decode_loop", 1),), keyed_spans=None if "image" not in output_modalities else keyed, ) + return {"main": main} # an edit generates at its input's size (see `get_initial_forward_pass_args`) - if images: - _, height, width = images[0].shape - latents = (height // cfg.latent_downsample) * (width // cfg.latent_downsample) - return LabelLayout( - spans=tuple(spans), growth=(("decode_loop", 1),) if think_mode else (), - scratch=latents + IMAGE_SENTINELS, keyed_spans=keyed, - ) + height, width = images[0].shape[-2:] if images else (params["height"], params["width"]) + scratch = (height // cfg.latent_downsample) * (width // cfg.latent_downsample) + IMAGE_SENTINELS + growth = (("decode_loop", 1),) if think_mode else () + main = LabelLayout(spans=tuple(spans), growth=growth, scratch=scratch, keyed_spans=keyed) + if not self._requires_cfg(target_output=target_output, **params): + return {"main": main, "cfg_text": LabelLayout(scratch=scratch), "cfg_img": LabelLayout(scratch=scratch)} + return { + "main": main, + "cfg_text": LabelLayout( + spans=tuple(spans[:-1] if spans and spans[-1].walk == "prefill_text" else spans), scratch=scratch, + ), + "cfg_img": LabelLayout( + spans=tuple(span for span in spans if span.walk == "prefill_text"), growth=growth, scratch=scratch, + ), + } @staticmethod def _resize_to_fit( @@ -837,10 +856,19 @@ def kv_labels(self) -> dict[str, dict[str, KVLabel]]: `_kv_layout`). A chat sends a reply back as its own assistant turn, so its pages match. """ - return {"kv": {"main": KVLabel( - growth=((DECODE, "decode_loop"),), - keys=PrefixKeys("text_inputs", ("prefill_text", "prefill_vit"), growth=True), - )}} + decode = ((DECODE, "decode_loop"),) + return {"kv": { + "main": KVLabel( + growth=decode, keys=PrefixKeys("text_inputs", ("prefill_text", "prefill_vit"), growth=True), + ), + # written by forks from main + "cfg_text": KVLabel(), + "cfg_img": KVLabel(growth=decode), + }} + + def get_loop_caps(self, **model_kwargs) -> dict[str, int]: + # The stop check allows max_tokens - 1 iterations, and a loop runs at least one. + return {"decode_loop": max(1, self.get_max_output_tokens(**model_kwargs) - 1)} def checkpoint_path(self) -> str: return snapshot_download( @@ -906,7 +934,10 @@ def get_request_resource_configs( from mstar.model.bagel.submodules import LLM_GRAPH_WALKS, active_labels model_kwargs = model_kwargs or {} - cfg = partition_fwd_args[DEFAULT_PARTITION].full_metadata.kwargs["requires_cfg"] + request = partition_fwd_args[DEFAULT_PARTITION].full_metadata.kwargs + cfg = request["requires_cfg"] + # the CFG nodes run only a generation's image_gen_cfg, so other requests open none of their labels + nodes = self._LLM_NODES if self._has_cfg_parallel and request["target_output"] == "image" else {"LLM"} sampling = self.get_sampling_config("LLM", model_kwargs) publish_labels: dict[tuple[str, str], list[str]] = {} final_publish_labels: dict[tuple[str, str], list[str]] = {} @@ -936,7 +967,7 @@ def get_request_resource_configs( "kv": KVReqConfig( needed_labels_per_node_walk={ (node, walk): active_labels(walk, cfg, node) - for node in self._LLM_NODES + for node in nodes for walk in LLM_GRAPH_WALKS }, # Only the main LLM instance creates cache branches consumed @@ -1251,6 +1282,32 @@ def _build_prefill_schedule( schedule.append(("prefill_vit", images[step.index])) return schedule + def _request_params( + self, input_modalities: list[str], output_modalities: list[str], model_kwargs: dict | None, + ) -> dict: + """The request's generation settings: its own where it set them, else the defaults for its input + and output modalities.""" + overridable_keys = [ + "cfg_text_scale", "cfg_img_scale", "cfg_interval", + "cfg_renorm_type", "cfg_renorm_min", "think_mode", + ] + params = {k: getattr(self.config, k) for k in overridable_keys} + params["width"] = 1024 + params["height"] = 1024 + params["image_preprocess"] = "default" + overridable_keys += ["width", "height", "image_preprocess"] + + if "image" in input_modalities and "image" in output_modalities: + params["cfg_img_scale"] = 2.0 + params["cfg_interval"] = [0.0, 1.0] + params["cfg_renorm_type"] = "text_channel" + + if model_kwargs: + for key in overridable_keys: + if key in model_kwargs: + params[key] = model_kwargs[key] + return params + def _requires_cfg( self, **kwargs, ): @@ -1373,27 +1430,7 @@ def get_initial_forward_pass_args( model_kwargs: dict | None = None, ) -> ForwardPassArgs: target_output = output_modalities[0] # "text" or "image" - - # Per-request overrides with config defaults - overridable_keys = [ - "cfg_text_scale", "cfg_img_scale", "cfg_interval", - "cfg_renorm_type", "cfg_renorm_min", "think_mode", - ] - params = {k: getattr(self.config, k) for k in overridable_keys} - params["width"] = 1024 - params["height"] = 1024 - params["image_preprocess"] = "default" - overridable_keys += ["width", "height", "image_preprocess"] - - if "image" in input_modalities and "image" in output_modalities: - params["cfg_img_scale"] = 2.0 - params["cfg_interval"] = [0.0, 1.0] - params["cfg_renorm_type"] = "text_channel" - - if model_kwargs: - for key in overridable_keys: - if key in model_kwargs: - params[key] = model_kwargs[key] + params = self._request_params(input_modalities, output_modalities, model_kwargs) if "image" in output_modalities and input_signals.get("image_inputs"): image = input_signals["image_inputs"][0] diff --git a/mstar/model/bagel/components/modeling_utils.py b/mstar/model/bagel/components/modeling_utils.py index 6e958fe282..5d3c0121ad 100644 --- a/mstar/model/bagel/components/modeling_utils.py +++ b/mstar/model/bagel/components/modeling_utils.py @@ -314,6 +314,16 @@ def vllm_vit_resize(image: torch.Tensor, size: int) -> torch.Tensor: return F.resize(image, (size, size), InterpolationMode.BICUBIC, antialias=True) +def vllm_vae_size(h: int, w: int, stride: int, max_img_size: int, min_img_size: int = 256) -> tuple[int, int]: + """The ``(height, width)`` `vllm_vae_resize` resizes an ``h`` x ``w`` image to.""" + scale = min(max_img_size / max(w, h), 1.0) + min_size = min(min_img_size, max_img_size) + scale = max(scale, min_size / min(w, h)) + new_w = max(stride, int(round(w * scale / stride) * stride)) + new_h = max(stride, int(round(h * scale / stride) * stride)) + return min(new_h, max_img_size), min(new_w, max_img_size) + + def vllm_vae_resize( image: torch.Tensor, stride: int, @@ -327,11 +337,5 @@ def vllm_vae_resize( normalization is applied separately by the caller. """ _, h, w = image.shape - scale = min(max_img_size / max(w, h), 1.0) - min_size = min(min_img_size, max_img_size) - scale = max(scale, min_size / min(w, h)) - new_w = max(stride, int(round(w * scale / stride) * stride)) - new_h = max(stride, int(round(h * scale / stride) * stride)) - new_w = min(new_w, max_img_size) - new_h = min(new_h, max_img_size) - return F.resize(image, (new_h, new_w), InterpolationMode.BICUBIC, antialias=True) + size = vllm_vae_size(h, w, stride, max_img_size, min_img_size) + return F.resize(image, size, InterpolationMode.BICUBIC, antialias=True) diff --git a/mstar/model/bagel/submodules.py b/mstar/model/bagel/submodules.py index eaa753f885..590305bcfd 100644 --- a/mstar/model/bagel/submodules.py +++ b/mstar/model/bagel/submodules.py @@ -31,6 +31,7 @@ get_flattened_position_ids_extrapolate, patchify, vllm_vae_resize, + vllm_vae_size, vllm_vit_resize, ) from mstar.model.bagel.components.vit_encoder import VIT_ATTN @@ -131,11 +132,7 @@ def vae_latents( ) -> int: """How many latents the VAE encodes an image of this size into. Its ``prepare_inputs`` asserts the same count.""" if image_preprocess == "vllm": - # the size `vllm_vae_resize` resizes to - stride, max_size = latent_downsample, max_latent_size * latent_downsample - scale = max(min(max_size / max(width, height), 1.0), min(256, max_size) / min(width, height)) - height = min(max(stride, round(height * scale / stride) * stride), max_size) - width = min(max(stride, round(width * scale / stride) * stride), max_size) + height, width = vllm_vae_size(height, width, latent_downsample, max_latent_size * latent_downsample) else: height, width = _VAE_TRANSFORM.resize_transform.target_size(height, width) return (height // latent_downsample) * (width // latent_downsample) diff --git a/test/modular/test_bagel_cfg_parallel.py b/test/modular/test_bagel_cfg_parallel.py index 0b51fcb616..481ebe545a 100644 --- a/test/modular/test_bagel_cfg_parallel.py +++ b/test/modular/test_bagel_cfg_parallel.py @@ -92,7 +92,7 @@ def test_pd_disaggregation_publishes_main_only_at_handoff_boundaries(): model._has_llm_disaggregation = True args = { "default": SimpleNamespace( - full_metadata=SimpleNamespace(kwargs={"requires_cfg": True}) + full_metadata=SimpleNamespace(kwargs={"requires_cfg": True, "target_output": "image"}) ) } diff --git a/test/modular/test_multimodal_ordering.py b/test/modular/test_multimodal_ordering.py index 7b8a32a276..21efd8c8c5 100644 --- a/test/modular/test_multimodal_ordering.py +++ b/test/modular/test_multimodal_ordering.py @@ -657,6 +657,8 @@ def _laid_out(bagel, modalities: list[str], output_modalities=("text",), tensors bagel.config.vit_config = SimpleNamespace(patch_size=PATCH_SIZE) bagel.config.vit_max_num_patch_per_side = MAX_PATCHES_PER_SIDE bagel.config.latent_downsample, bagel.config.max_latent_size = 16, 64 + bagel.config.cfg_text_scale, bagel.config.cfg_img_scale, bagel.config.cfg_interval = 4.0, 1.0, [0.4, 1.0] + bagel.config.cfg_renorm_type, bagel.config.cfg_renorm_min = "global", 0.0 texts = iter(["what is", "in this picture", "and in this one"]) parts, images = [], [] for modality in modalities: From fb6d7182eaa125971b10726ca77cdece68167004 Mon Sep 17 00:00:00 2001 From: Gaurav Shah Date: Wed, 21 Oct 2026 05:50:51 +0000 Subject: [PATCH 08/26] model: chatterbox declares its KV layout --- mstar/model/chatterbox/chatterbox_model.py | 56 ++++++++++++++++++---- mstar/model/chatterbox/submodules.py | 16 ++++++- test/modular/test_chatterbox_model.py | 10 ++-- 3 files changed, 66 insertions(+), 16 deletions(-) diff --git a/mstar/model/chatterbox/chatterbox_model.py b/mstar/model/chatterbox/chatterbox_model.py index 57788760a1..5450d6038a 100644 --- a/mstar/model/chatterbox/chatterbox_model.py +++ b/mstar/model/chatterbox/chatterbox_model.py @@ -59,7 +59,7 @@ ) from mstar.graph.base import GraphEdge, GraphNode, GraphSection, Loop, Sequential, TensorPointerInfo from mstar.graph.special_destinations import EMIT_TO_CLIENT, EMPTY_DESTINATION -from mstar.model.base import ForwardPassArgs, Model, TensorAndMetadata +from mstar.model.base import ForwardPassArgs, KVLabel, LabelLayout, Model, ProcessPromptOutput, Span, TensorAndMetadata from mstar.model.chatterbox.components.audio_frontend import resample, trim_silence from mstar.model.chatterbox.config import ( COND_LABEL, @@ -425,6 +425,14 @@ def resolve_generation_kwargs(self, model_kwargs: dict | None) -> dict[str, Any] def get_max_output_tokens(self, **model_kwargs: Any) -> int: return self.resolve_generation_kwargs(model_kwargs)["max_new_tokens"] + def get_loop_caps(self, **model_kwargs: Any) -> dict[str, int]: + # The stop check counts the prefill's token, so it allows max_new_tokens - 1 iterations. + return {"decode_loop": max(1, self.get_max_output_tokens(**model_kwargs) - 1)} + + def kv_labels(self) -> dict[str, dict[str, KVLabel]]: + t3 = KVLabel(growth=(("decode", "decode_loop"),)) + return {T3_KV: {COND_LABEL: t3, UNCOND_LABEL: t3}} + # ----------------------------------------------------------------------- # Graph # ----------------------------------------------------------------------- @@ -623,7 +631,7 @@ def process_prompt( output_modalities: list[str], tensors: NameToTensorList | None = None, **kwargs: Any, - ) -> NameToTensorList: + ) -> ProcessPromptOutput: if not prompt or not prompt.strip(): raise ValueError("Chatterbox requires a non-empty text prompt") if any(m not in ("text", "audio") for m in input_modalities): @@ -631,7 +639,7 @@ def process_prompt( if set(output_modalities) != {"audio"}: raise ValueError("Chatterbox produces audio output only") # a bad knob fails here as a 400; past this point it would hang the request - self.resolve_generation_kwargs(kwargs) + knobs = self.resolve_generation_kwargs(kwargs) text_ids = self._tokenize(prompt, kwargs.get("language_id")) if text_ids.numel() > self.config.max_text_tokens: @@ -645,16 +653,44 @@ def process_prompt( uploaded = (tensors or {}).get("audio_inputs") or [] if len(uploaded) > 1: raise ValueError("Give one reference clip per request") + wav = None if uploaded: wav = uploaded[0] - elif voice in _BUILTIN_VOICE_ALIASES: - return out - else: + elif voice not in _BUILTIN_VOICE_ALIASES: wav = self.load_audio(str(self._preset_voice_path(str(voice))), "cpu").data - wav = self._prepare_reference(wav) - out[REF_AUDIO] = [wav] - out[VOICE_KEY] = [voice_key_for(wav)] - return out + if wav is not None: + wav = self._prepare_reference(wav) + out[REF_AUDIO] = [wav] + out[VOICE_KEY] = [voice_key_for(wav)] + layout = self._t3_layout(text_ids, wav) + guided = knobs["cfg_weight"] > 0 + labels = (COND_LABEL, UNCOND_LABEL) if guided else (COND_LABEL,) + return ProcessPromptOutput(out, {"kv_layout": {T3_KV: dict.fromkeys(labels, layout)}}) + + def _t3_layout(self, text_ids: torch.Tensor, reference: torch.Tensor | None) -> LabelLayout: + """What T3 writes on each label the request opens: ``[cond | text | BOS]``, then a row per decode step.""" + from mstar.model.chatterbox.submodules import reference_tokens + + t3 = self.config.t3 + walk = "prefill" if reference is None else "prefill_voice" + cond = t3.cond_len + if not t3.use_perceiver_resampler: + # `cond_len` counts the most prompt tokens there can be, so count the raw prompt tokens written instead + prompt = ( + self._builtin_voice()["t3"]["cond_prompt_speech_tokens"].numel() if reference is None + else reference_tokens(reference.numel(), self.config) + ) + cond += prompt - t3.speech_cond_prompt_len + text = text_ids.numel() + bos = 2 if t3.duplicate_bos_in_prefill else 1 + return LabelLayout( + spans=( + Span("rows", walk, cond, cond, PROMPT_TOKENS), + Span("ids", walk, text, text, TEXT_INPUTS), + Span("rows", walk, bos, bos, PREV_TOKEN), + ), + growth=(("decode_loop", 1),), + ) def _tokenize(self, prompt: str, language_id: str | None) -> torch.Tensor: """Text ids for the variant's front end; ``language_id`` (a request knob) diff --git a/mstar/model/chatterbox/submodules.py b/mstar/model/chatterbox/submodules.py index bf19e96a1f..cdf85b21a4 100644 --- a/mstar/model/chatterbox/submodules.py +++ b/mstar/model/chatterbox/submodules.py @@ -41,6 +41,7 @@ SlotLease, SubmoduleStep, ) +from mstar.model.chatterbox.components.s3_tokenizer import AudioEncoderV2 from mstar.model.chatterbox.components.s3gen import FlowRow from mstar.model.chatterbox.config import ( CFG_LABEL, @@ -111,6 +112,14 @@ def put(self, key: int, value) -> None: # =========================================================================== +def reference_tokens(samples: int, config: ChatterboxConfig) -> int: + """How many S3 prompt tokens `VoiceEncoderSubmodule.condition` makes of a ``samples``-long + reference clip. It asserts the same count.""" + wav16 = min(-(-samples * S3_SR // config.sample_rate), round(config.enc_cond_seconds * S3_SR)) + frames = min(wav16 // config.s3_tokenizer.hop_size, 4 * config.t3.speech_cond_prompt_len) + return AudioEncoderV2._conv_out_len(AudioEncoderV2._conv_out_len(frames, 2), 2) + + class VoiceEncoderSubmodule(NodeSubmodule): """Speaker embedding (LSTM voice encoder over the whole clip) and the S3 prompt tokens (first ``enc_cond_seconds`` of the clip), as the reference @@ -144,7 +153,12 @@ def condition(self, wav24: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: tokens, lens = self.s3_tokenizer( [prompt_wav], max_len=self.config.t3.speech_cond_prompt_len, ) - return speaker_emb, tokens[0, : int(lens[0])] + count = int(lens[0]) + assert count == reference_tokens(wav24.numel(), self.config), ( + f"the S3 tokenizer made {count} prompt tokens of a {wav24.numel()}-sample clip, " + "not the count `reference_tokens` gives a layout to size it by" + ) + return speaker_emb, tokens[0, :count] def forward( self, graph_walk: str, engine_inputs: ModelInputsFromEngine, diff --git a/test/modular/test_chatterbox_model.py b/test/modular/test_chatterbox_model.py index 710575d3f6..783df99a17 100644 --- a/test/modular/test_chatterbox_model.py +++ b/test/modular/test_chatterbox_model.py @@ -296,10 +296,10 @@ def test_speech_adapter_maps_voice_and_reference_audio(tmp_path): def test_process_prompt_builtin_voice_emits_text_only(): model = _make_model() - out = model.process_prompt("hello", ["text"], ["audio"], voice="default") + out = model.process_prompt("hello", ["text"], ["audio"], voice="default").new_input_tensors assert set(out) == {TEXT_INPUTS} assert out[TEXT_INPUTS][0].tolist() == [1, 2, 3, 4, 5] - out = model.process_prompt("hello", ["text"], ["audio"]) + out = model.process_prompt("hello", ["text"], ["audio"]).new_input_tensors assert set(out) == {TEXT_INPUTS} @@ -308,7 +308,7 @@ def test_process_prompt_reference_audio_is_keyed_and_capped(): wav = torch.rand(40 * 24000) - 0.5 out = model.process_prompt( "hello", ["text", "audio"], ["audio"], tensors={"audio_inputs": [wav]}, - ) + ).new_input_tensors assert set(out) == {TEXT_INPUTS, REF_AUDIO, VOICE_KEY} assert out[REF_AUDIO][0].numel() == 30 * 24000 # capped at 30 s assert out[VOICE_KEY][0].dtype == torch.long @@ -353,7 +353,7 @@ def test_process_prompt_rejects_unusable_reference(wav, message): def test_process_prompt_accepts_short_but_usable_reference(): out = _make_model().process_prompt( "hello", ["text", "audio"], ["audio"], tensors={"audio_inputs": [_tone(0.6)]}, - ) + ).new_input_tensors assert out[REF_AUDIO][0].numel() == int(0.6 * 24000) @@ -522,7 +522,7 @@ def test_generation_kwargs_accept_boundary_values(): assert top_ks == [vocab - 1, vocab, vocab] # other request fields ride along in the same kwargs and are left alone out = model.process_prompt("hello", ["text"], ["audio"], voice="default", language_id=None, seed=7) - assert set(out) == {TEXT_INPUTS} + assert set(out.new_input_tensors) == {TEXT_INPUTS} def test_max_new_tokens_limit_bounds_requests_and_the_default(): From c62d8e79e5d378983ddfcd51cdb9035b887e0515 Mon Sep 17 00:00:00 2001 From: Gaurav Shah Date: Wed, 21 Oct 2026 05:50:51 +0000 Subject: [PATCH 09/26] model: cosmos3 declares its KV layout --- mstar/model/cosmos3/cosmos3_model.py | 88 ++++++++++++++++++++--- mstar/model/cosmos3/submodules.py | 74 ++++++++++++------- mstar/model/cosmos3/tests/test_edge.py | 4 +- mstar/model/cosmos3/tests/test_serving.py | 2 +- mstar/model/cosmos3/tests/test_sound.py | 14 ++-- 5 files changed, 134 insertions(+), 48 deletions(-) diff --git a/mstar/model/cosmos3/cosmos3_model.py b/mstar/model/cosmos3/cosmos3_model.py index f552a67b47..6fc14c7021 100644 --- a/mstar/model/cosmos3/cosmos3_model.py +++ b/mstar/model/cosmos3/cosmos3_model.py @@ -87,7 +87,7 @@ TensorPointerInfo, ) from mstar.graph.special_destinations import EMIT_TO_CLIENT, EMPTY_DESTINATION -from mstar.model.base import ForwardPassArgs, Model, ProcessPromptOutput +from mstar.model.base import ForwardPassArgs, KVLabel, LabelLayout, Model, ProcessPromptOutput, Span from mstar.model.cosmos3 import constants from mstar.model.cosmos3.components.packing import ACTION_MODES, resolve_action_domain_id from mstar.model.cosmos3.config import Cosmos3Config @@ -113,6 +113,9 @@ Cosmos3VAEDecoderSubmodule, Cosmos3VAEEncoderSubmodule, Cosmos3VisionEncoderSubmodule, + frame_tokens, + latent_frames, + sound_frames, ) from mstar.model.media_limits import VideoBudget from mstar.model.multimodal import TEXT, PromptPart, check_attachments, parts_from_modalities @@ -347,6 +350,17 @@ def get_node_resources(self) -> list[NodeResourceSpec]: ) return specs + def kv_labels(self) -> dict[str, dict[str, KVLabel]]: + uncond = KVLabel() + if not self._reasoner_enabled(): + return {KV_CACHE: {COND_LABEL: uncond, UNCOND_LABEL: uncond}} + main = KVLabel(growth=((self.REASONER_DECODE_WALK, REASONER_DECODE_LOOP),)) + return {KV_CACHE: {COND_LABEL: main, UNCOND_LABEL: uncond}} + + def get_loop_caps(self, **model_kwargs) -> dict[str, int]: + # The reasoner's stop check allows max_tokens - 1 iterations, and a loop runs at least one. + return {REASONER_DECODE_LOOP: max(1, self.get_max_output_tokens(**model_kwargs) - 1)} + def get_request_resource_configs( self, partition_fwd_args: dict[str, ForwardPassArgs], @@ -832,7 +846,9 @@ def process_prompt( prompt, input_modalities, tensors or {}, prompt_parts, input_metadata or {}, kwargs, ) if prompt is None: - return self._action_prompt_output({}, kwargs) + p = self._resolve_gen_params(kwargs, input_modalities, output_modalities) + layout = self._gen_layout(p, [], [], input_modalities, tensors or {}) + return self._action_prompt_output({}, kwargs, {"kv_layout": {KV_CACHE: layout}}) if self.tokenizer is None: # Tokenizer-less fallback used by structural unit tests. return self._action_prompt_output({ @@ -865,26 +881,70 @@ def process_prompt( add_duration_template=allow_templates and bool(kwargs.get("use_duration_template", False)), max_sequence_length=p["max_sequence_length"], ) + layout = self._gen_layout(p, cond_ids, uncond_ids, input_modalities, tensors or {}) return self._action_prompt_output({ "text_inputs": [ torch.tensor(cond_ids, dtype=torch.long), torch.tensor(uncond_ids, dtype=torch.long), ] - }, kwargs) + }, kwargs, {"kv_layout": {KV_CACHE: layout}}) + + def _gen_layout( + self, p: dict, cond_ids: list[int], uncond_ids: list[int], input_modalities: list[str], tensors: dict, + ) -> dict[str, LabelLayout]: + """Each guidance branch's writes: its prompt at prefill, one denoise step's + latents as scratch, and in window mode ``kv`` each window's latents once + they are clean, under the schedule's retention.""" + config = self.config + if "video" in input_modalities and "video_inputs" in tensors: + prefill = self.PREFILL_COND_VIDEO_WALK + elif p.get("has_image_condition") and "image_inputs" in tensors: + prefill = self.PREFILL_COND_WALK + else: + prefill = self.PREFILL_WALK + per_frame = frame_tokens(p["height"], p["width"], config) + schedule, windows = None, () + if p.get("window_mode") is not None: + schedule = WindowSchedule( + p["total_latent_units"], p["window_latent_units"], + context_units=p["context_latent_units"], overlap_units=p["overlap_latent_units"], + ) + scratch = max(plan.units for plan in schedule.windows()) * per_frame + if p["window_mode"] == "kv": + windows = tuple( + Span("rows", self.VIDEO_GEN_AR_WALK, length, length, "latents") + for length in ((plan.commit_end - plan.commit_start) * per_frame for plan in schedule.windows()) + ) + elif p.get("action_mode") is not None: + action_chunk = int(p["action_chunk_size"]) + scratch = latent_frames(int(p.get("num_frames") or action_chunk + 1), config) * per_frame + action_chunk + else: + scratch = latent_frames(p["num_frames"], config) * per_frame + if p.get("generate_sound"): + scratch += sound_frames(p, config)[1] + layouts = {UNCOND_LABEL: LabelLayout()} + branches = [(COND_LABEL, cond_ids), (UNCOND_LABEL, uncond_ids)] + for label, ids in branches if float(p["guidance_scale"]) != 1.0 else branches[:1]: + layouts[label] = LabelLayout( + spans=(Span("ids", prefill, len(ids), len(ids), "text_inputs"), *windows), scratch=scratch, + retention=schedule.retention(per_frame, len(ids)) if windows else None, + ) + return layouts @staticmethod - def _action_prompt_output(tensors: NameToTensorList, kwargs: dict): + def _action_prompt_output(tensors: NameToTensorList, kwargs: dict, metadata: dict | None = None): """Put conditioning values on a tensor edge; only their shape is metadata.""" + metadata = dict(metadata or {}) mode = str(kwargs.get("action_mode", "")).strip().lower() - if mode != "forward_dynamics" or kwargs.get("action") is None: - return tensors - action = torch.as_tensor(kwargs["action"]) - tensors["action_inputs"] = [action] - return ProcessPromptOutput(tensors, {"action_shape": list(action.shape)}) + if mode == "forward_dynamics" and kwargs.get("action") is not None: + action = torch.as_tensor(kwargs["action"]) + tensors["action_inputs"] = [action] + metadata["action_shape"] = list(action.shape) + return ProcessPromptOutput(tensors, metadata) if metadata else tensors def _process_reasoner_prompt( self, prompt, input_modalities, tensors, prompt_parts, input_metadata, model_kwargs, - ) -> NameToTensorList: + ) -> ProcessPromptOutput: """Render the chat prompt, preprocess its media, expand the placeholders and compute the mRoPE positions — everything the reasoner prefill needs besides the encoder pass. @@ -969,7 +1029,13 @@ def _process_reasoner_prompt( if ordered_patches: out["pixel_values"] = [torch.cat(ordered_patches, dim=0)] out["vision_grid_thw"] = [torch.tensor(ordered_grids, dtype=torch.long)] - return out + walk = self.REASONER_PREFILL_VISION_WALK if ordered_patches else self.REASONER_PREFILL_WALK + main = LabelLayout( + # media rows fill the expanded placeholders, so the span is rows, not ids + spans=(Span("rows", walk, ids.numel(), ids.numel(), "text_inputs"),), + growth=((REASONER_DECODE_LOOP, 1),), + ) + return ProcessPromptOutput(out, {"kv_layout": {KV_CACHE: {REASONER_LABEL: main, UNCOND_LABEL: LabelLayout()}}}) def load_video(self, filepath: str, device: str): """Decode a conditioning / reasoner video to ``[T, C, H, W]`` in [0, 1]. diff --git a/mstar/model/cosmos3/submodules.py b/mstar/model/cosmos3/submodules.py index e899794de2..0279f62b3e 100644 --- a/mstar/model/cosmos3/submodules.py +++ b/mstar/model/cosmos3/submodules.py @@ -160,6 +160,35 @@ REASONER_LABEL = "main" +def latent_frames(num_frames: int, config) -> int: + return 1 if num_frames == 1 else 1 + (num_frames - 1) // config.vae.scale_factor_temporal + + +def frame_tokens(height: int, width: int, config) -> int: + """Generation tokens per latent frame of a ``height`` x ``width`` request. `_build_static` asserts the count.""" + s, p = config.vae.scale_factor_spatial, config.latent_patch_size + return math.ceil(height // s / p) * math.ceil(width // s / p) + + +def sound_frames(md: dict, config) -> tuple[int, int]: + """Resolve a sound request's target sample count and latent frame count. + + The sound duration defaults to the video duration (``num_frames / fps``, + the request's ``sound_duration`` overrides), the target sample count is + that duration at the AVAE sample rate, and the latent frame count rounds + it up to whole tokenizer hops (the decoded tail past the target is + trimmed by the audio decode node).""" + num_frames = int(md.get("num_frames", 1)) + fps = float(md.get("fps", 24.0)) + duration = md.get("sound_duration") + if duration is None: + duration = num_frames / fps + duration = max(float(duration), 1.0 / max(fps, 1.0)) + target_samples = max(1, int(round(duration * config.sound_sample_rate))) + hop = int(round(config.sound_sample_rate / config.sound_latent_fps)) + return target_samples, max(1, math.ceil(target_samples / hop)) + + def _loop_step(fwd_info, graph_walk: str, time_index: torch.Tensor) -> int: """The denoise iteration a step runs, from the engine's loop counter. @@ -397,8 +426,7 @@ def _latent_shape( self, height: int, width: int, num_frames: int = 1 ) -> tuple[int, int, int, int, int]: s = self.config.vae.scale_factor_spatial - t = 1 if num_frames == 1 else 1 + (num_frames - 1) // self.config.vae.scale_factor_temporal - return (1, self.config.latent_channel, t, height // s, width // s) + return (1, self.config.latent_channel, latent_frames(num_frames, self.config), height // s, width // s) def _build_static( self, ids: list[int], height: int, width: int, num_frames: int, @@ -407,14 +435,19 @@ def _build_static( noisy_frames: list[int] | None = None, start_frame_offset: int = 0, ) -> dict: + latent_shape = self._latent_shape(height, width, num_frames) static = build_static_inputs( - list(ids), self._latent_shape(height, width, num_frames), self.config, + list(ids), latent_shape, self.config, self.config.vae.scale_factor_temporal, fps, device, has_image_condition=has_image_condition, sound_latent_frames=sound_latent_frames, noisy_frames=noisy_frames, start_frame_offset=start_frame_offset, ) + assert static["num_vision_tokens"] == latent_shape[2] * frame_tokens(height, width, self.config), ( + f"packing made {static['num_vision_tokens']} vision tokens of {latent_shape}, " + "not the count `frame_tokens` gives a layout to size it by" + ) # proj_out runs on the generation token block, so shift the joint-sequence # mse indexes to be relative to the generation tokens. static["mse_gen_indexes"] = static["vision_mse_loss_indexes"] - static["und_len"] @@ -422,24 +455,6 @@ def _build_static( static["sound_mse_gen_indexes"] = static["sound_mse_loss_indexes"] - static["und_len"] return static - def _resolve_sound_frames(self, md: dict) -> tuple[int, int]: - """Resolve a sound request's target sample count and latent frame count. - - The sound duration defaults to the video duration (``num_frames / fps``, - the request's ``sound_duration`` overrides), the target sample count is - that duration at the AVAE sample rate, and the latent frame count rounds - it up to whole tokenizer hops (the decoded tail past the target is - trimmed by the audio decode node).""" - num_frames = int(md.get("num_frames", 1)) - fps = float(md.get("fps", 24.0)) - duration = md.get("sound_duration") - if duration is None: - duration = num_frames / fps - duration = max(float(duration), 1.0 / max(fps, 1.0)) - target_samples = max(1, int(round(duration * self.config.sound_sample_rate))) - hop = int(round(self.config.sound_sample_rate / self.config.sound_latent_fps)) - return target_samples, max(1, math.ceil(target_samples / hop)) - def _new_scheduler(self, num_inference_steps: int, device, use_karras_sigma=None, flow_shift=None): if self.config.distilled_sigmas: return self._new_distilled_scheduler(device) @@ -529,10 +544,15 @@ def _build_action_static( self, ids: list[int], height: int, width: int, num_frames: int, action_chunk: int, mode: str, fps: float, action_fps: float, action_offset: int, device, ) -> dict: + latent_shape = self._latent_shape(height, width, num_frames) static = build_action_static_inputs( - list(ids), self._latent_shape(height, width, num_frames), action_chunk, mode, + list(ids), latent_shape, action_chunk, mode, self.config, self.config.vae.scale_factor_temporal, fps, action_fps, action_offset, device, ) + assert static["num_vision_tokens"] == latent_shape[2] * frame_tokens(height, width, self.config), ( + f"packing made {static['num_vision_tokens']} vision tokens of {latent_shape}, " + "not the count `frame_tokens` gives a layout to size it by" + ) # proj_out runs on the generation token block; shift the joint-sequence # mse indexes to be relative to the [vision | action] generation tokens. static["mse_gen_indexes"] = static["vision_mse_loss_indexes"] - static["und_len"] @@ -597,21 +617,21 @@ def _prepare_prefill(self, fwd_info, inputs, device) -> ARNodeInputs: # Opt-in sound: append a jointly denoised AVAE-latent band to the # generation block. Video-only (single-frame image and action requests # are rejected at request resolution). - sound_frames = None + num_sound = None if md.get("generate_sound"): if num_frames <= 1: raise ValueError("Cosmos3 sound generation requires a video request (num_frames > 1)") - _, sound_frames = self._resolve_sound_frames(md) + _, num_sound = sound_frames(md, self.config) cond = self._build_static( cond_ids, height, width, num_frames, fps, has_image_condition, device, - sound_latent_frames=sound_frames, noisy_frames=noisy_frames, + sound_latent_frames=num_sound, noisy_frames=noisy_frames, ) uncond = None if uncond_ids is not None: uncond = self._build_static( uncond_ids, height, width, num_frames, fps, has_image_condition, device, - sound_latent_frames=sound_frames, noisy_frames=noisy_frames, + sound_latent_frames=num_sound, noisy_frames=noisy_frames, ) node_inputs = self._get_prefill_node_inputs(cond, uncond) @@ -627,7 +647,7 @@ def _prepare_prefill(self, fwd_info, inputs, device) -> ARNodeInputs: use_karras_sigma=md.get("use_karras_sigma") ), latent_shape=self._latent_shape(height, width, num_frames), - num_sound=sound_frames, + num_sound=num_sound, ) # Video-to-video: pin the requested latent frames (the complement starts # from noise). Only the pin mask is derived here — the clean latents are diff --git a/mstar/model/cosmos3/tests/test_edge.py b/mstar/model/cosmos3/tests/test_edge.py index 42885b49ed..59b2396f13 100644 --- a/mstar/model/cosmos3/tests/test_edge.py +++ b/mstar/model/cosmos3/tests/test_edge.py @@ -595,7 +595,7 @@ def test_edge_reasoner_video_prompt() -> None: prompt_parts=[PromptPart("video", None, 0), PromptPart("text", "What happens?")], input_metadata={"video_inputs": [{"average_fps": 10.0, "num_frames": 30}]}, enable_thinking=False, - ) + ).new_input_tensors ids = out["text_inputs"][0] grid = out["vision_grid_thw"][0] assert grid.shape == (1, 3) and int(grid[0, 0]) == 6 @@ -616,7 +616,7 @@ def test_edge_reasoner_video_prompt() -> None: None, ["video", "text"], ["text"], tensors={"video_inputs": [video]}, prompt_parts=[PromptPart("video", None, 0), PromptPart("text", "What happens?")], input_metadata={"video_inputs": [{"average_fps": 10.0}]}, video_num_frames=4, - ) + ).new_input_tensors assert int(out2["vision_grid_thw"][0][0, 0]) == 4 diff --git a/mstar/model/cosmos3/tests/test_serving.py b/mstar/model/cosmos3/tests/test_serving.py index 1970fbe825..431edce1e7 100644 --- a/mstar/model/cosmos3/tests/test_serving.py +++ b/mstar/model/cosmos3/tests/test_serving.py @@ -182,7 +182,7 @@ def test_process_prompt_emits_cond_and_uncond() -> None: eos = model.tokenizer.eos_token_id out = model.process_prompt("a red cube on a table", ["text"], ["image"], tensors={}, size="256x256") - ti = out["text_inputs"] + ti = out.new_input_tensors["text_inputs"] assert len(ti) == 2, "t2i must emit a conditional and unconditional prompt" cond, uncond = ti[0].tolist(), ti[1].tolist() assert cond[-2:] == [eos, sog] diff --git a/mstar/model/cosmos3/tests/test_sound.py b/mstar/model/cosmos3/tests/test_sound.py index 9759ef33fe..d6a162bc98 100644 --- a/mstar/model/cosmos3/tests/test_sound.py +++ b/mstar/model/cosmos3/tests/test_sound.py @@ -23,7 +23,7 @@ from mstar.model.cosmos3.components.packing import build_static_inputs from mstar.model.cosmos3.config import Cosmos3Config from mstar.model.cosmos3.constants import VIDEO_SOUND_GEN_WALK -from mstar.model.cosmos3.submodules import Cosmos3DiTSubmodule +from mstar.model.cosmos3.submodules import sound_frames from mstar.model.cosmos3.tests.test_engine_cache import ( _engine_resources, _forward_step, @@ -70,18 +70,18 @@ def test_sound_segment_layout() -> None: def test_sound_frame_math() -> None: - dit = Cosmos3DiTSubmodule(transformer=None, config=Cosmos3Config()) + config = Cosmos3Config() # Video-duration default: 93 frames @ 16 fps = 5.8125 s -> 279000 samples # -> ceil(279000 / 1920) = 146 latent frames. - target, frames = dit._resolve_sound_frames({"num_frames": 93, "fps": 16.0}) + target, frames = sound_frames({"num_frames": 93, "fps": 16.0}, config) assert (target, frames) == (279000, 146) # 17 frames @ 24 fps = 0.7083 s -> 34000 samples -> 18 latent frames. - target, frames = dit._resolve_sound_frames({"num_frames": 17, "fps": 24.0}) + target, frames = sound_frames({"num_frames": 17, "fps": 24.0}, config) assert (target, frames) == (34000, 18) # Explicit sound_duration wins; sub-frame durations clamp to one video frame. - target, frames = dit._resolve_sound_frames({"num_frames": 17, "fps": 24.0, "sound_duration": 2.0}) + target, frames = sound_frames({"num_frames": 17, "fps": 24.0, "sound_duration": 2.0}, config) assert (target, frames) == (96000, 50) - target, frames = dit._resolve_sound_frames({"num_frames": 17, "fps": 24.0, "sound_duration": 0.0}) + target, frames = sound_frames({"num_frames": 17, "fps": 24.0, "sound_duration": 0.0}, config) assert target == round(48000 / 24.0) and frames == 2 @@ -290,7 +290,7 @@ def _sound_scenario(): use_system_prompt=False, add_resolution_template=False, add_duration_template=False, ) lat_t = 1 + (FRAMES - 1) // mpipe.vae_scale_temporal - _, t_s = model.get_submodule("dit")._resolve_sound_frames({"num_frames": FRAMES, "fps": 24.0}) + _, t_s = sound_frames({"num_frames": FRAMES, "fps": 24.0}, model.get_submodule("dit").config) gen = torch.Generator(device=device).manual_seed(SEED) init = torch.randn((1, 48, lat_t, H // 16, W // 16), generator=gen, device=device, dtype=dtype) # The fused pipeline draws its sound noise from `generator` after the video From ce8cd23943b29c20502da35b5c2cc584f0e2dfe7 Mon Sep 17 00:00:00 2001 From: Gaurav Shah Date: Wed, 21 Oct 2026 06:04:09 +0000 Subject: [PATCH 10/26] model: fish_s2_pro declares its KV layout --- mstar/model/fish_s2_pro/fish_s2_pro_model.py | 17 ++++++++++++++++- test/modular/test_fish_s2_pro_model.py | 6 ++++-- 2 files changed, 20 insertions(+), 3 deletions(-) diff --git a/mstar/model/fish_s2_pro/fish_s2_pro_model.py b/mstar/model/fish_s2_pro/fish_s2_pro_model.py index 301e647caa..1c33283b5c 100644 --- a/mstar/model/fish_s2_pro/fish_s2_pro_model.py +++ b/mstar/model/fish_s2_pro/fish_s2_pro_model.py @@ -56,7 +56,7 @@ ) from mstar.graph.base import GraphEdge, GraphNode, GraphSection, Loop, Sequential, TensorPointerInfo from mstar.graph.special_destinations import EMIT_TO_CLIENT, EMPTY_DESTINATION -from mstar.model.base import ForwardPassArgs, Model, ProcessPromptOutput, TensorAndMetadata +from mstar.model.base import ForwardPassArgs, KVLabel, LabelLayout, Model, ProcessPromptOutput, Span, TensorAndMetadata from mstar.model.fish_s2_pro.config import ( ALLOWED_TEXT_SPECIAL_TOKENS, AR_ATTN, @@ -382,6 +382,12 @@ def process_prompt( metadata: dict[str, Any] = {PROMPT_LAYOUT_KEY: [prefix_len, ref_frames]} if clips: metadata[REF_CLIP_KEY] = clip_digest(clips[0]) + # a clone prefills on its own walk, which splices the reference frames into the prompt's rows + clone = "audio" in input_modalities + rows = len(ids) + ref_frames + walk = "ar_prefill_clone" if clone else "ar_prefill" + prefill = Span("rows" if clone else "ids", walk, rows, rows, "prompt_ids") + metadata["kv_layout"] = {AR_KV: {"main": LabelLayout(spans=(prefill,), growth=(("ar_decode_loop", 1),))}} return ProcessPromptOutput({"prompt_ids": [torch.tensor(ids, dtype=torch.long)]}, metadata) # ----------------------------------------------------------------------- @@ -521,6 +527,15 @@ def get_max_output_tokens(self, **model_kwargs: Any) -> int: ) return requested + def get_loop_caps(self, **model_kwargs: Any) -> dict[str, int]: + # the stop check counts the prefill's frame, so decode runs one iteration fewer than max tokens + return {"ar_decode_loop": max(1, self.get_max_output_tokens(**model_kwargs) - 1)} + + def kv_labels(self) -> dict[str, dict[str, KVLabel]]: + return {AR_KV: {"main": KVLabel( + growth=(("ar_decode", "ar_decode_loop"),), + )}} + def get_output_sample_rate(self, modality: str = "audio") -> int: return self.config.codec.sample_rate diff --git a/test/modular/test_fish_s2_pro_model.py b/test/modular/test_fish_s2_pro_model.py index fc2e75cd61..f019d79d11 100644 --- a/test/modular/test_fish_s2_pro_model.py +++ b/test/modular/test_fish_s2_pro_model.py @@ -609,10 +609,12 @@ def test_process_prompt_lays_out_clones_and_validates_requests(): ids, prefix_len = model.build_prompt("hi", "r", 21) assert list(out.new_input_tensors) == ["prompt_ids"] assert out.new_input_tensors["prompt_ids"][0].tolist() == ids - assert out.metadata == {"prompt_layout": [prefix_len, 21], "ref_clip_key": clip_digest(clip)} + assert {key: value for key, value in out.metadata.items() if key != "kv_layout"} == { + "prompt_layout": [prefix_len, 21], "ref_clip_key": clip_digest(clip), + } plain = model.process_prompt("hi", ["text"], ["audio"]) - assert plain.metadata == {"prompt_layout": [0, 0]} + assert {key: value for key, value in plain.metadata.items() if key != "kv_layout"} == {"prompt_layout": [0, 0]} for prompt, inputs, outputs, tensors, match in [ (" ", ["text"], ["audio"], None, "non-empty"), ("hi", ["image"], ["audio"], None, "takes text input"), From 4fe59637205ab5b2deab3c8d50811d0d6fe5c3ea Mon Sep 17 00:00:00 2001 From: Gaurav Shah Date: Wed, 21 Oct 2026 05:50:51 +0000 Subject: [PATCH 11/26] model: higgs_audio declares its KV layout --- mstar/model/higgs_audio/higgs_audio_model.py | 26 +++++++++++++++++--- 1 file changed, 22 insertions(+), 4 deletions(-) diff --git a/mstar/model/higgs_audio/higgs_audio_model.py b/mstar/model/higgs_audio/higgs_audio_model.py index 167ed92c1c..e3a583fda7 100644 --- a/mstar/model/higgs_audio/higgs_audio_model.py +++ b/mstar/model/higgs_audio/higgs_audio_model.py @@ -49,7 +49,7 @@ ) from mstar.graph.base import GraphEdge, GraphNode, GraphSection, Loop, Sequential, TensorPointerInfo from mstar.graph.special_destinations import EMIT_TO_CLIENT -from mstar.model.base import ForwardPassArgs, Model +from mstar.model.base import ForwardPassArgs, KVLabel, LabelLayout, Model, ProcessPromptOutput, Span from mstar.model.higgs_audio.config import ATTN, KV_CACHE, ROPE, SAMPLER, HiggsAudioModelConfig from mstar.model.submodule_base import NodeSubmodule @@ -171,6 +171,15 @@ def get_max_output_tokens(self, **model_kwargs): # Reference transcribe.py generates up to 1024 new tokens. return model_kwargs.get("max_output_tokens", 1024) + def get_loop_caps(self, **model_kwargs) -> dict[str, int]: + # the decode stop check ends the loop after max_tokens iterations + return {"decode_loop": self.get_max_output_tokens(**model_kwargs)} + + def kv_labels(self) -> dict[str, dict[str, KVLabel]]: + return {KV_CACHE: {"main": KVLabel( + growth=(("decode", "decode_loop"),), + )}} + def get_graph_walk_graphs(self) -> dict[str, GraphSection]: llm_prefill_outputs = [ GraphEdge( @@ -374,7 +383,7 @@ def process_prompt( output_modalities: list[str], tensors: NameToTensorList | None = None, **kwargs, - ) -> NameToTensorList: + ) -> ProcessPromptOutput: """Chunk the audio, extract per-chunk mels, and split the ChatML prompt at the audio splice point. @@ -425,14 +434,23 @@ def enc(s: str) -> list[int]: if not kwargs.get("enable_thinking", False): post_ids += enc("\n\n\n\n") - return { + audio = sum(self.config.encoder_output_length(int(n)) for n in audio_feature_lens.tolist()) + layout = LabelLayout( + spans=( + Span("ids", "prefill_text", len(pre_ids), len(pre_ids), "text_inputs"), + Span("digest", "prefill_audio", audio, audio, ("audio", 0)), + Span("ids", "prefill_text", len(post_ids), len(post_ids), "text_inputs"), + ), + growth=(("decode_loop", 1),), + ) + return ProcessPromptOutput({ "text_inputs": [ torch.tensor(pre_ids, dtype=torch.long), torch.tensor(post_ids, dtype=torch.long), ], "audio_features": [audio_features], "audio_feature_lens": [audio_feature_lens], - } + }, {"kv_layout": {KV_CACHE: {"main": layout}}}) # ------------------------------------------------------------------- # Model ABC: postprocess From 06ccb0a11ebb8494667fa87f62e951941f3f291f Mon Sep 17 00:00:00 2001 From: Gaurav Shah Date: Wed, 21 Oct 2026 05:50:51 +0000 Subject: [PATCH 12/26] model: nemotron_duplex declares its KV layout --- .../nemotron_duplex/nemotron_duplex_model.py | 34 ++++++++++++++++--- test/modular/test_nemotron_duplex_model.py | 10 +++--- 2 files changed, 35 insertions(+), 9 deletions(-) diff --git a/mstar/model/nemotron_duplex/nemotron_duplex_model.py b/mstar/model/nemotron_duplex/nemotron_duplex_model.py index b351c85792..2e72c66c94 100644 --- a/mstar/model/nemotron_duplex/nemotron_duplex_model.py +++ b/mstar/model/nemotron_duplex/nemotron_duplex_model.py @@ -65,7 +65,7 @@ from mstar.engine.resources.recurrent import Mamba2Geometry, RecurrentStateConfig, RecurrentStateSpec from mstar.graph.base import GraphEdge, GraphNode, GraphSection, Loop, TensorPointerInfo from mstar.graph.special_destinations import EMIT_TO_CLIENT -from mstar.model.base import ForwardPassArgs, Model +from mstar.model.base import ForwardPassArgs, KVLabel, LabelLayout, Model, ProcessPromptOutput, Span from mstar.model.multimodal import PromptPart from mstar.model.nemotron_duplex.config import ( MAMBA, @@ -662,7 +662,7 @@ def process_prompt( tensors: NameToTensorList | None = None, prompt_parts: list[PromptPart] | None = None, **kwargs, - ) -> NameToTensorList: + ) -> ProcessPromptOutput: # Duplex request: raw user audio (-> Encoder partition) plus an optional # system-prompt text that primes the nano cache (prefill_text). out: NameToTensorList = {} @@ -726,7 +726,28 @@ def process_prompt( out["has_prompt"] = [torch.tensor([1], dtype=torch.long)] else: out["text_inputs"] = [torch.tensor([-1], dtype=torch.long)] - return out + # a session's first frame writes the prompt along with its own row, and the talker's first frame writes + # its speaker warm-up + prompt_len = ids.numel() if prompt else 0 + nano = LabelLayout( + spans=(Span("ids", "decode", prompt_len, prompt_len, "text_inputs"),) if prompt else (), + growth=(("decode_loop", 1),), + ) + warmup = self.TALKER_WARMUP_POSITIONS + talker = LabelLayout( + spans=(Span("rows", "talker_decode", warmup, warmup, "audio_prompt_latents"),), + growth=(("talker_decode_loop", 1),), + ) + return ProcessPromptOutput(out, {"kv_layout": { + NANO_KV: {"main": nano}, TALKER_KV: {"main": talker, "uncond": talker}, + }}) + + def kv_labels(self) -> dict[str, dict[str, KVLabel]]: + talker = KVLabel(growth=(("talker_decode", "talker_decode_loop"),)) + return { + NANO_KV: {"main": KVLabel(growth=(("decode", "decode_loop"),))}, + TALKER_KV: {"main": talker, "uncond": talker}, + } def load_audio(self, filepath: str, device: str): """Decode an uploaded audio file to 16 kHz mono via soundfile (the base uses @@ -1171,9 +1192,14 @@ def _create_talker_submodule( cv = build_char_vocab(tok) s2c, _ = subword_to_char_ids(tok, cv) # checkpoint config, so the talker samples with its inference_* knobs - return EarTTSTalkerSubmodule( + submodule = EarTTSTalkerSubmodule( talker=talker, config=cfg, subword_to_char=s2c, char_pad_idx=len(cv), ) + assert submodule.warmup_len == self.TALKER_WARMUP_POSITIONS, ( + f"the checkpoint's speaker warm-up is {submodule.warmup_len} positions, " + f"where the talker's layout counts {self.TALKER_WARMUP_POSITIONS}" + ) + return submodule def _create_codec_submodule( self, device: str, autocast_dtype: torch.dtype | None = None, diff --git a/test/modular/test_nemotron_duplex_model.py b/test/modular/test_nemotron_duplex_model.py index f8dadf4d31..3034108b54 100644 --- a/test/modular/test_nemotron_duplex_model.py +++ b/test/modular/test_nemotron_duplex_model.py @@ -227,7 +227,7 @@ def test_duplex_process_prompt_seeds_frame0_feedback(): out = model.process_prompt( None, ["audio"], ["audio", "text"], tensors={"audio_inputs": [torch.zeros(16000)]}, - ) + ).new_input_tensors assert "audio_features" in out assert int(out["prev_text"][0].item()) == model.config.text_bos_id assert int(out["prev_func"][0].item()) == model.config.text_pad_id @@ -238,7 +238,7 @@ def test_duplex_process_prompt_seeds_frame0_feedback(): out_p = model.process_prompt( "You are NVIDIA Voice Chat.", ["audio", "text"], ["audio", "text"], tensors={"audio_inputs": [torch.zeros(16000)]}, - ) + ).new_input_tensors assert int(out_p["prev_text"][0].item()) == model.config.text_pad_id assert out_p["text_inputs"][0].tolist() == [5, 6, 7] and "has_prompt" in out_p @@ -253,7 +253,7 @@ def test_duplex_process_prompt_rejects_audio_under_one_frame(): with pytest.raises(ValueError, match="need at least"): model.process_prompt(None, ["audio"], ["audio", "text"], tensors={"audio_inputs": [torch.zeros(n)]}) out = model.process_prompt(None, ["audio"], ["audio", "text"], tensors={"audio_inputs": [torch.zeros(frame)]}) - assert out["audio_features"][0].numel() == frame + assert out.new_input_tensors["audio_features"][0].numel() == frame def test_duplex_process_prompt_rejects_bad_requests(): @@ -283,10 +283,10 @@ def test_duplex_process_prompt_scales_down_over_full_scale_audio(): model = _make_model() wav = torch.sin(torch.arange(16000) / 10.0) out = model.process_prompt(None, ["audio"], ["audio", "text"], tensors={"audio_inputs": [wav * 1e6]}) - assert torch.allclose(out["audio_features"][0], wav, atol=1e-6) + assert torch.allclose(out.new_input_tensors["audio_features"][0], wav, atol=1e-6) quiet = wav * 0.5 out = model.process_prompt(None, ["audio"], ["audio", "text"], tensors={"audio_inputs": [quiet]}) - assert torch.equal(out["audio_features"][0], quiet) + assert torch.equal(out.new_input_tensors["audio_features"][0], quiet) def test_duplex_no_prompt_seeds_initial_decode_inputs(): From 0806f64bc8410aa39e6c360eafb3cf2e9188fe8f Mon Sep 17 00:00:00 2001 From: Gaurav Shah Date: Wed, 21 Oct 2026 05:50:51 +0000 Subject: [PATCH 13/26] model: pi05 declares its KV layout --- mstar/model/pi05/pi05_model.py | 19 ++++++++++++++++--- test/modular/test_pi05_model.py | 4 ++-- 2 files changed, 18 insertions(+), 5 deletions(-) diff --git a/mstar/model/pi05/pi05_model.py b/mstar/model/pi05/pi05_model.py index b596e76170..12290f1e5f 100644 --- a/mstar/model/pi05/pi05_model.py +++ b/mstar/model/pi05/pi05_model.py @@ -52,7 +52,7 @@ TensorPointerInfo, ) from mstar.graph.special_destinations import EMIT_TO_CLIENT -from mstar.model.base import ForwardPassArgs, Model +from mstar.model.base import ForwardPassArgs, KVLabel, LabelLayout, Model, ProcessPromptOutput, Span from mstar.model.loader import LLAMA_STACKED_PARAMS, load_hf_weights from mstar.model.pi05.components.action_expert import Pi05ActionExpert, Pi05TimeMLP from mstar.model.pi05.components.paligemma import Pi05PaliGemmaExpert @@ -489,7 +489,7 @@ def process_prompt( input_modalities: list[str], output_modalities: list[str], **kwargs, - ) -> NameToTensorList: + ) -> NameToTensorList | ProcessPromptOutput: """Tokenize the Pi0.5 prompt + robot state into a single token stream. Pi0.5's production preprocessor (lerobot's @@ -532,7 +532,20 @@ def process_prompt( full_prompt = cleaned text_ids = self.tokenizer.encode_prompt(full_prompt) - return {"text_inputs": [text_ids]} + cfg = self.config + # The most image rows a prefill writes: a captured SigLIP step encodes all num_cameras slots. + images = cfg.num_cameras * (cfg.vit_image_size // cfg.vit_patch_size) ** 2 + layout = LabelLayout( + spans=( + Span("rows", self.PREFILL_WALK, images, images, "img_emb"), + Span("ids", self.PREFILL_WALK, text_ids.numel(), text_ids.numel(), "text_inputs"), + ), + scratch=cfg.action_horizon, + ) + return ProcessPromptOutput({"text_inputs": [text_ids]}, {"kv_layout": {LLM_KV: {"main": layout}}}) + + def kv_labels(self) -> dict[str, dict[str, KVLabel]]: + return {LLM_KV: {"main": KVLabel()}} def postprocess(self, output: torch.Tensor, modality: str, request_kwargs: dict | None = None) -> bytes: if modality == "action": diff --git a/test/modular/test_pi05_model.py b/test/modular/test_pi05_model.py index 1fed65c568..3ea90af09c 100644 --- a/test/modular/test_pi05_model.py +++ b/test/modular/test_pi05_model.py @@ -334,7 +334,7 @@ def test_pi05_process_prompt_formats_state_into_text(): input_modalities=["image", "text"], output_modalities=["action"], robot_state=state, - ) + ).new_input_tensors # Result is a single text_inputs entry — no separate state_inputs. assert set(result.keys()) == {"text_inputs"} @@ -357,7 +357,7 @@ def test_pi05_process_prompt_without_state_uses_plain_text(): prompt="hello world", input_modalities=["text"], output_modalities=["action"], - ) + ).new_input_tensors assert "text_inputs" in result assert stub.last_prompt == "hello world" From 70e0701043ddbaea6ad8a664a2a8b3e299d53ec3 Mon Sep 17 00:00:00 2001 From: Gaurav Shah Date: Wed, 21 Oct 2026 05:50:51 +0000 Subject: [PATCH 14/26] model: qwen3_5 declares its KV layout through a shared prompt_spans --- mstar/model/bagel/bagel_model.py | 34 +++++++++++++-------------- mstar/model/base.py | 19 ++++++++++++++- mstar/model/qwen3_5/qwen3_5_model.py | 35 +++++++++++++++++++++++++++- test/modular/test_chat_templates.py | 2 +- 4 files changed, 69 insertions(+), 21 deletions(-) diff --git a/mstar/model/bagel/bagel_model.py b/mstar/model/bagel/bagel_model.py index cc7f2e1f31..a04e166e9f 100644 --- a/mstar/model/bagel/bagel_model.py +++ b/mstar/model/bagel/bagel_model.py @@ -94,6 +94,7 @@ PrefixKeys, ProcessPromptOutput, Span, + prompt_spans, ) from mstar.model.loader import iter_safetensors_file, load_hf_weights from mstar.model.loader.base import LLAMA_STACKED_PARAMS, StackedParamRule @@ -736,27 +737,24 @@ def _kv_layout( spans = [] if think_mode and not understanding and segments: spans.append(Span("ids", "prefill_text", segments[0].numel(), segments[0].numel(), "text_inputs")) - texts = segments[len(spans):] - for step in prefill_plan(parts, leading_text=understanding): - if step.modality == TEXT: - if step.index >= len(texts): - # no tensor for it: `_build_prefill_schedule` skips the step too - continue - length = texts[step.index].numel() - spans.append(Span("ids", "prefill_text", length, length, "text_inputs")) - continue - _, image_height, image_width = images[step.index].shape + + def image(step: PromptPart) -> list[Span]: + _, height, width = images[step.index].shape item = (("image", step.index), (image_preprocess,)) - if not understanding: - latents = vae_latents( - image_height, image_width, image_preprocess, cfg.latent_downsample, cfg.max_latent_size, - ) - spans.append(Span("digest", "prefill_vae", latents + IMAGE_SENTINELS, IMAGE_ADVANCE, *item)) patches = vit_patches( - image_height, image_width, image_preprocess, - cfg.vit_config.patch_size, cfg.vit_max_num_patch_per_side, + height, width, image_preprocess, cfg.vit_config.patch_size, cfg.vit_max_num_patch_per_side, ) - spans.append(Span("digest", "prefill_vit", patches + IMAGE_SENTINELS, IMAGE_ADVANCE, *item)) + vit = Span("digest", "prefill_vit", patches + IMAGE_SENTINELS, IMAGE_ADVANCE, *item) + if understanding: + return [vit] + latents = vae_latents(height, width, image_preprocess, cfg.latent_downsample, cfg.max_latent_size) + return [Span("digest", "prefill_vae", latents + IMAGE_SENTINELS, IMAGE_ADVANCE, *item), vit] + + texts = segments[len(spans):] + # skip a text step the prompt has no tensor for, as `_build_prefill_schedule` does + plan = prefill_plan(parts, leading_text=understanding) + plan = [step for step in plan if step.modality != TEXT or step.index < len(texts)] + spans += prompt_spans(plan, texts, "prefill_text", image) keyed = int(bool(spans) and spans[0].kind == "ids") if understanding: main = LabelLayout( diff --git a/mstar/model/base.py b/mstar/model/base.py index c95de6d258..49d04d6098 100644 --- a/mstar/model/base.py +++ b/mstar/model/base.py @@ -1,5 +1,6 @@ import itertools from abc import ABC, abstractmethod +from collections.abc import Callable, Iterable, Sequence from dataclasses import asdict, dataclass, field from fractions import Fraction from typing import NamedTuple, Type @@ -29,7 +30,7 @@ top_level_entities, ) from mstar.model.media_limits import VideoBudget, read_image -from mstar.model.multimodal import PromptPart +from mstar.model.multimodal import TEXT, PromptPart DECODE = "decode" MAX_OUTPUT_TOKENS = 2048 @@ -112,6 +113,22 @@ def row(self) -> list: ] +def prompt_spans( + plan: Iterable[PromptPart], texts: Sequence[torch.Tensor], text_walk: str, + item_spans: Callable[[PromptPart], Iterable[Span]], +) -> list[Span]: + """A prompt's prefill writes in ``plan`` order: a text step's ids from ``texts``, and an + attachment's from ``item_spans``.""" + spans = [] + for step in plan: + if step.modality != TEXT: + spans.extend(item_spans(step)) + continue + length = texts[step.index].numel() + spans.append(Span("ids", text_walk, length, length, "text_inputs")) + return spans + + class PrefixKeys(NamedTuple): """How a label's pages are keyed for cross-request prefix reuse. diff --git a/mstar/model/qwen3_5/qwen3_5_model.py b/mstar/model/qwen3_5/qwen3_5_model.py index 0bc17fb1aa..1914aa24a4 100644 --- a/mstar/model/qwen3_5/qwen3_5_model.py +++ b/mstar/model/qwen3_5/qwen3_5_model.py @@ -39,9 +39,13 @@ from mstar.graph.special_destinations import EMIT_TO_CLIENT from mstar.model.base import ( ForwardPassArgs, + KVLabel, + LabelLayout, Model, ProcessPromptOutput, + Span, TensorAndMetadata, + prompt_spans, ) from mstar.model.multimodal import ( TEXT, @@ -55,6 +59,7 @@ split_around_spans, ) from mstar.model.qwen3_5.components.language_model import Qwen3_5ForCausalLM +from mstar.model.qwen3_5.components.rope import vision_position_advance from mstar.model.qwen3_5.config import ( ATTN, GDN_STATE, @@ -549,8 +554,36 @@ def process_prompt( result["video_grid_thw"] = [video.grid for video in videos] # the decoded frames are spent; don't hold them for the request's life result["video_inputs"] = [] + # with an image or a video the whole prompt is one prefill_vision walk (see `_prefill_schedule`) + walk = "prefill_vision" if images or videos else "prefill_text" + merge = self.vision_config.spatial_merge_size + owners = group_owners(groups) + + def media(step: PromptPart) -> tuple[Span]: + # the encoder's merged tokens and their two sentinels, as the LLM's prefill_vision splices them. + # Each temporal group of a video is placed like a one-frame image. + video = step.modality == VIDEO + grid = result["video_grid_thw"][owners[step.index]] if video else result["image_grid_thw"][step.index] + grid = grid[0] if grid.dim() == 2 else grid + t, h, w = (int(v) for v in grid.tolist()) + grid = torch.tensor((1 if video else t, h, w)) + rows = int(grid.prod()) // merge**2 + 2 + source = (VIDEO, owners[step.index]) if video else ("image", step.index) + return (Span("rows" if video else "digest", walk, rows, vision_position_advance(grid, merge) + 2, source),) + + plan = prefill_plan(expand_video_parts(parts, groups)) + layout = LabelLayout(spans=tuple(prompt_spans(plan, segments, walk, media)), growth=(("decode_loop", 1),)) # always set, so a `video_groups` the client sent never reaches the schedule - return ProcessPromptOutput(result, {"video_groups": groups}) + return ProcessPromptOutput(result, {"video_groups": groups, "kv_layout": {KV_CACHE: {"main": layout}}}) + + def get_loop_caps(self, **model_kwargs) -> dict[str, int]: + # the decode stop check ends the loop after max_tokens iterations + return {"decode_loop": self.get_max_output_tokens(**model_kwargs)} + + def kv_labels(self) -> dict[str, dict[str, KVLabel]]: + return {KV_CACHE: {"main": KVLabel( + growth=(("decode", "decode_loop"),), + )}} def postprocess( diff --git a/test/modular/test_chat_templates.py b/test/modular/test_chat_templates.py index 1e2ebab002..2f68be5757 100644 --- a/test/modular/test_chat_templates.py +++ b/test/modular/test_chat_templates.py @@ -45,7 +45,7 @@ def _qwen3_5(path, text, in_mods, parts): from mstar.model.qwen3_5.qwen3_5_model import Qwen3_5DenseModel model = Qwen3_5DenseModel(path) - [ids] = model.process_prompt(text, in_mods, ["text"], prompt_parts=parts)["text_inputs"] + [ids] = model.process_prompt(text, in_mods, ["text"], prompt_parts=parts).new_input_tensors["text_inputs"] rendered = model.tokenizer.apply_chat_template( CHAT, tokenize=False, add_generation_prompt=True, enable_thinking=True, ) From 22d6f9b3d7177fb4093049ee51da959154ddb84c Mon Sep 17 00:00:00 2001 From: Gaurav Shah Date: Wed, 21 Oct 2026 05:50:51 +0000 Subject: [PATCH 15/26] model: qwen3_omni declares its KV layout --- mstar/model/qwen3_omni/qwen3_omni_model.py | 86 ++++++++++++++++++- mstar/model/qwen3_omni/submodules.py | 46 +++++----- test/modular/test_chat_templates.py | 2 +- test/modular/test_qwen3_omni_prompt_parity.py | 2 +- 4 files changed, 109 insertions(+), 27 deletions(-) diff --git a/mstar/model/qwen3_omni/qwen3_omni_model.py b/mstar/model/qwen3_omni/qwen3_omni_model.py index 15c338b864..9212ba4305 100644 --- a/mstar/model/qwen3_omni/qwen3_omni_model.py +++ b/mstar/model/qwen3_omni/qwen3_omni_model.py @@ -55,7 +55,16 @@ ) from mstar.graph.base import GraphEdge, GraphNode, Loop, Sequential, TensorPointerInfo from mstar.graph.special_destinations import EMIT_TO_CLIENT, EMPTY_DESTINATION -from mstar.model.base import MAX_OUTPUT_TOKENS, ForwardPassArgs, Model, TensorAndMetadata +from mstar.model.base import ( + MAX_OUTPUT_TOKENS, + ForwardPassArgs, + KVLabel, + LabelLayout, + Model, + ProcessPromptOutput, + Span, + TensorAndMetadata, +) from mstar.model.media_limits import VideoBudget, probe_video from mstar.model.multimodal import ( TEXT, @@ -68,6 +77,7 @@ prefill_plan, split_around_spans, ) +from mstar.model.qwen3_omni.components.rope import get_rope_index_vision from mstar.model.qwen3_omni.components.talker import Qwen3OmniCodePredictor from mstar.model.qwen3_omni.config import ( CODE_PRED_SAMPLER, @@ -1078,7 +1088,7 @@ def process_prompt( input_metadata: dict[str, dict] = {}, prompt_parts: list[PromptPart] | None = None, **kwargs, - ) -> NameToTensorList: + ) -> ProcessPromptOutput: """Build the full ChatML prompt + derived multimodal tensors. Uses HF's full ``AutoProcessor`` (combines tokenizer + image_processor @@ -1180,7 +1190,8 @@ def process_prompt( input_ids = self.tokenizer(text, return_tensors="pt")["input_ids"][0] spans = find_media_spans(input_ids, self._placeholder_specs()) segments = split_around_spans(input_ids, spans) - check_plan(prefill_plan(parts), spans, len(segments)) + plan = prefill_plan(parts) + check_plan(plan, spans, len(segments)) result["text_inputs"] = segments result["pixel_values"] = [] @@ -1240,7 +1251,74 @@ def process_prompt( result["video_grid_thw"] += vid_out["video_grid_thw"] result["pixel_values_videos"].append(vid_out["pixel_values_videos"]) - return result + if not attached: + # a call without tensors has no attachments to lay out + return ProcessPromptOutput(result, {}) + return ProcessPromptOutput(result, {"kv_layout": self._kv_layout(plan, result, "audio" in output_modalities)}) + + def _kv_layout(self, plan: list[PromptPart], tensors: NameToTensorList, audio_output: bool) -> dict: + """What each Thinker prefill step writes and, with audio out, what the Talker writes of it. + + The Talker prefills every step's rows except those of a system or earlier + assistant turn. It then prefills its assistant prefix, and writes a row per decode step. + """ + from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import _get_feat_extract_output_lengths + + from mstar.model.qwen3_omni.submodules import TalkerSubmodule, talker_text_mask + + merge = self.config.vision.spatial_merge_size + grids = {"image": tensors["image_grid_thw"], "video": tensors["video_grid_thw"]} + thinker, talker, position = [], [], 0 + for step in plan: + if step.modality == TEXT: + ids = tensors["text_inputs"][step.index] + span = Span("ids", "prefill_text", len(ids), len(ids), "text_inputs") + kept = int(talker_text_mask(ids, self.config).sum()) + elif step.modality == "audio": + # the encoder's rows, as it slices them, between two sentinels + rows = int(_get_feat_extract_output_lengths(tensors["audio_seqlens"][step.index])) + 2 + span = Span("digest", "prefill_audio", rows, rows, ("audio", step.index)) + kept = rows + else: + grid = grids[step.modality][step.index] + per_grid = tensors["video_second_per_grid"][step.index].item() if step.modality == "video" else None + rows = int(grid.prod()) // merge ** 2 + 2 + # the walk moves the counter one past its end sentinel, which follows the grid's largest position + positions = get_rope_index_vision( + grid, position + 1, self.config.thinker.position_id_per_seconds, + spatial_merge_size=merge, seconds_per_grid=per_grid, + ) + advance = int(positions.max()) + 2 - position + span = Span("digest", "prefill_vision", rows, advance, (step.modality, step.index)) + kept = rows + thinker.append(span) + talker.append(Span("rows", "talker_prefill", kept, kept, "thinker_states")) + position += span.advance + prefix = TalkerSubmodule.TALKER_LAST_PREFILL_TOKENS_PER_REQ + return { + THINKER_KV: {"main": LabelLayout(spans=tuple(thinker), growth=(("thinker_decode_loop", 1),))}, + TALKER_KV: {"main": LabelLayout( + spans=(*talker, Span("rows", "talker_last_prefill", prefix, prefix, "thinker_states")), + growth=(("talker_decode_loop", 1),), + ) if audio_output else LabelLayout()}, + } + + def get_loop_caps(self, **model_kwargs) -> dict[str, int]: + # the Thinker's and the Talker's stop checks end their decode loops after their max-token counts + return { + "thinker_decode_loop": self.get_max_output_tokens(**model_kwargs), + "talker_decode_loop": max(1, self._get_max_talker_output_tokens(**model_kwargs)), + } + + def kv_labels(self) -> dict[str, dict[str, KVLabel]]: + return { + THINKER_KV: {"main": KVLabel( + growth=(("thinker_decode", "thinker_decode_loop"),), + )}, + TALKER_KV: {"main": KVLabel( + growth=(("talker_decode", "talker_decode_loop"),), + )}, + } # ----------------------------------------------------------------------- # Model ABC: postprocess diff --git a/mstar/model/qwen3_omni/submodules.py b/mstar/model/qwen3_omni/submodules.py index f1932a93a4..663cc648bc 100644 --- a/mstar/model/qwen3_omni/submodules.py +++ b/mstar/model/qwen3_omni/submodules.py @@ -81,6 +81,27 @@ def _accepts_kwarg(fn: Any, name: str) -> bool: ) +def talker_text_mask(input_ids: torch.Tensor, config: Qwen3OmniModelConfig) -> torch.Tensor: + """ + Cut system prompt and previous assistant parts out of the talker input + """ + im_start_indexes = ( + input_ids == config.im_start_token_id + ).nonzero(as_tuple=True)[0] + mask = torch.ones(input_ids.shape, dtype=torch.bool, device=input_ids.device) + + for i in range(len(im_start_indexes) - 1): + im_start_index = im_start_indexes[i] + segment_end_index = im_start_indexes[i + 1] + role_token = input_ids[im_start_index + 1] + # Talker should ignore thinker system prompt + if role_token == config.system_token_id: + mask[im_start_index:segment_end_index] = 0 + elif role_token == config.assistant_token_id: + mask[im_start_index:segment_end_index] = 0 + return mask + + # =================================================================== # 1. AudioEncoderSubmodule (enc_dec engine) # =================================================================== @@ -321,26 +342,6 @@ def _get_decode_thinker_mask(self, device: torch.device) -> torch.Tensor: ) return self._decode_thinker_mask - def _get_talker_text_mask(self, input_ids: torch.Tensor) -> torch.Tensor: - """ - Cut system prompt and previous assistant parts out of the talker input - """ - im_start_indexes = ( - input_ids == self.config.im_start_token_id - ).nonzero(as_tuple=True)[0] - mask = torch.ones(input_ids.shape, dtype=torch.bool, device=input_ids.device) - - for i in range(len(im_start_indexes) - 1): - im_start_index = im_start_indexes[i] - segment_end_index = im_start_indexes[i + 1] - role_token = input_ids[im_start_index + 1] - # Talker should ignore thinker system prompt - if role_token == self.config.system_token_id: - mask[im_start_index:segment_end_index] = 0 - elif role_token == self.config.assistant_token_id: - mask[im_start_index:segment_end_index] = 0 - return mask - def _wrap_audio_input(self, audio_embeds: torch.Tensor): # Wrap the audio span in ``<|audio_bos|>`` / ``<|audio_eos|>`` # sentinel token embeddings so the Thinker sees the same @@ -437,7 +438,7 @@ def prepare_inputs( pos_ids = get_rope_index_text(seq_len, start_pos, device) masks_for_talker = torch.stack([ torch.zeros(text_ids.shape, dtype=torch.bool, device=device), # multimodal - self._get_talker_text_mask(text_ids) # text inclusion + talker_text_mask(text_ids, self.config) # text inclusion ]) return ARNodeInputs( input_seq_len=seq_len, @@ -507,6 +508,9 @@ def prepare_inputs( # ``image_grid_thw`` is available, use ``get_rope_index_vision``; # otherwise fall back to a 1-D sequence (test path) grid_thw = inputs.get("image_grid_thw", [None])[0] + assert vision_len == int(grid_thw.prod()) // self.config.vision.spatial_merge_size ** 2, ( + f"the vision encoder gave {vision_len} rows, where the layout counts the grid's merged patches" + ) seconds_per_grid = inputs.get("video_second_per_grid", []) seconds_per_grid = seconds_per_grid[0].item() if seconds_per_grid else None vision_pos_ids = get_rope_index_vision( diff --git a/test/modular/test_chat_templates.py b/test/modular/test_chat_templates.py index 2f68be5757..ad9cc4f39a 100644 --- a/test/modular/test_chat_templates.py +++ b/test/modular/test_chat_templates.py @@ -35,7 +35,7 @@ def _qwen3_omni(path, text, in_mods, parts): model = Qwen3OmniModel.__new__(Qwen3OmniModel) model._processor = AutoProcessor.from_pretrained(path, trust_remote_code=False) model.tokenizer = model._processor.tokenizer - [ids] = model.process_prompt(text, in_mods, ["text"], prompt_parts=parts)["text_inputs"] + [ids] = model.process_prompt(text, in_mods, ["text"], prompt_parts=parts).new_input_tensors["text_inputs"] rendered = model._processor.apply_chat_template(CHAT, tokenize=False, add_generation_prompt=True) return ids.tolist(), model.tokenizer(rendered)["input_ids"] diff --git a/test/modular/test_qwen3_omni_prompt_parity.py b/test/modular/test_qwen3_omni_prompt_parity.py index e88a6eb2fc..6c8cce2952 100644 --- a/test/modular/test_qwen3_omni_prompt_parity.py +++ b/test/modular/test_qwen3_omni_prompt_parity.py @@ -161,7 +161,7 @@ def _old_process_prompt(model, prompt, input_modalities): def _new_text_inputs(model, prompt, input_modalities, prompt_parts=None): return model.process_prompt( prompt, input_modalities, ["text"], prompt_parts=prompt_parts, - )["text_inputs"] + ).new_input_tensors["text_inputs"] def _ids(segments): From 609a51b0f7bd4222a841edb6e6ca52c94f9f7e17 Mon Sep 17 00:00:00 2001 From: Gaurav Shah Date: Wed, 21 Oct 2026 05:50:51 +0000 Subject: [PATCH 16/26] model: qwen3_tts declares its KV layout --- mstar/model/qwen3_tts/qwen3_tts_model.py | 18 +++++++++++++++++- mstar/model/qwen3_tts/submodules.py | 21 +++++++++++++++++++++ test/modular/test_model_internal_kwargs.py | 5 +++-- test/modular/test_qwen3_tts_model.py | 4 ++-- 4 files changed, 43 insertions(+), 5 deletions(-) diff --git a/mstar/model/qwen3_tts/qwen3_tts_model.py b/mstar/model/qwen3_tts/qwen3_tts_model.py index 15f55b5475..8bf297ffc2 100644 --- a/mstar/model/qwen3_tts/qwen3_tts_model.py +++ b/mstar/model/qwen3_tts/qwen3_tts_model.py @@ -68,7 +68,7 @@ ) from mstar.graph.special_destinations import EMIT_TO_CLIENT, EMPTY_DESTINATION from mstar.model.audio_watermark import AudioSealWatermarker, AudioWatermarkConfig, import_audioseal -from mstar.model.base import ForwardPassArgs, Model, ProcessPromptOutput, TensorAndMetadata +from mstar.model.base import ForwardPassArgs, KVLabel, LabelLayout, Model, ProcessPromptOutput, Span, TensorAndMetadata from mstar.model.qwen3_tts.config import ( CHATML_ASSISTANT_PREFIX_TOKEN_IDS, CHATML_ASSISTANT_SUFFIX_TOKEN_IDS, @@ -874,6 +874,15 @@ def process_prompt( clips = (tensors or {}).get("audio_inputs", []) if clips: metadata[REF_CLIP_KEY] = clip_digest(clips[0]) + from mstar.model.qwen3_tts.submodules import prefill_rows + + # a clone prefills on its own walk, and its x-vector takes the speaker slot + clone = "audio" in input_modalities + rows = prefill_rows(metadata[PROMPT_LAYOUT_KEY], language_id, speaker_id >= 0 or clone) + metadata["kv_layout"] = {TALKER_KV: {"main": LabelLayout( + spans=(Span("rows", "talker_prefill_clone" if clone else "talker_prefill", rows, rows, "text_inputs"),), + growth=(("talker_decode_loop", 1),), + )}} return ProcessPromptOutput( {"text_inputs": [torch.cat([instruct_ids, assistant_ids, ref_ids])]}, metadata ) @@ -1079,6 +1088,13 @@ def get_max_output_tokens(self, **model_kwargs: Any) -> int: model_kwargs.get("max_new_tokens", self.config.generation.max_new_tokens), ) + def get_loop_caps(self, **model_kwargs: Any) -> dict[str, int]: + # the stop check counts the prefill's frame, so decode runs one iteration fewer than max tokens + return {"talker_decode_loop": max(1, self.get_max_output_tokens(**model_kwargs) - 1)} + + def kv_labels(self) -> dict[str, dict[str, KVLabel]]: + return {TALKER_KV: {"main": KVLabel(growth=(("talker_decode", "talker_decode_loop"),))}} + def get_output_sample_rate(self, modality: str = "audio") -> int: return self.config.codec.output_sample_rate diff --git a/mstar/model/qwen3_tts/submodules.py b/mstar/model/qwen3_tts/submodules.py index 55ef9622b9..cc9aed97d3 100644 --- a/mstar/model/qwen3_tts/submodules.py +++ b/mstar/model/qwen3_tts/submodules.py @@ -96,6 +96,23 @@ # the CodePredictor depth loop, as a piecewise capture region DEPTH_LOOP_REGION = "code_predictor_loop" + +def prefill_rows(prompt_layout: Sequence[int], language_id: int, speaker: bool) -> int: + """The rows `TalkerSubmodule._build_prefill` writes, where ``speaker`` is True when a + speaker tag or x-vector takes a row.""" + instruct_len, text_len, stream_text, ref_text_len, ref_frames = prompt_layout + prefix = len(CHATML_ASSISTANT_PREFIX_TOKEN_IDS) + # the role, then every codec tag but codec BOS + rows = instruct_len + prefix + (3 if language_id < 0 else 4) + int(speaker) + 1 + if ref_frames: + # codec BOS and the reference frames. When the text isn't streamed, both transcripts and TTS EOS come first + text = ref_text_len - prefix - 2 + text_len + 1 + return rows + 1 + ref_frames + (0 if stream_text else text) + # a streamed prompt adds its first text token over codec BOS. Any other adds its text and TTS EOS, + # then TTS PAD over codec BOS + return rows + (1 if stream_text else text_len + 2) + + # =========================================================================== # 1. TalkerSubmodule - autoregressive 12 Hz codec-frame generation # =========================================================================== @@ -433,6 +450,10 @@ def prepare_inputs( speaker_embed=inputs["speaker_embed"][0] if "speaker_embed" in inputs else None, ref_codes=inputs["ref_codes"][0] if "ref_codes" in inputs else None, ) + speaker = int(step_metadata[SPEAKER_ID_KEY]) >= 0 or "speaker_embed" in inputs + assert input_embeds.shape[0] == prefill_rows( + step_metadata[PROMPT_LAYOUT_KEY], int(step_metadata[LANGUAGE_ID_KEY]), speaker, + ), f"the prefill built {input_embeds.shape[0]} rows, not the count `prefill_rows` lays out" state = self.request_state(fwd_info.rid_handle) elif graph_walk == "talker_decode": state = self.request_state(fwd_info.rid_handle) diff --git a/test/modular/test_model_internal_kwargs.py b/test/modular/test_model_internal_kwargs.py index f5f6db9d84..72bedb1e70 100644 --- a/test/modular/test_model_internal_kwargs.py +++ b/test/modular/test_model_internal_kwargs.py @@ -72,8 +72,9 @@ def test_qwen3_tts_a_client_cannot_set_the_codec_stream_lead(variant, mode, lead "Codec", modalities, ["audio"], {}, model_kwargs={**client, **result.metadata} ) assert codec.stream_lead_items == ({"codec_tokens": lead} if lead else {}) - # every key the conductor reads back is one the data worker strips from the client - assert set(result.metadata) <= set(model.INTERNAL_KWARGS) + # every key the conductor reads back is one the data worker strips from the client. + # The data worker strips kv_layout for every model. + assert set(result.metadata) - {"kv_layout"} <= set(model.INTERNAL_KWARGS) class _RecordingCommunicator: diff --git a/test/modular/test_qwen3_tts_model.py b/test/modular/test_qwen3_tts_model.py index 4090cf6180..2e6ec90943 100644 --- a/test/modular/test_qwen3_tts_model.py +++ b/test/modular/test_qwen3_tts_model.py @@ -351,7 +351,7 @@ def test_qwen3_tts_process_prompt_matches_official_template(): assert result.new_input_tensors["text_inputs"][0].tolist() == ASSISTANT_PREFIX + [1000] + ASSISTANT_SUFFIX # CustomVoice default: whole text in the prefill (stream_text = 0); no # reference transcript or frames. Host values, for the prefill's step metadata. - assert result.metadata == { + assert {key: value for key, value in result.metadata.items() if key != "kv_layout"} == { PROMPT_LAYOUT_KEY: [0, 1, 0, 0, 0], SPEAKER_ID_KEY: 3065, LANGUAGE_ID_KEY: 2055, Qwen3TTSModel.CODEC_STREAM_LEAD_KEY: 0, } @@ -450,7 +450,7 @@ def test_qwen3_tts_base_process_prompt_builds_in_context_clone(): assert model.tokenizer.texts[-1] == "<|im_start|>assistant\nthe reference says<|im_end|>\n" assert list(result.new_input_tensors) == ["text_inputs"] assert result.new_input_tensors["text_inputs"][0].tolist() == assistant_ids + ref_ids - assert result.metadata == { + assert {key: value for key, value in result.metadata.items() if key != "kv_layout"} == { # [instruct_len, text_len, stream_text (Base default: streaming), ref_text_len, ref_frames] PROMPT_LAYOUT_KEY: [0, 3, 1, len(ref_ids), 13], SPEAKER_ID_KEY: -1, From 26b36057e2d902cb1f3d498830150326b7dc766f Mon Sep 17 00:00:00 2001 From: Gaurav Shah Date: Wed, 21 Oct 2026 05:50:51 +0000 Subject: [PATCH 17/26] model: vjepa2 declares its KV layout --- mstar/model/vjepa2/vjepa2_model.py | 40 ++++++++++++++++++++++++++---- 1 file changed, 35 insertions(+), 5 deletions(-) diff --git a/mstar/model/vjepa2/vjepa2_model.py b/mstar/model/vjepa2/vjepa2_model.py index 033346ca88..2d2620f001 100644 --- a/mstar/model/vjepa2/vjepa2_model.py +++ b/mstar/model/vjepa2/vjepa2_model.py @@ -53,7 +53,15 @@ TensorPointerInfo, ) from mstar.graph.special_destinations import EMIT_TO_CLIENT -from mstar.model.base import ForwardPassArgs, Model, TensorAndMetadata, video_metadata_dict +from mstar.model.base import ( + ForwardPassArgs, + KVLabel, + LabelLayout, + Model, + ProcessPromptOutput, + TensorAndMetadata, + video_metadata_dict, +) from mstar.model.media_limits import VideoBudget from mstar.model.submodule_base import NodeSubmodule from mstar.model.vjepa2.components.ac_predictor import VisionTransformerPredictorAC @@ -650,7 +658,7 @@ def process_prompt( output_modalities: list[str], tensors: NameToTensorList | None = None, **kwargs, - ) -> NameToTensorList: + ) -> ProcessPromptOutput: """Pre-process raw video (and optional AC inputs) for the first forward pass. Expects the data worker to have populated ``tensors["video_inputs"]`` @@ -785,7 +793,30 @@ def process_prompt( if m is not None: out[mask_name] = [torch.as_tensor(m, dtype=torch.long)] - return out + kv_layout = {} + if self.config.predictor_kind == "ac": + # only a rollout runs the AC predictor, one frame per iteration: its + # action, state (and extrinsics) tokens, then its spatial tokens + frame = (3 if self.config.ac_predictor.use_extrinsics else 2) + self.config.grid_size ** 2 + rollout = self._initial_walk(kwargs) in (self.PREFILL_VIDEO_ROLLOUT, self.PREFILL_VIDEO_ROLLOUT_STREAMING) + kv_layout[KV_CACHE] = {"main": LabelLayout(growth=((self.ROLLOUT_LOOP_NAME, frame),) if rollout else ())} + return ProcessPromptOutput(out, {"kv_layout": kv_layout}) + + def get_loop_caps(self, **model_kwargs) -> dict[str, int]: + # the rollout's stop check ends it at its horizon + return {self.ROLLOUT_LOOP_NAME: self._rollout_horizon(model_kwargs)} + + def kv_labels(self) -> dict[str, dict[str, KVLabel]]: + if self.config.predictor_kind != "ac": + return {} + return {KV_CACHE: {"main": KVLabel(growth=( + (self.PREFILL_VIDEO_ROLLOUT, self.ROLLOUT_LOOP_NAME), + (self.PREFILL_VIDEO_ROLLOUT_STREAMING, self.ROLLOUT_LOOP_NAME), + ))}} + + def _rollout_horizon(self, model_kwargs: dict) -> int: + requested = int(model_kwargs.get("rollout_horizon", 2) or 2) + return max(1, min(requested, self.config.max_rollout_horizon)) def postprocess(self, output: torch.Tensor, modality: str, request_kwargs: dict | None = None) -> bytes: if modality == "video": @@ -901,8 +932,7 @@ def get_initial_forward_pass_args( # graph's Loop is always built with ``max_iters=config.max_rollout_horizon``. # Streaming variant uses the same horizon logic — the submodule # doesn't distinguish walks. - requested = int((model_kwargs or {}).get("rollout_horizon", 2) or 2) - step_metadata["rollout_horizon"] = max(1, min(requested, self.config.max_rollout_horizon)) + step_metadata["rollout_horizon"] = self._rollout_horizon(model_kwargs or {}) unpersist_tensors = sum([inp.tensor_info for inp in inputs], start=[]) return ForwardPassArgs( From d7dac61ec337b0f6ab8bd0674703119f93937a83 Mon Sep 17 00:00:00 2001 From: Gaurav Shah Date: Wed, 21 Oct 2026 05:50:51 +0000 Subject: [PATCH 18/26] model: whisper declares its KV layout --- mstar/model/whisper/whisper_model.py | 27 ++++++++++++++++++++++++--- 1 file changed, 24 insertions(+), 3 deletions(-) diff --git a/mstar/model/whisper/whisper_model.py b/mstar/model/whisper/whisper_model.py index 83f8660767..b880a7b8fd 100644 --- a/mstar/model/whisper/whisper_model.py +++ b/mstar/model/whisper/whisper_model.py @@ -48,7 +48,7 @@ ) from mstar.graph.base import GraphEdge, GraphNode, GraphSection, Loop, Sequential, TensorPointerInfo from mstar.graph.special_destinations import EMIT_TO_CLIENT -from mstar.model.base import ForwardPassArgs, Model +from mstar.model.base import ForwardPassArgs, KVLabel, LabelLayout, Model, ProcessPromptOutput, Span from mstar.model.submodule_base import NodeSubmodule from mstar.model.whisper.config import ( ATTN, @@ -383,7 +383,7 @@ def process_prompt( output_modalities: list[str], tensors: NameToTensorList | None = None, **kwargs, - ) -> NameToTensorList: + ) -> ProcessPromptOutput: """Extract the log-mel spectrogram and build the forced decoder prompt. The text ``prompt`` is unused — Whisper is conditioned via the @@ -411,9 +411,30 @@ def process_prompt( task=kwargs.get("task", "transcribe"), ) - return { + # the encoder takes the padded window whole, so its context is always this long + context = self.config.max_source_positions + kv_layout = { + KV_CACHE: {"main": LabelLayout( + spans=(Span("ids", "prefill", len(prompt_ids), len(prompt_ids), "text_inputs"),), + growth=(("decode_loop", 1),), + )}, + CROSS_KV_CACHE: {CONTEXT_LABEL: LabelLayout( + spans=(Span("rows", "prefill", context, context, "encoder_states"),), + )}, + } + return ProcessPromptOutput({ "audio_features": [audio_features], "text_inputs": [torch.tensor(prompt_ids, dtype=torch.long)], + }, {"kv_layout": kv_layout}) + + def get_loop_caps(self, **model_kwargs) -> dict[str, int]: + # the decode stop check ends the loop after max_tokens iterations + return {"decode_loop": self.get_max_output_tokens(**model_kwargs)} + + def kv_labels(self) -> dict[str, dict[str, KVLabel]]: + return { + KV_CACHE: {"main": KVLabel(growth=(("decode", "decode_loop"),))}, + CROSS_KV_CACHE: {CONTEXT_LABEL: KVLabel()}, } # ------------------------------------------------------------------- From 7fcf0678d349f7a040d81eb9d8c737055e39de8f Mon Sep 17 00:00:00 2001 From: Gaurav Shah Date: Wed, 21 Oct 2026 05:50:51 +0000 Subject: [PATCH 19/26] worker, conductor: refuse a paged label left undeclared --- mstar/conductor/conductor.py | 8 +++-- mstar/engine/resources/kv/config.py | 15 +++++++++ mstar/worker/engine_manager.py | 34 +++++++++++++++++++-- test/modular/test_kv_prefix_positions.py | 4 ++- test/modular/test_kv_prefix_request_keys.py | 3 +- 5 files changed, 57 insertions(+), 7 deletions(-) diff --git a/mstar/conductor/conductor.py b/mstar/conductor/conductor.py index c17e835a19..a3ea27a408 100644 --- a/mstar/conductor/conductor.py +++ b/mstar/conductor/conductor.py @@ -32,7 +32,7 @@ from mstar.distributed.base import ShardingConfig from mstar.distributed.communication import GlobalParallelConfig, WorkerParallelGroups from mstar.engine.resources import KVReqConfig, ResourceReqConfig -from mstar.engine.resources.kv.config import without_layouts +from mstar.engine.resources.kv.config import refuse_undeclared_labels, without_layouts from mstar.graph.base import GraphEdge, NodeAndGraphWalk, TensorPointerInfo from mstar.model.base import ForwardPassArgs, Model, WorkerGraph from mstar.observability import catalog @@ -603,7 +603,7 @@ def _get_resource_configs( configs = self.model.get_request_resource_configs( partition_fwd_args=partition_fwd_args, model_kwargs=model_kwargs ) - for key, labels in self.model.kv_labels().items(): + for key, labels in self._kv_labels.items(): configs.setdefault(key, KVReqConfig(needed_labels=list(labels))) return configs @@ -647,6 +647,9 @@ def _derive_worker_info(self): for worker_graph in self.worker_graphs.values(): for name, loop in worker_graph.section.get_loops().items(): self._loop_max_iters[name] = max(loop.max_iters, self._loop_max_iters.get(name, 0)) + self._kv_labels: dict[str, tuple[str, ...]] = { + key: tuple(labels) for key, labels in self.model.kv_labels().items() + } # set each group's _tp_rank (the worker's instance rank) per worker self.per_worker_sharding_config: dict[str, ShardingConfig] = {} @@ -1212,6 +1215,7 @@ def _do_ingest_request( prefix_decode = kwargs.get("prefix_decode") or {} prefix_cache = kwargs.get("prefix_cache") kv_layout = kwargs.get("kv_layout") or {} + refuse_undeclared_labels(type(self.model).__name__, self._kv_labels, kv_layout, request_data.resource_configs) for key, cfg in request_data.resource_configs.items(): cfg.apply_conductor_config( seed=seed, diff --git a/mstar/engine/resources/kv/config.py b/mstar/engine/resources/kv/config.py index eca5d1366c..0fee55f4de 100644 --- a/mstar/engine/resources/kv/config.py +++ b/mstar/engine/resources/kv/config.py @@ -328,6 +328,21 @@ def get_publish_labels( return mapping.get((node, walk), []) + +def refuse_undeclared_labels( + model: str, declared: Mapping[str, tuple[str, ...]], kv_layout: Mapping[str, Mapping], + configs: Mapping[str, ResourceReqConfig], +) -> None: + """Refuse a request whose layout leaves out a label the request opens, or lays + out one ``model`` doesn't declare. The model is at fault and not the client, so + this raises RuntimeError, which fails the request with a 500.""" + for key in declared.keys() | kv_layout.keys(): + laid_out = set(kv_layout.get(key, ())) + if extra := sorted(laid_out - set(declared.get(key, ()))): + raise RuntimeError(f"{model} returned a kv_layout for {key!r}/{extra[0]!r}, which it does not declare") + if missing := sorted(configs[key].opened_labels() - laid_out): + raise RuntimeError(f"{model} returned no kv_layout for {key!r}/{missing[0]!r}, which the request opens") + def without_layouts(configs: Mapping[str, ResourceReqConfig]) -> dict[str, ResourceReqConfig]: """``configs`` with each KV config's layouts and caps left out. Only the KV manager's ingest reads them, and the request's info that carries the configs rides every frame between workers.""" diff --git a/mstar/worker/engine_manager.py b/mstar/worker/engine_manager.py index 6132d62916..4eab7e3a3f 100644 --- a/mstar/worker/engine_manager.py +++ b/mstar/worker/engine_manager.py @@ -11,6 +11,7 @@ ResourceReqConfig, apply_yaml_overrides, ) +from mstar.engine.resources.kv.config import KVSpec, PagedKVConfig from mstar.engine.resources.kv.transfer import TransferEngineInfo from mstar.engine.resources.position.config import PositionSpec, PosScheme from mstar.graph.runtime.base import GraphRuntime @@ -120,12 +121,16 @@ def refuse_combined_walks_across_tp_groups( def _refuse_unknown_walks(model: Model) -> None: - """Refuse a declared label that names a walk the model never runs. + """Refuse a declared label that names a walk the model never runs, or a + growth loop its walk doesn't run. Only the named walks may file pages or be probed, so a name that matches - nothing leaves that cache empty for good, and nothing says why. + nothing leaves that cache empty for good, and nothing says why. A loop takes + its cap by name, and a loop built without a name gets a different one in every + process, so no declaration can refer to it. """ - walks = set(model.get_graph_walk_graphs()) + graphs = model.get_graph_walk_graphs() + walks = set(graphs) for key, by_label in model.kv_labels().items(): for label, kv_label in by_label.items(): named = {walk for walk, _ in kv_label.growth} @@ -138,6 +143,28 @@ def _refuse_unknown_walks(model: Model) -> None: f"{unknown}, which it never runs; its walks are " f"{sorted(walks)}" ) + for walk, loop in kv_label.growth: + loops = graphs[walk].get_loops() + if loop not in loops: + raise ValueError( + f"{type(model).__name__} declares {key!r}/{label!r} growing on " + f"{loop!r}, which {walk!r} doesn't run; its loops are {sorted(loops)}" + ) + + +def _refuse_undeclared_pools(specs: list[NodeResourceSpec], model: Model) -> None: + """Refuse a paged pool the model declares no labels for, since no request + could lay out its pages there. A ring pool needs none: its size is fixed + per request. + """ + declared = model.kv_labels() + for spec in specs: + if isinstance(spec, KVSpec) and isinstance(spec.config, PagedKVConfig) and not declared.get(spec.resource_key): + raise ValueError( + f"{type(model).__name__} declares no labels for its paged pool " + f"{spec.resource_key!r}; add them to `kv_labels`, and return each " + "request's layout for them from `process_prompt`" + ) def _refuse_unservable_walks(specs: list[NodeResourceSpec], model: Model) -> None: @@ -212,6 +239,7 @@ def build( _refuse_uncacheable_positions(specs, model) _refuse_unskippable_resources(specs, model) _refuse_unknown_walks(model) + _refuse_undeclared_pools(specs, model) _refuse_unservable_walks(specs, model) # Resolve autocast dtype: explicit YAML config wins; otherwise we diff --git a/test/modular/test_kv_prefix_positions.py b/test/modular/test_kv_prefix_positions.py index 2dc3f7bcdb..78133d78c3 100644 --- a/test/modular/test_kv_prefix_positions.py +++ b/test/modular/test_kv_prefix_positions.py @@ -34,6 +34,7 @@ from mstar.engine.resources.position.manager import RopeManager from mstar.engine.resources.spec import NodeResourceSpec from mstar.engine.resources.step import Segment, SubmoduleStep +from mstar.graph.base import GraphNode, Loop from mstar.model.base import KVLabel, PrefixKeys from mstar.worker.engine_manager import ( _refuse_uncacheable_positions, @@ -318,7 +319,8 @@ def kv_labels(self): ))}} def get_graph_walk_graphs(self): - return {walk: None for walk in self._runs} + llm = GraphNode(name="llm", input_names=["text_inputs"], outputs=[]) + return {walk: Loop(name="decode_loop", max_iters=8, outputs=[], section=llm) for walk in self._runs} def test_a_stream_naming_a_walk_the_model_never_runs_is_refused_at_load(): diff --git a/test/modular/test_kv_prefix_request_keys.py b/test/modular/test_kv_prefix_request_keys.py index 62dde9156b..846175ccd8 100644 --- a/test/modular/test_kv_prefix_request_keys.py +++ b/test/modular/test_kv_prefix_request_keys.py @@ -364,8 +364,9 @@ def _conductor_configs(model, model_kwargs: dict | None = None) -> dict: """The configs a request is opened with, as the conductor resolves them: the model's own, a default for anything it declared a stream for, and then each handed its own chain.""" + declared = {key: tuple(labels) for key, labels in model.kv_labels().items()} configs = Conductor._get_resource_configs( - SimpleNamespace(model=model), model_kwargs or {}, {}, + SimpleNamespace(model=model, _kv_labels=declared), model_kwargs or {}, {}, ) for key, cfg in configs.items(): _apply_chain(cfg, key, model_kwargs) From 4da8d7294bddba1042c7d20f8e8656e3636983e4 Mon Sep 17 00:00:00 2001 From: Gaurav Shah Date: Wed, 21 Oct 2026 05:50:58 +0000 Subject: [PATCH 20/26] test: footprints, span checks, loop caps and model layouts --- test/modular/test_kv_footprint.py | 115 +++ test/modular/test_kv_footprint_checks.py | 214 ++++++ test/modular/test_kv_layout_intake.py | 282 ++++++++ test/modular/test_model_kv_layouts.py | 851 +++++++++++++++++++++++ test/modular/test_qwen3_5_video.py | 6 +- test/rust/test_runtime_parity.py | 118 +++- 6 files changed, 1583 insertions(+), 3 deletions(-) create mode 100644 test/modular/test_kv_footprint.py create mode 100644 test/modular/test_kv_footprint_checks.py create mode 100644 test/modular/test_kv_layout_intake.py create mode 100644 test/modular/test_model_kv_layouts.py diff --git a/test/modular/test_kv_footprint.py b/test/modular/test_kv_footprint.py new file mode 100644 index 0000000000..f643b88601 --- /dev/null +++ b/test/modular/test_kv_footprint.py @@ -0,0 +1,115 @@ +"""A stream's footprint is never below the pages the manager gives it. + +The footprint is computed from the layout alone, before the request runs, so +the arithmetic copies the manager's: per-label rounding in `_alloc`, pages a +non-committing step leaves behind, and the page-floored release under a +retention policy. Each test runs the real manager and compares what the stream +held with what the layout predicted. A footprint below that would let a pool +promise room it doesn't have. +""" + +from __future__ import annotations + +import torch + +from mstar.communication.tensors import LocalTransferEngine +from mstar.engine.resources.kv.config import ( + KVStep, + PagedKVConfig, + PrefixSpan, + RetentionPolicy, + StreamLayout, +) +from mstar.engine.resources.kv.manager import KVManager +from mstar.engine.resources.kv.transfer import TransferEngineInfo +from mstar.engine.resources.step import Segment, StepContext + +PS = 8 + + +def _manager() -> KVManager: + return KVManager( + cfg=PagedKVConfig( + num_layers=1, num_kv_heads=1, head_dim=4, max_seq_len=512 * PS, + max_num_pages=512, page_size=PS, + ), + name="kv", joint_comm_group=None, + transfer_engine_info=TransferEngineInfo("h", "h", LocalTransferEngine("h")), + device=torch.device("cpu"), dtype=torch.float32, + ) + + +def _grow( + kv: KVManager, rid: str, span: int, label: str = "main", commit: bool = True, + retention: RetentionPolicy | None = None, +) -> int: + """Run one step and return the pages ``label`` holds at its peak, right after admit.""" + step = KVStep( + segments=(Segment(rid, label, span),), commit=commit, + retention={} if retention is None else {(rid, label): retention}, + ) + ctx = StepContext(request_ids=(rid,), graph_walk="w", slot=0, capture=False) + assert kv.admit(step, ctx).ok + held = len(kv._streams[rid][label].page_indices) + kv.plan(step, ctx) + kv.commit(step, ctx) + return held + + +def _spans(*lengths: int) -> tuple[PrefixSpan, ...]: + return tuple(PrefixSpan(length, length, "w") for length in lengths) + + +def test_each_label_rounds_up_on_its_own(): + kv = _manager() + kv.ingest_request("r") + held = _grow(kv, "r", 3, "main") + _grow(kv, "r", 3, "uncond") + layout = StreamLayout(spans=_spans(3)) + + assert held == 2 * layout.pages({}, PS), "rounding the labels' sum once is a page short" + + +def test_growth_is_its_cap_times_its_slots(): + kv = _manager() + kv.ingest_request("r") + held = _grow(kv, "r", 5) + for _ in range(7): + held = _grow(kv, "r", 2) + layout = StreamLayout(spans=_spans(5), growth=(("decode_loop", 2),)) + + assert held == layout.pages({"decode_loop": 7}, PS) == 3 + + +def test_scratch_counts_once_however_often_it_runs(): + """A non-committing step writes past the committed length and leaves its pages with the stream.""" + kv = _manager() + kv.ingest_request("r") + held = _grow(kv, "r", 6) + for _ in range(3): + held = max(held, _grow(kv, "r", 11, commit=False)) + layout = StreamLayout(spans=_spans(6), scratch=11) + + assert held == layout.pages({}, PS) == 3 + + +def test_a_window_holds_its_peak_not_its_sum(): + """Window commits under a retention policy, each denoised first by a step that commits nothing. + + The footprint must cover every admit. Once the window fills before the largest + step runs, the footprint may exceed the peak by one page at most. + """ + for prefix, budget, window in [(13, 24, 9), (16, 16, 16), (5, 40, 7), (0, 3, 12)]: + kv = _manager() + kv.ingest_request("r") + policy = RetentionPolicy(context_budget=budget, protected_prefix=prefix) + peak = _grow(kv, "r", prefix) if prefix else 0 + for _ in range(12): + peak = max(peak, _grow(kv, "r", window, commit=False)) + peak = max(peak, _grow(kv, "r", window, retention=policy)) + layout = StreamLayout( + spans=_spans(*([prefix] if prefix else []), *[window] * 12), + scratch=window, retention=policy, + ) + + assert peak <= layout.pages({}, PS), f"{prefix=} {budget=} {window=}: a window outgrew its footprint" + assert layout.pages({}, PS) - peak <= 1, f"{prefix=} {budget=} {window=}: the footprint over-counts" diff --git a/test/modular/test_kv_footprint_checks.py b/test/modular/test_kv_footprint_checks.py new file mode 100644 index 0000000000..99216e5683 --- /dev/null +++ b/test/modular/test_kv_footprint_checks.py @@ -0,0 +1,214 @@ +"""A write past a request's layout is logged, and the debug checks report where. + +The always-on line names the label and its layout the first time a stream takes +a page past its footprint. Under ``MSTAR_KV_DEBUG_ASSERTS`` each commit is +checked against the layout itself: the walk that wrote each span, the growth +after the spans, the scratch of a step that commits nothing, and the retention a +step declares. A span may fill over several steps, from a chunked prompt or the +tail a cache hit leaves, and neither is a violation. +""" + +from __future__ import annotations + +import logging +import sys + +sys.path.insert(0, ".") + +import pytest +from test_kv_prefix_lease import NODE, WALK, _ctx, _grow, _keys, _layout, _manager, _seed, _StubTransfer + +from mstar.engine.resources.kv import manager as manager_mod +from mstar.engine.resources.kv.config import KVReqConfig, KVStep, PrefixSpan, RetentionPolicy, StreamLayout +from mstar.engine.resources.step import Segment, StepContext +from mstar.model.submodule_base import ARNodeInputs + +LOGGER = "mstar.engine.resources.kv.manager" + + +@pytest.fixture(autouse=True) +def _stub_transfer(monkeypatch): + monkeypatch.setattr(manager_mod, "KVTransferManager", _StubTransfer) + + +@pytest.fixture +def debug(monkeypatch): + monkeypatch.setattr(manager_mod, "_DEBUG_ASSERTS", True) + + +@pytest.fixture +def always_on(monkeypatch): + """Turn the debug checks off. These tests write past a layout on purpose, and the checks would fail that commit.""" + monkeypatch.setattr(manager_mod, "_DEBUG_ASSERTS", False) + + +def _laid_out(kv, rid: str, spans=(PrefixSpan(20, 20, WALK),), **layout) -> None: + """Ingest ``rid`` with one layout on ``main``, whose ``decode_loop`` runs 3 times.""" + kv.ingest_request(rid, KVReqConfig( + kv_layout={"main": StreamLayout(spans=spans, **layout)}, loop_caps={"decode_loop": 3}, + )) + + +def test_a_label_past_its_footprint_logs_once(always_on, caplog): + kv = _manager() + _laid_out(kv, "r0") + _laid_out(kv, "r1") + + with caplog.at_level(logging.ERROR, logger=LOGGER): + _grow(kv, "r0", 20) + assert not caplog.records, "a stream inside its footprint was reported" + _grow(kv, "r0", 20) + _grow(kv, "r0", 20) + _grow(kv, "r1", 60) + + assert len(caplog.records) == 1, "an overrun on one label was reported more than once" + assert "footprint of 2" in caplog.text and "PrefixSpan(length=20" in caplog.text, ( + "the line doesn't name the layout that under-counted" + ) + + +def test_a_write_from_another_walk_fails_its_commit(debug): + kv = _manager() + _laid_out(kv, "r0") + step = KVStep(segments=(Segment("r0", "main", 20),)) + ctx = StepContext(request_ids=("r0",), graph_walk="decode", slot=0, capture=False) + assert kv.admit(step, ctx).ok + + with pytest.raises(AssertionError, match="decode wrote into span 0"): + kv.commit(step, ctx) + + +def test_growth_past_its_bound_fails_its_commit(debug): + kv = _manager() + _laid_out(kv, "r0", growth=(("decode_loop", 2),)) + _grow(kv, "r0", 20) + for _ in range(3): + _grow(kv, "r0", 2) + + with pytest.raises(AssertionError, match="past its layout's spans and growth"): + _grow(kv, "r0", 1) + + +def test_scratch_past_its_declaration_fails_its_admit(debug): + kv = _manager() + _laid_out(kv, "r0", scratch=4) + _grow(kv, "r0", 20) + + with pytest.raises(AssertionError, match="past its layout's scratch of 4"): + kv.admit(KVStep(segments=(Segment("r0", "main", 5),), commit=False), _ctx("r0")) + + +def test_a_retention_other_than_the_layouts_fails_its_commit(debug): + kv = _manager() + _laid_out(kv, "r0", retention=RetentionPolicy(context_budget=16, protected_prefix=4)) + + with pytest.raises(AssertionError, match="a step declares"): + _grow(kv, "r0", 20, retention={("r0", "main"): RetentionPolicy(context_budget=32, protected_prefix=4)}) + + +def test_a_prompt_chunked_over_three_steps_fires_nothing(debug): + kv = _manager() + _laid_out(kv, "r0", spans=(PrefixSpan(30, 30, WALK), PrefixSpan(5, 5, "decode"))) + + for _ in range(3): + _grow(kv, "r0", 10) + + assert kv._streams["r0"]["main"].span_cursor.done == 1, "the chunks never completed their span" + + +def test_the_tail_after_a_cache_hit_fires_nothing(debug, monkeypatch): + checked = [] + check = manager_mod.KVManager._assert_layout + monkeypatch.setattr(manager_mod.KVManager, "_assert_layout", lambda self, segment, *args: ( + checked.append(segment.request_id), check(self, segment, *args), + )) + kv = _manager() + tokens = list(range(100)) + keys = _seed(kv, tokens) + kv.ingest_request("r1", KVReqConfig(prefix_keys={"main": _keys(tokens)}, kv_layout=_layout(100), loop_caps={})) + inputs = ARNodeInputs(input_seq_len=100) + prefix = kv.resolve_cached_prefix("r1", NODE, WALK, inputs) + kv.apply_cached_prefix("r1", NODE, WALK, inputs, prefix) + + _grow(kv, "r1", 100 - prefix.tokens) + + assert "r1" in checked, "the tail's commit was never checked against its layout" + assert keys and kv._streams["r1"]["main"].span_cursor.done == 1, "the tail never completed the span the hit began" + + +def test_intake_drops_the_layout_of_a_label_the_request_does_not_open(): + """BAGEL lays out its guidance labels for every request, for the deployments where the + request opens them. A request is held only to the labels it opens.""" + config = KVReqConfig(needed_labels=["main"]) + config.apply_conductor_config(kv_layout={ + "main": [[[20, 20, WALK]], [["decode_loop", 1]], 0, None, 0], + "cfg_img": [[], [["decode_loop", 1]], 40, None, 0], + }, loop_caps={"decode_loop": 3}) + + assert list(config.kv_layout) == ["main"], "the request would be held to a label it never opens" + + +def test_a_label_the_request_does_not_open_holds_no_pages(debug, caplog, monkeypatch): + """A batch can write a label for a row that never reads it. Those pages are outside the + request's footprint, which the always-on line says.""" + checked = [] + check = manager_mod.KVManager._assert_layout + monkeypatch.setattr(manager_mod.KVManager, "_assert_layout", lambda self, segment, *args: ( + checked.append(segment.label), check(self, segment, *args), + )) + kv = _manager() + _laid_out(kv, "r0") + + with caplog.at_level(logging.WARNING, logger=LOGGER): + _grow(kv, "r0", 20) + _grow(kv, "r0", 60, label="cfg_img") + kv.remove_request("r0") + + assert checked == ["main"], "a label the request doesn't open was checked against its layout" + assert [r.levelno for r in caplog.records] == [logging.ERROR], "expected one overrun line and no under-use line" + assert "label cfg_img against a footprint of 0" in caplog.text, ( + "pages a batch took for the request on a label it doesn't open left no line" + ) + + +def test_a_label_the_request_does_not_open_does_not_hide_a_real_overrun(always_on, caplog): + kv = _manager() + _laid_out(kv, "r0") + kv.ingest_request("r1", KVReqConfig(kv_layout={ + "main": StreamLayout(spans=(PrefixSpan(20, 20, WALK),)), + "cfg_img": StreamLayout(spans=(PrefixSpan(20, 20, WALK),)), + }, loop_caps={})) + + with caplog.at_level(logging.ERROR, logger=LOGGER): + _grow(kv, "r0", 60, label="cfg_img") + _grow(kv, "r1", 60, label="cfg_img") + + assert "request r1 took 4 pages on label cfg_img against a footprint of 2" in caplog.text, ( + "a batch's write to a label one request doesn't open hid another's overrun there" + ) + + +def test_a_keyed_label_the_request_does_not_open_gets_no_chain(): + """A prefix chain reads its stream's span cursor, which a label the request doesn't open never gets.""" + kv = _manager() + kv.ingest_request("r0", KVReqConfig( + prefix_keys={"cfg_img": _keys(list(range(40)))}, kv_layout=_layout(40), loop_caps={}, + )) + + _grow(kv, "r0", 20, label="cfg_img") + + stream = kv._streams["r0"]["cfg_img"] + assert stream.span_cursor is None and stream.chain is None, "a chain without a cursor fails the next probe" + + +def test_a_label_left_under_its_footprint_is_reported_at_removal(debug, caplog): + kv = _manager() + _laid_out(kv, "r0", growth=(("decode_loop", 16),)) + _grow(kv, "r0", 20) + + with caplog.at_level(logging.WARNING, logger=LOGGER): + kv.remove_request("r0") + + assert "held at most 2 pages on label main against a footprint of 5" in caplog.text, ( + "an over-counted layout left no line to find it by" + ) diff --git a/test/modular/test_kv_layout_intake.py b/test/modular/test_kv_layout_intake.py new file mode 100644 index 0000000000..0ea7f5649f --- /dev/null +++ b/test/modular/test_kv_layout_intake.py @@ -0,0 +1,282 @@ +"""A request's KV layout travels from `process_prompt` to the stream that holds its pages. + +`process_prompt` is where every length a request writes is known: it tokenizes +and runs the media processors. Its layout crosses the preprocess worker and the +conductor as plain rows, the conductor resolves each growth loop's iterations, +and the manager turns each label's layout into a page footprint when the stream +opens. Keys are optional on that path, so a request with the prefix cache off +still arrives with its footprint. A client can't send a layout of its own. +""" + +from __future__ import annotations + +import queue +import sys +import threading +from collections import deque +from types import SimpleNamespace + +sys.path.insert(0, ".") + +import pytest +import torch +import yaml +from test_kv_prefix_request_keys import _StubTensorManager + +import mstar.conductor.conductor as conductor_mod +from mstar.api_server.data_worker import PreprocessWorkerThread +from mstar.api_server.request_types import PreprocessInput +from mstar.communication import wire +from mstar.communication.tensors import LocalTransferEngine +from mstar.conductor.conductor import Conductor +from mstar.conductor.request_info import CurrentForwardConductorMetadata +from mstar.engine.resources.kv.config import KVReqConfig, KVSpec, PagedKVConfig, RetentionPolicy, StreamLayout +from mstar.engine.resources.kv.manager import KVManager +from mstar.engine.resources.kv.transfer import TransferEngineInfo +from mstar.graph.base import GraphEdge, GraphNode, Loop, Sequential +from mstar.graph.special_destinations import EMIT_TO_CLIENT +from mstar.model.base import ForwardPassArgs, KVLabel, LabelLayout, Model, PrefixKeys, ProcessPromptOutput, Span +from mstar.utils.ipc_format import ConductorMessage, ConductorMessageType, NewRequestConductor, SetupDone + +PS = 16 +PROMPT = 100 +LAYOUT = {"kv": { + "main": LabelLayout( + spans=(Span("ids", "prefill", PROMPT, PROMPT, "text_inputs"),), growth=(("decode_loop", 1),), + ), + "uncond": LabelLayout( + spans=(Span("rows", "prefill", 10, 10, "cond"),), growth=(("talk_loop", 2),), scratch=7, + ), +}} + + +class _Comm: + """Records sends, and replays one batch of inbound messages per call.""" + + def __init__(self): + self.sent = [] + self.batches = deque() + + def send(self, entity_id, msg): + self.sent.append(msg) + + def get_all_new_messages(self): + return self.batches.popleft() if self.batches else [] + + +class _ScriptDoneError(Exception): + pass + + +class _Model(Model): + """Lays out two labels of one pool, and caps ``decode_loop`` at the request's ``max_output_tokens``.""" + + def __init__(self, layout=LAYOUT): + self._layout = layout + + def get_node_resources(self): + return [KVSpec(resource_key="kv", nodes={"llm"}, config=PagedKVConfig( + num_layers=1, num_kv_heads=1, head_dim=8, max_seq_len=4096, max_num_pages=64, page_size=PS, + ))] + + def kv_labels(self): + return {"kv": { + "main": KVLabel( + growth=(("decode", "decode_loop"),), keys=PrefixKeys("text_inputs", ("prefill",)), + ), + "uncond": KVLabel(growth=(("decode", "talk_loop"),)), + }} + + def get_graph_walk_graphs(self): + llm = GraphNode(name="llm", input_names=["text_inputs"], outputs=[ + GraphEdge(next_node=EMIT_TO_CLIENT, name="new_token", output_modality="text"), + ]) + return {"prefill": llm, "decode": Sequential(sections=[ + Loop(name="decode_loop", max_iters=50, outputs=[], section=llm), + Loop(name="talk_loop", max_iters=40, outputs=[], section=GraphNode( + name="talker", input_names=["new_token"], outputs=[], + )), + ])} + + def get_loop_caps(self, **model_kwargs): + return {"decode_loop": model_kwargs["max_output_tokens"]} + + def get_request_resource_configs(self, partition_fwd_args, model_kwargs=None): + return {"kv": KVReqConfig(needed_labels=["main", "uncond"])} + + def get_initial_forward_pass_args(self, partition_name, input_modalities, output_modalities, + input_signals, model_kwargs=None): + return ForwardPassArgs( + full_metadata=CurrentForwardConductorMetadata(graph_walk="prefill", is_prefill=True), + inputs=[], unpersist_tensors=[], + ) + + def get_partition_forward_pass_args(self, *args, **kwargs): + raise AssertionError("the request never leaves its first walk") + + def max_input_files(self): + return {} + + def process_prompt(self, *args, **kwargs): + tensors = {"text_inputs": [torch.arange(PROMPT)]} + return tensors if self._layout is None else ProcessPromptOutput(tensors, {"kv_layout": self._layout}) + + def postprocess(self, *args, **kwargs): + return b"" + + def get_submodule(self, *args, **kwargs): + return None + + +def _preprocess(model: _Model, model_kwargs: dict) -> dict: + """The kwargs the preprocess worker sends the conductor.""" + worker = PreprocessWorkerThread( + *(queue.Queue() for _ in range(8)), stop_event=threading.Event(), communicator=_Comm(), + tensor_manager=_StubTensorManager(), model=model, + model_config={"model": "stub", "resources": {"kv": {"page_size": PS}}}, + ) + worker._process_input(PreprocessInput( + request_id="r", text="hello", file_paths=None, input_modalities=["text"], + output_modalities=["text"], model_kwargs=model_kwargs, + )) + return worker.communicator.sent[-1].body.model_kwargs + + +def _conduct(model: _Model, model_kwargs: dict, tmp_path, monkeypatch) -> _Comm: + """Run the real `Conductor.run` to ingest the request. Return the communicator, which holds the + conductor and every message the conductor sent.""" + path = tmp_path / "cfg.yaml" + path.write_text(yaml.safe_dump({ + "max_seq_len": 4096, "node_groups": [{"node_names": ["llm", "talker"], "ranks": [0]}], + })) + comm = _Comm() + monkeypatch.setattr(Conductor, "_launch_workers", lambda self: None) + monkeypatch.setattr(conductor_mod, "make_communicator", lambda **kw: comm) + conductor = comm.conductor = Conductor(model, str(path), socket_path_prefix=str(tmp_path / "s")) + comm.batches.extend([ + [ConductorMessage(ConductorMessageType.SETUP_DONE, SetupDone(w)) for w in conductor.worker_ids], + [ConductorMessage(ConductorMessageType.NEW_REQUEST, NewRequestConductor( + request_id="r", initial_signals={}, initial_input_modalities=["text"], + initial_output_modalities=["text"], input_metadata={}, model_kwargs=model_kwargs, + ))], + ]) + + def poll(): + if not comm.batches: + raise _ScriptDoneError + + conductor._poll_worker_liveness = poll + with pytest.raises(_ScriptDoneError): + conductor.run() + return comm + + +def _ingest(model: _Model, model_kwargs: dict, tmp_path, monkeypatch): + """The NewRequest the conductor sent its worker.""" + sent = _conduct(model, model_kwargs, tmp_path, monkeypatch).sent + return next(msg.body for msg in sent if getattr(msg.message_type, "name", "") == "NEW_REQUEST") + + +def _manager() -> KVManager: + return KVManager( + cfg=PagedKVConfig( + num_layers=1, num_kv_heads=1, head_dim=8, max_seq_len=4096, max_num_pages=64, page_size=PS, + ), + name="kv", joint_comm_group=None, + transfer_engine_info=TransferEngineInfo("h", "h", LocalTransferEngine("h")), + device=torch.device("cpu"), dtype=torch.float32, + ) + + +@pytest.mark.parametrize("prefix_cache", [True, False]) +def test_a_layout_reaches_its_streams_as_footprints(tmp_path, monkeypatch, prefix_cache): + model = _Model() + kwargs = _preprocess(model, {"max_output_tokens": 30, "prefix_cache": prefix_cache}) + new_request = _ingest(model, kwargs, tmp_path, monkeypatch) + cfg = new_request.request_info.resource_configs["kv"] + kv = _manager() + + kv.ingest_request("r", cfg) + + assert ("prefix_keys" in kwargs) == prefix_cache + assert new_request.loop_caps == cfg.loop_caps == {"decode_loop": 30, "talk_loop": 40}, ( + "an uncapped loop runs to its max_iters" + ) + assert kv._streams["r"]["main"].footprint == -(-(PROMPT + 30) // PS) + assert kv._ensure_label("r", "uncond").footprint == -(-(10 + 2 * 40 + 7) // PS) + + +def test_the_new_request_crosses_the_wire_unchanged(tmp_path, monkeypatch): + """The worker sizes its pages from the frame the conductor encodes, retention included.""" + uncond = LAYOUT["kv"]["uncond"]._replace(retention=RetentionPolicy(context_budget=64, protected_prefix=10)) + model = _Model({"kv": {**LAYOUT["kv"], "uncond": uncond}}) + new_request = _ingest(model, _preprocess(model, {"max_output_tokens": 30}), tmp_path, monkeypatch) + + decoded = wire.decode(wire.encode(new_request)) + + sent, got = (msg.request_info.resource_configs["kv"] for msg in (new_request, decoded)) + assert decoded.loop_caps == {"decode_loop": 30, "talk_loop": 40} + assert got.kv_layout == sent.kv_layout and got.loop_caps == sent.loop_caps, ( + "the worker would size its pages from another layout than the conductor's" + ) + + +def test_a_cap_past_its_loop_reaches_the_workers_as_its_max_iters(tmp_path, monkeypatch): + """Rust's graph runtime holds a cap as a u32, so a cap from a large ``max_output_tokens`` would fail + the request there but not in Python.""" + model = _Model() + + new_request = _ingest(model, _preprocess(model, {"max_output_tokens": 2**40}), tmp_path, monkeypatch) + + assert new_request.loop_caps == {"decode_loop": 50, "talk_loop": 40}, "a cap went out past its loop's max_iters" + + +def test_only_the_new_request_carries_the_layout(tmp_path, monkeypatch): + """Each worker's KV manager ingests the layout from its NewRequest. The conductor's later messages + carry the request's configs to the workers again, so its own copy drops the layout and caps.""" + model = _Model() + + comm = _conduct(model, _preprocess(model, {"max_output_tokens": 30}), tmp_path, monkeypatch) + + [new_request] = [msg.body for msg in comm.sent if getattr(msg.message_type, "name", "") == "NEW_REQUEST"] + sent, kept = new_request.request_info.resource_configs["kv"], comm.conductor.requests["r"].resource_configs["kv"] + assert sent.kv_layout and sent.loop_caps, "the worker's KV manager would ingest no layout" + assert kept.kv_layout is None and kept.loop_caps is None, "the conductor's later messages would carry the layout" + + +def test_a_worker_keeps_the_layout_out_of_the_frames_it_sends(): + """The KV manager keeps the config it ingests. The request's info the worker stores rides every + frame it sends, so that copy drops the layout and caps.""" + from test_rid_handles import _worker + + from mstar.worker.worker import Worker + + w, _ = _worker() + ingested = {} + w.request_state.add_request = lambda request_id, request_info: None + w.request_state.per_request_info = {0: SimpleNamespace(stream_buffers={})} + w.engine_manager = SimpleNamespace(add_request=lambda rid, cfgs: ingested.update(cfgs), evictable_nodes=list) + w.tensor_manager = SimpleNamespace( + register_request=lambda rid, sharding: None, start_read_tensors=lambda rid, inputs, graph_walk=None: [], + ) + w._draining_rids, w._last_active, w._my_consumer_connections, w._unprocessed_messages = set(), {}, [], {} + w.wakeup_event = SimpleNamespace(register_futures=lambda futures: None) + config = KVReqConfig(kv_layout={"main": StreamLayout()}, loop_caps={"decode_loop": 30}) + info = SimpleNamespace(rid_handle=-1, resource_configs={"kv": config}, graph_walk="prefill", partition_name="p") + + Worker._add_new_request(w, SimpleNamespace( + request_id="r", request_info=info, partition_worker_graph_ids={}, worker_graph_to_workers={}, + initial_inputs=[], loop_caps={}, + )) + + assert ingested["kv"].kv_layout, "the KV manager was handed a config without its layout" + stored = info.resource_configs["kv"] + assert stored.kv_layout is None and stored.loop_caps is None, "the worker's frames would carry the layout" + + +def test_a_layout_the_client_sent_is_dropped(tmp_path, monkeypatch): + model = _Model() + kwargs = _preprocess(model, {"max_output_tokens": 30, "kv_layout": {"kv": {"main": [[[1, 1, "x"]], [], 0, None]}}}) + cfg = _ingest(model, kwargs, tmp_path, monkeypatch).request_info.resource_configs["kv"] + + assert cfg.kv_layout["main"].spans[0].length == PROMPT, "a client's layout would size another request's pages" diff --git a/test/modular/test_model_kv_layouts.py b/test/modular/test_model_kv_layouts.py new file mode 100644 index 0000000000..1235170795 --- /dev/null +++ b/test/modular/test_model_kv_layouts.py @@ -0,0 +1,851 @@ +"""Each model lays out every label it declares, and caps each growth loop where its stop check ends it. + +A footprint is computed from the layout, so a request's layout names every +label it opens, empty where the request writes nothing. A growth loop's cap must +end the loop on the iteration the model's own stop check does: an earlier cap +would cut an output short, and a later one would let speculation write a slot +the footprint doesn't count. Each test calls the +model's real `check_stop` with the loop's count, as the worker hands it over. +""" + +from __future__ import annotations + +import os +from types import SimpleNamespace + +import pytest +import torch + +from mstar.conductor.request_info import CurrentForwardPassInfo +from mstar.engine.engine import ForwardPassInfoWrapper +from mstar.engine.resources import SamplingReqConfig +from mstar.model.base import LabelLayout +from mstar.model.submodule_base import InputMetadata + +NOT_EOS = 42 + + +def _stops_after(check_stop, loop: str, walk: str, max_iters: int, outputs: dict, **info) -> int: + """How many iterations of ``loop`` ``check_stop`` allows: the first it stops at, plus one.""" + resource_configs = info.pop("resource_configs", {"sampler": SamplingReqConfig(ignore_eos=True)}) + for k in range(max_iters): + request_info = ForwardPassInfoWrapper( + info=CurrentForwardPassInfo( + request_id="r", graph_walk=walk, fwd_index=0, random_seed=0, + resource_configs=resource_configs, **info, + ), + metadata=InputMetadata(dynamic_loop_iter_counts={loop: k}), + ) + if loop in check_stop("r", request_info, outputs): + return k + 1 + return max_iters + + +# ── Orpheus ───────────────────────────────────────────────────────────── + + +def _orpheus(): + from mstar.model.orpheus.config import OrpheusModelConfig + from mstar.model.orpheus.orpheus_model import OrpheusModel + + model = object.__new__(OrpheusModel) + model.config = OrpheusModelConfig() + model.tokenizer = lambda text, return_tensors=None: SimpleNamespace( + input_ids=torch.tensor([[ord(c) for c in text]], dtype=torch.long), + ) + return model + + +@pytest.mark.parametrize("max_tokens", [1, 64]) +def test_orpheus_caps_decode_where_its_stop_check_ends_it(max_tokens): + from mstar.model.orpheus.submodules import OrpheusLLMSubmodule + + model = _orpheus() + llm = OrpheusLLMSubmodule( + language_model=SimpleNamespace(model=SimpleNamespace(embed_tokens=None), lm_head=None), config=model.config, + ) + + cap = model.get_loop_caps(max_output_tokens=max_tokens)["decode_loop"] + + assert cap == _stops_after( + llm.check_stop, "decode_loop", "decode", 2048, + {"new_token": [torch.tensor([NOT_EOS])]}, max_tokens=max_tokens, + ), "the cap and the stop check end decode on different iterations" + + +# ── Bagel ─────────────────────────────────────────────────────────────── + + +def _bagel(): + from test_multimodal_ordering import _StubTokenizer + + from mstar.model.bagel.bagel_model import BagelModel + + model = object.__new__(BagelModel) + model.config = SimpleNamespace( + think_mode=False, cfg_text_scale=4.0, cfg_img_scale=1.0, cfg_interval=[0.4, 1.0], + cfg_renorm_type="global", cfg_renorm_min=0.0, vit_config=SimpleNamespace(patch_size=14), + vit_max_num_patch_per_side=70, latent_downsample=16, max_latent_size=64, + temperature=0.0, top_k=-1, top_p=1.0, repetition_penalty=1.0, ignore_eos=False, vocab_size=1000, + ) + model.tokenizer = _StubTokenizer() + model.boi_token_id = _StubTokenizer.SPECIALS["<|vision_start|>"] + model.eoi_token_id = _StubTokenizer.SPECIALS["<|vision_end|>"] + model._has_cfg_parallel = model._has_llm_disaggregation = False + return model + + +@pytest.mark.parametrize("max_tokens", [1, 64]) +def test_bagel_caps_decode_where_its_stop_check_ends_it(max_tokens): + from test_bagel_guidance_sink import _llm + + llm = _llm() + + cap = _bagel().get_loop_caps(max_output_tokens=max_tokens)["decode_loop"] + + assert cap == _stops_after( + llm.check_stop, "decode_loop", "decode", 2048, + {"new_token": [torch.tensor([NOT_EOS])]}, max_tokens=max_tokens, + ), "the cap and the stop check end decode on different iterations" + + +_BAGEL_SHAPES = { + "text chat": (["text"], ["text"], {}), + "understanding an image": (["image", "text"], ["text"], {}), + "text to image": (["text"], ["image"], {"think_mode": True}), + "text to image without guidance": (["text"], ["image"], {"cfg_text_scale": 1.0}), + "an edit": (["image", "text"], ["image"], {}), + "an edit without guidance": (["image", "text"], ["image"], {"cfg_text_scale": 1.0, "cfg_img_scale": 1.0}), +} + + +@pytest.mark.parametrize("cfg_parallel", [False, True], ids=["one LLM node", "CFG-parallel"]) +@pytest.mark.parametrize(("input_modalities", "output_modalities", "kwargs"), _BAGEL_SHAPES.values(), ids=_BAGEL_SHAPES) +def test_bagel_lays_out_every_label_its_request_opens(input_modalities, output_modalities, kwargs, cfg_parallel): + """The layout is built in the API process, where the deployment isn't known, and the + labels a request opens come from the conductor's config, where it is.""" + model = _bagel() + images = [torch.rand(3, 480, 640)] if "image" in input_modalities else [] + output = model.process_prompt( + "a cat", input_modalities, output_modalities, tensors={"image_inputs": images}, **kwargs, + ) + params = model._request_params(input_modalities, output_modalities, kwargs) + target = output_modalities[0] + request = {"requires_cfg": model._requires_cfg(target_output=target, **params), "target_output": target} + + model._has_cfg_parallel = cfg_parallel + config = model.get_request_resource_configs({"default": SimpleNamespace(full_metadata=SimpleNamespace( + kwargs=request, + ))}, kwargs)["kv"] + + laid_out = set(output.metadata["kv_layout"]["kv"]) + assert config.opened_labels() <= laid_out, "intake would refuse the request for a label its layout leaves out" + + +_BAGEL_RUNS = { + "text chat": (["text"], ["text"], {}, "Hello there"), + "text chat, think": (["text"], ["text"], {"think_mode": True}, "Hello there"), + "understanding an image": (["image", "text"], ["text"], {}, "What is this?"), + "understanding, text first": (["text", "image"], ["text"], {}, "What is this?"), + "text to image": (["text"], ["image"], {}, "A cat on a mat"), + "text to image, think": (["text"], ["image"], {"think_mode": True}, "A cat on a mat"), + "text to image without guidance": (["text"], ["image"], {"cfg_text_scale": 1.0}, "A cat on a mat"), + "an edit": (["image", "text"], ["image"], {}, "Make it night"), + "an edit, text first": (["text", "image"], ["image"], {}, "Make it night"), + "an edit without text": (["image"], ["image"], {}, None), + "an edit without guidance": (["image", "text"], ["image"], {"cfg_text_scale": 1.0, "cfg_img_scale": 1.0}, "Night"), +} + + +def _pointer(tensor: torch.Tensor): + from mstar.graph.base import TensorPointerInfo + + return TensorPointerInfo( + dims=list(tensor.shape), dtype=tensor.dtype, nbytes=0, address=0, stride=[], uuid=id(tensor), + source_session_id="", source_entity="", + ) + + +def _bagel_request(model, prompt, input_modalities, output_modalities, kwargs): + """A request's KV config, laid out and capped, whether it is guided, and its walks in order with the + rows each writes: Bagel's own schedule, the ViT count from `vit_patches`, the VAE count from + `VAEEncoderSubmodule.prepare_inputs`, a decode row per cap, and three denoise steps.""" + from mstar.conductor.request_info import DEFAULT_PARTITION + from mstar.model.bagel.submodules import IMAGE_SENTINELS, VAEEncoderSubmodule, vit_patches, vit_preprocess + + cfg = model.config + images = [torch.rand(3, 480, 640)] if "image" in input_modalities else [] + out = model.process_prompt(prompt, input_modalities, output_modalities, tensors={"image_inputs": images}, **kwargs) + tensors = {"text_inputs": out.new_input_tensors.get("text_inputs", []), + "image_inputs": out.new_input_tensors.get("image_inputs", images)} + signals = {name: [_pointer(t) for t in by_name] for name, by_name in tensors.items()} + fwd = model.get_initial_forward_pass_args(DEFAULT_PARTITION, input_modalities, output_modalities, signals, kwargs) + metadata = fwd.full_metadata + by_uuid = {id(t): t for by_name in tensors.values() for t in by_name} + later = [] + while not model.get_partition_forward_pass_args(DEFAULT_PARTITION, metadata, {}).request_done: + if not metadata.is_prefill: + later.append(metadata.graph_walk) + if metadata.graph_walk in ("image_gen", "image_gen_cfg"): + break + config = model.get_request_resource_configs({DEFAULT_PARTITION: fwd}, kwargs)["kv"] + cap = model.get_loop_caps(max_output_tokens=6)["decode_loop"] + config.apply_conductor_config( + kv_layout={label: layout.row() for label, layout in out.metadata["kv_layout"]["kv"].items()}, + loop_caps={"decode_loop": cap}, + ) + vae = VAEEncoderSubmodule( + SimpleNamespace(posterior_noise=lambda *args, **kw: None), None, None, None, latent_patch_size=1, + latent_channel=16, latent_downsample=cfg.latent_downsample, max_latent_size=cfg.max_latent_size, + ) + preprocess = metadata.kwargs["image_preprocess"] + walks = [] + for walk, info in metadata.kwargs["prefill_schedule"]: + tensor = by_uuid[info.uuid] + if walk == "prefill_text": + walks.append((walk, tensor.numel())) + elif walk == "prefill_vit": + patches = vit_patches( + *tensor.shape[-2:], vit_preprocess(preprocess), cfg.vit_config.patch_size, + cfg.vit_max_num_patch_per_side, + ) + walks.append((walk, patches + IMAGE_SENTINELS)) + else: + info = SimpleNamespace(step_metadata={"image_preprocess": preprocess}, random_seed=0) + encoded = vae.prepare_inputs(walk, info, {"image_inputs": [tensor]}) + walks.append((walk, encoded.kwargs["h"] * encoded.kwargs["w"] + IMAGE_SENTINELS)) + height, width = metadata.kwargs["height"], metadata.kwargs["width"] + latents = (height // cfg.latent_downsample) * (width // cfg.latent_downsample) + for walk in later: + walks += [(walk, 1)] * cap if walk == "decode" else [(walk, latents + IMAGE_SENTINELS)] * 3 + return config, metadata.kwargs["requires_cfg"], walks + + +class _BagelPool: + """A CPU `KVManager` under a `StepRunner`, stepped by Bagel's own `LLMSubmodule.declare_step`.""" + + def __init__(self, configs: dict): + from test_bagel_guidance_sink import _llm + + from mstar.engine.resources import PagedKVConfig, StepRunner + from mstar.engine.resources.kv.manager import KVManager + + self.kv = KVManager( + cfg=PagedKVConfig( + num_layers=1, num_kv_heads=1, head_dim=4, max_seq_len=32768, max_num_pages=4096, page_size=16, + ), + name="kv", joint_comm_group=None, transfer_engine_info=None, device=torch.device("cpu"), + dtype=torch.float32, + ) + self.runner = StepRunner({"kv": self.kv}, node_resources={"LLM": ["kv"]}) + for rid, config in configs.items(): + self.runner.ingest_request(rid, {"kv": config}) + self.llm = _llm() + + def step(self, walk: str, rows: dict, guided: dict) -> None: + """One step of ``walk`` over ``rows`` (request -> rows it writes), each row guided or not.""" + from mstar.engine.resources import StepContext + from mstar.engine.resources.step import SubmoduleStep + from mstar.model.submodule_base import ARNodeInputs + + inputs = [] + for rid, count in rows.items(): + inputs.append(ARNodeInputs(input_seq_len=count, resource_step_info=guided[rid])) + if walk not in {"prefill_text", "decode"}: + positions = torch.zeros(count, dtype=torch.int32) + inputs[-1].custom_pos_ids = dict.fromkeys(("main", "cfg_text", "cfg_img"), positions) + kv_step = self.llm.declare_step(walk, list(rows), inputs, per_request_info={}).steps["kv"] + submodule_step = SubmoduleStep(steps={"kv": kv_step}, segments=list(kv_step.segments)) + submodule_step.set_ctx(StepContext(request_ids=tuple(rows), graph_walk=walk, slot=0, capture=False)) + assert self.runner.admit(submodule_step).outcome.ok + self.runner.plan(submodule_step) + self.runner.commit(submodule_step) + + def held(self, rid) -> dict: + return {label: (stream.peak, stream.footprint) for label, stream in self.kv._streams[rid].items()} + + +@pytest.fixture +def _bagel_checks(monkeypatch): + from mstar.engine.resources.kv import manager as manager_mod + + monkeypatch.setattr(manager_mod, "_DEBUG_ASSERTS", True) + monkeypatch.setattr( + manager_mod, "KVTransferManager", lambda *a, **kw: SimpleNamespace(remove_request=lambda rid: None), + ) + model = _bagel() + model.config.num_timesteps = 3 + return model + + +@pytest.mark.parametrize( + ("input_modalities", "output_modalities", "kwargs", "prompt"), _BAGEL_RUNS.values(), ids=_BAGEL_RUNS, +) +def test_bagel_peaks_at_its_footprint_on_every_label( + input_modalities, output_modalities, kwargs, prompt, _bagel_checks, +): + """Bagel's own schedule, steps and forks, run through a `KVManager`, take each label to its footprint. + ``cfg_text`` forks before a text walk and after an image walk, so what it holds depends on what the + prompt ends on.""" + config, guided, walks = _bagel_request(_bagel_checks, prompt, input_modalities, output_modalities, kwargs) + pool = _BagelPool({1: config}) + + for walk, rows in walks: + pool.step(walk, {1: rows}, {1: guided}) + + held = pool.held(1) + assert all(peak == footprint for peak, footprint in held.values()), f"(peak, footprint) by label: {held}" + + +def test_an_unguided_bagel_request_batched_with_a_guided_one_stays_inside_its_footprint(_bagel_checks, caplog): + """A guided batch runs the guidance labels for every row. The unguided chat runs them against + SINK_PAGE, so it holds only ``main`` and the always-on line stays silent.""" + import logging + + model = _bagel_checks + guided, _, guided_walks = _bagel_request(model, "A cat on a mat", ["text"], ["image"], {"think_mode": True}) + chat, _, chat_walks = _bagel_request(model, "Hello there", ["text"], ["text"], {}) + pool = _BagelPool({1: guided, 2: chat}) + flags = {1: True, 2: False} + + with caplog.at_level(logging.ERROR, logger="mstar.engine.resources.kv.manager"): + while guided_walks or chat_walks: + walk = (guided_walks or chat_walks)[0][0] + rows = {rid: walks.pop(0)[1] for rid, walks in ((1, guided_walks), (2, chat_walks)) + if walks and walks[0][0] == walk} + pool.step(walk, rows, flags) + + assert not caplog.records, f"a page went past a footprint: {caplog.text}" + assert set(pool.kv._streams[2]) == {"main"}, "the unguided chat took pages on a guidance label" + for rid in (1, 2): + held = pool.held(rid) + assert all(peak <= footprint for peak, footprint in held.values()), f"(peak, footprint): {held}" + + +@pytest.mark.parametrize("output_modalities", [["text"], ["image"]], ids=["understanding", "an edit"]) +def test_bagel_lays_out_a_request_without_text(output_modalities): + """The walks prefill the attachments alone, so the layout holds their spans, on every label the request opens.""" + model = _bagel() + images = [torch.rand(3, 480, 640)] + output = model.process_prompt(None, ["image"], output_modalities, tensors={"image_inputs": images}) + + layout = output.metadata["kv_layout"]["kv"] + assert [span.walk for span in layout["main"].spans] == ( + ["prefill_vit"] if output_modalities == ["text"] else ["prefill_vae", "prefill_vit"] + ), "the layout lists a text span no walk writes" + + +# ── Chatterbox ────────────────────────────────────────────────────────── + + +@pytest.mark.parametrize("max_tokens", [1, 64]) +def test_chatterbox_caps_decode_where_its_stop_check_ends_it(max_tokens): + """Chatterbox stops on a count its own postprocess keeps. The count starts at the prefill's token.""" + from test_chatterbox_model import _make_model + + from mstar.model.chatterbox.submodules import SPEECH_TOKENS, T3Submodule + + model = _make_model() + t3 = T3Submodule(None, model.config, builtin_voice=None) + outputs = {SPEECH_TOKENS: [torch.tensor([model.config.t3.stop_speech_token + 1])]} + t3.request_state("r").add_all(generated=0, requires_cfg=False) + t3.postprocess("r", None, outputs) + + def check_stop(request_id, request_info, outputs): + t3.postprocess(request_id, request_info, outputs) + return t3.check_stop(request_id, request_info, outputs) + + cap = model.get_loop_caps(max_new_tokens=max_tokens)["decode_loop"] + + assert cap == _stops_after( + check_stop, "decode_loop", "decode", model.config.t3.max_speech_tokens, outputs, + max_tokens=max_tokens, step_metadata={"max_new_tokens": max_tokens}, + ), "the cap and the stop check end decode on different iterations" + + +@pytest.mark.parametrize("variant", ["chatterbox", "turbo"]) +def test_a_reference_clip_is_laid_out_at_the_tokens_the_tokenizer_makes_of_it(variant): + """Turbo writes the clip's S3 tokens raw, so its layout counts them, and the voice encoder asserts the count.""" + from mstar.model.chatterbox.components.audio_frontend import resample, trim_reference + from mstar.model.chatterbox.components.s3_tokenizer import S3Tokenizer + from mstar.model.chatterbox.config import S3_SR, ChatterboxConfig + from mstar.model.chatterbox.submodules import reference_tokens + + config = ChatterboxConfig.from_variant(variant) + # the count is the encoder's length arithmetic, which the weights don't change + tokenizer = S3Tokenizer(config.s3_tokenizer).eval() + for samples in (14_400, 144_001, 720_000): + wav16 = trim_reference(resample(torch.rand(samples), config.sample_rate, S3_SR), S3_SR, config.enc_cond_seconds) + _, lens = tokenizer([wav16], max_len=config.t3.speech_cond_prompt_len) + + assert int(lens[0]) == reference_tokens(samples, config), f"a {samples}-sample clip is laid out wrong" + + +@pytest.mark.parametrize(("cfg_weight", "clip"), [(0.0, False), (0.5, False), (0.5, True)]) +def test_chatterbox_lays_out_every_label_it_declares(cfg_weight, clip): + from test_chatterbox_model import _make_model, _tone + + model = _make_model() + tensors = {"audio_inputs": [_tone(2.0)]} if clip else None + + output = model.process_prompt("hello", ["text", "audio"] if clip else ["text"], ["audio"], + tensors=tensors, cfg_weight=cfg_weight) + + opened = model.get_request_resource_configs({}, {"cfg_weight": cfg_weight})["t3_kv"].opened_labels() + assert set(output.metadata["kv_layout"]["t3_kv"]) == opened, "the layout doesn't name the labels the request opens" + + +# ── Cosmos3 ───────────────────────────────────────────────────────────── + + +class _Cosmos3Tokenizer: + """The three calls `tokenize_prompt` makes, with one token per character.""" + + eos_token_id = 1 + + def apply_chat_template(self, conversations, **kwargs): + return {"input_ids": [ord(c) for c in conversations[-1]["content"]]} + + def convert_tokens_to_ids(self, token): + return 2 + + +@pytest.mark.parametrize("max_tokens", [1, 64]) +def test_cosmos3_caps_the_reasoners_decode_where_its_stop_check_ends_it(max_tokens): + from mstar.model.cosmos3.cosmos3_model import Cosmos3Model + from mstar.model.cosmos3.submodules import REASONER_DECODE_LOOP, Cosmos3ReasonerSubmodule + + reasoner = SimpleNamespace(eos_token_id=None) + + def check_stop(request_id, request_info, outputs): + return Cosmos3ReasonerSubmodule.check_stop(reasoner, request_id, request_info, outputs) + + cap = Cosmos3Model.get_loop_caps(object.__new__(Cosmos3Model), max_output_tokens=max_tokens)[REASONER_DECODE_LOOP] + + assert cap == _stops_after( + check_stop, REASONER_DECODE_LOOP, "reasoner_decode", 2048, + {"new_token": [torch.tensor([NOT_EOS])]}, max_tokens=max_tokens, + ), "the cap and the stop check end the reasoner's decode on different iterations" + + +_COSMOS3_COMMITS = { + "text to image": ("a fox in the snow", ["image"], {}), + "text to image without guidance": ("a fox in the snow", ["image"], {"guidance_scale": 1.0}), + "text to video": ("a fox in the snow", ["video"], {"num_frames": 33}), + "kv windows": ("a fox in the snow", ["video"], {"window_mode": "kv", "num_frames": 121}), + "kv windows without context": ("a fox", ["video"], {"window_mode": "kv", "num_frames": 121, "context_frames": 0}), +} + + +@pytest.mark.parametrize(("prompt", "outputs", "kwargs"), _COSMOS3_COMMITS.values(), ids=_COSMOS3_COMMITS) +def test_cosmos3_lays_out_what_its_dit_commits(tmp_path, prompt, outputs, kwargs): + """Each label's prompt span is the DiT's text prefix, and its scratch is one denoise step's tokens. + In kv window mode, its later spans and its retention follow the windows the DiT commits.""" + from test_cosmos3_edge_model import _model + + from mstar.model.cosmos3.submodules import COND_LABEL, UNCOND_LABEL, Cosmos3DiTSubmodule + + model = _model(tmp_path, enable_windowed_video=True) + model.tokenizer = _Cosmos3Tokenizer() + dit = Cosmos3DiTSubmodule(None, model.config, scheduler=None) + dit.get_device = lambda: torch.device("cpu") + dit._new_scheduler = lambda *args, **kw: SimpleNamespace(timesteps=torch.arange(4)) + + output = model.process_prompt(prompt, ["text"], outputs, tensors={}, **kwargs) + args = model.get_initial_forward_pass_args("default", ["text"], outputs, {"text_inputs": []}, kwargs) + walk = args.full_metadata.graph_walk + info = CurrentForwardPassInfo(request_id="r", graph_walk=walk, fwd_index=0, random_seed=0, max_tokens=1, + rid_handle="r", step_metadata=args.step_metadata) + dit.prepare_inputs(walk, info, {"text_inputs": output.new_input_tensors.get("text_inputs", [])}) + + state, layout = dit.request_states["r"], output.metadata["kv_layout"]["kv"] + for label, static in ((COND_LABEL, state["cond"]), (UNCOND_LABEL, state["uncond"])): + laid_out = layout[label] + if static is None: + assert laid_out == LabelLayout(), f"{label} is laid out for a branch the DiT doesn't run" + continue + assert (laid_out.spans[0].walk, laid_out.spans[0].length) == (walk, static["und_len"]), label + if state.get("ar_schedule") is None: + step = static["num_vision_tokens"] + static.get("num_action_tokens", 0) + (state.get("num_sound") or 0) + assert laid_out.scratch == step, f"{label}'s scratch isn't one denoise step" + continue + schedule, per_unit = state["ar_schedule"], state["ar_tokens_per_unit"] + commits = [(plan.commit_end - plan.commit_start) * per_unit for plan in schedule.windows()] + assert [span.length for span in laid_out.spans[1:]] == commits, f"{label}'s windows aren't the commits" + assert laid_out.scratch == max(plan.units for plan in schedule.windows()) * per_unit + assert laid_out.retention == dit._window_retention(state).get(label), f"{label}'s retention isn't the DiT's" + + +# ── Fish S2-Pro ───────────────────────────────────────────────────────── + + +@pytest.mark.parametrize("max_tokens", [1, 64]) +def test_fish_caps_decode_where_its_stop_check_ends_it(max_tokens): + """The AR stops on a frame count its own stop check keeps. The count starts at the prefill's frame.""" + from test_fish_s2_pro_model import tiny_ar, tiny_model + + ar = tiny_ar() + outputs = {"new_token": [torch.tensor([4])]} + ar.request_state("r").add("generated_frames", 0) + ar.postprocess("r", None, outputs) + prefill = SimpleNamespace(step_metadata={"ar_max_tokens": max_tokens}, max_tokens=max_tokens, resource_configs={}) + ar.check_stop("r", prefill, outputs) + + cap = tiny_model().get_loop_caps(max_output_tokens=max_tokens)["ar_decode_loop"] + + assert cap == _stops_after( + ar.check_stop, "ar_decode_loop", "ar_decode", 4096, outputs, + max_tokens=max_tokens, step_metadata={"ar_max_tokens": max_tokens}, + ), "the cap and the stop check end decode on different iterations" + + +@pytest.mark.parametrize("clone", [False, True]) +def test_fish_lays_out_the_rows_its_prefill_builds(clone): + from test_fish_s2_pro_model import tiny_ar, tiny_model + + from mstar.model.fish_s2_pro.config import AR_KV, PROMPT_LAYOUT_KEY + + model = tiny_model() + # 21 codec frames + tensors = {"audio_inputs": [torch.zeros(20 * model.config.codec.hop_length + 1)]} if clone else None + + output = model.process_prompt("hi", ["text", "audio"] if clone else ["text"], ["audio"], tensors, ref_text="r") + + layout = output.metadata["kv_layout"] + [span] = layout[AR_KV]["main"].spans + assert span.walk == ("ar_prefill_clone" if clone else "ar_prefill") + prefix_len, frames = output.metadata[PROMPT_LAYOUT_KEY] + # the tiny AR's vocabulary stands in for the prompt's ids + prefill = tiny_ar()._build_prefill( + torch.zeros_like(output.new_input_tensors["prompt_ids"][0]), prefix_len, frames, + torch.zeros(frames, 4, dtype=torch.long) if frames else None, + ) + assert span.length == prefill.shape[0], "the layout and the prefill disagree on its rows" + +# ── Higgs Audio ───────────────────────────────────────────────────────── + + +def _higgs(): + from mstar.model.higgs_audio.config import HiggsAudioModelConfig + from mstar.model.higgs_audio.higgs_audio_model import HiggsAudioModel + + model = object.__new__(HiggsAudioModel) + model.config = HiggsAudioModelConfig() + model.tokenizer = SimpleNamespace(encode=lambda text, add_special_tokens=False: [ord(c) for c in text]) + # 100 mel frames a second, padded to the longest chunk, as the whisper extractor gives + model.feature_extractor = lambda chunks, **kwargs: { + "input_features": torch.zeros(len(chunks), 128, 3000), + "attention_mask": torch.stack([torch.arange(3000) < len(chunk) // 160 for chunk in chunks]), + } + return model + + +@pytest.mark.parametrize("max_tokens", [1, 64]) +def test_higgs_caps_decode_where_its_stop_check_ends_it(max_tokens): + from mstar.model.higgs_audio.submodules import HiggsAudioLLMSubmodule + + model = _higgs() + llm = HiggsAudioLLMSubmodule(None, model.config) + + cap = model.get_loop_caps(max_output_tokens=max_tokens)["decode_loop"] + + assert cap == _stops_after( + llm.check_stop, "decode_loop", "decode", 1024, + {"new_token": [torch.tensor([NOT_EOS])]}, max_tokens=max_tokens, + ), "the cap and the stop check end decode on different iterations" + + +def test_higgs_lays_out_every_label_it_declares(): + model = _higgs() + + output = model.process_prompt(None, ["audio"], ["text"], tensors={"audio_inputs": [torch.zeros(10 * 16000)]}) + + layout = output.metadata["kv_layout"] + assert [span.walk for span in layout["kv_cache"]["main"].spans] == ["prefill_text", "prefill_audio", "prefill_text"] + + +# ── Qwen3.5 ───────────────────────────────────────────────────────────── + + +@pytest.mark.parametrize("max_tokens", [1, 64]) +def test_qwen3_5_caps_decode_where_its_stop_check_ends_it(max_tokens): + from mstar.model.qwen3_5.qwen3_5_model import Qwen3_5DenseModel + from mstar.model.qwen3_5.submodules import LLMSubmodule + + llm = SimpleNamespace(config=SimpleNamespace(stop_token_ids=())) + + def check_stop(request_id, request_info, outputs): + return LLMSubmodule.check_stop(llm, request_id, request_info, outputs) + + cap = Qwen3_5DenseModel.get_loop_caps(object.__new__(Qwen3_5DenseModel), max_output_tokens=max_tokens) + + assert cap["decode_loop"] == _stops_after( + check_stop, "decode_loop", "decode", 2048, + {"new_token": [torch.tensor([NOT_EOS])]}, max_tokens=max_tokens, + ), "the cap and the stop check end decode on different iterations" + + +@pytest.mark.skipif(not os.environ.get("QWEN3_5_CKPT"), reason="set QWEN3_5_CKPT to a Qwen3.5 checkpoint") +@pytest.mark.parametrize("modalities", [["text"], ["image", "text"], ["text", "image", "text"]]) +def test_qwen3_5_lays_out_every_label_it_declares(modalities): + """An image prompt is one prefill_vision walk, so every span of it names that walk.""" + from mstar.model.qwen3_5.qwen3_5_model import Qwen3_5DenseModel + + model = Qwen3_5DenseModel(os.environ["QWEN3_5_CKPT"]) + images = [torch.rand(3, 224, 224)] * modalities.count("image") + + output = model.process_prompt("Describe this.", modalities, ["text"], tensors={"image_inputs": images}) + + layout = output.metadata["kv_layout"] + walks = {span.walk for span in layout["kv_cache"]["main"].spans} + assert walks == {"prefill_vision" if images else "prefill_text"} + + +# ── Qwen3-Omni ────────────────────────────────────────────────────────── + + +def _omni_submodule(): + """What the Thinker's and the Talker's stop checks read off their submodule.""" + return SimpleNamespace(config=SimpleNamespace(im_end_token_id=0, talker=SimpleNamespace(codec_eos_token_id=0))) + + +@pytest.mark.parametrize("max_tokens", [1, 64]) +def test_qwen3_omni_caps_the_thinkers_decode_where_its_stop_check_ends_it(max_tokens): + pytest.importorskip("flashinfer") # imported by Qwen3OmniModel's submodules + from mstar.model.qwen3_omni.config import THINKER_SAMPLER + from mstar.model.qwen3_omni.qwen3_omni_model import Qwen3OmniModel + from mstar.model.qwen3_omni.submodules import ThinkerSubmodule + + cap = Qwen3OmniModel.get_loop_caps(object.__new__(Qwen3OmniModel), max_output_tokens=max_tokens) + + assert cap["thinker_decode_loop"] == _stops_after( + lambda *args: ThinkerSubmodule.check_stop(_omni_submodule(), *args), "thinker_decode_loop", "thinker_decode", + 2048, {"new_token": [torch.tensor([NOT_EOS])]}, max_tokens=max_tokens, + resource_configs={THINKER_SAMPLER: SamplingReqConfig(ignore_eos=True)}, + ), "the cap and the stop check end decode on different iterations" + + +@pytest.mark.parametrize("talker_max_tokens", [0, 1, 64]) +def test_qwen3_omni_caps_the_talkers_decode_where_its_stop_check_ends_it(talker_max_tokens): + pytest.importorskip("flashinfer") + from mstar.model.qwen3_omni.qwen3_omni_model import Qwen3OmniModel + from mstar.model.qwen3_omni.submodules import TalkerSubmodule + + cap = Qwen3OmniModel.get_loop_caps(object.__new__(Qwen3OmniModel), talker_max_output_tokens=talker_max_tokens) + + assert cap["talker_decode_loop"] == _stops_after( + lambda *args: TalkerSubmodule.check_stop(_omni_submodule(), *args), "talker_decode_loop", "talker_decode", + 2048, {"layer0_codes": [torch.tensor([NOT_EOS])]}, max_tokens=2048, + step_metadata={"talker_max_tokens": talker_max_tokens}, + ), "the cap and the stop check end decode on different iterations" + + +@pytest.mark.skipif( + not os.environ.get("MSTAR_QWEN3_OMNI_PATH"), reason="set MSTAR_QWEN3_OMNI_PATH to the checkpoint's directory", +) +@pytest.mark.parametrize("output_modalities", [["text"], ["text", "audio"]]) +def test_qwen3_omni_lays_out_every_label_it_declares(output_modalities): + """With audio out, the Talker writes every Thinker prefill row but the system turn's, then its prefix.""" + pytest.importorskip("flashinfer") + from mstar.model.qwen3_omni.config import TALKER_KV, THINKER_KV + from mstar.model.qwen3_omni.qwen3_omni_model import Qwen3OmniModel + + model = Qwen3OmniModel(os.environ["MSTAR_QWEN3_OMNI_PATH"]) + tensors = {"image_inputs": [torch.rand(3, 224, 224)], "audio_inputs": [torch.zeros(16000)]} + + output = model.process_prompt("Describe both.", ["image", "audio", "text"], output_modalities, tensors=tensors) + + layout = output.metadata["kv_layout"] + thinker, talker = layout[THINKER_KV]["main"].spans, layout[TALKER_KV]["main"].spans + if "audio" not in output_modalities: + assert talker == (), "the Talker is laid out for a request with no audio out" + return + # the system turn runs to where the user's turn starts + system = output.new_input_tensors["text_inputs"][0].tolist().index(model.config.im_start_token_id, 1) + assert [span.length for span in talker] == [ + thinker[0].length - system, *(span.length for span in thinker[1:]), 6, + ], "the Talker's prefill rows aren't the Thinker's past the system turn, then its prefix" + + +# ── Qwen3-TTS ─────────────────────────────────────────────────────────── + + +def _tiny_qwen3_tts_talker(): + from test_qwen3_tts_model import _tiny_model_config + + from mstar.model.qwen3_tts.components.talker import Qwen3TTSCodePredictor, Qwen3TTSTalkerModel + from mstar.model.qwen3_tts.submodules import TalkerSubmodule + + config = _tiny_model_config() + return TalkerSubmodule(Qwen3TTSTalkerModel(config), Qwen3TTSCodePredictor(config), config) + + +@pytest.mark.parametrize("max_tokens", [1, 64]) +def test_qwen3_tts_caps_decode_where_its_stop_check_ends_it(max_tokens): + """The Talker stops on a frame count its own postprocess keeps. The count starts at the prefill's frame.""" + from test_qwen3_tts_model import _make_model + + talker = _tiny_qwen3_tts_talker() + outputs = {"new_token": [torch.tensor(NOT_EOS)]} + talker.request_state("r").add_all(generated_frames=0) + talker.postprocess("r", None, outputs) + + def check_stop(request_id, request_info, outputs): + talker.postprocess(request_id, request_info, outputs) + return talker.check_stop(request_id, request_info, outputs) + + cap = _make_model().get_loop_caps(max_output_tokens=max_tokens)["talker_decode_loop"] + + assert cap == _stops_after( + check_stop, "talker_decode_loop", "talker_decode", 2048, outputs, + max_tokens=max_tokens, step_metadata={"talker_max_tokens": max_tokens}, + ), "the cap and the stop check end decode on different iterations" + + +LONG_TEXT = "one two three four five six seven eight nine ten eleven twelve" +_QWEN3_TTS_SHAPES = { + "custom voice": ("custom_voice", "hello big world", {"language": "English"}), + "custom voice, streamed": ("custom_voice", "hello big world", {"non_streaming_mode": False}), + "voice design": ("voice_design", "hello big world", {"instruct": "a calm low voice"}), + "x-vector clone": ("base", "hello big world", {"x_vector_only_mode": True}), + "in-context clone, text past the clip": ("base", LONG_TEXT, {"ref_text": "the reference says"}), + "in-context clone, text short of the clip": ("base", "hi", {"ref_text": "ok"}), + "in-context clone, not streamed": ("base", "hello big world", {"ref_text": "ok", "non_streaming_mode": True}), +} + + +@pytest.mark.parametrize(("variant", "text", "kwargs"), _QWEN3_TTS_SHAPES.values(), ids=_QWEN3_TTS_SHAPES) +def test_qwen3_tts_lays_out_the_rows_its_prefill_builds(variant, text, kwargs): + from test_qwen3_tts_model import _variant_model + + from mstar.model.qwen3_tts.config import LANGUAGE_ID_KEY, PROMPT_LAYOUT_KEY, SPEAKER_ID_KEY, TALKER_KV + + model = _variant_model(variant) + clone = variant == "base" + # 1 s and a sample: 13 codec frames + tensors = {"audio_inputs": [torch.zeros(24000 + 1)]} if clone else None + + output = model.process_prompt( + text, ["audio", "text"] if clone else ["text"], ["audio"], tensors=tensors, **kwargs, + ) + + layout, metadata = output.metadata["kv_layout"], output.metadata + [span] = layout[TALKER_KV]["main"].spans + assert span.walk == ("talker_prefill_clone" if clone else "talker_prefill") + frames = metadata[PROMPT_LAYOUT_KEY][4] + # the tiny Talker's codec vocabulary stands in for the speaker and language tags + prefill = _tiny_qwen3_tts_talker()._build_prefill( + "r", torch.zeros_like(output.new_input_tensors["text_inputs"][0]), metadata[PROMPT_LAYOUT_KEY], + speaker_id=40 if metadata[SPEAKER_ID_KEY] >= 0 else -1, + language_id=41 if metadata[LANGUAGE_ID_KEY] >= 0 else -1, + speaker_embed=torch.randn(16) if clone else None, + ref_codes=torch.randint(0, 32, (frames, 4)) if frames else None, + ) + assert span.length == prefill.shape[0], "the layout and the prefill disagree on its rows" + + +# ── V-JEPA2-AC ────────────────────────────────────────────────────────── + + +def _vjepa2(predictor_kind: str = "ac", extrinsics: bool = False): + from mstar.model.vjepa2.vjepa2_model import VJepa2Model + + model = VJepa2Model("unused", skip_weight_loading=True, predictor_kind=predictor_kind) + if extrinsics: + model.config.ac_predictor.use_extrinsics = True + return model + + +@pytest.mark.parametrize("horizon", [2, 5, 16, 40]) +def test_vjepa2_caps_the_rollout_where_its_stop_check_ends_it(horizon): + from mstar.model.vjepa2.submodules import VJepa2ACRolloutPredictorSubmodule + + model = _vjepa2() + kwargs = {"rollout_horizon": horizon} + step_metadata = model.get_initial_forward_pass_args("p0", ["video"], ["video"], {}, kwargs).step_metadata + + assert model.get_loop_caps(**kwargs)["rollout_loop"] == _stops_after( + VJepa2ACRolloutPredictorSubmodule(None, model.config).check_stop, "rollout_loop", "prefill_video_rollout", + 64, {}, max_tokens=1, step_metadata=step_metadata, + ), "the cap and the stop check end the rollout on different iterations" + + +@pytest.mark.parametrize("extrinsics", [False, True]) +def test_vjepa2_lays_out_a_frame_per_rollout_iteration(extrinsics): + """Each iteration writes what the rollout step declares for one frame of context. No other walk writes.""" + from mstar.model.vjepa2.config import KV_CACHE + from mstar.model.vjepa2.submodules import VJepa2ACRolloutPredictorSubmodule + + model = _vjepa2(extrinsics=extrinsics) + kwargs = {"actions": [[0.1] * 7] * 5, "states": [[0.2] * 7] * 5} + if extrinsics: + kwargs["extrinsics"] = [[0.3] * 6] * 5 + + rollout = model.process_prompt(None, ["video"], ["video"], rollout_horizon=5, **kwargs).metadata["kv_layout"] + once = model.process_prompt(None, ["video"], ["video"], **kwargs).metadata["kv_layout"] + + step = VJepa2ACRolloutPredictorSubmodule(None, model.config).prepare_inputs( + "prefill_video_rollout", SimpleNamespace(dynamic_loop_iter_counts={"rollout_loop": 0}), { + "encoder_hidden": [torch.zeros(1, model.config.grid_size ** 2, model.config.hidden_size)], + **{name: [torch.tensor(value)] for name, value in kwargs.items()}, + }, + ) + assert rollout[KV_CACHE]["main"].growth == (("rollout_loop", step.input_seq_len),) + assert once[KV_CACHE]["main"] == LabelLayout(), "a walk that never runs the rollout is laid out" + assert _vjepa2("masked").process_prompt(None, ["video"], ["video"]).metadata["kv_layout"] == {} + + +# ── Whisper ───────────────────────────────────────────────────────────── + + +def _whisper(): + from transformers import WhisperFeatureExtractor + + from mstar.model.whisper.config import WhisperModelConfig + from mstar.model.whisper.whisper_model import WhisperModel + + model = object.__new__(WhisperModel) + model.config = WhisperModelConfig(lang_to_id={"<|en|>": 50259}, task_to_id={"transcribe": 50360}) + model.feature_extractor = WhisperFeatureExtractor(feature_size=model.config.num_mel_bins) + return model + + +@pytest.mark.parametrize("max_tokens", [1, 5, 64, 1000]) +def test_whisper_caps_decode_where_its_stop_check_ends_it(max_tokens): + """The cap is the request's own bound, which the learned positions hold to 444.""" + from mstar.model.whisper.config import SAMPLER + from mstar.model.whisper.submodules import WhisperDecoderSubmodule + + model = _whisper() + bound = model.get_max_output_tokens(max_output_tokens=max_tokens) + + assert model.get_loop_caps(max_output_tokens=max_tokens)["decode_loop"] == _stops_after( + WhisperDecoderSubmodule(None, model.config).check_stop, "decode_loop", "decode", 2048, + {"new_token": [torch.tensor([NOT_EOS])]}, max_tokens=bound, + resource_configs={SAMPLER: SamplingReqConfig(ignore_eos=True)}, + ), "the cap and the stop check end decode on different iterations" + + +@pytest.mark.parametrize("seconds", [1.0, 29.0, 45.0]) +def test_whisper_lays_out_the_forced_prompt_and_the_encoders_whole_window(seconds): + """Any clip is padded or cut to the window, so the context is the encoder's rows for any length.""" + from transformers import WhisperConfig + from transformers.models.whisper.modeling_whisper import WhisperEncoder + + from mstar.model.whisper.config import CONTEXT_LABEL, CROSS_KV_CACHE, KV_CACHE + + model = _whisper() + clip = torch.zeros(int(seconds * 16000)) + + output = model.process_prompt(None, ["audio"], ["text"], tensors={"audio_inputs": [clip]}) + + layout = output.metadata["kv_layout"] + assert [span.length for span in layout[KV_CACHE]["main"].spans] == [len(output.new_input_tensors["text_inputs"][0])] + # the encoder's length arithmetic, which neither its width nor its weights change + encoder = WhisperEncoder(WhisperConfig( + d_model=16, encoder_layers=1, encoder_attention_heads=2, encoder_ffn_dim=32, + num_mel_bins=model.config.num_mel_bins, max_source_positions=model.config.max_source_positions, + )).eval() + with torch.no_grad(): + rows = encoder(output.new_input_tensors["audio_features"][0][None]).last_hidden_state.shape[1] + assert [span.length for span in layout[CROSS_KV_CACHE][CONTEXT_LABEL].spans] == [rows] diff --git a/test/modular/test_qwen3_5_video.py b/test/modular/test_qwen3_5_video.py index 6a67b41c3b..3de1ec1e7c 100644 --- a/test/modular/test_qwen3_5_video.py +++ b/test/modular/test_qwen3_5_video.py @@ -383,7 +383,7 @@ def __init__(self, vocab_size: int): def test_vision_walk_matches_get_rope_index(model, processor, rope_model, layout): """The schedule and the LLM's `prefill_vision` row, against transformers: the row's token order and its 3D positions, plus where decode resumes.""" - from mstar.model.qwen3_5.config import ROPE + from mstar.model.qwen3_5.config import KV_CACHE, ROPE from mstar.model.qwen3_5.submodules import LLMSubmodule _, out, ref = run_both(model, processor, layout) @@ -435,6 +435,10 @@ def test_vision_walk_matches_get_rope_index(model, processor, rope_model, layout assert torch.equal(row.custom_pos_ids.long(), positions[:, 0]) # decode resumes where transformers' rope delta puts it assert row.tensor_inputs["mrope_advance"] == int(positions.max()) + 1 + # the layout's spans cover the rows the walk writes and advance the position as far as the walk does + spans = out.metadata["kv_layout"][KV_CACHE]["main"].spans + assert sum(span.length for span in spans) == row.input_seq_len + assert sum(span.advance for span in spans) == row.tensor_inputs["mrope_advance"] def test_mp4_matches_processor(model, processor, tmp_path): diff --git a/test/rust/test_runtime_parity.py b/test/rust/test_runtime_parity.py index e925f05af9..10382d4117 100644 --- a/test/rust/test_runtime_parity.py +++ b/test/rust/test_runtime_parity.py @@ -171,12 +171,12 @@ def _info(uuid): ) -def _admit(rt, rid="r1"): +def _admit(rt, rid="r1", loop_caps=None): return rt.add_request( request_id=rid, partition="default", graph_walk=WALK, partition_worker_graph_ids=[WG_ID], worker_graph_to_workers=ParallelList([WG_ID], [[WORKER]]), - loop_caps={}, + loop_caps=loop_caps or {}, ) @@ -349,6 +349,120 @@ def test_loop_iters_agree(pair): assert got.values == [{"ar_loop": 0}] +def _decode_graph(): + """prefill, then ``ar_loop`` for at most 4 iterations, emitting each token.""" + return Sequential(sections=[ + GraphNode( + name="prefill", input_names={"prompt"}, + outputs=[GraphEdge(name="token", next_node="ar_decode")], + ), + Loop( + name="ar_loop", + section=GraphNode( + name="ar_decode", input_names={"token"}, + outputs=[GraphEdge(name="token", next_node="ar_decode")], + ), + outputs=[GraphEdge(name="token", next_node=EMIT_TO_CLIENT, output_modality="text")], + max_iters=4, + ), + ]) + + +@pytest.fixture(params=["python", "rust"]) +def decode(request): + wg = WorkerGraph( + section=_decode_graph(), graph_walks={WALK}, ranks=[0], + worker_graph_id=WG_ID, + ) + common = dict( + my_worker_id=WORKER, my_worker_graphs=[wg], + all_wg_ids_to_graph_walks={WG_ID: {WALK}}, + all_wg_ids_to_dyn_loops={WG_ID: {"ar_loop"}}, + all_wg_ids_to_nodes={WG_ID: NODES}, + node_to_partition=dict.fromkeys(NODES, "default"), + sharding_config=_sharding(), + ) + if request.param == "python": + book = PythonTensorBookkeeping() + rt = PythonGraphRuntime( + tensor_manager=_StubTensorManager(book), communicator=None, **common, + ) + return rt, book + book = RustTensorBookkeeping() + return rust_runtime.RustGraphRuntime(bookkeeping=book, **common), book + + +def _run_pass(rt, book, rid, uuids, on_decode=None) -> list[str]: + """Feed ``rid`` a prompt and run each ready node until none is ready, calling + ``on_decode`` before each decode iteration completes.""" + def fresh(): + u = next(uuids) + book.put_tensor(u, _info(u)) + book.increment_ref(u, 1) # the safety hold + return u + + rt.ingest_inputs_batch( + _ingest_block([rid], [_spec("prompt", "prefill", uuids=[fresh()])]) + ) + log = [] + for _ in range(40): + ready = [r.node_name for r in rt.get_ready_nodes(set()) if rid in r.rids] + if not ready: + break + node = ready[0] + rt.pop_rids(node, WALK, [rid]) + log.append(node) + if node == "ar_decode" and on_decode is not None: + on_decode() + signals = rt.get_output_signals(node, WALK) + rt.complete_and_route_batch( + RouteInput( + partition="default", graph_walk=WALK, node_name=node, + output_signals=signals, + wg_ids=ParallelList([rid], [WG_ID]), + tensors=[fresh() for _ in signals], + num_tensors=[1] * len(signals), + ), + ) + return log + + +def test_a_capped_loop_ends_at_its_cap_on_every_pass(decode): + """A request's cap lowers only its own copy of the loop, and survives the + reset at the end of a pass. Another request keeps the graph's 4.""" + rt, book = decode + uuids = iter(range(1, 500)) + capped = _admit(rt, "r1", loop_caps={"ar_loop": 2}) + free = _admit(rt, "r2") + + first = _run_pass(rt, book, capped, uuids) + second = _run_pass(rt, book, capped, uuids) + other = _run_pass(rt, book, free, uuids) + + assert first.count("ar_decode") == second.count("ar_decode") == 2 + assert other.count("ar_decode") == 4, "the cap leaked to the next request" + + +def test_speculation_stops_at_the_cap(decode): + """Speculating the iteration past a stop check writes KV the request + then drops, so a capped loop never speculates past its cap.""" + rt, book = decode + rid = _admit(rt, loop_caps={"ar_loop": 2}) + speculated = [] + + def prep(): + out = rt.prep_spec_rids(SpeculationPrepInput( + spec_node_name="ar_decode", curr_node_name="ar_decode", + graph_walk=WALK, rids=[rid], room_for_continuing=None, + streaming_edges=[], streaming_edges_per_rid=[0], + )) + speculated.append(out.ready_rids == [rid]) + + _run_pass(rt, book, rid, iter(range(1, 500)), on_decode=prep) + + assert speculated == [True, False] + + def test_routing_agrees(pair): rt, book, store = pair rid = _admit(rt) From 0a5234b4118265c35adb3e9436183ef54d03b6d9 Mon Sep 17 00:00:00 2001 From: Gaurav Shah Date: Wed, 21 Oct 2026 05:50:58 +0000 Subject: [PATCH 21/26] test: refusals of an undeclared paged label --- test/modular/test_kv_layout_refusals.py | 92 +++++++++++++++++++++++++ 1 file changed, 92 insertions(+) create mode 100644 test/modular/test_kv_layout_refusals.py diff --git a/test/modular/test_kv_layout_refusals.py b/test/modular/test_kv_layout_refusals.py new file mode 100644 index 0000000000..11d0f46cae --- /dev/null +++ b/test/modular/test_kv_layout_refusals.py @@ -0,0 +1,92 @@ +"""A paged label the model leaves undeclared is refused, at load or at intake. + +A request's pages come from its layout, so nothing can size a label that nobody +declared or that a request writes without laying out. So a paged pool with no +labels fails the worker's load, and so does a declaration naming a walk or a +loop the graphs don't run, since a loop takes its cap by name. A request that +opens a label its layout leaves out, or lays out a label the model doesn't +declare, fails at the conductor with a 500, because the model got it wrong, not +the client. +""" + +from __future__ import annotations + +import sys + +sys.path.insert(0, ".") + +import pytest +from test_kv_layout_intake import LAYOUT, _conduct, _Model, _preprocess + +from mstar.engine.resources.kv.config import KVSpec, PagedKVConfig, RingKVConfig, RingKVLayerConfig +from mstar.graph.base import GraphNode, Loop +from mstar.model.base import KVLabel, LabelLayout +from mstar.worker.engine_manager import _refuse_undeclared_pools, _refuse_unknown_walks + +LLM = GraphNode(name="llm", input_names=["text_inputs"], outputs=[]) + + +class _Declaring: + """Declares ``labels`` on ``kv``, over a prefill and a decode that loops on ``decode_loop``.""" + + def __init__(self, labels: dict, decode: Loop | None = None): + self._labels = labels + self._decode = decode or Loop(name="decode_loop", max_iters=8, outputs=[], section=LLM) + + def kv_labels(self): + return {"kv": self._labels} if self._labels else {} + + def get_graph_walk_graphs(self): + return {"prefill": LLM, "decode": self._decode} + + +def _paged(key: str = "kv") -> KVSpec: + return KVSpec(resource_key=key, nodes={"llm"}, config=PagedKVConfig( + num_layers=1, num_kv_heads=1, head_dim=8, max_seq_len=64, max_num_pages=8, page_size=16, + )) + + +def test_a_paged_pool_with_no_labels_is_refused_at_load(): + with pytest.raises(ValueError, match="'kv'") as refusal: + _refuse_undeclared_pools([_paged()], _Declaring({})) + + assert "kv_labels" in str(refusal.value), "the refusal doesn't say what to add" + + +def test_a_ring_pool_needs_no_labels(): + ring = KVSpec(resource_key="ring", nodes={"llm"}, config=RingKVConfig( + num_layers=1, num_kv_heads=1, num_qo_heads=1, head_dim=8, tokens_per_frame=4, + layers=(RingKVLayerConfig(ring_frames=2, ring_buckets=1, pinned_dilation=1),), + )) + + _refuse_undeclared_pools([ring], _Declaring({})) + + +def test_growth_on_an_unnamed_loop_is_refused_at_load(): + unnamed = Loop(max_iters=8, outputs=[], section=LLM) + label = KVLabel(growth=(("decode", "decode_loop"),)) + + with pytest.raises(ValueError, match="'decode_loop', which 'decode' doesn't run") as refusal: + _refuse_unknown_walks(_Declaring({"main": label}, decode=unnamed)) + + assert unnamed.name in str(refusal.value), "the refusal doesn't name the loops the walk runs" + + +@pytest.mark.parametrize(("layout", "message"), [ + (None, "returned no kv_layout for 'kv'/'main', which the request opens"), + ({"kv": {"main": LAYOUT["kv"]["main"]}}, "returned no kv_layout for 'kv'/'uncond', which the request opens"), + ({"kv": {**LAYOUT["kv"], "x": LabelLayout()}}, "returned a kv_layout for 'kv'/'x', which it does not declare"), +], ids=["no layout", "a label left out", "an undeclared label"]) +def test_a_request_laid_out_wrong_fails_with_a_500(tmp_path, monkeypatch, layout, message): + model = _Model(layout=layout) + kwargs = _preprocess(model, {"max_output_tokens": 30}) + + sent = _conduct(model, kwargs, tmp_path, monkeypatch).sent + + [failure] = [msg.body for msg in sent if getattr(msg, "message_type", None) == "request_failed"] + assert failure.status == 500 and message in failure.error_message, ( + f"the request failed with {failure.status}: {failure.error_message}" + ) + assert not [msg for msg in sent if getattr(msg.message_type, "name", "") == "NEW_REQUEST"], ( + "a worker heard of a request the conductor refused" + ) From 69f8872ce1b981001fc378fc246cefcb59b31ac6 Mon Sep 17 00:00:00 2001 From: Gaurav Shah Date: Wed, 21 Oct 2026 18:58:41 +0000 Subject: [PATCH 22/26] docs: paged KV layouts, and prefix keys as a mark on a label --- docs/adding_models.rst | 157 +++++++++++++++++++++++++-------- docs/environment_variables.rst | 6 +- 2 files changed, 124 insertions(+), 39 deletions(-) diff --git a/docs/adding_models.rst b/docs/adding_models.rst index 26670b566d..1ff18df577 100644 --- a/docs/adding_models.rst +++ b/docs/adding_models.rst @@ -287,45 +287,130 @@ 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. +- ``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. It adds the scratch once on top of every span, so it +counts a step that runs before the last spans as if it ran after them. 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 ''/'