diff --git a/Cargo.lock b/Cargo.lock index ac8cca04..3c64ad5e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3872,6 +3872,7 @@ dependencies = [ "hex", "log", "memmap2", + "nvtx", "openinfer-bench", "openinfer-core", "openinfer-cupti", diff --git a/docs/index.md b/docs/index.md index e4896f29..fa15253f 100644 --- a/docs/index.md +++ b/docs/index.md @@ -30,6 +30,8 @@ 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/accuracy-gate.md` | Qwen3-4B instance of the logits golden gate (`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 `openinfer-kernels`; release build, test-target compile, accuracy gate, and bench snapshot pass. | +| `models/qwen3/dflash-model-download.md` | Download record for `z-lab/Qwen3-4B-DFlash-b16`: local `/data/models/Qwen3-4B-DFlash-b16` and 5090 `/data/Qwen3-4B-DFlash-b16`, the drafter artifact for Qwen3-4B speculative decoding bring-up. | +| `models/qwen3/dflash-speculative-decoding.md` | Qwen3-4B DFlash is wired for greedy TP1 serving with native drafter/verifier/scheduler, INFO acceptance logs with `committed_tokens`, verifier-span/full-draft HF regret gate, config-derived DFlash reserve, side-state byte-budget admission, transactional speculative KV rollback, per-request prefill capture, chunked-prefill continuity checks, and DFlash small-N cublasLt tuning. Draft PR #380. Latest local 5070 Ti PR-head results: Spec-Bench bs=1 149.32 tok/s (1.66x), Spec-Bench c4 330.42 tok/s (1.09x), random 1024/128 c4 349.50 tok/s (1.29x); post-hardening Spec-Bench c4 n=12 smoke completed 12/12 at 368.71 tok/s with four concurrent DFlash requests logged in one wave; 5090 OpenInfer Spec-Bench bs=1 is 251.48 tok/s and vLLM 0.22.1 DFlash reaches 289.57 tok/s with native acceptance metrics. | | `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. | | `models/qwen3/kv-pressure-hang.md` | Issue #85 Qwen3-4B KV pressure hang fixed by full-lifetime scheduler KV admission, waiting-queue deferral, cleanup on disconnect/error, impossible-request errors, scheduler/bridge gates, and real `vllm bench serve` QPS=2 `500/500` pass with post-pressure completion healthy. | diff --git a/docs/models/qwen3/dflash-model-download.md b/docs/models/qwen3/dflash-model-download.md new file mode 100644 index 00000000..618700a6 --- /dev/null +++ b/docs/models/qwen3/dflash-model-download.md @@ -0,0 +1,72 @@ +# DFlash Model Download + +> **TL;DR:** `z-lab/Qwen3-4B-DFlash-b16` is downloaded and verified locally at `/data/models/Qwen3-4B-DFlash-b16` and on the 5090 box at `/data/Qwen3-4B-DFlash-b16`; it is the drafter artifact for Qwen3-4B speculative decoding bring-up. +> +> **Last touched:** 2026-06 + +## Preparation + +- **Read**: + - `docs/index.md` - Qwen3 model-line docs live under `docs/models/qwen3/`. + - `docs/models/qwen3/model-crate.md` - Qwen3 runtime is model-crate owned; local model artifacts under `/data/models` are used by executor/tests through model paths. +- **Relevant history**: + - No existing DFlash or speculative-decoding task doc found. +- **Plan**: + 1. Check `/data/models` capacity and whether the target directory already exists. + 2. Download `z-lab/Qwen3-4B-DFlash-b16` with the Hugging Face Hub CLI into `/data/models/Qwen3-4B-DFlash-b16`. + 3. List the downloaded files and verify the safetensors/config metadata is present. +- **Risks / open questions**: + - The Hugging Face repo may require auth or may use custom code files that are not part of the plain safetensors load path yet. + +## Execution Log + +### Step 1: Check destination and capacity +- `/data/models` already contains the base `/data/models/Qwen3-4B` artifact. +- `df -h /data/models /data` reported `753G` available on `/data`, enough for the DFlash drafter. +- Target path chosen: `/data/models/Qwen3-4B-DFlash-b16`. + +### Step 2: Download from Hugging Face +- Command: + ```bash + uvx --from huggingface_hub hf download z-lab/Qwen3-4B-DFlash-b16 --local-dir /data/models/Qwen3-4B-DFlash-b16 + ``` +- Result: fetched `9` files, `1.08GB` total, into `/data/models/Qwen3-4B-DFlash-b16`. + +### Step 3: Verify local files +- `find /data/models/Qwen3-4B-DFlash-b16 -maxdepth 2 -type f` shows: + - `config.json` + - `model.safetensors` + - `modeling_dflash.py` + - `dflash.py` + - `utils.py` + - `README.md` + - `.gitattributes` + - `assets/dflash_system.png` + - `assets/speedup.png` +- `du -sh /data/models/Qwen3-4B-DFlash-b16` reports `1.1G`. +- `jq` parsed `config.json`; `architectures = ["DFlashDraftModel"]`, `hidden_size = 2560`, `num_hidden_layers = 5`, `vocab_size = 151936`. +- `safetensors` opened `model.safetensors` successfully and reported `58` tensors. + +### Step 4: Place artifact on 5090 +- User explicitly approved placing a copy under `/data` on the 5090 box. +- Remote path: `/data/Qwen3-4B-DFlash-b16`. +- Download command used the 5090 proxy from root's `.bashrc`: + ```bash + export http_proxy=http://172.17.0.1:1081 + export https_proxy=http://172.17.0.1:1081 + hf download z-lab/Qwen3-4B-DFlash-b16 --local-dir /data/Qwen3-4B-DFlash-b16 --max-workers 8 + ``` +- The 5090 copy contains the same core files as the local copy; `model.safetensors` is `1074860568` bytes. +- Real-weight validation now passes on 5090 with: + ```bash + OPENINFER_TEST_MODEL_PATH=/data/Qwen3-4B OPENINFER_DFLASH_TEST_MODEL_PATH=/data/Qwen3-4B-DFlash-b16 cargo test --release -p openinfer-qwen3-4b dflash::tests::downloaded_dflash_config_matches_qwen3_4b --lib -- --nocapture + ``` + +## Debrief + +- **Outcome**: The DFlash drafter model is present at `/data/models/Qwen3-4B-DFlash-b16` locally and `/data/Qwen3-4B-DFlash-b16` on 5090. +- **Pitfalls encountered**: + - None during download. The repo includes Python custom-code files, so runtime integration still needs a native Rust loader/forward path rather than relying on `trust_remote_code`. +- **Lessons learned**: + - The artifact is small enough (`1.1G`) to keep alongside the base Qwen3-4B model. + - `config.json` does not set `torch_dtype`; integration should infer/check tensor dtype from safetensors rather than trusting that config field. diff --git a/docs/models/qwen3/dflash-speculative-decoding.md b/docs/models/qwen3/dflash-speculative-decoding.md new file mode 100644 index 00000000..e562c842 --- /dev/null +++ b/docs/models/qwen3/dflash-speculative-decoding.md @@ -0,0 +1,649 @@ +# DFlash Speculative Decoding + +> **TL;DR:** Qwen3-4B DFlash is wired end-to-end for greedy TP1 serving with native Rust/CUDA drafter, target verifier, INFO acceptance logs with `committed_tokens`, profiling-only timing, verifier-span/full-draft HF regret gate, request-local draft scratch, reusable pending-context buffer, config-derived DFlash memory reserve, DFlash side-state byte-budget admission, transactional speculative KV rollback, per-request DFlash prefill capture, chunked-prefill DFlash context continuity checks, and DFlash small-N cublasLt tuning. Speculative protocol types live in `openinfer-qwen3-4b/src/speculative.rs`, and DFlash lane state lives in `openinfer-qwen3-4b/src/executor/dflash_lane.rs`. Draft PR: #380. Latest local 5070 Ti PR-head `vllm bench serve` results: Spec-Bench bs=1 149.32 tok/s (1.66x), Spec-Bench c4 330.42 tok/s (1.09x), random 1024/128 bs=1 136.09 tok/s (1.57x), random c4 349.50 tok/s (1.29x). Post-hardening smoke: Spec-Bench c4 n=12 completed 12/12 at 368.71 tok/s and logged four concurrent DFlash requests in one wave. 5090 OpenInfer bs=1 reaches 251.48 tok/s on Spec-Bench (1.50x); upstream vLLM 0.22.1 supports Qwen3 DFlash and reaches 289.57 tok/s on the same Spec-Bench (1.78x vs vLLM baseline), with native `vllm bench serve` acceptance metrics. The current OpenInfer DFlash path is not CUDA-graph captured: local nsys shows draft eager at ~2.7-2.9 ms/step, ~98-106 kernel launches/step, dominated by GEMM/GEMV. Multi-active DFlash is enabled for all eligible greedy requests, but the draft side is still per-request serial; target verification is batched. +> +> **Last touched:** 2026-06 + +## Preparation + +- **Read**: + - `docs/index.md` - Qwen3 model-line docs and scheduler/runtime docs are the relevant routing points. + - `docs/models/qwen3/dflash-model-download.md` - confirms `z-lab/Qwen3-4B-DFlash-b16` is local at `/data/models/Qwen3-4B-DFlash-b16`. + - `docs/models/qwen3/roadmap.md` - Qwen3 is the mature line; new decode behavior must sit under correctness/perf gates. + - `docs/subsystems/scheduler/scheduler.md` - current runtime is single GPU-owner thread, FCFS prefill priority, paged KV, bucket CUDA Graph decode, unified prefill+decode. + - `/data/models/Qwen3-4B-DFlash-b16/config.json` - DFlash drafter shape: `block_size=16`, 5 layers, hidden 2560, 32 Q heads, 8 KV heads, selected target layers `[1, 9, 17, 25, 33]`, mask token `151669`. + - `/data/models/Qwen3-4B-DFlash-b16/dflash.py` - generation contract: prefill target, use target hidden states, draft a block, verify with target, accept matching prefix plus one posterior token, crop target/drafter KV. + - `openinfer-qwen3-4b/src/scheduler/plan.rs` - current plan space is `Prefill`, `Decode`, `Unified`; pure decode calls `execute_decode`. + - `openinfer-qwen3-4b/src/unified_forward.rs` - decode rows already use prefill-style attention metadata as `qo_len=1`, which is the shape to generalize for verifier spans. + - `kvbm/kvbm-logical/src/integrations/scheduled.rs` - lower layer already supports `schedule_speculative` and `apply_speculative`, including partial accept and LIFO release of excess capacity. +- **Relevant history**: + - `docs/subsystems/scheduler/scheduler.md` records the FlashInfer paged metadata invariants; speculative views must continue to expose exact page rows, never raw assigned blocks with surplus generation capacity. + - `docs/models/qwen3/accuracy-gate.md` is the existing truth pattern for Qwen3 logits; DFlash should add a verifier-span logits/regret gate rather than rely on exact spec-on/spec-off greedy text identity. +- **Plan**: + 1. Add a shared speculative KV lifecycle surface in `openinfer-kv-cache`, backed by kvbm's existing `schedule_speculative` / `apply_speculative`. + 2. Add a model-owned speculative verifier API in Qwen3 executor that forwards verify spans and returns accepted tokens plus logits/hidden features needed by drafters. + 3. Port the DFlash drafter loader/forward path natively from the downloaded safetensors/custom Python, reusing target embeddings/lm head. + 4. Wire scheduler policy conservatively: enable DFlash first for pure greedy decode requests with no LoRA and no pending prefill, then broaden only after gates prove correctness and throughput. + 5. Add gates: EOS/max_tokens/context-window behavior, prefix-cache interaction, acceptance histogram, target/drafter/verifier timing, and serving throughput vs baseline. +- **Risks / open questions**: + - DFlash's Python implementation depends on target hidden states for selected layers; Qwen3 forward currently only returns logits, so hidden-state taps must be added without polluting normal decode hot paths. + - DFlash drafter has its own KV cache. The first correct native path may run eager before CUDA Graph capture; graphing belongs after verifier-span correctness and acceptance are proven. + - The README speedups are for specific stacks/workloads. Local expected speedup must be measured against current OpenInfer Qwen3-4B baseline, not copied from the paper. + - Exact spec-on/spec-off greedy text parity is not a stable gate for this implementation because the verifier uses multi-token target prefill while baseline decode uses the single-token decode path. Use target logits/regret gates plus scheduler/KV lifecycle gates for correctness; exact text mismatches at low-margin choices need separate investigation before becoming a release blocker. + +## Design Notes + +- **Verify span**: a speculative verifier step forwards `N` token positions for a request. In DFlash, `N=block_size`: the current dangling token plus `block_size - 1` draft candidates. The target posterior at the first mismatch supplies one bonus token. +- **Accepted tokens**: the sequence passed to `apply_speculative` is the newly produced output tokens for that step: matched draft tokens followed by the posterior bonus token. Its length equals the number of target KV positions that are now valid, leaving the last token dangling for the next step. +- **Abstraction boundary**: KV cache exposes generic speculative schedule/view/apply. Qwen3 executor owns verifier forward and sampling/acceptance. DFlash owns draft generation. Scheduler owns policy. +- **Code boundary**: + - `openinfer-qwen3-4b/src/speculative.rs` owns speculative draft/verify protocol types and acceptance-prefix construction. + - `openinfer-qwen3-4b/src/executor/dflash_lane.rs` owns DFlash per-lane request state, target-hidden context append, draft execution, INFO acceptance logs, and profiling timing. + - `openinfer-qwen3-4b/src/executor/dflash_prefill.rs` owns DFlash prefill eligibility and chunk-finalization policy. + - `openinfer-qwen3-4b/src/executor/speculative_exec.rs` owns Qwen3 executor-side speculative draft/verify scheduling and KV apply. + - `openinfer-qwen3-4b/src/executor/{lifecycle,model_executor,worker}.rs` split executor construction/offload lifecycle, runtime trait implementation, and worker-lane plumbing; every `src/executor/*.rs` file is now under 1k lines. + - `openinfer-qwen3-4b/src/dflash.rs` owns the DFlash model weights, request-local draft scratch, reusable pending-context buffer, DFlash-specific small-N cublasLt tuning, and GPU forward kernels. +- **Initial policy**: DFlash is opt-in and greedy-only. Unsupported combinations fail closed or route clearly to baseline, not silently change semantics. +- **Concurrency boundary**: when DFlash is configured, every eligible active greedy request may enter speculative decode. Pending prefill chunks still run before the next draft step so new requests can capture DFlash hidden context and join the speculative set. DFlash side-state admission is byte-budgeted from the config-derived reserve, so short concurrent requests can all enter while long requests queue instead of OOMing. The current multi-active implementation drafts serially per request, then batches the target verifier; this is correct and measurable, but not the final throughput shape. + +## Execution Log + +### Step 1: Download and inspect DFlash artifact +- Artifact: `/data/models/Qwen3-4B-DFlash-b16`. +- `config.json` confirms Qwen3-compatible geometry and DFlash-specific fields: + - `block_size = 16` + - `dflash_config.mask_token_id = 151669` + - `dflash_config.target_layer_ids = [1, 9, 17, 25, 33]` + - `architectures = ["DFlashDraftModel"]` +- `model.safetensors` has 58 BF16 tensors: 5 Qwen3-like layers plus `fc.weight`, `hidden_norm.weight`, and `norm.weight`. It has no lm-head; DFlash uses the target lm-head. +- Python `spec_generate` verifies the contract: target prefill → DFlash block draft → target block verify → prefix match → posterior bonus → crop caches. + +### Step 2: Expose speculative KV lifecycle +- Added `RequestKv::schedule_speculative`, `RequestKv::speculative_view`, and `RequestKv::apply_speculative` in `openinfer-kv-cache`. +- `speculative_view(num_draft_tokens)` uses the same exact page-row contract as `prefill_view` / `decode_view`: it covers `kv_position + num_draft_tokens` and does not expose extra eagerly-held generation blocks. +- Added lifecycle tests: + - `speculative_view_covers_verify_span` + - `speculative_partial_accept_releases_excess_capacity` +- Verification: + ```bash + cargo fmt --all --check + cargo test --release -p openinfer-kv-cache --test lifecycle speculative -- --nocapture + ``` + Both commands passed. + +### Step 3: Add Qwen3 target verifier API +- Added low-level runtime types: + - `SpeculativeVerifyStepItem` + - `SpeculativeVerifyPlan` + - `SpeculativeVerifyRequestResult` + - `SpeculativeVerifyResult` +- Added `Qwen3Executor::execute_speculative_verify`. +- Implementation: + - schedules each request with `RequestKv::schedule_speculative(verify_span_len)` + - builds exact `speculative_view` rows + - runs a prefill-style target forward with `echo=true` to get all-position logits + - greedily samples target posterior tokens + - computes the accepted draft prefix plus posterior bonus + - commits via `RequestKv::apply_speculative(&accepted_tokens)` +- Added unit tests for acceptance semantics: + - `speculative_verify_accepts_matching_prefix_plus_posterior_bonus` + - `speculative_verify_all_match_still_adds_block_end_posterior` +- Verification: + ```bash + cargo fmt --all --check + cargo test --release -p openinfer-qwen3-4b executor::tests::speculative_verify --lib -- --nocapture + cargo check --release -p openinfer-qwen3-4b --tests + ``` + All commands passed. + +### Step 4: Native DFlash drafter and scheduler wiring +- Added `openinfer-qwen3-4b/src/dflash.rs`: + - loads `fc`, `hidden_norm`, `norm`, and the 5 DFlash transformer layers from `model.safetensors` + - keeps a per-request DFlash KV cache and pending target hidden-state context + - reuses the target embedding/lm-head path for block logits +- Added selected hidden-state capture in Qwen3 prefill/verify for DFlash target layer IDs. +- Added scheduler/executor speculative decode stages: + - DFlash-ready request tracking after prefill hidden capture + - `SpeculativeDraft` to build the draft block + - `SpeculativeVerify` to run the target verifier and apply accepted tokens + - `DecodeEffect::{EmitManyAndContinue, EmitManyAndFinish}` for multi-token emission +- Server CLI: + - `--dflash-draft-model-path ` + - rejects tensor parallel use with DFlash + - disables prefix cache while DFlash is enabled, because hidden-state capture currently needs the full target suffix. +- Memory reservation: DFlash weights plus a config-derived request-state/scratch reserve are held out before Qwen3 KV budgeting. The reserve-polish local startup reserved 4034.1 MB total for DFlash and left 1028 Qwen3 KV blocks on the 16 GiB 5070 Ti. + +### Step 5: Edge-case fixes and local gates +- Fixed speculative verification near `max_tokens`: draft token IDs are clamped to the remaining output budget before scheduling the target verify span. The crash signature before the fix was `KvView pages must exactly cover seq_len=49`. +- Tightened DFlash startup/config invariants: + - DFlash draft/target config checks now reject mismatched hidden/vocab/head geometry, non-increasing target layer IDs, too-small block sizes, and out-of-vocab mask token IDs. + - The downloaded-model config test is gated by `OPENINFER_TEST_MODEL_PATH` and `OPENINFER_DFLASH_TEST_MODEL_PATH`, and skips cleanly on machines without local weights. +- Tightened serving behavior: + - DFlash prefill context capture requires exactly one supported request in the prefill step. + - DFlash hidden-state append now checks that captured rows exactly match the scheduled token span. + - Speculative verify context recording now checks request/result count, request IDs, and missing DFlash state before returning a successful worker result. + - Multi-request prefill can capture DFlash hidden context per request. Unsupported rows are dropped from the DFlash side path without poisoning supported greedy rows in the same prefill batch. + - Pending prefill chunks take priority over speculative decode, so a long speculative active request does not block new request prefill. + - DFlash request drop removes side-channel metric/timing maps as well as drafter request state. +- Fixed chunked-prefill DFlash context handling: + - `dflash_prefill_supported` now rejects prefix-cache hits (`cached_tokens > 0`) even though the server disables prefix cache when DFlash is enabled. This keeps the path fail-closed if a future config accidentally re-enables cache hits without full hidden-state replay. + - Non-final prefill chunks keep the DFlash request state pending instead of dropping it. The final chunk marks the request DFlash-ready. + - DFlash pending context append validates `pending_context_len == chunk_start` before accepting a captured hidden-state chunk. A discontinuous chunk sequence now fails early instead of silently drafting from an incomplete target context. + - The DFlash HF golden gate now also forces a chunked-prefill path and checks 144 verifier positions. +- Added/ran gates: + ```bash + cargo fmt --all --check + cargo check --release -p openinfer-server + cargo test --release -p openinfer-qwen3-4b executor::tests::speculative_verify --lib -- --nocapture + cargo test --release -p openinfer-qwen3-4b scheduler::plan::tests::speculative_verify_items_clamp_to_remaining_output_budget --lib -- --nocapture + cargo test --release -p openinfer-qwen3-4b dflash_prefill_capture_requires_single_supported_request --lib -- --nocapture + cargo test --release -p openinfer-qwen3-4b speculative_plan_runs_only_after_pending_prefill_is_drained --lib -- --nocapture + cargo test --release -p openinfer-kv-cache --test lifecycle speculative -- --nocapture + OPENINFER_TEST_MODEL_PATH=/data/models/Qwen3-4B OPENINFER_DFLASH_TEST_MODEL_PATH=/data/models/Qwen3-4B-DFlash-b16 cargo test --release -p openinfer-qwen3-4b dflash::tests::downloaded_dflash_config_matches_qwen3_4b --lib -- --nocapture + OPENINFER_TEST_MODEL_PATH=/data/models/Qwen3-4B OPENINFER_DFLASH_TEST_MODEL_PATH=/data/models/Qwen3-4B-DFlash-b16 cargo test --release -p openinfer-qwen3-4b --test hf_golden_gate dflash_speculative_verify_matches_hf_argmax_regret_gate -- --nocapture + cargo test --release -p openinfer-qwen3-4b --lib -- --nocapture + cargo build --release -p openinfer-server + ``` + All commands passed locally. +- Latest local DFlash HF regret gate after the chunked-prefill fix: + ```bash + OPENINFER_TEST_MODEL_PATH=/data/models/Qwen3-4B OPENINFER_DFLASH_TEST_MODEL_PATH=/data/models/Qwen3-4B-DFlash-b16 cargo test --release -p openinfer-qwen3-4b --test hf_golden_gate dflash_speculative_verify_matches_hf_argmax_regret_gate -- --nocapture + ``` + It passed and checked 144 target positions. +- Added a core sampling equivalence gate for the all-greedy contiguous argmax fast path: + ```bash + cargo test --release -p openinfer-core all_greedy_contiguous_argmax_matches_indexed_greedy_rows --lib -- --nocapture + ``` + It builds 8193-vocab logits, places maxima across multiple 4096-token tiles plus a cross-tile tie, runs the all-greedy contiguous path, then forces the mixed indexed path and checks greedy-row equality. + +### Step 5.5: Speculative abstraction cleanup +- Moved speculative draft/verify protocol types and acceptance construction out of `executor.rs` into `openinfer-qwen3-4b/src/speculative.rs`. +- Moved DFlash lane state and per-request lifecycle into `openinfer-qwen3-4b/src/executor/dflash_lane.rs`. +- `SpeculativeDraftStepItem` has a public constructor for external runtime/integration callers, while `execute_speculative_draft` and `execute_speculative_verify` fail early when DFlash is not loaded or the request has not completed DFlash prefill context capture. +- `execute_speculative_verify` validates every request before any speculative KV scheduling, so an unsupported mixed request cannot dirty KV state before failing. +- Split the executor surface after the DFlash integration: + - `executor.rs` now keeps shared request/result types, worker step dispatch, low-level sampling/logprob helpers, and tuning helpers. + - `executor/lifecycle.rs` holds constructors, public wrappers, KV-offload save/load helpers, and `run_step`. + - `executor/model_executor.rs` holds the `ModelExecutor` impl and LoRA load/unload policy. + - `executor/worker.rs` holds `LocalQwen3Lane`, `StepCommand`, `WorkerStepOutcome`, and `RankWorker`. + - `executor/dflash_prefill.rs` and `executor/speculative_exec.rs` hold the DFlash-specific policy added by this work. + - File sizes after the split: `executor.rs` 830 lines, `model_executor.rs` 668, `lifecycle.rs` 584, `worker.rs` 469, `dflash_lane.rs` 355. +- Added a DFlash verifier-span HF regret gate in `openinfer-qwen3-4b/tests/hf_golden_gate.rs`: + - Builds a real DFlash-enabled executor. + - Runs prefill hidden capture for fixed HF-golden prompts. + - Runs `execute_speculative_verify` with 16-token teacher-forced spans. + - Checks 144 target verifier positions against HF top-K using the same regret tolerance as the main Qwen3 logits gate, including one forced chunked-prefill request. + - Also runs 8 native DFlash draft → target verify steps, checks the first verifier posterior for each against HF top-K, validates `accepted_tokens` against the produced draft span, and requires at least one accepted draft candidate. This guards the DFlash draft column-offset contract without freezing exact draft token IDs across GPUs. +- Verification: + ```bash + cargo fmt --all --check + cargo check --release -p openinfer-qwen3-4b --lib + cargo test --release -p openinfer-qwen3-4b --lib -- --nocapture + OPENINFER_TEST_MODEL_PATH=/data/models/Qwen3-4B OPENINFER_DFLASH_TEST_MODEL_PATH=/data/models/Qwen3-4B-DFlash-b16 cargo test --release -p openinfer-qwen3-4b --test hf_golden_gate dflash_speculative_verify_matches_hf_argmax_regret_gate -- --nocapture + ``` + All commands passed locally. +- After adding the native full draft → verify assertions, the same DFlash HF gate passed and checked 152 target positions. +- The full serial HF gate also passed with the full-path DFlash assertions: + ```bash + OPENINFER_TEST_MODEL_PATH=/data/models/Qwen3-4B OPENINFER_DFLASH_TEST_MODEL_PATH=/data/models/Qwen3-4B-DFlash-b16 cargo test --release -p openinfer-qwen3-4b --test hf_golden_gate -- --nocapture --test-threads=1 + ``` + It covered DFlash full draft → verify, DFlash verifier-span, Qwen3 eager/cached replay, and Qwen3 CUDA Graph padded/cached replay buckets. +- The same full serial HF gate passed again after the config-derived DFlash memory reserve and effective-context admission changes; the DFlash verifier checked 152 target positions, and Qwen3 eager/cached/CUDA-graph padded replay gates all remained within tolerance. + +### Step 6: bs=1 serving performance on local 5070 Ti +- Baseline server command: + ```bash + ./target/release/openinfer --model-path /data/models/Qwen3-4B --served-model-name Qwen3-4B --port 8010 --max-prefill-tokens 1024 --no-prefix-cache + ``` +- DFlash server command: + ```bash + ./target/release/openinfer --model-path /data/models/Qwen3-4B --served-model-name Qwen3-4B --port 8010 --dflash-draft-model-path /data/models/Qwen3-4B-DFlash-b16 --max-prefill-tokens 1024 + ``` +- `vllm bench serve`, greedy bs=1 / `--max-concurrency 1`, measured output throughput: + +| Dataset | Prompt count | Output len | Baseline tok/s | DFlash tok/s | Speedup | Baseline mean TPOT | DFlash mean TPOT | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| Spec-Bench | 20 | 64 | 89.97 | 144.69-147.02 | 1.61-1.63x | 10.81 ms | 6.42-6.47 ms | +| ShareGPT sample | 20 | 64 | 88.23 | 126.00 | 1.43x | 10.99 ms | 7.50 ms | +| Random | 20 | 128 | 85.69 | 130.78 | 1.53x | 10.99 ms | 6.93 ms | + +- Final current Spec-Bench run after the executor split and 8193-vocab argmax equivalence gate (`target/dflash-bench/dflash-final-current-specbench-n20-out64.json`, server log `target/dflash-bench/dflash-server-final-current.log`): + - completed requests: 20 / 20 + - output throughput: 146.46 tok/s + - mean/median/p99 TPOT: 6.41 / 6.66 / 8.88 ms + - mean TTFT: 32.92 ms + - total input/output tokens: 3918 / 1280 + - this matches the previous chunkfix run within run-to-run noise, so the later cleanup and gate hardening did not regress bs=1 Spec-Bench throughput. +- Hotpath-cleanup Spec-Bench run after removing per-draft params/random allocation and reusing the DFlash block-token host buffer (`target/dflash-bench/dflash-hotpath-cleanup-specbench-n20-out64.json`, server log `target/dflash-bench/dflash-server-hotpath-cleanup.log`): + - completed requests: 20 / 20 + - output throughput: 147.02 tok/s + - mean/median/p99 TPOT: 6.42 / 6.67 / 8.90 ms + - mean TTFT: 30.59 ms + - total input/output tokens: 3918 / 1280 + - DFlash verify logs and acceptance were byte-for-byte aligned at the aggregate level with the final-current run: 491 rows, 769 / 6524 accepted/verified draft tokens, 1260 committed tokens. +- Host-buffer-cleanup Spec-Bench run after reusing `SamplingScratch::out_host` for DFlash greedy argmax D2H (`target/dflash-bench/dflash-host-buffer-cleanup-specbench-n20-out64.json`, server log `target/dflash-bench/dflash-server-host-buffer-cleanup.log`): + - completed requests: 20 / 20 + - output throughput: 145.75 tok/s + - mean/median/p99 TPOT: 6.42 / 6.67 / 8.90 ms + - mean TTFT: 34.51 ms + - total input/output tokens: 3918 / 1280 + - DFlash verify logs: 491 rows, accepted/verified draft tokens 769 / 6524 = 0.1179, committed 1260 tokens, timing fields stayed `-1.000`. +- Reserve-polish Spec-Bench run after deriving DFlash request-state memory reserve from the draft config and reducing the effective admission context limit by one DFlash block (`target/dflash-bench/dflash-reserve-polish-specbench-n20-out64.json`, server log `target/dflash-bench/dflash-server-reserve-polish.log`): + - startup reserve: weights 1074.9 MB, request/scratch extra 2959.3 MB, total 4034.1 MB. + - resulting local KV cache: 1028 blocks / 2313 MB; frontend `max_model_len=16432`. + - DFlash effective target context limit: `40944` tokens (`40960 - block_size 16`), so requests that would overflow the drafter cache are rejected at admission instead of failing mid-prefill. + - completed requests: 20 / 20 + - output throughput: 146.33 tok/s + - mean/median/p99 TPOT: 6.42 / 6.67 / 8.89 ms + - mean TTFT: 32.51 ms + - total input/output tokens: 3918 / 1280 + - DFlash verify logs: 491 rows, accepted/verified draft tokens 769 / 6524 = 0.1179, committed 1260 tokens, timing fields stayed `-1.000`. +- Latest post chunked-prefill-fix Spec-Bench run (`target/dflash-bench/dflash-chunkfix-specbench-n20-out64.json`, server log `target/dflash-bench/dflash-server-chunkfix.log`): + - completed requests: 20 / 20 + - output throughput: 146.90 tok/s + - mean/median/p99 TPOT: 6.42 / 6.67 / 8.90 ms + - request 19 used a 1188-token prompt context in DFlash after chunked prefill and completed 64 output tokens. + +- Result files live under `target/dflash-bench/`, including: + - `baseline-specbench-n20-out64.json` + - `dflash-specbench-n20-out64.json` + - `dflash-final-specbench-n20-out64.json` + - `dflash-nosync-specbench-n20-out64.json` + - `dflash-post-abstraction-specbench-n20-out64.json` + - `dflash-scratch-specbench-n20-out64.json` + - `dflash-tuned-specbench-n20-out64.json` + - `dflash-pending-buffer-specbench-n20-out64.json` + - `dflash-contiguous-argmax-specbench-n20-out64.json` + - `dflash-chunkfix-specbench-n20-out64.json` + - `dflash-final-current-specbench-n20-out64.json` + - `dflash-hotpath-cleanup-specbench-n20-out64.json` + - `dflash-host-buffer-cleanup-specbench-n20-out64.json` + - `dflash-reserve-polish-specbench-n20-out64.json` + - `baseline-sharegpt-n20-out64.json` + - `dflash-sharegpt-n20-out64.json` + - `baseline-random-in1024-out128-n20.json` + - `dflash-random-in1024-out128-n20.json` + +### Step 7: Acceptance, timing, and graph/profiling status +- Added INFO logs for every DFlash verify step: + - accepted draft tokens + - verified draft tokens + - committed tokens (`committed_tokens`, before scheduler stop-token suppression) + - cumulative acceptance rate + - last/average draft ms in profiling mode, otherwise `-1.000` + - last/average verify ms in profiling mode, otherwise `-1.000` + - pending and committed DFlash context lengths +- Default serving does not synchronize for DFlash timing. Profiling timings and NVTX ranges require: + ```bash + OPENINFER_QWEN3_DFLASH_NVTX=1 + ``` +- No-sync local Spec-Bench bs=1 run (`target/dflash-bench/dflash-nosync-specbench-n20-out64.json` and `target/dflash-bench/dflash-server-nosync.log`) showed: + - completed requests: 20 / 20 + - output throughput: 144.69 tok/s + - mean/median/p99 TPOT: 6.47 / 6.70 / 9.71 ms + - DFlash verify log rows: 494 + - accepted/verified draft tokens: 766 / 6544 = 0.117 + - committed tokens: 1260, or 2.55 tokens/speculative step + - timing fields: `draft_ms=-1.000`, `verify_ms=-1.000`, confirming default INFO logs do not force sync. +- Post-abstraction no-sync Spec-Bench bs=1 run (`target/dflash-bench/dflash-post-abstraction-specbench-n20-out64.json` and `target/dflash-bench/dflash-server-post-abstraction.log`) showed no throughput regression: + - completed requests: 20 / 20 + - output throughput: 144.60 tok/s + - mean/median/p99 TPOT: 6.47 / 6.69 / 9.72 ms + - DFlash verify log rows: 494 + - accepted/verified draft tokens: 766 / 6544 = 0.117 + - timing fields remained `draft_ms=-1.000`, `verify_ms=-1.000`. +- Draft scratch cleanup: + - `DFlashRequestState` now owns reusable draft buffers for context projection, block hidden states, tail K/V, MLP intermediates, and output logits. + - The first draft for a long prompt can grow the scratch to that prompt-context length; steady speculative steps reuse the same device buffers and only update `seq_len` metadata. + - The scratch-only Spec-Bench bs=1 run (`target/dflash-bench/dflash-scratch-specbench-n20-out64.json`) completed 20/20 requests at 144.89 tok/s with mean/median/p99 TPOT 6.47 / 6.69 / 9.71 ms, 494 log rows, acceptance 766 / 6544 = 0.117, and no timing sync. +- DFlash cublasLt tuning: + - Worker bind now extends target decode GEMM tuning with DFlash-specific shapes: `fc` `[M=2560, K=12800, N=1..16]` and DFlash K/V projection tails `[M=1024, K=2560, N=17..32]`. + - Startup logs the tuned ranges: + `Qwen3 DFlash cublasLt tuned: fc M=2560 K=12800 N=1..16, kv M=1024 K=2560 N=17..32`. + - Tuned Spec-Bench bs=1 run (`target/dflash-bench/dflash-tuned-specbench-n20-out64.json` and `target/dflash-bench/dflash-server-tuned.log`) completed 20/20 requests at 146.05 tok/s with mean/median/p99 TPOT 6.47 / 6.68 / 9.68 ms. + - Tuned DFlash verify logs: 495 rows, accepted/verified draft tokens 765 / 6559 = 0.1166, committed 1260 tokens, timing fields `-1.000`. +- Pending-context buffer cleanup: + - `DFlashRequestState` now keeps target hidden context in a reusable `DFlashPendingContext` instead of allocating a fresh `HiddenStates` on each append and allocating a merged buffer when chunks accumulate. + - The buffer tracks active length separately from capacity, grows to the first long prompt/chunk sequence, clears by resetting active length, and reuses the same device allocation for steady speculative steps. + - This removes a request-local allocation/merge source and makes the pending-context pointer stable after its first growth. Other graph blockers remain: K/V cache write offsets, RoPE positions, and attention `kv_len` still vary by step. + - Pending-buffer Spec-Bench bs=1 run (`target/dflash-bench/dflash-pending-buffer-specbench-n20-out64.json` and `target/dflash-bench/dflash-server-pending-buffer.log`) completed 20/20 requests at 145.97 tok/s with mean/median/p99 TPOT 6.46 / 6.67 / 9.67 ms. + - Pending-buffer DFlash verify logs: 495 rows, accepted/verified draft tokens 765 / 6559 = 0.1166, committed 1260 tokens, timing fields stayed `-1.000`. +- Contiguous greedy argmax cleanup: + - Added a Rust wrapper for the existing `argmax_batch_bf16_split_cuda` non-indexed path and taught `select_batch_tokens_into` to use it whenever every row is greedy and logits rows are contiguous. + - This removes the per-call row-index vector construction and H2D row-index copy from all-greedy paths, including DFlash draft and DFlash verifier token selection. It does not change the two argmax kernels themselves, so the launch-heavy draft profile is expected to remain dominated by GEMM/GEMV. + - Added `openinfer-core` GPU equivalence coverage comparing contiguous all-greedy output against indexed greedy rows from the mixed path at 8193 vocab, including cross-tile maxima and a cross-tile tie. + - Contiguous-argmax Spec-Bench bs=1 run (`target/dflash-bench/dflash-contiguous-argmax-specbench-n20-out64.json` and `target/dflash-bench/dflash-server-contiguous-argmax.log`) completed 20/20 requests at 145.98 tok/s with mean/median/p99 TPOT 6.47 / 6.68 / 9.69 ms. + - Contiguous-argmax DFlash verify logs: 495 rows, accepted/verified draft tokens 765 / 6559 = 0.1166, committed 1260 tokens, timing fields stayed `-1.000`. +- Chunked-prefill fix no-sync run: + - `target/dflash-bench/dflash-chunkfix-specbench-n20-out64.json` and `target/dflash-bench/dflash-server-chunkfix.log` completed 20/20 Spec-Bench requests at 146.90 tok/s with mean/median/p99 TPOT 6.42 / 6.67 / 8.90 ms. + - DFlash verify logs: 491 rows, accepted/verified draft tokens 769 / 6524 = 0.1179, committed 1260 tokens, or 2.57 tokens/speculative step. + - The longest prompt row had `draft_context_tokens=1188` followed by `draft_committed_context=1188` on the next speculative step, confirming the chunked prefill context was preserved across chunks. + - Timing fields stayed `draft_ms=-1.000`, `verify_ms=-1.000`. +- Final current no-sync run: + - `target/dflash-bench/dflash-final-current-specbench-n20-out64.json` and `target/dflash-bench/dflash-server-final-current.log` completed 20/20 Spec-Bench requests at 146.46 tok/s with mean/median/p99 TPOT 6.41 / 6.66 / 8.88 ms. + - DFlash verify logs: 491 rows, accepted/verified draft tokens 769 / 6524 = 0.1179, committed 1260 tokens, or 2.57 tokens/speculative step. + - Max observed `draft_context_tokens=1188`; max `draft_committed_context=1245`. + - Timing fields stayed `draft_ms=-1.000`, `verify_ms=-1.000`, so default INFO logging remains no-sync. +- Hotpath-cleanup no-sync run: + - `target/dflash-bench/dflash-hotpath-cleanup-specbench-n20-out64.json` and `target/dflash-bench/dflash-server-hotpath-cleanup.log` completed 20/20 Spec-Bench requests at 147.02 tok/s with mean/median/p99 TPOT 6.42 / 6.67 / 8.90 ms. + - DFlash verify logs: 491 rows, accepted/verified draft tokens 769 / 6524 = 0.1179, committed 1260 tokens, or 2.57 tokens/speculative step. + - Max observed `draft_context_tokens=1188`; max `draft_committed_context=1245`. + - Timing fields stayed `draft_ms=-1.000`, `verify_ms=-1.000`. +- Host-buffer-cleanup no-sync run: + - Reused `SamplingScratch::out_host` for the greedy contiguous argmax D2H path, removing the last per-DFlash-draft host allocation in token selection. + - `target/dflash-bench/dflash-host-buffer-cleanup-specbench-n20-out64.json` and `target/dflash-bench/dflash-server-host-buffer-cleanup.log` completed 20/20 Spec-Bench requests at 145.75 tok/s with mean/median/p99 TPOT 6.42 / 6.67 / 8.90 ms. + - DFlash verify logs: 491 rows, accepted/verified draft tokens 769 / 6524 = 0.1179, committed 1260 tokens, or 2.57 tokens/speculative step. + - Max observed `draft_context_tokens=1188`; max `draft_committed_context=1245`. + - Timing fields stayed `draft_ms=-1.000`, `verify_ms=-1.000`. +- DFlash memory/admission polish: + - The old fixed 1 GiB DFlash extra reserve only covered short serving probes. The reserve is now derived from the DFlash config and covers one request's worst-case drafter KV cache, pending target-hidden context, long-context projection scratch, tail K/V scratch, block scratch, logits scratch, and a 256 MiB allocator margin, with 1 GiB as a floor. + - Local startup after this change logged `weights=1074.9 MB, extra=2959.3 MB, total=4034.1 MB`; the resulting KV pool has 1028 Qwen3 blocks and still runs the standard bs=1 Spec-Bench profile at 146.33 tok/s. + - Qwen3 executor metadata now exposes an effective max context of `target_max_position_embeddings - dflash_block_size` when DFlash is loaded. This keeps DFlash cache overflow as an admission-time rejection instead of a GPU prefill failure. +- Toxic-reviewer follow-up caught a real shared-path regression in the first hotpath-cleanup patch: `select_batch_tokens_into` briefly required `params.len() == logits.seq_len`, which conflicts with Qwen3 CUDA Graph decode where logits are bucket-padded (`bs=5 -> bucket 8`, `bs=9 -> bucket 16`). + - Fixed the sampling contract so `logits.seq_len` may include padding rows while `params.len()` names the real rows to consume; `params.len() > logits.seq_len` still fails early. + - `LocalQwen3Lane::select_step_tokens` now sizes sampling scratch from `max(logits.seq_len, params.len())`, preserving padded-bucket capacity. + - Added `all_greedy_contiguous_argmax_ignores_cuda_graph_padding_rows`, which uses 8193-vocab logits with 8 padded rows and 5 real greedy params. + - Verification: + ```bash + cargo test --release -p openinfer-core all_greedy_contiguous_argmax --lib -- --nocapture + OPENINFER_TEST_MODEL_PATH=/data/models/Qwen3-4B OPENINFER_DFLASH_TEST_MODEL_PATH=/data/models/Qwen3-4B-DFlash-b16 cargo test --release -p openinfer-qwen3-4b --test hf_golden_gate -- --nocapture --test-threads=1 + ``` + Both passed. The serial HF gate covered the DFlash verifier-span gate plus Qwen3 `batched cuda-graph (9 padded)`, `batched cuda-graph (5 padded)`, and `batched cuda-graph cached replay (5)`. +- Local timing smoke (`target/dflash-bench/dflash-server-timing.log`, Spec-Bench bs=1) showed: + - 94 DFlash verify steps + - accepted/verified draft tokens: 158 / 1297 = 0.1218 + - committed tokens: 252, or 2.68 tokens/speculative step + - draft mean/median/p90/max: 2.740 / 2.692 / 2.740 / 4.907 ms + - verify mean/median/p90/max: 13.546 / 13.414 / 13.870 / 14.271 ms + - steady draft rows (`draft_context_tokens <= 16`) averaged 2.697 ms; cold prompt-context rows averaged 3.728 ms. +- Added optional NVTX ranges with `OPENINFER_QWEN3_DFLASH_NVTX=1`: + - `qwen3.dflash.draft` + - `qwen3.dflash.verify` +- Local nsys trace (`target/dflash-bench/nsys-dflash-full-with-request.nsys-rep`) on 2 Spec-Bench prompts: + - 47 draft ranges: mean 2.874 ms, min 2.717 ms, max 5.020 ms + - 47 verify ranges: mean 13.938 ms, min 13.616 ms, max 14.375 ms + - all captured kernels had null `graphId`; the DFlash speculative path is eager, not CUDA-graph replayed. + - draft range kernel sum: 117.6 ms total, 2.50 ms/range, ~106 kernel launches/range + - draft range kernel time split: + - GEMM/GEMV: 110.7 ms, 94.1% + - single-prefill attention: 2.8 ms, 2.4% + - norm/MLP elementwise: 1.5 ms, 1.3% + - hidden copy/embedding: 0.8 ms, 0.7% + - cublasLt split-K reduce: 0.8 ms, 0.7% + - DFlash qk-norm+RoPE: 0.5 ms, 0.4% + - argmax: 0.5 ms, 0.4% +- Pending-buffer nsys trace (`target/dflash-bench/nsys-dflash-pending-buffer.nsys-rep` / `.sqlite`) on 2 Spec-Bench prompts: + - 47 draft ranges: mean/median/min/max 2.758 / 2.677 / 2.616 / 5.137 ms by NVTX range. + - 47 verify ranges: mean/median/min/max 13.953 / 14.100 / 13.606 / 14.424 ms by NVTX range. + - all kernels inside both DFlash draft and DFlash verify ranges had null `graphId`; DFlash speculative draft and target verify are still eager. + - draft range kernel sum: 117.0 ms total, 2.49 ms/range, ~98 kernel launches/range. + - draft range kernel time split: + - CUTLASS/cublasLt GEMM/GEMV: 104.4 ms, 89.3% + - FlashInfer single-prefill attention: 2.8 ms, 2.4% + - SM120 nvjet GEMM variants: 2.4 ms, 2.0% + - DFlash qk-norm+RoPE: 0.5 ms, 0.4% + - hidden/context copies: 0.8 ms, 0.7% + - argmax: 0.5 ms, 0.4% + - profiling-mode server timings from the same run: 47 rows, draft mean/median/min/max 2.895 / 2.796 / 2.742 / 4.854 ms, verify mean/median/min/max 13.760 / 13.751 / 13.636 / 13.953 ms. Steady rows with `draft_context_tokens <= 16` averaged 2.852 ms; the 2 cold prompt-context rows averaged 3.858 ms. +- KernelWiki references checked for the graph/fill question: + - `sources/prs/vllm/PR-17484.md` and `sources/prs/vllm/PR-17668.md` show the relevant upstream pattern: move q/o projections and rotary work into the CUDA-graph region to reduce CPU overhead. + - `sources/prs/vllm/PR-25984.md` records a speculative-decode attention path that refactors metadata for optimized spec-as-decode FlashInfer-MLA kernels. + - `sources/prs/flashinfer/PR-969.md` and `sources/prs/flashinfer/PR-2244.md` reinforce the same constraint for our DFlash path: graph readiness requires removing host API/sync overhead and making all per-step metadata capturable or graph-updatable. +- Local GPU metrics/occupancy could not be collected: nsys reported `ERR_NVGPUCTRPERM` on the local 5070 Ti, and `ncu` is not installed locally. This means local evidence proves eager launch structure and time split, but not SM occupancy percentage. + +### Step 8: 5090 validation and vLLM comparison +- Remote machine: `5090` (`host-192-168-172-86`), 8x RTX 5090, idle at probe time. +- Applied the patch to remote branch `feat/qwen3-dflash-mtp-codex` in `~/develop/xingming/pegainfer`. +- Remote build/test commands: + ```bash + export PATH=/root/.cargo/bin:/usr/local/cuda-13.1/bin:$PATH + export CUDA_HOME=/usr/local/cuda-13.1 + cargo fmt --all --check + cargo check --release -p openinfer-server + cargo test --release -p openinfer-qwen3-4b --lib -- --nocapture + cargo test --release -p openinfer-kv-cache --test lifecycle speculative -- --nocapture + cargo build --release -p openinfer-server + ``` + All commands passed on 5090. +- After adding DFlash draft scratch and DFlash-specific cublasLt tuning, the same patch was re-synced to 5090 and the same fmt/check/lib/KV/build command set passed again. The DFlash config/model-weight gate skipped there because `OPENINFER_DFLASH_TEST_MODEL_PATH` is not available on the box. +- After adding the pending-context reusable buffer and contiguous greedy argmax fast path, the final patch was re-synced to 5090 and this command set passed: + ```bash + export PATH=/root/.cargo/bin:/usr/local/cuda-13.1/bin:$PATH + export CUDA_HOME=/usr/local/cuda-13.1 + cargo fmt --all --check + cargo check --release -p openinfer-server + cargo test --release -p openinfer-kernels ops::sampling --lib -- --nocapture + cargo test --release -p openinfer-qwen3-4b --lib -- --nocapture + cargo test --release -p openinfer-kv-cache --test lifecycle speculative -- --nocapture + cargo build --release -p openinfer-server + git diff --check + ``` + The remote DFlash config/model-weight gate still skipped because `OPENINFER_DFLASH_TEST_MODEL_PATH` is not available on the box. +- After the chunked-prefill DFlash context fix, the patch was re-synced to 5090 again and the same fmt/check/sampling/Qwen3-lib/KV/build/diff command set passed on `host-192-168-172-86` with CUDA 13.1. At that pre-download point, real-weight DFlash checks skipped. +- After splitting the executor files under 1k lines and adding the core contiguous-vs-indexed greedy argmax equivalence test, the patch was re-synced to 5090 again and this command set passed: + ```bash + export PATH=/root/.cargo/bin:/usr/local/cuda-13.1/bin:$PATH + export CUDA_HOME=/usr/local/cuda-13.1 + cargo fmt --all --check + cargo check --release -p openinfer-server + cargo test --release -p openinfer-core all_greedy_contiguous_argmax_matches_indexed_greedy_rows --lib -- --nocapture + cargo test --release -p openinfer-kernels ops::sampling --lib -- --nocapture + cargo test --release -p openinfer-qwen3-4b --lib -- --nocapture + cargo test --release -p openinfer-kv-cache --test lifecycle speculative -- --nocapture + cargo build --release -p openinfer-server + git diff --check + ``` + At that pre-download point, real-weight DFlash checks skipped. +- After removing per-DFlash-draft CPU allocations for params/random vectors and reusing the block-token host buffer, the patch was re-synced to 5090 again and this command set passed: + ```bash + export PATH=/root/.cargo/bin:/usr/local/cuda-13.1/bin:$PATH + export CUDA_HOME=/usr/local/cuda-13.1 + cargo fmt --all --check + cargo check --release -p openinfer-server + cargo test --release -p openinfer-qwen3-4b --lib -- --nocapture + cargo build --release -p openinfer-server + git diff --check + ``` + At that pre-download point, real-weight DFlash checks skipped. +- After fixing the padded-logits sampling contract from the toxic-reviewer follow-up, the patch was re-synced to 5090 again and this command set passed: + ```bash + export PATH=/root/.cargo/bin:/usr/local/cuda-13.1/bin:$PATH + export CUDA_HOME=/usr/local/cuda-13.1 + cargo fmt --all --check + cargo test --release -p openinfer-core all_greedy_contiguous_argmax --lib -- --nocapture + cargo test --release -p openinfer-qwen3-4b --lib -- --nocapture + cargo build --release -p openinfer-server + git diff --check + ``` + At that pre-download point, real-weight DFlash checks skipped. +- After adding the full native draft → verify assertions to the DFlash HF gate, the patch was re-synced to 5090 again and this command set passed: + ```bash + export PATH=/root/.cargo/bin:/usr/local/cuda-13.1/bin:$PATH + export CUDA_HOME=/usr/local/cuda-13.1 + cargo fmt --all --check + cargo test --release -p openinfer-qwen3-4b --test hf_golden_gate dflash_speculative_verify_matches_hf_argmax_regret_gate -- --nocapture + cargo test --release -p openinfer-qwen3-4b --lib -- --nocapture + cargo build --release -p openinfer-server + git diff --check + ``` + The remote integration test compiled; at that pre-download point, real-weight checks skipped because the default relative model paths did not exist on 5090. +- After reusing `SamplingScratch::out_host` for DFlash greedy argmax D2H, the patch was re-synced to 5090 again and this command set passed: + ```bash + export PATH=/root/.cargo/bin:/usr/local/cuda-13.1/bin:$PATH + export CUDA_HOME=/usr/local/cuda-13.1 + cargo fmt --all --check + cargo test --release -p openinfer-qwen3-4b --lib -- --nocapture + cargo build --release -p openinfer-server + git diff --check + ``` + At that pre-download point, real-weight DFlash checks skipped. +- After adding config-derived DFlash memory reserve and DFlash-aware effective context admission, the patch was re-synced to 5090 again and this command set passed: + ```bash + export PATH=/root/.cargo/bin:/usr/local/cuda-13.1/bin:$PATH + export CUDA_HOME=/usr/local/cuda-13.1 + cargo fmt --all --check + cargo test --release -p openinfer-qwen3-4b dflash_context_limit_reserves_one_draft_block --lib -- --nocapture + cargo test --release -p openinfer-qwen3-4b dflash_request_state_reserve_covers_long_context_scratch --lib -- --nocapture + cargo test --release -p openinfer-qwen3-4b --lib -- --nocapture + cargo build --release -p openinfer-server + git diff --check + ``` + At that pre-download point, real-weight DFlash checks skipped. +- After the user explicitly approved placing the DFlash artifact on 5090, downloaded `z-lab/Qwen3-4B-DFlash-b16` to `/data/Qwen3-4B-DFlash-b16` using the box proxy from `.bashrc` (`http://172.17.0.1:1081`). + - `model.safetensors`: 1074860568 bytes + - `config.json`, `README.md`, `dflash.py`, `modeling_dflash.py`, `utils.py`, and `assets/` are present. +- Real-weight 5090 gates now pass: + ```bash + OPENINFER_TEST_MODEL_PATH=/data/Qwen3-4B OPENINFER_DFLASH_TEST_MODEL_PATH=/data/Qwen3-4B-DFlash-b16 cargo test --release -p openinfer-qwen3-4b dflash::tests::downloaded_dflash_config_matches_qwen3_4b --lib -- --nocapture + OPENINFER_TEST_MODEL_PATH=/data/Qwen3-4B OPENINFER_DFLASH_TEST_MODEL_PATH=/data/Qwen3-4B-DFlash-b16 cargo test --release -p openinfer-qwen3-4b --test hf_golden_gate dflash_speculative_verify_matches_hf_argmax_regret_gate -- --nocapture + ``` + Both passed; the DFlash HF gate checked 152 target positions. +- 5090 OpenInfer bs=1 `vllm bench serve`, greedy, `--max-concurrency 1`: + +| Dataset | Prompt count | Output len | Baseline tok/s | DFlash tok/s | Speedup | Baseline mean TPOT | DFlash mean TPOT | Result files | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | --- | +| Spec-Bench | 20 | 64 | 167.34 | 251.48 | 1.50x | 5.83 ms | 3.79 ms | `/data/dflash-bench/{baseline,dflash}-5090-specbench-n20-out64-c1.json` | +| ShareGPT sample | 20 | 64 | 144.56 | 144.50 | 1.00x | 6.67 ms | 6.67 ms | `/data/dflash-bench/{baseline,dflash}-5090-sharegpt-n20-out64-c1.json` | +| Random 1024/128 | 20 | 128 | 154.00 | 168.83 | 1.10x | 6.18 ms | 5.59 ms | `/data/dflash-bench/{baseline,dflash}-5090-random-in1024-out128-n20-c1.json` | + +- OpenInfer concurrency probe: Spec-Bench `--max-concurrency 4`, 40 prompts, output 64: + - baseline: 558.18 tok/s, mean TPOT 6.83 ms + - DFlash: 540.23 tok/s, mean TPOT 6.91 ms + - This was measured before the multi-active default policy landed, and is retained as historical evidence that the first bs=1-only policy left batch-concurrency behavior under-optimized. +- 5090 DFlash INFO logs: + - `/data/dflash-bench/dflash-server-5090.log`: 485 rows, accepted/verified draft tokens 775 / 6478 = 0.1196, committed 1260 tokens, 2.60 committed tokens/step, no old `emitted=` field. + - `/data/dflash-bench/dflash-server-5090-r2.log`: 477 rows, accepted/verified draft tokens 669 / 6795 = 0.0985, committed 1146 tokens, 2.40 committed tokens/step, no old `emitted=` field. +- 5090 nsys DFlash Spec-Bench c1 trace: + - report: `/data/dflash-bench/nsys-5090-dflash-specbench.nsys-rep` + - sqlite: `/data/dflash-bench/nsys-5090-dflash-specbench.sqlite` + - result: `/data/dflash-bench/dflash-5090-nsys-specbench-n4-out64-c1.json` + - nsys-wrapped run completed 4/4 at 253.73 tok/s, mean TPOT 3.71 ms. + - Profiling-mode INFO rows show steady draft around 1.52-1.59 ms and verify around 8.0-8.6 ms, with tail verify rows near 9.3-9.9 ms. +- Upstream vLLM support check on 5090: + - `.venv/bin/vllm --version` reports `0.22.1`. + - `vllm serve --help=all` lists `--spec-method dflash`. + - Python import finds `vllm.model_executor.models.qwen3_dflash`. + - vLLM DFlash startup logs resolve `DFlashDraftModel`, build `SpeculativeConfig(method='dflash', model='/data/Qwen3-4B-DFlash-b16', num_spec_tokens=16)`, and warn that DFlash uses the v1 model runner because runner v2 does not yet support DFlash parallel drafting. +- 5090 vLLM 0.22.1 bs=1 `vllm bench serve`, greedy, `--max-num-seqs 1`, `--max-model-len 40944`, `--max-concurrency 1`: + +| Dataset | Prompt count | Output len | vLLM baseline tok/s | vLLM DFlash tok/s | Speedup | Baseline mean TPOT | DFlash mean TPOT | DFlash acceptance length | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| Spec-Bench | 20 | 64 | 162.48 | 289.57 | 1.78x | 5.97 ms | 3.07 ms | 2.62 | +| ShareGPT sample | 20 | 64 | 159.96 | 226.71 | 1.42x | 5.98 ms | 4.01 ms | 2.06 | +| Random 1024/128 | 20 | 128 | 156.85 | 258.13 | 1.65x | 6.03 ms | 3.49 ms | 2.64 | + + Result files: + - `/data/dflash-bench/vllm-baseline-5090-specbench-n20-out64-c1-r2.json` + - `/data/dflash-bench/vllm-dflash-5090-specbench-n20-out64-c1-r2.json` + - `/data/dflash-bench/vllm-baseline-5090-sharegpt-n20-out64-c1.json` + - `/data/dflash-bench/vllm-dflash-5090-sharegpt-n20-out64-c1.json` + - `/data/dflash-bench/vllm-baseline-5090-random-in1024-out128-n20-c1.json` + - `/data/dflash-bench/vllm-dflash-5090-random-in1024-out128-n20-c1.json` + Server logs: + - `/data/dflash-bench/vllm-dflash-5090-server.log` + - `/data/dflash-bench/vllm-baseline-5090-server.log` +- 5090 profiler state: + - `nsys` is installed and supports GPU metrics on all 8 GPUs. + - `ncu` is not installed. + +### Step 9: Review follow-up fixes +- Toxic reviewer follow-up found two real issues after the reserve/admission polish: + - A mid-prompt chunk could start a fresh DFlash state after an earlier mixed/chunked prefill step had dropped state, then fail on the hidden-context continuity check. Worker capture eligibility now uses real lane state: `chunk_start == 0` may start DFlash state, while `chunk_start > 0` captures hidden state only when a pending DFlash request state already exists. The worker returns `PrefillResult::dflash_context_captured_requests`, and the main executor uses the per-request result when deciding `MarkReady` / `KeepPending` / `Drop`. + - DFlash INFO logs used the field name `emitted` for `accepted_tokens.len()`. That value is committed to request state before scheduler stop-token suppression, so the log field is now `committed_tokens`. +- Local validation after this fix: + ```bash + cargo fmt --all --check + cargo test --release -p openinfer-qwen3-4b dflash_prefill --lib -- --nocapture + cargo test --release -p openinfer-qwen3-4b servable_len_reports_the_tighter_context_or_kv_limit --lib -- --nocapture + cargo test --release -p openinfer-qwen3-4b --lib -- --nocapture + cargo build --release -p openinfer-server + OPENINFER_TEST_MODEL_PATH=/data/models/Qwen3-4B OPENINFER_DFLASH_TEST_MODEL_PATH=/data/models/Qwen3-4B-DFlash-b16 cargo test --release -p openinfer-qwen3-4b --test hf_golden_gate dflash_speculative_verify_matches_hf_argmax_regret_gate -- --nocapture + git diff --check + ``` + All commands passed locally. The DFlash HF gate checked 152 target positions after this fix. +- Local Step 9 Spec-Bench bs=1 run (`target/dflash-bench/dflash-step9-specbench-n20-out64.json`, server log `target/dflash-bench/dflash-server-step9.log`) completed 20/20 requests: + - output throughput: 148.34 tok/s + - mean/median/p99 TPOT: 6.34 / 6.60 / 8.80 ms + - mean TTFT: 31.65 ms + - total input/output tokens: 3918 / 1280 + - DFlash verify logs: 491 rows, accepted/verified draft tokens 769 / 6524 = 0.1179, committed 1260 tokens, or 2.57 committed tokens/speculative step. + - Default timing fields stayed `draft_ms=-1.000`, `verify_ms=-1.000`, and no DFlash verify row used the old `emitted=` field. + +### Step 10: Draft PR and local multi-active validation +- Draft PR opened: . +- After user direction to remove the single-active restriction, scheduler policy now lets all eligible active greedy requests enter DFlash when the draft model is configured. The executor drafts each active request and batches target verification. +- Local Spec-Bench dataset materialized for vLLM 0.21: + ```bash + curl -L --fail --show-error https://raw.githubusercontent.com/hemingkx/Spec-Bench/refs/heads/main/data/spec_bench/question.jsonl -o target/dflash-bench/spec_bench_question.jsonl + wc -l target/dflash-bench/spec_bench_question.jsonl + ``` + The file has 480 JSONL rows with the expected `turns` field. +- Latest local PR-head serving results on RTX 5070 Ti, greedy `vllm bench serve`, `--temperature 0`, OpenAI `/v1/completions` endpoint: + +| Dataset | Concurrency | Prompts | Output len | Baseline tok/s | DFlash tok/s | Speedup | Baseline mean TPOT | DFlash mean TPOT | Result files | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | --- | +| Spec-Bench | 1 | 20 | 64 | 89.72 | 149.32 | 1.66x | 10.77 ms | 6.33 ms | `baseline-local-pr380-specbench-n20-out64-c1.json`, `dflash-local-pr380-specbench-n20-out64-c1-r2.json` | +| Spec-Bench | 4 | 40 | 64 | 303.92 | 330.42 | 1.09x | 12.09 ms | 10.95 ms | `baseline-local-pr380-specbench-n40-out64-c4.json`, `dflash-local-pr380-specbench-n40-out64-c4-r2.json` | +| Random 1024/128 | 1 | 20 | 128 | 86.61 | 136.09 | 1.57x | 10.86 ms | 6.63 ms | `baseline-local-pr380-random-in1024-out128-n20-c1.json`, `dflash-local-pr380-random-in1024-out128-n20-c1.json` | +| Random 1024/128 | 4 | 40 | 128 | 270.93 | 349.50 | 1.29x | 13.41 ms | 10.18 ms | `baseline-local-pr380-random-in1024-out128-n40-c4.json`, `dflash-local-pr380-random-in1024-out128-n40-c4.json` | + +- DFlash server log for the local multi-active session: `target/dflash-bench/dflash-local-pr380-multiactive-server.log`. + - The log contains DFlash rows for 180 internal request IDs (`0..179`) and shows multiple request IDs logging DFlash acceptance in the same decode wave, e.g. request `0/1/2/3` immediately after the first c4 batch. + - Aggregate over the session log: 5529 DFlash rows, accepted/verified draft tokens 9651 / 75682 = 0.1275, committed 15180 tokens, or 2.75 committed tokens/speculative row. + - The aggregate includes the discarded overlapping Spec-Bench c1/c4 attempt, so it is only evidence of multi-active coverage and acceptance logging, not a per-dataset performance number. + +### Step 11: Multi-active hardening after review +- Toxic review of the multi-active change found two release blockers: + - `Prefill` failure while DFlash-ready active requests exist only targeted pending requests, then cleared `active`, so active requests could disappear without an error event or executor cleanup. + - Removing the single-active restriction let each eligible request allocate DFlash side state without aggregate memory admission. +- Fixes: + - `failure_targets_for(Prefill)` now includes active requests as well as pending requests, so a failed active+pending prefill step errors and drops every touched request. + - `RequestKv::revert_schedule` exposes kvbm's scheduled-state rollback. `execute_speculative_verify_impl` now rolls back every request whose speculative KV was scheduled if later scheduling, worker execution, or result-shape validation fails. + - DFlash memory reserve is split into total reserved bytes and request-state budget bytes. Scheduler admission estimates each DFlash-supported request's true side-state footprint from the DFlash config and `prompt + max_tokens + block_size`; the global allocator margin/floor is reserved once, not charged per request. Requests that exceed the remaining side-state budget stay deferred instead of OOMing after admission. + - DFlash prefill hidden capture is now per request. A mixed prefill batch can still capture target hidden states; unsupported rows drop only their own DFlash side state, while supported greedy rows keep progressing toward `dflash_ready`. +- Local validation: + ```bash + cargo fmt --all --check + cargo test --release -p openinfer-kv-cache --test lifecycle speculative -- --nocapture + cargo test --release -p openinfer-qwen3-4b dflash_prefill --lib -- --nocapture + cargo test --release -p openinfer-qwen3-4b admission_respects_speculative_state_budget_for_supported_requests --lib -- --nocapture + cargo test --release -p openinfer-qwen3-4b admission_counts_only_active_requests_with_speculative_state --lib -- --nocapture + cargo test --release -p openinfer-qwen3-4b dflash_short_request_footprint_does_not_include_global_allocator_floor --lib -- --nocapture + cargo test --release -p openinfer-qwen3-4b prefill_failure_with_active_targets_active_and_pending_requests --lib -- --nocapture + cargo test --release -p openinfer-qwen3-4b --lib -- --nocapture + cargo build --release -p openinfer-server + OPENINFER_TEST_MODEL_PATH=/data/models/Qwen3-4B OPENINFER_DFLASH_TEST_MODEL_PATH=/data/models/Qwen3-4B-DFlash-b16 cargo test --release -p openinfer-qwen3-4b --test hf_golden_gate dflash_speculative_verify_matches_hf_argmax_regret_gate -- --nocapture + ``` + These passed locally; the HF gate checked 152 DFlash verifier target positions. +- Post-hardening serving smoke on RTX 5070 Ti after fixing per-request side-state footprint: + - Server log: `target/dflash-bench/dflash-pr380-hardening-smoke-r2-server.log` + - Result file: `target/dflash-bench/dflash-pr380-hardening-smoke-r2-specbench-n12-c4.json` + - Command shape: Spec-Bench, greedy, output 64, `--num-prompts 12`, `--max-concurrency 4`, OpenAI `/v1/completions`, served model `Qwen3-4B`. + - Startup log: `Qwen3 DFlash memory reserve: weights=1074.9 MB, extra=2959.3 MB, state_budget=2690.8 MB, total=4034.1 MB`. + - Result: 12/12 successful, 768 generated tokens, 368.71 tok/s, mean/median/p99 TPOT 9.54 / 9.44 / 11.36 ms, mean TTFT 70.73 ms. + - Log evidence: DFlash INFO rows include four live request IDs in the same wave (for example `0/1/2/3` at the start and `8/9/10/11` later) with cumulative acceptance around 0.12. + +## Debrief + +- **Outcome**: DFlash is integrated end-to-end for the Qwen3-4B TP1 greedy path, with native drafter loading/forward, reusable draft scratch, per-draft CPU allocation cleanup, reusable D2H host output buffer for greedy token selection, config-derived request-state memory reserve, DFlash-aware effective context admission, DFlash side-state byte admission, transactional speculative KV rollback, per-request DFlash prefill capture, DFlash-specific small-N cublasLt tuning, speculative scheduler stages, hidden-state capture, acceptance/timing logs, NVTX profiling ranges, local and 5090 standard bs=1 serving measurements, local multi-active serving measurements, verifier-span plus full draft→verify HF gates, executor files split under 1k lines, 5090 real-weight gates, a vLLM 0.22.1 DFlash baseline for comparison, and draft PR #380. +- **Review**: toxic reviewer previously returned Ready for review after the executor split and 8193-vocab argmax equivalence test, then follow-ups caught the chunked-prefill lifecycle / `emitted` log-field issues fixed in Step 9 and the multi-active failure/admission issues fixed in Step 11. Non-blockers called out in review are intentionally documented here: DFlash is not in CUDA Graph, draft execution is not batched/fused, and HTTP/scheduler-level DFlash E2E coverage can be broadened later. +- **Pitfalls encountered**: + - Draft tokens must be clamped to remaining `max_tokens` before building the verify span; otherwise speculative KV views can request a span larger than the committed output budget. + - `vllm bench serve` no longer forces greedy requests by default. DFlash only engages when requests are greedy, so benchmark commands must include `--temperature 0` when the server default is not enough. + - Per-step timing must not synchronize the hot path by default. Use no-sync INFO logs for acceptance/shape observability and `OPENINFER_QWEN3_DFLASH_NVTX=1` when profiling timing/graphs. + - Multi-active DFlash is correct enough to benchmark, but the draft path is still per-request serial. This preserves the user's "DFlash on means eligible requests use DFlash" semantics, while leaving real batched/fused draft execution as the next performance target. + - Chunked prefill needs continuity checks on the hidden-state side channel. Dropping state on a non-final chunk or accepting a mismatched `chunk_start` corrupts the drafter context without touching the normal target KV path. + - DFlash request memory is not just the 1.1 GB safetensors file. Long prompts can grow drafter KV, pending hidden-state context, and projection scratch by multiple GiB, so the KV pool reserve must be config-derived, the admission context limit must leave one draft block of headroom, and scheduler admission must budget DFlash side state across concurrent requests. + - Local NVTX capture-range triggering did not match reliably; a bounded full nsys capture plus sqlite post-processing produced the useful draft/verify breakdown. + - Local nsys GPU metrics are blocked by permissions. The 5090 supports nsys GPU metrics and now has the draft model, so exact graph/kernel attribution should be collected there before final bs=1 optimization claims. + - vLLM DFlash materially outperforms the current OpenInfer DFlash path on ShareGPT and random bs=1, not just Spec-Bench. That points to a real implementation/policy gap rather than benchmark noise. +- **Lessons learned**: + - DFlash's `acceptance_length + 1` maps naturally to kvbm's `apply_speculative(&accepted_tokens)`: accepted draft prefix plus posterior bonus. + - The current speedup is from amortizing one target verifier step over ~2.7 committed tokens, not from a fully optimized drafter. The drafter is eager, launch-heavy, and GEMM/GEMV dominated. + - DFlash draft graph capture requires more than wrapping the current function in `CudaGraphState`: pending context pointers, K/V cache write offsets, RoPE positions, and attention `kv_len` are dynamic today. They need stable scratch/device metadata or graph-exec parameter updates before capture would be correct. + - DFlash verifier correctness must be tested through a DFlash-enabled executor. A bare Qwen3 executor does not own DFlash hidden-context state and should fail closed on speculative verify. + - A production-quality next step is DFlash graph capture or a fused/smaller-launch draft path, plus aligning OpenInfer's request-state policy with vLLM's proven bs=1 behavior. Tuning acceptance policy alone will not answer the draft-kernel efficiency question. +- **Follow-ups**: + - Optimize bs=1 first against the 5090 vLLM DFlash reference: inspect why OpenInfer ShareGPT barely enters or barely benefits from DFlash, then profile draft/verify graph capture and kernel fill. + - Re-run the latest multi-active PR head on 5090 once the host is convenient again; local 5070 Ti shows positive c4 throughput, but 5090 is still the release performance gate. + - Replace the per-request serial DFlash draft loop with a batched/fused draft path or graph-captured draft replay, then re-run the same Spec-Bench/random c1/c4 gates. diff --git a/openinfer-core/src/ops.rs b/openinfer-core/src/ops.rs index 25e7049c..27de33eb 100644 --- a/openinfer-core/src/ops.rs +++ b/openinfer-core/src/ops.rs @@ -15,7 +15,8 @@ pub use attention::{ }; pub use openinfer_kernels::ops::{ GEMM_LT_MAX_N, LoraDecodeGroupedProjection, accumulate_bf16_token_scaled_to_f32_into, - add_batch, add_batch_into, bf16_hidden_to_f32_into, embedding_decode_into, extract_vec, + add_batch, add_batch_into, bf16_hidden_to_f32_into, copy_hidden_rows_into, + copy_hidden_token_range_into, dflash_qk_norm_rope_into, embedding_decode_into, extract_vec, extract_vec_into, extract_vec_ref, extract_vec_ref_into, f32_to_bf16_hidden_into, fused_add_rms_norm_into, gather_hidden_tokens_into, gemm, gemm_into_checked, gemm_lt_tune, gemm_per_token, gemv, linear, lora_decode_fused_delta_group3_into, @@ -23,7 +24,8 @@ pub use openinfer_kernels::ops::{ qk_norm_partial_rope_batched_decode_hd256_into, rms_norm, rms_norm_batch_offset_into, rms_norm_gated_batch_into, rms_norm_into, rms_norm_offset_into, scale_f32_in_place, scaled_add_batch_into, scaled_add_rows_indexed_into, scaled_add_rows_into, - scaled_add_rows_token_range_into, silu_mul_batch, silu_mul_batch_into, write_vec_into, + scaled_add_rows_token_range_into, silu_mul_batch, silu_mul_batch_into, + single_prefill_nhd_noncausal_into, write_vec_into, }; #[cfg(not(feature = "kernel-call-trace"))] pub use openinfer_kernels::ops::{ @@ -34,8 +36,8 @@ pub use openinfer_kernels::ops::{ pub use paged_plan::PrefillPagedPlan; pub use sampling::{ argmax, argmax_batch_bf16_into, argmax_batch_bf16_split_indexed_into, - argmax_batch_bf16_split_partials_len, flashinfer_topk_row_states_bytes, gpu_sample, - gpu_sample_into, select_batch_tokens_into, + argmax_batch_bf16_split_into, argmax_batch_bf16_split_partials_len, + flashinfer_topk_row_states_bytes, gpu_sample, gpu_sample_into, select_batch_tokens_into, }; #[cfg(feature = "kernel-call-trace")] pub use traced::{ diff --git a/openinfer-core/src/ops/sampling.rs b/openinfer-core/src/ops/sampling.rs index 73f85e64..0b261000 100644 --- a/openinfer-core/src/ops/sampling.rs +++ b/openinfer-core/src/ops/sampling.rs @@ -6,7 +6,8 @@ use crate::tensor::{DeviceContext, DeviceVec, HiddenStates}; pub use openinfer_kernels::ops::{ argmax, argmax_batch_bf16_into, argmax_batch_bf16_split_indexed_into, - argmax_batch_bf16_split_partials_len, flashinfer_topk_row_states_bytes, + argmax_batch_bf16_split_into, argmax_batch_bf16_split_partials_len, + flashinfer_topk_row_states_bytes, }; /// GPU sampling: temperature -> softmax -> top-k -> top-p -> multinomial. @@ -62,11 +63,12 @@ pub fn gpu_sample_into( ) } -/// Pick the next token for each row in a decode batch. +/// Pick the next token for each real row in a decode batch. /// -/// Greedy rows are selected together with indexed batched argmax. Non-greedy -/// rows still use the existing per-row sampler because each row may have its -/// own random value and sampling parameters. +/// `logits` may contain CUDA Graph padding rows after the real batch. Greedy +/// rows are selected together with batched argmax. Non-greedy rows still use +/// the existing per-row sampler because each row may have its own random value +/// and sampling parameters. #[allow(clippy::too_many_arguments)] pub fn select_batch_tokens_into( ctx: &DeviceContext, @@ -83,7 +85,44 @@ pub fn select_batch_tokens_into( out: &mut CudaSlice, ) -> Result> { let batch_size = params.len(); + if batch_size > logits.seq_len { + return Err(anyhow!( + "sampling params len {} exceeds logits seq_len {}", + batch_size, + logits.seq_len + )); + } + if random_vals.len() < batch_size { + return Err(anyhow!( + "sampling random_vals len {} is smaller than params len {}", + random_vals.len(), + batch_size + )); + } let mut tokens = vec![0; batch_size]; + if batch_size == 0 { + return Ok(tokens); + } + if params.iter().all(|params_i| params_i.is_greedy()) { + argmax_batch_bf16_split_into( + ctx, + logits, + argmax_partial_values_scratch, + argmax_partial_indices_scratch, + top1_value_scratch, + out, + )?; + let out_host = ctx + .stream + .clone_dtoh(out) + .map_err(|e| anyhow!("D2H batch argmax read failed: {}", e))?; + ctx.sync()?; + for (token, sampled) in tokens.iter_mut().zip(out_host) { + *token = sampled as u32; + } + return Ok(tokens); + } + let greedy_rows = params .iter() .enumerate() @@ -147,3 +186,148 @@ pub fn select_batch_tokens_into( Ok(tokens) } + +#[cfg(test)] +mod tests { + use half::bf16; + + use super::*; + + fn logits_with_argmax_rows( + ctx: &DeviceContext, + hidden_dim: usize, + max_tokens: &[usize], + tied_tokens: &[(usize, usize)], + ) -> HiddenStates { + let seq_len = max_tokens.len(); + let mut host = vec![bf16::from_f32(-1.0); seq_len * hidden_dim]; + for (row, &token) in max_tokens.iter().enumerate() { + host[row * hidden_dim + token] = bf16::from_f32(10.0 + row as f32); + } + for &(row, token) in tied_tokens { + host[row * hidden_dim + token] = host[row * hidden_dim + max_tokens[row]]; + } + HiddenStates { + data: ctx.stream.clone_htod(&host).expect("copy logits"), + hidden_dim, + seq_len, + } + } + + struct SelectScratch { + row_indices: CudaSlice, + argmax_values: CudaSlice, + argmax_indices: CudaSlice, + probs: CudaSlice, + top1_values: CudaSlice, + row_states: CudaSlice, + valid: CudaSlice, + out: CudaSlice, + } + + impl SelectScratch { + fn new(ctx: &DeviceContext, rows: usize, vocab: usize) -> Self { + let partials = argmax_batch_bf16_split_partials_len(rows, vocab); + Self { + row_indices: ctx.stream.alloc_zeros(rows).expect("row indices"), + argmax_values: ctx.stream.alloc_zeros(partials).expect("argmax values"), + argmax_indices: ctx.stream.alloc_zeros(partials).expect("argmax indices"), + probs: ctx.stream.alloc_zeros(vocab).expect("probs"), + top1_values: ctx.stream.alloc_zeros(rows).expect("top1 values"), + row_states: ctx + .stream + .alloc_zeros(flashinfer_topk_row_states_bytes()) + .expect("row states"), + valid: ctx.stream.alloc_zeros(1).expect("valid"), + out: ctx.stream.alloc_zeros(rows).expect("out"), + } + } + } + + #[test] + fn all_greedy_contiguous_argmax_matches_indexed_greedy_rows() { + let ctx = DeviceContext::new().expect("create CUDA context"); + let logits = logits_with_argmax_rows(&ctx, 8193, &[17, 5001, 8192, 4095], &[(3, 4096)]); + let greedy = SamplingParams::default(); + let mut sampled = SamplingParams::default(); + sampled.temperature = 1.0; + let random_vals = [0.0; 4]; + + let mut contiguous = SelectScratch::new(&ctx, logits.seq_len, logits.hidden_dim); + let all_greedy = [&greedy, &greedy, &greedy, &greedy]; + let contiguous_tokens = select_batch_tokens_into( + &ctx, + &logits, + &all_greedy, + &random_vals, + &mut contiguous.row_indices, + &mut contiguous.argmax_values, + &mut contiguous.argmax_indices, + &mut contiguous.probs, + &mut contiguous.top1_values, + &mut contiguous.row_states, + &mut contiguous.valid, + &mut contiguous.out, + ) + .expect("contiguous greedy select"); + + let mut indexed = SelectScratch::new(&ctx, logits.seq_len, logits.hidden_dim); + let mixed = [&greedy, &sampled, &greedy, &greedy]; + let indexed_tokens = select_batch_tokens_into( + &ctx, + &logits, + &mixed, + &random_vals, + &mut indexed.row_indices, + &mut indexed.argmax_values, + &mut indexed.argmax_indices, + &mut indexed.probs, + &mut indexed.top1_values, + &mut indexed.row_states, + &mut indexed.valid, + &mut indexed.out, + ) + .expect("mixed select"); + + for row in [0usize, 2, 3] { + assert_eq!( + contiguous_tokens[row], indexed_tokens[row], + "greedy row {row} differs between contiguous and indexed argmax" + ); + } + assert_eq!(contiguous_tokens, vec![17, 5001, 8192, 4095]); + } + + #[test] + fn all_greedy_contiguous_argmax_ignores_cuda_graph_padding_rows() { + let ctx = DeviceContext::new().expect("create CUDA context"); + let logits = logits_with_argmax_rows( + &ctx, + 8193, + &[17, 5001, 8192, 4095, 7, 123, 456, 789], + &[(3, 4096)], + ); + let greedy = SamplingParams::default(); + let real_rows = [&greedy, &greedy, &greedy, &greedy, &greedy]; + let random_vals = [0.0; 5]; + let mut scratch = SelectScratch::new(&ctx, logits.seq_len, logits.hidden_dim); + + let tokens = select_batch_tokens_into( + &ctx, + &logits, + &real_rows, + &random_vals, + &mut scratch.row_indices, + &mut scratch.argmax_values, + &mut scratch.argmax_indices, + &mut scratch.probs, + &mut scratch.top1_values, + &mut scratch.row_states, + &mut scratch.valid, + &mut scratch.out, + ) + .expect("padded greedy select"); + + assert_eq!(tokens, vec![17, 5001, 8192, 4095, 7]); + } +} diff --git a/openinfer-kernels/csrc/shared/elementwise.cu b/openinfer-kernels/csrc/shared/elementwise.cu index 92de04eb..7193b145 100644 --- a/openinfer-kernels/csrc/shared/elementwise.cu +++ b/openinfer-kernels/csrc/shared/elementwise.cu @@ -63,6 +63,43 @@ __global__ void gather_hidden_tokens_kernel( } } +__global__ void copy_hidden_rows_kernel( + const __nv_bfloat16 *__restrict__ src, + __nv_bfloat16 *__restrict__ dst, + int src_hidden_dim, + int dst_hidden_dim, + int row_offset, + int rows, + int seq_len) { + int total = rows * seq_len; + for (int idx = blockIdx.x * blockDim.x + threadIdx.x; + idx < total; + idx += gridDim.x * blockDim.x) { + int token = idx / rows; + int row = idx % rows; + dst[(size_t)token * dst_hidden_dim + row_offset + row] = + src[(size_t)token * src_hidden_dim + row]; + } +} + +__global__ void copy_hidden_token_range_kernel( + const __nv_bfloat16 *__restrict__ src, + __nv_bfloat16 *__restrict__ dst, + int hidden_dim, + int src_token_offset, + int dst_token_offset, + int token_count) { + int total = hidden_dim * token_count; + for (int idx = blockIdx.x * blockDim.x + threadIdx.x; + idx < total; + idx += gridDim.x * blockDim.x) { + int token = idx / hidden_dim; + int row = idx % hidden_dim; + dst[(size_t)(dst_token_offset + token) * hidden_dim + row] = + src[(size_t)(src_token_offset + token) * hidden_dim + row]; + } +} + __global__ void scaled_add_rows_indexed_kernel( const __nv_bfloat16 *__restrict__ delta, float scale, @@ -305,6 +342,53 @@ CUresult gather_hidden_tokens_cuda( return (CUresult)cudaGetLastError(); } +CUresult copy_hidden_rows_cuda( + const __nv_bfloat16 *src, + __nv_bfloat16 *dst, + int src_hidden_dim, + int dst_hidden_dim, + int row_offset, + int rows, + int seq_len, + cudaStream_t stream) { + if (src == nullptr || dst == nullptr || src_hidden_dim <= 0 || + dst_hidden_dim <= 0 || row_offset < 0 || rows <= 0 || seq_len <= 0 || + rows > src_hidden_dim || row_offset + rows > dst_hidden_dim) { + return CUDA_ERROR_INVALID_VALUE; + } + int total = rows * seq_len; + int block = 256; + int grid = (total + block - 1) / block; + copy_hidden_rows_kernel<<>>( + src, dst, src_hidden_dim, dst_hidden_dim, row_offset, rows, seq_len); + return (CUresult)cudaGetLastError(); +} + +CUresult copy_hidden_token_range_cuda( + const __nv_bfloat16 *src, + __nv_bfloat16 *dst, + int hidden_dim, + int src_token_offset, + int dst_token_offset, + int token_count, + int src_seq_len, + int dst_seq_len, + cudaStream_t stream) { + if (src == nullptr || dst == nullptr || hidden_dim <= 0 || + src_token_offset < 0 || dst_token_offset < 0 || token_count <= 0 || + src_seq_len <= 0 || dst_seq_len <= 0 || + src_token_offset + token_count > src_seq_len || + dst_token_offset + token_count > dst_seq_len) { + return CUDA_ERROR_INVALID_VALUE; + } + int total = hidden_dim * token_count; + int block = 256; + int grid = (total + block - 1) / block; + copy_hidden_token_range_kernel<<>>( + src, dst, hidden_dim, src_token_offset, dst_token_offset, token_count); + return (CUresult)cudaGetLastError(); +} + CUresult scaled_add_rows_indexed_cuda( const __nv_bfloat16 *delta, float scale, diff --git a/openinfer-kernels/csrc/shared/paged_attention.cu b/openinfer-kernels/csrc/shared/paged_attention.cu index 4506a60f..01f4e82c 100644 --- a/openinfer-kernels/csrc/shared/paged_attention.cu +++ b/openinfer-kernels/csrc/shared/paged_attention.cu @@ -607,6 +607,70 @@ int single_prefill_cuda( reinterpret_cast(stream))); } +int single_prefill_nhd_noncausal_cuda( + // Q and output (HiddenStates token-major: [seq_len, q_dim]) + void* q, + void* output, + // Contiguous KV cache (HiddenStates token-major: [max_seq_len, kv_dim]) + void* k_cache, + void* v_cache, + int32_t num_qo_heads, + int32_t num_kv_heads, + int32_t head_dim, + int32_t seq_len, + int32_t kv_len, + int32_t max_seq_len, + float sm_scale, + void* stream) +{ + if (q == nullptr || output == nullptr || k_cache == nullptr || v_cache == nullptr || + num_qo_heads <= 0 || num_kv_heads <= 0 || head_dim != 128 || + seq_len <= 0 || kv_len <= 0 || max_seq_len < kv_len) { + return static_cast(cudaErrorInvalidValue); + } + + uint32_t q_stride_n = num_qo_heads * head_dim; + uint32_t q_stride_h = head_dim; + uint32_t kv_stride_n = num_kv_heads * head_dim; + uint32_t kv_stride_h = head_dim; + + PrefillParamsT params( + reinterpret_cast(q), + reinterpret_cast(k_cache), + reinterpret_cast(v_cache), + /*maybe_custom_mask=*/nullptr, + reinterpret_cast(output), + /*lse=*/nullptr, + /*maybe_alibi_slopes=*/nullptr, + num_qo_heads, + num_kv_heads, + static_cast(seq_len), + static_cast(kv_len), + q_stride_n, + q_stride_h, + kv_stride_n, + kv_stride_h, + static_cast(head_dim), + /*window_left=*/-1, + /*logits_soft_cap=*/0.0f, + sm_scale, + /*rope_scale=*/1.0f, + /*rope_theta=*/1e6f); + + return static_cast( + SinglePrefillWithKVCacheDispatched< + /*HEAD_DIM_QK=*/128, + /*HEAD_DIM_VO=*/128, + PosEncodingMode::kNone, + /*USE_FP16_QK_REDUCTION=*/false, + MaskMode::kNone, + Variant, + PrefillParamsT>( + params, + /*tmp=*/nullptr, + reinterpret_cast(stream))); +} + // --------------------------------------------------------------------------- // Single-request prefill for HEAD_DIM=256 — wraps FlashInfer SinglePrefillWithKVCache. // diff --git a/openinfer-kernels/csrc/shared/prefill_attention.cu b/openinfer-kernels/csrc/shared/prefill_attention.cu index a7b24b66..0c69522b 100644 --- a/openinfer-kernels/csrc/shared/prefill_attention.cu +++ b/openinfer-kernels/csrc/shared/prefill_attention.cu @@ -93,6 +93,90 @@ __global__ void prefill_qk_norm_rope_kernel( data[offset] = result; } +__global__ void dflash_qk_norm_rope_kernel( + __nv_bfloat16* __restrict__ q, + __nv_bfloat16* __restrict__ k, + const __nv_bfloat16* __restrict__ q_norm_weight, + const __nv_bfloat16* __restrict__ k_norm_weight, + const __nv_bfloat16* __restrict__ cos_cache, + const __nv_bfloat16* __restrict__ sin_cache, + int num_q_heads, + int num_kv_heads, + int head_dim, + int q_len, + int k_len, + int q_start_pos, + int k_start_pos, + float eps, + int cos_max_pos +) { + int head_global = blockIdx.x; + int token = blockIdx.y; + int d = threadIdx.x; + + bool is_q = (head_global < num_q_heads); + int local_heads = is_q ? num_q_heads : num_kv_heads; + int seq_len = is_q ? q_len : k_len; + if (token >= seq_len) return; + + int head_local = is_q ? head_global : (head_global - num_q_heads); + if (head_local >= local_heads) return; + + __nv_bfloat16* data = is_q ? q : k; + int dim_stride = local_heads * head_dim; + const __nv_bfloat16* norm_w = is_q ? q_norm_weight : k_norm_weight; + int pos = (is_q ? q_start_pos : k_start_pos) + token; + if (pos < 0 || pos >= cos_max_pos) __trap(); + + int offset = token * dim_stride + head_local * head_dim + d; + float val = __bfloat162float(data[offset]); + + float sq = warp_reduce_sum(val * val); + int warp_id = d / WARP_SIZE; + int lane_id = d % WARP_SIZE; + __shared__ float warp_sums[4]; + if (lane_id == 0) warp_sums[warp_id] = sq; + __syncthreads(); + + __shared__ float s_inv_rms; + if (warp_id == 0) { + float v = (lane_id < 4) ? warp_sums[lane_id] : 0.0f; + float total = warp_reduce_sum(v); + if (lane_id == 0) s_inv_rms = rsqrtf(total / head_dim + eps); + } + __syncthreads(); + + __nv_bfloat16 normed = __float2bfloat16(val * s_inv_rms); + float normed_f = __bfloat162float(normed) * __bfloat162float(norm_w[d]); + + __shared__ __nv_bfloat16 smem[HEAD_DIM]; + smem[d] = __float2bfloat16(normed_f); + __syncthreads(); + + int half = head_dim / 2; + __nv_bfloat16 result; + if (d < half) { + float lo = __bfloat162float(smem[d]); + float hi = __bfloat162float(smem[d + half]); + float c = __bfloat162float(cos_cache[pos * head_dim + d]); + float s = __bfloat162float(sin_cache[pos * head_dim + d]); + float lo_cos = __bfloat162float(__float2bfloat16(lo * c)); + float hi_sin = __bfloat162float(__float2bfloat16(hi * s)); + result = __float2bfloat16(lo_cos - hi_sin); + } else { + int pair_d = d - half; + float lo = __bfloat162float(smem[pair_d]); + float hi = __bfloat162float(smem[d]); + float c = __bfloat162float(cos_cache[pos * head_dim + pair_d]); + float s = __bfloat162float(sin_cache[pos * head_dim + pair_d]); + float lo_sin = __bfloat162float(__float2bfloat16(lo * s)); + float hi_cos = __bfloat162float(__float2bfloat16(hi * c)); + result = __float2bfloat16(lo_sin + hi_cos); + } + + data[offset] = result; +} + extern "C" { // ============================================================================ @@ -136,4 +220,38 @@ void qk_norm_rope_batched_decode_cuda( ); } +int dflash_qk_norm_rope_cuda( + __nv_bfloat16* q, + __nv_bfloat16* k, + const __nv_bfloat16* q_norm_weight, + const __nv_bfloat16* k_norm_weight, + const __nv_bfloat16* cos_cache, + const __nv_bfloat16* sin_cache, + int num_q_heads, + int num_kv_heads, + int head_dim, + int q_len, + int k_len, + int q_start_pos, + int k_start_pos, + float rms_eps, + int cos_max_pos, + cudaStream_t stream +) { + if (q == nullptr || k == nullptr || q_norm_weight == nullptr || + k_norm_weight == nullptr || cos_cache == nullptr || sin_cache == nullptr || + num_q_heads <= 0 || num_kv_heads <= 0 || head_dim != HEAD_DIM || + q_len <= 0 || k_len <= 0 || q_start_pos < 0 || k_start_pos < 0 || + q_start_pos + q_len > cos_max_pos || k_start_pos + k_len > cos_max_pos) { + return static_cast(cudaErrorInvalidValue); + } + + dim3 grid(num_q_heads + num_kv_heads, q_len > k_len ? q_len : k_len); + dflash_qk_norm_rope_kernel<<>>( + q, k, q_norm_weight, k_norm_weight, cos_cache, sin_cache, + num_q_heads, num_kv_heads, head_dim, q_len, k_len, + q_start_pos, k_start_pos, rms_eps, cos_max_pos); + return static_cast(cudaGetLastError()); +} + } // extern "C" diff --git a/openinfer-kernels/src/ffi/shared.rs b/openinfer-kernels/src/ffi/shared.rs index 03748e9b..1c43b54c 100644 --- a/openinfer-kernels/src/ffi/shared.rs +++ b/openinfer-kernels/src/ffi/shared.rs @@ -30,6 +30,29 @@ unsafe extern "C" { stream: CUstream, ) -> CUresult; + pub fn copy_hidden_rows_cuda( + src: *const Half, + dst: *mut Half, + src_hidden_dim: i32, + dst_hidden_dim: i32, + row_offset: i32, + rows: i32, + seq_len: i32, + stream: CUstream, + ) -> CUresult; + + pub fn copy_hidden_token_range_cuda( + src: *const Half, + dst: *mut Half, + hidden_dim: i32, + src_token_offset: i32, + dst_token_offset: i32, + token_count: i32, + src_seq_len: i32, + dst_seq_len: i32, + stream: CUstream, + ) -> CUresult; + pub fn fused_add_rms_norm_cuda( hidden: *mut Half, residual: *const Half, @@ -240,6 +263,25 @@ unsafe extern "C" { stream: CUstream, ); + pub fn dflash_qk_norm_rope_cuda( + q: *mut Half, + k: *mut Half, + q_norm_weight: *const Half, + k_norm_weight: *const Half, + cos_cache: *const Half, + sin_cache: *const Half, + num_q_heads: i32, + num_kv_heads: i32, + head_dim: i32, + q_len: i32, + k_len: i32, + q_start_pos: i32, + k_start_pos: i32, + rms_eps: f32, + cos_max_pos: i32, + stream: CUstream, + ) -> i32; + // Scatter contiguous KV → paged layout (one layer, FlashInfer prefill append). pub fn paged_kv_scatter_cuda( kv_data: *const Half, @@ -366,6 +408,21 @@ unsafe extern "C" { stream: CUstream, ) -> i32; + pub fn single_prefill_nhd_noncausal_cuda( + q: *const Half, + output: *mut Half, + k_cache: *const Half, + v_cache: *const Half, + num_qo_heads: i32, + num_kv_heads: i32, + head_dim: i32, + seq_len: i32, + kv_len: i32, + max_seq_len: i32, + sm_scale: f32, + stream: CUstream, + ) -> i32; + // Paged attention decode (FlashInfer BatchDecode, no partition-KV). pub fn paged_attention_decode_cuda( q: *const Half, diff --git a/openinfer-kernels/src/ops.rs b/openinfer-kernels/src/ops.rs index 1ecf0fe1..5a57ecd7 100644 --- a/openinfer-kernels/src/ops.rs +++ b/openinfer-kernels/src/ops.rs @@ -13,9 +13,10 @@ mod norm; mod sampling; pub use attention::{ - PrefillPagedPlan, paged_attention_batch_decode_hd256_into, paged_attention_batch_decode_into, - paged_attention_batch_decode_split_kv_into, prefill_attention_paged_into, - qk_norm_partial_rope_batched_decode_hd256_into, qk_norm_rope_batch_decode_into, + PrefillPagedPlan, dflash_qk_norm_rope_into, paged_attention_batch_decode_hd256_into, + paged_attention_batch_decode_into, paged_attention_batch_decode_split_kv_into, + prefill_attention_paged_into, qk_norm_partial_rope_batched_decode_hd256_into, + qk_norm_rope_batch_decode_into, single_prefill_nhd_noncausal_into, }; #[cfg(feature = "kimi-k2")] pub use deepep::{ @@ -23,11 +24,11 @@ pub use deepep::{ }; pub use elementwise::{ accumulate_bf16_token_scaled_to_f32_into, add_batch, add_batch_into, bf16_hidden_to_f32_into, - extract_vec, extract_vec_into, extract_vec_ref, extract_vec_ref_into, f32_to_bf16_hidden_into, - gather_hidden_tokens_into, repeat_f32_for_reduce_scatter_into, scale_f32_in_place, - scaled_add_batch_into, scaled_add_rows_indexed_into, scaled_add_rows_into, - scaled_add_rows_token_range_into, silu_mul_batch, silu_mul_batch_into, - silu_mul_fused_batch_into, write_vec_into, + copy_hidden_rows_into, copy_hidden_token_range_into, extract_vec, extract_vec_into, + extract_vec_ref, extract_vec_ref_into, f32_to_bf16_hidden_into, gather_hidden_tokens_into, + repeat_f32_for_reduce_scatter_into, scale_f32_in_place, scaled_add_batch_into, + scaled_add_rows_indexed_into, scaled_add_rows_into, scaled_add_rows_token_range_into, + silu_mul_batch, silu_mul_batch_into, silu_mul_fused_batch_into, write_vec_into, }; pub use embedding::{embedding_batch, embedding_batch_vocab_shard, embedding_decode_into}; #[cfg(feature = "kimi-k2")] @@ -48,7 +49,7 @@ pub use norm::{ }; pub use sampling::{ BatchSamplingRow, BatchSamplingScratch, argmax, argmax_batch_bf16_into, - argmax_batch_bf16_split_indexed_into, argmax_batch_bf16_split_partials_len, - flashinfer_top1_batch_into, flashinfer_topk_row_states_bytes, gpu_sample, - gpu_sample_batch_into, gpu_sample_into, + argmax_batch_bf16_split_indexed_into, argmax_batch_bf16_split_into, + argmax_batch_bf16_split_partials_len, flashinfer_top1_batch_into, + flashinfer_topk_row_states_bytes, gpu_sample, gpu_sample_batch_into, gpu_sample_into, }; diff --git a/openinfer-kernels/src/ops/attention.rs b/openinfer-kernels/src/ops/attention.rs index e3ae955a..53a5a02e 100644 --- a/openinfer-kernels/src/ops/attention.rs +++ b/openinfer-kernels/src/ops/attention.rs @@ -497,6 +497,106 @@ pub fn qk_norm_rope_batch_decode_into( } } +#[allow(clippy::too_many_arguments)] +pub fn dflash_qk_norm_rope_into( + ctx: &DeviceContext, + q: &mut HiddenStates, + k: &mut HiddenStates, + q_norm_weight: &DeviceVec, + k_norm_weight: &DeviceVec, + cos_cache: &DeviceVec, + sin_cache: &DeviceVec, + num_q_heads: usize, + num_kv_heads: usize, + head_dim: usize, + q_start_pos: usize, + k_start_pos: usize, + rms_eps: f32, +) -> Result<()> { + assert_eq!(q.hidden_dim, num_q_heads * head_dim); + assert_eq!(k.hidden_dim, num_kv_heads * head_dim); + assert_eq!(q_norm_weight.len, head_dim); + assert_eq!(k_norm_weight.len, head_dim); + + let (q_ptr, _gq) = q.data.device_ptr_mut(&ctx.stream); + let (k_ptr, _gk) = k.data.device_ptr_mut(&ctx.stream); + let (qn_ptr, _gqn) = q_norm_weight.data.device_ptr(&ctx.stream); + let (kn_ptr, _gkn) = k_norm_weight.data.device_ptr(&ctx.stream); + let (cos_ptr, _gc) = cos_cache.data.device_ptr(&ctx.stream); + let (sin_ptr, _gs) = sin_cache.data.device_ptr(&ctx.stream); + + let result = unsafe { + ffi::dflash_qk_norm_rope_cuda( + q_ptr as *mut ffi::Half, + k_ptr as *mut ffi::Half, + qn_ptr as *const ffi::Half, + kn_ptr as *const ffi::Half, + cos_ptr as *const ffi::Half, + sin_ptr as *const ffi::Half, + num_q_heads as i32, + num_kv_heads as i32, + head_dim as i32, + q.seq_len as i32, + k.seq_len as i32, + q_start_pos as i32, + k_start_pos as i32, + rms_eps, + (cos_cache.data.len() / head_dim) as i32, + ctx.stream.cu_stream(), + ) + }; + if result != 0 { + anyhow::bail!("dflash_qk_norm_rope_cuda failed with error {result}"); + } + Ok(()) +} + +#[allow(clippy::too_many_arguments)] +pub fn single_prefill_nhd_noncausal_into( + ctx: &DeviceContext, + q: &HiddenStates, + k_cache: &HiddenStates, + v_cache: &HiddenStates, + output: &mut HiddenStates, + num_q_heads: usize, + num_kv_heads: usize, + head_dim: usize, + kv_len: usize, +) -> Result<()> { + assert_eq!(q.hidden_dim, num_q_heads * head_dim); + assert_eq!(output.hidden_dim, q.hidden_dim); + assert_eq!(output.seq_len, q.seq_len); + assert_eq!(k_cache.hidden_dim, num_kv_heads * head_dim); + assert_eq!(v_cache.hidden_dim, k_cache.hidden_dim); + assert_eq!(v_cache.seq_len, k_cache.seq_len); + assert!(kv_len <= k_cache.seq_len); + + let (q_ptr, _gq) = q.data.device_ptr(&ctx.stream); + let (k_ptr, _gk) = k_cache.data.device_ptr(&ctx.stream); + let (v_ptr, _gv) = v_cache.data.device_ptr(&ctx.stream); + let (out_ptr, _go) = output.data.device_ptr_mut(&ctx.stream); + let result = unsafe { + ffi::single_prefill_nhd_noncausal_cuda( + q_ptr as *const ffi::Half, + out_ptr as *mut ffi::Half, + k_ptr as *const ffi::Half, + v_ptr as *const ffi::Half, + num_q_heads as i32, + num_kv_heads as i32, + head_dim as i32, + q.seq_len as i32, + kv_len as i32, + k_cache.seq_len as i32, + 1.0f32 / (head_dim as f32).sqrt(), + ctx.stream.cu_stream(), + ) + }; + if result != 0 { + anyhow::bail!("single_prefill_nhd_noncausal_cuda failed with error {result}"); + } + Ok(()) +} + /// Batched QK RMSNorm + partial RoPE for Qwen3.5 HD256 decode. /// /// Reads Q from interleaved `q_full` ([q, gate] per head), writes prepared Q into `q`, diff --git a/openinfer-kernels/src/ops/elementwise.rs b/openinfer-kernels/src/ops/elementwise.rs index 1e651e97..7564f439 100644 --- a/openinfer-kernels/src/ops/elementwise.rs +++ b/openinfer-kernels/src/ops/elementwise.rs @@ -176,6 +176,90 @@ pub fn gather_hidden_tokens_into( Ok(()) } +pub fn copy_hidden_rows_into( + ctx: &DeviceContext, + src: &HiddenStates, + dst: &mut HiddenStates, + row_offset: usize, +) -> Result<()> { + assert!( + row_offset + src.hidden_dim <= dst.hidden_dim, + "row range [{}..{}) exceeds destination hidden_dim {}", + row_offset, + row_offset + src.hidden_dim, + dst.hidden_dim + ); + assert_eq!( + src.seq_len, dst.seq_len, + "copy_hidden_rows_into seq_len mismatch: src {}, dst {}", + src.seq_len, dst.seq_len + ); + + let (src_ptr, _gs) = src.data.device_ptr(&ctx.stream); + let (dst_ptr, _gd) = dst.data.device_ptr_mut(&ctx.stream); + let result = unsafe { + ffi::copy_hidden_rows_cuda( + src_ptr as *const ffi::Half, + dst_ptr as *mut ffi::Half, + src.hidden_dim as i32, + dst.hidden_dim as i32, + row_offset as i32, + src.hidden_dim as i32, + src.seq_len as i32, + ctx.stream.cu_stream(), + ) + }; + result.result()?; + Ok(()) +} + +pub fn copy_hidden_token_range_into( + ctx: &DeviceContext, + src: &HiddenStates, + src_token_offset: usize, + dst: &mut HiddenStates, + dst_token_offset: usize, + token_count: usize, +) -> Result<()> { + assert_eq!( + src.hidden_dim, dst.hidden_dim, + "copy_hidden_token_range_into hidden_dim mismatch: src {}, dst {}", + src.hidden_dim, dst.hidden_dim + ); + assert!( + src_token_offset + token_count <= src.seq_len, + "source token range [{}..{}) exceeds seq_len {}", + src_token_offset, + src_token_offset + token_count, + src.seq_len + ); + assert!( + dst_token_offset + token_count <= dst.seq_len, + "destination token range [{}..{}) exceeds seq_len {}", + dst_token_offset, + dst_token_offset + token_count, + dst.seq_len + ); + + let (src_ptr, _gs) = src.data.device_ptr(&ctx.stream); + let (dst_ptr, _gd) = dst.data.device_ptr_mut(&ctx.stream); + let result = unsafe { + ffi::copy_hidden_token_range_cuda( + src_ptr as *const ffi::Half, + dst_ptr as *mut ffi::Half, + src.hidden_dim as i32, + src_token_offset as i32, + dst_token_offset as i32, + token_count as i32, + src.seq_len as i32, + dst.seq_len as i32, + ctx.stream.cu_stream(), + ) + }; + result.result()?; + Ok(()) +} + pub fn scaled_add_rows_indexed_into( ctx: &DeviceContext, delta: &HiddenStates, diff --git a/openinfer-kernels/src/ops/sampling.rs b/openinfer-kernels/src/ops/sampling.rs index 23bb6a09..4dd5b2df 100644 --- a/openinfer-kernels/src/ops/sampling.rs +++ b/openinfer-kernels/src/ops/sampling.rs @@ -278,6 +278,60 @@ pub fn argmax_batch_bf16_split_partials_len(rows: usize, vocab: usize) -> usize rows * vocab.div_ceil(TILE_ELEMS) } +/// Two-stage batched argmax for contiguous rows: tile-parallel partials then +/// one finalize per row. Lowest index wins ties. +pub fn argmax_batch_bf16_split_into( + ctx: &DeviceContext, + logits: &HiddenStates, + partial_values: &mut CudaSlice, + partial_indices: &mut CudaSlice, + values: &mut CudaSlice, + out: &mut CudaSlice, +) -> Result<()> { + let rows = logits.seq_len; + if rows == 0 { + return Err(anyhow!("argmax split batch requires at least one row")); + } + let needed_partials = argmax_batch_bf16_split_partials_len(rows, logits.hidden_dim); + if partial_values.len() < needed_partials || partial_indices.len() < needed_partials { + return Err(anyhow!( + "argmax split partials scratch too small: have {}/{}, need {}", + partial_values.len(), + partial_indices.len(), + needed_partials + )); + } + if values.len() < rows || out.len() < rows { + return Err(anyhow!( + "argmax split outputs too small: have {}/{}, need {}", + values.len(), + out.len(), + rows + )); + } + + let (logits_ptr, _gl) = logits.data.device_ptr(&ctx.stream); + let (pv_ptr, _gpv) = partial_values.device_ptr_mut(&ctx.stream); + let (pi_ptr, _gpi) = partial_indices.device_ptr_mut(&ctx.stream); + let (values_ptr, _gv) = values.device_ptr_mut(&ctx.stream); + let (out_ptr, _go) = out.device_ptr_mut(&ctx.stream); + + unsafe { + ffi::argmax_batch_bf16_split_cuda( + logits_ptr as *const ffi::Half, + values_ptr as *mut ffi::Half, + out_ptr as *mut i32, + pv_ptr as *mut f32, + pi_ptr as *mut i32, + rows as i32, + logits.hidden_dim as i32, + ctx.stream.cu_stream(), + ); + } + + Ok(()) +} + /// Two-stage indexed batched argmax: tile-parallel partials then a per-row /// finalize. Lowest index wins ties; each vocab row spreads over many CTAs /// instead of one. diff --git a/openinfer-kv-cache/src/pool.rs b/openinfer-kv-cache/src/pool.rs index 77d08256..e2154683 100644 --- a/openinfer-kv-cache/src/pool.rs +++ b/openinfer-kv-cache/src/pool.rs @@ -262,7 +262,9 @@ impl LoadReservation { /// Per-request KV state wrapping `SchedulableSequence`. /// /// Lifecycle: `schedule_prefill → prefill_view/pages → forward → apply_prefill`, -/// then `schedule_decode → decode_view/pages → forward → apply_decode` in a loop. +/// then either `schedule_decode → decode_view/pages → forward → apply_decode` +/// or `schedule_speculative → speculative_view/pages → forward → +/// apply_speculative` in a loop. pub struct RequestKv { seq: SchedulableSequence<()>, } @@ -299,6 +301,15 @@ impl RequestKv { self.seq.schedule_decode(&pool.block_manager) } + pub fn schedule_speculative( + &mut self, + num_draft_tokens: usize, + pool: &BlockPool, + ) -> Result<(), ScheduleError> { + self.seq + .schedule_speculative(num_draft_tokens, &pool.block_manager) + } + // ── Views (for forward pass) ─────────────────────────────────────── /// Build an immutable `KvView` for prefill. @@ -327,6 +338,20 @@ impl RequestKv { ) } + /// Build an immutable `KvView` for speculative verification. + /// + /// The verifier forwards `num_draft_tokens` consecutive token positions: + /// the request's current dangling token followed by draft candidates. The + /// view covers the post-step KV extent; [`Self::apply_speculative`] later + /// commits only the accepted prefix and releases excess draft capacity. + pub fn speculative_view(&self, num_draft_tokens: usize) -> KvView { + KvView::new( + self.step_page_indices(num_draft_tokens), + self.seq.kv_position() + num_draft_tokens, + self.seq.block_size(), + ) + } + // ── Apply (register blocks, advance position) ────────────────────── pub fn apply_prefill(&mut self, token: u32, pool: &BlockPool) -> anyhow::Result<()> { @@ -350,6 +375,24 @@ impl RequestKv { .map_err(|e| anyhow::anyhow!("apply_decode: {e}")) } + pub fn apply_speculative( + &mut self, + accepted_tokens: &[u32], + pool: &BlockPool, + ) -> anyhow::Result { + self.seq + .apply_speculative(accepted_tokens, &pool.block_manager) + .map_err(|e| anyhow::anyhow!("apply_speculative: {e}")) + } + + /// Undo a scheduled-but-unapplied KV operation and return any blocks + /// reserved by that schedule to the pool. + pub fn revert_schedule(&mut self) -> anyhow::Result<()> { + self.seq + .revert_schedule() + .map_err(|e| anyhow::anyhow!("revert_schedule: {e}")) + } + pub fn release(&mut self) -> anyhow::Result<()> { self.seq .release() diff --git a/openinfer-kv-cache/tests/lifecycle.rs b/openinfer-kv-cache/tests/lifecycle.rs index 17ee5271..81f44817 100644 --- a/openinfer-kv-cache/tests/lifecycle.rs +++ b/openinfer-kv-cache/tests/lifecycle.rs @@ -215,6 +215,135 @@ fn decode_view_seq_len_invariant() { req.release().unwrap(); } +#[test] +fn speculative_view_covers_verify_span() { + let mgr = make_manager(10); + let mut req = mgr.pool().new_request(vec![1; 10], 8, None); + + req.schedule_prefill(10, mgr.pool()).unwrap(); + req.apply_prefill(42, mgr.pool()).unwrap(); + assert_eq!(req.kv_position(), 10); + + req.schedule_speculative(4, mgr.pool()).unwrap(); + let view = req.speculative_view(4); + assert_eq!(view.seq_len(), req.kv_position() + 4); + assert_eq!(view.num_pages(), 1); + assert_eq!(view.last_page_len(), 14); + + req.apply_speculative(&[100, 101], mgr.pool()).unwrap(); + assert_eq!( + req.kv_position(), + 12, + "accepted token count advances KV while preserving one dangling token" + ); + assert_eq!(req.generated_tokens(), 3); + + req.release().unwrap(); +} + +#[test] +fn speculative_partial_accept_releases_excess_capacity() { + let mgr = make_manager(5); // 4 usable blocks + let initial_avail = mgr.pool().available_blocks(); + + let mut req = mgr.pool().new_request(vec![1; 4], 64, None); + req.schedule_prefill(4, mgr.pool()).unwrap(); + req.apply_prefill(42, mgr.pool()).unwrap(); + + req.schedule_speculative(32, mgr.pool()).unwrap(); + let after_schedule = mgr.pool().available_blocks(); + assert!( + after_schedule < initial_avail, + "speculative scheduling should reserve draft capacity" + ); + + req.apply_speculative(&[100], mgr.pool()).unwrap(); + assert!( + mgr.pool().available_blocks() > after_schedule, + "partial accept should release excess draft capacity" + ); + assert_eq!(req.kv_position(), 5); + assert_eq!(req.generated_tokens(), 2); + + req.release().unwrap(); + assert_eq!(mgr.pool().available_blocks(), initial_avail); +} + +#[test] +fn speculative_revert_returns_reserved_blocks() { + let mgr = make_manager(5); // 4 usable blocks + let initial_avail = mgr.pool().available_blocks(); + + let mut req = mgr.pool().new_request(vec![1; 4], 64, None); + req.schedule_prefill(4, mgr.pool()).unwrap(); + req.apply_prefill(42, mgr.pool()).unwrap(); + let before_schedule = mgr.pool().available_blocks(); + + req.schedule_speculative(32, mgr.pool()).unwrap(); + assert!( + mgr.pool().available_blocks() < before_schedule, + "speculative scheduling should reserve draft capacity" + ); + + req.revert_schedule().unwrap(); + assert_eq!( + mgr.pool().available_blocks(), + before_schedule, + "reverting a speculative schedule should return reserved blocks" + ); + assert_eq!(req.kv_position(), 4); + assert_eq!(req.generated_tokens(), 1); + + req.schedule_decode(mgr.pool()).unwrap(); + req.apply_decode(100, mgr.pool()).unwrap(); + assert_eq!(req.kv_position(), 5); + + req.release().unwrap(); + assert_eq!(mgr.pool().available_blocks(), initial_avail); +} + +#[test] +fn speculative_partial_accept_keeps_cross_page_tail_visible() { + let mgr = make_manager(16); + let mut req = mgr.pool().new_request(vec![1; 10], 64, None); + + req.schedule_prefill(10, mgr.pool()).unwrap(); + req.apply_prefill(42, mgr.pool()).unwrap(); + + let accepted_lens = [1usize, 1, 2, 1, 2, 1, 4, 2, 2, 1, 4, 1, 1]; + let mut next_token = 100u32; + for accepted_len in accepted_lens { + req.schedule_speculative(16, mgr.pool()).unwrap(); + let view = req.speculative_view(16); + assert_eq!( + view.num_pages(), + view.seq_len().div_ceil(16), + "verify view must exactly cover seq_len={}", + view.seq_len() + ); + let accepted: Vec = (0..accepted_len) + .map(|_| { + let token = next_token; + next_token += 1; + token + }) + .collect(); + req.apply_speculative(&accepted, mgr.pool()).unwrap(); + } + + req.schedule_speculative(16, mgr.pool()).unwrap(); + let view = req.speculative_view(16); + assert_eq!(view.seq_len(), 49); + assert_eq!( + view.num_pages(), + 4, + "33 committed KV tokens + 16-token verify span must expose the fourth page" + ); + req.apply_speculative(&[next_token], mgr.pool()).unwrap(); + + req.release().unwrap(); +} + /// prefill_view seq_len == kv_position + prompt_len, page count covers /// that seq_len. #[test] diff --git a/openinfer-qwen3-4b/Cargo.toml b/openinfer-qwen3-4b/Cargo.toml index 1dfb795f..39a759fe 100644 --- a/openinfer-qwen3-4b/Cargo.toml +++ b/openinfer-qwen3-4b/Cargo.toml @@ -22,6 +22,7 @@ half = { workspace = true } hex = { workspace = true, optional = true } log = { workspace = true } memmap2 = { workspace = true } +nvtx = { workspace = true } rand = { workspace = true } safetensors = { workspace = true } serde = { workspace = true } diff --git a/openinfer-qwen3-4b/src/config.rs b/openinfer-qwen3-4b/src/config.rs index c02b140e..abffe4d7 100644 --- a/openinfer-qwen3-4b/src/config.rs +++ b/openinfer-qwen3-4b/src/config.rs @@ -38,6 +38,29 @@ pub(crate) struct Config { pub(crate) stop_token_ids: Vec, } +#[derive(Clone, Debug, Deserialize)] +pub(crate) struct DFlashConfig { + pub(crate) 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, + pub(crate) 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, + pub(crate) block_size: usize, + pub(crate) dflash_config: DFlashInnerConfig, +} + +#[derive(Clone, Debug, Deserialize)] +pub(crate) struct DFlashInnerConfig { + pub(crate) mask_token_id: u32, + pub(crate) target_layer_ids: Vec, +} + fn default_max_position_embeddings() -> usize { 40960 } @@ -117,6 +140,85 @@ impl Config { } } +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)?; + Ok(serde_json::from_str(&content)?) + } + + pub(crate) fn validate_for_target(&self, target: &Config) -> Result<()> { + anyhow::ensure!( + self.hidden_size == target.hidden_size, + "DFlash hidden_size {} does not match target {}", + self.hidden_size, + 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 + ); + anyhow::ensure!( + self.num_attention_heads == target.num_attention_heads + && self.num_key_value_heads == target.num_key_value_heads + && self.head_dim == target.head_dim, + "DFlash attention geometry does not match target" + ); + anyhow::ensure!( + self.vocab_size == target.vocab_size, + "DFlash vocab_size {} does not match target {}", + self.vocab_size, + target.vocab_size + ); + anyhow::ensure!( + self.rope_theta == target.rope_theta, + "DFlash rope_theta {} does not match target {}", + self.rope_theta, + target.rope_theta + ); + anyhow::ensure!( + self.max_position_embeddings >= target.max_position_embeddings, + "DFlash max_position_embeddings {} is smaller than target {}", + self.max_position_embeddings, + target.max_position_embeddings + ); + anyhow::ensure!( + self.block_size >= 2, + "DFlash block_size must be >= 2, got {}", + self.block_size + ); + anyhow::ensure!( + self.dflash_config.mask_token_id < target.vocab_size as u32, + "DFlash mask_token_id {} is outside target vocab_size {}", + self.dflash_config.mask_token_id, + target.vocab_size + ); + anyhow::ensure!( + self.dflash_config.target_layer_ids.len() == self.num_hidden_layers, + "DFlash target_layer_ids length {} does not match draft layers {}", + self.dflash_config.target_layer_ids.len(), + self.num_hidden_layers + ); + anyhow::ensure!( + self.dflash_config + .target_layer_ids + .iter() + .all(|&layer| layer < target.num_hidden_layers), + "DFlash target_layer_ids must be within target layer count" + ); + anyhow::ensure!( + self.dflash_config + .target_layer_ids + .windows(2) + .all(|pair| pair[0] < pair[1]), + "DFlash target_layer_ids must be strictly increasing" + ); + Ok(()) + } +} + impl TensorParallelConfig { pub(crate) fn validate_for(self, config: &Config) -> Result<()> { if self.world_size == 0 { diff --git a/openinfer-qwen3-4b/src/dflash.rs b/openinfer-qwen3-4b/src/dflash.rs new file mode 100644 index 00000000..df285d0a --- /dev/null +++ b/openinfer-qwen3-4b/src/dflash.rs @@ -0,0 +1,791 @@ +use anyhow::{Context, Result}; +use cudarc::driver::CudaSlice; +use log::debug; + +use crate::config::DFlashConfig; +use crate::weights::{Attention, MLP, Qwen3Model, TransformerBlock}; +use openinfer_core::ops; +use openinfer_core::tensor::HiddenStates; +use openinfer_core::tensor::{DeviceContext, DeviceMatrix, DeviceVec}; +use openinfer_core::weight_loader::{ + deserialize_shards, load_shard_info, load_tensor_1d, load_tensor_2d, mmap_shards, + precompute_rope, +}; + +pub(crate) struct DFlashDraftModel { + config: DFlashConfig, + layers: Vec, + norm: DeviceVec, + hidden_norm: DeviceVec, + fc: DeviceMatrix, + cos_cache: DeviceVec, + sin_cache: DeviceVec, +} + +pub(crate) struct DFlashRequestState { + layers: Vec, + pending_context: DFlashPendingContext, + scratch: DFlashDraftScratch, + committed_len: usize, + max_cache_len: usize, +} + +struct DFlashLayerCache { + k: HiddenStates, + v: HiddenStates, +} + +struct DFlashPendingContext { + buffer: HiddenStates, + len: usize, + capacity: usize, +} + +pub(crate) struct DFlashDraftOutput<'a> { + pub(crate) logits: &'a HiddenStates, + pub(crate) context_len: usize, + pub(crate) committed_len: usize, +} + +struct DFlashDraftScratch { + max_context_len: usize, + block_token_ids_h: Vec, + token_ids_d: CudaSlice, + context_projected: HiddenStates, + context_hidden: HiddenStates, + hidden: HiddenStates, + hidden_out: HiddenStates, + normed: HiddenStates, + tail_input: HiddenStates, + q_batch: HiddenStates, + k_tail: HiddenStates, + v_tail: HiddenStates, + attn_output: HiddenStates, + o_buf: HiddenStates, + gate_out: HiddenStates, + up_out: HiddenStates, + act_out: HiddenStates, + logits_normed: HiddenStates, + logits: HiddenStates, +} + +impl DFlashRequestState { + pub(crate) fn pending_context_len(&self) -> Option { + (self.pending_context.len > 0).then_some(self.pending_context.len) + } +} + +impl DFlashPendingContext { + fn new(ctx: &DeviceContext, hidden_dim: usize, capacity: usize) -> Result { + anyhow::ensure!( + capacity > 0, + "DFlash pending context capacity must be non-zero" + ); + let mut buffer = HiddenStates::zeros(ctx, hidden_dim, capacity)?; + buffer.seq_len = 0; + Ok(Self { + buffer, + len: 0, + capacity, + }) + } + + fn append_from( + &mut self, + ctx: &DeviceContext, + src: &HiddenStates, + src_token_offset: usize, + token_count: usize, + max_capacity: usize, + ) -> Result<()> { + let required_len = self + .len + .checked_add(token_count) + .context("DFlash pending context length overflow")?; + anyhow::ensure!( + required_len <= max_capacity, + "DFlash pending context length {} exceeds request capacity {}", + required_len, + max_capacity + ); + self.ensure_capacity(ctx, required_len, max_capacity)?; + self.buffer.seq_len = self.capacity; + ops::copy_hidden_token_range_into( + ctx, + src, + src_token_offset, + &mut self.buffer, + self.len, + token_count, + )?; + self.len = required_len; + self.buffer.seq_len = self.len; + Ok(()) + } + + fn ensure_capacity( + &mut self, + ctx: &DeviceContext, + required_len: usize, + max_capacity: usize, + ) -> Result<()> { + if required_len <= self.capacity { + return Ok(()); + } + let doubled = self + .capacity + .checked_mul(2) + .context("DFlash pending context capacity overflow")?; + let new_capacity = required_len.max(doubled).min(max_capacity); + anyhow::ensure!( + new_capacity >= required_len, + "DFlash pending context capacity {} cannot fit {} tokens", + new_capacity, + required_len + ); + let mut next = HiddenStates::zeros(ctx, self.buffer.hidden_dim, new_capacity)?; + if self.len > 0 { + self.buffer.seq_len = self.capacity; + ops::copy_hidden_token_range_into(ctx, &self.buffer, 0, &mut next, 0, self.len)?; + } + next.seq_len = self.len; + self.buffer = next; + self.capacity = new_capacity; + Ok(()) + } + + fn activate_for_read(&mut self) { + self.buffer.seq_len = self.len; + } + + fn clear(&mut self) { + self.len = 0; + self.buffer.seq_len = 0; + } +} + +impl DFlashDraftScratch { + fn new(ctx: &DeviceContext, config: &DFlashConfig, max_context_len: usize) -> Result { + let block_size = config.block_size; + let hidden_size = config.hidden_size; + let q_dim = config.num_attention_heads * config.head_dim; + let kv_dim = config.num_key_value_heads * config.head_dim; + let inter_dim = config.intermediate_size; + let tail_capacity = max_context_len + block_size; + + Ok(Self { + max_context_len, + block_token_ids_h: vec![config.dflash_config.mask_token_id; block_size], + token_ids_d: ctx.stream.alloc_zeros(block_size)?, + context_projected: HiddenStates::zeros(ctx, hidden_size, max_context_len)?, + context_hidden: HiddenStates::zeros(ctx, hidden_size, max_context_len)?, + hidden: HiddenStates::zeros(ctx, hidden_size, block_size)?, + hidden_out: HiddenStates::zeros(ctx, hidden_size, block_size)?, + normed: HiddenStates::zeros(ctx, hidden_size, block_size)?, + tail_input: HiddenStates::zeros(ctx, hidden_size, tail_capacity)?, + q_batch: HiddenStates::zeros(ctx, q_dim, block_size)?, + k_tail: HiddenStates::zeros(ctx, kv_dim, tail_capacity)?, + v_tail: HiddenStates::zeros(ctx, kv_dim, tail_capacity)?, + attn_output: HiddenStates::zeros(ctx, q_dim, block_size)?, + o_buf: HiddenStates::zeros(ctx, hidden_size, block_size)?, + gate_out: HiddenStates::zeros(ctx, inter_dim, block_size)?, + up_out: HiddenStates::zeros(ctx, inter_dim, block_size)?, + act_out: HiddenStates::zeros(ctx, inter_dim, block_size)?, + logits_normed: HiddenStates::zeros(ctx, hidden_size, block_size)?, + logits: HiddenStates::zeros(ctx, config.vocab_size, block_size)?, + }) + } + + fn ensure_context_capacity( + &mut self, + ctx: &DeviceContext, + config: &DFlashConfig, + context_len: usize, + ) -> Result<()> { + if context_len > self.max_context_len { + *self = Self::new(ctx, config, context_len)?; + } + let block_size = config.block_size; + let tail_len = context_len + block_size; + + self.context_projected.seq_len = context_len; + self.context_hidden.seq_len = context_len; + self.hidden.seq_len = block_size; + self.hidden_out.seq_len = block_size; + self.normed.seq_len = block_size; + self.tail_input.seq_len = tail_len; + self.q_batch.seq_len = block_size; + self.k_tail.seq_len = tail_len; + self.v_tail.seq_len = tail_len; + self.attn_output.seq_len = block_size; + self.o_buf.seq_len = block_size; + self.gate_out.seq_len = block_size; + self.up_out.seq_len = block_size; + self.act_out.seq_len = block_size; + self.logits_normed.seq_len = block_size; + self.logits.seq_len = block_size; + Ok(()) + } +} + +impl DFlashDraftModel { + pub(crate) fn from_safetensors_for_target( + ctx: &DeviceContext, + model_path: &str, + target: &Qwen3Model, + ) -> Result { + let config = DFlashConfig::from_file(model_path) + .with_context(|| format!("load DFlash config from {model_path}"))?; + config.validate_for_target(target.config())?; + + let (shard_paths, weight_map) = load_shard_info(model_path)?; + debug!( + "Loading DFlash drafter from {model_path}: {} shard(s)", + shard_paths.len() + ); + let mmaps = mmap_shards(&shard_paths)?; + let shards = deserialize_shards(&mmaps)?; + + 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}"); + + let q_proj = load_tensor_2d( + ctx, + &shards, + &weight_map, + &format!("{prefix}.self_attn.q_proj.weight"), + )?; + let k_proj = load_tensor_2d( + ctx, + &shards, + &weight_map, + &format!("{prefix}.self_attn.k_proj.weight"), + )?; + let v_proj = load_tensor_2d( + ctx, + &shards, + &weight_map, + &format!("{prefix}.self_attn.v_proj.weight"), + )?; + let q_dim = q_proj.rows; + let kv_dim = k_proj.rows; + let qkv_proj = DeviceMatrix::vstack(ctx, &[&q_proj, &k_proj, &v_proj])?; + drop(q_proj); + drop(k_proj); + drop(v_proj); + + let gate_proj = load_tensor_2d( + ctx, + &shards, + &weight_map, + &format!("{prefix}.mlp.gate_proj.weight"), + )?; + let up_proj = load_tensor_2d( + ctx, + &shards, + &weight_map, + &format!("{prefix}.mlp.up_proj.weight"), + )?; + let gate_up_proj = DeviceMatrix::vstack(ctx, &[&gate_proj, &up_proj])?; + drop(gate_proj); + drop(up_proj); + + layers.push(TransformerBlock { + input_layernorm: load_tensor_1d( + ctx, + &shards, + &weight_map, + &format!("{prefix}.input_layernorm.weight"), + )?, + attention: Attention { + qkv_proj, + o_proj: load_tensor_2d( + ctx, + &shards, + &weight_map, + &format!("{prefix}.self_attn.o_proj.weight"), + )?, + q_norm: load_tensor_1d( + ctx, + &shards, + &weight_map, + &format!("{prefix}.self_attn.q_norm.weight"), + )?, + k_norm: load_tensor_1d( + ctx, + &shards, + &weight_map, + &format!("{prefix}.self_attn.k_norm.weight"), + )?, + q_dim, + kv_dim, + }, + post_attention_layernorm: load_tensor_1d( + ctx, + &shards, + &weight_map, + &format!("{prefix}.post_attention_layernorm.weight"), + )?, + mlp: MLP { + gate_up_proj, + down_proj: load_tensor_2d( + ctx, + &shards, + &weight_map, + &format!("{prefix}.mlp.down_proj.weight"), + )?, + }, + }); + } + + let norm = load_tensor_1d(ctx, &shards, &weight_map, "norm.weight")?; + let hidden_norm = load_tensor_1d(ctx, &shards, &weight_map, "hidden_norm.weight")?; + let fc = load_tensor_2d(ctx, &shards, &weight_map, "fc.weight")?; + let (cos_cache, sin_cache) = precompute_rope( + ctx, + config.head_dim, + config.max_position_embeddings, + config.rope_theta, + )?; + ctx.sync()?; + + Ok(Self { + config, + layers, + norm, + hidden_norm, + fc, + cos_cache, + sin_cache, + }) + } + + pub(crate) fn block_size(&self) -> usize { + self.config.block_size + } + + pub(crate) fn config(&self) -> &DFlashConfig { + &self.config + } + + pub(crate) fn mask_token_id(&self) -> u32 { + self.config.dflash_config.mask_token_id + } + + pub(crate) fn target_layer_ids(&self) -> &[usize] { + &self.config.dflash_config.target_layer_ids + } + + pub(crate) fn tune_gemm_algos(&self, target: &Qwen3Model) -> Result<()> { + let ctx = target.device_ctx(); + let block_size = self.block_size().min(ops::GEMM_LT_MAX_N); + let hidden = self.config.hidden_size; + let q_dim = self.config.num_attention_heads * self.config.head_dim; + let kv_dim = self.config.num_key_value_heads * self.config.head_dim; + let context_dim = self.context_feature_dim(); + + let fc_samples = [(&self.fc, 0)]; + for n in 1..=block_size { + ops::gemm_lt_tune(ctx, &fc_samples, hidden, n)?; + } + + let kv_samples: Vec<_> = self + .layers + .iter() + .flat_map(|layer| { + [ + (&layer.attention.qkv_proj, q_dim), + (&layer.attention.qkv_proj, q_dim + kv_dim), + ] + }) + .collect(); + let min_tail_n = self.block_size() + 1; + let max_tail_n = (self.block_size() * 2).min(ops::GEMM_LT_MAX_N); + for n in min_tail_n..=max_tail_n { + ops::gemm_lt_tune(ctx, &kv_samples, kv_dim, n)?; + } + + log::info!( + "Qwen3 DFlash cublasLt tuned: fc M={} K={} N=1..{}, kv M={} K={} N={}..{}", + hidden, + context_dim, + block_size, + kv_dim, + hidden, + min_tail_n, + max_tail_n, + ); + Ok(()) + } + + pub(crate) fn new_request_state( + &self, + ctx: &DeviceContext, + max_cache_len: usize, + ) -> Result { + anyhow::ensure!( + max_cache_len <= self.config.max_position_embeddings, + "DFlash request cache length {} exceeds max_position_embeddings {}", + max_cache_len, + self.config.max_position_embeddings + ); + let kv_dim = self.config.num_key_value_heads * self.config.head_dim; + let mut layers = Vec::with_capacity(self.layers.len()); + for _ in 0..self.layers.len() { + layers.push(DFlashLayerCache { + k: HiddenStates::zeros(ctx, kv_dim, max_cache_len)?, + v: HiddenStates::zeros(ctx, kv_dim, max_cache_len)?, + }); + } + Ok(DFlashRequestState { + layers, + pending_context: DFlashPendingContext::new( + ctx, + self.context_feature_dim(), + self.config.block_size.min(max_cache_len), + )?, + scratch: DFlashDraftScratch::new(ctx, &self.config, self.config.block_size)?, + committed_len: 0, + max_cache_len, + }) + } + + pub(crate) fn append_pending_context( + &self, + ctx: &DeviceContext, + state: &mut DFlashRequestState, + captured_hidden: &HiddenStates, + token_offset: usize, + token_count: usize, + ) -> Result<()> { + anyhow::ensure!(token_count > 0, "DFlash context append needs tokens"); + anyhow::ensure!( + captured_hidden.hidden_dim == self.context_feature_dim(), + "DFlash captured hidden dim {} does not match expected {}", + captured_hidden.hidden_dim, + self.context_feature_dim() + ); + anyhow::ensure!( + token_offset + token_count <= captured_hidden.seq_len, + "DFlash captured hidden token range exceeds source" + ); + let required_committed_len = state + .committed_len + .checked_add(state.pending_context.len) + .and_then(|len| len.checked_add(token_count)) + .and_then(|len| len.checked_add(self.block_size())) + .context("DFlash pending context cache length overflow")?; + anyhow::ensure!( + required_committed_len <= state.max_cache_len, + "DFlash pending context would exceed cache: committed={}, pending={}, append={}, block={}, max={}", + state.committed_len, + state.pending_context.len, + token_count, + self.block_size(), + state.max_cache_len + ); + state.pending_context.append_from( + ctx, + captured_hidden, + token_offset, + token_count, + state.max_cache_len, + )?; + Ok(()) + } + + pub(crate) fn draft_logits<'a>( + &self, + target: &Qwen3Model, + state: &'a mut DFlashRequestState, + current_token: u32, + ) -> Result> { + let ctx = target.device_ctx(); + let Some(context_len) = state.pending_context_len() else { + anyhow::bail!("DFlash draft requested before target hidden context is available"); + }; + let block_size = self.block_size(); + let tail_len = context_len + block_size; + anyhow::ensure!( + state.committed_len + tail_len <= state.max_cache_len, + "DFlash draft cache overflow: committed={}, tail={}, max={}", + state.committed_len, + tail_len, + state.max_cache_len + ); + + state + .scratch + .ensure_context_capacity(ctx, &self.config, context_len)?; + state.scratch.block_token_ids_h.fill(self.mask_token_id()); + state.scratch.block_token_ids_h[0] = current_token; + ctx.stream.memcpy_htod( + &state.scratch.block_token_ids_h, + &mut state.scratch.token_ids_d, + )?; + target.get_embeddings_batch_into(&state.scratch.token_ids_d, &mut state.scratch.hidden)?; + + state.pending_context.activate_for_read(); + self.project_context_into(ctx, &state.pending_context.buffer, &mut state.scratch)?; + state.pending_context.clear(); + + let hidden_size = self.config.hidden_size; + let q_dim = self.config.num_attention_heads * self.config.head_dim; + let kv_dim = self.config.num_key_value_heads * self.config.head_dim; + let inter_dim = self.config.intermediate_size; + debug_assert_eq!(state.scratch.hidden.hidden_dim, hidden_size); + debug_assert_eq!(state.scratch.q_batch.hidden_dim, q_dim); + debug_assert_eq!(state.scratch.k_tail.hidden_dim, kv_dim); + debug_assert_eq!(state.scratch.gate_out.hidden_dim, inter_dim); + + for (layer_idx, layer) in self.layers.iter().enumerate() { + ops::rms_norm_batch_into( + ctx, + &state.scratch.hidden, + &layer.input_layernorm, + self.config.rms_norm_eps, + &mut state.scratch.normed, + ); + + ops::copy_hidden_token_range_into( + ctx, + &state.scratch.context_hidden, + 0, + &mut state.scratch.tail_input, + 0, + context_len, + )?; + ops::copy_hidden_token_range_into( + ctx, + &state.scratch.normed, + 0, + &mut state.scratch.tail_input, + context_len, + block_size, + )?; + + ops::gemm_rows_into( + ctx, + &layer.attention.qkv_proj, + 0, + q_dim, + &state.scratch.normed, + &mut state.scratch.q_batch, + ); + ops::gemm_rows_into( + ctx, + &layer.attention.qkv_proj, + q_dim, + kv_dim, + &state.scratch.tail_input, + &mut state.scratch.k_tail, + ); + ops::gemm_rows_into( + ctx, + &layer.attention.qkv_proj, + q_dim + kv_dim, + kv_dim, + &state.scratch.tail_input, + &mut state.scratch.v_tail, + ); + + ops::dflash_qk_norm_rope_into( + ctx, + &mut state.scratch.q_batch, + &mut state.scratch.k_tail, + &layer.attention.q_norm, + &layer.attention.k_norm, + &self.cos_cache, + &self.sin_cache, + self.config.num_attention_heads, + self.config.num_key_value_heads, + self.config.head_dim, + state.committed_len + context_len, + state.committed_len, + self.config.rms_norm_eps, + )?; + + let cache = &mut state.layers[layer_idx]; + ops::copy_hidden_token_range_into( + ctx, + &state.scratch.k_tail, + 0, + &mut cache.k, + state.committed_len, + tail_len, + )?; + ops::copy_hidden_token_range_into( + ctx, + &state.scratch.v_tail, + 0, + &mut cache.v, + state.committed_len, + tail_len, + )?; + ops::single_prefill_nhd_noncausal_into( + ctx, + &state.scratch.q_batch, + &cache.k, + &cache.v, + &mut state.scratch.attn_output, + self.config.num_attention_heads, + self.config.num_key_value_heads, + self.config.head_dim, + state.committed_len + tail_len, + )?; + + ops::gemm_into( + ctx, + &layer.attention.o_proj, + &state.scratch.attn_output, + &mut state.scratch.o_buf, + ); + openinfer_kernels::ops::fused_add_rms_norm_round_batch_into( + ctx, + &mut state.scratch.hidden, + &state.scratch.o_buf, + &layer.post_attention_layernorm, + self.config.rms_norm_eps, + &mut state.scratch.normed, + )?; + + ops::gemm_rows_into( + ctx, + &layer.mlp.gate_up_proj, + 0, + inter_dim, + &state.scratch.normed, + &mut state.scratch.gate_out, + ); + ops::gemm_rows_into( + ctx, + &layer.mlp.gate_up_proj, + inter_dim, + inter_dim, + &state.scratch.normed, + &mut state.scratch.up_out, + ); + ops::silu_mul_batch_into( + ctx, + &state.scratch.gate_out, + &state.scratch.up_out, + &mut state.scratch.act_out, + )?; + ops::gemm_into( + ctx, + &layer.mlp.down_proj, + &state.scratch.act_out, + &mut state.scratch.o_buf, + ); + ops::add_batch_into( + ctx, + &state.scratch.hidden, + &state.scratch.o_buf, + &mut state.scratch.hidden_out, + )?; + std::mem::swap(&mut state.scratch.hidden, &mut state.scratch.hidden_out); + } + + let committed_len = state.committed_len; + state.committed_len += context_len; + self.compute_logits_with_target_head_into(target, &mut state.scratch)?; + Ok(DFlashDraftOutput { + logits: &state.scratch.logits, + context_len, + committed_len, + }) + } + + fn context_feature_dim(&self) -> usize { + self.config.hidden_size * self.target_layer_ids().len() + } + + fn project_context_into( + &self, + ctx: &DeviceContext, + context_features: &HiddenStates, + scratch: &mut DFlashDraftScratch, + ) -> Result<()> { + ops::gemm_into( + ctx, + &self.fc, + context_features, + &mut scratch.context_projected, + ); + ops::rms_norm_batch_into( + ctx, + &scratch.context_projected, + &self.hidden_norm, + self.config.rms_norm_eps, + &mut scratch.context_hidden, + ); + Ok(()) + } + + fn compute_logits_with_target_head_into( + &self, + target: &Qwen3Model, + scratch: &mut DFlashDraftScratch, + ) -> Result<()> { + let ctx = target.device_ctx(); + ops::rms_norm_batch_into( + ctx, + &scratch.hidden, + &self.norm, + self.config.rms_norm_eps, + &mut scratch.logits_normed, + ); + ops::gemm_into( + ctx, + target.output_projection(), + &scratch.logits_normed, + &mut scratch.logits, + ); + Ok(()) + } +} + +#[cfg(test)] +pub(crate) fn validate_dflash_config_for_target( + dflash_path: &str, + target_config: &crate::config::Config, +) -> Result { + let config = DFlashConfig::from_file(dflash_path)?; + config.validate_for_target(target_config)?; + Ok(config) +} + +#[cfg(test)] +mod tests { + use super::validate_dflash_config_for_target; + use crate::config::Config; + use std::path::Path; + + #[test] + fn downloaded_dflash_config_matches_qwen3_4b() { + let target_path = std::env::var("OPENINFER_TEST_MODEL_PATH") + .unwrap_or_else(|_| "/data/models/Qwen3-4B".to_string()); + let dflash_path = std::env::var("OPENINFER_DFLASH_TEST_MODEL_PATH") + .unwrap_or_else(|_| "/data/models/Qwen3-4B-DFlash-b16".to_string()); + if !Path::new(&target_path).join("config.json").exists() + || !Path::new(&dflash_path).join("config.json").exists() + { + eprintln!( + "skipping DFlash config test; set OPENINFER_TEST_MODEL_PATH and OPENINFER_DFLASH_TEST_MODEL_PATH" + ); + return; + } + + let target = Config::from_file(&target_path).expect("target config"); + let dflash = validate_dflash_config_for_target(&dflash_path, &target) + .expect("DFlash config should match target"); + + assert_eq!(dflash.block_size, 16); + assert_eq!(dflash.dflash_config.mask_token_id, 151669); + assert_eq!( + dflash.dflash_config.target_layer_ids, + vec![1, 9, 17, 25, 33] + ); + } +} diff --git a/openinfer-qwen3-4b/src/executor.rs b/openinfer-qwen3-4b/src/executor.rs index e7cc9ac6..4aedc363 100644 --- a/openinfer-qwen3-4b/src/executor.rs +++ b/openinfer-qwen3-4b/src/executor.rs @@ -1,22 +1,34 @@ +use anyhow::{Context, Result}; use std::collections::{HashMap, HashSet}; -use std::thread; -use anyhow::Result; -use crossbeam_channel as channel; - -use crate::batch_decode_buffers::{BATCH_BUCKETS, BatchDecodeBuffers}; -use crate::config::{Config, TensorParallelConfig}; -use crate::weights::{ModelRuntimeConfig, Qwen3Model}; -use crate::{Qwen3LoraOptions, Qwen3OffloadOptions}; +use crate::batch_decode_buffers::BATCH_BUCKETS; +use crate::config::{Config, DFlashConfig}; +use crate::speculative::{ + DraftPlan as SpeculativeDraftPlan, DraftResult as SpeculativeDraftResult, + VerifyPlan as SpeculativeVerifyPlan, VerifyResult as SpeculativeVerifyResult, + VerifyStepItem as SpeculativeVerifyStepItem, build_verify_results, +}; +use crate::weights::Qwen3Model; +use crate::{Qwen3LoraOptions, Qwen3SpeculativeOptions}; use openinfer_core::engine::{LoadLoraAdapterRequest, TokenLogprob, UnloadLoraAdapterRequest}; -use openinfer_core::kv_pool::KvLayout; use openinfer_core::ops; use openinfer_core::sampler::SamplingParams; -use openinfer_core::tensor::{DeviceContext, DeviceVec, HiddenStates}; -use openinfer_kv_cache::{ - KvBlockGuard, KvBuffer, KvCacheManager, KvView, LoadReservation, PrefixProbe, -}; -use openinfer_kv_offload::{LoadHandle, OffloadConfig, OffloadEngine}; +use openinfer_core::tensor::{DeviceContext, HiddenStates}; +use openinfer_kv_cache::{KvCacheManager, LoadReservation, PrefixProbe}; +use openinfer_kv_offload::{LoadHandle, OffloadEngine}; + +mod dflash_lane; +mod dflash_prefill; +mod lifecycle; +mod model_executor; +mod speculative_exec; +mod worker; + +use self::worker::{LocalQwen3Lane, RankWorker, StepCommand, WorkerStepOutcome}; + +const BF16_BYTES: usize = 2; +const DFLASH_MIN_EXTRA_MEMORY_RESERVE_BYTES: usize = 1024 * 1024 * 1024; +const DFLASH_ALLOCATOR_MARGIN_BYTES: usize = 256 * 1024 * 1024; #[derive(Clone, Copy, Debug, Eq, PartialEq, Hash, Ord, PartialOrd)] pub struct RequestId(pub(crate) u64); @@ -83,6 +95,11 @@ impl PrefillStepItem { } } + pub fn with_chunk_budget(mut self, chunk_budget: usize) -> Self { + self.chunk_budget = chunk_budget; + self + } + /// Prompt tokens forwarded this step. fn as_slice(&self) -> &[u32] { &self.prompt_tokens[self.chunk_start..self.chunk_start + self.chunk_tokens] @@ -248,6 +265,154 @@ fn build_batch_decode_request_results( Ok(outputs) } +pub(super) struct DFlashMemoryReserve { + total_bytes: usize, + request_state_budget_bytes: usize, +} + +fn dflash_memory_reserve(options: &Qwen3SpeculativeOptions) -> Result { + let Some(dflash) = options.dflash.as_ref() else { + return Ok(DFlashMemoryReserve { + total_bytes: 0, + request_state_budget_bytes: 0, + }); + }; + let path = dflash + .model_path + .to_str() + .ok_or_else(|| anyhow::anyhow!("DFlash model path must be valid UTF-8"))?; + let config = DFlashConfig::from_file(path)?; + let (shard_paths, _) = openinfer_core::weight_loader::load_shard_info(path)?; + let weight_bytes = shard_paths.iter().try_fold(0usize, |acc, path| { + let len = std::fs::metadata(path) + .with_context(|| format!("stat DFlash shard {path}"))? + .len() as usize; + Ok::(acc + len) + })?; + let extra_bytes = dflash_request_state_reserve_bytes(&config)?; + let state_budget_bytes = extra_bytes.saturating_sub(DFLASH_ALLOCATOR_MARGIN_BYTES); + let reserved = weight_bytes + .checked_add(extra_bytes) + .context("DFlash memory reserve overflow")?; + log::info!( + "Qwen3 DFlash memory reserve: weights={:.1} MB, extra={:.1} MB, state_budget={:.1} MB, total={:.1} MB", + weight_bytes as f64 / 1e6, + extra_bytes as f64 / 1e6, + state_budget_bytes as f64 / 1e6, + reserved as f64 / 1e6, + ); + Ok(DFlashMemoryReserve { + total_bytes: reserved, + request_state_budget_bytes: state_budget_bytes, + }) +} + +fn checked_product(values: &[usize], context: &'static str) -> Result { + values.iter().try_fold(1usize, |acc, &value| { + acc.checked_mul(value).context(context) + }) +} + +fn checked_sum(values: &[usize], context: &'static str) -> Result { + values.iter().try_fold(0usize, |acc, &value| { + acc.checked_add(value).context(context) + }) +} + +fn bf16_tensor_bytes(rows: usize, cols: usize, context: &'static str) -> Result { + checked_product(&[rows, cols, BF16_BYTES], context) +} + +fn dflash_request_state_reserve_bytes(config: &DFlashConfig) -> Result { + let footprint = + dflash_request_state_footprint_bytes_for_len(config, config.max_position_embeddings)?; + let with_margin = footprint + .checked_add(DFLASH_ALLOCATOR_MARGIN_BYTES) + .context("DFlash request-state reserve overflow")?; + Ok(with_margin.max(DFLASH_MIN_EXTRA_MEMORY_RESERVE_BYTES)) +} + +pub(super) fn dflash_request_state_footprint_bytes_for_len( + config: &DFlashConfig, + max_cache_len: usize, +) -> Result { + let max_tokens = max_cache_len; + let block = config.block_size; + let hidden = config.hidden_size; + let q_dim = checked_product( + &[config.num_attention_heads, config.head_dim], + "DFlash q dimension overflow", + )?; + let kv_dim = checked_product( + &[config.num_key_value_heads, config.head_dim], + "DFlash KV dimension overflow", + )?; + let inter = config.intermediate_size; + let context_feature_dim = hidden + .checked_mul(config.dflash_config.target_layer_ids.len()) + .context("DFlash context feature dimension overflow")?; + + let kv_cache = checked_product( + &[config.num_hidden_layers, 2, max_tokens, kv_dim, BF16_BYTES], + "DFlash KV cache reserve overflow", + )?; + let pending_context = bf16_tensor_bytes( + max_tokens, + context_feature_dim, + "DFlash pending-context reserve overflow", + )?; + let context_projected = bf16_tensor_bytes( + max_tokens, + hidden, + "DFlash projected-context reserve overflow", + )?; + let context_hidden = + bf16_tensor_bytes(max_tokens, hidden, "DFlash context-hidden reserve overflow")?; + let tail_input = bf16_tensor_bytes(max_tokens, hidden, "DFlash tail-input reserve overflow")?; + let k_tail = bf16_tensor_bytes(max_tokens, kv_dim, "DFlash K-tail reserve overflow")?; + let v_tail = bf16_tensor_bytes(max_tokens, kv_dim, "DFlash V-tail reserve overflow")?; + + let block_hidden_count = checked_product( + &[5, block, hidden, BF16_BYTES], + "DFlash block hidden scratch reserve overflow", + )?; + let block_q_count = checked_product( + &[2, block, q_dim, BF16_BYTES], + "DFlash block attention scratch reserve overflow", + )?; + let block_mlp_count = checked_product( + &[3, block, inter, BF16_BYTES], + "DFlash block MLP scratch reserve overflow", + )?; + let block_logits = bf16_tensor_bytes( + block, + config.vocab_size, + "DFlash block logits reserve overflow", + )?; + let block_token_ids = checked_product( + &[block, std::mem::size_of::()], + "DFlash token-id scratch reserve overflow", + )?; + + checked_sum( + &[ + kv_cache, + pending_context, + context_projected, + context_hidden, + tail_input, + k_tail, + v_tail, + block_hidden_count, + block_q_count, + block_mlp_count, + block_logits, + block_token_ids, + ], + "DFlash request-state reserve overflow", + ) +} + fn execute_step_on_lane( lane: &mut LocalQwen3Lane, step: &StepCommand, @@ -264,8 +429,24 @@ fn execute_step_on_lane( .iter() .map(|req| req.lora_adapter.as_deref()) .collect(); - let (logits, all_position_logits) = - lane.execute_prefill(&prompts, kv_views, &lora_adapters, *echo)?; + let capture_dflash_context = lane.should_capture_dflash_prefill_context(requests); + let capture_layer_ids = if capture_dflash_context { + lane.dflash_capture_layer_ids() + } else { + None + }; + let (logits, all_position_logits, _captured_hidden) = lane.execute_prefill( + &prompts, + kv_views, + &lora_adapters, + *echo, + capture_layer_ids.as_deref(), + )?; + let dflash_context_captured_requests = lane.record_prefill_dflash_context( + requests, + capture_dflash_context, + _captured_hidden.as_ref(), + )?; if collect_result { let params: Vec<&SamplingParams> = requests.iter().map(|r| &r.params).collect(); let random_vals: Vec = requests.iter().map(|r| r.random_val).collect(); @@ -279,11 +460,77 @@ fn execute_step_on_lane( all_position_logits.as_ref(), *echo, )?, + dflash_context_captured_requests, })) } else { Ok(WorkerStepOutcome::Ack) } } + StepCommand::SpeculativeVerify { requests, kv_views } => { + for req in requests { + anyhow::ensure!( + req.params.is_greedy(), + "speculative verification currently supports greedy sampling only" + ); + } + let prompts: Vec<&[u32]> = requests + .iter() + .map(SpeculativeVerifyStepItem::as_slice) + .collect(); + let lora_adapters: Vec> = requests + .iter() + .map(|req| req.lora_adapter.as_deref()) + .collect(); + let capture_layer_ids = lane.dflash_capture_layer_ids(); + let verify_start = lane.start_dflash_timing()?; + let verify_nvtx_range = if lane.dflash_nvtx_enabled() { + Some(nvtx::range!("qwen3.dflash.verify")) + } else { + None + }; + let (_last_logits, all_position_logits, captured_hidden) = lane.execute_prefill( + &prompts, + kv_views, + &lora_adapters, + true, + capture_layer_ids.as_deref(), + )?; + if collect_result { + let all_position_logits = all_position_logits.ok_or_else(|| { + anyhow::anyhow!("speculative verification did not return all-position logits") + })?; + let params: Vec<&SamplingParams> = requests + .iter() + .flat_map(|req| std::iter::repeat_n(&req.params, req.token_ids.len())) + .collect(); + let random_vals = vec![0.0; params.len()]; + let target_tokens = + lane.select_step_tokens(&all_position_logits, ¶ms, &random_vals)?; + let verify_ms = lane.finish_dflash_timing(verify_start)?; + drop(verify_nvtx_range); + let results = build_verify_results(requests, &target_tokens)?; + lane.record_verify_dflash_context( + requests, + &results, + captured_hidden.as_ref(), + verify_ms, + )?; + Ok(WorkerStepOutcome::SpeculativeVerify( + SpeculativeVerifyResult { requests: results }, + )) + } else { + Ok(WorkerStepOutcome::Ack) + } + } + StepCommand::SpeculativeDraft { requests } => { + if collect_result { + Ok(WorkerStepOutcome::SpeculativeDraft( + lane.execute_dflash_draft(requests)?, + )) + } else { + Ok(WorkerStepOutcome::Ack) + } + } StepCommand::Decode { requests, kv_views } => { let token_ids: Vec = requests.iter().map(|req| req.token_id).collect(); let lora_adapters: Vec> = requests @@ -382,6 +629,7 @@ struct SamplingScratch { row_states: cudarc::driver::CudaSlice, valid: cudarc::driver::CudaSlice, out: cudarc::driver::CudaSlice, + out_host: Vec, } impl SamplingScratch { @@ -399,6 +647,7 @@ impl SamplingScratch { .alloc_zeros(openinfer_core::ops::flashinfer_topk_row_states_bytes())?, valid: ctx.stream.alloc_zeros(1)?, out: ctx.stream.alloc_zeros(max_batch_bucket)?, + out_host: vec![0; max_batch_bucket], }) } } @@ -550,6 +799,7 @@ pub struct DecodeRequestResult { pub struct PrefillResult { pub requests: Vec, + pub dflash_context_captured_requests: Vec, } pub struct DecodeResult { @@ -572,6 +822,30 @@ pub(crate) trait ModelExecutor: Send { fn execute_prefill(&mut self, plan: PrefillPlan<'_>) -> Result; fn execute_decode(&mut self, plan: DecodePlan<'_>) -> Result; + fn execute_speculative_verify( + &mut self, + _plan: SpeculativeVerifyPlan<'_>, + ) -> Result { + anyhow::bail!("speculative verification is not implemented for this executor") + } + fn execute_speculative_draft( + &mut self, + _plan: SpeculativeDraftPlan<'_>, + ) -> Result { + anyhow::bail!("speculative draft is not implemented for this executor") + } + fn speculative_enabled(&self) -> bool { + false + } + fn speculative_request_ready(&self, _request_id: RequestId) -> bool { + self.speculative_enabled() + } + fn speculative_state_budget_bytes(&self) -> Option { + None + } + fn speculative_request_state_bytes(&self, _prompt_len: usize, _max_tokens: usize) -> usize { + 0 + } fn execute_unified(&mut self, plan: UnifiedPlan<'_>) -> Result; fn load_lora_adapter(&mut self, request: &LoadLoraAdapterRequest) -> Result<()> { @@ -639,6 +913,9 @@ struct Qwen3ExecutorMetadata { block_size: usize, stop_token_ids: Vec, config: Config, + max_context_tokens: usize, + dflash_config: Option, + dflash_state_budget_bytes: usize, } pub struct Qwen3Executor { @@ -668,6 +945,8 @@ pub struct Qwen3Executor { /// so prefix matching itself stays enabled). Set via /// [`Self::set_no_prefix_cache`]. l1_retention_disabled: bool, + speculative_enabled: bool, + dflash_ready_requests: HashSet, } /// One request's in-flight CPU-tier KV prefetch. @@ -684,1092 +963,6 @@ struct PrefetchState { handle: Option, } -impl Qwen3Executor { - pub(crate) fn single(model: Qwen3Model, offload_opts: &Qwen3OffloadOptions) -> Result { - let budget = model.kv_budget(); - let kv_mgr = KvCacheManager::new( - &model.device_ctx().stream, - budget.num_layers, - budget.num_kv_heads, - budget.head_dim, - budget.block_size, - budget.num_blocks, - )?; - let metadata = Qwen3ExecutorMetadata { - block_size: budget.block_size, - stop_token_ids: model.config().stop_token_ids.clone(), - config: model.config().clone(), - }; - let kv_buffer = kv_mgr.buffer().clone(); - // Build the offload engine while the model's stream is still in hand - // (it moves into the RankWorker below). Registers the fused KV buffer. - let offload = build_offload(offload_opts, &kv_mgr, model.device_ctx())?; - let total_blocks = kv_mgr.pool().total_blocks(); - let padding_block_id = kv_mgr.pool().padding_block_id(); - Ok(Self { - metadata, - kv_mgr, - request_kvs: HashMap::new(), - primary: RankWorker::spawn( - 0, - LocalQwen3Lane::new(model, kv_buffer, total_blocks, padding_block_id)?, - )?, - workers: Vec::new(), - loaded_lora_adapters: HashSet::new(), - prefix_cache_enabled: true, - lora_options: Qwen3LoraOptions::default(), - offload, - saved_cursor: HashMap::new(), - prefetch: HashMap::new(), - l1_retention_disabled: false, - }) - } - - pub fn from_runtime( - model_path: &str, - enable_cuda_graph: bool, - device_ordinals: &[usize], - ) -> Result { - Self::from_runtime_with_lora_options( - model_path, - enable_cuda_graph, - device_ordinals, - Qwen3LoraOptions::default(), - Qwen3OffloadOptions::disabled(), - ) - } - - pub fn from_runtime_with_lora_options( - model_path: &str, - enable_cuda_graph: bool, - device_ordinals: &[usize], - lora_options: Qwen3LoraOptions, - offload_options: Qwen3OffloadOptions, - ) -> Result { - let lora_options = lora_options.validate()?; - anyhow::ensure!( - !device_ordinals.is_empty(), - "Qwen3 executor requires at least one device" - ); - anyhow::ensure!( - !offload_options.enabled || device_ordinals.len() == 1, - "KV offload is only supported on the single-GPU path (tensor parallel \ - shards KV per rank); got {} devices", - device_ordinals.len() - ); - if device_ordinals.len() == 1 { - let model = Qwen3Model::from_safetensors_with_runtime( - model_path, - ModelRuntimeConfig { - enable_cuda_graph, - tensor_parallel: None, - device_ordinal: device_ordinals[0], - max_loras: lora_options.max_loras, - max_lora_rank: lora_options.max_lora_rank, - }, - )?; - let mut executor = Self::single(model, &offload_options)?; - executor.lora_options = lora_options; - return Ok(executor); - } - - let world_size = device_ordinals.len(); - let mut models = Vec::with_capacity(world_size); - for (rank, &device_ordinal) in device_ordinals.iter().enumerate() { - models.push(Qwen3Model::from_safetensors_with_runtime( - model_path, - ModelRuntimeConfig { - enable_cuda_graph, - tensor_parallel: Some(TensorParallelConfig { rank, world_size }), - device_ordinal, - max_loras: lora_options.max_loras, - max_lora_rank: lora_options.max_lora_rank, - }, - )?); - } - - // Compute budget from first model (all ranks share geometry). - let budget = models[0].kv_budget(); - - // Create the centralized KvCacheManager on rank 0's stream. - let kv_mgr = KvCacheManager::new( - &models[0].device_ctx().stream, - budget.num_layers, - budget.num_kv_heads, - budget.head_dim, - budget.block_size, - budget.num_blocks, - )?; - - let metadata = Qwen3ExecutorMetadata { - block_size: budget.block_size, - stop_token_ids: models[0].config().stop_token_ids.clone(), - config: models[0].config().clone(), - }; - - // Create extra KvBuffers for ranks 1+ on their respective streams. - let mut extra_kv_buffers = Vec::with_capacity(world_size - 1); - for model in &models[1..] { - extra_kv_buffers.push(KvBuffer::new( - &model.device_ctx().stream, - budget.num_layers, - budget.num_kv_heads, - budget.head_dim, - budget.block_size, - budget.num_blocks, - )?); - } - - let streams = models - .iter() - .map(|m| m.device_ctx().stream.clone()) - .collect(); - let comms = cudarc::nccl::safe::Comm::from_devices(streams) - .map_err(|e| anyhow::anyhow!("failed to initialize NCCL comms: {e:?}"))?; - for (model, comm) in models.iter_mut().zip(comms) { - model.attach_tp_comm(comm); - } - - let total_blocks = kv_mgr.pool().total_blocks(); - let padding_block_id = kv_mgr.pool().padding_block_id(); - - // Primary rank gets the KvBuffer from the centralized manager. - let primary_buffer = kv_mgr.buffer().clone(); - let mut models_iter = models.into_iter(); - let primary_model = models_iter.next().unwrap(); - let primary = RankWorker::spawn( - 0, - LocalQwen3Lane::new( - primary_model, - primary_buffer, - total_blocks, - padding_block_id, - )?, - )?; - - // Worker ranks get their own extra KvBuffers. - let workers = models_iter - .zip(extra_kv_buffers) - .enumerate() - .map(|(index, (model, buffer))| { - let lane = LocalQwen3Lane::new(model, buffer, total_blocks, padding_block_id)?; - RankWorker::spawn(index + 1, lane) - }) - .collect::>>()?; - - Ok(Self { - metadata, - kv_mgr, - request_kvs: HashMap::new(), - primary, - workers, - loaded_lora_adapters: HashSet::new(), - prefix_cache_enabled: true, - lora_options, - // Offload is single-GPU only (asserted above); never built here. - offload: None, - saved_cursor: HashMap::new(), - prefetch: HashMap::new(), - l1_retention_disabled: false, - }) - } - - pub fn block_size(&self) -> usize { - ::block_size(self) - } - - pub fn max_request_blocks(&self) -> usize { - ::max_request_blocks(self) - } - - pub fn available_blocks(&self) -> usize { - ::available_blocks(self) - } - - pub fn is_stop_token(&self, token_id: u32) -> bool { - ::is_stop_token(self, token_id) - } - - pub fn drop_request(&mut self, request_id: RequestId) -> Result<()> { - ::drop_request(self, request_id) - } - - pub fn execute_prefill(&mut self, plan: PrefillPlan<'_>) -> Result { - ::execute_prefill(self, plan) - } - - pub fn execute_decode(&mut self, plan: DecodePlan<'_>) -> Result { - ::execute_decode(self, plan) - } - - pub fn execute_unified(&mut self, plan: UnifiedPlan<'_>) -> Result { - ::execute_unified(self, plan) - } - - /// Prefix caching is on by default; tests that assert bit-identical - /// replay disable it (a cache hit changes prefill GEMM shapes, which - /// drifts logits by bf16 ULPs). - pub fn set_prefix_cache_enabled(&mut self, enabled: bool) { - self.prefix_cache_enabled = enabled; - } - - /// vLLM-style `--no-prefix-cache`. Behaviour depends on whether offload is - /// active: - /// * **No offload** — classic: disable prefix matching outright, so every - /// prefill recomputes the full prompt. - /// * **With offload** — pure-L2 mode: keep matching on (the host-tier - /// restore registers blocks and relies on `match_and_add_prefix` to pick - /// them up) but stop retaining completed blocks in HBM, so no request - /// ever serves its prefix from a cross-request L1 hit. Every reuse then - /// comes from the host tier, which is the point of the L2 benchmark. - /// - /// A resident HBM block and its host-tier copy share one content hash, so - /// the cache cannot be told to prefer L2 for a block still in HBM — the only - /// way to force the bytes from L2 is to not keep the HBM copy around. - pub fn set_no_prefix_cache(&mut self, on: bool) { - if self.offload.is_some() { - self.l1_retention_disabled = on; - } else { - self.prefix_cache_enabled = !on; - } - } - - /// Whether KV offload is active on this executor. - pub fn offload_enabled(&self) -> bool { - self.offload.is_some() - } - - /// Flush pending offload saves into the host read cache so a following - /// query can see them. A persistence barrier for handoff and tests; no-op - /// without offload. - pub fn flush_offload_saves(&self) { - if let Some(offload) = &self.offload { - offload.flush_saves(); - } - } - - /// Drop every cached-but-unused GPU prefix block. With offload on, this - /// forces a cold prefix to be restored from the host tier on its next - /// request (rather than served from HBM). - pub fn evict_cached_blocks(&self) { - self.kv_mgr.pool().evict_inactive(); - } - - /// Begin an async CPU-tier KV prefetch for `request_id`; see the - /// [`ModelExecutor`] hook. Public so admission drivers and tests can park a - /// request on its load. Returns `true` when a load is in flight. - pub fn begin_kv_prefetch( - &mut self, - request_id: RequestId, - prompt_tokens: &[u32], - lora_adapter: Option<&str>, - reserve_floor: usize, - ) -> bool { - ::begin_kv_prefetch( - self, - request_id, - prompt_tokens, - lora_adapter, - reserve_floor, - ) - } - - /// Block until at least one in-flight prefetch settles, then sweep the - /// rest; returns the settled request ids (now prefill-eligible). - pub fn wait_ready_prefetch(&mut self) -> Vec { - ::wait_ready_prefetch(self) - } - - // ── KV-offload SAVE ──────────────────────────────────────────────── - - /// Save every block that sealed since this request's last save to the host - /// tier (fire-and-forget). Safe to call right after `apply_prefill`/ - /// `apply_decode`: the producing step's token read-back has already - /// synchronized the compute stream, so the sealed KV is fully written. - fn save_sealed_blocks(&mut self, request_id: RequestId) { - if self.offload.is_none() { - return; - } - let Some(rkv) = self.request_kvs.get(&request_id) else { - return; - }; - // `assigned_block_hashes` lists only sealed (registered) blocks; the - // partial tail block has no hash and never appears here. - let assigned = rkv.assigned_block_hashes(); - let prefix_matched = rkv.prefix_matched_blocks(); - let cursor = self - .saved_cursor - .entry(request_id) - .or_insert(prefix_matched); - if assigned.len() <= *cursor { - return; - } - let fresh = &assigned[*cursor..]; - let block_ids: Vec = fresh.iter().map(|(id, _)| *id).collect(); - let block_hashes: Vec> = fresh.iter().map(|(_, h)| h.to_vec()).collect(); - // Pin exactly the blocks being saved (aligned 1:1 with `assigned`) for - // the duration of the async D2H, so a finished request can't hand the - // slot to a new request that overwrites it before the copy lands. - let pins: Vec = rkv - .assigned_block_guards() - .into_iter() - .skip(*cursor) - .collect(); - *cursor = assigned.len(); - self.offload - .as_ref() - .expect("offload present") - .save(&block_ids, &block_hashes, pins); - } - - // ── Chunked prefill ──────────────────────────────────────────────── - - /// Prepare one prefill step for `req`: create its `RequestKv` on the - /// first chunk (matching the prefix cache), then clamp the scheduler's - /// chunk budget to the prompt tokens actually remaining and allocate KV - /// for them. Sets `chunk_start`/`chunk_tokens` on the item. - fn schedule_prefill_chunk(&mut self, req: &mut PrefillStepItem) -> Result<()> { - if !self.request_kvs.contains_key(&req.request_id) { - let mut rkv = self.kv_mgr.pool().new_request( - req.prompt_tokens.clone(), - req.max_output_tokens, - req.lora_adapter.as_deref(), - ); - // Echo needs logits for every prompt position; cached positions - // are never forwarded, so echo requests prefill from scratch. - if self.prefix_cache_enabled && !req.echo { - req.cached_tokens = rkv.match_and_add_prefix(self.kv_mgr.pool())?; - } - self.request_kvs.insert(req.request_id, rkv); - // match_and_add_prefix above already absorbed any CPU-prefetched - // blocks (now held by the request's sequence), so release the - // prefetch's separate hold. - self.prefetch.remove(&req.request_id); - } - let rkv = self - .request_kvs - .get_mut(&req.request_id) - .expect("inserted above"); - req.chunk_start = rkv.kv_position(); - let remaining = req.prompt_tokens.len() - req.chunk_start; - // Echo must produce all-position logits in a single forward, so it is - // exempt from chunking (the scheduler never splits echo requests). - req.chunk_tokens = if req.echo { - remaining - } else { - remaining.min(req.chunk_budget) - }; - assert!( - req.chunk_tokens > 0, - "zero-token prefill chunk for {:?} (budget {})", - req.request_id, - req.chunk_budget - ); - rkv.schedule_prefill(req.chunk_tokens, self.kv_mgr.pool()) - .map_err(|e| anyhow::anyhow!("schedule_prefill failed for {:?}: {e}", req.request_id)) - } - - /// Register a finished prefill step on the request's KV: the final chunk - /// carries the first generated token, non-final chunks only advance the - /// KV position. - fn apply_prefill_result(&mut self, result: &PrefillRequestResult) -> Result<()> { - let rkv = self - .request_kvs - .get_mut(&result.request_id) - .expect("request must exist after prefill"); - if result.completed { - rkv.apply_prefill(result.first_token, self.kv_mgr.pool()) - } else { - rkv.apply_prefill_chunk(self.kv_mgr.pool()) - } - } - - // ── KV-offload LOAD (async CPU-tier prefetch) ────────────────────── - // The trait-facing prefetch hooks (`begin_kv_prefetch`, - // `drain_ready_prefetch`, `wait_ready_prefetch`, `has_pending_prefetch`) - // live in the `ModelExecutor` impl below; `settle_prefetch` is their shared - // helper. - - /// Finalize one prefetch whose load returned `result`. On success the - /// reserved blocks are staged + registered (held by the probe until the - /// request prefills); on failure the state is dropped so the request - /// prefills from scratch. - fn settle_prefetch( - &mut self, - id: RequestId, - result: Result<(), openinfer_kv_offload::EngineError>, - ) { - if let Some(st) = self.prefetch.get_mut(&id) { - st.handle = None; - } - match result { - Ok(()) => { - let reservation = self - .prefetch - .get_mut(&id) - .and_then(|st| st.reservation.take()) - .expect("reservation present until commit"); - let st = self.prefetch.get_mut(&id).expect("prefetch present"); - self.kv_mgr - .pool() - .commit_loaded_blocks(&mut st.probe, reservation); - } - Err(e) => { - log::warn!("KV offload load failed for {id:?} (prefill from scratch): {e}"); - self.prefetch.remove(&id); - } - } - } - - fn wait_for_step_ack( - pending: Vec>>, - op_name: &'static str, - ) -> Result<()> { - for recv in pending { - match recv - .recv() - .map_err(|_| anyhow::anyhow!("tensor-parallel {op_name} worker dropped"))?? - { - WorkerStepOutcome::Ack => {} - other => { - return Err(anyhow::anyhow!( - "tensor-parallel {op_name} worker returned unexpected payload: {}", - other.kind() - )); - } - } - } - Ok(()) - } - - fn run_step(&self, step: &StepCommand) -> Result { - let primary = self.primary.run_step(step.clone(), true)?; - let mut pending = Vec::with_capacity(self.workers.len()); - for worker in &self.workers { - pending.push(worker.run_step(step.clone(), false)?); - } - let primary_result = primary - .recv() - .map_err(|_| anyhow::anyhow!("primary worker dropped step response"))??; - Self::wait_for_step_ack(pending, step.kind())?; - Ok(primary_result) - } -} - -/// Build the KV-offload engine for the single-GPU path, or `None` when offload -/// is disabled. Registers the fused KV buffer with pegaflow against the model's -/// device/stream — must be called while that stream is still owned by the model -/// (before it moves into the `RankWorker`). -fn build_offload( - opts: &Qwen3OffloadOptions, - kv_mgr: &KvCacheManager, - ctx: &DeviceContext, -) -> Result> { - if !opts.enabled { - return Ok(None); - } - let device_id = ctx.device_ordinal as i32; - let config = OffloadConfig::new( - format!("qwen3-4b-dev{device_id}"), - device_id, - opts.pinned_pool_bytes, - ); - let engine = OffloadEngine::new(config, kv_mgr.buffer(), &ctx.stream) - .map_err(|e| anyhow::anyhow!("KV offload engine init failed: {e}"))?; - log::info!( - "KV offload enabled on device {device_id} ({} MiB host tier)", - opts.pinned_pool_bytes >> 20 - ); - Ok(Some(engine)) -} - -fn ensure_lora_capacity( - loaded_lora_adapters: &HashSet, - lora_name: &str, - max_loras: usize, - load_inplace: bool, -) -> Result<()> { - if loaded_lora_adapters.contains(lora_name) { - anyhow::ensure!( - load_inplace, - "Qwen3 LoRA adapter {lora_name} is already loaded" - ); - return Ok(()); - } - anyhow::ensure!( - loaded_lora_adapters.len() < max_loras, - "Qwen3 LoRA adapter capacity exceeded: max_loras={}, loaded_adapters={}, requested={}", - max_loras, - loaded_lora_adapters.len(), - lora_name - ); - Ok(()) -} - -impl ModelExecutor for Qwen3Executor { - fn block_size(&self) -> usize { - self.metadata.block_size - } - - fn max_request_blocks(&self) -> usize { - self.kv_mgr.pool().max_request_blocks() - } - - fn max_context_tokens(&self) -> usize { - self.metadata.config.max_position_embeddings - } - - fn max_decode_batch_size(&self) -> usize { - *BATCH_BUCKETS.last().unwrap() - } - - fn available_blocks(&self) -> usize { - self.kv_mgr.pool().available_blocks() - } - - fn is_stop_token(&self, token_id: u32) -> bool { - self.metadata.stop_token_ids.contains(&token_id) - } - - fn prefetched_blocks(&self, request_id: RequestId) -> usize { - self.prefetch - .get(&request_id) - .map_or(0, |st| st.probe.held_blocks()) - } - - fn drop_request(&mut self, request_id: RequestId) -> Result<()> { - // Remove and drop — RAII on SchedulableSequence's block guards - // returns all allocated blocks regardless of lifecycle state. The same - // RAII frees any parked prefetch's reserved/held blocks. - self.request_kvs.remove(&request_id); - // A parked prefetch may still have a load in flight: pegaflow's worker - // is writing the reserved GPU blocks (H2D). Dropping the reservation now - // frees those physical pages for immediate reuse while the DMA keeps - // landing on them — silent KV corruption, the load-side mirror of the - // SAVE keep-alive pin. Block until the copy finishes before the - // reservation drops. The scheduler is a dedicated synchronous thread, so - // this brief wait costs nothing it could spend elsewhere. - if let Some(mut state) = self.prefetch.remove(&request_id) { - if let Some(handle) = state.handle.take() { - let _ = handle.wait(); - } - } - self.saved_cursor.remove(&request_id); - Ok(()) - } - - fn begin_kv_prefetch( - &mut self, - request_id: RequestId, - prompt_tokens: &[u32], - lora_adapter: Option<&str>, - reserve_floor: usize, - ) -> bool { - let Some(offload) = self.offload.as_ref() else { - return false; - }; - if !self.prefix_cache_enabled { - return false; - } - if self.l1_retention_disabled { - // Pure-L2 mode: drop any cross-request HBM retention so the probe - // sees gpu_hit == 0 and queries the whole cacheable prefix from the - // host tier. Only inactive (completed, unheld) blocks are drained — - // the current request holds nothing yet, and in-flight prefetches - // keep their reserved blocks, so this never touches live KV. - self.kv_mgr.pool().evict_inactive(); - } - let probe = self - .kv_mgr - .pool() - .probe_prefix(prompt_tokens.to_vec(), lora_adapter); - let query_hashes = probe.cpu_query_hashes(); - if query_hashes.is_empty() { - return false; - } - let hit = match offload.query(&request_id.0.to_string(), &query_hashes) { - Ok(hit) => hit, - Err(e) => { - log::warn!("KV offload query failed for {request_id:?} (skipping): {e}"); - return false; - } - }; - let (Some(lease), num_blocks) = (hit.lease, hit.num_blocks) else { - return false; // miss - }; - // Blocks promised to admitted requests are off-limits: reserving into - // them makes a later prefill chunk or decode growth fail allocation. - if self - .kv_mgr - .pool() - .available_blocks() - .saturating_sub(reserve_floor) - < num_blocks - { - offload.release_query_lease(lease); - return false; - } - let Some(reservation) = self.kv_mgr.pool().reserve_loaded_blocks(num_blocks) else { - // Block pressure: release the lease so its pinned host blocks aren't - // held for the full lease TTL, and prefill from scratch rather than - // stall. - offload.release_query_lease(lease); - return false; - }; - let page_ids = reservation.page_ids(); - let handle = match offload.load(lease, page_ids) { - Ok(handle) => handle, - Err(e) => { - log::warn!("KV offload load submit failed for {request_id:?} (skipping): {e}"); - // `load` consumes the lease only past its early validation; a - // submit error may leave it pinned, so release it (no-op if it - // was already consumed). - offload.release_query_lease(lease); - return false; - } - }; - self.prefetch.insert( - request_id, - PrefetchState { - probe, - reservation: Some(reservation), - handle: Some(handle), - }, - ); - true - } - - fn drain_ready_prefetch(&mut self) -> Vec { - let ids: Vec = self.prefetch.keys().copied().collect(); - let mut done = Vec::new(); - for id in ids { - let poll = match self.prefetch.get_mut(&id).and_then(|st| st.handle.as_mut()) { - Some(handle) => handle.poll(), - None => continue, // already settled, awaiting prefill - }; - if let Some(result) = poll { - self.settle_prefetch(id, result); - done.push(id); - } - } - done - } - - fn wait_ready_prefetch(&mut self) -> Vec { - let mut done = Vec::new(); - if let Some(id) = self - .prefetch - .iter() - .find(|(_, st)| st.handle.is_some()) - .map(|(id, _)| *id) - { - let handle = self - .prefetch - .get_mut(&id) - .and_then(|st| st.handle.take()) - .expect("in-flight handle present"); - let result = handle.wait(); - self.settle_prefetch(id, result); - // `settle_prefetch` clears the handle, so the drain below skips it; - // record it here as the one we blocked on. - done.push(id); - } - // Sweep any others that completed concurrently. - for id in self.drain_ready_prefetch() { - if !done.contains(&id) { - done.push(id); - } - } - done - } - - fn execute_prefill(&mut self, plan: PrefillPlan<'_>) -> Result { - // 1. Create RequestKvs (first chunk only), clamp chunk budgets, - // schedule KV for this step's tokens - let mut requests = plan.requests.to_vec(); - for req in &mut requests { - self.schedule_prefill_chunk(req)?; - } - - // 2. Build KvViews (seq_len = chunk_start + this chunk) - let kv_views: Vec = requests - .iter() - .map(|req| self.request_kvs[&req.request_id].prefill_view(req.chunk_tokens)) - .collect(); - - // 3. Execute forward - let step = StepCommand::Prefill { - requests, - kv_views, - echo: plan.echo, - }; - let outcome = self.run_step(&step)?; - - // 4. Apply prefill - let result = match outcome { - WorkerStepOutcome::Prefill(result) => result, - other => { - return Err(anyhow::anyhow!( - "prefill returned unexpected: {}", - other.kind() - )); - } - }; - for req_result in &result.requests { - self.apply_prefill_result(req_result)?; - } - // 5. Offload the blocks this prefill just sealed (post-step-sync). - for req_result in &result.requests { - self.save_sealed_blocks(req_result.request_id); - } - - Ok(result) - } - - fn execute_decode(&mut self, plan: DecodePlan<'_>) -> Result { - // 1. Schedule decode for all active requests - for req in plan.requests { - let rkv = self - .request_kvs - .get_mut(&req.request_id) - .ok_or_else(|| anyhow::anyhow!("missing RequestKv for {:?}", req.request_id))?; - rkv.schedule_decode(self.kv_mgr.pool()).map_err(|e| { - anyhow::anyhow!("schedule_decode failed for {:?}: {e}", req.request_id) - })?; - } - - // 2. Build KvViews - let kv_views: Vec = plan - .requests - .iter() - .map(|req| self.request_kvs[&req.request_id].decode_view()) - .collect(); - - // 3. Execute forward - let step = StepCommand::Decode { - requests: plan.requests.to_vec(), - kv_views, - }; - let outcome = self.run_step(&step)?; - - // 4. Apply decode - let result = match outcome { - WorkerStepOutcome::Decode(result) => result, - other => { - return Err(anyhow::anyhow!( - "decode returned unexpected: {}", - other.kind() - )); - } - }; - for req_result in &result.requests { - let rkv = self - .request_kvs - .get_mut(&req_result.request_id) - .expect("request must exist after decode"); - rkv.apply_decode(req_result.token, self.kv_mgr.pool())?; - } - // 5. Offload any block this decode step just sealed (post-step-sync). - for req_result in &result.requests { - self.save_sealed_blocks(req_result.request_id); - } - - Ok(result) - } - - fn execute_unified(&mut self, plan: UnifiedPlan<'_>) -> Result { - // 1. Create RequestKvs for prefill requests (first chunk only), clamp - // chunk budgets, schedule KV for this step's tokens - let mut prefill_requests = plan.prefill_requests.to_vec(); - for req in &mut prefill_requests { - self.schedule_prefill_chunk(req)?; - } - - // Schedule decode for active requests - for req in plan.decode_requests { - let rkv = self - .request_kvs - .get_mut(&req.request_id) - .ok_or_else(|| anyhow::anyhow!("missing RequestKv for {:?}", req.request_id))?; - rkv.schedule_decode(self.kv_mgr.pool()).map_err(|e| { - anyhow::anyhow!("schedule_decode failed for {:?}: {e}", req.request_id) - })?; - } - - // 2. Build KvViews - let prefill_kv_views: Vec = prefill_requests - .iter() - .map(|req| self.request_kvs[&req.request_id].prefill_view(req.chunk_tokens)) - .collect(); - let decode_kv_views: Vec = plan - .decode_requests - .iter() - .map(|req| self.request_kvs[&req.request_id].decode_view()) - .collect(); - - // 3. Execute forward - let step = StepCommand::Unified { - prefill_requests, - prefill_kv_views, - decode_requests: plan.decode_requests.to_vec(), - decode_kv_views, - }; - let outcome = self.run_step(&step)?; - - // 4. Apply both prefill and decode - let result = match outcome { - WorkerStepOutcome::Unified(result) => result, - other => { - return Err(anyhow::anyhow!( - "unified returned unexpected: {}", - other.kind() - )); - } - }; - for req_result in &result.prefill_requests { - self.apply_prefill_result(req_result)?; - } - for req_result in &result.decode_requests { - let rkv = self - .request_kvs - .get_mut(&req_result.request_id) - .expect("request must exist after unified decode"); - rkv.apply_decode(req_result.token, self.kv_mgr.pool())?; - } - // 5. Offload sealed blocks from both halves (post-step-sync). - for req_result in &result.prefill_requests { - self.save_sealed_blocks(req_result.request_id); - } - for req_result in &result.decode_requests { - self.save_sealed_blocks(req_result.request_id); - } - - Ok(result) - } - - fn load_lora_adapter(&mut self, request: &LoadLoraAdapterRequest) -> Result<()> { - ensure_lora_capacity( - &self.loaded_lora_adapters, - &request.lora_name, - self.lora_options.max_loras, - request.load_inplace, - )?; - let adapter = crate::lora::load_lora_adapter( - &request.lora_path, - &self.metadata.config, - self.lora_options.max_lora_rank, - )?; - let world_size = self.workers.len() + 1; - let projection_count: usize = adapter - .layers - .iter() - .map(|layer| layer.projections.len()) - .sum(); - let element_count: usize = adapter - .layers - .iter() - .flat_map(|layer| layer.projections.values()) - .map(|projection| projection.a.data.len() + projection.b.data.len()) - .sum(); - let shape_elems: usize = adapter - .layers - .iter() - .flat_map(|layer| layer.projections.values()) - .map(|projection| { - projection.a.rows * projection.a.cols + projection.b.rows * projection.b.cols - }) - .sum(); - debug_assert_eq!(element_count, shape_elems); - let rank = adapter.manifest.rank; - let targets = adapter.manifest.target_modules.join(", "); - let path = adapter.manifest.path.display().to_string(); - let mut sharded_adapters = Vec::with_capacity(world_size); - for rank in 0..world_size { - sharded_adapters.push(adapter.shard_for_tensor_parallel( - &self.metadata.config, - TensorParallelConfig { rank, world_size }, - )?); - } - - let mut sharded_adapters = sharded_adapters.into_iter(); - let primary_adapter = sharded_adapters - .next() - .expect("rank 0 adapter must exist for nonzero world_size"); - let primary_response = self.primary.load_lora_adapter( - request.lora_name.clone(), - primary_adapter, - request.load_inplace, - )?; - let mut pending = Vec::with_capacity(self.workers.len()); - let mut errors = Vec::new(); - for (index, worker) in self.workers.iter().enumerate() { - let rank = index + 1; - let rank_adapter = sharded_adapters - .next() - .expect("worker adapter must exist for every tensor-parallel rank"); - match worker.load_lora_adapter( - request.lora_name.clone(), - rank_adapter, - request.load_inplace, - ) { - Ok(response) => pending.push((rank, response)), - Err(err) => errors.push(format!("rank {rank} dispatch: {err:#}")), - } - } - - match primary_response.recv() { - Ok(Ok(())) => {} - Ok(Err(err)) => errors.push(format!("rank 0: {err:#}")), - Err(_) => errors.push("rank 0: dropped LoRA load response".to_string()), - } - for (rank, response) in pending { - match response.recv() { - Ok(Ok(())) => {} - Ok(Err(err)) => errors.push(format!("rank {rank}: {err:#}")), - Err(_) => errors.push(format!("rank {rank}: dropped LoRA load response")), - } - } - if !errors.is_empty() { - let mut cleanup_errors = Vec::new(); - match self.primary.discard_lora_adapter(request.lora_name.clone()) { - Ok(response) => match response.recv() { - Ok(Ok(())) => {} - Ok(Err(err)) => cleanup_errors.push(format!("rank 0 cleanup: {err:#}")), - Err(_) => cleanup_errors - .push("rank 0 cleanup: dropped LoRA discard response".to_string()), - }, - Err(err) => cleanup_errors.push(format!("rank 0 cleanup dispatch: {err:#}")), - } - for (index, worker) in self.workers.iter().enumerate() { - let rank = index + 1; - match worker.discard_lora_adapter(request.lora_name.clone()) { - Ok(response) => match response.recv() { - Ok(Ok(())) => {} - Ok(Err(err)) => { - cleanup_errors.push(format!("rank {rank} cleanup: {err:#}")); - } - Err(_) => cleanup_errors.push(format!( - "rank {rank} cleanup: dropped LoRA discard response" - )), - }, - Err(err) => { - cleanup_errors.push(format!("rank {rank} cleanup dispatch: {err:#}")); - } - } - } - if cleanup_errors.is_empty() { - self.loaded_lora_adapters.remove(&request.lora_name); - } - let cleanup_suffix = if cleanup_errors.is_empty() { - String::new() - } else { - format!("; cleanup errors: {}", cleanup_errors.join("; ")) - }; - anyhow::bail!( - "failed to load Qwen3 LoRA adapter {} on tensor-parallel ranks: {}{}", - request.lora_name, - errors.join("; "), - cleanup_suffix - ); - } - - log::info!( - "Loaded Qwen3 LoRA adapter {} from {} (rank={}, targets={}, projections={}, bf16_elements={}, tp_world_size={}, load_inplace={})", - request.lora_name, - path, - rank, - targets, - projection_count, - element_count, - world_size, - request.load_inplace - ); - self.loaded_lora_adapters.insert(request.lora_name.clone()); - Ok(()) - } - - fn unload_lora_adapter(&mut self, request: &UnloadLoraAdapterRequest) -> Result<()> { - let primary_response = self - .primary - .unload_lora_adapter(request.lora_name.clone())?; - let mut pending = Vec::with_capacity(self.workers.len()); - for (index, worker) in self.workers.iter().enumerate() { - pending.push(( - index + 1, - worker.unload_lora_adapter(request.lora_name.clone())?, - )); - } - - let mut errors = Vec::new(); - match primary_response.recv() { - Ok(Ok(())) => {} - Ok(Err(err)) => errors.push(format!("rank 0: {err:#}")), - Err(_) => errors.push("rank 0: dropped LoRA unload response".to_string()), - } - for (rank, response) in pending { - match response.recv() { - Ok(Ok(())) => {} - Ok(Err(err)) => errors.push(format!("rank {rank}: {err:#}")), - Err(_) => errors.push(format!("rank {rank}: dropped LoRA unload response")), - } - } - if !errors.is_empty() { - anyhow::bail!( - "failed to unload Qwen3 LoRA adapter {} on tensor-parallel ranks: {}", - request.lora_name, - errors.join("; ") - ); - } - - log::info!("Unloaded Qwen3 LoRA adapter {}", request.lora_name); - self.loaded_lora_adapters.remove(&request.lora_name); - Ok(()) - } - - fn list_lora_adapters(&self) -> Vec { - let mut names: Vec<_> = self.loaded_lora_adapters.iter().cloned().collect(); - names.sort(); - names - } -} - -#[cfg(test)] -mod tests { - use super::ensure_lora_capacity; - use std::collections::HashSet; - - #[test] - fn lora_capacity_rejects_new_adapter_at_limit() { - let loaded = HashSet::from(["adapter-a".to_string()]); - - let error = ensure_lora_capacity(&loaded, "adapter-b", 1, false) - .expect_err("new adapter should exceed capacity") - .to_string(); - - assert!(error.contains("max_loras=1")); - assert!(error.contains("requested=adapter-b")); - } - - #[test] - fn lora_capacity_allows_existing_adapter_replacement_at_limit_with_load_inplace() { - let loaded = HashSet::from(["adapter-a".to_string()]); - - ensure_lora_capacity(&loaded, "adapter-a", 1, true) - .expect("existing adapter should fit with load_inplace"); - } - - #[test] - fn lora_capacity_rejects_duplicate_without_load_inplace() { - let loaded = HashSet::from(["adapter-a".to_string()]); - - let error = ensure_lora_capacity(&loaded, "adapter-a", 1, false) - .expect_err("duplicate without load_inplace should fail") - .to_string(); - - assert!(error.contains("already loaded")); - } -} - impl Drop for Qwen3Executor { fn drop(&mut self) { self.primary.shutdown(); @@ -1779,387 +972,69 @@ impl Drop for Qwen3Executor { } } -struct LocalQwen3Lane { - model: Qwen3Model, - kv_buffer: KvBuffer, - layout: KvLayout, - bufs: BatchDecodeBuffers, - sample_scratch: SamplingScratch, -} - -impl LocalQwen3Lane { - fn new( - model: Qwen3Model, - kv_buffer: KvBuffer, - total_blocks: usize, - padding_block_id: i32, - ) -> Result { - let buf_layout = kv_buffer.layout(); - let layout = KvLayout::new( - buf_layout.num_layers, - buf_layout.num_kv_heads, - buf_layout.head_dim, - buf_layout.page_size, - ); - let max_bucket = *BATCH_BUCKETS.last().unwrap(); - let bufs = BatchDecodeBuffers::new( - model.device_ctx(), - model.config().hidden_size, - model.local_q_dim(), - model.local_kv_dim(), - model.local_intermediate_size(), - model.config().vocab_size, - max_bucket, - total_blocks, - padding_block_id, - model.local_num_attention_heads(), - )?; - let sample_scratch = - SamplingScratch::new(model.device_ctx(), model.config().vocab_size, max_bucket)?; - Ok(Self { - model, - kv_buffer, - layout, - bufs, - sample_scratch, - }) - } - - fn bind(&self) -> Result { - bind_model_thread(&self.model)?; - tune_decode_gemm_algos(&self.model)?; - Ok(CublasThreadGuard) - } - - /// Pick one token per logits column (batched argmax for greedy rows, - /// per-row sampler otherwise). Grows the sampling scratch when a step - /// is wider than the decode bucket it was sized for. - fn select_step_tokens( - &mut self, - logits: &HiddenStates, - params: &[&SamplingParams], - random_vals: &[f32], - ) -> Result> { - if params.len() > self.sample_scratch.row_indices.len() { - self.sample_scratch = SamplingScratch::new( - self.model.device_ctx(), - self.model.config().vocab_size, - params.len(), - )?; - } - openinfer_core::ops::select_batch_tokens_into( - self.model.device_ctx(), - logits, - params, - random_vals, - &mut self.sample_scratch.row_indices, - &mut self.sample_scratch.argmax_partial_values, - &mut self.sample_scratch.argmax_partial_indices, - &mut self.sample_scratch.probs, - &mut self.sample_scratch.top1_values, - &mut self.sample_scratch.row_states, - &mut self.sample_scratch.valid, - &mut self.sample_scratch.out, - ) - } - - fn extract_logprobs( - &self, - logits: &DeviceVec, - sampled_token: u32, - top_k: usize, - ) -> Result { - let logits_f32 = logits.to_host(self.model.device_ctx())?; - compute_logprobs_from_cpu(&logits_f32, sampled_token, top_k) - .ok_or_else(|| anyhow::anyhow!("logprobs computation failed")) - } - - fn extract_prompt_logprobs( - &self, - all_logits: &HiddenStates, - prev_pos: usize, - target_token: u32, - top_k: usize, - ) -> Option { - openinfer_core::ops::extract_vec(self.model.device_ctx(), all_logits, prev_pos) - .ok() - .and_then(|logits_vec| { - let logits_f32 = logits_vec.to_host(self.model.device_ctx()).ok()?; - compute_logprobs_from_cpu(&logits_f32, target_token, top_k) - }) - } - - fn execute_prefill( - &mut self, - prompts: &[&[u32]], - kv_views: &[KvView], - lora_adapters: &[Option<&str>], - echo: bool, - ) -> Result<(HiddenStates, Option)> { - self.model.batch_prefill( - prompts, - kv_views, - lora_adapters, - self.kv_buffer.buffer(), - &self.layout, - echo, - ) - } - - fn execute_decode( - &mut self, - token_ids: &[u32], - kv_views: &[KvView], - lora_adapters: &[Option<&str>], - ) -> Result<()> { - self.model.batch_decode( - token_ids, - kv_views, - lora_adapters, - self.kv_buffer.buffer(), - &self.layout, - &mut self.bufs, - ) - } - - fn execute_unified( - &mut self, - prefill_prompts: &[&[u32]], - prefill_views: &[KvView], - prefill_lora_adapters: &[Option<&str>], - decode_tokens: &[u32], - decode_views: &[KvView], - decode_lora_adapters: &[Option<&str>], - ) -> Result { - self.model.unified_step( - prefill_prompts, - prefill_views, - prefill_lora_adapters, - decode_tokens, - decode_views, - decode_lora_adapters, - self.kv_buffer.buffer(), - &self.layout, - ) - } - - fn load_lora_adapter( - &mut self, - name: String, - adapter: crate::lora::LoraAdapter, - load_inplace: bool, - ) -> Result<()> { - let device_adapter = - crate::lora::load_device_lora_adapter(self.model.device_ctx(), name, adapter)?; - self.model - .install_lora_adapter(device_adapter, load_inplace) - } - - fn unload_lora_adapter(&mut self, name: &str) -> Result<()> { - self.model.uninstall_lora_adapter(name) - } - - fn discard_lora_adapter(&mut self, name: &str) -> Result<()> { - self.model.discard_lora_adapter(name) - } -} - -#[derive(Clone)] -enum StepCommand { - Prefill { - requests: Vec, - kv_views: Vec, - echo: bool, - }, - Decode { - requests: Vec, - kv_views: Vec, - }, - Unified { - prefill_requests: Vec, - prefill_kv_views: Vec, - decode_requests: Vec, - decode_kv_views: Vec, - }, -} - -impl StepCommand { - fn kind(&self) -> &'static str { - match self { - Self::Prefill { .. } => "prefill", - Self::Decode { .. } => "decode", - Self::Unified { .. } => "unified", - } - } -} - -enum WorkerCommand { - RunStep { - step: StepCommand, - collect_result: bool, - resp: channel::Sender>, - }, - LoadLoraAdapter { - name: String, - adapter: crate::lora::LoraAdapter, - load_inplace: bool, - resp: channel::Sender>, - }, - UnloadLoraAdapter { - name: String, - resp: channel::Sender>, - }, - DiscardLoraAdapter { - name: String, - resp: channel::Sender>, - }, - Shutdown, -} - -enum WorkerStepOutcome { - Ack, - Prefill(PrefillResult), - Decode(DecodeResult), - Unified(UnifiedResult), -} - -impl WorkerStepOutcome { - fn kind(&self) -> &'static str { - match self { - Self::Ack => "ack", - Self::Prefill(_) => "prefill", - Self::Decode(_) => "decode", - Self::Unified(_) => "unified", - } - } -} - -struct RankWorker { - tx: channel::Sender, - handle: Option>, -} +#[cfg(test)] +mod tests { + use super::{ + DFLASH_MIN_EXTRA_MEMORY_RESERVE_BYTES, dflash_request_state_footprint_bytes_for_len, + dflash_request_state_reserve_bytes, + }; + use crate::config::{DFlashConfig, DFlashInnerConfig}; -impl RankWorker { - fn spawn(rank: usize, mut lane: LocalQwen3Lane) -> Result { - let (tx, rx) = channel::unbounded(); - let (startup_tx, startup_rx) = channel::bounded(1); - let handle = thread::Builder::new() - .name(format!("qwen3-tp-rank-{rank}")) - .spawn(move || { - let startup = lane.bind(); - match startup { - Ok(_guard) => { - let _ = startup_tx.send(Ok(())); - while let Ok(cmd) = rx.recv() { - match cmd { - WorkerCommand::RunStep { - step, - collect_result, - resp, - } => { - let result = - execute_step_on_lane(&mut lane, &step, collect_result); - let _ = resp.send(result); - } - WorkerCommand::LoadLoraAdapter { - name, - adapter, - load_inplace, - resp, - } => { - let result = - lane.load_lora_adapter(name, adapter, load_inplace); - let _ = resp.send(result); - } - WorkerCommand::UnloadLoraAdapter { name, resp } => { - let result = lane.unload_lora_adapter(&name); - let _ = resp.send(result); - } - WorkerCommand::DiscardLoraAdapter { name, resp } => { - let result = lane.discard_lora_adapter(&name); - let _ = resp.send(result); - } - WorkerCommand::Shutdown => break, - } - } - } - Err(err) => { - let _ = startup_tx.send(Err(err)); - } - } - }) - .map_err(|e| anyhow::anyhow!("failed to spawn tensor-parallel worker {rank}: {e}"))?; - startup_rx.recv().map_err(|_| { - anyhow::anyhow!("tensor-parallel worker {rank} exited during startup") - })??; - Ok(Self { - tx, - handle: Some(handle), - }) - } + #[test] + fn dflash_request_state_reserve_covers_long_context_scratch() { + let config = DFlashConfig { + hidden_size: 2560, + intermediate_size: 9728, + num_hidden_layers: 5, + num_attention_heads: 32, + num_key_value_heads: 8, + num_target_layers: 5, + head_dim: 128, + vocab_size: 151936, + rms_norm_eps: 1e-6, + rope_theta: 1_000_000.0, + max_position_embeddings: 40960, + block_size: 16, + dflash_config: DFlashInnerConfig { + mask_token_id: 151669, + target_layer_ids: vec![1, 9, 17, 25, 33], + }, + }; - fn run_step( - &self, - step: StepCommand, - collect_result: bool, - ) -> Result>> { - let (resp_tx, resp_rx) = channel::bounded(1); - self.tx - .send(WorkerCommand::RunStep { - step, - collect_result, - resp: resp_tx, - }) - .map_err(|_| anyhow::anyhow!("tensor-parallel worker step channel closed"))?; - Ok(resp_rx) - } + let reserve = dflash_request_state_reserve_bytes(&config).unwrap(); - fn load_lora_adapter( - &self, - name: String, - adapter: crate::lora::LoraAdapter, - load_inplace: bool, - ) -> Result>> { - let (resp_tx, resp_rx) = channel::bounded(1); - self.tx - .send(WorkerCommand::LoadLoraAdapter { - name, - adapter, - load_inplace, - resp: resp_tx, - }) - .map_err(|_| anyhow::anyhow!("tensor-parallel worker channel closed on LoRA load"))?; - Ok(resp_rx) + assert!( + reserve > DFLASH_MIN_EXTRA_MEMORY_RESERVE_BYTES * 2, + "Qwen3-4B DFlash reserve should account for long-context request state, got {reserve}" + ); } - fn unload_lora_adapter(&self, name: String) -> Result>> { - let (resp_tx, resp_rx) = channel::bounded(1); - self.tx - .send(WorkerCommand::UnloadLoraAdapter { - name, - resp: resp_tx, - }) - .map_err(|_| anyhow::anyhow!("tensor-parallel worker channel closed on LoRA unload"))?; - Ok(resp_rx) - } + #[test] + fn dflash_short_request_footprint_does_not_include_global_allocator_floor() { + let config = DFlashConfig { + hidden_size: 2560, + intermediate_size: 9728, + num_hidden_layers: 5, + num_attention_heads: 32, + num_key_value_heads: 8, + num_target_layers: 5, + head_dim: 128, + vocab_size: 151936, + rms_norm_eps: 1e-6, + rope_theta: 1_000_000.0, + max_position_embeddings: 40960, + block_size: 16, + dflash_config: DFlashInnerConfig { + mask_token_id: 151669, + target_layer_ids: vec![1, 9, 17, 25, 33], + }, + }; - fn discard_lora_adapter(&self, name: String) -> Result>> { - let (resp_tx, resp_rx) = channel::bounded(1); - self.tx - .send(WorkerCommand::DiscardLoraAdapter { - name, - resp: resp_tx, - }) - .map_err(|_| { - anyhow::anyhow!("tensor-parallel worker channel closed on LoRA discard") - })?; - Ok(resp_rx) - } + let footprint = dflash_request_state_footprint_bytes_for_len(&config, 80).unwrap(); - fn shutdown(&mut self) { - let _ = self.tx.send(WorkerCommand::Shutdown); - if let Some(handle) = self.handle.take() { - let _ = handle.join(); - } + assert!( + footprint < DFLASH_MIN_EXTRA_MEMORY_RESERVE_BYTES / 8, + "short-request admission must charge true DFlash state footprint, not the global floor: {footprint}" + ); } } diff --git a/openinfer-qwen3-4b/src/executor/dflash_lane.rs b/openinfer-qwen3-4b/src/executor/dflash_lane.rs new file mode 100644 index 00000000..852186da --- /dev/null +++ b/openinfer-qwen3-4b/src/executor/dflash_lane.rs @@ -0,0 +1,392 @@ +use std::collections::HashMap; +use std::env; +use std::time::Instant; + +use anyhow::Result; + +use crate::dflash::{DFlashDraftModel, DFlashRequestState}; +use crate::speculative::{ + DraftRequestResult as SpeculativeDraftRequestResult, DraftResult as SpeculativeDraftResult, + DraftStepItem as SpeculativeDraftStepItem, + VerifyRequestResult as SpeculativeVerifyRequestResult, + VerifyStepItem as SpeculativeVerifyStepItem, +}; +use openinfer_core::tensor::HiddenStates; + +use super::dflash_prefill::{ + dflash_prefill_can_capture, + should_capture_dflash_prefill_context as should_capture_dflash_prefill_context_with_state, +}; +use super::worker::LocalQwen3Lane; +use super::{PrefillStepItem, RequestId}; + +pub(super) struct DFlashLaneState { + pub(super) model: DFlashDraftModel, + pub(super) requests: HashMap, + verified_tokens: usize, + accepted_tokens: usize, + draft_steps: usize, + draft_ms_total: f64, + verify_steps: usize, + verify_ms_total: f64, + nvtx_enabled: bool, + last_draft_ms: HashMap, + last_draft_context_tokens: HashMap, + last_draft_committed_context: HashMap, +} + +impl DFlashLaneState { + pub(super) fn new(model: DFlashDraftModel) -> Self { + let nvtx_enabled = dflash_nvtx_enabled_from_env(); + if nvtx_enabled { + log::info!("Qwen3 DFlash NVTX ranges enabled"); + } + Self { + model, + requests: HashMap::new(), + verified_tokens: 0, + accepted_tokens: 0, + draft_steps: 0, + draft_ms_total: 0.0, + verify_steps: 0, + verify_ms_total: 0.0, + nvtx_enabled, + last_draft_ms: HashMap::new(), + last_draft_context_tokens: HashMap::new(), + last_draft_committed_context: HashMap::new(), + } + } +} + +impl LocalQwen3Lane { + pub(super) fn dflash_capture_layer_ids(&self) -> Option> { + self.dflash + .as_ref() + .map(|dflash| dflash.model.target_layer_ids().to_vec()) + } + + pub(super) fn dflash_nvtx_enabled(&self) -> bool { + self.dflash + .as_ref() + .is_some_and(|dflash| dflash.nvtx_enabled) + } + + pub(super) fn should_capture_dflash_prefill_context( + &self, + requests: &[PrefillStepItem], + ) -> bool { + let Some(dflash) = self.dflash.as_ref() else { + return false; + }; + should_capture_dflash_prefill_context_with_state(requests, |request_id| { + dflash.requests.contains_key(&request_id) + }) + } + + pub(super) fn start_dflash_timing(&self) -> Result> { + if !self.dflash_nvtx_enabled() { + return Ok(None); + } + self.model.device_ctx().sync()?; + Ok(Some(Instant::now())) + } + + pub(super) fn finish_dflash_timing(&self, start: Option) -> Result> { + let Some(start) = start else { + return Ok(None); + }; + self.model.device_ctx().sync()?; + Ok(Some(start.elapsed().as_secs_f64() * 1000.0)) + } + + pub(super) fn record_prefill_dflash_context( + &mut self, + requests: &[PrefillStepItem], + capture_requested: bool, + captured_hidden: Option<&HiddenStates>, + ) -> Result> { + let Some(captured_hidden) = captured_hidden else { + anyhow::ensure!( + !capture_requested, + "DFlash prefill context capture was requested but no hidden states were returned" + ); + return Ok(Vec::new()); + }; + anyhow::ensure!( + capture_requested, + "DFlash prefill hidden states were returned without a capture request" + ); + let Some(dflash) = self.dflash.as_mut() else { + anyhow::bail!("DFlash prefill context record requested without DFlash"); + }; + let expected_tokens: usize = requests.iter().map(|req| req.chunk_tokens).sum(); + anyhow::ensure!( + captured_hidden.seq_len == expected_tokens, + "DFlash prefill captured {} hidden rows for {} scheduled tokens", + captured_hidden.seq_len, + expected_tokens + ); + let ctx = self.model.device_ctx().clone(); + let mut captured_requests = Vec::new(); + let mut token_offset = 0usize; + for req in requests { + let pending_exists = dflash.requests.contains_key(&req.request_id); + if dflash_prefill_can_capture(req, pending_exists) { + let max_cache_len = + req.prompt_tokens.len() + req.max_output_tokens + dflash.model.block_size(); + let mut state = match dflash.requests.remove(&req.request_id) { + Some(state) => state, + None => dflash.model.new_request_state(&ctx, max_cache_len)?, + }; + let pending_len = state.pending_context_len().unwrap_or(0); + anyhow::ensure!( + pending_len == req.chunk_start, + "DFlash prefill context for {:?} is discontinuous: pending={}, chunk_start={}", + req.request_id, + pending_len, + req.chunk_start + ); + dflash.model.append_pending_context( + &ctx, + &mut state, + captured_hidden, + token_offset, + req.chunk_tokens, + )?; + dflash.requests.insert(req.request_id, state); + captured_requests.push(req.request_id); + } else { + dflash.requests.remove(&req.request_id); + } + token_offset += req.chunk_tokens; + } + Ok(captured_requests) + } + + pub(super) fn record_verify_dflash_context( + &mut self, + requests: &[SpeculativeVerifyStepItem], + results: &[SpeculativeVerifyRequestResult], + captured_hidden: Option<&HiddenStates>, + verify_ms: Option, + ) -> Result<()> { + let Some(captured_hidden) = captured_hidden else { + anyhow::bail!( + "DFlash speculative verify context capture was requested but no hidden states were returned" + ); + }; + let Some(dflash) = self.dflash.as_mut() else { + anyhow::bail!("DFlash speculative verify context record requested without DFlash"); + }; + anyhow::ensure!( + requests.len() == results.len(), + "DFlash speculative verify result count {} does not match request count {}", + results.len(), + requests.len() + ); + let expected_tokens: usize = requests.iter().map(|req| req.token_ids.len()).sum(); + anyhow::ensure!( + captured_hidden.seq_len == expected_tokens, + "DFlash speculative verify captured {} hidden rows for {} scheduled tokens", + captured_hidden.seq_len, + expected_tokens + ); + for (req, result) in requests.iter().zip(results) { + anyhow::ensure!( + req.request_id == result.request_id, + "DFlash speculative verify result {:?} does not match request {:?}", + result.request_id, + req.request_id + ); + anyhow::ensure!( + dflash.requests.contains_key(&req.request_id), + "missing DFlash state after speculative verify for {:?}", + req.request_id + ); + } + if let Some(ms) = verify_ms { + dflash.verify_steps += 1; + dflash.verify_ms_total += ms; + } + let ctx = self.model.device_ctx().clone(); + let mut token_offset = 0usize; + for (req, result) in requests.iter().zip(results) { + let Some(mut state) = dflash.requests.remove(&req.request_id) else { + dflash.last_draft_ms.remove(&req.request_id); + dflash.last_draft_context_tokens.remove(&req.request_id); + dflash.last_draft_committed_context.remove(&req.request_id); + anyhow::bail!( + "missing DFlash state after speculative verify for {:?}", + req.request_id + ); + }; + dflash.model.append_pending_context( + &ctx, + &mut state, + captured_hidden, + token_offset, + result.accepted_tokens.len(), + )?; + dflash.requests.insert(req.request_id, state); + dflash.verified_tokens += req.token_ids.len().saturating_sub(1); + dflash.accepted_tokens += result.matched_draft_tokens; + let rate = if dflash.verified_tokens == 0 { + 0.0 + } else { + dflash.accepted_tokens as f64 / dflash.verified_tokens as f64 + }; + let draft_ms = dflash.last_draft_ms.remove(&req.request_id).unwrap_or(-1.0); + let draft_context_tokens = dflash + .last_draft_context_tokens + .remove(&req.request_id) + .unwrap_or(0); + let draft_committed_context = dflash + .last_draft_committed_context + .remove(&req.request_id) + .unwrap_or(0); + let avg_draft_ms = if dflash.draft_steps == 0 { + -1.0 + } else { + dflash.draft_ms_total / dflash.draft_steps as f64 + }; + let avg_verify_ms = if dflash.verify_steps == 0 { + -1.0 + } else { + dflash.verify_ms_total / dflash.verify_steps as f64 + }; + log::info!( + "Qwen3 DFlash request={} accepted_draft={} verified_draft={} committed_tokens={} cumulative_accept_rate={:.3} draft_ms={:.3} verify_ms={:.3} avg_draft_ms={:.3} avg_verify_ms={:.3} draft_context_tokens={} draft_committed_context={}", + req.request_id.get(), + result.matched_draft_tokens, + req.token_ids.len().saturating_sub(1), + result.accepted_tokens.len(), + rate, + draft_ms, + verify_ms.unwrap_or(-1.0), + avg_draft_ms, + avg_verify_ms, + draft_context_tokens, + draft_committed_context, + ); + token_offset += req.token_ids.len(); + } + Ok(()) + } + + pub(super) fn execute_dflash_draft( + &mut self, + requests: &[SpeculativeDraftStepItem], + ) -> Result { + anyhow::ensure!( + !requests.is_empty(), + "DFlash draft requested without active requests" + ); + for req in requests { + anyhow::ensure!( + req.params.is_greedy(), + "DFlash draft currently supports greedy sampling only" + ); + } + + let ctx = self.model.device_ctx().clone(); + let profiling_enabled = self.dflash_nvtx_enabled(); + let draft_nvtx_range = if self.dflash_nvtx_enabled() { + Some(nvtx::range!("qwen3.dflash.draft")) + } else { + None + }; + let Some(mut dflash) = self.dflash.take() else { + anyhow::bail!("DFlash draft requested but DFlash is not loaded"); + }; + let draft_results = + (|| -> Result)>> { + let mut outputs = Vec::with_capacity(requests.len()); + for req in requests { + if profiling_enabled { + ctx.sync()?; + } + let draft_start = profiling_enabled.then(Instant::now); + let mut state = dflash.requests.remove(&req.request_id).ok_or_else(|| { + anyhow::anyhow!("missing DFlash state for {:?}", req.request_id) + })?; + state.pending_context_len().ok_or_else(|| { + anyhow::anyhow!( + "DFlash draft requested before target hidden context is available" + ) + })?; + let draft = + dflash + .model + .draft_logits(&self.model, &mut state, req.current_token)?; + let draft_len = draft.logits.seq_len; + let sampled = self.select_greedy_contiguous_tokens(draft.logits)?; + let context_tokens = draft.context_len; + let committed_context = draft.committed_len; + dflash.requests.insert(req.request_id, state); + if profiling_enabled { + ctx.sync()?; + } + let draft_ms = draft_start.map(|start| start.elapsed().as_secs_f64() * 1000.0); + anyhow::ensure!( + sampled.len() == draft_len && sampled.len() >= 2, + "DFlash draft sampled {} tokens from {} logits columns", + sampled.len(), + draft_len + ); + let mut token_ids = Vec::with_capacity(sampled.len()); + token_ids.push(req.current_token); + token_ids.extend(sampled.into_iter().skip(1)); + outputs.push(( + SpeculativeDraftRequestResult { + request_id: req.request_id, + token_ids, + }, + context_tokens, + committed_context, + draft_ms, + )); + } + Ok(outputs) + })(); + self.dflash = Some(dflash); + let draft_results = draft_results?; + drop(draft_nvtx_range); + if let Some(dflash) = self.dflash.as_mut() { + for (draft, context_tokens, committed_context, draft_ms) in &draft_results { + if let Some(draft_ms) = *draft_ms { + dflash.draft_steps += 1; + dflash.draft_ms_total += draft_ms; + dflash.last_draft_ms.insert(draft.request_id, draft_ms); + } + dflash + .last_draft_context_tokens + .insert(draft.request_id, *context_tokens); + dflash + .last_draft_committed_context + .insert(draft.request_id, *committed_context); + } + } + Ok(SpeculativeDraftResult { + requests: draft_results + .into_iter() + .map(|(draft, _, _, _)| draft) + .collect(), + }) + } + + pub(super) fn drop_dflash_request(&mut self, request_id: RequestId) { + if let Some(dflash) = self.dflash.as_mut() { + dflash.requests.remove(&request_id); + dflash.last_draft_ms.remove(&request_id); + dflash.last_draft_context_tokens.remove(&request_id); + dflash.last_draft_committed_context.remove(&request_id); + } + } +} + +fn dflash_nvtx_enabled_from_env() -> bool { + matches!( + env::var("OPENINFER_QWEN3_DFLASH_NVTX").ok().as_deref(), + Some("1" | "true" | "TRUE" | "yes" | "YES" | "on" | "ON") + ) +} diff --git a/openinfer-qwen3-4b/src/executor/dflash_prefill.rs b/openinfer-qwen3-4b/src/executor/dflash_prefill.rs new file mode 100644 index 00000000..bf2b99a5 --- /dev/null +++ b/openinfer-qwen3-4b/src/executor/dflash_prefill.rs @@ -0,0 +1,162 @@ +use super::{PrefillStepItem, RequestId}; + +pub(super) fn dflash_prefill_supported(req: &PrefillStepItem) -> bool { + req.lora_adapter.is_none() + && req.cached_tokens == 0 + && req.logprobs == 0 + && !req.echo + && req.params.is_greedy() +} + +pub(super) fn dflash_prefill_can_capture( + req: &PrefillStepItem, + pending_state_exists: bool, +) -> bool { + dflash_prefill_supported(req) && (req.chunk_start == 0 || pending_state_exists) +} + +pub(super) fn should_capture_dflash_prefill_context( + requests: &[PrefillStepItem], + pending_state_exists: impl Fn(RequestId) -> bool, +) -> bool { + !requests.is_empty() + && requests + .iter() + .any(|req| dflash_prefill_can_capture(req, pending_state_exists(req.request_id))) +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(super) enum DFlashPrefillAction { + MarkReady, + KeepPending, + Drop, +} + +pub(super) fn dflash_prefill_action( + captured_context: bool, + completed: bool, +) -> DFlashPrefillAction { + match (captured_context, completed) { + (true, true) => DFlashPrefillAction::MarkReady, + (true, false) => DFlashPrefillAction::KeepPending, + (false, _) => DFlashPrefillAction::Drop, + } +} + +#[cfg(test)] +mod tests { + use super::{ + DFlashPrefillAction, dflash_prefill_action, dflash_prefill_can_capture, + should_capture_dflash_prefill_context, + }; + use crate::executor::{PrefillStepItem, RequestId}; + use openinfer_core::sampler::SamplingParams; + + #[test] + fn dflash_prefill_capture_runs_when_any_request_can_capture() { + let greedy = SamplingParams::default(); + let mut non_greedy_params = greedy; + non_greedy_params.temperature = 0.7; + + let supported = PrefillStepItem { + request_id: RequestId::new(1), + prompt_tokens: vec![1, 2, 3], + max_output_tokens: 8, + params: greedy, + logprobs: 0, + echo: false, + lora_adapter: None, + random_val: 0.0, + cached_tokens: 0, + chunk_budget: 3, + chunk_start: 0, + chunk_tokens: 0, + }; + let mut second = supported.clone(); + second.request_id = RequestId::new(2); + let mut non_greedy = supported.clone(); + non_greedy.params = non_greedy_params; + + assert!(should_capture_dflash_prefill_context( + std::slice::from_ref(&supported), + |_| false + )); + let mut cached = supported.clone(); + cached.cached_tokens = 1; + assert!(should_capture_dflash_prefill_context( + &[supported.clone(), second], + |_| false + )); + assert!(should_capture_dflash_prefill_context( + &[supported.clone(), non_greedy.clone()], + |_| false + )); + assert!(!should_capture_dflash_prefill_context( + &[non_greedy], + |_| false + )); + assert!(!should_capture_dflash_prefill_context(&[cached], |_| false)); + + let mut mid_chunk = supported.clone(); + mid_chunk.chunk_start = 2; + mid_chunk.chunk_tokens = 1; + assert!(!dflash_prefill_can_capture(&mid_chunk, false)); + assert!(dflash_prefill_can_capture(&mid_chunk, true)); + } + + #[test] + fn dflash_prefill_capture_allows_multi_request_start_chunks() { + let supported = PrefillStepItem { + request_id: RequestId::new(1), + prompt_tokens: vec![1, 2, 3], + max_output_tokens: 8, + params: SamplingParams::default(), + logprobs: 0, + echo: false, + lora_adapter: None, + random_val: 0.0, + cached_tokens: 0, + chunk_budget: 3, + chunk_start: 0, + chunk_tokens: 0, + }; + let mut second = supported.clone(); + second.request_id = RequestId::new(2); + + assert!(should_capture_dflash_prefill_context( + &[supported.clone(), second.clone()], + |_| false + )); + + let mut mid_chunk = second; + mid_chunk.chunk_start = 3; + assert!(should_capture_dflash_prefill_context( + &[supported.clone(), mid_chunk.clone()], + |_| false + )); + assert!(should_capture_dflash_prefill_context( + &[supported, mid_chunk], + |id| id == RequestId::new(2) + )); + } + + #[test] + fn dflash_prefill_keeps_state_until_final_chunk() { + assert_eq!( + dflash_prefill_action(true, false), + DFlashPrefillAction::KeepPending + ); + assert_eq!( + dflash_prefill_action(true, true), + DFlashPrefillAction::MarkReady + ); + assert_eq!( + dflash_prefill_action(false, false), + DFlashPrefillAction::Drop + ); + assert_eq!( + dflash_prefill_action(false, true), + DFlashPrefillAction::Drop + ); + } +} diff --git a/openinfer-qwen3-4b/src/executor/lifecycle.rs b/openinfer-qwen3-4b/src/executor/lifecycle.rs new file mode 100644 index 00000000..564d9099 --- /dev/null +++ b/openinfer-qwen3-4b/src/executor/lifecycle.rs @@ -0,0 +1,627 @@ +use anyhow::Result; +use crossbeam_channel as channel; + +use crate::config::TensorParallelConfig; +use crate::dflash::DFlashDraftModel; +use crate::weights::{ModelRuntimeConfig, Qwen3Model}; +use crate::{Qwen3LoraOptions, Qwen3OffloadOptions, Qwen3SpeculativeOptions}; +use openinfer_core::tensor::DeviceContext; +use openinfer_kv_cache::{KvBlockGuard, KvBuffer, KvCacheManager}; +use openinfer_kv_offload::{OffloadConfig, OffloadEngine}; + +use super::worker::{LocalQwen3Lane, RankWorker, StepCommand, WorkerStepOutcome}; +use super::{ + DecodePlan, DecodeResult, ModelExecutor, PrefillPlan, PrefillRequestResult, PrefillResult, + PrefillStepItem, Qwen3Executor, Qwen3ExecutorMetadata, RequestId, SpeculativeDraftPlan, + SpeculativeDraftResult, SpeculativeVerifyPlan, SpeculativeVerifyResult, UnifiedPlan, + UnifiedResult, dflash_memory_reserve, +}; + +impl Qwen3Executor { + fn single( + model: Qwen3Model, + offload_opts: &Qwen3OffloadOptions, + speculative_options: Qwen3SpeculativeOptions, + ) -> Result { + let dflash_reserve = dflash_memory_reserve(&speculative_options)?; + let budget = model.kv_budget_with_reserved_bytes(dflash_reserve.total_bytes); + let kv_mgr = KvCacheManager::new( + &model.device_ctx().stream, + budget.num_layers, + budget.num_kv_heads, + budget.head_dim, + budget.block_size, + budget.num_blocks, + )?; + let kv_buffer = kv_mgr.buffer().clone(); + // Build the offload engine while the model's stream is still in hand + // (it moves into the RankWorker below). Registers the fused KV buffer. + let offload = build_offload(offload_opts, &kv_mgr, model.device_ctx())?; + let total_blocks = kv_mgr.pool().total_blocks(); + let padding_block_id = kv_mgr.pool().padding_block_id(); + let dflash = match speculative_options.dflash { + Some(options) => { + let path = options + .model_path + .to_str() + .ok_or_else(|| anyhow::anyhow!("DFlash model path must be valid UTF-8"))?; + Some(DFlashDraftModel::from_safetensors_for_target( + model.device_ctx(), + path, + &model, + )?) + } + None => None, + }; + let speculative_enabled = dflash.is_some(); + let max_context_tokens = effective_max_context_tokens( + model.config().max_position_embeddings, + dflash.as_ref().map(DFlashDraftModel::block_size), + )?; + let metadata = Qwen3ExecutorMetadata { + block_size: budget.block_size, + stop_token_ids: model.config().stop_token_ids.clone(), + config: model.config().clone(), + max_context_tokens, + dflash_config: dflash.as_ref().map(|dflash| dflash.config().clone()), + dflash_state_budget_bytes: dflash_reserve.request_state_budget_bytes, + }; + if speculative_enabled { + log::info!( + "Qwen3 DFlash loaded; disabling prefix cache for hidden-state capture; max_context_tokens={}", + max_context_tokens + ); + } + Ok(Self { + metadata, + kv_mgr, + request_kvs: Default::default(), + primary: RankWorker::spawn( + 0, + LocalQwen3Lane::new(model, kv_buffer, total_blocks, padding_block_id, dflash)?, + )?, + workers: Vec::new(), + loaded_lora_adapters: Default::default(), + prefix_cache_enabled: !speculative_enabled, + lora_options: Qwen3LoraOptions::default(), + offload, + saved_cursor: Default::default(), + prefetch: Default::default(), + l1_retention_disabled: false, + speculative_enabled, + dflash_ready_requests: Default::default(), + }) + } + + pub fn from_runtime( + model_path: &str, + enable_cuda_graph: bool, + device_ordinals: &[usize], + ) -> Result { + Self::from_runtime_with_lora_options( + model_path, + enable_cuda_graph, + device_ordinals, + Qwen3LoraOptions::default(), + Qwen3OffloadOptions::disabled(), + ) + } + + pub fn from_runtime_with_lora_options( + model_path: &str, + enable_cuda_graph: bool, + device_ordinals: &[usize], + lora_options: Qwen3LoraOptions, + offload_options: Qwen3OffloadOptions, + ) -> Result { + Self::from_runtime_with_options( + model_path, + enable_cuda_graph, + device_ordinals, + lora_options, + offload_options, + Qwen3SpeculativeOptions::disabled(), + ) + } + + pub fn from_runtime_with_options( + model_path: &str, + enable_cuda_graph: bool, + device_ordinals: &[usize], + lora_options: Qwen3LoraOptions, + offload_options: Qwen3OffloadOptions, + speculative_options: Qwen3SpeculativeOptions, + ) -> Result { + let lora_options = lora_options.validate()?; + anyhow::ensure!( + !device_ordinals.is_empty(), + "Qwen3 executor requires at least one device" + ); + anyhow::ensure!( + speculative_options.dflash.is_none() + || (device_ordinals.len() == 1 + && lora_options == Qwen3LoraOptions::default() + && !offload_options.enabled), + "DFlash speculative decoding currently requires the single-GPU base Qwen3 path without LoRA or KV offload" + ); + anyhow::ensure!( + !offload_options.enabled || device_ordinals.len() == 1, + "KV offload is only supported on the single-GPU path (tensor parallel \ + shards KV per rank); got {} devices", + device_ordinals.len() + ); + if device_ordinals.len() == 1 { + let model = Qwen3Model::from_safetensors_with_runtime( + model_path, + ModelRuntimeConfig { + enable_cuda_graph, + tensor_parallel: None, + device_ordinal: device_ordinals[0], + max_loras: lora_options.max_loras, + max_lora_rank: lora_options.max_lora_rank, + }, + )?; + let mut executor = Self::single(model, &offload_options, speculative_options)?; + executor.lora_options = lora_options; + return Ok(executor); + } + + let world_size = device_ordinals.len(); + let mut models = Vec::with_capacity(world_size); + for (rank, &device_ordinal) in device_ordinals.iter().enumerate() { + models.push(Qwen3Model::from_safetensors_with_runtime( + model_path, + ModelRuntimeConfig { + enable_cuda_graph, + tensor_parallel: Some(TensorParallelConfig { rank, world_size }), + device_ordinal, + max_loras: lora_options.max_loras, + max_lora_rank: lora_options.max_lora_rank, + }, + )?); + } + + // Compute budget from first model (all ranks share geometry). + let budget = models[0].kv_budget(); + + // Create the centralized KvCacheManager on rank 0's stream. + let kv_mgr = KvCacheManager::new( + &models[0].device_ctx().stream, + budget.num_layers, + budget.num_kv_heads, + budget.head_dim, + budget.block_size, + budget.num_blocks, + )?; + + let metadata = Qwen3ExecutorMetadata { + block_size: budget.block_size, + stop_token_ids: models[0].config().stop_token_ids.clone(), + config: models[0].config().clone(), + max_context_tokens: models[0].config().max_position_embeddings, + dflash_config: None, + dflash_state_budget_bytes: 0, + }; + + // Create extra KvBuffers for ranks 1+ on their respective streams. + let mut extra_kv_buffers = Vec::with_capacity(world_size - 1); + for model in &models[1..] { + extra_kv_buffers.push(KvBuffer::new( + &model.device_ctx().stream, + budget.num_layers, + budget.num_kv_heads, + budget.head_dim, + budget.block_size, + budget.num_blocks, + )?); + } + + let streams = models + .iter() + .map(|m| m.device_ctx().stream.clone()) + .collect(); + let comms = cudarc::nccl::safe::Comm::from_devices(streams) + .map_err(|e| anyhow::anyhow!("failed to initialize NCCL comms: {e:?}"))?; + for (model, comm) in models.iter_mut().zip(comms) { + model.attach_tp_comm(comm); + } + + let total_blocks = kv_mgr.pool().total_blocks(); + let padding_block_id = kv_mgr.pool().padding_block_id(); + + // Primary rank gets the KvBuffer from the centralized manager. + let primary_buffer = kv_mgr.buffer().clone(); + let mut models_iter = models.into_iter(); + let primary_model = models_iter.next().unwrap(); + let primary = RankWorker::spawn( + 0, + LocalQwen3Lane::new( + primary_model, + primary_buffer, + total_blocks, + padding_block_id, + None, + )?, + )?; + + // Worker ranks get their own extra KvBuffers. + let workers = models_iter + .zip(extra_kv_buffers) + .enumerate() + .map(|(index, (model, buffer))| { + let lane = + LocalQwen3Lane::new(model, buffer, total_blocks, padding_block_id, None)?; + RankWorker::spawn(index + 1, lane) + }) + .collect::>>()?; + + Ok(Self { + metadata, + kv_mgr, + request_kvs: Default::default(), + primary, + workers, + loaded_lora_adapters: Default::default(), + prefix_cache_enabled: true, + lora_options, + // Offload is single-GPU only (asserted above); never built here. + offload: None, + saved_cursor: Default::default(), + prefetch: Default::default(), + l1_retention_disabled: false, + speculative_enabled: false, + dflash_ready_requests: Default::default(), + }) + } + + pub fn block_size(&self) -> usize { + ::block_size(self) + } + + pub fn max_request_blocks(&self) -> usize { + ::max_request_blocks(self) + } + + pub fn available_blocks(&self) -> usize { + ::available_blocks(self) + } + + pub fn is_stop_token(&self, token_id: u32) -> bool { + ::is_stop_token(self, token_id) + } + + pub fn drop_request(&mut self, request_id: RequestId) -> Result<()> { + ::drop_request(self, request_id) + } + + pub fn execute_prefill(&mut self, plan: PrefillPlan<'_>) -> Result { + ::execute_prefill(self, plan) + } + + pub fn execute_decode(&mut self, plan: DecodePlan<'_>) -> Result { + ::execute_decode(self, plan) + } + + pub fn execute_speculative_verify( + &mut self, + plan: SpeculativeVerifyPlan<'_>, + ) -> Result { + ::execute_speculative_verify(self, plan) + } + + pub fn execute_speculative_draft( + &mut self, + plan: SpeculativeDraftPlan<'_>, + ) -> Result { + ::execute_speculative_draft(self, plan) + } + + pub fn execute_unified(&mut self, plan: UnifiedPlan<'_>) -> Result { + ::execute_unified(self, plan) + } + + /// Prefix caching is on by default; tests that assert bit-identical + /// replay disable it (a cache hit changes prefill GEMM shapes, which + /// drifts logits by bf16 ULPs). + pub fn set_prefix_cache_enabled(&mut self, enabled: bool) { + self.prefix_cache_enabled = enabled && !self.speculative_enabled; + } + + /// vLLM-style `--no-prefix-cache`. Behaviour depends on whether offload is + /// active: + /// * **No offload** — classic: disable prefix matching outright, so every + /// prefill recomputes the full prompt. + /// * **With offload** — pure-L2 mode: keep matching on (the host-tier + /// restore registers blocks and relies on `match_and_add_prefix` to pick + /// them up) but stop retaining completed blocks in HBM, so no request + /// ever serves its prefix from a cross-request L1 hit. Every reuse then + /// comes from the host tier, which is the point of the L2 benchmark. + /// + /// A resident HBM block and its host-tier copy share one content hash, so + /// the cache cannot be told to prefer L2 for a block still in HBM — the only + /// way to force the bytes from L2 is to not keep the HBM copy around. + pub fn set_no_prefix_cache(&mut self, on: bool) { + if self.speculative_enabled { + self.prefix_cache_enabled = false; + self.l1_retention_disabled = false; + return; + } + if self.offload.is_some() { + self.l1_retention_disabled = on; + } else { + self.prefix_cache_enabled = !on; + } + } + + /// Whether KV offload is active on this executor. + pub fn offload_enabled(&self) -> bool { + self.offload.is_some() + } + + /// Flush pending offload saves into the host read cache so a following + /// query can see them. A persistence barrier for handoff and tests; no-op + /// without offload. + pub fn flush_offload_saves(&self) { + if let Some(offload) = &self.offload { + offload.flush_saves(); + } + } + + /// Drop every cached-but-unused GPU prefix block. With offload on, this + /// forces a cold prefix to be restored from the host tier on its next + /// request (rather than served from HBM). + pub fn evict_cached_blocks(&self) { + self.kv_mgr.pool().evict_inactive(); + } + + /// Begin an async CPU-tier KV prefetch for `request_id`; see the + /// [`ModelExecutor`] hook. Public so admission drivers and tests can park a + /// request on its load. Returns `true` when a load is in flight. + pub fn begin_kv_prefetch( + &mut self, + request_id: RequestId, + prompt_tokens: &[u32], + lora_adapter: Option<&str>, + reserve_floor: usize, + ) -> bool { + ::begin_kv_prefetch( + self, + request_id, + prompt_tokens, + lora_adapter, + reserve_floor, + ) + } + + /// Block until at least one in-flight prefetch settles, then sweep the + /// rest; returns the settled request ids (now prefill-eligible). + pub fn wait_ready_prefetch(&mut self) -> Vec { + ::wait_ready_prefetch(self) + } + + // ── KV-offload SAVE ──────────────────────────────────────────────── + + /// Save every block that sealed since this request's last save to the host + /// tier (fire-and-forget). Safe to call right after `apply_prefill`/ + /// `apply_decode`: the producing step's token read-back has already + /// synchronized the compute stream, so the sealed KV is fully written. + pub(super) fn save_sealed_blocks(&mut self, request_id: RequestId) { + if self.offload.is_none() { + return; + } + let Some(rkv) = self.request_kvs.get(&request_id) else { + return; + }; + // `assigned_block_hashes` lists only sealed (registered) blocks; the + // partial tail block has no hash and never appears here. + let assigned = rkv.assigned_block_hashes(); + let prefix_matched = rkv.prefix_matched_blocks(); + let cursor = self + .saved_cursor + .entry(request_id) + .or_insert(prefix_matched); + if assigned.len() <= *cursor { + return; + } + let fresh = &assigned[*cursor..]; + let block_ids: Vec = fresh.iter().map(|(id, _)| *id).collect(); + let block_hashes: Vec> = fresh.iter().map(|(_, h)| h.to_vec()).collect(); + // Pin exactly the blocks being saved (aligned 1:1 with `assigned`) for + // the duration of the async D2H, so a finished request can't hand the + // slot to a new request that overwrites it before the copy lands. + let pins: Vec = rkv + .assigned_block_guards() + .into_iter() + .skip(*cursor) + .collect(); + *cursor = assigned.len(); + self.offload + .as_ref() + .expect("offload present") + .save(&block_ids, &block_hashes, pins); + } + + // ── Chunked prefill ──────────────────────────────────────────────── + + /// Prepare one prefill step for `req`: create its `RequestKv` on the + /// first chunk (matching the prefix cache), then clamp the scheduler's + /// chunk budget to the prompt tokens actually remaining and allocate KV + /// for them. Sets `chunk_start`/`chunk_tokens` on the item. + pub(super) fn schedule_prefill_chunk(&mut self, req: &mut PrefillStepItem) -> Result<()> { + if !self.request_kvs.contains_key(&req.request_id) { + let mut rkv = self.kv_mgr.pool().new_request( + req.prompt_tokens.clone(), + req.max_output_tokens, + req.lora_adapter.as_deref(), + ); + // Echo needs logits for every prompt position; cached positions + // are never forwarded, so echo requests prefill from scratch. + if self.prefix_cache_enabled && !req.echo { + req.cached_tokens = rkv.match_and_add_prefix(self.kv_mgr.pool())?; + } + self.request_kvs.insert(req.request_id, rkv); + // match_and_add_prefix above already absorbed any CPU-prefetched + // blocks (now held by the request's sequence), so release the + // prefetch's separate hold. + self.prefetch.remove(&req.request_id); + } + let rkv = self + .request_kvs + .get_mut(&req.request_id) + .expect("inserted above"); + req.chunk_start = rkv.kv_position(); + let remaining = req.prompt_tokens.len() - req.chunk_start; + // Echo must produce all-position logits in a single forward, so it is + // exempt from chunking (the scheduler never splits echo requests). + req.chunk_tokens = if req.echo { + remaining + } else { + remaining.min(req.chunk_budget) + }; + assert!( + req.chunk_tokens > 0, + "zero-token prefill chunk for {:?} (budget {})", + req.request_id, + req.chunk_budget + ); + rkv.schedule_prefill(req.chunk_tokens, self.kv_mgr.pool()) + .map_err(|e| anyhow::anyhow!("schedule_prefill failed for {:?}: {e}", req.request_id)) + } + + /// Register a finished prefill step on the request's KV: the final chunk + /// carries the first generated token, non-final chunks only advance the + /// KV position. + pub(super) fn apply_prefill_result(&mut self, result: &PrefillRequestResult) -> Result<()> { + let rkv = self + .request_kvs + .get_mut(&result.request_id) + .expect("request must exist after prefill"); + if result.completed { + rkv.apply_prefill(result.first_token, self.kv_mgr.pool()) + } else { + rkv.apply_prefill_chunk(self.kv_mgr.pool()) + } + } + + /// Finalize one prefetch whose load returned `result`. On success the + /// reserved blocks are staged + registered (held by the probe until the + /// request prefills); on failure the state is dropped so the request + /// prefills from scratch. + pub(super) fn settle_prefetch( + &mut self, + id: RequestId, + result: Result<(), openinfer_kv_offload::EngineError>, + ) { + if let Some(st) = self.prefetch.get_mut(&id) { + st.handle = None; + } + match result { + Ok(()) => { + let reservation = self + .prefetch + .get_mut(&id) + .and_then(|st| st.reservation.take()) + .expect("reservation present until commit"); + let st = self.prefetch.get_mut(&id).expect("prefetch present"); + self.kv_mgr + .pool() + .commit_loaded_blocks(&mut st.probe, reservation); + } + Err(e) => { + log::warn!("KV offload load failed for {id:?} (prefill from scratch): {e}"); + self.prefetch.remove(&id); + } + } + } + + fn wait_for_step_ack( + pending: Vec>>, + op_name: &'static str, + ) -> Result<()> { + for recv in pending { + match recv + .recv() + .map_err(|_| anyhow::anyhow!("tensor-parallel {op_name} worker dropped"))?? + { + WorkerStepOutcome::Ack => {} + other => { + return Err(anyhow::anyhow!( + "tensor-parallel {op_name} worker returned unexpected payload: {}", + other.kind() + )); + } + } + } + Ok(()) + } + + pub(super) fn run_step(&self, step: &StepCommand) -> Result { + let primary = self.primary.run_step(step.clone(), true)?; + let mut pending = Vec::with_capacity(self.workers.len()); + for worker in &self.workers { + pending.push(worker.run_step(step.clone(), false)?); + } + let primary_result = primary + .recv() + .map_err(|_| anyhow::anyhow!("primary worker dropped step response"))??; + Self::wait_for_step_ack(pending, step.kind())?; + Ok(primary_result) + } +} + +fn effective_max_context_tokens( + target_max_context_tokens: usize, + speculative_block_size: Option, +) -> Result { + match speculative_block_size { + Some(block_size) => { + anyhow::ensure!( + target_max_context_tokens > block_size, + "DFlash block_size {} leaves no usable context in target max_position_embeddings {}", + block_size, + target_max_context_tokens + ); + Ok(target_max_context_tokens - block_size) + } + None => Ok(target_max_context_tokens), + } +} + +/// Build the KV-offload engine for the single-GPU path, or `None` when offload +/// is disabled. Registers the fused KV buffer with pegaflow against the model's +/// device/stream — must be called while that stream is still owned by the model +/// (before it moves into the `RankWorker`). +fn build_offload( + opts: &Qwen3OffloadOptions, + kv_mgr: &KvCacheManager, + ctx: &DeviceContext, +) -> Result> { + if !opts.enabled { + return Ok(None); + } + let device_id = ctx.device_ordinal as i32; + let config = OffloadConfig::new( + format!("qwen3-4b-dev{device_id}"), + device_id, + opts.pinned_pool_bytes, + ); + let engine = OffloadEngine::new(config, kv_mgr.buffer(), &ctx.stream) + .map_err(|e| anyhow::anyhow!("KV offload engine init failed: {e}"))?; + log::info!( + "KV offload enabled on device {device_id} ({} MiB host tier)", + opts.pinned_pool_bytes >> 20 + ); + Ok(Some(engine)) +} + +#[cfg(test)] +mod tests { + use super::effective_max_context_tokens; + + #[test] + fn dflash_context_limit_reserves_one_draft_block() { + assert_eq!(effective_max_context_tokens(128, None).unwrap(), 128); + assert_eq!(effective_max_context_tokens(128, Some(16)).unwrap(), 112); + assert!(effective_max_context_tokens(16, Some(16)).is_err()); + } +} diff --git a/openinfer-qwen3-4b/src/executor/model_executor.rs b/openinfer-qwen3-4b/src/executor/model_executor.rs new file mode 100644 index 00000000..40d0ea07 --- /dev/null +++ b/openinfer-qwen3-4b/src/executor/model_executor.rs @@ -0,0 +1,691 @@ +use std::collections::HashSet; + +use anyhow::Result; + +use crate::batch_decode_buffers::BATCH_BUCKETS; +use crate::config::TensorParallelConfig; +use openinfer_core::engine::{LoadLoraAdapterRequest, UnloadLoraAdapterRequest}; +use openinfer_kv_cache::KvView; + +use super::dflash_prefill::{DFlashPrefillAction, dflash_prefill_action}; +use super::worker::{StepCommand, WorkerStepOutcome}; +use super::{ + DecodePlan, DecodeResult, ModelExecutor, PrefillPlan, PrefillResult, Qwen3Executor, RequestId, + SpeculativeDraftPlan, SpeculativeDraftResult, SpeculativeVerifyPlan, SpeculativeVerifyResult, + UnifiedPlan, UnifiedResult, dflash_request_state_footprint_bytes_for_len, +}; + +impl ModelExecutor for Qwen3Executor { + fn block_size(&self) -> usize { + self.metadata.block_size + } + + fn max_request_blocks(&self) -> usize { + self.kv_mgr.pool().max_request_blocks() + } + + fn max_context_tokens(&self) -> usize { + self.metadata.max_context_tokens + } + + fn max_decode_batch_size(&self) -> usize { + *BATCH_BUCKETS.last().unwrap() + } + + fn available_blocks(&self) -> usize { + self.kv_mgr.pool().available_blocks() + } + + fn is_stop_token(&self, token_id: u32) -> bool { + self.metadata.stop_token_ids.contains(&token_id) + } + + fn prefetched_blocks(&self, request_id: RequestId) -> usize { + self.prefetch + .get(&request_id) + .map_or(0, |st| st.probe.held_blocks()) + } + + fn drop_request(&mut self, request_id: RequestId) -> Result<()> { + // Remove and drop — RAII on SchedulableSequence's block guards + // returns all allocated blocks regardless of lifecycle state. The same + // RAII frees any parked prefetch's reserved/held blocks. + self.request_kvs.remove(&request_id); + // A parked prefetch may still have a load in flight: pegaflow's worker + // is writing the reserved GPU blocks (H2D). Dropping the reservation now + // frees those physical pages for immediate reuse while the DMA keeps + // landing on them — silent KV corruption, the load-side mirror of the + // SAVE keep-alive pin. Block until the copy finishes before the + // reservation drops. The scheduler is a dedicated synchronous thread, so + // this brief wait costs nothing it could spend elsewhere. + if let Some(mut state) = self.prefetch.remove(&request_id) { + if let Some(handle) = state.handle.take() { + let _ = handle.wait(); + } + } + self.saved_cursor.remove(&request_id); + if self.speculative_enabled { + self.dflash_ready_requests.remove(&request_id); + self.primary.drop_dflash_request(request_id)?; + } + Ok(()) + } + + fn begin_kv_prefetch( + &mut self, + request_id: RequestId, + prompt_tokens: &[u32], + lora_adapter: Option<&str>, + reserve_floor: usize, + ) -> bool { + let Some(offload) = self.offload.as_ref() else { + return false; + }; + if !self.prefix_cache_enabled { + return false; + } + if self.l1_retention_disabled { + // Pure-L2 mode: drop any cross-request HBM retention so the probe + // sees gpu_hit == 0 and queries the whole cacheable prefix from the + // host tier. Only inactive (completed, unheld) blocks are drained — + // the current request holds nothing yet, and in-flight prefetches + // keep their reserved blocks, so this never touches live KV. + self.kv_mgr.pool().evict_inactive(); + } + let probe = self + .kv_mgr + .pool() + .probe_prefix(prompt_tokens.to_vec(), lora_adapter); + let query_hashes = probe.cpu_query_hashes(); + if query_hashes.is_empty() { + return false; + } + let hit = match offload.query(&request_id.0.to_string(), &query_hashes) { + Ok(hit) => hit, + Err(e) => { + log::warn!("KV offload query failed for {request_id:?} (skipping): {e}"); + return false; + } + }; + let (Some(lease), num_blocks) = (hit.lease, hit.num_blocks) else { + return false; // miss + }; + // Blocks promised to admitted requests are off-limits: reserving into + // them makes a later prefill chunk or decode growth fail allocation. + if self + .kv_mgr + .pool() + .available_blocks() + .saturating_sub(reserve_floor) + < num_blocks + { + offload.release_query_lease(lease); + return false; + } + let Some(reservation) = self.kv_mgr.pool().reserve_loaded_blocks(num_blocks) else { + // Block pressure: release the lease so its pinned host blocks aren't + // held for the full lease TTL, and prefill from scratch rather than + // stall. + offload.release_query_lease(lease); + return false; + }; + let page_ids = reservation.page_ids(); + let handle = match offload.load(lease, page_ids) { + Ok(handle) => handle, + Err(e) => { + log::warn!("KV offload load submit failed for {request_id:?} (skipping): {e}"); + // `load` consumes the lease only past its early validation; a + // submit error may leave it pinned, so release it (no-op if it + // was already consumed). + offload.release_query_lease(lease); + return false; + } + }; + self.prefetch.insert( + request_id, + super::PrefetchState { + probe, + reservation: Some(reservation), + handle: Some(handle), + }, + ); + true + } + + fn drain_ready_prefetch(&mut self) -> Vec { + let ids: Vec = self.prefetch.keys().copied().collect(); + let mut done = Vec::new(); + for id in ids { + let poll = match self.prefetch.get_mut(&id).and_then(|st| st.handle.as_mut()) { + Some(handle) => handle.poll(), + None => continue, // already settled, awaiting prefill + }; + if let Some(result) = poll { + self.settle_prefetch(id, result); + done.push(id); + } + } + done + } + + fn wait_ready_prefetch(&mut self) -> Vec { + let mut done = Vec::new(); + if let Some(id) = self + .prefetch + .iter() + .find(|(_, st)| st.handle.is_some()) + .map(|(id, _)| *id) + { + let handle = self + .prefetch + .get_mut(&id) + .and_then(|st| st.handle.take()) + .expect("in-flight handle present"); + let result = handle.wait(); + self.settle_prefetch(id, result); + // `settle_prefetch` clears the handle, so the drain below skips it; + // record it here as the one we blocked on. + done.push(id); + } + // Sweep any others that completed concurrently. + for id in self.drain_ready_prefetch() { + if !done.contains(&id) { + done.push(id); + } + } + done + } + + fn execute_prefill(&mut self, plan: PrefillPlan<'_>) -> Result { + // 1. Create RequestKvs (first chunk only), clamp chunk budgets, + // schedule KV for this step's tokens + let mut requests = plan.requests.to_vec(); + for req in &mut requests { + self.schedule_prefill_chunk(req)?; + } + + // 2. Build KvViews (seq_len = chunk_start + this chunk) + let kv_views: Vec = requests + .iter() + .map(|req| self.request_kvs[&req.request_id].prefill_view(req.chunk_tokens)) + .collect(); + + // 3. Execute forward + let scheduled_requests = requests.clone(); + let step = StepCommand::Prefill { + requests, + kv_views, + echo: plan.echo, + }; + let outcome = self.run_step(&step)?; + + // 4. Apply prefill + let result = match outcome { + WorkerStepOutcome::Prefill(result) => result, + other => { + return Err(anyhow::anyhow!( + "prefill returned unexpected: {}", + other.kind() + )); + } + }; + for req_result in &result.requests { + self.apply_prefill_result(req_result)?; + } + if self.speculative_enabled { + for (req, req_result) in scheduled_requests.iter().zip(&result.requests) { + assert_eq!(req.request_id, req_result.request_id); + let captured = result + .dflash_context_captured_requests + .contains(&req.request_id); + match dflash_prefill_action(captured, req_result.completed) { + DFlashPrefillAction::MarkReady => { + self.dflash_ready_requests.insert(req.request_id); + } + DFlashPrefillAction::KeepPending => { + self.dflash_ready_requests.remove(&req.request_id); + } + DFlashPrefillAction::Drop => { + self.dflash_ready_requests.remove(&req.request_id); + self.primary.drop_dflash_request(req.request_id)?; + } + } + } + } + // 5. Offload the blocks this prefill just sealed (post-step-sync). + for req_result in &result.requests { + self.save_sealed_blocks(req_result.request_id); + } + + Ok(result) + } + + fn execute_decode(&mut self, plan: DecodePlan<'_>) -> Result { + // 1. Schedule decode for all active requests + for req in plan.requests { + let rkv = self + .request_kvs + .get_mut(&req.request_id) + .ok_or_else(|| anyhow::anyhow!("missing RequestKv for {:?}", req.request_id))?; + rkv.schedule_decode(self.kv_mgr.pool()).map_err(|e| { + anyhow::anyhow!("schedule_decode failed for {:?}: {e}", req.request_id) + })?; + } + + // 2. Build KvViews + let kv_views: Vec = plan + .requests + .iter() + .map(|req| self.request_kvs[&req.request_id].decode_view()) + .collect(); + + // 3. Execute forward + let scheduled_requests = plan.requests.to_vec(); + let step = StepCommand::Decode { + requests: scheduled_requests.clone(), + kv_views, + }; + let outcome = self.run_step(&step)?; + + // 4. Apply decode + let result = match outcome { + WorkerStepOutcome::Decode(result) => result, + other => { + return Err(anyhow::anyhow!( + "decode returned unexpected: {}", + other.kind() + )); + } + }; + for req_result in &result.requests { + let rkv = self + .request_kvs + .get_mut(&req_result.request_id) + .expect("request must exist after decode"); + rkv.apply_decode(req_result.token, self.kv_mgr.pool())?; + } + if self.speculative_enabled { + for req in &scheduled_requests { + self.dflash_ready_requests.remove(&req.request_id); + self.primary.drop_dflash_request(req.request_id)?; + } + } + // 5. Offload any block this decode step just sealed (post-step-sync). + for req_result in &result.requests { + self.save_sealed_blocks(req_result.request_id); + } + + Ok(result) + } + + fn execute_speculative_verify( + &mut self, + plan: SpeculativeVerifyPlan<'_>, + ) -> Result { + self.execute_speculative_verify_impl(plan) + } + + fn execute_speculative_draft( + &mut self, + plan: SpeculativeDraftPlan<'_>, + ) -> Result { + self.execute_speculative_draft_impl(plan) + } + + fn speculative_enabled(&self) -> bool { + self.speculative_enabled + } + + fn speculative_request_ready(&self, request_id: RequestId) -> bool { + self.dflash_ready_requests.contains(&request_id) + } + + fn speculative_state_budget_bytes(&self) -> Option { + self.speculative_enabled + .then_some(self.metadata.dflash_state_budget_bytes) + } + + fn speculative_request_state_bytes(&self, prompt_len: usize, max_tokens: usize) -> usize { + let Some(config) = self.metadata.dflash_config.as_ref() else { + return 0; + }; + let max_cache_len = prompt_len + .saturating_add(max_tokens) + .saturating_add(config.block_size) + .min(config.max_position_embeddings); + match dflash_request_state_footprint_bytes_for_len(config, max_cache_len) { + Ok(bytes) => bytes, + Err(error) => { + log::warn!( + "failed to estimate DFlash request-state memory for prompt_len={prompt_len} max_tokens={max_tokens}: {error}" + ); + usize::MAX + } + } + } + + fn execute_unified(&mut self, plan: UnifiedPlan<'_>) -> Result { + // 1. Create RequestKvs for prefill requests (first chunk only), clamp + // chunk budgets, schedule KV for this step's tokens + let mut prefill_requests = plan.prefill_requests.to_vec(); + for req in &mut prefill_requests { + self.schedule_prefill_chunk(req)?; + } + + // Schedule decode for active requests + let decode_requests = plan.decode_requests.to_vec(); + for req in &decode_requests { + let rkv = self + .request_kvs + .get_mut(&req.request_id) + .ok_or_else(|| anyhow::anyhow!("missing RequestKv for {:?}", req.request_id))?; + rkv.schedule_decode(self.kv_mgr.pool()).map_err(|e| { + anyhow::anyhow!("schedule_decode failed for {:?}: {e}", req.request_id) + })?; + } + + // 2. Build KvViews + let prefill_kv_views: Vec = prefill_requests + .iter() + .map(|req| self.request_kvs[&req.request_id].prefill_view(req.chunk_tokens)) + .collect(); + let decode_kv_views: Vec = decode_requests + .iter() + .map(|req| self.request_kvs[&req.request_id].decode_view()) + .collect(); + + // 3. Execute forward + let step = StepCommand::Unified { + prefill_requests, + prefill_kv_views, + decode_requests, + decode_kv_views, + }; + let outcome = self.run_step(&step)?; + + // 4. Apply both prefill and decode + let result = match outcome { + WorkerStepOutcome::Unified(result) => result, + other => { + return Err(anyhow::anyhow!( + "unified returned unexpected: {}", + other.kind() + )); + } + }; + for req_result in &result.prefill_requests { + self.apply_prefill_result(req_result)?; + } + for req_result in &result.decode_requests { + let rkv = self + .request_kvs + .get_mut(&req_result.request_id) + .expect("request must exist after unified decode"); + rkv.apply_decode(req_result.token, self.kv_mgr.pool())?; + } + if self.speculative_enabled { + for req_result in &result.prefill_requests { + self.dflash_ready_requests.remove(&req_result.request_id); + self.primary.drop_dflash_request(req_result.request_id)?; + } + for req_result in &result.decode_requests { + self.dflash_ready_requests.remove(&req_result.request_id); + self.primary.drop_dflash_request(req_result.request_id)?; + } + } + // 5. Offload sealed blocks from both halves (post-step-sync). + for req_result in &result.prefill_requests { + self.save_sealed_blocks(req_result.request_id); + } + for req_result in &result.decode_requests { + self.save_sealed_blocks(req_result.request_id); + } + + Ok(result) + } + + fn load_lora_adapter(&mut self, request: &LoadLoraAdapterRequest) -> Result<()> { + ensure_lora_capacity( + &self.loaded_lora_adapters, + &request.lora_name, + self.lora_options.max_loras, + request.load_inplace, + )?; + let adapter = crate::lora::load_lora_adapter( + &request.lora_path, + &self.metadata.config, + self.lora_options.max_lora_rank, + )?; + let world_size = self.workers.len() + 1; + let projection_count: usize = adapter + .layers + .iter() + .map(|layer| layer.projections.len()) + .sum(); + let element_count: usize = adapter + .layers + .iter() + .flat_map(|layer| layer.projections.values()) + .map(|projection| projection.a.data.len() + projection.b.data.len()) + .sum(); + let shape_elems: usize = adapter + .layers + .iter() + .flat_map(|layer| layer.projections.values()) + .map(|projection| { + projection.a.rows * projection.a.cols + projection.b.rows * projection.b.cols + }) + .sum(); + debug_assert_eq!(element_count, shape_elems); + let rank = adapter.manifest.rank; + let targets = adapter.manifest.target_modules.join(", "); + let path = adapter.manifest.path.display().to_string(); + let mut sharded_adapters = Vec::with_capacity(world_size); + for rank in 0..world_size { + sharded_adapters.push(adapter.shard_for_tensor_parallel( + &self.metadata.config, + TensorParallelConfig { rank, world_size }, + )?); + } + + let mut sharded_adapters = sharded_adapters.into_iter(); + let primary_adapter = sharded_adapters + .next() + .expect("rank 0 adapter must exist for nonzero world_size"); + let primary_response = self.primary.load_lora_adapter( + request.lora_name.clone(), + primary_adapter, + request.load_inplace, + )?; + let mut pending = Vec::with_capacity(self.workers.len()); + let mut errors = Vec::new(); + for (index, worker) in self.workers.iter().enumerate() { + let rank = index + 1; + let rank_adapter = sharded_adapters + .next() + .expect("worker adapter must exist for every tensor-parallel rank"); + match worker.load_lora_adapter( + request.lora_name.clone(), + rank_adapter, + request.load_inplace, + ) { + Ok(response) => pending.push((rank, response)), + Err(err) => errors.push(format!("rank {rank} dispatch: {err:#}")), + } + } + + match primary_response.recv() { + Ok(Ok(())) => {} + Ok(Err(err)) => errors.push(format!("rank 0: {err:#}")), + Err(_) => errors.push("rank 0: dropped LoRA load response".to_string()), + } + for (rank, response) in pending { + match response.recv() { + Ok(Ok(())) => {} + Ok(Err(err)) => errors.push(format!("rank {rank}: {err:#}")), + Err(_) => errors.push(format!("rank {rank}: dropped LoRA load response")), + } + } + if !errors.is_empty() { + let mut cleanup_errors = Vec::new(); + match self.primary.discard_lora_adapter(request.lora_name.clone()) { + Ok(response) => match response.recv() { + Ok(Ok(())) => {} + Ok(Err(err)) => cleanup_errors.push(format!("rank 0 cleanup: {err:#}")), + Err(_) => cleanup_errors + .push("rank 0 cleanup: dropped LoRA discard response".to_string()), + }, + Err(err) => cleanup_errors.push(format!("rank 0 cleanup dispatch: {err:#}")), + } + for (index, worker) in self.workers.iter().enumerate() { + let rank = index + 1; + match worker.discard_lora_adapter(request.lora_name.clone()) { + Ok(response) => match response.recv() { + Ok(Ok(())) => {} + Ok(Err(err)) => { + cleanup_errors.push(format!("rank {rank} cleanup: {err:#}")); + } + Err(_) => cleanup_errors.push(format!( + "rank {rank} cleanup: dropped LoRA discard response" + )), + }, + Err(err) => { + cleanup_errors.push(format!("rank {rank} cleanup dispatch: {err:#}")); + } + } + } + if cleanup_errors.is_empty() { + self.loaded_lora_adapters.remove(&request.lora_name); + } + let cleanup_suffix = if cleanup_errors.is_empty() { + String::new() + } else { + format!("; cleanup errors: {}", cleanup_errors.join("; ")) + }; + anyhow::bail!( + "failed to load Qwen3 LoRA adapter {} on tensor-parallel ranks: {}{}", + request.lora_name, + errors.join("; "), + cleanup_suffix + ); + } + + log::info!( + "Loaded Qwen3 LoRA adapter {} from {} (rank={}, targets={}, projections={}, bf16_elements={}, tp_world_size={}, load_inplace={})", + request.lora_name, + path, + rank, + targets, + projection_count, + element_count, + world_size, + request.load_inplace + ); + self.loaded_lora_adapters.insert(request.lora_name.clone()); + Ok(()) + } + + fn unload_lora_adapter(&mut self, request: &UnloadLoraAdapterRequest) -> Result<()> { + let primary_response = self + .primary + .unload_lora_adapter(request.lora_name.clone())?; + let mut pending = Vec::with_capacity(self.workers.len()); + for (index, worker) in self.workers.iter().enumerate() { + pending.push(( + index + 1, + worker.unload_lora_adapter(request.lora_name.clone())?, + )); + } + + let mut errors = Vec::new(); + match primary_response.recv() { + Ok(Ok(())) => {} + Ok(Err(err)) => errors.push(format!("rank 0: {err:#}")), + Err(_) => errors.push("rank 0: dropped LoRA unload response".to_string()), + } + for (rank, response) in pending { + match response.recv() { + Ok(Ok(())) => {} + Ok(Err(err)) => errors.push(format!("rank {rank}: {err:#}")), + Err(_) => errors.push(format!("rank {rank}: dropped LoRA unload response")), + } + } + if !errors.is_empty() { + anyhow::bail!( + "failed to unload Qwen3 LoRA adapter {} on tensor-parallel ranks: {}", + request.lora_name, + errors.join("; ") + ); + } + + log::info!("Unloaded Qwen3 LoRA adapter {}", request.lora_name); + self.loaded_lora_adapters.remove(&request.lora_name); + Ok(()) + } + + fn list_lora_adapters(&self) -> Vec { + let mut names: Vec<_> = self.loaded_lora_adapters.iter().cloned().collect(); + names.sort(); + names + } +} + +fn ensure_lora_capacity( + loaded_lora_adapters: &HashSet, + lora_name: &str, + max_loras: usize, + load_inplace: bool, +) -> Result<()> { + if loaded_lora_adapters.contains(lora_name) { + anyhow::ensure!( + load_inplace, + "Qwen3 LoRA adapter {lora_name} is already loaded" + ); + return Ok(()); + } + anyhow::ensure!( + loaded_lora_adapters.len() < max_loras, + "Qwen3 LoRA adapter capacity exceeded: max_loras={}, loaded_adapters={}, requested={}", + max_loras, + loaded_lora_adapters.len(), + lora_name + ); + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::ensure_lora_capacity; + use std::collections::HashSet; + + #[test] + fn lora_capacity_rejects_new_adapter_at_limit() { + let loaded = HashSet::from(["adapter-a".to_string()]); + + let error = ensure_lora_capacity(&loaded, "adapter-b", 1, false) + .expect_err("new adapter should exceed capacity") + .to_string(); + + assert!(error.contains("max_loras=1")); + assert!(error.contains("requested=adapter-b")); + } + + #[test] + fn lora_capacity_allows_existing_adapter_replacement_at_limit_with_load_inplace() { + let loaded = HashSet::from(["adapter-a".to_string()]); + + ensure_lora_capacity(&loaded, "adapter-a", 1, true) + .expect("existing adapter should fit with load_inplace"); + } + + #[test] + fn lora_capacity_rejects_duplicate_without_load_inplace() { + let loaded = HashSet::from(["adapter-a".to_string()]); + + let error = ensure_lora_capacity(&loaded, "adapter-a", 1, false) + .expect_err("duplicate without load_inplace should fail") + .to_string(); + + assert!(error.contains("already loaded")); + } +} diff --git a/openinfer-qwen3-4b/src/executor/speculative_exec.rs b/openinfer-qwen3-4b/src/executor/speculative_exec.rs new file mode 100644 index 00000000..107b6980 --- /dev/null +++ b/openinfer-qwen3-4b/src/executor/speculative_exec.rs @@ -0,0 +1,176 @@ +use anyhow::Result; + +use crate::speculative::{ + DraftPlan as SpeculativeDraftPlan, DraftResult as SpeculativeDraftResult, + VerifyPlan as SpeculativeVerifyPlan, VerifyResult as SpeculativeVerifyResult, +}; + +use super::{Qwen3Executor, RequestId, StepCommand, WorkerStepOutcome}; + +impl Qwen3Executor { + pub(super) fn execute_speculative_verify_impl( + &mut self, + plan: SpeculativeVerifyPlan<'_>, + ) -> Result { + anyhow::ensure!( + self.speculative_enabled, + "speculative verification requested but no draft model is loaded" + ); + for req in plan.requests { + anyhow::ensure!( + !req.token_ids.is_empty(), + "speculative verify request {:?} has an empty verify span", + req.request_id + ); + anyhow::ensure!( + req.params.is_greedy(), + "speculative verification currently supports greedy sampling only" + ); + anyhow::ensure!( + self.dflash_ready_requests.contains(&req.request_id), + "speculative verification requested before DFlash state is ready for {:?}", + req.request_id + ); + anyhow::ensure!( + self.request_kvs.contains_key(&req.request_id), + "missing RequestKv for {:?}", + req.request_id + ); + } + + let mut scheduled = Vec::with_capacity(plan.requests.len()); + for req in plan.requests { + let rkv = self + .request_kvs + .get_mut(&req.request_id) + .expect("RequestKv was validated before speculative scheduling"); + if let Err(e) = rkv.schedule_speculative(req.token_ids.len(), self.kv_mgr.pool()) { + self.revert_speculative_schedules(&scheduled); + return Err(anyhow::anyhow!( + "schedule_speculative failed for {:?}: {e}", + req.request_id + )); + } + scheduled.push(req.request_id); + } + + let kv_views = plan + .requests + .iter() + .map(|req| self.request_kvs[&req.request_id].speculative_view(req.token_ids.len())) + .collect(); + + let step = StepCommand::SpeculativeVerify { + requests: plan.requests.to_vec(), + kv_views, + }; + let outcome = match self.run_step(&step) { + Ok(outcome) => outcome, + Err(e) => { + self.revert_speculative_schedules(&scheduled); + return Err(e); + } + }; + + let result = match outcome { + WorkerStepOutcome::SpeculativeVerify(result) => result, + other => { + self.revert_speculative_schedules(&scheduled); + return Err(anyhow::anyhow!( + "speculative verify returned unexpected: {}", + other.kind() + )); + } + }; + if result.requests.len() != plan.requests.len() { + self.revert_speculative_schedules(&scheduled); + return Err(anyhow::anyhow!( + "speculative verify returned {} request results for {} requests", + result.requests.len(), + plan.requests.len() + )); + } + for (req, req_result) in plan.requests.iter().zip(&result.requests) { + if req.request_id != req_result.request_id { + self.revert_speculative_schedules(&scheduled); + return Err(anyhow::anyhow!( + "speculative verify returned request {:?} for {:?}", + req_result.request_id, + req.request_id + )); + } + } + + let mut applied = Vec::with_capacity(result.requests.len()); + for req_result in &result.requests { + let rkv = self + .request_kvs + .get_mut(&req_result.request_id) + .expect("request must exist after speculative verify"); + if let Err(e) = rkv.apply_speculative(&req_result.accepted_tokens, self.kv_mgr.pool()) { + let unapplied = scheduled + .iter() + .copied() + .filter(|request_id| !applied.contains(request_id)) + .collect::>(); + self.revert_speculative_schedules(&unapplied); + return Err(anyhow::anyhow!( + "apply_speculative failed for {:?}: {e}", + req_result.request_id + )); + } + applied.push(req_result.request_id); + } + for req_result in &result.requests { + self.save_sealed_blocks(req_result.request_id); + } + + Ok(result) + } + + pub(super) fn execute_speculative_draft_impl( + &mut self, + plan: SpeculativeDraftPlan<'_>, + ) -> Result { + anyhow::ensure!( + self.speculative_enabled, + "speculative draft requested but no draft model is loaded" + ); + for req in plan.requests { + anyhow::ensure!( + req.params.is_greedy(), + "speculative draft currently supports greedy sampling only" + ); + anyhow::ensure!( + self.dflash_ready_requests.contains(&req.request_id), + "speculative draft requested before DFlash state is ready for {:?}", + req.request_id + ); + } + let step = StepCommand::SpeculativeDraft { + requests: plan.requests.to_vec(), + }; + let outcome = self.run_step(&step)?; + match outcome { + WorkerStepOutcome::SpeculativeDraft(result) => Ok(result), + other => Err(anyhow::anyhow!( + "speculative draft returned unexpected: {}", + other.kind() + )), + } + } + + fn revert_speculative_schedules(&mut self, request_ids: &[RequestId]) { + for request_id in request_ids.iter().rev().copied() { + let Some(rkv) = self.request_kvs.get_mut(&request_id) else { + log::warn!( + "missing RequestKv while reverting speculative schedule for {request_id:?}" + ); + continue; + }; + if let Err(error) = rkv.revert_schedule() { + log::warn!("failed to revert speculative schedule for {request_id:?}: {error}"); + } + } + } +} diff --git a/openinfer-qwen3-4b/src/executor/worker.rs b/openinfer-qwen3-4b/src/executor/worker.rs new file mode 100644 index 00000000..ddf49318 --- /dev/null +++ b/openinfer-qwen3-4b/src/executor/worker.rs @@ -0,0 +1,507 @@ +use std::thread; + +use anyhow::Result; +use crossbeam_channel as channel; + +use crate::batch_decode_buffers::{BATCH_BUCKETS, BatchDecodeBuffers}; +use crate::dflash::DFlashDraftModel; +use crate::speculative::{ + DraftResult as SpeculativeDraftResult, DraftStepItem as SpeculativeDraftStepItem, + VerifyResult as SpeculativeVerifyResult, VerifyStepItem as SpeculativeVerifyStepItem, +}; +use crate::weights::Qwen3Model; +use openinfer_core::engine::TokenLogprob; +use openinfer_core::kv_pool::KvLayout; +use openinfer_core::sampler::SamplingParams; +use openinfer_core::tensor::{DeviceVec, HiddenStates}; +use openinfer_kv_cache::{KvBuffer, KvView}; + +use super::dflash_lane::DFlashLaneState; +use super::{ + CublasThreadGuard, DecodeResult, DecodeStepItem, PrefillResult, PrefillStepItem, RequestId, + SamplingScratch, UnifiedResult, bind_model_thread, compute_logprobs_from_cpu, + execute_step_on_lane, tune_decode_gemm_algos, +}; + +pub(super) struct LocalQwen3Lane { + pub(super) model: Qwen3Model, + kv_buffer: KvBuffer, + layout: KvLayout, + pub(super) bufs: BatchDecodeBuffers, + pub(super) sample_scratch: SamplingScratch, + pub(super) dflash: Option, +} + +impl LocalQwen3Lane { + pub(super) fn new( + model: Qwen3Model, + kv_buffer: KvBuffer, + total_blocks: usize, + padding_block_id: i32, + dflash_model: Option, + ) -> Result { + let buf_layout = kv_buffer.layout(); + let layout = KvLayout::new( + buf_layout.num_layers, + buf_layout.num_kv_heads, + buf_layout.head_dim, + buf_layout.page_size, + ); + let max_bucket = *BATCH_BUCKETS.last().unwrap(); + let bufs = BatchDecodeBuffers::new( + model.device_ctx(), + model.config().hidden_size, + model.local_q_dim(), + model.local_kv_dim(), + model.local_intermediate_size(), + model.config().vocab_size, + max_bucket, + total_blocks, + padding_block_id, + model.local_num_attention_heads(), + )?; + let sample_scratch = + SamplingScratch::new(model.device_ctx(), model.config().vocab_size, max_bucket)?; + Ok(Self { + model, + kv_buffer, + layout, + bufs, + sample_scratch, + dflash: dflash_model.map(DFlashLaneState::new), + }) + } + + fn bind(&self) -> Result { + bind_model_thread(&self.model)?; + tune_decode_gemm_algos(&self.model)?; + if let Some(dflash) = &self.dflash { + dflash.model.tune_gemm_algos(&self.model)?; + } + Ok(CublasThreadGuard) + } + + /// Pick one token per logits column (batched argmax for greedy rows, + /// per-row sampler otherwise). Grows the sampling scratch when a step + /// is wider than the decode bucket it was sized for. + pub(super) fn select_step_tokens( + &mut self, + logits: &HiddenStates, + params: &[&SamplingParams], + random_vals: &[f32], + ) -> Result> { + let scratch_rows = logits.seq_len.max(params.len()); + if scratch_rows > self.sample_scratch.row_indices.len() { + self.sample_scratch = SamplingScratch::new( + self.model.device_ctx(), + self.model.config().vocab_size, + scratch_rows, + )?; + } + openinfer_core::ops::select_batch_tokens_into( + self.model.device_ctx(), + logits, + params, + random_vals, + &mut self.sample_scratch.row_indices, + &mut self.sample_scratch.argmax_partial_values, + &mut self.sample_scratch.argmax_partial_indices, + &mut self.sample_scratch.probs, + &mut self.sample_scratch.top1_values, + &mut self.sample_scratch.row_states, + &mut self.sample_scratch.valid, + &mut self.sample_scratch.out, + ) + } + + pub(super) fn select_greedy_contiguous_tokens( + &mut self, + logits: &HiddenStates, + ) -> Result> { + if logits.seq_len > self.sample_scratch.row_indices.len() { + self.sample_scratch = SamplingScratch::new( + self.model.device_ctx(), + self.model.config().vocab_size, + logits.seq_len, + )?; + } + openinfer_core::ops::argmax_batch_bf16_split_into( + self.model.device_ctx(), + logits, + &mut self.sample_scratch.argmax_partial_values, + &mut self.sample_scratch.argmax_partial_indices, + &mut self.sample_scratch.top1_values, + &mut self.sample_scratch.out, + )?; + self.sample_scratch + .out_host + .resize(self.sample_scratch.out.len(), 0); + self.model + .device_ctx() + .stream + .memcpy_dtoh(&self.sample_scratch.out, &mut self.sample_scratch.out_host) + .map_err(|e| anyhow::anyhow!("D2H greedy argmax read failed: {}", e))?; + self.model.device_ctx().sync()?; + Ok(self + .sample_scratch + .out_host + .iter() + .take(logits.seq_len) + .map(|&token| token as u32) + .collect()) + } + + pub(super) fn extract_logprobs( + &self, + logits: &DeviceVec, + sampled_token: u32, + top_k: usize, + ) -> Result { + let logits_f32 = logits.to_host(self.model.device_ctx())?; + compute_logprobs_from_cpu(&logits_f32, sampled_token, top_k) + .ok_or_else(|| anyhow::anyhow!("logprobs computation failed")) + } + + pub(super) fn extract_prompt_logprobs( + &self, + all_logits: &HiddenStates, + prev_pos: usize, + target_token: u32, + top_k: usize, + ) -> Option { + openinfer_core::ops::extract_vec(self.model.device_ctx(), all_logits, prev_pos) + .ok() + .and_then(|logits_vec| { + let logits_f32 = logits_vec.to_host(self.model.device_ctx()).ok()?; + compute_logprobs_from_cpu(&logits_f32, target_token, top_k) + }) + } + + pub(super) fn execute_prefill( + &mut self, + prompts: &[&[u32]], + kv_views: &[KvView], + lora_adapters: &[Option<&str>], + echo: bool, + capture_layer_ids: Option<&[usize]>, + ) -> Result<(HiddenStates, Option, Option)> { + if let Some(capture_layer_ids) = capture_layer_ids { + self.model.batch_prefill_with_hidden_capture( + prompts, + kv_views, + lora_adapters, + self.kv_buffer.buffer(), + &self.layout, + echo, + Some(capture_layer_ids), + ) + } else { + let (logits, all_position_logits) = self.model.batch_prefill( + prompts, + kv_views, + lora_adapters, + self.kv_buffer.buffer(), + &self.layout, + echo, + )?; + Ok((logits, all_position_logits, None)) + } + } + + pub(super) fn execute_decode( + &mut self, + token_ids: &[u32], + kv_views: &[KvView], + lora_adapters: &[Option<&str>], + ) -> Result<()> { + self.model.batch_decode( + token_ids, + kv_views, + lora_adapters, + self.kv_buffer.buffer(), + &self.layout, + &mut self.bufs, + ) + } + + pub(super) fn execute_unified( + &mut self, + prefill_prompts: &[&[u32]], + prefill_views: &[KvView], + prefill_lora_adapters: &[Option<&str>], + decode_tokens: &[u32], + decode_views: &[KvView], + decode_lora_adapters: &[Option<&str>], + ) -> Result { + self.model.unified_step( + prefill_prompts, + prefill_views, + prefill_lora_adapters, + decode_tokens, + decode_views, + decode_lora_adapters, + self.kv_buffer.buffer(), + &self.layout, + ) + } + + fn load_lora_adapter( + &mut self, + name: String, + adapter: crate::lora::LoraAdapter, + load_inplace: bool, + ) -> Result<()> { + let device_adapter = + crate::lora::load_device_lora_adapter(self.model.device_ctx(), name, adapter)?; + self.model + .install_lora_adapter(device_adapter, load_inplace) + } + + fn unload_lora_adapter(&mut self, name: &str) -> Result<()> { + self.model.uninstall_lora_adapter(name) + } + + fn discard_lora_adapter(&mut self, name: &str) -> Result<()> { + self.model.discard_lora_adapter(name) + } +} + +#[derive(Clone)] +pub(super) enum StepCommand { + Prefill { + requests: Vec, + kv_views: Vec, + echo: bool, + }, + Decode { + requests: Vec, + kv_views: Vec, + }, + SpeculativeVerify { + requests: Vec, + kv_views: Vec, + }, + SpeculativeDraft { + requests: Vec, + }, + Unified { + prefill_requests: Vec, + prefill_kv_views: Vec, + decode_requests: Vec, + decode_kv_views: Vec, + }, +} + +impl StepCommand { + pub(super) fn kind(&self) -> &'static str { + match self { + Self::Prefill { .. } => "prefill", + Self::Decode { .. } => "decode", + Self::SpeculativeVerify { .. } => "speculative_verify", + Self::SpeculativeDraft { .. } => "speculative_draft", + Self::Unified { .. } => "unified", + } + } +} + +enum WorkerCommand { + RunStep { + step: StepCommand, + collect_result: bool, + resp: channel::Sender>, + }, + LoadLoraAdapter { + name: String, + adapter: crate::lora::LoraAdapter, + load_inplace: bool, + resp: channel::Sender>, + }, + UnloadLoraAdapter { + name: String, + resp: channel::Sender>, + }, + DiscardLoraAdapter { + name: String, + resp: channel::Sender>, + }, + DropDFlashRequest { + request_id: RequestId, + resp: channel::Sender>, + }, + Shutdown, +} + +pub(super) enum WorkerStepOutcome { + Ack, + Prefill(PrefillResult), + Decode(DecodeResult), + SpeculativeDraft(SpeculativeDraftResult), + SpeculativeVerify(SpeculativeVerifyResult), + Unified(UnifiedResult), +} + +impl WorkerStepOutcome { + pub(super) fn kind(&self) -> &'static str { + match self { + Self::Ack => "ack", + Self::Prefill(_) => "prefill", + Self::Decode(_) => "decode", + Self::SpeculativeDraft(_) => "speculative_draft", + Self::SpeculativeVerify(_) => "speculative_verify", + Self::Unified(_) => "unified", + } + } +} + +pub(super) struct RankWorker { + tx: channel::Sender, + handle: Option>, +} + +impl RankWorker { + pub(super) fn spawn(rank: usize, mut lane: LocalQwen3Lane) -> Result { + let (tx, rx) = channel::unbounded(); + let (startup_tx, startup_rx) = channel::bounded(1); + let handle = thread::Builder::new() + .name(format!("qwen3-tp-rank-{rank}")) + .spawn(move || { + let startup = lane.bind(); + match startup { + Ok(_guard) => { + let _ = startup_tx.send(Ok(())); + while let Ok(cmd) = rx.recv() { + match cmd { + WorkerCommand::RunStep { + step, + collect_result, + resp, + } => { + let result = + execute_step_on_lane(&mut lane, &step, collect_result); + let _ = resp.send(result); + } + WorkerCommand::LoadLoraAdapter { + name, + adapter, + load_inplace, + resp, + } => { + let result = + lane.load_lora_adapter(name, adapter, load_inplace); + let _ = resp.send(result); + } + WorkerCommand::UnloadLoraAdapter { name, resp } => { + let result = lane.unload_lora_adapter(&name); + let _ = resp.send(result); + } + WorkerCommand::DiscardLoraAdapter { name, resp } => { + let result = lane.discard_lora_adapter(&name); + let _ = resp.send(result); + } + WorkerCommand::DropDFlashRequest { request_id, resp } => { + lane.drop_dflash_request(request_id); + let _ = resp.send(Ok(())); + } + WorkerCommand::Shutdown => break, + } + } + } + Err(err) => { + let _ = startup_tx.send(Err(err)); + } + } + }) + .map_err(|e| anyhow::anyhow!("failed to spawn tensor-parallel worker {rank}: {e}"))?; + startup_rx.recv().map_err(|_| { + anyhow::anyhow!("tensor-parallel worker {rank} exited during startup") + })??; + Ok(Self { + tx, + handle: Some(handle), + }) + } + + pub(super) fn run_step( + &self, + step: StepCommand, + collect_result: bool, + ) -> Result>> { + let (resp_tx, resp_rx) = channel::bounded(1); + self.tx + .send(WorkerCommand::RunStep { + step, + collect_result, + resp: resp_tx, + }) + .map_err(|_| anyhow::anyhow!("tensor-parallel worker step channel closed"))?; + Ok(resp_rx) + } + + pub(super) fn load_lora_adapter( + &self, + name: String, + adapter: crate::lora::LoraAdapter, + load_inplace: bool, + ) -> Result>> { + let (resp_tx, resp_rx) = channel::bounded(1); + self.tx + .send(WorkerCommand::LoadLoraAdapter { + name, + adapter, + load_inplace, + resp: resp_tx, + }) + .map_err(|_| anyhow::anyhow!("tensor-parallel worker channel closed on LoRA load"))?; + Ok(resp_rx) + } + + pub(super) fn unload_lora_adapter( + &self, + name: String, + ) -> Result>> { + let (resp_tx, resp_rx) = channel::bounded(1); + self.tx + .send(WorkerCommand::UnloadLoraAdapter { + name, + resp: resp_tx, + }) + .map_err(|_| anyhow::anyhow!("tensor-parallel worker channel closed on LoRA unload"))?; + Ok(resp_rx) + } + + pub(super) fn discard_lora_adapter( + &self, + name: String, + ) -> Result>> { + let (resp_tx, resp_rx) = channel::bounded(1); + self.tx + .send(WorkerCommand::DiscardLoraAdapter { + name, + resp: resp_tx, + }) + .map_err(|_| { + anyhow::anyhow!("tensor-parallel worker channel closed on LoRA discard") + })?; + Ok(resp_rx) + } + + pub(super) fn drop_dflash_request(&self, request_id: RequestId) -> Result<()> { + let (resp_tx, resp_rx) = channel::bounded(1); + self.tx + .send(WorkerCommand::DropDFlashRequest { + request_id, + resp: resp_tx, + }) + .map_err(|_| anyhow::anyhow!("tensor-parallel worker channel closed on DFlash drop"))?; + resp_rx + .recv() + .map_err(|_| anyhow::anyhow!("tensor-parallel worker dropped DFlash drop response"))? + } + + pub(super) fn shutdown(&mut self) { + let _ = self.tx.send(WorkerCommand::Shutdown); + if let Some(handle) = self.handle.take() { + let _ = handle.join(); + } + } +} diff --git a/openinfer-qwen3-4b/src/lib.rs b/openinfer-qwen3-4b/src/lib.rs index 92243de2..55e61e5a 100644 --- a/openinfer-qwen3-4b/src/lib.rs +++ b/openinfer-qwen3-4b/src/lib.rs @@ -5,15 +5,17 @@ mod batch_decode_buffers; mod batch_decode_dag; pub mod batch_decode_trace; mod config; +mod dflash; mod executor; pub mod kernel_bench; mod lora; mod prefill; mod scheduler; +mod speculative; mod unified_forward; mod weights; -use std::path::Path; +use std::path::{Path, PathBuf}; use anyhow::Result; use openinfer_core::engine::{EngineHandle, EngineLoadOptions, ModelInfo}; @@ -103,6 +105,30 @@ impl Default for Qwen3OffloadOptions { } } +#[derive(Clone, Debug, Default, Eq, PartialEq)] +pub struct Qwen3SpeculativeOptions { + pub dflash: Option, +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct Qwen3DFlashOptions { + pub model_path: PathBuf, +} + +impl Qwen3SpeculativeOptions { + pub fn disabled() -> Self { + Self { dflash: None } + } + + pub fn dflash(model_path: impl Into) -> Self { + Self { + dflash: Some(Qwen3DFlashOptions { + model_path: model_path.into(), + }), + } + } +} + /// Low-level Qwen3 execution interface. /// /// This is the production phase boundary used by the Qwen3 scheduler and by @@ -113,6 +139,12 @@ pub mod runtime { PrefillRequestResult, PrefillResult, PrefillStepItem, Qwen3Executor, RequestId, UnifiedPlan, UnifiedResult, }; + pub use crate::speculative::{ + DraftPlan as SpeculativeDraftPlan, DraftRequestResult as SpeculativeDraftRequestResult, + DraftResult as SpeculativeDraftResult, DraftStepItem as SpeculativeDraftStepItem, + VerifyPlan as SpeculativeVerifyPlan, VerifyRequestResult as SpeculativeVerifyRequestResult, + VerifyResult as SpeculativeVerifyResult, VerifyStepItem as SpeculativeVerifyStepItem, + }; } pub fn probe_model(model_path: &Path) -> Result> { @@ -166,6 +198,24 @@ pub fn start_engine_with_offload( offload_options: Qwen3OffloadOptions, no_prefix_cache: bool, max_prefill_tokens: usize, +) -> Result { + start_engine_with_offload_and_speculative( + model_path, + options, + offload_options, + Qwen3SpeculativeOptions::disabled(), + no_prefix_cache, + max_prefill_tokens, + ) +} + +pub fn start_engine_with_offload_and_speculative( + model_path: &Path, + options: EngineLoadOptions, + offload_options: Qwen3OffloadOptions, + speculative_options: Qwen3SpeculativeOptions, + no_prefix_cache: bool, + max_prefill_tokens: usize, ) -> Result { let EngineLoadOptions { enable_cuda_graph, @@ -182,6 +232,7 @@ pub fn start_engine_with_offload( &device_ordinals, seed, offload_options, + speculative_options, no_prefix_cache, max_prefill_tokens, ) @@ -211,6 +262,7 @@ pub fn start_engine_with_lora_control( seed, lora_options.validate()?, offload_options, + Qwen3SpeculativeOptions::disabled(), no_prefix_cache, max_prefill_tokens, ) diff --git a/openinfer-qwen3-4b/src/prefill.rs b/openinfer-qwen3-4b/src/prefill.rs index a0d65cd5..687c48ae 100644 --- a/openinfer-qwen3-4b/src/prefill.rs +++ b/openinfer-qwen3-4b/src/prefill.rs @@ -75,6 +75,20 @@ impl Qwen3Model { Ok(out) } + pub(super) fn get_embeddings_batch_into( + &self, + token_ids_gpu: &cudarc::driver::CudaSlice, + out: &mut HiddenStates, + ) -> Result<()> { + anyhow::ensure!( + out.hidden_dim == self.config.hidden_size, + "embedding output hidden_dim {} does not match model hidden_size {}", + out.hidden_dim, + self.config.hidden_size + ); + ops::embedding_batch(&self.ctx, &self.embed_tokens, token_ids_gpu, out) + } + #[allow(clippy::too_many_arguments)] fn forward_layer_batch_paged( &self, @@ -322,6 +336,29 @@ impl Qwen3Model { layout: &KvLayout, echo: bool, ) -> Result<(HiddenStates, Option)> { + let (logits, all_logits, _) = self.batch_prefill_with_hidden_capture( + prompts, + kv_views, + lora_adapters, + kv_buffer, + layout, + echo, + None, + )?; + Ok((logits, all_logits)) + } + + #[allow(clippy::too_many_arguments)] + pub(crate) fn batch_prefill_with_hidden_capture( + &self, + prompts: &[&[u32]], + kv_views: &[KvView], + lora_adapters: &[Option<&str>], + kv_buffer: &CudaSlice, + layout: &KvLayout, + echo: bool, + capture_layer_ids: Option<&[usize]>, + ) -> Result<(HiddenStates, Option, Option)> { let batch_size = prompts.len(); assert_eq!(batch_size, kv_views.len()); assert_eq!(batch_size, lora_adapters.len()); @@ -360,8 +397,14 @@ impl Qwen3Model { )?; // Forward through all layers - let hidden = - self.process_all_layers_batch_multi(hidden, layout, kv_buffer, &plan, &lora_groups)?; + let (hidden, captured_hidden) = self.process_all_layers_batch_multi( + hidden, + layout, + kv_buffer, + &plan, + &lora_groups, + capture_layer_ids, + )?; // All-position logits for echo (before we extract last-token logits) let all_logits = if echo { @@ -379,7 +422,7 @@ impl Qwen3Model { } let logits = self.batch_token_logits(&hidden, &last_indices)?; - Ok((logits, all_logits)) + Ok((logits, all_logits, captured_hidden)) } fn process_all_layers_batch_multi( @@ -389,11 +432,33 @@ impl Qwen3Model { kv_buffer: &cudarc::driver::CudaSlice, plan: &PrefillPagedPlan, lora_groups: &[DeviceLoraTokenGroup<'_>], - ) -> Result { + capture_layer_ids: Option<&[usize]>, + ) -> Result<(HiddenStates, Option)> { let total_tokens = hidden.seq_len; let inter_dim = self.local_intermediate_size(); let q_dim = self.local_q_dim(); let kv_dim = self.local_kv_dim(); + let capture_layer_ids = capture_layer_ids.unwrap_or(&[]); + anyhow::ensure!( + capture_layer_ids.windows(2).all(|pair| pair[0] < pair[1]), + "target hidden capture layer ids must be strictly increasing" + ); + anyhow::ensure!( + capture_layer_ids + .iter() + .all(|&layer| layer < self.layers.len()), + "target hidden capture layer id out of range" + ); + let mut captured_hidden = if capture_layer_ids.is_empty() { + None + } else { + Some(HiddenStates::zeros( + &self.ctx, + self.config.hidden_size * capture_layer_ids.len(), + total_tokens, + )?) + }; + let mut next_capture = 0usize; let mut bufs = PrefillBuffers::new( &self.ctx, @@ -415,8 +480,20 @@ impl Qwen3Model { lora_groups, &mut bufs, )?; + if capture_layer_ids.get(next_capture) == Some(&layer_idx) { + let out = captured_hidden + .as_mut() + .expect("capture buffer exists when ids are non-empty"); + ops::copy_hidden_rows_into( + &self.ctx, + &hidden, + out, + next_capture * self.config.hidden_size, + )?; + next_capture += 1; + } } - Ok(hidden) + Ok((hidden, captured_hidden)) } } diff --git a/openinfer-qwen3-4b/src/scheduler.rs b/openinfer-qwen3-4b/src/scheduler.rs index a5697de6..827939cc 100644 --- a/openinfer-qwen3-4b/src/scheduler.rs +++ b/openinfer-qwen3-4b/src/scheduler.rs @@ -19,14 +19,14 @@ use rand::rngs::StdRng; use tokio::sync::mpsc; use crate::executor::{ModelExecutor, Qwen3Executor, RequestId}; -use crate::{Qwen3LoraOptions, Qwen3OffloadOptions}; +use crate::{Qwen3LoraOptions, Qwen3OffloadOptions, Qwen3SpeculativeOptions}; use openinfer_core::engine::{ EngineCommand, EngineControlRequest, EngineHandle, GenerateRequest, KvCapacity, TokenEvent, }; use openinfer_core::sampler::SamplingParams; use self::effects::apply_effects; -use self::plan::{build_next_plan, execute_plan}; +use self::plan::{ExecutionPlan, build_next_plan, execute_plan, should_speculative_decode}; use self::resolve::resolve_step; // ── Internal types ────────────────────────────────────────────────────── @@ -138,15 +138,17 @@ pub(crate) fn start_qwen3( device_ordinals: &[usize], seed: u64, offload_options: Qwen3OffloadOptions, + speculative_options: Qwen3SpeculativeOptions, no_prefix_cache: bool, max_prefill_tokens: usize, ) -> Result { - let mut executor = Qwen3Executor::from_runtime_with_lora_options( + let mut executor = Qwen3Executor::from_runtime_with_options( model_path, enable_cuda_graph, device_ordinals, Qwen3LoraOptions::default(), offload_options, + speculative_options, )?; executor.set_no_prefix_cache(no_prefix_cache); Ok(start_with_executor(executor, seed, max_prefill_tokens)) @@ -159,15 +161,17 @@ pub(crate) fn start_qwen3_with_lora_control( seed: u64, lora_options: Qwen3LoraOptions, offload_options: Qwen3OffloadOptions, + speculative_options: Qwen3SpeculativeOptions, no_prefix_cache: bool, max_prefill_tokens: usize, ) -> Result { - let mut executor = Qwen3Executor::from_runtime_with_lora_options( + let mut executor = Qwen3Executor::from_runtime_with_options( model_path, enable_cuda_graph, device_ordinals, lora_options, offload_options, + speculative_options, )?; executor.set_no_prefix_cache(no_prefix_cache); Ok(start_with_executor_with_lora_control( @@ -430,6 +434,11 @@ fn scheduler_loop( executor.max_request_blocks(), executor.max_context_tokens(), executor.max_decode_batch_size(), + executor.speculative_state_budget_bytes(), + |id| executor.speculative_request_ready(id), + |prompt_len, max_tokens| { + executor.speculative_request_state_bytes(prompt_len, max_tokens) + }, |id| executor.prefetched_blocks(id), ); for (rejected, reason) in &admission.rejected { @@ -438,11 +447,8 @@ fn scheduler_loop( } prefilling.extend(admission.pending); deferred = admission.deferred; - let pending = take_prefill_chunks(&mut prefilling, max_prefill_tokens); - - let Some(plan) = build_next_plan(!active.is_empty(), pending) else { - continue; - }; + let plan = build_runtime_plan(&executor, &active, &mut prefilling, max_prefill_tokens); + let Some(plan) = plan else { continue }; let failure_targets = failure_targets_for(&active, &plan); let artifacts = match execute_plan(&mut executor, &mut active, plan, &mut rng) { Ok(v) => v, @@ -558,6 +564,11 @@ fn scheduler_loop_with_lora_control( executor.max_request_blocks(), executor.max_context_tokens(), executor.max_decode_batch_size(), + executor.speculative_state_budget_bytes(), + |id| executor.speculative_request_ready(id), + |prompt_len, max_tokens| { + executor.speculative_request_state_bytes(prompt_len, max_tokens) + }, |id| executor.prefetched_blocks(id), ); for (rejected, reason) in &admission.rejected { @@ -566,9 +577,8 @@ fn scheduler_loop_with_lora_control( } prefilling.extend(admission.pending); deferred = admission.deferred; - let pending = take_prefill_chunks(&mut prefilling, max_prefill_tokens); - - if active.is_empty() && pending.is_empty() { + let plan = build_runtime_plan(&executor, &active, &mut prefilling, max_prefill_tokens); + if plan.is_none() && active.is_empty() { // A parked load must still be polled to completion before we block. if !loading.is_empty() { block_on_loading(&mut executor, &mut deferred, &mut loading); @@ -589,9 +599,7 @@ fn scheduler_loop_with_lora_control( continue; } - let Some(plan) = build_next_plan(!active.is_empty(), pending) else { - continue; - }; + let Some(plan) = plan else { continue }; let failure_targets = failure_targets_for(&active, &plan); let artifacts = match execute_plan(&mut executor, &mut active, plan, &mut rng) { Ok(v) => v, @@ -606,6 +614,24 @@ fn scheduler_loop_with_lora_control( } } +fn build_runtime_plan( + executor: &impl ModelExecutor, + active: &[ActiveRequestState], + prefilling: &mut Vec, + max_prefill_tokens: usize, +) -> Option { + let pending = take_prefill_chunks(prefilling, max_prefill_tokens); + if should_speculative_decode(executor, active) { + if pending.is_empty() { + Some(ExecutionPlan::SpeculativeDecode) + } else { + Some(ExecutionPlan::Prefill { pending }) + } + } else { + build_next_plan(!active.is_empty(), pending) + } +} + fn enqueue_engine_command( command: EngineCommand, deferred: &mut Vec, @@ -813,6 +839,41 @@ fn prefilling_future_blocks( .sum() } +fn active_speculative_state_candidate(req: &ActiveRequestState) -> bool { + req.lora_adapter.is_none() && req.logprobs == 0 && req.params.is_greedy() +} + +fn pending_speculative_state_candidate(req: &PendingRequest) -> bool { + req.lora_adapter.is_none() + && req.cached_tokens == 0 + && req.logprobs == 0 + && !req.echo + && req.params.is_greedy() +} + +fn active_speculative_state_bytes( + active: &[ActiveRequestState], + state_exists: &impl Fn(RequestId) -> bool, + estimate: &impl Fn(usize, usize) -> usize, +) -> usize { + active + .iter() + .filter(|req| active_speculative_state_candidate(req) && state_exists(req.request_id)) + .map(|req| estimate(req.prompt_len, req.max_tokens)) + .fold(0usize, usize::saturating_add) +} + +fn prefilling_speculative_state_bytes( + prefilling: &[PendingRequest], + estimate: &impl Fn(usize, usize) -> usize, +) -> usize { + prefilling + .iter() + .filter(|req| pending_speculative_state_candidate(req)) + .map(|req| estimate(req.prompt_tokens.len(), req.max_tokens)) + .fold(0usize, usize::saturating_add) +} + /// Default for `max_prefill_tokens`: prompt tokens forwarded in a single step /// (chunked prefill). Prefill activation scratch scales with the step's total /// prompt tokens (~22 KB/token measured on Qwen3-4B), so an unbounded prefill @@ -842,6 +903,9 @@ fn admit_deferred_requests( max_request_blocks: usize, max_context_tokens: usize, max_decode_batch_size: usize, + speculative_state_budget_bytes: Option, + active_speculative_state_exists: impl Fn(RequestId) -> bool, + speculative_request_state_bytes: impl Fn(usize, usize) -> usize, // Blocks a request already holds from a settled prefetch. These are already // out of `available_blocks`, so they must be credited against the request's // need or admission double-counts them and can wedge a near-budget CPU-hit @@ -858,6 +922,18 @@ fn admit_deferred_requests( let mut decode_slots = max_decode_batch_size .saturating_sub(active.len()) .saturating_sub(prefilling.len()); + let mut speculative_state_budget = speculative_state_budget_bytes.map(|budget| { + budget + .saturating_sub(active_speculative_state_bytes( + active, + &active_speculative_state_exists, + &speculative_request_state_bytes, + )) + .saturating_sub(prefilling_speculative_state_bytes( + prefilling, + &speculative_request_state_bytes, + )) + }); let mut pending = Vec::new(); let mut still_deferred = Vec::new(); let mut rejected = Vec::new(); @@ -885,9 +961,22 @@ fn admit_deferred_requests( // …but only the blocks not already held by this request's prefetch must // come from the free-pool budget. let fresh_needed = footprint.saturating_sub(prefetch_credit(req.request_id)); + let speculative_state_needed = if pending_speculative_state_candidate(&req) { + speculative_request_state_bytes(req.prompt_tokens.len(), req.max_tokens) + } else { + 0 + }; + if speculative_state_budget.is_some_and(|budget| speculative_state_needed > budget) { + still_deferred.push(req); + continue; + } if fresh_needed <= budget && decode_slots > 0 { budget -= fresh_needed; decode_slots -= 1; + if let Some(speculative_state_budget) = speculative_state_budget.as_mut() { + *speculative_state_budget = + speculative_state_budget.saturating_sub(speculative_state_needed); + } pending.push(req); } else { still_deferred.push(req); @@ -939,11 +1028,15 @@ fn failure_targets_for( let mut targets = Vec::new(); match plan { self::plan::ExecutionPlan::Prefill { pending } => { + targets.extend(active.iter().map(active_failure_target)); targets.extend(pending.iter().map(pending_failure_target)); } self::plan::ExecutionPlan::Decode => { targets.extend(active.iter().map(active_failure_target)); } + self::plan::ExecutionPlan::SpeculativeDecode => { + targets.extend(active.iter().map(active_failure_target)); + } self::plan::ExecutionPlan::Unified { pending } => { targets.extend(active.iter().map(active_failure_target)); targets.extend(pending.iter().map(pending_failure_target)); @@ -1010,6 +1103,12 @@ mod tests { PrefillStepItem, UnifiedPlan, UnifiedResult, }; + #[test] + fn servable_len_reports_the_tighter_context_or_kv_limit() { + assert_eq!(servable_len(40944, 4096, 16), 40944); + assert_eq!(servable_len(40944, 1027, 16), 16432); + } + struct FakeExecutor { block_size: usize, cached_tokens: usize, @@ -1022,6 +1121,7 @@ mod tests { prefill_positions: HashMap, fail_decode_once: bool, decode_delay: Duration, + speculative_ready: HashSet, loaded_lora_adapters: HashSet, dropped: Arc>>, prefetch_offers: Arc>>, @@ -1043,6 +1143,7 @@ mod tests { prefill_positions: HashMap::new(), fail_decode_once: false, decode_delay: Duration::ZERO, + speculative_ready: HashSet::new(), loaded_lora_adapters: HashSet::new(), dropped, prefetch_offers: Arc::new(Mutex::new(Vec::new())), @@ -1073,6 +1174,11 @@ mod tests { self } + fn with_speculative_ready(mut self, request_id: RequestId) -> Self { + self.speculative_ready.insert(request_id); + self + } + fn with_lora_adapters(mut self, names: &[&str]) -> Self { self.loaded_lora_adapters = names.iter().map(|name| (*name).to_string()).collect(); self @@ -1195,6 +1301,14 @@ mod tests { names } + fn speculative_enabled(&self) -> bool { + !self.speculative_ready.is_empty() + } + + fn speculative_request_ready(&self, request_id: RequestId) -> bool { + self.speculative_ready.contains(&request_id) + } + fn unload_lora_adapter(&mut self, request: &UnloadLoraAdapterRequest) -> Result<()> { anyhow::ensure!( self.loaded_lora_adapters.remove(&request.lora_name), @@ -1226,6 +1340,7 @@ mod tests { .iter() .map(|req| self.fake_prefill_result(req)) .collect(), + dflash_context_captured_requests: Vec::new(), }) } @@ -1407,8 +1522,20 @@ mod tests { ]; // available 4 blocks - 2 reserved for active growth = budget of 2. - let outcome = - admit_deferred_requests(deferred, &active, &[], 16, 4, 4, usize::MAX, 64, |_| 0); + let outcome = admit_deferred_requests( + deferred, + &active, + &[], + 16, + 4, + 4, + usize::MAX, + 64, + None, + |_| false, + |_, _| 0, + |_| 0, + ); let ids = |reqs: &[PendingRequest]| reqs.iter().map(|r| r.request_id.get()).collect::>(); @@ -1447,8 +1574,20 @@ mod tests { mk(3, 40, 1), // request 3: 40 prompt + 1 max = 41 total: overflows by 9 tokens → rejected ]; - let outcome = - admit_deferred_requests(deferred, &active, &[], 16, 1000, 1000, 32, 64, |_| 0); + let outcome = admit_deferred_requests( + deferred, + &active, + &[], + 16, + 1000, + 1000, + 32, + 64, + None, + |_| false, + |_, _| 0, + |_| 0, + ); let pending_ids = outcome .pending @@ -1503,6 +1642,9 @@ mod tests { 1024, usize::MAX, 64, + None, + |_| false, + |_, _| 0, |_| 0, ); @@ -1518,6 +1660,98 @@ mod tests { assert!(outcome.rejected.is_empty()); } + #[test] + fn admission_respects_speculative_state_budget_for_supported_requests() { + let active: [ActiveRequestState; 0] = []; + let mk = |id: u64| PendingRequest::from_scheduler_request(RequestId(id), request(8, 2).0); + let mut unsupported = mk(3); + unsupported.params.temperature = 0.7; + let deferred = vec![mk(1), mk(2), unsupported]; + + let outcome = admit_deferred_requests( + deferred, + &active, + &[], + 16, + 1024, + 1024, + usize::MAX, + 64, + Some(10), + |_| false, + |prompt_len, max_tokens| prompt_len + max_tokens, + |_| 0, + ); + + let ids = + |reqs: &[PendingRequest]| reqs.iter().map(|r| r.request_id.get()).collect::>(); + assert_eq!( + ids(&outcome.pending), + vec![1, 3], + "one supported request consumes the speculative side-state budget; unsupported fallback does not" + ); + assert_eq!( + ids(&outcome.deferred), + vec![2], + "extra supported requests wait for DFlash state memory instead of OOMing later" + ); + assert!(outcome.rejected.is_empty()); + } + + #[test] + fn admission_counts_only_active_requests_with_speculative_state() { + let (token_tx, _rx) = mpsc::unbounded_channel(); + let active = [ActiveRequestState { + request_id: RequestId(1), + lora_adapter: None, + token_tx, + last_token: 1, + generated_count: 1, + max_tokens: 2, + prompt_len: 8, + params: SamplingParams::default(), + logprobs: 0, + }]; + let mk_pending = || PendingRequest::from_scheduler_request(RequestId(2), request(8, 2).0); + + let fallback_active = admit_deferred_requests( + vec![mk_pending()], + &active, + &[], + 16, + 1024, + 1024, + usize::MAX, + 64, + Some(10), + |_| false, + |prompt_len, max_tokens| prompt_len + max_tokens, + |_| 0, + ); + assert_eq!( + fallback_active.pending[0].request_id, + RequestId(2), + "a greedy active request that already fell back to baseline must not consume DFlash side-state budget" + ); + + let dflash_active = admit_deferred_requests( + vec![mk_pending()], + &active, + &[], + 16, + 1024, + 1024, + usize::MAX, + 64, + Some(10), + |id| id == RequestId(1), + |prompt_len, max_tokens| prompt_len + max_tokens, + |_| 0, + ); + assert!(dflash_active.pending.is_empty()); + assert_eq!(dflash_active.deferred[0].request_id, RequestId(2)); + } + #[test] fn prefill_chunking_caps_step_tokens_and_keeps_fifo_progress() { let mk = |id: u64, prompt_len, max_tokens| { @@ -1597,6 +1831,147 @@ mod tests { assert_eq!(prefilling[0].request_id, RequestId(4)); } + #[test] + fn dflash_ready_active_prefills_waiting_request_before_next_draft() { + let dropped = Arc::new(Mutex::new(Vec::new())); + let executor = FakeExecutor::new(64, dropped).with_speculative_ready(RequestId(7)); + let (token_tx, _rx) = mpsc::unbounded_channel(); + let active = vec![ActiveRequestState { + request_id: RequestId(7), + lora_adapter: None, + token_tx, + last_token: 42, + generated_count: 1, + max_tokens: 32, + prompt_len: 16, + params: SamplingParams::default(), + logprobs: 0, + }]; + let mut prefilling = vec![PendingRequest::from_scheduler_request( + RequestId(8), + request(16, 1).0, + )]; + + match build_runtime_plan(&executor, &active, &mut prefilling, 1024) { + Some(ExecutionPlan::Prefill { pending }) => { + assert_eq!(pending.len(), 1); + assert_eq!(pending[0].request_id, RequestId(8)); + } + _ => panic!("DFlash should prefill waiting requests so they become DFlash-ready"), + } + assert!(prefilling.is_empty()); + } + + #[test] + fn dflash_ready_active_prefills_waiting_batch_before_next_draft() { + let dropped = Arc::new(Mutex::new(Vec::new())); + let executor = FakeExecutor::new(64, dropped).with_speculative_ready(RequestId(7)); + let (token_tx, _rx) = mpsc::unbounded_channel(); + let active = vec![ActiveRequestState { + request_id: RequestId(7), + lora_adapter: None, + token_tx, + last_token: 42, + generated_count: 1, + max_tokens: 32, + prompt_len: 16, + params: SamplingParams::default(), + logprobs: 0, + }]; + let mut prefilling = vec![ + PendingRequest::from_scheduler_request(RequestId(8), request(16, 1).0), + PendingRequest::from_scheduler_request(RequestId(9), request(16, 1).0), + ]; + + match build_runtime_plan(&executor, &active, &mut prefilling, 1024) { + Some(ExecutionPlan::Prefill { pending }) => { + assert_eq!( + pending.iter().map(|req| req.request_id).collect::>(), + vec![RequestId(8), RequestId(9)] + ); + } + _ => panic!("DFlash should prefill the waiting batch before the next draft"), + } + assert!(prefilling.is_empty()); + + let mut prefilling = Vec::new(); + match build_runtime_plan(&executor, &active, &mut prefilling, 1024) { + Some(ExecutionPlan::SpeculativeDecode) => {} + _ => panic!("ready single active request should speculate once pending is drained"), + } + } + + #[test] + fn prefill_failure_with_active_targets_active_and_pending_requests() { + let (active_tx, _active_rx) = mpsc::unbounded_channel(); + let active = vec![ActiveRequestState { + request_id: RequestId(7), + lora_adapter: None, + token_tx: active_tx, + last_token: 42, + generated_count: 1, + max_tokens: 32, + prompt_len: 16, + params: SamplingParams::default(), + logprobs: 0, + }]; + let pending = vec![PendingRequest::from_scheduler_request( + RequestId(8), + request(16, 1).0, + )]; + let plan = ExecutionPlan::Prefill { pending }; + + let targets = failure_targets_for(&active, &plan); + assert_eq!( + targets + .iter() + .map(|target| target.request_id) + .collect::>(), + vec![RequestId(7), RequestId(8)], + "a prefill step while active requests exist can dirty the shared executor step and must fail both sides" + ); + } + + #[test] + fn dflash_ready_multi_active_speculates_as_one_batch() { + let dropped = Arc::new(Mutex::new(Vec::new())); + let executor = FakeExecutor::new(64, dropped) + .with_speculative_ready(RequestId(7)) + .with_speculative_ready(RequestId(8)); + let (tx0, _rx0) = mpsc::unbounded_channel(); + let (tx1, _rx1) = mpsc::unbounded_channel(); + let active = vec![ + ActiveRequestState { + request_id: RequestId(7), + lora_adapter: None, + token_tx: tx0, + last_token: 42, + generated_count: 1, + max_tokens: 32, + prompt_len: 16, + params: SamplingParams::default(), + logprobs: 0, + }, + ActiveRequestState { + request_id: RequestId(8), + lora_adapter: None, + token_tx: tx1, + last_token: 43, + generated_count: 1, + max_tokens: 32, + prompt_len: 16, + params: SamplingParams::default(), + logprobs: 0, + }, + ]; + let mut prefilling = Vec::new(); + + match build_runtime_plan(&executor, &active, &mut prefilling, 1024) { + Some(ExecutionPlan::SpeculativeDecode) => {} + _ => panic!("all-ready active requests should run DFlash together"), + } + } + #[test] fn long_prompt_chunks_across_steps_and_still_completes() { let dropped = Arc::new(Mutex::new(Vec::new())); @@ -1745,6 +2120,9 @@ mod tests { 2, usize::MAX, 64, + None, + |_| false, + |_, _| 0, |_| 0, ); assert!( @@ -1766,6 +2144,9 @@ mod tests { 3, usize::MAX, 64, + None, + |_| false, + |_, _| 0, |_| 0, ); assert_eq!( diff --git a/openinfer-qwen3-4b/src/scheduler/effects.rs b/openinfer-qwen3-4b/src/scheduler/effects.rs index c64c5ccb..aa41a2e9 100644 --- a/openinfer-qwen3-4b/src/scheduler/effects.rs +++ b/openinfer-qwen3-4b/src/scheduler/effects.rs @@ -69,6 +69,17 @@ pub(super) enum DecodeEffect { logprob: Option, completion_tokens: usize, }, + EmitManyAndContinue { + request_id: RequestId, + tokens: Vec, + completion_tokens: usize, + }, + EmitManyAndFinish { + request_id: RequestId, + tokens: Vec, + finish_reason: FinishReason, + completion_tokens: usize, + }, } pub(super) struct StepEffects { @@ -180,6 +191,73 @@ pub(super) fn apply_effects( req.generated_count = completion_tokens; } } + DecodeEffect::EmitManyAndContinue { + request_id, + tokens, + completion_tokens, + } => { + let Some(index) = active.iter().position(|req| req.request_id == request_id) else { + continue; + }; + let req = &mut active[index]; + let mut sent = true; + for token in &tokens { + if req + .token_tx + .send(TokenEvent::Token { + id: *token, + logprob: None, + }) + .is_err() + { + sent = false; + break; + } + } + if sent { + req.last_token = *tokens + .last() + .expect("EmitManyAndContinue must carry at least one token"); + req.generated_count = completion_tokens; + } else { + let _ = executor.drop_request(request_id); + to_retire.push(index); + } + } + DecodeEffect::EmitManyAndFinish { + request_id, + tokens, + finish_reason, + completion_tokens, + } => { + let Some(index) = active.iter().position(|req| req.request_id == request_id) else { + continue; + }; + let req = &active[index]; + let mut sent = true; + for token in &tokens { + if req + .token_tx + .send(TokenEvent::Token { + id: *token, + logprob: None, + }) + .is_err() + { + sent = false; + break; + } + } + if sent { + let _ = req.token_tx.send(TokenEvent::Finished { + finish_reason, + prompt_tokens: req.prompt_len, + completion_tokens, + }); + } + let _ = executor.drop_request(request_id); + to_retire.push(index); + } } } to_retire.sort_unstable(); diff --git a/openinfer-qwen3-4b/src/scheduler/plan.rs b/openinfer-qwen3-4b/src/scheduler/plan.rs index 7a3ae694..96a8d8db 100644 --- a/openinfer-qwen3-4b/src/scheduler/plan.rs +++ b/openinfer-qwen3-4b/src/scheduler/plan.rs @@ -5,12 +5,18 @@ use crate::executor::{ DecodePlan, DecodeResult, DecodeStepItem, ModelExecutor, PrefillPlan, PrefillResult, PrefillStepItem, UnifiedPlan, UnifiedResult, }; +use crate::speculative::{ + DraftPlan as SpeculativeDraftPlan, DraftRequestResult as SpeculativeDraftRequestResult, + DraftStepItem as SpeculativeDraftStepItem, VerifyPlan as SpeculativeVerifyPlan, + VerifyResult as SpeculativeVerifyResult, VerifyStepItem as SpeculativeVerifyStepItem, +}; use super::{ActiveRequestState, PendingRequest}; pub(super) enum ExecutionPlan { Prefill { pending: Vec }, Decode, + SpeculativeDecode, Unified { pending: Vec }, } @@ -26,6 +32,9 @@ pub(super) enum ExecutionArtifacts { Decode { result: DecodeResult, }, + SpeculativeDecode { + verify: SpeculativeVerifyResult, + }, Unified { pending: Vec, result: UnifiedResult, @@ -80,6 +89,19 @@ pub(super) fn execute_plan( sort_decode_results(&mut result.requests); Ok(ExecutionArtifacts::Decode { result }) } + ExecutionPlan::SpeculativeDecode => { + let draft_requests = build_speculative_draft_items(active); + let mut draft = executor.execute_speculative_draft(SpeculativeDraftPlan { + requests: &draft_requests, + })?; + draft.requests.sort_by_key(|result| result.request_id); + let verify_requests = build_speculative_verify_items(active, &draft.requests); + let mut verify = executor.execute_speculative_verify(SpeculativeVerifyPlan { + requests: &verify_requests, + })?; + verify.requests.sort_by_key(|result| result.request_id); + Ok(ExecutionArtifacts::SpeculativeDecode { verify }) + } ExecutionPlan::Unified { pending } => { let scheduled_at_unix_s = openinfer_core::engine::unix_now_s(); let pending_indices: Vec = (0..pending.len()).collect(); @@ -101,6 +123,20 @@ pub(super) fn execute_plan( } } +pub(super) fn should_speculative_decode( + executor: &impl ModelExecutor, + active: &[ActiveRequestState], +) -> bool { + executor.speculative_enabled() + && !active.is_empty() + && active.iter().all(|req| { + executor.speculative_request_ready(req.request_id) + && req.lora_adapter.is_none() + && req.logprobs == 0 + && req.params.is_greedy() + }) +} + fn build_prefill_items( pending: &[PendingRequest], indices: &[usize], @@ -149,6 +185,44 @@ fn build_decode_items( .collect() } +fn build_speculative_draft_items(active: &[ActiveRequestState]) -> Vec { + active + .iter() + .map(|r| SpeculativeDraftStepItem { + request_id: r.request_id, + current_token: r.last_token, + params: r.params, + }) + .collect() +} + +fn build_speculative_verify_items( + active: &[ActiveRequestState], + draft_results: &[SpeculativeDraftRequestResult], +) -> Vec { + draft_results + .iter() + .map(|draft| { + let active = active + .iter() + .find(|req| req.request_id == draft.request_id) + .expect("draft request_id must exist in active set"); + SpeculativeVerifyStepItem { + request_id: draft.request_id, + token_ids: { + let remaining = active.max_tokens.saturating_sub(active.generated_count); + assert!(remaining > 0, "active request must have output budget"); + let mut token_ids = draft.token_ids.clone(); + token_ids.truncate(remaining); + token_ids + }, + params: active.params, + lora_adapter: None, + } + }) + .collect() +} + fn sort_prefill_results(results: &mut [crate::executor::PrefillRequestResult]) { results.sort_by_key(|result| result.request_id); } @@ -182,6 +256,21 @@ mod tests { } } + fn active(generated_count: usize, max_tokens: usize) -> ActiveRequestState { + let (token_tx, _rx) = tokio::sync::mpsc::unbounded_channel(); + ActiveRequestState { + request_id: RequestId::new(7), + lora_adapter: None, + token_tx, + last_token: 42, + generated_count, + max_tokens, + prompt_len: 10, + params: SamplingParams::default(), + logprobs: 0, + } + } + // The plan selector is the whole batch-formation policy: what the scheduler // does each tick is fully determined by (have_active, has_pending). Pin the // 2×2 truth table so a policy regression can't slip through silently. @@ -210,4 +299,19 @@ mod tests { "active + pending fuses prefill and decode into one unified step" ); } + + #[test] + fn speculative_verify_items_clamp_to_remaining_output_budget() { + let active = [active(24, 32)]; + let draft = SpeculativeDraftRequestResult { + request_id: RequestId::new(7), + token_ids: (0..16).collect(), + }; + + let verify = build_speculative_verify_items(&active, &[draft]); + + assert_eq!(verify.len(), 1); + assert_eq!(verify[0].token_ids.len(), 8); + assert_eq!(verify[0].token_ids, (0..8).collect::>()); + } } diff --git a/openinfer-qwen3-4b/src/scheduler/resolve.rs b/openinfer-qwen3-4b/src/scheduler/resolve.rs index 1f91dcd5..7688e6ad 100644 --- a/openinfer-qwen3-4b/src/scheduler/resolve.rs +++ b/openinfer-qwen3-4b/src/scheduler/resolve.rs @@ -1,4 +1,5 @@ use crate::executor::{DecodeRequestResult, ModelExecutor, PrefillRequestResult}; +use crate::speculative::VerifyRequestResult as SpeculativeVerifyRequestResult; use openinfer_core::engine::FinishReason; use super::effects::{DecodeEffect, PendingEffect, PromptEchoEffect, ScheduledEffect, StepEffects}; @@ -22,6 +23,12 @@ pub(super) fn resolve_step( pending: Vec::new(), decode: resolve_decode_outputs(executor, active, &result.requests), }, + ExecutionArtifacts::SpeculativeDecode { verify } => StepEffects { + scheduled: Vec::new(), + prompt_echoes: Vec::new(), + pending: Vec::new(), + decode: resolve_speculative_outputs(executor, active, &verify.requests), + }, ExecutionArtifacts::Unified { pending, result, @@ -39,6 +46,50 @@ pub(super) fn resolve_step( } } +fn resolve_speculative_outputs( + executor: &impl ModelExecutor, + active: &[ActiveRequestState], + request_results: &[SpeculativeVerifyRequestResult], +) -> Vec { + request_results + .iter() + .map(|result| { + let req = active + .iter() + .find(|req| req.request_id == result.request_id) + .expect("speculative request_id must exist in active set"); + let mut emitted = Vec::new(); + let mut completion_tokens = req.generated_count; + for &token in &result.accepted_tokens { + completion_tokens += 1; + let is_eos = !req.params.ignore_eos && executor.is_stop_token(token); + if is_eos { + return DecodeEffect::EmitManyAndFinish { + request_id: result.request_id, + tokens: emitted, + finish_reason: FinishReason::Stop, + completion_tokens, + }; + } + emitted.push(token); + if completion_tokens >= req.max_tokens { + return DecodeEffect::EmitManyAndFinish { + request_id: result.request_id, + tokens: emitted, + finish_reason: FinishReason::Length, + completion_tokens, + }; + } + } + DecodeEffect::EmitManyAndContinue { + request_id: result.request_id, + tokens: emitted, + completion_tokens, + } + }) + .collect() +} + fn resolve_prefill_outputs( executor: &impl ModelExecutor, pending: Vec, diff --git a/openinfer-qwen3-4b/src/speculative.rs b/openinfer-qwen3-4b/src/speculative.rs new file mode 100644 index 00000000..c195bd5c --- /dev/null +++ b/openinfer-qwen3-4b/src/speculative.rs @@ -0,0 +1,163 @@ +use anyhow::Result; + +use crate::executor::RequestId; +use openinfer_core::sampler::SamplingParams; + +pub struct VerifyPlan<'a> { + pub requests: &'a [VerifyStepItem], +} + +pub struct DraftPlan<'a> { + pub requests: &'a [DraftStepItem], +} + +#[derive(Clone)] +pub struct VerifyStepItem { + pub(crate) request_id: RequestId, + /// Verify-span token ids: current dangling token first, then draft + /// candidates. DFlash uses a fixed 16-token span. + pub(crate) token_ids: Vec, + pub(crate) params: SamplingParams, + pub(crate) lora_adapter: Option, +} + +impl VerifyStepItem { + pub fn new(request_id: RequestId, token_ids: Vec, params: SamplingParams) -> Self { + Self { + request_id, + token_ids, + params, + lora_adapter: None, + } + } + + pub(crate) fn as_slice(&self) -> &[u32] { + &self.token_ids + } +} + +#[derive(Clone, Debug)] +pub struct VerifyRequestResult { + pub request_id: RequestId, + /// Number of draft candidates accepted before the posterior bonus. + pub matched_draft_tokens: usize, + /// Tokens committed to request state: accepted draft prefix followed by + /// the target posterior token at the first mismatch (or at the block end). + /// The scheduler still owns stop-token suppression before client emission. + pub accepted_tokens: Vec, + /// Target-selected posterior token for each position in the verify span. + pub target_tokens: Vec, +} + +pub struct VerifyResult { + pub requests: Vec, +} + +#[derive(Clone)] +pub struct DraftStepItem { + pub(crate) request_id: RequestId, + pub(crate) current_token: u32, + pub(crate) params: SamplingParams, +} + +impl DraftStepItem { + pub fn new(request_id: RequestId, current_token: u32, params: SamplingParams) -> Self { + Self { + request_id, + current_token, + params, + } + } +} + +#[derive(Clone, Debug)] +pub struct DraftRequestResult { + pub request_id: RequestId, + /// Verify-span tokens: current dangling token first, then draft candidates. + pub token_ids: Vec, +} + +pub struct DraftResult { + pub requests: Vec, +} + +pub(crate) fn build_verify_results( + requests: &[VerifyStepItem], + target_tokens: &[u32], +) -> Result> { + let mut outputs = Vec::with_capacity(requests.len()); + let mut offset = 0usize; + for req in requests { + let span_len = req.token_ids.len(); + anyhow::ensure!( + span_len > 0, + "speculative verify request {:?} has an empty verify span", + req.request_id + ); + let end = offset + span_len; + anyhow::ensure!( + end <= target_tokens.len(), + "speculative target-token result is shorter than the verify span" + ); + let posterior = &target_tokens[offset..end]; + + let mut matched = 0usize; + while matched + 1 < span_len && req.token_ids[matched + 1] == posterior[matched] { + matched += 1; + } + + let mut accepted_tokens = req.token_ids[1..1 + matched].to_vec(); + accepted_tokens.push(posterior[matched]); + outputs.push(VerifyRequestResult { + request_id: req.request_id, + matched_draft_tokens: matched, + accepted_tokens, + target_tokens: posterior.to_vec(), + }); + offset = end; + } + anyhow::ensure!( + offset == target_tokens.len(), + "unused speculative target-token result columns: used {offset}, total {}", + target_tokens.len() + ); + Ok(outputs) +} + +#[cfg(test)] +mod tests { + use super::{VerifyStepItem, build_verify_results}; + use crate::executor::RequestId; + use openinfer_core::sampler::SamplingParams; + + #[test] + fn speculative_verify_accepts_matching_prefix_plus_posterior_bonus() { + let req = VerifyStepItem::new( + RequestId::new(7), + vec![10, 11, 12, 13], + SamplingParams::default(), + ); + + let results = build_verify_results(&[req], &[11, 12, 99, 100]).expect("verify results"); + + assert_eq!(results.len(), 1); + assert_eq!(results[0].request_id, RequestId::new(7)); + assert_eq!(results[0].matched_draft_tokens, 2); + assert_eq!(results[0].accepted_tokens, vec![11, 12, 99]); + assert_eq!(results[0].target_tokens, vec![11, 12, 99, 100]); + } + + #[test] + fn speculative_verify_all_match_still_adds_block_end_posterior() { + let req = VerifyStepItem::new( + RequestId::new(8), + vec![20, 21, 22], + SamplingParams::default(), + ); + + let results = build_verify_results(&[req], &[21, 22, 23]).expect("verify results"); + + assert_eq!(results[0].matched_draft_tokens, 2); + assert_eq!(results[0].accepted_tokens, vec![21, 22, 23]); + } +} diff --git a/openinfer-qwen3-4b/src/weights.rs b/openinfer-qwen3-4b/src/weights.rs index c7549aff..e200a5dc 100644 --- a/openinfer-qwen3-4b/src/weights.rs +++ b/openinfer-qwen3-4b/src/weights.rs @@ -833,6 +833,10 @@ impl Qwen3Model { /// KV cache geometry and budget for the executor to create a KvCacheManager. pub(crate) fn kv_budget(&self) -> KvBudget { + self.kv_budget_with_reserved_bytes(0) + } + + pub(crate) fn kv_budget_with_reserved_bytes(&self, reserved_bytes: usize) -> KvBudget { let page_size = 16; let num_kv_heads = self.local_num_key_value_heads(); let layout = openinfer_kv_cache::KvLayout::new( @@ -843,14 +847,30 @@ impl Qwen3Model { ); let bytes_per_block = layout.page_stride * std::mem::size_of::(); let (free_bytes, _) = cudarc::driver::result::mem_get_info().expect("cuMemGetInfo failed"); - let kv_budget_bytes = (free_bytes as f64 * 0.85) as usize; - let num_blocks = (kv_budget_bytes / bytes_per_block).max(64); - let kv_mb = num_blocks * bytes_per_block / (1024 * 1024); - log::info!( - "KV cache: {num_blocks} blocks ({kv_mb} MB, {:.0}% of {:.0} MB free)", - kv_budget_bytes as f64 / free_bytes as f64 * 100.0, - free_bytes as f64 / 1024.0 / 1024.0 + let usable_free_bytes = free_bytes.saturating_sub(reserved_bytes); + let max_blocks = usable_free_bytes / bytes_per_block; + assert!( + max_blocks >= 2, + "not enough GPU memory for Qwen3 KV cache after reserving {} MB", + reserved_bytes / (1024 * 1024) ); + let kv_budget_bytes = (usable_free_bytes as f64 * 0.85) as usize; + let num_blocks = (kv_budget_bytes / bytes_per_block).max(64).min(max_blocks); + let kv_mb = num_blocks * bytes_per_block / (1024 * 1024); + if reserved_bytes == 0 { + log::info!( + "KV cache: {num_blocks} blocks ({kv_mb} MB, {:.0}% of {:.0} MB free)", + kv_budget_bytes as f64 / free_bytes as f64 * 100.0, + free_bytes as f64 / 1024.0 / 1024.0 + ); + } else { + log::info!( + "KV cache: {num_blocks} blocks ({kv_mb} MB, {:.0}% of {:.0} MB free after reserving {:.0} MB)", + kv_budget_bytes as f64 / free_bytes as f64 * 100.0, + free_bytes as f64 / 1024.0 / 1024.0, + reserved_bytes as f64 / 1024.0 / 1024.0 + ); + } KvBudget { num_layers: self.config.num_hidden_layers, num_kv_heads, diff --git a/openinfer-qwen3-4b/tests/hf_golden_gate.rs b/openinfer-qwen3-4b/tests/hf_golden_gate.rs index 105fb553..7c5b639c 100644 --- a/openinfer-qwen3-4b/tests/hf_golden_gate.rs +++ b/openinfer-qwen3-4b/tests/hf_golden_gate.rs @@ -58,7 +58,10 @@ use openinfer_core::engine::TokenLogprob; use openinfer_core::sampler::SamplingParams; use openinfer_qwen3_4b::runtime::{ DecodePlan, DecodeStepItem, PrefillPlan, PrefillStepItem, Qwen3Executor, RequestId, + SpeculativeDraftPlan, SpeculativeDraftStepItem, SpeculativeVerifyPlan, + SpeculativeVerifyStepItem, }; +use openinfer_qwen3_4b::{Qwen3LoraOptions, Qwen3OffloadOptions, Qwen3SpeculativeOptions}; use safetensors::{Dtype, SafeTensors}; const MODEL_PATH: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/../models/Qwen3-4B"); @@ -126,6 +129,25 @@ fn model_path_or_skip() -> Option { } } +fn dflash_model_path_or_skip() -> Option { + match std::env::var("OPENINFER_DFLASH_TEST_MODEL_PATH") { + Ok(path) => Some(path), + Err(_) + if Path::new("/data/models/Qwen3-4B-DFlash-b16") + .join("config.json") + .exists() => + { + Some("/data/models/Qwen3-4B-DFlash-b16".to_string()) + } + Err(_) => { + eprintln!( + "skipping qwen3 DFlash verifier-span golden gate: set OPENINFER_DFLASH_TEST_MODEL_PATH" + ); + None + } + } +} + fn as_i32(st: &SafeTensors, name: &str) -> (Vec, Vec) { let t = st .tensor(name) @@ -289,6 +311,26 @@ fn prefill_item(id: RequestId, prompt: Vec) -> PrefillStepItem { ) } +fn dflash_prefill_item(id: RequestId, prompt: Vec) -> PrefillStepItem { + PrefillStepItem::new( + id, + prompt, + MAX_OUTPUT_TOKENS, + SamplingParams::default(), + 0, + false, + 0.0, + ) +} + +fn dflash_chunked_prefill_item( + id: RequestId, + prompt: Vec, + chunk_budget: usize, +) -> PrefillStepItem { + dflash_prefill_item(id, prompt).with_chunk_budget(chunk_budget) +} + fn decode_item(id: RequestId, fed: u32) -> DecodeStepItem { DecodeStepItem::new(id, fed, SamplingParams::default(), LOGPROBS, 0.0) } @@ -442,6 +484,21 @@ fn report_and_assert(label: &str, stats: &Stats) { let _ = max; // reported above, but not asserted: the worst single delta grows with coverage } +fn assert_token_within_hf_regret(label: &str, token: u32, hf: &[(u32, f32)]) { + let hf_top = hf[0].1; + let Some((_, token_lp)) = hf.iter().find(|(id, _)| *id == token) else { + panic!( + "{label}: verifier selected token {token}, which is absent from HF top-{}", + hf.len() + ); + }; + let regret = hf_top - *token_lp; + assert!( + regret <= MARGIN_TOL, + "{label}: verifier selected token {token}, which HF scores {regret:.4} nat below its best (> {MARGIN_TOL})" + ); +} + /// Number of visible CUDA devices; 0 when the driver cannot be queried. fn cuda_device_count() -> usize { cudarc::driver::CudaContext::device_count().map_or(0, |n| n.max(0) as usize) @@ -568,3 +625,195 @@ fn pega_logprobs_match_hf_golden_within_bf16_tolerance() { } run_eager_suite(&golden, &model_path, &[0, 1], "tp2 "); } + +#[test] +fn dflash_speculative_verify_matches_hf_argmax_regret_gate() { + let Some(model_path) = model_path_or_skip() else { + return; + }; + let Some(dflash_path) = dflash_model_path_or_skip() else { + return; + }; + let golden = Golden::load(); + let span_len = 16.min(golden.decode_len); + assert!(span_len > 0, "golden decode span must not be empty"); + + let mut ex = Qwen3Executor::from_runtime_with_options( + &model_path, + false, + &[0], + Qwen3LoraOptions::default(), + Qwen3OffloadOptions::disabled(), + Qwen3SpeculativeOptions::dflash(dflash_path), + ) + .expect("build DFlash-enabled executor"); + ex.set_prefix_cache_enabled(false); + + let seqs = 0..golden.num_seqs.min(8); + let mut checked = 0usize; + for seq in seqs { + let id = RequestId::new(50_000 + seq as u64); + ex.execute_prefill(PrefillPlan { + requests: &[dflash_prefill_item(id, golden.prompt(seq))], + echo: false, + }) + .expect("DFlash prefill hidden capture"); + + let verify_tokens: Vec = (0..span_len).map(|step| golden.decode(seq, step)).collect(); + let result = ex + .execute_speculative_verify(SpeculativeVerifyPlan { + requests: &[SpeculativeVerifyStepItem::new( + id, + verify_tokens, + SamplingParams::default(), + )], + }) + .expect("DFlash speculative verify"); + let request = &result.requests[0]; + assert_eq!(request.request_id, id); + assert_eq!(request.target_tokens.len(), span_len); + for (pos, &token) in request.target_tokens.iter().enumerate() { + assert_token_within_hf_regret( + &format!("seq {seq} verify pos {pos}"), + token, + &golden.topk(seq, pos + 1), + ); + checked += 1; + } + ex.drop_request(id).expect("drop DFlash golden request"); + } + + let mut full_path_matched = 0usize; + for seq in 0..golden.num_seqs.min(8) { + let id = RequestId::new(70_000 + seq as u64); + ex.execute_prefill(PrefillPlan { + requests: &[dflash_prefill_item(id, golden.prompt(seq))], + echo: false, + }) + .expect("DFlash full-path prefill hidden capture"); + + let current_token = golden.decode(seq, 0); + let draft = ex + .execute_speculative_draft(SpeculativeDraftPlan { + requests: &[SpeculativeDraftStepItem::new( + id, + current_token, + SamplingParams::default(), + )], + }) + .expect("DFlash native draft"); + let draft_request = &draft.requests[0]; + assert_eq!(draft_request.request_id, id); + assert_eq!(draft_request.token_ids[0], current_token); + assert_eq!( + draft_request.token_ids.len(), + span_len, + "DFlash draft span should match the verifier block size" + ); + + let result = ex + .execute_speculative_verify(SpeculativeVerifyPlan { + requests: &[SpeculativeVerifyStepItem::new( + id, + draft_request.token_ids.clone(), + SamplingParams::default(), + )], + }) + .expect("DFlash native draft verify"); + let request = &result.requests[0]; + assert_eq!(request.request_id, id); + assert_eq!(request.target_tokens.len(), draft_request.token_ids.len()); + assert_token_within_hf_regret( + &format!("full-path seq {seq} verify pos 0"), + request.target_tokens[0], + &golden.topk(seq, 1), + ); + assert_eq!( + request.accepted_tokens.len(), + request.matched_draft_tokens + 1 + ); + assert_eq!( + request.accepted_tokens.last().copied(), + request + .target_tokens + .get(request.matched_draft_tokens) + .copied() + ); + for matched in 0..request.matched_draft_tokens { + assert_eq!( + request.accepted_tokens[matched], + draft_request.token_ids[matched + 1], + "full-path accepted draft token mismatch at seq {seq} matched {matched}" + ); + } + full_path_matched += request.matched_draft_tokens; + checked += 1; + ex.drop_request(id) + .expect("drop DFlash full-path golden request"); + } + assert!( + full_path_matched > 0, + "DFlash full draft→verify gate accepted no draft candidates" + ); + + let chunked_seq = 0usize; + let chunked_prompt = golden.prompt(chunked_seq); + assert!( + chunked_prompt.len() >= 2, + "chunked DFlash gate needs a prompt with at least two tokens" + ); + let chunk_budget = (chunked_prompt.len() / 2).max(1); + let chunked_id = RequestId::new(60_000); + let mut chunked_completed = false; + let mut chunked_steps = 0usize; + while !chunked_completed { + let result = ex + .execute_prefill(PrefillPlan { + requests: &[dflash_chunked_prefill_item( + chunked_id, + chunked_prompt.clone(), + chunk_budget, + )], + echo: false, + }) + .expect("DFlash chunked prefill hidden capture"); + let request = &result.requests[0]; + assert_eq!(request.request_id, chunked_id); + chunked_completed = request.completed; + chunked_steps += 1; + assert!( + chunked_steps <= chunked_prompt.len(), + "chunked DFlash prefill failed to make progress" + ); + } + assert!( + chunked_steps > 1, + "DFlash chunked prefill gate did not force multiple chunks" + ); + let verify_tokens: Vec = (0..span_len) + .map(|step| golden.decode(chunked_seq, step)) + .collect(); + let result = ex + .execute_speculative_verify(SpeculativeVerifyPlan { + requests: &[SpeculativeVerifyStepItem::new( + chunked_id, + verify_tokens, + SamplingParams::default(), + )], + }) + .expect("DFlash speculative verify after chunked prefill"); + let request = &result.requests[0]; + assert_eq!(request.request_id, chunked_id); + assert_eq!(request.target_tokens.len(), span_len); + for (pos, &token) in request.target_tokens.iter().enumerate() { + assert_token_within_hf_regret( + &format!("chunked seq {chunked_seq} verify pos {pos}"), + token, + &golden.topk(chunked_seq, pos + 1), + ); + checked += 1; + } + ex.drop_request(chunked_id) + .expect("drop chunked DFlash golden request"); + eprintln!("dflash verifier-span golden gate checked {checked} target positions"); +} diff --git a/openinfer-server/src/main.rs b/openinfer-server/src/main.rs index 5ca4eec2..0db55f27 100644 --- a/openinfer-server/src/main.rs +++ b/openinfer-server/src/main.rs @@ -11,7 +11,7 @@ use openinfer_core::engine::{EngineHandle, EngineLoadOptions, EpBackend}; #[cfg(feature = "kimi-k2")] use openinfer_core::parallel::ParallelConfig; #[cfg(feature = "qwen3-4b")] -use openinfer_qwen3_4b::{Qwen3LoraOptions, Qwen3OffloadOptions}; +use openinfer_qwen3_4b::{Qwen3LoraOptions, Qwen3OffloadOptions, Qwen3SpeculativeOptions}; #[cfg(not(target_env = "msvc"))] #[global_allocator] @@ -98,6 +98,11 @@ struct Args { #[arg(long, default_value_t = false)] no_prefix_cache: bool, + /// Enable Qwen3 DFlash speculative decoding with this drafter model path. + #[cfg(feature = "qwen3-4b")] + #[arg(long = "dflash-draft-model-path")] + dflash_draft_model_path: Option, + /// Cap on total prompt tokens forwarded in one Qwen3 scheduler step /// (chunked prefill). Prefill activation scratch scales with the step's /// prompt tokens, so this bounds peak VRAM under request bursts; prompts @@ -320,6 +325,24 @@ fn load_engine( } else { Qwen3OffloadOptions::disabled() }; + let speculative = if let Some(path) = args.dflash_draft_model_path.as_ref() { + if args.enable_lora { + bail!("--dflash-draft-model-path is not supported with --enable-lora"); + } + if args.kv_offload { + bail!("--dflash-draft-model-path is not supported with --kv-offload"); + } + if args.tp_size != 1 { + bail!("--dflash-draft-model-path currently requires --tp-size=1"); + } + info!( + "Qwen3 DFlash enabled: draft_model_path={}, prefix cache disabled for hidden-state capture", + path.display() + ); + Qwen3SpeculativeOptions::dflash(path.clone()) + } else { + Qwen3SpeculativeOptions::disabled() + }; if args.enable_lora { let lora_options = Qwen3LoraOptions { @@ -339,10 +362,11 @@ fn load_engine( args.max_prefill_tokens, ) } else { - openinfer_qwen3_4b::start_engine_with_offload( + openinfer_qwen3_4b::start_engine_with_offload_and_speculative( &args.model_path, options, offload, + speculative, args.no_prefix_cache, args.max_prefill_tokens, )