Repository navigation
Integrate the DFlash2 device pipeline with Qwen graph execution - #2771
Akshay Sonawane (apsonawane) wants to merge 6 commits into
Conversation
Include native Engine and multi-row verification prerequisites, persistent staging, shared-stream GPU lattice selection, DFlash-only memory reductions, per-emission buffered logits, and a portable BF16/FP8 drafter exporter. Honor explicit ORT_LIB_PATH for custom runtime builds. Preserve published graph/resource behavior and device-only shared-weight placement. Fresh CUDA build and focused C++ checks: 147 passed, one skipped. Exact 64-token Engine/Generator parity at 2048 and 8193 tokens with repeated per-emission logits checks. Exporter FP8 smoke and rebuilt plane tests passed. Combined historical experiment and accuracy caveats are documented in the PR description. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
There was a problem hiding this comment.
🟡 Changes recommended
Supported-platform build failures and rejection of headless DFlash-only configurations remain unresolved.
4 open issues
What changed in this PR
Integrates three-token DFlash2 drafting with four-row Qwen graph verification through GenAI’s native Engine and Generator. This remains a draft integration dependent on companion ORT changes, with combined-stack numerical qualification pending.
Changes:
- Adds a single-slot Engine executor, multi-row device contract, and request lifecycle tests.
- Introduces asynchronous device drafting, unified prefill, and per-emission buffered logits.
- Adds a local-checkpoint drafter exporter and honors explicit Python
ORT_LIB_PATHselection.
| File | Description |
|---|---|
| tools/qwen38_graph/README.md | Documents integration, export, configuration, and qualification limits. |
| tools/qwen38_graph/export_dflash2.py | Adds BF16/FP8 drafter export using target embeddings and head. |
| test/cpp/engine/qwen38_graph_test_doubles.h | Adds a scripted multi-draft device. |
| test/cpp/engine/qwen38_graph_executor_tests.cpp | Tests execution, limits, cancellation, continuation, and failures. |
| src/qwen38_plane/qwen38_plane.cu | Adds device lattice selection, argmax completion, and CUDA helpers. |
| src/python/py/_dll_directory.py | Gives explicit ORT_LIB_PATH precedence. |
| src/models/qwen38_suffix_drafter.h | Adds historical suffix-based draft proposals. |
| src/models/qwen38_graph.h | Declares multi-row execution and DFlash state. |
| src/models/qwen38_graph.cpp | Integrates DFlash drafting and multi-row verification. |
| src/models/qwen38_graph_prefill.cpp | Adds unified prefill and lazy logits readback. |
| src/models/qwen38_decode_device.h | Defines the portable device contract and state ledger. |
| src/engine/scheduler.cpp | Honors backend-specific static batch limits. |
| src/engine/scheduled_requests.h | Declares executor-selected token staging. |
| src/engine/scheduled_requests.cpp | Commits staged tokens through requests. |
| src/engine/request.h | Declares external token commitment. |
| src/engine/request.cpp | Applies stopping rules to externally selected tokens. |
| src/engine/qwen38_graph_executor.h | Defines the Qwen session, cache manager, and executor. |
| src/engine/qwen38_graph_executor.cpp | Implements single-slot Engine execution and device leasing. |
| src/engine/model_executor.h | Adds token-selection, validation, and statistics hooks. |
| src/engine/engine.cpp | Routes Qwen models through the new executor. |
| src/engine/cache_manager.h | Adds static backend capacity hooks. |
| src/dflash2_drafter.h | Exposes device-selection outputs and persistent staging. |
| src/dflash2_drafter.cpp | Reuses staging buffers and supports asynchronous proposals. |
| src/decoding/qwen38_graph_strategy.h | Consolidates speculative statistics storage. |
| src/decoding/qwen38_graph_strategy.cpp | Handles multi-token queues and buffered logits. |
| src/config.h | Adds prefill and suffix-lookup settings. |
| src/config.cpp | Parses and validates the new settings. |
Review effort: Balanced
| if (auto graph_model = std::dynamic_pointer_cast<Qwen38GraphModel>(model)) { | ||
| return CreateQwen38GraphDependencies(std::move(graph_model)); | ||
| } |
|
|
||
| #include <cstddef> | ||
| #include <cstdint> | ||
| #include <span> |
| throw std::runtime_error("qwen38_graph: invalid numerical shadow cache " + name); | ||
| shape[0] = 1; | ||
| if (scoring) { | ||
| if (!name.ends_with(".scoring")) |
| if (q.fold.head_r2.filename.empty() && config_->model.dflash2.filename.empty()) | ||
| throw std::runtime_error("model.qwen38_graph.fold.head_r2.filename is required"); |
Expose aliased control/context outputs and order tail publication after the mid-stage state update; document custom-op removal and retained runtime helper. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Pack target weights and PLE into one external file, omit inactive head graphs, support headless folded configuration, and extend cache/rotary/pool contracts with four internal guard slots. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
|
Published consolidation and native 262K support in 7b22453. Includes out-of-place single-target-data packing, removal of inactive graph references, native-context exporter and rotary/KV/index/shared-pool updates. Validation: exact 512-token short-context parity, full 262112-prompt + 32-output boundary twice, zero fallback/DFlash failure/disable counters, EOS/logits checks, 148 native tests passed (1 skipped, 1 disabled). Peak sampled GPU memory at full window: 66.68 GiB on H200. Target KV/index allocation remains full-capacity rather than input-sized; long-context task accuracy is not qualified. Model assets and binaries are not committed. Data-plane integration is the next separate change. |
Replace model-local helper loading with a version/size-checked function table from the existing GenAI CUDA library. Remove the standalone helper target and reject obsolete library settings with a migration diagnostic. Update packaging, documentation and native/CUDA regression tests. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Pin the published combined ORT source and integrated GenAI data plane. Cover source checkout, CUDA plugin builds, model conversion prerequisites, packaging, native context, runtime tuning and validation limits. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
|
Latest integration source is now fully committed and published. GenAI runtime db154bc integrates the Qwen data plane into the standard CUDA add-on: no model-local libdflash_plane_v2.so, no separate helper target, and an explicit migration diagnostic for legacy plane_library configs. Documentation-only tip 65424b4 adds the complete pinned-source build/conversion/run guide at tools/qwen38_graph/BUILD_AND_RUN.md. This supersedes older notes saying a separate runtime helper remains required. Companion complete combined ORT integration branch: asonawane/remaining-standard-op-fusions at 13aa2a1e428a1e31ebc6f027e774674bc27d0aef (https://github.com/microsoft/onnxruntime/tree/asonawane/remaining-standard-op-fusions). Incremental operator PRs remain separate; no cherry-pick recipe onto main is implied. Pre-publication checks: combined ORT core/CUDA-plugin/provider-test build passed; 5 fold GPU tests passed. GenAI build and 150 native tests passed (1 skipped, 1 disabled), plus 15 source-tool tests and 4 integrated CUDA tests. Earlier exact source-content model validation preserved 512 token IDs at 2048/8192/16384/32768/8193, passed the 262112+32 boundary twice, EOS/logits and no-helper process mappings, with zero fallback/failure/disable counters. Full-window sampled peak: 67.13 GiB. No new full-model run or from-empty build was performed during this publication because GPUs are occupied. Accuracy caveats remain; model assets are not uploaded or committed. |
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
## Summary Replace the existing `com.microsoft` speculative-fold operators at opset 1 with the multi-row implementation required by Qwen/DFlash verification. No opset-2 schemas or kernels are added; model imports remain compatible with the supported `[1, 1]` `com.microsoft` domain range. - Support multi-row verification and muted prefill, including row counts greater than `head_rows`. - Preserve the combined verifier/downstream error status through `FoldTail`. - Add an InferenceSession load/run regression and CUDA launcher coverage. - Update the operator references and migration guide. ## Related GenAI changes - [GenAI #2771](microsoft/onnxruntime-genai#2771) normalizes converted graphs to the version-1 replacement schema (latest commit `efb26517`). - [GenAI #2776](microsoft/onnxruntime-genai#2776) emits `com.microsoft` opset 1 for the procedural graph builder (latest commit `2e63ae0c`). ## Base Incremental PR against `asonawane/main-ops-mtp-fold-state`. ## Validation - `lintrunner -a` passed for all touched ORT and GenAI paths; `git diff --check` passed. - GenAI graph-builder tests: **10 passed, 3 subtests passed**. - GenAI converter and target-graph tests: **14 passed**. - ORT C++/CUDA tests and generated-doc validation could not be run locally: this worktree has no configured ORT build directory or `nvcc`. These require CI/configured-build verification. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Summary
Integrate three-token DFlash2 drafting with four-row Qwen graph verification through native Engine and Generator.
ORT_LIB_PATHfrom Python instead of silently replacing it with the installed Python package's ORT library. This is required to load the custom API32/contrib stack.Base and dependencies
Incremental integration PR against
asonawane/main-qwen-graph-runtime, not main. Upstream main and the published feature branch do not contain all native Engine/multi-row prerequisites, so those required C++ pieces are included here explicitly.Requires the companion ORT branches based on:
asonawane/main-ops-sparse-attention-indexer(grouped INT8 S4)asonawane/main-ops-quantized-moe-fp8(INT2 S4 GEMV/fused routing)asonawane/main-ops-hyperconnection-tensor(S4 work-group tuning)No benchmark/evaluation scripts, downloaded checkpoints, model assets, binaries, or build helpers are committed. Added standalone tooling includes the drafter exporter and the native version-1 replacement fold converter/test. No Mobius PR is needed because the exporter/runtime changes are in GenAI.
Validation of this rebased GenAI branch
*Qwen38*:*Dflash2*:*DFlash2*): 147 passed, 1 skipped, 1 disabled test not executed.git diff --checkand changed-line C++ formatting passed.The three companion ORT branches have source/test translation-unit compilation only at this publication stage. The original combined experiment passed 24 numeric reference tests, but that is not exact-rebase operator qualification. A clean combined-branch ORT build and full numerical/accuracy qualification remain pending; therefore this is a draft integration PR.
Historical combined-stack H200 results
INT2 Qwen target + FP8 DFlash2; greedy batch one; six-run means:
Peak sampled GPU memory: 59.350 GiB, versus 64.322 GiB fresh old-runtime control. Exact 512-token parity with the prior DFlash candidate at all four contexts. Steady-state profiling observed no
cudaDeviceSynchronize/cudaHostAlloc, with lattice selection and fused INT2 FC2 kernels present.These are historical results for the previously tested combined experimental integration, not fresh six-run measurements of this rebased branch or per-component ablations. Synthetic prompts have unusually high draft acceptance; real-code throughput/acceptance varies.
Accuracy caveats
abc400_dregressed to a capacity-limited incomplete result;abc400_eandabc400_galso hit session capacity. Historical MTP is not a matched same-S4 control.Companion ORT draft PRs
Native fold version-1 replacement update
Commit
efb26517updates the graph converter, tests, and migration documentation to normalize legacycom.microsoft::FoldMid2/FoldTailgraphs to the native version-1 replacement schema. Companion ORT PR: microsoft/onnxruntime#33550 (baseasonawane/main-ops-mtp-fold-state).The converter preserves external-weight paths, exposes aliased control/context outputs and adds the mid-to-tail state dependency. It rejects repeat conversion and mismatched legacy control buffers. After conversion, the cleaned package no longer registers or distributes
libqwen_fold.soor other ONNX custom-op libraries, and no longer bundles provider/runtime directories. It still retainslibdflash_plane_v2.so, a required runtime helper, not an ONNX custom-op registration library. Source-built ORT core/provider and GenAI are required; stock upstream wheels are insufficient.Validation: all 4 converter tests passed on this publication branch; the combined converter/target-graph suite passed 14 tests. Previously reported combined-stack model parity and cleaned-package checks predate the version-1 replacement and are historical only. The five ORT GPU tests reported for the earlier v2 implementation do not validate this revision; ORT CUDA and InferenceSession tests remain pending a configured build/CI. Existing performance/accuracy caveats above still apply.