[0z5a]

Summary

A fixed-depth decode row now runs all its remaining loops in one RECURRENT submission (one graph replay per loop, no host round trip between loops) when every row of the batch has a fixed depth on the LAST_EXITED layout. Rows with an adaptive exit, trace exits or the SHARED layout keep per-loop boundaries. Fused submissions use the real row count under padded static buffers (the first loop's device KV view and the later loops' host views must agree; this caused a rotary row mismatch in the static+padded async GPU tests before the fix).

A fused step is L loops long, and the refill policy admitted max_prefill_batches_before_decode prompt batches per decode step, so prompt work got 1/L of its share and TTFT grew under load. After a decode submission that fuses L loops the policy now allows L times the configured prompt batches (SchedulerConfig.prefill_share_per_loop, default on).

Each row's arithmetic is unchanged. Batches can differ from unfused execution (rows no longer stagger across loop phases), so bf16 GEMMs can pick other algorithms at high load and a few greedy sequences change there.

Measurements

RTX 5090, Ouro-1.4B, 512-token prompts, 128 output tokens, async + graphs, one run per level:

• Submissions per output token: 6.0 → 3.0 at batch 1. Throughput 0.99× at batch 1, 1.09× at 32 and 128 requests. Batch 1 is bound by kernel count, not by submissions (see the exact-kernel PR of this series).
• Token identity: 160/169 requests; all requests at C ≤ 8 identical.
• Huginn-0125 (32 loops): 33.5 → 3.0 submissions per token, 1.10–1.12×.
• GSM8K 87-question subset (serial): 87/87 token sequences and exit depths identical; real-text prompts 41/41.

GH200, Ouro-1.4B, prefill share on vs off:

	share off	share on
burst, 32 requests (2 kept runs): tok/s	877	1001
TTFT p50 / p90 (s)	1.11 / 2.13	0.82 / 1.63
TPOT p50 (ms)	25.9	22.8
64 agent sessions, 6 turns (two pairs): TTFT p50 (s)	10.26, 9.69	2.64, 2.47
TPOT p50 (s)	0.50, 0.46	1.19, 1.18
tok/s	62.8, 67.4	60.2, 56.7

For bursts the share helps on every metric. For sessions it trades time per token for time to first token; prefill_share_per_loop=False restores the fused-only behaviour. Without fusion, main had TTFT p50 1.41 s and TPOT p50 1.30 s at 64 sessions (separate run).

Platforms. GH200: torch 2.14.0+cu130, Triton 3.8.0, NIXL 1.4.1. RTX 5090: torch 2.13.0+cu130, Triton 3.7.1. Models: Ouro-1.4B (24 layers, 4 loops, fixed depth) and Huginn-0125 (2 prelude + 4 recurrent + 2 coda layers, 32 loops). Burst workloads issue C requests at once with random-token prompts; session workloads replay multi-turn agent traces with the prefix cache on.

Tests

• CPU suite: 831 passed; ruff passes. New tests: fused-loop scheduling and identity against unfused execution, the per-loop prefill share (a late prompt gets its first token at step 12 instead of 24 with identical tokens), the async overlap probe, graph-staging contract for fused later loops.
• GPU tests on this branch (GH200, head 9d277b7, flash-attn-3 3.0.1 and flash-attn-4 4.0.0b30 installed, pytest --run-gpu tests/test_engine.py tests/test_async_pipeline.py tests/test_input_preparation.py tests/test_cuda_graph.py): 137 passed. Without FlashAttention installed, the FlashAttention tests failed at import on this branch and on main (66920da) alike.

Notes

The GPU numbers come from a development branch that contains all changes of this series; most cells are single runs, and on the shared GH200 host two runs of main at 32 requests differed by up to 16%, so smaller differences are ties. This branch alone was checked with the CPU suite, ruff and the GPU tests above.

Series. Draft PRs from one study of serving looped models. They are independent except #152, which is stacked on #147; several touch the same files, so later ones may need a rebase after earlier ones merge.

• [Kernels] Bitwise-exact fused RMSNorm, RoPE and SiLU*up for CUDA graph decode #146 [Kernels] Bitwise-exact fused RMSNorm, RoPE and SiLU·up for graph decode: bitwise, 1.14–1.44× on Ouro
• [KV] Pack recurrent depth planes for partially recurrent models #147 [KV] Pack recurrent depth planes: bitwise, 1.88× Huginn KV capacity
• [Scheduler] Fuse fixed-depth decode loops into one submission, keeping the per-loop prefill share #148 [Scheduler] Fused decode loops with the per-loop prefill share: per-row arithmetic unchanged ← this PR
• [Engine] Packed prefill resolves automatically, for Huginn too #149 [Engine] Automatic packed prefill, Huginn included: accuracy-gated, 2.0–2.4× on 2048-token prompts
• [PD] Coalesced transfers, UCX error-handling mode, larger default budget #150 [PD] Huginn coordinator config, coalesced transfers, UCX error handling, 2 GiB default budget
• [Perf] Host-path costs: prefix-key caching, O(1) free-block count, per-run KV bookkeeping, one argmax per batch #151 [Perf] Host-path costs in admission, KV bookkeeping and sampling
• [Huginn] Opt-in prefix caching with position-keyed recurrent-state noise (stacked on #147) #152 [Huginn] Opt-in prefix caching with position-keyed noise (stacked on [KV] Pack recurrent depth planes for partially recurrent models #147): 1.91× session throughput

[LyxWxj]

Two additional points:

1. Lines 802–807 in scheduler.py fix a semantic issue when wavefront_chunked_prefill is enabled. If a batch contains only RECURRENT requests, _planned_loop_count() could previously return 0, after which ExecutionBatch.__post_init__() implicitly changed it to 1. This usually works, but it makes the meaning of _planned_loop_count()`` ambiguous. Calling _fused_decode_loops()` directly makes the scheduler’s output explicit: loop_count represents the actual number of recurrent loops planned for the batch.
2. It may also be worth folding:

if RequestPhase.RECURRENT in phases:
    return self._fused_decode_loops(scheduled_requests)

into the mixed-mode branch:

if self.config.mode == "mixed" and (
    RequestPhase.RECURRENT in phases or self.has_ready(RequestPhase.RECURRENT)
):
    return 1

This would allow mixed batches containing both RECURRENT and PREFILLING_RECURRENT requests to execute multiple consecutive recurrent loops. For example, if both the decode rows and the prefill rows can safely execute at least two loops, the batch could use loop_count=2 instead of executing one loop and rebuilding the batch immediately afterward.

The existing safety conditions should still apply: fusion should only be used when decode exit depths are known, the KV layout supports fused execution, and every prefill row covers the same number of contiguous depth frontiers. Adaptive exits, trace-dependent decisions, the SHARED layout, or any prefill row limited to one loop should continue to force loop_count=1.