Performance 1/7: shared engine, conductor and component changes

.claude/skills/diffusion-scaffold/SKILL.md

@@ -41,7 +41,7 @@ A recipe such as LTX-2.5's two-stage distilled one (half-resolution denoise →
 Every attention in the DiT should go through a ragged resource, so the model gets the engine's kernels and its backend can change in config.
 
 - `attention_segments(bucket_key)` declares every span a request's layers attend over, as `(label, length)`; the default is one `"main"` span of all its tokens. A layer that attends over a sub-span (a refiner over the image tokens alone, a joint model's per-modality self-attention) gets its own label, and `ragged_for(label)` gives that layer its callable.
-- `ragged_for` memoizes one callable per label for the life of a binding. The compiled transformer guards on the identity of the callables it is handed; a fresh one per step recompiles every step.
+- `ragged_for` memoizes one callable per label (the resource's `run` with the label bound) for the life of a binding, and a rebinding drops them so none attends through the old resource.
 - A model with more than one head geometry (LTX-2.5's video 32×128 and audio 32×64 heads) declares one resource per geometry.
 - Cross-attention between spans of the same request (video queries over text keys) is a separate cross-attention resource kind (`RaggedCrossAttentionSpec`, added in #372) whose step names `(q_label, kv_label)` pairs; see "Cross-attention between spans" in `docs/adding_models.rst`. Don't fall back to SDPA for it because the shapes are fixed inside a bucket.
 - `joint_attention` in `attention.py` packs per-row q/k/v into the ragged layout (or runs SDPA when no resource is bound, which is what a parity backend uses).
@@ -50,7 +50,7 @@ Every attention in the DiT should go through a ragged resource, so the model get
 
 - The engine's blanket compile is off for the node (`disable_torch_compile = True`); the model compiles its transformer region with `compile_utils.compile_transformer_forward`, which keeps the exact-op exclusions that preserve eager rounding. Compiling the DiT mostly buys the memory-bound elementwise work (modulation, fp32 norms, RoPE, gating), which on LTX-2.5 was ~45% of an eager step.
 - `capture_buckets` lists `(graph_walk, bucket_key)` pairs to capture; `capture_request_inputs` supplies one placeholder row per bucket. A bucket is captured per batch size in `capture_batch_sizes`, and each one costs startup, so capture the shapes you serve, not every shape you accept. `replay_walks` lets another walk reuse a captured bucket.
-- The captured graph removes the gaps the per-layer attention graph breaks leave in a compiled-but-eager step; measure both before deciding a model doesn't need capture.
+- The captured graph removes the launch gaps of a compiled-but-eager step. Ragged attention breaks that compiled graph per layer unless the resource sets `compile_op=True` (then `run` traces as the `mstar::ragged_attention` custom op), so measure both before deciding a model needs capture.
 
 ## The overshoot veto
 

benchmark/waypoint/benchmark_streaming.py

@@ -558,7 +558,9 @@ def _dit_step_timestamps(log_text: str, request_ids: set[str]) -> list[float]:
             batch = ast.literal_eval(line.split(marker, 1)[1].strip())
         except (SyntaxError, ValueError):
             continue
-        if not isinstance(batch, (list, tuple)) or not any(rid in request_ids for rid in batch):
+        if not isinstance(batch, (list, tuple)) or not any(
+            str(rid).split("#", 1)[0] in request_ids for rid in batch
+        ):
             continue
         match = _LOG_TIMESTAMP_RE.match(line)
         if match is None:

docs/adding_models.rst

@@ -682,17 +682,21 @@ The spec types are:
        ``head_dim`` (all pre-sharding, as in ``KVConfig``), an ``sm_scale`` that
        defaults to ``head_dim ** -0.5``, ``backend`` (``RaggedAttnBackend.AUTO``, the
        default, ``RaggedAttnBackend.FLASHINFER``, or ``RaggedAttnBackend.DENSE``),
-       ``flashinfer_backend``, and the
+       ``flashinfer_backend``, the
        two CUDA-graph ceilings ``max_segments_per_request`` and
-       ``max_tokens_per_request``. ``DENSE`` runs FA3 ``flash_attn_varlen`` (AITER on ROCm)
-       and falls back to FlashInfer if that kernel is missing. By default the heads are split
-       across all TP × SP ranks of the node; ``shard_over_sp=False`` splits them across the
-       TP ranks only, for an attention each sequence-parallel rank runs whole over its own
-       tokens (one over a replicated stream, or a cross-attention from the rank's query
-       shard). ``workspace_mb`` sizes each label's FlashInfer workspace per slot (default
-       ``MSTAR_WORKSPACE_BUFFER_MB``, 512); a node with many labels can declare less, since a
-       ragged prefill only uses it for split-KV partials. See `Cacheless attention (encoder
-       towers)`_.
+       ``max_tokens_per_request``, ``compile_op``, ``shard_over_sp`` and
+       ``workspace_mb``. ``DENSE`` runs FA3
+       ``flash_attn_varlen`` (AITER on ROCm) and falls back to FlashInfer if that kernel is
+       missing. By default the heads are split across all TP × SP ranks of the node;
+       ``shard_over_sp=False`` splits them across the TP ranks only, for an attention
+       each sequence-parallel rank runs whole over its own tokens (one over a
+       replicated stream, or a cross-attention from the rank's query shard).
+       ``workspace_mb`` sizes each label's FlashInfer workspace per slot (default
+       ``MSTAR_WORKSPACE_BUFFER_MB``, 512); a node with many labels can declare less,
+       since a ragged prefill only uses it for split-KV partials. ``compile_op``,
+       ``shard_over_sp`` and ``workspace_mb`` are the model's declaration; the serving
+       config's ``resources:`` overrides take only the backends and the two ceilings.
+       See `Cacheless attention (encoder towers)`_ for ``compile_op``.
    * - ``RaggedCrossAttentionSpec(config=RaggedAttentionConfig(...))``
      - The cross-attention counterpart: one span of each request attends another
        (``RaggedCrossAttentionStep(pairs=...)``). Same config. See `Cross-attention
@@ -939,7 +943,11 @@ encoder node only, and names no cache:
 
 A layer binds the key as it binds any other resource, then attends through the resource's
 ``run(q, k, v, label=None)``. There is no KV write to pair it with, so no
-``AttentionCallable`` is involved.
+``AttentionCallable`` is involved. A compiled caller breaks at ``run``; with
+``RaggedAttentionConfig(compile_op=True)``, ``run`` traces as the ``mstar::ragged_attention``
+custom op instead, so the region keeps one graph across the attention (worth it for a
+compiled, uncaptured region; it costs startup time and capture memory where the region
+is also captured). An eager call goes straight to the kernel.
 
 The step is an ordinary ``AttentionStep``, and its ``segments`` are the entire layout.
 There is no cache to read the layout off, so a label that carries no segment in the
@@ -967,14 +975,15 @@ step names the ``(q_label, kv_label)`` pairs:
 
 Each pair is planned as its own layout, pairing each request's ``q_label`` span with that
 request's ``kv_label`` span; cross-attention is never causal. A layer attends through
-``RaggedAttentionCallable(resource, cross_label(q_label, kv_label))``, passing the key
+``resource.run(q, k, v, label=cross_label(q_label, kv_label))``, passing the key
 span's tokens as ``k`` and ``v``. One resource has one head geometry, so a model whose
 attentions differ in head count or head dim declares one spec per (kind, geometry).
 LTX-2.5 is the reference (``mstar/model/ltx2_5/submodules.py``).
 
-Called with ``return_lse=True``, a ragged or ragged cross-attention callable on the FlashInfer
-backend also returns each query row's log-sum-exp (``[total_tokens, H]``, fp32, natural log)
-from the same plan, which is what merging attentions over disjoint key sets needs. Under sequence parallelism
+Called with ``return_lse=True``, a ragged or ragged cross-attention resource's ``run`` on the
+FlashInfer backend also returns each query row's log-sum-exp (``[total_tokens, H]``, fp32,
+natural log) from the same plan, which is what merging attentions over disjoint key sets
+needs; a compiled caller breaks there even with ``compile_op``. Under sequence parallelism
 LTX-2.5's audio queries attend each rank's share of the video tokens, and
 ``attention_over_sharded_keys`` (``mstar/model/components/distributed/sequence_parallel.py``)
 merges them.
@@ -1173,15 +1182,22 @@ methods are:
 ``check_stop(...) -> set[str] | StopCheck`` (optional)
    Runs off the GPU thread, and may read tensor values. Return the names of the ``Loop``
    objects to stop, for example after seeing the EOS token. This is how a decode loop
-   terminates. A prefill samples the first token before its decode loop exists, so a stop
-   there is returned as ``END_PARTITION`` (from ``mstar.model.submodule_base``) instead:
-   the worker reports the partition done with that pass, and the conductor ends it rather
-   than start decode. Return it only from a row outside any loop; from inside one the
-   request fails. ``TokenStepNode`` does this for its ``PREFILL_WALK``. It suits a partition
-   with nothing left to run once generation stops: a model that runs another walk after its
-   loop (an ``align`` walk, or image generation after thinking) or publishes final-only
-   state at the loop's stop needs its own transition instead. Keep ``check_stop_batched``,
-   if you write one, in agreement.
+   terminates. An AR submodule whose loop ends on a stop token sets
+   ``self.token_stop = TokenStop(...)`` in ``__init__`` instead of writing one; that also
+   gives it ``check_stop_batched``, which checks a whole step off one host copy of its
+   tokens. The token budget is the loop's cap from ``get_loop_caps``, not a stop check.
+   A prefill samples the first token before its decode loop exists, so a stop there is
+   returned as ``END_PARTITION`` (from ``mstar.model.submodule_base``) instead: the worker
+   reports the partition done with that pass, and the conductor ends it rather than start
+   decode. Return it only from a row outside any loop; from inside one the request fails.
+   ``TokenStop`` does this for the walks in its ``prefill_walks``, whose token counts toward
+   the budget, so ``get_loop_caps`` gives the loop ``max(1, max_output_tokens - 1)``. A
+   model that prefills a prompt over several passes sets ``last_prefill_flag`` to the
+   ``step_metadata`` flag of its last pass (Qwen3.5: ``"last_prefill"``), or a stop id
+   sampled on its first pass ends the request. It suits a partition with nothing left to run once generation stops: a model that runs
+   another walk after its loop (an ``align`` walk, or image generation after thinking) or
+   publishes final-only state at the loop's stop needs its own transition instead. Keep
+   ``check_stop_batched``, if you write one, in agreement.
    To also keep some of this step's outputs from being routed, return
    ``StopCheck(stops, withheld=(names...))``, for example a frame streamed alongside the
    EOS token: its consumer then never receives it, instead of finding and dropping it with
@@ -1732,6 +1748,14 @@ Both types share the base ``PiecewiseCudaGraphConfig`` fields:
   when the region was leased, for example when the outer step is ``None``, and not
   whenever ``runner.can_run`` is true: a refused batch can still fit a bucket, and
   replaying it would declare the region's resources twice.
+- ``one_row_per_step`` makes the runner lease, check (``can_run``) and replay the region as
+  one row, whatever the step's batch size, under the step's first request. It is for a
+  packed region that runs every request's segments as one ragged batch, such as a vision
+  tower, and needs ``lease_before_step`` and ``capture_batch_sizes=[1]``. The region's
+  ``declare_step`` then receives ``[first_rid]`` with every request's ``seq_lens``, so the
+  whole step's segments count against that one request's ``max_segments_per_request``,
+  which for this region is a per-step ceiling. Pair it with a ``can_lease`` that refuses a
+  step with more.
 - ``capture_batch_sizes`` lists the batch sizes to record. ``None`` uses the runner
   default.
 - ``compile`` runs ``torch.compile`` on ``capture_fn`` before capture. The default is
@@ -1966,7 +1990,9 @@ default:
    preprocess_torch_threads: 4
 
 The setting applies only to the API server process. The conductor and workers keep
-torch's default.
+torch's default. It does not size video decoding: each torchcodec decode that asks
+``media_limits.video_decode_threads`` (``Model.load_video`` does) runs FFmpeg with up to 8
+threads, no more than the CPUs the process's CPU plan leaves unpinned (``CpuPlan.left``).
 
 A pool that small is slow for a large copy, and a 121-frame 832x480 video is 145 MB. A
 video decoder should therefore emit its clip in the layout the encoder takes, transposed

docs/environment_variables.rst

@@ -544,7 +544,8 @@ Worker scheduling
        last core of the mask. A number outside the mask, or too near its top
        for every worker, stops the server at startup. Threads and pools a
        pinned thread starts keep the process mask (``mstar.utils.threads``),
-       and the API server's intra-op and encode pools are sized from the
+       and the API server's intra-op and encode pools and FFmpeg's per-decode
+       threads (``media_limits.video_decode_threads``) are sized from the
        CPUs the pins leave (``CpuPlan`` in ``mstar.utils.cpus``). Each worker
        and the API server log their CPUs and why. Read by
        ``knobs.step_thread_core``; no effect where ``os.sched_setaffinity``

docs/serving.rst

@@ -305,7 +305,13 @@ A config maps the model's computation-graph nodes to physical GPU ranks. The key
        ``hf://<owner>/<repo>/<path>`` (a file in a model's HF repo); its modality
        comes from the extension, or give ``{path: ..., modality: ...}``.
        ``repeat: N`` sends ``text`` N times over, a long prompt for a path only
-       long prompts take (a long-context prefill). A malformed or failing warmup
+       long prompts take (a long-context prefill). Each entry is sent as many
+       times at once as the most data-parallel replicas any of the model's worker
+       graphs is served by (the replicas the conductor picks from; a node group no
+       walk places adds none), so every replica warms. When some worker graph has
+       fewer replicas than that, such as one DiT instance beside a replicated VAE,
+       the entry first runs once alone, so that graph also warms its one-request
+       shape before the copies reach it together. A malformed or failing warmup
        is logged and skipped.
    * - ``dist_timeout_s``
      - *(optional)* Timeout in seconds for the NCCL world group and its parallel