From 29ab0bef3529e46ecff807932a79214710526a4b Mon Sep 17 00:00:00 2001 From: RicardoMin <17879681016@163.com> Date: Wed, 2 Sep 2026 10:36:43 +0800 Subject: [PATCH 1/3] feat(qwen3): add DFlash2 Phase 1 path selector Part of #930 Signed-off-by: RicardoMin <17879681016@163.com> --- docs/index.md | 1 + docs/models/qwen3/dflash2-phase1-930.md | 114 ++++ .../csrc/shared/dflash2_selector.cu | 222 +++++++ pegainfer-kernels/src/ffi/shared.rs | 32 + pegainfer-kernels/src/ops.rs | 5 + pegainfer-kernels/src/ops/dflash2.rs | 171 +++++ pegainfer-qwen3/src/config.rs | 611 +++++++++++++++--- pegainfer-qwen3/src/dflash.rs | 49 +- pegainfer-qwen3/src/dflash/loading.rs | 49 +- pegainfer-qwen3/src/dflash/manifest.rs | 91 +++ pegainfer-qwen3/src/dflash/reservation.rs | 18 +- pegainfer-qwen3/src/dflash/selector.rs | 247 +++++++ pegainfer-qwen3/src/dspark.rs | 4 +- pegainfer-qwen3/src/executor/dflash_lane.rs | 5 +- 14 files changed, 1527 insertions(+), 92 deletions(-) create mode 100644 docs/models/qwen3/dflash2-phase1-930.md create mode 100644 pegainfer-kernels/csrc/shared/dflash2_selector.cu create mode 100644 pegainfer-kernels/src/ops/dflash2.rs create mode 100644 pegainfer-qwen3/src/dflash/manifest.rs create mode 100644 pegainfer-qwen3/src/dflash/selector.rs diff --git a/docs/index.md b/docs/index.md index 021ca864a..7241a9fed 100644 --- a/docs/index.md +++ b/docs/index.md @@ -35,6 +35,7 @@ Organized by domain (model line / subsystem / playbook / lesson) instead of by l | `models/qwen3/prefix-cache.md` | Prefix caching on by default for Qwen3-4B: full-block kvbm radix matching at the executor, suffix-only prefill. Repeated ~1900-token prompt TTFT 141.8 → 16.3ms p50 (8.7×); warm TTFT ≈ TPOT + ~5ms setup. Includes the RoPE scalar-path corruption fix and the drain-the-stream TTFT measurement pitfall. | | `models/qwen3/dspark-integration.md` | DeepSeek **DSpark** Phase 1 is implemented for Qwen3-4B: DFlash backbone + rank-256 Markov head, anchor-first DeepSpec layout, one strided argmax-with-bias kernel, PDL polish, and one D2H per draft block. Greedy losslessness passes; 5090 block7 A/B vs matched DFlash shows DSpark +3.6% geomean output tok/s overall (+3–16% on text/code, random synthetic exception) and better accepted-draft distribution (2.52 vs 2.30 draft tokens/round). | | `models/qwen3/dflash-speculative-decoding.md` | DFlash speculative decoding behind `--dflash-draft-model-path`, modelled as an optimistic transaction (propose K → verify K+1 span → accept longest argmax prefix + 1 bonus → commit/roll back KV). Lossless up to bf16 tie-flips (bit-identical multi-token accepts; lm-eval gsm8k strict-match identical spec on/off). Single-stream decode 1.82× on 5070 Ti, 1.56× on 5090. Concurrent throughput fixed by batching the draft forward, then a piecewise verify CUDA Graph (dense ops captured, attention eager) closed single-stream: 5090 greedy c1 274 ≈ vLLM 278, c8 1525 > 1240, c16 1834 ≈ 1846 — all batch sizes now ≥ vLLM. Accept measured equal (9.1% vs 8.85%, same drafter); draft-side piecewise graph tracked next. Proposer trait deferred to EAGLE. | +| `models/qwen3/dflash2-phase1-930.md` | Issue #930 Phase 1 record: bounded top-16 DFlash2 candidate selection; dynamic convolution, sliding-window execution, and sampled rejection remain out of scope. | | `models/qwen3/accuracy-gate.md` | Qwen3 size-keyed logits golden gate (all six sizes 0.6B–32B committed) (`tests/hf_golden_gate.rs`): 48 teacher-forced sequences / 816 positions vs a stored HF bf16 golden, replayed over bs=1 / batched eager / CUDA-graph. Strict guards: regret check + mean ≤ 0.06 + p99 ≤ 0.20; absolute max printed but not asserted (coverage-unstable). Methodology in `subsystems/correctness/`. | | `models/qwen3/kernels-crate.md` | Phase 1 split implemented and 5090-verified: Qwen3-4B kernel surface lives in `pegainfer-kernels`; release build, test-target compile, accuracy gate, and bench snapshot pass. | | `models/qwen3/tp-design.md` | Qwen3 tensor-parallel design: `TP=2` milestone scope plus the controller/worker broadcast execution model, request identity, and coarse-grained step protocol for future TP/MoE work. | diff --git a/docs/models/qwen3/dflash2-phase1-930.md b/docs/models/qwen3/dflash2-phase1-930.md new file mode 100644 index 000000000..48e33cbb4 --- /dev/null +++ b/docs/models/qwen3/dflash2-phase1-930.md @@ -0,0 +1,114 @@ +# DFlash2 Phase 1 Selector + +> **TL;DR:** Issue #930 Phase 1 adds a bounded, deterministic top-16 candidate selector on top of the existing Qwen3 DFlash backbone; dynamic convolution, sliding-window execution, and sampled rejection are deliberately out of scope. +> +> **Last touched:** 2026-09 + +## Preparation + +- **Read**: + - `docs/index.md` - routes Qwen3 model and kernel design records. + - `docs/models/qwen3/dflash-speculative-decoding.md` - defines the existing proposer/verify/KV transaction contract and the batch layout. + - `docs/models/qwen3/dspark-integration.md` - documents the legacy anchor-first/Markov path that must remain unchanged. + - `docs/models/qwen3/kernels-crate.md` - assigns CUDA primitives and FFI ownership to `pegainfer-kernels`. + - `docs/models/qwen3/model-crate.md` - documents Qwen3 model-crate boundaries and single-GPU speculative decoding. + - `docs/conventions/coding-style.md` - requires focused tests and project logging conventions. + - `CLAUDE.md` - defines build, branch, and AI-assisted contribution requirements. +- **Relevant history**: + - `docs/models/qwen3/dflash-speculative-decoding.md` - the existing DFlash lane owns proposal while the shared verify and KV transaction contracts stay method-agnostic. + - No prior DFlash2 Phase 1 task record exists in this checkout. +- **Plan**: + 1. Audit the current configuration and loader scaffold; keep legacy DFlash and DSpark behavior unchanged and reject native hybrid DFlash2 capabilities that Phase 1 cannot execute. + 2. Load and validate the selector projection/codebooks, add a fixed-size GPU selector primitive and Rust wrapper, and account for its persistent and scratch allocations. + 3. Dispatch `TopKSelector` from the DFlash draft lane without changing draft span, verify, KV transaction, or CUDA-Graph shapes. + 4. Run formatting, compile, focused selector/reference checks, GPU-vs-reference checks, and legacy DFlash/DSpark regression checks; record actual results and limitations. +- **Risks / open questions**: + - The only discovered DFlash2 checkpoint also declares Phase 2 convolution and sliding-window capabilities; it must fail closed until those execution paths exist. + - Selector tie-breaking, anchor mapping, request-major row offsets, and scratch reservation must be deterministic and shape-safe. + +## Execution Log + +### Step 1: Normalize the DFlash2 capability contract + +- Added a `DFlashProposal::TopKSelector` capability and an explicit + `DFlashLayout` in `pegainfer-qwen3/src/config.rs`. +- Legacy DFlash and DSpark schemas remain on their existing proposal paths. +- Native DFlash2 configurations are parsed, but Phase 2 convolution, + sliding-window attention, anchor-first selector layouts, and an untied draft + output head fail closed before GPU weight allocation. + +### Step 2: Load selector weights and wire the proposer + +- Added SafeTensors manifest checks for the hidden projection and predecessor / + successor codebooks. +- Added persistent selector scratch and a two-launch CUDA implementation: + deterministic top-16 candidate extraction followed by a request-local path + walk using the predecessor/successor codebooks. +- Kept the existing full-block draft result contract, verify span, KV updates, + and CUDA-Graph shapes unchanged. + +### Step 3: Fix anchor-drop row mapping + +- The DFlash backbone emits an anchor-inclusive block. For the current + anchor-drop layout, row 0 is discarded by the executor and rows 1..N-1 are + the real proposal positions. +- The selector now uses compact candidate/output rows for those real positions, + while reading the corresponding rows from the original anchor-inclusive + logits/hidden buffers. Every request-local walk starts from the verified + anchor token, so no draft depends on a candidate from the discarded row 0. +- The host wrapper reconstructs `[anchor, selected_1, ..., selected_N-1]` for + the unchanged executor contract and rejects an invalid GPU token id before + it can reach token lookup. + +### Step 4: Lightweight verification + +Commands were run in the Linux feature checkout +`/database/ricardo.zheng/projects/open-access/pegainfer` +with `/usr/bin` present in `PATH` (the build script invokes `git`): + +| Command | Result | +| --- | --- | +| `cargo fmt --all -- --check` | Passed | +| `git diff --check` | Passed | +| `cargo check --release -p pegainfer-qwen3 --tests` | Passed; CUDA `sm_89` build | +| `cargo test --release -p pegainfer-qwen3 --lib` | 88 passed, 0 failed | +| `cargo test --release -p pegainfer-build --lib` | 8 passed, 0 failed | + +### Step 5: Remove redundant scaffolding + +- Kept the selector tensor preflight because the shared loader does not check + SafeTensors dtype or malformed rank; removed its unused `SelectorManifest` + wrapper and duplicate positive-value checks. +- Removed the unused selector scratch accessor and ABI-only bf16 assertion. +- Shortened comments to the anchor mapping, two-launch dependency, and + unsupported-capability boundaries. +- Re-ran formatting, Qwen3/build tests, and the server release build. + +| Cleanup verification | Result | +| --- | --- | +| `cargo fmt --all -- --check` | Passed | +| `git diff --check` | Passed | +| `cargo check --release -p pegainfer-qwen3 --tests` | Passed | +| `cargo test --release -p pegainfer-qwen3 --lib` | 88 passed, 0 failed | +| `cargo test --release -p pegainfer-build --lib` | 8 passed, 0 failed | +| `cargo build --release -p pegainfer-server --bin pegainfer` | Passed | + +## Debrief + +- **Outcome:** Phase 1 selector wiring, anchor-drop mapping, and a focused + cleanup of redundant scaffolding are complete in the feature branch. The + checkout is fast-forwarded to upstream main; no changes are staged or + committed. +- **Pitfalls encountered:** The first verification command omitted system + directories from `PATH`, so `pegainfer-kernels/build.rs` could not spawn + `git`. Re-running with `/usr/bin:/bin` succeeded. The row mapping bug was a + real semantic issue that compilation alone could not detect. +- **Lessons learned:** Selector buffers must distinguish the input block shape + from the compact set of positions actually proposed. The anchor is a + request-level predecessor, not a selector candidate when the executor drops + row 0. +- **Follow-ups:** Run the GPU selector/checkpoint import, Qwen3 greedy + losslessness, legacy DFlash/DSpark regressions, HTTP serving, and performance + A/B on the Linux GPU host. Phase 2 remains responsible for dynamic + convolution and sliding-window execution; Phase 3 remains responsible for + sampled losslessness/rejection sampling. diff --git a/pegainfer-kernels/csrc/shared/dflash2_selector.cu b/pegainfer-kernels/csrc/shared/dflash2_selector.cu new file mode 100644 index 000000000..db7252bd2 --- /dev/null +++ b/pegainfer-kernels/csrc/shared/dflash2_selector.cu @@ -0,0 +1,222 @@ +#include "common.cuh" + +#include +#include +#include +#include +#include +#include + +// DFlash2 keeps a small candidate set per draft position. The selector is +// deliberately split into two kernels: top-k is embarrassingly parallel over +// logits, while the path walk has a request-local dependency on the previously +// selected token. Keeping that dependency out of the top-k kernel makes both +// launches easy to reason about and keeps the temporary layout graph-safe. +namespace { + +constexpr int SELECTOR_TOP_K = 16; +constexpr int SELECTOR_TOPK_THREADS = 256; +constexpr int SELECTOR_WALK_THREADS = 512; + +__device__ __forceinline__ bool selector_better(float lhs_value, int lhs_id, + float rhs_value, int rhs_id) { + return lhs_value > rhs_value || + (lhs_value == rhs_value && lhs_id < rhs_id); +} + +__device__ __forceinline__ void selector_insert(float value, int id, float* values, int* ids) { + if (!selector_better(value, id, values[SELECTOR_TOP_K - 1], + ids[SELECTOR_TOP_K - 1])) { + return; + } + int slot = SELECTOR_TOP_K - 1; + while (slot > 0 && + selector_better(value, id, values[slot - 1], ids[slot - 1])) { + values[slot] = values[slot - 1]; + ids[slot] = ids[slot - 1]; + --slot; + } + values[slot] = value; + ids[slot] = id; +} + +__global__ void dflash2_selector_topk_kernel( + const __nv_bfloat16* __restrict__ logits, uint32_t* __restrict__ ids, + float* __restrict__ scores, int rows, int input_block_size, + int position_offset, int positions_per_request, int vocab) { + // Output rows are compact, while the source logits retain the anchor row. + // Translate each compact row back to its request-major input row. + const int row = blockIdx.x; + if (row >= rows || positions_per_request <= 0) { + return; + } + const int request = row / positions_per_request; + const int position = row % positions_per_request; + const size_t source_row = static_cast(request) * input_block_size + + position_offset + position; + + // Each thread keeps a private top-16 list. The lists occupy 32 KiB of + // shared memory and are merged by thread zero in canonical score/id order. + __shared__ float local_values[SELECTOR_TOPK_THREADS][SELECTOR_TOP_K]; + __shared__ int local_ids[SELECTOR_TOPK_THREADS][SELECTOR_TOP_K]; + float* my_values = local_values[threadIdx.x]; + int* my_ids = local_ids[threadIdx.x]; + for (int j = 0; j < SELECTOR_TOP_K; ++j) { + my_values[j] = -INFINITY; + my_ids[j] = INT_MAX; + } + + const __nv_bfloat16* row_logits = logits + source_row * vocab; + for (int token = threadIdx.x; token < vocab; + token += SELECTOR_TOPK_THREADS) { + selector_insert(__bfloat162float(row_logits[token]), token, my_values, + my_ids); + } + __syncthreads(); + + if (threadIdx.x == 0) { + float best_values[SELECTOR_TOP_K]; + int best_ids[SELECTOR_TOP_K]; + for (int j = 0; j < SELECTOR_TOP_K; ++j) { + best_values[j] = -INFINITY; + best_ids[j] = INT_MAX; + } + for (int thread = 0; thread < SELECTOR_TOPK_THREADS; ++thread) { + for (int j = 0; j < SELECTOR_TOP_K; ++j) { + selector_insert(local_values[thread][j], local_ids[thread][j], + best_values, best_ids); + } + } + for (int j = 0; j < SELECTOR_TOP_K; ++j) { + ids[static_cast(row) * SELECTOR_TOP_K + j] = + static_cast(best_ids[j]); + scores[static_cast(row) * SELECTOR_TOP_K + j] = best_values[j]; + } + } +} + +__global__ void dflash2_selector_walk_kernel( + const __nv_bfloat16* __restrict__ projected_hidden, + const __nv_bfloat16* __restrict__ predecessor, + const __nv_bfloat16* __restrict__ successor, + const uint32_t* __restrict__ anchor_tokens, + const uint32_t* __restrict__ candidate_ids, + const float* __restrict__ candidate_unary, uint32_t* __restrict__ output, + int requests, int input_block_size, int position_offset, + int positions_per_request, int vocab, int rank) { + const int request = blockIdx.x; + if (request >= requests) { + return; + } + + __shared__ float edge_scores[SELECTOR_TOP_K]; + __shared__ uint32_t edge_ids[SELECTOR_TOP_K]; + __shared__ uint32_t previous; + if (threadIdx.x == 0) { + previous = anchor_tokens[request]; + } + __syncthreads(); + + const int lane = threadIdx.x & 31; + const int candidate = threadIdx.x >> 5; + for (int position = 0; position < positions_per_request; ++position) { + // Candidate/output rows are compact; hidden rows retain the anchor slot. + const int row = request * positions_per_request + position; + const size_t source_row = static_cast(request) * input_block_size + + position_offset + position; + if (candidate < SELECTOR_TOP_K) { + const uint32_t candidate_id = + candidate_ids[static_cast(row) * SELECTOR_TOP_K + candidate]; + float dot = 0.0f; + if (candidate_id < static_cast(vocab) && + previous < static_cast(vocab)) { + const __nv_bfloat16* hidden_row = + projected_hidden + source_row * rank; + const __nv_bfloat16* predecessor_row = + predecessor + static_cast(previous) * rank; + const __nv_bfloat16* successor_row = + successor + static_cast(candidate_id) * rank; + for (int component = lane; component < rank; component += 32) { + dot += __bfloat162float(predecessor_row[component]) * + __bfloat162float(hidden_row[component]) * + __bfloat162float(successor_row[component]); + } + } + dot = warp_reduce_sum(dot); + if (lane == 0) { + edge_ids[candidate] = candidate_id; + edge_scores[candidate] = + candidate_unary[static_cast(row) * SELECTOR_TOP_K + + candidate] + + dot; + } + } + __syncthreads(); + + if (threadIdx.x == 0) { + uint32_t best_id = edge_ids[0]; + float best_score = edge_scores[0]; + for (int j = 1; j < SELECTOR_TOP_K; ++j) { + if (selector_better(edge_scores[j], static_cast(edge_ids[j]), + best_score, static_cast(best_id))) { + best_score = edge_scores[j]; + best_id = edge_ids[j]; + } + } + output[row] = best_id; + previous = best_id; + } + __syncthreads(); + } +} + +} // namespace + +extern "C" int dflash2_selector_topk_cuda( + const __nv_bfloat16* logits, uint32_t* candidate_ids, + float* candidate_scores, int rows, int input_block_size, + int position_offset, int positions_per_request, int vocab, + cudaStream_t stream) { + if (logits == nullptr || candidate_ids == nullptr || candidate_scores == nullptr || + rows <= 0 || input_block_size <= 0 || position_offset < 0 || + positions_per_request <= 0 || + position_offset > input_block_size - positions_per_request || + vocab < SELECTOR_TOP_K) { + return static_cast(cudaErrorInvalidValue); + } + dflash2_selector_topk_kernel<<>>( + logits, candidate_ids, candidate_scores, rows, input_block_size, + position_offset, positions_per_request, vocab); + return static_cast(cudaGetLastError()); +} + +extern "C" int dflash2_selector_walk_cuda( + const __nv_bfloat16* projected_hidden, + const __nv_bfloat16* predecessor, + const __nv_bfloat16* successor, + const uint32_t* anchor_tokens, + const uint32_t* candidate_ids, + const float* candidate_unary, + uint32_t* output, + int requests, + int input_block_size, + int position_offset, + int positions_per_request, + int vocab, + int rank, + cudaStream_t stream) { + if (projected_hidden == nullptr || predecessor == nullptr || + successor == nullptr || anchor_tokens == nullptr || candidate_ids == nullptr || + candidate_unary == nullptr || output == nullptr || requests <= 0 || + input_block_size <= 0 || position_offset < 0 || + positions_per_request <= 0 || + position_offset > input_block_size - positions_per_request || + vocab < SELECTOR_TOP_K || rank <= 0) { + return static_cast(cudaErrorInvalidValue); + } + dflash2_selector_walk_kernel<<>>( + projected_hidden, predecessor, successor, anchor_tokens, + candidate_ids, candidate_unary, output, requests, input_block_size, + position_offset, positions_per_request, vocab, rank); + return static_cast(cudaGetLastError()); +} diff --git a/pegainfer-kernels/src/ffi/shared.rs b/pegainfer-kernels/src/ffi/shared.rs index 08f159707..731fc802d 100644 --- a/pegainfer-kernels/src/ffi/shared.rs +++ b/pegainfer-kernels/src/ffi/shared.rs @@ -940,6 +940,38 @@ unsafe extern "C" { stream: CUstream, ); + // DFlash2 proposer-side selector. The first launch writes a deterministic + // top-k candidate list for each logits row; the second launch walks those + // candidates request by request and writes one selected token per row. + pub fn dflash2_selector_topk_cuda( + logits: *const Half, + candidate_ids: *mut u32, + candidate_scores: *mut f32, + rows: i32, + input_block_size: i32, + position_offset: i32, + positions_per_request: i32, + vocab: i32, + stream: CUstream, + ) -> i32; + + pub fn dflash2_selector_walk_cuda( + projected_hidden: *const Half, + predecessor: *const Half, + successor: *const Half, + anchor_tokens: *const u32, + candidate_ids: *const u32, + candidate_unary: *const f32, + output: *mut u32, + requests: i32, + input_block_size: i32, + position_offset: i32, + positions_per_request: i32, + vocab: i32, + rank: i32, + stream: CUstream, + ) -> i32; + pub fn bf16_to_f32_cuda( input: *const Half, output: *mut f32, diff --git a/pegainfer-kernels/src/ops.rs b/pegainfer-kernels/src/ops.rs index a144c65cf..ebe2dbf44 100644 --- a/pegainfer-kernels/src/ops.rs +++ b/pegainfer-kernels/src/ops.rs @@ -5,6 +5,7 @@ mod attention; mod deepep; #[cfg(feature = "deepseek-v2-lite")] mod deepseek_v2_lite; +mod dflash2; mod elementwise; mod embedding; #[cfg(feature = "gemma4")] @@ -83,6 +84,10 @@ pub use deepep::glm52_deepep_info; pub use deepep::glm52_ep_deepep_unique_id; #[cfg(feature = "deepseek-v2-lite")] pub use deepseek_v2_lite::*; +pub use dflash2::DFlash2SelectorScratch; +pub use dflash2::dflash2_selector_into; +pub use dflash2::dflash2_selector_scratch_bytes; +pub use dflash2::dflash2_selector_selected_host; pub use elementwise::SuppressIds; pub use elementwise::accumulate_bf16_token_scaled_to_f32_into; pub use elementwise::add_batch; diff --git a/pegainfer-kernels/src/ops/dflash2.rs b/pegainfer-kernels/src/ops/dflash2.rs new file mode 100644 index 000000000..5115f20d9 --- /dev/null +++ b/pegainfer-kernels/src/ops/dflash2.rs @@ -0,0 +1,171 @@ +//! CUDA-backed DFlash2 candidate selection. +//! +//! Runs deterministic top-k extraction and request-local path walks. + +use anyhow::Result; +use anyhow::anyhow; +use anyhow::ensure; +use cudarc::driver::CudaSlice; +use cudarc::driver::DevicePtr; +use cudarc::driver::DevicePtrMut; + +use crate::ffi; +use crate::tensor::DeviceContext; +use crate::tensor::DeviceMatrix; +use crate::tensor::HiddenStates; + +const DFLASH2_SELECTOR_TOP_K: usize = 16; + +/// Persistent buffers for the two selector launches. +pub struct DFlash2SelectorScratch { + pub candidate_ids: CudaSlice, + pub candidate_scores: CudaSlice, + pub selected: CudaSlice, + rows: usize, +} + +impl DFlash2SelectorScratch { + pub fn new(ctx: &DeviceContext, rows: usize) -> Result { + ensure!(rows > 0, "DFlash2 selector rows must be positive"); + let candidate_rows = rows + .checked_mul(DFLASH2_SELECTOR_TOP_K) + .ok_or_else(|| anyhow!("DFlash2 selector scratch size overflow"))?; + Ok(Self { + candidate_ids: ctx.stream.alloc_zeros(candidate_rows)?, + candidate_scores: ctx.stream.alloc_zeros(candidate_rows)?, + selected: ctx.stream.alloc_zeros(rows)?, + rows, + }) + } +} + +/// Run top-16 extraction and request-local path walks. +/// +/// Inputs are anchor-inclusive request-major rows; output rows are compact. +/// Ties are resolved by score descending, then token id ascending. +#[allow(clippy::too_many_arguments)] +pub fn dflash2_selector_into( + ctx: &DeviceContext, + logits: &HiddenStates, + projected_hidden: &HiddenStates, + predecessor: &DeviceMatrix, + successor: &DeviceMatrix, + anchors: &CudaSlice, + input_block_size: usize, + position_offset: usize, + positions_per_request: usize, + scratch: &mut DFlash2SelectorScratch, +) -> Result<()> { + ensure!( + input_block_size > 0, + "DFlash2 selector input block size must be positive" + ); + ensure!( + positions_per_request > 0, + "DFlash2 selector positions must be positive" + ); + ensure!(position_offset <= input_block_size); + ensure!( + positions_per_request <= input_block_size - position_offset, + "DFlash2 selector positions exceed the input block" + ); + ensure!(logits.seq_len == projected_hidden.seq_len); + ensure!(logits.seq_len.is_multiple_of(input_block_size)); + let requests = logits.seq_len / input_block_size; + let compact_rows = requests + .checked_mul(positions_per_request) + .ok_or_else(|| anyhow!("DFlash2 selector compact row count overflow"))?; + ensure!( + anchors.len() >= requests, + "selector anchor buffer is too small" + ); + ensure!( + scratch.rows >= compact_rows, + "selector scratch is too small" + ); + ensure!(projected_hidden.hidden_dim == predecessor.cols); + ensure!(predecessor.cols == successor.cols); + ensure!(predecessor.rows == successor.rows); + ensure!(logits.hidden_dim >= DFLASH2_SELECTOR_TOP_K); + ensure!(predecessor.rows == logits.hidden_dim); + let rows_i32 = i32::try_from(compact_rows) + .map_err(|_| anyhow!("DFlash2 selector row count exceeds i32"))?; + let requests_i32 = i32::try_from(requests) + .map_err(|_| anyhow!("DFlash2 selector request count exceeds i32"))?; + let input_block_size_i32 = i32::try_from(input_block_size) + .map_err(|_| anyhow!("DFlash2 selector input block size exceeds i32"))?; + let position_offset_i32 = i32::try_from(position_offset) + .map_err(|_| anyhow!("DFlash2 selector position offset exceeds i32"))?; + let positions_per_request_i32 = i32::try_from(positions_per_request) + .map_err(|_| anyhow!("DFlash2 selector position count exceeds i32"))?; + let vocab_i32 = i32::try_from(logits.hidden_dim) + .map_err(|_| anyhow!("DFlash2 selector vocabulary size exceeds i32"))?; + let rank_i32 = i32::try_from(predecessor.cols) + .map_err(|_| anyhow!("DFlash2 selector rank exceeds i32"))?; + + let (logits_ptr, _gl) = logits.data.device_ptr(&ctx.stream); + let (candidate_ids_ptr, _gi) = scratch.candidate_ids.device_ptr_mut(&ctx.stream); + let (candidate_scores_ptr, _gs) = scratch.candidate_scores.device_ptr_mut(&ctx.stream); + let status = unsafe { + ffi::dflash2_selector_topk_cuda( + logits_ptr as *const ffi::Half, + candidate_ids_ptr as *mut u32, + candidate_scores_ptr as *mut f32, + rows_i32, + input_block_size_i32, + position_offset_i32, + positions_per_request_i32, + vocab_i32, + crate::tensor::active_cu_stream(ctx), + ) + }; + ensure!( + status == 0, + "DFlash2 selector top-k launch failed: {status}" + ); + + let (hidden_ptr, _gh) = projected_hidden.data.device_ptr(&ctx.stream); + let (predecessor_ptr, _gp) = predecessor.data.device_ptr(&ctx.stream); + let (successor_ptr, _gs) = successor.data.device_ptr(&ctx.stream); + let (anchors_ptr, _ga) = anchors.device_ptr(&ctx.stream); + let (selected_ptr, _go) = scratch.selected.device_ptr_mut(&ctx.stream); + let status = unsafe { + ffi::dflash2_selector_walk_cuda( + hidden_ptr as *const ffi::Half, + predecessor_ptr as *const ffi::Half, + successor_ptr as *const ffi::Half, + anchors_ptr as *const u32, + candidate_ids_ptr as *const u32, + candidate_scores_ptr as *const f32, + selected_ptr as *mut u32, + requests_i32, + input_block_size_i32, + position_offset_i32, + positions_per_request_i32, + vocab_i32, + rank_i32, + crate::tensor::active_cu_stream(ctx), + ) + }; + ensure!(status == 0, "DFlash2 selector path launch failed: {status}"); + Ok(()) +} + +/// Copy selected token ids to host after the caller synchronizes. +pub fn dflash2_selector_selected_host( + ctx: &DeviceContext, + scratch: &DFlash2SelectorScratch, + rows: usize, +) -> Result> { + ensure!(rows <= scratch.rows, "selector output rows exceed scratch"); + ctx.stream + .clone_dtoh(&scratch.selected.slice(..rows)) + .map_err(|e| anyhow!("DFlash2 selector D2H failed: {e}")) +} + +/// Bytes required by [`DFlash2SelectorScratch`] for `rows` active rows. +pub const fn dflash2_selector_scratch_bytes(rows: usize) -> usize { + rows * (DFLASH2_SELECTOR_TOP_K * std::mem::size_of::() + + DFLASH2_SELECTOR_TOP_K * std::mem::size_of::() + + std::mem::size_of::()) +} diff --git a/pegainfer-qwen3/src/config.rs b/pegainfer-qwen3/src/config.rs index 0e43a5140..bf514ac63 100644 --- a/pegainfer-qwen3/src/config.rs +++ b/pegainfer-qwen3/src/config.rs @@ -8,6 +8,8 @@ use serde::Deserialize; use serde_json::Value; pub(crate) const PREFILL_ATTENTION_CTA_TILE_Q: i32 = 64; +pub(crate) const DFLASH2_SELECTOR_TOP_K: usize = 16; +const DEFAULT_MARKOV_HEAD_TYPE: &str = "vanilla"; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub(crate) struct TensorParallelConfig { @@ -43,46 +45,63 @@ pub(crate) struct Config { pub(crate) stop_token_ids: Vec, } -/// Resolved drafter config shared by DFlash and DSpark. DSpark extends the -/// DFlash backbone with a Markov head and an optional confidence head; -/// `markov_rank == 0` is plain DFlash. Two on-disk schemas are normalized into -/// this in `from_file`: our `Qwen3-4B-DFlash-b16` nests -/// `dflash_config: {mask_token_id, target_layer_ids}` and puts `rope_theta` at -/// the top level, while DeepSpec's `dflash_/dspark_*_block7` put those fields -/// flat and nest `rope_theta` under `rope_parameters`. +/// Normalized DFlash/DSpark configuration for legacy and native DFlash2 schemas. +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) enum DFlashProposal { + PlainArgmax, + Markov { rank: usize, head_type: String }, + TopKSelector { rank: usize, top_k: usize }, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) enum DFlashLayout { + /// Position zero is an anchor slot and is removed before verification. + AnchorDrop, + /// Position zero is already the first proposed token. + AnchorFirst, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) struct DynamicConv { + pub(crate) kernel_size: usize, + pub(crate) group_size: usize, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) struct SlidingWindow { + pub(crate) window_size: usize, + pub(crate) non_causal: bool, + pub(crate) enabled: bool, + pub(crate) layer_types: Vec, +} + #[derive(Clone, Debug)] pub(crate) struct DFlashConfig { pub(crate) hidden_size: usize, + pub(crate) target_hidden_size: usize, pub(crate) intermediate_size: usize, pub(crate) num_hidden_layers: usize, pub(crate) num_attention_heads: usize, pub(crate) num_key_value_heads: usize, - num_target_layers: usize, pub(crate) head_dim: usize, pub(crate) vocab_size: usize, pub(crate) rms_norm_eps: f32, pub(crate) rope_theta: f32, pub(crate) max_position_embeddings: usize, + /// Legacy target-layer count; absent in native DFlash2. + num_target_layers: Option, pub(crate) block_size: usize, pub(crate) mask_token_id: u32, pub(crate) target_layer_ids: Vec, - /// DSpark Markov head low-rank size; 0 disables the head (= plain DFlash). - pub(crate) markov_rank: usize, - markov_head_type: String, - /// Block draft layout. DeepSpec `Qwen3DSparkModel` checkpoints (both the - /// markov and the `markov_rank == 0` ones) are *anchor-first*: block position - /// 0 is already the first real prediction, so all `block_size` positions - /// draft (verify span `block_size + 1`). Our native `DFlashDraftModel` (b16) - /// is *anchor-drop*: position 0 is a throwaway anchor slot, so only positions - /// `1..block_size` draft (verify span `block_size`). This is a property of the - /// checkpoint, NOT of the markov head — keying it on the markov head silently - /// mis-drafts a no-markov DeepSpec checkpoint (accept rate collapses to ~0). - anchor_first: bool, - /// Whether the checkpoint carries a confidence head. Phase 1 does not use it - /// (full-block verify, no confidence-scheduled truncation); surfaced at load - /// so the operator knows that capability is being ignored. See - /// docs/models/qwen3/dspark-integration.md (Phase 2). + pub(crate) draft_vocab_size: usize, + /// Proposal capability selected by the checkpoint schema. + pub(crate) proposal: DFlashProposal, + pub(crate) layout: DFlashLayout, + /// Whether the target output head can be reused by the drafter. + pub(crate) reuse_target_head: bool, pub(crate) enable_confidence_head: bool, + pub(crate) dynamic_convolution: Option, + pub(crate) sliding_window: Option, } #[derive(Clone, Debug, Deserialize)] @@ -97,14 +116,15 @@ struct RopeParameters { } fn default_markov_head_type() -> String { - "vanilla".to_string() + DEFAULT_MARKOV_HEAD_TYPE.to_owned() } -/// On-disk drafter config tolerant of both the nested (`b16`) and flat -/// (DeepSpec) schemas; `from_file` resolves it into `DFlashConfig`. +/// Legacy nested and flat drafter schemas. #[derive(Deserialize)] struct RawDFlashConfig { hidden_size: usize, + #[serde(default)] + target_hidden_size: Option, intermediate_size: usize, num_hidden_layers: usize, num_attention_heads: usize, @@ -112,6 +132,8 @@ struct RawDFlashConfig { num_target_layers: usize, head_dim: usize, vocab_size: usize, + #[serde(default)] + draft_vocab_size: Option, rms_norm_eps: f32, #[serde(default)] rope_theta: Option, @@ -132,12 +154,77 @@ struct RawDFlashConfig { markov_head_type: String, #[serde(default)] enable_confidence_head: bool, - /// DeepSpec `Qwen3DSparkModel` checkpoints declare this (anchor-first - /// drafting); native `DFlashDraftModel` checkpoints omit it (anchor-drop). + #[serde(default)] + selector_rank: Option, + #[serde(default)] + selector_top_k: Option, + /// DeepSpec declares anchors; native DFlash2 omits them. #[serde(default)] num_anchors: Option, } +#[derive(Debug, Deserialize)] +struct RawDFlash2TransformerConfig { + hidden_size: usize, + intermediate_size: usize, + num_hidden_layers: usize, + num_attention_heads: usize, + num_key_value_heads: usize, + head_dim: usize, + vocab_size: usize, + rms_norm_eps: f32, + rope_parameters: RopeParameters, + #[serde(default = "default_max_position_embeddings")] + max_position_embeddings: usize, + #[serde(default)] + sliding_window: Option, + #[serde(default)] + use_sliding_window: bool, + #[serde(default)] + layer_types: Vec, +} + +#[derive(Debug, Deserialize)] +struct RawDFlash2ProposalMethod { + proposal_type: String, + speculative_tokens: usize, + verifier_accept_k: usize, + #[serde(default)] + accept_tolerance: f32, +} + +#[derive(Debug, Deserialize)] +struct RawDFlash2SpeculatorsConfig { + algorithm: String, + default_proposal_method: String, + #[serde(default)] + proposal_methods: Vec, +} + +#[derive(Debug, Deserialize)] +struct RawDFlash2Config { + transformer_layer_config: RawDFlash2TransformerConfig, + aux_hidden_state_layer_ids: Vec, + block_size: usize, + mask_token_id: u32, + draft_vocab_size: usize, + selector_rank: usize, + selector_top_k: usize, + sample_from_anchor: bool, + #[serde(default)] + target_hidden_size: Option, + + #[serde(default)] + conv_kernel_size: Option, + #[serde(default)] + conv_group_size: Option, + + #[serde(default)] + sliding_window_non_causal: Option, + + speculators_config: RawDFlash2SpeculatorsConfig, +} + /// EAGLE-3 drafter config (e.g. `AngelSlim/Qwen3-4B_eagle3`). /// /// A single-layer (`midlayer`) head whose attention takes @@ -269,61 +356,435 @@ impl DFlashConfig { pub(crate) fn from_file(model_path: &str) -> Result { let config_path = format!("{}/config.json", model_path); let content = fs::read_to_string(&config_path)?; - let raw: RawDFlashConfig = serde_json::from_str(&content)?; + let json: Value = serde_json::from_str(&content)?; + + let model_type = match json.get("speculators_model_type") { + None => None, + Some(value) => Some( + value + .as_str() + .context("speculators_model_type must be a string")? + .to_owned(), + ), + }; - // rope_theta: flat (b16) or nested under rope_parameters (DeepSpec). - let rope_theta = raw - .rope_theta - .or(raw.rope_parameters.map(|r| r.rope_theta)) - .context("drafter config missing rope_theta / rope_parameters.rope_theta")?; + let architecture_marker = match json.get("architectures") { + None => false, + Some(value) => value + .as_array() + .context("architectures must be an array")? + .iter() + .any(|value| value.as_str() == Some("DFlash2DraftModel")), + }; - // mask_token_id + target_layer_ids: nested dflash_config (b16) or flat (DeepSpec). - let (mask_token_id, target_layer_ids) = match raw.dflash_config { - Some(inner) => (inner.mask_token_id, inner.target_layer_ids), + let config = match (model_type.as_deref(), architecture_marker) { + (Some("dflash2"), true) => { + let raw: RawDFlash2Config = serde_json::from_value(json)?; + Self::from_dflash2_raw(raw) + } + (None | Some("dflash" | "dspark"), false) => { + let raw: RawDFlashConfig = serde_json::from_value(json)?; + Self::from_legacy_raw(&raw) + } + (Some("dflash2"), false) => { + bail!( + "DFlash2 config declares speculators_model_type=dflash2 but not DFlash2DraftModel" + ) + } + (None, true) => { + bail!( + "DFlash2 config declares DFlash2DraftModel but missing speculators_model_type" + ) + } + (Some(kind), true) => { + bail!("unsupported speculator type {kind:?} for DFlash2DraftModel") + } + (Some(kind), false) => { + bail!("unsupported speculator model type {kind:?}") + } + }?; + config.validate_proposal()?; + Ok(config) + } + + fn from_legacy_raw(raw: &RawDFlashConfig) -> Result { + let nested_rope_theta = raw.rope_parameters.as_ref().map(|value| value.rope_theta); + let rope_theta = match (raw.rope_theta, nested_rope_theta) { + (Some(flat), Some(nested)) => { + ensure!( + flat.to_bits() == nested.to_bits(), + "legacy DFlash rope_theta and rope_parameters.rope_theta disagree" + ); + flat + } + (Some(value), None) | (None, Some(value)) => value, + (None, None) => { + bail!("drafter config missing rope_theta / rope_parameters.rope_theta") + } + }; + + let (mask_token_id, target_layer_ids) = match raw.dflash_config.as_ref() { + Some(inner) => { + if let Some(flat) = raw.mask_token_id { + ensure!( + flat == inner.mask_token_id, + "legacy DFlash mask_token_id and dflash_config.mask_token_id disagree" + ); + } + if let Some(flat) = raw.target_layer_ids.as_ref() { + ensure!( + flat == &inner.target_layer_ids, + "legacy DFlash target_layer_ids and dflash_config.target_layer_ids disagree" + ); + } + (inner.mask_token_id, inner.target_layer_ids.clone()) + } None => ( raw.mask_token_id - .context("drafter config missing mask_token_id (no dflash_config block)")?, + .context("drafter config missing mask_token_id")?, raw.target_layer_ids - .context("drafter config missing target_layer_ids (no dflash_config block)")?, + .clone() + .context("drafter config missing target_layer_ids")?, ), }; + let markov_rank = raw.markov_rank; + let markov_head_type = raw.markov_head_type.clone(); + ensure!( + markov_rank > 0 || markov_head_type == default_markov_head_type(), + "legacy DFlash markov_head_type must be \"vanilla\" when markov_rank is zero" + ); + if let Some(num_anchors) = raw.num_anchors { + ensure!( + num_anchors > 0, + "legacy DFlash num_anchors must be positive when declared" + ); + } + ensure!( + markov_rank == 0 || raw.num_anchors.is_some(), + "legacy DFlash Markov proposals must declare num_anchors" + ); + + let proposal = match (markov_rank, raw.selector_rank, raw.selector_top_k) { + (0, None, None) => DFlashProposal::PlainArgmax, + (rank, None, None) => DFlashProposal::Markov { + rank, + head_type: markov_head_type, + }, + (0, Some(rank), Some(top_k)) => DFlashProposal::TopKSelector { rank, top_k }, + (0, _, _) => { + bail!("legacy DFlash selector_rank and selector_top_k must be declared together") + } + (_, Some(_), _) | (_, _, Some(_)) => { + bail!("legacy DFlash cannot enable both Markov and top-k selector proposals") + } + }; + Ok(Self { hidden_size: raw.hidden_size, + target_hidden_size: raw.target_hidden_size.unwrap_or(raw.hidden_size), intermediate_size: raw.intermediate_size, num_hidden_layers: raw.num_hidden_layers, num_attention_heads: raw.num_attention_heads, num_key_value_heads: raw.num_key_value_heads, - num_target_layers: raw.num_target_layers, head_dim: raw.head_dim, vocab_size: raw.vocab_size, rms_norm_eps: raw.rms_norm_eps, rope_theta, max_position_embeddings: raw.max_position_embeddings, + num_target_layers: Some(raw.num_target_layers), block_size: raw.block_size, mask_token_id, target_layer_ids, - markov_rank: raw.markov_rank, - markov_head_type: raw.markov_head_type, + draft_vocab_size: raw.draft_vocab_size.unwrap_or(raw.vocab_size), + proposal, + layout: if raw.num_anchors.is_some() { + DFlashLayout::AnchorFirst + } else { + DFlashLayout::AnchorDrop + }, + reuse_target_head: true, enable_confidence_head: raw.enable_confidence_head, - // A markov head only ever ships on a DeepSpec (anchor-first) - // checkpoint, so num_anchors is always present alongside it; treat - // markov as an independent corroborating signal so a future flat - // schema can't accidentally route a markov checkpoint anchor-drop. - anchor_first: raw.num_anchors.is_some() || raw.markov_rank > 0, + dynamic_convolution: None, + sliding_window: None, }) } - /// DSpark Markov head is active (`markov_rank > 0`); the draft uses the - /// semi-autoregressive sample loop instead of independent argmax. + fn from_dflash2_raw(raw: RawDFlash2Config) -> Result { + let RawDFlash2Config { + transformer_layer_config, + aux_hidden_state_layer_ids, + block_size, + mask_token_id, + draft_vocab_size, + selector_rank, + selector_top_k, + sample_from_anchor, + target_hidden_size, + conv_kernel_size, + conv_group_size, + sliding_window_non_causal, + speculators_config, + } = raw; + + let RawDFlash2TransformerConfig { + hidden_size, + intermediate_size, + num_hidden_layers, + num_attention_heads, + num_key_value_heads, + head_dim, + vocab_size, + rms_norm_eps, + rope_parameters, + max_position_embeddings, + sliding_window, + use_sliding_window, + layer_types, + } = transformer_layer_config; + + ensure!( + speculators_config.algorithm == "dflash2", + "DFlash2 speculators_config.algorithm must be \"dflash2\", got {:?}", + speculators_config.algorithm + ); + ensure!( + speculators_config.default_proposal_method == "greedy", + "DFlash2 proposal method {:?} is not supported in Phase 1", + speculators_config.default_proposal_method + ); + ensure!( + speculators_config.proposal_methods.len() == 1, + "DFlash2 Phase 1 requires exactly one proposal method, got {}", + speculators_config.proposal_methods.len() + ); + let proposal_method = &speculators_config.proposal_methods[0]; + ensure!( + proposal_method.proposal_type == "greedy", + "DFlash2 proposal type {:?} is not supported in Phase 1", + proposal_method.proposal_type + ); + ensure!( + proposal_method.accept_tolerance == 0.0, + "DFlash2 accept_tolerance {} is not supported in Phase 1", + proposal_method.accept_tolerance + ); + let expected_speculative_tokens = if sample_from_anchor { + block_size + } else { + block_size.saturating_sub(1) + }; + ensure!( + proposal_method.speculative_tokens == expected_speculative_tokens, + "DFlash2 proposal speculative_tokens {} does not match block/layout expectation {}", + proposal_method.speculative_tokens, + expected_speculative_tokens + ); + ensure!( + proposal_method.verifier_accept_k == 1, + "DFlash2 verifier_accept_k {} is not supported in Phase 1", + proposal_method.verifier_accept_k + ); + ensure!( + draft_vocab_size == vocab_size, + "DFlash2 draft_vocab_size {} must match transformer vocab_size {} for the full-vocabulary selector", + draft_vocab_size, + vocab_size + ); + let dynamic_convolution = match (conv_kernel_size, conv_group_size) { + (Some(kernel_size), Some(group_size)) => { + ensure!( + kernel_size > 0 && group_size > 0, + "DFlash2 convolution kernel_size and group_size must be positive" + ); + Some(DynamicConv { + kernel_size, + group_size, + }) + } + (None, None) => None, + _ => bail!("DFlash2 convolution fields must be declared together"), + }; + + ensure!( + layer_types.is_empty() || layer_types.len() == num_hidden_layers, + "DFlash2 transformer layer_types length {} does not match num_hidden_layers {}", + layer_types.len(), + num_hidden_layers + ); + for (layer_idx, kind) in layer_types.iter().enumerate() { + ensure!( + matches!(kind.as_str(), "full_attention" | "sliding_attention"), + "DFlash2 transformer layer_types[{}] has unsupported value {:?}", + layer_idx, + kind + ); + } + let has_non_full_layer = layer_types.iter().any(|kind| kind == "sliding_attention"); + let declares_sliding = use_sliding_window + || sliding_window.is_some() + || sliding_window_non_causal.is_some() + || has_non_full_layer; + let sliding_window = if declares_sliding { + let window_size = sliding_window + .context("DFlash2 sliding-window capability is missing sliding_window")?; + ensure!( + window_size > 0, + "DFlash2 sliding_window must be positive, got {}", + window_size + ); + let non_causal = sliding_window_non_causal.context( + "DFlash2 sliding-window capability is missing sliding_window_non_causal", + )?; + Some(SlidingWindow { + window_size, + non_causal, + enabled: use_sliding_window, + layer_types, + }) + } else { + None + }; + + let proposal = DFlashProposal::TopKSelector { + rank: selector_rank, + top_k: selector_top_k, + }; + + let layout = if sample_from_anchor { + DFlashLayout::AnchorFirst + } else { + DFlashLayout::AnchorDrop + }; + + Ok(Self { + hidden_size, + target_hidden_size: target_hidden_size.unwrap_or(hidden_size), + intermediate_size, + num_hidden_layers, + num_attention_heads, + num_key_value_heads, + head_dim, + vocab_size, + rms_norm_eps, + rope_theta: rope_parameters.rope_theta, + max_position_embeddings, + + num_target_layers: None, + + block_size, + mask_token_id, + target_layer_ids: aux_hidden_state_layer_ids, + draft_vocab_size, + proposal, + layout, + reuse_target_head: false, + + enable_confidence_head: false, + + dynamic_convolution, + sliding_window, + }) + } + + /// Whether this config selects the DSpark Markov path. pub(crate) fn uses_markov_head(&self) -> bool { - self.markov_rank > 0 + matches!(&self.proposal, DFlashProposal::Markov { .. }) } - /// Anchor-first block layout (see [`DFlashConfig::anchor_first`]). The markov - /// head implies it, but a `markov_rank == 0` DeepSpec checkpoint needs it too. + pub(crate) fn uses_selector(&self) -> bool { + matches!(&self.proposal, DFlashProposal::TopKSelector { .. }) + } + + pub(crate) fn markov_rank(&self) -> usize { + match &self.proposal { + DFlashProposal::Markov { rank, .. } => *rank, + _ => 0, + } + } + + /// Whether position zero is a real draft token in this checkpoint layout. pub(crate) fn anchor_first(&self) -> bool { - self.anchor_first + matches!(&self.layout, DFlashLayout::AnchorFirst) + } + + /// Validate normalized proposal and vocabulary invariants. + fn validate_proposal(&self) -> Result<()> { + ensure!( + self.block_size >= 2, + "DFlash block_size must be at least 2, got {}", + self.block_size + ); + match &self.proposal { + DFlashProposal::PlainArgmax => {} + DFlashProposal::Markov { rank, head_type } => { + ensure!(*rank > 0, "Markov proposal rank must be positive"); + ensure!( + head_type == DEFAULT_MARKOV_HEAD_TYPE, + "DSpark markov_head_type {:?} not supported (only \"vanilla\")", + head_type + ); + } + DFlashProposal::TopKSelector { rank, top_k } => { + ensure!(*rank > 0, "DFlash selector rank must be positive"); + ensure!( + *top_k == DFLASH2_SELECTOR_TOP_K, + "DFlash selector_top_k must be {}, got {}", + DFLASH2_SELECTOR_TOP_K, + top_k + ); + ensure!( + self.draft_vocab_size >= *top_k, + "DFlash selector_top_k {} exceeds draft vocabulary size {}", + top_k, + self.draft_vocab_size + ); + } + } + ensure!( + self.draft_vocab_size > 0 && self.draft_vocab_size <= self.vocab_size, + "DFlash draft_vocab_size {} must be in 1..={}", + self.draft_vocab_size, + self.vocab_size, + ); + Ok(()) + } + + /// Reject capabilities outside the Phase 1 execution path. + pub(crate) fn validate_runtime_capabilities(&self) -> Result<()> { + if let Some(conv) = &self.dynamic_convolution { + bail!( + "DFlash2 dynamic convolution (kernel_size={}, group_size={}) is not supported in Phase 1", + conv.kernel_size, + conv.group_size + ); + } + + if let Some(window) = &self.sliding_window { + bail!( + "DFlash2 sliding-window attention (window_size={}, non_causal={}, use_sliding_window={}, layer_types={:?}) is not supported in Phase 1", + window.window_size, + window.non_causal, + window.enabled, + window.layer_types + ); + } + + if matches!( + (&self.proposal, &self.layout), + ( + DFlashProposal::TopKSelector { .. }, + DFlashLayout::AnchorFirst + ) + ) { + bail!("DFlash Phase 1 supports only anchor-drop selector checkpoints"); + } + + ensure!( + self.reuse_target_head, + "DFlash checkpoint uses an independent draft output head, which the current runtime does not load" + ); + Ok(()) } pub(crate) fn validate_for_target(&self, target: &Config) -> Result<()> { @@ -334,11 +795,19 @@ impl DFlashConfig { target.hidden_size ); anyhow::ensure!( - self.num_target_layers == target.num_hidden_layers, - "DFlash num_target_layers {} does not match target layers {}", - self.num_target_layers, - target.num_hidden_layers + self.target_hidden_size == target.hidden_size, + "DFlash target_hidden_size {} does not match target {}", + self.target_hidden_size, + target.hidden_size ); + if let Some(num_target_layers) = self.num_target_layers { + anyhow::ensure!( + num_target_layers == target.num_hidden_layers, + "DFlash num_target_layers {} does not match target layers {}", + num_target_layers, + target.num_hidden_layers + ); + } anyhow::ensure!( self.num_attention_heads == target.num_attention_heads && self.num_key_value_heads == target.num_key_value_heads @@ -351,6 +820,12 @@ impl DFlashConfig { self.vocab_size, target.vocab_size ); + anyhow::ensure!( + self.draft_vocab_size == target.vocab_size, + "DFlash draft_vocab_size {} must match target vocab_size {} for full-vocabulary logits", + self.draft_vocab_size, + target.vocab_size + ); anyhow::ensure!( self.rope_theta.to_bits() == target.rope_theta.to_bits(), "DFlash rope_theta {} does not match target {}", @@ -364,12 +839,7 @@ impl DFlashConfig { target.max_position_embeddings ); anyhow::ensure!( - self.block_size >= 2, - "DFlash block_size must be >= 2, got {}", - self.block_size - ); - anyhow::ensure!( - self.mask_token_id < target.vocab_size as u32, + u64::from(self.mask_token_id) < target.vocab_size as u64, "DFlash mask_token_id {} is outside target vocab_size {}", self.mask_token_id, target.vocab_size @@ -392,19 +862,6 @@ impl DFlashConfig { .all(|pair| pair[0] < pair[1]), "DFlash target_layer_ids must be strictly increasing" ); - - // DSpark Markov head: only the released `vanilla` low-rank head is - // implemented; reject the gated/rnn variants loudly rather than silently - // mis-drafting. The confidence head is intentionally ignored in Phase 1 - // (full-block verify, no confidence-scheduled truncation) — its weights - // are simply not loaded; see docs/models/qwen3/dspark-integration.md. - if self.uses_markov_head() { - anyhow::ensure!( - self.markov_head_type == "vanilla", - "DSpark markov_head_type {:?} not supported (only \"vanilla\")", - self.markov_head_type - ); - } Ok(()) } } diff --git a/pegainfer-qwen3/src/dflash.rs b/pegainfer-qwen3/src/dflash.rs index 6cdeb1ef9..778cdfadb 100644 --- a/pegainfer-qwen3/src/dflash.rs +++ b/pegainfer-qwen3/src/dflash.rs @@ -6,6 +6,8 @@ use pegainfer_core::tensor::DeviceContext; use pegainfer_core::tensor::DeviceMatrix; use pegainfer_core::tensor::DeviceVec; use pegainfer_core::tensor::HiddenStates; +use selector::SelectorScratch; +use selector::SelectorWeights; use crate::config::DFlashConfig; use crate::dspark::MarkovHead; @@ -14,7 +16,9 @@ use crate::weights::Qwen3Model; use crate::weights::TransformerBlock; mod loading; +mod manifest; mod reservation; +pub(crate) mod selector; pub(crate) use reservation::DFlashMemoryReservation; @@ -30,6 +34,8 @@ pub(crate) struct DFlashDraftModel { /// draft proposes via [`MarkovHead::sample_block`] (anchor-first, all /// `block_size` positions) instead of an independent per-position argmax. markov: Option, + /// DFlash2 selector, mutually exclusive with the DSpark Markov head. + selector: Option, } pub(crate) struct DFlashRequestState { @@ -95,6 +101,8 @@ pub(crate) struct DFlashBatchScratch { v_tail: HiddenStates, // DSpark Markov sample-loop scratch; `None` for plain DFlash drafters. markov: Option, + // DFlash2 selector scratch; absent for legacy drafters. + selector: Option, } impl DFlashRequestState { @@ -260,6 +268,10 @@ impl DFlashBatchScratch { .uses_markov_head() .then(|| MarkovScratch::new(ctx, config, max_decode_batch_size)) .transpose()?, + selector: config + .uses_selector() + .then(|| SelectorScratch::new(ctx, config, max_decode_batch_size)) + .transpose()?, }) } @@ -340,6 +352,10 @@ impl DFlashDraftModel { self.markov.is_some() } + pub(crate) fn uses_selector(&self) -> bool { + self.selector.is_some() + } + /// Anchor-first block layout (a checkpoint property, see /// [`DFlashConfig::anchor_first`]) — drives both the verify span and the /// draft-block slice start, independently of the markov head. @@ -794,8 +810,39 @@ impl DFlashDraftModel { ) } + /// Select one bounded path per request from the draft outputs. + pub(crate) fn selector_draft_tokens( + &self, + ctx: &DeviceContext, + current_tokens: &[u32], + scratch: &mut DFlashBatchScratch, + ) -> Result> { + let selector = self + .selector + .as_ref() + .context("selector_draft_tokens called on a non-selector drafter")?; + let DFlashBatchScratch { + logits, + logits_normed, + selector: selector_scratch, + .. + } = scratch; + let selector_scratch = selector_scratch + .as_mut() + .context("selector scratch was not allocated for this drafter")?; + selector.select_block( + ctx, + logits, + logits_normed, + current_tokens, + self.block_size(), + self.anchor_first(), + selector_scratch, + ) + } + fn context_feature_dim(&self) -> usize { - self.config.hidden_size * self.target_layer_ids().len() + self.config.target_hidden_size * self.target_layer_ids().len() } fn project_context_into( diff --git a/pegainfer-qwen3/src/dflash/loading.rs b/pegainfer-qwen3/src/dflash/loading.rs index b935619e5..b1241a5db 100644 --- a/pegainfer-qwen3/src/dflash/loading.rs +++ b/pegainfer-qwen3/src/dflash/loading.rs @@ -12,7 +12,13 @@ use pegainfer_core::weight_loader::load_tensor_2d; use pegainfer_core::weight_loader::mmap_shards; use super::DFlashDraftModel; +use super::manifest::HIDDEN_PROJECTION_TENSOR; +use super::manifest::PREDECESSOR_CODEBOOK_TENSOR; +use super::manifest::SUCCESSOR_CODEBOOK_TENSOR; +use super::manifest::validate_selector_tensors; +use super::selector::SelectorWeights; use crate::config::DFlashConfig; +use crate::config::DFlashProposal; use crate::dspark::MARKOV_W1_TENSOR; use crate::dspark::MARKOV_W2_TENSOR; use crate::dspark::MarkovHead; @@ -30,6 +36,8 @@ impl DFlashDraftModel { let config = DFlashConfig::from_file(model_path) .with_context(|| format!("load DFlash config from {model_path}"))?; config.validate_for_target(target.config())?; + // Reject unsupported capabilities before GPU allocation. + config.validate_runtime_capabilities()?; let (shard_paths, weight_map) = load_shard_info(model_path)?; debug!( @@ -39,6 +47,32 @@ impl DFlashDraftModel { let mmaps = mmap_shards(&shard_paths)?; let shards = deserialize_shards(&mmaps)?; + let selector = if let DFlashProposal::TopKSelector { rank, .. } = &config.proposal { + validate_selector_tensors( + &shards, + &weight_map, + *rank, + config.hidden_size, + config.draft_vocab_size, + )?; + let hidden_projection = + load_tensor_2d(ctx, &shards, &weight_map, HIDDEN_PROJECTION_TENSOR)?; + let predecessor_codebook = + load_tensor_2d(ctx, &shards, &weight_map, PREDECESSOR_CODEBOOK_TENSOR)?; + let successor_codebook = + load_tensor_2d(ctx, &shards, &weight_map, SUCCESSOR_CODEBOOK_TENSOR)?; + Some(SelectorWeights::new( + *rank, + config.hidden_size, + config.draft_vocab_size, + hidden_projection, + predecessor_codebook, + successor_codebook, + )?) + } else { + None + }; + let mut layers = Vec::with_capacity(config.num_hidden_layers); for layer_idx in 0..config.num_hidden_layers { let prefix = format!("layers.{layer_idx}"); @@ -139,16 +173,16 @@ impl DFlashDraftModel { // DSpark Markov head (Phase 1). The confidence head and the tied // embed_tokens/lm_head are intentionally skipped: the head is byte-identical // to the target's, which we reuse for the verify-equivalent logits. + if config.enable_confidence_head { + log::info!( + "DFlash confidence head present in {model_path} but unused in Phase 1 \ + (full-block verify, no confidence-scheduled truncation)" + ); + } let markov = if config.uses_markov_head() { let w1 = load_tensor_2d(ctx, &shards, &weight_map, MARKOV_W1_TENSOR)?; let w2 = load_tensor_2d(ctx, &shards, &weight_map, MARKOV_W2_TENSOR)?; - if config.enable_confidence_head { - log::info!( - "DSpark confidence head present in {model_path} but unused in Phase 1 \ - (full-block verify, no confidence-scheduled truncation)" - ); - } - Some(MarkovHead::new(config.markov_rank, w1, w2)?) + Some(MarkovHead::new(config.markov_rank(), w1, w2)?) } else { None }; @@ -173,6 +207,7 @@ impl DFlashDraftModel { cos_cache, sin_cache, markov, + selector, }) } } diff --git a/pegainfer-qwen3/src/dflash/manifest.rs b/pegainfer-qwen3/src/dflash/manifest.rs new file mode 100644 index 000000000..b6e1590b8 --- /dev/null +++ b/pegainfer-qwen3/src/dflash/manifest.rs @@ -0,0 +1,91 @@ +//! Preflight validation for DFlash2 selector tensors. + +use std::collections::HashMap; + +use anyhow::Context; +use anyhow::Result; +use anyhow::ensure; +use safetensors::Dtype; +use safetensors::SafeTensors; +use safetensors::tensor::TensorView; + +pub(crate) const HIDDEN_PROJECTION_TENSOR: &str = "candidate_selector.hidden_projection.weight"; +pub(crate) const PREDECESSOR_CODEBOOK_TENSOR: &str = "candidate_selector.predecessor_codebook"; +pub(crate) const SUCCESSOR_CODEBOOK_TENSOR: &str = "candidate_selector.successor_codebook"; + +/// Validate selector dtype and shapes before GPU upload. +pub(crate) fn validate_selector_tensors( + shards: &[SafeTensors<'_>], + weight_map: &HashMap, + rank: usize, + hidden_size: usize, + vocab_size: usize, +) -> Result<()> { + validate_matrix( + shards, + weight_map, + HIDDEN_PROJECTION_TENSOR, + rank, + hidden_size, + )?; + validate_matrix( + shards, + weight_map, + PREDECESSOR_CODEBOOK_TENSOR, + vocab_size, + rank, + )?; + validate_matrix( + shards, + weight_map, + SUCCESSOR_CODEBOOK_TENSOR, + vocab_size, + rank, + )?; + + Ok(()) +} + +fn validate_matrix( + shards: &[SafeTensors<'_>], + weight_map: &HashMap, + name: &str, + expected_rows: usize, + expected_cols: usize, +) -> Result<()> { + let tensor = find_tensor(shards, weight_map, name)?; + ensure!( + tensor.dtype() == Dtype::BF16, + "DFlash selector tensor {name:?} must be BF16, got {:?}", + tensor.dtype() + ); + ensure!( + tensor.shape() == [expected_rows, expected_cols], + "DFlash selector tensor {name:?} has shape {:?}, expected [{expected_rows}, {expected_cols}]", + tensor.shape() + ); + Ok(()) +} + +fn find_tensor<'a>( + shards: &'a [SafeTensors<'a>], + weight_map: &HashMap, + name: &str, +) -> Result> { + if let Some(&shard_idx) = weight_map.get(name) { + let shard = shards.get(shard_idx).with_context(|| { + format!("DFlash selector tensor {name:?} references missing shard index {shard_idx}") + })?; + return shard + .tensor(name) + .with_context(|| format!("load DFlash selector tensor {name:?}")); + } + + for shard in shards { + if let Ok(tensor) = shard.tensor(name) { + return Ok(tensor); + } + } + + anyhow::bail!("DFlash selector tensor {name:?} is missing") +} diff --git a/pegainfer-qwen3/src/dflash/reservation.rs b/pegainfer-qwen3/src/dflash/reservation.rs index c71a9932c..942dea669 100644 --- a/pegainfer-qwen3/src/dflash/reservation.rs +++ b/pegainfer-qwen3/src/dflash/reservation.rs @@ -1,6 +1,8 @@ use anyhow::Result; use crate::config::DFlashConfig; +use crate::config::DFlashProposal; +use crate::dflash::selector::SelectorScratch; use crate::dspark::MarkovHead; /// GPU memory DFlash needs on top of the target KV pool, derived from the draft @@ -29,12 +31,14 @@ pub(crate) struct DFlashMemoryReservation { impl DFlashMemoryReservation { pub(crate) fn from_path(draft_path: &str, max_decode_batch_size: usize) -> Result { let config = DFlashConfig::from_file(draft_path)?; + config.validate_runtime_capabilities()?; Ok(Self::from_config(&config, max_decode_batch_size)) } pub(crate) fn from_config(config: &DFlashConfig, max_decode_batch_size: usize) -> Self { const BF16: usize = 2; let hidden = config.hidden_size; + let target_hidden = config.target_hidden_size; let kv_dim = config.num_key_value_heads * config.head_dim; let q_dim = config.num_attention_heads * config.head_dim; let inter = config.intermediate_size; @@ -47,7 +51,7 @@ impl DFlashMemoryReservation { // tail, which is one block past the prefix. let context_scratch = 2 * hidden * BF16; // context_projected + context_hidden let tail_scratch = (hidden + 2 * kv_dim) * BF16; // tail_input + k_tail + v_tail - let pending = hidden * capture_layers * BF16; // context_feature_dim + let pending = target_hidden * capture_layers * BF16; // context_feature_dim let kv_bytes_per_token = draft_kv + context_scratch + tail_scratch + pending; // Lane-level batched dense scratch: every dense buffer is sized for the @@ -65,7 +69,7 @@ impl DFlashMemoryReservation { + q_dim * hidden // o_proj + hidden * 2 * inter // gate_up_proj + inter * hidden); // down_proj - let fc = BF16 * hidden * (hidden * capture_layers); // context projection + let fc = BF16 * hidden * (target_hidden * capture_layers); // context projection let weights = per_layer * config.num_hidden_layers + fc; let weights = weights + weights / 10; @@ -79,10 +83,18 @@ impl DFlashMemoryReservation { // DSpark Markov head: weights (2 × vocab × rank) + sample scratch (the // per-step bias is the dominant term). Zero for plain DFlash drafters. let markov = MarkovHead::reservation_bytes(config, max_decode_batch_size); + let selector = match &config.proposal { + DFlashProposal::TopKSelector { rank, .. } => { + // Selector weights and persistent launch scratch. + let weights = BF16 * (rank * hidden + 2 * config.draft_vocab_size * rank); + weights + SelectorScratch::bytes(config, max_decode_batch_size) + } + _ => 0, + }; Self { kv_bytes_per_token, - fixed_bytes: weights + scratch_total + block_headroom + markov, + fixed_bytes: weights + scratch_total + block_headroom + markov + selector, } } } diff --git a/pegainfer-qwen3/src/dflash/selector.rs b/pegainfer-qwen3/src/dflash/selector.rs new file mode 100644 index 000000000..0d6a98e66 --- /dev/null +++ b/pegainfer-qwen3/src/dflash/selector.rs @@ -0,0 +1,247 @@ +//! CUDA-backed DFlash2 candidate selection. +//! +//! DFlash2 selector weights and persistent scratch. +//! +//! CUDA performs bounded candidate and path selection. + +use anyhow::Context; +use anyhow::Result; +use anyhow::ensure; +use cudarc::driver::CudaSlice; +use pegainfer_core::ops; +use pegainfer_core::tensor::DeviceContext; +use pegainfer_core::tensor::DeviceMatrix; +use pegainfer_core::tensor::HiddenStates; +use pegainfer_kernels::ops::DFlash2SelectorScratch; +use pegainfer_kernels::ops::dflash2_selector_into; +use pegainfer_kernels::ops::dflash2_selector_scratch_bytes; +use pegainfer_kernels::ops::dflash2_selector_selected_host; + +use crate::config::DFlashConfig; + +/// DFlash2 selector weights. +pub(crate) struct SelectorWeights { + hidden_projection: DeviceMatrix, + predecessor_codebook: DeviceMatrix, + successor_codebook: DeviceMatrix, +} + +impl SelectorWeights { + pub(crate) fn new( + rank: usize, + hidden_size: usize, + vocab_size: usize, + hidden_projection: DeviceMatrix, + predecessor_codebook: DeviceMatrix, + successor_codebook: DeviceMatrix, + ) -> Result { + ensure!(rank > 0, "DFlash selector rank must be positive"); + ensure!( + hidden_size > 0, + "DFlash selector hidden size must be positive" + ); + ensure!( + vocab_size > 0, + "DFlash selector vocabulary must be positive" + ); + ensure!( + hidden_projection.rows == rank && hidden_projection.cols == hidden_size, + "DFlash selector projection shape {}x{} does not match {}x{}", + hidden_projection.rows, + hidden_projection.cols, + rank, + hidden_size + ); + ensure!( + predecessor_codebook.rows == vocab_size && predecessor_codebook.cols == rank, + "DFlash selector predecessor shape {}x{} does not match {}x{}", + predecessor_codebook.rows, + predecessor_codebook.cols, + vocab_size, + rank + ); + ensure!( + successor_codebook.rows == vocab_size && successor_codebook.cols == rank, + "DFlash selector successor shape {}x{} does not match {}x{}", + successor_codebook.rows, + successor_codebook.cols, + vocab_size, + rank + ); + Ok(Self { + hidden_projection, + predecessor_codebook, + successor_codebook, + }) + } + + pub(crate) fn select_block( + &self, + ctx: &DeviceContext, + logits: &HiddenStates, + logits_normed: &HiddenStates, + current_tokens: &[u32], + block_size: usize, + anchor_first: bool, + scratch: &mut SelectorScratch, + ) -> Result> { + ensure!(!current_tokens.is_empty(), "selector needs active requests"); + ensure!(block_size > 0, "selector block size must be positive"); + let rows = current_tokens + .len() + .checked_mul(block_size) + .context("selector row count overflow")?; + ensure!( + logits.seq_len == rows, + "selector logits rows {} != requests {} * block {}", + logits.seq_len, + current_tokens.len(), + block_size + ); + ensure!( + logits_normed.seq_len == rows, + "selector hidden rows {} != requests {} * block {}", + logits_normed.seq_len, + current_tokens.len(), + block_size + ); + ensure!( + logits.hidden_dim == self.successor_codebook.rows, + "selector logits vocabulary {} != codebook vocabulary {}", + logits.hidden_dim, + self.successor_codebook.rows + ); + ensure!( + logits_normed.hidden_dim == self.hidden_projection.cols, + "selector hidden size {} != projection input {}", + logits_normed.hidden_dim, + self.hidden_projection.cols + ); + scratch.activate(rows, current_tokens.len(), self.hidden_projection.rows)?; + + let mut anchor_dst = scratch.anchor_tokens.slice_mut(..current_tokens.len()); + ctx.stream.memcpy_htod(current_tokens, &mut anchor_dst)?; + ops::gemm_into( + ctx, + &self.hidden_projection, + logits_normed, + &mut scratch.projected_hidden, + ); + let (position_offset, positions_per_request) = if anchor_first { + (0, block_size) + } else { + let positions = block_size + .checked_sub(1) + .context("anchor-drop selector block size must exceed one")?; + (1, positions) + }; + dflash2_selector_into( + ctx, + logits, + &scratch.projected_hidden, + &self.predecessor_codebook, + &self.successor_codebook, + &scratch.anchor_tokens, + block_size, + position_offset, + positions_per_request, + &mut scratch.selector, + )?; + + let selected_rows = current_tokens + .len() + .checked_mul(positions_per_request) + .context("selector selected row count overflow")?; + let selected = dflash2_selector_selected_host(ctx, &scratch.selector, selected_rows)?; + ctx.sync()?; + ensure!( + selected + .iter() + .all(|&token_id| u64::from(token_id) < logits.hidden_dim as u64), + "DFlash2 selector produced a token outside vocabulary size {}", + logits.hidden_dim + ); + if anchor_first { + ensure!(selected.len() == rows); + return Ok(selected); + } + + // Restore the existing anchor-inclusive draft contract. + let mut output = Vec::with_capacity(rows); + for (request_idx, &anchor) in current_tokens.iter().enumerate() { + output.push(anchor); + let start = request_idx * positions_per_request; + output.extend_from_slice(&selected[start..start + positions_per_request]); + } + Ok(output) + } +} + +/// Persistent GPU scratch for a selector-enabled lane. +pub(crate) struct SelectorScratch { + max_rows: usize, + projected_hidden: HiddenStates, + anchor_tokens: CudaSlice, + selector: DFlash2SelectorScratch, +} + +impl SelectorScratch { + pub(crate) fn new( + ctx: &DeviceContext, + config: &DFlashConfig, + max_decode_batch_size: usize, + ) -> Result { + let rank = match &config.proposal { + crate::config::DFlashProposal::TopKSelector { rank, .. } => *rank, + _ => { + return Err(anyhow::anyhow!( + "selector scratch requested without selector" + )); + } + }; + let max_rows = max_decode_batch_size + .checked_mul(config.block_size) + .context("selector scratch row count overflow")?; + Ok(Self { + max_rows, + projected_hidden: HiddenStates::zeros(ctx, rank, max_rows)?, + anchor_tokens: ctx.stream.alloc_zeros(max_decode_batch_size)?, + selector: DFlash2SelectorScratch::new(ctx, max_rows)?, + }) + } + + fn activate(&mut self, rows: usize, requests: usize, rank: usize) -> Result<()> { + ensure!( + rows <= self.max_rows, + "selector rows {} exceed scratch capacity {}", + rows, + self.max_rows + ); + ensure!( + requests <= self.anchor_tokens.len(), + "selector requests {} exceed anchor capacity {}", + requests, + self.anchor_tokens.len() + ); + ensure!( + self.projected_hidden.hidden_dim == rank, + "selector scratch rank {} != projection rank {}", + self.projected_hidden.hidden_dim, + rank + ); + self.projected_hidden.seq_len = rows; + Ok(()) + } + + pub(crate) fn bytes(config: &DFlashConfig, max_decode_batch_size: usize) -> usize { + const BF16: usize = 2; + let rank = match &config.proposal { + crate::config::DFlashProposal::TopKSelector { rank, .. } => *rank, + _ => return 0, + }; + let rows = max_decode_batch_size * config.block_size; + rows * rank * BF16 + + max_decode_batch_size * std::mem::size_of::() + + dflash2_selector_scratch_bytes(rows) + } +} diff --git a/pegainfer-qwen3/src/dspark.rs b/pegainfer-qwen3/src/dspark.rs index ed132e8cf..c21e08d6b 100644 --- a/pegainfer-qwen3/src/dspark.rs +++ b/pegainfer-qwen3/src/dspark.rs @@ -147,7 +147,7 @@ impl MarkovHead { return 0; } let vocab = config.vocab_size; - let rank = config.markov_rank; + let rank = config.markov_rank(); let weights = 2 * vocab * rank * BF16; let scratch = MarkovScratch::bytes(vocab, rank, config.block_size, max_decode_batch_size); weights + scratch @@ -180,7 +180,7 @@ impl MarkovScratch { "DSpark markov scratch needs a non-zero batch size" ); let vocab = config.vocab_size; - let rank = config.markov_rank; + let rank = config.markov_rank(); let partials = markov_step_argmax_partials_len(max_decode_batch_size, vocab); let sampled = max_decode_batch_size * config.block_size; Ok(Self { diff --git a/pegainfer-qwen3/src/executor/dflash_lane.rs b/pegainfer-qwen3/src/executor/dflash_lane.rs index d064a5278..c123a5de7 100644 --- a/pegainfer-qwen3/src/executor/dflash_lane.rs +++ b/pegainfer-qwen3/src/executor/dflash_lane.rs @@ -277,8 +277,9 @@ impl LocalQwen3Lane { // Propose tokens from the base logits. DFlash takes an independent // greedy argmax per position; DSpark adds the Markov bias and samples // the block left-to-right (anchor-first, all `block_size` positions). - let markov = model.uses_markov_head(); - let sampled = if markov { + let sampled = if model.uses_selector() { + model.selector_draft_tokens(self.model.device_ctx(), ¤t_tokens, scratch)? + } else if model.uses_markov_head() { model.markov_draft_tokens(self.model.device_ctx(), ¤t_tokens, scratch)? } else { let greedy = SamplingParams::default(); From 9ce23575683af3cfe1c15e9ac86197aba11a4212 Mon Sep 17 00:00:00 2001 From: RicardoMin <17879681016@163.com> Date: Wed, 2 Sep 2026 19:31:02 +0800 Subject: [PATCH 2/3] fix(qwen3): support native DFlash2 head variants Signed-off-by: RicardoMin <17879681016@163.com> --- docs/models/qwen3/dflash2-phase1-930.md | 62 ++++++++-- pegainfer-qwen3/src/config.rs | 33 +++-- pegainfer-qwen3/src/dflash.rs | 16 +-- pegainfer-qwen3/src/dflash/loading.rs | 27 ++++- pegainfer-qwen3/src/dflash/manifest.rs | 30 ++++- pegainfer-qwen3/src/dflash/reservation.rs | 7 +- .../tests/dflash_speculative_gate.rs | 114 +++++++++++++++++- 7 files changed, 248 insertions(+), 41 deletions(-) diff --git a/docs/models/qwen3/dflash2-phase1-930.md b/docs/models/qwen3/dflash2-phase1-930.md index 48e33cbb4..26e91905a 100644 --- a/docs/models/qwen3/dflash2-phase1-930.md +++ b/docs/models/qwen3/dflash2-phase1-930.md @@ -18,7 +18,7 @@ - `docs/models/qwen3/dflash-speculative-decoding.md` - the existing DFlash lane owns proposal while the shared verify and KV transaction contracts stay method-agnostic. - No prior DFlash2 Phase 1 task record exists in this checkout. - **Plan**: - 1. Audit the current configuration and loader scaffold; keep legacy DFlash and DSpark behavior unchanged and reject native hybrid DFlash2 capabilities that Phase 1 cannot execute. + 1. Audit the current configuration and loader scaffold; keep legacy DFlash and DSpark behavior unchanged, load an independent native output head when the checkpoint declares untied embeddings, and reject hybrid capabilities that Phase 1 cannot execute. 2. Load and validate the selector projection/codebooks, add a fixed-size GPU selector primitive and Rust wrapper, and account for its persistent and scratch allocations. 3. Dispatch `TopKSelector` from the DFlash draft lane without changing draft span, verify, KV transaction, or CUDA-Graph shapes. 4. Run formatting, compile, focused selector/reference checks, GPU-vs-reference checks, and legacy DFlash/DSpark regression checks; record actual results and limitations. @@ -33,9 +33,12 @@ - Added a `DFlashProposal::TopKSelector` capability and an explicit `DFlashLayout` in `pegainfer-qwen3/src/config.rs`. - Legacy DFlash and DSpark schemas remain on their existing proposal paths. -- Native DFlash2 configurations are parsed, but Phase 2 convolution, - sliding-window attention, anchor-first selector layouts, and an untied draft - output head fail closed before GPU weight allocation. +- Native DFlash2 configurations parse their root or nested + `tie_word_embeddings` field into a head-source contract. Legacy and tied + checkpoints reuse the verifier embedding/output projection; untied native + checkpoints load only their separate `lm_head.weight`. +- Phase 2 convolution, sliding-window attention, and anchor-first selector + layouts still fail closed before GPU weight allocation. ### Step 2: Load selector weights and wire the proposer @@ -46,6 +49,10 @@ walk using the predecessor/successor codebooks. - Kept the existing full-block draft result contract, verify span, KV updates, and CUDA-Graph shapes unchanged. +- The native verifier embedding is intentionally reused instead of loading a + duplicate `embed_tokens.weight`; the downloaded Qwen3-4B DFlash2 checkpoint's + embedding bytes match the verifier exactly, while its `lm_head.weight` is + loaded when the schema is untied. ### Step 3: Fix anchor-drop row mapping @@ -74,6 +81,26 @@ with `/usr/bin` present in `PATH` (the build script invokes `git`): | `cargo test --release -p pegainfer-qwen3 --lib` | 88 passed, 0 failed | | `cargo test --release -p pegainfer-build --lib` | 8 passed, 0 failed | +The native selector gate was also run with the real Qwen3-4B DFlash2 tensors +and a selector-only config overlay. The overlay removes only the checkpoint's +Phase 2 convolution/sliding-window declarations; it does not replace selector, +backbone, embedding, or output-head weights: + +```bash +export PATH=/home/ricardo.zheng/.cargo/bin:/usr/local/cuda/bin:/usr/bin:/bin:$PATH +export CUDA_HOME=/usr/local/cuda +export PEGAINFER_CUDA_SM=89 +PEGAINFER_TEST_MODEL_PATH=/database/ricardo.zheng/models/Qwen3/Qwen3-4B \ +PEGAINFER_DFLASH2_TEST_MODEL_PATH=/tmp/dflash2-phase1-native-overlay \ +cargo test --release -p pegainfer-qwen3 \ + --test dflash_speculative_gate \ + dflash2_native_selector_untied_head_greedy_gate \ + -- --ignored --nocapture --test-threads=1 +``` + +Result: `1 passed, 0 failed`; the selector-only native launch dispatched the +CUDA top-k/path-walk path and matched the plain Qwen3 greedy continuation. + ### Step 5: Remove redundant scaffolding - Kept the selector tensor preflight because the shared loader does not check @@ -93,12 +120,22 @@ with `/usr/bin` present in `PATH` (the build script invokes `git`): | `cargo test --release -p pegainfer-build --lib` | 8 passed, 0 failed | | `cargo build --release -p pegainfer-server --bin pegainfer` | Passed | +### Step 6: Address review feedback + +- Removed the native-schema-wide head rejection. The loader now distinguishes + verifier-owned tied heads from an untied native `lm_head.weight`, so a + selector-only DFlash2 checkpoint can reach the selector path. +- Added an ignored GPU gate that imports a native untied checkpoint, runs the + selector CUDA launches with real weights, and checks greedy losslessness. + The public Qwen3-4B DFlash2 artifact is currently Phase 2-capable, so the + gate uses a config-only overlay while preserving every model tensor. + ## Debrief -- **Outcome:** Phase 1 selector wiring, anchor-drop mapping, and a focused - cleanup of redundant scaffolding are complete in the feature branch. The - checkout is fast-forwarded to upstream main; no changes are staged or - committed. +- **Outcome:** Phase 1 selector wiring, anchor-drop mapping, native untied-head + loading, and a focused cleanup of redundant scaffolding are complete in the + feature branch. The checkout is based on upstream main; no changes are + staged or committed. - **Pitfalls encountered:** The first verification command omitted system directories from `PATH`, so `pegainfer-kernels/build.rs` could not spawn `git`. Re-running with `/usr/bin:/bin` succeeded. The row mapping bug was a @@ -107,8 +144,7 @@ with `/usr/bin` present in `PATH` (the build script invokes `git`): from the compact set of positions actually proposed. The anchor is a request-level predecessor, not a selector candidate when the executor drops row 0. -- **Follow-ups:** Run the GPU selector/checkpoint import, Qwen3 greedy - losslessness, legacy DFlash/DSpark regressions, HTTP serving, and performance - A/B on the Linux GPU host. Phase 2 remains responsible for dynamic - convolution and sliding-window execution; Phase 3 remains responsible for - sampled losslessness/rejection sampling. +- **Follow-ups:** Run the full native checkpoint import only after Phase 2 + convolution and sliding-window execution land, because the public checkpoint + intentionally advertises those capabilities and Phase 1 rejects them. Phase 3 + remains responsible for sampled losslessness/rejection sampling. diff --git a/pegainfer-qwen3/src/config.rs b/pegainfer-qwen3/src/config.rs index bf514ac63..bf4f7c34b 100644 --- a/pegainfer-qwen3/src/config.rs +++ b/pegainfer-qwen3/src/config.rs @@ -61,6 +61,14 @@ pub(crate) enum DFlashLayout { AnchorFirst, } +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum DFlashHeadSource { + /// Use the verifier's embedding and output projection. + Target, + /// Native DFlash2 provides a separate draft output projection. + DraftOutput, +} + #[derive(Clone, Debug, PartialEq, Eq)] pub(crate) struct DynamicConv { pub(crate) kernel_size: usize, @@ -97,8 +105,7 @@ pub(crate) struct DFlashConfig { /// Proposal capability selected by the checkpoint schema. pub(crate) proposal: DFlashProposal, pub(crate) layout: DFlashLayout, - /// Whether the target output head can be reused by the drafter. - pub(crate) reuse_target_head: bool, + pub(crate) head_source: DFlashHeadSource, pub(crate) enable_confidence_head: bool, pub(crate) dynamic_convolution: Option, pub(crate) sliding_window: Option, @@ -182,6 +189,8 @@ struct RawDFlash2TransformerConfig { use_sliding_window: bool, #[serde(default)] layer_types: Vec, + #[serde(default)] + tie_word_embeddings: Option, } #[derive(Debug, Deserialize)] @@ -212,6 +221,8 @@ struct RawDFlash2Config { selector_top_k: usize, sample_from_anchor: bool, #[serde(default)] + tie_word_embeddings: Option, + #[serde(default)] target_hidden_size: Option, #[serde(default)] @@ -503,7 +514,7 @@ impl DFlashConfig { } else { DFlashLayout::AnchorDrop }, - reuse_target_head: true, + head_source: DFlashHeadSource::Target, enable_confidence_head: raw.enable_confidence_head, dynamic_convolution: None, sliding_window: None, @@ -520,6 +531,7 @@ impl DFlashConfig { selector_rank, selector_top_k, sample_from_anchor, + tie_word_embeddings, target_hidden_size, conv_kernel_size, conv_group_size, @@ -541,8 +553,13 @@ impl DFlashConfig { sliding_window, use_sliding_window, layer_types, + tie_word_embeddings: transformer_tie_word_embeddings, } = transformer_layer_config; + let tie_word_embeddings = tie_word_embeddings + .or(transformer_tie_word_embeddings) + .unwrap_or(false); + ensure!( speculators_config.algorithm == "dflash2", "DFlash2 speculators_config.algorithm must be \"dflash2\", got {:?}", @@ -678,7 +695,11 @@ impl DFlashConfig { draft_vocab_size, proposal, layout, - reuse_target_head: false, + head_source: if tie_word_embeddings { + DFlashHeadSource::Target + } else { + DFlashHeadSource::DraftOutput + }, enable_confidence_head: false, @@ -780,10 +801,6 @@ impl DFlashConfig { bail!("DFlash Phase 1 supports only anchor-drop selector checkpoints"); } - ensure!( - self.reuse_target_head, - "DFlash checkpoint uses an independent draft output head, which the current runtime does not load" - ); Ok(()) } diff --git a/pegainfer-qwen3/src/dflash.rs b/pegainfer-qwen3/src/dflash.rs index 778cdfadb..4ef8cdc5c 100644 --- a/pegainfer-qwen3/src/dflash.rs +++ b/pegainfer-qwen3/src/dflash.rs @@ -36,6 +36,8 @@ pub(crate) struct DFlashDraftModel { markov: Option, /// DFlash2 selector, mutually exclusive with the DSpark Markov head. selector: Option, + /// Native untied DFlash2 output head; tied checkpoints reuse the target. + draft_lm_head: Option, } pub(crate) struct DFlashRequestState { @@ -773,7 +775,7 @@ impl DFlashDraftModel { for (i, state) in states.iter_mut().enumerate() { state.committed_len += context_lens[i]; } - self.compute_logits_with_target_head_into(target, scratch); + self.compute_logits_into(target, scratch); Ok(&scratch.logits) } @@ -866,11 +868,7 @@ impl DFlashDraftModel { ); } - fn compute_logits_with_target_head_into( - &self, - target: &Qwen3Model, - scratch: &mut DFlashBatchScratch, - ) { + fn compute_logits_into(&self, target: &Qwen3Model, scratch: &mut DFlashBatchScratch) { let ctx = target.device_ctx(); ops::rms_norm_batch_into( ctx, @@ -879,9 +877,13 @@ impl DFlashDraftModel { self.config.rms_norm_eps, &mut scratch.logits_normed, ); + let output_projection = self + .draft_lm_head + .as_ref() + .unwrap_or_else(|| target.output_projection()); ops::gemm_into( ctx, - target.output_projection(), + output_projection, &scratch.logits_normed, &mut scratch.logits, ); diff --git a/pegainfer-qwen3/src/dflash/loading.rs b/pegainfer-qwen3/src/dflash/loading.rs index b1241a5db..42c366529 100644 --- a/pegainfer-qwen3/src/dflash/loading.rs +++ b/pegainfer-qwen3/src/dflash/loading.rs @@ -12,12 +12,15 @@ use pegainfer_core::weight_loader::load_tensor_2d; use pegainfer_core::weight_loader::mmap_shards; use super::DFlashDraftModel; +use super::manifest::DRAFT_LM_HEAD_TENSOR; use super::manifest::HIDDEN_PROJECTION_TENSOR; use super::manifest::PREDECESSOR_CODEBOOK_TENSOR; use super::manifest::SUCCESSOR_CODEBOOK_TENSOR; +use super::manifest::validate_native_output_head; use super::manifest::validate_selector_tensors; use super::selector::SelectorWeights; use crate::config::DFlashConfig; +use crate::config::DFlashHeadSource; use crate::config::DFlashProposal; use crate::dspark::MARKOV_W1_TENSOR; use crate::dspark::MARKOV_W2_TENSOR; @@ -73,6 +76,24 @@ impl DFlashDraftModel { None }; + let draft_lm_head = match config.head_source { + DFlashHeadSource::Target => None, + DFlashHeadSource::DraftOutput => { + validate_native_output_head( + &shards, + &weight_map, + config.draft_vocab_size, + config.hidden_size, + )?; + Some(load_tensor_2d( + ctx, + &shards, + &weight_map, + DRAFT_LM_HEAD_TENSOR, + )?) + } + }; + let mut layers = Vec::with_capacity(config.num_hidden_layers); for layer_idx in 0..config.num_hidden_layers { let prefix = format!("layers.{layer_idx}"); @@ -170,9 +191,8 @@ impl DFlashDraftModel { let hidden_norm = load_tensor_1d(ctx, &shards, &weight_map, "hidden_norm.weight")?; let fc = load_tensor_2d(ctx, &shards, &weight_map, "fc.weight")?; - // DSpark Markov head (Phase 1). The confidence head and the tied - // embed_tokens/lm_head are intentionally skipped: the head is byte-identical - // to the target's, which we reuse for the verify-equivalent logits. + // The confidence head is outside Phase 1; native DFlash2 heads were + // loaded above according to the checkpoint's tie policy. if config.enable_confidence_head { log::info!( "DFlash confidence head present in {model_path} but unused in Phase 1 \ @@ -208,6 +228,7 @@ impl DFlashDraftModel { sin_cache, markov, selector, + draft_lm_head, }) } } diff --git a/pegainfer-qwen3/src/dflash/manifest.rs b/pegainfer-qwen3/src/dflash/manifest.rs index b6e1590b8..23538a5c8 100644 --- a/pegainfer-qwen3/src/dflash/manifest.rs +++ b/pegainfer-qwen3/src/dflash/manifest.rs @@ -1,4 +1,4 @@ -//! Preflight validation for DFlash2 selector tensors. +//! Preflight validation for DFlash2 tensors. use std::collections::HashMap; @@ -12,6 +12,7 @@ use safetensors::tensor::TensorView; pub(crate) const HIDDEN_PROJECTION_TENSOR: &str = "candidate_selector.hidden_projection.weight"; pub(crate) const PREDECESSOR_CODEBOOK_TENSOR: &str = "candidate_selector.predecessor_codebook"; pub(crate) const SUCCESSOR_CODEBOOK_TENSOR: &str = "candidate_selector.successor_codebook"; +pub(crate) const DRAFT_LM_HEAD_TENSOR: &str = "lm_head.weight"; /// Validate selector dtype and shapes before GPU upload. pub(crate) fn validate_selector_tensors( @@ -46,6 +47,23 @@ pub(crate) fn validate_selector_tensors( Ok(()) } +/// Validate the native DFlash2 output head before uploading it. +pub(crate) fn validate_native_output_head( + shards: &[SafeTensors<'_>], + weight_map: &HashMap, + vocab_size: usize, + hidden_size: usize, +) -> Result<()> { + validate_matrix( + shards, + weight_map, + DRAFT_LM_HEAD_TENSOR, + vocab_size, + hidden_size, + )?; + Ok(()) +} + fn validate_matrix( shards: &[SafeTensors<'_>], weight_map: &HashMap, @@ -56,12 +74,12 @@ fn validate_matrix( let tensor = find_tensor(shards, weight_map, name)?; ensure!( tensor.dtype() == Dtype::BF16, - "DFlash selector tensor {name:?} must be BF16, got {:?}", + "DFlash2 tensor {name:?} must be BF16, got {:?}", tensor.dtype() ); ensure!( tensor.shape() == [expected_rows, expected_cols], - "DFlash selector tensor {name:?} has shape {:?}, expected [{expected_rows}, {expected_cols}]", + "DFlash2 tensor {name:?} has shape {:?}, expected [{expected_rows}, {expected_cols}]", tensor.shape() ); Ok(()) @@ -74,11 +92,11 @@ fn find_tensor<'a>( ) -> Result> { if let Some(&shard_idx) = weight_map.get(name) { let shard = shards.get(shard_idx).with_context(|| { - format!("DFlash selector tensor {name:?} references missing shard index {shard_idx}") + format!("DFlash2 tensor {name:?} references missing shard index {shard_idx}") })?; return shard .tensor(name) - .with_context(|| format!("load DFlash selector tensor {name:?}")); + .with_context(|| format!("load DFlash2 tensor {name:?}")); } for shard in shards { @@ -87,5 +105,5 @@ fn find_tensor<'a>( } } - anyhow::bail!("DFlash selector tensor {name:?} is missing") + anyhow::bail!("DFlash2 tensor {name:?} is missing") } diff --git a/pegainfer-qwen3/src/dflash/reservation.rs b/pegainfer-qwen3/src/dflash/reservation.rs index 942dea669..3f4e240b8 100644 --- a/pegainfer-qwen3/src/dflash/reservation.rs +++ b/pegainfer-qwen3/src/dflash/reservation.rs @@ -1,6 +1,7 @@ use anyhow::Result; use crate::config::DFlashConfig; +use crate::config::DFlashHeadSource; use crate::config::DFlashProposal; use crate::dflash::selector::SelectorScratch; use crate::dspark::MarkovHead; @@ -91,10 +92,14 @@ impl DFlashMemoryReservation { } _ => 0, }; + let native_head = match config.head_source { + DFlashHeadSource::DraftOutput => BF16 * config.draft_vocab_size * hidden, + DFlashHeadSource::Target => 0, + }; Self { kv_bytes_per_token, - fixed_bytes: weights + scratch_total + block_headroom + markov + selector, + fixed_bytes: weights + scratch_total + block_headroom + markov + selector + native_head, } } } diff --git a/pegainfer-qwen3/tests/dflash_speculative_gate.rs b/pegainfer-qwen3/tests/dflash_speculative_gate.rs index 8803bcd68..7db870d4a 100644 --- a/pegainfer-qwen3/tests/dflash_speculative_gate.rs +++ b/pegainfer-qwen3/tests/dflash_speculative_gate.rs @@ -35,9 +35,10 @@ //! Runs the two engines sequentially (baseline dropped before the speculative //! engine loads) so only one Qwen3-4B is resident at a time. //! -//! Requires a CUDA GPU, Qwen3-4B weights, and the DFlash drafter. Set -//! `PEGAINFER_TEST_MODEL_PATH` (target) and `PEGAINFER_DFLASH_TEST_MODEL_PATH` -//! (drafter); skips cleanly when either is absent. +//! Requires a CUDA GPU, Qwen3-4B weights, and a drafter. Legacy gates use +//! `PEGAINFER_DFLASH_TEST_MODEL_PATH` and skip when weights are absent. The +//! ignored native-selector gate requires `PEGAINFER_DFLASH2_TEST_MODEL_PATH` +//! and fails when its explicitly requested checkpoint is not usable. use std::path::Path; use std::path::PathBuf; @@ -107,6 +108,52 @@ fn draft_path_or_skip() -> Option { } } +fn required_model_path(name: &str) -> String { + let path = std::env::var(name).unwrap_or_else(|_| panic!("set {name} to run this test")); + assert!( + Path::new(&path).join("config.json").is_file(), + "{name}={path:?} has no config.json" + ); + path +} + +fn native_dflash2_path() -> String { + let path = required_model_path("PEGAINFER_DFLASH2_TEST_MODEL_PATH"); + let config: serde_json::Value = serde_json::from_str( + &std::fs::read_to_string(Path::new(&path).join("config.json")) + .expect("read native DFlash2 config"), + ) + .expect("parse native DFlash2 config"); + assert_eq!( + config + .get("speculators_model_type") + .and_then(serde_json::Value::as_str), + Some("dflash2"), + "native selector gate requires speculators_model_type=dflash2" + ); + assert!( + config + .get("architectures") + .and_then(serde_json::Value::as_array) + .is_some_and(|values| values.iter().any(|value| value == "DFlash2DraftModel")), + "native selector gate requires DFlash2DraftModel" + ); + let tied = config + .get("tie_word_embeddings") + .and_then(serde_json::Value::as_bool) + .or_else(|| { + config + .pointer("/transformer_layer_config/tie_word_embeddings") + .and_then(serde_json::Value::as_bool) + }) + .unwrap_or(false); + assert!( + !tied, + "this gate must exercise native DFlash2's independent draft lm_head" + ); + path +} + fn launch_options(draft: Option) -> Qwen3LaunchOptions { Qwen3LaunchOptions { device_ordinal: 0, @@ -420,6 +467,67 @@ fn dflash_speculative_greedy_matches_plain_greedy() { ); } +/// Production gate for the native DFlash2 selector and its independent draft +/// output head. The native schema guarantees that the speculative launch +/// dispatches through the CUDA top-k and path-walk kernels; target verification +/// must still produce a lossless greedy continuation. +#[test] +#[ignore = "requires CUDA, Qwen3-4B, and a selector-only native DFlash2 checkpoint"] +fn dflash2_native_selector_untied_head_greedy_gate() { + const NATIVE_GENERATED_TOKENS: usize = 32; + + let model_path = required_model_path("PEGAINFER_TEST_MODEL_PATH"); + let draft_path = native_dflash2_path(); + let _gpu = GPU + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + + let prompt = "def fibonacci(n):"; + let tokenizer = common::load_tokenizer(&model_path); + let prompt_tokens = tokenizer.encode(prompt, false).expect("encode failed"); + + let baseline = { + let engine = EngineHarness::new( + pegainfer_qwen3::launch(Path::new(&model_path), launch_options(None)) + .expect("failed to start baseline engine"), + ); + let output = generate( + &engine, + prompt_tokens.clone(), + LOGPROBS, + NATIVE_GENERATED_TOKENS, + ); + drop(engine); + std::thread::sleep(Duration::from_secs(2)); + output + }; + + let engine = EngineHarness::new( + pegainfer_qwen3::launch( + Path::new(&model_path), + launch_options(Some(PathBuf::from(&draft_path))), + ) + .expect("failed to load native DFlash2 selector checkpoint"), + ); + let speculative = generate(&engine, prompt_tokens.clone(), 0, NATIVE_GENERATED_TOKENS); + let result = check_lossless( + &engine, + &tokenizer, + 0, + prompt, + &prompt_tokens, + &baseline, + &speculative, + ); + drop(engine); + + assert!( + result.is_ok(), + "native DFlash2 selector is not greedy-lossless:\n{}", + result.unwrap_err() + ); +} + /// Verify-graph capture-shape regression (heterogeneous `max_tokens`). /// /// The piecewise verify CUDA Graph keys its captured dense segments by From 65e3a8c286b6605b7060adc4f90aba15936a8d45 Mon Sep 17 00:00:00 2001 From: RicardoMin <17879681016@163.com> Date: Sat, 5 Sep 2026 13:48:06 +0800 Subject: [PATCH 3/3] fix(qwen3): satisfy CUDA Clippy gate Signed-off-by: RicardoMin <17879681016@163.com> --- docs/models/qwen3/dflash2-phase1-930.md | 24 ++++++++++++++++--- .../tests/dflash_speculative_gate.rs | 2 +- 2 files changed, 22 insertions(+), 4 deletions(-) diff --git a/docs/models/qwen3/dflash2-phase1-930.md b/docs/models/qwen3/dflash2-phase1-930.md index 0281c8dce..043474b13 100644 --- a/docs/models/qwen3/dflash2-phase1-930.md +++ b/docs/models/qwen3/dflash2-phase1-930.md @@ -150,12 +150,30 @@ CUDA top-k/path-walk path and matched the plain Qwen3 greedy continuation. | Linux `cargo test --release -p pegainfer-build --lib` | 7 passed, 0 failed | | Linux Qwen3 unit-test link | Blocked by existing `cudaLaunchKernelExC` linker mismatch | +### Step 8: Fix the Qwen3 CUDA Clippy gate + +- Reproduced the failing CI command on Linux at commit `7d04ca5a`. +- Clippy reported one `redundant_clone` in + `pegainfer-qwen3/tests/dflash_speculative_gate.rs:586`; the native gate + consumed `dflash2_view.path`, so cloning it was unnecessary. +- Removed only that clone. The same CI package set now passes with + `--all-targets -- -D warnings`. +- Local Windows Clippy could not reach project diagnostics because its + environment lacks OpenSSL and builds `esaxx-rs` with exceptions disabled; + this is an environment limitation, not a source failure. + +| CI follow-up verification | Result | +| --- | --- | +| Linux Qwen3 CUDA Clippy package set, `sm_80` | Passed | +| `cargo fmt --all -- --check` | Passed | +| `git diff --check` | Passed | + ## Debrief - **Outcome:** Phase 1 selector wiring, anchor-drop mapping, native tied/untied - head loading, an executable selector gate, and the merge with current - upstream main are complete in the feature branch. The merge result is staged - but not committed or pushed. + head loading, an executable selector gate, the merge with current upstream + main, and the Qwen3 CUDA Clippy cleanup are complete in the feature branch. + The source fix is unstaged and not committed or pushed. - **Pitfalls encountered:** The first verification command omitted system directories from `PATH`, so `pegainfer-kernels/build.rs` could not spawn `git`. Re-running with `/usr/bin:/bin` succeeded. The row mapping bug was a diff --git a/pegainfer-qwen3/tests/dflash_speculative_gate.rs b/pegainfer-qwen3/tests/dflash_speculative_gate.rs index d1c214cb2..a54c2084e 100644 --- a/pegainfer-qwen3/tests/dflash_speculative_gate.rs +++ b/pegainfer-qwen3/tests/dflash_speculative_gate.rs @@ -583,7 +583,7 @@ fn dflash2_native_selector_greedy_gate() { let engine = EngineHarness::new( pegainfer_qwen3::launch( Path::new(&model_path), - launch_options(Some(dflash2_view.path.clone())), + launch_options(Some(dflash2_view.path)), ) .expect("failed to load native DFlash2 selector checkpoint"), );