Agentic multi-turn rollout: kernel, model wiring, engine resume API

[KJLdefeated]

Build the cross-turn KV-reuse path described in docs/AGENTIC_ROLLOUT_SPEC.md end-to-end, so a sequence can pause at turn end, accept an observation, and resume with only the new tokens prefilled (attending over the cached prefix).

Kernel

• New launch_flash_attention_prefill_paged: a tile of new-chunk queries attends over the paged KV cache causally by absolute position. Same online-softmax math as paged decode; per-query (q_seq_idx, q_abs_pos) metadata with q_abs_pos<0 marking padding rows.
• Fix Flash-Decode split-K NaN when num_splits > num_blocks (every short- context decode): empty splits computed expf(-inf - (-inf)) = NaN and poisoned the reduce. Empty splits now write neutral partials and the reduce skips them. Was blocking 100% of generation on the H100 box and is structurally on the path multi-turn resume hits every turn.

Model

• qwen3_prefill detects any seq with num_cached_tokens>0 and routes the whole batch through the paged kernel via a cont_prefill flag; pure- fresh batches keep dense WMMA FA2. New d_q_seq_idx / d_q_abs_pos scratch buffers + host mirrors.

Engine

• SeqStatus::PAUSED, SamplingParams::agentic; Sequence tracks next_turn_start / turn_count / asst_spans for per-turn token slicing and trainer mask building.
• BlockManager::append_blocks_for grows a resumed seq's block table by delta only.
• Scheduler::paused_ deque + resume(); postprocess routes agentic stops to PAUSED (KV + batch_slot retained) instead of FINISHED+delete; schedule_prefill takes a separate branch for resumed seqs that skips alloc_slot() and allocates only delta blocks.
• LLMEngine::rollout_add / rollout_continue / rollout_run_turns / rollout_finish drive the full lifecycle.

Tests

• test_chunked_prefill: weights-free spike + regression guard for the paged kernel (GQA, block boundaries, empty prefix, multi-seq+padding).
• test_continued_prefill_model: 2/3-chunk + 1-token-resume prefill through the full 28-layer model matches a single full prefill; argmax identical, cos ~0.99998.
• test_rollout_engine: end-to-end multi-turn rollout via the engine vs single-shot reference. Turn-1 tokens match generate(P, G1) exactly; all 12 turn-2 greedy tokens match generate(P+t1+O, G2).

Docs

• agentic_rl_plan.md (roadmap) + docs/AGENTIC_ROLLOUT_SPEC.md (live spec with progress table; calls out the split-K decode bug and why it was load-bearing for the agentic path).

[Copilot]

Pull request overview

Implements an end-to-end “agentic” multi-turn rollout flow with cross-turn KV reuse, enabling sequences to pause at turn boundaries, accept observation tokens, and resume generation using a continued-prefill path that attends over the cached prefix.

Changes:

• Adds a paged, cache-reading causal prefill attention kernel (launch_flash_attention_prefill_paged) and wires it into Qwen3 prefill when any sequence resumes with num_cached_tokens > 0.
• Introduces engine/scheduler support for pausing/resuming sequences (SeqStatus::PAUSED, SamplingParams::agentic, rollout_* API, resumed scheduling + delta block growth).
• Adds kernel/model/engine regression tests plus supporting build plumbing and documentation/spec updates, and fixes split-K decode NaNs for empty splits.

Reviewed changes

Copilot reviewed 16 out of 16 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
src/kernels/attention.cu	Adds paged continued-prefill kernel; fixes split-K empty-split NaN behavior; adds launcher for cache-reading prefill.
include/kernels/attention.cuh	Declares the new cache-reading continued-prefill launcher API.
src/model/qwen3.cu	Routes prefill through paged continued-prefill path when resuming; adds/query metadata buffers; integrates new kernel at layer forward.
include/model/qwen3.h	Adds continued-prefill metadata buffers + cont_prefill flag to the model struct.
include/model/sampling_parmas.h	Adds SamplingParams::agentic, SeqStatus::PAUSED, and per-turn bookkeeping on Sequence.
include/engine/scheduler.h	Adds paused queue + resume/take_paused plumbing; continued-prefill scheduling branch for resumed sequences.
include/engine/block_manager.h	Adds delta-only block growth helpers for resumed sequences (blocks_needed_for, append_blocks_for).
include/engine/llm_engine.h	Adds rollout_add/continue/run_turns/finish API driving multi-turn lifecycle.
tests/kernels/test_chunked_prefill.cu	Weights-free kernel-level regression validating continued-prefill behavior vs golden reference (incl. multi-seq + padding).
tests/models/test_continued_prefill_model.cu	Model-level validation that chunked/continued prefill matches full prefill logits (argmax/cosine).
tests/models/test_rollout_engine.cu	End-to-end engine test for multi-turn rollout equivalence to single-shot reference generation.
tests/models/test_llmengine.cu	Tweaks benchmark parameters/comments related to realistic bench data dependency.
Makefile	Registers new kernel/model tests in the make-based test targets.
CMakeLists.txt	Registers new kernel/model tests in the CMake build.
docs/AGENTIC_ROLLOUT_SPEC.md	Adds/updates implementation spec and progress tracking for the agentic rollout path and prerequisites.
agentic_rl_plan.md	Adds roadmap document describing the agentic RL pivot plan and sequencing.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.