[TRTLLM-16460][feat] agent-flow: trtllm_modeling_v2 workflow for gated ModelingV2 onboarding and perf tuning

[Wanli-Jiang]

@coderabbitai summary

Description

What trtllm_modeling_v2 is

trtllm_modeling_v2 (previously staircase-onboard; renamed following review) is an agent-flow workflow with two stages:

• onboard takes one Hugging Face checkpoint to a gated ModelingV2 target: a flat, self-contained forward at tensorrt_llm/_torch/_experimental/modeling_v2/models/<family>/<checkpoint>__<arch>__<parallel>/, built from catalog entries and the target's own ops/. The target is trusted because it passes accuracy gates, not because it shares abstractions with other models.
• tune makes a delivered target fast without making it less correct: a config ladder, then a campaign of agent-written candidates, each measured end to end against a same-host control and kept only if it beats the measured noise floor and passes the accuracy gate.

The modeling_v2 catalog, targets and router are already in tree (#19977); this PR is the agent half. An onboard run has five phases:

1. Preflight. Survey the checkpoint and the launch environment.
2. Plan. A drafter proposes a staircase of units (core first, then one feature per unit) and a cross-model reviewer reviews it.
3. Assemble. For each unit, an assembler writes the target, and kernel and modeling smiths fill catalog gaps.
4. Gate. Boot, accuracy, acceptance and witness gates run at every parallel member.
5. Deliver. An auditor runs a bounded remediation cycle, then delivery checks run.

Everything is driven by one task.yaml; perf.enabled: true chains tune after delivery, and trtllm-modeling-v2 tune runs it standalone on a delivered target.

Why: onboard

The workflow has onboarded two checkpoints end to end. A MiMo-V2.6 run produced findings F001–F011 and a retrospective; a DeepSeek-V4.1-Flash-NVFP4 run (8 units, about 51 h) produced F001–F051 and a reflection.

In the second run, 9 defects escaped every gate and were found later by hand:

• Silent accuracy shift. A position defect shifted accuracy on every member. The witnesses compared the engine with itself, and the floor was anchor − 5, about 7 standard errors (F050).
• Crashes outside the tested shapes. An import-time capacity constant, an int32 overflow, and a recompute past 16K escaped because the tuple the gates ran was treated as the certified set (F045, F047).
• Untested deployment path. HTTP 400 on every image request through trtllm-serve (F040).
• No performance gate. A +19% prefill regression went through (F048).
• Unrepresentative gate environment. One defect was hidden by an autotuner cache the gates never had; one OOM/hang only appeared right after another engine exited.
• Noisy acceptance gate. 4 prompts reported findings that a 32-prompt paired test showed were noise (F051).

The run also cost 6 session deaths, lock starvation on the shared GPU slots, and about 30 manual claim/dispatch/unlock interventions (F010–F012, F025, F031–F032, F044, F046). Records contradicted each other (F033), and commits carried sandbox placeholders and DCO identity mismatches (F002, F022, F043).

This PR ships the workflow together with the fixes for those findings, condensed into LESSONS.md: each lesson with its symptom, its cause, the in-tree mechanism that now prevents it, and any gap still open.

Why: tune, and what it did on one model

A correct target is not a fast one: the onboarded Qwen3.8-Flash-Next-NVFP4 target (1× GB300, ISL 8192 / OSL 1024, CUDA graphs on, MTP off) started at 0.83× the in-tree model at concurrency 1 (both on their default kernels). The tune stage then ran on it for about 28 h: 30 candidates, 19 kept, 2 delivered as kernel entries. Every keep passed the target's test suite, the modeling_v2 contract tests and gsm8k (anchor 96.36; kept scores 95.91–96.66).

A concurrency sweep afterwards, all three variants on each of 3 hosts (output tok/s/GPU, max_batch_size 128):

concurrency	in-tree model (modeling_v2: off, default kernels)	target before tune	target after tune	after / in-tree
1	119.3	99.2	272.0	2.28×
8	712.9	546.5	1,206.2	1.69×
32	1,628.5	1,345.2	2,589.4	1.59×
128	2,675.3	2,243.6	4,318.4	1.61×

The in-tree model in this sweep runs the engine's default kernel switches: moe_config.backend: AUTO, which picks the CUTLASS MoE for this checkpoint, and TRTLLM_LOW_M_GEMM_BACKEND off. With moe_config.backend: TRTLLM and TRTLLM_LOW_M_GEMM_BACKEND=auto, the kernel families the tuned target also uses, the in-tree model reaches 245.6 tok/s/user at concurrency 1 and 3,342 tok/s/GPU at 128 (gsm8k 95.94). Measured on 3 hosts with the tuned target on each (the last row on 2 of them); the ratios are same-host:

in-tree model	concurrency 1, tok/s/user	concurrency 128, tok/s/GPU	after / in-tree
default kernels	119.6	2,672	2.28× / 1.60×
moe_config.backend: TRTLLM and TRTLLM_LOW_M_GEMM_BACKEND=auto	245.6	3,342	1.11× / 1.28×

So most of the gap to the default in-tree model is its kernel selection; what the in-tree model cannot switch to (the target's fusions and kernels below) is worth 1.11× at concurrency 1 and 1.28× at 128. The tune stage does not yet search the reference model's own switches; that is a follow-up below.

Most of the gain is fusion (hyper-connection sites, MoE routing/finalize into neighbouring kernels, PLE), kernel sourcing (FlashInfer GDN and low-M GEMMs, a row-tiled sparse-attention prefill kernel) and one config-ladder rung. The tuned target code is not part of this PR; it lives on its own branch.

What changed (by area)

#	Area	Main points
1	Workflow core	The orchestrator, its six roles plus runner, prompts, the materialized skills, and the gates and witnesses. It also adds the agent-flow runtime package: turn executors (in-process, file dispatch) and command executors (local, slurm, spool, delegated).
2	Envelope and certification	serving_envelope and feature_matrix in task.yaml; the certified set is derived from them, not from the gate tuples. A capacity-constant audit (F045); integer-width probes in kernel-verify and a smith brief per call site (F042, F047); a typed certification_gap decision fanned out to verify smiths (F019, F020, F026, F028); early delivery checks as warnings (F041); a preflight refusal of any speculation algorithm the launch path cannot run.
3	Upstream adaptation, no model specifics	Per-role backend and model routing through an agents: block (#19522), with Codex roles fixed in-process. Prompt snapshots (#19715), the spill check (#19766), duplicate request ids in file dispatch fixed. Model specifics live in examples/, which a run never materializes.
4	Commit hygiene, pre-merge, embargo	commit.sh commits and signs off as the repo user and refuses symlinks, placeholders and out-of-policy paths (F002, F022, F043). trtllm-modeling-v2 premerge lint|ci|copyright|hygiene|split|resign|pr-open|pack; source_embargo enforced with PreToolUse/PostToolUse hooks and an attempt log; opt-in --overlap-units (F027); a delivery tail that audits the run's commits.
5	Ledger and reports	An append-only evidence ledger; a verdict goes stale when a later commit touches the paths it covered, and stale verdicts are re-run. The final regression is bound to the delivered tree. Progress and report tables come from the ledger; decision files are never overwritten (F033); journal ordinals are atomic (F013, F014).
6	Gates	Accuracy floors of anchor − k·SE and a cross-member finding (F050); a per-prompt paired acceptance gate on natural text (F051, F040, F029, F024); a scheduler-equivalence witness; boot workloads matched on generated text, with a real image for multimodal (F034, F037, F049); a launch-path accuracy smoke through trtllm-serve and the disagg proxy; witnesses at every member; trtllm-modeling-v2 gate <name>.
7	Execution substrate	A spool scheduler (priority queue, device slots, wall and memory limits, MPI environment scrubbed) replaces the flock protocol (F046, F011, F012, F031, F032); --tree jobs with mixed-checkout results refused (F044); atomic native rebuilds (F025); sessions declared in task.yaml with planned hand-over; deadlines from job start plus a watchdog; multi-node mn_* launchers; the external-transport driver trtllm-modeling-v2 driver (F010, F017, F030).
8	modeling_v2 as merged (#19977)	Catalog entries are OpWrapper classes called through __call__ with every argument at the call site; CELLS and ARCHS are the executable contract, read by the survey and checked by delivery (shared catalog and a target's own ops/). The catalog is its entry files (no index). Delivery checks the core's _select_target over plain PrefillTarget / DecodeTarget classes, a W-row weight table with MODEL_WEIGHTS, and routing that returns the target class. Target selection is the LLM-API argument modeling_v2; gates pass it through the harness and set TRTLLM_MODELING_V2_VALIDATE=1 for the step-contract check.
9	FlashInfer as a sourcing tier	Tier 3 of the sourcing ladder is a FlashInfer callable called directly from an entry's __call__. The survey reads the installed package's public API off disk (no import, no GPU) and records which callables the engine already references. Against flashinfer-python 0.6.18: 558 public callables, 510 referenced nowhere in the engine.
10	tune stage	tune/: ground truth (baseline, in-tree reference, identity repeats as the noise floor), the config ladder, and the candidate campaign (one hypothesis = one worktree = one tuner turn; same-host controls; stacked evaluation of compatible winners and of near misses; speculative rebases; the keep gate = target suite + contract tests + gsm8k). bench.py (LLM-API yaml configs, --nsys session-mode capture), nsys_summary.py (per step class, CUDA-graph replays placed at their launch), microbench.py, a FlashInfer/kernel-ledger roadmap requirement, and the perf tooling CLI.
11	Docs	README (including the tune stage), TROUBLESHOOTING, BUILD_AND_TEST and LESSONS.md; a short "environment failures that read as model failures" table in AGENTS.md.

Outside agent-flow/agent_flow/workflows/trtllm_modeling_v2/ and its tests, the PR touches only:

• agent-flow/agent_flow/runtime/ (new package; the old runtime.py becomes runtime/portal.py) and agent-flow/tests/runtime/;
• agent-flow/pyproject.toml: the trtllm-modeling-v2 entry point, package-data, and an E501 per-file ignore for prompt tables;
• agent-flow/README.md: one sentence;
• AGENTS.md: one table.

There are no changes under tensorrt_llm/, cpp/ or tests/, and no new dependencies.

Naming

Following review, the workflow is trtllm_modeling_v2 (package), trtllm-modeling-v2 (CLI), TrtllmModelingV2Workflow, and skills trtllm-modeling-v2-* (the repo's trtllm- skill prefix). The workflow's own environment variables use MV2FLOW_*, so they never share a namespace with the engine's TRTLLM_MODELING_V2* switches (and stay clear of MVAPICH2's MV2_*). In the code, "staircase" now appears only in two quoted historical error messages; in prose it is just the shape of a plan (core first, then one feature per unit).

How to run

pip install -e agent-flow                       # provides the trtllm-modeling-v2 CLI
cp agent-flow/agent_flow/workflows/trtllm_modeling_v2/task.example.yaml task.yaml   # then edit it

# 1. Cheapest check first: preflight, envelope derivation and environment probe; no agent turn
trtllm-modeling-v2 --task task.yaml --workspace ws/ --dry-run
trtllm-modeling-v2 --task task.yaml --workspace ws/      # the run (resumable: re-run the same command)

# 2. Performance: chained (perf.enabled: true in task.yaml), or on a delivered target
trtllm-modeling-v2 tune --task task.yaml --workspace ws/
trtllm-modeling-v2 perf ladder run --spec tune.yaml      # tooling: the ladder alone
trtllm-modeling-v2 perf table --baseline A A B C         # a curve table of label dirs

# 3. Where the orchestrator cannot authenticate an agent (sandbox, rotating tokens):
trtllm-modeling-v2 --task task.yaml --workspace ws/ --agent-transport external
trtllm-modeling-v2 driver --workspace ws/ wait        # in the driver session: NEW <id> | ORCHESTRATOR_EXITED | ...
trtllm-modeling-v2 driver --workspace ws/ prompt <id> --claim

# 4. One gate or witness on its own, against a tree
trtllm-modeling-v2 gate <name> --task task.yaml --unit <unit> --member <member> --tree <checkout>

# 5. Before a delivered branch becomes a PR (nothing in the pack pushes)
trtllm-modeling-v2 premerge pack --base origin/main   # lint x2, DCO hygiene, copyright, CI lines as one report

Test Coverage

All tests are CPU-only and need neither GPUs nor model weights. Pass PYTHONPATH=agent-flow when an editable install points at another checkout.

• pytest agent-flow/tests/workflows/trtllm_modeling_v2: 1870 passed, 2 skipped, including the tune suites (test_tune_stage, test_tune_ladder, test_tune_bench, test_tune_nsys_summary, test_tune_worktree_parallel, ...), the gate, ledger and substrate suites (with the substrate's fault injection), test_delivery against the in-tree gpt_oss and DeepSeek-V3 targets, and test_prompts, which refuses the pre-merge modeling_v2 names in any role prompt.
• pytest agent-flow/tests (all agent-flow workflows, guarding the shared runtime package): 3524 passed, 5 skipped.
• Lint: the repo's .pre-commit-config.yaml over every file this PR changes (origin/main..HEAD): all hooks pass.

GPU validation. The tune stage ran end to end on Qwen3.8-Flash-Next-NVFP4 (above) and was started on gpt-oss-120b (one keep, gsm8k pass, before it was stopped). The onboard gates were exercised end to end by the two onboarding runs that produced the findings; the historical-defect replay and a pilot onboarding run on this branch are follow-ups.

Known limitations and follow-ups

• Historical-defect replay and pilot run of the onboard stage on this branch (targets at least 7 of the 9 escaped defects; pilot compared with the 51 h DeepSeek-V4.1 run).
• Performance is measured, not gated. The per-commit prefill/decode regression gate (F048) is a follow-up; tune measures but never blocks delivery.
• The reference model's switches are not tuned. The config ladder skips module knobs for a target, which pins its kernels, but it also never tries them, or engine environment switches such as TRTLLM_LOW_M_GEMM_BACKEND, for the in-tree reference. A reference-side ladder, gated on accuracy, is a follow-up.
• tune substrate. The spool runs pending jobs in name order, not submission order; long campaigns depend on 4 h interactive sessions being re-submitted; a stopped run has no tune --stop that cancels its jobs and writes the partial summary.
• No concurrent witness leg yet (batch 1 only); envelope probes use the aggregated quickstart boot.
• Embargo enforcement applies only to in-process roles whose backend runs SDK hooks; Codex roles and external-transport turns get it as a prompt instruction.
• Unverifiable backend on external transport (F008).
• Opt-in unit overlap (--overlap-units, off by default).
• No DSpark branch in main's quickstart. A dspark spec stops at preflight with an explicit message.

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

[litaotju]

Thanks for putting all these together.

1. Lets try to rename staircase to "trtllm_modeling_v2"?

2. One open question. Shall each feature have its own skill or it should be part of the main skill with reference doc on how each feature shall be implemented?

[Wanli-Jiang]

updated codes and comments in #20145 (comment)

[WeiHaocheng]

I am still understanding and reviewing the code in the workflow section.

Regarding the "skills" section, I do not recommend incorporating these skills into the repo. The previous experience of the modeling agent tells us that, first, detailed skills may limit the model's performance. Second, the large amount of information carried in the skills is difficult to maintain, and errors and outdated information will become difficult to manage. We can first place them in an internal small repo. If there are clear AB tests proving that this is effective, then we can incorporate them into the repo.

[Wanli-Jiang]

Thanks both for the reviews. I've force-pushed an updated branch: it is rebased onto main now that #19977 has merged (so only this PR's own commits remain), aligned with modeling_v2 as merged, and it adds the tune stage. The description is updated to match.

@litaotju

1. Rename: done. The workflow is now trtllm_modeling_v2: package agent_flow.workflows.trtllm_modeling_v2, CLI trtllm-modeling-v2, TrtllmModelingV2Workflow, and skills trtllm-modeling-v2-* (following the repo's trtllm- skill prefix). The workflow's own env vars became MV2FLOW_* rather than TRTLLM_MODELING_V2_*, so they can't be mistaken for the engine's TRTLLM_MODELING_V2 / _VALIDATE / _VALUE_DUMP switches. In the code, "staircase" now appears only in two quoted historical error messages.
2. One skill per feature, or one main skill with per-feature references? Today it is one core skill plus one skill per feature. The reason: a run assembles one unit per feature, and each assembler turn is told to load the core skill plus only its own unit's feature skill (FEATURE_SKILLS in task_schema.py), so a turn never carries the other features' text. A main skill whose references are named per feature would keep that property as long as the turn prompt names the reference to read, so I'm happy to restructure it that way if you prefer it; it would also shrink the surface @WeiHaocheng is concerned about (below).

@WeiHaocheng

Both concerns are fair, and this week gave a concrete example of the second: aligning with modeling_v2 as merged meant rewriting 32 skill documents. Some context on how these skills differ from .claude/skills, then a proposal.

• They are inputs of this workflow rather than general-purpose skills: they are versioned with the prompts and gates that reference them, packaged with the workflow, and copied into each run's workspace, so a run never depends on whatever a checkout's .claude/skills happens to contain.
• Part of their correctness is tested: test_skills checks that every repo path a skill names exists, and test_prompts now refuses the pre-merge modeling_v2 names. That catches path-level drift, not semantic drift, which is exactly the staleness you describe.

Proposal, and I'd like your view before I change the code:

• (a) Slim in tree. Fold the per-feature skills into the core skill's references (answering @litaotju's question 2 at the same time) and cut the content to the rules a gate or delivery check can't enforce on its own. The volume drops and what remains stays versioned with the workflow.
• (b) Out of tree. Add a skills_dir option to task.yaml, ship only the minimal set the prompts require, and keep the full skills in an internal repo until an A/B (the same onboarding with and without them) shows they help.

I lean towards (a) for this PR plus the A/B as a follow-up, but would go with (b) if you'd rather keep the detailed text out of the repo until there is data. Which do you prefer?