Run a DFlash2 block drafter in the batch engine

[Tianlei Wu (tianleiwu)]

Description

Adds the DFlash2 block-drafter runtime to the paged batch engine on top of #2495. The target decoder exposes packed auxiliary hidden states, and the drafter ingests committed rows into its own bounded sliding-window KV cache before proposing a coherent token path through the exported candidate lattice.

This replaces #2460 while excluding builder changes already covered by #2487/#2489, diagnostics churn, unsafe mixed-checkpoint experiments, and ONNX Runtime dependency changes.

Summary of Changes

Packed target output

• Binds an optional packed aux_hidden_states decoder output sized to the active token rows.
• Forwards the output through hybrid decoder IO and scheduled requests without capturing dense recurrent checkpoints.
• Rejects CUDA graph capture for this dynamically sized intermediate output.

DFlash2 drafter

• Parses runtime-only model.dflash2 session, geometry, tensor-name, and sliding-window configuration.
• Creates the DFlash2 session before sizing the target cache so its fixed batch-scaled footprint is included in available-memory accounting.
• Maintains a bounded per-request paged KV ring and releases its blocks with request teardown.
• Packs committed context and query-block rows into one drafter execution and walks the candidate/edge-score lattice greedily.

Engine integration

• Drops hidden-state rows for rejected target drafts before feeding the drafter.
• Proposes only for eligible greedy requests within per-request and maximum-length draft limits.
• Captures feed metadata before request commit clears accepted-prefix counters, then publishes drafts after the target/MTP transaction commits.
• Records DFlash2 forward passes in the shared speculative telemetry.

Testing

• build/Linux/Debug/unit_tests (238 passed, 22 skipped)
• build/Linux/Debug/engine_unit_tests (295/295 passed)
• cmake --build build/Linux/Debug (passed)
• Added 6 focused tests for DFlash2 geometry validation, config projection, exact KV bytes, bounded ring sizing, and rejection of unbounded cache pools.

Performance validation

Measured on one NVIDIA H200 with the CUDA Release build from stack head dabb4103, ORT CUDA 1.30.0, compact-final native-GDN artifacts, K=7, 1,024 generated tokens, CUDA graphs off, 2 warmups, and 3 measured runs:

Prompt tokens	MTP decode tokens/s	DFlash2 decode tokens/s	DFlash2/MTP
512	139.92	166.40	1.189x
2,048	138.59	155.92	1.125x
8,192	111.26	129.43	1.163x

DFlash2's geometric-mean decode speedup was 1.159x (+15.89%). Its evaluated-draft acceptance was 75.13%, 74.60%, and 72.86% respectively. All measured runs produced deterministic output hashes and speculative counters, with no dense-checkpoint compatibility warnings.

Dependencies

• Based on Run chained MTP drafting in the batch engine #2495.
• Uses auxiliary hidden-state exporter support from Tap auxiliary hidden states for speculative block drafters #2487/Export DFlash 2 block drafters for Qwen 3.5 and 3.8 #2489 when producing models; builder code is not duplicated here.
• Does not require an ONNX Runtime package or submodule bump.

Checklist

• Focused unit coverage added
• Full C++ and engine suites pass
• No dense recurrent checkpoints
• No ONNX Runtime dependency bump
• Real-model DFlash2 proposal and throughput validation

[Copilot]

Pull request overview

This PR integrates a DFlash2 block drafter into the paged dynamic-batching Engine, including plumbing to surface a packed aux_hidden_states decoder output and a new Dflash2Drafter runtime that maintains its own windowed paged KV ring and proposes draft token blocks post-commit.

Changes:

• Add an optional packed aux_hidden_states output to the decoder IO chain (and forward it through ScheduledRequests) for downstream speculative drafters.
• Introduce the Dflash2Model/Dflash2Drafter runtime and wire it into Engine::StepDynamic() (feed capture pre-commit, publish post-commit, per-request cache release on teardown).
• Extend config parsing for model.dflash2 and add focused C++ unit coverage for config projection, geometry validation, and cache sizing.

Reviewed changes

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

Show a summary per file

File	Description
test/cpp/dflash2_config_test.cpp	Adds unit tests for DFlash2 config validation/projection and cache sizing helpers.
src/engine/scheduled_requests.h	Exposes AuxHiddenStates() accessor for scheduled request batches.
src/engine/scheduled_requests.cpp	Implements ScheduledRequests::AuxHiddenStates() forwarding to decoder state.
src/engine/engine.h	Adds DFlash2 drafter ownership + step hooks for feed capture and draft publication.
src/engine/engine.cpp	Creates/owns the drafter, captures feeds pre-commit, publishes drafts post-commit, releases per-request cache blocks.
src/engine/decoders/decoder.h	Extends DecoderIO interface with AuxHiddenStates().
src/engine/decoders/hybrid_decoder_io.h	Forwards AuxHiddenStates() from varlen IO through hybrid IO.
src/engine/decoders/varlen_decoder_io.h	Stores optional aux_hidden_states_ output tensor and exposes it via AuxHiddenStates().
src/engine/decoders/varlen_decoder_io.cpp	Binds optional aux_hidden_states output; rejects CUDA graph capture when present.
src/dflash2_drafter.h	Defines Dflash2Model, CreateDflash2Config, and the Dflash2Drafter runtime API.
src/dflash2_drafter.cpp	Implements config projection, drafter cache sizing/allocation, packed layout construction, and greedy lattice walk.
src/config.h	Adds decoder aux_hidden_states output name and introduces the model.dflash2 config section.
src/config.cpp	Parses decoder aux_hidden_states output and the new model.dflash2 JSON section (including tensor-name overrides).

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.