Reuse the block drafter runtime for DSpark

[Tianlei Wu (tianleiwu)]

Description

Reuses the DFlash2 block-drafter runtime from #2496 for DSpark artifacts. DSpark emits the same candidate-and-edge-score lattice with a Markov head, so the runtime only needs a config alias, support for one prediction per block row, and full-attention KV budgeting against the target paged-cache size.

This replaces #2462 while excluding DSpark builder changes already covered by #2490, historical diagnostics, dense-checkpoint experiments, and ONNX Runtime dependency changes.

Summary of Changes

Shared drafter configuration

• Parses model.dspark through the existing block-drafter runtime schema.
• Accepts num_draft_tokens == block_size for DSpark while preserving DFlash2's block_size - 1 layout.
• Continues to consume the same packed auxiliary hidden-state output and candidate lattice contract.

Full-attention cache budgeting

• Keeps the existing fixed batch-scaled ring for windowed DFlash2 artifacts.
• Bills full-attention DSpark KV bytes per target paged-cache block before target allocation.
• Creates a matching drafter pool after target sizing, with one query-spill block per configured batch row.

Scope

• Reuses the Run a DFlash2 block drafter in the batch engine #2496 engine lifecycle, accepted-prefix filtering, request teardown, and telemetry path.
• Does not duplicate Python builder/export code from Export DSpark block drafters for Qwen 3.5 and 3.8 #2490.
• Does not add dense recurrent checkpoints or mixed-checkpoint scheduling.

Testing

• build/Linux/Debug/unit_tests (240 passed, 22 skipped)
• build/Linux/Debug/engine_unit_tests (295/295 passed)
• cmake --build build/Linux/Debug (passed)
• Focused DFlash2/DSpark config and cache-planning tests (8/8 passed), including public JSON parsing of model.dspark.

Real-model DSpark proposal correctness and throughput remain follow-up validation for this draft.

Dependencies

• Based on Run a DFlash2 block drafter in the batch engine #2496.
• Uses DSpark exporter support from Export DSpark block drafters for Qwen 3.5 and 3.8 #2490; 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 DSpark proposal and throughput validation

[Copilot]

Pull request overview

Adds DSpark support by reusing the existing Engine-hosted block-drafter runtime (originally for DFlash2), including config aliasing, stricter contract validation, and updated paged-KV cache budgeting so full-attention drafters are accounted for against the target cache size.

Changes:

• Parse model.dspark as an alias of the existing DFlash2 block-drafter config/runtime and enforce mutual exclusivity with model.dflash2.
• Add full-attention KV budgeting support by charging auxiliary bytes-per-target-block plus fixed “query spill” reserved bytes before selecting target paged-cache capacity.
• Expand validation and test coverage for overflow handling, DSpark geometry, cache tensor bindings/types, and lattice output shapes.

Reviewed changes

Copilot reviewed 12 out of 12 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
src/config.cpp	Adds model.dspark alias parsing and rejects configuring both dspark and dflash2.
src/config.h	Extends block-drafter config with an is_dspark flag and clarifies DSpark vs DFlash2 geometry semantics.
src/dflash2_drafter.h	Generalizes the interface to cover DSpark and adds helpers for full-attention cache sizing/billing.
src/dflash2_drafter.cpp	Implements full target/drafter contract validation (cache tensors + lattice outputs) and adds full-attention sizing helpers.
src/engine/engine.cpp	Integrates full-attention drafter budgeting into engine cache planning and conditionally constructs fixed-window vs mirrored full-attention drafter pools.
src/engine/cache_manager.{h,cpp} / src/engine/paged_key_value_cache.{h,cpp}	Threads auxiliary reserved-memory into paged-cache capacity computation and adds overflow-safe sizing.
test/cpp/dflash2_config_test.cpp	Adds DSpark alias/geometry tests and broader compatibility validation coverage.
test/cpp/engine/paged_key_value_cache_tests.cpp	Adds explicit overflow coverage for paged-cache capacity math.
docs/paged_attention_engine.md	Updates documentation for DSpark aliasing and full-attention cache budgeting behavior.

---

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