diff --git a/harness/benchmarks/QWEN38_DSPARK_ADAPTIVE_SELECTION.md b/harness/benchmarks/QWEN38_DSPARK_ADAPTIVE_SELECTION.md new file mode 100644 index 000000000..ff685211f --- /dev/null +++ b/harness/benchmarks/QWEN38_DSPARK_ADAPTIVE_SELECTION.md @@ -0,0 +1,118 @@ +# Qwen3.8 adaptive-selection prompts (historical DSpark baseline) + +This six-request workload is a prompt-selection fixture for the concurrent +adaptive speculation gate. It is deliberately balanced rather than +representative: + +- two prompts where dense DSpark is a strong win; +- two prompts where dense DSpark is only a marginal win; +- two prompts where dense autoregressive decode wins. + +The machine-readable source is +`prompts/qwen38_dspark_adaptive_selection.jsonl`. Every row records its +selection class, expected dense oracle, and the R9700 screening measurement +that justified the label. +The exact completion lengths and ordered output hashes are retained locally in +`QWEN38_DSPARK_ADAPTIVE_SELECTION_R9700.json`. + +The dense labels below were collected with DSpark and remain useful priors for +request selection. They do not prescribe the current proposal mechanism. +Concurrent runs use whichever single `Speculator` adapter the server registers; +currently that is DFlash2. The activation engine owns scoring, ranking, cost +comparison, sticky AR fallback, and telemetry. + +## Dense screening baseline + +Measured 2026-08-19 on the Radeon AI PRO R9700 (`gfx1201`, ROCm 7.2) with: + +- target: Qwen3.8-27B PR #625 requantization, IQ4_XS body, Q5_K output, + Q6_K `attn_v` and `ssm_out`; +- drafter: RadixArk Qwen3.8-27B DSpark, no YaRN, Q8_0, block width 8; +- build: Release, `gfx1201`, HIP graphs; +- dense cache/attention: Q8_0 K/V, `--fa-window 2048`; +- greedy chat requests, up to 256 output tokens; +- fresh AR process without a drafter; +- forced DSpark with `DFLASH_QWEN35_SPEC_STEP_RATIO=0`, so the dense + decoder could not hide weak speculation behind its own AR bursts. + +The throughput columns are model-side decode rates. This was a one-repeat +selection screen, not a publication measurement. + +| Prompt | Class | Dense oracle | AR tok/s | Forced spec tok/s | Spec/AR | Accept | Commit/step | +| :--- | :--- | :--- | ---: | ---: | ---: | ---: | ---: | +| `he_09 sum_product` | strong win | speculation | 34.89 | 46.35 | 1.328 | 35.8% | 2.51 | +| `he_10 rolling_max` | strong win | speculation | 34.89 | 50.30 | 1.442 | 38.9% | 2.72 | +| `he_02 separate_paren_groups` | marginal win | speculation | 34.95 | 37.69 | 1.078 | 29.0% | 2.03 | +| `he_03 truncate_number` | marginal win | speculation | 35.15 | 37.44 | 1.065 | 28.9% | 2.02 | +| `he_08 filter_by_substring` | loss | AR | 35.04 | 33.95 | 0.969 | 26.2% | 1.84 | +| `prose-01 reproducibility` | loss | AR | 34.90 | 27.10 | 0.777 | 20.8% | 1.45 | + +All six selected rows produced identical ordered content hashes under dense AR +and forced dense speculation. Prompts with mismatched hashes were rejected +from the performance fixture: HumanEval 01, 05, 06, and 07, plus prose 02, +03, and 04. Keep those rejected prompts as correctness diagnostics; do not +interpret their throughput as a valid speculative win or loss. + +## Concurrent activation benchmark + +The smallest adversarial cohort is fixed at C=3: the two dense strong-win +prompts (`sum_product` and `rolling_max`) plus the strongest dense AR win +(`reproducibility`). The wider boundary uses all six prompts at C=6. + +The removed DSpark-only matrix runner must not be used for new measurements. +Launch paired fresh AR, forced-speculation, and adaptive server processes with +the active DFlash2 drafter, and drive each process with +`harness/benchmarks/concurrency/concurrent_benchmark.py`. The checked-in JSONL +is already suitable as `--prompt-file`; select the C=3 IDs explicitly when +running the smaller cohort. `run_qwen38_dflash2_subsets.sh` remains the core +forced-mode/refill control. + +The dense oracle labels are priors, not hard assertions about the concurrent +executor. Adaptive mode evaluates each request once at its first target-decode +boundary. The active adapter returns proposal tokens, an activation score, +expected yield, and optional conditional hazards. The generic engine ranks +those scores against the profiled cost table and commits each request to AR or +speculation through retirement. Evaluation failure emits +`activation_evaluation_failed` and commits sticky AR. Accepted-token history +and user identity are not scoring inputs. + +The prefill logits produce the first sampled output token before a target-decode +step exists. A speculative request reuses its activation proposal immediately; +a request that retires directly from prefill has no target-decode mode to +activate. The useful behavior is still to preserve ordered greedy output hashes, +keep known losses in AR, and let measured concurrent costs decide the marginal +pair. + +## Profiling outputs + +`DFLASH_STEP_TIMING=1` is enabled by default. Every case retains: + +- `server.log`: startup, warmup, and measured request evidence; +- `benchmark-server.log`: only the measured window; +- `bench.json`: request and aggregate throughput; +- `feature-proof.json`: request-correlated execution/correctness proof. + +The matrix root adds: + +- `profiling.json`: complete machine-readable gate, request, shape, and phase + distributions; +- `profiling.md`: concise activation-regret and bottleneck tables; +- `summary.md`: oracle-gated aggregate summary, written only when output + stability and adaptive criteria pass. + +Read the report in this order: + +1. Require identical outputs between AR, forced speculation, and adaptive. + A mismatch is a correctness failure, not a throughput result. +2. Use paired concurrent AR/speculation measurements as the empirical oracle; + never use the dense prompt label as the final activation answer. +3. Inspect `Activation outcome against matched pure AR`. It compares each + `(live, k, path)` shape to pure AR at the same live concurrency. +4. Inspect `Initial prediction accuracy` for predicted-versus-realized + goodput and realized-versus-AR regret. +5. Use phase attribution to choose the next optimization: unexpected draft + work on k=0, verify, or the structural replay forward. + +The profiler's `total_us` uses one common decode-round origin for pure AR, +adaptive k=0, and speculative rounds. `draft_us` is a subset of that wall +time; do not add it a second time. diff --git a/harness/benchmarks/QWEN38_PR625_BASELINE.md b/harness/benchmarks/QWEN38_PR625_BASELINE.md new file mode 100644 index 000000000..207825bf4 --- /dev/null +++ b/harness/benchmarks/QWEN38_PR625_BASELINE.md @@ -0,0 +1,83 @@ +# Qwen3.8-27B PR #625 baseline + +Establish this dense, single-request baseline before measuring PR #626's +concurrent paged path. The two paths intentionally do not share cache or +attention settings. + +## Models + +Sources: + +- target: `bartowski/Qwen3.8-27B-GGUF`, `Qwen3.8-27B-IQ4_XS.gguf` +- drafter: `RadixArk/Qwen3.8-27B-DSpark`, `model.safetensors` + +Prepare the permanent local pair with: + +```bash +TARGET_SOURCE=/path/Qwen3.8-27B-IQ4_XS.gguf \ +DRAFT_SOURCE=/path/RadixArk-Qwen3.8-27B-DSpark/model.safetensors \ +LLAMA_QUANTIZE=/path/llama-quantize \ +server/scripts/prepare_qwen38_pr625_models.sh +``` + +This produces and validates: + +- target: pure IQ4_XS body, Q5_K `output.weight`, Q6_K `attn_v` and + `ssm_out` +- drafter: no YaRN, Q8_0, capture layers `4,16,28,40,52`, mask token + `248077`. PR #625 does not spell out the DSpark quantization command, but + Q8_0 reproduces its reported compute regime: this setup measured a 1.78x + speculative/plain step-time ratio on the R9700, versus the PR's ~1.8x. + The unquantized F16 drafter measured 4.20x and is therefore not the + benchmark artifact. Set `DRAFT_SCHEME=f16` or `DRAFT_SCHEME=q4-mix` only + for an explicit ablation. + +## Build on Radeon AI PRO R9700 + +Use ROCm 7.2, `gfx1201`, Release, and HIP graphs: + +```bash +cmake -S server -B server/build-pr625-r9700 -G Ninja \ + -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_HIP_COMPILER=/opt/rocm/llvm/bin/clang++ \ + -DDFLASH27B_GPU_BACKEND=hip \ + -DDFLASH27B_HIP_ARCHITECTURES=gfx1201 \ + -DDFLASH27B_HIP_SM80_EQUIV=ON \ + -DGGML_HIP_GRAPHS=ON \ + -DDFLASH27B_FA_ALL_QUANTS=ON \ + -DDFLASH27B_SERVER=ON -DDFLASH27B_TESTS=OFF +cmake --build server/build-pr625-r9700 -j +``` + +## Dense single-request launch + +Use `HIP_VISIBLE_DEVICES=0` when the R9700 is the first physical GPU. Inside +the process it remains `hip:0`. + +```bash +HIP_VISIBLE_DEVICES=0 \ +DFLASH_SINGLE_CHAIN_CHECKPOINT_F32=1 \ +DFLASH_FAST_ROLLBACK_THRESHOLD=1 \ +LUCE_Q8_MEMO=1 \ +DFLASH_KV_ROTATE=0 \ +server/build-pr625-r9700/dflash_server \ + models/.lucebox/qwen38-pr625/Qwen3.8-27B-PR625-IQ4_XS.gguf \ + --draft models/.lucebox/qwen38-pr625/Qwen3.8-27B-DSpark-RadixArk-no-yarn-q8_0.gguf \ + --target-device hip:0 --draft-device hip:0 \ + --fa-window 2048 --cache-type-k q8_0 --cache-type-v q8_0 \ + --max-ctx 8192 --prefix-cache-slots 0 --prefill-cache-slots 0 \ + --decode-mode speculation --host 127.0.0.1 --port 18140 +``` + +Do not pass `--paged-attention` or `--max-concurrency` for this baseline. For +the AR control, start a fresh process **without `--draft`** and use +`--decode-mode ar`; the dense backend otherwise has a loaded drafter and can +enter its original speculative loop. All target, cache, and attention settings +stay identical. Use greedy 300-token generations with +`harness/benchmarks/prompts/qwen38_pr625.jsonl`, and reject a measurement that +ends before the 300-token cap. + +PR #625 reported R9700 decode throughput of 34.3/34.4 tok/s for AR and +45.6/32.4 tok/s for DSpark on its code/prose prompts. Exact numeric parity +requires the original unpublished prompts; the structural check is that code +benefits while prose can remain below AR. diff --git a/harness/benchmarks/concurrency/FEATURE_MATRIX.md b/harness/benchmarks/concurrency/FEATURE_MATRIX.md new file mode 100644 index 000000000..03f7e1a2b --- /dev/null +++ b/harness/benchmarks/concurrency/FEATURE_MATRIX.md @@ -0,0 +1,147 @@ +# Qwen3.6 concurrent feature matrix + +The bounded Strix Halo measurements collected for the draft implementation are +recorded in [`STRIX_HALO_RESULTS.md`](STRIX_HALO_RESULTS.md). + +## Qwen3.8 adaptive speculation + +The DSpark-only matrix runner was removed with the concurrent DSpark execution +path. The checked-in Qwen3.8 selection prompts remain useful workload fixtures, +but proposal generation and activation scoring now come from the one active +`Speculator` adapter (currently DFlash2). + +Use `concurrent_benchmark.py` as the common request client for AR, forced +speculation, and adaptive server processes. Use +`run_qwen38_dflash2_subsets.sh` for the forced AR/speculation and refill +controls. Adaptive runs must retain the fail-closed `[spec-activation]` proof: +one scored or failed decision per measured request, opaque `score_kind`, +optional hazards, and `activation_evaluation_failed` for sticky-AR fallback. + +## Qwen3.6 DDTree/PFlash/KVFlash matrix + +`run_qwen36_feature_matrix.sh` extends the PR #596 protocol with feature +ablations for the complete Strix Halo configuration: + +- `ar`: concurrent paged autoregressive control. +- `ddtree`: adds the decode draft, DDTree, and the recorded budget. +- `pflash`: adds auto prefill compression, its drafter, and persistent draft + residency. +- `kvflash`: adds bounded KV residency in auto mode and explicitly supplies + the hashed prefill drafter for relevance-scored page selection; prefill + compression remains off in this ablation. +- `full`: enables DDTree, PFlash, and KVFlash together with both devices on + `hip:0`. +- `llama`: optional external comparison; its binary is required only when this + variant is explicitly requested. + +Run the default bounded C4 screening repeat (seven applicable fresh-server cases): + +```bash +MODEL=/opt/models/Qwen3.6-27B-Q4_K_M.gguf \ +DRAFT_MODEL=/opt/models/draft/dflash-draft-3.6-q8_0.gguf \ +PREFILL_DRAFTER=/opt/models/Qwen3-0.6B-BF16.gguf \ +REPEATS=1 \ +harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh +``` + +For the canonical AMD Strix Halo recipe, use the published Q8_0 3.6 drafter and make the tuning explicit: + +```bash +MODEL=/opt/models/Qwen3.6-27B-Q4_K_M.gguf \ +DRAFT_MODEL=/opt/models/draft/dflash-draft-3.6-q8_0.gguf \ +PREFILL_DRAFTER=/opt/models/Qwen3-0.6B-BF16.gguf \ +DRAFT_SWA=2048 PREFILL_UBATCH=512 DDTREE_ADAPTIVE=0 \ +VARIANTS=ddtree,pflash,kvflash,full CLIENTS=1,4,8,16 REPEATS=5 \ +harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh +``` + +`DDTREE_ADAPTIVE=0` matches the blog's continuous DDTree probe policy; leave it at the default `1` when measuring the concurrent engine's adaptive fallback policy. The concurrent path now records a startup `[parallel-ddtree]` marker and per-request `ddtree_steps`; these are the proof that DDTree actually ran. + +On a 128 GiB Strix Halo host, budget roughly 45–90 minutes for this smoke run; +the long-context AR controls dominate and actual time depends on the build. +Every row remains independently selectable through `VARIANTS`. For example: + +```bash +WORKLOADS=short CLIENTS=4 VARIANTS=ar,ddtree MAX_TOKENS=256 REPEATS=1 \ +harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh +``` + +Use five fresh-process, paired repeats for published measurements: + +```bash +WORKLOADS=short,compression CLIENTS=1,4,8,16 \ +VARIANTS=ar,ddtree,pflash,kvflash,full MAX_TOKENS=256 REPEATS=5 \ +harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh +``` + +That full matrix can take roughly 15–30 hours on Strix Halo; keep the generated +case directory so interrupted or suspect rows can be diagnosed rather than +quoted. + +## Activation workloads + +Auto features cannot be validated with the original 400–4,000 word prompts. +The extension adds two deterministic, disjoint 29-prompt cohorts: + +- `compression`: 34K–40K words, chosen after observing 38,130–44,856 + tokens with the development Qwen GGUF tokenizer. +- `kv-pressure`: 12K–18K words, which produced 13,463–20,190 tokens with + that tokenizer against the runner's explicitly recorded 8K pool cap. + +The observed counts are a fixture sanity check, not a claim derived from word +count and not publication evidence. Each row records the model hash; the proof +cross-checks logged raw PFlash input against `usage.prompt_tokens`, requires +it to meet the recorded auto threshold, and uses server-reported effective +tokens plus actual page traffic for KVFlash. + +The bounded default uses `short,compression` at C4. It runs PFlash and the +full configuration only on `compression`; KVFlash can also be selected on +`kv-pressure`. Inapplicable pairs are printed as skips and never appear as +successful rows. AR controls use the same prompts, so feature deltas remain +paired. + +## Fail-closed feature proof + +The server must write one JSON object per completed request with this prefix: + +```text +[concurrency-metrics] {"request_id":"...", ...} +``` + +Required fields are `effective_prompt_tokens`, `ddtree_steps`, +`ddtree_suspensions`, `ddtree_accepted_tokens`, `target_forwards`, +`kvflash_page_ins`, `kvflash_page_outs`, `kvflash_resident_blocks`, +`kvflash_reselects`, +`pflash_applied`, `pflash_input_tokens`, and `pflash_output_tokens`. + +The proof tool correlates log objects with measured SSE request IDs and also +checks the log's effective token count against +`usage.timings.effective_prompt_tokens`. A requested feature fails the case +unless: + +- DDTree has positive step and target-forward counts. Acceptance may be zero. + The required per-request suspension counter must be either zero or one. It + records adaptive fallback activation, but does not prove whether AR work ran + before or after the suspension; temporal claims require direct ordered-log + evidence. +- PFlash reports `pflash_applied=true`, a smaller output prompt, and (in auto + mode) an input token count at or above the recorded activation threshold. +- KVFlash always records an explicit hashed scorer drafter, reports its + startup-observed physical pool and enabled metadata, and has a positive + resident-block count. `kvflash`-only and `kv-pressure` rows must also show page-in or + page-out traffic. For a `full` row, traffic is required only when a + server-reported `effective_prompt_tokens` value exceeds that observed pool + token limit; zero traffic is valid when PFlash compression fits in the pool. + +After health succeeds, the runner fail-closed parses the backend's +`[parallel-kvflash] physical resident pool ...` and `[paged-attention] ...` +startup markers into `runtime_observed`. This distinguishes the actual resident +pool from both the requested auto cap and `--kv-pool-tokens`, which concurrent +KVFlash intentionally does not use to expand VRAM. + +Each case retains the exact shell-escaped command, controlled launch +environment, literal client process arguments and client-script hash, +binary/shared-library/target/draft/PFlash-and-KV-scorer hashes, the ordered +`literal_screenshot_flags` array, all feature values, raw request report, +server log, and `feature-proof.json`. The summary refuses to +include a Lucebox row whose proof is missing or invalid. diff --git a/harness/benchmarks/concurrency/README.md b/harness/benchmarks/concurrency/README.md new file mode 100644 index 000000000..ae9dfdae2 --- /dev/null +++ b/harness/benchmarks/concurrency/README.md @@ -0,0 +1,157 @@ +# Qwen3.6 concurrency benchmark + +This protocol measures the serving behavior targeted by packed continuous +prefill and concurrent decode. It is intentionally small: one streaming client, +one fresh-process runner, one deterministic prompt generator, and one summary +script. + +Run a quick screening repeat: + +```bash +MODEL=/path/to/Qwen3.6-27B-Q4_K_M.gguf \ +LUCE_SERVER_BIN=server/build-hip/dflash_server \ +LLAMA_SERVER_BIN=/path/to/llama-server \ +harness/benchmarks/concurrency/run_qwen36_concurrency.sh +``` + +Run a decode-heavy comparison with the same harness: + +```bash +MODEL=/path/to/Qwen3.6-27B-Q4_K_M.gguf \ +LUCE_SERVER_BIN=server/build-hip/dflash_server \ +LLAMA_SERVER_BIN=/path/to/llama-server \ +WORKLOADS=short MAX_TOKENS=256 VARIANTS=luce-k8,llama REPEATS=3 \ +harness/benchmarks/concurrency/run_qwen36_concurrency.sh +``` + +The short ragged prompts keep admission realistic while 256 forced output +tokens make generation dominate the measured window. Use `REPEATS=5` for +publication. Every measured case starts a fresh server and first runs a +discarded warmup at the same concurrency. The variants are: + +- `luce-k8`: packed prefill with up to eight concurrent prefills. +- `luce-k1`: the same binary/configuration with packing width limited to one. +- `llama`: llama.cpp continuous batching with fixed `-b 2048 -ub 512`. + +The 29 generated prompts are disjoint cohorts for C1/C4/C8/C16. C4 and above +contain four substantial length strata while holding the mean target length +constant. The default short, medium, and long profiles target mean lengths of +400, 1,000, and 3,000 words per request. Those generator targets are not token +counts; reports retain the exact server-observed token counts for the selected +model and tokenizer. The client refuses to wrap or reuse a prompt. + +The headline metric is aggregate output goodput: exact server-reported +completion tokens divided by level wall time. It includes queueing, prefill, +and decode and must not be called decode throughput. + +`Output-window tok/s` divides exact completion tokens by the interval from the +earliest observed first output to the final request completion. It removes the +initial all-prefill interval and is decode-facing, but it can still contain +staggered prefill while later requests await their first token. +`Request decode tok/s` is the median per-request estimate +`(completion_tokens - 1) / (end - first_output)`; it assumes the first +observed streaming event accounts for one token. Neither metric is pure kernel +decode throughput. + +`Prompt tok/s to first` is the sum of server-reported prompt tokens divided by +the latest first-token arrival; it is a useful prefill-facing metric but still +includes admission, queueing, and transport. Report TTFT median/max alongside +all throughput metrics. + +The K8-vs-K1 comparison is the causal packing ablation. The K8-vs-llama +comparison is the product comparison. Five paired repeats, the exact command +and hashes recorded in each case, zero failures, and a fixed declared output +length are required before using results in a post. The standard prefill-facing +protocol uses 64 output tokens; the decode-heavy protocol above uses 256. +Variant gains are computed as the median of same-repeat ratios, not as a ratio +of independently aggregated medians. The summarizer rejects mismatched repeat +sets. It also marks whether each variant produced the same ordered output +hashes across at least two repeats; a one-repeat screen reports stability as +`n/a`, and an unstable result is a correctness warning, not a performance win. + +## Forced DFlash2 subset/depth diagnostics + +`forced_subset_benchmark.py` is the fail-closed client for the DFlash2 +concurrency bring-up. It is intentionally limited to forced controls: every +request receives an explicit `decode_mode` of `ar` or `speculation`, and the +artifact explicitly forbids interpreting the result as adaptive activation. +The repository-owned runner generates DFlash2-specific metadata and keeps one +server process alive across the complete positional mask set: + +```bash +MODEL=/path/Qwen3.8-27B-target.gguf \ +DRAFT_MODEL=/path/Qwen3.8-27B-DFlash2-q8_0.gguf \ +PROMPT_FILE=/path/prompts.jsonl \ +CLIENTS=2 MASKS=AA,AS,SA,SS SPEC_DEPTH=4 \ +harness/benchmarks/concurrency/run_qwen38_dflash2_subsets.sh +``` + +Run a fresh process/output directory for each depth in a `{2,4,8}` screen. +The runner performs all-AR and all-SPEC warmups before measuring any mask. +It starts the backend in global `speculation` mode so the chain is allocated; +each `A` request then overrides that default to request-local AR. The host +default is `VISIBLE_DEVICES=0` (the R9700), and can be overridden explicitly. + +Launch the server with `DFLASH_SPEC_CHAIN_DEPTH=`, +`DFLASH_STEP_TIMING=1`, and, whenever any lane is speculative, +`DFLASH_DFLASH2_SELECTOR_LOG=1`. Those values must also appear in the metadata +file's `launch_environment`; the client rejects a mismatch. The server log may +contain startup and warmup output: the client hashes and parses only bytes +appended after its synchronized request cohort is ready to run. + +The JSON retains each request, its positional forced mode, prompt/payload and +exact content/reasoning hashes, every measured `[step-timing]`, +`[spec-selector]`, `[spec-activation]`, and `[concurrency-metrics]` record, plus +an exact hash and byte range for the measured server-log span. A case fails if +the requested modes did not execute, the inferred tree depth differs, output +or token accounting is incomplete, request starts exceed the configured skew, +or live concurrency `C` is not sustained for at least two consecutive decode +rounds. Use `--min-full-live-rounds` only to state a different threshold +explicitly; do not waive the requirement for at least one full-live round. + +### Refill/saturated service diagnostic + +Set `REFILL_WAVES` to at least 3 to run `refill_subset_benchmark.py` through +the same persistent-server runner: + +```bash +MODEL=/path/Qwen3.8-27B-target.gguf \ +DRAFT_MODEL=/path/Qwen3.8-27B-DFlash2-q8_0.gguf \ +PROMPT_FILE=/path/prompts.jsonl \ +CLIENTS=6 MASKS=AAAAAA,SSSSSS SPEC_DEPTH=8 REFILL_WAVES=4 \ +MAX_TOKENS=64 \ +harness/benchmarks/concurrency/run_qwen38_dflash2_subsets.sh +``` + +Each positional lane reuses its deterministic prompt and forced A/S mode, and +the first `C * (waves - 1)` successful completions submit another request on +the lane that completed. Faster lanes therefore receive more requests instead +of exhausting a fixed per-lane quota, and the active A/S mask is preserved +until the final drain. This closed-loop workload answers whether the server can +profit while work is replenished. The default `REFILL_WAVES=1` instead measures +a single closed cohort whose faster lanes create a terminal tail. Refill mode +requires at least three waves: the last C scheduled refills form a guard cohort +around the saturation proof. Run both protocols when comparing scheduler policy. + +The refill report deliberately separates two metrics: + +- `aggregate_refill_tok_s` is exact completion tokens divided by the entire + multi-wave wall interval. It includes initial prefill, client/server handoff + gaps, later prefills, and the final drain. Increasing the declared wave count + amortizes the one-time boundaries; never relabel it as a closed-cohort rate. +- `validation.full_live.engine_round_goodput_tok_s` is emitted tokens divided by + summed server `total_us` for only `live=C` timed engine rounds. It isolates + saturated round economics and includes work accounted inside `total_us`, but + excludes time between timing records such as client handoff gaps and is not + end-to-end throughput. + +The client retains every raw timing, selector, activation, and per-request +metric record and its exact measured log byte range. It also retains each +request payload/prompt/output hash. It fails closed unless all `C * waves` +requests finish the fixed token count, each forced mode and DFlash2 selector +mapping is proven by per-request telemetry, depth matches, identical repeated +lane inputs produce identical output hashes, refill handoff gaps meet the +stated bound, and a later `live=C` round occurs after at least +`C * (waves - 2)` completions while one C-request guard cohort remains before +the final drain. The artifact explicitly permits neither an adaptive activation +claim nor a closed-cohort makespan claim. diff --git a/harness/benchmarks/concurrency/STRIX_HALO_RESULTS.md b/harness/benchmarks/concurrency/STRIX_HALO_RESULTS.md new file mode 100644 index 000000000..c42eea4f9 --- /dev/null +++ b/harness/benchmarks/concurrency/STRIX_HALO_RESULTS.md @@ -0,0 +1,134 @@ +# Qwen3.6 concurrent feature results — Strix Halo + +- Date: 2026-08-13 +- Implementation: `568fbac03b498d53d6efc0b2ab5893044543a321` +- Stack base: PR #595 head `a90ffe45c1d4ad58f5f73c4107571d3cf6c51bfd` + +These are bounded engineering measurements for the draft PR, not the +five-repeat publication matrix described in `FEATURE_MATRIX.md`. The paired +AR/DDTree screen has three fresh-process repeats. The long-context activation +rows have one fresh-process repeat per concurrency level because they are much +more expensive; treat their throughput as screening data. + +## System and artifacts + +- AMD Ryzen AI MAX+ 395 with Radeon 8060S (`gfx1151`), 128 GiB unified memory. +- ROCm runtime 7.2.4. +- Release HIP build for `gfx1151` with + `DFLASH27B_HIP_SM80_EQUIV=ON`. +- Server SHA-256: + `c77b4d2c7d1505fcc751600a6603cd65e51514b685bc66c4f7d33cd64a87c8a6`. +- Target SHA-256: + `5ed60d0af4650a854b1755bd392f9aef4872643dc25a254bc68043fa638392a0`. +- Decode draft SHA-256: + `e2500e90165a0f8e7b52c9882c29ed1fa391c60b300ff11b817bf10e31fa092e`. +- PFlash/KV scorer drafter SHA-256: + `f9c9f1d3c1e21755b82d4e165f88dbbbd4355646d632fb5d6cef7c66ed4ee04e`. + +Every case started a fresh server, discarded an 8-token same-concurrency +warmup, then requested exactly 64 output tokens per request with temperature +zero, seed one, and EOS ignored. Prompts were deterministic, disjoint across +concurrency levels, and identical between paired variants. The runner rotated +variant order across repeats. + +`Output-window` counts all completion tokens from the earliest first output to +the last completion. `Goodput` counts completion tokens over the whole level, +including TTFT. Every reported row passed exact token accounting and the +request-ID-correlated feature proof. + +The retained screening artifacts contain maximum TTFT but not median TTFT. +Their max-only columns below are an explicit screening exception, not a +protocol-complete publication result; a publication rerun must report both. + +## Paired AR and adaptive DDTree + +The DDTree configuration adds the local decode draft, budget 22, and target and +draft placement on `hip:0`. Values are medians over three fresh-process +repeats. + +| C | Variant | N | Goodput tok/s | Output-window tok/s | vs AR goodput | Accepted/step | Steps/suspensions | Max TTFT s | Output hashes stable | +| ---: | :--- | ---: | ---: | ---: | ---: | ---: | :--- | ---: | :---: | +| 1 | AR | 3 | 9.41 | 12.57 | — | — | 0/0 | 1.707 | yes | +| 1 | DDTree | 3 | 9.00 | 11.85 | -4.4% | 1.00 | 1/1 | 1.715 | yes | +| 4 | AR | 3 | 20.46 | 36.21 | — | — | 0/0 | 5.508 | no | +| 4 | DDTree | 3 | 19.43 | 33.61 | n/a | 3.08 | 4/4 | 5.533 | no | +| 8 | AR | 3 | 27.44 | 65.93 | — | — | 0/0 | 11.008 | no | +| 8 | DDTree | 3 | 25.63 | 56.32 | n/a | 2.79 | 8/8 | 11.076 | no | +| 16 | AR | 3 | 31.82 | 58.79 | — | — | 0/0 | 22.487 | no | +| 16 | DDTree | 3 | 29.36 | 54.63 | n/a | 2.00 | 16/16 | 22.514 | no | + +The supplied draft had weak acceptance on this cohort. The adaptive policy +sampled one real packed-tree step, then suspended the whole cohort because its +aggregate emitted yield was below six tokens per request. At C4 and above, +the raw timings are retained only to diagnose this fallback behavior; unstable +outputs do not support a performance comparison with AR. + +At C4 and above, greedy text hashes varied across fresh repeats in both the AR +control and DDTree. C1 was byte-stable. These measurements therefore establish +exact token accounting and feature execution, but do not claim bitwise text +reproducibility for concurrent batches. + +## Full screenshot configuration + +These rows enable the complete requested product configuration: + +```text +--target-device hip:0 +--draft-device hip:0 +--ddtree +--ddtree-budget 22 +--draft-residency persistent +--prefill-compression auto +--prefill-drafter /opt/models/Qwen3-0.6B-BF16.gguf +--kvflash auto +``` + +The controlled runner sets the auto PFlash threshold to 32K tokens, the keep +ratio to 0.05, and the KVFlash resident cap to 8,192 tokens. Startup telemetry +confirmed 512 physical blocks of 16 tokens, 16 configured slots, and a 65,536 +logical-token bound per slot. + +| C | N | Goodput tok/s | Output-window tok/s | Request decode tok/s | Raw prompt range | Effective prompt range | Max TTFT s | DDTree steps/susp. | KV page in/out | PFlash requests | +| ---: | ---: | ---: | ---: | ---: | :--- | :--- | ---: | :--- | :--- | ---: | +| 1 | 1 | 2.61 | 12.39 | 12.19 | 41,504 | 2,021 | 19.324 | 1/1 | 0/0 | 1 | +| 4 | 1 | 3.00 | 31.61 | 7.95 | 38,142–44,866 | 1,870–2,235 | 77.495 | 4/4 | 1/18 | 4 | +| 8 | 1 | 3.13 | 52.60 | 6.56 | 38,141–44,870 | 1,869–2,237 | 153.782 | 8/8 | 245/792 | 8 | +| 16 | 1 | 3.19 | 18.73 | 2.34 | 38,140–44,872 | 1,867–2,238 | 307.360 | 16/16 | 293/1,880 | 16 | + +All four rows proved DDTree, PFlash, and KVFlash active. PFlash retained about +4.9% of raw prompt tokens. Output-window throughput scaled through C8, then +dropped at C16 while roughly 32K effective prompt tokens shared the 8K resident +pool; the concurrent page traffic rose accordingly. + +## Feature ablations + +| Workload | C | Variant | N | Goodput tok/s | Output-window tok/s | Request decode tok/s | Effective/raw | Max TTFT s | Activation evidence | +| :--- | ---: | :--- | ---: | ---: | ---: | ---: | ---: | ---: | :--- | +| compression | 4 | PFlash | 1 | 3.15 | 35.54 | 8.89 | 0.049 | 74.261 | 4/4 prompts compressed, 166,016 -> 8,179 tokens | +| kv-pressure | 4 | KVFlash | 1 | 1.24 | 8.56 | 5.32 | 1.000 | 197.859 | 129 resident blocks max, 0 page-ins / 3,714 page-outs | + +The PFlash-only row uses the same C4 prompts as the full row; adding DDTree and +KVFlash reduced output-window throughput from 35.54 to 31.61 tok/s in this +single screening repeat. The KVFlash-only row deliberately disables PFlash and +uses 13,474–20,203-token histories against the 8K pool. It is an activation and +pressure test, not a recommended latency configuration. + +## Reproduction + +The exact per-case command, controlled environment, startup-observed pool, +binary/shared-library/model hashes, raw request report, server log, and +`feature-proof.json` are retained by the runner. The principal invocations were: + +```bash +WORKLOADS=short CLIENTS=1,4,8,16 VARIANTS=ar,ddtree MAX_TOKENS=64 REPEATS=3 SLOTS=16 harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh + +WORKLOADS=compression CLIENTS=1,4,8,16 VARIANTS=full MAX_TOKENS=64 REPEATS=1 SLOTS=16 harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh + +WORKLOADS=compression CLIENTS=4 VARIANTS=pflash MAX_TOKENS=64 REPEATS=1 SLOTS=16 harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh + +WORKLOADS=kv-pressure CLIENTS=4 VARIANTS=kvflash MAX_TOKENS=64 REPEATS=1 SLOTS=16 harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh +``` + +Set `MODEL`, `DRAFT_MODEL`, `PREFILL_DRAFTER`, and `LUCE_SERVER_BIN` as shown +in `FEATURE_MATRIX.md`. For publication-quality claims, rerun the documented +five-repeat 256-token matrix. diff --git a/harness/benchmarks/concurrency/analyze_dflash2_selector.py b/harness/benchmarks/concurrency/analyze_dflash2_selector.py new file mode 100644 index 000000000..760a5f5d0 --- /dev/null +++ b/harness/benchmarks/concurrency/analyze_dflash2_selector.py @@ -0,0 +1,817 @@ +#!/usr/bin/env python3 +"""Offline analysis for forced DFlash2 subset/depth artifacts. + +The forced-subset client retains exact JSON records from ``[spec-selector]``, +``[concurrency-metrics]``, and ``[step-timing]`` inside each ``bench.json``. +This analyzer validates those records, joins selector engine IDs to wire +requests, and keeps two questions separate: + +* can proposal-local raw signals predict accepted yield; and +* can independent request rankings predict the best concurrent cohort? + +It deliberately does not fit or emit a runtime activation threshold. A +calibrator is only justified when held-out data contains both positive and +negative acceptance/benefit labels. +""" + +from __future__ import annotations + +import argparse +import itertools +import json +import math +import statistics +from collections import defaultdict +from pathlib import Path +from typing import Any, Iterable + + +FEATURE_DIRECTIONS = { + "chain_lm_logp": 1, + "chain_lm_probability": 1, + "mean_selected_logp": 1, + "min_selected_logp": 1, + "mean_lm_margin": 1, + "min_lm_margin": 1, + "mean_topk_mass": 1, + "min_topk_mass": 1, + "lm_top1_fraction": 1, + "mean_selector_margin": 1, + "min_selector_margin": 1, + "selector_chain_probability": 1, + "mean_selector_mass": 1, + "min_selector_mass": 1, + "mean_selector_entropy": -1, + "max_selector_entropy": -1, + "mean_rank": -1, + "max_rank": -1, +} + + +def _finite_number(value: Any, label: str) -> float: + if type(value) not in (int, float) or not math.isfinite(value): + raise ValueError(f"{label} must be a finite number") + return float(value) + + +def _non_negative_int(value: Any, label: str) -> int: + if type(value) is not int or value < 0: + raise ValueError(f"{label} must be a non-negative integer") + return value + + +def _record(wrapper: Any, label: str) -> dict[str, Any]: + if not isinstance(wrapper, dict) or not isinstance(wrapper.get("record"), dict): + raise ValueError(f"{label} must wrap a JSON object in record") + record = wrapper["record"] + raw = wrapper.get("raw_json") + if not isinstance(raw, str): + raise ValueError(f"{label} must retain raw_json") + try: + parsed = json.loads(raw) + except json.JSONDecodeError as exc: + raise ValueError(f"{label} contains invalid raw_json: {exc}") from exc + if parsed != record: + raise ValueError(f"{label} raw_json and parsed record disagree") + return record + + +def parse_profile_lines(data: bytes) -> dict[str, list[dict[str, Any]]]: + """Parse the three profiling records used by this analyzer. + + This helper is intentionally small; the benchmark client remains the + authority for log-span capture. It is useful for parser tests and for + diagnosing an artifact whose retained record is malformed. + """ + prefixes = { + "selectors": b"[spec-selector] ", + "requests": b"[concurrency-metrics] ", + "rounds": b"[step-timing] ", + } + out = {key: [] for key in prefixes} + for line_index, line in enumerate(data.splitlines(), 1): + for key, prefix in prefixes.items(): + position = line.find(prefix) + if position < 0: + continue + raw = line[position + len(prefix):].decode("utf-8", errors="strict") + try: + record = json.loads(raw) + except json.JSONDecodeError as exc: + raise ValueError( + f"line {line_index}: invalid {prefix.decode().strip()} JSON: {exc}" + ) from exc + if not isinstance(record, dict): + raise ValueError(f"line {line_index}: profiling JSON must be an object") + out[key].append({ + "line_index": line_index, + "raw_json": raw, + "record": record, + }) + break + return out + + +def first_block_features(selector: dict[str, Any]) -> dict[str, float]: + """Collapse pre-verification raw signals from one proposal block.""" + depths = selector.get("depths") + if not isinstance(depths, list) or not depths: + raise ValueError("selector depths must be a non-empty array") + numeric: dict[str, list[float]] = defaultdict(list) + top1: list[float] = [] + expected_depth = 1 + for row in depths: + if not isinstance(row, dict): + raise ValueError("selector depth entry must be an object") + if row.get("depth") != expected_depth: + raise ValueError("selector depths must be contiguous and one-based") + expected_depth += 1 + for key in ( + "selected_logp", "lm_margin", "topk_mass", "selector_margin", + "selector_mass", "selector_entropy", "rank", + ): + numeric[key].append(_finite_number(row.get(key), f"selector {key}")) + if type(row.get("lm_top1")) is not bool: + raise ValueError("selector lm_top1 must be boolean") + top1.append(float(row["lm_top1"])) + + logp = numeric["selected_logp"] + selector_mass = numeric["selector_mass"] + chain_lm_logp = sum(logp) + selector_chain_logp = sum(math.log(max(value, 1e-300)) for value in selector_mass) + return { + "chain_lm_logp": chain_lm_logp, + "chain_lm_probability": math.exp(chain_lm_logp), + "mean_selected_logp": statistics.fmean(logp), + "min_selected_logp": min(logp), + "mean_lm_margin": statistics.fmean(numeric["lm_margin"]), + "min_lm_margin": min(numeric["lm_margin"]), + "mean_topk_mass": statistics.fmean(numeric["topk_mass"]), + "min_topk_mass": min(numeric["topk_mass"]), + "lm_top1_fraction": statistics.fmean(top1), + "mean_selector_margin": statistics.fmean(numeric["selector_margin"]), + "min_selector_margin": min(numeric["selector_margin"]), + "selector_chain_probability": math.exp(selector_chain_logp), + "mean_selector_mass": statistics.fmean(selector_mass), + "min_selector_mass": min(selector_mass), + "mean_selector_entropy": statistics.fmean(numeric["selector_entropy"]), + "max_selector_entropy": max(numeric["selector_entropy"]), + "mean_rank": statistics.fmean(numeric["rank"]), + "max_rank": max(numeric["rank"]), + } + + +def _validate_selector(selector: dict[str, Any], label: str) -> None: + _non_negative_int(selector.get("request_id"), f"{label} request_id") + _non_negative_int(selector.get("slot"), f"{label} slot") + _non_negative_int(selector.get("generated"), f"{label} generated") + accepted_depth = _non_negative_int( + selector.get("accepted_depth"), f"{label} accepted_depth", + ) + features = first_block_features(selector) + del features + depths = selector["depths"] + if accepted_depth > len(depths): + raise ValueError(f"{label} accepted_depth exceeds proposal depth") + for index, row in enumerate(depths, 1): + if type(row.get("accepted")) is not bool: + raise ValueError(f"{label} depth {index} accepted must be boolean") + if row["accepted"] != (index <= accepted_depth): + raise ValueError(f"{label} accepted flags are not a prefix") + + +def _round_summary(rows: list[dict[str, Any]]) -> dict[str, Any]: + total_us = sum(float(row["total_us"]) for row in rows) + emitted = sum(int(row["emitted_tokens"]) for row in rows) + accepted = sum(int(row["accepted_tokens"]) for row in rows) + lane_steps = sum(int(row["k"]) for row in rows) + live_histogram: dict[int, int] = defaultdict(int) + for row in rows: + live_histogram[int(row["live"])] += 1 + return { + "rounds": len(rows), + "total_us": total_us, + "emitted_tokens": emitted, + "accepted_tokens": accepted, + "spec_lane_steps": lane_steps, + "goodput_tok_s": emitted * 1e6 / total_us if total_us else None, + "accepted_per_spec_lane_step": ( + accepted / lane_steps if lane_steps else None + ), + "live_histogram": dict(sorted(live_histogram.items())), + } + + +def analyze_artifact(path: Path) -> dict[str, Any]: + """Validate and join one forced-subset ``bench.json`` artifact.""" + bench = json.loads(path.read_text(encoding="utf-8")) + if bench.get("kind") != "dflash2-forced-subset-diagnostic": + raise ValueError(f"{path}: not a DFlash2 forced-subset artifact") + validation = bench.get("validation") + if not isinstance(validation, dict) or validation.get("passed") is not True: + raise ValueError(f"{path}: benchmark evidence did not pass validation") + level = bench.get("level") + if not isinstance(level, dict): + raise ValueError(f"{path}: level must be an object") + clients = _non_negative_int(level.get("clients"), f"{path} clients") + if clients < 1: + raise ValueError(f"{path}: clients must be positive") + spec_depth = _non_negative_int(bench.get("spec_depth"), f"{path} spec_depth") + if spec_depth < 2: + raise ValueError(f"{path}: spec_depth must be at least two") + mask = level.get("request_mode_mask") + if not isinstance(mask, str) or len(mask) != clients or set(mask) - {"A", "S"}: + raise ValueError(f"{path}: invalid request mode mask") + + details = level.get("requests_detail") + if not isinstance(details, list) or len(details) != clients: + raise ValueError(f"{path}: requests_detail count does not match clients") + detail_by_wire: dict[str, dict[str, Any]] = {} + for detail in details: + if not isinstance(detail, dict) or not isinstance(detail.get("request_id"), str): + raise ValueError(f"{path}: request detail lacks request_id") + wire = detail["request_id"] + if wire in detail_by_wire: + raise ValueError(f"{path}: duplicate wire request {wire}") + detail_by_wire[wire] = detail + + retained = bench.get("server_records") + if not isinstance(retained, dict): + raise ValueError(f"{path}: server_records must be an object") + timings = [] + for index, wrapper in enumerate(retained.get("rounds") or []): + row = _record(wrapper, f"{path} timing record {index}") + for key in ("live", "k", "accepted_tokens", "emitted_tokens"): + _non_negative_int(row.get(key), f"{path} timing {key}") + if row.get("path") not in ("ar", "spec", "spec-direct"): + raise ValueError( + f"{path}: timing path must be ar, spec, or spec-direct" + ) + total_us = _finite_number(row.get("total_us"), f"{path} timing total_us") + if total_us <= 0.0: + raise ValueError(f"{path}: timing total_us must be positive") + if int(row["live"]) < 1 or int(row["live"]) > clients: + raise ValueError(f"{path}: timing live exceeds request cohort") + timings.append(row) + + full_live = [row for row in timings if int(row["live"]) == clients] + tail = [row for row in timings if int(row["live"]) < clients] + round_timing = { + "all": _round_summary(timings), + "full_live": _round_summary(full_live), + "tail": _round_summary(tail), + } + + metric_by_engine: dict[int, dict[str, Any]] = {} + metric_by_wire: dict[str, dict[str, Any]] = {} + for index, wrapper in enumerate(retained.get("requests") or []): + metric = _record(wrapper, f"{path} request record {index}") + wire = metric.get("request_id") + engine = metric.get("engine_request_id") + if not isinstance(wire, str) or wire not in detail_by_wire: + raise ValueError(f"{path}: metric references unknown wire request {wire!r}") + engine = _non_negative_int(engine, f"{path} engine_request_id") + if engine in metric_by_engine or wire in metric_by_wire: + raise ValueError(f"{path}: duplicate request metric mapping") + metric_by_engine[engine] = metric + metric_by_wire[wire] = metric + if set(metric_by_wire) != set(detail_by_wire): + missing = sorted(set(detail_by_wire) - set(metric_by_wire)) + raise ValueError(f"{path}: missing concurrency metrics for {missing}") + + selectors_by_engine: dict[int, list[dict[str, Any]]] = defaultdict(list) + for index, wrapper in enumerate(retained.get("selectors") or []): + selector = _record(wrapper, f"{path} selector record {index}") + _validate_selector(selector, f"{path} selector record {index}") + engine = int(selector["request_id"]) + if engine not in metric_by_engine: + raise ValueError(f"{path}: selector references unknown engine request {engine}") + if len(selector["depths"]) != spec_depth - 1: + raise ValueError(f"{path}: selector proposal depth disagrees with spec_depth") + selectors_by_engine[engine].append(selector) + + requests = [] + for position, detail in enumerate(details): + wire = detail["request_id"] + metric = metric_by_wire[wire] + engine = int(metric["engine_request_id"]) + spec_steps = _non_negative_int(metric.get("spec_steps"), f"{path} spec_steps") + accepted = _non_negative_int( + metric.get("spec_accepted_tokens"), f"{path} spec_accepted_tokens", + ) + selectors = sorted( + selectors_by_engine.get(engine, []), key=lambda row: row["generated"], + ) + expected_mode = "speculation" if mask[position] == "S" else "ar" + if detail.get("decode_mode") != expected_mode: + raise ValueError(f"{path}: request mode mask and detail disagree") + if expected_mode == "ar": + if spec_steps or accepted or selectors: + raise ValueError(f"{path}: forced AR request contains speculation") + first = None + lifetime_yield = None + yield_fraction = None + else: + if spec_steps < 1 or len(selectors) != spec_steps: + raise ValueError(f"{path}: selector count does not match spec_steps") + if sum(int(row["accepted_depth"]) for row in selectors) != accepted: + raise ValueError(f"{path}: selector acceptance does not match request metric") + generated = [int(row["generated"]) for row in selectors] + if len(set(generated)) != len(generated) or generated[0] != 0: + raise ValueError(f"{path}: selector sequence lacks a unique first block") + first = first_block_features(selectors[0]) + lifetime_yield = accepted / spec_steps + yield_fraction = lifetime_yield / (spec_depth - 1) + requests.append({ + "position": position, + "wire_request_id": wire, + "engine_request_id": engine, + "mode": expected_mode, + "prompt_index": detail.get("prompt_index"), + "prompt_sha256": detail.get("prompt_sha256"), + "request_decode_tok_s": detail.get("request_decode_tok_s"), + "spec_steps": spec_steps, + "accepted_tokens": accepted, + "lifetime_accepted_yield": lifetime_yield, + "lifetime_yield_fraction": yield_fraction, + "first_accepted_depth": selectors[0]["accepted_depth"] if selectors else None, + "first_features": first, + "selectors": selectors, + }) + + return { + "path": str(path), + "clients": clients, + "spec_depth": spec_depth, + "mask": mask, + "repeat": (bench.get("server_metadata") or {}).get("repeat"), + "prompt_set_sha256": level.get("selected_prompt_set_sha256"), + "aggregate_tok_s": _finite_number( + level.get("aggregate_tok_s"), f"{path} aggregate_tok_s", + ), + "wall_s": _finite_number(level.get("wall_s"), f"{path} wall_s"), + "round_timing": round_timing, + "requests": requests, + } + + +def _average_ranks(values: list[float]) -> list[float]: + ordered = sorted(range(len(values)), key=values.__getitem__) + ranks = [0.0] * len(values) + offset = 0 + while offset < len(ordered): + end = offset + 1 + while end < len(ordered) and values[ordered[end]] == values[ordered[offset]]: + end += 1 + rank = (offset + end - 1) / 2.0 + 1.0 + for index in ordered[offset:end]: + ranks[index] = rank + offset = end + return ranks + + +def spearman_correlation(x: list[float], y: list[float]) -> float | None: + if len(x) != len(y) or len(x) < 2: + return None + rx = _average_ranks(x) + ry = _average_ranks(y) + mx = statistics.fmean(rx) + my = statistics.fmean(ry) + numerator = sum((a - mx) * (b - my) for a, b in zip(rx, ry)) + dx = sum((a - mx) ** 2 for a in rx) + dy = sum((b - my) ** 2 for b in ry) + if dx == 0.0 or dy == 0.0: + return None + return numerator / math.sqrt(dx * dy) + + +def _correlations( + rows: list[dict[str, Any]], outcome: str, +) -> dict[str, Any]: + eligible = [row for row in rows if row.get(outcome) is not None] + grouped: dict[str, list[dict[str, Any]]] = defaultdict(list) + for index, row in enumerate(eligible): + prompt = row.get("prompt_sha256") + grouped[str(prompt) if prompt is not None else f"row-{index}"].append(row) + + independent = [] + for prompt, prompt_rows in sorted(grouped.items()): + independent.append({ + "prompt_sha256": prompt, + outcome: statistics.fmean(float(row[outcome]) for row in prompt_rows), + "first_features": { + feature: statistics.fmean( + float(row["first_features"][feature]) for row in prompt_rows + ) + for feature in FEATURE_DIRECTIONS + }, + }) + + def correlations(sample: list[dict[str, Any]]) -> dict[str, float | None]: + y = [float(row[outcome]) for row in sample] + values = {} + for feature in FEATURE_DIRECTIONS: + x = [float(row["first_features"][feature]) for row in sample] + raw = spearman_correlation(x, y) + values[feature] = ( + raw * FEATURE_DIRECTIONS[feature] if raw is not None else None + ) + return values + + independent_y = [float(row[outcome]) for row in independent] + full_yield = None + non_full_yield = None + if outcome == "lifetime_yield_fraction" and independent_y: + maximum = max(independent_y) + full_yield = sum(abs(value - maximum) <= 1e-12 for value in independent_y) + non_full_yield = len(independent_y) - full_yield + label_support = full_yield >= 3 and non_full_yield >= 3 + else: + label_support = len(set(independent_y)) >= 3 + identifiable = len(independent) >= 10 and label_support + observation_y = [float(row[outcome]) for row in eligible] + return { + "observations": len(eligible), + "unique_prompts": len(independent), + "outcome_min": min(independent_y) if independent_y else None, + "outcome_max": max(independent_y) if independent_y else None, + "identifiable": identifiable, + "minimum_recommended_unique_prompts": 10, + "full_yield_prompts": full_yield, + "non_full_yield_prompts": non_full_yield, + "observation_outcome_min": min(observation_y) if observation_y else None, + "observation_outcome_max": max(observation_y) if observation_y else None, + "observation_weighted_spearman": correlations(eligible), + "spearman_higher_is_better": correlations(independent), + } + + +def _calibration( + selectors: Iterable[tuple[int, int, str, dict[str, Any]]], +) -> list[dict[str, Any]]: + grouped: dict[ + tuple[int, int, int], + list[tuple[bool, float, float, float, float, str]], + ] = defaultdict(list) + for clients, spec_depth, prompt, selector in selectors: + lm_chain = 1.0 + selector_chain = 1.0 + for depth in selector["depths"]: + lm_chain *= math.exp(float(depth["selected_logp"])) + selector_chain *= float(depth["selector_mass"]) + grouped[(clients, spec_depth, int(depth["depth"]))].append(( + bool(depth["accepted"]), + math.exp(float(depth["selected_logp"])), + float(depth["selector_mass"]), + lm_chain, + selector_chain, + prompt, + )) + out = [] + for (clients, spec_depth, depth), rows in sorted(grouped.items()): + labels = [float(row[0]) for row in rows] + observed = statistics.fmean(labels) + by_prompt: dict[str, list[bool]] = defaultdict(list) + for row in rows: + by_prompt[row[5]].append(bool(row[0])) + + def metric(index: int) -> dict[str, float]: + predicted = [row[index] for row in rows] + mean = statistics.fmean(predicted) + return { + "mean_raw_probability": mean, + "observed_minus_raw": observed - mean, + "brier": statistics.fmean( + (prediction - label) ** 2 + for prediction, label in zip(predicted, labels) + ), + } + + rejected_prompts = sum( + any(not label for label in prompt_labels) + for prompt_labels in by_prompt.values() + ) + out.append({ + "clients": clients, + "spec_depth": spec_depth, + "proposal_depth": depth, + "observations": len(rows), + "unique_prompts": len(by_prompt), + "prompts_with_rejection": rejected_prompts, + "prompts_always_accepted": len(by_prompt) - rejected_prompts, + "accepted": sum(int(value) for value in labels), + "observed_survival": observed, + "selected_token_probability": metric(1), + "selector_mass": metric(2), + "lm_chain_probability": metric(3), + "selector_chain_probability": metric(4), + "has_both_labels": len(set(labels)) > 1, + }) + return out + + +def _mask_for_positions(clients: int, positions: Iterable[int]) -> str: + chosen = set(positions) + return "".join("S" if index in chosen else "A" for index in range(clients)) + + +def _shapley_values(values: dict[str, float], clients: int) -> list[float]: + denominator = math.factorial(clients) + out = [] + for player in range(clients): + contribution = 0.0 + others = [index for index in range(clients) if index != player] + for size in range(clients): + weight = ( + math.factorial(size) * math.factorial(clients - size - 1) + / denominator + ) + for subset in itertools.combinations(others, size): + without = _mask_for_positions(clients, subset) + with_player = _mask_for_positions(clients, (*subset, player)) + contribution += weight * (values[with_player] - values[without]) + out.append(contribution) + return out + + +def _relative_regret(oracle: float, value: float | None) -> float | None: + if value is None or oracle <= 0: + return None + return max(0.0, (oracle - value) / oracle) + + +def _subset_group(cases: list[dict[str, Any]]) -> dict[str, Any]: + clients = int(cases[0]["clients"]) + depth = int(cases[0]["spec_depth"]) + by_mask: dict[str, list[float]] = defaultdict(list) + case_by_mask: dict[str, list[dict[str, Any]]] = defaultdict(list) + for case in cases: + by_mask[case["mask"]].append(float(case["aggregate_tok_s"])) + case_by_mask[case["mask"]].append(case) + values = {mask: statistics.median(rows) for mask, rows in by_mask.items()} + walls = { + mask: statistics.median(float(case["wall_s"]) for case in rows) + for mask, rows in case_by_mask.items() + } + + def timing_values(scope: str) -> dict[str, float]: + out = {} + for mask, rows in case_by_mask.items(): + samples = [ + case["round_timing"][scope]["goodput_tok_s"] for case in rows + if case["round_timing"][scope]["goodput_tok_s"] is not None + ] + if samples: + out[mask] = statistics.median(float(value) for value in samples) + return dict(sorted(out.items())) + + expected = { + "".join(bits) for bits in itertools.product("AS", repeat=clients) + } + missing = sorted(expected - set(values)) + result: dict[str, Any] = { + "clients": clients, + "spec_depth": depth, + "prompt_set_sha256": cases[0].get("prompt_set_sha256"), + "mask_goodput_tok_s": dict(sorted(values.items())), + "mask_makespan_s": dict(sorted(walls.items())), + "mask_all_round_goodput_tok_s": timing_values("all"), + "mask_full_live_goodput_tok_s": timing_values("full_live"), + "mask_tail_goodput_tok_s": timing_values("tail"), + "complete_exhaustive": not missing, + "missing_masks": missing, + } + if missing: + return result + + oracle_mask = max(values, key=values.get) + oracle = values[oracle_mask] + all_ar_mask = "A" * clients + all_spec_mask = "S" * clients + mixed = {mask: value for mask, value in values.items() if "A" in mask and "S" in mask} + best_mixed_mask = max(mixed, key=mixed.get) if mixed else None + homogeneous_mask = max((all_ar_mask, all_spec_mask), key=values.get) + homogeneous = values[homogeneous_mask] + result.update({ + "oracle_mask": oracle_mask, + "oracle_goodput_tok_s": oracle, + "all_ar_ratio": values[all_ar_mask] / oracle, + "all_spec_ratio": values[all_spec_mask] / oracle, + "best_homogeneous_mask": homogeneous_mask, + "best_homogeneous_regret": _relative_regret(oracle, homogeneous), + "best_mixed_mask": best_mixed_mask, + "best_mixed_regret": _relative_regret( + oracle, values.get(best_mixed_mask) if best_mixed_mask else None, + ), + "oracle_is_homogeneous": oracle_mask in (all_ar_mask, all_spec_mask), + "homogeneous_dominates_every_mixed": ( + bool(mixed) and homogeneous > max(mixed.values()) + ), + }) + + shapley = _shapley_values(values, clients) + marginals = [] + sign_flips = 0 + for player in range(clients): + deltas = [] + others = [index for index in range(clients) if index != player] + for size in range(clients): + for subset in itertools.combinations(others, size): + without = _mask_for_positions(clients, subset) + with_player = _mask_for_positions(clients, (*subset, player)) + deltas.append(values[with_player] - values[without]) + flip = min(deltas) < 0.0 < max(deltas) + sign_flips += int(flip) + marginals.append({ + "position": player, + "shapley_goodput_tok_s": shapley[player], + "marginal_min_tok_s": min(deltas), + "marginal_max_tok_s": max(deltas), + "marginal_sign_flips_with_peer_modes": flip, + }) + result["per_request_marginals"] = marginals + result["requests_with_contextual_sign_flip"] = sign_flips + + source_cases = case_by_mask[all_spec_mask] + source = source_cases[0] + by_position = {row["position"]: row for row in source["requests"]} + prefix_results = {} + benefit_rows = [] + for position in range(clients): + request = by_position[position] + benefit_rows.append({ + **request, + "shapley_goodput_tok_s": shapley[position], + }) + for feature, direction in FEATURE_DIRECTIONS.items(): + ranking = sorted( + range(clients), + key=lambda position: direction * float( + by_position[position]["first_features"][feature] + ), + reverse=True, + ) + by_k = [] + for count in range(clients + 1): + prefix_mask = _mask_for_positions(clients, ranking[:count]) + candidates = { + mask: value for mask, value in values.items() if mask.count("S") == count + } + best_mask = max(candidates, key=candidates.get) + by_k.append({ + "spec_requests": count, + "ranked_prefix_mask": prefix_mask, + "best_same_size_mask": best_mask, + "same_size_regret": _relative_regret( + candidates[best_mask], values[prefix_mask], + ), + }) + prefix_results[feature] = { + "ranked_positions": ranking, + "by_subset_size": by_k, + } + result["raw_feature_prefix_rankings"] = prefix_results + result["first_feature_vs_shapley"] = _correlations( + benefit_rows, "shapley_goodput_tok_s", + ) + return result + + +def analyze(paths: Iterable[Path]) -> dict[str, Any]: + bench_paths: set[Path] = set() + for path in paths: + if path.is_file(): + bench_paths.add(path) + elif path.is_dir(): + bench_paths.update(path.rglob("bench.json")) + else: + raise ValueError(f"input does not exist: {path}") + cases = [] + skipped = [] + for path in sorted(bench_paths): + try: + header = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise ValueError(f"cannot read {path}: {exc}") from exc + if header.get("kind") != "dflash2-forced-subset-diagnostic": + skipped.append(str(path)) + continue + # Warmups have the same artifact kind but live below a warmup directory. + if "warmup" in path.parts: + skipped.append(str(path)) + continue + cases.append(analyze_artifact(path)) + if not cases: + raise ValueError("no measured DFlash2 forced-subset artifacts found") + + request_rows = [ + {**row, "clients": case["clients"], "spec_depth": case["spec_depth"], + "mask": case["mask"], "artifact": case["path"]} + for case in cases for row in case["requests"] if row["mode"] == "speculation" + ] + all_selectors = [ + (int(row["clients"]), int(row["spec_depth"]), + str(row.get("prompt_sha256")), selector) + for row in request_rows for selector in row["selectors"] + ] + first_selectors = [ + (int(row["clients"]), int(row["spec_depth"]), + str(row.get("prompt_sha256")), row["selectors"][0]) + for row in request_rows + ] + + grouped_cases: dict[tuple[int, int, Any], list[dict[str, Any]]] = defaultdict(list) + for case in cases: + grouped_cases[( + int(case["clients"]), int(case["spec_depth"]), + case.get("prompt_set_sha256"), + )].append(case) + subset_groups = [ + _subset_group(rows) for _key, rows in sorted(grouped_cases.items()) + ] + + correlation_groups = [] + grouped_requests: dict[tuple[int, int], list[dict[str, Any]]] = defaultdict(list) + for row in request_rows: + grouped_requests[(int(row["clients"]), int(row["spec_depth"]))].append(row) + for (clients, depth), rows in sorted(grouped_requests.items()): + correlation_groups.append({ + "clients": clients, + "spec_depth": depth, + **_correlations(rows, "lifetime_yield_fraction"), + }) + + complete = [row for row in subset_groups if row["complete_exhaustive"]] + sign_flips = sum(row.get("requests_with_contextual_sign_flip", 0) for row in complete) + homogeneous_wins = sum( + int(row.get("homogeneous_dominates_every_mixed", False)) for row in complete + ) + calibration_all = _calibration(all_selectors) + calibration_first = _calibration(first_selectors) + has_negative_acceptance = any( + row["observed_survival"] < 1.0 for row in calibration_all + ) + calibration_label_support = any( + row["prompts_with_rejection"] >= 3 + and row["prompts_always_accepted"] >= 3 + for row in calibration_all + ) + benefit_identifiable = any( + bool(row.get("first_feature_vs_shapley", {}).get("identifiable")) + for row in complete + ) + yield_identifiable = any(row["identifiable"] for row in correlation_groups) + return { + "schema_version": 1, + "artifact_count": len(cases), + "skipped_artifacts": skipped, + "spec_request_observations": len(request_rows), + "unique_spec_prompts": len({row.get("prompt_sha256") for row in request_rows}), + "acceptance_calibration": { + "all_blocks": calibration_all, + "first_blocks": calibration_first, + "contains_negative_acceptance_labels": has_negative_acceptance, + "independent_label_support": calibration_label_support, + }, + "first_feature_vs_lifetime_yield": correlation_groups, + "subset_oracle": subset_groups, + "evidence_assessment": { + "yield_rank_identifiable": yield_identifiable, + "benefit_rank_identifiable": benefit_identifiable, + "complete_exhaustive_groups": len(complete), + "groups_where_homogeneous_beats_every_mixed_subset": homogeneous_wins, + "request_positions_with_context_dependent_marginal_sign": sign_flips, + "runtime_calibrator_ready": ( + calibration_label_support and yield_identifiable + and benefit_identifiable + ), + "interpretation": ( + "Raw selector values are proposal features, not confidence. " + "Fit only on held-out request-level benefit labels, stratified " + "by executed concurrency/depth shape; keep cohort utility in the " + "activation objective." + ), + }, + } + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("inputs", nargs="+", type=Path) + parser.add_argument("--out", type=Path) + parser.add_argument("--compact", action="store_true") + args = parser.parse_args(argv) + report = analyze(args.inputs) + text = json.dumps( + report, sort_keys=True, indent=None if args.compact else 2, + separators=(",", ":") if args.compact else None, + ) + "\n" + if args.out is None: + print(text, end="") + else: + args.out.write_text(text, encoding="utf-8") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/harness/benchmarks/concurrency/analyze_gate_decisions.py b/harness/benchmarks/concurrency/analyze_gate_decisions.py new file mode 100644 index 000000000..bef65f9d9 --- /dev/null +++ b/harness/benchmarks/concurrency/analyze_gate_decisions.py @@ -0,0 +1,1195 @@ +#!/usr/bin/env python3 +"""Analyze adaptive gate decisions and per-phase decode timing. + +Inputs may be individual matrix case directories or a completed matrix root. +The analyzer joins benchmark request IDs to engine IDs, preserves prompt +selection labels, compares paired AR/speculation controls, and summarizes +machine-readable [step-timing] records. When benchmark-server.log exists it is +preferred so warmup rounds cannot contaminate the measured distributions. +""" + +from __future__ import annotations + +import argparse +import json +import math +import re +import statistics +from collections import defaultdict +from pathlib import Path +from typing import Any + + +NUMBER = r"(?:[0-9]+(?:\.[0-9]+)?|\.[0-9]+)(?:[eE][+-]?[0-9]+)?" +GATE_RE = re.compile( + rf"\[spec-gate\] C=(?P\d+) k=(?P\d+) " + rf"scores=\[(?P[^\]]*)\].*?" + rf"G\(k\)=(?P{NUMBER}) G\(0\)=(?P{NUMBER}).*?" + rf"predicted_cost=(?P{NUMBER})us " + rf"measured=(?:(?P{NUMBER})us|ar-path)" +) +SCORE_RE = re.compile( + rf"(?P\d+):(?P{NUMBER}|nan)/" + r"(?P[a-z_-]+)" + r"(?:/(?P[a-z0-9_-]+))?(?P\*?)" +) +METRIC_RE = re.compile(r"\[concurrency-metrics\] (?P\{.*\})") +TIMING_RE = re.compile(r"\[step-timing\] (?P\{.*\})") +ACTIVATION_RE = re.compile(r"\[spec-activation\]\s+(?P.*)$") +TIMING_COUNT_FIELDS = ( + "live", "k", "emitted_tokens", "accepted_tokens", "target_forwards", +) +TYPED_REASON_RE = re.compile(r"^[a-z][a-z0-9]*(?:_[a-z0-9]+)*$") + + +def _json_object(match: re.Match[str], path: Path, line_no: int) -> dict[str, Any]: + try: + value = json.loads(match.group("json")) + except json.JSONDecodeError as exc: + raise ValueError(f"{path}:{line_no}: invalid profiling JSON: {exc}") from exc + if not isinstance(value, dict): + raise ValueError(f"{path}:{line_no}: profiling record must be an object") + return value + + +def _validate_timing(row: dict[str, Any], path: Path, line_no: int) -> None: + if row.get("path") not in ("ar", "spec", "spec-direct"): + raise ValueError(f"{path}:{line_no}: invalid step-timing path") + for key in TIMING_COUNT_FIELDS: + value = row.get(key) + if type(value) is not int or value < 0: + raise ValueError( + f"{path}:{line_no}: step-timing {key} must be non-negative int" + ) + total = row.get("total_us") + if type(total) not in (int, float) or not math.isfinite(total) or total <= 0: + raise ValueError(f"{path}:{line_no}: step-timing total_us must be positive") + for key, value in row.items(): + if key.endswith("_us") and ( + type(value) not in (int, float) + or not math.isfinite(value) + or value < 0 + ): + raise ValueError( + f"{path}:{line_no}: step-timing {key} must be non-negative" + ) + + +def _validate_activation(row: dict[str, Any], path: Path, line_no: int) -> None: + required = ( + "request_id", "slot", "activation_score", "score_kind", + "expected_yield", "evaluation", "fallback_reason", + "decision_reason", "decision", + ) + for key in required: + if key not in row: + raise ValueError( + f"{path}:{line_no}: spec-activation {key} is required" + ) + + for key in ("request_id", "slot"): + value = row[key] + if type(value) is not int or value < 0: + raise ValueError( + f"{path}:{line_no}: spec-activation {key} must be a " + "non-negative int" + ) + if row["decision"] not in ("ar", "speculation"): + raise ValueError( + f"{path}:{line_no}: spec-activation decision must be ar or " + "speculation" + ) + if row["evaluation"] not in ("scored", "failed"): + raise ValueError( + f"{path}:{line_no}: spec-activation evaluation must be scored " + "or failed" + ) + + kind = row["score_kind"] + if not isinstance(kind, str) or not kind: + raise ValueError( + f"{path}:{line_no}: spec-activation score_kind must be a " + "nonempty string" + ) + reason = row["decision_reason"] + if not isinstance(reason, str) or not TYPED_REASON_RE.fullmatch(reason): + raise ValueError( + f"{path}:{line_no}: spec-activation decision_reason must be a " + "nonempty snake-case reason" + ) + + hazards = row.get("hazards") + if hazards is not None: + if not isinstance(hazards, list) or any( + type(value) not in (int, float) + or not math.isfinite(value) + or not 0.0 <= value <= 1.0 + for value in hazards + ): + raise ValueError( + f"{path}:{line_no}: spec-activation hazards must be an " + "array of finite probabilities" + ) + + score_fields = ("activation_score", "expected_yield") + if row["evaluation"] == "scored": + if kind == "unspecified": + raise ValueError( + f"{path}:{line_no}: scored spec-activation score_kind must " + "identify its scoring model" + ) + for key in score_fields: + value = row[key] + if ( + type(value) not in (int, float) + or not math.isfinite(value) + or value < 1.0 + ): + raise ValueError( + f"{path}:{line_no}: spec-activation {key} must be finite " + "and at least 1 for a scored evaluation" + ) + if row["fallback_reason"] is not None: + raise ValueError( + f"{path}:{line_no}: scored spec-activation fallback_reason " + "must be null" + ) + return + + if any(row[key] is not None for key in score_fields) or hazards is not None: + raise ValueError( + f"{path}:{line_no}: failed spec-activation scores must be null" + ) + if row["decision"] != "ar": + raise ValueError( + f"{path}:{line_no}: failed spec-activation decision must be ar" + ) + if row["fallback_reason"] != "activation_evaluation_failed": + raise ValueError( + f"{path}:{line_no}: failed spec-activation fallback_reason must " + "be activation_evaluation_failed" + ) + + +def parse_server_log( + path: Path, +) -> tuple[ + list[dict[str, Any]], + dict[str, dict[str, Any]], + list[dict[str, Any]], + list[dict[str, Any]], +]: + rounds: list[dict[str, Any]] = [] + metrics: dict[str, dict[str, Any]] = {} + timings: list[dict[str, Any]] = [] + activations: list[dict[str, Any]] = [] + for line_no, line in enumerate( + path.read_text(encoding="utf-8", errors="replace").splitlines(), 1, + ): + gate = GATE_RE.search(line) + if gate: + entries = [] + for score in SCORE_RE.finditer(gate.group("scores")): + raw_score = score.group("score") + entries.append({ + "request": int(score.group("rid")), + "score": ( + float(raw_score) if raw_score != "nan" + else math.nan + ), + "source": score.group("source"), + "score_kind": score.group("score_kind"), + "admitted": score.group("admitted") == "*", + }) + rounds.append({ + "concurrency": int(gate.group("c")), + "k": int(gate.group("k")), + "predicted_goodput_tok_s": float(gate.group("gk")) * 1e6, + "ar_goodput_tok_s": float(gate.group("g0")) * 1e6, + "predicted_cost_us": float(gate.group("predicted")), + "measured_cost_us": ( + float(gate.group("measured")) + if gate.group("measured") is not None else None + ), + "entries": entries, + }) + continue + activation = ACTIVATION_RE.search(line) + if activation: + row = _json_object(activation, path, line_no) + _validate_activation(row, path, line_no) + activations.append(row) + continue + metric = METRIC_RE.search(line) + if metric: + row = _json_object(metric, path, line_no) + request_id = row.get("request_id") + if isinstance(request_id, str): + if request_id in metrics: + raise ValueError( + f"{path}:{line_no}: duplicate request metric {request_id}" + ) + metrics[request_id] = row + continue + timing = TIMING_RE.search(line) + if timing: + row = _json_object(timing, path, line_no) + _validate_timing(row, path, line_no) + timings.append(row) + return rounds, metrics, timings, activations + + +def _prompt_records(prompt_file: Path | None) -> list[dict[str, Any]]: + if prompt_file is None or not prompt_file.exists(): + return [] + records = [] + for line_no, line in enumerate( + prompt_file.read_text(encoding="utf-8").splitlines(), 1, + ): + if not line.strip(): + continue + value = json.loads(line) + if not isinstance(value, dict): + raise ValueError(f"{prompt_file}:{line_no}: prompt must be an object") + records.append(value) + return records + + +def prompt_names(bench: dict[str, Any], prompt_file: Path | None) -> dict[str, dict]: + """Map wire request_id to prompt identity and selection metadata.""" + records = _prompt_records(prompt_file) + out: dict[str, dict] = {} + for level in bench.get("levels", []): + for detail in level.get("requests_detail", []): + request_id = detail.get("request_id") + index = detail.get("prompt_index") + if not isinstance(request_id, str) or type(index) is not int: + continue + source = records[index] if index < len(records) else {} + out[request_id] = { + "prompt_index": index, + "prompt_id": str(source.get("id", f"prompt-{index}")), + "selection_class": source.get("selection_class"), + "expected_dense_oracle": source.get("expected_dense_oracle"), + "dense_r9700_baseline": source.get("dense_r9700_baseline"), + "decode_tok_s": detail.get("request_decode_tok_s"), + "output_sha256": detail.get("content_sha256"), + } + return out + + +def _find_prompt_file( + case_dir: Path, workload: str, explicit: Path | None, +) -> Path | None: + if explicit is not None: + return explicit + for parent in (case_dir, *case_dir.parents): + candidate = parent / "prompts" / f"{workload}.jsonl" + if candidate.is_file(): + return candidate + return None + + +def _percentile(values: list[float], fraction: float) -> float | None: + if not values: + return None + ordered = sorted(values) + position = (len(ordered) - 1) * fraction + lower = math.floor(position) + upper = math.ceil(position) + if lower == upper: + return ordered[lower] + weight = position - lower + return ordered[lower] * (1.0 - weight) + ordered[upper] * weight + + +def _timing_summary(rows: list[dict[str, Any]]) -> dict[str, Any]: + emitted = sum(int(row["emitted_tokens"]) for row in rows) + accepted = sum(int(row["accepted_tokens"]) for row in rows) + forwards = sum(int(row["target_forwards"]) for row in rows) + total_us = sum(float(row["total_us"]) for row in rows) + spec_lane_steps = sum(int(row["k"]) for row in rows) + phase_keys = sorted({ + key for row in rows for key in row + if key.endswith("_us") and key != "total_us" + }) + means = { + key: statistics.fmean(float(row.get(key, 0.0)) for row in rows) + for key in phase_keys + } + medians = { + key: statistics.median(float(row.get(key, 0.0)) for row in rows) + for key in phase_keys + } + p95 = { + key: _percentile( + [float(row.get(key, 0.0)) for row in rows], 0.95, + ) + for key in phase_keys + } + return { + "rounds": len(rows), + "total_wall_us": total_us, + "emitted_tokens": emitted, + "accepted_tokens": accepted, + "target_forwards": forwards, + "emitted_tokens_per_target_forward": ( + emitted / forwards if forwards else None + ), + "accepted_tokens_per_spec_lane_step": ( + accepted / spec_lane_steps if spec_lane_steps else None + ), + "round_goodput_tok_s": ( + emitted * 1e6 / total_us if total_us else None + ), + "mean_total_us": statistics.fmean( + float(row["total_us"]) for row in rows + ) if rows else None, + "p50_total_us": statistics.median( + float(row["total_us"]) for row in rows + ) if rows else None, + "p95_total_us": _percentile( + [float(row["total_us"]) for row in rows], 0.95, + ), + "phase_mean_us": means, + "phase_p50_us": medians, + "phase_p95_us": p95, + } + + +def _gate_summary( + rounds: list[dict[str, Any]], timings: list[dict[str, Any]], +) -> tuple[list[dict[str, Any]], int]: + gate_by_k: dict[int, list[dict[str, Any]]] = defaultdict(list) + timing_by_k: dict[int, list[dict[str, Any]]] = defaultdict(list) + for row in rounds: + gate_by_k[int(row["k"])].append(row) + for row in timings: + timing_by_k[int(row["k"])].append(row) + + summaries = [] + mismatch = 0 + for k in sorted(gate_by_k): + gate_rows = gate_by_k[k] + timing_rows = timing_by_k.get(k, []) + paired = min(len(gate_rows), len(timing_rows)) + mismatch += abs(len(gate_rows) - len(timing_rows)) + paired_timing = timing_rows[:paired] + realized_tokens = sum( + int(row["emitted_tokens"]) for row in paired_timing + ) + realized_us = sum(float(row["total_us"]) for row in paired_timing) + predicted = statistics.fmean( + float(row["predicted_goodput_tok_s"]) for row in gate_rows + ) + realized = ( + realized_tokens * 1e6 / realized_us if realized_us else None + ) + measured_costs = [ + float(row["measured_cost_us"]) for row in gate_rows + if row["measured_cost_us"] is not None + ] + summaries.append({ + "k": k, + "gate_rounds": len(gate_rows), + "timing_rounds": len(timing_rows), + "paired_rounds": paired, + "predicted_goodput_mean_tok_s": predicted, + "predicted_ar_goodput_mean_tok_s": statistics.fmean( + float(row["ar_goodput_tok_s"]) for row in gate_rows + ), + "realized_goodput_tok_s": realized, + "realized_over_predicted": ( + realized / predicted if realized is not None and predicted else None + ), + "realized_over_predicted_ar": ( + realized / statistics.fmean( + float(row["ar_goodput_tok_s"]) for row in gate_rows + ) + if realized is not None else None + ), + "predicted_cost_mean_us": statistics.fmean( + float(row["predicted_cost_us"]) for row in gate_rows + ), + "gate_measured_cost_mean_us": ( + statistics.fmean(measured_costs) if measured_costs else None + ), + "timed_wall_mean_us": ( + statistics.fmean(float(row["total_us"]) for row in paired_timing) + if paired_timing else None + ), + }) + return summaries, mismatch + + +def _class_summary(requests: list[dict[str, Any]]) -> list[dict[str, Any]]: + grouped: dict[str, list[dict[str, Any]]] = defaultdict(list) + for row in requests: + selection_class = row.get("selection_class") + if isinstance(selection_class, str): + grouped[selection_class].append(row) + out = [] + for selection_class, rows in sorted(grouped.items()): + gate_rounds = sum(int(row["gate_rounds"]) for row in rows) + admitted = sum(int(row["admitted_rounds"]) for row in rows) + spec_steps = sum(int(row["spec_steps"]) for row in rows) + accepted = sum(int(row["spec_accepted_tokens"]) for row in rows) + out.append({ + "selection_class": selection_class, + "requests": len(rows), + "gate_rounds": gate_rounds, + "admitted_fraction": admitted / gate_rounds if gate_rounds else 0.0, + "spec_steps": spec_steps, + "accepted_tokens": accepted, + "commit_per_spec_step": ( + (spec_steps + accepted) / spec_steps if spec_steps else None + ), + }) + return out + + +def _activation_proof( + variant: str, + activations: list[dict[str, Any]], + metrics: dict[str, dict[str, Any]], + measured_wire_ids: set[str], + log_path: Path, +) -> tuple[dict[str, Any], dict[int, dict[str, Any]]]: + expected_ids: set[int] = set() + engine_to_wire: dict[int, str] = {} + mapping_errors: list[str] = [] + for wire_id, metric in metrics.items(): + engine_id = metric.get("engine_request_id") + if type(engine_id) is not int or engine_id < 0: + mapping_errors.append( + f"metric {wire_id} has no non-negative integer engine_request_id" + ) + continue + previous = engine_to_wire.get(engine_id) + if previous is not None and previous != wire_id: + mapping_errors.append( + f"engine request {engine_id} maps both {previous} and {wire_id}" + ) + continue + engine_to_wire[engine_id] = wire_id + expected_ids.add(engine_id) + + by_id: dict[int, list[dict[str, Any]]] = defaultdict(list) + for activation in activations: + by_id[int(activation["request_id"])].append(activation) + + required = variant == "adaptive-on" + if required: + missing_metric_ids = sorted(measured_wire_ids - set(metrics)) + unknown_metric_ids = sorted(set(metrics) - measured_wire_ids) + duplicate_ids = sorted( + request_id for request_id, rows in by_id.items() if len(rows) != 1 + ) + unknown_ids = sorted(set(by_id) - expected_ids) + missing_ids = sorted(expected_ids - set(by_id)) + errors = list(mapping_errors) + if missing_metric_ids: + errors.append( + "missing engine metrics for measured wire requests " + + ",".join(missing_metric_ids) + ) + if unknown_metric_ids: + errors.append( + "engine metrics reference unknown measured wire requests " + + ",".join(unknown_metric_ids) + ) + if duplicate_ids: + errors.append( + "duplicate activations for engine requests " + + ",".join(str(value) for value in duplicate_ids) + ) + if unknown_ids: + errors.append( + "activations reference unknown engine requests " + + ",".join(str(value) for value in unknown_ids) + ) + if missing_ids: + errors.append( + "missing activations for engine requests " + + ",".join(str(value) for value in missing_ids) + ) + for request_id in sorted(expected_ids & set(by_id)): + rows = by_id[request_id] + if len(rows) != 1: + continue + metric = metrics[engine_to_wire[request_id]] + spec_steps = metric.get("spec_steps") + service_steps = metric.get("spec_service_ar_steps", 0) + target_forwards = metric.get("target_forwards") + if ( + type(spec_steps) is not int + or spec_steps < 0 + or type(service_steps) is not int + or service_steps < 0 + or type(target_forwards) is not int + or target_forwards < 0 + ): + errors.append( + f"engine request {request_id} has invalid execution " + "counters" + ) + continue + decision = rows[0]["decision"] + if decision == "ar" and ( + spec_steps != 0 or service_steps != 0 + ): + errors.append( + f"AR activation for engine request {request_id} " + f"executed {spec_steps} speculation steps and " + f"{service_steps} service AR steps" + ) + if decision == "speculation" and ( + spec_steps == 0 + or target_forwards != 2 * spec_steps + service_steps + ): + errors.append( + f"Spec activation for engine request {request_id} " + "contains a non-speculative target step " + f"(spec_steps={spec_steps}, " + f"spec_service_ar_steps={service_steps}, " + f"target_forwards={target_forwards})" + ) + if errors: + raise ValueError( + f"{log_path}: adaptive-on activation proof failed: " + + "; ".join(errors) + ) + + unique = {request_id: rows[0] for request_id, rows in by_id.items()} + decisions = { + decision: sum( + 1 for row in activations if row["decision"] == decision + ) + for decision in ("ar", "speculation") + } + evaluations = { + evaluation: sum( + 1 for row in activations if row["evaluation"] == evaluation + ) + for evaluation in ("scored", "failed") + } + score_kind_counts: dict[str, int] = defaultdict(int) + fallback_reason_counts: dict[str, int] = defaultdict(int) + decision_reason_counts: dict[str, int] = defaultdict(int) + for row in activations: + score_kind_counts[row["score_kind"]] += 1 + if isinstance(row.get("fallback_reason"), str): + fallback_reason_counts[row["fallback_reason"]] += 1 + if isinstance(row.get("decision_reason"), str): + decision_reason_counts[row["decision_reason"]] += 1 + return ({ + "required": required, + "validation": "passed" if required else "not-required", + "execution_validation": "passed" if required else "not-required", + "measured_engine_requests": len(expected_ids), + "records": len(activations), + "unique_requests": len(by_id), + "matched_requests": len(set(by_id) & expected_ids), + "decision_counts": decisions, + "evaluation_counts": evaluations, + "score_kind_counts": dict(sorted(score_kind_counts.items())), + "fallback_reason_counts": dict(sorted(fallback_reason_counts.items())), + "decision_reason_counts": dict(sorted(decision_reason_counts.items())), + }, unique) + + +def analyze_case(case_dir: Path, prompt_file: Path | None = None) -> dict[str, Any]: + bench_path = case_dir / "bench.json" + bench = json.loads(bench_path.read_text(encoding="utf-8")) + metadata = bench.get("server_metadata") or {} + workload = str(metadata.get("workload") or "") + variant = str(metadata.get("variant") or case_dir.name) + repeat = metadata.get("repeat") + clients = metadata.get("clients") + if type(clients) is not int: + levels = bench.get("levels") or [] + clients = levels[0].get("clients") if levels else None + + log_path = case_dir / "benchmark-server.log" + if not log_path.is_file(): + log_path = case_dir / "server.log" + rounds, metrics, timings, activations = parse_server_log(log_path) + resolved_prompt_file = _find_prompt_file( + case_dir, workload, prompt_file, + ) + by_wire = prompt_names(bench, resolved_prompt_file) + activation_summary, activation_by_engine = _activation_proof( + variant, activations, metrics, set(by_wire), log_path, + ) + + engine_to_prompt: dict[int, dict[str, Any]] = {} + for wire_id, row in metrics.items(): + info = by_wire.get(wire_id) + engine_id = row.get("engine_request_id") + if info is None or type(engine_id) is not int: + continue + engine_to_prompt[engine_id] = { + **info, + "spec_accepted_tokens": row.get("spec_accepted_tokens", 0), + "spec_steps": row.get("spec_steps", 0), + "target_forwards": row.get("target_forwards", 0), + "output_tokens": row.get("output_tokens", 0), + } + + per_request: dict[int, dict[str, Any]] = defaultdict( + lambda: { + "rounds": 0, "admitted": 0, + "activation_score_sum": 0.0, "activation_scored": 0, + } + ) + k_histogram: dict[int, int] = defaultdict(int) + for entry in rounds: + k_histogram[int(entry["k"])] += 1 + for score in entry["entries"]: + stats = per_request[int(score["request"])] + stats["rounds"] += 1 + if score["admitted"]: + stats["admitted"] += 1 + if ( + math.isfinite(score["score"]) + and score["source"] in ("fresh", "initial") + ): + stats["activation_score_sum"] += score["score"] + stats["activation_scored"] += 1 + + requests = [] + for engine_id, prompt in sorted( + engine_to_prompt.items(), key=lambda item: item[1]["prompt_index"], + ): + stats = per_request[engine_id] + steps = prompt.get("spec_steps", 0) + accepted = prompt.get("spec_accepted_tokens", 0) + activation = activation_by_engine.get(engine_id) + requests.append({ + "engine_request_id": engine_id, + **prompt, + "gate_rounds": stats["rounds"], + "admitted_rounds": stats["admitted"], + "admitted_fraction": ( + stats["admitted"] / stats["rounds"] if stats["rounds"] else 0.0 + ), + "mean_activation_score": ( + stats["activation_score_sum"] / stats["activation_scored"] + if stats["activation_scored"] else None + ), + "activation_slot": ( + activation.get("slot") if activation is not None else None + ), + "activation_score": ( + activation.get("activation_score") + if activation is not None else None + ), + "activation_score_kind": ( + activation.get("score_kind") + if activation is not None else None + ), + "activation_hazards": ( + activation.get("hazards") if activation is not None else None + ), + "expected_yield": ( + activation.get("expected_yield") + if activation is not None else None + ), + "activation_decision": ( + activation.get("decision") if activation is not None else None + ), + "activation_evaluation": ( + activation.get("evaluation") + if activation is not None else None + ), + "activation_fallback_reason": ( + activation.get("fallback_reason") + if activation is not None else None + ), + "activation_decision_reason": ( + activation.get("decision_reason") + if activation is not None else None + ), + "commit_per_spec_step": ( + (steps + accepted) / steps if steps else None + ), + }) + + levels = bench.get("levels") or [] + aggregate_tok_s = ( + levels[0].get("aggregate_tok_s") if len(levels) == 1 else None + ) + by_path = { + path: _timing_summary([ + row for row in timings if row["path"] == path + ]) + for path in ("ar", "spec", "spec-direct") + if any(row["path"] == path for row in timings) + } + shape_groups: dict[tuple[str, int, int], list[dict[str, Any]]] = defaultdict(list) + for row in timings: + shape_groups[ + (str(row["path"]), int(row["live"]), int(row["k"])) + ].append(row) + by_shape = [ + { + "path": path, + "live": live, + "k": k, + **_timing_summary(rows), + } + for (path, live, k), rows in sorted(shape_groups.items()) + ] + ar_timings = [row for row in timings if row["path"] == "ar"] + drafted_ar = [row for row in ar_timings if float(row.get("draft_us", 0)) > 0] + ar_wall = sum(float(row["total_us"]) for row in ar_timings) + ar_draft = sum(float(row.get("draft_us", 0)) for row in ar_timings) + gate_by_k, gate_timing_mismatch = _gate_summary(rounds, timings) + + return { + "case": str(case_dir), + "log": str(log_path), + "workload": workload, + "clients": clients, + "repeat": repeat, + "variant": variant, + "aggregate_tok_s": aggregate_tok_s, + "activation": activation_summary, + "gate_rounds": len(rounds), + "k_histogram": dict(sorted(k_histogram.items())), + "gate_by_k": gate_by_k, + "gate_timing_count_mismatch": gate_timing_mismatch, + "timing": { + "records": len(timings), + "by_path": by_path, + "by_shape": by_shape, + "draft_tax_on_ar": { + "ar_rounds": len(ar_timings), + "drafted_ar_rounds": len(drafted_ar), + "draft_us": ar_draft, + "ar_wall_us": ar_wall, + "fraction_of_ar_wall": ar_draft / ar_wall if ar_wall else None, + }, + }, + "selection_classes": _class_summary(requests), + "requests": requests, + } + + +def discover_case_dirs(paths: list[Path]) -> list[Path]: + case_dirs: set[Path] = set() + for path in paths: + if path.name == "bench.json" and path.is_file(): + case_dirs.add(path.parent) + elif (path / "bench.json").is_file(): + case_dirs.add(path) + elif path.is_dir(): + case_dirs.update(item.parent for item in path.rglob("bench.json")) + else: + raise ValueError(f"{path}: not a benchmark case or matrix root") + if not case_dirs: + raise ValueError("no bench.json files found") + return sorted(case_dirs) + + +def compare_prompts(cases: list[dict[str, Any]]) -> list[dict[str, Any]]: + grouped: dict[tuple[Any, ...], dict[str, dict[str, Any]]] = defaultdict(dict) + metadata: dict[tuple[Any, ...], dict[str, Any]] = {} + for case in cases: + for request in case["requests"]: + key = ( + case["workload"], case["clients"], case["repeat"], + request["prompt_id"], + ) + grouped[key][case["variant"]] = request + metadata[key] = { + name: request.get(name) for name in ( + "selection_class", "expected_dense_oracle", + "dense_r9700_baseline", + ) + } + + comparisons = [] + for key in sorted(grouped, key=lambda value: tuple(str(item) for item in value)): + variants = grouped[key] + if "ar" not in variants or "speculation" not in variants: + continue + ar_rate = variants["ar"].get("decode_tok_s") + spec_rate = variants["speculation"].get("decode_tok_s") + ratio = ( + spec_rate / ar_rate + if type(ar_rate) in (int, float) and ar_rate > 0 + and type(spec_rate) in (int, float) else None + ) + empirical_oracle = ( + "speculation" if ratio is not None and ratio > 1.0 else "ar" + ) + row = { + "workload": key[0], + "clients": key[1], + "repeat": key[2], + "prompt_id": key[3], + **metadata[key], + "ar_decode_tok_s": ar_rate, + "speculation_decode_tok_s": spec_rate, + "speculation_over_ar": ratio, + "empirical_concurrent_oracle": empirical_oracle, + "matches_expected_dense_oracle": ( + empirical_oracle == metadata[key].get("expected_dense_oracle") + if metadata[key].get("expected_dense_oracle") is not None else None + ), + "output_stable": ( + variants["ar"].get("output_sha256") + == variants["speculation"].get("output_sha256") + and variants["ar"].get("output_sha256") is not None + ), + "adaptive": {}, + } + for name, request in sorted(variants.items()): + if name.startswith("adaptive-"): + row["adaptive"][name] = { + "decode_tok_s": request.get("decode_tok_s"), + "admitted_fraction": request.get("admitted_fraction"), + "activation_decision": request.get( + "activation_decision" + ), + "activation_evaluation": request.get( + "activation_evaluation" + ), + "activation_fallback_reason": request.get( + "activation_fallback_reason" + ), + "activation_decision_reason": request.get( + "activation_decision_reason" + ), + "activation_score_kind": request.get( + "activation_score_kind" + ), + "activation_score": request.get("activation_score"), + "mean_activation_score": request.get("mean_activation_score"), + "expected_yield": request.get("expected_yield"), + "activation_hazards": request.get("activation_hazards"), + "commit_per_spec_step": request.get( + "commit_per_spec_step" + ), + } + comparisons.append(row) + return comparisons + + +def compare_activation_shapes( + cases: list[dict[str, Any]], +) -> list[dict[str, Any]]: + """Compare every non-control round shape with pure AR at the same live C.""" + grouped: dict[tuple[Any, ...], dict[str, dict[str, Any]]] = defaultdict(dict) + for case in cases: + grouped[ + (case["workload"], case["clients"], case["repeat"]) + ][case["variant"]] = case + + comparisons = [] + for key, variants in sorted( + grouped.items(), key=lambda item: tuple(str(value) for value in item[0]), + ): + ar_case = variants.get("ar") + if ar_case is None: + continue + ar_shapes = { + int(shape["live"]): shape + for shape in ar_case["timing"]["by_shape"] + if shape["path"] == "ar" and int(shape["k"]) == 0 + } + ar_aggregate = ar_case.get("aggregate_tok_s") + for variant, case in sorted(variants.items()): + if variant == "ar": + continue + aggregate = case.get("aggregate_tok_s") + case["aggregate_over_ar"] = ( + aggregate / ar_aggregate + if type(aggregate) in (int, float) + and type(ar_aggregate) in (int, float) + and ar_aggregate > 0 else None + ) + for shape in case["timing"]["by_shape"]: + baseline = ar_shapes.get(int(shape["live"])) + if baseline is None: + continue + realized = shape.get("round_goodput_tok_s") + ar_rate = baseline.get("round_goodput_tok_s") + ratio = ( + realized / ar_rate + if type(realized) in (int, float) + and type(ar_rate) in (int, float) + and ar_rate > 0 else None + ) + path = str(shape["path"]) + comparisons.append({ + "workload": key[0], + "clients": key[1], + "repeat": key[2], + "variant": variant, + "path": path, + "live": shape["live"], + "k": shape["k"], + "rounds": shape["rounds"], + "realized_goodput_tok_s": realized, + "pure_ar_goodput_tok_s": ar_rate, + "realized_over_pure_ar": ratio, + "mean_wall_us": shape.get("mean_total_us"), + "mean_draft_us": ( + shape.get("phase_mean_us") or {} + ).get("draft_us"), + "activation_outcome": ( + "profitable" if path in ("spec", "spec-direct") + and ratio is not None and ratio > 1.0 + else "unprofitable" if path in ("spec", "spec-direct") + else "ar-with-draft-tax" if ( + (shape.get("phase_mean_us") or {}).get( + "draft_us", 0.0 + ) > 0 + ) else "ar" + ), + }) + for case in cases: + if case["variant"] == "ar": + case["aggregate_over_ar"] = 1.0 + else: + case.setdefault("aggregate_over_ar", None) + return comparisons + + +def _fmt(value: Any, digits: int = 2) -> str: + return f"{value:.{digits}f}" if type(value) in (int, float) else "n/a" + + +def render_markdown(report: dict[str, Any]) -> str: + lines = [ + "# Adaptive speculation profiling", + "", + "Round goodput is emitted decode tokens divided by the common measured " + "decode-round wall. Draft time is already inside that wall and is also " + "reported separately as attribution.", + "", + "## Case overview", + "", + "| Workload | C | Repeat | Variant | Benchmark tok/s | AR rounds | " + "Spec rounds | Timed tok/s | vs AR | AR draft tax |", + "| :--- | ---: | ---: | :--- | ---: | ---: | ---: | ---: | ---: | " + "---: |", + ] + for case in report["cases"]: + paths = case["timing"]["by_path"] + ar = paths.get("ar", {}) + spec = paths.get("spec-direct") or paths.get("spec", {}) + total_tokens = sum( + value.get("emitted_tokens", 0) for value in paths.values() + ) + total_wall = sum( + value.get("total_wall_us", 0.0) for value in paths.values() + ) + timed_rate = total_tokens * 1e6 / total_wall if total_wall else None + tax = case["timing"]["draft_tax_on_ar"]["fraction_of_ar_wall"] + tax_text = f"{tax * 100:.1f}%" if tax is not None else "n/a" + lines.append( + f"| {case['workload']} | {case['clients']} | {case['repeat']} | " + f"{case['variant']} | {_fmt(case['aggregate_tok_s'])} | " + f"{ar.get('rounds', 0)} | {spec.get('rounds', 0)} | " + f"{_fmt(timed_rate)} | {_fmt(case['aggregate_over_ar'], 3)} | " + f"{tax_text} |" + ) + + if report["activation_comparisons"]: + lines += [ + "", + "## Activation outcome against matched pure AR", + "", + "Each row compares the measured round shape with pure AR at the same " + "number of live requests. A speculative ratio below 1 is an " + "activation error for that concurrency shape.", + "", + "| Workload | Variant | Live | k | Path | Rounds | Actual tok/s | " + "Pure AR tok/s | Actual/AR | Draft mean us | Outcome |", + "| :--- | :--- | ---: | ---: | :--- | ---: | ---: | ---: | ---: | " + "---: | :--- |", + ] + for row in report["activation_comparisons"]: + lines.append( + f"| {row['workload']} | {row['variant']} | {row['live']} | " + f"{row['k']} | {row['path']} | {row['rounds']} | " + f"{_fmt(row['realized_goodput_tok_s'])} | " + f"{_fmt(row['pure_ar_goodput_tok_s'])} | " + f"{_fmt(row['realized_over_pure_ar'], 3)} | " + f"{_fmt(row['mean_draft_us'], 1)} | " + f"{row['activation_outcome']} |" + ) + + phase_rows = [ + (case, path, summary) + for case in report["cases"] + for path, summary in case["timing"]["by_path"].items() + ] + if phase_rows: + lines += [ + "", + "## Mean phase attribution", + "", + "Draft is a subset of pre-round time, so it must not be added to " + "the other columns a second time. Verify and replay are the two " + "target forwards on speculative rounds.", + "", + "| Variant | Path | Rounds | Wall us | Draft us | AR graph us | " + "Verify us | Replay us | Build us | Readback us |", + "| :--- | :--- | ---: | ---: | ---: | ---: | ---: | ---: | " + "---: | ---: |", + ] + for case, path, summary in phase_rows: + phase = summary["phase_mean_us"] + build_us = sum( + phase.get(key, 0.0) for key in ( + "graph_build_us", "graph_prepare_us", + "verify_build_us", "replay_build_us", + ) + ) + read_us = sum( + phase.get(key, 0.0) for key in ( + "posterior_read_us", "sample_read_us", + ) + ) + lines.append( + f"| {case['variant']} | {path} | {summary['rounds']} | " + f"{_fmt(summary['mean_total_us'], 1)} | " + f"{_fmt(phase.get('draft_us'), 1)} | " + f"{_fmt(phase.get('graph_exec_us'), 1)} | " + f"{_fmt(phase.get('verify_exec_us'), 1)} | " + f"{_fmt(phase.get('replay_exec_us'), 1)} | " + f"{_fmt(build_us, 1)} | {_fmt(read_us, 1)} |" + ) + + if report["prompt_comparisons"]: + lines += [ + "", + "## Per-prompt concurrent oracle", + "", + "| Workload | C | Prompt | Class | AR tok/s | Spec tok/s | Spec/AR | " + "Concurrent oracle | Dense label agrees | Stable |", + "| :--- | ---: | :--- | :--- | ---: | ---: | ---: | :--- | " + ":---: | :---: |", + ] + for row in report["prompt_comparisons"]: + lines.append( + f"| {row['workload']} | {row['clients']} | {row['prompt_id']} | " + f"{row.get('selection_class') or 'n/a'} | " + f"{_fmt(row['ar_decode_tok_s'])} | " + f"{_fmt(row['speculation_decode_tok_s'])} | " + f"{_fmt(row['speculation_over_ar'], 3)} | " + f"{row['empirical_concurrent_oracle']} | " + f"{row['matches_expected_dense_oracle']} | " + f"{row['output_stable']} |" + ) + + adaptive_classes = [ + (case, item) + for case in report["cases"] if case["variant"].startswith("adaptive-") + for item in case["selection_classes"] + ] + if adaptive_classes: + lines += [ + "", + "## Adaptive decisions by prompt class", + "", + "| Workload | C | Variant | Class | Gate rounds | Admitted | " + "Commit/spec step |", + "| :--- | ---: | :--- | :--- | ---: | ---: | ---: |", + ] + for case, item in adaptive_classes: + lines.append( + f"| {case['workload']} | {case['clients']} | {case['variant']} | " + f"{item['selection_class']} | {item['gate_rounds']} | " + f"{_fmt(item['admitted_fraction'] * 100, 1)}% | " + f"{_fmt(item['commit_per_spec_step'])} |" + ) + + gate_rows = [ + (case, item) + for case in report["cases"] + for item in case["gate_by_k"] + ] + if gate_rows: + lines += [ + "", + "## Initial prediction accuracy", + "", + "| Workload | C | Variant | k | Rounds | Pred tok/s | " + "Realized tok/s | Realized/pred | Realized/AR | Pred cost us | " + "Timed wall us |", + "| :--- | ---: | :--- | ---: | ---: | ---: | ---: | ---: | " + "---: | ---: | ---: |", + ] + for case, item in gate_rows: + lines.append( + f"| {case['workload']} | {case['clients']} | " + f"{case['variant']} | {item['k']} | {item['gate_rounds']} | " + f"{_fmt(item['predicted_goodput_mean_tok_s'])} | " + f"{_fmt(item['realized_goodput_tok_s'])} | " + f"{_fmt(item['realized_over_predicted'], 3)} | " + f"{_fmt(item['realized_over_predicted_ar'], 3)} | " + f"{_fmt(item['predicted_cost_mean_us'], 1)} | " + f"{_fmt(item['timed_wall_mean_us'], 1)} |" + ) + + lines.append("") + return "\n".join(lines) + + +def build_report( + paths: list[Path], prompt_file: Path | None = None, +) -> dict[str, Any]: + cases = [ + analyze_case(case_dir, prompt_file) + for case_dir in discover_case_dirs(paths) + ] + activation_comparisons = compare_activation_shapes(cases) + return { + "schema_version": 7, + "cases": cases, + "prompt_comparisons": compare_prompts(cases), + "activation_comparisons": activation_comparisons, + } + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("paths", nargs="+", type=Path) + parser.add_argument( + "--prompt-file", type=Path, default=None, + help="override prompt metadata (normally discovered from matrix root)", + ) + parser.add_argument("--out", type=Path) + parser.add_argument("--markdown-out", type=Path) + args = parser.parse_args() + report = build_report(args.paths, args.prompt_file) + rendered = json.dumps(report, indent=2, sort_keys=True) + "\n" + if args.out: + args.out.write_text(rendered, encoding="utf-8") + if args.markdown_out: + args.markdown_out.write_text( + render_markdown(report), encoding="utf-8", + ) + if args.out or args.markdown_out: + destinations = ", ".join( + str(path) for path in (args.out, args.markdown_out) + if path is not None + ) + print( + f"[profile] cases={len(report['cases'])} " + f"prompt_comparisons={len(report['prompt_comparisons'])} " + f"activation_comparisons={len(report['activation_comparisons'])} " + f"wrote {destinations}" + ) + else: + print(rendered, end="") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/harness/benchmarks/concurrency/concurrent_benchmark.py b/harness/benchmarks/concurrency/concurrent_benchmark.py new file mode 100755 index 000000000..e5004089a --- /dev/null +++ b/harness/benchmarks/concurrency/concurrent_benchmark.py @@ -0,0 +1,441 @@ +#!/usr/bin/env python3 +"""Measure concurrent goodput, TTFT, request IDs, and prompt telemetry.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import statistics +import sys +import threading +import time +import urllib.request +from pathlib import Path +from typing import Any, Iterable + +CLIENT_SCRIPT = Path(__file__).resolve() + + +def client_provenance(argv: list[str] | None = None) -> dict[str, Any]: + """Return the literal process argv and exact client source digest.""" + process_argv = list(sys.orig_argv if argv is None else argv) + if not process_argv or not all(isinstance(value, str) for value in process_argv): + raise ValueError("client process argv must be a non-empty string array") + return { + "client_argv": process_argv, + "client_script": str(CLIENT_SCRIPT), + "client_script_sha256": hashlib.sha256(CLIENT_SCRIPT.read_bytes()).hexdigest(), + } + + + +def sha256_text(text: str) -> str: + return hashlib.sha256(text.encode("utf-8")).hexdigest() + + +def load_prompts(path: Path) -> list[str]: + prompts = [] + for line_no, raw in enumerate(path.read_text(encoding="utf-8").splitlines(), 1): + line = raw.strip() + if not line: + continue + if line.startswith("{"): + value = json.loads(line).get("prompt") + if not isinstance(value, str) or not value: + raise ValueError(f"{path}:{line_no}: missing string 'prompt'") + prompts.append(value) + else: + prompts.append(line) + if not prompts: + raise ValueError(f"{path}: no prompts") + return prompts + + +def request_prompts(prompts: list[str], count: int, offset: int) -> list[str]: + if offset < 0: + raise ValueError("--prompt-offset must be >= 0") + if offset + count > len(prompts): + raise ValueError( + f"need prompts [{offset}, {offset + count}), but only " + f"{len(prompts)} were supplied; refusing to reuse prompts" + ) + return prompts[offset:offset + count] + + +def iter_sse_data(lines: Iterable[bytes]) -> Iterable[str]: + data: list[str] = [] + for raw in lines: + line = raw.decode("utf-8", errors="replace").rstrip("\r\n") + if not line: + if data: + yield "\n".join(data) + data.clear() + elif line.startswith("data:"): + data.append(line[5:].lstrip()) + if data: + yield "\n".join(data) + + +def stream_request(args: argparse.Namespace, prompt: str) -> dict[str, Any]: + started = time.perf_counter() + first = None + request_id = None + content: list[str] = [] + reasoning: list[str] = [] + completion_tokens = None + prompt_tokens = None + finish_reason = None + done_received = False + timings: dict[str, Any] = {} + wire_metrics: dict[str, Any] = {} + error = None + payload = { + "model": args.model, + "messages": [{"role": "user", "content": prompt}], + "stream": True, + "stream_options": {"include_usage": True}, + "max_tokens": args.max_tokens, + "temperature": args.temperature, + "seed": args.seed, + } + if args.ignore_eos: + payload["ignore_eos"] = True + headers = {"Content-Type": "application/json"} + if args.api_key: + headers["Authorization"] = f"Bearer {args.api_key}" + request = urllib.request.Request( + args.base_url.rstrip("/") + "/chat/completions", + data=json.dumps(payload).encode("utf-8"), headers=headers, method="POST", + ) + try: + with urllib.request.urlopen(request, timeout=args.timeout) as response: + for data in iter_sse_data(response): + if data == "[DONE]": + done_received = True + break + event = json.loads(data) + if isinstance(event.get("id"), str): + request_id = event["id"] + usage = event.get("usage") or {} + if type(usage.get("completion_tokens")) is int: + completion_tokens = usage["completion_tokens"] + if type(usage.get("prompt_tokens")) is int: + prompt_tokens = usage["prompt_tokens"] + if isinstance(usage.get("timings"), dict): + timings = dict(usage["timings"]) + if isinstance(usage.get("concurrency_metrics"), dict): + wire_metrics = dict(usage["concurrency_metrics"]) + for choice in event.get("choices") or []: + if choice.get("finish_reason") is not None: + finish_reason = choice["finish_reason"] + delta = choice.get("delta") or {} + piece = delta.get("content") + thought = delta.get("reasoning_content") + if isinstance(piece, str) and piece: + first = first or time.perf_counter() + content.append(piece) + if isinstance(thought, str) and thought: + first = first or time.perf_counter() + reasoning.append(thought) + except Exception as exc: # retain partial data for diagnosis + error = f"{type(exc).__name__}: {exc}" + if error is None and not done_received: + error = "ProtocolError: stream ended before [DONE]" + elif error is None and finish_reason is None: + error = "ProtocolError: stream ended without a terminal finish_reason" + ended = time.perf_counter() + output = "".join(content) + reasoning_output = "".join(reasoning) + decode_duration = ended - first if first is not None and ended > first else None + request_decode_tok_s = ( + (completion_tokens - 1) / decode_duration + if type(completion_tokens) is int and completion_tokens > 0 + and decode_duration is not None else None + ) + return { + "request_id": request_id, + "t_start": started, "t_first": first, "t_end": ended, + "duration_s": ended - started, + "ttft_s": first - started if first is not None else None, + "decode_duration_s": decode_duration, + "completion_tokens": completion_tokens, "prompt_tokens": prompt_tokens, + "effective_prompt_tokens": timings.get("effective_prompt_tokens"), + "prefilled_tokens": timings.get("prefilled_tokens"), + "cached_prefix_tokens": timings.get("cached_prefix_tokens"), + "cache_hit": timings.get("cache_hit"), + "server_prefill_ms": timings.get("prefill_ms"), + "server_decode_ms": timings.get("decode_ms"), + "server_decode_tokens_per_sec": timings.get("decode_tokens_per_sec"), + "server_timings": timings, + "wire_concurrency_metrics": wire_metrics, + "finish_reason": finish_reason, "done_received": done_received, "error": error, + "content_sha256": sha256_text(output), + "reasoning_content_sha256": sha256_text(reasoning_output), + "content_chars": len(output), "reasoning_content_chars": len(reasoning_output), + "request_output_tok_s": ( + completion_tokens / (ended - started) + if completion_tokens is not None and ended > started else None + ), + "request_decode_tok_s": request_decode_tok_s, + } + + +def run_level( + clients: int, args: argparse.Namespace, prompts: list[str], offset: int, +) -> dict[str, Any]: + selected = request_prompts(prompts, clients, offset) + barrier = threading.Barrier(clients) + records: list[dict[str, Any] | None] = [None] * clients + + def worker(index: int) -> None: + barrier.wait() + record = stream_request(args, selected[index]) + record["prompt_index"] = offset + index + record["prompt_sha256"] = sha256_text(selected[index]) + records[index] = record + + threads = [threading.Thread(target=worker, args=(i,), daemon=True) for i in range(clients)] + for thread in threads: + thread.start() + deadline = time.monotonic() + args.timeout + 30 + for thread in threads: + thread.join(max(0.0, deadline - time.monotonic())) + hung = sum(thread.is_alive() for thread in threads) + if hung: + raise TimeoutError( + f"{hung} request worker(s) exceeded the level deadline" + ) + completed = [record for record in records if record is not None] + failures = sum(record["error"] is not None for record in completed) + ok = [record for record in completed if record["error"] is None] + starts = [record["t_start"] for record in completed] + ends = [record["t_end"] for record in completed] + level_start = min(starts) if starts else time.perf_counter() + wall = max(ends) - level_start if ends else 0.0 + for record in completed: + record["start_offset_s"] = record["t_start"] - level_start + + completion_counts = [r["completion_tokens"] for r in ok] + prompt_counts = [r["prompt_tokens"] for r in ok] + completion_complete = bool(ok) and all(isinstance(v, int) for v in completion_counts) + prompt_complete = bool(ok) and all(isinstance(v, int) for v in prompt_counts) + ttfts = [r["ttft_s"] for r in ok if r["ttft_s"] is not None] + first_window = ( + max(r["start_offset_s"] + r["ttft_s"] for r in ok) + if len(ttfts) == len(ok) and ok else None + ) + first_times = [r["t_first"] for r in ok if r["t_first"] is not None] + output_window = ( + max(r["t_end"] for r in ok) - min(first_times) + if len(first_times) == len(ok) and ok else None + ) + request_decode_rates = [ + r["request_decode_tok_s"] for r in ok + if r.get("request_decode_tok_s") is not None + ] + fixed_valid = ( + failures == 0 and len(ok) == clients + and completion_complete + and all(v == args.max_tokens for v in completion_counts) + ) if args.ignore_eos else None + prompt_hashes = [r["prompt_sha256"] for r in ok] + output_hashes = [ + [r["content_sha256"], r["reasoning_content_sha256"]] for r in ok + ] + digest = lambda value: sha256_text(json.dumps(value, separators=(",", ":"))) + return { + "clients": clients, "requests": clients, "requests_ok": len(ok), + "failures": failures, "wall_s": wall, + "start_skew_s": max(starts) - min(starts) if starts else None, + "completion_tokens_total": sum(completion_counts) if completion_complete else None, + "token_count_complete": completion_complete, + "fixed_token_workload_valid": fixed_valid, + "aggregate_tok_s": ( + sum(completion_counts) / wall if completion_complete and wall > 0 else None + ), + "aggregate_metric": "completion_tokens_per_level_wall_second", + "output_window_s": output_window, + "output_window_tok_s": ( + sum(completion_counts) / output_window + if completion_complete and output_window is not None and output_window > 0 + else None + ), + "output_window_metric": "completion_tokens_per_first_output_to_final_completion_second", + "request_decode_tok_s_median": ( + statistics.median(request_decode_rates) + if len(request_decode_rates) == len(ok) and ok else None + ), + "prompt_tokens_total": sum(prompt_counts) if prompt_complete else None, + "prompt_tokens_min": min(prompt_counts) if prompt_complete else None, + "prompt_tokens_max": max(prompt_counts) if prompt_complete else None, + "prompt_tokens_distinct": len(set(prompt_counts)) if prompt_complete else None, + "prompt_token_count_complete": prompt_complete, + "prompt_to_first_token_s": first_window, + "prompt_tokens_per_s_to_first_token": ( + sum(prompt_counts) / first_window + if prompt_complete and first_window is not None and first_window > 0 else None + ), + "ttft_median_s": statistics.median(ttfts) if ttfts else None, + "ttft_max_s": max(ttfts) if ttfts else None, + "selected_prompt_set_sha256": digest(prompt_hashes), + "selected_output_set_sha256": digest(output_hashes), + "requests_detail": completed, + } + + +def enrich_level(level: dict[str, Any]) -> None: + ok = [row for row in level["requests_detail"] if row.get("error") is None] + effective = [row.get("effective_prompt_tokens") for row in ok] + effective_complete = bool(ok) and all(type(value) is int for value in effective) + request_ids = [row.get("request_id") for row in ok] + request_ids_complete = ( + bool(ok) + and all(isinstance(value, str) and value for value in request_ids) + and len(set(request_ids)) == len(request_ids) + ) + level.update({ + "request_ids_complete": request_ids_complete, + "effective_prompt_token_count_complete": effective_complete, + "effective_prompt_tokens_total": sum(effective) if effective_complete else None, + "effective_prompt_tokens_min": min(effective) if effective_complete else None, + "effective_prompt_tokens_max": max(effective) if effective_complete else None, + "effective_to_wire_prompt_ratio": ( + sum(effective) / level["prompt_tokens_total"] + if effective_complete and level.get("prompt_tokens_total") else None + ), + }) + for key in ( + "server_prefill_ms", "server_decode_ms", "server_decode_tokens_per_sec", + ): + values = [ + row.get(key) for row in ok + if type(row.get(key)) in (int, float) + ] + level[f"{key}_median"] = ( + statistics.median(values) if len(values) == len(ok) and ok else None + ) + + +def fmt(value: Any, spec: str = ".2f") -> str: + return format(value, spec) if isinstance(value, (int, float)) else "n/a" + + +def markdown(report: dict[str, Any]) -> str: + lines = [ + f"# Concurrent benchmark — {report['label']}", "", + "| C | Ok | Output goodput tok/s | Output-window tok/s | " + "Request decode tok/s | Prompt tok/s to first | Wire prompt range | " + "Effective prompt range | Effective/wire | TTFT median/max s |", + "| ---: | ---: | ---: | ---: | ---: | ---: | :--- | :--- | ---: | :--- |", + ] + for level in report["levels"]: + lines.append( + f"| {level['clients']} | {level['requests_ok']}/{level['requests']} | " + f"{fmt(level['aggregate_tok_s'])} | " + f"{fmt(level['output_window_tok_s'])} | " + f"{fmt(level['request_decode_tok_s_median'])} | " + f"{fmt(level['prompt_tokens_per_s_to_first_token'])} | " + f"{fmt(level['prompt_tokens_min'], '.0f')}–" + f"{fmt(level['prompt_tokens_max'], '.0f')} | " + f"{fmt(level['effective_prompt_tokens_min'], '.0f')}–" + f"{fmt(level['effective_prompt_tokens_max'], '.0f')} | " + f"{fmt(level['effective_to_wire_prompt_ratio'], '.3f')} | " + f"{fmt(level['ttft_median_s'], '.3f')}/" + f"{fmt(level['ttft_max_s'], '.3f')} |" + ) + return "\n".join(lines) + "\n" + + +def level_failed(level: dict[str, Any], ignore_eos: bool) -> bool: + return bool( + level["failures"] + or not level["token_count_complete"] + or not level["prompt_token_count_complete"] + or (ignore_eos and level["fixed_token_workload_valid"] is not True) + ) + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base-url", default="http://127.0.0.1:18080/v1") + parser.add_argument("--api-key", default="") + parser.add_argument("--model", default="luce-dflash") + parser.add_argument("--clients", type=int, action="append", dest="client_levels") + parser.add_argument("--prompt-file", type=Path, required=True) + parser.add_argument("--prompt-offset", type=int, default=0) + parser.add_argument("--require-distinct-prompts", action="store_true", + help="Compatibility flag; this client always refuses reuse") + parser.add_argument("--max-tokens", type=int, default=64) + parser.add_argument("--temperature", type=float, default=0.0) + parser.add_argument("--seed", type=int, default=1) + parser.add_argument("--ignore-eos", action="store_true") + parser.add_argument("--timeout", type=float, default=1200.0) + parser.add_argument("--cooldown", type=float, default=0.0) + parser.add_argument("--server-metadata-json", type=Path) + parser.add_argument("--out", type=Path, required=True) + parser.add_argument("--label", default="") + parser.add_argument( + "--require-effective-prompt-telemetry", action="store_true", + help="fail when usage.timings.effective_prompt_tokens is absent", + ) + return parser + + +def run(args: argparse.Namespace) -> int: + levels = args.client_levels or [1, 4, 8, 16] + if any(level < 1 for level in levels): + raise ValueError("--clients must be positive") + if args.prompt_offset < 0 or args.max_tokens < 1 or args.timeout <= 0: + raise ValueError("invalid offset, max-tokens, or timeout") + prompts = load_prompts(args.prompt_file) + results = [] + offset = args.prompt_offset + for index, clients in enumerate(levels): + if index and args.cooldown > 0: + time.sleep(args.cooldown) + print(f"[bench] C={clients} max_tokens={args.max_tokens}", flush=True) + level = run_level(clients, args, prompts, offset) + enrich_level(level) + results.append(level) + offset += clients + metadata = ( + json.loads(args.server_metadata_json.read_text(encoding="utf-8")) + if args.server_metadata_json else {} + ) + report = { + "schema_version": 3, "label": args.label, "base_url": args.base_url, + "model": args.model, "max_tokens": args.max_tokens, + "temperature": args.temperature, "seed": args.seed, + "ignore_eos": args.ignore_eos, "prompt_offset": args.prompt_offset, + "prompt_file_sha256": hashlib.sha256(args.prompt_file.read_bytes()).hexdigest(), + "server_metadata": metadata, "levels": results, + **client_provenance(), + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text( + json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8", + ) + print(markdown(report), end="") + bad = any( + level_failed(level, args.ignore_eos) + or not level["request_ids_complete"] + or ( + args.require_effective_prompt_telemetry + and not level["effective_prompt_token_count_complete"] + ) + for level in results + ) + return 1 if bad else 0 + + +def main() -> int: + try: + return run(build_parser().parse_args()) + except Exception as exc: + print(f"[bench] error: {exc}", file=sys.stderr) + return 2 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/harness/benchmarks/concurrency/forced_subset_benchmark.py b/harness/benchmarks/concurrency/forced_subset_benchmark.py new file mode 100755 index 000000000..78824f109 --- /dev/null +++ b/harness/benchmarks/concurrency/forced_subset_benchmark.py @@ -0,0 +1,618 @@ +#!/usr/bin/env python3 +"""Forced AR/speculation subset diagnostic for concurrent DFlash2. + +This is deliberately not an adaptive benchmark. Every request carries an +explicit ``decode_mode`` and the report fails closed unless server telemetry +proves that the requested mixed mode, active chain depth, and live concurrency +actually executed during the measured window. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import statistics +import sys +import threading +import time +import urllib.request +from pathlib import Path +from typing import Any, Iterable + +import concurrent_benchmark as base + + +CLIENT_SCRIPT = Path(__file__).resolve() +PROFILE_PREFIXES = { + "rounds": "[step-timing] ", + "selectors": "[spec-selector] ", + "activations": "[spec-activation] ", + "requests": "[concurrency-metrics] ", +} +MODES = ("ar", "speculation") + + +def digest_bytes(value: bytes) -> str: + return hashlib.sha256(value).hexdigest() + + +def canonical_digest(value: Any) -> str: + wire = json.dumps(value, ensure_ascii=False, separators=(",", ":")) + return base.sha256_text(wire) + + +def client_provenance(argv: list[str] | None = None) -> dict[str, Any]: + process_argv = list(sys.orig_argv if argv is None else argv) + if not process_argv or not all(isinstance(item, str) for item in process_argv): + raise ValueError("client process argv must be a non-empty string array") + return { + "client_argv": process_argv, + "client_script": str(CLIENT_SCRIPT), + "client_script_sha256": digest_bytes(CLIENT_SCRIPT.read_bytes()), + } + + +def parse_request_modes(raw: str, clients: int) -> list[str]: + modes = [item.strip() for item in raw.split(",")] + if len(modes) != clients or any(item not in MODES for item in modes): + raise ValueError( + "--request-modes must provide exactly one ar/speculation mode " + "per client; adaptive is intentionally out of scope" + ) + return modes + + +def validate_server_metadata( + metadata: dict[str, Any], clients: int, spec_depth: int, + prompt_offset: int, require_selector: bool, +) -> None: + if type(metadata.get("clients")) is not int or metadata["clients"] != clients: + raise ValueError("server metadata clients does not match --clients") + launch = metadata.get("launch_environment") + if not isinstance(launch, dict): + raise ValueError("server metadata lacks launch_environment") + expected = { + "DFLASH_SPEC_CHAIN_DEPTH": str(spec_depth), + "DFLASH_STEP_TIMING": "1", + "PROMPT_OFFSET": str(prompt_offset), + } + if require_selector: + expected["DFLASH_DFLASH2_SELECTOR_LOG"] = "1" + for key, value in expected.items(): + if launch.get(key) != value: + raise ValueError( + f"server metadata must record {key}={value}; got " + f"{launch.get(key)!r}" + ) + + +def stream_request( + args: argparse.Namespace, prompt: str, decode_mode: str, +) -> dict[str, Any]: + started = time.perf_counter() + first = None + second = None + request_id = None + content: list[str] = [] + reasoning: list[str] = [] + completion_tokens = None + prompt_tokens = None + finish_reason = None + done_received = False + timings: dict[str, Any] = {} + wire_metrics: dict[str, Any] = {} + error = None + payload = { + "model": args.model, + "messages": [{"role": "user", "content": prompt}], + "stream": True, + "stream_options": {"include_usage": True}, + "max_tokens": args.max_tokens, + "temperature": 0.0, + "seed": args.seed, + "ignore_eos": True, + "decode_mode": decode_mode, + } + headers = {"Content-Type": "application/json"} + if args.api_key: + headers["Authorization"] = f"Bearer {args.api_key}" + request = urllib.request.Request( + args.base_url.rstrip("/") + "/chat/completions", + data=json.dumps(payload).encode("utf-8"), headers=headers, method="POST", + ) + try: + with urllib.request.urlopen(request, timeout=args.timeout) as response: + for data in base.iter_sse_data(response): + if data == "[DONE]": + done_received = True + break + event = json.loads(data) + if isinstance(event.get("id"), str): + request_id = event["id"] + usage = event.get("usage") or {} + if type(usage.get("completion_tokens")) is int: + completion_tokens = usage["completion_tokens"] + if type(usage.get("prompt_tokens")) is int: + prompt_tokens = usage["prompt_tokens"] + if isinstance(usage.get("timings"), dict): + timings = dict(usage["timings"]) + if isinstance(usage.get("concurrency_metrics"), dict): + wire_metrics = dict(usage["concurrency_metrics"]) + for choice in event.get("choices") or []: + if choice.get("finish_reason") is not None: + finish_reason = choice["finish_reason"] + delta = choice.get("delta") or {} + piece = delta.get("content") + thought = delta.get("reasoning_content") + output_event = ( + isinstance(piece, str) and bool(piece) + ) or ( + isinstance(thought, str) and bool(thought) + ) + if output_event: + now = time.perf_counter() + if first is None: + first = now + elif second is None: + second = now + if isinstance(piece, str) and piece: + content.append(piece) + if isinstance(thought, str) and thought: + reasoning.append(thought) + except Exception as exc: # retain partial evidence for diagnosis + error = f"{type(exc).__name__}: {exc}" + if error is None and not done_received: + error = "ProtocolError: stream ended before [DONE]" + elif error is None and finish_reason is None: + error = "ProtocolError: stream ended without a terminal finish_reason" + ended = time.perf_counter() + output = "".join(content) + reasoning_output = "".join(reasoning) + decode_duration = ended - first if first is not None and ended > first else None + request_decode_tok_s = ( + (completion_tokens - 1) / decode_duration + if type(completion_tokens) is int and completion_tokens > 0 + and decode_duration is not None else None + ) + return { + "request_id": request_id, + "decode_mode": decode_mode, + "request_payload_sha256": canonical_digest(payload), + "t_start": started, "t_first": first, "t_end": ended, + "duration_s": ended - started, + "ttft_s": first - started if first is not None else None, + "first_to_second_output_event_s": ( + second - first if first is not None and second is not None else None + ), + "decode_duration_s": decode_duration, + "completion_tokens": completion_tokens, + "prompt_tokens": prompt_tokens, + "effective_prompt_tokens": timings.get("effective_prompt_tokens"), + "server_timings": timings, + "wire_concurrency_metrics": wire_metrics, + "finish_reason": finish_reason, + "done_received": done_received, + "error": error, + "content_sha256": base.sha256_text(output), + "reasoning_content_sha256": base.sha256_text(reasoning_output), + "combined_output_sha256": canonical_digest([output, reasoning_output]), + "content_chars": len(output), + "reasoning_content_chars": len(reasoning_output), + "request_output_tok_s": ( + completion_tokens / (ended - started) + if type(completion_tokens) is int and ended > started else None + ), + "request_decode_tok_s": request_decode_tok_s, + } + + +def run_level( + args: argparse.Namespace, prompts: list[str], modes: list[str], +) -> dict[str, Any]: + selected = base.request_prompts(prompts, args.clients, args.prompt_offset) + barrier = threading.Barrier(args.clients + 1) + records: list[dict[str, Any] | None] = [None] * args.clients + worker_errors: list[BaseException | None] = [None] * args.clients + + def worker(index: int) -> None: + try: + barrier.wait(timeout=min(args.timeout, 60.0)) + record = stream_request(args, selected[index], modes[index]) + record["request_index"] = index + record["prompt_index"] = args.prompt_offset + index + record["prompt_sha256"] = base.sha256_text(selected[index]) + records[index] = record + except BaseException as exc: # surfaced in the main thread + worker_errors[index] = exc + + threads = [ + threading.Thread(target=worker, args=(index,), daemon=True) + for index in range(args.clients) + ] + for thread in threads: + thread.start() + barrier.wait(timeout=min(args.timeout, 60.0)) + barrier_released = time.perf_counter() + deadline = time.monotonic() + args.timeout + 30.0 + for thread in threads: + thread.join(max(0.0, deadline - time.monotonic())) + hung = sum(thread.is_alive() for thread in threads) + if hung: + raise TimeoutError(f"{hung} request worker(s) exceeded the level deadline") + first_worker_error = next((error for error in worker_errors if error), None) + if first_worker_error is not None: + raise RuntimeError(f"request worker failed: {first_worker_error}") + completed = [record for record in records if record is not None] + if len(completed) != args.clients: + raise RuntimeError("not every synchronized request worker returned a record") + + starts = [float(record["t_start"]) for record in completed] + ends = [float(record["t_end"]) for record in completed] + level_start = min(starts) + wall = max(ends) - level_start + for record in completed: + record["start_offset_s"] = float(record["t_start"]) - level_start + record["barrier_release_offset_s"] = ( + float(record["t_start"]) - barrier_released + ) + ok = [record for record in completed if record["error"] is None] + completion = [record["completion_tokens"] for record in ok] + prompt_counts = [record["prompt_tokens"] for record in ok] + complete_tokens = bool(ok) and all(type(value) is int for value in completion) + complete_prompts = bool(ok) and all(type(value) is int for value in prompt_counts) + decode_rates = [ + record["request_decode_tok_s"] for record in ok + if type(record.get("request_decode_tok_s")) in (int, float) + ] + prompt_hashes = [record["prompt_sha256"] for record in completed] + output_hashes = [ + [record["content_sha256"], record["reasoning_content_sha256"]] + for record in completed + ] + return { + "clients": args.clients, + "request_modes": modes, + "request_mode_mask": "".join( + "A" if mode == "ar" else "S" for mode in modes + ), + "requests": args.clients, + "requests_ok": len(ok), + "failures": args.clients - len(ok), + "wall_s": wall, + "start_skew_s": max(starts) - min(starts), + "completion_tokens_total": sum(completion) if complete_tokens else None, + "token_count_complete": complete_tokens, + "prompt_token_count_complete": complete_prompts, + "fixed_token_workload_valid": ( + len(ok) == args.clients and complete_tokens + and all(value == args.max_tokens for value in completion) + ), + "aggregate_tok_s": ( + sum(completion) / wall if complete_tokens and wall > 0 else None + ), + "request_decode_tok_s_median": ( + statistics.median(decode_rates) + if len(decode_rates) == len(ok) and ok else None + ), + "prompt_tokens_total": sum(prompt_counts) if complete_prompts else None, + "selected_prompt_set_sha256": canonical_digest(prompt_hashes), + "selected_output_set_sha256": canonical_digest(output_hashes), + "requests_detail": completed, + } + + +def parse_profile_records(data: bytes) -> dict[str, list[dict[str, Any]]]: + records = {key: [] for key in PROFILE_PREFIXES} + text = data.decode("utf-8", errors="replace") + for line_index, line in enumerate(text.splitlines(), 1): + for key, prefix in PROFILE_PREFIXES.items(): + marker = line.find(prefix) + if marker < 0: + continue + raw = line[marker + len(prefix):] + try: + value = json.loads(raw) + except json.JSONDecodeError as exc: + raise ValueError( + f"measured server log line {line_index}: invalid {prefix.strip()} " + f"JSON: {exc}" + ) from exc + if not isinstance(value, dict): + raise ValueError( + f"measured server log line {line_index}: {prefix.strip()} " + "record must be an object" + ) + records[key].append({ + "line_index": line_index, + "raw_json": raw, + "record": value, + }) + break + return records + + +def read_log_span(path: Path, start: int) -> tuple[bytes, int]: + size = path.stat().st_size + if size < start: + raise ValueError("server log was truncated or rotated during the benchmark") + with path.open("rb") as handle: + handle.seek(start) + return handle.read(size - start), size + + +def _longest_full_live_streak(rounds: Iterable[dict[str, Any]], clients: int) -> int: + longest = 0 + current = 0 + for wrapped in rounds: + row = wrapped["record"] + if type(row.get("live")) is int and row["live"] == clients: + current += 1 + longest = max(longest, current) + else: + current = 0 + return longest + + +def validate_evidence( + level: dict[str, Any], records: dict[str, list[dict[str, Any]]], + clients: int, modes: list[str], spec_depth: int, + max_start_skew_ms: float, min_full_live_rounds: int, +) -> dict[str, Any]: + errors: list[str] = [] + if level["failures"] or level["requests_ok"] != clients: + errors.append("one or more requests failed") + if not level["token_count_complete"] or not level["prompt_token_count_complete"]: + errors.append("wire token accounting is incomplete") + if level["fixed_token_workload_valid"] is not True: + errors.append("ignore-eos fixed-token workload was not completed exactly") + if float(level["start_skew_s"]) * 1000.0 > max_start_skew_ms: + errors.append( + f"request start skew exceeds {max_start_skew_ms:g} ms" + ) + request_ids = [row.get("request_id") for row in level["requests_detail"]] + if ( + any(not isinstance(value, str) or not value for value in request_ids) + or len(set(request_ids)) != clients + ): + errors.append("wire request IDs are missing or not unique") + + rounds = records["rounds"] + full_live_rounds = sum( + wrapped["record"].get("live") == clients for wrapped in rounds + ) + longest_streak = _longest_full_live_streak(rounds, clients) + if longest_streak < min_full_live_rounds: + errors.append( + f"sustained live=C proof absent: longest live={clients} streak " + f"is {longest_streak}, need {min_full_live_rounds}" + ) + invalid_live = [ + wrapped["record"].get("live") for wrapped in rounds + if type(wrapped["record"].get("live")) is not int + or wrapped["record"]["live"] < 1 + or wrapped["record"]["live"] > clients + ] + if invalid_live: + errors.append(f"step-timing contains invalid live values: {invalid_live}") + + spec_requested = any(mode == "speculation" for mode in modes) + spec_rounds = [ + wrapped["record"] for wrapped in rounds + if wrapped["record"].get("path") in ("spec", "spec-direct") + and type(wrapped["record"].get("k")) is int + and wrapped["record"]["k"] > 0 + ] + inferred_depths: list[int] = [] + for row in spec_rounds: + bucket = row.get("tree_bucket") + tree_rows = row.get("tree_rows") + ar_lanes = row.get("ar_lanes", 0) + spec_rows = ( + tree_rows - ar_lanes + if row.get("path") == "spec-direct" + and type(tree_rows) is int and type(ar_lanes) is int + else tree_rows + ) + if ( + type(bucket) is not int or bucket <= 0 + or type(tree_rows) is not int or tree_rows <= 0 + or type(spec_rows) is not int or spec_rows <= 0 + or spec_rows % bucket != 0 + ): + errors.append("spec step-timing lacks a valid tree_rows/tree_bucket shape") + continue + inferred_depths.append(spec_rows // bucket) + if spec_requested: + if not spec_rounds: + errors.append("speculation was requested but no spec round executed") + if not records["selectors"]: + errors.append("speculation was requested but no DFlash2 selector records exist") + wrong_depths = sorted(set(depth for depth in inferred_depths if depth != spec_depth)) + if wrong_depths: + errors.append( + f"executed chain depths {wrong_depths} do not match requested " + f"depth {spec_depth}" + ) + elif spec_rounds: + errors.append("all-AR mask unexpectedly executed speculation") + + metrics_by_id: dict[str, dict[str, Any]] = {} + for wrapped in records["requests"]: + row = wrapped["record"] + request_id = row.get("request_id") + if not isinstance(request_id, str) or not request_id: + errors.append("concurrency metric has no wire request_id") + elif request_id in metrics_by_id: + errors.append(f"duplicate concurrency metric for {request_id}") + else: + metrics_by_id[request_id] = row + missing = sorted(set(request_ids) - set(metrics_by_id)) + extra = sorted(set(metrics_by_id) - set(request_ids)) + if missing: + errors.append(f"missing per-request concurrency metrics: {missing}") + if extra: + errors.append(f"unmatched per-request concurrency metrics: {extra}") + for request, requested_mode in zip(level["requests_detail"], modes): + metric = metrics_by_id.get(request.get("request_id")) + if metric is None: + continue + steps = metric.get("spec_steps") + if type(steps) is not int or steps < 0: + errors.append(f"{request.get('request_id')}: invalid spec_steps") + elif requested_mode == "speculation" and steps == 0: + errors.append(f"{request.get('request_id')}: forced speculation never executed") + elif requested_mode == "ar" and steps != 0: + errors.append(f"{request.get('request_id')}: forced AR executed speculation") + + output_hashes_complete = all( + isinstance(request.get(key), str) and len(request[key]) == 64 + for request in level["requests_detail"] + for key in ( + "content_sha256", "reasoning_content_sha256", + "combined_output_sha256", + ) + ) + if not output_hashes_complete: + errors.append("exact request output hashes are incomplete") + return { + "passed": not errors, + "errors": errors, + "adaptive_claims_permitted": False, + "full_live_rounds": full_live_rounds, + "longest_full_live_streak": longest_streak, + "min_full_live_rounds": min_full_live_rounds, + "executed_spec_depths": sorted(set(inferred_depths)), + "round_records": len(rounds), + "selector_records": len(records["selectors"]), + "request_metric_records": len(records["requests"]), + } + + +def markdown(report: dict[str, Any]) -> str: + level = report["level"] + validation = report["validation"] + status = "PASS" if validation["passed"] else "FAIL" + return ( + f"# Forced DFlash2 subset diagnostic — {report['label']}\n\n" + "This is a forced-control diagnostic; it makes no adaptive result claim.\n\n" + "| C | Mask | Depth | Ok | Goodput tok/s | Full-live streak | Status |\n" + "| ---: | :--- | ---: | ---: | ---: | ---: | :--- |\n" + f"| {level['clients']} | {level['request_mode_mask']} | " + f"{report['spec_depth']} | {level['requests_ok']}/{level['requests']} | " + f"{base.fmt(level['aggregate_tok_s'])} | " + f"{validation['longest_full_live_streak']} | {status} |\n" + ) + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base-url", default="http://127.0.0.1:18080/v1") + parser.add_argument("--api-key", default="") + parser.add_argument("--model", default="luce-dflash") + parser.add_argument("--clients", type=int, required=True) + parser.add_argument("--request-modes", required=True) + parser.add_argument("--spec-depth", type=int, required=True) + parser.add_argument("--prompt-file", type=Path, required=True) + parser.add_argument("--prompt-offset", type=int, default=0) + parser.add_argument("--max-tokens", type=int, default=256) + parser.add_argument("--seed", type=int, default=1) + parser.add_argument("--timeout", type=float, default=1800.0) + parser.add_argument("--max-start-skew-ms", type=float, default=100.0) + parser.add_argument("--min-full-live-rounds", type=int, default=2) + parser.add_argument("--log-settle-ms", type=float, default=100.0) + parser.add_argument("--server-metadata-json", type=Path, required=True) + parser.add_argument("--server-log", type=Path, required=True) + parser.add_argument("--out", type=Path, required=True) + parser.add_argument("--label", default="") + return parser + + +def run(args: argparse.Namespace) -> int: + if args.clients < 1: + raise ValueError("--clients must be positive") + if args.spec_depth < 2: + raise ValueError("--spec-depth must be at least 2; depth 1 is AR-equivalent") + if args.prompt_offset < 0 or args.max_tokens < 1 or args.timeout <= 0: + raise ValueError("invalid prompt offset, max tokens, or timeout") + if args.max_start_skew_ms < 0 or args.min_full_live_rounds < 1: + raise ValueError("invalid synchronization proof threshold") + if args.log_settle_ms < 0: + raise ValueError("--log-settle-ms must be non-negative") + modes = parse_request_modes(args.request_modes, args.clients) + prompts = base.load_prompts(args.prompt_file) + metadata_bytes = args.server_metadata_json.read_bytes() + metadata = json.loads(metadata_bytes) + if not isinstance(metadata, dict): + raise ValueError("server metadata must be a JSON object") + validate_server_metadata( + metadata, args.clients, args.spec_depth, args.prompt_offset, + require_selector=any(mode == "speculation" for mode in modes), + ) + log_start = args.server_log.stat().st_size + level = run_level(args, prompts, modes) + if args.log_settle_ms: + time.sleep(args.log_settle_ms / 1000.0) + log_span, log_end = read_log_span(args.server_log, log_start) + records = parse_profile_records(log_span) + validation = validate_evidence( + level, records, args.clients, modes, args.spec_depth, + args.max_start_skew_ms, args.min_full_live_rounds, + ) + report = { + "schema_version": 1, + "kind": "dflash2-forced-subset-diagnostic", + "label": args.label, + "scope": { + "forced_controls_only": True, + "adaptive_evaluation": False, + "interpretation": ( + "This artifact may compare forced AR/speculation subsets and " + "depths; it is not an adaptive activation result." + ), + }, + "base_url": args.base_url, + "model": args.model, + "max_tokens": args.max_tokens, + "temperature": 0.0, + "seed": args.seed, + "ignore_eos": True, + "spec_depth": args.spec_depth, + "prompt_offset": args.prompt_offset, + "prompt_file": str(args.prompt_file.resolve()), + "prompt_file_sha256": digest_bytes(args.prompt_file.read_bytes()), + "server_metadata": metadata, + "server_metadata_sha256": digest_bytes(metadata_bytes), + "server_log": { + "path": str(args.server_log.resolve()), + "start_offset": log_start, + "end_offset": log_end, + "span_bytes": len(log_span), + "span_sha256": digest_bytes(log_span), + }, + "level": level, + "server_records": records, + "validation": validation, + **client_provenance(), + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text( + json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8", + ) + print(markdown(report), end="") + if not validation["passed"]: + for error in validation["errors"]: + print(f"[forced-subset] validation: {error}", file=sys.stderr) + return 0 if validation["passed"] else 1 + + +def main() -> int: + try: + return run(build_parser().parse_args()) + except Exception as exc: + print(f"[forced-subset] error: {exc}", file=sys.stderr) + return 2 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/harness/benchmarks/concurrency/generate_ragged_prompts.py b/harness/benchmarks/concurrency/generate_ragged_prompts.py new file mode 100755 index 000000000..3595dd4d8 --- /dev/null +++ b/harness/benchmarks/concurrency/generate_ragged_prompts.py @@ -0,0 +1,124 @@ +#!/usr/bin/env python3 +"""Generate a small deterministic ragged-prompt manifest for concurrency runs.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Callable, Mapping + + +PROFILES = { + "tiny": (64, 96, 128, 160), + "short": (250, 350, 450, 550), + "medium": (650, 850, 1150, 1350), + "long": (2000, 2600, 3400, 4000), + # Long-context profiles. Runtime telemetry, not word count, proves that + # PFlash compression or KVFlash pressure actually activated. + "compression": (34000, 36000, 38000, 40000), + "kv-pressure": (12000, 14000, 16000, 18000), +} + +WORD_BANK = ( + "systems engineers compare latency throughput scheduling memory kernels queues " + "batches requests tokens caches pages attention arithmetic bandwidth occupancy " + "profiling measurement fairness reproducibility workloads concurrency admission " + "prefill decoding evidence tradeoffs implementation validation production service" +).split() + + +def prompt_text(profile: str, cohort: str, index: int, target_words: int) -> str: + prefix = ( + f"Ragged benchmark {profile} cohort {cohort} request {index}. " + "Write a structured engineering analysis of the following observations, " + "including assumptions, likely bottlenecks, and a concise conclusion." + ).split() + words = list(prefix) + cursor = (index * 7 + target_words) % len(WORD_BANK) + while len(words) < target_words: + words.append(WORD_BANK[cursor % len(WORD_BANK)]) + cursor += 1 + return " ".join(words[:target_words]) + + +ExtraFields = Callable[[str], Mapping[str, object]] + + +def build_profile_records( + profile: str, + profiles: Mapping[str, tuple[int, ...]], + extra_fields: ExtraFields | None = None, +) -> list[dict[str, object]]: + """Build the standard disjoint C1/C4/C8/C16 cohort layout.""" + strata = profiles[profile] + if not strata: + raise ValueError(f"profile {profile!r} has no length strata") + layout = [ + ("c1", [sum(strata) // len(strata)]), + ("c4", list(strata)), + ("c8", list(strata) * 2), + ("c16", list(strata) * 4), + ] + records: list[dict[str, object]] = [] + for cohort, targets in layout: + for target in targets: + index = len(records) + record: dict[str, object] = { + "id": f"{profile}-{index:02d}", + "cohort": cohort, + "stratum": strata.index(target) if target in strata else "mean", + "target_words": target, + "prompt": prompt_text(profile, cohort, index, target), + } + if extra_fields is not None: + additions = dict(extra_fields(profile)) + overlap = record.keys() & additions.keys() + if overlap: + raise ValueError( + "extra profile fields may not replace standard fields: " + f"{sorted(overlap)}" + ) + record.update(additions) + records.append(record) + return records + + +def build_records(profile: str) -> list[dict[str, object]]: + return build_profile_records(profile, PROFILES) + +def build_records(profile: str) -> list[dict[str, object]]: + activation_target = ( + "pflash-auto" if profile == "compression" + else "kvflash-pressure" if profile == "kv-pressure" + else "none" + ) + return build_profile_records( + profile, PROFILES, + lambda _profile: {"activation_target": activation_target}, + ) + if path.exists(): + raise FileExistsError(f"refusing to overwrite {path}") + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + "".join(json.dumps(row, sort_keys=True) + "\n" for row in records), + encoding="utf-8", + ) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--profile", choices=sorted(PROFILES), required=True) + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + records = build_records(args.profile) + try: + write_records(args.out, records) + except FileExistsError as exc: + parser.error(str(exc)) + print(f"wrote {len(records)} prompts to {args.out}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/harness/benchmarks/concurrency/record_feature_runtime.py b/harness/benchmarks/concurrency/record_feature_runtime.py new file mode 100644 index 000000000..afb4fea6e --- /dev/null +++ b/harness/benchmarks/concurrency/record_feature_runtime.py @@ -0,0 +1,119 @@ +#!/usr/bin/env python3 +"""Add startup-observed pool dimensions to one feature-case metadata file.""" + +from __future__ import annotations + +import argparse +import json +import re +import tempfile +from pathlib import Path +from typing import Any + + +KVFLASH_POOL_RE = re.compile( + r"\[parallel-kvflash\] physical resident pool (?P\d+) tokens; " + r"logical per-slot cap (?P\d+) across (?P\d+) slots" +) +PAGED_POOL_RE = re.compile( + r"\[paged-attention\] (?P\d+) physical blocks x " + r"(?P\d+) tokens \((?P\d+) pool tokens, " + r"per-sequence max_ctx (?P\d+)\)" +) + + +def _one_consistent(matches: list[dict[str, int]], label: str) -> dict[str, int] | None: + if not matches: + return None + first = matches[0] + if any(row != first for row in matches[1:]): + raise ValueError(f"conflicting {label} startup markers: {matches}") + return first + + +def observe_startup(log_text: str) -> dict[str, Any]: + kvflash = _one_consistent( + [ + {key: int(value) for key, value in match.groupdict().items()} + for match in KVFLASH_POOL_RE.finditer(log_text) + ], + "KVFlash pool", + ) + paged = _one_consistent( + [ + {key: int(value) for key, value in match.groupdict().items()} + for match in PAGED_POOL_RE.finditer(log_text) + ], + "paged pool", + ) + if paged and paged["blocks"] * paged["block_size"] != paged["tokens"]: + raise ValueError("paged-attention startup marker has inconsistent dimensions") + if kvflash and paged: + if kvflash["tokens"] != paged["tokens"]: + raise ValueError("KVFlash and paged-attention startup pool sizes disagree") + if kvflash["max_ctx"] != paged["max_ctx"]: + raise ValueError("KVFlash and paged-attention logical max_ctx values disagree") + + return { + "kvflash_active": kvflash is not None, + "physical_kv_pool_tokens": ( + kvflash["tokens"] if kvflash else paged["tokens"] if paged else None + ), + "physical_kv_pool_blocks": paged["blocks"] if paged else None, + "kv_block_size_tokens": paged["block_size"] if paged else None, + "logical_per_slot_max_ctx": ( + kvflash["max_ctx"] if kvflash else paged["max_ctx"] if paged else None + ), + "configured_slots": kvflash["slots"] if kvflash else None, + "proof_sources": { + "kvflash_pool_startup_marker": kvflash is not None, + "paged_pool_startup_marker": paged is not None, + }, + } + + +def update_metadata(metadata: dict[str, Any], log_text: str) -> dict[str, Any]: + observed = observe_startup(log_text) + feature_config = metadata.get("feature_config") or {} + kvflash_mode = feature_config.get("kvflash") + kvflash_requested = isinstance(kvflash_mode, str) and kvflash_mode not in ( + "", "off", "0", + ) + if kvflash_requested and not observed["kvflash_active"]: + raise ValueError( + "KVFlash metadata is enabled but its physical-pool startup marker is missing" + ) + if not observed["proof_sources"]["paged_pool_startup_marker"]: + raise ValueError("paged physical-pool startup marker is missing") + + result = dict(metadata) + result["schema_version"] = max(3, int(result.get("schema_version", 0))) + result["runtime_observed"] = observed + return result + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--metadata", type=Path, required=True) + parser.add_argument("--server-log", type=Path, required=True) + args = parser.parse_args() + + metadata = json.loads(args.metadata.read_text(encoding="utf-8")) + updated = update_metadata( + metadata, + args.server_log.read_text(encoding="utf-8", errors="replace"), + ) + # Replace atomically so a killed run never leaves half-written metadata. + with tempfile.NamedTemporaryFile( + "w", encoding="utf-8", dir=args.metadata.parent, + prefix=f".{args.metadata.name}.", delete=False, + ) as handle: + json.dump(updated, handle, indent=2, sort_keys=True) + handle.write("\n") + temporary = Path(handle.name) + temporary.replace(args.metadata) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/harness/benchmarks/concurrency/refill_subset_benchmark.py b/harness/benchmarks/concurrency/refill_subset_benchmark.py new file mode 100644 index 000000000..ba5e3ca2f --- /dev/null +++ b/harness/benchmarks/concurrency/refill_subset_benchmark.py @@ -0,0 +1,1014 @@ +#!/usr/bin/env python3 +"""Sustained/refill AR/speculation/adaptive diagnostic for concurrent DFlash2. + +Each positional lane keeps one request in flight. The first +``clients * (waves - 1)`` completions immediately refill the lane that +completed, preserving each positional prompt and requested mode until the +final C-request drain. This measures a closed-loop saturated service workload. +Forced A/S controls remain supported; adaptive mode additionally fails closed +unless every activation, expected initial route, execution counter, chain +depth, refill recovery, token count, and exact output hash is proven. +""" + + +from __future__ import annotations + +import argparse +import json +import statistics +import sys +import threading +import time +from pathlib import Path +from typing import Any + +import analyze_gate_decisions as gate_analysis +import concurrent_benchmark as base +import forced_subset_benchmark as forced + + +CLIENT_SCRIPT = Path(__file__).resolve() +MODES = ("ar", "speculation", "adaptive") +MODE_CHARS = {"ar": "A", "speculation": "S", "adaptive": "D"} + + +def parse_request_modes(raw: str, clients: int) -> list[str]: + modes = [item.strip() for item in raw.split(",")] + if len(modes) != clients or any(item not in MODES for item in modes): + raise ValueError( + "--request-modes must provide exactly one ar/speculation/adaptive " + "mode per client" + ) + adaptive = [mode == "adaptive" for mode in modes] + if any(adaptive) and not all(adaptive): + raise ValueError( + "adaptive refill must use adaptive mode for every positional lane" + ) + return modes + + +def latency_summary(values: list[float]) -> dict[str, float | int | None]: + return { + "count": len(values), + "median_s": statistics.median(values) if values else None, + "max_s": max(values) if values else None, + } + + +def client_provenance(argv: list[str] | None = None) -> dict[str, Any]: + process_argv = list(sys.orig_argv if argv is None else argv) + if not process_argv or not all(isinstance(item, str) for item in process_argv): + raise ValueError("client process argv must be a non-empty string array") + return { + "client_argv": process_argv, + "client_script": str(CLIENT_SCRIPT), + "client_script_sha256": forced.digest_bytes(CLIENT_SCRIPT.read_bytes()), + } + + +def run_refill( + args: argparse.Namespace, prompts: list[str], modes: list[str], +) -> dict[str, Any]: + """Keep C positional lanes full for a fixed global request budget.""" + selected = base.request_prompts(prompts, args.clients, args.prompt_offset) + barrier = threading.Barrier(args.clients + 1) + lane_records: list[list[dict[str, Any]]] = [ + [] for _ in range(args.clients) + ] + worker_errors: list[BaseException | None] = [None] * args.clients + refill_lock = threading.Lock() + refill_budget = args.clients * (args.waves - 1) + completed_ok = 0 + next_request_index = args.clients + abort_refills = False + + def worker(lane_index: int) -> None: + nonlocal completed_ok, next_request_index, abort_refills + try: + barrier.wait(timeout=min(args.timeout, 60.0)) + previous_end: float | None = None + lane_request_index = 0 + request_index: int | None = lane_index + while request_index is not None: + record = forced.stream_request( + args, selected[lane_index], modes[lane_index], + ) + record["lane_index"] = lane_index + record["lane_request_index"] = lane_request_index + record["request_index"] = request_index + record["admission_group_index"] = request_index // args.clients + record["prompt_index"] = args.prompt_offset + lane_index + record["prompt_sha256"] = base.sha256_text(selected[lane_index]) + record["refill_gap_s"] = ( + float(record["t_start"]) - previous_end + if previous_end is not None else None + ) + lane_records[lane_index].append(record) + previous_end = float(record["t_end"]) + lane_request_index += 1 + with refill_lock: + if record["error"] is not None: + abort_refills = True + request_index = None + else: + completed_ok += 1 + if not abort_refills and completed_ok <= refill_budget: + request_index = next_request_index + next_request_index += 1 + else: + request_index = None + except BaseException as exc: # surfaced in the main thread + with refill_lock: + abort_refills = True + worker_errors[lane_index] = exc + + threads = [ + threading.Thread(target=worker, args=(index,), daemon=True) + for index in range(args.clients) + ] + for thread in threads: + thread.start() + barrier.wait(timeout=min(args.timeout, 60.0)) + barrier_released = time.perf_counter() + deadline = time.monotonic() + args.timeout * args.clients * args.waves + 30.0 + for thread in threads: + thread.join(max(0.0, deadline - time.monotonic())) + hung = sum(thread.is_alive() for thread in threads) + if hung: + raise TimeoutError(f"{hung} refill lane(s) exceeded the workload deadline") + first_worker_error = next((error for error in worker_errors if error), None) + if first_worker_error is not None: + raise RuntimeError(f"refill request worker failed: {first_worker_error}") + + completed = sorted( + (record for lane in lane_records for record in lane), + key=lambda record: int(record["request_index"]), + ) + expected_requests = args.clients * args.waves + starts = [float(record["t_start"]) for record in completed] + ends = [float(record["t_end"]) for record in completed] + if not starts: + raise RuntimeError("no refill request returned a record") + workload_start = min(starts) + workload_end = max(ends) + wall = workload_end - workload_start + for record in completed: + record["start_offset_s"] = float(record["t_start"]) - workload_start + record["end_offset_s"] = float(record["t_end"]) - workload_start + if record["lane_request_index"] == 0: + record["barrier_release_offset_s"] = ( + float(record["t_start"]) - barrier_released + ) + + ok = [record for record in completed if record["error"] is None] + completion = [record["completion_tokens"] for record in ok] + prompt_counts = [record["prompt_tokens"] for record in ok] + complete_tokens = bool(ok) and all(type(value) is int for value in completion) + complete_prompts = bool(ok) and all( + type(value) is int for value in prompt_counts + ) + first_outputs = [ + float(record["t_first"]) for record in ok + if type(record.get("t_first")) in (int, float) + ] + output_wall = ( + workload_end - min(first_outputs) + if len(first_outputs) == len(ok) and first_outputs else None + ) + ttfts = [ + float(record["ttft_s"]) for record in ok + if type(record.get("ttft_s")) in (int, float) + ] + refill_ttfts = [ + float(record["ttft_s"]) for record in ok + if record["lane_request_index"] > 0 + and type(record.get("ttft_s")) in (int, float) + ] + first_decode_gaps = [ + float(record["first_to_second_output_event_s"]) for record in ok + if type(record.get("first_to_second_output_event_s")) in (int, float) + ] + refill_first_decode_gaps = [ + float(record["first_to_second_output_event_s"]) for record in ok + if record["lane_request_index"] > 0 + and type(record.get("first_to_second_output_event_s")) in (int, float) + ] + first_wave_starts = [ + float(lane[0]["t_start"]) for lane in lane_records if lane + ] + refill_gaps = [ + float(record["refill_gap_s"]) + for record in completed + if type(record.get("refill_gap_s")) in (int, float) + ] + output_hashes = [ + [ + record["request_index"], record["lane_index"], + record["lane_request_index"], + record["content_sha256"], record["reasoning_content_sha256"], + ] + for record in completed + ] + lane_summaries = [] + for lane_index, lane in enumerate(lane_records): + lane_completion = [ + record["completion_tokens"] for record in lane + if type(record.get("completion_tokens")) is int + ] + signatures = { + ( + record.get("content_sha256"), + record.get("reasoning_content_sha256"), + ) + for record in lane + } + lane_summaries.append({ + "lane_index": lane_index, + "decode_mode": modes[lane_index], + "prompt_index": args.prompt_offset + lane_index, + "requests": len(lane), + "requests_ok": sum(record["error"] is None for record in lane), + "completion_tokens_total": sum(lane_completion), + "exact_output_stable": bool(lane) and len(signatures) == 1, + "ordered_output_set_sha256": forced.canonical_digest([ + [record["content_sha256"], record["reasoning_content_sha256"]] + for record in lane + ]), + }) + return { + "clients": args.clients, + "waves": args.waves, + "scheduled_refills": refill_budget, + "terminal_guard_requests": args.clients, + "request_modes": modes, + "request_mode_mask": "".join( + MODE_CHARS[mode] for mode in modes + ), + "expected_requests": expected_requests, + "requests": len(completed), + "requests_ok": len(ok), + "failures": len(completed) - len(ok), + "missing_requests": expected_requests - len(completed), + "wall_s": wall, + "initial_start_skew_s": ( + max(first_wave_starts) - min(first_wave_starts) + if len(first_wave_starts) == args.clients else None + ), + "completion_tokens_total": sum(completion) if complete_tokens else None, + "token_count_complete": complete_tokens, + "prompt_token_count_complete": complete_prompts, + "fixed_token_workload_valid": ( + len(ok) == expected_requests and complete_tokens + and all(value == args.max_tokens for value in completion) + ), + "aggregate_refill_tok_s": ( + sum(completion) / wall if complete_tokens and wall > 0 else None + ), + "output_window_s": output_wall, + "output_window_tok_s": ( + sum(completion) / output_wall + if complete_tokens and output_wall is not None and output_wall > 0 + else None + ), + "ttft": latency_summary(ttfts), + "refill_ttft": latency_summary(refill_ttfts), + "first_to_second_output_event": latency_summary(first_decode_gaps), + "refill_first_to_second_output_event": latency_summary(refill_first_decode_gaps), + "prompt_tokens_total": sum(prompt_counts) if complete_prompts else None, + "refill_handoffs": len(refill_gaps), + "refill_gap_s_median": ( + statistics.median(refill_gaps) if refill_gaps else None + ), + "refill_gap_s_max": max(refill_gaps) if refill_gaps else None, + "selected_prompt_set_sha256": forced.canonical_digest([ + base.sha256_text(prompt) for prompt in selected + ]), + "ordered_output_set_sha256": forced.canonical_digest(output_hashes), + "exact_output_stable_per_lane": all( + lane["exact_output_stable"] for lane in lane_summaries + ), + "lanes": lane_summaries, + "requests_detail": completed, + } + + +def _full_live_round_summary( + rounds: list[dict[str, Any]], clients: int, +) -> tuple[dict[str, Any], list[str]]: + errors: list[str] = [] + full_live: list[dict[str, Any]] = [] + all_timed_us = 0.0 + all_emitted = 0 + longest = 0 + current = 0 + for wrapped in rounds: + row = wrapped["record"] + live = row.get("live") + if type(live) is not int or live < 1 or live > clients: + errors.append(f"step-timing line {wrapped['line_index']} has invalid live={live!r}") + current = 0 + continue + total_us = row.get("total_us") + emitted = row.get("emitted_tokens") + if ( + type(total_us) not in (int, float) or total_us <= 0 + or type(emitted) is not int or emitted < 1 + ): + errors.append( + f"step-timing line {wrapped['line_index']} lacks positive " + "total_us/emitted_tokens" + ) + current = 0 + continue + all_timed_us += float(total_us) + all_emitted += emitted + if live == clients: + full_live.append(wrapped) + current += 1 + longest = max(longest, current) + else: + current = 0 + full_us = sum(float(row["record"]["total_us"]) for row in full_live) + full_emitted = sum(int(row["record"]["emitted_tokens"]) for row in full_live) + return ({ + "rounds": len(full_live), + "longest_streak": longest, + "timed_us": full_us, + "emitted_tokens": full_emitted, + "engine_round_goodput_tok_s": ( + full_emitted * 1_000_000.0 / full_us if full_us > 0 else None + ), + "timed_fraction": full_us / all_timed_us if all_timed_us > 0 else None, + "emitted_fraction": ( + full_emitted / all_emitted if all_emitted > 0 else None + ), + "first_line_index": full_live[0]["line_index"] if full_live else None, + "last_line_index": full_live[-1]["line_index"] if full_live else None, + "path_counts": { + path: sum(row["record"].get("path") == path for row in full_live) + for path in ("ar", "spec", "spec-direct") + }, + }, errors) + + +def validate_adaptive_activations( + records: dict[str, list[dict[str, Any]]], + requests: list[dict[str, Any]], + metrics_by_id: dict[str, dict[str, Any]], + expected_mask: str, + errors: list[str], +) -> dict[str, Any]: + direct_executor = any( + wrapped["record"].get("path") == "spec-direct" + for wrapped in records["rounds"] + ) + by_engine: dict[int, list[dict[str, Any]]] = {} + valid_rows: list[dict[str, Any]] = [] + for wrapped in records["activations"]: + row = wrapped["record"] + try: + gate_analysis._validate_activation( # shared fail-closed schema + row, Path(""), wrapped["line_index"], + ) + except ValueError as exc: + errors.append(str(exc)) + continue + engine_id = int(row["request_id"]) + by_engine.setdefault(engine_id, []).append(row) + valid_rows.append(row) + + expected_engine_ids: set[int] = set() + initial_routes: dict[int, str] = {} + matched = 0 + for request in requests: + wire_id = request.get("request_id") + metric = metrics_by_id.get(wire_id) + if metric is None: + continue + engine_id = metric.get("engine_request_id") + if type(engine_id) is not int or engine_id < 0: + continue + expected_engine_ids.add(engine_id) + rows = by_engine.get(engine_id, []) + if len(rows) != 1: + errors.append( + f"{wire_id}: expected exactly one adaptive activation for " + f"engine request {engine_id}, got {len(rows)}" + ) + continue + matched += 1 + activation = rows[0] + lane = int(request["lane_index"]) + if activation["evaluation"] != "scored": + errors.append( + f"{wire_id}: adaptive activation evaluation was " + f"{activation['evaluation']!r}, expected 'scored'" + ) + steps = metric.get("spec_steps") + service_steps = metric.get("spec_service_ar_steps", 0) + target_forwards = metric.get("target_forwards") + if type(steps) is not int or steps < 0: + errors.append(f"{wire_id}: adaptive request has invalid spec_steps") + elif type(service_steps) is not int or service_steps < 0: + errors.append( + f"{wire_id}: adaptive request has invalid " + "spec_service_ar_steps" + ) + elif activation["decision"] == "ar" and ( + steps != 0 or service_steps != 0 + ): + errors.append( + f"{wire_id}: AR activation executed spec_steps={steps}, " + f"spec_service_ar_steps={service_steps}" + ) + elif activation["decision"] == "speculation" and ( + steps == 0 + or type(target_forwards) is not int + or target_forwards < + (1 if direct_executor else 2) * steps + service_steps + ): + errors.append( + f"{wire_id}: speculation activation has invalid execution " + f"counters spec_steps={steps!r}, " + f"spec_service_ar_steps={service_steps!r}, " + f"target_forwards={target_forwards!r}" + ) + if request.get("lane_request_index") == 0: + initial_routes[lane] = ( + "S" if activation["decision"] == "speculation" else "A" + ) + expected_decision = ( + "speculation" if expected_mask[lane] == "S" else "ar" + ) + if activation["decision"] != expected_decision: + errors.append( + f"{wire_id}: initial adaptive decision " + f"{activation['decision']!r} does not match lane {lane} " + f"expected {expected_decision!r}" + ) + + duplicate_ids = sorted( + engine_id for engine_id, rows in by_engine.items() if len(rows) != 1 + ) + unknown_ids = sorted(set(by_engine) - expected_engine_ids) + missing_ids = sorted(expected_engine_ids - set(by_engine)) + if duplicate_ids: + errors.append( + "duplicate adaptive activations for engine requests " + + ",".join(str(value) for value in duplicate_ids) + ) + if unknown_ids: + errors.append( + "adaptive activations reference unknown engine requests " + + ",".join(str(value) for value in unknown_ids) + ) + if missing_ids: + errors.append( + "missing adaptive activations for engine requests " + + ",".join(str(value) for value in missing_ids) + ) + + decision_counts = { + decision: sum(row["decision"] == decision for row in valid_rows) + for decision in ("ar", "speculation") + } + evaluation_counts = { + evaluation: sum(row["evaluation"] == evaluation for row in valid_rows) + for evaluation in ("scored", "failed") + } + score_kinds = sorted({str(row["score_kind"]) for row in valid_rows}) + return { + "required": True, + "records": len(records["activations"]), + "valid_records": len(valid_rows), + "matched_requests": matched, + "decision_counts": decision_counts, + "evaluation_counts": evaluation_counts, + "score_kinds": score_kinds, + "expected_initial_route_mask": expected_mask, + "observed_initial_route_mask": "".join( + initial_routes.get(lane, "?") for lane in range(len(expected_mask)) + ), + } + + +def validate_evidence( + workload: dict[str, Any], records: dict[str, list[dict[str, Any]]], + clients: int, modes: list[str], spec_depth: int, waves: int, + max_start_skew_ms: float, max_refill_gap_ms: float, + min_full_live_rounds: int, + expected_adaptive_mask: str | None = None, + adaptive_scoring_enabled: bool = True, +) -> dict[str, Any]: + errors: list[str] = [] + adaptive_requested = all(mode == "adaptive" for mode in modes) + adaptive_activation_required = adaptive_requested and adaptive_scoring_enabled + expected_requests = clients * waves + adaptive_mask_valid = ( + isinstance(expected_adaptive_mask, str) + and len(expected_adaptive_mask) == clients + and set(expected_adaptive_mask) <= {"A", "S"} + ) + if adaptive_requested and not adaptive_mask_valid: + errors.append( + "adaptive refill requires an expected A/S route for every lane" + ) + elif not adaptive_requested and expected_adaptive_mask is not None: + errors.append("expected adaptive mask supplied for a forced refill") + if ( + adaptive_requested and not adaptive_scoring_enabled + and expected_adaptive_mask != "A" * clients + ): + errors.append("adaptive scoring-off control must expect all-AR routing") + if workload["requests"] != expected_requests or workload["missing_requests"] != 0: + errors.append( + f"completed request records {workload['requests']} do not match " + f"clients*waves={expected_requests}" + ) + if workload["failures"] or workload["requests_ok"] != expected_requests: + errors.append("one or more refill requests failed") + if not workload["token_count_complete"] or not workload["prompt_token_count_complete"]: + errors.append("wire token accounting is incomplete") + if workload["fixed_token_workload_valid"] is not True: + errors.append("ignore-eos fixed-token refill workload was not completed exactly") + start_skew = workload.get("initial_start_skew_s") + if type(start_skew) not in (int, float): + errors.append("initial synchronized request start skew is unavailable") + elif float(start_skew) * 1000.0 > max_start_skew_ms: + errors.append(f"initial request start skew exceeds {max_start_skew_ms:g} ms") + expected_handoffs = clients * (waves - 1) + if workload.get("refill_handoffs") != expected_handoffs: + errors.append( + f"refill handoff count {workload.get('refill_handoffs')} does not " + f"match {expected_handoffs}" + ) + max_gap = workload.get("refill_gap_s_max") + if type(max_gap) not in (int, float): + errors.append("refill handoff latency is unavailable") + elif float(max_gap) < 0: + errors.append("a refill request started before its predecessor completed") + elif float(max_gap) * 1000.0 > max_refill_gap_ms: + errors.append(f"refill handoff gap exceeds {max_refill_gap_ms:g} ms") + if workload.get("exact_output_stable_per_lane") is not True: + errors.append("deterministic exact output hashes changed across refill waves") + + requests = workload["requests_detail"] + request_ids = [row.get("request_id") for row in requests] + if ( + any(not isinstance(value, str) or not value for value in request_ids) + or len(set(request_ids)) != expected_requests + ): + errors.append("wire request IDs are missing or not unique") + for lane_index, lane in enumerate(workload.get("lanes") or []): + if lane.get("lane_index") != lane_index or not lane.get("requests"): + errors.append(f"lane {lane_index} did not execute an initial request") + if lane.get("decode_mode") != modes[lane_index]: + errors.append(f"lane {lane_index} mode changed during refill") + output_hashes_complete = all( + isinstance(request.get(key), str) and len(request[key]) == 64 + for request in requests + for key in ( + "content_sha256", "reasoning_content_sha256", + "combined_output_sha256", + ) + ) + if not output_hashes_complete: + errors.append("exact request output hashes are incomplete") + + rounds = records["rounds"] + direct_executor = any( + wrapped["record"].get("path") == "spec-direct" + for wrapped in rounds + ) + full_live, timing_errors = _full_live_round_summary(rounds, clients) + errors.extend(timing_errors) + if full_live["longest_streak"] < min_full_live_rounds: + errors.append( + f"sustained live=C proof absent: longest live={clients} streak is " + f"{full_live['longest_streak']}, need {min_full_live_rounds}" + ) + + spec_requested = ( + any(mode == "speculation" for mode in modes) + or adaptive_activation_required + and expected_adaptive_mask is not None + and "S" in expected_adaptive_mask + ) + spec_rounds = [ + wrapped for wrapped in rounds + if wrapped["record"].get("path") in ("spec", "spec-direct") + and type(wrapped["record"].get("k")) is int + and wrapped["record"]["k"] > 0 + ] + inferred_depths: list[int] = [] + for wrapped in spec_rounds: + row = wrapped["record"] + bucket = row.get("tree_bucket") + tree_rows = row.get("tree_rows") + ar_lanes = row.get("ar_lanes", 0) + spec_rows = ( + tree_rows - ar_lanes + if row.get("path") == "spec-direct" + and type(tree_rows) is int and type(ar_lanes) is int + else tree_rows + ) + if ( + type(bucket) is not int or bucket <= 0 + or type(tree_rows) is not int or tree_rows <= 0 + or type(spec_rows) is not int or spec_rows <= 0 + or spec_rows % bucket != 0 + ): + errors.append("spec step-timing lacks a valid tree_rows/tree_bucket shape") + continue + inferred_depths.append(spec_rows // bucket) + if spec_requested: + if not spec_rounds: + errors.append("speculation was requested but no spec round executed") + wrong_depths = sorted(set(depth for depth in inferred_depths if depth != spec_depth)) + if wrong_depths: + errors.append( + f"executed chain depths {wrong_depths} do not match requested " + f"depth {spec_depth}" + ) + elif spec_rounds: + errors.append("all-AR refill mask unexpectedly executed speculation") + + metrics_by_id: dict[str, dict[str, Any]] = {} + metric_wrappers_by_id: dict[str, dict[str, Any]] = {} + for wrapped in records["requests"]: + row = wrapped["record"] + request_id = row.get("request_id") + if not isinstance(request_id, str) or not request_id: + errors.append("concurrency metric has no wire request_id") + elif request_id in metrics_by_id: + errors.append(f"duplicate concurrency metric for {request_id}") + else: + metrics_by_id[request_id] = row + metric_wrappers_by_id[request_id] = wrapped + missing = sorted(set(request_ids) - set(metrics_by_id)) + extra = sorted(set(metrics_by_id) - set(request_ids)) + if missing: + errors.append(f"missing per-request concurrency metrics: {missing}") + if extra: + errors.append(f"unmatched per-request concurrency metrics: {extra}") + + selector_engine_ids = { + wrapped["record"].get("request_id") for wrapped in records["selectors"] + if type(wrapped["record"].get("request_id")) is int + } + for request in requests: + request_id = request.get("request_id") + metric = metrics_by_id.get(request_id) + if metric is None: + continue + expected_mode = modes[int(request["lane_index"])] + steps = metric.get("spec_steps") + service_steps = metric.get("spec_service_ar_steps", 0) + target_forwards = metric.get("target_forwards") + if type(steps) is not int or steps < 0: + errors.append(f"{request_id}: invalid spec_steps") + elif type(service_steps) is not int or service_steps < 0: + errors.append(f"{request_id}: invalid spec_service_ar_steps") + elif expected_mode == "speculation" and steps == 0: + errors.append(f"{request_id}: forced speculation never executed") + elif expected_mode == "speculation" and ( + type(target_forwards) is not int + or target_forwards < + (1 if direct_executor else 2) * steps + service_steps + ): + errors.append( + f"{request_id}: forced speculation has invalid execution " + f"counters spec_steps={steps!r}, " + f"spec_service_ar_steps={service_steps!r}, " + f"target_forwards={target_forwards!r}" + ) + elif expected_mode == "ar" and (steps != 0 or service_steps != 0): + errors.append( + f"{request_id}: forced AR executed speculation/service" + ) + elif ( + expected_mode == "adaptive" + and not adaptive_scoring_enabled + and (steps != 0 or service_steps != 0) + ): + errors.append( + f"{request_id}: scoring-off control executed speculation/service" + ) + engine_id = metric.get("engine_request_id") + if type(engine_id) is not int: + errors.append(f"{request_id}: missing integer engine_request_id") + elif expected_mode == "speculation" and engine_id not in selector_engine_ids: + errors.append(f"{request_id}: no DFlash2 selector evidence") + elif expected_mode == "ar" and engine_id in selector_engine_ids: + errors.append(f"{request_id}: forced AR emitted DFlash2 selector evidence") + + activation_summary: dict[str, Any] = { + "required": adaptive_activation_required, + "records": len(records["activations"]), + } + if adaptive_activation_required and adaptive_mask_valid: + assert expected_adaptive_mask is not None + activation_summary = validate_adaptive_activations( + records, requests, metrics_by_id, expected_adaptive_mask, errors, + ) + elif adaptive_activation_required: + activation_summary["validation"] = "invalid-expected-mask" + elif adaptive_requested: + activation_summary.update({ + "scoring_disabled": True, + "expected_initial_route_mask": expected_adaptive_mask, + "observed_initial_route_mask": "A" * clients, + }) + if records["activations"]: + errors.append( + "adaptive scoring-off control emitted activation records" + ) + elif not adaptive_requested and records["activations"]: + errors.append( + "forced refill unexpectedly emitted adaptive activation records" + ) + + last_full_live_line = full_live["last_line_index"] + completed_before_last_full_live = ( + sum( + wrapped["line_index"] < last_full_live_line + for wrapped in metric_wrappers_by_id.values() + ) + if type(last_full_live_line) is int else 0 + ) + scheduled_refills = expected_requests - clients + required_refill_recoveries = expected_requests - 2 * clients + refill_recovery_proved = ( + completed_before_last_full_live >= required_refill_recoveries + ) + if not refill_recovery_proved: + errors.append( + "full live=C was not recovered inside the guarded refill window: " + f"completions: proved {completed_before_last_full_live}, need " + f"{required_refill_recoveries}" + ) + return { + "passed": not errors, + "errors": errors, + "adaptive_claims_permitted": ( + adaptive_activation_required and not errors + ), + "adaptive_stack_control_permitted": ( + adaptive_requested + and not adaptive_scoring_enabled + and not errors + ), + "activation": activation_summary, + "closed_cohort_claims_permitted": False, + "full_live": full_live, + "min_full_live_rounds": min_full_live_rounds, + "executed_spec_depths": sorted(set(inferred_depths)), + "round_records": len(rounds), + "selector_records": len(records["selectors"]), + "request_metric_records": len(records["requests"]), + "required_refill_recoveries": required_refill_recoveries, + "scheduled_refills": scheduled_refills, + "terminal_guard_requests": clients, + "completed_before_last_full_live": completed_before_last_full_live, + "refill_recovery_proved": refill_recovery_proved, + } + + +def markdown(report: dict[str, Any]) -> str: + workload = report["workload"] + validation = report["validation"] + full_live = validation["full_live"] + status = "PASS" if validation["passed"] else "FAIL" + adaptive_requests = bool(report["scope"].get("adaptive_requests")) + adaptive_evaluation = bool(report["scope"]["adaptive_evaluation"]) + if adaptive_evaluation: + title = "Adaptive" + scope = "This is a fail-closed adaptive closed-loop refill diagnostic." + elif adaptive_requests: + title = "Adaptive OFF" + scope = ( + "This is a matched-stack adaptive-scoring-off refill control." + ) + else: + title = "Forced" + scope = "This is a forced closed-loop refill diagnostic." + activation = validation["activation"] + route_mask = ( + activation.get("observed_initial_route_mask") + if adaptive_requests else workload["request_mode_mask"] + ) + refill_ttft = workload["refill_ttft"] + refill_decode_gap = workload["refill_first_to_second_output_event"] + return ( + f"# {title} DFlash2 refill diagnostic — {report['label']}\n\n" + f"{scope} It makes no closed-cohort makespan claim.\n\n" + "| C | Route | Waves | Depth | Requests | Refill tok/s | " + "Full-live tok/s | Refill TTFT med/max s | " + "First-output gap med/max s | Recovery | Status |\n" + "| ---: | :--- | ---: | ---: | ---: | ---: | ---: | :--- | :--- | :--- | :--- |\n" + f"| {workload['clients']} | {route_mask} | " + f"{workload['waves']} | {report['spec_depth']} | " + f"{workload['requests_ok']}/{workload['expected_requests']} | " + f"{base.fmt(workload['aggregate_refill_tok_s'])} | " + f"{base.fmt(full_live['engine_round_goodput_tok_s'])} | " + f"{base.fmt(refill_ttft['median_s'])}/{base.fmt(refill_ttft['max_s'])} | " + f"{base.fmt(refill_decode_gap['median_s'])}/" + f"{base.fmt(refill_decode_gap['max_s'])} | " + f"{validation['completed_before_last_full_live']}/" + f"{validation['required_refill_recoveries']} | {status} |\n" + ) + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base-url", default="http://127.0.0.1:18080/v1") + parser.add_argument("--api-key", default="") + parser.add_argument("--model", default="luce-dflash") + parser.add_argument("--clients", type=int, required=True) + parser.add_argument("--request-modes", required=True) + parser.add_argument( + "--expected-adaptive-mask", + help=( + "Required initial per-lane A/S route for adaptive requests" + ), + ) + parser.add_argument( + "--adaptive-scoring-disabled", + action="store_true", + help=( + "Validate the DFLASH_SPEC_ACTIVATION_SCORE=0 matched-stack control" + ), + ) + parser.add_argument("--waves", type=int, required=True) + parser.add_argument("--spec-depth", type=int, required=True) + parser.add_argument("--prompt-file", type=Path, required=True) + parser.add_argument("--prompt-offset", type=int, default=0) + parser.add_argument("--max-tokens", type=int, default=256) + parser.add_argument("--seed", type=int, default=1) + parser.add_argument("--timeout", type=float, default=1800.0) + parser.add_argument("--max-start-skew-ms", type=float, default=100.0) + parser.add_argument("--max-refill-gap-ms", type=float, default=100.0) + parser.add_argument("--min-full-live-rounds", type=int, default=2) + parser.add_argument("--log-settle-ms", type=float, default=100.0) + parser.add_argument("--server-metadata-json", type=Path, required=True) + parser.add_argument("--server-log", type=Path, required=True) + parser.add_argument("--out", type=Path, required=True) + parser.add_argument("--label", default="") + return parser + + +def run(args: argparse.Namespace) -> int: + if args.clients < 1: + raise ValueError("--clients must be positive") + if args.waves < 3: + raise ValueError("--waves must be at least 3 for a guarded refill workload") + if args.spec_depth < 2: + raise ValueError("--spec-depth must be at least 2; depth 1 is AR-equivalent") + if args.prompt_offset < 0 or args.max_tokens < 1 or args.timeout <= 0: + raise ValueError("invalid prompt offset, max tokens, or timeout") + if ( + args.max_start_skew_ms < 0 or args.max_refill_gap_ms < 0 + or args.min_full_live_rounds < 1 + ): + raise ValueError("invalid synchronization/refill proof threshold") + if args.log_settle_ms < 0: + raise ValueError("--log-settle-ms must be non-negative") + modes = parse_request_modes(args.request_modes, args.clients) + adaptive_requested = all(mode == "adaptive" for mode in modes) + adaptive_scoring_enabled = not bool( + getattr(args, "adaptive_scoring_disabled", False) + ) + if not adaptive_requested and not adaptive_scoring_enabled: + raise ValueError( + "--adaptive-scoring-disabled requires adaptive request modes" + ) + raw_expected_mask = getattr(args, "expected_adaptive_mask", None) + expected_adaptive_mask = ( + raw_expected_mask.strip().upper() + if isinstance(raw_expected_mask, str) and raw_expected_mask.strip() + else None + ) + if adaptive_requested and ( + expected_adaptive_mask is None + or len(expected_adaptive_mask) != args.clients + or set(expected_adaptive_mask) - {"A", "S"} + ): + raise ValueError( + "--expected-adaptive-mask must provide one A/S route per client" + ) + if not adaptive_requested and expected_adaptive_mask is not None: + raise ValueError("--expected-adaptive-mask requires adaptive request modes") + prompts = base.load_prompts(args.prompt_file) + metadata_bytes = args.server_metadata_json.read_bytes() + metadata = json.loads(metadata_bytes) + if not isinstance(metadata, dict): + raise ValueError("server metadata must be a JSON object") + forced.validate_server_metadata( + metadata, args.clients, args.spec_depth, args.prompt_offset, + require_selector=any( + mode in ("speculation", "adaptive") for mode in modes + ), + ) + declared_decode_mode = (metadata.get("feature_config") or {}).get("decode_mode") + if adaptive_requested and declared_decode_mode != "adaptive": + raise ValueError( + "adaptive refill requires server metadata decode_mode=adaptive" + ) + if adaptive_requested: + expected_score_switch = "1" if adaptive_scoring_enabled else "0" + actual_score_switch = (metadata.get("launch_environment") or {}).get( + "DFLASH_SPEC_ACTIVATION_SCORE" + ) + if actual_score_switch != expected_score_switch: + raise ValueError( + "adaptive refill metadata must record " + f"DFLASH_SPEC_ACTIVATION_SCORE={expected_score_switch}" + ) + log_start = args.server_log.stat().st_size + workload = run_refill(args, prompts, modes) + if args.log_settle_ms: + time.sleep(args.log_settle_ms / 1000.0) + log_span, log_end = forced.read_log_span(args.server_log, log_start) + records = forced.parse_profile_records(log_span) + validation = validate_evidence( + workload, records, args.clients, modes, args.spec_depth, args.waves, + args.max_start_skew_ms, args.max_refill_gap_ms, + args.min_full_live_rounds, expected_adaptive_mask, + adaptive_scoring_enabled, + ) + report = { + "schema_version": 2, + "kind": ( + "dflash2-adaptive-refill-diagnostic" + if adaptive_requested and adaptive_scoring_enabled else + "dflash2-adaptive-scoring-off-refill-control" + if adaptive_requested else + "dflash2-forced-refill-diagnostic" + ), + "label": args.label, + "scope": { + "forced_controls_only": not adaptive_requested, + "adaptive_requests": adaptive_requested, + "adaptive_evaluation": ( + adaptive_requested and adaptive_scoring_enabled + ), + "adaptive_stack_control": ( + adaptive_requested and not adaptive_scoring_enabled + ), + "closed_cohort_makespan": False, + "interpretation": ( + "aggregate_refill_tok_s is end-to-end closed-loop goodput " + "across persistent client lanes and includes startup, request " + "handoffs, prefill, and terminal drain. full_live decode-round " + "goodput uses only timed live=C decode rounds and excludes " + "handoff/prefill gaps. Adaptive claims require fail-closed " + "activation validation. Neither metric is a single " + "closed-cohort makespan result." + ), + }, + "base_url": args.base_url, + "model": args.model, + "max_tokens": args.max_tokens, + "temperature": 0.0, + "seed": args.seed, + "ignore_eos": True, + "spec_depth": args.spec_depth, + "expected_adaptive_mask": expected_adaptive_mask, + "prompt_offset": args.prompt_offset, + "prompt_file": str(args.prompt_file.resolve()), + "prompt_file_sha256": forced.digest_bytes(args.prompt_file.read_bytes()), + "server_metadata": metadata, + "server_metadata_sha256": forced.digest_bytes(metadata_bytes), + "server_log": { + "path": str(args.server_log.resolve()), + "start_offset": log_start, + "end_offset": log_end, + "span_bytes": len(log_span), + "span_sha256": forced.digest_bytes(log_span), + }, + "workload": workload, + "server_records": records, + "validation": validation, + **client_provenance(), + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text( + json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8", + ) + print(markdown(report), end="") + if not validation["passed"]: + for error in validation["errors"]: + print(f"[forced-refill] validation: {error}", file=sys.stderr) + return 0 if validation["passed"] else 1 + + +def main() -> int: + try: + return run(build_parser().parse_args()) + except Exception as exc: + print(f"[forced-refill] error: {exc}", file=sys.stderr) + return 2 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/harness/benchmarks/concurrency/run_qwen36_concurrency.sh b/harness/benchmarks/concurrency/run_qwen36_concurrency.sh new file mode 100755 index 000000000..5583bacf4 --- /dev/null +++ b/harness/benchmarks/concurrency/run_qwen36_concurrency.sh @@ -0,0 +1,274 @@ +#!/usr/bin/env bash +# Paired, fresh-process Qwen3.6 concurrency benchmark for Lucebox and llama.cpp. +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)" +REPO="${REPO:-$(cd -- "$SCRIPT_DIR/../../.." && pwd -P)}" +CLIENT="${CLIENT:-$SCRIPT_DIR/concurrent_benchmark.py}" +GENERATOR="${GENERATOR:-$SCRIPT_DIR/generate_ragged_prompts.py}" +SUMMARIZER="${SUMMARIZER:-$SCRIPT_DIR/summarize_concurrency.py}" + +MODEL="${MODEL:-}" +LUCE_SERVER_BIN="${LUCE_SERVER_BIN:-$REPO/server/build-hip/dflash_server}" +LLAMA_SERVER_BIN="${LLAMA_SERVER_BIN:-$(command -v llama-server 2>/dev/null || true)}" +OUT="${OUT:-$REPO/.harness-runs/qwen36-concurrency-$(date -u +%Y%m%dT%H%M%SZ)}" +REPEATS="${REPEATS:-1}" +WORKLOADS="${WORKLOADS:-short,medium,long}" +VARIANTS="${VARIANTS:-luce-k8,luce-k1,llama}" +CLIENTS="${CLIENTS:-1,4,8,16}" +PORT="${PORT:-18114}" +COOLDOWN_SECONDS="${COOLDOWN_SECONDS:-3}" +HEALTH_TIMEOUT_SECONDS="${HEALTH_TIMEOUT_SECONDS:-600}" +MAX_TOKENS="${MAX_TOKENS:-64}" +WARMUP_TOKENS="${WARMUP_TOKENS:-8}" +SLOTS=16 + +usage() { + cat <<'EOF' +Usage: MODEL=/path/model.gguf [REPEATS=5] run_qwen36_concurrency.sh + +Runs fresh-server, same-concurrency warmup + measurement cases for luce-k8, +luce-k1, and llama at C=1/4/8/16. Defaults to one repeat for screening; use at +least five paired repeats for publication. For a decode-heavy comparison, set +WORKLOADS=short MAX_TOKENS=256 VARIANTS=luce-k8,llama. OUT must not already +exist. +EOF +} + +if [[ "${1:-}" == "--help" ]]; then usage; exit 0; fi +if [[ $# -ne 0 ]]; then usage >&2; exit 2; fi +for cmd in python3 curl sha256sum; do command -v "$cmd" >/dev/null || { echo "missing $cmd" >&2; exit 2; }; done +[[ -r "$MODEL" ]] || { echo "set MODEL to a readable GGUF" >&2; exit 2; } +[[ -x "$LUCE_SERVER_BIN" ]] || { echo "missing Lucebox server: $LUCE_SERVER_BIN" >&2; exit 2; } +if [[ ",$VARIANTS," == *,llama,* ]]; then + [[ -x "$LLAMA_SERVER_BIN" ]] || { echo "missing llama.cpp server: $LLAMA_SERVER_BIN" >&2; exit 2; } +fi +[[ "$REPEATS" =~ ^[1-9][0-9]*$ ]] || { echo "REPEATS must be positive" >&2; exit 2; } +[[ ! -e "$OUT" ]] || { echo "refusing to overwrite $OUT" >&2; exit 2; } +ambient_tuning="$(env | grep -E '^(GGML_|DFLASH_|LUCE_|HIP_|ROCR_|HSA_|LD_PRELOAD=|LD_LIBRARY_PATH=)' \ + | grep -v '^LUCE_SERVER_BIN=' || true)" +if [[ -n "$ambient_tuning" ]]; then + echo "refusing ambient GPU/backend tuning variables:" >&2 + echo "$ambient_tuning" >&2 + exit 2 +fi +MODEL_SHA256="$(sha256sum "$MODEL" | awk '{print $1}')" + +IFS=, read -r -a workload_list <<< "$WORKLOADS" +IFS=, read -r -a variant_list <<< "$VARIANTS" +IFS=, read -r -a client_list <<< "$CLIENTS" +reject_duplicates() { + local list_name="$1" value + shift + local -A seen=() + for value in "$@"; do + if [[ -n "${seen[$value]+yes}" ]]; then + echo "$list_name contains duplicate entry: $value" >&2 + return 1 + fi + seen["$value"]=1 + done +} +reject_duplicates CLIENTS "${client_list[@]}" || exit 2 +reject_duplicates VARIANTS "${variant_list[@]}" || exit 2 +declare -A prompt_offsets=([1]=0 [4]=1 [8]=5 [16]=13) +for c in "${client_list[@]}"; do + [[ -n "${prompt_offsets[$c]+yes}" ]] || { echo "supported CLIENTS are 1,4,8,16" >&2; exit 2; } +done +for v in "${variant_list[@]}"; do + [[ "$v" == luce-k8 || "$v" == luce-k1 || "$v" == luce-k16-b2 || + "$v" == luce-k16-b4 || "$v" == luce-k16-dyn || + "$v" == llama ]] || { echo "unknown variant $v" >&2; exit 2; } +done + +mkdir -p "$OUT/prompts" +for workload in "${workload_list[@]}"; do + python3 "$GENERATOR" --profile "$workload" --out "$OUT/prompts/$workload.jsonl" +done + +server_pid="" +stop_server() { + if [[ -n "$server_pid" ]] && kill -0 "$server_pid" 2>/dev/null; then + kill "$server_pid" 2>/dev/null || true + for _ in $(seq 1 30); do + kill -0 "$server_pid" 2>/dev/null || break + sleep 1 + done + kill -9 "$server_pid" 2>/dev/null || true + wait "$server_pid" 2>/dev/null || true + fi + server_pid="" +} +trap stop_server EXIT +trap 'exit 130' INT TERM + +served_model_matches() { + python3 - "$PORT" "$1" <<'PY' +import json +import sys +import urllib.request + +port, expected = sys.argv[1:] +with urllib.request.urlopen( + f"http://127.0.0.1:{port}/v1/models", timeout=2, +) as response: + payload = json.load(response) +matches = any( + isinstance(row, dict) and row.get("id") == expected + for row in payload.get("data", []) +) +raise SystemExit(0 if matches else 1) +PY +} + +wait_health() { + local expected_model="$1" + local deadline=$((SECONDS + HEALTH_TIMEOUT_SECONDS)) + while (( SECONDS < deadline )); do + kill -0 "$server_pid" 2>/dev/null || return 1 + if curl -fsS --max-time 2 "http://127.0.0.1:$PORT/health" >/dev/null 2>&1 && + served_model_matches "$expected_model" >/dev/null 2>&1; then + return 0 + fi + sleep 1 + done + return 1 +} + +port_is_available() { + python3 - "$PORT" <<'PY' +import socket +import sys + +port = int(sys.argv[1]) +with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + try: + sock.bind(("127.0.0.1", port)) + except OSError as exc: + print(f"PORT {port} is unavailable: {exc}", file=sys.stderr) + raise SystemExit(1) +PY +} + +write_metadata() { + local path="$1" variant="$2" workload="$3" clients="$4" repeat="$5" binary="$6" max_prefills="$7" mixed_budget="$8" idle_budget="$9" quantum="${10}" command_file="${11}" + python3 -c 'import hashlib,json,pathlib,subprocess,sys +p,variant,workload,clients,repeat,binary,max_prefills,mixed,idle,quantum,cmd_file,model_sha,prompts,repo=sys.argv[1:] +digest=lambda x: hashlib.sha256(pathlib.Path(x).read_bytes()).hexdigest() +libs={} +for line in subprocess.run(["ldd",binary],text=True,capture_output=True).stdout.splitlines(): + fields=line.replace("=>"," ").split() + paths=[x for x in fields if x.startswith("/") and pathlib.Path(x).is_file()] + for lib in paths: libs[str(pathlib.Path(lib).resolve())]=digest(lib) +lucebox_git_head=subprocess.run(["git","-C",repo,"rev-parse","HEAD"],text=True,capture_output=True).stdout.strip() or None +server_version=None +if variant == "llama": + version=subprocess.run([binary,"--version"],text=True,capture_output=True,timeout=30) + server_version="\n".join(x.strip() for x in (version.stdout,version.stderr) if x.strip()) or None + if version.returncode != 0 or server_version is None: + raise RuntimeError(f"cannot identify llama.cpp source version from {binary} --version") +obj={"variant":variant,"workload":workload,"clients":int(clients),"repeat":int(repeat), + "max_concurrent_prefills":int(max_prefills),"server_binary":str(pathlib.Path(binary).resolve()), + "mixed_prefill_tokens":int(mixed),"idle_prefill_tokens":int(idle), + "prefill_allocation_quantum":int(quantum), + "server_binary_sha256":digest(binary),"model_sha256":model_sha, + "prompt_file_sha256":digest(prompts),"server_command":pathlib.Path(cmd_file).read_text().strip(), + "resolved_shared_library_sha256":libs, + "lucebox_git_head":lucebox_git_head if variant != "llama" else None, + "server_version":server_version} +pathlib.Path(p).write_text(json.dumps(obj,indent=2,sort_keys=True)+"\n")' \ + "$path" "$variant" "$workload" "$clients" "$repeat" "$binary" "$max_prefills" "$mixed_budget" "$idle_budget" "$quantum" "$command_file" "$MODEL_SHA256" "$OUT/prompts/$workload.jsonl" "$REPO" +} + +run_case() { + local repeat="$1" workload="$2" clients="$3" variant="$4" + local max_ctx timeout capacity max_prefills mixed_budget idle_budget quantum binary model_id + if [[ "$workload" == long ]]; then + max_ctx=8192; timeout=1800 + else + max_ctx=4096; timeout=1200 + fi + capacity=$((SLOTS * max_ctx)) + local case_dir="$OUT/$workload/c$clients/r$repeat/$variant" + mkdir -p "$case_dir" + local -a command launch_env + if [[ "$variant" == llama ]]; then + binary="$LLAMA_SERVER_BIN"; model_id=qwen36-llama; max_prefills=0 + mixed_budget=0; idle_budget=0; quantum=0 + command=("$binary" -m "$MODEL" -ngl all --parallel "$SLOTS" -c "$capacity" + -b 2048 -ub 512 --cont-batching --no-context-shift --no-mmap -fa on + -ctk q4_0 -ctv q4_0 --no-cache-prompt --host 127.0.0.1 --port "$PORT" --alias "$model_id") + launch_env=() + else + binary="$LUCE_SERVER_BIN"; model_id=qwen36-luce + case "$variant" in + luce-k8) max_prefills=8; mixed_budget=2048; idle_budget=4096; quantum=512 ;; + luce-k16-b2) max_prefills=16; mixed_budget=2048; idle_budget=2048; quantum=128 ;; + luce-k16-b4) max_prefills=16; mixed_budget=4096; idle_budget=4096; quantum=256 ;; + luce-k16-dyn) max_prefills=16; mixed_budget=2048; idle_budget=4096; quantum=256 ;; + *) max_prefills=1; mixed_budget=2048; idle_budget=4096; quantum=512 ;; + esac + command=("$binary" "$MODEL" --target-device hip:0 --paged-attention + --max-concurrency "$SLOTS" --kv-pool-tokens "$capacity" --max-ctx "$max_ctx" + --cache-type-k q4_0 --cache-type-v q4_0 --fa-window 0 + --prefix-cache-slots 0 --prefill-cache-slots 0 --admission-coalesce-ms 5 + --host 127.0.0.1 --port "$PORT" --model-name "$model_id") + launch_env=("DFLASH_MIN_TOKENS=$WARMUP_TOKENS" + "DFLASH_MAX_CONCURRENT_PREFILLS=$max_prefills" + "DFLASH_MIXED_PREFILL_TOKENS=$mixed_budget" + "DFLASH_IDLE_PREFILL_TOKENS=$idle_budget" + "DFLASH_PREFILL_ALLOCATION_QUANTUM=$quantum") + fi + if ((${#launch_env[@]})); then + printf 'env ' > "$case_dir/server-command.txt" + printf '%q ' "${launch_env[@]}" "${command[@]}" >> "$case_dir/server-command.txt" + else + printf '%q ' "${command[@]}" > "$case_dir/server-command.txt" + fi + printf '\n' >> "$case_dir/server-command.txt" + write_metadata "$case_dir/server-metadata.json" "$variant" "$workload" "$clients" "$repeat" "$binary" "$max_prefills" "$mixed_budget" "$idle_budget" "$quantum" "$case_dir/server-command.txt" + + echo "[run] $workload C=$clients repeat=$repeat variant=$variant" + port_is_available || return 1 + if ((${#launch_env[@]})); then + env "${launch_env[@]}" "${command[@]}" > "$case_dir/server.log" 2>&1 & + else + "${command[@]}" > "$case_dir/server.log" 2>&1 & + fi + server_pid=$! + if ! wait_health "$model_id"; then tail -n 80 "$case_dir/server.log" >&2 || true; return 1; fi + + local offset="${prompt_offsets[$clients]}" prompts="$OUT/prompts/$workload.jsonl" + python3 "$CLIENT" --base-url "http://127.0.0.1:$PORT/v1" --model "$model_id" \ + --clients "$clients" --prompt-file "$prompts" --prompt-offset "$offset" \ + --require-distinct-prompts --max-tokens "$WARMUP_TOKENS" --temperature 0 \ + --ignore-eos --timeout "$timeout" --cooldown 0 --out "$case_dir/warmup.json" \ + --label "$variant $workload C=$clients warmup" > "$case_dir/warmup.txt" + sleep 1 + python3 "$CLIENT" --base-url "http://127.0.0.1:$PORT/v1" --model "$model_id" \ + --clients "$clients" --prompt-file "$prompts" --prompt-offset "$offset" \ + --require-distinct-prompts --max-tokens "$MAX_TOKENS" --temperature 0 \ + --ignore-eos --timeout "$timeout" --cooldown 0 \ + --server-metadata-json "$case_dir/server-metadata.json" --out "$case_dir/bench.json" \ + --label "$variant $workload C=$clients repeat=$repeat" | tee "$case_dir/bench.txt" + stop_server + sleep "$COOLDOWN_SECONDS" +} + +for ((repeat=1; repeat<=REPEATS; repeat++)); do + for workload in "${workload_list[@]}"; do + for c_index in "${!client_list[@]}"; do + clients="${client_list[$c_index]}" + # Rotate start variant by case so one engine is not always hot or cold. + shift_by=$(((repeat + c_index) % ${#variant_list[@]})) + for ((i=0; i<${#variant_list[@]}; i++)); do + variant="${variant_list[$(((i + shift_by) % ${#variant_list[@]}))]}" + run_case "$repeat" "$workload" "$clients" "$variant" + done + done + done +done + +python3 "$SUMMARIZER" "$OUT" --out "$OUT/summary.md" +echo "[run] complete: $OUT" diff --git a/harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh b/harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh new file mode 100755 index 000000000..aeb8c1d50 --- /dev/null +++ b/harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh @@ -0,0 +1,395 @@ +#!/usr/bin/env bash +# Fresh-process Qwen3.6 concurrent feature ablations with fail-closed proof. +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)" +REPO="${REPO:-$(cd -- "$SCRIPT_DIR/../../.." && pwd -P)}" +CLIENT="${CLIENT:-$SCRIPT_DIR/concurrent_benchmark.py}" +GENERATOR="${GENERATOR:-$SCRIPT_DIR/generate_ragged_prompts.py}" +SUMMARIZER="${SUMMARIZER:-$SCRIPT_DIR/summarize_feature_matrix.py}" +PROOF_TOOL="${PROOF_TOOL:-$SCRIPT_DIR/verify_feature_metrics.py}" +METADATA_TOOL="${METADATA_TOOL:-$SCRIPT_DIR/write_feature_metadata.py}" +RUNTIME_METADATA_TOOL="${RUNTIME_METADATA_TOOL:-$SCRIPT_DIR/record_feature_runtime.py}" + +MODEL="${MODEL:-}" +DRAFT_MODEL="${DRAFT_MODEL:-}" +PREFILL_DRAFTER="${PREFILL_DRAFTER:-}" +LUCE_SERVER_BIN="${LUCE_SERVER_BIN:-$REPO/server/build-hip/dflash_server}" +LLAMA_SERVER_BIN="${LLAMA_SERVER_BIN:-$(command -v llama-server 2>/dev/null || true)}" +OUT="${OUT:-$REPO/.harness-runs/qwen36-feature-matrix-$(date -u +%Y%m%dT%H%M%SZ)}" +REPEATS="${REPEATS:-1}" +WORKLOADS="${WORKLOADS:-short,compression}" +VARIANTS="${VARIANTS:-ar,ddtree,pflash,kvflash,full}" +CLIENTS="${CLIENTS:-4}" +PORT="${PORT:-18114}" +COOLDOWN_SECONDS="${COOLDOWN_SECONDS:-3}" +HEALTH_TIMEOUT_SECONDS="${HEALTH_TIMEOUT_SECONDS:-900}" +MAX_TOKENS="${MAX_TOKENS:-64}" +WARMUP_TOKENS="${WARMUP_TOKENS:-8}" +SLOTS="${SLOTS:-16}" +MAX_CONCURRENT_PREFILLS="${MAX_CONCURRENT_PREFILLS:-8}" +DRAFT_SWA="${DRAFT_SWA:-2048}" +PREFILL_UBATCH="${PREFILL_UBATCH:-512}" +DDTREE_ADAPTIVE="${DDTREE_ADAPTIVE:-1}" + +# The requested Strix Halo configuration. Every value is serialized into case +# metadata; no performance-affecting DFLASH variable is inherited implicitly. +TARGET_DEVICE="${TARGET_DEVICE:-hip:0}" +DRAFT_DEVICE="${DRAFT_DEVICE:-hip:0}" +DDTREE_BUDGET="${DDTREE_BUDGET:-22}" +DRAFT_RESIDENCY="${DRAFT_RESIDENCY:-persistent}" +PREFILL_COMPRESSION="${PREFILL_COMPRESSION:-auto}" +PREFILL_THRESHOLD="${PREFILL_THRESHOLD:-32000}" +PREFILL_KEEP_RATIO="${PREFILL_KEEP_RATIO:-0.05}" +KVFLASH_MODE="${KVFLASH_MODE:-auto}" +KVFLASH_MAX_POOL_TOKENS="${KVFLASH_MAX_POOL_TOKENS:-8192}" + +usage() { + cat <<'EOF' +Usage: + MODEL=/path/Qwen3.6-27B-Q4_K_M.gguf \ + DRAFT_MODEL=/path/dflash-draft-3.6-q8_0.gguf \ + PREFILL_DRAFTER=/path/Qwen3-0.6B-BF16.gguf \ + DRAFT_SWA=2048 PREFILL_UBATCH=512 \ + harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh + +The default is a bounded C4 smoke matrix with independently selectable +ar, ddtree, pflash, kvflash, and full rows. The full row is the requested +Strix Halo configuration: target/draft hip:0, Q8 DFlash SWA=2048, DDTree +budget 22, persistent PFlash auto, and KVFlash auto. The long-context profiles +are intended to cross the recorded 32K PFlash and 8K KV-residency thresholds; +word count is not treated as proof. Per-request wire/log token counts and +activation telemetry fail the case if an "auto" feature did not execute. + +llama is optional: include it explicitly with VARIANTS=ar,ddtree,llama and set +LLAMA_SERVER_BIN. For publication, set CLIENTS=1,4,8,16 and REPEATS=5. +For the AMD blog recipe, use the Q8_0 3.6 drafter, DRAFT_SWA=2048, PREFILL_UBATCH=512, and set DDTREE_ADAPTIVE=0 to keep DDTree active for every eligible step. +OUT must not already exist. +EOF +} + +if [[ "${1:-}" == "--help" ]]; then usage; exit 0; fi +if [[ $# -ne 0 ]]; then usage >&2; exit 2; fi +for cmd in python3 curl sha256sum awk; do + command -v "$cmd" >/dev/null || { echo "missing $cmd" >&2; exit 2; } +done +[[ -r "$MODEL" ]] || { echo "set MODEL to a readable target GGUF" >&2; exit 2; } +[[ "$REPEATS" =~ ^[1-9][0-9]*$ ]] || { echo "REPEATS must be positive" >&2; exit 2; } +[[ "$SLOTS" =~ ^[1-9][0-9]*$ ]] || { echo "SLOTS must be positive" >&2; exit 2; } +[[ "$MAX_CONCURRENT_PREFILLS" =~ ^[1-9][0-9]*$ ]] || { echo "MAX_CONCURRENT_PREFILLS must be positive" >&2; exit 2; } +[[ "$DDTREE_BUDGET" =~ ^[1-9][0-9]*$ ]] || { echo "DDTREE_BUDGET must be positive" >&2; exit 2; } +[[ "$PREFILL_THRESHOLD" =~ ^[1-9][0-9]*$ ]] || { echo "PREFILL_THRESHOLD must be positive" >&2; exit 2; } +[[ "$KVFLASH_MAX_POOL_TOKENS" =~ ^[1-9][0-9]*$ ]] || { echo "KVFLASH_MAX_POOL_TOKENS must be positive" >&2; exit 2; } +[[ "$DRAFT_SWA" =~ ^[0-9]+$ ]] || { echo "DRAFT_SWA must be a non-negative integer" >&2; exit 2; } +[[ "$PREFILL_UBATCH" =~ ^[1-9][0-9]*$ ]] || { echo "PREFILL_UBATCH must be positive" >&2; exit 2; } +[[ "$DDTREE_ADAPTIVE" == 0 || "$DDTREE_ADAPTIVE" == 1 ]] || { echo "DDTREE_ADAPTIVE must be 0 or 1" >&2; exit 2; } +[[ ! -e "$OUT" ]] || { echo "refusing to overwrite $OUT" >&2; exit 2; } + +ambient_tuning="$(env | grep -E '^(GGML_|DFLASH_|LUCE_|HIP_|ROCR_|HSA_|LD_PRELOAD=|LD_LIBRARY_PATH=)' \ + | grep -v '^LUCE_SERVER_BIN=' || true)" +if [[ -n "$ambient_tuning" ]]; then + echo "refusing ambient GPU/backend tuning variables:" >&2 + echo "$ambient_tuning" >&2 + exit 2 +fi + +IFS=, read -r -a workload_list <<< "$WORKLOADS" +IFS=, read -r -a variant_list <<< "$VARIANTS" +IFS=, read -r -a client_list <<< "$CLIENTS" +reject_duplicates() { + local list_name="$1" value + shift + local -A seen=() + for value in "$@"; do + if [[ -n "${seen[$value]+yes}" ]]; then + echo "$list_name contains duplicate entry: $value" >&2 + return 1 + fi + seen["$value"]=1 + done +} +reject_duplicates CLIENTS "${client_list[@]}" || exit 2 +reject_duplicates VARIANTS "${variant_list[@]}" || exit 2 +declare -A prompt_offsets=([1]=0 [4]=1 [8]=5 [16]=13) +for c in "${client_list[@]}"; do + [[ -n "${prompt_offsets[$c]+yes}" ]] || { echo "supported CLIENTS are 1,4,8,16" >&2; exit 2; } + (( c <= SLOTS )) || { echo "CLIENTS=$c exceeds SLOTS=$SLOTS" >&2; exit 2; } +done +luce_requested=0 +for v in "${variant_list[@]}"; do + case "$v" in + llama) ;; + ar|ddtree|pflash|kvflash|full) luce_requested=1 ;; + *) echo "unknown variant $v" >&2; exit 2 ;; + esac +done +if (( luce_requested )); then + [[ -x "$LUCE_SERVER_BIN" ]] || { echo "missing Lucebox server: $LUCE_SERVER_BIN" >&2; exit 2; } +fi + +contains_variant() { + local needle="$1" value + for value in "${variant_list[@]}"; do [[ "$value" == "$needle" ]] && return 0; done + return 1 +} +if contains_variant ddtree || contains_variant full; then + [[ -r "$DRAFT_MODEL" ]] || { echo "DRAFT_MODEL is required for DDTree/full" >&2; exit 2; } +fi +if contains_variant pflash || contains_variant kvflash || contains_variant full; then + [[ -r "$PREFILL_DRAFTER" ]] || { + echo "PREFILL_DRAFTER is required for PFlash and drafter-scored KVFlash" >&2 + exit 2 + } +fi +if contains_variant llama; then + [[ -x "$LLAMA_SERVER_BIN" ]] || { echo "llama requested but LLAMA_SERVER_BIN is missing" >&2; exit 2; } +fi + +MODEL_SHA256="$(sha256sum "$MODEL" | awk '{print $1}')" +DRAFT_MODEL_SHA256="" +PREFILL_DRAFTER_SHA256="" +[[ -n "$DRAFT_MODEL" ]] && DRAFT_MODEL_SHA256="$(sha256sum "$DRAFT_MODEL" | awk '{print $1}')" +[[ -n "$PREFILL_DRAFTER" ]] && PREFILL_DRAFTER_SHA256="$(sha256sum "$PREFILL_DRAFTER" | awk '{print $1}')" + +mkdir -p "$OUT/prompts" +for workload in "${workload_list[@]}"; do + python3 "$GENERATOR" --profile "$workload" --out "$OUT/prompts/$workload.jsonl" +done + +server_pid="" +stop_server() { + if [[ -n "$server_pid" ]] && kill -0 "$server_pid" 2>/dev/null; then + kill "$server_pid" 2>/dev/null || true + for _ in $(seq 1 30); do + kill -0 "$server_pid" 2>/dev/null || break + sleep 1 + done + kill -9 "$server_pid" 2>/dev/null || true + wait "$server_pid" 2>/dev/null || true + fi + server_pid="" +} +trap stop_server EXIT +trap 'exit 130' INT +trap 'exit 143' TERM + +wait_health() { + local deadline=$((SECONDS + HEALTH_TIMEOUT_SECONDS)) + while (( SECONDS < deadline )); do + kill -0 "$server_pid" 2>/dev/null || return 1 + curl -fsS --max-time 2 "http://127.0.0.1:$PORT/health" >/dev/null 2>&1 && return 0 + sleep 1 + done + return 1 +} + +port_is_available() { + python3 - "$PORT" <<'PY' +import socket +import sys + +port = int(sys.argv[1]) +with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + try: + sock.bind(("127.0.0.1", port)) + except OSError as exc: + print(f"PORT {port} is unavailable: {exc}", file=sys.stderr) + sys.exit(1) +PY +} + +# Auto modes cannot be proven on sub-threshold inputs. These skips are listed +# explicitly and are not emitted as successful benchmark rows. +case_applicable() { + local variant="$1" workload="$2" + case "$variant" in + pflash|full) [[ "$workload" == compression ]] ;; + kvflash) [[ "$workload" == compression || "$workload" == kv-pressure ]] ;; + *) return 0 ;; + esac +} + +workload_limits() { + case "$1" in + short) echo "4096 1200" ;; + medium) echo "8192 1800" ;; + long) echo "16384 2400" ;; + kv-pressure) echo "32768 3600" ;; + compression) echo "65536 5400" ;; + *) echo "unknown workload $1" >&2; return 1 ;; + esac +} + +run_case() { + local repeat="$1" workload="$2" clients="$3" variant="$4" + local max_ctx timeout + read -r max_ctx timeout <<< "$(workload_limits "$workload")" + local capacity=$((SLOTS * max_ctx)) + local case_dir="$OUT/$workload/c$clients/r$repeat/$variant" + mkdir -p "$case_dir" + + local binary model_id max_prefills + local -a command launch_env metadata expected + if [[ "$variant" == llama ]]; then + binary="$LLAMA_SERVER_BIN"; model_id=qwen36-llama; max_prefills=0 + command=("$binary" -m "$MODEL" -ngl all --parallel "$SLOTS" -c "$capacity" + -b 2048 -ub 512 --cont-batching --no-context-shift --no-mmap -fa on + -ctk q4_0 -ctv q4_0 --no-cache-prompt --host 127.0.0.1 --port "$PORT" --alias "$model_id") + launch_env=() + else + binary="$LUCE_SERVER_BIN"; model_id=qwen36-luce; max_prefills="$MAX_CONCURRENT_PREFILLS" + command=("$binary" "$MODEL" --target-device "$TARGET_DEVICE" --paged-attention + --max-concurrency "$SLOTS" --kv-pool-tokens "$capacity" --max-ctx "$max_ctx" + --cache-type-k q4_0 --cache-type-v q4_0 --fa-window 0 + --prefix-cache-slots 0 --prefill-cache-slots 0 --admission-coalesce-ms 5 + --host 127.0.0.1 --port "$PORT" --model-name "$model_id") + launch_env=("DFLASH_MIN_TOKENS=$WARMUP_TOKENS" "DFLASH_MAX_CONCURRENT_PREFILLS=$max_prefills" + "DFLASH27B_DRAFT_SWA=$DRAFT_SWA" "DFLASH27B_PREFILL_UBATCH=$PREFILL_UBATCH") + if [[ "$DDTREE_ADAPTIVE" == 0 ]]; then + launch_env+=("DFLASH_DDTREE_ADAPTIVE=0") + fi + if [[ "$variant" == ddtree || "$variant" == full ]]; then + command+=(--draft "$DRAFT_MODEL" --draft-device "$DRAFT_DEVICE" + --ddtree --ddtree-budget "$DDTREE_BUDGET" --fast-rollback) + expected+=(--expect ddtree) + if [[ "$variant" == ddtree ]]; then command+=(--draft-residency "$DRAFT_RESIDENCY"); fi + fi + if [[ "$variant" == pflash || "$variant" == full ]]; then + if [[ "$variant" == pflash ]]; then command+=(--draft-device "$DRAFT_DEVICE"); fi + command+=(--prefill-compression "$PREFILL_COMPRESSION" + --prefill-threshold "$PREFILL_THRESHOLD" + --prefill-keep-ratio "$PREFILL_KEEP_RATIO" + --prefill-drafter "$PREFILL_DRAFTER" + --draft-residency "$DRAFT_RESIDENCY") + expected+=(--expect pflash) + fi + if [[ "$variant" == kvflash || "$variant" == full ]]; then + # --prefill-drafter also selects the default drafter-scored KVFlash + # residency policy; it does not enable PFlash when compression is off. + if [[ "$variant" == kvflash ]]; then + command+=(--prefill-drafter "$PREFILL_DRAFTER") + fi + command+=(--kvflash "$KVFLASH_MODE") + launch_env+=("DFLASH_KVFLASH_MAX_POOL=$KVFLASH_MAX_POOL_TOKENS") + expected+=(--expect kvflash) + fi + fi + + if ((${#launch_env[@]})); then + printf 'env ' > "$case_dir/server-command.txt" + printf '%q ' "${launch_env[@]}" "${command[@]}" >> "$case_dir/server-command.txt" + else + printf '%q ' "${command[@]}" > "$case_dir/server-command.txt" + fi + printf '\n' >> "$case_dir/server-command.txt" + + metadata=(python3 "$METADATA_TOOL" --out "$case_dir/server-metadata.json" + --variant "$variant" --workload "$workload" --clients "$clients" --repeat "$repeat" + --binary "$binary" --model "$MODEL" --model-sha256 "$MODEL_SHA256" + --prompt-file "$OUT/prompts/$workload.jsonl" + --command-file "$case_dir/server-command.txt" --repo "$REPO" + --max-concurrent-prefills "$max_prefills") + if [[ "$variant" != llama ]]; then + metadata+=(--target-device "$TARGET_DEVICE") + local item + for item in "${launch_env[@]}"; do metadata+=(--launch-env "$item"); done + fi + if [[ "$variant" == ddtree || "$variant" == full ]]; then + metadata+=(--draft-device "$DRAFT_DEVICE" --draft-model "$DRAFT_MODEL" + --draft-model-sha256 "$DRAFT_MODEL_SHA256" --ddtree --ddtree-budget "$DDTREE_BUDGET" --fast-rollback) + if [[ "$variant" == ddtree ]]; then metadata+=(--draft-residency "$DRAFT_RESIDENCY"); fi + fi + if [[ "$variant" == pflash || "$variant" == full ]]; then + metadata+=(--draft-device "$DRAFT_DEVICE" --prefill-compression "$PREFILL_COMPRESSION" + --prefill-threshold "$PREFILL_THRESHOLD" --prefill-keep-ratio "$PREFILL_KEEP_RATIO" + --prefill-drafter "$PREFILL_DRAFTER" + --prefill-drafter-sha256 "$PREFILL_DRAFTER_SHA256" + --draft-residency "$DRAFT_RESIDENCY") + fi + if [[ "$variant" == kvflash || "$variant" == full ]]; then + metadata+=(--kvflash "$KVFLASH_MODE" + --kvflash-max-pool-tokens "$KVFLASH_MAX_POOL_TOKENS" + --kvflash-scorer-drafter "$PREFILL_DRAFTER" + --kvflash-scorer-drafter-sha256 "$PREFILL_DRAFTER_SHA256") + if [[ "$variant" == kvflash ]]; then + # Record the literal server flag too, while keeping compression off. + metadata+=(--prefill-drafter "$PREFILL_DRAFTER" + --prefill-drafter-sha256 "$PREFILL_DRAFTER_SHA256") + fi + fi + "${metadata[@]}" + + echo "[run] $workload C=$clients repeat=$repeat variant=$variant" + port_is_available || return 1 + if ((${#launch_env[@]})); then + env "${launch_env[@]}" "${command[@]}" > "$case_dir/server.log" 2>&1 & + else + "${command[@]}" > "$case_dir/server.log" 2>&1 & + fi + server_pid=$! + if ! wait_health; then tail -n 120 "$case_dir/server.log" >&2 || true; return 1; fi + if [[ "$variant" != llama ]]; then + python3 "$RUNTIME_METADATA_TOOL" --metadata "$case_dir/server-metadata.json" \ + --server-log "$case_dir/server.log" + fi + + local offset="${prompt_offsets[$clients]}" prompts="$OUT/prompts/$workload.jsonl" + local -a common_client=(--base-url "http://127.0.0.1:$PORT/v1" --model "$model_id" + --clients "$clients" --prompt-file "$prompts" --prompt-offset "$offset" + --require-distinct-prompts --temperature 0 --ignore-eos --timeout "$timeout" --cooldown 0) + local -a telemetry_arg=() + [[ "$variant" != llama ]] && telemetry_arg+=(--require-effective-prompt-telemetry) + local -a warmup_cmd=( + python3 "$CLIENT" "${common_client[@]}" + --max-tokens "$WARMUP_TOKENS" "${telemetry_arg[@]}" + --out "$case_dir/warmup.json" + --label "$variant $workload C=$clients warmup" + ) + "${warmup_cmd[@]}" > "$case_dir/warmup.txt" + + local -a benchmark_cmd=( + python3 "$CLIENT" "${common_client[@]}" + --max-tokens "$MAX_TOKENS" "${telemetry_arg[@]}" + --server-metadata-json "$case_dir/server-metadata.json" + --out "$case_dir/bench.json" + --label "$variant $workload C=$clients repeat=$repeat" + ) + "${benchmark_cmd[@]}" | tee "$case_dir/bench.txt" + stop_server + + if [[ "$variant" != llama ]]; then + local -a proof_cmd=( + python3 "$PROOF_TOOL" + --bench "$case_dir/bench.json" + --server-log "$case_dir/server.log" + "${expected[@]}" + --out "$case_dir/feature-proof.json" + ) + "${proof_cmd[@]}" + fi + sleep "$COOLDOWN_SECONDS" +} + +active_cases=0 +for ((repeat=1; repeat<=REPEATS; repeat++)); do + for workload in "${workload_list[@]}"; do + for c_index in "${!client_list[@]}"; do + clients="${client_list[$c_index]}" + shift_by=$(((repeat + c_index) % ${#variant_list[@]})) + for ((i=0; i<${#variant_list[@]}; i++)); do + variant="${variant_list[$(((i + shift_by) % ${#variant_list[@]}))]}" + if ! case_applicable "$variant" "$workload"; then + echo "[skip] $variant requires an activation workload; workload=$workload" + continue + fi + active_cases=$((active_cases + 1)) + run_case "$repeat" "$workload" "$clients" "$variant" + done + done + done +done +(( active_cases > 0 )) || { echo "no applicable benchmark cases" >&2; exit 2; } + +python3 "$SUMMARIZER" "$OUT" --out "$OUT/summary.md" +echo "[run] complete: $OUT" diff --git a/harness/benchmarks/concurrency/run_qwen38_dflash2_subsets.sh b/harness/benchmarks/concurrency/run_qwen38_dflash2_subsets.sh new file mode 100755 index 000000000..1c34a22c8 --- /dev/null +++ b/harness/benchmarks/concurrency/run_qwen38_dflash2_subsets.sh @@ -0,0 +1,283 @@ +#!/usr/bin/env bash +# Persistent-server forced AR/speculation subset screen for concurrent DFlash2. +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)" +REPO="${REPO:-$(cd -- "$SCRIPT_DIR/../../.." && pwd -P)}" +CLIENT="${CLIENT:-$SCRIPT_DIR/forced_subset_benchmark.py}" +REFILL_CLIENT="${REFILL_CLIENT:-$SCRIPT_DIR/refill_subset_benchmark.py}" +METADATA_TOOL="${METADATA_TOOL:-$SCRIPT_DIR/write_feature_metadata.py}" +RUNTIME_METADATA_TOOL="${RUNTIME_METADATA_TOOL:-$SCRIPT_DIR/record_feature_runtime.py}" + +MODEL="${MODEL:-}" +DRAFT_MODEL="${DRAFT_MODEL:-}" +PROMPT_FILE="${PROMPT_FILE:-}" +PROMPT_OFFSET="${PROMPT_OFFSET:-0}" +LUCE_SERVER_BIN="${LUCE_SERVER_BIN:-$REPO/server/build-hip/dflash_server}" +OUT="${OUT:-$REPO/.harness-runs/qwen38-dflash2-subsets-$(date -u +%Y%m%dT%H%M%SZ)}" +CLIENTS="${CLIENTS:-2}" +MASKS="${MASKS:-AA,AS,SA,SS}" +SPEC_DEPTH="${SPEC_DEPTH:-4}" +REPEATS="${REPEATS:-1}" +REFILL_WAVES="${REFILL_WAVES:-1}" +MAX_TOKENS="${MAX_TOKENS:-256}" +WARMUP_TOKENS="${WARMUP_TOKENS:-16}" +MIN_FULL_LIVE_ROUNDS="${MIN_FULL_LIVE_ROUNDS:-2}" +MAX_START_SKEW_MS="${MAX_START_SKEW_MS:-100}" +MAX_REFILL_GAP_MS="${MAX_REFILL_GAP_MS:-100}" +SLOTS="${SLOTS:-8}" +MAX_CTX="${MAX_CTX:-8192}" +MAX_CONCURRENT_PREFILLS="${MAX_CONCURRENT_PREFILLS:-8}" +CACHE_TYPE_K="${CACHE_TYPE_K:-q8_0}" +CACHE_TYPE_V="${CACHE_TYPE_V:-q8_0}" +FA_WINDOW="${FA_WINDOW:-0}" +PORT="${PORT:-18139}" +HEALTH_TIMEOUT_SECONDS="${HEALTH_TIMEOUT_SECONDS:-900}" +REQUEST_TIMEOUT_SECONDS="${REQUEST_TIMEOUT_SECONDS:-1800}" +COOLDOWN_SECONDS="${COOLDOWN_SECONDS:-1}" +TARGET_DEVICE="${TARGET_DEVICE:-hip:0}" +DRAFT_DEVICE="${DRAFT_DEVICE:-hip:0}" +VISIBLE_DEVICES="${VISIBLE_DEVICES:-0}" + +usage() { + cat <<'EOF' +Usage: + MODEL=/path/Qwen3.8-27B-target.gguf \ + DRAFT_MODEL=/path/Qwen3.8-27B-DFlash2-q8_0.gguf \ + PROMPT_FILE=/path/prompts.jsonl \ + CLIENTS=2 MASKS=AA,AS,SA,SS SPEC_DEPTH=4 PROMPT_OFFSET=0 \ + REFILL_WAVES=1 \ + harness/benchmarks/concurrency/run_qwen38_dflash2_subsets.sh + +Mask characters are positional: A forces request-local AR and S forces +request-local speculation. Adaptive is intentionally unsupported. One server +is launched for the selected depth, warmed once in all-AR and all-SPEC modes, +then kept alive for every mask and repeat. Invoke a fresh OUT/process for each +depth in a {2,4,8} screen. REFILL_WAVES=1 runs the original synchronized, +closed-cohort diagnostic. REFILL_WAVES>=3 instead keeps each positional lane +full with repeated deterministic requests and reports refill goodput separately +from full-live engine-round goodput. The report fails unless at least +MIN_FULL_LIVE_ROUNDS consecutive engine rounds execute at live=CLIENTS; refill +also reserves one guard cohort to prove saturation away from the final drain. + +The runner records the target, DFlash2 draft, binary, libraries, git revision, +exact command, depth/timing/selector environment, prompt hash, request masks, +per-request output hashes, and measured server records. These are forced +controls only and cannot be reported as adaptive activation results. +EOF +} + +if [[ "${1:-}" == "--help" ]]; then usage; exit 0; fi +if [[ $# -ne 0 ]]; then usage >&2; exit 2; fi +for cmd in python3 curl sha256sum awk; do + command -v "$cmd" >/dev/null || { echo "missing $cmd" >&2; exit 2; } +done +[[ -r "$MODEL" ]] || { echo "set MODEL to a readable Qwen3.8 target GGUF" >&2; exit 2; } +[[ -r "$DRAFT_MODEL" ]] || { echo "set DRAFT_MODEL to a readable DFlash2 GGUF" >&2; exit 2; } +[[ -r "$PROMPT_FILE" ]] || { echo "set PROMPT_FILE to a readable JSONL/text prompt file" >&2; exit 2; } +[[ -x "$LUCE_SERVER_BIN" ]] || { echo "missing Lucebox server: $LUCE_SERVER_BIN" >&2; exit 2; } +[[ -r "$CLIENT" && -r "$REFILL_CLIENT" && -r "$METADATA_TOOL" && -r "$RUNTIME_METADATA_TOOL" ]] || { + echo "missing DFlash2 subset harness tool" >&2; exit 2; +} +for value_name in CLIENTS REPEATS REFILL_WAVES MAX_TOKENS WARMUP_TOKENS MIN_FULL_LIVE_ROUNDS SLOTS MAX_CTX MAX_CONCURRENT_PREFILLS HEALTH_TIMEOUT_SECONDS REQUEST_TIMEOUT_SECONDS; do + value="${!value_name}" + [[ "$value" =~ ^[1-9][0-9]*$ ]] || { echo "$value_name must be positive" >&2; exit 2; } +done +(( REFILL_WAVES == 1 || REFILL_WAVES >= 3 )) || { echo "REFILL_WAVES must be 1 (closed cohort) or at least 3 (guarded refill)" >&2; exit 2; } +[[ "$SPEC_DEPTH" =~ ^[2-9][0-9]*$ ]] || { echo "SPEC_DEPTH must be at least 2" >&2; exit 2; } +[[ "$PROMPT_OFFSET" =~ ^[0-9]+$ ]] || { echo "PROMPT_OFFSET must be non-negative" >&2; exit 2; } +[[ "$MAX_REFILL_GAP_MS" =~ ^[0-9]+([.][0-9]+)?$ ]] || { echo "MAX_REFILL_GAP_MS must be non-negative" >&2; exit 2; } +[[ "$PORT" =~ ^[1-9][0-9]*$ ]] || { echo "PORT must be positive" >&2; exit 2; } +[[ "$COOLDOWN_SECONDS" =~ ^[0-9]+$ ]] || { echo "COOLDOWN_SECONDS must be non-negative" >&2; exit 2; } +[[ "$MAX_START_SKEW_MS" =~ ^[0-9]+([.][0-9]+)?$ ]] || { echo "MAX_START_SKEW_MS must be non-negative" >&2; exit 2; } +(( CLIENTS <= SLOTS )) || { echo "CLIENTS exceeds SLOTS" >&2; exit 2; } +[[ "$FA_WINDOW" == 0 ]] || { echo "paged concurrency requires FA_WINDOW=0" >&2; exit 2; } +[[ ! -e "$OUT" ]] || { echo "refusing to overwrite OUT=$OUT" >&2; exit 2; } + +ambient_tuning="$(env | grep -E '^(GGML_|DFLASH_|LUCE_|HIP_|ROCR_|HSA_|LD_PRELOAD=|LD_LIBRARY_PATH=)' \ + | grep -v '^LUCE_SERVER_BIN=' || true)" +if [[ -n "$ambient_tuning" ]]; then + echo "refusing ambient GPU/backend tuning variables:" >&2 + echo "$ambient_tuning" >&2 + exit 2 +fi + +IFS=, read -r -a mask_list <<< "$MASKS" +(( ${#mask_list[@]} > 0 )) || { echo "MASKS must not be empty" >&2; exit 2; } +declare -A seen_masks=() +for mask in "${mask_list[@]}"; do + [[ ${#mask} -eq CLIENTS && "$mask" =~ ^[AS]+$ ]] || { + echo "mask $mask must contain exactly CLIENTS=$CLIENTS A/S characters" >&2 + exit 2 + } + [[ -z "${seen_masks[$mask]+present}" ]] || { echo "duplicate mask $mask" >&2; exit 2; } + seen_masks[$mask]=1 +done + +mode_csv() { + local mask="$1" result="" mode index + for ((index=0; index/dev/null; then + kill "$server_pid" 2>/dev/null || true + for _ in $(seq 1 30); do + kill -0 "$server_pid" 2>/dev/null || break + sleep 1 + done + kill -9 "$server_pid" 2>/dev/null || true + wait "$server_pid" 2>/dev/null || true + fi + server_pid="" +} +trap stop_server EXIT +trap 'exit 130' INT +trap 'exit 143' TERM + +wait_health() { + local deadline=$((SECONDS + HEALTH_TIMEOUT_SECONDS)) + while (( SECONDS < deadline )); do + kill -0 "$server_pid" 2>/dev/null || return 1 + curl -fsS --max-time 2 "http://127.0.0.1:$PORT/health" >/dev/null 2>&1 && return 0 + sleep 1 + done + return 1 +} + +MODEL_SHA256="$(sha256sum "$MODEL" | awk '{print $1}')" +DRAFT_MODEL_SHA256="$(sha256sum "$DRAFT_MODEL" | awk '{print $1}')" +capacity=$((SLOTS * MAX_CTX)) +model_id=qwen38-dflash2 +mkdir -p "$OUT" + +command=( + "$LUCE_SERVER_BIN" "$MODEL" --draft "$DRAFT_MODEL" + --target-device "$TARGET_DEVICE" --draft-device "$DRAFT_DEVICE" + --paged-attention --max-concurrency "$SLOTS" + --kv-pool-tokens "$capacity" --max-ctx "$MAX_CTX" + --cache-type-k "$CACHE_TYPE_K" --cache-type-v "$CACHE_TYPE_V" + --fa-window "$FA_WINDOW" --prefix-cache-slots 0 --prefill-cache-slots 0 + --admission-coalesce-ms 20 --draft-residency persistent + --decode-mode speculation --host 127.0.0.1 --port "$PORT" --model-name "$model_id" +) +launch_env=( + "HIP_VISIBLE_DEVICES=$VISIBLE_DEVICES" + "DFLASH_MAX_CONCURRENT_PREFILLS=$MAX_CONCURRENT_PREFILLS" + "DFLASH_SPEC_BATCHED_DRAFT=1" + "DFLASH_SPEC_CHAIN_DEPTH=$SPEC_DEPTH" + "DFLASH_STEP_TIMING=1" + "DFLASH_DFLASH2_SELECTOR_LOG=1" + "PROMPT_OFFSET=$PROMPT_OFFSET" +) +printf 'env ' > "$OUT/server-command.txt" +printf '%q ' "${launch_env[@]}" "${command[@]}" >> "$OUT/server-command.txt" +printf '\n' >> "$OUT/server-command.txt" + +write_case_metadata() { + local case_dir="$1" variant="$2" repeat="$3" workload="${4:-dflash2-forced-subsets}" + mkdir -p "$case_dir" + metadata=( + python3 "$METADATA_TOOL" --out "$case_dir/server-metadata.json" + --variant "$variant" --workload "$workload" + --clients "$CLIENTS" --repeat "$repeat" + --binary "$LUCE_SERVER_BIN" --model "$MODEL" + --model-sha256 "$MODEL_SHA256" --prompt-file "$PROMPT_FILE" + --command-file "$OUT/server-command.txt" --repo "$REPO" + --max-concurrent-prefills "$MAX_CONCURRENT_PREFILLS" + --target-device "$TARGET_DEVICE" --draft-device "$DRAFT_DEVICE" + --draft-model "$DRAFT_MODEL" --draft-model-sha256 "$DRAFT_MODEL_SHA256" + --decode-mode speculation --cache-type-k "$CACHE_TYPE_K" --cache-type-v "$CACHE_TYPE_V" + --fa-window "$FA_WINDOW" --draft-residency persistent + ) + local item + for item in "${launch_env[@]}"; do metadata+=(--launch-env "$item"); done + "${metadata[@]}" + python3 "$RUNTIME_METADATA_TOOL" \ + --metadata "$case_dir/server-metadata.json" --server-log "$OUT/server.log" +} + +run_client_case() { + local case_dir="$1" mask="$2" max_tokens="$3" repeat="$4" min_rounds="$5" waves="${6:-1}" + local modes client="$CLIENT" workload=dflash2-forced-subsets variant + modes="$(mode_csv "$mask")" + variant="dflash2-depth-$SPEC_DEPTH-mask-$mask" + if (( waves > 1 )); then + client="$REFILL_CLIENT" + workload=dflash2-forced-refill + variant+="-refill-w$waves" + fi + write_case_metadata "$case_dir" "$variant" "$repeat" "$workload" + client_cmd=( + python3 "$client" --base-url "http://127.0.0.1:$PORT/v1" --model "$model_id" + --clients "$CLIENTS" --request-modes "$modes" --spec-depth "$SPEC_DEPTH" + --prompt-file "$PROMPT_FILE" --prompt-offset "$PROMPT_OFFSET" --max-tokens "$max_tokens" + --timeout "$REQUEST_TIMEOUT_SECONDS" --max-start-skew-ms "$MAX_START_SKEW_MS" + --min-full-live-rounds "$min_rounds" + --server-metadata-json "$case_dir/server-metadata.json" + --server-log "$OUT/server.log" --out "$case_dir/bench.json" + --label "DFlash2 depth=$SPEC_DEPTH mask=$mask repeat=$repeat" + ) + if (( waves > 1 )); then + client_cmd+=(--waves "$waves" --max-refill-gap-ms "$MAX_REFILL_GAP_MS") + fi + printf '%q ' "${client_cmd[@]}" > "$case_dir/client-command.txt" + printf '\n' >> "$case_dir/client-command.txt" + "${client_cmd[@]}" | tee "$case_dir/bench.txt" +} + +port_is_available +env "${launch_env[@]}" "${command[@]}" > "$OUT/server.log" 2>&1 & +server_pid=$! +if ! wait_health; then + tail -n 160 "$OUT/server.log" >&2 || true + exit 1 +fi + +all_ar="$(repeat_char A)" +all_spec="$(repeat_char S)" +run_client_case "$OUT/warmup/ar" "$all_ar" "$WARMUP_TOKENS" 0 1 +run_client_case "$OUT/warmup/speculation" "$all_spec" "$WARMUP_TOKENS" 0 1 +sleep "$COOLDOWN_SECONDS" + +for ((repeat=1; repeat<=REPEATS; repeat++)); do + shift_by=$(((repeat - 1) % ${#mask_list[@]})) + for ((index=0; index<${#mask_list[@]}; index++)); do + mask="${mask_list[$(((index + shift_by) % ${#mask_list[@]}))]}" + case_dir="$OUT/c$CLIENTS/depth$SPEC_DEPTH/r$repeat/$mask" + echo "[run] C=$CLIENTS depth=$SPEC_DEPTH repeat=$repeat mask=$mask refill_waves=$REFILL_WAVES" + run_client_case "$case_dir" "$mask" "$MAX_TOKENS" "$repeat" "$MIN_FULL_LIVE_ROUNDS" "$REFILL_WAVES" + sleep "$COOLDOWN_SECONDS" + done +done + +stop_server +echo "[run] complete: $OUT" diff --git a/harness/benchmarks/concurrency/summarize_concurrency.py b/harness/benchmarks/concurrency/summarize_concurrency.py new file mode 100755 index 000000000..0f4374bcb --- /dev/null +++ b/harness/benchmarks/concurrency/summarize_concurrency.py @@ -0,0 +1,221 @@ +#!/usr/bin/env python3 +"""Summarize paired Lucebox/llama.cpp concurrency benchmark reports.""" + +from __future__ import annotations + +import argparse +import json +import statistics +from collections import defaultdict +from pathlib import Path + + +def load_reports(root: Path) -> list[dict]: + reports = [] + for path in sorted(root.rglob("bench.json")): + report = json.loads(path.read_text(encoding="utf-8")) + meta = report.get("server_metadata") or {} + if len(report.get("levels", [])) != 1: + raise ValueError(f"{path}: expected exactly one client level") + level = report["levels"][0] + if ( + level.get("failures") + or not level.get("token_count_complete") + or not level.get("prompt_token_count_complete") + ): + raise ValueError(f"{path}: failed or incomplete token accounting") + if report.get("ignore_eos") and level.get("fixed_token_workload_valid") is not True: + raise ValueError(f"{path}: fixed-token validation failed") + reports.append({"path": path, "report": report, "level": level, "meta": meta}) + if not reports: + raise ValueError(f"{root}: no bench.json files found") + return reports + + +def median(values: list[float]) -> float: + return statistics.median(values) + + +def complete_median(values: list[float | None]) -> float | None: + """Return a median only when every repeat measured the metric.""" + if not values or any(value is None for value in values): + return None + return median([value for value in values if value is not None]) + + +def output_stability(items: list[dict]) -> str: + output_digests = [ + item["level"].get("selected_output_set_sha256") for item in items + ] + complete = all(isinstance(value, str) and bool(value) for value in output_digests) + hashes = {value for value in output_digests if isinstance(value, str)} + return ( + "n/a" if len(items) < 2 or not complete + else "yes" if len(hashes) == 1 + else "NO" + ) + + +def run_signature(item: dict) -> tuple[object, ...]: + report, meta = item["report"], item["meta"] + max_tokens = report.get("max_tokens") + ignore_eos = report.get("ignore_eos") + temperature = report.get("temperature") + seed = report.get("seed") + model_sha256 = meta.get("model_sha256") + if ( + type(max_tokens) is not int or max_tokens <= 0 + or not isinstance(ignore_eos, bool) + or type(temperature) not in (int, float) + or type(seed) is not int + or not isinstance(model_sha256, str) or not model_sha256 + ): + raise ValueError("incomplete run metadata") + return max_tokens, ignore_eos, temperature, seed, model_sha256 + + +def report_key(item: dict) -> tuple[str, int, str]: + meta, level = item["meta"], item["level"] + workload = meta.get("workload") + variant = meta.get("variant") + clients = level.get("clients") + if not isinstance(workload, str) or not workload: + raise ValueError("incomplete report metadata: missing workload") + if not isinstance(variant, str) or not variant: + raise ValueError("incomplete report metadata: missing variant") + if type(clients) is not int or clients <= 0: + raise ValueError("incomplete report metadata: invalid clients") + return workload, clients, variant + + +def summarize(reports: list[dict]) -> str: + grouped: dict[tuple[str, int, str], list[dict]] = defaultdict(list) + for item in reports: + grouped[report_key(item)].append(item) + for key, items in grouped.items(): + repeats = [item["meta"].get("repeat") for item in items] + if any(type(repeat) is not int or repeat <= 0 for repeat in repeats): + raise ValueError(f"{key}: invalid or missing repeat") + if len(repeats) != len(set(repeats)): + raise ValueError(f"{key}: duplicate repeat") + if len({run_signature(item) for item in items}) != 1: + raise ValueError(f"{key}: incompatible run metadata") + + lines = [ + "# Concurrency benchmark summary", "", + "Aggregate output goodput includes queueing, prefill, and decode. " + "Output-window goodput starts at the first observed output and is decode-facing, " + "but it can include staggered prefill. Prompt tok/s to first token includes " + "admission and TTFT.", "", + "| Workload | C | Variant | Repeats | Output goodput tok/s | " + "Output-window tok/s | Request decode tok/s | Prompt tok/s to first | " + "TTFT max s | Stable output | vs llama | Decode vs llama | K8 vs K1 |", + "| :--- | ---: | :--- | ---: | ---: | ---: | ---: | ---: | ---: | " + ":---: | ---: | ---: | ---: |", + ] + for workload, clients, variant in sorted(grouped): + items = grouped[(workload, clients, variant)] + prompt_digests = [ + item["level"].get("selected_prompt_set_sha256") for item in items + ] + if not all(isinstance(value, str) and value for value in prompt_digests): + raise ValueError( + f"{workload} C={clients} {variant}: missing selected prompt set hash" + ) + hashes = set(prompt_digests) + if len(hashes) != 1: + raise ValueError(f"{workload} C={clients} {variant}: prompt sets differ") + goodput_values = [item["level"].get("aggregate_tok_s") for item in items] + if any(type(value) not in (int, float) for value in goodput_values): + raise ValueError( + f"{workload} C={clients} {variant}: missing aggregate token rate" + ) + goodput = median(goodput_values) + output_window = complete_median([ + item["level"].get("output_window_tok_s") for item in items + ]) + request_decode = complete_median([ + item["level"].get("request_decode_tok_s_median") for item in items + ]) + prompt_rate = complete_median([ + item["level"].get("prompt_tokens_per_s_to_first_token") for item in items + ]) + ttft = complete_median([ + item["level"].get("ttft_max_s") for item in items + ]) + stable = output_stability(items) + + def delta(other: str, metric: str) -> str: + peers = grouped.get((workload, clients, other), []) + if not peers: + return "n/a" + peer_hashes = {p["level"]["selected_prompt_set_sha256"] for p in peers} + if peer_hashes != hashes: + raise ValueError(f"{workload} C={clients}: {variant}/{other} prompts differ") + if {run_signature(item) for item in items} != { + run_signature(peer) for peer in peers + }: + raise ValueError( + f"{workload} C={clients}: {variant}/{other} run metadata differs" + ) + by_repeat = {int(item["meta"]["repeat"]): item for item in items} + peers_by_repeat = {int(item["meta"]["repeat"]): item for item in peers} + if by_repeat.keys() != peers_by_repeat.keys(): + raise ValueError( + f"{workload} C={clients}: {variant}/{other} repeat sets differ" + ) + if stable == "NO" or output_stability(peers) == "NO": + return "n/a" + ratios = [] + for repeat in sorted(by_repeat): + value = by_repeat[repeat]["level"].get(metric) + base = peers_by_repeat[repeat]["level"].get(metric) + if value is None or base is None: + return "n/a" + if base <= 0: + raise ValueError( + f"{workload} C={clients} repeat={repeat}: " + f"non-positive {other} {metric}" + ) + ratios.append(value / base - 1.0) + return f"{median(ratios) * 100:+.1f}%" + + vs_llama = ( + delta("llama", "aggregate_tok_s") + if variant == "luce-k8" else "—" + ) + decode_vs_llama = ( + delta("llama", "output_window_tok_s") + if variant == "luce-k8" else "—" + ) + vs_k1 = ( + delta("luce-k1", "aggregate_tok_s") + if variant == "luce-k8" else "—" + ) + output_window_text = f"{output_window:.2f}" if output_window is not None else "n/a" + request_decode_text = f"{request_decode:.2f}" if request_decode is not None else "n/a" + prompt_rate_text = f"{prompt_rate:.2f}" if prompt_rate is not None else "n/a" + ttft_text = f"{ttft:.3f}" if ttft is not None else "n/a" + lines.append( + f"| {workload} | {clients} | {variant} | {len(items)} | {goodput:.2f} | " + f"{output_window_text} | {request_decode_text} | {prompt_rate_text} | " + f"{ttft_text} | {stable} | {vs_llama} | {decode_vs_llama} | {vs_k1} |" + ) + lines.append("") + return "\n".join(lines) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("root", type=Path) + parser.add_argument("--out", type=Path) + args = parser.parse_args() + text = summarize(load_reports(args.root)) + if args.out: + args.out.write_text(text + "\n", encoding="utf-8") + print(text) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/harness/benchmarks/concurrency/summarize_feature_matrix.py b/harness/benchmarks/concurrency/summarize_feature_matrix.py new file mode 100755 index 000000000..17076a62e --- /dev/null +++ b/harness/benchmarks/concurrency/summarize_feature_matrix.py @@ -0,0 +1,255 @@ +#!/usr/bin/env python3 +"""Summarize Qwen3.6 concurrent feature ablations.""" + +from __future__ import annotations + +import argparse +import json +import statistics +from collections import defaultdict +from pathlib import Path + + +def load_reports(root: Path) -> list[dict]: + reports = [] + for path in sorted(root.rglob("bench.json")): + report = json.loads(path.read_text(encoding="utf-8")) + meta = report.get("server_metadata") or {} + if len(report.get("levels", [])) != 1: + raise ValueError(f"{path}: expected exactly one client level") + level = report["levels"][0] + variant = str(meta.get("variant", "")) + if ( + level.get("failures") + or not level.get("token_count_complete") + or not level.get("prompt_token_count_complete") + or (variant != "llama" + and not level.get("effective_prompt_token_count_complete")) + ): + raise ValueError(f"{path}: failed or incomplete token accounting") + if report.get("ignore_eos") and level.get("fixed_token_workload_valid") is not True: + raise ValueError(f"{path}: fixed-token validation failed") + proof_path = path.with_name("feature-proof.json") + proof = None + if variant != "llama": + if not proof_path.is_file(): + raise ValueError(f"{path}: missing feature-proof.json") + proof = json.loads(proof_path.read_text(encoding="utf-8")) + if proof.get("valid") is not True: + raise ValueError(f"{proof_path}: activation proof failed") + expected_by_variant = { + "ar": [], "ddtree": ["ddtree"], "pflash": ["pflash"], + "kvflash": ["kvflash"], "full": ["ddtree", "kvflash", "pflash"], + "speculation": ["chain"], + "adaptive-on": ["chain"], + "adaptive-confidence-off": ["chain"], + } + if variant not in expected_by_variant: + raise ValueError(f"{path}: unknown Lucebox variant {variant!r}") + if proof.get("expected_features") != expected_by_variant[variant]: + raise ValueError( + f"{proof_path}: expected_features does not match variant {variant}" + ) + mode_by_variant = { + "ar": "ar", "speculation": "speculation", + "adaptive-on": "adaptive", + "adaptive-confidence-off": "adaptive", + } + if variant in mode_by_variant: + mode = (meta.get("feature_config") or {}).get("decode_mode") + # `ar` is shared with the older Qwen3.6 matrix, whose metadata + # intentionally has no decode_mode field. + if ((variant != "ar" or mode is not None) + and (mode != mode_by_variant[variant] + or proof.get("decode_mode") != mode)): + raise ValueError( + f"{path}: decode_mode proof does not match variant {variant}" + ) + reports.append({ + "path": path, "report": report, "level": level, + "meta": meta, "proof": proof, + }) + if not reports: + raise ValueError(f"{root}: no bench.json files found") + return reports + + +def median(values: list[float]) -> float: + return statistics.median(values) + + +def fmt(value: float | None, digits: int = 2) -> str: + return f"{value:.{digits}f}" if value is not None else "n/a" + + +def complete_median(values: list[float | None]) -> float | None: + """Return a median only when every repeat measured the metric.""" + if not values or any(value is None for value in values): + return None + return median([value for value in values if value is not None]) + + +def run_signature(item: dict) -> tuple[object, ...]: + report, meta = item["report"], item["meta"] + max_tokens = report.get("max_tokens") + ignore_eos = report.get("ignore_eos") + temperature = report.get("temperature") + seed = report.get("seed") + model_sha256 = meta.get("model_sha256") + if ( + type(max_tokens) is not int or max_tokens <= 0 + or not isinstance(ignore_eos, bool) + or type(temperature) not in (int, float) + or type(seed) is not int + or not isinstance(model_sha256, str) or not model_sha256 + ): + raise ValueError("incomplete run metadata") + return max_tokens, ignore_eos, temperature, seed, model_sha256 + + +def output_stability(items: list[dict]) -> str: + output_digests = [ + item["level"].get("selected_output_set_sha256") for item in items + ] + complete = all(isinstance(value, str) and bool(value) for value in output_digests) + hashes = {value for value in output_digests if isinstance(value, str)} + return ( + "n/a" if len(items) < 2 or not complete + else "yes" if len(hashes) == 1 + else "NO" + ) + + +def summarize_qwen36(reports: list[dict]) -> str: + grouped: dict[tuple[str, int, str], list[dict]] = defaultdict(list) + for item in reports: + meta, level = item["meta"], item["level"] + key = (str(meta["workload"]), int(level["clients"]), str(meta["variant"])) + grouped[key].append(item) + for key, items in grouped.items(): + repeats = [int(item["meta"]["repeat"]) for item in items] + if len(repeats) != len(set(repeats)): + raise ValueError(f"{key}: duplicate repeat") + if len({run_signature(item) for item in items}) != 1: + raise ValueError(f"{key}: incompatible run metadata") + + lines = [ + "# Qwen3.6 concurrent feature matrix", "", + "Every Lucebox row is included only after request-correlated server telemetry " + "proves its requested features executed. Throughput is the median across fresh-process repeats.", + "", + "| Workload | C | Variant | N | Output goodput | Output-window | vs AR | " + "Effective/wire | DDTree accepted/step | DDTree steps/susp. | " + "Target forwards | KV in/out | PFlash requests | TTFT max s | " + "Stable output |", + "| :--- | ---: | :--- | ---: | ---: | ---: | ---: | ---: | ---: | " + ":--- | ---: | :--- | ---: | ---: | :---: |", + ] + for workload, clients, variant in sorted(grouped): + items = grouped[(workload, clients, variant)] + prompt_digests = [ + item["level"].get("selected_prompt_set_sha256") for item in items + ] + if not all(isinstance(value, str) and value for value in prompt_digests): + raise ValueError( + f"{workload} C={clients} {variant}: missing selected prompt set hash" + ) + prompt_hashes = set(prompt_digests) + if len(prompt_hashes) != 1: + raise ValueError(f"{workload} C={clients} {variant}: prompt sets differ") + goodput = median([item["level"]["aggregate_tok_s"] for item in items]) + window = complete_median([ + item["level"].get("output_window_tok_s") for item in items + ]) + ratio = complete_median([ + item["level"].get("effective_to_wire_prompt_ratio") for item in items + ]) + ttft = complete_median([ + item["level"].get("ttft_max_s") for item in items + ]) + stable = output_stability(items) + + peers = grouped.get((workload, clients, "ar"), []) + vs_ar = "—" if variant == "ar" else "n/a" + if variant != "ar" and peers: + peer_hashes = {p["level"]["selected_prompt_set_sha256"] for p in peers} + if peer_hashes != prompt_hashes: + raise ValueError(f"{workload} C={clients}: {variant}/ar prompts differ") + if {run_signature(item) for item in items} != { + run_signature(peer) for peer in peers + }: + raise ValueError( + f"{workload} C={clients}: {variant}/ar run metadata differs" + ) + by_repeat = {int(item["meta"]["repeat"]): item for item in items} + peers_by_repeat = {int(item["meta"]["repeat"]): item for item in peers} + if by_repeat.keys() != peers_by_repeat.keys(): + raise ValueError( + f"{workload} C={clients}: {variant}/ar repeat sets differ" + ) + if stable != "NO" and output_stability(peers) != "NO": + ratios = [] + for repeat in sorted(by_repeat): + value = by_repeat[repeat]["level"].get("aggregate_tok_s") + base = peers_by_repeat[repeat]["level"].get("aggregate_tok_s") + if value is None or base is None or base <= 0: + raise ValueError( + f"{workload} C={clients} repeat={repeat}: invalid AR goodput" + ) + ratios.append(value / base - 1.0) + vs_ar = f"{median(ratios) * 100:+.1f}%" + + proofs = [item["proof"] for item in items if item["proof"] is not None] + aggregates = [p["aggregate"] for p in proofs] + steps = sum(a["ddtree_steps"] for a in aggregates) + accepted = sum(a["ddtree_accepted_tokens"] for a in aggregates) + accepted_per_step = accepted / steps if steps else None + median_steps = median( + [a["ddtree_steps"] for a in aggregates] + ) if aggregates else None + median_suspensions = median( + [a["ddtree_suspensions"] for a in aggregates] + ) if aggregates else None + ddtree_activity = ( + f"{median_steps:.0f}/{median_suspensions:.0f}" + if median_steps is not None and median_suspensions is not None + else "n/a" + ) + target_forwards = median([a["target_forwards"] for a in aggregates]) if aggregates else None + page_ins = median([a["kvflash_page_ins"] for a in aggregates]) if aggregates else None + page_outs = median([a["kvflash_page_outs"] for a in aggregates]) if aggregates else None + pflash_requests = median([a["pflash_applied_requests"] for a in aggregates]) if aggregates else None + kv_text = ( + f"{page_ins:.0f}/{page_outs:.0f}" + if page_ins is not None and page_outs is not None else "n/a" + ) + lines.append( + f"| {workload} | {clients} | {variant} | {len(items)} | {goodput:.2f} | " + f"{fmt(window)} | {vs_ar} | {fmt(ratio, 3)} | {fmt(accepted_per_step)} | " + f"{ddtree_activity} | {fmt(target_forwards, 0)} | {kv_text} | " + f"{fmt(pflash_requests, 0)} | {fmt(ttft, 3)} | {stable} |" + ) + lines.append("") + return "\n".join(lines) + + + + +def summarize(reports: list[dict]) -> str: + return summarize_qwen36(reports) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("root", type=Path) + parser.add_argument("--out", type=Path) + args = parser.parse_args() + text = summarize(load_reports(args.root)) + if args.out: + args.out.write_text(text + "\n", encoding="utf-8") + print(text) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/harness/benchmarks/concurrency/test_analyze_dflash2_selector.py b/harness/benchmarks/concurrency/test_analyze_dflash2_selector.py new file mode 100644 index 000000000..d84475889 --- /dev/null +++ b/harness/benchmarks/concurrency/test_analyze_dflash2_selector.py @@ -0,0 +1,252 @@ +#!/usr/bin/env python3 +"""Tests for the offline DFlash2 selector/subset analyzer.""" + +from __future__ import annotations + +import importlib.util +import json +import math +import tempfile +import unittest +from pathlib import Path + + +HERE = Path(__file__).parent +SCRIPT = HERE / "analyze_dflash2_selector.py" +SPEC = importlib.util.spec_from_file_location("analyze_dflash2_selector", SCRIPT) +assert SPEC is not None and SPEC.loader is not None +analyzer = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(analyzer) + + +def wrapped(record: dict, line: int = 1) -> dict: + return { + "line_index": line, + "raw_json": json.dumps(record, separators=(",", ":")), + "record": record, + } + + +def depth_row(depth: int = 1, accepted: bool = True) -> dict: + return { + "depth": depth, + "accepted": accepted, + "selected_logp": -0.2 * depth, + "lm_margin": 2.0 / depth, + "topk_mass": 0.99, + "rank": 0, + "lm_top1": True, + "selector_margin": 3.0 / depth, + "selector_mass": 0.9, + "selector_entropy": 0.1, + } + + +def selector(engine: int, generated: int, accepted: int = 1) -> dict: + return { + "request_id": engine, + "slot": 1, + "generated": generated, + "accepted_depth": accepted, + "depths": [depth_row(1, accepted >= 1)], + } + + +def artifact() -> dict: + details = [ + { + "request_id": "wire-a", "request_index": 0, "prompt_index": 0, + "prompt_sha256": "a" * 64, "decode_mode": "ar", + "request_decode_tok_s": 20.0, + }, + { + "request_id": "wire-s", "request_index": 1, "prompt_index": 1, + "prompt_sha256": "b" * 64, "decode_mode": "speculation", + "request_decode_tok_s": 30.0, + }, + ] + metrics = [ + { + "request_id": "wire-a", "engine_request_id": 10, + "spec_steps": 0, "spec_accepted_tokens": 0, + }, + { + "request_id": "wire-s", "engine_request_id": 11, + "spec_steps": 2, "spec_accepted_tokens": 2, + }, + ] + selectors = [selector(11, 0), selector(11, 2)] + timings = [ + {"path": "spec", "live": 2, "k": 1, "accepted_tokens": 1, + "emitted_tokens": 3, "total_us": 1000.0}, + {"path": "ar", "live": 1, "k": 0, "accepted_tokens": 0, + "emitted_tokens": 1, "total_us": 1000.0}, + ] + return { + "kind": "dflash2-forced-subset-diagnostic", + "spec_depth": 2, + "validation": {"passed": True}, + "server_metadata": {"repeat": 1}, + "level": { + "clients": 2, + "request_mode_mask": "AS", + "requests_detail": details, + "selected_prompt_set_sha256": "c" * 64, + "aggregate_tok_s": 50.0, + "wall_s": 2.0, + }, + "server_records": { + "requests": [wrapped(row, index + 1) for index, row in enumerate(metrics)], + "selectors": [wrapped(row, index + 3) for index, row in enumerate(selectors)], + "rounds": [wrapped(row, index + 5) for index, row in enumerate(timings)], + "activations": [], + }, + } + + +def all_features(value: float) -> dict[str, float]: + return {key: value for key in analyzer.FEATURE_DIRECTIONS} + + +def subset_case(mask: str, goodput: float) -> dict: + requests = [] + for position, mode in enumerate(mask): + requests.append({ + "position": position, + "mode": "speculation" if mode == "S" else "ar", + "prompt_sha256": str(position), + "first_features": all_features(float(2 - position)) if mode == "S" else None, + "lifetime_yield_fraction": 1.0 if mode == "S" else None, + }) + return { + "clients": 2, + "spec_depth": 8, + "prompt_set_sha256": "set", + "wall_s": 1.0, + "round_timing": { + "all": {"goodput_tok_s": goodput}, + "full_live": {"goodput_tok_s": goodput}, + "tail": {"goodput_tok_s": None}, + }, + "mask": mask, + "aggregate_tok_s": goodput, + "requests": requests, + } + + +class DFlash2SelectorAnalyzerTests(unittest.TestCase): + def test_profile_parser_retains_selector_metric_and_timing_json(self) -> None: + data = ( + b'noise [spec-selector] {"request_id":7,"depths":[]}\n' + b'[concurrency-metrics] {"request_id":"wire","engine_request_id":7}\n' + b'[step-timing] {"path":"spec","live":2}\n' + ) + parsed = analyzer.parse_profile_lines(data) + self.assertEqual(parsed["selectors"][0]["record"]["request_id"], 7) + self.assertEqual( + parsed["requests"][0]["record"]["engine_request_id"], 7, + ) + self.assertEqual(parsed["rounds"][0]["record"]["live"], 2) + self.assertIn('"request_id":7', parsed["selectors"][0]["raw_json"]) + + def test_profile_parser_rejects_malformed_selector_json(self) -> None: + with self.assertRaisesRegex(ValueError, "invalid \\[spec-selector\\]"): + analyzer.parse_profile_lines(b"[spec-selector] {bad}\n") + + def test_artifact_joins_engine_selector_to_wire_request(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "bench.json" + path.write_text(json.dumps(artifact()), encoding="utf-8") + case = analyzer.analyze_artifact(path) + request = case["requests"][1] + self.assertEqual(request["wire_request_id"], "wire-s") + self.assertEqual(request["engine_request_id"], 11) + self.assertEqual(request["spec_steps"], 2) + self.assertEqual(request["lifetime_accepted_yield"], 1.0) + self.assertEqual(request["lifetime_yield_fraction"], 1.0) + self.assertAlmostEqual( + request["first_features"]["chain_lm_probability"], math.exp(-0.2), + ) + self.assertEqual(case["round_timing"]["full_live"]["goodput_tok_s"], 3000.0) + self.assertEqual(case["round_timing"]["tail"]["goodput_tok_s"], 1000.0) + self.assertEqual(case["round_timing"]["all"]["goodput_tok_s"], 2000.0) + + def test_artifact_accepts_direct_spec_timing(self) -> None: + value = artifact() + value["server_records"]["rounds"][0]["record"]["path"] = "spec-direct" + value["server_records"]["rounds"][0]["raw_json"] = json.dumps( + value["server_records"]["rounds"][0]["record"], + separators=(",", ":"), + ) + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "bench.json" + path.write_text(json.dumps(value), encoding="utf-8") + case = analyzer.analyze_artifact(path) + self.assertEqual( + case["round_timing"]["full_live"]["goodput_tok_s"], 3000.0, + ) + + def test_artifact_rejects_unknown_selector_engine_id(self) -> None: + value = artifact() + bad = selector(99, 0) + value["server_records"]["selectors"][0] = wrapped(bad) + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "bench.json" + path.write_text(json.dumps(value), encoding="utf-8") + with self.assertRaisesRegex(ValueError, "unknown engine request 99"): + analyzer.analyze_artifact(path) + + def test_artifact_rejects_acceptance_counter_mismatch(self) -> None: + value = artifact() + metric = value["server_records"]["requests"][1]["record"] + metric["spec_accepted_tokens"] = 1 + value["server_records"]["requests"][1] = wrapped(metric) + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "bench.json" + path.write_text(json.dumps(value), encoding="utf-8") + with self.assertRaisesRegex(ValueError, "acceptance does not match"): + analyzer.analyze_artifact(path) + + def test_calibration_reports_raw_probability_gap_per_depth(self) -> None: + rows = analyzer._calibration([(1, 2, "prompt", selector(1, 0, accepted=1))]) + self.assertEqual(len(rows), 1) + row = rows[0] + self.assertEqual(row["observed_survival"], 1.0) + self.assertFalse(row["has_both_labels"]) + self.assertEqual(row["unique_prompts"], 1) + self.assertEqual(row["prompts_with_rejection"], 0) + self.assertAlmostEqual( + row["selected_token_probability"]["mean_raw_probability"], + math.exp(-0.2), + ) + self.assertGreater( + row["selected_token_probability"]["observed_minus_raw"], 0.0, + ) + + def test_subset_oracle_exposes_homogeneous_synergy_and_sign_flip(self) -> None: + result = analyzer._subset_group([ + subset_case("AA", 100.0), + subset_case("AS", 70.0), + subset_case("SA", 80.0), + subset_case("SS", 130.0), + ]) + self.assertTrue(result["complete_exhaustive"]) + self.assertEqual(result["oracle_mask"], "SS") + self.assertTrue(result["oracle_is_homogeneous"]) + self.assertTrue(result["homogeneous_dominates_every_mixed"]) + self.assertEqual(result["requests_with_contextual_sign_flip"], 2) + self.assertAlmostEqual( + result["per_request_marginals"][0]["shapley_goodput_tok_s"], 20.0, + ) + self.assertAlmostEqual( + result["per_request_marginals"][1]["shapley_goodput_tok_s"], 10.0, + ) + ranked = result["raw_feature_prefix_rankings"]["chain_lm_logp"] + self.assertEqual(ranked["ranked_positions"], [0, 1]) + self.assertEqual(ranked["by_subset_size"][1]["ranked_prefix_mask"], "SA") + self.assertEqual(ranked["by_subset_size"][1]["same_size_regret"], 0.0) + self.assertFalse(result["first_feature_vs_shapley"]["identifiable"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/harness/benchmarks/concurrency/test_concurrency_tools.py b/harness/benchmarks/concurrency/test_concurrency_tools.py new file mode 100644 index 000000000..60bc925a9 --- /dev/null +++ b/harness/benchmarks/concurrency/test_concurrency_tools.py @@ -0,0 +1,266 @@ +#!/usr/bin/env python3 +"""Unit tests for the deterministic prompt generator and compact summarizer.""" + +from __future__ import annotations + +import importlib.util +import json +import tempfile +import unittest +from pathlib import Path + + +HERE = Path(__file__).parent + + +def load(name: str): + path = HERE / f"{name}.py" + spec = importlib.util.spec_from_file_location(name, path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +generator = load("generate_ragged_prompts") +summarizer = load("summarize_concurrency") + + +class PromptGeneratorTests(unittest.TestCase): + def test_cohorts_are_disjoint_ragged_and_mean_matched(self) -> None: + records = generator.build_records("short") + self.assertEqual(len(records), 29) + self.assertEqual( + [row["cohort"] for row in records], + ["c1"] + ["c4"] * 4 + ["c8"] * 8 + ["c16"] * 16, + ) + self.assertEqual(len({row["prompt"] for row in records}), 29) + by_cohort = { + cohort: [row for row in records if row["cohort"] == cohort] + for cohort in ("c1", "c4", "c8", "c16") + } + means = { + cohort: sum(row["target_words"] for row in rows) / len(rows) + for cohort, rows in by_cohort.items() + } + self.assertEqual(len(set(means.values())), 1) + for cohort in ("c4", "c8", "c16"): + self.assertEqual(len({row["target_words"] for row in by_cohort[cohort]}), 4) + for row in records: + self.assertEqual(len(row["prompt"].split()), row["target_words"]) + + +class RunnerShellTests(unittest.TestCase): + def test_runner_guards_case_identity_and_records_launch_environment(self) -> None: + runner = (HERE / "run_qwen36_concurrency.sh").read_text( + encoding="utf-8", + ) + self.assertIn( + 'reject_duplicates CLIENTS "${client_list[@]}"', runner, + ) + self.assertIn( + 'reject_duplicates VARIANTS "${variant_list[@]}"', runner, + ) + self.assertIn('printf \'env \' > "$case_dir/server-command.txt"', runner) + self.assertIn('"${launch_env[@]}" "${command[@]}"', runner) + self.assertIn("port_is_available || return 1", runner) + self.assertIn('wait_health "$model_id"', runner) + + +class SummarizerTests(unittest.TestCase): + @staticmethod + def item( + variant: str, + goodput: float, + output_window: float | None = None, + *, + repeat: int = 1, + output_hash: str | None = "same-outputs", + ) -> dict: + return { + "report": { + "max_tokens": 256, "ignore_eos": True, + "temperature": 0.0, "seed": 1, + }, + "meta": { + "workload": "short", "variant": variant, "repeat": repeat, + "model_sha256": "a" * 64, + }, + "level": { + "clients": 8, + "aggregate_tok_s": goodput, + "output_window_tok_s": output_window if output_window is not None else goodput, + "request_decode_tok_s_median": goodput / 8, + "prompt_tokens_per_s_to_first_token": 100.0, + "ttft_max_s": 2.0, + "selected_prompt_set_sha256": "same-prompts", + "selected_output_set_sha256": output_hash, + }, + } + + def test_summary_reports_product_and_packing_deltas(self) -> None: + text = summarizer.summarize([ + self.item("luce-k8", 20.0), + self.item("luce-k1", 10.0), + self.item("llama", 8.0), + ]) + self.assertIn("+150.0%", text) + self.assertIn("+100.0%", text) + self.assertIn("Decode vs llama", text) + + def test_summary_uses_same_repeat_ratios(self) -> None: + reports = [] + for repeat, luce, llama in ( + (1, 10.0, 1.0), + (2, 20.0, 90.0), + (3, 100.0, 50.0), + ): + reports.extend([ + self.item("luce-k8", luce, repeat=repeat), + self.item("llama", llama, repeat=repeat), + ]) + text = summarizer.summarize(reports) + luce_row = next(line for line in text.splitlines() if "| luce-k8 |" in line) + self.assertIn("+100.0%", luce_row) + self.assertNotIn("-60.0%", luce_row) + + def test_summary_rejects_mismatched_repeat_sets(self) -> None: + reports = [ + self.item("luce-k8", 20.0, repeat=1), + self.item("luce-k8", 22.0, repeat=2), + self.item("llama", 10.0, repeat=1), + ] + with self.assertRaisesRegex(ValueError, "repeat sets differ"): + summarizer.summarize(reports) + + def test_single_repeat_does_not_claim_stability(self) -> None: + text = summarizer.summarize([self.item("llama", 8.0)]) + row = next(line for line in text.splitlines() if "| llama |" in line) + self.assertEqual(row.split("|")[10].strip(), "n/a") + + def test_multiple_repeats_report_output_stability(self) -> None: + stable = summarizer.summarize([ + self.item("llama", 8.0, repeat=1), + self.item("llama", 9.0, repeat=2), + ]) + stable_row = next(line for line in stable.splitlines() if "| llama |" in line) + self.assertEqual(stable_row.split("|")[10].strip(), "yes") + + unstable = summarizer.summarize([ + self.item("llama", 8.0, repeat=1, output_hash="first"), + self.item("llama", 9.0, repeat=2, output_hash="second"), + ]) + unstable_row = next(line for line in unstable.splitlines() if "| llama |" in line) + self.assertEqual(unstable_row.split("|")[10].strip(), "NO") + + def test_unstable_current_variant_suppresses_deltas(self) -> None: + reports = [ + self.item("llama", 10.0, repeat=1), + self.item("llama", 10.0, repeat=2), + self.item("luce-k8", 20.0, repeat=1, output_hash="first"), + self.item("luce-k8", 22.0, repeat=2, output_hash="second"), + ] + text = summarizer.summarize(reports) + row = next(line for line in text.splitlines() if "| luce-k8 |" in line) + self.assertEqual(row.split("|")[10].strip(), "NO") + self.assertEqual(row.split("|")[11].strip(), "n/a") + self.assertEqual(row.split("|")[12].strip(), "n/a") + + def test_unstable_peer_variant_suppresses_deltas(self) -> None: + reports = [ + self.item( + "llama", 10.0, repeat=1, output_hash="first", + ), + self.item( + "llama", 10.0, repeat=2, output_hash="second", + ), + self.item("luce-k8", 20.0, repeat=1), + self.item("luce-k8", 22.0, repeat=2), + ] + text = summarizer.summarize(reports) + row = next(line for line in text.splitlines() if "| luce-k8 |" in line) + self.assertEqual(row.split("|")[10].strip(), "yes") + self.assertEqual(row.split("|")[11].strip(), "n/a") + self.assertEqual(row.split("|")[12].strip(), "n/a") + + + def test_missing_output_digest_does_not_claim_stability(self) -> None: + reports = [ + self.item("llama", 8.0, repeat=1, output_hash=None), + self.item("llama", 9.0, repeat=2, output_hash=None), + ] + text = summarizer.summarize(reports) + row = next(line for line in text.splitlines() if "| llama |" in line) + self.assertEqual(row.split("|")[10].strip(), "n/a") + + def test_incomplete_repeat_metrics_are_reported_as_na(self) -> None: + first = self.item("llama", 8.0, repeat=1) + second = self.item("llama", 9.0, repeat=2) + for metric in ( + "output_window_tok_s", "request_decode_tok_s_median", + "prompt_tokens_per_s_to_first_token", "ttft_max_s", + ): + second["level"][metric] = None + text = summarizer.summarize([first, second]) + row = next(line for line in text.splitlines() if "| llama |" in line) + self.assertEqual( + [row.split("|")[column].strip() for column in range(6, 10)], + ["n/a"] * 4, + ) + + def test_comparison_rejects_incompatible_run_metadata(self) -> None: + fields = ( + ("report", "max_tokens", 64), + ("report", "ignore_eos", False), + ("report", "temperature", 0.5), + ("report", "seed", 2), + ("meta", "model_sha256", "b" * 64), + ) + for container, key, value in fields: + with self.subTest(key=key): + luce = self.item("luce-k8", 20.0) + llama = self.item("llama", 10.0) + luce[container][key] = value + with self.assertRaisesRegex(ValueError, "run metadata differs"): + summarizer.summarize([luce, llama]) + + def test_summary_rejects_incomplete_run_metadata(self) -> None: + item = self.item("llama", 10.0) + del item["report"]["max_tokens"] + with self.assertRaisesRegex(ValueError, "incomplete run metadata"): + summarizer.summarize([item]) + + def test_summary_reports_descriptive_errors_for_missing_fields(self) -> None: + fields = ( + ("meta", "workload", "workload"), + ("level", "clients", "clients"), + ("level", "aggregate_tok_s", "aggregate token rate"), + ("level", "selected_prompt_set_sha256", "prompt set hash"), + ) + for container, key, message in fields: + with self.subTest(key=key): + item = self.item("llama", 10.0) + del item[container][key] + with self.assertRaisesRegex(ValueError, message): + summarizer.summarize([item]) + + def test_load_reports_rejects_missing_prompt_usage(self) -> None: + report = { + "ignore_eos": True, + "server_metadata": {"workload": "short", "variant": "llama", "repeat": 1}, + "levels": [{ + "failures": 0, + "token_count_complete": True, + "prompt_token_count_complete": False, + "fixed_token_workload_valid": True, + }], + } + with tempfile.TemporaryDirectory() as root: + path = Path(root) / "bench.json" + path.write_text(json.dumps(report), encoding="utf-8") + with self.assertRaisesRegex(ValueError, "incomplete token accounting"): + summarizer.load_reports(Path(root)) + + +if __name__ == "__main__": + unittest.main() diff --git a/harness/benchmarks/concurrency/test_concurrent_benchmark.py b/harness/benchmarks/concurrency/test_concurrent_benchmark.py new file mode 100644 index 000000000..da1251c79 --- /dev/null +++ b/harness/benchmarks/concurrency/test_concurrent_benchmark.py @@ -0,0 +1,147 @@ +#!/usr/bin/env python3 +"""Focused tests for concurrent_benchmark.py.""" + +from __future__ import annotations + +import argparse +import importlib.util +import json +import unittest +from pathlib import Path +from unittest import mock + + +SCRIPT = Path(__file__).with_name("concurrent_benchmark.py") +SPEC = importlib.util.spec_from_file_location("concurrent_benchmark", SCRIPT) +assert SPEC is not None and SPEC.loader is not None +benchmark = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(benchmark) + + +class BenchmarkTests(unittest.TestCase): + def test_sse_parser_handles_events_and_done(self) -> None: + lines = [ + b'data: {"choices":[{"delta":{"content":"hi"}}]}\n', b"\n", + b"data: [DONE]\n", b"\n", + ] + self.assertEqual( + list(benchmark.iter_sse_data(lines)), + ['{"choices":[{"delta":{"content":"hi"}}]}', "[DONE]"], + ) + + def test_prompt_selection_never_wraps(self) -> None: + self.assertEqual(benchmark.request_prompts(["a", "b", "c"], 2, 1), ["b", "c"]) + with self.assertRaisesRegex(ValueError, "refusing to reuse"): + benchmark.request_prompts(["a", "b"], 2, 1) + + def test_level_uses_exact_usage_and_first_token_window(self) -> None: + def fake_request(_args: argparse.Namespace, prompt: str) -> dict: + start, first, end, prompt_tokens = { + "first": (10.0, 12.0, 14.0, 10), + "second": (10.25, 11.25, 15.0, 30), + }[prompt] + return { + "t_start": start, "t_first": first, "t_end": end, + "duration_s": end - start, "ttft_s": first - start, + "decode_duration_s": end - first, + "completion_tokens": 8, "prompt_tokens": prompt_tokens, + "finish_reason": "length", "error": None, + "content_sha256": benchmark.sha256_text(prompt + " output"), + "reasoning_content_sha256": benchmark.sha256_text(""), + "content_chars": 6, "reasoning_content_chars": 0, + "request_output_tok_s": 8 / (end - start), + "request_decode_tok_s": 7 / (end - first), + } + + args = argparse.Namespace(max_tokens=8, ignore_eos=True, timeout=2.0) + with mock.patch.object(benchmark, "stream_request", side_effect=fake_request): + level = benchmark.run_level(2, args, ["first", "second"], 0) + self.assertEqual(level["completion_tokens_total"], 16) + self.assertEqual(level["prompt_tokens_total"], 40) + self.assertTrue(level["fixed_token_workload_valid"]) + # Assert independently known windows before their derived rates. + self.assertEqual(level["wall_s"], 5.0) + self.assertEqual(level["output_window_s"], 3.75) + self.assertEqual(level["prompt_to_first_token_s"], 2.0) + self.assertAlmostEqual(level["aggregate_tok_s"], 3.2) + self.assertAlmostEqual( + level["output_window_tok_s"], 16 / 3.75, + ) + self.assertAlmostEqual( + level["request_decode_tok_s_median"], (3.5 + 7 / 3.75) / 2, + ) + self.assertAlmostEqual(level["prompt_tokens_per_s_to_first_token"], 20.0) + + def test_stream_request_keeps_usage_separate_from_sse_chunks(self) -> None: + class Response: + def __enter__(self): return self + def __exit__(self, *_args): return None + def __iter__(self): + return iter([ + b'data: {"choices":[{"delta":{"content":"one chunk"}}]}\n', b"\n", + b'data: {"choices":[{"delta":{},"finish_reason":"length"}]}\n', b"\n", + b'data: {"choices":[],"usage":{"prompt_tokens":12,"completion_tokens":64}}\n', b"\n", + b"data: [DONE]\n", b"\n", + ]) + + args = argparse.Namespace( + model="m", max_tokens=64, temperature=0.0, seed=1, ignore_eos=True, + api_key="", base_url="http://localhost/v1", timeout=2.0, + ) + with mock.patch.object(benchmark.urllib.request, "urlopen", return_value=Response()): + record = benchmark.stream_request(args, "prompt") + self.assertEqual(record["completion_tokens"], 64) + self.assertEqual(record["prompt_tokens"], 12) + self.assertTrue(record["done_received"]) + self.assertIsNone(record["error"]) + self.assertIsNotNone(record["request_decode_tok_s"]) + self.assertEqual(record["content_sha256"], benchmark.sha256_text("one chunk")) + + def test_stream_request_rejects_clean_eof_without_done(self) -> None: + class Response: + def __enter__(self): return self + def __exit__(self, *_args): return None + def __iter__(self): + return iter([ + b'data: {"choices":[{"delta":{"content":"partial"}}]}\n', b"\n", + b'data: {"choices":[{"delta":{},"finish_reason":"length"}]}\n', b"\n", + b'data: {"choices":[],"usage":{"prompt_tokens":12,"completion_tokens":64}}\n', b"\n", + ]) + + args = argparse.Namespace( + model="m", max_tokens=64, temperature=0.0, seed=1, ignore_eos=True, + api_key="", base_url="http://localhost/v1", timeout=2.0, + ) + with mock.patch.object(benchmark.urllib.request, "urlopen", return_value=Response()): + record = benchmark.stream_request(args, "prompt") + self.assertFalse(record["done_received"]) + self.assertIn("before [DONE]", record["error"]) + + def test_missing_prompt_usage_fails_level(self) -> None: + level = { + "failures": 0, + "token_count_complete": True, + "prompt_token_count_complete": False, + "fixed_token_workload_valid": True, + } + self.assertTrue(benchmark.level_failed(level, ignore_eos=True)) + + def test_hung_worker_aborts_level_instead_of_overlapping_next(self) -> None: + thread = mock.Mock() + thread.is_alive.return_value = True + args = argparse.Namespace(timeout=1.0) + with ( + mock.patch.object(benchmark.threading, "Thread", return_value=thread), + mock.patch.object( + benchmark.time, "monotonic", side_effect=(10.0, 50.0), + ), + ): + with self.assertRaisesRegex(TimeoutError, "exceeded the level deadline"): + benchmark.run_level(1, args, ["prompt"], 0) + thread.start.assert_called_once_with() + thread.join.assert_called_once_with(0.0) + + + +if __name__ == "__main__": + unittest.main() diff --git a/harness/benchmarks/concurrency/test_feature_concurrent_benchmark.py b/harness/benchmarks/concurrency/test_feature_concurrent_benchmark.py new file mode 100644 index 000000000..d21be4450 --- /dev/null +++ b/harness/benchmarks/concurrency/test_feature_concurrent_benchmark.py @@ -0,0 +1,149 @@ +#!/usr/bin/env python3 +"""Focused tests for request-correlated feature benchmark telemetry.""" + +from __future__ import annotations + +import argparse +import hashlib +import importlib.util +import sys +import unittest +from pathlib import Path +from unittest import mock + + +HERE = Path(__file__).parent +sys.path.insert(0, str(HERE)) +SCRIPT = HERE / "concurrent_benchmark.py" +SPEC = importlib.util.spec_from_file_location("concurrent_benchmark", SCRIPT) +assert SPEC is not None and SPEC.loader is not None +benchmark = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(benchmark) + + +class FeatureBenchmarkTests(unittest.TestCase): + def test_stream_request_captures_id_and_effective_prompt(self) -> None: + class Response: + def __enter__(self): return self + def __exit__(self, *_args): return None + def __iter__(self): + return iter([ + b'data: {"id":"chatcmpl-42","choices":[{"delta":{"content":"x"}}]}\n', b"\n", + b'data: {"id":"chatcmpl-42","choices":[{"delta":{},' + b'"finish_reason":"length"}]}\n', b"\n", + b'data: {"id":"chatcmpl-42","choices":[],"usage":' + b'{"prompt_tokens":40000,"completion_tokens":8,"timings":' + b'{"effective_prompt_tokens":2000,"prefilled_tokens":2000,' + b'"cached_prefix_tokens":0,"cache_hit":false,"prefill_ms":12.5,' + b'"decode_ms":20.0,"decode_tokens_per_sec":400.0}}}\n', b"\n", + b"data: [DONE]\n", b"\n", + ]) + + args = argparse.Namespace( + model="m", max_tokens=8, temperature=0.0, seed=1, ignore_eos=True, + api_key="", base_url="http://localhost/v1", timeout=2.0, + ) + with mock.patch.object(benchmark.urllib.request, "urlopen", return_value=Response()): + row = benchmark.stream_request(args, "prompt") + self.assertEqual(row["request_id"], "chatcmpl-42") + self.assertEqual(row["prompt_tokens"], 40000) + self.assertEqual(row["effective_prompt_tokens"], 2000) + self.assertEqual(row["server_prefill_ms"], 12.5) + self.assertTrue(row["done_received"]) + self.assertIsNone(row["error"]) + + def test_clean_eof_without_done_is_rejected(self) -> None: + class Response: + def __enter__(self): return self + def __exit__(self, *_args): return None + def __iter__(self): + return iter([ + b'data: {"id":"chatcmpl-cut","choices":[{"delta":{"content":"x"},' + b'"finish_reason":"length"}]}\n', b"\n", + ]) + + args = argparse.Namespace( + model="m", max_tokens=8, temperature=0.0, seed=1, ignore_eos=True, + api_key="", base_url="http://localhost/v1", timeout=2.0, + ) + with mock.patch.object(benchmark.urllib.request, "urlopen", return_value=Response()): + row = benchmark.stream_request(args, "prompt") + self.assertIn("before [DONE]", row["error"]) + + def test_enrich_level_reports_compression_ratio(self) -> None: + level = { + "prompt_tokens_total": 40000, + "requests_detail": [{ + "request_id": "r1", "error": None, + "effective_prompt_tokens": 2000, + "server_prefill_ms": 10.0, "server_decode_ms": 20.0, + "server_decode_tokens_per_sec": 100.0, + }], + } + benchmark.enrich_level(level) + self.assertTrue(level["request_ids_complete"]) + self.assertTrue(level["effective_prompt_token_count_complete"]) + self.assertEqual(level["effective_prompt_tokens_total"], 2000) + self.assertEqual(level["effective_to_wire_prompt_ratio"], 0.05) + + def test_duplicate_request_ids_are_not_complete(self) -> None: + level = { + "prompt_tokens_total": 20, + "requests_detail": [ + {"request_id": "same", "error": None, "effective_prompt_tokens": 10}, + {"request_id": "same", "error": None, "effective_prompt_tokens": 10}, + ], + } + benchmark.enrich_level(level) + self.assertFalse(level["request_ids_complete"]) + + def test_boolean_wire_counts_are_not_accepted_as_integers(self) -> None: + level = { + "prompt_tokens_total": 1, + "requests_detail": [{ + "request_id": "r1", "error": None, + "effective_prompt_tokens": True, + }], + } + benchmark.enrich_level(level) + self.assertFalse(level["effective_prompt_token_count_complete"]) + + def test_client_provenance_records_exact_argv_and_source_digest(self) -> None: + argv = ["python3", "concurrent_benchmark.py", "--clients", "4"] + result = benchmark.client_provenance(argv) + self.assertEqual(result["client_argv"], argv) + self.assertEqual(result["client_script"], str(SCRIPT.resolve())) + self.assertEqual( + result["client_script_sha256"], + hashlib.sha256(SCRIPT.read_bytes()).hexdigest(), + ) + + def test_markdown_reports_prompt_rate_and_ttft_median_max(self) -> None: + text = benchmark.markdown({ + "label": "feature", + "levels": [{ + "clients": 4, + "requests_ok": 4, + "requests": 4, + "aggregate_tok_s": 12.5, + "output_window_tok_s": 20.0, + "request_decode_tok_s_median": 5.0, + "prompt_tokens_per_s_to_first_token": 123.4, + "prompt_tokens_min": 100, + "prompt_tokens_max": 400, + "effective_prompt_tokens_min": 50, + "effective_prompt_tokens_max": 200, + "effective_to_wire_prompt_ratio": 0.5, + "ttft_median_s": 1.25, + "ttft_max_s": 2.5, + }], + }) + self.assertIn("Prompt tok/s to first", text) + self.assertIn("TTFT median/max s", text) + self.assertIn("| 4 | 4/4 | 12.50 | 20.00 | 5.00 | 123.40 |", text) + self.assertIn("| 0.500 | 1.250/2.500 |", text) + + + +if __name__ == "__main__": + unittest.main() diff --git a/harness/benchmarks/concurrency/test_feature_metadata.py b/harness/benchmarks/concurrency/test_feature_metadata.py new file mode 100644 index 000000000..229685030 --- /dev/null +++ b/harness/benchmarks/concurrency/test_feature_metadata.py @@ -0,0 +1,184 @@ +#!/usr/bin/env python3 +"""Tests for literal feature flags and reproducibility metadata.""" + +from __future__ import annotations + +import hashlib +import importlib.util +import json +import sys +import tempfile +import unittest +from pathlib import Path +from unittest import mock + + +HERE = Path(__file__).parent +SCRIPT = HERE / "write_feature_metadata.py" +SPEC = importlib.util.spec_from_file_location("write_feature_metadata", SCRIPT) +assert SPEC is not None and SPEC.loader is not None +metadata = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(metadata) +RUNTIME_SCRIPT = HERE / "record_feature_runtime.py" +RUNTIME_SPEC = importlib.util.spec_from_file_location( + "record_feature_runtime", RUNTIME_SCRIPT, +) +assert RUNTIME_SPEC is not None and RUNTIME_SPEC.loader is not None +runtime_metadata = importlib.util.module_from_spec(RUNTIME_SPEC) +RUNTIME_SPEC.loader.exec_module(runtime_metadata) + + +class FeatureMetadataTests(unittest.TestCase): + def test_full_row_records_literal_screenshot_flags_and_hashes(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + target = root / "target.gguf" + draft = root / "draft.gguf" + prefill = root / "prefill.gguf" + prompts = root / "prompts.jsonl" + command = root / "command.txt" + out = root / "metadata.json" + target.write_bytes(b"target") + draft.write_bytes(b"draft") + prefill.write_bytes(b"prefill") + draft_sha = hashlib.sha256(draft.read_bytes()).hexdigest() + prefill_sha = hashlib.sha256(prefill.read_bytes()).hexdigest() + prompts.write_text('{"prompt":"p"}\n', encoding="utf-8") + command.write_text("server --target-device hip:0\n", encoding="utf-8") + argv = [ + str(SCRIPT), "--out", str(out), "--variant", "full", + "--workload", "compression", "--clients", "4", "--repeat", "1", + "--binary", "/bin/true", "--model", str(target), + "--model-sha256", hashlib.sha256(target.read_bytes()).hexdigest(), + "--prompt-file", str(prompts), "--command-file", str(command), + "--repo", str(HERE.parents[2]), "--max-concurrent-prefills", "8", + "--target-device", "hip:0", "--draft-device", "hip:0", + "--draft-model", str(draft), "--draft-model-sha256", draft_sha, + "--ddtree", "--ddtree-budget", "22", "--fast-rollback", + "--prefill-compression", "auto", "--prefill-threshold", "32000", + "--prefill-keep-ratio", "0.05", "--prefill-drafter", str(prefill), + "--prefill-drafter-sha256", prefill_sha, + "--draft-residency", "persistent", "--kvflash", "auto", + "--kvflash-max-pool-tokens", "8192", + "--kvflash-scorer-drafter", str(prefill), + "--kvflash-scorer-drafter-sha256", prefill_sha, + ] + with mock.patch.object(sys, "argv", argv): + self.assertEqual(metadata.main(), 0) + result = json.loads(out.read_text(encoding="utf-8")) + self.assertEqual(result["model_sha256"], hashlib.sha256(b"target").hexdigest()) + self.assertTrue(result["server_binary_sha256"]) + self.assertTrue(result["git_head"]) + self.assertEqual(result["feature_config"]["draft_model_sha256"], draft_sha) + self.assertEqual(result["feature_config"]["prefill_drafter_sha256"], prefill_sha) + self.assertEqual( + result["feature_config"]["kvflash_scorer_drafter"], + str(prefill.resolve()), + ) + self.assertEqual( + result["feature_config"]["kvflash_scorer_drafter_sha256"], + prefill_sha, + ) + self.assertEqual(result["literal_screenshot_flags"], [ + "--target-device", "hip:0", + "--draft-device", "hip:0", + "--ddtree", + "--ddtree-budget", "22", + "--fast-rollback", + "--draft-residency", "persistent", + "--prefill-compression", "auto", + "--prefill-drafter", str(prefill.resolve()), + "--kvflash", "auto", + ]) + self.assertIsNone(result["runtime_observed"]) + + def test_model_digest_claim_must_match_referenced_file(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + model = Path(tmp) / "model.gguf" + model.write_bytes(b"model") + expected = hashlib.sha256(b"model").hexdigest() + cache: dict[Path, str] = {} + self.assertEqual( + metadata.validated_digest(model, expected, "model", cache), + expected, + ) + with self.assertRaisesRegex(ValueError, "does not match"): + metadata.validated_digest(model, "stale", "model", cache) + with self.assertRaisesRegex(ValueError, "is required"): + metadata.validated_digest(model, None, "model", cache) + with self.assertRaisesRegex(ValueError, "without a model file"): + metadata.validated_digest(None, expected, "model", cache) + + + def test_ldd_failure_is_fatal(self) -> None: + failed = mock.Mock(returncode=1, stdout="", stderr="not a dynamic executable") + with mock.patch.object(metadata.subprocess, "run", return_value=failed): + with self.assertRaisesRegex(RuntimeError, "ldd failed"): + metadata.resolved_libraries(Path("/tmp/server")) + + def test_unresolved_shared_library_is_fatal(self) -> None: + unresolved = mock.Mock( + returncode=0, stdout="libmissing.so => not found\n", stderr="", + ) + with mock.patch.object(metadata.subprocess, "run", return_value=unresolved): + with self.assertRaisesRegex(RuntimeError, "unresolved libraries"): + metadata.resolved_libraries(Path("/tmp/server")) + + def test_git_revision_failure_and_empty_output_are_fatal(self) -> None: + for result, message in ( + (mock.Mock(returncode=128, stdout="", stderr="not a repository"), + "git rev-parse failed"), + (mock.Mock(returncode=0, stdout="\n", stderr=""), + "empty revision"), + ): + with self.subTest(message=message): + with mock.patch.object(metadata.subprocess, "run", return_value=result): + with self.assertRaisesRegex(RuntimeError, message): + metadata.repository_head(Path("/tmp/repo")) + + def test_runtime_records_actual_kvflash_pool_from_startup(self) -> None: + original = { + "schema_version": 3, + "feature_config": { + "kvflash": "auto", + "kvflash_max_pool_tokens": 16384, + }, + } + log = "\n".join(( + "[parallel-kvflash] physical resident pool 8192 tokens; " + "logical per-slot cap 65536 across 16 slots " + "(--kv-pool-tokens does not expand resident VRAM)", + "[paged-attention] 512 physical blocks x 16 tokens " + "(8192 pool tokens, per-sequence max_ctx 65536)", + )) + result = runtime_metadata.update_metadata(original, log) + observed = result["runtime_observed"] + self.assertTrue(observed["kvflash_active"]) + self.assertEqual(observed["physical_kv_pool_tokens"], 8192) + self.assertEqual(observed["physical_kv_pool_blocks"], 512) + self.assertEqual(observed["kv_block_size_tokens"], 16) + self.assertEqual(observed["logical_per_slot_max_ctx"], 65536) + self.assertEqual(observed["configured_slots"], 16) + + def test_runtime_rejects_enabled_kvflash_without_marker(self) -> None: + original = {"feature_config": {"kvflash": "auto"}} + with self.assertRaisesRegex(ValueError, "startup marker is missing"): + runtime_metadata.update_metadata( + original, + "[paged-attention] 512 physical blocks x 16 tokens " + "(8192 pool tokens, per-sequence max_ctx 65536)", + ) + + def test_runtime_rejects_kvflash_marker_without_paged_marker(self) -> None: + original = {"feature_config": {"kvflash": "auto"}} + log = ( + "[parallel-kvflash] physical resident pool 8192 tokens; " + "logical per-slot cap 65536 across 16 slots " + "(--kv-pool-tokens does not expand resident VRAM)" + ) + with self.assertRaisesRegex(ValueError, "paged physical-pool"): + runtime_metadata.update_metadata(original, log) + + +if __name__ == "__main__": + unittest.main() diff --git a/harness/benchmarks/concurrency/test_feature_tools.py b/harness/benchmarks/concurrency/test_feature_tools.py new file mode 100644 index 000000000..c93207dda --- /dev/null +++ b/harness/benchmarks/concurrency/test_feature_tools.py @@ -0,0 +1,1130 @@ +#!/usr/bin/env python3 +"""Tests for pressure prompts, activation proof, and feature summaries.""" + +from __future__ import annotations + +import importlib.util +import json +import math +import os +import subprocess +import sys +import tempfile +import unittest +from pathlib import Path + + +HERE = Path(__file__).parent +sys.path.insert(0, str(HERE)) + + +def load(name: str): + path = HERE / f"{name}.py" + spec = importlib.util.spec_from_file_location(name, path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +generator = load("generate_ragged_prompts") +proof = load("verify_feature_metrics") +summary = load("summarize_feature_matrix") +gate_analysis = load("analyze_gate_decisions") + + +def report( + variant: str = "full", workload: str = "compression", + effective: tuple[int, int] = (2000, 2100), +) -> dict: + return { + "server_metadata": { + "variant": variant, + "workload": workload, + "feature_config": { + "prefill_compression": "auto", + "prefill_threshold": 32000, + "kvflash": "auto", + "kvflash_max_pool_tokens": 8192, + "kvflash_scorer_drafter": "/models/Qwen3-0.6B-BF16.gguf", + "kvflash_scorer_drafter_sha256": "a" * 64, + }, + "runtime_observed": { + "kvflash_active": True, + "physical_kv_pool_tokens": 8192, + }, + }, + "levels": [{ + "requests_detail": [ + {"request_id": "r1", "error": None, + "prompt_tokens": 40000, + "effective_prompt_tokens": effective[0]}, + {"request_id": "r2", "error": None, + "prompt_tokens": 40000, + "effective_prompt_tokens": effective[1]}, + ], + }], + } + + +def metric(request_id: str, effective: int, page_outs: int = 1) -> dict: + return { + "request_id": request_id, + "effective_prompt_tokens": effective, + "ddtree_steps": 3, + "ddtree_suspensions": 0, + # Zero acceptance is legitimate and must not invalidate execution proof. + "ddtree_accepted_tokens": 0, + "spec_steps": 0, + "spec_accepted_tokens": 0, + "target_forwards": 3, + "kvflash_page_ins": 0, + "kvflash_page_outs": page_outs, + "kvflash_resident_blocks": 8, + "kvflash_reselects": 1, + "pflash_applied": True, + "pflash_input_tokens": 40000, + "pflash_output_tokens": effective, + } + + +class FeaturePromptTests(unittest.TestCase): + def test_activation_profiles_are_disjoint_and_above_thresholds(self) -> None: + compression = generator.build_records("compression") + pressure = generator.build_records("kv-pressure") + self.assertEqual(len(compression), 29) + self.assertEqual(len(pressure), 29) + self.assertEqual(len({row["prompt"] for row in compression}), 29) + self.assertEqual(len({row["prompt"] for row in pressure}), 29) + self.assertTrue( + {row["prompt"] for row in compression}.isdisjoint( + row["prompt"] for row in pressure + ) + ) + self.assertGreaterEqual(min(row["target_words"] for row in compression), 34000) + self.assertGreaterEqual(min(row["target_words"] for row in pressure), 12000) + self.assertTrue(all(row["activation_target"] == "pflash-auto" for row in compression)) + +class FeatureRunnerShellTests(unittest.TestCase): + def run_invalid_matrix( + self, tmp: str, **overrides: str, + ) -> subprocess.CompletedProcess[str]: + model = Path(tmp) / "model.gguf" + model.touch() + env = { + key: value for key, value in os.environ.items() + if not key.startswith(("GGML_", "DFLASH_", "LUCE_", "HIP_", "ROCR_", "HSA_")) + and key not in ("LD_PRELOAD", "LD_LIBRARY_PATH") + } + env.update({ + "MODEL": str(model), + "LUCE_SERVER_BIN": "/bin/true", + "OUT": str(Path(tmp) / "out"), + "VARIANTS": "ar", + **overrides, + }) + return subprocess.run( + ["bash", str(HERE / "run_qwen36_feature_matrix.sh")], + env=env, capture_output=True, text=True, check=False, + ) + + def test_client_and_proof_invocations_are_array_backed(self) -> None: + runner = (HERE / "run_qwen36_feature_matrix.sh").read_text( + encoding="utf-8", + ) + client_lines = [ + line.strip() for line in runner.splitlines() + if 'python3 "$CLIENT"' in line + ] + self.assertEqual(client_lines, [ + 'python3 "$CLIENT" "${common_client[@]}"', + 'python3 "$CLIENT" "${common_client[@]}"', + ]) + self.assertIn('local -a warmup_cmd=(', runner) + self.assertIn('local -a benchmark_cmd=(', runner) + self.assertIn( + '"${warmup_cmd[@]}" > "$case_dir/warmup.txt"', + runner, + ) + self.assertIn( + '"${benchmark_cmd[@]}" | tee "$case_dir/bench.txt"', + runner, + ) + + proof_lines = [ + line.strip() for line in runner.splitlines() + if 'python3 "$PROOF_TOOL"' in line + ] + self.assertEqual(proof_lines, ['python3 "$PROOF_TOOL"']) + self.assertIn('local -a proof_cmd=(', runner) + self.assertIn('"${proof_cmd[@]}"', runner) + + def test_signals_exit_and_launch_environment_is_recorded(self) -> None: + runner = (HERE / "run_qwen36_feature_matrix.sh").read_text(encoding="utf-8") + self.assertIn("trap stop_server EXIT", runner) + self.assertIn("trap 'exit 130' INT", runner) + self.assertIn("trap 'exit 143' TERM", runner) + self.assertNotIn("trap stop_server EXIT INT TERM", runner) + self.assertIn("'env ' > \"$case_dir/server-command.txt\"", runner) + self.assertIn('"${launch_env[@]}" "${command[@]}"', runner) + + def test_duplicate_clients_are_rejected_before_artifacts_are_created(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + result = self.run_invalid_matrix(tmp, CLIENTS="4,4") + self.assertEqual(result.returncode, 2, result.stderr) + self.assertIn("CLIENTS contains duplicate entry: 4", result.stderr) + self.assertFalse((Path(tmp) / "out").exists()) + + def test_duplicate_variants_are_rejected_before_artifacts_are_created(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + result = self.run_invalid_matrix(tmp, VARIANTS="ar,ar") + self.assertEqual(result.returncode, 2, result.stderr) + self.assertIn("VARIANTS contains duplicate entry: ar", result.stderr) + self.assertFalse((Path(tmp) / "out").exists()) + + def test_llama_only_does_not_require_lucebox_binary(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + result = self.run_invalid_matrix( + tmp, VARIANTS="llama", LUCE_SERVER_BIN="/does/not/exist", + LLAMA_SERVER_BIN="/bin/true", REPEATS="0", + ) + self.assertEqual(result.returncode, 2, result.stderr) + self.assertIn("REPEATS must be positive", result.stderr) + self.assertNotIn("missing Lucebox server", result.stderr) + + + +class GateAnalysisTests(unittest.TestCase): + @staticmethod + def _scored( + request_id: int, + slot: int, + score: float, + expected_yield: float, + decision: str, + score_kind: str = "test_score", + hazards: list[float] | None = None, + ) -> dict: + return { + "request_id": request_id, + "slot": slot, + "activation_score": score, + "score_kind": score_kind, + "expected_yield": expected_yield, + "hazards": [] if hazards is None else hazards, + "evaluation": "scored", + "fallback_reason": None, + "decision_reason": ( + "selected_by_joint_goodput" + if decision == "speculation" else "ar_counterfactual_won" + ), + "decision": decision, + } + + @staticmethod + def _failed( + request_id: int, slot: int, score_kind: str = "unspecified", + ) -> dict: + return { + "request_id": request_id, + "slot": slot, + "activation_score": None, + "score_kind": score_kind, + "expected_yield": None, + "hazards": None, + "evaluation": "failed", + "fallback_reason": "activation_evaluation_failed", + "decision_reason": "evaluation_failed", + "decision": "ar", + } + + @staticmethod + def _write_activation_case( + root: Path, + variant: str, + activations: list[dict], + engine_ids: tuple[int, ...] = (7, 8), + execution: dict[int, tuple[int, int]] | None = None, + ) -> Path: + case = root / "selection" / "c2" / "r1" / variant + case.mkdir(parents=True) + details = [ + { + "request_id": f"wire-{index}", + "prompt_index": index - 1, + "request_decode_tok_s": 8.0, + "content_sha256": f"hash-{index}", + } + for index in range(1, len(engine_ids) + 1) + ] + (case / "bench.json").write_text(json.dumps({ + "server_metadata": { + "workload": "selection", "variant": variant, + "clients": len(engine_ids), "repeat": 1, + }, + "levels": [{ + "aggregate_tok_s": 9.0, + "clients": len(engine_ids), + "requests_detail": details, + }], + }), encoding="utf-8") + metric_rows = [ + { + "request_id": f"wire-{index}", + "engine_request_id": engine_id, + "spec_accepted_tokens": 0, + "spec_steps": (execution or {}).get(engine_id, (0, 1))[0], + "target_forwards": (execution or {}).get( + engine_id, (0, 1) + )[1], + "output_tokens": 1, + } + for index, engine_id in enumerate(engine_ids, 1) + ] + lines = [ + f"[spec-activation] {json.dumps(row)}" for row in activations + ] + [ + f"[concurrency-metrics] {json.dumps(row)}" for row in metric_rows + ] + (case / "benchmark-server.log").write_text( + "\n".join(lines) + "\n", encoding="utf-8", + ) + return case + + def test_measured_step_timing_is_joined_and_summarized(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + prompt_dir = root / "prompts" + prompt_dir.mkdir() + (prompt_dir / "selection.jsonl").write_text( + json.dumps({ + "id": "friendly", + "prompt": "hello", + "selection_class": "speculation_strong_win", + "expected_dense_oracle": "speculation", + }) + "\n", + encoding="utf-8", + ) + case = root / "selection" / "c1" / "r1" / "adaptive-on" + case.mkdir(parents=True) + (case / "bench.json").write_text(json.dumps({ + "server_metadata": { + "workload": "selection", + "variant": "adaptive-on", + "clients": 1, + "repeat": 1, + }, + "levels": [{ + "aggregate_tok_s": 9.0, + "clients": 1, + "requests_detail": [{ + "request_id": "wire-1", + "prompt_index": 0, + "request_decode_tok_s": 8.0, + "content_sha256": "abc", + }], + }], + }), encoding="utf-8") + timing = { + "path": "ar", "live": 1, "k": 0, "decode_bucket": 1, + "n_prefill": 0, "max_kv_len": 10, + "draft_us": 20.0, "draft_lanes": 1, + "pre_us": 30.0, "graph_build_us": 5.0, + "graph_prepare_us": 5.0, "graph_exec_us": 40.0, + "sample_read_us": 10.0, "finish_us": 10.0, + "total_us": 100.0, "accepted_tokens": 0, + "emitted_tokens": 1, "target_forwards": 1, + } + metric_row = { + "request_id": "wire-1", "engine_request_id": 7, + "spec_accepted_tokens": 0, "spec_steps": 0, + "target_forwards": 1, "output_tokens": 1, + } + activation = self._scored(7, 0, 3.25, 1.75, "ar") + (case / "benchmark-server.log").write_text( + "[spec-gate] C=1 k=0 scores=[7:1.000/fresh/test_score] " + "sources=fresh:1,initial:0,unavailable:0 " + "G(k)=0.010000 G(0)=0.020000 " + "predicted_cost=50.0us measured=ar-path\n" + f"[spec-activation] {json.dumps(activation)}\n" + f"[step-timing] {json.dumps(timing)}\n" + f"[concurrency-metrics] {json.dumps(metric_row)}\n", + encoding="utf-8", + ) + + report = gate_analysis.analyze_case(case) + self.assertEqual(report["log"], str(case / "benchmark-server.log")) + self.assertEqual(report["requests"][0]["prompt_id"], "friendly") + self.assertEqual(report["timing"]["by_path"]["ar"]["rounds"], 1) + self.assertAlmostEqual( + report["timing"]["draft_tax_on_ar"]["fraction_of_ar_wall"], + 0.2, + ) + self.assertAlmostEqual( + report["gate_by_k"][0]["realized_goodput_tok_s"], 10000.0, + ) + self.assertEqual(report["gate_timing_count_mismatch"], 0) + self.assertEqual(report["activation"]["validation"], "passed") + self.assertEqual(report["activation"]["records"], 1) + request = report["requests"][0] + self.assertEqual(request["activation_slot"], 0) + self.assertEqual(request["activation_score"], 3.25) + self.assertEqual(request["expected_yield"], 1.75) + self.assertEqual(request["activation_decision"], "ar") + self.assertEqual(request["activation_evaluation"], "scored") + self.assertIsNone(request["activation_fallback_reason"]) + + def test_adaptive_on_activation_proof_fails_closed(self) -> None: + activation_7 = self._scored(7, 0, 2.0, 1.25, "speculation") + activation_8 = self._scored(8, 1, 1.0, 1.0, "ar") + activation_9 = {**activation_8, "request_id": 9} + cases = ( + ("missing", [activation_7], "missing activations.*8"), + ( + "duplicate", [activation_7, activation_7, activation_8], + "duplicate activations.*7", + ), + ( + "unknown", [activation_7, activation_8, activation_9], + "unknown engine requests.*9", + ), + ) + for label, activations, message in cases: + with self.subTest(label=label), tempfile.TemporaryDirectory() as tmp: + case = self._write_activation_case( + Path(tmp), "adaptive-on", activations, + ) + with self.assertRaisesRegex(ValueError, message): + gate_analysis.analyze_case(case) + + def test_confidence_off_is_exempt_from_activation_coverage(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + case = self._write_activation_case( + Path(tmp), "adaptive-confidence-off", [], + ) + report = gate_analysis.analyze_case(case) + self.assertFalse(report["activation"]["required"]) + self.assertEqual(report["activation"]["validation"], "not-required") + self.assertEqual(report["activation"]["records"], 0) + self.assertTrue(all( + row["activation_decision"] is None + for row in report["requests"] + )) + + def test_adaptive_on_activation_proof_enforces_sticky_execution(self) -> None: + activation = self._scored(7, 0, 2.0, 1.5, "speculation") + with tempfile.TemporaryDirectory() as tmp: + case = self._write_activation_case( + Path(tmp), "adaptive-on", [activation], + engine_ids=(7,), execution={7: (2, 4)}, + ) + report = gate_analysis.analyze_case(case) + self.assertEqual( + report["activation"]["execution_validation"], "passed", + ) + bad_cases = ( + ( + {**activation, "decision": "ar"}, {7: (1, 2)}, + "AR activation.*executed 1 speculation steps", + ), + ( + activation, {7: (2, 5)}, + "Spec activation.*contains a non-speculative target step", + ), + ( + activation, {7: (0, 1)}, + "Spec activation.*contains a non-speculative target step", + ), + ) + for row, execution, message in bad_cases: + with self.subTest(message=message), tempfile.TemporaryDirectory() as tmp: + case = self._write_activation_case( + Path(tmp), "adaptive-on", [row], + engine_ids=(7,), execution=execution, + ) + with self.assertRaisesRegex(ValueError, message): + gate_analysis.analyze_case(case) + + def test_failed_evaluation_activation_is_request_local_and_explicit( + self, + ) -> None: + failed = self._failed(7, 0) + scored = self._scored(8, 1, 2.0, 1.5, "speculation") + with tempfile.TemporaryDirectory() as tmp: + case = self._write_activation_case( + Path(tmp), "adaptive-on", [failed, scored], + execution={7: (0, 1), 8: (2, 4)}, + ) + report = gate_analysis.analyze_case(case) + self.assertEqual( + report["activation"]["evaluation_counts"], + {"scored": 1, "failed": 1}, + ) + by_id = { + row["engine_request_id"]: row for row in report["requests"] + } + self.assertEqual(by_id[7]["activation_decision"], "ar") + self.assertEqual(by_id[7]["activation_evaluation"], "failed") + self.assertIsNone(by_id[7]["activation_score"]) + self.assertEqual( + by_id[7]["activation_fallback_reason"], + "activation_evaluation_failed", + ) + self.assertEqual(by_id[8]["activation_evaluation"], "scored") + + + def test_dflash2_activation_fields_and_gate_score_kind_are_preserved( + self, + ) -> None: + score_kind = "qwen38-dflash2-selector-benefit-v1" + scored = self._scored( + 7, 0, 5.3306, 5.3306, "speculation", score_kind, [0.1, 0.2], + ) + failed = self._failed(8, 1, score_kind) + with tempfile.TemporaryDirectory() as tmp: + case = self._write_activation_case( + Path(tmp), "adaptive-on", [scored, failed], + execution={7: (2, 4), 8: (0, 1)}, + ) + log_path = case / "benchmark-server.log" + log_path.write_text( + "[spec-gate] C=2 k=1 " + "scores=[7:5.331/fresh/qwen38-dflash2-selector-benefit-v1*," + "8:2.685/initial/qwen38-dflash2-selector-benefit-v1] " + "G(k)=0.010 G(0)=0.009 predicted_cost=1us measured=2us\n" + + log_path.read_text(encoding="utf-8"), + encoding="utf-8", + ) + report = gate_analysis.analyze_case(case) + self.assertEqual( + report["activation"]["score_kind_counts"], {score_kind: 2}, + ) + self.assertEqual( + report["activation"]["fallback_reason_counts"], + {"activation_evaluation_failed": 1}, + ) + by_id = { + row["engine_request_id"]: row for row in report["requests"] + } + request = by_id[7] + self.assertAlmostEqual(request["activation_score"], 5.3306) + self.assertEqual(request["activation_hazards"], [0.1, 0.2]) + self.assertEqual(request["activation_score_kind"], score_kind) + self.assertEqual( + request["activation_decision_reason"], + "selected_by_joint_goodput", + ) + self.assertAlmostEqual(request["mean_activation_score"], 5.331) + rounds, _, _, _ = gate_analysis.parse_server_log(log_path) + self.assertEqual(rounds[0]["entries"][0]["score_kind"], score_kind) + + + + def test_failed_activation_uses_one_generic_fallback_reason(self) -> None: + base = self._failed(7, 0) + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "server.log" + path.write_text( + f"[spec-activation] {json.dumps(base)}\n", encoding="utf-8", + ) + _, _, _, activations = gate_analysis.parse_server_log(path) + self.assertEqual(activations, [base]) + no_adapter = {**base, "decision_reason": "no_speculator_adapter"} + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "server.log" + path.write_text( + f"[spec-activation] {json.dumps(no_adapter)}\n", + encoding="utf-8", + ) + _, _, _, activations = gate_analysis.parse_server_log(path) + self.assertEqual(activations, [no_adapter]) + for reason in ("", "benefit_evaluation_failed", "unknown_failure"): + with self.subTest(reason=reason), tempfile.TemporaryDirectory() as tmp: + row = {**base, "fallback_reason": reason} + path = Path(tmp) / "server.log" + path.write_text( + f"[spec-activation] {json.dumps(row)}\n", encoding="utf-8", + ) + with self.assertRaisesRegex(ValueError, "fallback_reason"): + gate_analysis.parse_server_log(path) + + + + def test_activation_record_fields_are_strictly_validated(self) -> None: + valid = self._scored(7, 0, 1.0, 1.0, "ar") + failed = self._failed(7, 0) + missing_evaluation = dict(valid) + missing_evaluation.pop("evaluation") + missing_fallback_reason = dict(valid) + missing_fallback_reason.pop("fallback_reason") + missing_failed_score = dict(failed) + missing_failed_score.pop("activation_score") + cases = ( + ({**valid, "request_id": True}, "request_id"), + ({**valid, "slot": -1}, "slot"), + ({**valid, "activation_score": 0.99}, "activation_score"), + ({**valid, "expected_yield": math.nan}, "expected_yield"), + ({**valid, "hazards": [1.1]}, "hazards"), + ({**valid, "score_kind": ""}, "score_kind"), + ({**valid, "decision": "undecided"}, "decision"), + (missing_evaluation, "evaluation"), + (missing_fallback_reason, "fallback_reason"), + (missing_failed_score, "activation_score"), + ({**valid, "fallback_reason": "failure"}, "fallback_reason"), + ({**failed, "activation_score": 1.0}, "scores must be null"), + ({**failed, "hazards": []}, "scores must be null"), + ({**failed, "decision": "speculation"}, "decision must be ar"), + ({**failed, "fallback_reason": "draft_failed"}, "fallback_reason"), + ) + for row, message in cases: + with self.subTest(field=message), tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "server.log" + path.write_text( + f"[spec-activation] {json.dumps(row)}\n", encoding="utf-8", + ) + with self.assertRaisesRegex(ValueError, message): + gate_analysis.parse_server_log(path) + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "server.log" + path.write_text( + f"[spec-activation] {json.dumps(failed)}\n", encoding="utf-8", + ) + _, _, _, activations = gate_analysis.parse_server_log(path) + self.assertEqual(activations, [failed]) + + + def test_paired_controls_produce_a_per_prompt_concurrent_oracle(self) -> None: + base_request = { + "prompt_id": "p1", + "selection_class": "speculation_strong_win", + "expected_dense_oracle": "speculation", + "dense_r9700_baseline": {"spec_over_ar": 1.4}, + "output_sha256": "same", + } + cases = [] + for variant, rate in ( + ("ar", 10.0), ("speculation", 15.0), ("adaptive-on", 14.0), + ): + cases.append({ + "workload": "selection", "clients": 3, "repeat": 1, + "variant": variant, + "requests": [{ + **base_request, + "decode_tok_s": rate, + "admitted_fraction": 0.75, + "mean_confidence_yield": 2.5, + "commit_per_spec_step": 2.0, + }], + }) + comparison = gate_analysis.compare_prompts(cases)[0] + self.assertEqual( + comparison["empirical_concurrent_oracle"], "speculation", + ) + self.assertAlmostEqual(comparison["speculation_over_ar"], 1.5) + self.assertTrue(comparison["matches_expected_dense_oracle"]) + self.assertTrue(comparison["output_stable"]) + self.assertEqual( + comparison["adaptive"]["adaptive-on"]["admitted_fraction"], 0.75, + ) + + def test_activation_shapes_are_compared_at_the_same_live_concurrency(self) -> None: + def shape(path: str, live: int, k: int, rate: float, draft: float) -> dict: + return { + "path": path, "live": live, "k": k, "rounds": 4, + "round_goodput_tok_s": rate, "mean_total_us": 100.0, + "phase_mean_us": {"draft_us": draft}, + } + + cases = [ + { + "workload": "selection", "clients": 3, "repeat": 1, + "variant": "ar", "aggregate_tok_s": 60.0, + "timing": {"by_shape": [ + shape("ar", 2, 0, 50.0, 0.0), + shape("ar", 3, 0, 70.0, 0.0), + ]}, + }, + { + "workload": "selection", "clients": 3, "repeat": 1, + "variant": "adaptive-on", "aggregate_tok_s": 42.0, + "timing": {"by_shape": [ + shape("ar", 2, 0, 40.0, 10.0), + shape("spec-direct", 3, 2, 49.0, 15.0), + ]}, + }, + ] + rows = gate_analysis.compare_activation_shapes(cases) + self.assertEqual(len(rows), 2) + self.assertAlmostEqual(cases[1]["aggregate_over_ar"], 0.7) + self.assertEqual(rows[0]["activation_outcome"], "ar-with-draft-tax") + self.assertEqual(rows[1]["activation_outcome"], "unprofitable") + self.assertAlmostEqual(rows[1]["realized_over_pure_ar"], 0.7) + + +class FeatureProofTests(unittest.TestCase): + def test_full_below_pool_passes_without_page_traffic(self) -> None: + rows = [ + metric("warmup", 50), + metric("r1", 2000, page_outs=0), + metric("r2", 2100, page_outs=0), + ] + result = proof.verify( + report(), rows, {"ddtree", "pflash", "kvflash"}, + ) + self.assertTrue(result["valid"], result["errors"]) + self.assertFalse(result["kvflash_page_traffic_required"]) + self.assertEqual( + result["kvflash_page_traffic_reason"], + "compressed-prompt-fits-pool", + ) + self.assertEqual(result["aggregate"]["ddtree_accepted_tokens"], 0) + self.assertEqual(result["aggregate"]["ddtree_suspensions"], 0) + self.assertEqual(result["matched_metric_count"], 2) + + def test_ddtree_suspensions_required_and_aggregated(self) -> None: + missing = metric("r1", 2000) + del missing["ddtree_suspensions"] + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "server.log" + path.write_text( + proof.PREFIX + json.dumps(missing) + "\n", + encoding="utf-8", + ) + with self.assertRaisesRegex( + ValueError, "missing telemetry keys.*ddtree_suspensions", + ): + proof.parse_markers(path) + + rows = [metric("r1", 2000), metric("r2", 2100)] + rows[1]["ddtree_suspensions"] = 1 + result = proof.verify( + report(), rows, {"ddtree", "pflash", "kvflash"}, + ) + self.assertTrue(result["valid"], result["errors"]) + self.assertEqual(result["aggregate"]["ddtree_suspensions"], 1) + self.assertEqual( + [row["ddtree_suspensions"] for row in result["requests"]], + [0, 1], + ) + + def test_ddtree_suspensions_must_be_binary_per_request(self) -> None: + raw_two = metric("r1", 2000) + raw_two["ddtree_suspensions"] = 2 + with self.assertRaisesRegex( + ValueError, "ddtree_suspensions must be 0 or 1 per request", + ): + proof.aggregate_rows([raw_two]) + + def test_duplicate_terminal_telemetry_is_rejected(self) -> None: + row = metric("r1", 2000) + with self.assertRaisesRegex(ValueError, "duplicate telemetry request ID r1"): + proof.aggregate_rows([row, dict(row)]) + + def test_boolean_telemetry_counts_are_rejected(self) -> None: + integer_fields = (*proof.COUNTERS, "effective_prompt_tokens", + "kvflash_resident_blocks", "pflash_input_tokens", + "pflash_output_tokens") + for key in integer_fields: + with self.subTest(key=key): + row = metric("r1", 2000) + row[key] = True + with self.assertRaisesRegex(ValueError, key): + proof.aggregate_rows([row]) + + def test_boolean_wire_and_metadata_counts_do_not_prove_features(self) -> None: + rows = [metric("r1", 2000), metric("r2", 2100)] + wire = report() + wire["levels"][0]["requests_detail"][0]["effective_prompt_tokens"] = True + result = proof.verify(wire, rows, set()) + self.assertFalse(result["valid"]) + self.assertIn("missing usage.timings.effective_prompt_tokens", result["errors"][0]) + + threshold = report() + threshold["server_metadata"]["feature_config"]["prefill_threshold"] = True + result = proof.verify(threshold, rows, {"pflash"}) + self.assertFalse(result["valid"]) + self.assertIn("positive recorded token threshold", "\n".join(result["errors"])) + + pool = report() + pool["server_metadata"]["runtime_observed"]["physical_kv_pool_tokens"] = True + result = proof.verify(pool, rows, {"kvflash"}) + self.assertFalse(result["valid"]) + self.assertIn("positive startup-observed physical pool", "\n".join(result["errors"])) + + requested = report() + requested["server_metadata"]["feature_config"]["kvflash_max_pool_tokens"] = True + result = proof.verify(requested, rows, {"kvflash"}) + self.assertFalse(result["valid"]) + self.assertIn("recorded pool-token cap", "\n".join(result["errors"])) + + def test_full_above_pool_requires_page_traffic(self) -> None: + rows = [ + metric("r1", 9000, page_outs=0), + metric("r2", 9100, page_outs=0), + ] + result = proof.verify( + report(effective=(9000, 9100)), rows, + {"ddtree", "pflash", "kvflash"}, + ) + self.assertFalse(result["valid"]) + self.assertTrue(result["kvflash_page_traffic_required"]) + self.assertIn( + "no page-in/page-out", "\n".join(result["errors"]) + ) + + def test_full_aggregate_effective_demand_requires_page_traffic(self) -> None: + rows = [ + metric("r1", 5000, page_outs=0), + metric("r2", 5000, page_outs=0), + ] + result = proof.verify( + report(effective=(5000, 5000)), rows, + {"ddtree", "pflash", "kvflash"}, + ) + self.assertFalse(result["valid"]) + self.assertTrue(result["kvflash_page_traffic_required"]) + self.assertEqual( + result["kvflash_page_traffic_reason"], + "effective-prompt-exceeds-pool", + ) + self.assertIn("no page-in/page-out", "\n".join(result["errors"])) + + def test_full_effective_demand_is_evaluated_per_client_level(self) -> None: + input_report = report(effective=(5000, 5000)) + requests = input_report["levels"][0]["requests_detail"] + input_report["levels"] = [ + {"requests_detail": [requests[0]]}, + {"requests_detail": [requests[1]]}, + ] + rows = [ + metric("r1", 5000, page_outs=0), + metric("r2", 5000, page_outs=0), + ] + result = proof.verify( + input_report, rows, {"ddtree", "pflash", "kvflash"}, + ) + self.assertTrue(result["valid"], result["errors"]) + self.assertFalse(result["kvflash_page_traffic_required"]) + self.assertEqual( + result["kvflash_page_traffic_reason"], + "compressed-prompt-fits-pool", + ) + + def test_unknown_kvflash_variant_fails_closed(self) -> None: + rows = [ + metric("r1", 2000, page_outs=0), + metric("r2", 2100, page_outs=0), + ] + result = proof.verify( + report(variant="typo"), rows, {"kvflash"}, + ) + self.assertFalse(result["valid"]) + self.assertTrue(result["kvflash_page_traffic_required"]) + self.assertEqual(result["kvflash_page_traffic_reason"], "unknown-variant") + self.assertIn("no page-in/page-out", "\n".join(result["errors"])) + + def test_pflash_auto_requires_measured_input_above_threshold(self) -> None: + rows = [metric("r1", 2000), metric("r2", 2100)] + rows[0]["pflash_input_tokens"] = 31999 + result = proof.verify(report(), rows, {"pflash"}) + self.assertFalse(result["valid"]) + self.assertIn( + "did not reach its recorded token threshold", + "\n".join(result["errors"]), + ) + + def test_kvflash_only_cannot_succeed_without_page_traffic(self) -> None: + rows = [ + metric("r1", 2000, page_outs=0), + metric("r2", 2100, page_outs=0), + ] + result = proof.verify( + report(variant="kvflash"), rows, {"kvflash"}, + ) + self.assertFalse(result["valid"]) + self.assertTrue(result["kvflash_page_traffic_required"]) + self.assertEqual( + result["kvflash_page_traffic_reason"], + "kvflash-only-ablation", + ) + self.assertIn( + "no page-in/page-out", "\n".join(result["errors"]) + ) + + def test_kvflash_requires_explicit_hashed_scorer(self) -> None: + input_report = report(variant="kvflash") + config = input_report["server_metadata"]["feature_config"] + config["kvflash_scorer_drafter"] = None + config["kvflash_scorer_drafter_sha256"] = None + rows = [metric("r1", 9000), metric("r2", 9100)] + result = proof.verify(input_report, rows, {"kvflash"}) + self.assertFalse(result["valid"]) + text = "\n".join(result["errors"]) + self.assertIn("no explicit scorer drafter", text) + self.assertIn("not a valid SHA-256 digest", text) + + def test_kvflash_rejects_malformed_scorer_hash(self) -> None: + rows = [metric("r1", 9000), metric("r2", 9100)] + for digest in ("not-a-hash", "g" * 64, "a" * 63, "a" * 65): + with self.subTest(digest=digest): + input_report = report(variant="kvflash") + input_report["server_metadata"]["feature_config"][ + "kvflash_scorer_drafter_sha256" + ] = digest + result = proof.verify(input_report, rows, {"kvflash"}) + self.assertFalse(result["valid"]) + self.assertIn( + "not a valid SHA-256 digest", "\n".join(result["errors"]) + ) + + def test_kvflash_requires_startup_observed_pool(self) -> None: + input_report = report(variant="kvflash") + input_report["server_metadata"]["runtime_observed"] = None + rows = [metric("r1", 9000), metric("r2", 9100)] + result = proof.verify(input_report, rows, {"kvflash"}) + self.assertFalse(result["valid"]) + text = "\n".join(result["errors"]) + self.assertIn("startup marker was not recorded", text) + self.assertIn("no positive startup-observed physical pool", text) + + def test_requested_features_cannot_succeed_silently(self) -> None: + rows = [metric("r1", 9000, page_outs=0), metric("r2", 9100, page_outs=0)] + rows[0]["ddtree_steps"] = 0 + rows[1]["pflash_applied"] = False + result = proof.verify( + report(effective=(9000, 9100)), rows, + {"ddtree", "pflash", "kvflash"}, + ) + self.assertFalse(result["valid"]) + text = "\n".join(result["errors"]) + self.assertIn("ddtree_steps is zero", text) + self.assertIn("pflash_applied is false", text) + self.assertIn("no page-in/page-out", text) + + def test_pflash_input_mismatch_with_wire_tokens_fails(self) -> None: + rows = [metric("r1", 2000), metric("r2", 2100)] + rows[0]["pflash_input_tokens"] = 39999 + result = proof.verify(report(), rows, {"pflash"}) + self.assertFalse(result["valid"]) + self.assertIn( + "pflash_input_tokens=39999 does not match wire value 40000", + "\n".join(result["errors"]), + ) + + def test_effective_prompt_mismatch_fails(self) -> None: + rows = [metric("r1", 1999), metric("r2", 2100)] + result = proof.verify(report(), rows, set()) + self.assertFalse(result["valid"]) + self.assertIn("does not match wire value", result["errors"][0]) + + def test_measured_request_requires_explicit_error_status(self) -> None: + input_report = report() + del input_report["levels"][0]["requests_detail"][0]["error"] + with self.assertRaisesRegex(ValueError, "explicit error status"): + proof.measured_requests(input_report) + + def test_forced_chain_requires_steps_and_matching_startup_proof(self) -> None: + input_report = report(variant="speculation") + input_report["server_metadata"]["feature_config"]["decode_mode"] = ( + "speculation" + ) + rows = [metric("r1", 2000), metric("r2", 2100)] + for row in rows: + row["ddtree_steps"] = 0 + row["ddtree_accepted_tokens"] = 0 + row["spec_steps"] = 3 + row["spec_accepted_tokens"] = 2 + startup = ( + "[parallel-chain] speculator=test-score " + "decode_mode=speculation draft=q4-mix-compatible" + ) + result = proof.verify(input_report, rows, {"chain"}, startup) + self.assertTrue(result["valid"], result["errors"]) + self.assertEqual(result["decode_mode"], "speculation") + self.assertEqual(result["aggregate"]["spec_steps"], 6) + self.assertEqual(result["aggregate"]["spec_accepted_tokens"], 4) + + rows[0]["spec_steps"] = 0 + rows[0]["spec_accepted_tokens"] = 0 + result = proof.verify(input_report, rows, {"chain"}, startup) + self.assertFalse(result["valid"]) + self.assertIn("spec_steps is zero", "\n".join(result["errors"])) + + result = proof.verify(input_report, rows[1:], {"chain"}, "") + self.assertFalse(result["valid"]) + self.assertIn("startup proof is missing", "\n".join(result["errors"])) + + def test_adaptive_chain_allows_ar_argmax_but_requires_profile(self) -> None: + input_report = report(variant="adaptive-on") + input_report["server_metadata"]["feature_config"]["decode_mode"] = ( + "adaptive" + ) + rows = [metric("r1", 2000), metric("r2", 2100)] + for row in rows: + row["ddtree_steps"] = 0 + row["ddtree_accepted_tokens"] = 0 + startup = "\n".join(( + "[spec-profile] context=4096 reps=5 mode=batched-draft", + "[parallel-chain] speculator=test-score " + "", + )) + result = proof.verify(input_report, rows, {"chain"}, startup) + self.assertTrue(result["valid"], result["errors"]) + + no_profile = startup.splitlines()[1] + result = proof.verify(input_report, rows, {"chain"}, no_profile) + self.assertFalse(result["valid"]) + self.assertIn("cost-profile proof is missing", "\n".join(result["errors"])) + + def test_ar_decode_mode_rejects_speculation_activity(self) -> None: + input_report = report(variant="ar") + input_report["server_metadata"]["feature_config"]["decode_mode"] = "ar" + rows = [metric("r1", 2000), metric("r2", 2100)] + rows[0]["spec_steps"] = 1 + result = proof.verify(input_report, rows, set()) + self.assertFalse(result["valid"]) + self.assertIn( + "AR decode_mode emitted chain speculation", "\n".join(result["errors"]) + ) + + def test_spec_acceptance_without_step_is_rejected(self) -> None: + row = metric("r1", 2000) + row["spec_accepted_tokens"] = 1 + with self.assertRaisesRegex( + ValueError, "spec_accepted_tokens requires positive spec_steps", + ): + proof.aggregate_rows([row]) + + +class FeatureSummaryTests(unittest.TestCase): + @staticmethod + def item( + variant: str, goodput: float, *, repeat: int = 1, + output_hash: str | None = "same-output", + ) -> dict: + return { + "report": { + "max_tokens": 256, "ignore_eos": True, + "temperature": 0.0, "seed": 1, + }, + "meta": { + "workload": "compression", "variant": variant, "repeat": repeat, + "model_sha256": "a" * 64, + }, + "level": { + "clients": 8, + "aggregate_tok_s": goodput, + "output_window_tok_s": goodput, + "effective_to_wire_prompt_ratio": 0.05 if variant == "full" else 1.0, + "ttft_max_s": 2.0, + "selected_prompt_set_sha256": "same-prompts", + "selected_output_set_sha256": output_hash, + }, + "proof": { + "aggregate": { + "ddtree_steps": 4 if variant == "full" else 0, + "ddtree_suspensions": 1 if variant == "full" else 0, + "ddtree_accepted_tokens": 8 if variant == "full" else 0, + "target_forwards": 4, + "kvflash_page_ins": 2 if variant == "full" else 0, + "kvflash_page_outs": 3 if variant == "full" else 0, + "pflash_applied_requests": 8 if variant == "full" else 0, + }, + }, + } + def test_summary_compares_feature_row_to_ar(self) -> None: + text = summary.summarize([self.item("ar", 10.0), self.item("full", 12.0)]) + self.assertIn("+20.0%", text) + self.assertIn("2.00", text) + self.assertIn("DDTree steps/susp.", text) + self.assertIn("| 4/1 | 4 | 2/3 |", text) + + def test_feature_row_without_ar_control_reports_na(self) -> None: + text = summary.summarize([self.item("full", 12.0)]) + row = next(line for line in text.splitlines() if "| full |" in line) + self.assertEqual(row.split("|")[7].strip(), "n/a") + + def test_missing_output_digest_does_not_claim_stability(self) -> None: + first = self.item("full", 12.0, output_hash=None) + second = self.item("full", 13.0, repeat=2, output_hash=None) + text = summary.summarize([first, second]) + row = next(line for line in text.splitlines() if "| full |" in line) + self.assertEqual(row.split("|")[15].strip(), "n/a") + + def test_incomplete_repeat_metrics_are_reported_as_na(self) -> None: + first = self.item("full", 12.0) + second = self.item("full", 13.0, repeat=2) + second["level"]["output_window_tok_s"] = None + second["level"]["effective_to_wire_prompt_ratio"] = None + text = summary.summarize([first, second]) + row = next(line for line in text.splitlines() if "| full |" in line) + self.assertEqual(row.split("|")[6].strip(), "n/a") + self.assertEqual(row.split("|")[8].strip(), "n/a") + + def test_incomplete_repeat_ttft_is_reported_as_na(self) -> None: + first = self.item("full", 12.0) + second = self.item("full", 13.0, repeat=2) + second["level"]["ttft_max_s"] = None + text = summary.summarize([first, second]) + row = next(line for line in text.splitlines() if "| full |" in line) + self.assertEqual(row.split("|")[14].strip(), "n/a") + + def test_summary_rejects_incompatible_repeat_metadata(self) -> None: + first = self.item("full", 12.0) + second = self.item("full", 13.0, repeat=2) + second["report"]["max_tokens"] = 128 + with self.assertRaisesRegex(ValueError, "incompatible run metadata"): + summary.summarize([first, second]) + + def test_summary_rejects_incompatible_ar_control_metadata(self) -> None: + ar = self.item("ar", 10.0) + feature = self.item("full", 12.0) + feature["meta"]["model_sha256"] = "b" * 64 + with self.assertRaisesRegex(ValueError, "run metadata differs"): + summary.summarize([ar, feature]) + + def test_summary_rejects_sampling_mismatch(self) -> None: + for field, value in (("temperature", 0.5), ("seed", 2)): + with self.subTest(field=field): + ar = self.item("ar", 10.0) + feature = self.item("full", 12.0) + feature["report"][field] = value + with self.assertRaisesRegex(ValueError, "run metadata differs"): + summary.summarize([ar, feature]) + + def test_summary_rejects_missing_prompt_hash(self) -> None: + for value in (None, ""): + with self.subTest(value=value): + feature = self.item("full", 12.0) + feature["level"]["selected_prompt_set_sha256"] = value + with self.assertRaisesRegex( + ValueError, "missing selected prompt set hash", + ): + summary.summarize([feature]) + + + def test_unstable_feature_row_suppresses_ar_delta(self) -> None: + reports = [ + self.item("ar", 10.0, repeat=1), + self.item("ar", 10.0, repeat=2), + self.item("full", 12.0, repeat=1, output_hash="first"), + self.item("full", 13.0, repeat=2, output_hash="second"), + ] + text = summary.summarize(reports) + row = next(line for line in text.splitlines() if "| full |" in line) + self.assertEqual(row.split("|")[7].strip(), "n/a") + self.assertEqual(row.split("|")[15].strip(), "NO") + + def test_unstable_ar_control_suppresses_feature_delta(self) -> None: + reports = [ + self.item("ar", 10.0, repeat=1, output_hash="first"), + self.item("ar", 10.0, repeat=2, output_hash="second"), + self.item("full", 12.0, repeat=1), + self.item("full", 13.0, repeat=2), + ] + text = summary.summarize(reports) + row = next(line for line in text.splitlines() if "| full |" in line) + self.assertEqual(row.split("|")[7].strip(), "n/a") + +if __name__ == "__main__": + unittest.main() diff --git a/harness/benchmarks/concurrency/test_forced_subset_benchmark.py b/harness/benchmarks/concurrency/test_forced_subset_benchmark.py new file mode 100644 index 000000000..6f18e6f03 --- /dev/null +++ b/harness/benchmarks/concurrency/test_forced_subset_benchmark.py @@ -0,0 +1,284 @@ +#!/usr/bin/env python3 +"""Tests for the forced DFlash2 subset/depth diagnostic client.""" + +from __future__ import annotations + +import argparse +import importlib.util +import json +import sys +import tempfile +import unittest +from pathlib import Path +from unittest import mock + + +HERE = Path(__file__).parent +sys.path.insert(0, str(HERE)) +SCRIPT = HERE / "forced_subset_benchmark.py" +SPEC = importlib.util.spec_from_file_location("forced_subset_benchmark", SCRIPT) +assert SPEC is not None and SPEC.loader is not None +benchmark = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(benchmark) + + +def wrapped(record: dict, line: int = 1) -> dict: + raw = json.dumps(record, separators=(",", ":")) + return {"line_index": line, "raw_json": raw, "record": record} + + +def request(request_id: str, mode: str) -> dict: + return { + "request_id": request_id, + "decode_mode": mode, + "error": None, + "content_sha256": "a" * 64, + "reasoning_content_sha256": "b" * 64, + "combined_output_sha256": "c" * 64, + } + + +def valid_level() -> dict: + return { + "requests": 2, + "requests_ok": 2, + "failures": 0, + "token_count_complete": True, + "prompt_token_count_complete": True, + "fixed_token_workload_valid": True, + "start_skew_s": 0.001, + "requests_detail": [request("wire-ar", "ar"), request("wire-spec", "speculation")], + } + + +def valid_records() -> dict: + return { + "rounds": [ + wrapped({"path": "ar", "live": 2, "k": 0}, 1), + wrapped({ + "path": "spec", "live": 2, "k": 1, + "tree_bucket": 2, "tree_rows": 8, + }, 2), + ], + "selectors": [wrapped({"request_id": 7, "depths": []}, 3)], + "activations": [], + "requests": [ + wrapped({"request_id": "wire-ar", "spec_steps": 0}, 4), + wrapped({"request_id": "wire-spec", "spec_steps": 3}, 5), + ], + } + + +class ForcedSubsetTests(unittest.TestCase): + def test_request_payload_forces_mode_zero_temperature_and_ignore_eos(self) -> None: + captured: dict = {} + + class Response: + def __enter__(self): return self + def __exit__(self, *_args): return None + def __iter__(self): + return iter([ + b'data: {"id":"wire-1","choices":[{"delta":{"content":"x"}}]}\n', + b"\n", + b'data: {"id":"wire-1","choices":[{"delta":{"content":"y"}}]}\n', + b"\n", + b'data: {"id":"wire-1","choices":[{"delta":{},' + b'"finish_reason":"length"}]}\n', + b"\n", + b'data: {"id":"wire-1","choices":[],"usage":' + b'{"prompt_tokens":3,"completion_tokens":8}}\n', + b"\n", + b"data: [DONE]\n", b"\n", + ]) + + def open_request(http_request, timeout): + captured.update(json.loads(http_request.data)) + self.assertEqual(timeout, 2.0) + return Response() + + args = argparse.Namespace( + model="m", max_tokens=8, seed=1, api_key="", + base_url="http://localhost/v1", timeout=2.0, + ) + with mock.patch.object( + benchmark.urllib.request, "urlopen", side_effect=open_request, + ): + row = benchmark.stream_request(args, "prompt", "speculation") + self.assertEqual(captured["decode_mode"], "speculation") + self.assertEqual(captured["temperature"], 0.0) + self.assertIs(captured["ignore_eos"], True) + self.assertEqual(row["request_id"], "wire-1") + self.assertEqual(row["decode_mode"], "speculation") + self.assertEqual(len(row["combined_output_sha256"]), 64) + self.assertIsNotNone(row["first_to_second_output_event_s"]) + self.assertGreaterEqual(row["first_to_second_output_event_s"], 0.0) + self.assertIsNone(row["error"]) + + def test_modes_are_positional_and_reject_adaptive(self) -> None: + self.assertEqual( + benchmark.parse_request_modes("ar,speculation", 2), + ["ar", "speculation"], + ) + with self.assertRaisesRegex(ValueError, "adaptive is intentionally out of scope"): + benchmark.parse_request_modes("adaptive", 1) + with self.assertRaisesRegex(ValueError, "exactly one"): + benchmark.parse_request_modes("ar", 2) + + def test_depth_and_profiling_environment_are_required_in_metadata(self) -> None: + metadata = { + "clients": 4, + "launch_environment": { + "DFLASH_SPEC_CHAIN_DEPTH": "8", + "DFLASH_STEP_TIMING": "1", + "DFLASH_DFLASH2_SELECTOR_LOG": "1", + "PROMPT_OFFSET": "5", + }, + } + benchmark.validate_server_metadata(metadata, 4, 8, 5, True) + metadata["launch_environment"]["DFLASH_SPEC_CHAIN_DEPTH"] = "4" + with self.assertRaisesRegex(ValueError, "DFLASH_SPEC_CHAIN_DEPTH=8"): + benchmark.validate_server_metadata(metadata, 4, 8, 5, True) + metadata["launch_environment"]["DFLASH_SPEC_CHAIN_DEPTH"] = "8" + with self.assertRaisesRegex(ValueError, "PROMPT_OFFSET=4"): + benchmark.validate_server_metadata(metadata, 4, 8, 4, True) + + def test_ar_only_metadata_does_not_require_selector_logging(self) -> None: + metadata = { + "clients": 1, + "launch_environment": { + "DFLASH_SPEC_CHAIN_DEPTH": "4", + "DFLASH_STEP_TIMING": "1", + "PROMPT_OFFSET": "0", + }, + } + benchmark.validate_server_metadata(metadata, 1, 4, 0, False) + + def test_profile_parser_retains_raw_round_and_request_records(self) -> None: + data = ( + b'prefix [step-timing] {"path":"spec","live":2,"k":1}\n' + b'[spec-selector] {"request_id":7,"depths":[]}\n' + b'[concurrency-metrics] {"request_id":"wire","spec_steps":1}\n' + ) + records = benchmark.parse_profile_records(data) + self.assertEqual(records["rounds"][0]["record"]["live"], 2) + self.assertEqual(records["selectors"][0]["record"]["request_id"], 7) + self.assertEqual(records["requests"][0]["record"]["request_id"], "wire") + self.assertIn('"path":"spec"', records["rounds"][0]["raw_json"]) + + def test_profile_parser_rejects_malformed_measured_json(self) -> None: + with self.assertRaisesRegex(ValueError, r"invalid \[step-timing\]"): + benchmark.parse_profile_records(b"[step-timing] {bad}\n") + + def test_valid_mixed_subset_proves_depth_and_sustained_full_live(self) -> None: + result = benchmark.validate_evidence( + valid_level(), valid_records(), 2, ["ar", "speculation"], 4, + max_start_skew_ms=100.0, min_full_live_rounds=2, + ) + self.assertTrue(result["passed"], result["errors"]) + self.assertEqual(result["executed_spec_depths"], [4]) + self.assertEqual(result["longest_full_live_streak"], 2) + self.assertIs(result["adaptive_claims_permitted"], False) + + def test_direct_mixed_subset_infers_depth_without_ar_prefix_rows(self) -> None: + records = valid_records() + records["rounds"][1]["record"].update({ + "path": "spec-direct", "ar_lanes": 1, "tree_rows": 9, + }) + result = benchmark.validate_evidence( + valid_level(), records, 2, ["ar", "speculation"], 4, + max_start_skew_ms=100.0, min_full_live_rounds=2, + ) + self.assertTrue(result["passed"], result["errors"]) + self.assertEqual(result["executed_spec_depths"], [4]) + + def test_missing_requested_live_concurrency_fails_closed(self) -> None: + records = valid_records() + for row in records["rounds"]: + row["record"]["live"] = 1 + result = benchmark.validate_evidence( + valid_level(), records, 2, ["ar", "speculation"], 4, + max_start_skew_ms=100.0, min_full_live_rounds=1, + ) + self.assertFalse(result["passed"]) + self.assertTrue(any("sustained live=C proof absent" in error for error in result["errors"])) + + def test_executed_depth_mismatch_fails_closed(self) -> None: + records = valid_records() + records["rounds"][1]["record"]["tree_rows"] = 16 + result = benchmark.validate_evidence( + valid_level(), records, 2, ["ar", "speculation"], 4, + max_start_skew_ms=100.0, min_full_live_rounds=2, + ) + self.assertFalse(result["passed"]) + self.assertTrue(any("do not match requested depth 4" in error for error in result["errors"])) + + def test_request_mode_execution_mismatch_fails_closed(self) -> None: + records = valid_records() + records["requests"][0]["record"]["spec_steps"] = 2 + result = benchmark.validate_evidence( + valid_level(), records, 2, ["ar", "speculation"], 4, + max_start_skew_ms=100.0, min_full_live_rounds=2, + ) + self.assertFalse(result["passed"]) + self.assertTrue(any("forced AR executed speculation" in error for error in result["errors"])) + + def test_log_span_is_exact_and_rejects_truncation(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "server.log" + path.write_bytes(b"warmup\nmeasured\n") + data, end = benchmark.read_log_span(path, len(b"warmup\n")) + self.assertEqual(data, b"measured\n") + self.assertEqual(end, path.stat().st_size) + with self.assertRaisesRegex(ValueError, "truncated or rotated"): + benchmark.read_log_span(path, end + 1) + + def test_ar_only_evidence_passes_without_selector_records(self) -> None: + level = { + "requests": 1, "requests_ok": 1, "failures": 0, + "token_count_complete": True, + "prompt_token_count_complete": True, + "fixed_token_workload_valid": True, + "start_skew_s": 0.0, + "requests_detail": [request("wire-ar", "ar")], + } + records = { + "rounds": [ + wrapped({"path": "ar", "live": 1, "k": 0}, 1), + wrapped({"path": "ar", "live": 1, "k": 0}, 2), + ], + "selectors": [], "activations": [], + "requests": [wrapped({"request_id": "wire-ar", "spec_steps": 0}, 3)], + } + result = benchmark.validate_evidence( + level, records, 1, ["ar"], 4, + max_start_skew_ms=100.0, min_full_live_rounds=2, + ) + self.assertTrue(result["passed"], result["errors"]) + + def test_runner_pins_one_server_and_dflash2_metadata_across_masks(self) -> None: + runner = (HERE / "run_qwen38_dflash2_subsets.sh").read_text( + encoding="utf-8", + ) + self.assertEqual( + runner.count( + 'env "${launch_env[@]}" "${command[@]}" > "$OUT/server.log"' + ), + 1, + ) + self.assertIn('"DFLASH_SPEC_CHAIN_DEPTH=$SPEC_DEPTH"', runner) + self.assertIn('"DFLASH_DFLASH2_SELECTOR_LOG=1"', runner) + self.assertIn('PROMPT_OFFSET="${PROMPT_OFFSET:-0}"', runner) + self.assertIn('"PROMPT_OFFSET=$PROMPT_OFFSET"', runner) + self.assertIn('--prompt-offset "$PROMPT_OFFSET"', runner) + self.assertNotIn('--prompt-offset 0', runner) + self.assertIn('--draft-model "$DRAFT_MODEL"', runner) + self.assertIn('--decode-mode speculation --host', runner) + self.assertIn('--decode-mode speculation --cache-type-k', runner) + self.assertIn('VISIBLE_DEVICES="${VISIBLE_DEVICES:-0}"', runner) + self.assertNotIn('--decode-mode ar', runner) + self.assertIn('run_client_case "$case_dir" "$mask"', runner) + self.assertNotIn("dspark", runner.lower()) + + +if __name__ == "__main__": + unittest.main() diff --git a/harness/benchmarks/concurrency/test_refill_subset_benchmark.py b/harness/benchmarks/concurrency/test_refill_subset_benchmark.py new file mode 100644 index 000000000..b4e30884e --- /dev/null +++ b/harness/benchmarks/concurrency/test_refill_subset_benchmark.py @@ -0,0 +1,452 @@ +#!/usr/bin/env python3 +"""Tests for the fail-closed DFlash2 refill diagnostic.""" + +from __future__ import annotations + +import argparse +import importlib.util +import json +import sys +import threading +import time +import unittest +from pathlib import Path +from unittest import mock + + +HERE = Path(__file__).parent +sys.path.insert(0, str(HERE)) +SCRIPT = HERE / "refill_subset_benchmark.py" +SPEC = importlib.util.spec_from_file_location("refill_subset_benchmark", SCRIPT) +assert SPEC is not None and SPEC.loader is not None +benchmark = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(benchmark) + + +def wrapped(record: dict, line: int) -> dict: + raw = json.dumps(record, separators=(",", ":")) + return {"line_index": line, "raw_json": raw, "record": record} + + +def valid_workload() -> dict: + requests = [] + for wave in range(3): + for lane in range(2): + digest = ("a" if lane == 0 else "b") * 64 + requests.append({ + "request_id": f"wire-{wave}-{lane}", + "lane_index": lane, + "wave_index": wave, + "decode_mode": "ar" if lane == 0 else "speculation", + "error": None, + "content_sha256": digest, + "reasoning_content_sha256": "c" * 64, + "combined_output_sha256": "d" * 64, + }) + return { + "clients": 2, + "waves": 3, + "requests": 6, + "requests_ok": 6, + "failures": 0, + "missing_requests": 0, + "token_count_complete": True, + "prompt_token_count_complete": True, + "fixed_token_workload_valid": True, + "initial_start_skew_s": 0.001, + "refill_handoffs": 4, + "refill_gap_s_max": 0.001, + "exact_output_stable_per_lane": True, + "lanes": [ + {"lane_index": 0, "decode_mode": "ar", "requests": 3}, + {"lane_index": 1, "decode_mode": "speculation", "requests": 3}, + ], + "requests_detail": requests, + } + + +def valid_records() -> dict: + rounds = [ + wrapped({ + "path": "spec", "live": 2, "k": 1, + "tree_bucket": 1, "tree_rows": 4, + "total_us": 20_000.0, "emitted_tokens": 2, + }, 1), + wrapped({ + "path": "spec", "live": 2, "k": 1, + "tree_bucket": 1, "tree_rows": 4, + "total_us": 20_000.0, "emitted_tokens": 2, + }, 4), + wrapped({ + "path": "spec", "live": 2, "k": 1, + "tree_bucket": 1, "tree_rows": 4, + "total_us": 20_000.0, "emitted_tokens": 2, + }, 7), + ] + metrics = [] + selectors = [] + line_by_wave = (2, 5, 8) + for wave, line in enumerate(line_by_wave): + for lane in range(2): + engine_id = wave * 2 + lane + 1 + metrics.append(wrapped({ + "request_id": f"wire-{wave}-{lane}", + "engine_request_id": engine_id, + "spec_steps": 0 if lane == 0 else 3, + "target_forwards": 3 if lane == 0 else 6, + }, line + lane)) + if lane == 1: + selectors.append(wrapped({ + "request_id": engine_id, + "accepted_depth": 3, + "depths": [], + }, line + lane)) + return { + "rounds": rounds, + "selectors": selectors, + "activations": [], + "requests": metrics, + } + + +def valid_adaptive_case() -> tuple[dict, dict]: + workload = valid_workload() + for request in workload["requests_detail"]: + request["decode_mode"] = "adaptive" + request["lane_request_index"] = request["wave_index"] + for lane in workload["lanes"]: + lane["decode_mode"] = "adaptive" + + records = valid_records() + records["selectors"] = [] + records["activations"] = [] + for wrapped_metric in records["requests"]: + metric = wrapped_metric["record"] + wire_parts = metric["request_id"].split("-") + lane = int(wire_parts[-1]) + engine_id = metric["engine_request_id"] + decision = "speculation" if lane == 0 else "ar" + metric["spec_steps"] = 3 if lane == 0 else 0 + metric["target_forwards"] = 6 if lane == 0 else 3 + if lane == 0: + records["selectors"].append(wrapped({ + "request_id": engine_id, + "accepted_depth": 3, + "depths": [], + }, wrapped_metric["line_index"])) + records["activations"].append(wrapped({ + "request_id": engine_id, + "slot": lane, + "activation_score": 6.0 if lane == 0 else 2.0, + "score_kind": "qwen38-dflash2-selector-benefit-v1", + "expected_yield": 6.0 if lane == 0 else 2.0, + "evaluation": "scored", + "fallback_reason": None, + "decision_reason": ( + "selected_by_joint_goodput" + if lane == 0 else "ar_counterfactual_won" + ), + "decision": decision, + "hazards": [0.1, 0.2], + }, wrapped_metric["line_index"] - 1)) + return workload, records + + +class RefillSubsetTests(unittest.TestCase): + def test_refill_client_keeps_positional_modes_and_sequences_each_lane(self) -> None: + calls: dict[str, int] = {"prompt-a": 0, "prompt-b": 0} + lock = threading.Lock() + + def fake_request(args, prompt, mode): + with lock: + wave = calls[prompt] + calls[prompt] += 1 + started = time.perf_counter() + digest = benchmark.base.sha256_text(prompt) + return { + "request_id": f"{prompt}-{wave}", + "decode_mode": mode, + "t_start": started, + "t_first": started, + "t_end": time.perf_counter(), + "completion_tokens": args.max_tokens, + "prompt_tokens": 4, + "error": None, + "content_sha256": digest, + "reasoning_content_sha256": "e" * 64, + "combined_output_sha256": "f" * 64, + } + + args = argparse.Namespace( + clients=2, waves=3, prompt_offset=0, timeout=2.0, max_tokens=8, + ) + with mock.patch.object( + benchmark.forced, "stream_request", side_effect=fake_request, + ): + workload = benchmark.run_refill( + args, ["prompt-a", "prompt-b"], ["ar", "speculation"], + ) + self.assertEqual(sum(calls.values()), 6) + self.assertTrue(all(count >= 1 for count in calls.values())) + self.assertEqual(workload["requests"], 6) + self.assertEqual(workload["refill_handoffs"], 4) + self.assertEqual(workload["request_mode_mask"], "AS") + self.assertTrue(workload["exact_output_stable_per_lane"]) + self.assertEqual( + [row["request_index"] for row in workload["requests_detail"]], + list(range(6)), + ) + self.assertEqual( + [row["admission_group_index"] for row in workload["requests_detail"]], + [0, 0, 1, 1, 2, 2], + ) + + def test_valid_refill_proves_every_handoff_and_full_live_round_rate(self) -> None: + result = benchmark.validate_evidence( + valid_workload(), valid_records(), 2, ["ar", "speculation"], + 4, 3, max_start_skew_ms=100.0, max_refill_gap_ms=100.0, + min_full_live_rounds=2, + ) + self.assertTrue(result["passed"], result["errors"]) + self.assertTrue(result["refill_recovery_proved"]) + self.assertEqual(result["completed_before_last_full_live"], 4) + self.assertEqual(result["required_refill_recoveries"], 2) + self.assertEqual(result["scheduled_refills"], 4) + self.assertEqual(result["terminal_guard_requests"], 2) + self.assertEqual(result["executed_spec_depths"], [4]) + self.assertAlmostEqual( + result["full_live"]["engine_round_goodput_tok_s"], 100.0, + ) + self.assertIs(result["adaptive_claims_permitted"], False) + self.assertIs(result["closed_cohort_claims_permitted"], False) + + def test_direct_refill_validates_compact_depth_and_one_pass_counters(self) -> None: + workload, records = valid_adaptive_case() + for wrapped_round in records["rounds"]: + wrapped_round["record"].update({ + "path": "spec-direct", "ar_lanes": 1, "tree_rows": 5, + }) + for wrapped_metric in records["requests"]: + metric = wrapped_metric["record"] + if metric["spec_steps"]: + # One forward per direct Spec step plus an ordinary AR step + # before the sticky activation was committed. + metric["target_forwards"] = metric["spec_steps"] + 1 + result = benchmark.validate_evidence( + workload, records, 2, ["adaptive", "adaptive"], 4, 3, + max_start_skew_ms=100.0, max_refill_gap_ms=100.0, + min_full_live_rounds=2, expected_adaptive_mask="SA", + ) + self.assertTrue(result["passed"], result["errors"]) + self.assertEqual(result["executed_spec_depths"], [4]) + self.assertEqual(result["full_live"]["path_counts"]["spec-direct"], 3) + + def test_direct_forced_refill_validates_one_pass_counters(self) -> None: + workload = valid_workload() + records = valid_records() + for wrapped_round in records["rounds"]: + wrapped_round["record"].update({ + "path": "spec-direct", "ar_lanes": 1, "tree_rows": 5, + }) + for wrapped_metric in records["requests"]: + metric = wrapped_metric["record"] + if metric["spec_steps"]: + metric["target_forwards"] = metric["spec_steps"] + result = benchmark.validate_evidence( + workload, records, 2, ["ar", "speculation"], 4, 3, + max_start_skew_ms=100.0, max_refill_gap_ms=100.0, + min_full_live_rounds=2, + ) + self.assertTrue(result["passed"], result["errors"]) + + def test_adaptive_refill_proves_route_and_execution_for_every_request( + self, + ) -> None: + workload, records = valid_adaptive_case() + result = benchmark.validate_evidence( + workload, records, 2, ["adaptive", "adaptive"], 4, 3, + max_start_skew_ms=100.0, max_refill_gap_ms=100.0, + min_full_live_rounds=2, expected_adaptive_mask="SA", + ) + self.assertTrue(result["passed"], result["errors"]) + self.assertIs(result["adaptive_claims_permitted"], True) + self.assertEqual(result["activation"]["matched_requests"], 6) + self.assertEqual( + result["activation"]["observed_initial_route_mask"], "SA", + ) + self.assertEqual( + result["activation"]["decision_counts"], + {"ar": 3, "speculation": 3}, + ) + + def test_adaptive_refill_allows_later_joint_gate_route_changes(self) -> None: + workload, records = valid_adaptive_case() + later_engine_id = 3 + for wrapped_activation in records["activations"]: + activation = wrapped_activation["record"] + if activation["request_id"] == later_engine_id: + activation.update({ + "decision": "ar", + "decision_reason": "ar_counterfactual_won", + }) + for wrapped_metric in records["requests"]: + metric = wrapped_metric["record"] + if metric["engine_request_id"] == later_engine_id: + metric.update({"spec_steps": 0, "target_forwards": 3}) + result = benchmark.validate_evidence( + workload, records, 2, ["adaptive", "adaptive"], 4, 3, + 100.0, 100.0, 2, expected_adaptive_mask="SA", + ) + self.assertTrue(result["passed"], result["errors"]) + self.assertEqual( + result["activation"]["observed_initial_route_mask"], "SA", + ) + self.assertEqual( + result["activation"]["decision_counts"], + {"ar": 4, "speculation": 2}, + ) + + def test_adaptive_spec_service_rounds_are_explicit_and_valid(self) -> None: + workload, records = valid_adaptive_case() + for wrapped_metric in records["requests"]: + metric = wrapped_metric["record"] + lane = int(metric["request_id"].split("-")[-1]) + if lane == 0: + metric["spec_service_ar_steps"] = 1 + metric["target_forwards"] = 2 * metric["spec_steps"] + 1 + result = benchmark.validate_evidence( + workload, records, 2, ["adaptive", "adaptive"], 4, 3, + max_start_skew_ms=100.0, max_refill_gap_ms=100.0, + min_full_live_rounds=2, expected_adaptive_mask="SA", + ) + self.assertTrue(result["passed"], result["errors"]) + + records["requests"][1]["record"]["spec_service_ar_steps"] = 1 + result = benchmark.validate_evidence( + workload, records, 2, ["adaptive", "adaptive"], 4, 3, + 100.0, 100.0, 2, expected_adaptive_mask="SA", + ) + self.assertFalse(result["passed"]) + + def test_adaptive_refill_route_mismatch_fails_closed(self) -> None: + workload, records = valid_adaptive_case() + result = benchmark.validate_evidence( + workload, records, 2, ["adaptive", "adaptive"], 4, 3, + 100.0, 100.0, 2, expected_adaptive_mask="AS", + ) + self.assertFalse(result["passed"]) + self.assertIs(result["adaptive_claims_permitted"], False) + self.assertTrue(any( + "does not match lane" in error for error in result["errors"] + )) + + def test_adaptive_scoring_off_is_proven_as_all_ar_control(self) -> None: + workload, records = valid_adaptive_case() + records["activations"] = [] + records["selectors"] = [] + for wrapped_round in records["rounds"]: + wrapped_round["record"].update({ + "path": "ar", "k": 0, "emitted_tokens": 2, + }) + for wrapped_metric in records["requests"]: + wrapped_metric["record"].update({ + "spec_steps": 0, "target_forwards": 3, + }) + result = benchmark.validate_evidence( + workload, records, 2, ["adaptive", "adaptive"], 4, 3, + 100.0, 100.0, 2, expected_adaptive_mask="AA", + adaptive_scoring_enabled=False, + ) + self.assertTrue(result["passed"], result["errors"]) + self.assertIs(result["adaptive_claims_permitted"], False) + self.assertIs(result["adaptive_stack_control_permitted"], True) + self.assertEqual(result["activation"]["records"], 0) + self.assertEqual( + result["activation"]["observed_initial_route_mask"], "AA", + ) + + def test_refill_mode_parser_requires_uniform_adaptive_lanes(self) -> None: + self.assertEqual( + benchmark.parse_request_modes("adaptive,adaptive", 2), + ["adaptive", "adaptive"], + ) + with self.assertRaisesRegex(ValueError, "every positional lane"): + benchmark.parse_request_modes("adaptive,ar", 2) + + def test_missing_post_refill_full_live_proof_fails_closed(self) -> None: + records = valid_records() + del records["rounds"][1:] + result = benchmark.validate_evidence( + valid_workload(), records, 2, ["ar", "speculation"], 4, 3, + 100.0, 100.0, 2, + ) + self.assertFalse(result["passed"]) + self.assertTrue(any( + "full live=C was not recovered" in error for error in result["errors"] + )) + + def test_every_request_mode_and_selector_mapping_are_checked(self) -> None: + records = valid_records() + records["requests"][0]["record"]["spec_steps"] = 1 + records["selectors"].append(wrapped({"request_id": 1}, 3)) + result = benchmark.validate_evidence( + valid_workload(), records, 2, ["ar", "speculation"], 4, 3, + 100.0, 100.0, 2, + ) + self.assertFalse(result["passed"]) + self.assertTrue(any( + "forced AR executed speculation" in error for error in result["errors"] + )) + self.assertTrue(any( + "forced AR emitted DFlash2 selector" in error for error in result["errors"] + )) + + def test_exact_output_instability_and_slow_handoff_fail_closed(self) -> None: + workload = valid_workload() + workload["exact_output_stable_per_lane"] = False + workload["refill_gap_s_max"] = 0.2 + result = benchmark.validate_evidence( + workload, valid_records(), 2, ["ar", "speculation"], 4, 3, + 100.0, 100.0, 2, + ) + self.assertFalse(result["passed"]) + self.assertTrue(any("exact output hashes changed" in e for e in result["errors"])) + self.assertTrue(any("handoff gap exceeds" in e for e in result["errors"])) + + def test_timing_rows_must_have_exact_positive_fields(self) -> None: + records = valid_records() + records["rounds"][0]["record"]["emitted_tokens"] = 0 + result = benchmark.validate_evidence( + valid_workload(), records, 2, ["ar", "speculation"], 4, 3, + 100.0, 100.0, 2, + ) + self.assertFalse(result["passed"]) + self.assertTrue(any("positive total_us/emitted_tokens" in e for e in result["errors"])) + + def test_runner_refill_hook_is_opt_in_and_warmups_remain_closed(self) -> None: + runner = (HERE / "run_qwen38_dflash2_subsets.sh").read_text( + encoding="utf-8", + ) + self.assertIn('REFILL_WAVES="${REFILL_WAVES:-1}"', runner) + self.assertIn('REFILL_CLIENT="${REFILL_CLIENT:-$SCRIPT_DIR/refill_subset_benchmark.py}"', runner) + self.assertIn('client_cmd+=(--waves "$waves" --max-refill-gap-ms "$MAX_REFILL_GAP_MS")', runner) + self.assertIn('workload=dflash2-forced-refill', runner) + self.assertIn( + 'run_client_case "$OUT/warmup/ar" "$all_ar" "$WARMUP_TOKENS" 0 1', + runner, + ) + self.assertIn( + 'run_client_case "$case_dir" "$mask" "$MAX_TOKENS" "$repeat" ' + '"$MIN_FULL_LIVE_ROUNDS" "$REFILL_WAVES"', + runner, + ) + self.assertIn('> "$case_dir/client-command.txt"', runner) + + def test_refill_requires_three_waves_for_a_guard_cohort(self) -> None: + args = argparse.Namespace(clients=1, waves=2) + with self.assertRaisesRegex(ValueError, "at least 3"): + benchmark.run(args) + + +if __name__ == "__main__": + unittest.main() diff --git a/harness/benchmarks/concurrency/verify_feature_metrics.py b/harness/benchmarks/concurrency/verify_feature_metrics.py new file mode 100644 index 000000000..0bd98cc27 --- /dev/null +++ b/harness/benchmarks/concurrency/verify_feature_metrics.py @@ -0,0 +1,392 @@ +#!/usr/bin/env python3 +"""Fail closed unless server telemetry proves requested features executed.""" + +from __future__ import annotations + +import argparse +import json +import re +import sys +from pathlib import Path +from typing import Any + + +PREFIX = "[concurrency-metrics] " +COUNTERS = ( + "ddtree_steps", "ddtree_suspensions", "ddtree_accepted_tokens", + "spec_steps", "spec_accepted_tokens", + "target_forwards", "kvflash_page_ins", "kvflash_page_outs", + "kvflash_reselects", +) +DECODE_MODES = ("ar", "speculation", "adaptive") +SPECULATOR_STARTUP_PREFIX = "[parallel-chain] speculator=" +SPEC_PROFILE_PREFIX = "[spec-profile] context=" +REQUIRED_KEYS = ( + "request_id", "effective_prompt_tokens", *COUNTERS, + "kvflash_resident_blocks", "pflash_applied", "pflash_input_tokens", + "pflash_output_tokens", +) + + +def parse_markers(path: Path) -> list[dict[str, Any]]: + rows: list[dict[str, Any]] = [] + for line_no, line in enumerate(path.read_text(encoding="utf-8", errors="replace").splitlines(), 1): + marker = line.find(PREFIX) + if marker < 0: + continue + raw = line[marker + len(PREFIX):] + try: + row = json.loads(raw) + except json.JSONDecodeError as exc: + raise ValueError(f"{path}:{line_no}: invalid concurrency metric JSON: {exc}") from exc + if not isinstance(row, dict): + raise ValueError(f"{path}:{line_no}: concurrency metric must be an object") + missing = [key for key in REQUIRED_KEYS if key not in row] + if missing: + raise ValueError(f"{path}:{line_no}: missing telemetry keys {missing}") + rows.append(row) + return rows + + +def measured_requests(report: dict[str, Any]) -> dict[str, dict[str, Any]]: + requests: dict[str, dict[str, Any]] = {} + for level in report.get("levels") or []: + for row in level.get("requests_detail") or []: + if "error" not in row: + raise ValueError("bench report request lacks an explicit error status") + if row["error"] is not None: + continue + request_id = row.get("request_id") + if not isinstance(request_id, str) or not request_id: + raise ValueError("bench report lacks a request ID for a successful request") + if request_id in requests: + raise ValueError(f"duplicate measured request ID {request_id}") + requests[request_id] = row + if not requests: + raise ValueError("bench report has no successful measured requests") + return requests + + +def aggregate_rows(rows: list[dict[str, Any]]) -> dict[str, dict[str, Any]]: + aggregate: dict[str, dict[str, Any]] = {} + for row in rows: + request_id = row.get("request_id") + if not isinstance(request_id, str) or not request_id: + raise ValueError("telemetry request_id must be a non-empty string") + if request_id in aggregate: + raise ValueError(f"duplicate telemetry request ID {request_id}") + dst = { + "request_id": request_id, + "effective_prompt_tokens": row["effective_prompt_tokens"], + "kvflash_resident_blocks": row["kvflash_resident_blocks"], + "pflash_applied": False, + "pflash_input_tokens": row["pflash_input_tokens"], + "pflash_output_tokens": row["pflash_output_tokens"], + **{key: 0 for key in COUNTERS}, + } + aggregate[request_id] = dst + for key in COUNTERS: + value = row[key] + if key == "ddtree_suspensions": + if type(value) is not int or value not in (0, 1): + raise ValueError( + f"{request_id}: ddtree_suspensions must be 0 or 1 " + "per request" + ) + elif type(value) is not int or value < 0: + raise ValueError(f"{request_id}: {key} must be a non-negative integer") + dst[key] += value + for key in ( + "effective_prompt_tokens", "kvflash_resident_blocks", + "pflash_input_tokens", "pflash_output_tokens", + ): + value = row[key] + if type(value) is not int or value < 0: + raise ValueError(f"{request_id}: {key} must be a non-negative integer") + dst[key] = value + if not isinstance(row["pflash_applied"], bool): + raise ValueError(f"{request_id}: pflash_applied must be boolean") + dst["pflash_applied"] = dst["pflash_applied"] or row["pflash_applied"] + for request_id, row in aggregate.items(): + if row["ddtree_suspensions"] not in (0, 1): + raise ValueError( + f"{request_id}: ddtree_suspensions must be 0 or 1 per request" + ) + if row["spec_accepted_tokens"] > 0 and row["spec_steps"] == 0: + raise ValueError( + f"{request_id}: spec_accepted_tokens requires positive spec_steps" + ) + return aggregate + + +def verify( + report: dict[str, Any], markers: list[dict[str, Any]], expected: set[str], + server_log_text: str | None = None, +) -> dict[str, Any]: + measured = measured_requests(report) + all_rows = aggregate_rows(markers) + rows = {request_id: all_rows[request_id] for request_id in measured if request_id in all_rows} + errors: list[str] = [] + missing = sorted(set(measured) - set(rows)) + if missing: + errors.append(f"missing concurrency telemetry for {len(missing)} measured request(s): {missing}") + + metadata = report.get("server_metadata") or {} + feature_config = metadata.get("feature_config") or {} + decode_mode = feature_config.get("decode_mode") + if decode_mode is not None and decode_mode not in DECODE_MODES: + errors.append(f"invalid recorded decode_mode {decode_mode!r}") + + for request_id, measured_row in measured.items(): + metric = rows.get(request_id) + if metric is None: + continue + wire_effective = measured_row.get("effective_prompt_tokens") + if type(wire_effective) is not int or wire_effective < 0: + errors.append(f"{request_id}: missing usage.timings.effective_prompt_tokens") + elif metric["effective_prompt_tokens"] != wire_effective: + errors.append( + f"{request_id}: log effective_prompt_tokens={metric['effective_prompt_tokens']} " + f"does not match wire value {wire_effective}" + ) + if "pflash" in expected: + wire_input = measured_row.get("prompt_tokens") + if type(wire_input) is not int or wire_input < 0: + errors.append(f"{request_id}: missing usage.prompt_tokens") + elif metric["pflash_input_tokens"] != wire_input: + errors.append( + f"{request_id}: log pflash_input_tokens=" + f"{metric['pflash_input_tokens']} does not match wire value " + f"{wire_input}" + ) + if "ddtree" in expected: + if metric["ddtree_steps"] <= 0: + errors.append(f"{request_id}: DDTree requested but ddtree_steps is zero") + if metric["target_forwards"] <= 0: + errors.append(f"{request_id}: DDTree requested but target_forwards is zero") + if "chain" in expected: + if metric["target_forwards"] <= 0: + errors.append( + f"{request_id}: chain decode requested but target_forwards is zero" + ) + if metric["ddtree_steps"] != 0 or metric["ddtree_accepted_tokens"] != 0: + errors.append( + f"{request_id}: chain run must keep DDTree counters at zero" + ) + if metric["ddtree_suspensions"] != 0: + errors.append( + f"{request_id}: chain run must keep ddtree_suspensions at zero" + ) + if decode_mode == "speculation" and metric["spec_steps"] <= 0: + errors.append( + f"{request_id}: forced chain speculation requested but " + "spec_steps is zero" + ) + if "pflash" in expected: + if metric["pflash_applied"] is not True: + errors.append(f"{request_id}: PFlash requested but pflash_applied is false") + if not (0 < metric["pflash_output_tokens"] < metric["pflash_input_tokens"]): + errors.append( + f"{request_id}: PFlash did not reduce prompt tokens " + f"({metric['pflash_input_tokens']} -> {metric['pflash_output_tokens']})" + ) + + if "chain" in expected: + if decode_mode not in ("speculation", "adaptive"): + errors.append( + "chain requested but metadata decode_mode is not speculation or adaptive" + ) + log_text = server_log_text or "" + startup_mode = ( + isinstance(decode_mode, str) + and re.search( + rf"^.*{re.escape(SPECULATOR_STARTUP_PREFIX)}\S+.*$", + log_text, + flags=re.MULTILINE, + ) + ) + if not startup_mode: + errors.append( + "chain requested but matching speculator adapter startup proof is missing" + ) + if decode_mode == "adaptive" and SPEC_PROFILE_PREFIX not in log_text: + errors.append( + "adaptive chain requested but startup cost-profile proof is missing" + ) + elif decode_mode == "ar": + active_spec = sorted( + request_id for request_id, row in rows.items() + if row["spec_steps"] != 0 or row["spec_accepted_tokens"] != 0 + ) + if active_spec: + errors.append( + f"AR decode_mode emitted chain speculation for request(s): {active_spec}" + ) + + totals = {key: sum(row[key] for row in rows.values()) for key in COUNTERS} + resident = [row["kvflash_resident_blocks"] for row in rows.values()] + variant = str(metadata.get("variant") or "") + workload = str(metadata.get("workload") or "") + pflash_mode = feature_config.get("prefill_compression") + pflash_threshold = feature_config.get("prefill_threshold") + if "pflash" in expected: + if not isinstance(pflash_mode, str) or pflash_mode in ("", "off", "0"): + errors.append("PFlash requested but metadata does not prove it was enabled") + if pflash_mode == "auto": + if type(pflash_threshold) is not int or pflash_threshold <= 0: + errors.append( + "PFlash auto row lacks a positive recorded token threshold" + ) + else: + below_threshold = sorted( + request_id for request_id, row in rows.items() + if row["pflash_input_tokens"] < pflash_threshold + ) + if below_threshold: + errors.append( + "PFlash auto input did not reach its recorded token threshold " + f"for request(s): {below_threshold}" + ) + kvflash_mode = feature_config.get("kvflash") + requested_pool_tokens = feature_config.get("kvflash_max_pool_tokens") + runtime_observed = metadata.get("runtime_observed") or {} + pool_tokens = runtime_observed.get("physical_kv_pool_tokens") + kvflash_page_traffic_required = False + kvflash_page_traffic_reason = "not-requested" + if "kvflash" in expected: + if not isinstance(kvflash_mode, str) or kvflash_mode in ("", "off", "0"): + errors.append("KVFlash requested but metadata does not prove it was enabled") + if (type(requested_pool_tokens) is not int + or requested_pool_tokens <= 0): + errors.append( + "KVFlash requested but its recorded pool-token cap is not a positive integer") + scorer_drafter = feature_config.get("kvflash_scorer_drafter") + scorer_sha256 = feature_config.get("kvflash_scorer_drafter_sha256") + if not isinstance(scorer_drafter, str) or not scorer_drafter: + errors.append( + "KVFlash requested but no explicit scorer drafter was recorded" + ) + if ( + not isinstance(scorer_sha256, str) + or re.fullmatch(r"[0-9a-fA-F]{64}", scorer_sha256) is None + ): + errors.append( + "KVFlash requested but the scorer drafter hash is not a valid SHA-256 digest" + ) + if runtime_observed.get("kvflash_active") is not True: + errors.append( + "KVFlash requested but its physical-pool startup marker was not recorded" + ) + if type(pool_tokens) is not int or pool_tokens <= 0: + errors.append( + "KVFlash requested but no positive startup-observed physical pool was recorded" + ) + if not resident or max(resident) <= 0: + errors.append("KVFlash requested but resident block count never became positive") + + effective_demand_by_level = [] + for level in report.get("levels") or []: + demand = 0 + for request in level.get("requests_detail") or []: + request_id = request.get("request_id") + if request.get("error") is None and request_id in rows: + demand += rows[request_id]["effective_prompt_tokens"] + effective_demand_by_level.append(demand) + if variant not in ("kvflash", "full"): + # A report whose declared variant is inconsistent with a KVFlash + # proof must not pass merely because its current prompts fit. + kvflash_page_traffic_required = True + kvflash_page_traffic_reason = "unknown-variant" + elif variant == "kvflash": + kvflash_page_traffic_required = True + kvflash_page_traffic_reason = "kvflash-only-ablation" + elif workload == "kv-pressure": + kvflash_page_traffic_required = True + kvflash_page_traffic_reason = "kv-pressure-workload" + elif variant == "full": + if type(pool_tokens) is not int or pool_tokens <= 0: + errors.append( + "full KVFlash row lacks a positive recorded pool-token limit" + ) + kvflash_page_traffic_reason = "missing-pool-limit" + else: + kvflash_page_traffic_required = any( + demand > pool_tokens + for demand in effective_demand_by_level + ) + kvflash_page_traffic_reason = ( + "effective-prompt-exceeds-pool" + if kvflash_page_traffic_required + else "compressed-prompt-fits-pool" + ) + + if ( + kvflash_page_traffic_required + and totals["kvflash_page_ins"] + totals["kvflash_page_outs"] <= 0 + ): + errors.append( + "KVFlash paging was required but no page-in/page-out was observed" + ) + + return { + "schema_version": 5, + "decode_mode": decode_mode, + "expected_features": sorted(expected), + "valid": not errors, + "errors": errors, + "measured_request_count": len(measured), + "matched_metric_count": len(rows), + "ignored_marker_count": len(markers) - len(rows), + "kvflash_page_traffic_required": kvflash_page_traffic_required, + "kvflash_page_traffic_reason": kvflash_page_traffic_reason, + "kvflash_pool_tokens": pool_tokens if type(pool_tokens) is int else None, + "kvflash_requested_max_pool_tokens": ( + requested_pool_tokens if type(requested_pool_tokens) is int else None + ), + "pflash_threshold_tokens": ( + pflash_threshold if type(pflash_threshold) is int else None + ), + "aggregate": { + **totals, + "kvflash_resident_blocks_max": max(resident) if resident else None, + "pflash_applied_requests": sum(row["pflash_applied"] is True for row in rows.values()), + "pflash_input_tokens": sum(row["pflash_input_tokens"] for row in rows.values()), + "pflash_output_tokens": sum(row["pflash_output_tokens"] for row in rows.values()), + }, + "requests": [rows[key] for key in sorted(rows)], + } + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--bench", type=Path, required=True) + parser.add_argument("--server-log", type=Path, required=True) + parser.add_argument( + "--expect", action="append", + choices=("ddtree", "chain", "pflash", "kvflash"), default=[], + ) + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + try: + report = json.loads(args.bench.read_text(encoding="utf-8")) + log_text = args.server_log.read_text(encoding="utf-8", errors="replace") + result = verify( + report, parse_markers(args.server_log), set(args.expect), log_text, + ) + except Exception as exc: + print(f"[proof] error: {exc}", file=sys.stderr) + return 2 + args.out.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8") + if not result["valid"]: + for error in result["errors"]: + print(f"[proof] {error}", file=sys.stderr) + return 1 + print( + f"[proof] valid features={','.join(result['expected_features']) or 'ar'} " + f"requests={result['matched_metric_count']}" + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/harness/benchmarks/concurrency/write_feature_metadata.py b/harness/benchmarks/concurrency/write_feature_metadata.py new file mode 100644 index 000000000..e6c48073b --- /dev/null +++ b/harness/benchmarks/concurrency/write_feature_metadata.py @@ -0,0 +1,225 @@ +#!/usr/bin/env python3 +"""Write reproducible server and feature configuration metadata for one case.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import pathlib +import subprocess + + +def digest(path: pathlib.Path) -> str: + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def validated_digest( + path: pathlib.Path | None, + claimed: str | None, + label: str, + cache: dict[pathlib.Path, str], +) -> str | None: + if path is None: + if claimed is not None: + raise ValueError(f"{label} SHA-256 supplied without a model file") + return None + if not isinstance(claimed, str) or not claimed: + raise ValueError(f"{label} SHA-256 is required") + resolved = path.resolve() + if resolved not in cache: + cache[resolved] = digest(resolved) + actual = cache[resolved] + if claimed != actual: + raise ValueError(f"{label} SHA-256 does not match {resolved}") + return actual + + +def resolved_libraries(binary: pathlib.Path) -> dict[str, str]: + result = subprocess.run( + ["ldd", str(binary)], text=True, capture_output=True, + ) + if result.returncode != 0: + detail = result.stderr.strip() or result.stdout.strip() or "no diagnostic" + raise RuntimeError(f"ldd failed for {binary}: {detail}") + unresolved = [ + line.strip() for line in result.stdout.splitlines() + if "=> not found" in line + ] + if unresolved: + raise RuntimeError(f"ldd found unresolved libraries for {binary}: {unresolved}") + libraries: dict[str, str] = {} + for line in result.stdout.splitlines(): + fields = line.replace("=>", " ").split() + paths = [pathlib.Path(value) for value in fields if value.startswith("/")] + for path in paths: + if path.is_file(): + libraries[str(path.resolve())] = digest(path) + return libraries + + +def repository_head(repo: pathlib.Path) -> str: + result = subprocess.run( + ["git", "-C", str(repo), "rev-parse", "HEAD"], + text=True, capture_output=True, + ) + if result.returncode != 0: + detail = result.stderr.strip() or result.stdout.strip() or "no diagnostic" + raise RuntimeError(f"git rev-parse failed for {repo}: {detail}") + head = result.stdout.strip() + if not head: + raise RuntimeError(f"git rev-parse returned an empty revision for {repo}") + return head + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--out", type=pathlib.Path, required=True) + parser.add_argument("--variant", required=True) + parser.add_argument("--workload", required=True) + parser.add_argument("--clients", type=int, required=True) + parser.add_argument("--repeat", type=int, required=True) + parser.add_argument("--binary", type=pathlib.Path, required=True) + parser.add_argument("--model", type=pathlib.Path, required=True) + parser.add_argument("--model-sha256", required=True) + parser.add_argument("--prompt-file", type=pathlib.Path, required=True) + parser.add_argument("--command-file", type=pathlib.Path, required=True) + parser.add_argument("--repo", type=pathlib.Path, required=True) + parser.add_argument("--max-concurrent-prefills", type=int, required=True) + parser.add_argument("--target-device", default=None) + parser.add_argument("--draft-device", default=None) + parser.add_argument("--draft-model", type=pathlib.Path) + parser.add_argument("--draft-model-sha256") + parser.add_argument("--decode-mode", choices=("ar", "speculation", "adaptive")) + parser.add_argument("--cache-type-k") + parser.add_argument("--cache-type-v") + parser.add_argument("--fa-window", type=int) + parser.add_argument("--draft-always", choices=("on", "off")) + parser.add_argument("--confidence", choices=("on", "off")) + parser.add_argument("--ddtree", action="store_true") + parser.add_argument("--fast-rollback", action="store_true") + parser.add_argument("--ddtree-budget", type=int) + parser.add_argument("--prefill-compression", default="off") + parser.add_argument("--prefill-threshold", type=int) + parser.add_argument("--prefill-keep-ratio", type=float) + parser.add_argument("--prefill-drafter", type=pathlib.Path) + parser.add_argument("--prefill-drafter-sha256") + parser.add_argument("--draft-residency", default=None) + parser.add_argument("--kvflash", default="off") + parser.add_argument("--kvflash-max-pool-tokens", type=int) + parser.add_argument("--kvflash-scorer-drafter", type=pathlib.Path) + parser.add_argument("--kvflash-scorer-drafter-sha256") + parser.add_argument("--launch-env", action="append", default=[]) + args = parser.parse_args() + + launch_env: dict[str, str] = {} + for item in args.launch_env: + key, sep, value = item.partition("=") + if not sep or not key: + parser.error(f"bad --launch-env {item!r}; expected KEY=VALUE") + launch_env[key] = value + + digest_cache: dict[pathlib.Path, str] = {} + model_sha256 = validated_digest( + args.model, args.model_sha256, "target model", digest_cache, + ) + draft_model_sha256 = validated_digest( + args.draft_model, args.draft_model_sha256, "draft model", digest_cache, + ) + prefill_drafter_sha256 = validated_digest( + args.prefill_drafter, args.prefill_drafter_sha256, + "prefill drafter", digest_cache, + ) + kvflash_scorer_drafter_sha256 = validated_digest( + args.kvflash_scorer_drafter, + args.kvflash_scorer_drafter_sha256, + "KVFlash scorer drafter", + digest_cache, + ) + + libraries = resolved_libraries(args.binary) + git_head = repository_head(args.repo) + literal_flags: list[str] = [] + if args.target_device: + literal_flags += ["--target-device", args.target_device] + if args.draft_device: + literal_flags += ["--draft-device", args.draft_device] + if args.decode_mode: + literal_flags += ["--decode-mode", args.decode_mode] + if args.cache_type_k: + literal_flags += ["--cache-type-k", args.cache_type_k] + if args.cache_type_v: + literal_flags += ["--cache-type-v", args.cache_type_v] + if args.fa_window is not None: + literal_flags += ["--fa-window", str(args.fa_window)] + if args.ddtree: + literal_flags += ["--ddtree"] + if args.ddtree_budget is not None: + literal_flags += ["--ddtree-budget", str(args.ddtree_budget)] + if args.fast_rollback: + literal_flags += ["--fast-rollback"] + if args.draft_residency: + literal_flags += ["--draft-residency", args.draft_residency] + if args.prefill_compression != "off": + literal_flags += ["--prefill-compression", args.prefill_compression] + if args.prefill_drafter: + literal_flags += ["--prefill-drafter", str(args.prefill_drafter.resolve())] + if args.kvflash != "off": + literal_flags += ["--kvflash", args.kvflash] + + obj = { + "schema_version": 4, + "variant": args.variant, + "workload": args.workload, + "clients": args.clients, + "repeat": args.repeat, + "max_concurrent_prefills": args.max_concurrent_prefills, + "server_binary": str(args.binary.resolve()), + "server_binary_sha256": digest(args.binary), + "model": str(args.model.resolve()), + "model_sha256": model_sha256, + "prompt_file_sha256": digest(args.prompt_file), + "server_command": args.command_file.read_text(encoding="utf-8").strip(), + "launch_environment": launch_env, + "resolved_shared_library_sha256": libraries, + "git_head": git_head, + "literal_screenshot_flags": literal_flags, + # Populated from fail-closed startup markers after the server is healthy. + "runtime_observed": None, + "feature_config": { + "target_device": args.target_device, + "draft_device": args.draft_device, + "draft_model": str(args.draft_model.resolve()) if args.draft_model else None, + "draft_model_sha256": draft_model_sha256, + "decode_mode": args.decode_mode, + "cache_type_k": args.cache_type_k, + "cache_type_v": args.cache_type_v, + "fa_window": args.fa_window, + "draft_always": args.draft_always, + "confidence": args.confidence, + "ddtree": args.ddtree, + "fast_rollback": args.fast_rollback, + "ddtree_budget": args.ddtree_budget, + "prefill_compression": args.prefill_compression, + "prefill_threshold": args.prefill_threshold, + "prefill_keep_ratio": args.prefill_keep_ratio, + "prefill_drafter": ( + str(args.prefill_drafter.resolve()) if args.prefill_drafter else None + ), + "prefill_drafter_sha256": prefill_drafter_sha256, + "draft_residency": args.draft_residency, + "kvflash": args.kvflash, + "kvflash_max_pool_tokens": args.kvflash_max_pool_tokens, + "kvflash_scorer_drafter": ( + str(args.kvflash_scorer_drafter.resolve()) + if args.kvflash_scorer_drafter else None + ), + "kvflash_scorer_drafter_sha256": kvflash_scorer_drafter_sha256, + }, + } + args.out.write_text(json.dumps(obj, indent=2, sort_keys=True) + "\n", encoding="utf-8") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/harness/benchmarks/prompts/qwen38_dspark_adaptive_selection.jsonl b/harness/benchmarks/prompts/qwen38_dspark_adaptive_selection.jsonl new file mode 100644 index 000000000..a2721c465 --- /dev/null +++ b/harness/benchmarks/prompts/qwen38_dspark_adaptive_selection.jsonl @@ -0,0 +1,6 @@ +{"dense_r9700_baseline":{"ar_decode_tok_s":34.89,"lossless":true,"spec_accept_pct":35.8,"spec_avg_commit":2.51,"spec_decode_tok_s":46.35,"spec_over_ar":1.328},"expected_dense_oracle":"speculation","id":"adaptive-he-09-sum-product","prompt":"Complete the following Python function.\n\nfrom typing import List, Tuple\n\ndef sum_product(numbers: List[int]) -> Tuple[int, int]:\n \"\"\" For a given list of integers, return a tuple consisting of a sum and a product of all the integers in a list.\n Empty sum should be equal to 0 and empty product should be equal to 1.\n >>> sum_product([])\n (0, 1)\n >>> sum_product([1, 2, 3, 4])\n (10, 24)\n \"\"\"\n","selection_class":"speculation_strong_win","source_id":"he_09","suite":"humaneval"} +{"dense_r9700_baseline":{"ar_decode_tok_s":34.89,"lossless":true,"spec_accept_pct":38.9,"spec_avg_commit":2.72,"spec_decode_tok_s":50.3,"spec_over_ar":1.442},"expected_dense_oracle":"speculation","id":"adaptive-he-10-rolling-max","prompt":"Complete the following Python function.\n\nfrom typing import List\n\ndef rolling_max(numbers: List[int]) -> List[int]:\n \"\"\" From a given list of integers, generate a list of rolling maximum element found until given moment\n in the sequence.\n >>> rolling_max([1, 2, 3, 2, 3, 4, 2])\n [1, 2, 3, 3, 3, 4, 4]\n \"\"\"\n","selection_class":"speculation_strong_win","source_id":"he_10","suite":"humaneval"} +{"dense_r9700_baseline":{"ar_decode_tok_s":34.95,"lossless":true,"spec_accept_pct":29.0,"spec_avg_commit":2.03,"spec_decode_tok_s":37.69,"spec_over_ar":1.078},"expected_dense_oracle":"speculation","id":"adaptive-he-02-separate-paren-groups","prompt":"Complete the following Python function.\n\nfrom typing import List\n\ndef separate_paren_groups(paren_string: str) -> List[str]:\n \"\"\" Input to this function is a string containing multiple groups of nested parentheses. Your goal is to\n separate those group into separate strings and return the list of those.\n Separate groups are balanced (each open brace is properly closed) and not nested within each other\n Ignore any spaces in the input string.\n >>> separate_paren_groups('( ) (( )) (( )( ))')\n ['()', '(())', '(()())']\n \"\"\"\n","selection_class":"speculation_marginal_win","source_id":"he_02","suite":"humaneval"} +{"dense_r9700_baseline":{"ar_decode_tok_s":35.15,"lossless":true,"spec_accept_pct":28.9,"spec_avg_commit":2.02,"spec_decode_tok_s":37.44,"spec_over_ar":1.065},"expected_dense_oracle":"speculation","id":"adaptive-he-03-truncate-number","prompt":"Complete the following Python function.\n\ndef truncate_number(number: float) -> float:\n \"\"\" Given a positive floating point number, it can be decomposed into\n and integer part (largest integer smaller than given number) and decimals\n (leftover part always smaller than 1).\n\n Return the decimal part of the number.\n >>> truncate_number(3.5)\n 0.5\n \"\"\"\n","selection_class":"speculation_marginal_win","source_id":"he_03","suite":"humaneval"} +{"dense_r9700_baseline":{"ar_decode_tok_s":35.04,"lossless":true,"spec_accept_pct":26.2,"spec_avg_commit":1.84,"spec_decode_tok_s":33.95,"spec_over_ar":0.969},"expected_dense_oracle":"ar","id":"adaptive-he-08-filter-by-substring","prompt":"Complete the following Python function.\n\nfrom typing import List\n\ndef filter_by_substring(strings: List[str], substring: str) -> List[str]:\n \"\"\" Filter an input list of strings only for ones that contain given substring\n >>> filter_by_substring([], 'a')\n []\n >>> filter_by_substring(['abc', 'bacd', 'cde', 'array'], 'a')\n ['abc', 'bacd', 'array']\n \"\"\"\n","selection_class":"speculation_loss","source_id":"he_08","suite":"humaneval"} +{"dense_r9700_baseline":{"ar_decode_tok_s":34.9,"lossless":true,"spec_accept_pct":20.8,"spec_avg_commit":1.45,"spec_decode_tok_s":27.1,"spec_over_ar":0.777},"expected_dense_oracle":"ar","id":"adaptive-prose-01-reproducibility","prompt":"Write a clear, self-contained technical essay of about 500 words on why reproducible benchmarks need immutable inputs and explicit hardware metadata. Include one concrete example and end with a concise conclusion.","selection_class":"speculation_loss","source_id":"prose-01","suite":"prose"} diff --git a/harness/benchmarks/prompts/qwen38_pr625.jsonl b/harness/benchmarks/prompts/qwen38_pr625.jsonl new file mode 100644 index 000000000..635ec2eec --- /dev/null +++ b/harness/benchmarks/prompts/qwen38_pr625.jsonl @@ -0,0 +1,2 @@ +{"id":"code-long","suite":"code","prompt":"Write a single self-contained, production-quality Python module implementing an asynchronous bounded worker pool. Include type annotations, docstrings, graceful cancellation, backpressure, per-job timeouts, structured result objects, clean shutdown, an executable usage example, and comprehensive unittest tests. Return only Python code and make the module at least 250 lines long."} +{"id":"prose-long","suite":"prose","prompt":"Write a clear, self-contained technical essay of about 500 words on why reproducible benchmarks need immutable inputs and explicit hardware metadata. Include one concrete example and end with a concise conclusion."} diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 46ecfda81..06739ed43 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -439,6 +439,11 @@ add_library(dflash_common STATIC src/common/dynamic_backend.cpp src/common/domino_head.cpp src/common/dspark_head.cpp + src/common/dflash2_benefit.cpp + src/common/speculation/spec_cost_profile.cpp + src/common/speculation/adapters/dflash2_speculator.cpp + src/common/dflash2_head.cpp + src/common/dflash2_batch.cpp src/common/target_shard_ipc.cpp src/common/target_shard_ipc_daemon.cpp src/common/dflash_feature_ring.cpp @@ -450,6 +455,9 @@ add_library(dflash_common STATIC src/common/dflash_draft_kv.cpp src/common/dflash_spec_decode.cpp src/common/concurrency/paged_kv_pool.cpp + src/common/concurrency/paged_kv_residency.cpp + src/common/concurrency/qwen_paged_kv_transfer_layout.cpp + src/common/concurrency/qwen_paged_kv_transfer.cpp src/qwen35/concurrency/qwen35_slot_manager.cpp src/common/layer_split_backend.cpp src/common/layer_split_runtime.cpp @@ -1391,12 +1399,33 @@ if(DFLASH27B_TESTS) ${CMAKE_CURRENT_SOURCE_DIR}/src) list(APPEND _raw_unit_test_targets test_paged_kv_pool) endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_paged_kv_residency.cpp") + # Pure host-side multi-sequence residency policy + mock DMA test. + add_executable(test_paged_kv_residency + test/test_paged_kv_residency.cpp + src/common/concurrency/paged_kv_pool.cpp + src/common/concurrency/paged_kv_residency.cpp) + target_include_directories(test_paged_kv_residency PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src) + list(APPEND _raw_unit_test_targets test_paged_kv_residency) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_qwen_paged_kv_transfer_layout.cpp") + # Pure host-side validation of packed K/V block byte/stride math. + add_executable(test_qwen_paged_kv_transfer_layout + test/test_qwen_paged_kv_transfer_layout.cpp + src/common/concurrency/qwen_paged_kv_transfer_layout.cpp) + target_include_directories(test_qwen_paged_kv_transfer_layout PRIVATE + ${DFLASH27B_SRC_INCLUDE_DIRS} + ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/include) + list(APPEND _raw_unit_test_targets test_qwen_paged_kv_transfer_layout) + endif() if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_seq_slot_manager.cpp") # Host-side slot bookkeeping test (concurrent serving): no GPU. add_executable(test_seq_slot_manager test/test_seq_slot_manager.cpp src/qwen35/concurrency/qwen35_slot_manager.cpp - src/common/concurrency/paged_kv_pool.cpp) + src/common/concurrency/paged_kv_pool.cpp + src/common/concurrency/paged_kv_residency.cpp) target_include_directories(test_seq_slot_manager PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src) list(APPEND _raw_unit_test_targets test_seq_slot_manager) @@ -1409,6 +1438,84 @@ if(DFLASH27B_TESTS) ${CMAKE_CURRENT_SOURCE_DIR}/test) list(APPEND _raw_unit_test_targets test_seq_engine_contract) endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_ddtree_path.cpp") + # Pure host-side accepted-path/pending-token contract tests. + add_executable(test_ddtree_path + test/test_ddtree_path.cpp + src/common/ddtree.cpp) + target_include_directories(test_ddtree_path PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src) + list(APPEND _raw_unit_test_targets test_ddtree_path) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_chain_spec_shapes.cpp") + # Pure host-side DSpark chain topology and mixed-launch arithmetic. + add_executable(test_chain_spec_shapes + test/test_chain_spec_shapes.cpp + src/common/ddtree.cpp) + target_include_directories(test_chain_spec_shapes PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src + ${CMAKE_CURRENT_SOURCE_DIR}/test) + list(APPEND _raw_unit_test_targets test_chain_spec_shapes) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_dflash2_selector_validation.cpp") + # Pure host-side selector metadata/layout validation: no GPU. + add_executable(test_dflash2_selector_validation + test/test_dflash2_selector_validation.cpp) + target_include_directories(test_dflash2_selector_validation PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src + ${CMAKE_CURRENT_SOURCE_DIR}/test) + list(APPEND _raw_unit_test_targets test_dflash2_selector_validation) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_dflash2_benefit.cpp") + # Pure host-side, versioned selector-to-benefit adapter: no GPU. + add_executable(test_dflash2_benefit + test/test_dflash2_benefit.cpp + src/common/dflash2_benefit.cpp) + target_include_directories(test_dflash2_benefit PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src + ${CMAKE_CURRENT_SOURCE_DIR}/include + ${CMAKE_CURRENT_SOURCE_DIR}/test + ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/include) + list(APPEND _raw_unit_test_targets test_dflash2_benefit) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_delta_transition_journal.cpp") + # Pure host-side proof for replay-free DeltaNet transition commits. + add_executable(test_delta_transition_journal + test/test_delta_transition_journal.cpp + src/qwen35/delta_transition_journal.cpp) + target_include_directories(test_delta_transition_journal PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src + ${CMAKE_CURRENT_SOURCE_DIR}/test) + list(APPEND _raw_unit_test_targets test_delta_transition_journal) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_dspark_batched_head.cpp") + add_executable(test_dspark_batched_head + test/test_dspark_batched_head.cpp) + target_include_directories(test_dspark_batched_head PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src + ${CMAKE_CURRENT_SOURCE_DIR}/test) + target_link_libraries(test_dspark_batched_head PRIVATE + dflash_common ggml ggml-cpu + ${DFLASH27B_GGML_BACKEND_TARGET} ggml-base) + list(APPEND _raw_unit_test_targets test_dspark_batched_head) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_speculation_gate.cpp") + add_executable(test_speculation_gate + test/test_speculation_gate.cpp) + target_include_directories(test_speculation_gate PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src + ${CMAKE_CURRENT_SOURCE_DIR}/test) + list(APPEND _raw_unit_test_targets test_speculation_gate) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_spec_cost_profile.cpp") + add_executable(test_spec_cost_profile + test/test_spec_cost_profile.cpp + src/common/speculation/spec_cost_profile.cpp) + target_include_directories(test_spec_cost_profile PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src + ${CMAKE_CURRENT_SOURCE_DIR}/test) + list(APPEND _raw_unit_test_targets test_spec_cost_profile) + endif() if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_seq_batch_plan.cpp") # Pure-host tests for model-neutral token-budget/FIFO planning. add_executable(test_seq_batch_plan test/test_seq_batch_plan.cpp) @@ -1791,6 +1898,8 @@ if(DFLASH27B_TESTS) set(_unit_ctest_name recurrent_snapshot) elseif(_unit_target STREQUAL "test_paged_kv_pool") set(_unit_ctest_name paged_kv_pool) + elseif(_unit_target STREQUAL "test_paged_kv_residency") + set(_unit_ctest_name paged_kv_residency) endif() add_test(NAME "${_unit_ctest_name}" COMMAND ${_unit_target}) set_tests_properties("${_unit_ctest_name}" PROPERTIES SKIP_RETURN_CODE 77) @@ -1896,6 +2005,17 @@ if(DFLASH27B_TESTS) add_dependencies(check test_batched_gdn) endif() endif() + if((DFLASH27B_GPU_BACKEND STREQUAL "cuda" OR + DFLASH27B_GPU_BACKEND STREQUAL "hip") + AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_gdn_transition_journal.cpp") + dflash_add_ggml_gpu_executable( + test_gdn_transition_journal + test/test_gdn_transition_journal.cpp) + add_test(NAME gdn_transition_journal COMMAND test_gdn_transition_journal) + if(TARGET check) + add_dependencies(check test_gdn_transition_journal) + endif() + endif() if((DFLASH27B_GPU_BACKEND STREQUAL "cuda" OR DFLASH27B_GPU_BACKEND STREQUAL "hip") AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_concat_transpose.cpp") diff --git a/server/deps/llama.cpp/ggml/include/ggml-cuda.h b/server/deps/llama.cpp/ggml/include/ggml-cuda.h index bb850e2a9..a418a9ab7 100644 --- a/server/deps/llama.cpp/ggml/include/ggml-cuda.h +++ b/server/deps/llama.cpp/ggml/include/ggml-cuda.h @@ -93,6 +93,43 @@ GGML_BACKEND_API ggml_backend_reg_t ggml_backend_cuda_reg(void); GGML_BACKEND_API bool ggml_backend_cuda_topk_rows(const struct ggml_tensor * logits, int k, float * probs_out, int32_t * ids_out); +// Apply compact GDN journal prefixes to persistent F32 state. The journal is +// [J,H,T,B] (see ggml_gated_delta_net_set_transition_journal); accepted and +// active slots are contiguous I32 [B]. Negative/out-of-range slots are +// padding. Phase 1 is synchronous, single-device, and requires unique slots. +// Each live state slot must still contain the same base state from which its +// journal row was captured, because delta is state-dependent. Call only after +// the synchronous graph compute that produced the journal has returned. +GGML_BACKEND_API bool ggml_backend_cuda_gdn_transition_journal_commit( + const struct ggml_tensor * journal, + struct ggml_tensor * state, + const struct ggml_tensor * accepted_prefixes, + const struct ggml_tensor * active_slot_ids); + +// Batched concurrent-tree commit. Validation is fail-closed before any kernel +// launches; all layer journals and convolution windows commit on one device +// synchronization. +GGML_BACKEND_API bool ggml_backend_cuda_gdn_transition_journal_commit_many( + const struct ggml_tensor * const * journals, + struct ggml_tensor * const * states, + const struct ggml_tensor * const * conv_inputs, + struct ggml_tensor * const * conv_states, + int n_layers, + const struct ggml_tensor * accepted_prefixes, + const struct ggml_tensor * active_slot_ids); + +// Promote accepted packed-tree K/V scratch rows into pager-owned rows. +GGML_BACKEND_API bool ggml_backend_cuda_tree_cache_commit_many( + struct ggml_tensor * const * caches, int n_caches, + const struct ggml_tensor * commit_rows, + const struct ggml_tensor * active_slot_ids, + int tree_scratch_base, int tree_scratch_stride); + +// Promote accepted BF16 tree feature rows into slot-local feature rings. +GGML_BACKEND_API bool ggml_backend_cuda_tree_feature_commit( + const struct ggml_tensor * source, struct ggml_tensor * destination, + const struct ggml_tensor * destination_rows); + // Attach learned per-expert decode tables to a mixed-precision tensor. The // host variants copy the tables to the device that owns `base`. Call the // matching unregister function before releasing the tensor's backing buffer. diff --git a/server/deps/llama.cpp/ggml/include/ggml.h b/server/deps/llama.cpp/ggml/include/ggml.h index acac40c1f..fd01a3a17 100644 --- a/server/deps/llama.cpp/ggml/include/ggml.h +++ b/server/deps/llama.cpp/ggml/include/ggml.h @@ -221,7 +221,7 @@ #define GGML_MAX_DIMS 4 #define GGML_MAX_PARAMS 2048 -#define GGML_MAX_SRC 10 +#define GGML_MAX_SRC 12 #define GGML_MAX_N_THREADS 512 #define GGML_MAX_OP_PARAMS 64 @@ -2502,6 +2502,18 @@ extern "C" { // prefill chunks can attend the paged pool causally. A negative position // marks a padding row. NULL keeps the decode semantics (full cached // length per row). + // + // parent_ids/tree_sizes optionally enable packed tree verification. + // Queries are flattened sequence-major: tree sequence s occupies rows + // [s*tree_width, (s+1)*tree_width). parent_ids is contiguous I32 + // [tree_width, n_tree_seq] (root parent -1), and tree_sizes is contiguous + // I32 [n_tree_seq]. active_slot_ids is required and remains per query row; + // it selects the physical block-table column and scratch slab. Each live + // query attends its complete committed prefix from the block table plus + // its own candidate node and ancestors from physical K/V rows + // tree_scratch_base + slot*tree_scratch_stride + node. Siblings and rows + // at or beyond tree_sizes[s] are excluded. query_positions must be NULL + // in tree mode. Pass NULL/NULL/0/0/0 to retain standard paged attention. GGML_API struct ggml_tensor * ggml_paged_attn_ext( struct ggml_context * ctx, struct ggml_tensor * q, @@ -2513,7 +2525,12 @@ extern "C" { struct ggml_tensor * query_positions, float scale, int block_size, - int max_kv_seq_len); + int max_kv_seq_len, + struct ggml_tensor * parent_ids, + struct ggml_tensor * tree_sizes, + int tree_width, + int tree_scratch_base, + int tree_scratch_stride); // TurboQuant FWHT rotation. direction: 0 = forward, 1 = inverse. // Applies signs1 -> FWHT -> signs2 (forward) or signs2 -> FWHT -> signs1 (inverse). @@ -2741,6 +2758,25 @@ extern "C" { struct ggml_tensor * c, struct ggml_tensor * parent_ids); + // dflash extension: fused causal-conv step for recurrent decode/verify. + // Replaces transpose + concat(state, x) + ssm_conv + silu + state + // write-back with one kernel. + // x: [C, T, S] f32, rows contiguous (token stride may be + // larger than C, e.g. a row-slice of a stacked GEMV) + // c: [K, C] f32 depthwise conv weights + // conv_state: [K-1, C, S] f32 history; READ, then OVERWRITTEN in + // place with the last K-1 conv inputs + // conv_input_out: optional [>= K-1+T, C, S] f32; receives the full + // conv window (history rows then x rows) per channel, + // for speculative-decode rollback. May be a view. + // Returns silu(conv(x)) as [C, T, S]. CUDA/HIP only. + GGML_API struct ggml_tensor * ggml_ssm_conv_step( + struct ggml_context * ctx, + struct ggml_tensor * x, + struct ggml_tensor * c, + struct ggml_tensor * conv_state, + struct ggml_tensor * conv_input_out); + GGML_API struct ggml_tensor * ggml_ssm_scan( struct ggml_context * ctx, struct ggml_tensor * s, @@ -2887,6 +2923,25 @@ extern "C" { struct ggml_tensor * tensor, bool skip_intermediate); + // CUDA/HIP linear-chain journal in compact F32 [J,H,T,B] layout: + // scalar gate J=2*S_v+1 stores [g | k | delta], while KDA J=3*S_v + // stores [g[S_v] | k | delta]. Delta is captured after the + // state-dependent reduction. Tree mode is deliberately rejected. + GGML_API void ggml_gated_delta_net_set_transition_journal( + struct ggml_tensor * tensor, + struct ggml_tensor * journal); + + // dflash extension: let the kernel derive the gates from the raw + // projections instead of graph-side sigmoid/softplus ops: + // beta_val = sigmoid(beta_raw) + // g_val = exp(softplus(alpha_raw + dt_bias[h]) * A[h]) + // `g` then carries alpha_raw and `beta` carries beta_raw (both [1,H,T,S]); + // dt_bias and A are [H] f32. Only for the non-tree, non-KDA CUDA/HIP path. + GGML_API void ggml_gated_delta_net_set_raw_gates( + struct ggml_tensor * tensor, + struct ggml_tensor * dt_bias, + struct ggml_tensor * A); + // dflash extension: tree-mode gated delta net for DDTree-style // speculative decoding verify. `parent_ids` is an int32 tensor of shape // [n_tokens, n_seqs] where entry [t, s] is the index within sequence s of diff --git a/server/deps/llama.cpp/ggml/src/ggml-cpu/ops.cpp b/server/deps/llama.cpp/ggml/src/ggml-cpu/ops.cpp index b10e8c75d..2c05b85f2 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cpu/ops.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-cpu/ops.cpp @@ -9332,6 +9332,8 @@ void ggml_compute_forward_flash_attn_back( static void ggml_compute_forward_ssm_conv_f32( const ggml_compute_params * params, ggml_tensor * dst) { + // dflash: the fused step mode (ggml_ssm_conv_step) is CUDA/HIP only + GGML_ASSERT(ggml_get_op_params_i32(dst, 0) == 0 && "ggml_ssm_conv_step is not supported on CPU"); const ggml_tensor * src0 = dst->src[0]; // conv_x const ggml_tensor * src1 = dst->src[1]; // conv1d.weight diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn-vec.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn-vec.cuh index bcf1dd804..85a2af718 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn-vec.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn-vec.cuh @@ -534,7 +534,10 @@ void ggml_cuda_flash_attn_ext_vec_case_impl(ggml_backend_cuda_context & ctx, ggm const bool need_f16_K = type_K == GGML_TYPE_F16; const bool need_f16_V = type_V == GGML_TYPE_F16; constexpr size_t nbytes_shared = 0; - launch_fattn(ctx, dst, fattn_kernel, nwarps, nbytes_shared, D, need_f16_K, need_f16_V, false); + // The kernel walks the KV sequence in steps of nthreads (not D); telling + // launch_fattn so lets it split a short KV span (e.g. a 256-token window + // at head_dim 256) across two blocks per head instead of one. + launch_fattn(ctx, dst, fattn_kernel, nwarps, nbytes_shared, nthreads, need_f16_K, need_f16_V, false); } template diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/gated_delta_net.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/gated_delta_net.cu index 76f2de9da..50a7a4184 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/gated_delta_net.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/gated_delta_net.cu @@ -82,6 +82,7 @@ gated_delta_net_cuda(const float * q, float * state_out, const int * parent_ids, // TREE_MODE only; else ignored InterT * persist_inter, // optional external buffer for per-token intermediates + float * transition_journal, int64_t H, int64_t n_tokens, int64_t n_seqs, @@ -97,7 +98,9 @@ gated_delta_net_cuda(const float * q, int64_t sb3, const uint3 neqk1_magic, const uint3 rq3_magic, - float scale) { + float scale, + const float * gate_bias, // raw-gate mode: dt_bias[H], else nullptr + const float * gate_A) { // raw-gate mode: A[H] const uint32_t h_idx = blockIdx.x; const uint32_t sequence = blockIdx.y; // each warp owns one column, using warp-level primitives to reduce across rows @@ -196,7 +199,16 @@ gated_delta_net_cuda(const float * q, const float * beta_t = beta + gb_offset; const float * g_t = g + gb_offset * (KDA ? S_v : 1); - const float beta_val = *beta_t; + constexpr int journal_gate_values = KDA ? S_v : 1; + constexpr int journal_width = journal_gate_values + 2*S_v; + float * journal_t = transition_journal + ? transition_journal + + ((sequence * n_tokens + t) * H + h_idx) * journal_width + : nullptr; + + // raw-gate mode: beta = sigmoid(beta_raw); g = softplus(alpha_raw + bias) * A + const bool raw_gates = gate_bias != nullptr; + const float beta_val = raw_gates ? 1.0f / (1.0f + expf(-(*beta_t))) : *beta_t; // Cache k and q in registers float k_reg[rows_per_lane]; @@ -208,8 +220,27 @@ gated_delta_net_cuda(const float * q, q_reg[r] = q_t[i]; } + if (journal_t && blockIdx.z == 0 && threadIdx.y == 0) { +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { + const int i = r * warp_size + lane; + journal_t[journal_gate_values + i] = k_reg[r]; + if constexpr (KDA) { + journal_t[i] = expf(g_t[i]); + } + } + } + if constexpr (!KDA) { - const float g_val = expf(*g_t); + float g_log = *g_t; + if (raw_gates) { + const float a = g_log + gate_bias[h_idx]; + g_log = ((a > 20.0f) ? a : logf(1.0f + expf(a))) * gate_A[h_idx]; + } + const float g_val = expf(g_log); + if (journal_t && lane == 0 && col == 0) { + journal_t[0] = g_val; + } // kv[col] = (S^T @ k)[col] = sum_i S[i][col] * k[i] float kv_shard = 0.0f; @@ -221,6 +252,9 @@ gated_delta_net_cuda(const float * q, // delta[col] = (v[col] - g * kv[col]) * beta float delta_col = (v_t[col] - g_val * kv_col) * beta_val; + if (journal_t && lane == 0) { + journal_t[journal_gate_values + S_v + col] = delta_col; + } // fused: S[i][col] = g * S[i][col] + k[i] * delta[col] // attn[col] = (S^T @ q)[col] = sum_i S[i][col] * q[i] @@ -249,6 +283,9 @@ gated_delta_net_cuda(const float * q, // delta[col] = (v[col] - kv[col]) * beta float delta_col = (v_t[col] - kv_col) * beta_val; + if (journal_t && lane == 0) { + journal_t[journal_gate_values + S_v + col] = delta_col; + } // fused: S[i][col] = g[i] * S[i][col] + k[i] * delta[col] // attn[col] = (S^T @ q)[col] = sum_i S[i][col] * q[i] @@ -291,7 +328,7 @@ gated_delta_net_cuda(const float * q, } } -template +template __global__ void __launch_bounds__(WARP_THREADS * 8, 2) gated_delta_net_cuda_grouped_cols(const float * q, const float * k, @@ -300,9 +337,11 @@ gated_delta_net_cuda_grouped_cols(const float * q, const float * beta, const float * curr_state, const int * active_slot_ids, + const int * parent_ids, float * dst, float * state_out, InterT * persist_inter, + float * transition_journal, int64_t H, int64_t n_tokens, int64_t n_seqs, @@ -318,7 +357,9 @@ gated_delta_net_cuda_grouped_cols(const float * q, int64_t sb3, const uint3 neqk1_magic, const uint3 rq3_magic, - float scale) { + float scale, + const float * gate_bias, // raw-gate mode: dt_bias[H], else nullptr + const float * gate_A) { // raw-gate mode: A[H] static_assert(S_v == 128, "grouped GDN kernel is specialized for S_v=128"); static_assert(WIDTH == 16, "grouped GDN kernel expects 16-lane subgroups"); static_assert(COLS == 4, "grouped GDN kernel expects 4 columns per subgroup"); @@ -352,7 +393,7 @@ gated_delta_net_cuda_grouped_cols(const float * q, n_seqs, n_state_slots, physical_sequence, physical_state_offset); InterT * inter_states = nullptr; InterT * inter_base = nullptr; - if constexpr (WRITE_INTER) { + if constexpr (WRITE_INTER || TREE_MODE) { inter_states = persist_inter ? persist_inter : (InterT *)(dst + attn_score_elems + final_state_elems); @@ -362,6 +403,10 @@ gated_delta_net_cuda_grouped_cols(const float * q, const float * curr_state_seq = physical_sequence >= 0 ? curr_state + physical_state_offset : nullptr; + const int * parent_ids_seq = nullptr; + if constexpr (TREE_MODE) { + parent_ids_seq = parent_ids + sequence * n_tokens; + } attn_data += (sequence * n_tokens * H + h_idx) * S_v; float state_shard[COLS][rows_per_lane]; @@ -378,6 +423,39 @@ gated_delta_net_cuda_grouped_cols(const float * q, } for (int t = 0; t < n_tokens; ++t) { + if constexpr (TREE_MODE) { + if (t > 0) { + const int parent_t = parent_ids_seq[t]; + if (parent_t == GGML_GDN_TREE_ROOT_PARENT) { +#pragma unroll + for (int c = 0; c < COLS; ++c) { + const int col = col_base + c; +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { + const int row = r * WIDTH + lane; + state_shard[c][r] = curr_state_seq + ? curr_state_seq[col * S_v + row] + : 0.0f; + } + } + } else if (parent_t != t - 1) { + const InterT * parent_base = inter_states + + ((sequence * n_tokens + parent_t) * H + h_idx) * + S_v * S_v; +#pragma unroll + for (int c = 0; c < COLS; ++c) { + const int col = col_base + c; +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { + const int row = r * WIDTH + lane; + state_shard[c][r] = load_inter_state( + parent_base, col * S_v + row); + } + } + } + } + } + const float * q_t = q + iq3 * sq3 + t * sq2 + iq1 * sq1; const float * k_t = k + iq3 * sq3 + t * sq2 + iq1 * sq1; const float * v_t = v + sequence * sv3 + t * sv2 + h_idx * sv1; @@ -387,12 +465,26 @@ gated_delta_net_cuda_grouped_cols(const float * q, float g_val = 0.0f; float beta_val = 0.0f; if (threadIdx.x == 0) { - g_val = expf(g[gb_offset]); - beta_val = beta[gb_offset]; + if (gate_bias != nullptr) { + // raw-gate mode: g = exp(softplus(alpha_raw + bias) * A), beta = sigmoid(beta_raw) + const float a = g[gb_offset] + gate_bias[h_idx]; + const float sp = (a > 20.0f) ? a : logf(1.0f + expf(a)); + g_val = expf(sp * gate_A[h_idx]); + beta_val = 1.0f / (1.0f + expf(-beta[gb_offset])); + } else { + g_val = expf(g[gb_offset]); + beta_val = beta[gb_offset]; + } } g_val = __shfl_sync(0xffffffffU, g_val, 0); beta_val = __shfl_sync(0xffffffffU, beta_val, 0); + constexpr int journal_width = 2*S_v + 1; + float * journal_t = transition_journal + ? transition_journal + + ((sequence * n_tokens + t) * H + h_idx) * journal_width + : nullptr; + float k_reg[rows_per_lane]; float q_reg[rows_per_lane]; float kv_partial[COLS]; @@ -409,12 +501,20 @@ gated_delta_net_cuda_grouped_cols(const float * q, const float k_val = k_t[row]; q_reg[r] = q_val; k_reg[r] = k_val; + if (journal_t && blockIdx.z == 0 && threadIdx.y == 0 && + subgroup == 0) { + journal_t[1 + row] = k_val; + } #pragma unroll for (int c = 0; c < COLS; ++c) { kv_partial[c] += state_shard[c][r] * k_val; } } + if (journal_t && blockIdx.z == 0 && threadIdx.y == 0 && + subgroup == 0 && lane == 0) { + journal_t[0] = g_val; + } float delta[COLS]; #pragma unroll @@ -423,6 +523,9 @@ gated_delta_net_cuda_grouped_cols(const float * q, float delta_val = 0.0f; if (lane == 0) { delta_val = (v_t[col_base + c] - g_val * kv_col) * beta_val; + if (journal_t) { + journal_t[1 + S_v + col_base + c] = delta_val; + } } delta[c] = gdn_subgroup_broadcast_lane0(delta_val, WIDTH); } @@ -455,7 +558,7 @@ gated_delta_net_cuda_grouped_cols(const float * q, } } - if constexpr (WRITE_INTER) { + if constexpr (WRITE_INTER || TREE_MODE) { #pragma unroll for (int c = 0; c < COLS; ++c) { const int col = col_base + c; @@ -497,7 +600,9 @@ static void launch_gated_delta_net( int64_t sv1, int64_t sv2, int64_t sv3, int64_t sb1, int64_t sb2, int64_t sb3, int64_t neqk1, int64_t rq3, - float scale, cudaStream_t stream) { + float scale, cudaStream_t stream, + const float * gate_bias = nullptr, const float * gate_A = nullptr, + float * transition_journal_d = nullptr) { //TODO: Add chunked kernel for even faster pre-fill const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size; const int num_warps = 4; @@ -519,25 +624,25 @@ static void launch_gated_delta_net( switch (S_v) { case 16: gated_delta_net_cuda<16, KDA, TREE_MODE, WRITE_INTER, InterT><<>>( - q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, H, + q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, transition_journal_d, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, - sb1, sb2, sb3, neqk1_magic, rq3_magic, scale); + sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, gate_bias, gate_A); break; case 32: gated_delta_net_cuda<32, KDA, TREE_MODE, WRITE_INTER, InterT><<>>( - q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, H, + q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, transition_journal_d, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, - sb1, sb2, sb3, neqk1_magic, rq3_magic, scale); + sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, gate_bias, gate_A); break; case 64: { gated_delta_net_cuda<64, KDA, TREE_MODE, WRITE_INTER, InterT><<>>( - q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, H, + q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, transition_journal_d, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, - sb1, sb2, sb3, neqk1_magic, rq3_magic, scale); + sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, gate_bias, gate_A); break; } case 128: { - if constexpr (!KDA && !TREE_MODE) { + if constexpr (!KDA) { if (use_grouped_cols && ((GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_AMPERE) || GGML_CUDA_CC_IS_AMD(cc))) { @@ -549,35 +654,35 @@ static void launch_gated_delta_net( constexpr int groups_per_warp = 32 / width; dim3 grouped_grid_dims(H, n_seqs, (groups + column_groups_per_block * groups_per_warp - 1) / (column_groups_per_block * groups_per_warp)); dim3 grouped_block_dims(32, column_groups_per_block, 1); - gated_delta_net_cuda_grouped_cols<128, cols, width, 32, WRITE_INTER, InterT><<>>( - q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, persist_inter_d, H, + gated_delta_net_cuda_grouped_cols<128, cols, width, 32, TREE_MODE, WRITE_INTER, InterT><<>>( + q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, parent_ids_d, dst_d, state_out_d, persist_inter_d, transition_journal_d, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, - sb1, sb2, sb3, neqk1_magic, rq3_magic, scale); + sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, gate_bias, gate_A); } else if (warp_size == 64) { constexpr int groups_per_warp = 64 / width; dim3 grouped_grid_dims(H, n_seqs, (groups + column_groups_per_block * groups_per_warp - 1) / (column_groups_per_block * groups_per_warp)); dim3 grouped_block_dims(64, column_groups_per_block, 1); - gated_delta_net_cuda_grouped_cols<128, cols, width, 64, WRITE_INTER, InterT><<>>( - q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, persist_inter_d, H, + gated_delta_net_cuda_grouped_cols<128, cols, width, 64, TREE_MODE, WRITE_INTER, InterT><<>>( + q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, parent_ids_d, dst_d, state_out_d, persist_inter_d, transition_journal_d, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, - sb1, sb2, sb3, neqk1_magic, rq3_magic, scale); + sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, gate_bias, gate_A); } else { gated_delta_net_cuda<128, KDA, TREE_MODE, WRITE_INTER, InterT><<>>( - q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, H, + q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, transition_journal_d, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, - sb1, sb2, sb3, neqk1_magic, rq3_magic, scale); + sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, gate_bias, gate_A); } } else { gated_delta_net_cuda<128, KDA, TREE_MODE, WRITE_INTER, InterT><<>>( - q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, H, + q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, transition_journal_d, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, - sb1, sb2, sb3, neqk1_magic, rq3_magic, scale); + sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, gate_bias, gate_A); } } else { gated_delta_net_cuda<128, KDA, TREE_MODE, WRITE_INTER, InterT><<>>( - q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, H, + q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, transition_journal_d, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, - sb1, sb2, sb3, neqk1_magic, rq3_magic, scale); + sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, gate_bias, gate_A); } break; } @@ -605,6 +710,8 @@ void ggml_cuda_op_gated_delta_net(ggml_backend_cuda_context & ctx, ggml_tensor * // Optional 9th source maps compact sequence rows to physical recurrent // state slabs. Negative ids are graph-bucket padding rows. ggml_tensor * src_active_slots = dst->src[8]; + // Optional compact transition journal [J,H,T,B]. Linear-chain only. + ggml_tensor * src_transition_journal = dst->src[11]; GGML_TENSOR_LOCALS(int64_t, neq, src_q, ne); GGML_TENSOR_LOCALS(size_t , nbq, src_q, nb); @@ -647,6 +754,9 @@ void ggml_cuda_op_gated_delta_net(ggml_backend_cuda_context & ctx, ggml_tensor * void * persist_inter_d = src_persist_inter ? src_persist_inter->data : nullptr; + float * transition_journal_d = src_transition_journal + ? (float *) src_transition_journal->data + : nullptr; const bool persist_is_f16 = src_persist_inter && src_persist_inter->type == GGML_TYPE_F16; if (src_persist_inter) { @@ -675,6 +785,15 @@ void ggml_cuda_op_gated_delta_net(ggml_backend_cuda_context & ctx, ggml_tensor * GGML_ASSERT(ggml_is_contiguous(src_active_slots)); GGML_ASSERT(ggml_nelements(src_active_slots) == n_seqs); } + if (src_transition_journal) { + const int64_t journal_width = kda ? 3*S_v : 2*S_v + 1; + GGML_ASSERT(src_transition_journal->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(src_transition_journal)); + GGML_ASSERT(src_transition_journal->ne[0] == journal_width); + GGML_ASSERT(src_transition_journal->ne[1] == H); + GGML_ASSERT(src_transition_journal->ne[2] == n_tokens); + GGML_ASSERT(src_transition_journal->ne[3] == n_seqs); + } // strides in floats (beta strides used for both g and beta offset computation) const int64_t sq1 = nbq1 / sizeof(float); @@ -693,6 +812,18 @@ void ggml_cuda_op_gated_delta_net(ggml_backend_cuda_context & ctx, ggml_tensor * const bool tree_mode = (parent_ids_d != nullptr); const bool skip_intermediate = ggml_get_op_params_i32(dst, 0) != 0; + // dflash raw-gate mode: src[9] = dt_bias[H], src[10] = A[H]; src[8] + // remains available for the optional active-slot map. The kernel + // applies sigmoid / softplus+bias / A itself (see ggml_gated_delta_net_set_raw_gates). + const bool raw_gates = ggml_get_op_params_i32(dst, 2) != 0; + const float * gate_bias_d = nullptr; + const float * gate_A_d = nullptr; + if (raw_gates) { + GGML_ASSERT(dst->src[9] && dst->src[10]); + GGML_ASSERT(!kda && !tree_mode); + gate_bias_d = (const float *) dst->src[9]->data; + gate_A_d = (const float *) dst->src[10]->data; + } const bool write_intermediate = tree_mode || !skip_intermediate || persist_inter_d != nullptr; // Macro to expand KDA × TREE_MODE × WRITE_INTER for a given InterT. @@ -704,35 +835,35 @@ void ggml_cuda_op_gated_delta_net(ggml_backend_cuda_context & ctx, ggml_tensor * if (tree_mode) { \ launch_gated_delta_net( \ q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_typed, \ - S_v, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, \ - sb1, sb2, sb3, neqk1, rq3, scale, stream); \ + S_v, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, \ + sb1, sb2, sb3, neqk1, rq3, scale, stream, gate_bias_d, gate_A_d, transition_journal_d); \ } else if (write_intermediate) { \ launch_gated_delta_net( \ q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, nullptr, persist_typed, \ - S_v, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, \ - sb1, sb2, sb3, neqk1, rq3, scale, stream); \ + S_v, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, \ + sb1, sb2, sb3, neqk1, rq3, scale, stream, gate_bias_d, gate_A_d, transition_journal_d); \ } else { \ launch_gated_delta_net( \ q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, nullptr, persist_typed, \ - S_v, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, \ - sb1, sb2, sb3, neqk1, rq3, scale, stream); \ + S_v, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, \ + sb1, sb2, sb3, neqk1, rq3, scale, stream, gate_bias_d, gate_A_d, transition_journal_d); \ } \ } else { \ if (tree_mode) { \ launch_gated_delta_net( \ q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_typed, \ - S_v, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, \ - sb1, sb2, sb3, neqk1, rq3, scale, stream); \ + S_v, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, \ + sb1, sb2, sb3, neqk1, rq3, scale, stream, gate_bias_d, gate_A_d, transition_journal_d); \ } else if (write_intermediate) { \ launch_gated_delta_net( \ q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, nullptr, persist_typed, \ - S_v, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, \ - sb1, sb2, sb3, neqk1, rq3, scale, stream); \ + S_v, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, \ + sb1, sb2, sb3, neqk1, rq3, scale, stream, gate_bias_d, gate_A_d, transition_journal_d); \ } else { \ launch_gated_delta_net( \ q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, nullptr, persist_typed, \ - S_v, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, \ - sb1, sb2, sb3, neqk1, rq3, scale, stream); \ + S_v, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, \ + sb1, sb2, sb3, neqk1, rq3, scale, stream, gate_bias_d, gate_A_d, transition_journal_d); \ } \ } \ } while (0) diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/gdn-transition-journal.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/gdn-transition-journal.cu new file mode 100644 index 000000000..9ca1cdd9f --- /dev/null +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/gdn-transition-journal.cu @@ -0,0 +1,501 @@ +#include "common.cuh" +#include "ggml-cuda.h" + +#include +#include +#include + +#if defined(GGML_USE_HIP) +#ifndef cudaPointerAttributes +#define cudaPointerAttributes hipPointerAttribute_t +#define cudaPointerGetAttributes hipPointerGetAttributes +#define cudaMemoryTypeDevice hipMemoryTypeDevice +#define cudaMemoryTypeManaged hipMemoryTypeManaged +#endif +#endif + +namespace { + +__global__ void gdn_transition_journal_commit_kernel( + const float * journal, + float * state, + const int32_t * accepted_prefixes, + const int32_t * active_slot_ids, + int state_size, + int n_heads, + int n_tokens, + int n_seqs, + int n_state_slots, + int journal_width, + int gate_values) { + const int sequence = blockIdx.z; + const int head = blockIdx.y; + const int element = blockIdx.x * blockDim.x + threadIdx.x; + const int state_elements = state_size * state_size; + if (sequence >= n_seqs || head >= n_heads || + element >= state_elements) { + return; + } + + const int slot = active_slot_ids[sequence]; + const int accepted = accepted_prefixes[sequence]; + if (slot < 0 || slot >= n_state_slots || + accepted < 0 || accepted > n_tokens) { + return; + } + + const int row = element % state_size; + const int col = element / state_size; + const size_t state_offset = + (((size_t) slot*n_heads + head)*state_size + col)*state_size + row; + float current = state[state_offset]; + + for (int token = 0; token < accepted; ++token) { + const float * transition = journal + + (((size_t) sequence*n_tokens + token)*n_heads + head) * + journal_width; + const float gate = gate_values == 1 + ? transition[0] + : transition[row]; + const float key = transition[gate_values + row]; + const float delta = + transition[gate_values + state_size + col]; + current = fmaf(key, delta, gate * current); + } + + state[state_offset] = current; +} + +bool device_pointer(const void * pointer, int & device) { + if (pointer == nullptr) return false; + cudaPointerAttributes attributes{}; + if (cudaPointerGetAttributes(&attributes, pointer) != cudaSuccess) { + (void) cudaGetLastError(); + return false; + } + if (attributes.type != cudaMemoryTypeDevice && + attributes.type != cudaMemoryTypeManaged) { + return false; + } + device = attributes.device; + return true; +} + +} // namespace + +extern "C" bool ggml_backend_cuda_gdn_transition_journal_commit( + const ggml_tensor * journal, + ggml_tensor * state, + const ggml_tensor * accepted_prefixes, + const ggml_tensor * active_slot_ids) { + if (!journal || !state || !accepted_prefixes || !active_slot_ids || + journal->type != GGML_TYPE_F32 || + state->type != GGML_TYPE_F32 || + accepted_prefixes->type != GGML_TYPE_I32 || + active_slot_ids->type != GGML_TYPE_I32 || + !ggml_is_contiguous(journal) || + !ggml_is_contiguous(state) || + !ggml_is_contiguous(accepted_prefixes) || + !ggml_is_contiguous(active_slot_ids)) { + return false; + } + + const int64_t state_size = state->ne[0]; + const int64_t n_heads = state->ne[2]; + const int64_t n_state_slots = state->ne[3]; + const int64_t journal_width = journal->ne[0]; + const int64_t n_tokens = journal->ne[2]; + const int64_t n_seqs = journal->ne[3]; + const bool supported_state_size = + state_size == 16 || state_size == 32 || + state_size == 64 || state_size == 128; + if (!supported_state_size || state->ne[1] != state_size || + n_heads < 1 || journal->ne[1] != n_heads || + n_tokens < 1 || n_seqs < 1 || n_state_slots < 1 || + ggml_nelements(accepted_prefixes) != n_seqs || + ggml_nelements(active_slot_ids) != n_seqs || + (journal_width != 2*state_size + 1 && + journal_width != 3*state_size) || + state_size > std::numeric_limits::max() || + n_heads > 65535 || + n_tokens > std::numeric_limits::max() || + n_seqs > 65535 || + n_state_slots > std::numeric_limits::max() || + journal_width > std::numeric_limits::max()) { + return false; + } + + int device = -1; + int pointer_device = -1; + const void * pointers[] = { + journal->data, state->data, + accepted_prefixes->data, active_slot_ids->data, + }; + for (const void * pointer : pointers) { + if (!device_pointer(pointer, pointer_device)) return false; + if (device < 0) device = pointer_device; + if (pointer_device != device) return false; + } + ggml_cuda_set_device(device); + + std::vector accepted((size_t) n_seqs); + std::vector slots((size_t) n_seqs); + const size_t map_bytes = (size_t) n_seqs * sizeof(int32_t); + if (cudaMemcpy(accepted.data(), accepted_prefixes->data, map_bytes, + cudaMemcpyDeviceToHost) != cudaSuccess || + cudaMemcpy(slots.data(), active_slot_ids->data, map_bytes, + cudaMemcpyDeviceToHost) != cudaSuccess) { + return false; + } + + std::vector seen((size_t) n_state_slots, 0); + for (int64_t sequence = 0; sequence < n_seqs; ++sequence) { + if (accepted[(size_t) sequence] < 0 || + accepted[(size_t) sequence] > n_tokens) { + return false; + } + const int32_t slot = slots[(size_t) sequence]; + if (slot < 0 || slot >= n_state_slots) continue; + if (seen[(size_t) slot]) return false; + seen[(size_t) slot] = 1; + } + + constexpr int threads = 256; + const int64_t state_elements = state_size * state_size; + const dim3 grid( + (unsigned int) ((state_elements + threads - 1) / threads), + (unsigned int) n_heads, + (unsigned int) n_seqs); + (void) cudaGetLastError(); + gdn_transition_journal_commit_kernel<<>>( + (const float *) journal->data, + (float *) state->data, + (const int32_t *) accepted_prefixes->data, + (const int32_t *) active_slot_ids->data, + (int) state_size, + (int) n_heads, + (int) n_tokens, + (int) n_seqs, + (int) n_state_slots, + (int) journal_width, + journal_width == 2*state_size + 1 ? 1 : (int) state_size); + if (cudaGetLastError() != cudaSuccess) return false; + return cudaDeviceSynchronize() == cudaSuccess; +} + +namespace { + +__global__ void gdn_conv_journal_commit_kernel( + const float * conv_input, + float * conv_state, + const int32_t * accepted_prefixes, + const int32_t * active_slot_ids, + int window, + int channels, + int n_tokens, + int n_seqs, + int n_state_slots) { + const int sequence = blockIdx.y; + const int element = blockIdx.x * blockDim.x + threadIdx.x; + const int count = window * channels; + if (sequence >= n_seqs || element >= count) return; + const int slot = active_slot_ids[sequence]; + const int accepted = accepted_prefixes[sequence]; + if (slot < 0 || slot >= n_state_slots || + accepted < 0 || accepted > n_tokens) return; + const int k = element % window; + const int channel = element / window; + const size_t source = + ((size_t) sequence * channels + channel) * (window + n_tokens) + + accepted + k; + const size_t destination = + ((size_t) slot * channels + channel) * window + k; + conv_state[destination] = conv_input[source]; +} + +__global__ void tree_cache_commit_kernel( + uint8_t * cache, + const int64_t * commit_rows, + const int32_t * active_slot_ids, + size_t row_bytes, + size_t head_stride, + int n_heads, + int tree_width, + int n_seqs, + int n_cache_rows, + int scratch_base, + int scratch_stride) { + const int byte = blockIdx.x * blockDim.x + threadIdx.x; + const int flat = blockIdx.y; + const int head = blockIdx.z; + if ((size_t) byte >= row_bytes || flat >= tree_width*n_seqs || + head >= n_heads) return; + const int lane = flat / tree_width; + const int node = flat % tree_width; + const int slot = active_slot_ids[lane]; + const int64_t destination_row = commit_rows[flat]; + if (slot < 0 || destination_row < 0 || + destination_row >= n_cache_rows) return; + const int64_t source_row = + (int64_t) scratch_base + (int64_t) slot*scratch_stride + node; + if (source_row < 0 || source_row >= n_cache_rows) return; + const size_t source = + (size_t) head*head_stride + (size_t) source_row*row_bytes + byte; + const size_t destination = + (size_t) head*head_stride + (size_t) destination_row*row_bytes + byte; + cache[destination] = cache[source]; +} + +__global__ void tree_feature_commit_kernel( + const uint8_t * source, + uint8_t * destination, + const int32_t * destination_rows, + size_t row_bytes, + int n_rows, + int destination_capacity) { + const int byte = blockIdx.x * blockDim.x + threadIdx.x; + const int source_row = blockIdx.y; + if ((size_t) byte >= row_bytes || source_row >= n_rows) return; + const int destination_row = destination_rows[source_row]; + if (destination_row < 0 || destination_row >= destination_capacity) return; + destination[(size_t) destination_row*row_bytes + byte] = + source[(size_t) source_row*row_bytes + byte]; +} + +bool same_device_pointer(const void * pointer, int expected_device) { + int pointer_device = -1; + return device_pointer(pointer, pointer_device) && + pointer_device == expected_device; +} + +} // namespace + +extern "C" bool ggml_backend_cuda_gdn_transition_journal_commit_many( + const ggml_tensor * const * journals, + ggml_tensor * const * states, + const ggml_tensor * const * conv_inputs, + ggml_tensor * const * conv_states, + int n_layers, + const ggml_tensor * accepted_prefixes, + const ggml_tensor * active_slot_ids) { + if (!journals || !states || !conv_inputs || !conv_states || + n_layers <= 0 || !accepted_prefixes || !active_slot_ids || + accepted_prefixes->type != GGML_TYPE_I32 || + active_slot_ids->type != GGML_TYPE_I32 || + !ggml_is_contiguous(accepted_prefixes) || + !ggml_is_contiguous(active_slot_ids)) return false; + + const int64_t n_seqs = ggml_nelements(accepted_prefixes); + if (n_seqs < 1 || ggml_nelements(active_slot_ids) != n_seqs) return false; + int device = -1; + if (!device_pointer(accepted_prefixes->data, device) || + !same_device_pointer(active_slot_ids->data, device)) return false; + + std::vector accepted((size_t) n_seqs); + std::vector slots((size_t) n_seqs); + const size_t map_bytes = (size_t) n_seqs*sizeof(int32_t); + if (cudaMemcpy(accepted.data(), accepted_prefixes->data, map_bytes, + cudaMemcpyDeviceToHost) != cudaSuccess || + cudaMemcpy(slots.data(), active_slot_ids->data, map_bytes, + cudaMemcpyDeviceToHost) != cudaSuccess) return false; + + int common_tokens = -1; + int common_state_slots = -1; + std::vector seen; + for (int layer = 0; layer < n_layers; ++layer) { + const ggml_tensor * journal = journals[layer]; + ggml_tensor * state = states[layer]; + const ggml_tensor * conv_input = conv_inputs[layer]; + ggml_tensor * conv_state = conv_states[layer]; + if (!journal || !state || !conv_input || !conv_state || + journal->type != GGML_TYPE_F32 || state->type != GGML_TYPE_F32 || + conv_input->type != GGML_TYPE_F32 || conv_state->type != GGML_TYPE_F32 || + !ggml_is_contiguous(journal) || !ggml_is_contiguous(state) || + !ggml_is_contiguous(conv_input) || !ggml_is_contiguous(conv_state)) return false; + const int64_t state_size = state->ne[0]; + const int64_t heads = state->ne[2]; + const int64_t tokens = journal->ne[2]; + const int64_t state_slots = state->ne[3]; + if ((state_size != 16 && state_size != 32 && + state_size != 64 && state_size != 128) || + state->ne[1] != state_size || heads < 1 || + journal->ne[1] != heads || journal->ne[3] != n_seqs || + (journal->ne[0] != 2*state_size + 1 && + journal->ne[0] != 3*state_size) || tokens < 1 || + conv_state->ne[0] < 1 || conv_state->ne[1] < 1 || + conv_state->ne[2] != state_slots || conv_state->ne[3] != 1 || + conv_input->ne[0] != conv_state->ne[0] + tokens || + conv_input->ne[1] != conv_state->ne[1] || + conv_input->ne[2] != n_seqs || conv_input->ne[3] != 1) return false; + if (common_tokens < 0) { + common_tokens = (int) tokens; + common_state_slots = (int) state_slots; + seen.assign((size_t) state_slots, 0); + } else if (tokens != common_tokens || state_slots != common_state_slots) { + return false; + } + const void * pointers[] = { + journal->data, state->data, conv_input->data, conv_state->data, + }; + for (const void * pointer : pointers) { + if (!same_device_pointer(pointer, device)) return false; + } + } + for (int64_t lane = 0; lane < n_seqs; ++lane) { + if (accepted[(size_t) lane] < 0 || + accepted[(size_t) lane] > common_tokens) return false; + const int slot = slots[(size_t) lane]; + if (slot < 0 || slot >= common_state_slots) continue; + if (seen[(size_t) slot]) return false; + seen[(size_t) slot] = 1; + } + + ggml_cuda_set_device(device); + constexpr int threads = 256; + (void) cudaGetLastError(); + for (int layer = 0; layer < n_layers; ++layer) { + const ggml_tensor * journal = journals[layer]; + ggml_tensor * state = states[layer]; + const int state_size = (int) state->ne[0]; + const int heads = (int) state->ne[2]; + const int tokens = (int) journal->ne[2]; + const int journal_width = (int) journal->ne[0]; + const int64_t state_elements = (int64_t) state_size*state_size; + const dim3 state_grid( + (unsigned int) ((state_elements + threads - 1)/threads), + (unsigned int) heads, (unsigned int) n_seqs); + gdn_transition_journal_commit_kernel<<>>( + (const float *) journal->data, (float *) state->data, + (const int32_t *) accepted_prefixes->data, + (const int32_t *) active_slot_ids->data, + state_size, heads, tokens, (int) n_seqs, + (int) state->ne[3], journal_width, + journal_width == 2*state_size + 1 ? 1 : state_size); + + const ggml_tensor * conv_input = conv_inputs[layer]; + ggml_tensor * conv_state = conv_states[layer]; + const int conv_elements = + (int) (conv_state->ne[0]*conv_state->ne[1]); + const dim3 conv_grid( + (unsigned int) ((conv_elements + threads - 1)/threads), + (unsigned int) n_seqs, 1); + gdn_conv_journal_commit_kernel<<>>( + (const float *) conv_input->data, (float *) conv_state->data, + (const int32_t *) accepted_prefixes->data, + (const int32_t *) active_slot_ids->data, + (int) conv_state->ne[0], (int) conv_state->ne[1], tokens, + (int) n_seqs, (int) conv_state->ne[2]); + } + if (cudaGetLastError() != cudaSuccess) return false; + return cudaDeviceSynchronize() == cudaSuccess; +} + +extern "C" bool ggml_backend_cuda_tree_cache_commit_many( + ggml_tensor * const * caches, + int n_caches, + const ggml_tensor * commit_rows, + const ggml_tensor * active_slot_ids, + int tree_scratch_base, + int tree_scratch_stride) { + if (!caches || n_caches <= 0 || !commit_rows || !active_slot_ids || + commit_rows->type != GGML_TYPE_I64 || + active_slot_ids->type != GGML_TYPE_I32 || + !ggml_is_contiguous(commit_rows) || + !ggml_is_contiguous(active_slot_ids) || + commit_rows->ne[0] < 1 || commit_rows->ne[1] < 1 || + ggml_nelements(active_slot_ids) != commit_rows->ne[1] || + tree_scratch_base < 0 || tree_scratch_stride < commit_rows->ne[0]) return false; + const int tree_width = (int) commit_rows->ne[0]; + const int n_seqs = (int) commit_rows->ne[1]; + const int n_rows = tree_width*n_seqs; + int device = -1; + if (!device_pointer(commit_rows->data, device) || + !same_device_pointer(active_slot_ids->data, device)) return false; + + std::vector destinations((size_t) n_rows); + std::vector slots((size_t) n_seqs); + if (cudaMemcpy(destinations.data(), commit_rows->data, + destinations.size()*sizeof(int64_t), cudaMemcpyDeviceToHost) != cudaSuccess || + cudaMemcpy(slots.data(), active_slot_ids->data, + slots.size()*sizeof(int32_t), cudaMemcpyDeviceToHost) != cudaSuccess) return false; + int cache_rows = -1; + for (int index = 0; index < n_caches; ++index) { + ggml_tensor * cache = caches[index]; + if (!cache || !ggml_is_contiguous(cache) || cache->ne[0] < 1 || + cache->ne[1] < 1 || cache->ne[2] < 1 || cache->ne[3] != 1 || + cache->nb[1] < ggml_row_size(cache->type, cache->ne[0]) || + !same_device_pointer(cache->data, device)) return false; + if (cache_rows < 0) cache_rows = (int) cache->ne[1]; + else if (cache->ne[1] != cache_rows) return false; + } + for (int lane = 0; lane < n_seqs; ++lane) { + const int slot = slots[(size_t) lane]; + if (slot < 0) continue; + const int64_t source_end = (int64_t) tree_scratch_base + + (int64_t) slot*tree_scratch_stride + tree_width; + if (source_end > cache_rows) return false; + for (int node = 0; node < tree_width; ++node) { + const int64_t destination = + destinations[(size_t) lane*tree_width + node]; + if (destination < -1 || destination >= cache_rows) return false; + } + } + + ggml_cuda_set_device(device); + constexpr int threads = 256; + (void) cudaGetLastError(); + for (int index = 0; index < n_caches; ++index) { + ggml_tensor * cache = caches[index]; + const dim3 grid( + (unsigned int) ((cache->nb[1] + threads - 1)/threads), + (unsigned int) n_rows, (unsigned int) cache->ne[2]); + tree_cache_commit_kernel<<>>( + (uint8_t *) cache->data, + (const int64_t *) commit_rows->data, + (const int32_t *) active_slot_ids->data, + cache->nb[1], cache->nb[2], (int) cache->ne[2], + tree_width, n_seqs, cache_rows, + tree_scratch_base, tree_scratch_stride); + } + if (cudaGetLastError() != cudaSuccess) return false; + return cudaDeviceSynchronize() == cudaSuccess; +} + +extern "C" bool ggml_backend_cuda_tree_feature_commit( + const ggml_tensor * source, + ggml_tensor * destination, + const ggml_tensor * destination_rows) { + if (!source || !destination || !destination_rows || + source->type != destination->type || source->type != GGML_TYPE_BF16 || + destination_rows->type != GGML_TYPE_I32 || + !ggml_is_contiguous(source) || !ggml_is_contiguous(destination) || + !ggml_is_contiguous(destination_rows) || + source->ne[0] != destination->ne[0] || source->ne[2] != 1 || + source->ne[3] != 1 || destination->ne[2] != 1 || + destination->ne[3] != 1 || + ggml_nelements(destination_rows) != source->ne[1] || + source->nb[1] != destination->nb[1]) return false; + int device = -1; + if (!device_pointer(source->data, device) || + !same_device_pointer(destination->data, device) || + !same_device_pointer(destination_rows->data, device)) return false; + const int n_rows = (int) source->ne[1]; + std::vector rows((size_t) n_rows); + if (cudaMemcpy(rows.data(), destination_rows->data, + rows.size()*sizeof(int32_t), cudaMemcpyDeviceToHost) != cudaSuccess) return false; + for (int row : rows) { + if (row < -1 || row >= destination->ne[1]) return false; + } + ggml_cuda_set_device(device); + constexpr int threads = 256; + const dim3 grid( + (unsigned int) ((source->nb[1] + threads - 1)/threads), + (unsigned int) n_rows, 1); + (void) cudaGetLastError(); + tree_feature_commit_kernel<<>>( + (const uint8_t *) source->data, (uint8_t *) destination->data, + (const int32_t *) destination_rows->data, + source->nb[1], n_rows, (int) destination->ne[1]); + if (cudaGetLastError() != cudaSuccess) return false; + return cudaDeviceSynchronize() == cudaSuccess; +} diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu index 1c55d51d1..543417701 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu @@ -473,7 +473,11 @@ const ggml_cuda_device_info & ggml_cuda_info() { // buffer pool for cuda (legacy) struct ggml_cuda_pool_leg : public ggml_cuda_pool { - static const int MAX_BUFFERS = 256; + // 1024 (upstream 256): LUCE_Q8_MEMO keeps one pooled q8_1 activation + // buffer per quantized matmul alive across a whole graph evaluation + // (~300 on a 64-layer hybrid), and a full pool falls back to freeing + // in-flight buffers with cudaFree. + static const int MAX_BUFFERS = 1024; int device; struct ggml_cuda_buffer { @@ -4311,6 +4315,49 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, } } + // dflash: residual ADD + RMS_NORM + MUL. The add output stays live (it is + // the next residual), so this is a subgraph fusion with two outputs. + if (ops.size() == 3 && ops.begin()[0] == GGML_OP_ADD && ops.begin()[1] == GGML_OP_RMS_NORM && + ops.begin()[2] == GGML_OP_MUL) { + if (!ggml_can_fuse_subgraph(cgraph, node_idx, ops, { node_idx, node_idx + 2 })) { + return false; + } + const ggml_tensor * add = cgraph->nodes[node_idx]; + const ggml_tensor * rms = cgraph->nodes[node_idx + 1]; + const ggml_tensor * mul = cgraph->nodes[node_idx + 2]; + if (rms->src[0] != add) { + return false; + } + const ggml_tensor * w = nullptr; + if (mul->src[0] == rms) { + w = mul->src[1]; + } else if (mul->src[1] == rms) { + w = mul->src[0]; + } else { + return false; + } + const ggml_tensor * a = add->src[0]; + const ggml_tensor * b = add->src[1]; + if (a->type != GGML_TYPE_F32 || b->type != GGML_TYPE_F32 || w->type != GGML_TYPE_F32 || + add->type != GGML_TYPE_F32 || rms->type != GGML_TYPE_F32 || mul->type != GGML_TYPE_F32) { + return false; + } + if (!ggml_is_contiguous(a) || !ggml_is_contiguous(b) || !ggml_is_contiguous(w) || + !ggml_is_contiguous(add) || !ggml_is_contiguous(mul)) { + return false; + } + if (!ggml_are_same_shape(a, b) || !ggml_are_same_shape(a, add) || !ggml_are_same_shape(a, mul)) { + return false; + } + if (w->ne[0] != a->ne[0] || ggml_nelements(w) != a->ne[0]) { + return false; + } + if (ggml_backend_buft_is_cuda_split(a->buffer->buft) || ggml_backend_buft_is_cuda_split(b->buffer->buft)) { + return false; + } + return true; + } + if (!ggml_can_fuse(cgraph, node_idx, ops)) { return false; } @@ -4910,6 +4957,12 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud continue; } + if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_ADD, GGML_OP_RMS_NORM, GGML_OP_MUL }, {})) { + ggml_cuda_op_add_rms_norm_mul_fused(*cuda_ctx, node, cgraph->nodes[i+1], cgraph->nodes[i+2]); + i += 2; + continue; + } + if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ADD}, {})) { ggml_cuda_op_rms_norm_fused_add(*cuda_ctx, node, cgraph->nodes[i+1], cgraph->nodes[i+2]); i += 2; @@ -6157,6 +6210,10 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g } } case GGML_OP_SSM_CONV: { + // dflash fused step mode handles any channel count + if (ggml_get_op_params_i32(op, 0) == 1) { + return true; + } // assumes d_inner % threads == 0 return op->src[0]->ne[1] % 128 == 0; } diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh index 876a8ba45..29a1d4db2 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh @@ -107,9 +107,14 @@ struct tile_x_sizes { int sc; }; -// RDNA uses 128x128, eight-warp MMQ tiles by default. Q4_K narrows the row -// dimension to 128x64, while ROCmFPX uses 64x64 four-warp tiles. Their -// unpacking pressure makes the smaller tiles faster on gfx1151. +// RDNA uses 128x128, eight-warp MMQ tiles by default. Template instances +// compiled with GGML_CUDA_MMQ_SMALL_TILE use 64x64, four-warp tiles: +// - ROCmFPX formats: their unpacking pressure makes the smaller tile faster +// on gfx1151; +// - IQ4_XS / Q6_K / Q8_0 (dense hybrid targets): at spec-decode verify +// widths (N<=16) the 128-row tile leaves a 5120-row projection with only +// 40 blocks on a 64-CU gfx1201; the small tile measured +12-23% there +// (mmq_probe) at the cost of ~8% prefill throughput. #ifndef LUCEBOX_RDNA_MMQ_TILE_OVERRIDE #define LUCEBOX_RDNA_MMQ_TILE_OVERRIDE 1 #endif @@ -122,7 +127,7 @@ struct tile_x_sizes { static int get_mmq_x_max_host(const int cc) { if (LUCEBOX_RDNA_TILE_HOST(cc)) { -#if defined(GGML_CUDA_ROCMFPX_MMQ_TILE) +#if defined(GGML_CUDA_MMQ_SMALL_TILE) return 64; #else return 128; @@ -139,7 +144,7 @@ static int get_mmq_x_max_host(const int cc) { static constexpr __device__ int get_mmq_x_max_device() { #if LUCEBOX_RDNA_TILE_DEVICE -#if defined(GGML_CUDA_ROCMFPX_MMQ_TILE) +#if defined(GGML_CUDA_MMQ_SMALL_TILE) return 64; #else return 128; @@ -169,7 +174,7 @@ static constexpr __device__ int get_mmq_x_max_device() { static int get_mmq_y_host(const int cc) { if (LUCEBOX_RDNA_TILE_HOST(cc)) { -#if defined(GGML_CUDA_ROCMFPX_MMQ_TILE) +#if defined(GGML_CUDA_MMQ_SMALL_TILE) return 64; #elif defined(LUCEBOX_RDNA_MMQ_Y) return LUCEBOX_RDNA_MMQ_Y; @@ -191,7 +196,7 @@ static constexpr __device__ int get_iter_k([[maybe_unused]] const ggml_type type static constexpr __device__ int get_mmq_y_device() { #if LUCEBOX_RDNA_TILE_DEVICE -#if defined(GGML_CUDA_ROCMFPX_MMQ_TILE) +#if defined(GGML_CUDA_MMQ_SMALL_TILE) return 64; #elif defined(LUCEBOX_RDNA_MMQ_Y) return LUCEBOX_RDNA_MMQ_Y; @@ -346,7 +351,7 @@ static constexpr __device__ int mmq_get_granularity_device(const int /*mmq_x*/) #if defined(GGML_USE_HIP) static int mmq_get_nwarps_host(const int cc, const int warp_size) { if (LUCEBOX_RDNA_TILE_HOST(cc)) { -#if defined(GGML_CUDA_ROCMFPX_MMQ_TILE) +#if defined(GGML_CUDA_MMQ_SMALL_TILE) return 4; #elif defined(LUCEBOX_RDNA_MMQ_Y) return 4; @@ -364,7 +369,7 @@ static int mmq_get_nwarps_host(const int /*cc*/, const int warp_size) { static constexpr __device__ int mmq_get_nwarps_device() { #if LUCEBOX_RDNA_TILE_DEVICE -#if defined(GGML_CUDA_ROCMFPX_MMQ_TILE) +#if defined(GGML_CUDA_MMQ_SMALL_TILE) return 4; #elif defined(LUCEBOX_RDNA_MMQ_Y) return 4; @@ -4223,7 +4228,7 @@ template #if defined(GGML_USE_HIP) // RDNA4 is compute-bound on MMQ (WMMA path); allow compiler to use more VGPRs // (minBlocks=1 matches NVIDIA Volta+ behavior and reduces register spilling). -#if defined(RDNA4) && !defined(GGML_CUDA_ROCMFPX_MMQ_TILE) +#if defined(RDNA4) && !defined(GGML_CUDA_MMQ_SMALL_TILE) __launch_bounds__(ggml_cuda_get_physical_warp_size()*mmq_get_nwarps_device(), 1) #elif defined(RDNA3) || defined(RDNA2) || defined(CDNA) || defined(GCN) __launch_bounds__(ggml_cuda_get_physical_warp_size()*mmq_get_nwarps_device(), 2) @@ -4785,6 +4790,14 @@ void mul_mat_q_case(ggml_backend_cuda_context & ctx, const mmq_args & args, cuda if (mmq_x % granularity != 0 || mmq_get_nbytes_shared(mmq_x, mmq_y, cc, warp_size, nwarps) > smpbo) { continue; } +#if defined(GGML_CUDA_MMQ_SMALL_TILE) + // The 64-row/4-warp tile is pathological at mmq_x == 32 on gfx1201 + // (17408x5120 IQ4_XS: N=16 443 GB/s, N=24..32 180 GB/s, N=48 315 GB/s + // in mmq_probe); a wider tile with more padding is still faster. + if (LUCEBOX_RDNA_TILE_HOST(cc) && mmq_x == 32) { + continue; + } +#endif const int ntiles_x = (args.ncols_max + mmq_x - 1) / mmq_x; diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/norm.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/norm.cu index ef98f675a..696a6f441 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/norm.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/norm.cu @@ -150,6 +150,57 @@ static __global__ void rms_norm_f32(const float * x, } } +// dflash: residual add fused into the following rms_norm * weight. +// sum = a + b (written to sum_out; it is the next residual) +// dst = rms_norm(sum) * w +// All of a, b, sum_out, dst are contiguous [ncols, R]; w is [ncols]. +template +static __global__ void add_rms_norm_mul_f32(const float * __restrict__ a, + const float * __restrict__ b, + float * __restrict__ sum_out, + float * __restrict__ dst, + const float * __restrict__ w, + const int ncols, + const float eps) { + const int64_t row = blockIdx.x; + const int tid = threadIdx.x; + + a += row * ncols; + b += row * ncols; + sum_out += row * ncols; + dst += row * ncols; + + float tmp = 0.0f; + for (int col = tid; col < ncols; col += block_size) { + const float s = a[col] + b[col]; + sum_out[col] = s; + tmp += s * s; + } + + extern __shared__ float s_sum[]; + tmp = block_reduce(tmp, s_sum); + + const float mean = tmp / ncols; + const float scale = rsqrtf(mean + eps); + + for (int col = tid; col < ncols; col += block_size) { + dst[col] = scale * sum_out[col] * w[col]; + } +} + +static void add_rms_norm_mul_f32_cuda(const float * a, const float * b, float * sum_out, float * dst, + const float * w, const int ncols, const int64_t nrows, + const float eps, cudaStream_t stream) { + const dim3 blocks_num(nrows, 1, 1); + if (ncols < 1024) { + const dim3 block_dims(256, 1, 1); + add_rms_norm_mul_f32<256><<>>(a, b, sum_out, dst, w, ncols, eps); + } else { + const dim3 block_dims(1024, 1, 1); + add_rms_norm_mul_f32<1024><<>>(a, b, sum_out, dst, w, ncols, eps); + } +} + template static __global__ void rms_norm_back_f32( const float * grad, const float * xf, float * dst, const int ncols, const float eps) { @@ -533,6 +584,36 @@ void ggml_cuda_op_rms_norm_fused(ggml_backend_cuda_context & ctx, ggml_tensor * eps, stream); } +// dflash: ADD (residual) + RMS_NORM + MUL in one launch. `add_tensor` is the +// residual add node (its output is materialized), `rms_tensor` is elided, +// `mul_tensor` receives the normalized * weight result. +void ggml_cuda_op_add_rms_norm_mul_fused(ggml_backend_cuda_context & ctx, + ggml_tensor * add_tensor, + ggml_tensor * rms_tensor, + ggml_tensor * mul_tensor) { + const ggml_tensor * a = add_tensor->src[0]; + const ggml_tensor * b = add_tensor->src[1]; + const ggml_tensor * w = (mul_tensor->src[0] == rms_tensor) ? mul_tensor->src[1] : mul_tensor->src[0]; + + float eps = 0.0f; + memcpy(&eps, rms_tensor->op_params, sizeof(float)); + GGML_ASSERT(eps >= 0.0f); + + GGML_ASSERT(a->type == GGML_TYPE_F32 && b->type == GGML_TYPE_F32 && w->type == GGML_TYPE_F32); + GGML_ASSERT(add_tensor->type == GGML_TYPE_F32 && mul_tensor->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(a) && ggml_is_contiguous(b) && ggml_is_contiguous(w)); + GGML_ASSERT(ggml_is_contiguous(add_tensor) && ggml_is_contiguous(mul_tensor)); + GGML_ASSERT(ggml_are_same_shape(a, b) && ggml_are_same_shape(a, add_tensor) && ggml_are_same_shape(a, mul_tensor)); + GGML_ASSERT(w->ne[0] == a->ne[0] && ggml_nelements(w) == a->ne[0]); + + const int ncols = (int) a->ne[0]; + const int64_t nrows = ggml_nrows(a); + + add_rms_norm_mul_f32_cuda((const float *) a->data, (const float *) b->data, + (float *) add_tensor->data, (float *) mul_tensor->data, + (const float *) w->data, ncols, nrows, eps, ctx.stream()); +} + void ggml_cuda_op_rms_norm_fused_add(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * mul_tensor, diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/norm.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/norm.cuh index a74f63767..6313a98ce 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/norm.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/norm.cuh @@ -16,3 +16,6 @@ void ggml_cuda_op_rms_norm_fused_add(ggml_backend_cuda_context & ctx, void ggml_cuda_op_rms_norm_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst); void ggml_cuda_op_l2_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +// dflash: residual ADD + RMS_NORM + MUL fusion (see norm.cu) +void ggml_cuda_op_add_rms_norm_mul_fused(ggml_backend_cuda_context & ctx, ggml_tensor * add_tensor, ggml_tensor * rms_tensor, ggml_tensor * mul_tensor); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/paged-attn.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/paged-attn.cu index c76d97b36..ad321d094 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/paged-attn.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/paged-attn.cu @@ -29,6 +29,42 @@ static __host__ __device__ __forceinline__ int32_t paged_attn_partitions( return requested < available ? requested : available; } +// parent_ids is a sequence-major [tree_width, n_tree_seq] table. Walk only +// from the current query node toward the root; a candidate is visible iff it +// appears on that chain. The bounded walk also turns malformed cycles or +// out-of-range parents into invisible edges instead of an unsafe read. +static __device__ __forceinline__ bool paged_attn_tree_visible( + const char * __restrict__ parent_ids, + int64_t parent_nb0, + int64_t parent_nb1, + int32_t tree_seq, + int32_t query_node, + int32_t candidate, + int32_t tree_size) { + if (candidate < 0 || candidate >= tree_size || + query_node < 0 || query_node >= tree_size) { + return false; + } + + int32_t current = query_node; + for (int32_t depth = 0; depth < tree_size; ++depth) { + if (current == candidate) { + return true; + } + if (current < 0 || current >= tree_size) { + return false; + } + const int32_t parent = *(const int32_t *) ( + parent_ids + (int64_t) current * parent_nb0 + + (int64_t) tree_seq * parent_nb1); + if (parent == current) { + return false; + } + current = parent; + } + return false; +} + // All scores are computed in the log2 domain: log2(e) is folded into the same // Q prescale that already carries the 1/sqrt(D) attention scale, so every // softmax exponential uses the fast exp2f SFU path. @@ -180,6 +216,8 @@ static __global__ void paged_attn_decode( const char * __restrict__ kv_seq_lens, const char * __restrict__ active_slot_ids, const char * __restrict__ query_positions, + const char * __restrict__ parent_ids, + const char * __restrict__ tree_sizes, char * __restrict__ dst, half * __restrict__ partial_acc, float2 * __restrict__ partial_meta, @@ -189,6 +227,7 @@ static __global__ void paged_attn_decode( int64_t bt_nb0, int64_t bt_nb1, int64_t ksl_nb0, int64_t asi_nb0, int64_t qpos_nb0, + int64_t parent_nb0, int64_t parent_nb1, int64_t tree_size_nb0, int64_t dst_nb1, int64_t dst_nb2, int32_t n_table_seq, int32_t n_head, @@ -197,6 +236,10 @@ static __global__ void paged_attn_decode( int32_t max_blocks, int32_t block_size, int32_t min_partitions, + int32_t tree_width, + int32_t tree_row_offset, + int32_t tree_scratch_base, + int32_t tree_scratch_stride, float scale) { constexpr int nthreads = WARP_SIZE; constexpr int values_per_load = 4; @@ -222,29 +265,40 @@ static __global__ void paged_attn_decode( const int n_seq = gridDim.y; const int n_partitions = gridDim.z; + const bool tree_mode = parent_ids != nullptr; const int32_t physical_seq_raw = active_slot_ids ? *(const int32_t *) (active_slot_ids + (int64_t) seq * asi_nb0) : seq; const int32_t query_pos = query_positions ? *(const int32_t *) (query_positions + (int64_t) seq * qpos_nb0) : -1; + const bool tree_query = tree_mode && seq >= tree_row_offset; + const int32_t tree_seq = tree_query + ? (seq - tree_row_offset) / tree_width : 0; + const int32_t query_node = tree_query + ? seq - tree_row_offset - tree_seq * tree_width : -1; + const int32_t tree_size = tree_query + ? *(const int32_t *) ( + tree_sizes + (int64_t) tree_seq * tree_size_nb0) + : 0; // A row is live when its slot id selects a real block-table column and, - // for ragged batches, its causal position is non-negative. Dead rows are - // pinned to column 0 with kv_seq_len forced to 0, which routes every - // partition through the existing zero-output early path; the block table - // is then never read for them. + // for ragged batches, its causal position is non-negative. Tree padding + // rows are validated by tree_sizes. Dead rows are pinned to column 0 with + // an empty virtual context, so the block table and scratch are never read. const bool valid_query = physical_seq_raw >= 0 && physical_seq_raw < n_table_seq && - (!query_positions || query_pos >= 0); + (!query_positions || tree_query || query_pos >= 0) && + (!tree_query || + (tree_size >= 0 && tree_size <= tree_width && + query_node < tree_size)); const int32_t physical_seq = valid_query ? physical_seq_raw : 0; int32_t kv_seq_len_raw = valid_query ? *(const int32_t *) (kv_seq_lens + (int64_t) physical_seq * ksl_nb0) : 0; - // The inclusive clamp IS the causal mask: this row attends tokens - // [0, pos] only, and every downstream bound (partition count, token loop - // extents) already derives from kv_seq_len. - if (query_positions && query_pos < kv_seq_len_raw) { + // The inclusive clamp IS the causal mask for non-tree ragged rows. Tree + // rows always read the whole committed prefix carried by kv_seq_lens. + if (query_positions && !tree_query && query_pos < kv_seq_len_raw) { kv_seq_len_raw = query_pos + 1; } const int64_t table_capacity = @@ -254,8 +308,14 @@ static __global__ void paged_attn_decode( : (kv_seq_len_raw < table_capacity ? kv_seq_len_raw : (int32_t) table_capacity); + // Treat the candidate slab as a virtual tail of tree_width tokens. The + // normal partition split then covers prefix and tree candidates in one + // stable softmax; invisible siblings/padding resolve to no physical row. + const int32_t virtual_tokens = valid_query + ? kv_seq_len + (tree_query ? tree_width : 0) + : 0; const int32_t n_logical_blocks = - (kv_seq_len + block_size - 1) / block_size; + (virtual_tokens + block_size - 1) / block_size; const int32_t active_partitions = paged_attn_partitions(n_logical_blocks, min_partitions, n_partitions); @@ -289,7 +349,7 @@ static __global__ void paged_attn_decode( const int32_t token_begin = logical_block_begin * block_size; const int32_t token_end_blocks = logical_block_end * block_size; const int32_t token_end = - kv_seq_len < token_end_blocks ? kv_seq_len : token_end_blocks; + virtual_tokens < token_end_blocks ? virtual_tokens : token_end_blocks; constexpr bool quantize_q = type_K != GGML_TYPE_F16; constexpr int q_registers = (D / 2) / nthreads; @@ -363,7 +423,12 @@ static __global__ void paged_attn_decode( qk_sum[h] = 0.0f; } - const int32_t n_physical_blocks = pool_tokens / block_size; + // In tree mode the committed block table may address only the prefix + // pool before tree_scratch_base. Candidate rows are addressed directly + // below, keeping uncommitted nodes out of every sequence block table. + const int32_t prefix_pool_tokens = + tree_mode ? tree_scratch_base : pool_tokens; + const int32_t n_physical_blocks = prefix_pool_tokens / block_size; for (int32_t tile_begin = token_begin; tile_begin < token_end; @@ -380,7 +445,7 @@ static __global__ void paged_attn_decode( // read; their tokens contribute nothing, mirroring the CPU reference. int32_t phys_mine = -1; const int32_t my_token = tile_begin + lane; - if (my_token < token_end) { + if (my_token < token_end && my_token < kv_seq_len) { const int32_t logical_block = my_token / block_size; const int32_t physical_block = *(const int32_t *) (block_table + @@ -390,6 +455,19 @@ static __global__ void paged_attn_decode( phys_mine = physical_block * block_size + my_token % block_size; } + } else if (tree_query && my_token < token_end) { + const int32_t candidate = my_token - kv_seq_len; + if (paged_attn_tree_visible( + parent_ids, parent_nb0, parent_nb1, + tree_seq, query_node, candidate, tree_size)) { + const int64_t physical = + (int64_t) tree_scratch_base + + (int64_t) physical_seq * tree_scratch_stride + + candidate; + if (physical >= 0 && physical < pool_tokens) { + phys_mine = (int32_t) physical; + } + } } float score_mine[n_batch_heads]; @@ -619,6 +697,8 @@ bool ggml_cuda_paged_attn_supported(const ggml_tensor * dst) { const ggml_tensor * kv_seq_lens = dst->src[4]; const ggml_tensor * active_slot_ids = dst->src[5]; const ggml_tensor * query_positions = dst->src[6]; + const ggml_tensor * parent_ids = dst->src[7]; + const ggml_tensor * tree_sizes = dst->src[8]; if (!q || !k || !v || !block_table || !kv_seq_lens) { return false; @@ -628,6 +708,11 @@ bool ggml_cuda_paged_attn_supported(const ggml_tensor * dst) { if (query_positions && !active_slot_ids) { return false; } + const bool tree_mode = parent_ids || tree_sizes; + if ((parent_ids == nullptr) != (tree_sizes == nullptr) || + (tree_mode && !active_slot_ids)) { + return false; + } if (dst->type != GGML_TYPE_F32 || q->type != GGML_TYPE_F32 || @@ -636,7 +721,9 @@ bool ggml_cuda_paged_attn_supported(const ggml_tensor * dst) { block_table->type != GGML_TYPE_I32 || kv_seq_lens->type != GGML_TYPE_I32 || (active_slot_ids && active_slot_ids->type != GGML_TYPE_I32) || - (query_positions && query_positions->type != GGML_TYPE_I32)) { + (query_positions && query_positions->type != GGML_TYPE_I32) || + (parent_ids && parent_ids->type != GGML_TYPE_I32) || + (tree_sizes && tree_sizes->type != GGML_TYPE_I32)) { return false; } @@ -647,6 +734,8 @@ bool ggml_cuda_paged_attn_supported(const ggml_tensor * dst) { kv_seq_lens->nb[0] != sizeof(int32_t) || (active_slot_ids && active_slot_ids->nb[0] != sizeof(int32_t)) || (query_positions && query_positions->nb[0] != sizeof(int32_t)) || + (parent_ids && parent_ids->nb[0] != sizeof(int32_t)) || + (tree_sizes && tree_sizes->nb[0] != sizeof(int32_t)) || dst->nb[0] != sizeof(float)) { return false; } @@ -701,10 +790,49 @@ bool ggml_cuda_paged_attn_supported(const ggml_tensor * dst) { const int32_t block_size = ggml_get_op_params_i32(dst, 1); const int32_t max_kv_seq_len = ggml_get_op_params_i32(dst, 2); - return block_size > 0 && - max_kv_seq_len > 0 && - max_kv_seq_len <= k->ne[1] && - k->ne[1] % block_size == 0; + const int32_t tree_width = ggml_get_op_params_i32(dst, 3); + const int32_t tree_scratch_base = ggml_get_op_params_i32(dst, 4); + const int32_t tree_scratch_stride = ggml_get_op_params_i32(dst, 5); + if (block_size <= 0 || + max_kv_seq_len <= 0 || + (int64_t) max_kv_seq_len + tree_width > INT32_MAX || + k->ne[1] % block_size != 0) { + return false; + } + + if (!tree_mode) { + return tree_width == 0 && + tree_scratch_base == 0 && + tree_scratch_stride == 0; + } + + if (tree_width <= 0 || + tree_scratch_base <= 0 || + tree_scratch_base % block_size != 0 || + tree_scratch_stride < tree_width || + !ggml_is_contiguous(parent_ids) || + !ggml_is_contiguous(tree_sizes) || + parent_ids->ne[0] != tree_width || + parent_ids->ne[1] <= 0 || + parent_ids->ne[1] != tree_sizes->ne[0] || + parent_ids->ne[2] != 1 || + parent_ids->ne[3] != 1 || + tree_sizes->ne[1] != 1 || + tree_sizes->ne[2] != 1 || + tree_sizes->ne[3] != 1 || + parent_ids->ne[1] > INT64_MAX / tree_width || + q->ne[1] < parent_ids->ne[1] * tree_width || + (!query_positions && + q->ne[1] != parent_ids->ne[1] * tree_width) || + (int64_t) max_kv_seq_len + tree_width > INT32_MAX) { + return false; + } + + const int64_t scratch_end = + (int64_t) tree_scratch_base + + (block_table->ne[1] - 1) * (int64_t) tree_scratch_stride + + tree_width; + return scratch_end <= k->ne[1]; } // Cached max resident blocks/SM for this instantiation at the given block @@ -764,6 +892,12 @@ static bool try_launch_paged_attn( const ggml_tensor * kv_seq_lens = dst->src[4]; const ggml_tensor * active_slot_ids = dst->src[5]; const ggml_tensor * query_positions = dst->src[6]; + const ggml_tensor * parent_ids = dst->src[7]; + const ggml_tensor * tree_sizes = dst->src[8]; + + const int32_t tree_width = ggml_get_op_params_i32(dst, 3); + const int32_t tree_scratch_base = ggml_get_op_params_i32(dst, 4); + const int32_t tree_scratch_stride = ggml_get_op_params_i32(dst, 5); const int32_t n_head = (int32_t) q->ne[2]; const int32_t n_head_kv = (int32_t) k->ne[2]; @@ -837,15 +971,21 @@ static bool try_launch_paged_attn( if (min_partitions > partition_limit) { min_partitions = partition_limit; } - if (min_partitions > block_table->ne[0]) { - min_partitions = (int32_t) block_table->ne[0]; + const int32_t tree_blocks = + (tree_width + block_size - 1) / block_size; + const int64_t partitionable_blocks = + block_table->ne[0] + (parent_ids ? tree_blocks : 0); + if (min_partitions > partitionable_blocks) { + min_partitions = (int32_t) partitionable_blocks; } - // Size the launch from the live maximum sequence length carried in the - // graph op, not the block-table capacity. Ragged sequences still clamp - // their own active partition count from kv_seq_lens on device. + // Size the launch from the live maximum committed prefix plus the virtual + // tree tail. Ragged/tree rows still clamp their own active partition count + // from device metadata. + const int32_t live_tokens = + max_kv_seq_len + (parent_ids ? tree_width : 0); const int32_t live_blocks = - (max_kv_seq_len + block_size - 1) / block_size; + (live_tokens + block_size - 1) / block_size; int32_t n_partitions = paged_attn_partitions( live_blocks, min_partitions, PAGED_ATTN_MAX_PARTITIONS); @@ -861,7 +1001,7 @@ static bool try_launch_paged_attn( }(); if (forced_partitions >= 1 && forced_partitions <= PAGED_ATTN_MAX_PARTITIONS && - forced_partitions <= block_table->ne[0]) { + forced_partitions <= partitionable_blocks) { min_partitions = forced_partitions; n_partitions = forced_partitions; } @@ -914,6 +1054,8 @@ static bool try_launch_paged_attn( (const char *) kv_seq_lens->data, active_slot_ids ? (const char *) active_slot_ids->data : nullptr, query_positions ? (const char *) query_positions->data : nullptr, + parent_ids ? (const char *) parent_ids->data : nullptr, + tree_sizes ? (const char *) tree_sizes->data : nullptr, (char *) dst->data, partial_acc, partial_meta, @@ -924,6 +1066,9 @@ static bool try_launch_paged_attn( kv_seq_lens->nb[0], active_slot_ids ? active_slot_ids->nb[0] : 0, query_positions ? query_positions->nb[0] : 0, + parent_ids ? parent_ids->nb[0] : 0, + parent_ids ? parent_ids->nb[1] : 0, + tree_sizes ? tree_sizes->nb[0] : 0, dst->nb[1], dst->nb[2], (int32_t) block_table->ne[1], n_head, @@ -932,6 +1077,12 @@ static bool try_launch_paged_attn( (int32_t) block_table->ne[0], block_size, min_partitions, + tree_width, + parent_ids + ? (int32_t)(q->ne[1] - parent_ids->ne[1] * tree_width) + : 0, + tree_scratch_base, + tree_scratch_stride, scale); if (n_partitions > 1) { diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/ssm-conv.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/ssm-conv.cu index e6ce26f72..8c82cb2ac 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/ssm-conv.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/ssm-conv.cu @@ -244,7 +244,113 @@ static void ssm_conv_f32_cuda(const float * src0, const float * src1, const int } } +// dflash: fused conv step (see ggml_ssm_conv_step). One thread per channel +// walks the token loop with the K-1 history in registers, writes silu(conv), +// the optional rollback window and the new history in place. +template +static __global__ void ssm_conv_step_f32(const float * __restrict__ x, const int x_nb1, const int x_nb2, + const float * __restrict__ w, const int w_nb1, + float * state, const int st_nb1, const int st_nb2, + float * __restrict__ y, const int y_nb1, const int y_nb2, + float * ci, const int ci_nb1, const int ci_nb2, + const int C, const int T) { + const int c = blockIdx.x * blockDim.x + threadIdx.x; + const int s = blockIdx.y; + if (c >= C) return; + + const float * xs = (const float *) ((const char *) x + (size_t) s * x_nb2) + c; + float * st = (float *) ((char *) state + (size_t) s * st_nb2 + (size_t) c * st_nb1); + float * ys = (float *) ((char *) y + (size_t) s * y_nb2) + c; + float * cs = ci ? (float *) ((char *) ci + (size_t) s * ci_nb2 + (size_t) c * ci_nb1) : nullptr; + const float * wc = (const float *) ((const char *) w + (size_t) c * w_nb1); + + const int xs_stride = x_nb1 / sizeof(float); + const int ys_stride = y_nb1 / sizeof(float); + + float wt[K]; + float win[K]; // oldest first; win[K-1] is the current input +#pragma unroll + for (int k = 0; k < K; k++) { + wt[k] = wc[k]; + } +#pragma unroll + for (int j = 0; j < K - 1; j++) { + win[j] = st[j]; + if (cs) cs[j] = win[j]; + } + for (int t = 0; t < T; t++) { + const float xt = xs[(size_t) t * xs_stride]; + win[K - 1] = xt; + float acc = 0.0f; +#pragma unroll + for (int k = 0; k < K; k++) { + acc += win[k] * wt[k]; + } + ys[(size_t) t * ys_stride] = ggml_cuda_op_silu_single(acc); + if (cs) cs[K - 1 + t] = xt; +#pragma unroll + for (int j = 0; j < K - 1; j++) { + win[j] = win[j + 1]; + } + } +#pragma unroll + for (int j = 0; j < K - 1; j++) { + st[j] = win[j]; + } +} + +static void ggml_cuda_op_ssm_conv_step(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * x = dst->src[0]; + const ggml_tensor * w = dst->src[1]; + ggml_tensor * st = dst->src[2]; + ggml_tensor * ci = dst->src[3]; + + const int K = (int) w->ne[0]; + const int C = (int) w->ne[1]; + const int T = (int) dst->ne[1]; + const int S = (int) dst->ne[2]; + + GGML_ASSERT(x->type == GGML_TYPE_F32 && w->type == GGML_TYPE_F32 && st->type == GGML_TYPE_F32); + GGML_ASSERT(x->nb[0] == sizeof(float)); + GGML_ASSERT(w->nb[0] == sizeof(float)); + GGML_ASSERT(st->nb[0] == sizeof(float) && st->nb[1] == (size_t) (K - 1) * sizeof(float)); + GGML_ASSERT(dst->nb[0] == sizeof(float)); + if (ci) { + GGML_ASSERT(ci->type == GGML_TYPE_F32 && ci->nb[0] == sizeof(float)); + GGML_ASSERT(ci->ne[0] >= K - 1 + T); + } + + const int threads = 256; + const dim3 blocks((C + threads - 1) / threads, S, 1); + cudaStream_t stream = ctx.stream(); + + auto launch = [&](auto KK) { + constexpr int kK = decltype(KK)::value; + ssm_conv_step_f32<<>>( + (const float *) x->data, (int) x->nb[1], (int) x->nb[2], + (const float *) w->data, (int) w->nb[1], + (float *) st->data, (int) st->nb[1], (int) st->nb[2], + (float *) dst->data, (int) dst->nb[1], (int) dst->nb[2], + ci ? (float *) ci->data : nullptr, ci ? (int) ci->nb[1] : 0, ci ? (int) ci->nb[2] : 0, + C, T); + }; + switch (K) { + case 3: launch(std::integral_constant{}); break; + case 4: launch(std::integral_constant{}); break; + case 5: launch(std::integral_constant{}); break; + case 9: launch(std::integral_constant{}); break; + default: GGML_ABORT("ssm_conv_step only supports kernel sizes 3, 4, 5, 9."); + } +} + void ggml_cuda_op_ssm_conv(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * silu_dst) { + // dflash: fused step mode (silu already applied by the kernel) + if (ggml_get_op_params_i32(dst, 0) == 1) { + GGML_ASSERT(silu_dst == nullptr); + ggml_cuda_op_ssm_conv_step(ctx, dst); + return; + } + const struct ggml_tensor * src0 = dst->src[0]; // conv_x const struct ggml_tensor * src1 = dst->src[1]; // conv1d.weight // dflash27b_ggml: optional src[2] = parent_ids (i32) enables tree mode diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/generate_cu_files.py b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/generate_cu_files.py index f87396f5a..5e3a1f00d 100755 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/generate_cu_files.py +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/generate_cu_files.py @@ -102,10 +102,16 @@ def get_short_name(long_quant_name): "GGML_TYPE_Q2_1_ROCMFP2_MIX", "GGML_TYPE_Q3_0_ROCMFPX", "GGML_TYPE_Q3_1_ROCMFP3_MIX", + # Dense hybrid (Qwen3.5/3.8) verify widths N<=16 on gfx1201: the + # 128-row tile leaves a 5120-row projection with only 40 blocks; + # 64x64/4-warp tiles measured +12-23% on those shapes (mmq_probe). + "GGML_TYPE_IQ4_XS", + "GGML_TYPE_Q4_K", + "GGML_TYPE_Q5_K", + "GGML_TYPE_Q6_K", + "GGML_TYPE_Q8_0", }: - guard = "#define GGML_CUDA_ROCMFPX_MMQ_TILE 1\n" - if type == "GGML_TYPE_Q4_K": - guard = "#define LUCEBOX_RDNA_MMQ_Y 64\n" + guard = "#define GGML_CUDA_MMQ_SMALL_TILE 1\n" f.write(SOURCE_MMQ.format(type=type, guard=guard)) for type in range(1, 17): diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-iq4_xs.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-iq4_xs.cu index 1eb3b7430..5e2a1127a 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-iq4_xs.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-iq4_xs.cu @@ -1,5 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. +#define GGML_CUDA_MMQ_SMALL_TILE 1 #include "../mmq.cuh" DECL_MMQ_CASE(GGML_TYPE_IQ4_XS); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q2_0_rocmfp2.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q2_0_rocmfp2.cu index 8221e1d1e..b00cd9a0c 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q2_0_rocmfp2.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q2_0_rocmfp2.cu @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#define GGML_CUDA_ROCMFPX_MMQ_TILE 1 +#define GGML_CUDA_MMQ_SMALL_TILE 1 #include "../mmq.cuh" DECL_MMQ_CASE(GGML_TYPE_Q2_0_ROCMFP2); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q2_1_rocmfp2_mix.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q2_1_rocmfp2_mix.cu index 647b4572f..f73033e33 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q2_1_rocmfp2_mix.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q2_1_rocmfp2_mix.cu @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#define GGML_CUDA_ROCMFPX_MMQ_TILE 1 +#define GGML_CUDA_MMQ_SMALL_TILE 1 #include "../mmq.cuh" DECL_MMQ_CASE(GGML_TYPE_Q2_1_ROCMFP2_MIX); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q3_0_rocmfpx.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q3_0_rocmfpx.cu index 2380af75c..486782982 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q3_0_rocmfpx.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q3_0_rocmfpx.cu @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#define GGML_CUDA_ROCMFPX_MMQ_TILE 1 +#define GGML_CUDA_MMQ_SMALL_TILE 1 #include "../mmq.cuh" DECL_MMQ_CASE(GGML_TYPE_Q3_0_ROCMFPX); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q3_1_rocmfp3_mix.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q3_1_rocmfp3_mix.cu index 1873e073f..92197f871 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q3_1_rocmfp3_mix.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q3_1_rocmfp3_mix.cu @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#define GGML_CUDA_ROCMFPX_MMQ_TILE 1 +#define GGML_CUDA_MMQ_SMALL_TILE 1 #include "../mmq.cuh" DECL_MMQ_CASE(GGML_TYPE_Q3_1_ROCMFP3_MIX); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_0_rocmfp4_fast.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_0_rocmfp4_fast.cu index 94a2bb0f5..92cb4653d 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_0_rocmfp4_fast.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_0_rocmfp4_fast.cu @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#define GGML_CUDA_ROCMFPX_MMQ_TILE 1 +#define GGML_CUDA_MMQ_SMALL_TILE 1 #include "../mmq.cuh" DECL_MMQ_CASE(GGML_TYPE_Q4_0_ROCMFP4_FAST); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_k.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_k.cu index dcf47b2c2..f9a206d20 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_k.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_k.cu @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#define LUCEBOX_RDNA_MMQ_Y 64 +#define GGML_CUDA_MMQ_SMALL_TILE 1 #include "../mmq.cuh" DECL_MMQ_CASE(GGML_TYPE_Q4_K); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q5_k.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q5_k.cu index a2e90ffd5..7cf43f75e 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q5_k.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q5_k.cu @@ -1,5 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. +#define GGML_CUDA_MMQ_SMALL_TILE 1 #include "../mmq.cuh" DECL_MMQ_CASE(GGML_TYPE_Q5_K); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q6_k.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q6_k.cu index 470938fef..8bc6b7434 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q6_k.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q6_k.cu @@ -1,5 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. +#define GGML_CUDA_MMQ_SMALL_TILE 1 #include "../mmq.cuh" DECL_MMQ_CASE(GGML_TYPE_Q6_K); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q8_0.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q8_0.cu index 974477bbb..fb8fcf911 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q8_0.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q8_0.cu @@ -1,5 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. +#define GGML_CUDA_MMQ_SMALL_TILE 1 #include "../mmq.cuh" DECL_MMQ_CASE(GGML_TYPE_Q8_0); diff --git a/server/deps/llama.cpp/ggml/src/ggml.c b/server/deps/llama.cpp/ggml/src/ggml.c index 86e4520d6..d6290babd 100644 --- a/server/deps/llama.cpp/ggml/src/ggml.c +++ b/server/deps/llama.cpp/ggml/src/ggml.c @@ -5708,7 +5708,12 @@ struct ggml_tensor * ggml_paged_attn_ext( struct ggml_tensor * query_positions, float scale, int block_size, - int max_kv_seq_len) { + int max_kv_seq_len, + struct ggml_tensor * parent_ids, + struct ggml_tensor * tree_sizes, + int tree_width, + int tree_scratch_base, + int tree_scratch_stride) { GGML_ASSERT(q->type == GGML_TYPE_F32); GGML_ASSERT(k->type == GGML_TYPE_F16 || k->type == GGML_TYPE_Q4_0 || k->type == GGML_TYPE_Q8_0); GGML_ASSERT(v->type == GGML_TYPE_F16 || v->type == GGML_TYPE_Q4_0 || v->type == GGML_TYPE_Q8_0); @@ -5720,6 +5725,16 @@ struct ggml_tensor * ggml_paged_attn_ext( GGML_ASSERT(query_positions == NULL || active_slot_ids != NULL); GGML_ASSERT(query_positions == NULL || query_positions->type == GGML_TYPE_I32); + const bool tree_mode = parent_ids != NULL || tree_sizes != NULL; + GGML_ASSERT((parent_ids == NULL) == (tree_sizes == NULL)); + GGML_ASSERT(!tree_mode || active_slot_ids != NULL); + // Mixed direct-commit batches use causal positions for a compact AR + // prefix and -1 for the fixed-width tree tail. Pure trees keep this null. + GGML_ASSERT(!tree_mode || query_positions == NULL || + query_positions->ne[0] == q->ne[1]); + GGML_ASSERT(!tree_mode || parent_ids->type == GGML_TYPE_I32); + GGML_ASSERT(!tree_mode || tree_sizes->type == GGML_TYPE_I32); + GGML_ASSERT(q->ne[0] == k->ne[0] && q->ne[0] == v->ne[0]); GGML_ASSERT(k->ne[1] == v->ne[1]); GGML_ASSERT(k->ne[2] > 0); @@ -5749,13 +5764,51 @@ struct ggml_tensor * ggml_paged_attn_ext( GGML_ASSERT(block_size > 0); GGML_ASSERT(k->ne[1] % block_size == 0); GGML_ASSERT(max_kv_seq_len > 0); - GGML_ASSERT(max_kv_seq_len <= k->ne[1]); + // This is a padded logical launch bound, not a physical-cache extent. + // Each row clamps its actual sequence length to the block-table capacity + // and validates every resolved physical block before dereferencing K/V. + GGML_ASSERT((int64_t) max_kv_seq_len + tree_width <= INT32_MAX); + + if (tree_mode) { + GGML_ASSERT(tree_width > 0); + GGML_ASSERT(tree_scratch_base > 0); + GGML_ASSERT(tree_scratch_base % block_size == 0); + GGML_ASSERT(tree_scratch_stride >= tree_width); + GGML_ASSERT(ggml_is_contiguous(parent_ids)); + GGML_ASSERT(ggml_is_contiguous(tree_sizes)); + GGML_ASSERT(parent_ids->ne[0] == tree_width); + GGML_ASSERT(parent_ids->ne[1] == tree_sizes->ne[0]); + GGML_ASSERT(parent_ids->ne[2] == 1 && parent_ids->ne[3] == 1); + GGML_ASSERT(tree_sizes->ne[1] == 1 && tree_sizes->ne[2] == 1 && tree_sizes->ne[3] == 1); + GGML_ASSERT(parent_ids->ne[1] > 0); + GGML_ASSERT(parent_ids->ne[1] <= INT64_MAX / tree_width); + const int64_t tree_rows = parent_ids->ne[1] * tree_width; + GGML_ASSERT(q->ne[1] >= tree_rows); + GGML_ASSERT(query_positions || q->ne[1] == tree_rows); + + // Every physical sequence slot owns one non-overlapping scratch slab. + // Bound the largest address with int64 arithmetic before the GPU sees + // the int32 op parameters. + const int64_t scratch_end = + (int64_t) tree_scratch_base + + (block_table->ne[1] - 1) * (int64_t) tree_scratch_stride + + tree_width; + GGML_ASSERT(scratch_end <= k->ne[1]); + GGML_ASSERT((int64_t) max_kv_seq_len + tree_width <= INT32_MAX); + } else { + GGML_ASSERT(tree_width == 0); + GGML_ASSERT(tree_scratch_base == 0); + GGML_ASSERT(tree_scratch_stride == 0); + } struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, GGML_MAX_DIMS, q->ne); ggml_set_op_params_f32(result, 0, scale); ggml_set_op_params_i32(result, 1, block_size); ggml_set_op_params_i32(result, 2, max_kv_seq_len); + ggml_set_op_params_i32(result, 3, tree_width); + ggml_set_op_params_i32(result, 4, tree_scratch_base); + ggml_set_op_params_i32(result, 5, tree_scratch_stride); result->op = GGML_OP_PAGED_ATTN; result->src[0] = q; @@ -5765,6 +5818,8 @@ struct ggml_tensor * ggml_paged_attn_ext( result->src[4] = kv_seq_lens; result->src[5] = active_slot_ids; result->src[6] = query_positions; + result->src[7] = parent_ids; + result->src[8] = tree_sizes; return result; } @@ -5891,6 +5946,53 @@ struct ggml_tensor * ggml_ssm_conv_tree( return result; } +// dflash: fused conv step. Same op id as ggml_ssm_conv; op_params[0] = 1 +// marks step mode, srcs are (x, c, conv_state, conv_input_out). +struct ggml_tensor * ggml_ssm_conv_step( + struct ggml_context * ctx, + struct ggml_tensor * x, + struct ggml_tensor * c, + struct ggml_tensor * conv_state, + struct ggml_tensor * conv_input_out) { + GGML_ASSERT(x->type == GGML_TYPE_F32); + GGML_ASSERT(c->type == GGML_TYPE_F32); + GGML_ASSERT(conv_state->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_matrix(c)); + GGML_ASSERT(ggml_is_contiguous(c)); + GGML_ASSERT(x->nb[0] == sizeof(float)); + GGML_ASSERT(x->ne[3] == 1); + + const int64_t d_conv = c->ne[0]; + const int64_t d_inner = c->ne[1]; + const int64_t n_t = x->ne[1]; + const int64_t n_s = x->ne[2]; + + GGML_ASSERT(x->ne[0] == d_inner); + GGML_ASSERT(conv_state->ne[0] == d_conv - 1); + GGML_ASSERT(conv_state->ne[1] == d_inner); + GGML_ASSERT(conv_state->ne[2] == n_s); + GGML_ASSERT(conv_state->nb[0] == sizeof(float)); + GGML_ASSERT(conv_state->nb[1] == (size_t)(d_conv - 1) * sizeof(float)); + if (conv_input_out) { + GGML_ASSERT(conv_input_out->type == GGML_TYPE_F32); + GGML_ASSERT(conv_input_out->ne[0] >= d_conv - 1 + n_t); + GGML_ASSERT(conv_input_out->ne[1] == d_inner); + GGML_ASSERT(conv_input_out->ne[2] == n_s); + GGML_ASSERT(conv_input_out->nb[0] == sizeof(float)); + } + + struct ggml_tensor * result = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, d_inner, n_t, n_s); + ggml_set_op_params_i32(result, 0, 1); // step mode + + result->op = GGML_OP_SSM_CONV; + result->src[0] = x; + result->src[1] = c; + result->src[2] = conv_state; + result->src[3] = conv_input_out; + + return result; +} + // ggml_ssm_scan struct ggml_tensor * ggml_ssm_scan( @@ -6704,6 +6806,54 @@ void ggml_gated_delta_net_set_skip_intermediate( tensor->nb[3] = tensor->nb[2]*tensor->ne[2]; } +void ggml_gated_delta_net_set_transition_journal( + struct ggml_tensor * tensor, + struct ggml_tensor * journal) { + GGML_ASSERT(tensor != NULL && journal != NULL); + GGML_ASSERT(tensor->op == GGML_OP_GATED_DELTA_NET); + GGML_ASSERT(journal->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(journal)); + + const struct ggml_tensor * v = tensor->src[2]; + const struct ggml_tensor * g = tensor->src[3]; + GGML_ASSERT(v != NULL && g != NULL); + const int64_t S_v = v->ne[0]; + const int64_t H = v->ne[1]; + const int64_t n_tokens = v->ne[2]; + const int64_t n_seqs = v->ne[3]; + const bool kda = g->ne[0] == S_v; + const int64_t journal_width = kda ? 3*S_v : 2*S_v + 1; + GGML_ASSERT(journal->ne[0] == journal_width && + journal->ne[1] == H && + journal->ne[2] == n_tokens && + journal->ne[3] == n_seqs); + + tensor->src[11] = journal; +} + +// dflash: raw-gate mode (see ggml.h). src[8] is reserved for the optional +// active-slot map; dt_bias -> src[9], A -> src[10], +// op_params[2] = 1. +void ggml_gated_delta_net_set_raw_gates( + struct ggml_tensor * tensor, + struct ggml_tensor * dt_bias, + struct ggml_tensor * A) { + GGML_ASSERT(tensor != NULL); + GGML_ASSERT(tensor->op == GGML_OP_GATED_DELTA_NET); + GGML_ASSERT(dt_bias != NULL && A != NULL); + GGML_ASSERT(dt_bias->type == GGML_TYPE_F32 && A->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(dt_bias) && ggml_is_contiguous(A)); + const struct ggml_tensor * v = tensor->src[2]; + GGML_ASSERT(ggml_nelements(dt_bias) == v->ne[1]); + GGML_ASSERT(ggml_nelements(A) == v->ne[1]); + // scalar gate only (no KDA), no tree mode + GGML_ASSERT(tensor->src[3]->ne[0] == 1); + GGML_ASSERT(tensor->src[6] == NULL); + tensor->src[9] = dt_bias; + tensor->src[10] = A; + ggml_set_op_params_i32(tensor, 2, 1); +} + // dflash: tree-mode variant. Same op, with parent_ids plumbed into // src[6] so the CUDA kernel can branch-reload state at DFS transitions. struct ggml_tensor * ggml_gated_delta_net_tree( diff --git a/server/scripts/convert_dflash_to_gguf.py b/server/scripts/convert_dflash_to_gguf.py index b904d7ea5..74f71bce2 100644 --- a/server/scripts/convert_dflash_to_gguf.py +++ b/server/scripts/convert_dflash_to_gguf.py @@ -110,6 +110,37 @@ def pick(*keys): or c.get("aux_hidden_state_layer_ids")) if _tli: a["capture_layer_ids"] = [int(x) for x in _tli] + # Newer HF configs (transformers >= 5.x, e.g. the Qwen3.8 DSpark + # drafter) nest rope_theta / YaRN under rope_parameters and + # mask_token_id under dflash_config instead of top-level. + rp = c.get("rope_parameters") or c.get("rope_scaling") or {} + if isinstance(rp, dict): + if rp.get("rope_theta") is not None: + a["rope_theta"] = float(rp["rope_theta"]) + if str(rp.get("rope_type", "")).lower() == "yarn": + a["yarn_factor"] = float(rp.get("factor", 0.0)) + a["yarn_orig_ctx"] = int(rp.get("original_max_position_embeddings", 0)) + a["yarn_beta_fast"] = float(rp.get("beta_fast", 32.0)) + a["yarn_beta_slow"] = float(rp.get("beta_slow", 1.0)) + if dfc.get("mask_token_id") is not None: + a["mask_token_id"] = int(dfc["mask_token_id"]) + if dfc.get("block_size") is not None: + a["block_size"] = int(dfc["block_size"]) + # DFlash 2 (z-lab/inco): grouped dynamic convs + candidate selector. + if dfc.get("conv_kernel_size") is not None: + a["conv_kernel_size"] = int(dfc["conv_kernel_size"]) + a["conv_group_size"] = int(dfc.get("conv_group_size", 16)) + if dfc.get("selector_rank") is not None: + a["selector_rank"] = int(dfc["selector_rank"]) + a["selector_top_k"] = int(dfc.get("selector_top_k", 16)) + # Per-layer sliding-window / causal attention (Qwen3.6-style drafters + # and DFlash 2). HF: layer_types + sliding_window; a top-level + # is_causal=false (DFlash 2) makes every layer bidirectional, which is + # our default (no SWA pattern emitted). + lt = c.get("layer_types") + if lt and c.get("sliding_window") and c.get("is_causal", None) is not False: + a["swa_window"] = int(c["sliding_window"]) + a["swa_pattern"] = [str(x) == "sliding_attention" for x in lt] print(f"[info] read arch from {cfg_path}") else: print(f"[warn] no config.json next to safetensors; using 27B defaults") @@ -182,8 +213,17 @@ def map_name(name: str) -> str | None: "mlp.gate_proj.weight": f"blk.{i}.ffn_gate.weight", "mlp.up_proj.weight": f"blk.{i}.ffn_up.weight", "mlp.down_proj.weight": f"blk.{i}.ffn_down.weight", + # DFlash 2 grouped dynamic convs + "attention_conv.base_kernel": f"blk.{i}.attn_conv.base", + "attention_conv.kernel_projection.weight": f"blk.{i}.attn_conv.proj.weight", + "mlp_conv.base_kernel": f"blk.{i}.ffn_conv.base", + "mlp_conv.kernel_projection.weight": f"blk.{i}.ffn_conv.proj.weight", } return layer_map.get(rest) + # DFlash 2 candidate selector + if name == "candidate_selector.hidden_projection.weight": return "dflash.selector.hproj.weight" + if name == "candidate_selector.predecessor_codebook": return "dflash.selector.pred_cb" + if name == "candidate_selector.successor_codebook": return "dflash.selector.succ_cb" return None @@ -248,18 +288,30 @@ def bytes_to_np(raw: bytes, dtype: str, shape: list[int]) -> np.ndarray: } +# Alias sets per head tensor: SpecForge sidecar names, DS4 MTP-shard names, +# and single-file releases (e.g. RadixArk Qwen3.8-27B-DSpark) that carry the +# heads inline in the main model.safetensors. +DSPARK_MARKOV_W1_KEYS = ("dspark_markov_head.markov_w1.weight", + "mtp.2.markov_head.markov_w1.weight", + "markov_head.markov_w1.weight") +DSPARK_MARKOV_W2_KEYS = ("dspark_markov_head.markov_w2.weight", + "mtp.2.markov_head.markov_w2.weight", + "markov_head.markov_w2.weight") +DSPARK_CONF_W_KEYS = ("dspark_confidence_head.weight", + "mtp.2.confidence_head.proj.weight", + "confidence_head.proj.weight") +DSPARK_CONF_B_KEYS = ("dspark_confidence_head.bias", + "mtp.2.confidence_head.proj.bias", + "confidence_head.proj.bias") + DSPARK_TENSOR_MAP = { - ("dspark_markov_head.markov_w1.weight", - "mtp.2.markov_head.markov_w1.weight"): ("dflash.dspark.markov.w1", gguf.GGMLQuantizationType.F16), - ("dspark_markov_head.markov_w2.weight", - "mtp.2.markov_head.markov_w2.weight"): ("dflash.dspark.markov.w2", gguf.GGMLQuantizationType.F16), + DSPARK_MARKOV_W1_KEYS: ("dflash.dspark.markov.w1", gguf.GGMLQuantizationType.F16), + DSPARK_MARKOV_W2_KEYS: ("dflash.dspark.markov.w2", gguf.GGMLQuantizationType.F16), } DSPARK_CONFIDENCE_TENSOR_MAP = { - ("dspark_confidence_head.weight", - "mtp.2.confidence_head.proj.weight"): ("dflash.dspark.confidence.weight", gguf.GGMLQuantizationType.F16), - ("dspark_confidence_head.bias", - "mtp.2.confidence_head.proj.bias"): ("dflash.dspark.confidence.bias", gguf.GGMLQuantizationType.F32), + DSPARK_CONF_W_KEYS: ("dflash.dspark.confidence.weight", gguf.GGMLQuantizationType.F16), + DSPARK_CONF_B_KEYS: ("dflash.dspark.confidence.bias", gguf.GGMLQuantizationType.F32), } @@ -372,8 +424,8 @@ def add_dspark_aux_heads(writer, arch: str, aux_path: Path | None): return print(f"[info] reading DSpark aux heads from {aux_path}") - w1 = resolved[("dspark_markov_head.markov_w1.weight", "mtp.2.markov_head.markov_w1.weight")][1] - w2 = resolved[("dspark_markov_head.markov_w2.weight", "mtp.2.markov_head.markov_w2.weight")][1] + w1 = resolved[DSPARK_MARKOV_W1_KEYS][1] + w2 = resolved[DSPARK_MARKOV_W2_KEYS][1] vocab = int(w1.shape[0]) rank = int(w1.shape[1]) if tuple(w2.shape) != (vocab, rank): @@ -397,8 +449,8 @@ def add_dspark_aux_heads(writer, arch: str, aux_path: Path | None): conf_missing.append(names) continue conf_resolved[names] = (found_name, tensor, spec) - weight_names = ("dspark_confidence_head.weight", "mtp.2.confidence_head.proj.weight") - bias_names = ("dspark_confidence_head.bias", "mtp.2.confidence_head.proj.bias") + weight_names = DSPARK_CONF_W_KEYS + bias_names = DSPARK_CONF_B_KEYS if weight_names not in conf_resolved: if conf_missing: print("[warn] incomplete DSpark confidence head; Markov head will still load") @@ -440,6 +492,9 @@ def main(): help="optional Domino/DSpark aux-head .pt or DS4 MTP .safetensors file; defaults to dflash_aux_heads.pt next to the safetensors") ap.add_argument("--no-aux-heads", action="store_true", help="do not auto-embed Domino/DSpark aux-head tensors") + ap.add_argument("--no-yarn", action="store_true", + help="omit YaRN scaling metadata while retaining rope_theta; " + "this matches the PR #625 short-context Qwen3.8 DSpark artifact") args = ap.parse_args() if not args.safetensors.exists(): @@ -452,6 +507,9 @@ def main(): print(f"[info] {n_entries} tensor entries") a = load_arch(args.safetensors, header) + if args.no_yarn: + for key in ("yarn_factor", "yarn_orig_ctx", "yarn_beta_fast", "yarn_beta_slow"): + a.pop(key, None) writer = gguf.GGUFWriter(args.out_gguf, ARCH) @@ -470,6 +528,12 @@ def main(): writer.add_uint32(f"{ARCH}.vocab_size", a["vocab"]) writer.add_float32(f"{ARCH}.attention.layer_norm_rms_epsilon", a["rms_eps"]) writer.add_float32(f"{ARCH}.rope.freq_base", a["rope_theta"]) + if a.get("yarn_factor", 0.0) > 1.0: + writer.add_string(f"{ARCH}.rope.scaling.type", "yarn") + writer.add_float32(f"{ARCH}.rope.scaling.factor", a["yarn_factor"]) + writer.add_uint32(f"{ARCH}.rope.scaling.original_context_length", a["yarn_orig_ctx"]) + writer.add_float32(f"{ARCH}.rope.scaling.beta_fast", a["yarn_beta_fast"]) + writer.add_float32(f"{ARCH}.rope.scaling.beta_slow", a["yarn_beta_slow"]) # DFlash-specific hyperparameters writer.add_uint32(f"{ARCH}.dflash.n_target_layers", a["n_target_layers"]) @@ -484,6 +548,15 @@ def main(): elif _cap_ids: print(f"[warn] capture_layer_ids len {len(_cap_ids)} != n_target_layers " f"{a['n_target_layers']}; not embedding ids", file=sys.stderr) + if a.get("swa_pattern"): + writer.add_uint32(f"{ARCH}.attention.sliding_window", a["swa_window"]) + writer.add_array(f"{ARCH}.attention.sliding_window_pattern", [bool(x) for x in a["swa_pattern"]]) + if a.get("conv_kernel_size"): + writer.add_uint32(f"{ARCH}.dflash.dflash2.conv_kernel_size", a["conv_kernel_size"]) + writer.add_uint32(f"{ARCH}.dflash.dflash2.conv_group_size", a["conv_group_size"]) + if a.get("selector_rank"): + writer.add_uint32(f"{ARCH}.dflash.dflash2.selector_rank", a["selector_rank"]) + writer.add_uint32(f"{ARCH}.dflash.dflash2.selector_top_k", a["selector_top_k"]) # Walk + add tensors. Sort: dflash.* singletons first, then output_*, # then per-layer in numeric order — keeps the on-disk layout stable. @@ -522,7 +595,8 @@ def sort_key(t): is_norm = ( gguf_name.endswith("_norm.weight") or gguf_name == "output_norm.weight" or - gguf_name == "dflash.hidden_norm.weight" + gguf_name == "dflash.hidden_norm.weight" or + gguf_name.endswith("_conv.base") # DFlash 2 conv base kernels [2, K, hidden] ) if is_norm: arr = arr.astype("&2; exit 2; } +[[ -r "$DRAFT_SOURCE" ]] || { echo "unreadable DRAFT_SOURCE: $DRAFT_SOURCE" >&2; exit 2; } +[[ -x "$LLAMA_QUANTIZE" ]] || { echo "LLAMA_QUANTIZE is not executable: $LLAMA_QUANTIZE" >&2; exit 2; } +command -v "$PYTHON" >/dev/null || { echo "PYTHON is unavailable: $PYTHON" >&2; exit 2; } +[[ "$DRAFT_SCHEME" == f16 || "$DRAFT_SCHEME" == q8_0 || "$DRAFT_SCHEME" == q4-mix ]] || { + echo "DRAFT_SCHEME must be f16, q8_0, or q4-mix" >&2 + exit 2 +} + +mkdir -p "$OUT_DIR" +work_dir="$(mktemp -d "$OUT_DIR/.prepare.XXXXXX")" +cleanup() { rm -rf -- "$work_dir"; } +trap cleanup EXIT + +draft_f16="$work_dir/Qwen3.8-27B-DSpark-RadixArk-no-yarn-f16.gguf" +draft_final="$work_dir/Qwen3.8-27B-DSpark-RadixArk-no-yarn-$DRAFT_SCHEME.gguf" +target_final="$work_dir/Qwen3.8-27B-PR625-IQ4_XS.gguf" + +"$PYTHON" "$SCRIPT_DIR/convert_dflash_to_gguf.py" \ + "$DRAFT_SOURCE" "$draft_f16" --no-yarn +if [[ "$DRAFT_SCHEME" != f16 ]]; then + "$PYTHON" "$SCRIPT_DIR/quantize_dflash_draft.py" \ + "$draft_f16" "$draft_final" --scheme "$DRAFT_SCHEME" +fi + +# PR #625 target: pure IQ4_XS body, Q5_K output, Q6_K attn_v/ssm_out. +# The validator below deliberately catches quantizers that let --pure suppress +# explicit --tensor-type overrides. +"$LLAMA_QUANTIZE" \ + --allow-requantize --pure \ + --output-tensor-type q5_k \ + --tensor-type ssm_out=q6_k \ + --tensor-type attn_v=q6_k \ + "$TARGET_SOURCE" "$target_final" iq4_xs + +"$PYTHON" "$SCRIPT_DIR/validate_qwen38_pr625_models.py" \ + --target "$target_final" --draft "$draft_final" --draft-scheme "$DRAFT_SCHEME" + +mv -- "$target_final" "$OUT_DIR/Qwen3.8-27B-PR625-IQ4_XS.gguf" +mv -- "$draft_final" "$OUT_DIR/Qwen3.8-27B-DSpark-RadixArk-no-yarn-$DRAFT_SCHEME.gguf" + +echo "PR #625 model pair ready in $OUT_DIR" diff --git a/server/scripts/quantize_draft_q8.py b/server/scripts/quantize_draft_q8.py index b8f72fdbc..df1793065 100644 --- a/server/scripts/quantize_draft_q8.py +++ b/server/scripts/quantize_draft_q8.py @@ -19,35 +19,24 @@ """ import argparse -import json -import struct import sys from pathlib import Path import numpy as np -sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "deps" / "llama.cpp" / "gguf-py")) +SCRIPT_DIR = Path(__file__).resolve().parent +sys.path.insert(0, str(SCRIPT_DIR.parent / "deps" / "llama.cpp" / "gguf-py")) +sys.path.insert(0, str(SCRIPT_DIR)) import gguf +import convert_dflash_to_gguf as canonical # ────────────────────────────────────────────────────────────────────── -# DFlash 27B draft architecture constants (must match dflash27b.h) +# Legacy fallback constants. The canonical converter resolves these from +# config.json + tensor shapes when the source ships model metadata. # ────────────────────────────────────────────────────────────────────── -ARCH = "qwen35-dflash-draft" -HIDDEN = 5120 -N_LAYER = 5 -N_HEAD = 32 -N_HEAD_KV = 8 -HEAD_DIM = 128 -INTERMEDIATE = 17408 -VOCAB = 248320 -N_TARGET_LAYERS = 5 -ROPE_THETA = 1_000_000.0 -RMS_EPS = 1e-6 -MASK_TOKEN_ID = 248070 -BLOCK_SIZE = 16 -CTX_LEN = 32768 +ARCH = canonical.ARCH Q8_0_BLOCK_SIZE = 32 # elements per Q8_0 block @@ -67,60 +56,28 @@ def add_qwen36_swa_metadata(writer, enabled: bool) -> None: # ────────────────────────────────────────────────────────────────────── -# Tensor name mapping — DFlash safetensors -> llama.cpp GGUF -# (Identical to convert_dflash_to_gguf.py) +# Tensor name mapping — share the DFlash2-aware canonical converter. # ────────────────────────────────────────────────────────────────────── def map_name(name: str) -> str | None: - if name == "fc.weight": return "dflash.fc.weight" - if name == "hidden_norm.weight": return "dflash.hidden_norm.weight" - if name == "norm.weight": return "output_norm.weight" - if name.startswith("layers."): - parts = name.split(".", 2) - if len(parts) < 3: return None - i = int(parts[1]) - rest = parts[2] - layer_map = { - "input_layernorm.weight": f"blk.{i}.attn_norm.weight", - "post_attention_layernorm.weight": f"blk.{i}.ffn_norm.weight", - "self_attn.q_proj.weight": f"blk.{i}.attn_q.weight", - "self_attn.k_proj.weight": f"blk.{i}.attn_k.weight", - "self_attn.v_proj.weight": f"blk.{i}.attn_v.weight", - "self_attn.o_proj.weight": f"blk.{i}.attn_output.weight", - "self_attn.q_norm.weight": f"blk.{i}.attn_q_norm.weight", - "self_attn.k_norm.weight": f"blk.{i}.attn_k_norm.weight", - "mlp.gate_proj.weight": f"blk.{i}.ffn_gate.weight", - "mlp.up_proj.weight": f"blk.{i}.ffn_up.weight", - "mlp.down_proj.weight": f"blk.{i}.ffn_down.weight", - } - return layer_map.get(rest) - return None + return canonical.map_name(name) def is_norm_tensor(gguf_name: str) -> bool: return ( gguf_name.endswith("_norm.weight") or gguf_name == "output_norm.weight" or - gguf_name == "dflash.hidden_norm.weight" + gguf_name == "dflash.hidden_norm.weight" or + gguf_name.endswith("_conv.base") ) # ────────────────────────────────────────────────────────────────────── -# safetensors reader +# safetensors reader aliases — kept public for existing tests/importers. # ────────────────────────────────────────────────────────────────────── -def load_safetensors_header(path: Path): - with open(path, "rb") as f: - header_size = struct.unpack(" bytes: - start, end = info["data_offsets"] - with open(path, "rb") as f: - f.seek(8 + header_size + start) - return f.read(end - start) +load_safetensors_header = canonical.load_safetensors_header +read_tensor_bytes = canonical.read_tensor_bytes def bf16_bytes_to_f32(raw: bytes, shape: list[int]) -> np.ndarray: @@ -128,6 +85,44 @@ def bf16_bytes_to_f32(raw: bytes, shape: list[int]) -> np.ndarray: u32 = (u16.astype(np.uint32) << 16) return u32.view(" None: + """Write the same resolved architecture profile as the F16 converter.""" + writer.add_string("general.name", f"DFlash-Draft-{a['hidden']}h-{a['n_layer']}L-Q8_0") + writer.add_quantization_version(gguf.GGML_QUANT_VERSION) + writer.add_uint32(f"{ARCH}.context_length", a["ctx_len"]) + writer.add_uint32(f"{ARCH}.embedding_length", a["hidden"]) + writer.add_uint32(f"{ARCH}.block_count", a["n_layer"]) + writer.add_uint32(f"{ARCH}.feed_forward_length", a["intermediate"]) + writer.add_uint32(f"{ARCH}.attention.head_count", a["n_head"]) + writer.add_uint32(f"{ARCH}.attention.head_count_kv", a["n_head_kv"]) + writer.add_uint32(f"{ARCH}.attention.key_length", a["head_dim"]) + writer.add_uint32(f"{ARCH}.attention.value_length", a["head_dim"]) + writer.add_uint32(f"{ARCH}.vocab_size", a["vocab"]) + writer.add_float32(f"{ARCH}.attention.layer_norm_rms_epsilon", a["rms_eps"]) + writer.add_float32(f"{ARCH}.rope.freq_base", a["rope_theta"]) + + if qwen36_swa: + add_qwen36_swa_metadata(writer, True) + elif a.get("swa_pattern"): + writer.add_uint32(f"{ARCH}.attention.sliding_window", a["swa_window"]) + writer.add_array(f"{ARCH}.attention.sliding_window_pattern", [bool(x) for x in a["swa_pattern"]]) + + writer.add_uint32(f"{ARCH}.dflash.n_target_layers", a["n_target_layers"]) + writer.add_uint32(f"{ARCH}.dflash.block_size", a["block_size"]) + writer.add_uint32(f"{ARCH}.dflash.mask_token_id", a["mask_token_id"]) + capture_ids = a.get("capture_layer_ids") + if capture_ids and len(capture_ids) == a["n_target_layers"]: + writer.add_array(f"{ARCH}.dflash.target_layer_ids", [int(x) for x in capture_ids]) + elif capture_ids: + print(f"[warn] capture_layer_ids len {len(capture_ids)} != n_target_layers {a['n_target_layers']}; not embedding ids", file=sys.stderr) + + if a.get("conv_kernel_size"): + writer.add_uint32(f"{ARCH}.dflash.dflash2.conv_kernel_size", a["conv_kernel_size"]) + writer.add_uint32(f"{ARCH}.dflash.dflash2.conv_group_size", a["conv_group_size"]) + if a.get("selector_rank"): + writer.add_uint32(f"{ARCH}.dflash.dflash2.selector_rank", a["selector_rank"]) + writer.add_uint32(f"{ARCH}.dflash.dflash2.selector_top_k", a["selector_top_k"]) + # ────────────────────────────────────────────────────────────────────── # Main @@ -160,32 +155,14 @@ def main(): header_size, header = load_safetensors_header(args.safetensors) n_entries = sum(1 for k in header if k != "__metadata__") print(f"[info] {n_entries} tensor entries") + arch = canonical.load_arch(args.safetensors, header) writer = gguf.GGUFWriter(args.out_gguf, ARCH) - # Architecture metadata (identical to convert_dflash_to_gguf.py) - writer.add_string("general.name", "Qwen3.5-27B-DFlash-Draft-Q8_0") - writer.add_quantization_version(gguf.GGML_QUANT_VERSION) - writer.add_uint32(f"{ARCH}.context_length", CTX_LEN) - writer.add_uint32(f"{ARCH}.embedding_length", HIDDEN) - writer.add_uint32(f"{ARCH}.block_count", N_LAYER) - writer.add_uint32(f"{ARCH}.feed_forward_length", INTERMEDIATE) - writer.add_uint32(f"{ARCH}.attention.head_count", N_HEAD) - writer.add_uint32(f"{ARCH}.attention.head_count_kv", N_HEAD_KV) - writer.add_uint32(f"{ARCH}.attention.key_length", HEAD_DIM) - writer.add_uint32(f"{ARCH}.attention.value_length", HEAD_DIM) - writer.add_uint32(f"{ARCH}.vocab_size", VOCAB) - writer.add_float32(f"{ARCH}.attention.layer_norm_rms_epsilon", RMS_EPS) - writer.add_float32(f"{ARCH}.rope.freq_base", ROPE_THETA) - add_qwen36_swa_metadata(writer, args.qwen36_swa) + add_arch_metadata(writer, arch, args.qwen36_swa) if args.qwen36_swa: print("[info] Qwen3.6 draft SWA: layers 0-3 window=2048; layer 4 full attention") - # DFlash-specific hyperparameters - writer.add_uint32(f"{ARCH}.dflash.n_target_layers", N_TARGET_LAYERS) - writer.add_uint32(f"{ARCH}.dflash.block_size", BLOCK_SIZE) - writer.add_uint32(f"{ARCH}.dflash.mask_token_id", MASK_TOKEN_ID) - # Collect and sort tensors (same order as convert_dflash_to_gguf.py) pending = [] for st_name, info in header.items(): diff --git a/server/scripts/validate_qwen38_pr625_models.py b/server/scripts/validate_qwen38_pr625_models.py new file mode 100755 index 000000000..d3215f48d --- /dev/null +++ b/server/scripts/validate_qwen38_pr625_models.py @@ -0,0 +1,126 @@ +#!/usr/bin/env python3 +"""Fail closed unless target and drafter match the PR #625 Qwen3.8 recipe.""" + +import argparse +import sys +from collections import Counter +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "deps" / "llama.cpp" / "gguf-py")) + +from gguf import GGUFReader # noqa: E402 + + +DRAFT_ARCH = "qwen35-dflash-draft" + + +def require(condition: bool, message: str, errors: list[str]) -> None: + if not condition: + errors.append(message) + + +def field(reader: GGUFReader, name: str): + value = reader.fields.get(name) + return None if value is None else value.contents() + + +def validate_target(path: Path, errors: list[str]) -> None: + reader = GGUFReader(path) + tensors = {tensor.name: tensor for tensor in reader.tensors} + counts = Counter(tensor.tensor_type.name for tensor in reader.tensors) + + require(field(reader, "general.architecture") == "qwen35", + "target architecture must be qwen35", errors) + require(counts == Counter({"IQ4_XS": 440, "F32": 360, "Q6_K": 65, "Q5_K": 1}), + f"target tensor-type counts differ from PR #625: {dict(counts)}", errors) + require(tensors.get("output.weight") is not None and + tensors["output.weight"].tensor_type.name == "Q5_K", + "target output.weight must be Q5_K", errors) + + q6_names = { + tensor.name for tensor in reader.tensors if tensor.tensor_type.name == "Q6_K" + } + invalid_q6 = sorted( + name for name in q6_names + if not (name.endswith("ssm_out.weight") or name.endswith("attn_v.weight")) + ) + require(not invalid_q6, + f"target has unexpected Q6_K tensors: {invalid_q6}", errors) + require(all( + tensor.tensor_type.name == "Q6_K" + for name, tensor in tensors.items() + if name.endswith("ssm_out.weight") or name.endswith("attn_v.weight") + ), "every target ssm_out/attn_v tensor must be Q6_K", errors) + + +def validate_draft(path: Path, scheme: str, errors: list[str]) -> None: + reader = GGUFReader(path) + prefix = DRAFT_ARCH + "." + counts = Counter(tensor.tensor_type.name for tensor in reader.tensors) + + require(field(reader, "general.architecture") == DRAFT_ARCH, + f"drafter architecture must be {DRAFT_ARCH}", errors) + expected_counts = { + "f16": Counter({"F16": 39, "F32": 23}), + "q8_0": Counter({"Q8_0": 39, "F32": 23}), + "q4-mix": Counter({"Q4_0": 35, "F32": 23, "Q8_0": 4}), + }[scheme] + require(counts == expected_counts, + f"drafter tensor-type counts differ from {scheme}: {dict(counts)}", errors) + require(field(reader, prefix + "rope.freq_base") == 10_000_000.0, + "drafter rope.freq_base must be 10000000", errors) + require(not any("rope.scaling" in name for name in reader.fields), + "drafter must not contain YaRN/rope.scaling metadata", errors) + + expected = { + "dflash.n_target_layers": 5, + "dflash.block_size": 7, + "dflash.mask_token_id": 248077, + "dflash.target_layer_ids": [4, 16, 28, 40, 52], + "dflash.dspark.enabled": 1, + "dflash.dspark.markov_rank": 256, + "dflash.dspark.vocab_size": 248320, + "dflash.dspark.confidence_dim": 5376, + "dflash.dspark.confidence.enabled": 1, + } + for key, wanted in expected.items(): + actual = field(reader, prefix + key) + require(actual == wanted, + f"drafter {key} must be {wanted!r}, got {actual!r}", errors) + + if scheme == "q4-mix": + invalid_q8 = sorted( + tensor.name for tensor in reader.tensors + if tensor.tensor_type.name == "Q8_0" and not tensor.name.startswith("dflash.") + ) + require(not invalid_q8, + f"q4-mix has non-head Q8_0 tensors: {invalid_q8}", errors) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--target", required=True, type=Path) + parser.add_argument("--draft", required=True, type=Path) + parser.add_argument("--draft-scheme", choices=("f16", "q8_0", "q4-mix"), + default="q8_0") + args = parser.parse_args() + + errors: list[str] = [] + for label, path in (("target", args.target), ("draft", args.draft)): + if not path.is_file(): + errors.append(f"{label} is not a readable file: {path}") + if not errors: + validate_target(args.target, errors) + validate_draft(args.draft, args.draft_scheme, errors) + + if errors: + for error in errors: + print(f"error: {error}", file=sys.stderr) + return 1 + print(f"PR #625 target OK: {args.target}") + print(f"PR #625 no-YaRN {args.draft_scheme} drafter OK: {args.draft}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/server/src/common/backend_args.h b/server/src/common/backend_args.h index ee7b55d33..66b6feecd 100644 --- a/server/src/common/backend_args.h +++ b/server/src/common/backend_args.h @@ -10,6 +10,7 @@ #include "placement/remote_target_shard_config.h" #include "prefill_attention_mode.h" +#include "speculation_policy.h" namespace dflash::common { // Server-owned features that participate in backend admission even though @@ -81,6 +82,7 @@ struct BackendArgs { bool ddtree_chain_seed = true; int verify_width = 0; // chain spec verify width; 0 = adaptive bool use_feature_mirror = false; + SpeculationPolicy speculation_policy = SpeculationPolicy::Adaptive; }; } // namespace dflash::common diff --git a/server/src/common/backend_factory.cpp b/server/src/common/backend_factory.cpp index 0d1ccc61b..d49d2584a 100644 --- a/server/src/common/backend_factory.cpp +++ b/server/src/common/backend_factory.cpp @@ -284,6 +284,7 @@ std::unique_ptr create_backend( cfg.ddtree_chain_seed = args.ddtree_chain_seed; cfg.use_feature_mirror = args.use_feature_mirror; + cfg.speculation_policy = args.speculation_policy; auto backend = std::make_unique(cfg); if (!backend->init()) { std::fprintf(stderr, "[backend_factory] Qwen35Backend init failed\n"); @@ -310,6 +311,7 @@ std::unique_ptr create_backend( cfg.ddtree_chain_seed = args.ddtree_chain_seed; cfg.use_feature_mirror = args.use_feature_mirror; + cfg.speculation_policy = args.speculation_policy; auto backend = std::make_unique(cfg); if (!backend->init()) { std::fprintf(stderr, "[backend_factory] Qwen35MoeBackend init failed\n"); diff --git a/server/src/common/concurrency/chain_spec_shapes.h b/server/src/common/concurrency/chain_spec_shapes.h new file mode 100644 index 000000000..380babb61 --- /dev/null +++ b/server/src/common/concurrency/chain_spec_shapes.h @@ -0,0 +1,151 @@ +// Pure host-side shape helpers for path-shaped DSpark verification. + +#pragma once + +#include "common/ddtree.h" + +#include +#include +#include + +namespace dflash::common { + +inline int chain_decode_bucket_width(int lanes) { + static constexpr int buckets[] = { + 1, 2, 3, 4, 6, 8, 12, 16, 24, 32, 48, 64, + }; + if (lanes <= 0) return 0; + for (int bucket : buckets) { + if (bucket >= lanes) return bucket; + } + return 64; +} + +// draft_tokens[0] is the already-pending root; positions 1.. form the +// proposal. DDTree's flat indices then coincide with chain depth. +inline DDTree make_chain_verify_tree( + const std::vector & draft_tokens) { + DDTree tree; + if (draft_tokens.size() <= 1) return tree; + + tree.n_nodes = static_cast(draft_tokens.size()) - 1; + tree.token_ids.assign(draft_tokens.begin() + 1, draft_tokens.end()); + tree.depths.resize(static_cast(tree.n_nodes)); + tree.parents.resize(static_cast(tree.n_nodes) + 1); + tree.child_maps.resize(static_cast(tree.n_nodes) + 1); + tree.parents[0] = -1; + for (int node = 1; node <= tree.n_nodes; ++node) { + tree.depths[static_cast(node) - 1] = node; + tree.parents[static_cast(node)] = node - 1; + tree.child_maps[static_cast(node) - 1] + [tree.token_ids[static_cast(node) - 1]] = node; + } + + const int width = tree.n_nodes + 1; + tree.visibility.assign(static_cast(width) * width, 0); + for (int row = 0; row < width; ++row) { + for (int col = 0; col <= row; ++col) { + tree.visibility[static_cast(row) * width + col] = 1; + } + } + return tree; +} + +// A chain verify always includes the pending root. Depth 1 would therefore +// be an AR-equivalent target step, which is forbidden once a request has +// committed to sticky speculation. `requested == 0` means use the configured +// drafter maximum; any other invalid value fails closed. +inline int resolve_chain_verify_depth(int requested, int maximum) { + if (maximum < 2) return 0; + if (requested == 0) return maximum; + return requested >= 2 && requested <= maximum ? requested : 0; +} + +// Keep proposal generation at its configured maximum while allowing one +// common verify depth to be selected per round. Failure leaves the proposal +// untouched, which makes malformed controller/config output non-destructive. +inline bool truncate_chain_proposal( + std::vector & draft_tokens, int verify_depth) { + if (verify_depth < 2 || + verify_depth > static_cast(draft_tokens.size())) { + return false; + } + draft_tokens.resize(static_cast(verify_depth)); + return true; +} + +struct ChainLaunchShape { + int spec_lanes = 0; + int tree_bucket = 0; + int tree_rows = 0; + int ar_lanes = 0; + int ar_bucket = 0; + int accepted_rows = 0; + int commit_rows = 0; +}; + +inline ChainLaunchShape chain_launch_shape( + const std::vector & admitted, + const std::vector & accepted_lengths, + int tree_width) { + ChainLaunchShape shape; + const size_t count = admitted.size(); + for (size_t i = 0; i < count; ++i) { + if (admitted[i]) { + ++shape.spec_lanes; + if (i < accepted_lengths.size()) { + shape.accepted_rows += std::max(0, accepted_lengths[i]); + } + } + } + shape.ar_lanes = static_cast(count) - shape.spec_lanes; + shape.tree_bucket = chain_decode_bucket_width(shape.spec_lanes); + shape.tree_rows = shape.tree_bucket * std::max(0, tree_width); + shape.ar_bucket = chain_decode_bucket_width(shape.ar_lanes); + shape.commit_rows = shape.accepted_rows + shape.ar_bucket; + return shape; +} + +// Proposal preparation is independent per request. A requested speculative +// lane whose proposal failed must be removed from both executor cohorts: it is +// a lane-local failure, never an AR fallback for that decode step. +enum class ChainLaneDisposition : uint8_t { + AR, + Speculation, + Failed, +}; + +inline ChainLaneDisposition chain_lane_disposition( + bool requested_speculation, bool proposal_failed) { + if (proposal_failed) return ChainLaneDisposition::Failed; + return requested_speculation + ? ChainLaneDisposition::Speculation + : ChainLaneDisposition::AR; +} + +inline bool chain_lane_executes(ChainLaneDisposition disposition) { + return disposition != ChainLaneDisposition::Failed; +} + +// The pending root at path[0] was sampled by the preceding target step, so +// the ordinary sampler has already applied the min-token EOS floor to it. +// Accepted children would bypass that sampler. Stop before an EOS that is +// still below the floor so replay samples a replacement from the kept +// tip's exact logits. Once the floor is met, keep the EOS itself but discard +// deeper accepted tokens that the scheduler would hide after retirement. +template +inline size_t chain_min_tokens_safe_prefix( + const std::vector & path, + int generated_tokens_before_root, + int min_tokens, + IsEos is_eos) { + const int generated = std::max(0, generated_tokens_before_root); + for (size_t child = 1; child < path.size(); ++child) { + if (!is_eos(path[child])) continue; + return generated + static_cast(child) < min_tokens + ? child : child + 1; + } + return path.size(); +} + +} // namespace dflash::common diff --git a/server/src/common/concurrency/paged_kv_pool.cpp b/server/src/common/concurrency/paged_kv_pool.cpp index 6bdf0f3fe..741ace600 100644 --- a/server/src/common/concurrency/paged_kv_pool.cpp +++ b/server/src/common/concurrency/paged_kv_pool.cpp @@ -20,6 +20,12 @@ const char * paged_kv_status_string(PagedKvStatus status) { return "physical blocks exhausted"; case PagedKvStatus::StaleHandle: return "stale sequence handle"; + case PagedKvStatus::LogicalBlockOutOfRange: + return "logical block out of range"; + case PagedKvStatus::BlockNotResident: + return "logical block is not resident"; + case PagedKvStatus::BlockAlreadyResident: + return "logical block is already resident"; } return "unknown paged KV status"; } @@ -138,8 +144,35 @@ PagedKvAppendResult PagedKvPool::append(PagedKvSequenceHandle handle, const uint32_t old_kv_seq_len = sequence.kv_seq_len; const uint32_t new_kv_seq_len = old_kv_seq_len + token_count; - result.status = - extend_block_table(sequence, blocks_for_tokens(new_kv_seq_len)); + const uint32_t required_blocks = blocks_for_tokens(new_kv_seq_len); + const uint32_t additional_blocks = + required_blocks - static_cast(sequence.block_table.size()); + + // Appending into a partially-filled cold head needs one physical remap in + // addition to any newly opened logical blocks. Preflight the aggregate so + // append remains all-or-nothing on BlocksExhausted. + const bool remap_head = + old_kv_seq_len % block_size_ != 0 && + sequence.block_table[old_kv_seq_len / block_size_] == + PAGED_KV_COLD_BLOCK; + const uint32_t allocations = additional_blocks + (remap_head ? 1u : 0u); + const uint64_t available = + static_cast(sequence.reserved_blocks.size()) + + free_blocks_.size(); + if (allocations > available) { + result.status = PagedKvStatus::BlocksExhausted; + return result; + } + + sequence.block_table.reserve(required_blocks); + if (remap_head) { + const uint32_t logical_block = old_kv_seq_len / block_size_; + const uint32_t physical_block = take_append_block(sequence); + sequence.block_table[logical_block] = physical_block; + result.remapped_cold_blocks.push_back( + {logical_block, physical_block}); + } + result.status = extend_block_table(sequence, required_blocks); if (result.status != PagedKvStatus::Ok) return result; const auto make_slot = [&](uint32_t logical_position) { @@ -176,7 +209,9 @@ PagedKvStatus PagedKvPool::release(PagedKvSequenceHandle handle) { SequenceState & sequence = sequences_[handle.slot]; request_to_slot_.erase(sequence.request_id); for (uint32_t block : sequence.block_table) { - give_back(free_blocks_, block); + if (block != PAGED_KV_COLD_BLOCK) { + give_back(free_blocks_, block); + } } for (uint32_t block : sequence.reserved_blocks) { give_back(free_blocks_, block); @@ -226,11 +261,70 @@ PagedKvStatus PagedKvPool::owned_block_count( if (status != PagedKvStatus::Ok) return status; const SequenceState & sequence = sequences_[handle.slot]; + // Cold logical blocks remain appended sequence capacity even though they + // no longer own a physical page. Preserve the pre-residency API contract: + // appended logical blocks plus blocks reserved for future append. out_count = static_cast( sequence.block_table.size() + sequence.reserved_blocks.size()); return PagedKvStatus::Ok; } +PagedKvStatus PagedKvPool::page_out_block( + PagedKvSequenceHandle handle, uint32_t logical_block, + uint32_t & out_physical_block) { + const PagedKvStatus status = validate(handle); + if (status != PagedKvStatus::Ok) return status; + + SequenceState & sequence = sequences_[handle.slot]; + if (logical_block >= sequence.block_table.size()) { + return PagedKvStatus::LogicalBlockOutOfRange; + } + const uint32_t physical_block = sequence.block_table[logical_block]; + if (physical_block == PAGED_KV_COLD_BLOCK) { + return PagedKvStatus::BlockNotResident; + } + + sequence.block_table[logical_block] = PAGED_KV_COLD_BLOCK; + give_back(free_blocks_, physical_block); + out_physical_block = physical_block; + return PagedKvStatus::Ok; +} + +PagedKvStatus PagedKvPool::page_in_block( + PagedKvSequenceHandle handle, uint32_t logical_block, + uint32_t & out_physical_block) { + const PagedKvStatus status = validate(handle); + if (status != PagedKvStatus::Ok) return status; + + SequenceState & sequence = sequences_[handle.slot]; + if (logical_block >= sequence.block_table.size()) { + return PagedKvStatus::LogicalBlockOutOfRange; + } + if (sequence.block_table[logical_block] != PAGED_KV_COLD_BLOCK) { + return PagedKvStatus::BlockAlreadyResident; + } + if (free_blocks_.empty()) { + return PagedKvStatus::BlocksExhausted; + } + + const uint32_t physical_block = take_lowest(free_blocks_); + sequence.block_table[logical_block] = physical_block; + out_physical_block = physical_block; + return PagedKvStatus::Ok; +} + +PagedKvStatus PagedKvPool::resident_block_count( + PagedKvSequenceHandle handle, uint32_t & out_count) const { + const PagedKvStatus status = validate(handle); + if (status != PagedKvStatus::Ok) return status; + + const SequenceState & sequence = sequences_[handle.slot]; + out_count = static_cast(std::count_if( + sequence.block_table.begin(), sequence.block_table.end(), + [](uint32_t block) { return block != PAGED_KV_COLD_BLOCK; })); + return PagedKvStatus::Ok; +} + uint32_t PagedKvPool::blocks_for_tokens(uint32_t token_count) const { if (token_count == 0) return 0; return 1 + (token_count - 1) / block_size_; @@ -262,13 +356,17 @@ PagedKvStatus PagedKvPool::extend_block_table(SequenceState & sequence, sequence.block_table.reserve(required_blocks); for (uint32_t i = 0; i < additional_blocks; ++i) { - std::vector & source = sequence.reserved_blocks.empty() - ? free_blocks_ : sequence.reserved_blocks; - sequence.block_table.push_back(take_lowest(source)); + sequence.block_table.push_back(take_append_block(sequence)); } return PagedKvStatus::Ok; } +uint32_t PagedKvPool::take_append_block(SequenceState & sequence) { + std::vector & source = sequence.reserved_blocks.empty() + ? free_blocks_ : sequence.reserved_blocks; + return take_lowest(source); +} + void PagedKvPool::take_reserved_blocks(SequenceState & sequence, uint32_t additional_blocks) { sequence.reserved_blocks.reserve( diff --git a/server/src/common/concurrency/paged_kv_pool.h b/server/src/common/concurrency/paged_kv_pool.h index 9d745739f..3cf90964a 100644 --- a/server/src/common/concurrency/paged_kv_pool.h +++ b/server/src/common/concurrency/paged_kv_pool.h @@ -33,6 +33,9 @@ enum class PagedKvStatus : uint8_t { SequenceSlotsExhausted, // all max_sequences slots are in use BlocksExhausted, // not enough free physical blocks for the growth StaleHandle, // handle refers to a released or reused slot + LogicalBlockOutOfRange, // logical block is not materialized by the sequence + BlockNotResident, // page_out_block() targeted an already-cold block + BlockAlreadyResident, // page_in_block() targeted a resident block }; // Human-readable status name for logs and error messages. @@ -57,6 +60,23 @@ struct PagedKvWriteSlot { uint64_t physical_token_index = 0; }; +// Sentinel in PagedKvSequenceSnapshot::block_table for a logical block whose +// full-attention K/V bytes are host-backed by PagedKvResidencyManager. It is +// never returned as a PagedKvWriteSlot::physical_block. +inline constexpr uint32_t PAGED_KV_COLD_BLOCK = + std::numeric_limits::max(); + +// A formerly-cold append-head block that append() had to remap. The caller +// must restore the block's host-backed bytes before consuming or overwriting +// any row in the returned physical block. PagedKvResidencyManager::append() +// and prepare_append() do that before returning to the engine; this record +// primarily makes a direct pool append fail-obvious instead of losing the +// remap information. +struct PagedKvBlockRemap { + uint32_t logical_block = 0; + uint32_t physical_block = 0; +}; + // Outcome of append(). On success, `token_count` is the number of appended // tokens. By default, `write_slots` holds one entry per token in logical // order. With `only_first_last_slots`, it is empty and `first` and `last` @@ -68,6 +88,7 @@ struct PagedKvAppendResult { std::vector write_slots; PagedKvWriteSlot first; PagedKvWriteSlot last; + std::vector remapped_cold_blocks; explicit operator bool() const { return status == PagedKvStatus::Ok; } }; @@ -75,6 +96,8 @@ struct PagedKvAppendResult { // Copy of one sequence's bookkeeping state, as returned by sequence(). struct PagedKvSequenceSnapshot { uint32_t kv_seq_len = 0; + // Entries are physical block indices or PAGED_KV_COLD_BLOCK. Logical + // length and block-table length do not shrink when a block is paged out. std::vector block_table; // Physical blocks held for future append() calls by this sequence. They // are not visible in block_table until append consumes them. @@ -155,6 +178,26 @@ class PagedKvPool { PagedKvStatus owned_block_count(PagedKvSequenceHandle handle, uint32_t & out_count) const; + // Relinquish one materialized physical block while preserving its logical + // block-table position as PAGED_KV_COLD_BLOCK. The caller must first copy + // the complete block to host backing. On success, out_physical_block is + // the returned pool block and may be reused immediately. + PagedKvStatus page_out_block(PagedKvSequenceHandle handle, + uint32_t logical_block, + uint32_t & out_physical_block); + + // Allocate a physical block for one cold logical entry. The caller must + // restore its complete host-backed bytes before attention or append reads + // it. On failure, the cold mapping and output argument are unchanged. + PagedKvStatus page_in_block(PagedKvSequenceHandle handle, + uint32_t logical_block, + uint32_t & out_physical_block); + + // Number of materialized logical blocks that currently own physical + // storage. Reserved append capacity is deliberately excluded. + PagedKvStatus resident_block_count(PagedKvSequenceHandle handle, + uint32_t & out_count) const; + private: // Bookkeeping for one sequence slot. `generation` survives release so // the next acquire on this slot invalidates old handles. @@ -182,6 +225,11 @@ class PagedKvPool { PagedKvStatus extend_block_table(SequenceState & sequence, uint32_t required_blocks); + // Allocate a physical block from a sequence's reservation first, then the + // global free list. Used only by append(), where consuming promised future + // capacity is correct. + uint32_t take_append_block(SequenceState & sequence); + // Move exactly `additional_blocks` globally free blocks into a sequence's // private reservation. Caller must preflight availability. void take_reserved_blocks(SequenceState & sequence, diff --git a/server/src/common/concurrency/paged_kv_residency.cpp b/server/src/common/concurrency/paged_kv_residency.cpp new file mode 100644 index 000000000..c0af73b12 --- /dev/null +++ b/server/src/common/concurrency/paged_kv_residency.cpp @@ -0,0 +1,974 @@ +#include "paged_kv_residency.h" + +#include +#include +#include +#include + +namespace dflash::common { + +namespace { + +PagedKvResidencyStatus from_pool_status(PagedKvStatus status) { + switch (status) { + case PagedKvStatus::Ok: + return PagedKvResidencyStatus::Ok; + case PagedKvStatus::StaleHandle: + return PagedKvResidencyStatus::StaleHandle; + case PagedKvStatus::BlocksExhausted: + return PagedKvResidencyStatus::PoolExhausted; + case PagedKvStatus::InvalidArgument: + case PagedKvStatus::LogicalBlockOutOfRange: + case PagedKvStatus::BlockNotResident: + case PagedKvStatus::BlockAlreadyResident: + return PagedKvResidencyStatus::InvalidArgument; + case PagedKvStatus::DuplicateRequest: + case PagedKvStatus::SequenceSlotsExhausted: + return PagedKvResidencyStatus::InconsistentPoolState; + } + return PagedKvResidencyStatus::InconsistentPoolState; +} + +} // namespace + +const char * paged_kv_residency_status_string( + PagedKvResidencyStatus status) { + switch (status) { + case PagedKvResidencyStatus::Ok: + return "ok"; + case PagedKvResidencyStatus::InvalidArgument: + return "invalid argument"; + case PagedKvResidencyStatus::SequenceNotRegistered: + return "sequence not registered"; + case PagedKvResidencyStatus::StaleHandle: + return "stale sequence handle"; + case PagedKvResidencyStatus::PoolExhausted: + return "physical pool exhausted"; + case PagedKvResidencyStatus::NoEvictableBlock: + return "no evictable block"; + case PagedKvResidencyStatus::HostAllocationFailed: + return "pinned host allocation failed"; + case PagedKvResidencyStatus::TransferFailed: + return "K/V transfer failed"; + case PagedKvResidencyStatus::HostCopyMissing: + return "cold block has no valid host copy"; + case PagedKvResidencyStatus::InconsistentPoolState: + return "residency state disagrees with paged pool"; + } + return "unknown paged K/V residency status"; +} + +PagedKvResidencyManager::PagedKvResidencyManager( + PagedKvPool & pool, PagedKvResidencyConfig config, + PagedKvResidencyTransferOps transfers) + : pool_(pool), config_(config), transfers_(std::move(transfers)) { + if (config_.block_bytes == 0 || !transfers_.allocate_pinned || + !transfers_.free_pinned || !transfers_.copy_out_async || + !transfers_.copy_in_async || !transfers_.synchronize) { + throw std::invalid_argument("invalid paged K/V residency callbacks"); + } + if (config_.resident_budget_blocks == 0) { + config_.resident_budget_blocks = pool_.physical_block_count(); + } + if (config_.resident_budget_blocks > pool_.physical_block_count()) { + throw std::invalid_argument("resident budget exceeds paged K/V pool"); + } + sequences_.resize(pool_.max_sequences()); +} + +PagedKvResidencyManager::~PagedKvResidencyManager() { + try { + reset(); + } catch (...) { + // Destructors must not propagate callback failures. Production pinned + // allocators/free functions are non-throwing; this catch protects test + // and plugin callbacks from terminating teardown. + } +} + +PagedKvResidencyStatus PagedKvResidencyManager::validate_registered( + PagedKvSequenceHandle handle) const { + if (transfer_barrier_failed_) { + return PagedKvResidencyStatus::TransferFailed; + } + if (handle.slot >= sequences_.size()) { + return PagedKvResidencyStatus::StaleHandle; + } + const SequenceState & state = sequences_[handle.slot]; + if (!state.active) { + return PagedKvResidencyStatus::SequenceNotRegistered; + } + if (state.generation != handle.generation) { + return PagedKvResidencyStatus::StaleHandle; + } + return PagedKvResidencyStatus::Ok; +} + +PagedKvResidencyStatus PagedKvResidencyManager::register_sequence( + PagedKvSequenceHandle handle) { + if (transfer_barrier_failed_) { + return PagedKvResidencyStatus::TransferFailed; + } + PagedKvSequenceSnapshot snapshot; + const PagedKvStatus pool_status = pool_.sequence(handle, snapshot); + if (pool_status != PagedKvStatus::Ok) { + return from_pool_status(pool_status); + } + if (handle.slot >= sequences_.size()) { + return PagedKvResidencyStatus::StaleHandle; + } + if (std::find(snapshot.block_table.begin(), snapshot.block_table.end(), + PAGED_KV_COLD_BLOCK) != snapshot.block_table.end()) { + return PagedKvResidencyStatus::InconsistentPoolState; + } + + SequenceState & state = sequences_[handle.slot]; + if (state.active && state.generation == handle.generation) { + return PagedKvResidencyStatus::Ok; + } + if (state.active) { + const auto synced = synchronize_before_read(); + if (synced != PagedKvResidencyStatus::Ok) return synced; + free_sequence_buffers(state); + } + state.active = true; + state.generation = handle.generation; + state.blocks.resize(snapshot.block_table.size()); + for (BlockState & block : state.blocks) block.last_use = ++clock_; + return PagedKvResidencyStatus::Ok; +} + +PagedKvResidencyStatus PagedKvResidencyManager::forget_sequence( + PagedKvSequenceHandle handle) { + // Teardown is also the recovery path after a failed copy-stream barrier. + // validate_registered() deliberately rejects ordinary operations while a + // transfer is quarantined, so validate the generation directly here and + // allow synchronize_before_read() to retry the barrier. + if (handle.slot >= sequences_.size()) { + return PagedKvResidencyStatus::StaleHandle; + } + const SequenceState & state = sequences_[handle.slot]; + if (!state.active) { + return PagedKvResidencyStatus::SequenceNotRegistered; + } + if (state.generation != handle.generation) { + return PagedKvResidencyStatus::StaleHandle; + } + const auto synced = synchronize_before_read(); + if (synced != PagedKvResidencyStatus::Ok) return synced; + free_sequence_buffers(sequences_[handle.slot]); + return PagedKvResidencyStatus::Ok; +} + +void PagedKvResidencyManager::reset() { + // A failed barrier should not let teardown free memory still referenced by + // an async transfer. Leak those buffers rather than creating a use-after- + // free; the process/backend is already unhealthy in this case. + if (synchronize_before_read() != PagedKvResidencyStatus::Ok) return; + for (SequenceState & state : sequences_) free_sequence_buffers(state); + stats_ = {}; + clock_ = 0; +} + +void PagedKvResidencyManager::free_sequence_buffers( + SequenceState & state) noexcept { + for (BlockState & block : state.blocks) { + if (!block.host) continue; + try { + transfers_.free_pinned(block.host); + } catch (...) { + // The callback contract is non-throwing. Continue freeing other + // pages if a third-party implementation violates it. + } + block.host = nullptr; + if (stats_.host_bytes >= config_.block_bytes) { + stats_.host_bytes -= config_.block_bytes; + } + } + state = {}; +} + +PagedKvResidencyStatus PagedKvResidencyManager::refresh_sequence( + PagedKvSequenceHandle handle, bool appended_rows, + uint32_t append_first, uint32_t append_count) { + const auto registered = validate_registered(handle); + if (registered != PagedKvResidencyStatus::Ok) return registered; + PagedKvSequenceSnapshot snapshot; + const PagedKvStatus pool_status = pool_.sequence(handle, snapshot); + if (pool_status != PagedKvStatus::Ok) return from_pool_status(pool_status); + + SequenceState & state = sequences_[handle.slot]; + if (snapshot.block_table.size() < state.blocks.size()) { + return PagedKvResidencyStatus::InconsistentPoolState; + } + state.blocks.resize(snapshot.block_table.size()); + for (uint32_t logical = 0; logical < snapshot.block_table.size(); ++logical) { + if (snapshot.block_table[logical] == PAGED_KV_COLD_BLOCK && + !state.blocks[logical].host_valid) { + return PagedKvResidencyStatus::HostCopyMissing; + } + } + + if (appended_rows && append_count > 0) { + const uint64_t last = static_cast(append_first) + + append_count - 1; + const uint32_t first_block = append_first / pool_.block_size(); + const uint32_t last_block = + static_cast(last / pool_.block_size()); + if (last_block >= state.blocks.size()) { + return PagedKvResidencyStatus::InconsistentPoolState; + } + for (uint32_t logical = first_block; logical <= last_block; ++logical) { + // The target graph has not written the returned rows yet. Retain a + // restored partial page's old host image until commit, but make the + // physical page ineligible for eviction in every policy path. + state.blocks[logical].write_pending = true; + state.blocks[logical].last_use = ++clock_; + } + } + return PagedKvResidencyStatus::Ok; +} + +PagedKvResidencyStatus PagedKvResidencyManager::prepare_append( + PagedKvSequenceHandle handle, uint32_t token_count) { + const auto registered = validate_registered(handle); + if (registered != PagedKvResidencyStatus::Ok) return registered; + + PagedKvSequenceSnapshot snapshot; + const PagedKvStatus pool_status = pool_.sequence(handle, snapshot); + if (pool_status != PagedKvStatus::Ok) return from_pool_status(pool_status); + if (token_count > std::numeric_limits::max() - + snapshot.kv_seq_len) { + return PagedKvResidencyStatus::InvalidArgument; + } + if (token_count == 0) return finish_transfers(PagedKvResidencyStatus::Ok); + + BlockState * append_head = nullptr; + bool restore_partial_head = false; + uint32_t partial_head = 0; + if (snapshot.kv_seq_len % pool_.block_size() != 0) { + partial_head = snapshot.kv_seq_len / pool_.block_size(); + append_head = &sequences_[handle.slot].blocks[partial_head]; + append_head->reservation_pending = true; + restore_partial_head = + snapshot.block_table[partial_head] == PAGED_KV_COLD_BLOCK; + } + + const uint32_t new_length = snapshot.kv_seq_len + token_count; + const uint32_t required_blocks = + 1 + (new_length - 1) / pool_.block_size(); + const uint32_t additional_blocks = required_blocks - + static_cast(snapshot.block_table.size()); + const uint32_t globally_needed = additional_blocks > + snapshot.reserved_block_count + ? additional_blocks - snapshot.reserved_block_count : 0; + // Reserve the partial-head restoration and every new physical page as one + // transaction. Restoring first and making room later can otherwise select + // the just-restored append head as the next eviction victim. + const auto room = make_room( + handle, globally_needed + (restore_partial_head ? 1u : 0u), + additional_blocks + (restore_partial_head ? 1u : 0u)); + if (room != PagedKvResidencyStatus::Ok) { + if (append_head) append_head->reservation_pending = false; + return room; + } + if (restore_partial_head) { + const auto restored = restore_block_async( + handle, partial_head); + if (restored != PagedKvResidencyStatus::Ok) { + const auto finished = finish_transfers(restored); + if (append_head) append_head->reservation_pending = false; + return finished; + } + } + const auto finished = finish_transfers(PagedKvResidencyStatus::Ok); + if (append_head) append_head->reservation_pending = false; + return finished; +} + +PagedKvResidentAppendResult PagedKvResidencyManager::append( + PagedKvSequenceHandle handle, uint32_t token_count, + bool only_first_last_slots) { + PagedKvResidentAppendResult result; + result.status = prepare_append(handle, token_count); + if (result.status != PagedKvResidencyStatus::Ok) return result; + + result.pool_result = + pool_.append(handle, token_count, only_first_last_slots); + if (result.pool_result.status != PagedKvStatus::Ok) { + result.status = from_pool_status(result.pool_result.status); + return result; + } + result.status = observe_append(handle, result.pool_result); + return result; +} + +PagedKvResidencyStatus PagedKvResidencyManager::observe_append( + PagedKvSequenceHandle handle, const PagedKvAppendResult & result) { + const auto registered = validate_registered(handle); + if (registered != PagedKvResidencyStatus::Ok) return registered; + if (result.status != PagedKvStatus::Ok) return from_pool_status(result.status); + + SequenceState & state = sequences_[handle.slot]; + for (const PagedKvBlockRemap & remap : result.remapped_cold_blocks) { + if (remap.logical_block >= state.blocks.size() || + !state.blocks[remap.logical_block].host_valid) { + return finish_transfers(PagedKvResidencyStatus::HostCopyMissing); + } + BlockState & block = state.blocks[remap.logical_block]; + // A callback may reject only after queuing a prefix of a multi-tensor + // copy. Quarantine the destination before invoking it either way. + transfers_pending_ = true; + block.page_in_pending = true; + bool queued = false; + try { + queued = transfers_.copy_in_async( + handle, remap.logical_block, remap.physical_block, + block.host, config_.block_bytes); + } catch (...) { + queued = false; + } + const PendingTransfer transfer{ + handle, remap.logical_block, remap.physical_block}; + if (!queued) { + pending_page_in_rollbacks_.push_back(transfer); + return finish_transfers(PagedKvResidencyStatus::TransferFailed); + } + pending_page_ins_.push_back(transfer); + } + const auto synced = finish_transfers(PagedKvResidencyStatus::Ok); + if (synced != PagedKvResidencyStatus::Ok) return synced; + + PagedKvSequenceSnapshot snapshot; + const PagedKvStatus pool_status = pool_.sequence(handle, snapshot); + if (pool_status != PagedKvStatus::Ok) return from_pool_status(pool_status); + if (result.token_count > snapshot.kv_seq_len) { + return PagedKvResidencyStatus::InconsistentPoolState; + } + return refresh_sequence( + handle, /*appended_rows=*/true, + snapshot.kv_seq_len - result.token_count, result.token_count); +} + +PagedKvResidencyStatus PagedKvResidencyManager::commit_pending_writes( + PagedKvSequenceHandle handle) { + const auto registered = validate_registered(handle); + if (registered != PagedKvResidencyStatus::Ok) return registered; + + // The caller supplies the target-compute -> pager dependency by + // synchronizing the target backend before this call. Any old host image is + // stale only now, after the device rows have actually been overwritten. + SequenceState & state = sequences_[handle.slot]; + for (BlockState & block : state.blocks) { + if (!block.write_pending) continue; + block.write_pending = false; + block.host_valid = false; + block.last_use = ++clock_; + } + return PagedKvResidencyStatus::Ok; +} + +PagedKvResidencyStatus PagedKvResidencyManager::finish_transfers( + PagedKvResidencyStatus status) { + const auto synced = synchronize_before_read(); + return synced == PagedKvResidencyStatus::Ok ? status : synced; +} + +PagedKvResidencyStatus PagedKvResidencyManager::synchronize_before_read() { + if (!transfers_pending_) { + return transfer_barrier_failed_ + ? PagedKvResidencyStatus::TransferFailed + : PagedKvResidencyStatus::Ok; + } + bool ok = false; + try { + ok = transfers_.synchronize(); + } catch (...) { + ok = false; + } + if (!ok) { + // A failed barrier does not prove that the stream stopped using its + // source and destination blocks. Keep every transfer pending and every + // H2D destination mapped (therefore quarantined from the free list), + // and reject all manager operations until an explicit retry confirms + // that the stream has drained. + transfer_barrier_failed_ = true; + return PagedKvResidencyStatus::TransferFailed; + } + + transfers_pending_ = false; + transfer_barrier_failed_ = false; + + PagedKvResidencyStatus result = PagedKvResidencyStatus::Ok; + for (const PendingTransfer & transfer : pending_page_outs_) { + if (validate_registered(transfer.handle) != + PagedKvResidencyStatus::Ok) { + result = PagedKvResidencyStatus::InconsistentPoolState; + continue; + } + BlockState & block = + sequences_[transfer.handle.slot].blocks[transfer.logical_block]; + block.page_out_pending = false; + PagedKvSequenceSnapshot snapshot; + if (pool_.sequence(transfer.handle, snapshot) != PagedKvStatus::Ok || + transfer.logical_block >= snapshot.block_table.size() || + snapshot.block_table[transfer.logical_block] != + transfer.physical_block) { + block.host_valid = false; + result = PagedKvResidencyStatus::InconsistentPoolState; + continue; + } + uint32_t released = PAGED_KV_COLD_BLOCK; + if (pool_.page_out_block( + transfer.handle, transfer.logical_block, released) != + PagedKvStatus::Ok || released != transfer.physical_block) { + block.host_valid = false; + result = PagedKvResidencyStatus::InconsistentPoolState; + continue; + } + block.host_valid = true; + stats_.page_outs++; + stats_.moved_bytes += config_.block_bytes; + } + for (const PendingTransfer & transfer : pending_page_ins_) { + if (validate_registered(transfer.handle) != + PagedKvResidencyStatus::Ok) { + result = PagedKvResidencyStatus::InconsistentPoolState; + continue; + } + BlockState & block = + sequences_[transfer.handle.slot].blocks[transfer.logical_block]; + block.page_in_pending = false; + stats_.page_ins++; + stats_.moved_bytes += config_.block_bytes; + } + for (const PendingTransfer & transfer : pending_page_in_rollbacks_) { + if (validate_registered(transfer.handle) != + PagedKvResidencyStatus::Ok) { + result = PagedKvResidencyStatus::InconsistentPoolState; + continue; + } + BlockState & block = + sequences_[transfer.handle.slot].blocks[transfer.logical_block]; + PagedKvSequenceSnapshot snapshot; + if (pool_.sequence(transfer.handle, snapshot) != PagedKvStatus::Ok || + transfer.logical_block >= snapshot.block_table.size() || + snapshot.block_table[transfer.logical_block] != + transfer.physical_block) { + result = PagedKvResidencyStatus::InconsistentPoolState; + continue; + } + uint32_t released = PAGED_KV_COLD_BLOCK; + if (pool_.page_out_block( + transfer.handle, transfer.logical_block, released) != + PagedKvStatus::Ok || released != transfer.physical_block) { + result = PagedKvResidencyStatus::InconsistentPoolState; + continue; + } + block.page_in_pending = false; + } + pending_page_outs_.clear(); + pending_page_ins_.clear(); + pending_page_in_rollbacks_.clear(); + if (result != PagedKvResidencyStatus::Ok) return result; + return PagedKvResidencyStatus::Ok; +} + +bool PagedKvResidencyManager::is_protected( + PagedKvSequenceHandle handle, uint32_t logical_block) const { + PagedKvSequenceSnapshot snapshot; + if (pool_.sequence(handle, snapshot) != PagedKvStatus::Ok || + logical_block >= snapshot.block_table.size()) { + return true; + } + if (logical_block < config_.sink_blocks) return true; + const uint32_t tail_begin = snapshot.block_table.size() > config_.tail_blocks + ? static_cast(snapshot.block_table.size()) - + config_.tail_blocks + : 0; + return logical_block >= tail_begin; +} + +uint32_t PagedKvResidencyManager::sequence_resident_count( + PagedKvSequenceHandle handle) const { + uint32_t count = 0; + return pool_.resident_block_count(handle, count) == PagedKvStatus::Ok + ? count : 0; +} + +uint32_t PagedKvResidencyManager::sequence_pending_page_out_count( + PagedKvSequenceHandle handle) const { + if (handle.slot >= sequences_.size()) return 0; + const SequenceState & state = sequences_[handle.slot]; + if (!state.active || state.generation != handle.generation) return 0; + return static_cast(std::count_if( + state.blocks.begin(), state.blocks.end(), + [](const BlockState & block) { return block.page_out_pending; })); +} + +uint32_t PagedKvResidencyManager::sequence_protected_count( + PagedKvSequenceHandle handle) const { + PagedKvSequenceSnapshot snapshot; + if (pool_.sequence(handle, snapshot) != PagedKvStatus::Ok) return 0; + uint32_t count = 0; + for (uint32_t logical = 0; logical < snapshot.block_table.size(); ++logical) { + if (snapshot.block_table[logical] != PAGED_KV_COLD_BLOCK && + is_protected(handle, logical)) { + ++count; + } + } + return count; +} + +uint32_t PagedKvResidencyManager::total_resident_count() const { + uint32_t total = 0; + for (uint32_t slot = 0; slot < sequences_.size(); ++slot) { + const SequenceState & state = sequences_[slot]; + if (!state.active) continue; + total += sequence_resident_count({slot, state.generation}); + } + return total; +} + +uint32_t PagedKvResidencyManager::quota_for_slot(uint32_t slot) const { + uint32_t active = 0; + uint32_t rank = 0; + for (uint32_t i = 0; i < sequences_.size(); ++i) { + if (!sequences_[i].active) continue; + if (i < slot) ++rank; + ++active; + } + if (active == 0 || slot >= sequences_.size() || + !sequences_[slot].active) { + return 0; + } + const uint32_t base = config_.resident_budget_blocks / active; + const uint32_t remainder = config_.resident_budget_blocks % active; + const uint32_t fair = base + (rank < remainder ? 1u : 0u); + const PagedKvSequenceHandle handle{slot, sequences_[slot].generation}; + return std::max(fair, sequence_protected_count(handle)); +} + +uint32_t PagedKvResidencyManager::fair_quota( + PagedKvSequenceHandle handle) const { + return validate_registered(handle) == PagedKvResidencyStatus::Ok + ? quota_for_slot(handle.slot) : 0; +} + +PagedKvResidencyManager::Victim PagedKvResidencyManager::choose_victim( + PagedKvSequenceHandle requester) const { + Victim best; + for (uint32_t slot = 0; slot < sequences_.size(); ++slot) { + const SequenceState & state = sequences_[slot]; + if (!state.active) continue; + const PagedKvSequenceHandle owner{slot, state.generation}; + PagedKvSequenceSnapshot snapshot; + if (pool_.sequence(owner, snapshot) != PagedKvStatus::Ok) continue; + const uint32_t resident = sequence_resident_count(owner); + const uint32_t pending_page_outs = + sequence_pending_page_out_count(owner); + const uint32_t resident_after_pending = + resident > pending_page_outs ? resident - pending_page_outs : 0; + const uint32_t quota = quota_for_slot(slot); + + int class_rank = 2; + if (resident_after_pending > quota) class_rank = 0; + else if (slot == requester.slot && + state.generation == requester.generation) class_rank = 1; + + for (uint32_t logical = 0; logical < snapshot.block_table.size(); ++logical) { + if (snapshot.block_table[logical] == PAGED_KV_COLD_BLOCK || + state.blocks[logical].write_pending || + state.blocks[logical].page_out_pending || + state.blocks[logical].page_in_pending || + state.blocks[logical].reservation_pending || + is_protected(owner, logical)) { + continue; + } + const BlockState & block = state.blocks[logical]; + // An unscored page is the first eviction candidate once relevance + // scoring is active for only part of a sequence. + const float score = block.score_valid + ? block.score : -std::numeric_limits::infinity(); + const bool better = !best.found || + class_rank < best.class_rank || + (class_rank == best.class_rank && score < best.score) || + (class_rank == best.class_rank && score == best.score && + block.last_use < best.last_use) || + (class_rank == best.class_rank && score == best.score && + block.last_use == best.last_use && + (slot < best.handle.slot || + (slot == best.handle.slot && + logical < best.logical_block))); + if (better) { + best.found = true; + best.handle = owner; + best.logical_block = logical; + best.class_rank = class_rank; + best.score = score; + best.last_use = block.last_use; + } + } + } + return best; +} + +PagedKvResidencyStatus PagedKvResidencyManager::make_room( + PagedKvSequenceHandle requester, uint32_t pool_blocks_needed, + uint32_t future_resident_blocks) { + const uint32_t free = pool_.free_block_count(); + const uint32_t need_for_pool = pool_blocks_needed > free + ? pool_blocks_needed - free : 0; + const uint32_t resident = total_resident_count(); + const uint64_t projected = + static_cast(resident) + future_resident_blocks; + const uint32_t need_for_budget = + projected > config_.resident_budget_blocks + ? static_cast(projected - config_.resident_budget_blocks) + : 0; + const uint32_t evictions = std::max(need_for_pool, need_for_budget); + for (uint32_t i = 0; i < evictions; ++i) { + const Victim victim = choose_victim(requester); + if (!victim.found) { + return finish_transfers(PagedKvResidencyStatus::NoEvictableBlock); + } + const auto status = evict_block_async( + victim.handle, victim.logical_block, + /*allow_protected=*/false); + if (status != PagedKvResidencyStatus::Ok) return finish_transfers(status); + } + // A page is recyclable only after the complete D2H batch succeeds. + return finish_transfers(PagedKvResidencyStatus::Ok); +} + +PagedKvResidencyStatus PagedKvResidencyManager::evict_block_async( + PagedKvSequenceHandle handle, uint32_t logical_block, + bool allow_protected) { + const auto registered = validate_registered(handle); + if (registered != PagedKvResidencyStatus::Ok) return registered; + PagedKvSequenceSnapshot snapshot; + const PagedKvStatus pool_status = pool_.sequence(handle, snapshot); + if (pool_status != PagedKvStatus::Ok) return from_pool_status(pool_status); + if (logical_block >= snapshot.block_table.size()) { + return PagedKvResidencyStatus::InvalidArgument; + } + if (snapshot.block_table[logical_block] == PAGED_KV_COLD_BLOCK) { + return PagedKvResidencyStatus::InvalidArgument; + } + + BlockState & block = sequences_[handle.slot].blocks[logical_block]; + if (block.write_pending || block.page_out_pending || + block.page_in_pending || + (!allow_protected && is_protected(handle, logical_block))) { + return PagedKvResidencyStatus::NoEvictableBlock; + } + if (!block.host) { + try { + block.host = transfers_.allocate_pinned(config_.block_bytes); + } catch (...) { + block.host = nullptr; + } + if (!block.host) return PagedKvResidencyStatus::HostAllocationFailed; + stats_.host_bytes += config_.block_bytes; + } + + const uint32_t physical = snapshot.block_table[logical_block]; + bool queued = false; + // False may mean a multi-tensor callback queued only a prefix. Invalidate + // the old host image and require a barrier before any later operation. + transfers_pending_ = true; + block.host_valid = false; + try { + queued = transfers_.copy_out_async( + handle, logical_block, physical, block.host, + config_.block_bytes); + } catch (...) { + queued = false; + } + if (!queued) return PagedKvResidencyStatus::TransferFailed; + block.page_out_pending = true; + pending_page_outs_.push_back({handle, logical_block, physical}); + return PagedKvResidencyStatus::Ok; +} + +PagedKvResidencyStatus PagedKvResidencyManager::restore_block_async( + PagedKvSequenceHandle handle, uint32_t logical_block) { + const auto registered = validate_registered(handle); + if (registered != PagedKvResidencyStatus::Ok) return registered; + PagedKvSequenceSnapshot snapshot; + const PagedKvStatus pool_status = pool_.sequence(handle, snapshot); + if (pool_status != PagedKvStatus::Ok) return from_pool_status(pool_status); + if (logical_block >= snapshot.block_table.size()) { + return PagedKvResidencyStatus::InvalidArgument; + } + if (snapshot.block_table[logical_block] != PAGED_KV_COLD_BLOCK) { + sequences_[handle.slot].blocks[logical_block].last_use = ++clock_; + return PagedKvResidencyStatus::Ok; + } + BlockState & block = sequences_[handle.slot].blocks[logical_block]; + if (block.page_out_pending || block.page_in_pending) { + return PagedKvResidencyStatus::InconsistentPoolState; + } + if (!block.host || !block.host_valid) { + return PagedKvResidencyStatus::HostCopyMissing; + } + + uint32_t physical = PAGED_KV_COLD_BLOCK; + const PagedKvStatus page_status = + pool_.page_in_block(handle, logical_block, physical); + if (page_status != PagedKvStatus::Ok) return from_pool_status(page_status); + + bool queued = false; + transfers_pending_ = true; + block.page_in_pending = true; + try { + queued = transfers_.copy_in_async( + handle, logical_block, physical, block.host, + config_.block_bytes); + } catch (...) { + queued = false; + } + const PendingTransfer transfer{handle, logical_block, physical}; + if (!queued) { + // The callback may have queued a prefix. Keep the physical mapping + // quarantined until finish_transfers() proves the stream is drained. + pending_page_in_rollbacks_.push_back(transfer); + return PagedKvResidencyStatus::TransferFailed; + } + pending_page_ins_.push_back(transfer); + block.last_use = ++clock_; + return PagedKvResidencyStatus::Ok; +} + +PagedKvResidencyStatus PagedKvResidencyManager::ensure_resident( + PagedKvSequenceHandle handle, + const std::vector & logical_blocks) { + const auto registered = validate_registered(handle); + if (registered != PagedKvResidencyStatus::Ok) return registered; + + PagedKvSequenceSnapshot snapshot; + const PagedKvStatus pool_status = pool_.sequence(handle, snapshot); + if (pool_status != PagedKvStatus::Ok) return from_pool_status(pool_status); + std::vector cold; + std::vector seen(snapshot.block_table.size(), 0); + // Validate the complete request before reserving any member. Otherwise a + // later bad index leaves earlier blocks permanently ineligible as victims. + if (std::any_of(logical_blocks.begin(), logical_blocks.end(), + [&](uint32_t logical) { + return logical >= snapshot.block_table.size(); + })) { + return PagedKvResidencyStatus::InvalidArgument; + } + for (uint32_t logical : logical_blocks) { + if (seen[logical]) continue; + seen[logical] = 1; + sequences_[handle.slot].blocks[logical].reservation_pending = true; + if (snapshot.block_table[logical] == PAGED_KV_COLD_BLOCK) { + cold.push_back(logical); + } else { + sequences_[handle.slot].blocks[logical].last_use = ++clock_; + } + } + const auto clear_reservations = [&] { + for (uint32_t logical = 0; logical < seen.size(); ++logical) { + if (seen[logical]) { + sequences_[handle.slot].blocks[logical].reservation_pending = false; + } + } + }; + const auto room = make_room( + handle, static_cast(cold.size()), + static_cast(cold.size())); + if (room != PagedKvResidencyStatus::Ok) { + clear_reservations(); + return room; + } + + // Restore the requested set from one reserved capacity pool. No member can + // become the victim of a later member in this same operation. + for (uint32_t logical : cold) { + const auto status = restore_block_async( + handle, logical); + if (status != PagedKvResidencyStatus::Ok) { + const auto finished = finish_transfers(status); + clear_reservations(); + return finished; + } + } + const auto finished = finish_transfers(PagedKvResidencyStatus::Ok); + clear_reservations(); + return finished; +} + +PagedKvResidencyStatus PagedKvResidencyManager::evict_block( + PagedKvSequenceHandle handle, uint32_t logical_block, + bool allow_protected) { + return finish_transfers( + evict_block_async(handle, logical_block, allow_protected)); +} + +PagedKvResidencyStatus PagedKvResidencyManager::touch( + PagedKvSequenceHandle handle, uint32_t logical_block) { + const auto registered = validate_registered(handle); + if (registered != PagedKvResidencyStatus::Ok) return registered; + if (logical_block >= sequences_[handle.slot].blocks.size()) { + return PagedKvResidencyStatus::InvalidArgument; + } + sequences_[handle.slot].blocks[logical_block].last_use = ++clock_; + return PagedKvResidencyStatus::Ok; +} + +PagedKvResidencyStatus PagedKvResidencyManager::set_scores( + PagedKvSequenceHandle handle, const std::vector & scores) { + const auto registered = validate_registered(handle); + if (registered != PagedKvResidencyStatus::Ok) return registered; + SequenceState & state = sequences_[handle.slot]; + if (scores.empty()) { + for (BlockState & block : state.blocks) block.score_valid = false; + return PagedKvResidencyStatus::Ok; + } + if (scores.size() != state.blocks.size() || + std::any_of(scores.begin(), scores.end(), + [](float score) { return !std::isfinite(score); })) { + return PagedKvResidencyStatus::InvalidArgument; + } + for (uint32_t i = 0; i < scores.size(); ++i) { + state.blocks[i].score = scores[i]; + state.blocks[i].score_valid = true; + } + return PagedKvResidencyStatus::Ok; +} + +PagedKvResidencyStatus PagedKvResidencyManager::reselect( + PagedKvSequenceHandle handle) { + const auto registered = validate_registered(handle); + if (registered != PagedKvResidencyStatus::Ok) return registered; + stats_.reselects++; + + PagedKvSequenceSnapshot snapshot; + const PagedKvStatus pool_status = pool_.sequence(handle, snapshot); + if (pool_status != PagedKvStatus::Ok) return from_pool_status(pool_status); + SequenceState & state = sequences_[handle.slot]; + + std::vector wanted(snapshot.block_table.size(), 0); + uint32_t wanted_count = 0; + for (uint32_t logical = 0; logical < snapshot.block_table.size(); ++logical) { + if (state.blocks[logical].write_pending || + is_protected(handle, logical)) { + wanted[logical] = 1; + ++wanted_count; + } + } + const uint32_t target = std::max( + wanted_count, + std::min(quota_for_slot(handle.slot), + snapshot.block_table.size())); + + std::vector candidates; + for (uint32_t logical = 0; logical < snapshot.block_table.size(); ++logical) { + if (!wanted[logical]) candidates.push_back(logical); + } + std::sort(candidates.begin(), candidates.end(), + [&](uint32_t a, uint32_t b) { + const BlockState & lhs = state.blocks[a]; + const BlockState & rhs = state.blocks[b]; + if (lhs.score_valid != rhs.score_valid) return lhs.score_valid; + if (lhs.score_valid && lhs.score != rhs.score) return lhs.score > rhs.score; + if (lhs.last_use != rhs.last_use) return lhs.last_use > rhs.last_use; + return a < b; + }); + for (uint32_t logical : candidates) { + if (wanted_count >= target) break; + wanted[logical] = 1; + ++wanted_count; + } + + for (uint32_t logical = 0; logical < wanted.size(); ++logical) { + if (wanted[logical]) state.blocks[logical].reservation_pending = true; + } + const auto clear_reservations = [&] { + for (uint32_t logical = 0; logical < wanted.size(); ++logical) { + if (wanted[logical]) state.blocks[logical].reservation_pending = false; + } + }; + + // Out first, making one batch of free physical pages before recalls. + for (uint32_t logical = 0; logical < snapshot.block_table.size(); ++logical) { + if (!wanted[logical] && + snapshot.block_table[logical] != PAGED_KV_COLD_BLOCK) { + const auto status = evict_block_async( + handle, logical, /*allow_protected=*/false); + if (status != PagedKvResidencyStatus::Ok) { + const auto finished = finish_transfers(status); + clear_reservations(); + return finished; + } + } + } + auto finished = finish_transfers(PagedKvResidencyStatus::Ok); + if (finished != PagedKvResidencyStatus::Ok) { + clear_reservations(); + return finished; + } + + std::vector cold_wanted; + for (uint32_t logical = 0; logical < snapshot.block_table.size(); ++logical) { + if (wanted[logical] && !is_resident(handle, logical)) { + cold_wanted.push_back(logical); + } + } + const auto room = make_room( + handle, static_cast(cold_wanted.size()), + static_cast(cold_wanted.size())); + if (room != PagedKvResidencyStatus::Ok) { + clear_reservations(); + return room; + } + for (uint32_t logical : cold_wanted) { + const auto status = restore_block_async( + handle, logical); + if (status != PagedKvResidencyStatus::Ok) { + finished = finish_transfers(status); + clear_reservations(); + return finished; + } + } + finished = finish_transfers(PagedKvResidencyStatus::Ok); + clear_reservations(); + return finished; +} + +PagedKvResidencyStatus PagedKvResidencyManager::rebalance() { + while (true) { + // The dummy requester prevents any real sequence from getting the + // requester-swap class. A class-0 victim is necessarily over quota. + const PagedKvSequenceHandle none{ + std::numeric_limits::max(), 0}; + const Victim victim = choose_victim(none); + if (!victim.found || victim.class_rank != 0) break; + const auto status = evict_block_async( + victim.handle, victim.logical_block, + /*allow_protected=*/false); + if (status != PagedKvResidencyStatus::Ok) { + return finish_transfers(status); + } + } + return finish_transfers(PagedKvResidencyStatus::Ok); +} + +bool PagedKvResidencyManager::is_resident( + PagedKvSequenceHandle handle, uint32_t logical_block) const { + if (validate_registered(handle) != PagedKvResidencyStatus::Ok) return false; + PagedKvSequenceSnapshot snapshot; + return pool_.sequence(handle, snapshot) == PagedKvStatus::Ok && + logical_block < snapshot.block_table.size() && + snapshot.block_table[logical_block] != PAGED_KV_COLD_BLOCK; +} + +PagedKvResidencyStats PagedKvResidencyManager::stats() const { + PagedKvResidencyStats out = stats_; + out.resident_blocks = total_resident_count(); + return out; +} + +} // namespace dflash::common diff --git a/server/src/common/concurrency/paged_kv_residency.h b/server/src/common/concurrency/paged_kv_residency.h new file mode 100644 index 000000000..ac2d982b0 --- /dev/null +++ b/server/src/common/concurrency/paged_kv_residency.h @@ -0,0 +1,261 @@ +// Multi-sequence host residency for PagedKvPool full-attention K/V blocks. +// +// The policy and bookkeeping are backend-neutral. The Qwen engine supplies a +// pinned allocator plus async full-block D2H/H2D callbacks that capture its +// dedicated copy stream and know how to concatenate every full-attention K/V +// tensor into one host block. The manager batches transfers and synchronizes +// once before returning any operation that permits device K/V to be read or a +// recycled physical block to be written. + +#pragma once + +#include "paged_kv_pool.h" + +#include +#include +#include +#include + +namespace dflash::common { + +enum class PagedKvResidencyStatus : uint8_t { + Ok = 0, + InvalidArgument, + SequenceNotRegistered, + StaleHandle, + PoolExhausted, + NoEvictableBlock, + HostAllocationFailed, + TransferFailed, + HostCopyMissing, + InconsistentPoolState, +}; + +const char * paged_kv_residency_status_string(PagedKvResidencyStatus status); + +struct PagedKvResidencyConfig { + // Bytes copied for one physical pool block across all full-attention K/V + // tensors. Recurrent/DeltaNet state is intentionally outside this pager. + size_t block_bytes = 0; + + // Hard bound for materialized logical blocks. Zero uses the complete + // physical PagedKvPool. Reserved-but-unappended pool blocks still consume + // physical capacity and may temporarily make the effective bound smaller. + uint32_t resident_budget_blocks = 0; + + // Attention sinks and the trailing local window are never automatic + // eviction victims. Explicit evict_block(..., allow_protected=true) is + // available for teardown/tests only. + uint32_t sink_blocks = 1; + uint32_t tail_blocks = 4; +}; + +struct PagedKvResidencyTransferOps { + // Production implementations should use cudaMallocHost/hipHostMalloc (or + // the equivalent backend pinned allocator). Every allocation is exactly + // config.block_bytes and lives until forget_sequence()/reset(). + std::function allocate_pinned; + std::function free_pinned; + + // Queue a complete physical-block copy on one dedicated copy stream. + // handle/logical_block are supplied for diagnostics only; physical_block + // identifies the source/destination in the Qwen paged K/V tensors. + std::function copy_out_async; + std::function copy_in_async; + + // Synchronize that copy stream. The manager calls it once after a batch + // and before attention reads or append writes can observe remapped blocks. + std::function synchronize; +}; + +struct PagedKvResidencyStats { + uint64_t page_ins = 0; + uint64_t page_outs = 0; + uint64_t resident_blocks = 0; + uint64_t reselects = 0; + uint64_t host_bytes = 0; + uint64_t moved_bytes = 0; +}; + +struct PagedKvResidentAppendResult { + PagedKvResidencyStatus status = PagedKvResidencyStatus::Ok; + PagedKvAppendResult pool_result; + + explicit operator bool() const { + return status == PagedKvResidencyStatus::Ok && + pool_result.status == PagedKvStatus::Ok; + } +}; + +class PagedKvResidencyManager { +public: + // Throws std::invalid_argument for an empty transfer callback, zero block + // bytes, or a resident budget larger than the physical pool. + PagedKvResidencyManager(PagedKvPool & pool, + PagedKvResidencyConfig config, + PagedKvResidencyTransferOps transfers); + ~PagedKvResidencyManager(); + + PagedKvResidencyManager(const PagedKvResidencyManager &) = delete; + PagedKvResidencyManager & operator=(const PagedKvResidencyManager &) = delete; + + // Call immediately after PagedKvPool::acquire[_reserved](). Registration + // rejects pre-existing cold entries because their host backing is unknown. + PagedKvResidencyStatus register_sequence(PagedKvSequenceHandle handle); + + // Free this sequence's pinned backing. Call before pool.release, and do + // not release the pool handle unless this returns Ok: a failed transfer + // barrier keeps its physical pages and pinned buffers quarantined until + // forget_sequence is retried. The handle must still match the manager's + // registered generation. + PagedKvResidencyStatus forget_sequence(PagedKvSequenceHandle handle); + void reset(); + + // Recommended append integration: restores a cold partial append head, + // fairly evicts enough blocks, synchronizes the copy stream, appends in + // PagedKvPool, and marks every touched block pending until the engine has + // finished the target graph that writes the returned physical rows. + PagedKvResidentAppendResult append(PagedKvSequenceHandle handle, + uint32_t token_count, + bool only_first_last_slots = false); + + // Split integration for callers whose slot manager owns pool.append(): + // prepare_append(handle, n); pool.append(handle, n); observe_append(...) + // prepare_append always synchronizes before returning Ok. observe_append + // must be called before the next residency operation. + PagedKvResidencyStatus prepare_append(PagedKvSequenceHandle handle, + uint32_t token_count); + PagedKvResidencyStatus observe_append( + PagedKvSequenceHandle handle, const PagedKvAppendResult & result); + + // Clear all blocks staged by append()/observe_append() for this sequence. + // Call only after synchronizing the target compute that wrote every returned + // physical row. Pending blocks cannot be evicted or deselected, preventing a + // later slot staged in the same packed step from recycling an unwritten page. + PagedKvResidencyStatus commit_pending_writes( + PagedKvSequenceHandle handle); + + // Restore the requested host-backed logical blocks as one transfer batch. + // The operation synchronizes before returning Ok, so attention may read + // the returned pool block table immediately. + PagedKvResidencyStatus ensure_resident( + PagedKvSequenceHandle handle, + const std::vector & logical_blocks); + + // Explicit single-block eviction. Automatic policy never evicts sink/tail + // blocks; protected eviction requires an explicit opt-in. Pending target + // writes are never evictable, including with allow_protected=true. + PagedKvResidencyStatus evict_block(PagedKvSequenceHandle handle, + uint32_t logical_block, + bool allow_protected = false); + + // Update recency after a block participates in attention. + PagedKvResidencyStatus touch(PagedKvSequenceHandle handle, + uint32_t logical_block); + + // Optional relevance array, one score per materialized logical block. + // Higher values are retained. Without scores, reselection uses LRU. + PagedKvResidencyStatus set_scores(PagedKvSequenceHandle handle, + const std::vector & scores); + PagedKvResidencyStatus reselect(PagedKvSequenceHandle handle); + + // Evict unprotected blocks above each active sequence's deterministic + // fair share. Spare capacity remains borrowable; it is reclaimed first + // from borrowers when another sequence needs a block. + PagedKvResidencyStatus rebalance(); + + // Explicit barrier for engine paths that directly issue operations and + // then read the paged K/V tensors. Normally append/ensure/reselect/evict + // already provide this barrier. + PagedKvResidencyStatus synchronize_before_read(); + + uint32_t fair_quota(PagedKvSequenceHandle handle) const; + bool is_resident(PagedKvSequenceHandle handle, + uint32_t logical_block) const; + PagedKvResidencyStats stats() const; + +private: + struct BlockState { + void * host = nullptr; + bool host_valid = false; + bool write_pending = false; + bool page_out_pending = false; + bool page_in_pending = false; + bool reservation_pending = false; + bool score_valid = false; + float score = 0.0f; + uint64_t last_use = 0; + }; + + struct SequenceState { + bool active = false; + uint64_t generation = 0; + std::vector blocks; + }; + + struct Victim { + bool found = false; + PagedKvSequenceHandle handle; + uint32_t logical_block = 0; + int class_rank = 0; + float score = 0.0f; + uint64_t last_use = 0; + }; + + struct PendingTransfer { + PagedKvSequenceHandle handle; + uint32_t logical_block = 0; + uint32_t physical_block = PAGED_KV_COLD_BLOCK; + }; + + PagedKvResidencyStatus validate_registered( + PagedKvSequenceHandle handle) const; + PagedKvResidencyStatus refresh_sequence( + PagedKvSequenceHandle handle, bool appended_rows = false, + uint32_t append_first = 0, uint32_t append_count = 0); + PagedKvResidencyStatus finish_transfers( + PagedKvResidencyStatus status); + + bool is_protected(PagedKvSequenceHandle handle, + uint32_t logical_block) const; + uint32_t sequence_resident_count(PagedKvSequenceHandle handle) const; + uint32_t sequence_pending_page_out_count( + PagedKvSequenceHandle handle) const; + uint32_t sequence_protected_count(PagedKvSequenceHandle handle) const; + uint32_t total_resident_count() const; + uint32_t quota_for_slot(uint32_t slot) const; + + Victim choose_victim(PagedKvSequenceHandle requester) const; + PagedKvResidencyStatus make_room(PagedKvSequenceHandle requester, + uint32_t pool_blocks_needed, + uint32_t future_resident_blocks); + PagedKvResidencyStatus evict_block_async( + PagedKvSequenceHandle handle, uint32_t logical_block, + bool allow_protected); + PagedKvResidencyStatus restore_block_async( + PagedKvSequenceHandle handle, uint32_t logical_block); + + void free_sequence_buffers(SequenceState & state) noexcept; + + PagedKvPool & pool_; + PagedKvResidencyConfig config_; + PagedKvResidencyTransferOps transfers_; + std::vector sequences_; + PagedKvResidencyStats stats_; + uint64_t clock_ = 0; + bool transfers_pending_ = false; + bool transfer_barrier_failed_ = false; + std::vector pending_page_outs_; + std::vector pending_page_ins_; + std::vector pending_page_in_rollbacks_; +}; + +} // namespace dflash::common diff --git a/server/src/common/concurrency/qwen_paged_kv_transfer.cpp b/server/src/common/concurrency/qwen_paged_kv_transfer.cpp new file mode 100644 index 000000000..d9a2b74cf --- /dev/null +++ b/server/src/common/concurrency/qwen_paged_kv_transfer.cpp @@ -0,0 +1,405 @@ +#include "qwen_paged_kv_transfer.h" + +#include "common/gpu_runtime_compat.h" +#include "internal.h" + +#include "ggml.h" + +#include +#include +#include + +namespace dflash::common { +namespace { + +void set_error(std::string * error, const std::string & message) { + if (error) *error = message; +} + +bool runtime_pointer_is_device(const void * pointer, int device) { + cudaPointerAttributes attributes{}; + const cudaError_t status = cudaPointerGetAttributes(&attributes, pointer); + if (status != cudaSuccess) { + (void)cudaGetLastError(); + return false; + } +#if defined(DFLASH27B_BACKEND_HIP) || defined(GGML_USE_HIP) + return attributes.type == hipMemoryTypeDevice && + attributes.device == device; +#else + return attributes.type == cudaMemoryTypeDevice && + attributes.device == device; +#endif +} + +bool tensor_is_device_backed(const ggml_tensor * tensor, int device, + std::string * error) { + if (!tensor || !tensor->buffer || !tensor->data) { + set_error(error, "paged K/V tensor is null or unallocated"); + return false; + } + const ggml_backend_buffer_type_t buft = + ggml_backend_buffer_get_type(tensor->buffer); + if (!buft || ggml_backend_buft_is_meta(buft)) { + set_error(error, + "meta/tensor-parallel paged K/V buffers are unsupported"); + return false; + } + if (ggml_backend_buft_is_host(buft)) { + set_error(error, "host-backed paged K/V buffers are unsupported"); + return false; + } + const ggml_backend_dev_t tensor_device = + ggml_backend_buft_get_device(buft); + if (!tensor_device) { + set_error(error, "paged K/V buffer has no backend device"); + return false; + } + const enum ggml_backend_dev_type tensor_device_type = + ggml_backend_dev_type(tensor_device); + if (tensor_device_type != GGML_BACKEND_DEVICE_TYPE_GPU && + tensor_device_type != GGML_BACKEND_DEVICE_TYPE_IGPU) { + set_error(error, "paged K/V buffer is not GPU device memory"); + return false; + } + if (!runtime_pointer_is_device(tensor->data, device)) { + set_error(error, + "paged K/V tensor is not on the requested HIP/CUDA device"); + return false; + } + return true; +} + +} // namespace + +struct QwenPagedKvResidencyTransfer::State { + struct TensorCopy { + uint8_t * device_data = nullptr; + QwenPagedKvTensorLayout layout; + size_t host_offset = 0; + size_t host_head_bytes = 0; + }; + + int device = -1; + uint32_t block_size = 0; + size_t block_bytes = 0; + uint32_t physical_block_count = 0; + cudaStream_t stream = nullptr; + bool stream_may_reference_host = false; + std::vector tensors; + std::vector pinned_allocations; + + ~State() { + if (device >= 0 && cudaSetDevice(device) == cudaSuccess) { + if (stream && cudaStreamSynchronize(stream) != cudaSuccess) { + // Match PagedKvResidencyManager's fail-safe teardown rule: + // leak stream/backing rather than free host memory that a + // failed runtime may still reference. + pinned_allocations.clear(); + stream = nullptr; + return; + } + for (void * pointer : pinned_allocations) { + if (pointer) (void)cudaFreeHost(pointer); + } + pinned_allocations.clear(); + if (stream) (void)cudaStreamDestroy(stream); + } + stream = nullptr; + } + + bool select_device() const { + return device >= 0 && cudaSetDevice(device) == cudaSuccess; + } + + void * allocate_pinned(size_t bytes) { + if (bytes != block_bytes || !select_device()) return nullptr; + void * pointer = nullptr; + if (cudaMallocHost(&pointer, bytes) != cudaSuccess || !pointer) { + (void)cudaGetLastError(); + return nullptr; + } + try { + pinned_allocations.push_back(pointer); + } catch (...) { + (void)cudaFreeHost(pointer); + return nullptr; + } + return pointer; + } + + void free_pinned(void * pointer) { + if (!pointer) return; + const auto found = std::find( + pinned_allocations.begin(), pinned_allocations.end(), pointer); + if (found == pinned_allocations.end()) return; + if (stream_may_reference_host) { + if (!stream || !select_device() || + cudaStreamSynchronize(stream) != cudaSuccess) { + return; + } + stream_may_reference_host = false; + } + if (select_device() && cudaFreeHost(pointer) == cudaSuccess) { + pinned_allocations.erase(found); + } + } + + bool synchronize() { + if (!stream || !select_device()) return false; + if (cudaStreamSynchronize(stream) != cudaSuccess) { + (void)cudaGetLastError(); + stream_may_reference_host = true; + return false; + } + stream_may_reference_host = false; + return true; + } + + bool queue_copy(uint32_t physical_block, void * host, size_t bytes, + bool copy_out) { + if (!host || bytes != block_bytes || !stream || + physical_block >= physical_block_count || !select_device()) { + return false; + } + + const size_t first_row = + static_cast(physical_block) * block_size; + uint8_t * const host_bytes = static_cast(host); + for (const TensorCopy & tensor : tensors) { + uint8_t * const device_block = tensor.device_data + + first_row * tensor.layout.row_stride; + uint8_t * const host_tensor = + host_bytes + tensor.host_offset; + + cudaError_t status = cudaSuccess; + if (tensor.layout.row_stride == tensor.layout.row_bytes) { + // Common Qwen cache layout: block rows are contiguous inside + // each head plane, so one 2D copy covers every head. + status = copy_out + ? cudaMemcpy2DAsync( + host_tensor, tensor.host_head_bytes, + device_block, tensor.layout.head_stride, + tensor.host_head_bytes, tensor.layout.heads, + cudaMemcpyDeviceToHost, stream) + : cudaMemcpy2DAsync( + device_block, tensor.layout.head_stride, + host_tensor, tensor.host_head_bytes, + tensor.host_head_bytes, tensor.layout.heads, + cudaMemcpyHostToDevice, stream); + } else { + // Preserve uncommon row padding with one pitched async copy + // per head. Host images remain tightly packed payload bytes. + for (size_t head = 0; + head < static_cast(tensor.layout.heads); + ++head) { + uint8_t * const device_head = + device_block + head * tensor.layout.head_stride; + uint8_t * const host_head = + host_tensor + head * tensor.host_head_bytes; + status = copy_out + ? cudaMemcpy2DAsync( + host_head, tensor.layout.row_bytes, + device_head, tensor.layout.row_stride, + tensor.layout.row_bytes, block_size, + cudaMemcpyDeviceToHost, stream) + : cudaMemcpy2DAsync( + device_head, tensor.layout.row_stride, + host_head, tensor.layout.row_bytes, + tensor.layout.row_bytes, block_size, + cudaMemcpyHostToDevice, stream); + if (status != cudaSuccess) break; + } + } + if (status != cudaSuccess) { + // Drain any prefix already queued here. The residency manager + // also treats a false return as pending, so a failed drain is + // quarantined and retried before ownership can be released. + (void)cudaGetLastError(); + if (cudaStreamSynchronize(stream) != cudaSuccess) { + stream_may_reference_host = true; + } + return false; + } + } + return true; + } +}; + +QwenPagedKvResidencyTransfer::QwenPagedKvResidencyTransfer( + std::shared_ptr state) + : state_(std::move(state)) {} + +QwenPagedKvResidencyTransfer::~QwenPagedKvResidencyTransfer() = default; + +std::unique_ptr +QwenPagedKvResidencyTransfer::create( + const TargetCache & cache, + ggml_backend_t backend, + int device, + uint32_t block_size, + std::string * error) { + if (error) error->clear(); + if (!backend || device < 0 || block_size == 0) { + set_error(error, "invalid paged K/V transfer backend/device/block size"); + return nullptr; + } + const ggml_backend_dev_t backend_device = ggml_backend_get_device(backend); + if (!backend_device) { + set_error(error, "paged K/V transfer backend has no device"); + return nullptr; + } + const enum ggml_backend_dev_type backend_type = + ggml_backend_dev_type(backend_device); + if (backend_type == GGML_BACKEND_DEVICE_TYPE_META || + ggml_backend_buft_is_meta( + ggml_backend_get_default_buffer_type(backend))) { + set_error(error, + "meta/tensor-parallel paged K/V transfer is unsupported"); + return nullptr; + } + if (backend_type != GGML_BACKEND_DEVICE_TYPE_GPU && + backend_type != GGML_BACKEND_DEVICE_TYPE_IGPU) { + set_error(error, "paged K/V transfer requires a GPU backend"); + return nullptr; + } + if (cache.backend && cache.backend != backend) { + set_error(error, "paged K/V cache belongs to a different backend"); + return nullptr; + } + if (cache.attn_k.empty() || + cache.attn_k.size() != cache.attn_v.size()) { + set_error(error, "paged K/V cache has no complete attention layers"); + return nullptr; + } + + int device_count = 0; + if (cudaGetDeviceCount(&device_count) != cudaSuccess || + device >= device_count || cudaSetDevice(device) != cudaSuccess) { + (void)cudaGetLastError(); + set_error(error, "invalid HIP/CUDA device for paged K/V transfer"); + return nullptr; + } + + std::vector cache_tensors; + std::vector layouts; + try { + cache_tensors.reserve(cache.attn_k.size() * 2); + layouts.reserve(cache.attn_k.size() * 2); + } catch (...) { + set_error(error, "paged K/V transfer layout allocation failed"); + return nullptr; + } + + for (size_t layer = 0; layer < cache.attn_k.size(); ++layer) { + const ggml_tensor * pair[] = { + cache.attn_k[layer], cache.attn_v[layer], + }; + if (!pair[0] || !pair[1]) { + set_error(error, + "partial/tensor-parallel paged K/V cache is unsupported"); + return nullptr; + } + if (pair[0]->ne[1] != pair[1]->ne[1] || + pair[0]->ne[2] != pair[1]->ne[2]) { + set_error(error, "K/V cache pair dimensions do not match"); + return nullptr; + } + for (const ggml_tensor * tensor : pair) { + if (!tensor_is_device_backed(tensor, device, error)) return nullptr; + if (tensor->ne[0] <= 0 || tensor->ne[1] <= 0 || + tensor->ne[2] <= 0 || tensor->ne[3] != 1 || + tensor->ne[0] % ggml_blck_size(tensor->type) != 0) { + set_error(error, "unsupported paged K/V tensor dimensions"); + return nullptr; + } + const size_t row_bytes = + ggml_row_size(tensor->type, tensor->ne[0]); + cache_tensors.push_back(tensor); + layouts.push_back({ + row_bytes, + tensor->nb[1], + tensor->nb[2], + ggml_nbytes(tensor), + static_cast(tensor->ne[1]), + static_cast(tensor->ne[2]), + }); + } + } + + QwenPagedKvBlockLayout plan; + if (!plan_qwen_paged_kv_block_layout( + layouts, block_size, plan, error)) { + return nullptr; + } + + std::shared_ptr state; + try { + state = std::make_shared(); + state->device = device; + state->block_size = block_size; + state->block_bytes = plan.block_bytes; + state->physical_block_count = plan.physical_block_count; + state->tensors.reserve(cache_tensors.size()); + for (size_t i = 0; i < cache_tensors.size(); ++i) { + state->tensors.push_back({ + static_cast(cache_tensors[i]->data), + layouts[i], + plan.tensor_offsets[i], + plan.tensor_head_bytes[i], + }); + } + } catch (...) { + set_error(error, "paged K/V transfer state allocation failed"); + return nullptr; + } + + if (cudaStreamCreateWithFlags( + &state->stream, cudaStreamNonBlocking) != cudaSuccess) { + (void)cudaGetLastError(); + set_error(error, "failed to create paged K/V transfer stream"); + return nullptr; + } + + try { + return std::unique_ptr( + new QwenPagedKvResidencyTransfer(std::move(state))); + } catch (...) { + set_error(error, "paged K/V transfer wrapper allocation failed"); + return nullptr; + } +} + +size_t QwenPagedKvResidencyTransfer::block_bytes() const noexcept { + return state_ ? state_->block_bytes : 0; +} + +PagedKvResidencyTransferOps +QwenPagedKvResidencyTransfer::callbacks() const { + PagedKvResidencyTransferOps result; + const std::shared_ptr state = state_; + if (!state) return result; + result.allocate_pinned = [state](size_t bytes) { + return state->allocate_pinned(bytes); + }; + result.free_pinned = [state](void * pointer) { + state->free_pinned(pointer); + }; + result.copy_out_async = + [state](PagedKvSequenceHandle, uint32_t, uint32_t physical_block, + void * host, size_t bytes) { + return state->queue_copy( + physical_block, host, bytes, /*copy_out=*/true); + }; + result.copy_in_async = + [state](PagedKvSequenceHandle, uint32_t, uint32_t physical_block, + const void * host, size_t bytes) { + return state->queue_copy( + physical_block, const_cast(host), bytes, + /*copy_out=*/false); + }; + result.synchronize = [state] { return state->synchronize(); }; + return result; +} + +} // namespace dflash::common diff --git a/server/src/common/concurrency/qwen_paged_kv_transfer.h b/server/src/common/concurrency/qwen_paged_kv_transfer.h new file mode 100644 index 000000000..2aee297a4 --- /dev/null +++ b/server/src/common/concurrency/qwen_paged_kv_transfer.h @@ -0,0 +1,87 @@ +// Production HIP/CUDA block transfers for Qwen paged K/V residency. +// +// A host image packs the payload bytes for one physical block from every +// full-attention K and V tensor. Device row/head strides are retained in the +// copy plan, so quantized K/V types and padded tensor layouts are copied +// without reinterpretation. + +#pragma once + +#include "paged_kv_residency.h" + +#include "ggml-backend.h" + +#include +#include +#include +#include +#include + +namespace dflash::common { + +struct TargetCache; + +// GPU-independent input/output for block-layout validation. storage_bytes is +// the accessible span starting at the tensor data pointer (ggml_nbytes for a +// normal cache tensor). +struct QwenPagedKvTensorLayout { + size_t row_bytes = 0; + size_t row_stride = 0; + size_t head_stride = 0; + size_t storage_bytes = 0; + uint64_t physical_rows = 0; + uint64_t heads = 0; +}; + +struct QwenPagedKvBlockLayout { + size_t block_bytes = 0; + uint32_t physical_block_count = 0; + // One entry per input tensor. Tensor payloads are packed consecutively; + // each head occupies tensor_head_bytes[i] bytes in the host image. + std::vector tensor_offsets; + std::vector tensor_head_bytes; +}; + +// Validates all dimensions/strides and computes the packed host block image. +// Every tensor must cover the same physical row count and head count. Row +// padding is supported; padding bytes are not copied into the host image. +bool plan_qwen_paged_kv_block_layout( + const std::vector & tensors, + uint32_t block_size, + QwenPagedKvBlockLayout & out, + std::string * error = nullptr); + +// Owns the dedicated nonblocking transfer stream. callbacks() retains shared +// ownership of the stream/layout state, so this wrapper may be destroyed as +// soon as callbacks are handed to PagedKvResidencyManager. TargetCache and +// its K/V buffers must outlive those callbacks. Before asking the residency +// manager to evict, the engine must have completed the compute work that last +// wrote the source block; the manager's synchronize callback supplies the +// opposite copy-stream -> later-attention/append barrier. +class QwenPagedKvResidencyTransfer { +public: + static std::unique_ptr create( + const TargetCache & cache, + ggml_backend_t backend, + int device, + uint32_t block_size, + std::string * error = nullptr); + + ~QwenPagedKvResidencyTransfer(); + + QwenPagedKvResidencyTransfer( + const QwenPagedKvResidencyTransfer &) = delete; + QwenPagedKvResidencyTransfer & operator=( + const QwenPagedKvResidencyTransfer &) = delete; + + size_t block_bytes() const noexcept; + PagedKvResidencyTransferOps callbacks() const; + +private: + struct State; + explicit QwenPagedKvResidencyTransfer(std::shared_ptr state); + + std::shared_ptr state_; +}; + +} // namespace dflash::common diff --git a/server/src/common/concurrency/qwen_paged_kv_transfer_layout.cpp b/server/src/common/concurrency/qwen_paged_kv_transfer_layout.cpp new file mode 100644 index 000000000..d75a3f864 --- /dev/null +++ b/server/src/common/concurrency/qwen_paged_kv_transfer_layout.cpp @@ -0,0 +1,119 @@ +#include "qwen_paged_kv_transfer.h" + +#include + +namespace dflash::common { +namespace { + +bool fail(QwenPagedKvBlockLayout & out, std::string * error, + const char * message) { + out = {}; + if (error) *error = message; + return false; +} + +bool checked_mul(size_t lhs, size_t rhs, size_t & out) { + if (lhs != 0 && rhs > std::numeric_limits::max() / lhs) { + return false; + } + out = lhs * rhs; + return true; +} + +bool checked_add(size_t lhs, size_t rhs, size_t & out) { + if (rhs > std::numeric_limits::max() - lhs) return false; + out = lhs + rhs; + return true; +} + +} // namespace + +bool plan_qwen_paged_kv_block_layout( + const std::vector & tensors, + uint32_t block_size, + QwenPagedKvBlockLayout & out, + std::string * error) { + out = {}; + if (tensors.empty() || block_size == 0) { + return fail(out, error, "empty paged K/V layout or zero block size"); + } + + const uint64_t physical_rows = tensors.front().physical_rows; + const uint64_t heads = tensors.front().heads; + if (physical_rows == 0 || heads == 0 || + physical_rows % block_size != 0 || + physical_rows / block_size > std::numeric_limits::max()) { + return fail(out, error, "invalid paged K/V physical dimensions"); + } + out.physical_block_count = + static_cast(physical_rows / block_size); + + try { + out.tensor_offsets.reserve(tensors.size()); + out.tensor_head_bytes.reserve(tensors.size()); + } catch (...) { + return fail(out, error, "paged K/V layout allocation failed"); + } + + size_t host_offset = 0; + for (const QwenPagedKvTensorLayout & tensor : tensors) { + if (tensor.row_bytes == 0 || tensor.row_stride < tensor.row_bytes || + tensor.head_stride == 0 || tensor.storage_bytes == 0 || + tensor.physical_rows != physical_rows || tensor.heads != heads) { + return fail(out, error, "inconsistent paged K/V tensor layout"); + } + if (physical_rows > std::numeric_limits::max() || + heads > std::numeric_limits::max()) { + return fail(out, error, "paged K/V dimensions exceed host size_t"); + } + + size_t last_row_offset = 0; + size_t last_head_offset = 0; + size_t used_end = 0; + if (!checked_mul(static_cast(physical_rows - 1), + tensor.row_stride, last_row_offset) || + !checked_mul(static_cast(heads - 1), + tensor.head_stride, last_head_offset) || + !checked_add(last_head_offset, last_row_offset, used_end) || + !checked_add(used_end, tensor.row_bytes, used_end) || + used_end > tensor.storage_bytes) { + return fail(out, error, "paged K/V tensor strides exceed storage"); + } + + // Heads must not overlap. Larger head strides (alignment/padding) are + // fine because device copies retain the original pitch. + size_t plane_span = 0; + if (!checked_mul(static_cast(physical_rows - 1), + tensor.row_stride, plane_span) || + !checked_add(plane_span, tensor.row_bytes, plane_span) || + tensor.head_stride < plane_span) { + return fail(out, error, "overlapping paged K/V head planes"); + } + + size_t head_bytes = 0; + size_t tensor_bytes = 0; + size_t next_offset = 0; + if (!checked_mul(tensor.row_bytes, block_size, head_bytes) || + !checked_mul(head_bytes, static_cast(heads), + tensor_bytes) || + !checked_add(host_offset, tensor_bytes, next_offset)) { + return fail(out, error, "paged K/V host block size overflow"); + } + try { + out.tensor_offsets.push_back(host_offset); + out.tensor_head_bytes.push_back(head_bytes); + } catch (...) { + return fail(out, error, "paged K/V layout allocation failed"); + } + host_offset = next_offset; + } + + if (host_offset == 0) { + return fail(out, error, "zero-byte paged K/V host block"); + } + out.block_bytes = host_offset; + if (error) error->clear(); + return true; +} + +} // namespace dflash::common diff --git a/server/src/common/concurrency/seq_engine.h b/server/src/common/concurrency/seq_engine.h index fa4dbba35..ddd2e2cf6 100644 --- a/server/src/common/concurrency/seq_engine.h +++ b/server/src/common/concurrency/seq_engine.h @@ -6,11 +6,15 @@ // paged KV cache and execute a batched decode step. Any additional // per-sequence model state is owned by the concrete engine, not by this // interface. admit() claims a slot and queues its prompt without compute. -// Each step() then advances a scheduler-selected cohort of prompt slices -// alongside the complete live decode batch. Once a prefill completes, the -// scheduler advances that slot one token per step(), feeding each sampled -// token back as the next step's input — which is what lets it override a token -// (thinking-budget force-close) before it is committed to the cache. +// Each step() advances the complete live decode batch and may also advance a +// scheduler-selected cohort of prompt slices. A decode graph that cannot mix +// with prefill reports those slices as deferred and leaves their prompt state +// unchanged; the scheduler selects them again after the exclusive decode wave. +// Once a prefill completes, the scheduler feeds the final sampled token back as +// the next step's input. An engine may also return speculative children it +// already committed before that final pending token. The scheduler can disable +// that burst path per slot when it must preserve authority to substitute a +// thinking-budget close token. // // The split of duties is deliberate and is the reason this interface exists // apart from ModelBackend: @@ -47,6 +51,7 @@ // engine through it before wiring it up. #pragma once +#include "common/speculation_policy.h" #include #include @@ -181,20 +186,52 @@ class SeqEngine { struct StepInput { int slot = -1; int32_t token = -1; // token to commit at this slot's next position + // False when scheduler-side policy may replace the sampled token + // before it is committed (currently the thinking-budget close hook). + bool allow_speculation = true; + // Effective server default + per-request override. The hard safety + // check above always wins over this policy. + SpeculationPolicy speculation_policy = SpeculationPolicy::Adaptive; }; struct DecodeOutput { int slot = -1; - int32_t token = -1; // newly sampled token (pending until next step) + // Final newly sampled token, pending until the scheduler feeds it into + // the next step. `committed_tokens`, when non-empty, precede this token + // and are already present in backend state. + int32_t token = -1; bool failed = false; // Present when failed=true so the scheduler can report an honest // per-request error instead of silently truncating generation. std::string error; + std::vector committed_tokens; + + // Per-slot deltas for this engine step. The scheduler aggregates them + // until retirement and emits one machine-readable proof record. + uint64_t ddtree_steps = 0; + uint64_t ddtree_accepted_tokens = 0; + uint64_t ddtree_suspensions = 0; + uint64_t spec_steps = 0; + uint64_t spec_accepted_tokens = 0; + // Sticky-SPEC lane advances performed by an explicit packed + // AR+prefill service round. These are scheduling suspensions, not a + // routing-mode change, and are reported separately from spec_steps. + uint64_t spec_service_ar_steps = 0; + uint64_t target_forwards = 0; + uint64_t kvflash_page_ins = 0; + uint64_t kvflash_page_outs = 0; + uint64_t kvflash_resident_blocks = 0; + uint64_t kvflash_reselects = 0; }; struct PrefillOutput { enum class Status { advanced, completed, + // Selected work intentionally made no progress because the same + // engine call executed an incompatible decode graph. No prompt + // state changed; the scheduler keeps it pending and selects it + // again after the exclusive decode wave. + deferred, failed, }; @@ -237,8 +274,9 @@ class SeqEngine { virtual StepPlanLimits step_plan_limits(int decode_rows) const = 0; // A successful result returns one decode output for every decode input and - // one explicit advanced/completed/failed result for every selected - // prefill. Invalid plans return a fatal error without advancing state. + // one explicit advanced/completed/deferred/failed result for every + // selected prefill. Invalid plans return a fatal error without advancing + // state. // Runtime failures are terminal for the live cohort and may follow partial // backend mutation, but expose no consumable payload. virtual StepResult step(const StepPlan & plan) = 0; @@ -250,6 +288,33 @@ class SeqEngine { virtual bool token_is_eos(int32_t token) const = 0; }; +// Scheduler fairness advances only when selected prompt work actually moved. +// A deferred or failed slice remains non-progress even though both are valid, +// explicit answers for the selected row. +inline bool prefill_result_made_progress( + const SeqEngine::StepResult & result) { + using Status = SeqEngine::PrefillOutput::Status; + return std::any_of( + result.prefills.begin(), result.prefills.end(), + [](const SeqEngine::PrefillOutput & output) { + return output.status == Status::advanced || + output.status == Status::completed; + }); +} + +// Deliver a successful decode result in wire order. The visitor returns false +// after a stop/EOS/output-cap decision; in that case later committed children +// and the final pending token are intentionally hidden and the slot is retired. +template +inline bool consume_decode_output_tokens( + const SeqEngine::DecodeOutput & output, Advance advance) { + if (output.failed) return false; + for (int32_t token : output.committed_tokens) { + if (!advance(token)) return false; + } + return advance(output.token); +} + // Validate the model-neutral step protocol before the scheduler consumes any // output. Malformed row ownership is fatal because re-feeding a token after an // omitted output would silently corrupt that sequence. @@ -265,6 +330,7 @@ inline std::string validate_step_result( } std::vector decode_planned((size_t)slot_count, 0); + std::vector speculation_allowed((size_t)slot_count, 0); std::vector prefill_planned((size_t)slot_count, 0); for (const SeqEngine::StepInput & input : plan.decode) { if (input.slot < 0 || input.slot >= slot_count || input.token < 0) @@ -272,6 +338,8 @@ inline std::string validate_step_result( if (decode_planned[(size_t)input.slot]) return "decode plan contains a duplicate slot"; decode_planned[(size_t)input.slot] = 1; + speculation_allowed[(size_t)input.slot] = + input.allow_speculation ? 1 : 0; } for (const PrefillSlice & slice : plan.prefills) { if (slice.slot < 0 || slice.slot >= slot_count || @@ -290,10 +358,39 @@ inline std::string validate_step_result( return "decode output names an unplanned slot"; if (decode_seen[(size_t)output.slot]) return "step returned duplicate decode outputs"; - if (output.failed && (output.token >= 0 || output.error.empty())) - return "failed decode has invalid payload"; - if (!output.failed && (output.token < 0 || !output.error.empty())) - return "successful decode has invalid payload"; + if (output.failed) { + if (output.error.empty()) + return "failed decode has no diagnostic"; + if (output.token >= 0 || !output.committed_tokens.empty()) + return "failed decode exposes token payload"; + if (output.ddtree_suspensions != 0) + return "failed decode carries DDTree suspension telemetry"; + if (output.spec_steps != 0 || output.spec_accepted_tokens != 0) + return "failed decode carries chain speculation telemetry"; + if (output.spec_service_ar_steps != 0) + return "failed decode carries chain service telemetry"; + } else { + if (output.token < 0) + return "successful decode has no pending token"; + if (!output.error.empty()) + return "successful decode carries an error diagnostic"; + if (!speculation_allowed[(size_t)output.slot] && + !output.committed_tokens.empty()) + return "decode output burst violates disabled speculation"; + if (std::any_of( + output.committed_tokens.begin(), + output.committed_tokens.end(), + [](int32_t token) { return token < 0; })) + return "decode output burst contains an invalid token"; + if (output.ddtree_suspensions > output.ddtree_steps) + return "DDTree suspension has no successful DDTree step"; + if (output.spec_accepted_tokens != 0 && output.spec_steps == 0) + return "chain acceptance has no successful speculation step"; + if (output.spec_service_ar_steps > output.target_forwards) + return "chain service steps exceed target forwards"; + if (output.spec_service_ar_steps != 0 && output.spec_steps != 0) + return "step mixes chain speculation and AR service"; + } decode_seen[(size_t)output.slot] = 1; } @@ -307,6 +404,7 @@ inline std::string validate_step_result( return "step returned duplicate prefill outputs"; if (output.status != PrefillStatus::advanced && output.status != PrefillStatus::completed && + output.status != PrefillStatus::deferred && output.status != PrefillStatus::failed) return "prefill output has an unknown status"; if (output.status == PrefillStatus::advanced && @@ -315,6 +413,10 @@ inline std::string validate_step_result( if (output.status == PrefillStatus::completed && (output.token < 0 || !output.error.empty())) return "completed prefill has invalid payload"; + if (output.status == PrefillStatus::deferred && + (plan.decode.empty() || output.token >= 0 || + !output.error.empty())) + return "deferred prefill has invalid payload or no decode peer"; if (output.status == PrefillStatus::failed && (output.token >= 0 || output.error.empty())) return "failed prefill has invalid payload"; diff --git a/server/src/common/ddtree.cpp b/server/src/common/ddtree.cpp index 08ca33464..54402c2cc 100644 --- a/server/src/common/ddtree.cpp +++ b/server/src/common/ddtree.cpp @@ -223,4 +223,14 @@ std::vector follow_verified_tree(const DDTree & tree, return accepted; } +bool truncate_verified_path(std::vector & accepted, + std::size_t max_committed, + const int32_t * posterior, + int & out_next_token) { + if (accepted.size() <= max_committed) return false; + accepted.resize(max_committed); + out_next_token = accepted.empty() ? -1 : posterior[accepted.back()]; + return true; +} + } // namespace dflash::common diff --git a/server/src/common/ddtree.h b/server/src/common/ddtree.h index afe22f226..718c05b57 100644 --- a/server/src/common/ddtree.h +++ b/server/src/common/ddtree.h @@ -8,6 +8,7 @@ #pragma once +#include #include #include #include @@ -61,4 +62,13 @@ std::vector follow_verified_tree(const DDTree & tree, int & out_next_token, int * out_node_idx = nullptr); +// Bound a verified path to the number of tokens that can actually be +// committed. When truncation removes the old tip, the pending token must be +// recomputed from the posterior at the new tip; otherwise it describes model +// state that was never committed. +bool truncate_verified_path(std::vector & accepted, + std::size_t max_committed, + const int32_t * posterior, + int & out_next_token); + } // namespace dflash::common diff --git a/server/src/common/dflash2_batch.cpp b/server/src/common/dflash2_batch.cpp new file mode 100644 index 000000000..879564798 --- /dev/null +++ b/server/src/common/dflash2_batch.cpp @@ -0,0 +1,440 @@ +#include "dflash2_head.h" + +#include "dflash2_selector_validation.h" +#include "ddtree.h" +#include "geometric_draft_topk_cuda.h" +#include "ggml-alloc.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace dflash::common { +namespace { + +struct ProjectionGraph { + const DraftWeights * dw = nullptr; + ggml_backend_t backend = nullptr; + ggml_tensor * lm_head = nullptr; + int n_positions = 0; + std::vector arena; + ggml_context * ctx = nullptr; + ggml_cgraph * gf = nullptr; + ggml_gallocr_t galloc = nullptr; + ggml_tensor * inp_hidden = nullptr; + ggml_tensor * logits = nullptr; +}; + +struct BatchedSelectorGraph { + const DraftWeights * dw = nullptr; + ggml_backend_t backend = nullptr; + int n_lanes = 0; + int n_cand = 0; + int K = 0; + std::vector arena; + ggml_context * ctx = nullptr; + ggml_cgraph * gf = nullptr; + ggml_gallocr_t galloc = nullptr; + ggml_tensor * inp_hidden = nullptr; + ggml_tensor * inp_succ = nullptr; + ggml_tensor * inp_pred = nullptr; + ggml_tensor * hproj = nullptr; + ggml_tensor * succ = nullptr; + ggml_tensor * pred = nullptr; +}; + +ProjectionGraph & projection_graph() { + static thread_local ProjectionGraph graph; + return graph; +} + +BatchedSelectorGraph & batched_selector_graph() { + static thread_local BatchedSelectorGraph graph; + return graph; +} + +void free_projection_graph(ProjectionGraph & graph) { + if (graph.galloc) { + ggml_gallocr_free(graph.galloc); + graph.galloc = nullptr; + } + if (graph.ctx) { + ggml_free(graph.ctx); + graph.ctx = nullptr; + } + graph = {}; +} + +void free_selector_graph(BatchedSelectorGraph & graph) { + if (graph.galloc) { + ggml_gallocr_free(graph.galloc); + graph.galloc = nullptr; + } + if (graph.ctx) { + ggml_free(graph.ctx); + graph.ctx = nullptr; + } + graph = {}; +} + +bool ensure_projection_graph( + ProjectionGraph & graph, const DraftWeights & dw, + ggml_backend_t backend, ggml_tensor * lm_head, int n_positions) { + if (graph.ctx && graph.dw == &dw && graph.backend == backend && + graph.lm_head == lm_head && graph.n_positions == n_positions) { + return true; + } + free_projection_graph(graph); + if (!backend || !lm_head || n_positions <= 0 || dw.n_embd <= 0 || + lm_head->ne[0] != dw.n_embd || lm_head->ne[1] <= 0) { + return false; + } + + const size_t arena_size = + ggml_tensor_overhead() * 32 + + ggml_graph_overhead_custom(256, false) + 4096; + graph.arena.assign(arena_size, 0); + ggml_init_params params{}; + params.mem_size = graph.arena.size(); + params.mem_buffer = graph.arena.data(); + params.no_alloc = true; + graph.ctx = ggml_init(params); + if (!graph.ctx) return false; + graph.gf = ggml_new_graph_custom(graph.ctx, 256, false); + graph.inp_hidden = ggml_new_tensor_2d( + graph.ctx, GGML_TYPE_F32, dw.n_embd, n_positions); + ggml_set_input(graph.inp_hidden); + graph.logits = ggml_mul_mat(graph.ctx, lm_head, graph.inp_hidden); + ggml_set_output(graph.logits); + ggml_build_forward_expand(graph.gf, graph.logits); + graph.galloc = + ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!graph.galloc || !ggml_gallocr_alloc_graph(graph.galloc, graph.gf)) { + std::fprintf(stderr, + "dflash2_select_chains_batched: projection graph alloc failed\n"); + free_projection_graph(graph); + return false; + } + graph.dw = &dw; + graph.backend = backend; + graph.lm_head = lm_head; + graph.n_positions = n_positions; + return true; +} + +bool ensure_selector_graph( + BatchedSelectorGraph & graph, const DraftWeights & dw, + ggml_backend_t backend, int n_lanes, int n_cand, int K) { + if (graph.ctx && graph.dw == &dw && graph.backend == backend && + graph.n_lanes == n_lanes && graph.n_cand == n_cand && + graph.K == K) { + return true; + } + free_selector_graph(graph); + const DraftSelectorWeights & selector = dw.selector; + if (!backend || n_lanes <= 0 || n_cand <= 0 || K <= 0 || + dw.n_embd <= 0 || selector.rank <= 0 || !selector.hproj || + !selector.pred_cb || !selector.succ_cb) { + return false; + } + + const int n_positions = n_lanes * n_cand; + const int n_pred_rows = n_lanes + n_positions * K; + const size_t arena_size = + ggml_tensor_overhead() * 48 + + ggml_graph_overhead_custom(256, false) + 4096; + graph.arena.assign(arena_size, 0); + ggml_init_params params{}; + params.mem_size = graph.arena.size(); + params.mem_buffer = graph.arena.data(); + params.no_alloc = true; + graph.ctx = ggml_init(params); + if (!graph.ctx) return false; + graph.gf = ggml_new_graph_custom(graph.ctx, 256, false); + graph.inp_hidden = ggml_new_tensor_2d( + graph.ctx, GGML_TYPE_F32, dw.n_embd, n_positions); + graph.inp_succ = ggml_new_tensor_1d( + graph.ctx, GGML_TYPE_I32, n_positions * K); + graph.inp_pred = ggml_new_tensor_1d( + graph.ctx, GGML_TYPE_I32, n_pred_rows); + ggml_set_input(graph.inp_hidden); + ggml_set_input(graph.inp_succ); + ggml_set_input(graph.inp_pred); + graph.hproj = + ggml_mul_mat(graph.ctx, selector.hproj, graph.inp_hidden); + graph.succ = + ggml_get_rows(graph.ctx, selector.succ_cb, graph.inp_succ); + graph.pred = + ggml_get_rows(graph.ctx, selector.pred_cb, graph.inp_pred); + ggml_set_output(graph.hproj); + ggml_set_output(graph.succ); + ggml_set_output(graph.pred); + ggml_build_forward_expand(graph.gf, graph.hproj); + ggml_build_forward_expand(graph.gf, graph.succ); + ggml_build_forward_expand(graph.gf, graph.pred); + graph.galloc = + ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!graph.galloc || !ggml_gallocr_alloc_graph(graph.galloc, graph.gf)) { + std::fprintf(stderr, + "dflash2_select_chains_batched: selector graph alloc failed\n"); + free_selector_graph(graph); + return false; + } + graph.dw = &dw; + graph.backend = backend; + graph.n_lanes = n_lanes; + graph.n_cand = n_cand; + graph.K = K; + return true; +} + +DFlash2DepthSignal summarize_depth( + const float * log_probs, const std::vector & scores, + int K, int selected) { + DFlash2DepthSignal signal; + if (!log_probs || K <= 0 || selected < 0 || selected >= K || + static_cast(scores.size()) != K) { + return signal; + } + signal.selected_log_prob = log_probs[selected]; + signal.lm_top2_margin = K > 1 ? log_probs[0] - log_probs[1] + : std::numeric_limits::infinity(); + float top_k_mass = 0.0f; + for (int k = 0; k < K; ++k) top_k_mass += std::exp(log_probs[k]); + signal.top_k_mass = std::clamp(top_k_mass, 0.0f, 1.0f); + signal.selected_rank = selected; + signal.agrees_with_lm_top1 = selected == 0; + + float runner_up = -INFINITY; + for (int k = 0; k < K; ++k) { + if (k != selected) { + runner_up = std::max(runner_up, scores[(size_t) k]); + } + } + signal.selector_margin = K > 1 + ? scores[(size_t) selected] - runner_up + : std::numeric_limits::infinity(); + + const float max_score = + *std::max_element(scores.begin(), scores.end()); + float z = 0.0f; + for (float score : scores) z += std::exp(score - max_score); + if (z > 0.0f && std::isfinite(z)) { + signal.selector_winner_mass = + std::exp(scores[(size_t) selected] - max_score) / z; + float entropy = 0.0f; + for (float score : scores) { + const float p = std::exp(score - max_score) / z; + if (p > 0.0f) entropy -= p * std::log(p); + } + signal.selector_entropy = entropy; + } + return signal; +} + +} // namespace + +bool dflash2_select_chains_batched( + const DraftWeights & dw, + ggml_backend_t backend, + ggml_tensor * lm_head, + const std::vector & hidden_by_lane, + int q_len, + const std::vector & last_tokens, + std::vector> & draft_tokens, + std::vector * traces) { + draft_tokens.clear(); + if (traces) traces->clear(); + const DraftSelectorWeights & selector = dw.selector; + const int n_lanes = static_cast(hidden_by_lane.size()); + const int n_cand = q_len - 1; + const int K = selector.top_k; + const int rank = selector.rank; + const int hdim = dw.n_embd; + if (!selector.enabled || !selector.hproj || !selector.pred_cb || + !selector.succ_cb || !backend || !lm_head || n_lanes <= 0 || + static_cast(last_tokens.size()) != n_lanes || + n_cand <= 0 || K <= 0 || rank <= 0 || hdim <= 0) { + return false; + } + DFlash2SelectorLayout selector_layout; + selector_layout.rank = rank; + selector_layout.top_k = K; + selector_layout.hproj_rank = selector.hproj->ne[1]; + selector_layout.pred_rank = selector.pred_cb->ne[0]; + selector_layout.pred_vocab = selector.pred_cb->ne[1]; + selector_layout.succ_rank = selector.succ_cb->ne[0]; + selector_layout.succ_vocab = selector.succ_cb->ne[1]; + selector_layout.target_output_vocab = lm_head->ne[1]; + std::string selector_error; + if (!validate_dflash2_selector_layout( + selector_layout, selector_error)) { + std::fprintf(stderr, "dflash2_select_chains_batched: %s\n", + selector_error.c_str()); + return false; + } + for (const float * hidden : hidden_by_lane) { + if (!hidden) return false; + } + + const int n_positions = n_lanes * n_cand; + std::vector candidate_hidden( + (size_t) hdim * (size_t) n_positions); + for (int lane = 0; lane < n_lanes; ++lane) { + for (int depth = 0; depth < n_cand; ++depth) { + const int position = lane * n_cand + depth; + const float * source = hidden_by_lane[(size_t) lane] + + (size_t) (depth + 1) * (size_t) hdim; + std::memcpy( + candidate_hidden.data() + + (size_t) position * (size_t) hdim, + source, sizeof(float) * (size_t) hdim); + } + } + + ProjectionGraph & projection = projection_graph(); + if (!ensure_projection_graph( + projection, dw, backend, lm_head, n_positions)) { + return false; + } + ggml_backend_tensor_set( + projection.inp_hidden, candidate_hidden.data(), 0, + sizeof(float) * candidate_hidden.size()); + if (ggml_backend_graph_compute(backend, projection.gf) != + GGML_STATUS_SUCCESS) { + std::fprintf(stderr, + "dflash2_select_chains_batched: projection compute failed\n"); + return false; + } + + const int vocab = static_cast(lm_head->ne[1]); + std::vector candidate_log_probs( + (size_t) n_positions * (size_t) K); + std::vector candidate_ids( + (size_t) n_positions * (size_t) K); + bool have_top_k = false; +#ifdef DFLASH27B_HAVE_DRAFT_TOPK + static const bool gpu_top_k = []() { + const char * value = std::getenv("DFLASH_GPU_DRAFT_TOPK"); + return value == nullptr || value[0] != '0'; + }(); + if (gpu_top_k && projection.logits && projection.logits->data) { + have_top_k = geometric_extract_draft_topk_cuda( + projection.logits->data, n_positions, vocab, K, + candidate_log_probs.data(), candidate_ids.data(), 1.0f); + } +#endif + if (!have_top_k) { + std::vector logits( + (size_t) vocab * (size_t) n_positions); + ggml_backend_tensor_get( + projection.logits, logits.data(), 0, + sizeof(float) * logits.size()); + extract_draft_topk( + logits.data(), n_positions, vocab, K, + candidate_log_probs.data(), candidate_ids.data(), 1.0f); + } + + BatchedSelectorGraph & graph = batched_selector_graph(); + if (!ensure_selector_graph( + graph, dw, backend, n_lanes, n_cand, K)) { + return false; + } + std::vector predecessor_ids( + (size_t) n_lanes + candidate_ids.size()); + std::copy( + last_tokens.begin(), last_tokens.end(), predecessor_ids.begin()); + std::copy( + candidate_ids.begin(), candidate_ids.end(), + predecessor_ids.begin() + n_lanes); + ggml_backend_tensor_set( + graph.inp_hidden, candidate_hidden.data(), 0, + sizeof(float) * candidate_hidden.size()); + ggml_backend_tensor_set( + graph.inp_succ, candidate_ids.data(), 0, + sizeof(int32_t) * candidate_ids.size()); + ggml_backend_tensor_set( + graph.inp_pred, predecessor_ids.data(), 0, + sizeof(int32_t) * predecessor_ids.size()); + if (ggml_backend_graph_compute(backend, graph.gf) != + GGML_STATUS_SUCCESS) { + std::fprintf(stderr, + "dflash2_select_chains_batched: selector compute failed\n"); + return false; + } + + std::vector projected_hidden( + (size_t) rank * (size_t) n_positions); + std::vector successor_codes( + (size_t) rank * candidate_ids.size()); + std::vector predecessor_codes( + (size_t) rank * predecessor_ids.size()); + ggml_backend_tensor_get_async( + backend, graph.hproj, projected_hidden.data(), 0, + sizeof(float) * projected_hidden.size()); + ggml_backend_tensor_get_async( + backend, graph.succ, successor_codes.data(), 0, + sizeof(float) * successor_codes.size()); + ggml_backend_tensor_get_async( + backend, graph.pred, predecessor_codes.data(), 0, + sizeof(float) * predecessor_codes.size()); + ggml_backend_synchronize(backend); + + draft_tokens.assign( + (size_t) n_lanes, + std::vector((size_t) q_len)); + if (traces) traces->resize((size_t) n_lanes); + for (int lane = 0; lane < n_lanes; ++lane) { + draft_tokens[(size_t) lane][0] = last_tokens[(size_t) lane]; + if (traces) { + (*traces)[(size_t) lane].depths.reserve((size_t) n_cand); + } + int predecessor_row = lane; + for (int depth = 0; depth < n_cand; ++depth) { + const int position = lane * n_cand + depth; + const float * predecessor = predecessor_codes.data() + + (size_t) predecessor_row * (size_t) rank; + const float * hidden = projected_hidden.data() + + (size_t) position * (size_t) rank; + float best_score = -INFINITY; + int best_candidate = 0; + std::vector scores((size_t) K); + for (int candidate = 0; candidate < K; ++candidate) { + const int candidate_row = position * K + candidate; + const float * successor = successor_codes.data() + + (size_t) candidate_row * (size_t) rank; + float correction = 0.0f; + for (int r = 0; r < rank; ++r) { + correction += + predecessor[r] * hidden[r] * successor[r]; + } + const float score = + candidate_log_probs[(size_t) candidate_row] + + correction; + scores[(size_t) candidate] = score; + if (score > best_score) { + best_score = score; + best_candidate = candidate; + } + } + const int selected_row = position * K + best_candidate; + draft_tokens[(size_t) lane][(size_t) depth + 1] = + candidate_ids[(size_t) selected_row]; + if (traces) { + (*traces)[(size_t) lane].depths.push_back( + summarize_depth( + candidate_log_probs.data() + + (size_t) position * (size_t) K, + scores, K, best_candidate)); + } + predecessor_row = n_lanes + selected_row; + } + } + return true; +} + +} // namespace dflash::common diff --git a/server/src/common/dflash2_benefit.cpp b/server/src/common/dflash2_benefit.cpp new file mode 100644 index 000000000..98dc138ca --- /dev/null +++ b/server/src/common/dflash2_benefit.cpp @@ -0,0 +1,207 @@ +#include "dflash2_benefit.h" + +#include +#include +#include +#include +#include +#include + +namespace dflash::common { +namespace { + +// Seed artifacts (full hashes retained for offline provenance; startup uses the +// cheap size plus structural signature): target IQ4_XS sha256 +// 4e44edf892af6d57506fcd9eeaf5d0628f8737cfdce89652cb9d9bff82808eae, +// draft Q8_0 sha256 +// bb727abc583498aa4deea8b3cd0c34c2d96553954cbff25b5f7bdd469f0f1306. +constexpr DFlash2BenefitModelSignature kSeededQwen38DFlash2 = { + /*target_layers=*/64, + /*target_hidden=*/5120, + /*target_vocab=*/248320, + /*draft_layers=*/5, + /*draft_hidden=*/5120, + /*draft_block_size=*/8, + /*selector_rank=*/256, + /*selector_top_k=*/16, + /*selector_vocab=*/248320, + /*conv_kernel_size=*/2, + /*conv_group_size=*/16, + /*target_file_size=*/15195272800ULL, + /*draft_file_size=*/2045471776ULL, +}; + +bool same_signature(const DFlash2BenefitModelSignature & a, + const DFlash2BenefitModelSignature & b) { + return a.target_layers == b.target_layers && + a.target_hidden == b.target_hidden && + a.target_vocab == b.target_vocab && + a.draft_layers == b.draft_layers && + a.draft_hidden == b.draft_hidden && + a.draft_block_size == b.draft_block_size && + a.selector_rank == b.selector_rank && + a.selector_top_k == b.selector_top_k && + a.selector_vocab == b.selector_vocab && + a.conv_kernel_size == b.conv_kernel_size && + a.conv_group_size == b.conv_group_size && + a.target_file_size == b.target_file_size && + a.draft_file_size == b.draft_file_size; +} + +bool parse_finite_env(const char * name, double minimum, double maximum, + double & value, std::string & error) { + const char * text = std::getenv(name); + if (!text || !*text) return true; + errno = 0; + char * end = nullptr; + const double parsed = std::strtod(text, &end); + if (errno != 0 || end == text || !end || *end != '\0' || + !std::isfinite(parsed) || parsed < minimum || parsed > maximum) { + error = std::string(name) + " must be finite in [" + + std::to_string(minimum) + "," + std::to_string(maximum) + "]"; + return false; + } + value = parsed; + return true; +} + +void set_error(std::string * output, const std::string & value) { + if (output) *output = value; +} + +} // namespace + +std::string DFlash2BenefitModelSignature::str() const { + std::ostringstream out; + out << "target:l" << target_layers << ":h" << target_hidden + << ":v" << target_vocab + << "/draft:l" << draft_layers << ":h" << draft_hidden + << ":b" << draft_block_size + << "/selector:r" << selector_rank << ":k" << selector_top_k + << ":v" << selector_vocab + << "/conv:k" << conv_kernel_size << ":g" << conv_group_size + << "/files:t" << target_file_size << ":d" << draft_file_size; + return out.str(); +} + +DFlash2BenefitConfig DFlash2BenefitProvider::config_from_environment( + std::string & error) { + error.clear(); + DFlash2BenefitConfig config; + if (const char * version = + std::getenv("DFLASH_DFLASH2_BENEFIT_ADAPTER")) { + config.adapter_version = version; + } + if (!parse_finite_env( + "DFLASH_DFLASH2_BENEFIT_LM_WEIGHT", 0.0, 1.0, + config.lm_log_weight, error)) { + return config; + } + if (!parse_finite_env( + "DFLASH_DFLASH2_BENEFIT_HAZARD_SCALE", 0.0, 1.0, + config.hazard_scale, error) || config.hazard_scale <= 0.0) { + if (error.empty()) { + error = "DFLASH_DFLASH2_BENEFIT_HAZARD_SCALE must be in (0,1]"; + } + return config; + } + if (!parse_finite_env( + "DFLASH_DFLASH2_BENEFIT_YIELD_SCALE", 0.25, 4.0, + config.yield_scale, error)) { + return config; + } + return config; +} + +DFlash2BenefitProvider::DFlash2BenefitProvider( + DFlash2BenefitModelSignature model_signature, DFlash2BenefitConfig config) + : model_signature_(std::move(model_signature)), config_(std::move(config)) { + if (config_.adapter_version != kDFlash2BenefitAdapterVersion) { + error_ = "unsupported DFlash2 benefit adapter version '" + + config_.adapter_version + "'"; + return; + } + if (!std::isfinite(config_.lm_log_weight) || + config_.lm_log_weight < 0.0 || config_.lm_log_weight > 1.0 || + !std::isfinite(config_.hazard_scale) || + config_.hazard_scale <= 0.0 || config_.hazard_scale > 1.0 || + !std::isfinite(config_.yield_scale) || + config_.yield_scale < 0.25 || config_.yield_scale > 4.0) { + error_ = "invalid DFlash2 benefit coefficients"; + return; + } + if (!same_signature(model_signature_, kSeededQwen38DFlash2)) { + error_ = "unsupported DFlash2 model signature " + + model_signature_.str(); + } +} + +bool DFlash2BenefitProvider::estimate( + const DFlash2SelectorTrace & trace, int max_accept, + DFlash2BenefitEstimate & out, std::string * error) const { + out = {}; + if (!ready()) { + set_error(error, error_); + return false; + } + if (max_accept < 2 || max_accept > model_signature_.draft_block_size) { + set_error(error, "DFlash2 benefit depth is outside the seeded block"); + return false; + } + const size_t required = static_cast(max_accept - 1); + if (trace.depths.size() < required) { + set_error(error, "DFlash2 selector trace is missing required depths"); + return false; + } + + out.conditional_hazards.reserve(required); + double survival = 1.0; + double expected = 1.0; + const double selector_weight = 1.0 - config_.lm_log_weight; + for (size_t depth = 0; depth < required; ++depth) { + const DFlash2DepthSignal & signal = trace.depths[depth]; + if (!std::isfinite(signal.selected_log_prob) || + signal.selected_log_prob > 1e-6f || + !std::isfinite(signal.selector_winner_mass) || + signal.selector_winner_mass <= 0.0f || + signal.selector_winner_mass > 1.0f + 1e-6f) { + set_error(error, "DFlash2 selector trace contains invalid evidence"); + out = {}; + return false; + } + const double log_hazard = + config_.lm_log_weight * signal.selected_log_prob + + selector_weight * + std::log(std::min(1.0, signal.selector_winner_mass)); + const double hazard = std::clamp( + config_.hazard_scale * std::exp(log_hazard), 0.0, 1.0); + if (!std::isfinite(hazard)) { + set_error(error, "DFlash2 selector trace produced a nonfinite hazard"); + out = {}; + return false; + } + out.conditional_hazards.push_back(hazard); + survival *= hazard; + expected += survival; + } + out.expected_yield = std::clamp( + config_.yield_scale * expected, 1.0, + static_cast(max_accept)); + if (error) error->clear(); + return true; +} + +bool DFlash2BenefitProvider::publish_once( + const DFlash2SelectorTrace & trace, int max_accept, + double & destination, std::string * error) const { + if (std::isfinite(destination)) { + if (error) error->clear(); + return true; + } + DFlash2BenefitEstimate estimate_result; + if (!estimate(trace, max_accept, estimate_result, error)) return false; + destination = estimate_result.expected_yield; + return true; +} + +} // namespace dflash::common diff --git a/server/src/common/dflash2_benefit.h b/server/src/common/dflash2_benefit.h new file mode 100644 index 000000000..99b8b3d92 --- /dev/null +++ b/server/src/common/dflash2_benefit.h @@ -0,0 +1,94 @@ +#pragma once + +#include "dflash2_head.h" + +#include +#include +#include +#include + +namespace dflash::common { + +// Versioned structural signature of the Qwen3.8/DFlash2 pair on which +// the deliberately simple benefit heuristic below was empirically seeded. +// The architecture plus exact artifact sizes reject known-incompatible +// artifacts; they are not a cryptographic identity or a claim of held-out +// calibration. Deployment must preserve the seeded hashes recorded in the +// implementation, and any different target/selector needs its own offline fit. +struct DFlash2BenefitModelSignature { + int target_layers = 0; + int target_hidden = 0; + int target_vocab = 0; + int draft_layers = 0; + int draft_hidden = 0; + int draft_block_size = 0; + int selector_rank = 0; + int selector_top_k = 0; + int selector_vocab = 0; + int conv_kernel_size = 0; + int conv_group_size = 0; + uint64_t target_file_size = 0; + uint64_t draft_file_size = 0; + + std::string str() const; +}; + +inline constexpr const char * kDFlash2BenefitAdapterVersion = + "qwen38-dflash2-selector-benefit-v1"; + +struct DFlash2BenefitConfig { + std::string adapter_version = kDFlash2BenefitAdapterVersion; + + // Conditional acceptance hazard at each depth: + // exp(lm_log_weight * selected_log_prob + // + (1-lm_log_weight) * log(selector_winner_mass)) + // The selector is the better signal on the seed traces, while the + // LM term conservatively lowers a selector winner that has weak model + // probability. hazard_scale may only lower the estimate. + double lm_log_weight = 0.10; + double hazard_scale = 1.0; + // Offline per-adapter calibration; the generic gate never rescales yield. + double yield_scale = 1.0; +}; + +struct DFlash2BenefitEstimate { + double expected_yield = std::numeric_limits::quiet_NaN(); + std::vector conditional_hazards; +}; + +// Stateless request-local adapter. It does not learn from prior requests. +// Request lifetime is owned by the caller via publish_once(destination): the +// first valid trace fills an empty destination, and subsequent traces cannot +// overwrite that request's activation score. +class DFlash2BenefitProvider { +public: + DFlash2BenefitProvider( + DFlash2BenefitModelSignature model_signature, + DFlash2BenefitConfig config = {}); + + static DFlash2BenefitConfig config_from_environment( + std::string & error); + + bool ready() const { return error_.empty(); } + const std::string & error() const { return error_; } + const DFlash2BenefitModelSignature & model_signature() const { + return model_signature_; + } + const DFlash2BenefitConfig & config() const { return config_; } + const char * score_kind() const { return kDFlash2BenefitAdapterVersion; } + + bool estimate(const DFlash2SelectorTrace & trace, int max_accept, + DFlash2BenefitEstimate & out, + std::string * error = nullptr) const; + + bool publish_once(const DFlash2SelectorTrace & trace, int max_accept, + double & destination, + std::string * error = nullptr) const; + +private: + DFlash2BenefitModelSignature model_signature_; + DFlash2BenefitConfig config_; + std::string error_; +}; + +} // namespace dflash::common diff --git a/server/src/common/dflash2_head.cpp b/server/src/common/dflash2_head.cpp new file mode 100644 index 000000000..e44cda882 --- /dev/null +++ b/server/src/common/dflash2_head.cpp @@ -0,0 +1,227 @@ +#include "dflash2_head.h" + +#include "dflash2_selector_validation.h" +#include "ddtree.h" +#include "geometric_draft_topk_cuda.h" +#include "ggml-alloc.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace dflash::common { + +namespace { + +// Selector projection graph, built once per (drafter, backend, n_cand, K). +struct SelectorGraph { + const DraftWeights * dw = nullptr; + ggml_backend_t backend = nullptr; + int n_cand = 0; + int K = 0; + std::vector arena; + ggml_context * ctx = nullptr; + ggml_cgraph * gf = nullptr; + ggml_gallocr_t galloc = nullptr; + ggml_tensor * inp_hidden = nullptr; + ggml_tensor * inp_succ = nullptr; + ggml_tensor * inp_pred = nullptr; + ggml_tensor * hproj = nullptr; + ggml_tensor * succ = nullptr; + ggml_tensor * pred = nullptr; +}; + +SelectorGraph & selector_graph() { + static thread_local SelectorGraph g; + return g; +} + +void selector_graph_free(SelectorGraph & g) { + if (g.galloc) { ggml_gallocr_free(g.galloc); g.galloc = nullptr; } + if (g.ctx) { ggml_free(g.ctx); g.ctx = nullptr; } + g.gf = nullptr; + g.dw = nullptr; + g.n_cand = 0; + g.K = 0; +} + +DFlash2DepthSignal make_depth_signal( + const float * log_probs, const std::vector & scores, + int K, int selected) { + DFlash2DepthSignal signal; + if (!log_probs || K <= 0 || selected < 0 || selected >= K || + static_cast(scores.size()) != K) { + return signal; + } + signal.selected_log_prob = log_probs[selected]; + signal.lm_top2_margin = K > 1 ? log_probs[0] - log_probs[1] + : std::numeric_limits::infinity(); + float top_k_mass = 0.0f; + for (int k = 0; k < K; ++k) top_k_mass += std::exp(log_probs[k]); + signal.top_k_mass = std::clamp(top_k_mass, 0.0f, 1.0f); + signal.selected_rank = selected; + signal.agrees_with_lm_top1 = selected == 0; + + float runner_up = -INFINITY; + for (int k = 0; k < K; ++k) { + if (k != selected) runner_up = std::max(runner_up, scores[(size_t)k]); + } + signal.selector_margin = K > 1 + ? scores[(size_t)selected] - runner_up + : std::numeric_limits::infinity(); + + const float max_score = *std::max_element(scores.begin(), scores.end()); + float z = 0.0f; + for (float score : scores) z += std::exp(score - max_score); + if (z > 0.0f && std::isfinite(z)) { + signal.selector_winner_mass = + std::exp(scores[(size_t)selected] - max_score) / z; + float entropy = 0.0f; + for (float score : scores) { + const float p = std::exp(score - max_score) / z; + if (p > 0.0f) entropy -= p * std::log(p); + } + signal.selector_entropy = entropy; + } + return signal; +} + +} // namespace + +bool dflash2_select_chain(const DraftWeights & dw, + ggml_backend_t backend, + DFlashTarget & target, + const float * local_hidden, + int q_len, + int32_t last_tok, + std::vector & draft_tok, + DFlash2SelectorTrace * trace) { + const DraftSelectorWeights & sel = dw.selector; + if (!sel.enabled || !sel.hproj || !sel.pred_cb || !sel.succ_cb) return false; + if (q_len <= 1 || !local_hidden || !backend) return false; + const int hdim = dw.n_embd; + const int rank = sel.rank; + const int K = sel.top_k; + const int n_cand = q_len - 1; + if (hdim <= 0 || rank <= 0 || K <= 0) return false; + DFlash2SelectorLayout selector_layout; + selector_layout.rank = rank; + selector_layout.top_k = K; + selector_layout.hproj_rank = sel.hproj->ne[1]; + selector_layout.pred_rank = sel.pred_cb->ne[0]; + selector_layout.pred_vocab = sel.pred_cb->ne[1]; + selector_layout.succ_rank = sel.succ_cb->ne[0]; + selector_layout.succ_vocab = sel.succ_cb->ne[1]; + std::string selector_error; + if (!validate_dflash2_selector_layout( + selector_layout, selector_error)) { + std::fprintf(stderr, "dflash2_select_chain: %s\n", + selector_error.c_str()); + return false; + } + + // 1. Top-k candidates (log-probs) per block position through the target + // lm_head. Position 0 of local_hidden is the seed slot; candidates are + // rows 1 .. q_len-1. + std::vector cand_lp; + std::vector cand_ids; + if (!target.project_hidden_to_topk(local_hidden + (size_t)hdim, n_cand, K, /*temperature=*/1.0f, + cand_lp, cand_ids)) { + return false; + } + if (cand_lp.size() != (size_t)n_cand * K || cand_ids.size() != (size_t)n_cand * K) return false; + + // 2. One graph on the draft backend: hproj(h) for every candidate position, + // successor rows for every candidate, predecessor rows for the seed and + // every candidate (the path picks its predecessor among them). The + // graph shape only depends on (n_cand, K), so it is built once and + // reused across steps. + const int n_rows_pred = 1 + n_cand * K; + SelectorGraph & g = selector_graph(); + if (!g.ctx || g.dw != &dw || g.backend != backend || g.n_cand != n_cand || g.K != K) { + selector_graph_free(g); + const size_t arena_size = ggml_tensor_overhead() * 32 + ggml_graph_overhead() + 4096; + g.arena.assign(arena_size, 0); + ggml_init_params ip{}; + ip.mem_size = g.arena.size(); + ip.mem_buffer = g.arena.data(); + ip.no_alloc = true; + g.ctx = ggml_init(ip); + if (!g.ctx) return false; + g.gf = ggml_new_graph(g.ctx); + g.inp_hidden = ggml_new_tensor_2d(g.ctx, GGML_TYPE_F32, hdim, n_cand); + g.inp_succ = ggml_new_tensor_1d(g.ctx, GGML_TYPE_I32, n_cand * K); + g.inp_pred = ggml_new_tensor_1d(g.ctx, GGML_TYPE_I32, n_rows_pred); + ggml_set_input(g.inp_hidden); + ggml_set_input(g.inp_succ); + ggml_set_input(g.inp_pred); + g.hproj = ggml_mul_mat(g.ctx, sel.hproj, g.inp_hidden); // [rank, n_cand] + g.succ = ggml_get_rows(g.ctx, sel.succ_cb, g.inp_succ); // [rank, n_cand*K] f32 + g.pred = ggml_get_rows(g.ctx, sel.pred_cb, g.inp_pred); // [rank, 1+n_cand*K] f32 + ggml_set_output(g.hproj); + ggml_set_output(g.succ); + ggml_set_output(g.pred); + ggml_build_forward_expand(g.gf, g.hproj); + ggml_build_forward_expand(g.gf, g.succ); + ggml_build_forward_expand(g.gf, g.pred); + g.galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!g.galloc || !ggml_gallocr_alloc_graph(g.galloc, g.gf)) { + std::fprintf(stderr, "dflash2_select_chain: gallocr_alloc_graph failed\n"); + selector_graph_free(g); + return false; + } + g.dw = &dw; g.backend = backend; g.n_cand = n_cand; g.K = K; + } + + std::vector pred_ids((size_t)n_rows_pred); + pred_ids[0] = last_tok; + std::memcpy(pred_ids.data() + 1, cand_ids.data(), sizeof(int32_t) * (size_t)n_cand * K); + ggml_backend_tensor_set(g.inp_hidden, local_hidden + (size_t)hdim, 0, sizeof(float) * (size_t)hdim * n_cand); + ggml_backend_tensor_set(g.inp_succ, cand_ids.data(), 0, sizeof(int32_t) * (size_t)n_cand * K); + ggml_backend_tensor_set(g.inp_pred, pred_ids.data(), 0, sizeof(int32_t) * (size_t)n_rows_pred); + if (ggml_backend_graph_compute(backend, g.gf) != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "dflash2_select_chain: graph_compute failed\n"); + return false; + } + std::vector h_hproj((size_t)rank * n_cand); + std::vector h_succ((size_t)rank * n_cand * K); + std::vector h_pred((size_t)rank * n_rows_pred); + ggml_backend_tensor_get_async(backend, g.hproj, h_hproj.data(), 0, sizeof(float) * h_hproj.size()); + ggml_backend_tensor_get_async(backend, g.succ, h_succ.data(), 0, sizeof(float) * h_succ.size()); + ggml_backend_tensor_get_async(backend, g.pred, h_pred.data(), 0, sizeof(float) * h_pred.size()); + ggml_backend_synchronize(backend); + + // 3. Path search: greedy over the candidates, conditioned on the previous pick. + draft_tok.assign((size_t)q_len, last_tok); + if (trace) { + trace->depths.clear(); + trace->depths.reserve((size_t)n_cand); + } + int prev_row = 0; // row in h_pred: 0 = seed, 1 + i*K + k = candidate k of position i + for (int i = 0; i < n_cand; ++i) { + const float * pr = h_pred.data() + (size_t)prev_row * rank; + const float * hp = h_hproj.data() + (size_t)i * rank; + float best = -INFINITY; + int best_k = 0; + std::vector scores((size_t)K); + for (int k = 0; k < K; ++k) { + const float * sc = h_succ.data() + ((size_t)i * K + k) * rank; + float dot = 0.0f; + for (int r = 0; r < rank; ++r) dot += pr[r] * hp[r] * sc[r]; + const float score = cand_lp[(size_t)i * K + k] + dot; + scores[(size_t)k] = score; + if (score > best) { best = score; best_k = k; } + } + draft_tok[(size_t)i + 1] = cand_ids[(size_t)i * K + best_k]; + if (trace) trace->depths.push_back(make_depth_signal( + cand_lp.data() + (size_t)i * K, scores, K, best_k)); + prev_row = 1 + i * K + best_k; + } + return true; +} + +} // namespace dflash::common diff --git a/server/src/common/dflash2_head.h b/server/src/common/dflash2_head.h new file mode 100644 index 000000000..3b47bc2aa --- /dev/null +++ b/server/src/common/dflash2_head.h @@ -0,0 +1,65 @@ +#pragma once + +#include "dflash_target.h" +#include "internal.h" + +#include +#include + +namespace dflash::common { + +// Raw selector diagnostics for one proposed depth. These are deliberately +// not called confidence: the DFlash2 selector is trained to rank its top-K +// candidates, not to emit calibrated target-acceptance probabilities. An +// offline, model-specific adapter may later map these values to survival +// probabilities for the adaptive gate. +struct DFlash2DepthSignal { + float selected_log_prob = 0.0f; + float lm_top2_margin = 0.0f; + float top_k_mass = 0.0f; + int selected_rank = 0; + bool agrees_with_lm_top1 = false; + float selector_margin = 0.0f; + float selector_winner_mass = 0.0f; + float selector_entropy = 0.0f; +}; + +struct DFlash2SelectorTrace { + std::vector depths; +}; + +// DFlash 2 candidate selector for greedy chain drafting. +// +// For every drafted block position the target lm_head logits are reduced to +// the selector's top-k candidates (log-probs, so per-position constants do +// not matter for the argmax), then one path is traced through them: +// score(c) = logp(c) + < pred_cb[prev] * hproj(h_pos), succ_cb[c] > +// prev = argmax_c score(c) +// starting from the block seed `last_tok`. Runs the projections (hproj GEMV +// and codebook row gathers) in one small graph on `backend`, the k-way path +// search on the host. Fills draft_tok = [last_tok, tok_1 .. tok_{q_len-1}]. +bool dflash2_select_chain(const DraftWeights & dw, + ggml_backend_t backend, + DFlashTarget & target, + const float * local_hidden, + int q_len, + int32_t last_tok, + std::vector & draft_tok, + DFlash2SelectorTrace * trace = nullptr); + +// Same selector, batched over host-resident drafter hidden blocks and using a +// local target lm_head tensor. The expensive lm_head projection covers every +// (lane, depth) in one graph, GPU top-K is invoked once, and selector +// projections/readback are shared across the cohort. This is the concurrent +// paged-engine entry point; no non-paged DFlashTarget adapter is required. +bool dflash2_select_chains_batched( + const DraftWeights & dw, + ggml_backend_t backend, + ggml_tensor * lm_head, + const std::vector & hidden_by_lane, + int q_len, + const std::vector & last_tokens, + std::vector> & draft_tokens, + std::vector * traces = nullptr); + +} // namespace dflash::common diff --git a/server/src/common/dflash2_selector_validation.h b/server/src/common/dflash2_selector_validation.h new file mode 100644 index 000000000..6bcd50403 --- /dev/null +++ b/server/src/common/dflash2_selector_validation.h @@ -0,0 +1,93 @@ +#pragma once + +#include "geometric_draft_topk_cuda.h" + +#include +#include + +namespace dflash::common { + +// Host-only description of the selector tensors. Keeping validation in terms +// of dimensions makes it usable both while GGUF tensor descriptors are being +// loaded and when a concrete target lm_head is attached to the batched path. +struct DFlash2SelectorLayout { + int rank = 0; + int top_k = 0; + int64_t hproj_rank = 0; + int64_t pred_rank = 0; + int64_t pred_vocab = 0; + int64_t succ_rank = 0; + int64_t succ_vocab = 0; + // Zero means that source is unavailable at this validation point. A + // partial target shard, for example, can declare n_vocab without owning + // the final output tensor; the concrete lm_head is checked again at use. + int64_t target_output_vocab = 0; + int64_t target_declared_vocab = 0; +}; + +inline bool validate_dflash2_selector_layout( + const DFlash2SelectorLayout & layout, std::string & error) { + error.clear(); + if (layout.rank <= 0) { + error = "DFlash 2 selector rank must be positive (got " + + std::to_string(layout.rank) + ")"; + return false; + } + if (!geometric_draft_topk_cuda_supports_k(layout.top_k)) { + error = "DFlash 2 selector top_k=" + std::to_string(layout.top_k) + + " is unsupported; expected one of 1..8, 12, or 16"; + return false; + } + if (layout.hproj_rank != layout.rank || + layout.pred_rank != layout.rank || + layout.succ_rank != layout.rank) { + error = "DFlash 2 selector rank mismatch: metadata=" + + std::to_string(layout.rank) + " hproj=" + + std::to_string(layout.hproj_rank) + " pred_cb=" + + std::to_string(layout.pred_rank) + " succ_cb=" + + std::to_string(layout.succ_rank); + return false; + } + if (layout.pred_vocab <= 0 || layout.succ_vocab <= 0) { + error = "DFlash 2 selector codebook vocab must be positive: pred_cb=" + + std::to_string(layout.pred_vocab) + " succ_cb=" + + std::to_string(layout.succ_vocab); + return false; + } + if (layout.pred_vocab != layout.succ_vocab) { + error = "DFlash 2 selector codebook vocab mismatch: pred_cb=" + + std::to_string(layout.pred_vocab) + " succ_cb=" + + std::to_string(layout.succ_vocab); + return false; + } + if (layout.top_k > layout.pred_vocab) { + error = "DFlash 2 selector top_k=" + std::to_string(layout.top_k) + + " exceeds codebook vocab=" + std::to_string(layout.pred_vocab); + return false; + } + if (layout.target_output_vocab > 0 && + layout.target_declared_vocab > 0 && + layout.target_output_vocab != layout.target_declared_vocab) { + error = "DFlash 2 target vocab mismatch: output/lm_head=" + + std::to_string(layout.target_output_vocab) + " target.n_vocab=" + + std::to_string(layout.target_declared_vocab); + return false; + } + if (layout.target_output_vocab > 0 && + layout.pred_vocab != layout.target_output_vocab) { + error = "DFlash 2 selector vocab mismatch: codebook=" + + std::to_string(layout.pred_vocab) + " target output/lm_head=" + + std::to_string(layout.target_output_vocab); + return false; + } + if (layout.target_declared_vocab > 0 && + layout.pred_vocab != layout.target_declared_vocab) { + error = "DFlash 2 selector vocab mismatch: codebook=" + + std::to_string(layout.pred_vocab) + " target.n_vocab=" + + std::to_string(layout.target_declared_vocab); + return false; + } + return true; +} + +} // namespace dflash::common diff --git a/server/src/common/dflash_draft_kv.cpp b/server/src/common/dflash_draft_kv.cpp index 05ec25aca..1ff4cb531 100644 --- a/server/src/common/dflash_draft_kv.cpp +++ b/server/src/common/dflash_draft_kv.cpp @@ -96,9 +96,12 @@ bool draft_kv_init(DraftKvState & st, si.mask_swa = st.mask_swa; si.lm_head = lm_head; DraftGraphOutputs go = build_draft_kv_step(st.g_ctx, st.gf, dw, st.cache, si); - if (!go.hidden_states) return false; + if (!go.hidden_prenorm || !go.hidden_states) return false; + st.hidden_prenorm = go.hidden_prenorm; st.hidden_states = go.hidden_states; st.logits = go.logits; + ggml_set_output(st.hidden_prenorm); + ggml_build_forward_expand(st.gf, st.hidden_prenorm); ggml_set_output(st.hidden_states); ggml_build_forward_expand(st.gf, st.hidden_states); if (st.logits) { @@ -142,7 +145,7 @@ void draft_kv_free(DraftKvState & st) { if (st.mem_ctx) { ggml_free(st.mem_ctx); st.mem_ctx = nullptr; } st.meta_arena.clear(); st.meta_arena.shrink_to_fit(); - st.hidden_states = st.logits = nullptr; + st.hidden_prenorm = st.hidden_states = st.logits = nullptr; st.cache.k.clear(); st.cache.v.clear(); st.slot_pos.clear(); @@ -318,4 +321,173 @@ bool draft_kv_begin_step(DraftKvState & st, return true; } +void draft_kv_batch_free(DraftKvBatchGraph & batch) { + if (batch.galloc) { + ggml_gallocr_free(batch.galloc); + batch.galloc = nullptr; + } + if (batch.g_ctx) { + ggml_free(batch.g_ctx); + batch.g_ctx = nullptr; + } + batch.gf = nullptr; + batch.hidden_by_lane.clear(); + batch.prenorm_by_lane.clear(); + batch.lane_states.clear(); + batch.meta_arena.clear(); + batch.n_lanes = 0; + batch.q_len = 0; + batch.outputs_prenorm = false; + batch.built_for = nullptr; +} + +static bool draft_kv_batch_build( + DraftKvBatchGraph & batch, + const DraftWeights & dw, + ggml_backend_t backend, + const std::vector & lane_states, + bool need_prenorm) { + if (!backend || lane_states.empty() || dw.block_size <= 1) { + return false; + } + for (DraftKvState * state : lane_states) { + if (!state || !state->mem_buf || state->q_len != dw.block_size || + state->built_for != static_cast(&dw)) { + return false; + } + } + + draft_kv_batch_free(batch); + const int n_lanes = static_cast(lane_states.size()); + const size_t arena_size = + (32u + 16u * static_cast(n_lanes)) * 1024u * 1024u; + batch.meta_arena.resize(arena_size); + ggml_init_params params{}; + params.mem_size = batch.meta_arena.size(); + params.mem_buffer = batch.meta_arena.data(); + params.no_alloc = true; + batch.g_ctx = ggml_init(params); + if (!batch.g_ctx) { + draft_kv_batch_free(batch); + return false; + } + batch.gf = ggml_new_graph_custom( + batch.g_ctx, 4096 * n_lanes + 2048, false); + + batch.hidden_by_lane.reserve(static_cast(n_lanes)); + if (need_prenorm) { + batch.prenorm_by_lane.reserve(static_cast(n_lanes)); + } + for (DraftKvState * state : lane_states) { + DraftKvAppendInputs append{}; + append.n_rows = state->a_step; + append.feat = state->ap_feat; + append.positions = state->ap_pos; + append.rows = state->ap_rows; + if (!build_draft_kv_append( + batch.g_ctx, batch.gf, dw, state->cache, append)) { + draft_kv_batch_free(batch); + return false; + } + + DraftKvStepInputs step{}; + step.noise_embed = state->inp_embed; + step.positions_q = state->pos_q; + step.noise_rows = state->noise_rows; + step.mask_full = state->mask_full; + step.mask_swa = state->mask_swa; + DraftGraphOutputs output = build_draft_kv_step( + batch.g_ctx, batch.gf, dw, state->cache, step); + if (!output.hidden_states || + (need_prenorm && !output.hidden_prenorm)) { + draft_kv_batch_free(batch); + return false; + } + ggml_set_output(output.hidden_states); + ggml_build_forward_expand(batch.gf, output.hidden_states); + batch.hidden_by_lane.push_back(output.hidden_states); + if (need_prenorm) { + ggml_set_output(output.hidden_prenorm); + ggml_build_forward_expand(batch.gf, output.hidden_prenorm); + batch.prenorm_by_lane.push_back(output.hidden_prenorm); + } + } + + batch.galloc = + ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!batch.galloc || + !ggml_gallocr_alloc_graph(batch.galloc, batch.gf)) { + std::fprintf(stderr, + "[draft-kv-batch] graph alloc failed lanes=%d\n", n_lanes); + draft_kv_batch_free(batch); + return false; + } + + batch.n_lanes = n_lanes; + batch.q_len = dw.block_size; + batch.outputs_prenorm = need_prenorm; + batch.built_for = &dw; + batch.lane_states = lane_states; + std::fprintf(stderr, + "[draft-kv-batch] packed backbone ready lanes=%d q_len=%d " + "prenorm=%s\n", + n_lanes, dw.block_size, need_prenorm ? "on" : "off"); + return true; +} + +bool draft_kv_batch_compute( + DraftKvBatchGraph & batch, + const DraftWeights & dw, + ggml_backend_t backend, + const std::vector & lane_states, + bool need_prenorm, + std::vector> & hidden_by_lane, + std::vector> & prenorm_by_lane) { + hidden_by_lane.clear(); + prenorm_by_lane.clear(); + if (lane_states.empty()) return false; + + const bool reusable = + batch.gf && batch.built_for == static_cast(&dw) && + batch.lane_states == lane_states && + batch.outputs_prenorm == need_prenorm; + if (!reusable && + !draft_kv_batch_build( + batch, dw, backend, lane_states, need_prenorm)) { + return false; + } + if (ggml_backend_graph_compute(backend, batch.gf) != + GGML_STATUS_SUCCESS) { + std::fprintf(stderr, + "[draft-kv-batch] graph compute failed lanes=%d\n", + batch.n_lanes); + return false; + } + + const size_t elements = + static_cast(dw.n_embd) * static_cast(batch.q_len); + hidden_by_lane.assign( + static_cast(batch.n_lanes), + std::vector(elements)); + if (need_prenorm) { + prenorm_by_lane.assign( + static_cast(batch.n_lanes), + std::vector(elements)); + } + for (int lane = 0; lane < batch.n_lanes; ++lane) { + ggml_backend_tensor_get_async( + backend, batch.hidden_by_lane[static_cast(lane)], + hidden_by_lane[static_cast(lane)].data(), 0, + sizeof(float) * elements); + if (need_prenorm) { + ggml_backend_tensor_get_async( + backend, batch.prenorm_by_lane[static_cast(lane)], + prenorm_by_lane[static_cast(lane)].data(), 0, + sizeof(float) * elements); + } + } + ggml_backend_synchronize(backend); + return true; +} + } // namespace dflash::common diff --git a/server/src/common/dflash_draft_kv.h b/server/src/common/dflash_draft_kv.h index 888f9b46b..fb9c46f43 100644 --- a/server/src/common/dflash_draft_kv.h +++ b/server/src/common/dflash_draft_kv.h @@ -64,6 +64,7 @@ struct DraftKvState { ggml_context * g_ctx = nullptr; ggml_cgraph * gf = nullptr; ggml_gallocr_t galloc = nullptr; + ggml_tensor * hidden_prenorm = nullptr; ggml_tensor * hidden_states = nullptr; ggml_tensor * logits = nullptr; // iff lm_head passed at init @@ -99,4 +100,37 @@ bool draft_kv_begin_step(DraftKvState & st, const DraftFeatureMirror & ring, int committed); +struct DraftKvBatchGraph { + DraftKvBatchGraph() = default; + DraftKvBatchGraph(const DraftKvBatchGraph &) = delete; + DraftKvBatchGraph & operator=(const DraftKvBatchGraph &) = delete; + + int n_lanes = 0; + int q_len = 0; + bool outputs_prenorm = false; + const void * built_for = nullptr; + std::vector lane_states; + + std::vector meta_arena; + ggml_context * g_ctx = nullptr; + ggml_cgraph * gf = nullptr; + ggml_gallocr_t galloc = nullptr; + std::vector hidden_by_lane; + std::vector prenorm_by_lane; +}; + +void draft_kv_batch_free(DraftKvBatchGraph & batch); + +// All lane states must already have draft_kv_begin_step() inputs and +// inp_embed uploaded. The packed graph only computes the shared backbone; +// adapters consume the returned per-lane host views. +bool draft_kv_batch_compute( + DraftKvBatchGraph & batch, + const DraftWeights & dw, + ggml_backend_t backend, + const std::vector & lane_states, + bool need_prenorm, + std::vector> & hidden_by_lane, + std::vector> & prenorm_by_lane); + } // namespace dflash::common diff --git a/server/src/common/dspark_head.cpp b/server/src/common/dspark_head.cpp index f0df52c10..4146ceacc 100644 --- a/server/src/common/dspark_head.cpp +++ b/server/src/common/dspark_head.cpp @@ -307,6 +307,120 @@ bool build_markov_chain_graph(const DraftWeights & dw, } // namespace +bool build_dspark_markov_batched_chain( + ggml_context * ctx, + ggml_cgraph * gf, + const DraftWeights & dw, + ggml_tensor * lm_head, + const std::vector & hidden_by_lane, + const std::vector & prenorm_by_lane, + ggml_tensor * seed_tokens, + int q_len, + bool want_confidence, + DSparkBatchedChainOutputs & out) { + out = {}; + const int n_lanes = static_cast(hidden_by_lane.size()); + const int hdim = dw.n_embd; + if (!ctx || !gf || !lm_head || !seed_tokens || n_lanes <= 0 || + q_len <= 1 || hdim <= 0 || !dw.dspark.enabled || + !dw.dspark.markov_w1 || !dw.dspark.markov_w2 || + seed_tokens->ne[0] != n_lanes || + prenorm_by_lane.size() != hidden_by_lane.size()) { + return false; + } + const int vocab = static_cast(lm_head->ne[1]); + if (vocab <= 0 || + (dw.dspark.vocab_size > 0 && vocab != dw.dspark.vocab_size)) { + return false; + } + for (int lane = 0; lane < n_lanes; ++lane) { + ggml_tensor * hidden = hidden_by_lane[(size_t)lane]; + ggml_tensor * prenorm = prenorm_by_lane[(size_t)lane]; + if (!hidden || hidden->ne[0] != hdim || hidden->ne[1] < q_len || + !prenorm || prenorm->ne[0] != hdim || prenorm->ne[1] < q_len) { + return false; + } + } + + const bool have_confidence = want_confidence && + dw.dspark.confidence_w && dw.dspark.confidence_b && + (dw.dspark.confidence_dim == hdim || + dw.dspark.confidence_dim == hdim + dw.dspark.markov_rank); + const int n_depths = q_len - 1; + std::vector hidden_by_depth((size_t)n_depths); + std::vector confidence_by_depth((size_t)n_depths); + + auto concat_lane_column = [&](const std::vector & sources, + int column) -> ggml_tensor * { + ggml_tensor * packed = nullptr; + for (ggml_tensor * source : sources) { + ggml_tensor * lane = ggml_view_2d( + ctx, source, hdim, 1, source->nb[1], + (size_t)column * source->nb[1]); + packed = packed ? ggml_concat(ctx, packed, lane, 1) : lane; + } + return packed; + }; + + ggml_tensor * all_hidden = nullptr; + for (int depth = 0; depth < n_depths; ++depth) { + hidden_by_depth[(size_t)depth] = + concat_lane_column(hidden_by_lane, depth + 1); + confidence_by_depth[(size_t)depth] = + concat_lane_column(prenorm_by_lane, depth + 1); + if (!hidden_by_depth[(size_t)depth] || + !confidence_by_depth[(size_t)depth]) { + return false; + } + all_hidden = all_hidden + ? ggml_concat(ctx, all_hidden, hidden_by_depth[(size_t)depth], 1) + : hidden_by_depth[(size_t)depth]; + } + + // Depth-major layout makes each Markov step a contiguous [vocab, lanes] + // view while retaining one lm_head projection for the whole cohort. + ggml_tensor * base = ggml_mul_mat(ctx, lm_head, all_hidden); + ggml_tensor * prev_ids = seed_tokens; + out.n_lanes = n_lanes; + out.q_len = q_len; + out.tokens.assign((size_t)n_depths, nullptr); + out.confidence.assign((size_t)n_depths, nullptr); + + for (int depth = 0; depth < n_depths; ++depth) { + ggml_tensor * prev_emb = + ggml_get_rows(ctx, dw.dspark.markov_w1, prev_ids); + ggml_tensor * bias = + ggml_mul_mat(ctx, dw.dspark.markov_w2, prev_emb); + ggml_tensor * base_depth = ggml_view_2d( + ctx, base, vocab, n_lanes, base->nb[1], + (size_t)depth * (size_t)n_lanes * base->nb[1]); + ggml_tensor * corrected = ggml_add(ctx, base_depth, bias); + ggml_tensor * tok = ggml_argmax(ctx, corrected); + ggml_set_output(tok); + ggml_build_forward_expand(gf, tok); + out.tokens[(size_t)depth] = tok; + + if (have_confidence) { + ggml_tensor * conf_in = confidence_by_depth[(size_t)depth]; + if (dw.dspark.confidence_dim == + hdim + dw.dspark.markov_rank) { + conf_in = ggml_concat(ctx, conf_in, prev_emb, 0); + } + ggml_tensor * conf = + ggml_mul_mat(ctx, dw.dspark.confidence_w, conf_in); + conf = ggml_add( + ctx, conf, + ggml_reshape_2d(ctx, dw.dspark.confidence_b, 1, 1)); + conf = ggml_sigmoid(ctx, conf); + ggml_set_output(conf); + ggml_build_forward_expand(gf, conf); + out.confidence[(size_t)depth] = conf; + } + prev_ids = tok; + } + return true; +} + bool dspark_markov_correct_greedy_chain_fused(const DraftWeights & dw, ggml_backend_t backend, ggml_tensor * lm_head, diff --git a/server/src/common/dspark_head.h b/server/src/common/dspark_head.h index 9b97b261d..6259482e0 100644 --- a/server/src/common/dspark_head.h +++ b/server/src/common/dspark_head.h @@ -36,6 +36,31 @@ bool dspark_markov_correct_greedy_chain_fused(const DraftWeights & dw, std::vector * confidence_out = nullptr, const float * confidence_hidden = nullptr); +// Outputs embedded in a caller-owned graph for a lane-batched DSpark chain. +// Each depth tensor is shaped [n_lanes], with confidence in [1, n_lanes]. +struct DSparkBatchedChainOutputs { + int n_lanes = 0; + int q_len = 0; + std::vector tokens; + std::vector confidence; +}; + +// Append a depth-major, multi-lane Markov chain to an existing draft graph. +// The lane backbones remain independent; their hidden tensors stay on-device. +// One lm_head matmul covers every (depth, lane), then each depth performs a +// batched Markov lookup, correction, argmax, and calibrated confidence head. +bool build_dspark_markov_batched_chain( + ggml_context * ctx, + ggml_cgraph * gf, + const DraftWeights & dw, + ggml_tensor * lm_head, + const std::vector & hidden_by_lane, + const std::vector & prenorm_by_lane, + ggml_tensor * seed_tokens, + int q_len, + bool want_confidence, + DSparkBatchedChainOutputs & out); + // DDTree candidate generation with the Markov correction: base logits for // all n_tokens positions in ONE lm_head matmul; rows 1..n-1 get the low-rank // previous-token bias chained along the main (argmax) path; top-K extracted diff --git a/server/src/common/feature_gate.cpp b/server/src/common/feature_gate.cpp index f655f5d4f..dfbd1e604 100644 --- a/server/src/common/feature_gate.cpp +++ b/server/src/common/feature_gate.cpp @@ -163,6 +163,30 @@ std::string check_feature_compatibility( "' does not support PFlash compression"; } + const bool concurrent_paged_qwen = + arch == "qwen35" && args.paged_attention && + args.max_concurrency > 1 && + !args.device.is_layer_split() && + !args.remote_target_shard.enabled() && + args.fa_window == 0; + const bool concurrent_local_paged_qwen = + concurrent_paged_qwen && !args.remote_draft.enabled() && + target_backend == draft_backend && + args.device.gpu == args.draft_device.gpu && + !args.device.is_tensor_parallel(); + const bool concurrent_local_ddtree = + concurrent_local_paged_qwen && args.draft_path != nullptr && + args.ddtree_mode; + const bool concurrent_local_chain = + concurrent_local_paged_qwen && args.draft_path != nullptr && + !args.ddtree_mode && + args.speculation_policy != SpeculationPolicy::Never; + + if (args.ddtree_mode && + (args.ddtree_budget < 1 || args.ddtree_budget > 255)) { + return "--ddtree-budget must be in [1, 255]"; + } + // ── --paged-attention × architecture, placement, and decode features // Paged decode swaps the contiguous K/V cache for a block table owned by // the monolithic qwen35 backend, so every rule below is about reaching @@ -180,20 +204,30 @@ std::string check_feature_compatibility( args.remote_target_shard.enabled()) { return "--paged-attention requires one local target device"; } - if (args.draft_path != nullptr || args.remote_draft.enabled() || - args.ddtree_mode) { - return "--paged-attention requires autoregressive decode without a " - "draft or DDTree"; + if (args.remote_draft.enabled() && + args.speculation_policy != SpeculationPolicy::Never) { + return "concurrent paged DDTree requires a local draft on the target device"; + } + if (args.ddtree_mode && !concurrent_local_ddtree) { + return "paged DDTree requires concurrent local target/draft execution " + "on one device"; + } + if (args.draft_path != nullptr && !args.ddtree_mode && + args.speculation_policy != SpeculationPolicy::Never && + !concurrent_local_chain) { + return "paged chain speculation requires concurrent local target/draft " + "execution on one device"; } if (args.fa_window != 0) { return "--paged-attention requires full attention (--fa-window 0)"; } - if (features.pflash_enabled) { - return "--paged-attention cannot be combined with PFlash prefill " - "compression"; + if (features.pflash_enabled && !concurrent_local_paged_qwen) { + return "paged PFlash prefill compression requires concurrent local " + "Qwen3.5/Qwen3.6 serving on one target/draft device"; } - if (features.kvflash_enabled) { - return "--paged-attention cannot be combined with KVFlash"; + if (features.kvflash_enabled && !concurrent_paged_qwen) { + return "paged KVFlash requires concurrent local Qwen3.5/Qwen3.6 " + "serving with full attention"; } // The pool rounds max_ctx up to a whole number of blocks, so the top // of the range is what can be rounded without overflowing int. @@ -232,10 +266,18 @@ std::string check_feature_compatibility( if (args.max_concurrency <= 1) { return "--kv-pool-tokens requires --max-concurrency greater than 1"; } - // The cache appends one scratch block after the physical pool, and - // the requested pool itself is rounded up to a whole block. Cap the - // request at the largest aligned pool that leaves room for scratch. - const int64_t max_pool_tokens = paged_kv_address_cap(); + // Reserve the dead-row block plus one rounded DDTree candidate slab + // per slot, exactly matching Qwen35Backend's cache allocation. + const int64_t tree_scratch = concurrent_local_ddtree + ? (int64_t)args.max_concurrency * + paged_token_capacity(args.ddtree_budget + 1) + : concurrent_local_chain + ? (int64_t)args.max_concurrency * paged_token_capacity(16) + : 0; + const int64_t scratch_tokens = PAGED_BLOCK_SIZE + tree_scratch; + const int64_t max_pool_tokens = + ((int64_t)INT32_MAX - scratch_tokens) / PAGED_BLOCK_SIZE * + PAGED_BLOCK_SIZE; if (args.kv_pool_tokens < PAGED_BLOCK_SIZE || args.kv_pool_tokens > max_pool_tokens) { return "--kv-pool-tokens must be in [" + diff --git a/server/src/common/geometric_draft_topk_cuda.cu b/server/src/common/geometric_draft_topk_cuda.cu index 71086c98a..45a9477f2 100644 --- a/server/src/common/geometric_draft_topk_cuda.cu +++ b/server/src/common/geometric_draft_topk_cuda.cu @@ -13,7 +13,7 @@ namespace dflash::common { namespace { -constexpr int kMaxK = 8; // ddtree_K is 8 in practice; K>kMaxK → CPU fallback +constexpr int kMaxK = 16; // largest instantiated K; supported set is declared in the header constexpr int kBlock = 256; // threads per block (power of two for the reduction) constexpr int kMaxSplit = 128; // max vocab splits per position (combine-block cap) @@ -331,7 +331,13 @@ bool geometric_extract_draft_topk_cuda(const void * d_logits, float * out_log_probs, int32_t * out_token_ids, float temperature) { - if (!d_logits || n_positions <= 0 || vocab <= 0 || K <= 0 || K > kMaxK) return false; + // Reject before touching CUDA or scratch. In particular, K values in the + // holes between instantiated templates (9-11 and 13-15) must fall back to + // the CPU path instead of copying stale data from a previous invocation. + if (!d_logits || !out_log_probs || !out_token_ids || n_positions <= 0 || + vocab <= 0 || K > vocab || !geometric_draft_topk_cuda_supports_k(K)) { + return false; + } cudaPointerAttributes attr{}; if (cudaPointerGetAttributes(&attr, d_logits) != cudaSuccess) { @@ -362,9 +368,11 @@ bool geometric_extract_draft_topk_cuda(const void * d_logits, // the tensor base aligned and a vocab stride that is a multiple of 4. const bool use_vec = (vocab % 4 == 0) && (reinterpret_cast(lp_in) % 16 == 0); + bool dispatched = false; // K (and the vectorization flag) are compile-time template parameters // so the per-thread/per-partial top-K stays register-resident; dispatch - // the runtime K to its instantiation. K>kMaxK is already rejected above. + // the runtime K to its instantiation. Unsupported K is rejected before + // scratch allocation above, so the default is unreachable hardening. #define DFLASH_TOPK_LAUNCH(KV, VEC) \ geometric_draft_topk_partial<<>>( \ lp_in, vocab, inv_t, split, \ @@ -374,19 +382,22 @@ bool geometric_extract_draft_topk_cuda(const void * d_logits, split, g_scratch.d_lp, g_scratch.d_ids); #define DFLASH_TOPK_CASE(KV) \ case KV: \ + dispatched = true; \ if (use_vec) { DFLASH_TOPK_LAUNCH(KV, true) } \ else { DFLASH_TOPK_LAUNCH(KV, false) } \ break; switch (K) { DFLASH_TOPK_CASE(1) DFLASH_TOPK_CASE(2) DFLASH_TOPK_CASE(3) DFLASH_TOPK_CASE(4) DFLASH_TOPK_CASE(5) DFLASH_TOPK_CASE(6) DFLASH_TOPK_CASE(7) DFLASH_TOPK_CASE(8) + DFLASH_TOPK_CASE(12) DFLASH_TOPK_CASE(16) default: break; } #undef DFLASH_TOPK_CASE #undef DFLASH_TOPK_LAUNCH if (kProfile) cudaEventRecord(e_k1); - if (cudaGetLastError() == cudaSuccess && cudaDeviceSynchronize() == cudaSuccess) { + if (dispatched && cudaGetLastError() == cudaSuccess && + cudaDeviceSynchronize() == cudaSuccess) { const cudaError_t e1 = cudaMemcpy(out_log_probs, g_scratch.d_lp, n * sizeof(float), cudaMemcpyDeviceToHost); const cudaError_t e2 = cudaMemcpy(out_token_ids, g_scratch.d_ids, diff --git a/server/src/common/geometric_draft_topk_cuda.h b/server/src/common/geometric_draft_topk_cuda.h index b926dbefc..0edd48aa5 100644 --- a/server/src/common/geometric_draft_topk_cuda.h +++ b/server/src/common/geometric_draft_topk_cuda.h @@ -25,6 +25,15 @@ namespace dflash::common { +// Keep the public capability predicate in lockstep with the template +// instantiations dispatched by geometric_draft_topk_cuda.cu. Callers use this +// to choose the CPU fallback without entering CUDA, and loader validation uses +// it to reject selector metadata that cannot be executed consistently across +// the host and device paths. +inline constexpr bool geometric_draft_topk_cuda_supports_k(int K) noexcept { + return (K >= 1 && K <= 8) || K == 12 || K == 16; +} + // d_logits: device pointer to row-major [n_positions][vocab] f32 logits (the // position stride is `vocab` floats — pass an offset pointer to skip // leading positions). out_* are HOST buffers of size n_positions*K. diff --git a/server/src/common/gpu_runtime_compat.h b/server/src/common/gpu_runtime_compat.h index dba8eaa7c..8a21fe025 100644 --- a/server/src/common/gpu_runtime_compat.h +++ b/server/src/common/gpu_runtime_compat.h @@ -60,6 +60,7 @@ #define cudaPointerAttributes hipPointerAttribute_t #define cudaPointerGetAttributes hipPointerGetAttributes #define cudaStreamCreate hipStreamCreate +#define cudaStreamCreateWithFlags hipStreamCreateWithFlags #define cudaStreamDefault hipStreamDefault #define cudaStreamDestroy hipStreamDestroy #define cudaStreamNonBlocking hipStreamNonBlocking diff --git a/server/src/common/model_backend.h b/server/src/common/model_backend.h index 12445e6b2..422043ade 100644 --- a/server/src/common/model_backend.h +++ b/server/src/common/model_backend.h @@ -335,6 +335,14 @@ struct ModelBackend { // and stays valid until shutdown(). virtual SeqEngine * seq_engine() { return nullptr; } + // Request-level decode_mode is consumed by the concurrent scheduler. + // Report only capabilities whose resources were actually initialized; + // the scheduler rejects unsupported non-AR modes before claiming a slot. + virtual ConcurrentDecodeCapabilities concurrent_decode_capabilities() + const { + return {}; + } + // ── Snapshots ──────────────────────────────────────────────────── // With right-sized CPU-resident snapshots, each slot costs only // ~(cur_pos × 5 KB) of system RAM, so we can afford many slots. diff --git a/server/src/common/speculation/adapters/dflash2_speculator.cpp b/server/src/common/speculation/adapters/dflash2_speculator.cpp new file mode 100644 index 000000000..d7c7709b3 --- /dev/null +++ b/server/src/common/speculation/adapters/dflash2_speculator.cpp @@ -0,0 +1,102 @@ +#include "common/speculation/adapters/dflash2_speculator.h" + +#include "common/dflash2_head.h" + +#include +#include +#include + +namespace dflash::common { +namespace { + +std::string depth_debug_fields(const DFlash2DepthSignal & signal) { + std::ostringstream out; + out << std::setprecision(9) + << "\"selected_logp\":" << signal.selected_log_prob + << ",\"lm_margin\":" << signal.lm_top2_margin + << ",\"topk_mass\":" << signal.top_k_mass + << ",\"rank\":" << signal.selected_rank + << ",\"lm_top1\":" + << (signal.agrees_with_lm_top1 ? "true" : "false") + << ",\"selector_margin\":" << signal.selector_margin + << ",\"selector_mass\":" << signal.selector_winner_mass + << ",\"selector_entropy\":" << signal.selector_entropy; + return out.str(); +} + +} // namespace + +DFlash2Speculator::DFlash2Speculator( + const DraftWeights & weights, + ggml_backend_t backend, + ggml_tensor * lm_head, + DFlash2BenefitModelSignature signature, + DFlash2BenefitConfig config) + : weights_(weights), backend_(backend), lm_head_(lm_head), + benefit_(std::move(signature), std::move(config)), + score_kind_(benefit_.score_kind()) { + const DraftSelectorWeights & selector = weights_.selector; + if (!backend_ || !lm_head_ || !selector.enabled || !selector.hproj || + !selector.pred_cb || !selector.succ_cb || selector.rank <= 0 || + selector.top_k <= 0 || weights_.block_size <= 1) { + error_ = "DFlash2 selector inputs are unavailable"; + } else if (!benefit_.ready()) { + error_ = benefit_.error(); + } +} + +int DFlash2Speculator::max_block_size() const { + return weights_.block_size; +} + +bool DFlash2Speculator::propose( + const SpeculatorBatchInput & input, + std::vector & output) { + output.clear(); + if (!ready() || + !speculator_input_satisfies(input, input_requirements()) || + input.requested_depth > max_block_size()) { + return false; + } + + std::vector> draft_tokens; + std::vector traces; + if (!dflash2_select_chains_batched( + weights_, backend_, lm_head_, input.hidden_by_lane, + input.requested_depth, input.seed_tokens, + draft_tokens, &traces) || + static_cast(draft_tokens.size()) != input.lane_count || + static_cast(traces.size()) != input.lane_count) { + return false; + } + + output.resize(static_cast(input.lane_count)); + for (int lane = 0; lane < input.lane_count; ++lane) { + SpecProposal & proposal = output[static_cast(lane)]; + proposal.tokens = std::move(draft_tokens[static_cast(lane)]); + + DFlash2BenefitEstimate estimate; + std::string estimate_error; + if (!benefit_.estimate( + traces[static_cast(lane)], + input.requested_depth, estimate, &estimate_error)) { + proposal.error = estimate_error.empty() + ? "DFlash2 activation estimate failed" + : std::move(estimate_error); + continue; + } + proposal.estimate.expected_yield = estimate.expected_yield; + proposal.estimate.conditional_hazards = + std::move(estimate.conditional_hazards); + proposal.debug_depth_fields.reserve( + traces[static_cast(lane)].depths.size()); + for (const DFlash2DepthSignal & signal : + traces[static_cast(lane)].depths) { + proposal.debug_depth_fields.push_back( + depth_debug_fields(signal)); + } + } + return true; +} + +} // namespace dflash::common diff --git a/server/src/common/speculation/adapters/dflash2_speculator.h b/server/src/common/speculation/adapters/dflash2_speculator.h new file mode 100644 index 000000000..d9ba6dcec --- /dev/null +++ b/server/src/common/speculation/adapters/dflash2_speculator.h @@ -0,0 +1,42 @@ +#pragma once + +#include "common/dflash2_benefit.h" +#include "common/speculation/speculator.h" +#include "internal.h" + +#include "ggml-backend.h" + +#include + +namespace dflash::common { + +class DFlash2Speculator final : public Speculator { +public: + DFlash2Speculator( + const DraftWeights & weights, + ggml_backend_t backend, + ggml_tensor * lm_head, + DFlash2BenefitModelSignature signature, + DFlash2BenefitConfig config = {}); + + const std::string & score_kind() const override { return score_kind_; } + int max_block_size() const override; + uint32_t input_requirements() const override { + return SpeculatorInputHidden; + } + bool ready() const override { return error_.empty(); } + const std::string & error() const override { return error_; } + + bool propose(const SpeculatorBatchInput & input, + std::vector & output) override; + +private: + const DraftWeights & weights_; + ggml_backend_t backend_ = nullptr; + ggml_tensor * lm_head_ = nullptr; + DFlash2BenefitProvider benefit_; + std::string score_kind_; + std::string error_; +}; + +} // namespace dflash::common diff --git a/server/src/common/speculation/spec_cost_profile.cpp b/server/src/common/speculation/spec_cost_profile.cpp new file mode 100644 index 000000000..9e4cd9dca --- /dev/null +++ b/server/src/common/speculation/spec_cost_profile.cpp @@ -0,0 +1,292 @@ +#include "common/speculation/spec_cost_profile.h" +#include "common/sha1.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +namespace dflash::common { +namespace { + +void sort_unique_positive(std::vector & values) { + values.erase( + std::remove_if( + values.begin(), values.end(), + [](int value) { return value <= 0; }), + values.end()); + std::sort(values.begin(), values.end()); + values.erase(std::unique(values.begin(), values.end()), values.end()); +} + +struct SeriesResult { + SpecCostSeries table; + std::string error; +}; + +SeriesResult profile_monotonic_costs( + std::vector indices, + const SpecCostProfiler::Runner & runner, + int repetitions) { + SeriesResult result; + if (!runner) { + result.error = "profiling runner is missing"; + return result; + } + if (repetitions <= 0) { + result.error = "profiling repetitions must be positive"; + return result; + } + sort_unique_positive(indices); + if (indices.empty()) { + result.error = "profiling grid is empty"; + return result; + } + + result.table.indices = indices; + result.table.costs.reserve(indices.size()); + for (int index : indices) { + (void)runner(index); + std::vector samples; + samples.reserve(static_cast(repetitions)); + for (int rep = 0; rep < repetitions; ++rep) { + const double sample = runner(index); + if (!std::isfinite(sample) || sample <= 0.0) { + result.error = "profiling runner returned an invalid cost"; + result.table = {}; + return result; + } + samples.push_back(sample); + } + std::sort(samples.begin(), samples.end()); + double median = samples[static_cast(repetitions) / 2]; + if (repetitions % 2 == 0) { + median = 0.5 * ( + samples[static_cast(repetitions) / 2 - 1] + + median); + } + if (!result.table.costs.empty()) { + median = std::max(median, result.table.costs.back()); + } + result.table.costs.push_back(median); + } + return result; +} + + +constexpr int kProfileCacheVersion = 1; +constexpr size_t kMaxSeriesEntries = 4096; + +bool read_series( + std::istream & input, const char * expected, + SpecCostSeries & series) { + std::string name; + size_t count = 0; + if (!(input >> name >> count) || name != expected || + count == 0 || count > kMaxSeriesEntries) { + return false; + } + series.indices.resize(count); + series.costs.resize(count); + for (size_t i = 0; i < count; ++i) { + if (!(input >> series.indices[i] >> series.costs[i])) return false; + } + return true; +} + +void write_series( + std::ostream & output, const char * name, + const SpecCostSeries & series) { + output << name << ' ' << series.indices.size() << '\n'; + output << std::setprecision(17); + for (size_t i = 0; i < series.indices.size(); ++i) { + output << series.indices[i] << ' ' << series.costs[i] << '\n'; + } +} + +std::string hex_sha1(const std::string & value) { + uint8_t digest[20]; + sha1_hash(value.data(), value.size(), digest); + std::ostringstream out; + out << std::hex << std::setfill('0'); + for (uint8_t byte : digest) out << std::setw(2) << (unsigned)byte; + return out.str(); +} + +} // namespace + +std::string spec_cost_profile_cache_path(const std::string & identity) { + if (identity.empty()) return {}; + if (const char * configured = std::getenv("DFLASH_SPEC_PROFILE_PATH")) { + if (std::string(configured) == "0") return {}; + if (*configured) return configured; + } + std::filesystem::path root; + if (const char * xdg = std::getenv("XDG_CACHE_HOME"); xdg && *xdg) { + root = xdg; + } else if (const char * home = std::getenv("HOME"); home && *home) { + root = std::filesystem::path(home) / ".cache"; + } else { + return {}; + } + return (root / "lucebox" / + ("spec-cost-v1-" + hex_sha1(identity) + ".profile")).string(); +} + +bool load_spec_cost_profile( + const std::string & path, const std::string & identity, + SpecCostTables & tables, std::string & error) { + tables = {}; + error.clear(); + if (path.empty()) { + error = "profile cache disabled"; + return false; + } + std::ifstream input(path); + if (!input) { + error = "profile cache miss"; + return false; + } + std::string magic; + int version = 0; + std::string stored_identity; + SpecCostTables loaded; + if (!(input >> magic >> version) || magic != "dflash-spec-cost-profile" || + version != kProfileCacheVersion || + !(input >> std::quoted(stored_identity)) || + !(input >> std::quoted(loaded.speculator_id)) || + !read_series(input, "tree", loaded.tree_cost) || + !read_series(input, "step", loaded.step_cost) || + !read_series(input, "draft", loaded.draft_cost)) { + error = "invalid profile cache"; + return false; + } + input >> std::ws; + if (!input.eof() || stored_identity != identity || !loaded.valid()) { + error = stored_identity != identity + ? "profile cache identity mismatch" : "invalid profile cache"; + return false; + } + tables = std::move(loaded); + return true; +} + +bool save_spec_cost_profile( + const std::string & path, const std::string & identity, + const SpecCostTables & tables, std::string & error) { + error.clear(); + if (path.empty()) return true; + if (identity.empty() || !tables.valid()) { + error = "refusing to save invalid profile cache"; + return false; + } + const std::filesystem::path destination(path); + std::error_code ec; + if (!destination.parent_path().empty()) { + std::filesystem::create_directories(destination.parent_path(), ec); + if (ec) { + error = "could not create profile cache directory"; + return false; + } + } + const std::filesystem::path temporary = + destination.string() + ".tmp." + std::to_string(getpid()); + { + std::ofstream output(temporary, std::ios::trunc); + if (!output) { + error = "could not open temporary profile cache"; + return false; + } + output << "dflash-spec-cost-profile " << kProfileCacheVersion << '\n' + << std::quoted(identity) << '\n' + << std::quoted(tables.speculator_id) << '\n'; + write_series(output, "tree", tables.tree_cost); + write_series(output, "step", tables.step_cost); + write_series(output, "draft", tables.draft_cost); + output.flush(); + if (!output) { + error = "could not write profile cache"; + output.close(); + std::filesystem::remove(temporary, ec); + return false; + } + } + std::filesystem::rename(temporary, destination, ec); + if (ec) { + error = "could not publish profile cache"; + std::filesystem::remove(temporary, ec); + return false; + } + return true; +} + +SpecProfileGrid build_spec_profile_grid( + int max_concurrency, int tree_width, int max_accept, + const std::function & bucket) { + SpecProfileGrid grid; + if (max_concurrency <= 0 || tree_width <= 0 || max_accept <= 0) { + return grid; + } + auto bucketed = [&](int lanes) { + if (lanes <= 0) return 0; + return std::max(lanes, bucket ? bucket(lanes) : lanes); + }; + for (int lanes = 1; lanes <= max_concurrency; ++lanes) { + grid.tree_rows.push_back(bucketed(lanes) * tree_width); + grid.draft_lanes.push_back(lanes); + } + for (int concurrency = 1; concurrency <= max_concurrency; ++concurrency) { + grid.step_rows.push_back(bucketed(concurrency)); + for (int spec_lanes = 1; spec_lanes <= concurrency; ++spec_lanes) { + const int ar_rows = bucketed(concurrency - spec_lanes); + for (int accepted = spec_lanes; + accepted <= spec_lanes * max_accept; ++accepted) { + grid.step_rows.push_back(accepted + ar_rows); + } + } + } + sort_unique_positive(grid.tree_rows); + sort_unique_positive(grid.step_rows); + sort_unique_positive(grid.draft_lanes); + return grid; +} + +SpecCostProfileResult SpecCostProfiler::profile( + const SpecProfileGrid & grid, + Runner tree_runner, + Runner step_runner, + Runner draft_runner, + std::string speculator_id, + int repetitions) const { + SpecCostProfileResult result; + if (speculator_id.empty()) { + result.error = "speculator id is empty"; + return result; + } + + SeriesResult tree = profile_monotonic_costs( + grid.tree_rows, tree_runner, repetitions); + SeriesResult step = profile_monotonic_costs( + grid.step_rows, step_runner, repetitions); + SeriesResult draft = profile_monotonic_costs( + grid.draft_lanes, draft_runner, repetitions); + if (!tree.error.empty() || !step.error.empty() || !draft.error.empty()) { + result.error = tree.error + step.error + draft.error; + return result; + } + + result.tables.tree_cost = std::move(tree.table); + result.tables.step_cost = std::move(step.table); + result.tables.draft_cost = std::move(draft.table); + result.tables.speculator_id = std::move(speculator_id); + return result; +} + +} // namespace dflash::common diff --git a/server/src/common/speculation/spec_cost_profile.h b/server/src/common/speculation/spec_cost_profile.h new file mode 100644 index 000000000..924c527d4 --- /dev/null +++ b/server/src/common/speculation/spec_cost_profile.h @@ -0,0 +1,53 @@ +// Generic startup profiling protocol for monotone speculation cost tables. +#pragma once + +#include "common/speculation/speculation_gate.h" + +#include +#include +#include + +namespace dflash::common { + +struct SpecProfileGrid { + std::vector tree_rows; + std::vector step_rows; + std::vector draft_lanes; +}; + +SpecProfileGrid build_spec_profile_grid( + int max_concurrency, int tree_width, int max_accept, + const std::function & bucket); + +struct SpecCostProfileResult { + SpecCostTables tables; + std::string error; + + bool ok() const { return error.empty() && tables.valid(); } +}; + +// Returns an empty path when disk caching is disabled or no cache root exists. +std::string spec_cost_profile_cache_path(const std::string & identity); + +bool load_spec_cost_profile( + const std::string & path, const std::string & identity, + SpecCostTables & tables, std::string & error); + +bool save_spec_cost_profile( + const std::string & path, const std::string & identity, + const SpecCostTables & tables, std::string & error); + +class SpecCostProfiler { +public: + using Runner = std::function; + + SpecCostProfileResult profile( + const SpecProfileGrid & grid, + Runner tree_runner, + Runner step_runner, + Runner draft_runner, + std::string speculator_id, + int repetitions = 5) const; +}; + +} // namespace dflash::common diff --git a/server/src/common/speculation/speculation_gate.h b/server/src/common/speculation/speculation_gate.h new file mode 100644 index 000000000..1e4d4b9fa --- /dev/null +++ b/server/src/common/speculation/speculation_gate.h @@ -0,0 +1,708 @@ +// Cohort-planned adaptive speculation policy over startup-profiled costs. +// Request-local state contains only immutable activation knowledge. The engine +// owns the current cohort epoch and decides when to run plan() again. +// Pure host code: no graph, backend, or scheduler types belong here. + +#pragma once + +#include "common/speculation_policy.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace dflash::common { + +struct SpecGateConfig { + double cost_ema_alpha = 0.20; + double adaptive_gain_margin = 0.01; +}; + +struct SpecCostLookup { + double cost = std::numeric_limits::infinity(); + int requested_index = 0; + int profiled_index = 0; + bool clamped = false; + bool rounded_up = false; +}; + +struct SpecCostSeries { + std::vector indices; + std::vector costs; + + bool valid() const { + if (indices.empty() || indices.size() != costs.size()) return false; + for (size_t i = 0; i < indices.size(); ++i) { + if (indices[i] < 0 || !std::isfinite(costs[i]) || costs[i] <= 0.0) + return false; + if (i > 0 && (indices[i] <= indices[i - 1] || + costs[i] < costs[i - 1])) + return false; + } + return true; + } + + SpecCostLookup lookup(int index) const { + SpecCostLookup result; + result.requested_index = index; + if (!valid()) return result; + auto it = std::lower_bound(indices.begin(), indices.end(), index); + if (it == indices.end()) { + result.profiled_index = indices.back(); + result.cost = costs.back(); + result.clamped = true; + return result; + } + const size_t pos = static_cast(it - indices.begin()); + result.profiled_index = *it; + result.cost = costs[pos]; + result.clamped = index < indices.front() || index > indices.back(); + result.rounded_up = *it != index && !result.clamped; + return result; + } +}; + +struct SpecCostTables { + SpecCostSeries tree_cost; + SpecCostSeries step_cost; + SpecCostSeries draft_cost; + std::string speculator_id; + + bool valid() const { + return tree_cost.valid() && step_cost.valid() && draft_cost.valid(); + } +}; + +inline constexpr const char * kUnspecifiedScoreKind = "unspecified"; + +struct SpecCandidate { + uint64_t request_id = 0; + int slot = -1; + SpeculationPolicy policy = SpeculationPolicy::Adaptive; + // scoreable is a request-lifetime capability: the engine can produce the + // one-time activation score from the committed feature mirror. + bool scoreable = false; + // can_speculate is also request-lifetime (for example, false for an + // unsupported sampler or thinking hook). The min-token EOS policy is + // enforced inside the speculative path. An incompatible prefill graph may + // suspend speculative execution for one explicitly telemetered AR service + // round without changing this capability or the chosen request mode. + bool can_speculate = false; + // NaN requests a one-time bootstrap when no cached score exists. A finite + // value is the preferred activation measurement. Evaluation failure keeps + // the request AR without inventing a score. Adapter estimates are already + // calibrated; the gate only clamps executor bounds. + double activation_yield = std::numeric_limits::quiet_NaN(); + std::vector conditional_hazards; + std::string score_kind = kUnspecifiedScoreKind; +}; + +struct SpecStepGeometry { + int tree_width = 1; + std::function bucket = [](int lanes) { + return std::max(0, lanes); + }; + + int bucketed_lanes(int lanes) const { + return lanes <= 0 ? 0 : std::max(lanes, bucket ? bucket(lanes) : lanes); + } + int tree_rows(int spec_lanes) const { + return bucketed_lanes(spec_lanes) * std::max(1, tree_width); + } + double expected_step_rows(int concurrency, int spec_lanes, + double expected_spec_tokens) const { + const double accepted_rows = std::max( + static_cast(spec_lanes), expected_spec_tokens); + return accepted_rows + bucketed_lanes(concurrency - spec_lanes); + } +}; + +enum class SpecScoreSource : uint8_t { + Fresh, + Initial, + Unavailable, +}; + +inline const char * spec_score_source_name(SpecScoreSource source) { + switch (source) { + case SpecScoreSource::Fresh: return "fresh"; + case SpecScoreSource::Initial: return "initial"; + case SpecScoreSource::Unavailable: return "unavailable"; + } + return "unknown"; +} + +enum class SpecEvaluationAction : uint8_t { + Score, + FallbackAR, +}; + +struct SpecPendingEvaluation { + uint64_t request_id = 0; + int slot = -1; + SpecEvaluationAction action = SpecEvaluationAction::Score; +}; + +struct SpecPlanScore { + uint64_t request_id = 0; + int slot = -1; + double expected_yield = 1.0; + SpecScoreSource source = SpecScoreSource::Unavailable; + bool forced = false; + bool admitted = false; + std::string score_kind = kUnspecifiedScoreKind; + bool execution_unsupported = false; +}; + +// Discrete graph shape that actually ran. Unlike a SpecPlan's fractional +// expected replay rows, every field here comes from executor telemetry and is +// safe to use as an online timing key. +struct SpecExecutionShape { + int concurrency = 0; + int admitted_count = 0; + int tree_rows = 0; + int step_rows = 0; + int draft_lanes = 0; + + bool operator==(const SpecExecutionShape & other) const { + return concurrency == other.concurrency && + admitted_count == other.admitted_count && + tree_rows == other.tree_rows && + step_rows == other.step_rows && + draft_lanes == other.draft_lanes; + } +}; + +struct SpecPlan { + bool valid = true; + std::string error; + int concurrency = 0; + int admitted_count = 0; + int tree_rows = 0; + double expected_step_rows = 0.0; + int draft_lanes = 0; + double expected_tokens = 0.0; + // Startup-profiled cost before online correction, and the shape-local + // correction applied to it. predicted_cost is their product. + double profiled_cost = 0.0; + double cost_scale = 1.0; + double predicted_cost = 0.0; + // Fixed-scale expected yield for admitted activation-scored lanes. This is + // directly comparable with realized emitted tokens in telemetry. + double initial_predicted_tokens = 0.0; + double goodput = 0.0; + double ar_goodput = 0.0; + int unavailable_count = 0; + // The engine resolves pending actions, immediately replans, and caches + // only the completed result as the current cohort epoch. + bool cost_lookup_clamped = false; + std::vector ordered; + std::vector admitted_request_ids; + std::vector admitted_slots; + // Score actions are batched for one-time activation-score initialization. + // FallbackAR actions cannot attempt scoring and instead record an + // explicit failed evaluation. One tagged record keeps + // request identity and slot inseparable on all failure paths. + std::vector pending_evaluations; +}; + +struct SpecCohortEpoch { + uint64_t id = 0; + std::vector request_ids; + SpecPlan plan; + + bool matches(const std::vector & candidates) const { + return request_ids == ids(candidates); + } + + static std::vector ids( + const std::vector & candidates) { + std::vector out; + out.reserve(candidates.size()); + for (const SpecCandidate & candidate : candidates) + out.push_back(candidate.request_id); + std::sort(out.begin(), out.end()); + return out; + } +}; + +class SpeculationGate { +private: + struct RequestState { + double initial_score = + std::numeric_limits::quiet_NaN(); + bool evaluation_failed = false; + std::string score_kind = kUnspecifiedScoreKind; + std::vector conditional_hazards; + }; + + struct ExecutionShapeHash { + size_t operator()(const SpecExecutionShape & shape) const { + size_t seed = 0; + auto mix = [&](int value) { + seed ^= std::hash{}(value) + + static_cast(0x9e3779b9U) + + (seed << 6) + (seed >> 2); + }; + mix(shape.concurrency); + mix(shape.admitted_count); + mix(shape.tree_rows); + mix(shape.step_rows); + mix(shape.draft_lanes); + return seed; + } + }; + + struct CostState { + double scale = 1.0; + uint64_t observations = 0; + }; + + struct CostPrice { + double profiled = std::numeric_limits::infinity(); + double predicted = std::numeric_limits::infinity(); + }; + + struct CandidateScore { + double expected_yield = 1.0; + SpecScoreSource source = SpecScoreSource::Unavailable; + std::string score_kind = kUnspecifiedScoreKind; + }; + +public: + using ClampLogger = std::function; + + SpeculationGate(SpecCostTables costs, SpecStepGeometry geometry, + int max_accept, ClampLogger clamp_logger = {}, + SpecGateConfig config = {}, + bool direct_commit = false) + : config_(config), costs_(std::move(costs)), + geometry_(std::move(geometry)), + max_accept_(std::max(1, max_accept)), + direct_commit_(direct_commit), + clamp_logger_(std::move(clamp_logger)) {} + + SpeculationGate(SpecGateConfig config, SpecCostTables costs, + SpecStepGeometry geometry, int max_accept, + ClampLogger clamp_logger = {}, + bool direct_commit = false) + : SpeculationGate(std::move(costs), std::move(geometry), max_accept, + std::move(clamp_logger), config, direct_commit) {} + + bool valid() const { + auto valid_alpha = [](double value) { + return std::isfinite(value) && value > 0.0 && value <= 1.0; + }; + return valid_alpha(config_.cost_ema_alpha) && + std::isfinite(config_.adaptive_gain_margin) && + config_.adaptive_gain_margin >= 0.0 && + costs_.valid() && geometry_.tree_width >= 1 && + max_accept_ >= 1; + } + + // draft_lanes_override prices always-drafting. -1 means admitted-only. + // The engine calls this only for a new cohort epoch and once more after + // resolving any cold-score actions. + SpecPlan plan(int concurrency, + const std::vector & candidates, + int k_cap, int draft_lanes_override = -1) { + SpecPlan out; + out.concurrency = concurrency; + if (!valid() || concurrency < 0 || + static_cast(concurrency) != candidates.size() || + k_cap < 0) { + out.valid = false; + out.error = "invalid speculation gate inputs"; + return out; + } + + struct Ranked { + const SpecCandidate * candidate = nullptr; + double score = 1.0; + SpecScoreSource source = SpecScoreSource::Unavailable; + std::string score_kind = kUnspecifiedScoreKind; + bool forced = false; + }; + std::vector forced; + std::vector adaptive; + std::vector forced_ar; + forced.reserve(candidates.size()); + adaptive.reserve(candidates.size()); + forced_ar.reserve(candidates.size()); + + for (const SpecCandidate & candidate : candidates) { + if (candidate.policy == SpeculationPolicy::Never) continue; + if (candidate.policy == SpeculationPolicy::Adaptive && + evaluation_failed(candidate.request_id)) { + forced_ar.push_back({ + &candidate, 1.0, SpecScoreSource::Unavailable, + initial_score_kind(candidate.request_id), false}); + continue; + } + + const CandidateScore score = score_candidate(candidate); + if (candidate.policy == SpeculationPolicy::Adaptive && + score.source == SpecScoreSource::Unavailable) { + ++out.unavailable_count; + out.pending_evaluations.push_back({ + candidate.request_id, candidate.slot, + candidate.scoreable + ? SpecEvaluationAction::Score + : SpecEvaluationAction::FallbackAR, + }); + continue; + } + if (candidate.policy == SpeculationPolicy::Adaptive && + !candidate.can_speculate) { + forced_ar.push_back({ + &candidate, score.expected_yield, score.source, + score.score_kind, false}); + continue; + } + Ranked ranked{ + &candidate, score.expected_yield, score.source, + score.score_kind, + candidate.policy == SpeculationPolicy::Always}; + (ranked.forced ? forced : adaptive).push_back(ranked); + } + + // A cohort plan is publishable only after every cold adaptive request + // has resolved. Explicit Always lanes may still run in the bootstrap + // service plan, but the engine never caches that partial result. + if (!out.pending_evaluations.empty()) adaptive.clear(); + + auto request_order = [](const Ranked & a, const Ranked & b) { + return a.candidate->request_id < b.candidate->request_id; + }; + std::sort(forced.begin(), forced.end(), request_order); + std::sort(forced_ar.begin(), forced_ar.end(), request_order); + std::sort(adaptive.begin(), adaptive.end(), + [](const Ranked & a, const Ranked & b) { + if (a.score != b.score) return a.score > b.score; + return a.candidate->request_id < b.candidate->request_id; + }); + + if (static_cast(forced.size()) > k_cap) { + out.valid = false; + out.error = "forced speculation exceeds executor capacity"; + return out; + } + + std::vector ranked; + ranked.reserve(forced.size() + adaptive.size()); + ranked.insert(ranked.end(), forced.begin(), forced.end()); + ranked.insert(ranked.end(), adaptive.begin(), adaptive.end()); + for (const Ranked & item : ranked) { + out.ordered.push_back({ + item.candidate->request_id, + item.candidate->slot, + item.score, + item.source, + item.forced, + false, + item.score_kind, + false, + }); + } + + const int forced_count = static_cast(forced.size()); + const int max_k = std::min(k_cap, ranked.size()); + double expected_sum = 0.0; + for (int i = 0; i < forced_count; ++i) expected_sum += ranked[i].score; + + const SpecExecutionShape ar_shape{ + concurrency, 0, 0, geometry_.bucketed_lanes(concurrency), 0}; + const CostPrice ar_price = price_execution_shape(ar_shape, &out); + out.ar_goodput = concurrency == 0 ? 0.0 + : static_cast(concurrency) / ar_price.predicted; + + struct PlanPoint { + int k = 0; + double goodput = -1.0; + double profiled_cost = 0.0; + double predicted_cost = 0.0; + double cost_scale = 1.0; + double expected_tokens = 0.0; + int tree_rows = 0; + double expected_step_rows = 0.0; + int draft_lanes = 0; + }; + PlanPoint baseline; + PlanPoint best; + + for (int k = forced_count; k <= max_k; ++k) { + if (k > forced_count) expected_sum += ranked[k - 1].score; + const double expected_tokens = + static_cast(concurrency - k) + expected_sum; + int tree_rows = 0; + double expected_step_rows = geometry_.bucketed_lanes(concurrency); + const int draft_lanes = draft_lanes_override >= 0 + ? draft_lanes_override : k; + + if (k > 0) { + if (direct_commit_) { + tree_rows = geometry_.tree_rows(k) + concurrency - k; + expected_step_rows = 0.0; + } else { + tree_rows = geometry_.tree_rows(k); + expected_step_rows = geometry_.expected_step_rows( + concurrency, k, expected_sum); + } + } + const CostPrice price = price_expected_shape( + {concurrency, k, tree_rows, 0, draft_lanes}, + expected_step_rows, &out); + const double scale = price.predicted / price.profiled; + const double goodput = expected_tokens / price.predicted; + const PlanPoint point{ + k, goodput, price.profiled, price.predicted, scale, + expected_tokens, tree_rows, expected_step_rows, draft_lanes}; + if (k == forced_count) baseline = point; + if (goodput > best.goodput) best = point; + } + + // Explicit Always lanes establish the non-negotiable baseline. The + // safety margin applies to this epoch's adaptive subset selection. + if (best.k > forced_count && + best.goodput < baseline.goodput * + (1.0 + config_.adaptive_gain_margin)) { + best = baseline; + } + + out.admitted_count = best.k; + out.goodput = std::max(0.0, best.goodput); + out.profiled_cost = best.profiled_cost; + out.cost_scale = best.cost_scale; + out.predicted_cost = best.predicted_cost; + out.expected_tokens = best.expected_tokens; + out.tree_rows = best.tree_rows; + out.expected_step_rows = best.expected_step_rows; + out.draft_lanes = best.draft_lanes; + for (int i = 0; i < best.k; ++i) { + out.ordered[static_cast(i)].admitted = true; + out.admitted_request_ids.push_back( + ranked[static_cast(i)].candidate->request_id); + out.admitted_slots.push_back( + ranked[static_cast(i)].candidate->slot); + if (ranked[static_cast(i)].source != + SpecScoreSource::Unavailable) { + out.initial_predicted_tokens += + ranked[static_cast(i)].score; + } + } + for (const Ranked & item : forced_ar) { + out.ordered.push_back({ + item.candidate->request_id, + item.candidate->slot, + item.score, + item.source, + false, + false, + item.score_kind, + true, + }); + } + return out; + } + + void observe_cost(const SpecExecutionShape & executed, + double measured_us) { + if (!std::isfinite(measured_us) || measured_us <= 0.0 || + executed.concurrency < 0 || executed.admitted_count < 0 || + executed.admitted_count > executed.concurrency || + executed.tree_rows < 0 || executed.step_rows < 0 || + (executed.step_rows == 0 && + (!direct_commit_ || executed.admitted_count == 0)) || + executed.draft_lanes < 0) { + return; + } + const double profiled_cost = + profile_execution_shape(executed, nullptr); + if (!std::isfinite(profiled_cost) || profiled_cost <= 0.0) return; + const double ratio = std::clamp( + measured_us / profiled_cost, + kCostScaleMin, kCostScaleMax); + CostState & state = cost_states_[executed]; + update_ema(state.scale, state.observations, ratio, + config_.cost_ema_alpha); + } + + // Record a cold evaluation failure without inventing a score. The + // request remains AR in every later cohort because it cannot be ranked. + bool record_evaluation_failure(uint64_t request_id) { + RequestState & state = request_states_[request_id]; + if (state.evaluation_failed || + std::isfinite(state.initial_score)) return false; + state.evaluation_failed = true; + return true; + } + + void forget(uint64_t request_id) { + request_states_.erase(request_id); + } + + bool has_state(uint64_t request_id) const { + return request_states_.find(request_id) != request_states_.end(); + } + bool has_score(uint64_t request_id) const { + auto state = request_states_.find(request_id); + return state != request_states_.end() && + std::isfinite(state->second.initial_score); + } + bool evaluation_failed(uint64_t request_id) const { + auto state = request_states_.find(request_id); + return state != request_states_.end() && + state->second.evaluation_failed; + } + double initial_score(uint64_t request_id) const { + auto state = request_states_.find(request_id); + return state == request_states_.end() + ? std::numeric_limits::quiet_NaN() + : state->second.initial_score; + } + std::string initial_score_kind(uint64_t request_id) const { + auto state = request_states_.find(request_id); + return state == request_states_.end() + ? kUnspecifiedScoreKind : state->second.score_kind; + } + const std::vector & initial_hazards(uint64_t request_id) const { + static const std::vector empty; + auto state = request_states_.find(request_id); + return state == request_states_.end() + ? empty : state->second.conditional_hazards; + } + const SpecCostTables & costs() const { return costs_; } + +private: + CandidateScore score_candidate(const SpecCandidate & candidate) { + bool accepted_initial_score = false; + if (std::isfinite(candidate.activation_yield)) { + const double raw = std::clamp( + candidate.activation_yield, 1.0, + static_cast(max_accept_)); + RequestState & state = request_states_[candidate.request_id]; + if (!std::isfinite(state.initial_score)) { + state.initial_score = raw; + state.score_kind = candidate.score_kind; + state.conditional_hazards = candidate.conditional_hazards; + accepted_initial_score = true; + } + return { + state.initial_score, + accepted_initial_score ? SpecScoreSource::Fresh + : SpecScoreSource::Initial, + state.score_kind, + }; + } + auto state = request_states_.find(candidate.request_id); + if (state != request_states_.end() && + std::isfinite(state->second.initial_score)) { + return { + state->second.initial_score, + SpecScoreSource::Initial, + state->second.score_kind, + }; + } + return {1.0, SpecScoreSource::Unavailable, + kUnspecifiedScoreKind}; + } + + static void update_ema(double & value, uint64_t & observations, + double sample, double alpha) { + value = observations == 0 + ? sample : (1.0 - alpha) * value + alpha * sample; + ++observations; + } + + double cost_scale(const SpecExecutionShape & shape) const { + auto state = cost_states_.find(shape); + return state == cost_states_.end() || state->second.observations == 0 + ? 1.0 : state->second.scale; + } + + double profile_execution_shape(const SpecExecutionShape & shape, + SpecPlan * plan) const { + double cost = 0.0; + auto add = [&](const char * name, const SpecCostLookup & lookup) { + if (plan) report_clamp(name, lookup, *plan); + cost += lookup.cost; + }; + + if (shape.admitted_count > 0) { + add("tree", costs_.tree_cost.lookup(shape.tree_rows)); + if (!direct_commit_) { + add("step", costs_.step_cost.lookup(shape.step_rows)); + } + } else { + add("step", costs_.step_cost.lookup(shape.step_rows)); + } + if (shape.draft_lanes > 0) { + add("draft", costs_.draft_cost.lookup(shape.draft_lanes)); + } + return cost; + } + + CostPrice price_execution_shape(const SpecExecutionShape & shape, + SpecPlan * plan) const { + const double profiled = profile_execution_shape(shape, plan); + return {profiled, profiled * cost_scale(shape)}; + } + + CostPrice price_expected_shape(SpecExecutionShape shape, + double expected_step_rows, + SpecPlan * plan) const { + if (!std::isfinite(expected_step_rows) || expected_step_rows < 0.0) + return {}; + + // The executor can only launch an integer row count, while the gate + // owns an expected count. Price that expectation continuously across + // the two neighboring executable shapes so nonlinear profile cliffs + // retain their cost without an lround() decision discontinuity. Apply + // each neighbor's own online correction before interpolating it. + const int lower_rows = + static_cast(std::floor(expected_step_rows)); + const int upper_rows = + static_cast(std::ceil(expected_step_rows)); + shape.step_rows = lower_rows; + const CostPrice lower = price_execution_shape(shape, plan); + if (lower_rows == upper_rows) return lower; + + shape.step_rows = upper_rows; + const CostPrice upper = price_execution_shape(shape, plan); + const double upper_weight = expected_step_rows - lower_rows; + return { + lower.profiled + upper_weight * (upper.profiled - lower.profiled), + lower.predicted + + upper_weight * (upper.predicted - lower.predicted), + }; + } + + void report_clamp(const char * name, const SpecCostLookup & lookup, + SpecPlan & plan) const { + if (!lookup.clamped) return; + plan.cost_lookup_clamped = true; + if (clamp_logger_) + clamp_logger_(name, lookup.requested_index, lookup.profiled_index); + } + + SpecGateConfig config_; + SpecCostTables costs_; + SpecStepGeometry geometry_; + int max_accept_ = 1; + bool direct_commit_ = false; + static constexpr double kCostScaleMin = 0.25; + static constexpr double kCostScaleMax = 4.0; + std::unordered_map request_states_; + std::unordered_map + cost_states_; + ClampLogger clamp_logger_; +}; + +} // namespace dflash::common diff --git a/server/src/common/speculation/speculator.h b/server/src/common/speculation/speculator.h new file mode 100644 index 000000000..6a6d47fc4 --- /dev/null +++ b/server/src/common/speculation/speculator.h @@ -0,0 +1,91 @@ +// Model-agnostic speculation adapter contract. +#pragma once + +#include +#include +#include +#include + +namespace dflash::common { + +enum SpeculatorInputRequirement : uint32_t { + SpeculatorInputNone = 0, + SpeculatorInputHidden = 1u << 0, + SpeculatorInputPrenorm = 1u << 1, +}; + +struct ActivationEstimate { + double expected_yield = std::numeric_limits::quiet_NaN(); + std::vector conditional_hazards; +}; + +struct SpeculatorBatchInput { + int lane_count = 0; + int requested_depth = 0; + std::vector hidden_by_lane; + std::vector prenorm_by_lane; + std::vector seed_tokens; +}; + +struct SpecProposal { + std::vector tokens; + ActivationEstimate estimate; + std::string error; + // Optional per-depth JSON field fragments used only by debug telemetry. + std::vector debug_depth_fields; +}; + +inline bool speculator_input_satisfies( + const SpeculatorBatchInput & input, uint32_t requirements) { + if (input.lane_count <= 0 || input.requested_depth < 2 || + static_cast(input.seed_tokens.size()) != input.lane_count) { + return false; + } + auto has_lanes = [&](const std::vector & lanes) { + if (static_cast(lanes.size()) != input.lane_count) return false; + for (const float * lane : lanes) { + if (!lane) return false; + } + return true; + }; + if ((requirements & SpeculatorInputHidden) != 0 && + !has_lanes(input.hidden_by_lane)) { + return false; + } + if ((requirements & SpeculatorInputPrenorm) != 0 && + !has_lanes(input.prenorm_by_lane)) { + return false; + } + return true; +} + +class Speculator { +public: + virtual ~Speculator() = default; + + // Opaque, versioned identity. The activation engine never enumerates it. + virtual const std::string & score_kind() const = 0; + virtual int max_block_size() const = 0; + virtual uint32_t input_requirements() const = 0; + virtual bool ready() const = 0; + virtual const std::string & error() const = 0; + + // Draft and score each lane in one adapter call. A false return means the + // whole batch failed; lane-local failures use SpecProposal::error. + virtual bool propose(const SpeculatorBatchInput & input, + std::vector & output) = 0; +}; + +inline constexpr const char * kNoSpeculatorAdapterReason = + "no_speculator_adapter"; + +inline bool speculator_is_ready(const Speculator * speculator) { + return speculator != nullptr && speculator->ready(); +} + +inline const char * speculator_fallback_reason(const Speculator * speculator) { + return speculator_is_ready(speculator) + ? nullptr : kNoSpeculatorAdapterReason; +} + +} // namespace dflash::common diff --git a/server/src/common/speculation/survival_score.h b/server/src/common/speculation/survival_score.h new file mode 100644 index 000000000..4d1edf64e --- /dev/null +++ b/server/src/common/speculation/survival_score.h @@ -0,0 +1,49 @@ +// Reusable scoring contract for conditional acceptance hazards. +#pragma once + +#include "common/speculation/speculator.h" + +#include +#include + +namespace dflash::common { + +// Each hazard is a probability-like, monotone-in-acceptance score for one +// position conditioned on accepting its prefix. The yield includes the root. +inline double hazard_survival_yield( + const std::vector & hazards, int max_accept) { + if (max_accept <= 1) return 1.0; + double expected = 1.0; + double survival = 1.0; + const int depth = std::min( + static_cast(hazards.size()), max_accept - 1); + for (int i = 0; i < depth; ++i) { + const double hazard = + std::clamp(hazards[static_cast(i)], 0.0, 1.0); + survival *= hazard; + expected += survival; + } + return std::clamp(expected, 1.0, static_cast(max_accept)); +} + +class ConfidenceVectorScorer { +public: + ActivationEstimate score( + const std::vector & confidences, int max_accept) const { + ActivationEstimate estimate; + const int depth = std::max( + 0, std::min( + static_cast(confidences.size()), max_accept - 1)); + estimate.conditional_hazards.reserve(static_cast(depth)); + for (int i = 0; i < depth; ++i) { + estimate.conditional_hazards.push_back( + std::clamp( + confidences[static_cast(i)], 0.0, 1.0)); + } + estimate.expected_yield = hazard_survival_yield( + estimate.conditional_hazards, max_accept); + return estimate; + } +}; + +} // namespace dflash::common diff --git a/server/src/common/speculation_policy.h b/server/src/common/speculation_policy.h new file mode 100644 index 000000000..1f56d8ffe --- /dev/null +++ b/server/src/common/speculation_policy.h @@ -0,0 +1,67 @@ +// Decode policy shared by the HTTP layer, scheduler, and sequence engines. + +#pragma once + +#include +#include + +namespace dflash::common { + +enum class SpeculationPolicy { + Adaptive, + Always, + Never, +}; + +// Runtime capabilities for the concurrent decode path. Forced speculation +// needs an executable draft/verify chain. Adaptive means the server can honor +// the one-shot activation contract; a configured chain may satisfy that by +// recording a request-local sticky-AR fallback when scoring or profiling is +// unavailable. AR is always supported. +struct ConcurrentDecodeCapabilities { + bool forced_speculation = false; + bool adaptive = false; + + constexpr bool supports(SpeculationPolicy policy) const { + switch (policy) { + case SpeculationPolicy::Always: return forced_speculation; + case SpeculationPolicy::Adaptive: return adaptive; + case SpeculationPolicy::Never: return true; + } + return false; + } +}; + +inline const char * speculation_policy_name(SpeculationPolicy policy) { + switch (policy) { + case SpeculationPolicy::Adaptive: return "adaptive"; + case SpeculationPolicy::Always: return "speculation"; + case SpeculationPolicy::Never: return "ar"; + } + return "adaptive"; +} + +inline bool parse_speculation_policy( + std::string_view value, SpeculationPolicy & policy) { + if (value == "adaptive") { + policy = SpeculationPolicy::Adaptive; + return true; + } + if (value == "speculation") { + policy = SpeculationPolicy::Always; + return true; + } + if (value == "ar") { + policy = SpeculationPolicy::Never; + return true; + } + return false; +} + +inline SpeculationPolicy resolve_speculation_policy( + SpeculationPolicy server_default, + const std::optional & request_override) { + return request_override.value_or(server_default); +} + +} // namespace dflash::common diff --git a/server/src/common/step_graph.h b/server/src/common/step_graph.h index cdd4d0210..d11a2a208 100644 --- a/server/src/common/step_graph.h +++ b/server/src/common/step_graph.h @@ -20,6 +20,8 @@ struct StepGraph { ggml_context * ctx = nullptr; ggml_cgraph * gf = nullptr; ggml_gallocr_t alloc = nullptr; + ggml_context * commit_ctx = nullptr; + ggml_backend_buffer_t commit_buffer = nullptr; // Persistent metadata arena for the draft graph. Reusing the same arena // across rebuilds keeps every ggml_tensor at a stable address, which is @@ -35,7 +37,8 @@ struct StepGraph { ggml_tensor * inp_embed = nullptr; ggml_tensor * positions = nullptr; ggml_tensor * attn_mask = nullptr; // may be null - ggml_tensor * parent_ids = nullptr; // DDTree tree-mode; null for chain mode + ggml_tensor * parent_ids = nullptr; // DDTree [tree_width,n_tree_seqs] + ggml_tensor * tree_sizes = nullptr; // DDTree [n_tree_seqs], 0 = padding ggml_tensor * target_hidden_cat = nullptr; // draft only ggml_tensor * positions_k = nullptr; // draft only ggml_tensor * pad_mask_full = nullptr; // draft only; padded-ctx mask @@ -55,11 +58,21 @@ struct StepGraph { // state_slot_ids has the same shape but maps padding to a safe readable // slot for graph-level conv-state gathers. ggml_tensor * active_slot_ids = nullptr; + // Recurrent gather rows. Unlike active/paged IDs, padding must name a + // valid harmless slot (normally 0): ggml_get_rows does not mask -1. ggml_tensor * state_slot_ids = nullptr; // Ragged paged read (concurrent prefill): per-row block-table column and // inclusive causal position, [n_tokens] i32 each. Padding rows carry -1. ggml_tensor * paged_query_seq_ids = nullptr; ggml_tensor * paged_query_positions = nullptr; + // DFlash target-feature destination rows. Multi-slot replay maps each + // token to its slot-local ring; padding maps to the cache's dead row. + ggml_tensor * target_feat_rows = nullptr; + // Packed-tree direct-commit metadata uploaded after posterior selection. + ggml_tensor * accepted_prefixes = nullptr; // [n_tree_seqs] i32 + ggml_tensor * commit_slot_ids = nullptr; // [n_tree_seqs] i32 + ggml_tensor * commit_rows = nullptr; // [tree_width,n_tree_seqs] i64 + ggml_tensor * feature_commit_rows = nullptr; // same shape, i32 // Multi-prompt steps: i32 row indices gathered from the final norm // before the LM head (committing rows + decode rows). ggml_tensor * logits_row_indices = nullptr; @@ -77,12 +90,18 @@ struct StepGraph { // Per-delta-net-layer captures (verify only). std::vector delta_captures; + ggml_tensor * tree_features = nullptr; std::vector moe_selected; }; // Reset the per-call graph state (ctx + graph + tensor handles) but KEEP the // persistent CUDA buffer in `sg.alloc` alive across steps. inline void step_graph_free(StepGraph & sg) { + if (sg.commit_buffer) { + ggml_backend_buffer_free(sg.commit_buffer); + sg.commit_buffer = nullptr; + } + if (sg.commit_ctx) { ggml_free(sg.commit_ctx); sg.commit_ctx = nullptr; } if (sg.ctx) { ggml_free(sg.ctx); sg.ctx = nullptr; } sg.gf = nullptr; sg.inp_embed = sg.positions = sg.attn_mask = nullptr; @@ -92,11 +111,17 @@ inline void step_graph_free(StepGraph & sg) { sg.built_view = false; sg.hidden_input = nullptr; sg.parent_ids = nullptr; + sg.tree_sizes = nullptr; sg.kv_write_rows = nullptr; sg.active_slot_ids = nullptr; sg.state_slot_ids = nullptr; sg.paged_query_seq_ids = nullptr; sg.paged_query_positions = nullptr; + sg.target_feat_rows = nullptr; + sg.accepted_prefixes = nullptr; + sg.commit_slot_ids = nullptr; + sg.commit_rows = nullptr; + sg.feature_commit_rows = nullptr; sg.logits_row_indices = nullptr; sg.logits = nullptr; sg.hidden_states = nullptr; @@ -108,6 +133,7 @@ inline void step_graph_free(StepGraph & sg) { sg.hot_local_lut = nullptr; sg.valid_lut = nullptr; sg.delta_captures.clear(); + sg.tree_features = nullptr; sg.moe_selected.clear(); } diff --git a/server/src/draft/draft_gguf_loader.cpp b/server/src/draft/draft_gguf_loader.cpp index e5a04721c..43304f9f3 100644 --- a/server/src/draft/draft_gguf_loader.cpp +++ b/server/src/draft/draft_gguf_loader.cpp @@ -25,6 +25,7 @@ // blk..ffn_down.weight [hidden, intermediate] Q8_0 / F16 #include "internal.h" +#include "common/dflash2_selector_validation.h" #include "common/derived_scalars.h" #include "common/gguf_mmap.h" #include "common/gguf_bounds.h" @@ -75,7 +76,7 @@ int count_attn_gate_layers(const DraftWeights & w) { bool check_shape_1d(const ggml_tensor * t, int64_t ne0, const char * name, char * buf, size_t buf_sz) { if (!t || t->ne[0] != ne0) { - std::snprintf(buf, buf_sz, "draft GGUF: Domino tensor %s shape mismatch: got [%lld], expected [%lld]", + std::snprintf(buf, buf_sz, "draft GGUF: tensor %s shape mismatch: got [%lld], expected [%lld]", name, t ? (long long)t->ne[0] : -1LL, (long long)ne0); return false; } @@ -86,7 +87,7 @@ bool check_shape_2d(const ggml_tensor * t, int64_t ne0, int64_t ne1, const char * name, char * buf, size_t buf_sz) { if (!t || t->ne[0] != ne0 || t->ne[1] != ne1) { std::snprintf(buf, buf_sz, - "draft GGUF: Domino tensor %s shape mismatch: got [%lld,%lld], expected [%lld,%lld]", + "draft GGUF: tensor %s shape mismatch: got [%lld,%lld], expected [%lld,%lld]", name, t ? (long long)t->ne[0] : -1LL, t ? (long long)t->ne[1] : -1LL, @@ -96,6 +97,21 @@ bool check_shape_2d(const ggml_tensor * t, int64_t ne0, int64_t ne1, return true; } +bool check_shape_3d(const ggml_tensor * t, int64_t ne0, int64_t ne1, int64_t ne2, + const char * name, char * buf, size_t buf_sz) { + if (!t || t->ne[0] != ne0 || t->ne[1] != ne1 || t->ne[2] != ne2) { + std::snprintf(buf, buf_sz, + "draft GGUF: tensor %s shape mismatch: got [%lld,%lld,%lld], expected [%lld,%lld,%lld]", + name, + t ? (long long)t->ne[0] : -1LL, + t ? (long long)t->ne[1] : -1LL, + t ? (long long)t->ne[2] : -1LL, + (long long)ne0, (long long)ne1, (long long)ne2); + return false; + } + return true; +} + } // namespace bool load_draft_gguf(const std::string & path, @@ -209,6 +225,14 @@ bool load_draft_gguf(const std::string & path, if (target) { out.mask_token_id = target->mask_token_id; } + // The drafter's own MASK id wins over the family default: newer drafters + // (e.g. the Qwen3.8 DSpark release) are trained with a different mask + // token than the target-side default, and drafting with the wrong mask + // embedding silently destroys acceptance. + { + const uint32_t mask_meta = read_u32("dflash.mask_token_id", 0); + if (mask_meta != 0) out.mask_token_id = (int32_t)mask_meta; + } // Upper bounds on hparams. Guards against malformed/hostile GGUFs that // would otherwise trigger huge allocations or signed-int overflow when @@ -245,6 +269,24 @@ bool load_draft_gguf(const std::string & path, if (out.rope_theta == 0.0f) { fprintf(stderr, "[draft-gguf] WARNING: rope.freq_base not found in GGUF, draft RoPE will be wrong\n"); } + // YaRN rope scaling (optional). Drafters trained with YaRN (e.g. Qwen3.8 + // DSpark: factor 32, orig ctx 8192) apply it at every position; plain + // RoPE at inference silently degrades acceptance. + { + const float yarn_factor = read_f32("rope.scaling.factor", 0.0f); + if (yarn_factor > 1.0f) { + out.rope_freq_scale = 1.0f / yarn_factor; + out.rope_ext_factor = 1.0f; + out.rope_attn_factor = read_f32("rope.scaling.attn_factor", 1.0f); + out.rope_beta_fast = read_f32("rope.scaling.beta_fast", 32.0f); + out.rope_beta_slow = read_f32("rope.scaling.beta_slow", 1.0f); + out.rope_n_ctx_orig = (int)read_u32("rope.scaling.original_context_length", 0); + fprintf(stderr, + "[draft-gguf] YaRN rope: factor=%.1f orig_ctx=%d beta=%.1f/%.1f\n", + yarn_factor, out.rope_n_ctx_orig, + out.rope_beta_fast, out.rope_beta_slow); + } + } out.layers.assign((size_t)n_layer, DraftLayer{}); auto g = [&](const char * name) -> ggml_tensor * { @@ -301,6 +343,11 @@ bool load_draft_gguf(const std::string & path, L.w_gate = fnd("ffn_gate.weight"); L.w_up = fnd("ffn_up.weight"); L.w_down = fnd("ffn_down.weight"); + // DFlash 2 grouped dynamic convs (optional) + L.attn_conv.base = fnd("attn_conv.base"); + L.attn_conv.proj = fnd("attn_conv.proj.weight"); + L.mlp_conv.base = fnd("ffn_conv.base"); + L.mlp_conv.proj = fnd("ffn_conv.proj.weight"); if (!L.attn_norm || !L.ffn_norm || !L.wq || !L.wk || !L.wv || !L.wo || !L.q_norm || !L.k_norm || !L.w_gate || !L.w_up || !L.w_down) { char b[128]; @@ -451,6 +498,82 @@ bool load_draft_gguf(const std::string & path, out.dspark.confidence_dim); } + // DFlash 2: dynamic convs in every layer + candidate selector head. + { + const int conv_k = (int)read_u32("dflash.dflash2.conv_kernel_size", 0); + int n_conv = 0; + for (const DraftLayer & L : out.layers) { + if (L.attn_conv.present() && L.mlp_conv.present()) n_conv++; + } + if (n_conv > 0 || conv_k > 0) { + if (n_conv != out.n_layer || conv_k <= 0) { + set_last_error("draft GGUF: DFlash 2 conv tensors/metadata incomplete " + "(need attn_conv/ffn_conv base+proj in every layer and " + "dflash.dflash2.conv_kernel_size)"); + ggml_free(meta_ctx); out.ctx = nullptr; gguf_free(gctx); + return false; + } + out.conv_kernel_size = conv_k; + out.conv_group_size = (int)read_u32("dflash.dflash2.conv_group_size", 16); + const DraftLayer & L0 = out.layers[0]; + const int64_t groups = out.n_embd / out.conv_group_size; + char shape_err[192]; + if (!check_shape_3d(L0.attn_conv.base, out.n_embd, conv_k, 2, "attn_conv.base", shape_err, sizeof(shape_err)) || + !check_shape_2d(L0.attn_conv.proj, out.n_embd, 2 * conv_k * groups, "attn_conv.proj", shape_err, sizeof(shape_err))) { + set_last_error(shape_err); + ggml_free(meta_ctx); out.ctx = nullptr; gguf_free(gctx); + return false; + } + std::fprintf(stderr, "[draft GGUF] DFlash 2 dynamic convs: kernel=%d group=%d\n", + out.conv_kernel_size, out.conv_group_size); + } + out.selector = DraftSelectorWeights{}; + out.selector.hproj = g("dflash.selector.hproj.weight"); + out.selector.pred_cb = g("dflash.selector.pred_cb"); + out.selector.succ_cb = g("dflash.selector.succ_cb"); + const uint32_t sel_rank = read_u32("dflash.dflash2.selector_rank", 0); + if (out.selector.hproj || out.selector.pred_cb || out.selector.succ_cb || sel_rank) { + if (!out.selector.hproj || !out.selector.pred_cb || !out.selector.succ_cb) { + set_last_error("draft GGUF: DFlash 2 selector tensors incomplete " + "(hproj.weight, pred_cb, succ_cb)"); + ggml_free(meta_ctx); out.ctx = nullptr; gguf_free(gctx); + return false; + } + out.selector.rank = sel_rank ? (int)sel_rank : (int)out.selector.hproj->ne[1]; + out.selector.top_k = (int)read_u32("dflash.dflash2.selector_top_k", 16); + char shape_err[192]; + const int64_t R = out.selector.rank; + if (!check_shape_2d(out.selector.hproj, out.n_embd, R, + "selector.hproj", shape_err, sizeof(shape_err))) { + set_last_error(shape_err); + ggml_free(meta_ctx); out.ctx = nullptr; gguf_free(gctx); + return false; + } + DFlash2SelectorLayout selector_layout; + selector_layout.rank = out.selector.rank; + selector_layout.top_k = out.selector.top_k; + selector_layout.hproj_rank = out.selector.hproj->ne[1]; + selector_layout.pred_rank = out.selector.pred_cb->ne[0]; + selector_layout.pred_vocab = out.selector.pred_cb->ne[1]; + selector_layout.succ_rank = out.selector.succ_cb->ne[0]; + selector_layout.succ_vocab = out.selector.succ_cb->ne[1]; + selector_layout.target_output_vocab = + target && target->output ? target->output->ne[1] : 0; + selector_layout.target_declared_vocab = + target ? target->n_vocab : 0; + std::string selector_error; + if (!validate_dflash2_selector_layout( + selector_layout, selector_error)) { + set_last_error("draft GGUF: " + selector_error); + ggml_free(meta_ctx); out.ctx = nullptr; gguf_free(gctx); + return false; + } + out.selector.enabled = true; + std::fprintf(stderr, "[draft GGUF] DFlash 2 selector enabled: rank=%d top_k=%d vocab=%lld\n", + out.selector.rank, out.selector.top_k, (long long)out.selector.pred_cb->ne[1]); + } + } + // GGUF Qwen3.6 drafters carry SWA metadata emitted by the converter: // dflash-draft.attention.sliding_window = 2048 // dflash-draft.attention.sliding_window_pattern = [true,true,true,true,false] diff --git a/server/src/draft/draft_graph.cpp b/server/src/draft/draft_graph.cpp index 472c214c9..0178fd40b 100644 --- a/server/src/draft/draft_graph.cpp +++ b/server/src/draft/draft_graph.cpp @@ -40,6 +40,19 @@ namespace dflash::common { +// RoPE with the drafter's scaling config. YaRN-trained drafters (e.g. the +// Qwen3.8 DSpark release: factor 32, orig ctx 8192) apply the scaled rotary +// at every position, so plain-RoPE inference silently degrades acceptance. +static ggml_tensor * draft_rope(ggml_context * ctx, ggml_tensor * t, + ggml_tensor * positions, + const DraftWeights & w) { + return ggml_rope_ext(ctx, t, positions, /*freq_factors=*/nullptr, + w.head_dim, GGML_ROPE_TYPE_NEOX, w.rope_n_ctx_orig, + w.rope_theta, w.rope_freq_scale, + w.rope_ext_factor, w.rope_attn_factor, + w.rope_beta_fast, w.rope_beta_slow); +} + // Feature fusion shared by the legacy one-shot graph and the cached-KV // builders: optional per-capture RMSNorm slices, fc projection, hidden_norm. // Row-independent, so it is bit-identical whether run over the full window @@ -72,6 +85,69 @@ static ggml_tensor * draft_fuse_features( return target_feat; } +// ── DFlash 2 grouped dynamic causal conv ──────────────────────────── +// +// Two taps over the draft block (positions within the block; the block's +// first slot has no predecessor). For each tap k the coefficient is a +// per-element base kernel plus a per-group dynamic kernel projected from +// the block's normalized hidden state: +// dyn = proj @ x_norm [2*K*groups, q_len] +// coef_s_k = base[s][k] (per element) + dyn[s][k] (per group, broadcast) +// out = sum_k coef_s_k * shift_k(x) +// s = 0 ("prepare", applied to the sub-block input) or 1 ("finish", applied +// to the sub-block output); both use the dyn computed from the input. +struct DraftDynConv { + ggml_tensor * dyn = nullptr; // [2*K*groups, q_len] +}; + +static DraftDynConv draft_dyn_conv_kernel(ggml_context * ctx, + const DraftConvWeights & cw, + ggml_tensor * x_norm) { + DraftDynConv dc; + dc.dyn = ggml_mul_mat(ctx, cw.proj, x_norm); // [2*K*groups, q_len] + return dc; +} + +static ggml_tensor * draft_dyn_conv_apply(ggml_context * ctx, + const DraftWeights & w, + const DraftConvWeights & cw, + const DraftDynConv & dc, + int s, // 0 = prepare, 1 = finish + ggml_tensor * x) { // [hidden, q_len] + const int64_t hidden = x->ne[0]; + const int64_t q_len = x->ne[1]; + const int K = w.conv_kernel_size; + const int64_t gs = w.conv_group_size; + const int64_t groups = hidden / gs; + const size_t e = ggml_element_size(dc.dyn); + + ggml_tensor * out = nullptr; + for (int k = 0; k < K; ++k) { + // shift_k(x): column l takes x[:, l-k], zero for l < k + ggml_tensor * xs = x; + if (k > 0) { + if (q_len <= k) break; + ggml_tensor * head = ggml_view_2d(ctx, x, hidden, q_len - k, x->nb[1], 0); + xs = ggml_pad_ext(ctx, head, 0, 0, k, 0, 0, 0, 0, 0); // [hidden, q_len] + } + // per-group dynamic coefficient for (s, k): rows [(s*K+k)*groups, +groups) + ggml_tensor * dyn_sk = ggml_view_3d(ctx, dc.dyn, 1, groups, q_len, + e, dc.dyn->nb[1], + (size_t)((s * K + k) * groups) * e); + ggml_tensor * xs3 = ggml_reshape_3d(ctx, xs, gs, groups, q_len); + ggml_tensor * dyn3 = ggml_repeat(ctx, dyn_sk, xs3); // [gs, groups, q_len] + // per-element base coefficient base[s][k]: [hidden] at offset (s*K+k)*hidden + ggml_tensor * base_sk = ggml_view_3d(ctx, cw.base, gs, groups, 1, + cw.base->nb[0] * gs, cw.base->nb[0] * hidden, + (size_t)(s * K + k) * cw.base->nb[1]); + ggml_tensor * coef = ggml_add(ctx, dyn3, base_sk); // broadcast over q_len + ggml_tensor * term = ggml_mul(ctx, xs3, coef); + term = ggml_reshape_2d(ctx, term, hidden, q_len); + out = out ? ggml_add(ctx, out, term) : term; + } + return out; +} + DraftGraphOutputs build_draft_graph( ggml_context * ctx, const DraftWeights & w, @@ -83,7 +159,6 @@ DraftGraphOutputs build_draft_graph( const int n_kv = w.n_head_kv; const int head_dim = w.head_dim; const float eps = DFLASH27B_RMS_EPS; - const float rope_base = w.rope_theta; // ── 1. Feature fusion: target_feat = rms_norm(fc @ target_hidden_cat, hidden_norm) // fc: [5*hidden, hidden] (ggml: ne[0]=5*hidden, ne[1]=hidden) @@ -118,10 +193,18 @@ DraftGraphOutputs build_draft_graph( const int eff_total_k = eff_ctx + q_len; const int ctx_offset = use_swa ? (ctx_len - w.swa_window) : 0; + const bool dyn_conv = w.conv_kernel_size > 0 && L.attn_conv.present() && L.mlp_conv.present(); + if (!disable_attn) { // ── 2a. Attention pre-norm ggml_tensor * hn = ggml_rms_norm(ctx, h, eps); hn = ggml_mul(ctx, hn, L.attn_norm); + // DFlash 2: dynamic conv "prepare" on the attention input + DraftDynConv attn_dc; + if (dyn_conv) { + attn_dc = draft_dyn_conv_kernel(ctx, L.attn_conv, hn); + hn = draft_dyn_conv_apply(ctx, w, L.attn_conv, attn_dc, 0, hn); + } std::snprintf(probe_name, sizeof(probe_name), "draft_l%d_hn", il); ggml_set_name(hn, probe_name); @@ -185,14 +268,8 @@ DraftGraphOutputs build_draft_graph( pk = ggml_view_1d(ctx, in.positions_k, eff_total_k, ctx_offset * ggml_element_size(in.positions_k)); } - Q = ggml_rope_ext(ctx, Q, in.positions_q, /*freq_factors=*/nullptr, - head_dim, GGML_ROPE_TYPE_NEOX, /*n_ctx_orig=*/0, - rope_base, /*freq_scale=*/1.0f, - /*ext_factor=*/0.0f, /*attn_factor=*/1.0f, - /*beta_fast=*/0.0f, /*beta_slow=*/0.0f); - K = ggml_rope_ext(ctx, K, pk, nullptr, - head_dim, GGML_ROPE_TYPE_NEOX, 0, - rope_base, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + Q = draft_rope(ctx, Q, in.positions_q, w); + K = draft_rope(ctx, K, pk, w); // ── 2e. Permute into the layout flash_attn_ext wants // q: [n_embd_k=head_dim, n_batch=q_len, n_head, ne3] @@ -235,6 +312,9 @@ DraftGraphOutputs build_draft_graph( // ── 2g. Output projection + residual // wo: [q_dim, hidden] (ne[0]=q_dim, ne[1]=hidden) ggml_tensor * attn_out = ggml_mul_mat(ctx, L.wo, attn); // [hidden, q_len] + if (dyn_conv) { + attn_out = draft_dyn_conv_apply(ctx, w, L.attn_conv, attn_dc, 1, attn_out); + } std::snprintf(probe_name, sizeof(probe_name), "draft_l%d_attn_out", il); ggml_set_name(attn_out, probe_name); h = ggml_add(ctx, h, attn_out); @@ -246,6 +326,11 @@ DraftGraphOutputs build_draft_graph( // ── 2h. FFN pre-norm ggml_tensor * hf = ggml_rms_norm(ctx, h, eps); hf = ggml_mul(ctx, hf, L.ffn_norm); + DraftDynConv mlp_dc; + if (dyn_conv) { + mlp_dc = draft_dyn_conv_kernel(ctx, L.mlp_conv, hf); + hf = draft_dyn_conv_apply(ctx, w, L.mlp_conv, mlp_dc, 0, hf); + } // ── 2i. SwiGLU: down(silu(gate(x)) * up(x)) // w_gate, w_up: [hidden, intermediate] @@ -255,6 +340,9 @@ DraftGraphOutputs build_draft_graph( ggml_tensor * u = ggml_mul_mat(ctx, L.w_up, hf); // [inter, q_len] ggml_tensor * gu = ggml_mul(ctx, g, u); ggml_tensor * ffn_out = ggml_mul_mat(ctx, L.w_down, gu); // [hidden, q_len] + if (dyn_conv) { + ffn_out = draft_dyn_conv_apply(ctx, w, L.mlp_conv, mlp_dc, 1, ffn_out); + } h = ggml_add(ctx, h, ffn_out); std::snprintf(probe_name, sizeof(probe_name), "draft_l%d_h_after_ffn", il); @@ -262,12 +350,15 @@ DraftGraphOutputs build_draft_graph( } } - // ── 3. Final norm + // ── 3. Final norm. DSpark's confidence head is calibrated on h before + // this normalization, so retain both tensors as distinct graph outputs. + ggml_set_name(h, "draft_hidden_prenorm"); ggml_tensor * out = ggml_rms_norm(ctx, h, eps); out = ggml_mul(ctx, out, w.out_norm); ggml_set_name(out, "draft_hidden_out"); DraftGraphOutputs og{}; + og.hidden_prenorm = h; og.hidden_states = out; og.logits = nullptr; @@ -309,11 +400,7 @@ static void draft_ctx_kv_rows( K = ggml_reshape_3d(ctx, K, w.head_dim, w.n_head_kv, n); K = ggml_rms_norm(ctx, K, eps); K = ggml_mul (ctx, K, L.k_norm); - K = ggml_rope_ext(ctx, K, positions, /*freq_factors=*/nullptr, - w.head_dim, GGML_ROPE_TYPE_NEOX, /*n_ctx_orig=*/0, - w.rope_theta, /*freq_scale=*/1.0f, - /*ext_factor=*/0.0f, /*attn_factor=*/1.0f, - /*beta_fast=*/0.0f, /*beta_slow=*/0.0f); + K = draft_rope(ctx, K, positions, w); // rope output is contiguous [head_dim, n_kv, n] → head-major rows view *k_rows_out = ggml_view_2d(ctx, K, (int64_t)w.head_dim * w.n_head_kv, n, K->nb[2], 0); @@ -356,7 +443,6 @@ DraftGraphOutputs build_draft_kv_step( const int n_kv = w.n_head_kv; const int head_dim = w.head_dim; const float eps = DFLASH27B_RMS_EPS; - const float rope_base = w.rope_theta; const int kv_total = cache.kv_total; static const bool disable_attn_gate = @@ -370,28 +456,30 @@ DraftGraphOutputs build_draft_kv_step( for (int il = 0; il < w.n_layer; il++) { const DraftLayer & L = w.layers[il]; const bool layer_is_swa = L.is_swa && !disable_swa; + const bool dyn_conv = w.conv_kernel_size > 0 && L.attn_conv.present() && L.mlp_conv.present(); - // ── attention pre-norm + // ── attention pre-norm (+ DFlash 2 dynamic conv "prepare") ggml_tensor * hn = ggml_rms_norm(ctx, h, eps); hn = ggml_mul(ctx, hn, L.attn_norm); + DraftDynConv attn_dc; + if (dyn_conv) { + attn_dc = draft_dyn_conv_kernel(ctx, L.attn_conv, hn); + hn = draft_dyn_conv_apply(ctx, w, L.attn_conv, attn_dc, 0, hn); + } // ── Q from noise, per-head RMSNorm, RoPE at absolute positions ggml_tensor * Q = ggml_mul_mat(ctx, L.wq, hn); Q = ggml_reshape_3d(ctx, Q, head_dim, n_head, q_len); Q = ggml_rms_norm(ctx, Q, eps); Q = ggml_mul (ctx, Q, L.q_norm); - Q = ggml_rope_ext(ctx, Q, in.positions_q, nullptr, - head_dim, GGML_ROPE_TYPE_NEOX, 0, - rope_base, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + Q = draft_rope(ctx, Q, in.positions_q, w); // ── noise K/V into the scratch cache slots ggml_tensor * Kn = ggml_mul_mat(ctx, L.wk, hn); Kn = ggml_reshape_3d(ctx, Kn, head_dim, n_kv, q_len); Kn = ggml_rms_norm(ctx, Kn, eps); Kn = ggml_mul (ctx, Kn, L.k_norm); - Kn = ggml_rope_ext(ctx, Kn, in.positions_q, nullptr, - head_dim, GGML_ROPE_TYPE_NEOX, 0, - rope_base, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + Kn = draft_rope(ctx, Kn, in.positions_q, w); ggml_tensor * Kn_rows = ggml_view_2d(ctx, Kn, (int64_t)head_dim * n_kv, q_len, Kn->nb[2], 0); ggml_tensor * Vn_rows = ggml_mul_mat(ctx, L.wv, hn); // [kv_dim, q_len] @@ -436,24 +524,37 @@ DraftGraphOutputs build_draft_kv_step( attn = ggml_reshape_2d(ctx, attn, head_dim * n_head, q_len); ggml_tensor * attn_out = ggml_mul_mat(ctx, L.wo, attn); + if (dyn_conv) { + attn_out = draft_dyn_conv_apply(ctx, w, L.attn_conv, attn_dc, 1, attn_out); + } h = ggml_add(ctx, h, attn_out); - // ── FFN + // ── FFN (+ DFlash 2 dynamic conv prepare/finish) ggml_tensor * hf = ggml_rms_norm(ctx, h, eps); hf = ggml_mul(ctx, hf, L.ffn_norm); + DraftDynConv mlp_dc; + if (dyn_conv) { + mlp_dc = draft_dyn_conv_kernel(ctx, L.mlp_conv, hf); + hf = draft_dyn_conv_apply(ctx, w, L.mlp_conv, mlp_dc, 0, hf); + } ggml_tensor * g = ggml_mul_mat(ctx, L.w_gate, hf); g = ggml_silu(ctx, g); ggml_tensor * u = ggml_mul_mat(ctx, L.w_up, hf); ggml_tensor * gu = ggml_mul(ctx, g, u); ggml_tensor * ffn_out = ggml_mul_mat(ctx, L.w_down, gu); + if (dyn_conv) { + ffn_out = draft_dyn_conv_apply(ctx, w, L.mlp_conv, mlp_dc, 1, ffn_out); + } h = ggml_add(ctx, h, ffn_out); } ggml_tensor * out = ggml_rms_norm(ctx, h, eps); + ggml_set_name(h, "draft_kv_hidden_prenorm"); out = ggml_mul(ctx, out, w.out_norm); ggml_set_name(out, "draft_kv_hidden_out"); DraftGraphOutputs og{}; + og.hidden_prenorm = h; og.hidden_states = out; og.logits = nullptr; if (in.lm_head) { diff --git a/server/src/draft/draft_graph.h b/server/src/draft/draft_graph.h index b89429dac..1be162963 100644 --- a/server/src/draft/draft_graph.h +++ b/server/src/draft/draft_graph.h @@ -30,6 +30,7 @@ struct DraftGraphInputs { }; struct DraftGraphOutputs { + ggml_tensor * hidden_prenorm; // [hidden, q_len, 1] before final RMSNorm ggml_tensor * hidden_states; // [hidden, q_len, 1] (always set) ggml_tensor * logits; // [vocab, q_len, 1] (non-null iff lm_head was provided) }; diff --git a/server/src/internal.h b/server/src/internal.h index fffda41ac..7db863adf 100644 --- a/server/src/internal.h +++ b/server/src/internal.h @@ -76,6 +76,13 @@ struct TargetLayer { ggml_tensor * ssm_dt_bias = nullptr; // [dt_rank] per-head alpha bias ggml_tensor * ssm_norm = nullptr; // [head_v_dim] ggml_tensor * ssm_out = nullptr; // output projection after delta-net + // Zero-copy stacked projections (set by the loader when the two source + // tensors share a type and were placed back to back in the weight buffer): + // wqkv_z: rows [0, n_z) = wqkv_gate (z), rows [n_z, ...) = wqkv + // ssm_ba: rows [0, dt_rank) = ssm_beta, rows [dt_rank, ...) = ssm_alpha + // One GEMV each instead of two; nullptr when stacking was not possible. + ggml_tensor * wqkv_z = nullptr; + ggml_tensor * ssm_ba = nullptr; // MoE FFN (qwen35moe only; nullptr on dense qwen35) ggml_tensor * ffn_gate_inp = nullptr; // [hidden, n_expert] router @@ -147,6 +154,7 @@ struct CpuEmbedder { struct TargetWeights { ggml_context * ctx = nullptr; + ggml_context * stack_ctx = nullptr; // owns the stacked alias tensors ggml_backend_t backend = nullptr; ggml_backend_buffer_t buf = nullptr; @@ -234,6 +242,18 @@ void free_target_weights(TargetWeights & w); // ─── Draft weights (z-lab DFlash, bf16) ─────────────────────────── +// DFlash 2 grouped dynamic causal conv (two taps over the draft block, one +// instance before/after attention and one before/after the MLP): +// dyn = proj @ x_norm [2*K*groups, q_len] +// prepare = sum_k (base[0][k] + dyn[0][k]) * shift_k(x_norm) +// finish = sum_k (base[1][k] + dyn[1][k]) * shift_k(sub_block_out) +// base is per element, dyn per group of conv_group_size elements. +struct DraftConvWeights { + ggml_tensor * base = nullptr; // [hidden, K, 2] f32 + ggml_tensor * proj = nullptr; // [hidden, 2*K*groups] + bool present() const { return base != nullptr && proj != nullptr; } +}; + struct DraftLayer { ggml_tensor * attn_norm; ggml_tensor * ffn_norm; @@ -247,6 +267,8 @@ struct DraftLayer { ggml_tensor * w_gate; ggml_tensor * w_up; ggml_tensor * w_down; + DraftConvWeights attn_conv; // optional DFlash 2 conv around attention + DraftConvWeights mlp_conv; // optional DFlash 2 conv around the MLP bool is_swa = false; // true for SWA layers (Qwen3.6 pattern) bool attn_gate_per_head = false; }; @@ -280,6 +302,18 @@ struct DraftDSparkWeights { ggml_tensor * confidence_b = nullptr; // [1] f32 }; +// DFlash 2 candidate selector: top-k candidates per block position from the +// target lm_head logits, then one path through them scored by a low-rank +// bigram form unary[c] + . +struct DraftSelectorWeights { + bool enabled = false; + int rank = 0; + int top_k = 0; + ggml_tensor * hproj = nullptr; // [hidden, rank] + ggml_tensor * pred_cb = nullptr; // [rank, vocab] predecessor codebook + ggml_tensor * succ_cb = nullptr; // [rank, vocab] successor codebook +}; + struct DraftWeights { ggml_context * ctx = nullptr; ggml_backend_t backend = nullptr; @@ -324,6 +358,12 @@ struct DraftWeights { // Optional DSpark/DeepSpec-style Markov correction head. When present, // greedy chain decode adds a low-rank previous-token bias before argmax. DraftDSparkWeights dspark; + + // Optional DFlash 2 pieces: dynamic convs live in the layers, the + // selector replaces argmax/markov projection for the drafted chain. + int conv_kernel_size = 0; // 0 = no dynamic convs + int conv_group_size = 0; + DraftSelectorWeights selector; }; bool load_draft_safetensors(const std::string & path, @@ -414,13 +454,12 @@ struct TargetCache { std::vector conv_input_cache; // size = n_delta (48) // Rolling target layer features captured during target forward passes. - // Shape [5 * hidden, target_feat_cap] bf16. target_feat_cap is typically - // << max_ctx (e.g. 4096) so the buffer stays small at 128K context. The - // graph writes to slot `(kv_start + i) % target_feat_cap` so positions - // beyond the cap wrap and overwrite older entries. Readers (draft) only - // need the last DRAFT_CTX_MAX positions, so wrap is invisible in - // practice. Fed into the draft graph's fc projection after a bf16→f32 - // cast (ggml_get_to_fp32_cuda). + // Single-sequence shape: [5 * hidden, target_feat_cap] bf16. A multi-slot + // cache owns one ring per physical sequence slot and one final dead row: + // [5 * hidden, target_feat_cap * n_seq_slots + 1]. Live row P in slot S + // maps to S*target_feat_cap + P%target_feat_cap; bucket padding maps to the + // dead final row because ggml_set_rows does not accept a negative index. + // target_feat_cap remains the per-sequence ring width. ggml_tensor * target_feat = nullptr; int target_feat_cap = 0; @@ -559,6 +598,10 @@ bool restore_target_cache_chain(const PrefixSnapshot * thick, // decode is AR-only). With // n_seq_slots > 1 the attention K/V tensors are sized by ctx_alloc (the shared // pool capacity plus one scratch block) rather than one sequence's max_ctx. +// `concurrent_tree` declares that a paged multi-slot caller will build packed +// DDTree verification graphs. Those graphs are deliberately side-effect-free +// for recurrent state and commit accepted paths through a later replay, so no +// rollback snapshots/intermediates are allocated. bool create_target_cache(const TargetWeights & w, int max_ctx, int max_verify_tokens, @@ -567,7 +610,8 @@ bool create_target_cache(const TargetWeights & w, bool prefill_only = false, int ctx_alloc = 0, bool paged_attention = false, - int n_seq_slots = 1); + int n_seq_slots = 1, + bool concurrent_tree = false); // `f32_ssm_intermediates` enables exact per-token checkpoints for the opt-in // layer-split fast rollback path. The default preserves the established Q8_0 @@ -584,7 +628,8 @@ bool create_target_cache_partial(const TargetWeights & w, int ctx_alloc = 0, bool f32_ssm_intermediates = false, bool paged_attention = false, - int n_seq_slots = 1); + int n_seq_slots = 1, + bool concurrent_tree = false); void free_target_cache(TargetCache & c); @@ -633,6 +678,10 @@ bool migrate_prefill_cache(const TargetWeights & w, struct DeltaNetCapture { ggml_tensor * ssm_intermediate_states = nullptr; ggml_tensor * conv_input = nullptr; + // Concurrent tree direct-commit data. The compact journal plus the + // tree conv input can advance accepted recurrent prefixes without a + // second target-model forward. These are graph-owned outputs. + ggml_tensor * transition_journal = nullptr; }; // One contiguous prompt chunk on the flattened token axis of a concurrent @@ -654,10 +703,12 @@ struct QwenGraphInputs { int kv_start; // position where the new tokens begin bool capture_layers; // if true, write captured layer features into cache.target_feat bool capture_delta_intermediate = false; // if true, populate out_delta_captures + bool capture_tree_commit = false; // compact recurrent journal + tree features bool capture_moe_router = false; // if true, expose selected expert ids for MoE layers int fa_window = 0; // sliding window for FA layers: 0 = full attention int logits_tail_rows = 0; // compute logits only for last n rows; 0 = all - ggml_tensor * parent_ids = nullptr; // [n_tokens] i32; tree mode when non-null + ggml_tensor * parent_ids = nullptr; // tree: [tree_width,n_tree_seqs] i32 + ggml_tensor * tree_sizes = nullptr; // tree: [n_tree_seqs] i32; 0 = padding tree // [n_tokens,n_head_kv] i64 physical destination rows for the // ggml_set_rows KV write; step-invariant. ggml_tensor * kv_write_rows = nullptr; @@ -683,6 +734,10 @@ struct QwenGraphInputs { // last row plus the decode rows), which a tail view cannot express. // Non-null overrides logits_tail_rows. ggml_tensor * logits_row_indices = nullptr; + // Optional replay-stable DFlash capture destinations. When present, all + // captured layers are concatenated once and written with ggml_set_rows. + // Multi-slot callers provide per-slot ring rows (padding uses dead row). + ggml_tensor * target_feat_rows = nullptr; // [n_tokens] i32 // Prefill segments on the leading token axis (see QwenPrefillSegment). // n_prefill_tokens is their total row count. seq_slot is ignored when // segments are present. @@ -714,9 +769,19 @@ struct QwenGraphInputs { // Packed steps use logits_row_indices for scattered committing rows and // compact decode rows; logits_tail_rows remains the dense-path fallback. int n_seqs = 1; + // Mixed direct-commit tree graphs place this many one-token mapped AR + // sequences before the fixed-width speculative tree segment. Their slot + // IDs share active_slot_ids/state_slot_ids with the tree lanes. + int mapped_ar_seqs = 0; int seq_slot = 0; int paged_max_kv_len = 0; int n_prefill_tokens = 0; + // Packed paged-tree metadata. Tokens are flattened sequence-major: + // row = sequence*tree_width + node. tree_scratch_* describe the physical + // KV scratch slab owned by each physical sequence slot. + int tree_width = 0; + int tree_scratch_base = 0; + int tree_scratch_stride = 0; // Capture the LAST token's post-RoPE/post-rotation Q per full-attention // layer into cache.q_cap (KVFlash target-QK scorer). Step-invariant: // node properties depend only on n_tokens and the layer index. @@ -732,6 +797,8 @@ struct QwenGraphOutputs { // views marked as ggml_set_output() so their data persists after // graph_compute; the spec-decode loop reads them host-side for rollback. std::vector delta_captures; + // BF16 [n_capture_layers*n_embd, n_tokens], packed-tree only. + ggml_tensor * tree_features = nullptr; // One entry per target layer. Populated only when capture_moe_router is // true; qwen35 dense layers and non-MoE models leave entries null. std::vector moe_selected; diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index 0a54761f9..2c1c2af9d 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -12,11 +12,22 @@ #include "attn_masks.h" #include "prefill_helpers.h" #include "common/sampler.h" +#include "common/ddtree.h" +#include "common/geometric_draft_topk_cuda.h" +#include "common/speculation/adapters/dflash2_speculator.h" +#include "common/concurrency/chain_spec_shapes.h" +#include "common/speculation/spec_cost_profile.h" #include "internal.h" +#include "ggml-cuda.h" #include +#include +#include #include #include +#include +#include +#include #include #include @@ -33,8 +44,2726 @@ int decode_bucket_width(int live_count) { return 64; } +bool chain_direct_commit_enabled() { + const char * value = std::getenv("DFLASH_CHAIN_DURABLE_REPLAY"); + return !(value && std::atoi(value) != 0); +} + +double initial_prediction_realized_tokens( + const SpecPlan & plan, const SeqEngine::StepResult & result) { + double realized = 0.0; + for (const SpecPlanScore & score : plan.ordered) { + if (!score.admitted || + score.source == SpecScoreSource::Unavailable) { + continue; + } + const auto output = std::find_if( + result.decode.begin(), result.decode.end(), + [&](const SeqEngine::DecodeOutput & item) { + return item.slot == score.slot; + }); + if (output == result.decode.end() || output->failed) + return std::numeric_limits::quiet_NaN(); + realized += 1.0 + static_cast(output->spec_accepted_tokens); + } + return realized; +} + +void log_spec_gate_plan( + const SpecPlan & plan, + double initial_realized_tokens, + double measured_us) { + int fresh = 0; + int initial = 0; + int unavailable = plan.unavailable_count; + for (const SpecPlanScore & score : plan.ordered) { + switch (score.source) { + case SpecScoreSource::Fresh: ++fresh; break; + case SpecScoreSource::Initial: ++initial; break; + case SpecScoreSource::Unavailable: ++unavailable; break; + } + } + + std::fprintf(stderr, "[spec-gate] C=%d k=%d scores=[", + plan.concurrency, plan.admitted_count); + for (size_t i = 0; i < plan.ordered.size(); ++i) { + const SpecPlanScore & score = plan.ordered[i]; + std::fprintf(stderr, "%s%llu:%.3f/%s/%s%s", + i == 0 ? "" : ",", + static_cast(score.request_id), + score.expected_yield, + spec_score_source_name(score.source), + score.score_kind.c_str(), + score.admitted ? "*" : ""); + } + std::fprintf(stderr, "] decisions=["); + for (size_t i = 0; i < plan.ordered.size(); ++i) { + const SpecPlanScore & score = plan.ordered[i]; + std::fprintf(stderr, "%s%llu:%s", + i == 0 ? "" : ",", + static_cast(score.request_id), + score.admitted ? "speculation" : "ar"); + } + std::fprintf(stderr, + "] sources=fresh:%d,initial:%d,unavailable:%d " + "initial_tokens=%.3f/", + fresh, initial, unavailable, plan.initial_predicted_tokens); + if (std::isfinite(initial_realized_tokens)) { + std::fprintf(stderr, "%.3f", initial_realized_tokens); + } else { + std::fprintf(stderr, "n/a"); + } + std::fprintf(stderr, + " profiled_cost=%.1fus cost_scale=%.3f" + " G(k)=%.6f G(0)=%.6f predicted_cost=%.1fus", + plan.profiled_cost, plan.cost_scale, + plan.goodput, plan.ar_goodput, plan.predicted_cost); + if (std::isfinite(measured_us)) { + std::fprintf(stderr, " measured=%.1fus\n", measured_us); + } else { + std::fprintf(stderr, " measured=ar-path\n"); + } +} + +void log_spec_epoch( + const SpecCohortEpoch & epoch, + const SpeculationGate & gate) { + const SpecPlan & plan = epoch.plan; + std::fprintf(stderr, + "[spec-epoch] {\"epoch_id\":%llu,\"request_ids\":[", + static_cast(epoch.id)); + for (size_t i = 0; i < epoch.request_ids.size(); ++i) { + std::fprintf(stderr, "%s%llu", i == 0 ? "" : ",", + static_cast(epoch.request_ids[i])); + } + std::fprintf(stderr, "],\"selected_request_ids\":["); + for (size_t i = 0; i < plan.admitted_request_ids.size(); ++i) { + std::fprintf(stderr, "%s%llu", i == 0 ? "" : ",", + static_cast(plan.admitted_request_ids[i])); + } + std::fprintf(stderr, "],\"requests\":["); + for (size_t i = 0; i < plan.ordered.size(); ++i) { + const SpecPlanScore & score = plan.ordered[i]; + const bool failed = gate.evaluation_failed(score.request_id); + const double initial = gate.initial_score(score.request_id); + const std::string kind = gate.initial_score_kind(score.request_id); + const std::vector & hazards = + gate.initial_hazards(score.request_id); + std::fprintf(stderr, + "%s{\"request_id\":%llu,\"slot\":%d," + "\"activation_score\":", + i == 0 ? "" : ",", + static_cast(score.request_id), score.slot); + if (std::isfinite(initial)) std::fprintf(stderr, "%.6f", initial); + else std::fprintf(stderr, "null"); + std::fprintf(stderr, + ",\"score_kind\":\"%s\",\"expected_yield\":", + kind.c_str()); + if (std::isfinite(initial)) + std::fprintf(stderr, "%.6f", score.expected_yield); + else std::fprintf(stderr, "null"); + std::fprintf(stderr, ",\"hazards\":"); + if (std::isfinite(initial)) { + std::fprintf(stderr, "["); + for (size_t j = 0; j < hazards.size(); ++j) + std::fprintf(stderr, "%s%.8g", j == 0 ? "" : ",", hazards[j]); + std::fprintf(stderr, "]"); + } else { + std::fprintf(stderr, "null"); + } + const char * evaluation = failed ? "failed" + : std::isfinite(initial) ? "scored" : "unavailable"; + const char * reason = failed ? "activation_evaluation_failed" + : score.execution_unsupported ? "execution_unsupported" + : score.admitted ? "selected_by_joint_goodput" + : "ar_counterfactual_won"; + std::fprintf(stderr, + ",\"evaluation\":\"%s\",\"route\":\"%s\"," + "\"reason\":\"%s\"}", + evaluation, score.admitted ? "speculation" : "ar", reason); + } + std::fprintf(stderr, + "],\"profiled_cost_us\":%.1f,\"cost_scale\":%.6f," + "\"predicted_cost_us\":%.1f,\"goodput\":%.9f," + "\"ar_goodput\":%.9f}\n", + plan.profiled_cost, plan.cost_scale, plan.predicted_cost, + plan.goodput, plan.ar_goodput); +} + +void log_spec_evaluation_fallback( + uint64_t request_id, + int slot, + const std::string & kind, + const char * reason) { + const std::string cause = reason + ? reason : "activation_evaluation_failed"; + std::fprintf(stderr, + "[spec-evaluation] {\"request_id\":%llu,\"slot\":%d," + "\"score_kind\":\"%s\",\"evaluation\":\"failed\"," + "\"reason\":\"%s\"}\n", + static_cast(request_id), slot, kind.c_str(), + cause.c_str()); +} + +uint64_t file_size_or_zero(const char * path) { + if (!path || !*path) return 0; + std::ifstream file(path, std::ios::binary | std::ios::ate); + if (!file) return 0; + const std::streamoff size = file.tellg(); + return size > 0 ? static_cast(size) : 0; +} + +int configured_chain_verify_depth(int maximum) { + const char * value = std::getenv("DFLASH_SPEC_CHAIN_DEPTH"); + if (!value || !*value) { + return resolve_chain_verify_depth(0, maximum); + } + + char * end = nullptr; + const long parsed = std::strtol(value, &end, 10); + const bool integer = end != value && end && *end == '\0' && + parsed >= std::numeric_limits::min() && + parsed <= std::numeric_limits::max(); + const int resolved = integer + ? resolve_chain_verify_depth(static_cast(parsed), maximum) + : 0; + if (resolved != 0) return resolved; + + std::fprintf(stderr, + "[parallel-chain] ignoring invalid DFLASH_SPEC_CHAIN_DEPTH=%s; " + "expected root-inclusive depth 2..%d, using %d\n", + value, maximum, maximum); + return resolve_chain_verify_depth(0, maximum); +} + } // namespace +Qwen35SeqEngine::Qwen35SeqEngine( + Qwen35Backend & backend, PagedKvPool & pool, int max_ctx, + int64_t scratch_row, int tree_width, int tree_scratch_base, + int tree_scratch_stride, SpecMode spec_mode, int max_prefills, + int mixed_prefill_tokens, int long_mixed_prefill_tokens, + int long_prefill_threshold, int idle_prefill_tokens, + int prefill_quantum) + : max_prefills_(std::max(1, max_prefills)), + mixed_prefill_tokens_(std::max(1, mixed_prefill_tokens)), + long_mixed_prefill_tokens_(std::max(1, long_mixed_prefill_tokens)), + long_prefill_threshold_(std::max(1, long_prefill_threshold)), + idle_prefill_tokens_(std::max(1, idle_prefill_tokens)), + prefill_quantum_(std::max(1, prefill_quantum)), b_(backend), + slots_(pool, max_ctx, std::max(1, tree_width), + backend.paged_kv_residency_.get()), + scratch_row_(scratch_row), tree_width_(tree_width), + chain_verify_depth_(configured_chain_verify_depth(tree_width)), + tree_scratch_base_(tree_scratch_base), + tree_scratch_stride_(tree_scratch_stride), spec_mode_(spec_mode) { + if (spec_mode_ == SpecMode::chain && + chain_verify_depth_ >= 2 && chain_verify_depth_ != tree_width_) { + std::fprintf(stderr, + "[parallel-chain] verify_depth=%d draft_width=%d " + "(DFLASH_SPEC_CHAIN_DEPTH)\n", + chain_verify_depth_, tree_width_); + } + const int n_slots = slots_.slot_count(); + slot_draft_kv_.resize((size_t)n_slots); + prepared_chain_drafts_.resize((size_t)n_slots); + last_activation_estimate_.resize((size_t)n_slots); + + adaptive_fallback_ar_.assign((size_t)n_slots, 0); + if (spec_mode_ == SpecMode::chain && b_.dw_.selector.enabled) { + DFlash2BenefitModelSignature signature; + signature.target_layers = b_.w_.n_layer; + signature.target_hidden = b_.w_.n_embd; + signature.target_vocab = b_.w_.n_vocab; + signature.draft_layers = b_.dw_.n_layer; + signature.draft_hidden = b_.dw_.n_embd; + signature.draft_block_size = b_.dw_.block_size; + signature.selector_rank = b_.dw_.selector.rank; + signature.selector_top_k = b_.dw_.selector.top_k; + signature.selector_vocab = b_.dw_.selector.pred_cb + ? static_cast(b_.dw_.selector.pred_cb->ne[1]) : 0; + signature.conv_kernel_size = b_.dw_.conv_kernel_size; + signature.conv_group_size = b_.dw_.conv_group_size; + signature.target_file_size = file_size_or_zero(b_.cfg_.target_path); + signature.draft_file_size = file_size_or_zero(b_.cfg_.draft_path); + + std::string config_error; + DFlash2BenefitConfig config = + DFlash2BenefitProvider::config_from_environment(config_error); + if (config_error.empty()) { + speculator_ = std::make_unique( + b_.dw_, b_.draft_backend_, b_.w_.output, + signature, config); + } + if (!config_error.empty() || !speculator_is_ready(speculator_.get())) { + const std::string error = !config_error.empty() + ? config_error + : speculator_ ? speculator_->error() + : "adapter construction failed"; + speculator_.reset(); + adaptive_fallback_reason_ = + speculator_fallback_reason(speculator_.get()); + std::fprintf(stderr, + "[parallel-chain] no speculator adapter: %s; " + "adaptive requests will use AR fallback\n", + error.c_str()); + } else { + std::fprintf(stderr, + "[parallel-chain] speculator=%s lm_weight=%.3f " + "hazard_scale=%.3f yield_scale=%.3f signature=%s\n", + speculator_->score_kind().c_str(), + config.lm_log_weight, config.hazard_scale, + config.yield_scale, signature.str().c_str()); + } + } else if (spec_mode_ == SpecMode::chain) { + adaptive_fallback_reason_ = + speculator_fallback_reason(speculator_.get()); + std::fprintf(stderr, + "[parallel-chain] no speculator adapter for loaded drafter; " + "adaptive requests will use AR fallback\n"); + } + + // The concurrent DDTree stack is gated to a local same-device drafter. + // Build metadata-only BF16 views over each slot's disjoint target feature + // ring; draft_kv_begin_step converts only newly committed rows to its F32 + // append input instead of syncing the entire 200 MiB ring per round. + capture_features_ = tree_width_ > 0 && b_.cache_.target_feat && + b_.cache_.target_feat_cap > 0 && b_.cfg_.draft_path && + !b_.cfg_.remote_draft.enabled() && !b_.split_gpus_ && + b_.cfg_.draft_gpu == b_.cfg_.device.gpu; + if (!capture_features_) return; + + ggml_init_params ip{}; + ip.mem_size = ggml_tensor_overhead() * (size_t)(n_slots + 1); + ip.no_alloc = true; + feature_view_ctx_ = ggml_init(ip); + if (!feature_view_ctx_) { + capture_features_ = false; + return; + } + + const int cap = b_.cache_.target_feat_cap; + const int64_t fc_in = + (int64_t)b_.w_.n_capture_layers * b_.w_.n_embd; + slot_feature_mirrors_.resize((size_t)n_slots); + for (int slot = 0; slot < n_slots; ++slot) { + DraftFeatureMirror & mirror = slot_feature_mirrors_[(size_t)slot]; + mirror.target_feat = ggml_view_2d( + feature_view_ctx_, b_.cache_.target_feat, fc_in, cap, + b_.cache_.target_feat->nb[1], + (size_t)slot * (size_t)cap * b_.cache_.target_feat->nb[1]); + mirror.device = b_.cfg_.draft_gpu; + mirror.target_device = b_.cfg_.device.gpu; + mirror.cap = cap; + mirror.n_target_layers = b_.w_.n_capture_layers; + mirror.hidden_size = b_.w_.n_embd; + mirror.storage_type = b_.cache_.target_feat->type; + } +} + +Qwen35SeqEngine::~Qwen35SeqEngine() { + draft_kv_batch_free(batch_draft_graph_); + for (std::unique_ptr & state : dummy_draft_kv_) { + if (state) draft_kv_free(*state); + } + dummy_draft_kv_.clear(); + for (std::unique_ptr & state : slot_draft_kv_) { + if (state) draft_kv_free(*state); + } + slot_draft_kv_.clear(); + for (DraftFeatureMirror & mirror : slot_feature_mirrors_) { + draft_feature_mirror_free(mirror); + } + slot_feature_mirrors_.clear(); + if (feature_view_ctx_) { + ggml_free(feature_view_ctx_); + feature_view_ctx_ = nullptr; + } +} + +bool Qwen35SeqEngine::spec_gate_debug_enabled() const { + const char * value = std::getenv("DFLASH_SPEC_GATE_LOG"); + return value && std::atoi(value) != 0; +} + +bool Qwen35SeqEngine::step_timing_enabled() { + static const bool enabled = []() { + const char * value = std::getenv("DFLASH_STEP_TIMING"); + return value && std::atoi(value) != 0; + }(); + return enabled; +} + +bool Qwen35SeqEngine::profile_spec_costs(int context_tokens) { + adaptive_fallback_reason_ = "cost_profile_unavailable"; + speculation_gate_.reset(); + spec_cohort_epoch_.reset(); + next_spec_cohort_epoch_id_ = 1; + if (spec_mode_ != SpecMode::chain || !capture_features_ || + !activation_scoring_available() || tree_width_ <= 1 || tree_width_ > 16 || + resolve_chain_verify_depth(chain_verify_depth_, tree_width_) == 0 || + slots_.residency_active()) { + std::fprintf(stderr, + "[spec-profile] disabled: chain/features unavailable or " + "concurrent KVFlash residency active; adaptive capability " + "unavailable\n"); + return false; + } + + const int n_slots = slots_.slot_count(); + const int T = tree_width_; + const int V = chain_verify_depth_for_round(); + const int hidden = b_.w_.n_embd; + const int n_head_kv = b_.w_.n_head_kv; + const int max_profile_ctx = std::min( + slots_.max_context() - T, b_.cache_.target_feat_cap); + if (n_slots < 1 || max_profile_ctx < 1) return false; + const int ctx_tokens = std::clamp(context_tokens, 1, max_profile_ctx); + const SpecProfileGrid grid = build_spec_profile_grid( + n_slots, V, V, [](int lanes) { + return chain_decode_bucket_width(lanes); + }); + const bool profile_batched = batched_drafting_enabled(); + + std::ostringstream identity; + identity << "qwen35-spec-cost-v1" + << "|slots=" << n_slots + << "|tree=" << T + << "|verify=" << V + << "|ctx=" << ctx_tokens + << "|max_ctx=" << slots_.max_context() + << "|kq_pad=" << b_.cfg_.kq_stride_pad + << "|draft_mode=" << (profile_batched ? "batched" : "serial") + << "|commit_mode=" + << (chain_direct_commit_enabled() ? "direct" : "replay") + << "|speculator=" << speculator_->score_kind() + << "|target_device=" << placement_device_name(b_.cfg_.device) + << "|draft_device=" << b_.cfg_.draft_gpu; + if (ggml_backend_dev_t device = + ggml_backend_get_device(b_.target_backend_)) { + identity << "|device_name=" << ggml_backend_dev_name(device) + << "|device_description=" + << ggml_backend_dev_description(device); + } + auto append_file_identity = [&](const char * label, const char * path) { + identity << '|' << label << '=' << (path ? path : ""); + if (!path || !*path) return; + std::error_code ec; + const uintmax_t size = std::filesystem::file_size(path, ec); + identity << ":size=" << (ec ? 0 : size); + ec.clear(); + const auto modified = std::filesystem::last_write_time(path, ec); + identity << ":mtime=" << (ec ? 0 : modified.time_since_epoch().count()); + }; + append_file_identity("target", b_.cfg_.target_path); + append_file_identity("draft", b_.cfg_.draft_path); + auto append_grid = [&](const char * label, const std::vector & values) { + identity << '|' << label << '='; + for (size_t i = 0; i < values.size(); ++i) { + if (i) identity << ','; + identity << values[i]; + } + }; + append_grid("tree_rows", grid.tree_rows); + append_grid("step_rows", grid.step_rows); + append_grid("draft_lanes", grid.draft_lanes); + const std::string profile_identity = identity.str(); + const std::string profile_cache_path = + spec_cost_profile_cache_path(profile_identity); + + auto install_profile = [&](const SpecCostTables & tables) { + SpecStepGeometry geometry; + geometry.tree_width = V; + geometry.bucket = [](int lanes) { + return chain_decode_bucket_width(lanes); + }; + speculation_gate_ = std::make_unique( + tables, geometry, V, + [](const char * table, int requested, int profiled) { + std::fprintf(stderr, + "[spec-gate] %s_cost index %d outside profile; " + "clamped to %d\n", + table, requested, profiled); + }, SpecGateConfig{}, chain_direct_commit_enabled()); + if (!speculation_gate_->valid()) { + speculation_gate_.reset(); + return false; + } + adaptive_fallback_reason_.clear(); + return true; + }; + + SpecCostTables cached_tables; + std::string cache_error; + if (load_spec_cost_profile( + profile_cache_path, profile_identity, + cached_tables, cache_error) && + install_profile(cached_tables)) { + std::fprintf(stderr, + "[spec-profile] loaded %s context=%d mode=%s-draft " + "speculator=%s\n", + profile_cache_path.c_str(), ctx_tokens, + profile_batched ? "batched" : "serial", + cached_tables.speculator_id.c_str()); + return true; + } + if (!profile_cache_path.empty() && + cache_error != "profile cache miss") { + std::fprintf(stderr, "[spec-profile] cache ignored: %s\n", + cache_error.c_str()); + } + + std::string profile_error; + std::vector synthetic_slots; + synthetic_slots.reserve((size_t)n_slots); + auto cleanup = [&]() { + for (int slot : synthetic_slots) { + if (slot >= 0 && slot < (int)slot_draft_kv_.size() && + slot_draft_kv_[(size_t)slot]) { + draft_kv_reset(*slot_draft_kv_[(size_t)slot]); + } + if (slots_.is_active(slot)) slots_.retire(slot); + if (slot >= 0 && slot < (int)last_activation_estimate_.size()) { + last_activation_estimate_[(size_t)slot] = {}; + } + } + }; + + // Fabricate page tables and steady-state sequence lengths without a + // target prefill. Zero K/V and captured target features once, then every + // timed graph sees a deterministic context at the requested length. + const int32_t profile_token = b_.w_.mask_token_id >= 0 + ? b_.w_.mask_token_id : 0; + std::vector prompt((size_t)ctx_tokens, profile_token); + const SamplerCfg greedy{}; + for (int lane = 0; lane < n_slots; ++lane) { + AdmitResult admitted = admit( + std::numeric_limits::max() - (uint64_t)lane, + prompt, greedy); + if (admitted.status != AdmitResult::Status::admitted) { + profile_error = admitted.error.empty() + ? "synthetic slot admission failed" : admitted.error; + cleanup(); + std::fprintf(stderr, "[spec-profile] %s\n", profile_error.c_str()); + return false; + } + synthetic_slots.push_back(admitted.slot); + Qwen35SlotManager::PrefillChunk chunk = + slots_.append_prefill(admitted.slot, ctx_tokens); + if (!chunk.ok || + !upload_block_table_delta(admitted.slot, chunk.first_new_block, + chunk.new_blocks.data(), + chunk.new_blocks.size())) { + profile_error = "synthetic paged context allocation failed"; + cleanup(); + std::fprintf(stderr, "[spec-profile] %s\n", profile_error.c_str()); + return false; + } + slots_.commit_prefill(admitted.slot); + } + for (ggml_tensor * tensor : b_.cache_.attn_k) { + if (tensor) ggml_backend_tensor_memset(tensor, 0, 0, ggml_nbytes(tensor)); + } + for (ggml_tensor * tensor : b_.cache_.attn_v) { + if (tensor) ggml_backend_tensor_memset(tensor, 0, 0, ggml_nbytes(tensor)); + } + if (b_.cache_.target_feat) { + ggml_backend_tensor_memset( + b_.cache_.target_feat, 0, 0, ggml_nbytes(b_.cache_.target_feat)); + } + seq_lens_.assign((size_t)n_slots, ctx_tokens); + ggml_backend_tensor_set( + b_.cache_.paged_kv_seq_lens, seq_lens_.data(), 0, + sizeof(int32_t) * seq_lens_.size()); + ggml_backend_synchronize(b_.target_backend_); + + int prepared_tree_rows = -1; + auto tree_runner = [&](int total_rows) -> double { + if (!profile_error.empty()) + return std::numeric_limits::infinity(); + if (prepared_tree_rows != total_rows) { + if (total_rows <= 0 || total_rows % V != 0) { + profile_error = "invalid tree profiling shape"; + return std::numeric_limits::infinity(); + } + const int bucket = total_rows / V; + const int live = std::min(bucket, n_slots); + StepGraph & sg = b_.sg_; + if (!build_target_step_paged_tree( + sg, b_.w_, b_.cache_, b_.target_backend_, + V, bucket, ctx_tokens, + tree_scratch_base_, tree_scratch_stride_, + b_.cfg_.kq_stride_pad)) { + profile_error = "tree profiling graph build failed"; + return std::numeric_limits::infinity(); + } + + std::vector tokens((size_t)total_rows, profile_token); + std::vector embeddings((size_t)hidden * total_rows); + std::vector parents((size_t)total_rows, -1); + std::vector sizes((size_t)bucket, 0); + std::vector active((size_t)bucket, -1); + std::vector state((size_t)bucket, 0); + std::vector queries((size_t)total_rows, -1); + std::vector positions((size_t)4 * total_rows, 0); + std::vector rows( + (size_t)n_head_kv * total_rows, scratch_row_); + if (!b_.w_.embedder.embed( + tokens.data(), total_rows, embeddings.data())) { + profile_error = "tree profiling embedding failed"; + return std::numeric_limits::infinity(); + } + for (int lane = 0; lane < live; ++lane) { + const int slot = synthetic_slots[(size_t)lane]; + sizes[(size_t)lane] = V; + active[(size_t)lane] = slot; + state[(size_t)lane] = slot; + for (int node = 0; node < V; ++node) { + const int row = lane * V + node; + parents[(size_t)row] = node == 0 ? -1 : node - 1; + queries[(size_t)row] = slot; + for (int axis = 0; axis < 3; ++axis) { + positions[(size_t)axis * total_rows + row] = + ctx_tokens + node; + } + for (int head = 0; head < n_head_kv; ++head) { + rows[(size_t)head * total_rows + row] = + (int64_t)tree_scratch_base_ + + (int64_t)slot * tree_scratch_stride_ + node; + } + } + } + ggml_backend_tensor_set(sg.inp_embed, embeddings.data(), 0, + sizeof(float) * embeddings.size()); + ggml_backend_tensor_set(sg.positions, positions.data(), 0, + sizeof(int32_t) * positions.size()); + ggml_backend_tensor_set(sg.parent_ids, parents.data(), 0, + sizeof(int32_t) * parents.size()); + ggml_backend_tensor_set(sg.tree_sizes, sizes.data(), 0, + sizeof(int32_t) * sizes.size()); + if (detail::target_paged_tree_active_slots_need_upload(sg)) { + ggml_backend_tensor_set(sg.active_slot_ids, active.data(), 0, + sizeof(int32_t) * active.size()); + } + ggml_backend_tensor_set(sg.state_slot_ids, state.data(), 0, + sizeof(int32_t) * state.size()); + ggml_backend_tensor_set(sg.paged_query_seq_ids, queries.data(), 0, + sizeof(int32_t) * queries.size()); + ggml_backend_tensor_set(sg.kv_write_rows, rows.data(), 0, + sizeof(int64_t) * rows.size()); + prepared_tree_rows = total_rows; + } + const auto start = std::chrono::steady_clock::now(); + if (ggml_backend_graph_compute(b_.target_backend_, b_.sg_.gf) != + GGML_STATUS_SUCCESS) { + profile_error = "tree profiling compute failed"; + return std::numeric_limits::infinity(); + } + ggml_backend_synchronize(b_.target_backend_); + // The serving path rebuilds this graph for every decode iteration. + // Rebuild between samples too: replaying the same captured target + // graph is not a supported lifecycle on the HIP graph backend. + prepared_tree_rows = -1; + return std::chrono::duration( + std::chrono::steady_clock::now() - start).count(); + }; + + int prepared_step_rows = -1; + auto step_runner = [&](int total_rows) -> double { + if (!profile_error.empty()) + return std::numeric_limits::infinity(); + if (prepared_step_rows != total_rows) { + if (total_rows <= 0 || total_rows > n_slots * V) { + profile_error = "invalid durable-step profiling shape"; + return std::numeric_limits::infinity(); + } + std::vector segments; + int offset = 0; + for (int lane = 0; offset < total_rows; ++lane) { + const int length = std::min(V, total_rows - offset); + segments.push_back({ + offset, length, synthetic_slots[(size_t)lane]}); + offset += length; + } + StepGraph & sg = b_.sg_; + if (!build_target_step( + sg, b_.w_, b_.cache_, b_.target_backend_, + 0, total_rows, false, true, false, 0, 0, + b_.cfg_.kq_stride_pad, false, false, false, true, + 1, 0, ctx_tokens + V, + total_rows, segments.data(), (int)segments.size(), + (int)segments.size(), false) || + !sg.kv_write_rows || !sg.target_feat_rows || + !sg.paged_query_seq_ids || !sg.paged_query_positions || + !sg.logits_row_indices) { + profile_error = "durable-step profiling graph build failed"; + return std::numeric_limits::infinity(); + } + + std::vector tokens((size_t)total_rows, profile_token); + std::vector embeddings((size_t)hidden * total_rows); + std::vector positions((size_t)4 * total_rows, 0); + std::vector rows( + (size_t)n_head_kv * total_rows, scratch_row_); + std::vector queries((size_t)total_rows, -1); + std::vector feature_rows( + (size_t)total_rows, + b_.cache_.target_feat_cap * n_slots); + std::vector logits_rows; + logits_rows.reserve(segments.size()); + if (!b_.w_.embedder.embed( + tokens.data(), total_rows, embeddings.data())) { + profile_error = "durable-step profiling embedding failed"; + return std::numeric_limits::infinity(); + } + for (const QwenPrefillSegment & segment : segments) { + for (int j = 0; j < segment.n_tokens; ++j) { + const int row = segment.token_offset + j; + queries[(size_t)row] = segment.seq_slot; + for (int axis = 0; axis < 3; ++axis) { + positions[(size_t)axis * total_rows + row] = + ctx_tokens + j; + } + for (int head = 0; head < n_head_kv; ++head) { + rows[(size_t)head * total_rows + row] = + (int64_t)tree_scratch_base_ + + (int64_t)segment.seq_slot * tree_scratch_stride_ + j; + } + } + logits_rows.push_back(segment.token_offset + segment.n_tokens - 1); + } + ggml_backend_tensor_set(sg.inp_embed, embeddings.data(), 0, + sizeof(float) * embeddings.size()); + ggml_backend_tensor_set(sg.positions, positions.data(), 0, + sizeof(int32_t) * positions.size()); + ggml_backend_tensor_set(sg.kv_write_rows, rows.data(), 0, + sizeof(int64_t) * rows.size()); + ggml_backend_tensor_set(sg.paged_query_seq_ids, queries.data(), 0, + sizeof(int32_t) * queries.size()); + ggml_backend_tensor_set(sg.paged_query_positions, + positions.data(), 0, + sizeof(int32_t) * total_rows); + ggml_backend_tensor_set(sg.target_feat_rows, feature_rows.data(), 0, + sizeof(int32_t) * feature_rows.size()); + ggml_backend_tensor_set(sg.logits_row_indices, logits_rows.data(), 0, + sizeof(int32_t) * logits_rows.size()); + prepared_step_rows = total_rows; + } + const auto start = std::chrono::steady_clock::now(); + if (ggml_backend_graph_compute(b_.target_backend_, b_.sg_.gf) != + GGML_STATUS_SUCCESS) { + profile_error = "durable-step profiling compute failed"; + return std::numeric_limits::infinity(); + } + ggml_backend_synchronize(b_.target_backend_); + prepared_step_rows = -1; + return std::chrono::duration( + std::chrono::steady_clock::now() - start).count(); + }; + + auto draft_runner = [&](int lanes) -> double { + if (!profile_error.empty()) { + return std::numeric_limits::infinity(); + } + std::vector profile_inputs; + std::vector selected; + profile_inputs.reserve(static_cast(lanes)); + selected.assign(static_cast(lanes), 1); + for (int lane = 0; lane < lanes; ++lane) { + profile_inputs.push_back({ + synthetic_slots[static_cast(lane)], + profile_token, + true, + SpeculationPolicy::Always, + }); + } + const auto start = std::chrono::steady_clock::now(); + if (!prepare_chain_drafts( + profile_inputs, selected, + /*force_serial=*/!profile_batched, + /*fail_fast_batch=*/profile_batched)) { + profile_error = "draft profiling adapter proposal failed"; + return std::numeric_limits::infinity(); + } + for (const StepInput & input : profile_inputs) { + const PreparedChainDraft & prepared = + prepared_chain_drafts_[static_cast(input.slot)]; + if (!prepared.valid || + static_cast(prepared.tokens.size()) != T) { + profile_error = "draft profiling proposal shape failed"; + return std::numeric_limits::infinity(); + } + } + return std::chrono::duration( + std::chrono::steady_clock::now() - start).count(); + }; + + SpecCostProfileResult profiled = SpecCostProfiler{}.profile( + grid, tree_runner, step_runner, draft_runner, + speculator_->score_kind(), 5); + cleanup(); + if (!profiled.ok() || !profile_error.empty()) { + std::fprintf(stderr, "[spec-profile] failed: %s%s\n", + profile_error.c_str(), profiled.error.c_str()); + return false; + } + + SpecCostTables tables = std::move(profiled.tables); + if (!install_profile(tables)) return false; + if (!save_spec_cost_profile( + profile_cache_path, profile_identity, tables, cache_error)) { + std::fprintf(stderr, "[spec-profile] cache save failed: %s\n", + cache_error.c_str()); + } else if (!profile_cache_path.empty()) { + std::fprintf(stderr, "[spec-profile] saved %s\n", + profile_cache_path.c_str()); + } + + auto print_table = [](const char * name, const SpecCostSeries & series) { + std::fprintf(stderr, "[spec-profile] %s", name); + for (size_t i = 0; i < series.indices.size(); ++i) { + std::fprintf(stderr, "%s%d:%.1fus", + i == 0 ? " " : ",", series.indices[i], series.costs[i]); + } + std::fprintf(stderr, "\n"); + }; + std::fprintf(stderr, + "[spec-profile] context=%d reps=5 mode=%s-draft speculator=%s\n", + ctx_tokens, + profile_batched ? "batched" : "serial", + tables.speculator_id.c_str()); + print_table("tree_cost", tables.tree_cost); + print_table("step_cost", tables.step_cost); + print_table("draft_cost", tables.draft_cost); + return true; +} + +DraftFeatureMirror * Qwen35SeqEngine::slot_feature_mirror(int slot) { + if (!capture_features_ || slot < 0 || + slot >= (int)slot_feature_mirrors_.size()) { + return nullptr; + } + return &slot_feature_mirrors_[(size_t)slot]; +} + +DraftKvState * Qwen35SeqEngine::ensure_slot_draft_kv(int slot) { + DraftFeatureMirror * mirror = slot_feature_mirror(slot); + if (!mirror || slot < 0 || slot >= (int)slot_draft_kv_.size()) { + return nullptr; + } + std::unique_ptr & state = slot_draft_kv_[(size_t)slot]; + if (state && state->gf && state->built_for == (const void *)&b_.dw_) { + return state.get(); + } + if (state) draft_kv_free(*state); + state = std::make_unique(); + const int cap = std::min( + mirror->cap, std::max(1, b_.cfg_.draft_ctx_max)); + if (!draft_kv_init(*state, b_.dw_, b_.draft_backend_, cap, nullptr)) { + draft_kv_free(*state); + state.reset(); + return nullptr; + } + return state.get(); +} +bool Qwen35SeqEngine::batched_drafting_enabled() const { + const char * value = std::getenv("DFLASH_SPEC_BATCHED_DRAFT"); + return !value || std::atoi(value) != 0; +} + +bool Qwen35SeqEngine::activation_scoring_available() const { + return spec_mode_ == SpecMode::chain && capture_features_ && + speculator_is_ready(speculator_.get()); +} + +std::string Qwen35SeqEngine::chain_activation_score_kind() const { + return speculator_is_ready(speculator_.get()) + ? speculator_->score_kind() : kUnspecifiedScoreKind; +} + +bool Qwen35SeqEngine::activation_scoring_enabled() const { + const char * value = std::getenv("DFLASH_SPEC_ACTIVATION_SCORE"); + return !value || std::atoi(value) != 0; +} + +bool Qwen35SeqEngine::prepare_chain_drafts( + const std::vector & inputs, + const std::vector & selected, + bool force_serial, + bool fail_fast_batch) { + if (selected.size() != inputs.size() || + !speculator_is_ready(speculator_.get())) return false; + + // Accumulate the full drafting wall (draft graph compute + fused + // selector + readbacks) into the round's [step-timing] attribution, + // including early-failure paths. + struct DraftTimer { + Qwen35SeqEngine * engine; + std::chrono::steady_clock::time_point start; + int lanes; + ~DraftTimer() { + if (!engine) return; + engine->round_draft_us_ += + std::chrono::duration( + std::chrono::steady_clock::now() - start).count(); + engine->round_draft_lanes_ += lanes; + } + }; + DraftTimer draft_timer{ + this, + std::chrono::steady_clock::now(), + (int)std::count_if( + selected.begin(), selected.end(), + [](uint8_t value) { return value != 0; })}; + + const int T = tree_width_; + const int hidden = b_.w_.n_embd; + const uint32_t requirements = speculator_->input_requirements(); + const bool need_prenorm = + (requirements & SpeculatorInputPrenorm) != 0; + struct Lane { + size_t input_index = 0; + int slot = -1; + int32_t seed = -1; + DraftKvState * state = nullptr; + DraftFeatureMirror * mirror = nullptr; + }; + std::vector lanes; + lanes.reserve(inputs.size()); + std::vector noise((size_t)T, b_.w_.mask_token_id); + std::vector noise_embed((size_t)hidden * T); + + // Proposal validity is current-block-specific. Do not clear the last + // published activation score before the draft succeeds: bootstrap must either + // publish a finite score or fail, and a later draft failure must not erase + // the immutable activation score already owned by the gate. + for (size_t i = 0; i < inputs.size(); ++i) { + if (!selected[i]) continue; + const int slot = inputs[i].slot; + if (slot < 0 || slot >= (int)prepared_chain_drafts_.size()) continue; + prepared_chain_drafts_[(size_t)slot].valid = false; + } + + for (size_t i = 0; i < inputs.size(); ++i) { + if (!selected[i]) continue; + const StepInput & in = inputs[i]; + if (!chain_proposal_input_capable(in)) return false; + DraftKvState * state = ensure_slot_draft_kv(in.slot); + DraftFeatureMirror * mirror = slot_feature_mirror(in.slot); + if (!state || !mirror || + !draft_kv_begin_step( + *state, b_.dw_, b_.draft_backend_, *mirror, + slots_.slot(in.slot).cur_pos)) { + return false; + } + noise[0] = in.token; + std::fill( + noise.begin() + 1, noise.end(), b_.w_.mask_token_id); + if (!b_.w_.embedder.embed( + noise.data(), T, noise_embed.data())) { + return false; + } + ggml_backend_tensor_set( + state->inp_embed, noise_embed.data(), 0, + sizeof(float) * noise_embed.size()); + lanes.push_back({i, in.slot, in.token, state, mirror}); + } + if (lanes.empty()) return true; + + std::vector proposals; + auto invoke_adapter = [&]( + const std::vector> & hidden_blocks, + const std::vector> & prenorm_blocks, + const std::vector & seeds) { + SpeculatorBatchInput adapter_input; + adapter_input.lane_count = static_cast(seeds.size()); + adapter_input.requested_depth = T; + adapter_input.seed_tokens = seeds; + for (const std::vector & block : hidden_blocks) { + adapter_input.hidden_by_lane.push_back(block.data()); + } + if (need_prenorm) { + for (const std::vector & block : prenorm_blocks) { + adapter_input.prenorm_by_lane.push_back(block.data()); + } + } + return speculator_input_satisfies(adapter_input, requirements) && + speculator_->propose(adapter_input, proposals); + }; + bool used_batch = false; + const bool try_batched = !force_serial && batched_drafting_enabled(); + if (try_batched) { + const int bucket = + chain_decode_bucket_width((int)lanes.size()); + std::vector batch_states; + std::vector seeds; + batch_states.reserve((size_t)bucket); + seeds.reserve((size_t)bucket); + for (const Lane & lane : lanes) { + batch_states.push_back(lane.state); + seeds.push_back(lane.seed); + } + + const int dummy_count = bucket - (int)lanes.size(); + const int cap = std::min( + lanes[0].mirror->cap, + std::max(1, b_.cfg_.draft_ctx_max)); + while ((int)dummy_draft_kv_.size() < dummy_count) { + auto dummy = std::make_unique(); + if (!draft_kv_init( + *dummy, b_.dw_, b_.draft_backend_, cap, nullptr)) { + draft_kv_free(*dummy); + break; + } + dummy_draft_kv_.push_back(std::move(dummy)); + } + if ((int)dummy_draft_kv_.size() >= dummy_count) { + noise[0] = lanes[0].seed; + std::fill( + noise.begin() + 1, noise.end(), + b_.w_.mask_token_id); + bool dummy_ok = b_.w_.embedder.embed( + noise.data(), T, noise_embed.data()); + for (int i = 0; dummy_ok && i < dummy_count; ++i) { + DraftKvState * dummy = + dummy_draft_kv_[(size_t)i].get(); + dummy_ok = draft_kv_begin_step( + *dummy, b_.dw_, b_.draft_backend_, + *lanes[0].mirror, 1); + if (dummy_ok) { + ggml_backend_tensor_set( + dummy->inp_embed, noise_embed.data(), 0, + sizeof(float) * noise_embed.size()); + batch_states.push_back(dummy); + seeds.push_back(lanes[0].seed); + } + } + if (dummy_ok) { + std::vector> hidden_blocks; + std::vector> prenorm_blocks; + used_batch = draft_kv_batch_compute( + batch_draft_graph_, b_.dw_, b_.draft_backend_, + batch_states, need_prenorm, + hidden_blocks, prenorm_blocks) && + invoke_adapter(hidden_blocks, prenorm_blocks, seeds) && + proposals.size() >= lanes.size(); + } + } + if (!used_batch) { + static bool warned = false; + if (!warned) { + warned = true; + std::fprintf(stderr, + "[draft-kv-batch] unavailable; using serial fallback\n"); + } + } + } + + if (!used_batch && try_batched && fail_fast_batch) return false; + + if (!used_batch && try_batched) { + // A failed backend compute can leave a subset of the packed draft + // graph's cache writes visible. Rebuild every real lane from its + // captured target-feature ring before entering the serial fallback. + for (const Lane & lane : lanes) { + draft_kv_reset(*lane.state); + if (!draft_kv_begin_step( + *lane.state, b_.dw_, b_.draft_backend_, *lane.mirror, + slots_.slot(lane.slot).cur_pos)) { + return false; + } + noise[0] = lane.seed; + std::fill( + noise.begin() + 1, noise.end(), b_.w_.mask_token_id); + if (!b_.w_.embedder.embed( + noise.data(), T, noise_embed.data())) { + return false; + } + ggml_backend_tensor_set( + lane.state->inp_embed, noise_embed.data(), 0, + sizeof(float) * noise_embed.size()); + } + } + + if (!used_batch) { + std::vector> hidden_blocks( + lanes.size(), + std::vector( + static_cast(hidden) * static_cast(T))); + std::vector> prenorm_blocks; + if (need_prenorm) { + prenorm_blocks.assign( + lanes.size(), + std::vector( + static_cast(hidden) * static_cast(T))); + } + std::vector seeds; + seeds.reserve(lanes.size()); + for (size_t lane = 0; lane < lanes.size(); ++lane) { + DraftKvState * state = lanes[lane].state; + if (ggml_backend_graph_compute( + b_.draft_backend_, state->gf) != + GGML_STATUS_SUCCESS) { + return false; + } + ggml_backend_tensor_get_async( + b_.draft_backend_, state->hidden_states, + hidden_blocks[lane].data(), 0, + sizeof(float) * hidden_blocks[lane].size()); + if (need_prenorm) { + ggml_backend_tensor_get_async( + b_.draft_backend_, state->hidden_prenorm, + prenorm_blocks[lane].data(), 0, + sizeof(float) * prenorm_blocks[lane].size()); + } + seeds.push_back(lanes[lane].seed); + } + ggml_backend_synchronize(b_.draft_backend_); + if (!invoke_adapter(hidden_blocks, prenorm_blocks, seeds) || + proposals.size() != lanes.size()) { + return false; + } + } + + if (proposals.size() < lanes.size()) return false; + for (size_t lane = 0; lane < lanes.size(); ++lane) { + SpecProposal & proposal = proposals[lane]; + if (!proposal.error.empty() || + static_cast(proposal.tokens.size()) != T || + !std::isfinite(proposal.estimate.expected_yield) || + static_cast( + proposal.estimate.conditional_hazards.size()) < T - 1) { + if (!proposal.error.empty()) { + std::fprintf(stderr, + "[spec-gate] activation evaluation failed " + "request=%llu slot=%d kind=%s: %s\n", + static_cast( + slots_.slot(lanes[lane].slot).request_id), + lanes[lane].slot, speculator_->score_kind().c_str(), + proposal.error.c_str()); + } + return false; + } + + const Lane & info = lanes[lane]; + PreparedChainDraft & prepared = + prepared_chain_drafts_[static_cast(info.slot)]; + prepared.valid = true; + prepared.generated = slots_.slot(info.slot).generated_tokens(); + prepared.root = info.seed; + prepared.tokens = std::move(proposal.tokens); + prepared.estimate = proposal.estimate; + prepared.debug_depth_fields = + std::move(proposal.debug_depth_fields); + + ActivationEstimate & published = + last_activation_estimate_[static_cast(info.slot)]; + if (!std::isfinite(published.expected_yield)) { + published = std::move(proposal.estimate); + } + } + return true; +} + +bool Qwen35SeqEngine::ddtree_eligible(const StepPlan & plan) const { + if (spec_mode_ != SpecMode::ddtree || tree_width_ <= 1 || + !capture_features_ || !plan.prefills.empty() || + plan.decode.empty() || b_.dw_.block_size <= 1 || + b_.cfg_.ddtree_budget + 1 != tree_width_) { + return false; + } + const int min_floor = []() { + const char * value = std::getenv("DFLASH_MIN_TOKENS"); + return value ? std::max(0, std::atoi(value)) : 0; + }(); + for (const StepInput & in : plan.decode) { + if (!in.allow_speculation || + in.speculation_policy == SpeculationPolicy::Never || + in.slot < 0 || + in.slot >= slots_.slot_count() || + !slots_.slot(in.slot).decoding() || + (in.speculation_policy != SpeculationPolicy::Always && + !slots_.ddtree_speculation_allowed(in.slot)) || + slots_.slot(in.slot).sampler.needs_logit_processing() || + slots_.slot(in.slot).cur_pos < 1 || + slots_.slot(in.slot).cur_pos >= slots_.max_context()) { + return false; + } + const Qwen35Slot & seq = slots_.slot(in.slot); + const int generated = seq.generated_tokens(); + if (generated < min_floor) return false; + } + return true; +} +bool Qwen35SeqEngine::chain_proposal_input_capable( + const StepInput & in) const { + const uint32_t supported_inputs = + SpeculatorInputHidden | SpeculatorInputPrenorm; + const bool adapter_capable = + speculator_is_ready(speculator_.get()) && + speculator_->max_block_size() == tree_width_ && + (speculator_->input_requirements() & ~supported_inputs) == 0; + return spec_mode_ == SpecMode::chain && capture_features_ && + tree_width_ > 1 && tree_width_ <= 16 && + resolve_chain_verify_depth( + chain_verify_depth_for_round(), tree_width_) != 0 && + b_.dw_.block_size == tree_width_ && adapter_capable && + in.slot >= 0 && in.slot < slots_.slot_count() && + slots_.slot(in.slot).decoding() && + slots_.slot(in.slot).cur_pos >= 1 && + slots_.slot(in.slot).cur_pos < slots_.max_context(); +} + +bool Qwen35SeqEngine::chain_activation_input_scoreable( + const StepInput & in) const { + return chain_proposal_input_capable(in) && + activation_scoring_available(); +} + +bool Qwen35SeqEngine::chain_spec_request_capable( + const StepInput & in) const { + return chain_proposal_input_capable(in) && + in.allow_speculation && + in.speculation_policy != SpeculationPolicy::Never && + !slots_.slot(in.slot).sampler.needs_logit_processing(); +} + +bool Qwen35SeqEngine::chain_spec_input_eligible( + const StepInput & in) const { + return chain_spec_request_capable(in); +} +SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( + const StepPlan & plan, const std::vector & admitted, + std::chrono::steady_clock::time_point round_started) { + StepResult result; + const std::vector & inputs = plan.decode; + if (admitted.size() != inputs.size() || !plan.prefills.empty()) { + result.error = "invalid chain speculation admission plan"; + return result; + } + + const int T = tree_width_; + const int V = chain_verify_depth_for_round(); + if (resolve_chain_verify_depth(V, T) == 0) { + result.error = + "invalid root-inclusive speculative chain verify depth"; + return result; + } + const int hidden = b_.w_.n_embd; + const int n_head_kv = b_.w_.n_head_kv; + const int n_slots = slots_.slot_count(); + const int min_tokens = []() { + const char * value = std::getenv("DFLASH_MIN_TOKENS"); + return value ? std::max(0, std::atoi(value)) : 0; + }(); + const int requested_spec_count = static_cast(std::count_if( + admitted.begin(), admitted.end(), + [](uint8_t value) { return value != 0; })); + if (requested_spec_count == 0) { + result.error = "empty chain speculation admission plan"; + return result; + } + + // Optional per-phase wall attribution ([step-timing]). Timestamps mark + // phase boundaries; graph computes are synchronous on this backend, so + // each *_exec span covers upload + compute up to its trailing sync. + const bool timing = step_timing_enabled(); + using timing_clock = std::chrono::steady_clock; + const auto t_round_start = round_started; + timing_clock::time_point t_verify_build_start, t_verify_build_end, + t_verify_exec_end, t_posterior_end, t_commit_end, + t_replay_build_end, t_replay_exec_end, t_sample_end; + auto span_us = [](timing_clock::time_point from, + timing_clock::time_point to) { + return std::chrono::duration(to - from).count(); + }; + + struct Proposal { + size_t input_index = 0; + int slot = -1; + int32_t root = -1; + DDTree tree; + std::vector flat; + std::vector accepted; + std::vector path; + std::vector debug_depth_fields; + int32_t verify_bonus = -1; + int32_t pending = -1; + }; + struct ArLane { + size_t input_index = 0; + int slot = -1; + int32_t token = -1; + int position = -1; + int64_t physical_row = -1; + int32_t pending = -1; + }; + + std::vector proposals; + proposals.reserve(static_cast(requested_spec_count)); + std::vector proposal_for_input(inputs.size(), -1); + std::vector active_admitted = admitted; + std::vector retried(inputs.size(), 0); + std::vector proposal_errors(inputs.size()); + + auto clean_proposal_lane = [&](size_t i) { + const int slot = inputs[i].slot; + if (slot >= 0 && + slot < static_cast(prepared_chain_drafts_.size())) { + prepared_chain_drafts_[static_cast(slot)].valid = false; + } + if (slot >= 0 && slot < static_cast(slot_draft_kv_.size()) && + slot_draft_kv_[static_cast(slot)]) { + draft_kv_reset(*slot_draft_kv_[static_cast(slot)]); + } + }; + auto fail_proposal_lane = [&](size_t i, const char * error) { + clean_proposal_lane(i); + active_admitted[i] = 0; + proposal_errors[i] = error; + std::fprintf(stderr, + "[spec-proposal-failure] request_id=%llu slot=%d error=%s\n", + (unsigned long long)slots_.slot(inputs[i].slot).request_id, + inputs[i].slot, error); + }; + auto retry_proposal_lane = [&](size_t i) { + retried[i] = 1; + clean_proposal_lane(i); + std::vector selected(inputs.size(), 0); + selected[i] = 1; + if (prepare_chain_drafts(inputs, selected, /*force_serial=*/true)) { + return true; + } + fail_proposal_lane( + i, "chain proposal preparation failed after clean retry"); + return false; + }; + + std::vector need_prepare(inputs.size(), 0); + for (size_t i = 0; i < inputs.size(); ++i) { + const StepInput & in = inputs[i]; + const bool hard_eligible = chain_spec_input_eligible(in); + if (!admitted[i]) continue; + if (!hard_eligible) { + fail_proposal_lane( + i, "epoch-selected speculation request became ineligible"); + continue; + } + const PreparedChainDraft & prepared = + prepared_chain_drafts_[(size_t)in.slot]; + const Qwen35Slot & seq = slots_.slot(in.slot); + need_prepare[i] = + !prepared.valid || + prepared.generated != seq.generated_tokens() || + prepared.root != in.token || + (int)prepared.tokens.size() != T; + } + if (std::any_of( + need_prepare.begin(), need_prepare.end(), + [](uint8_t value) { return value != 0; }) && + !prepare_chain_drafts( + inputs, need_prepare, /*force_serial=*/false, + /*fail_fast_batch=*/true)) { + // The packed prepare has no target-side effects. Reset its drafter + // state, then retry each affected request once through the serial path + // so one broken lane cannot fail or demote healthy peers. + for (size_t i = 0; i < inputs.size(); ++i) { + if (need_prepare[i]) clean_proposal_lane(i); + } + for (size_t i = 0; i < inputs.size(); ++i) { + if (need_prepare[i] && active_admitted[i]) { + retry_proposal_lane(i); + } + } + } + + auto take_prepared_proposal = [&](size_t i, Proposal & proposal) { + const StepInput & in = inputs[i]; + PreparedChainDraft & prepared = + prepared_chain_drafts_[(size_t)in.slot]; + if (!prepared.valid || + prepared.generated != slots_.slot(in.slot).generated_tokens() || + prepared.root != in.token || + (int)prepared.tokens.size() != T) { + return false; + } + + Proposal next; + next.input_index = i; + next.slot = in.slot; + next.root = in.token; + next.debug_depth_fields = + std::move(prepared.debug_depth_fields); + next.flat = std::move(prepared.tokens); + prepared.valid = false; + if (!truncate_chain_proposal(next.flat, V)) return false; + const size_t verified_signal_depths = + static_cast(V - 1); + if (next.debug_depth_fields.size() > verified_signal_depths) { + next.debug_depth_fields.resize(verified_signal_depths); + } + next.tree = make_chain_verify_tree(next.flat); + if (next.tree.n_nodes + 1 != V) return false; + proposal = std::move(next); + return true; + }; + + for (size_t i = 0; i < inputs.size(); ++i) { + if (!active_admitted[i]) continue; + Proposal proposal; + if (!take_prepared_proposal(i, proposal)) { + if (!retried[i] && retry_proposal_lane(i) && + take_prepared_proposal(i, proposal)) { + // The clean retry repaired a stale or malformed proposal. + } else { + if (active_admitted[i]) { + fail_proposal_lane( + i, "chain proposal remained invalid after clean retry"); + } + continue; + } + } + proposal_for_input[i] = + static_cast(proposals.size()); + proposals.push_back(std::move(proposal)); + } + const int spec_count = static_cast(proposals.size()); + const bool direct_commit = chain_direct_commit_enabled(); + auto lane_disposition = [&](size_t i) { + return chain_lane_disposition( + active_admitted[i] != 0, !proposal_errors[i].empty()); + }; + + std::vector ar_lanes; + ar_lanes.reserve(inputs.size() - static_cast(spec_count)); + std::vector ar_for_input(inputs.size(), -1); + for (size_t i = 0; i < inputs.size(); ++i) { + if (lane_disposition(i) != ChainLaneDisposition::AR) continue; + ArLane lane; + lane.input_index = i; + lane.slot = inputs[i].slot; + lane.token = inputs[i].token; + lane.position = slots_.slot(lane.slot).cur_pos; + ar_for_input[i] = static_cast(ar_lanes.size()); + ar_lanes.push_back(lane); + } + const int ar_count = static_cast(ar_lanes.size()); + const int tree_lane_count = spec_count; + const int tree_bucket = tree_lane_count > 0 + ? chain_decode_bucket_width(tree_lane_count) : 0; + + // Compact direct graphs write AR K/V and recurrent state durably in the + // fused launch. Allocate their physical rows now, but keep history and + // cur_pos staged until the target graph and all promotions succeed. + if (direct_commit) { + for (ArLane & lane : ar_lanes) { + const Qwen35SlotManager::StepAppend app = + slots_.append_token(lane.slot, lane.token); + if (!app.ok) { + result.error = app.busy + ? "paged KV pool exhausted during compact AR staging" + : "compact AR K/V staging failed"; + return result; + } + const bool table_ok = slots_.residency_active() || + app.new_block < 0 || + upload_block_table_delta( + lane.slot, app.new_block_index, &app.new_block, 1); + if (!table_ok) { + result.error = "compact AR block-table update failed"; + return result; + } + lane.position = app.position; + lane.physical_row = app.physical_row; + } + } + StepGraph & tree_sg = b_.sg_; + std::vector posterior; + + int replay_total = 0; + if (tree_lane_count > 0) { + // Launch 1: scratch-only packed path-tree verification. + int max_prefix = 1; + for (const Proposal & proposal : proposals) { + max_prefix = std::max( + max_prefix, slots_.slot(proposal.slot).cur_pos); + } + if (direct_commit) { + for (const ArLane & ar : ar_lanes) { + max_prefix = std::max(max_prefix, ar.position + 1); + } + } + t_verify_build_start = timing_clock::now(); + if (!build_target_step_paged_tree( + tree_sg, b_.w_, b_.cache_, b_.target_backend_, + V, tree_bucket, max_prefix, + tree_scratch_base_, tree_scratch_stride_, + b_.cfg_.kq_stride_pad, direct_commit ? ar_count : 0, + direct_commit)) { + result.error = "packed chain speculation verify graph build failed"; + return result; + } + t_verify_build_end = timing_clock::now(); + + const int spec_tree_rows = V * tree_bucket; + const int spec_row_offset = direct_commit ? ar_count : 0; + const int total_tree = spec_row_offset + spec_tree_rows; + std::vector flat_tokens(static_cast(total_tree), 0); + std::vector parents( + static_cast(spec_tree_rows), -1); + std::vector sizes(static_cast(tree_bucket), 0); + const int mapped_slot_count = spec_row_offset + tree_bucket; + std::vector tree_slots( + static_cast(mapped_slot_count), -1); + std::vector tree_state_slots( + static_cast(mapped_slot_count), 0); + std::vector query_slots(static_cast(total_tree), -1); + std::vector query_positions( + direct_commit ? static_cast(total_tree) : 0, -1); + std::vector tree_rows( + static_cast(total_tree) * n_head_kv, scratch_row_); + std::vector tree_positions( + static_cast(4) * total_tree, 0); + std::vector tree_embed( + static_cast(hidden) * total_tree, 0.0f); + seq_lens_.assign(static_cast(n_slots), 0); + + if (direct_commit) { + for (int ar_index = 0; ar_index < ar_count; ++ar_index) { + const ArLane & ar = ar_lanes[static_cast(ar_index)]; + const int row = ar_index; + tree_slots[static_cast(ar_index)] = ar.slot; + tree_state_slots[static_cast(ar_index)] = ar.slot; + seq_lens_[static_cast(ar.slot)] = ar.position + 1; + flat_tokens[static_cast(row)] = ar.token; + query_slots[static_cast(row)] = ar.slot; + query_positions[static_cast(row)] = ar.position; + tree_positions[static_cast(0) * total_tree + row] = + ar.position; + tree_positions[static_cast(1) * total_tree + row] = + ar.position; + tree_positions[static_cast(2) * total_tree + row] = + ar.position; + for (int head = 0; head < n_head_kv; ++head) { + tree_rows[static_cast(head) * total_tree + row] = + ar.physical_row; + } + } + } + + for (int lane = 0; lane < spec_count; ++lane) { + const Proposal & proposal = proposals[static_cast(lane)]; + const int tree_base = lane * V; + const int row_base = spec_row_offset + tree_base; + const int mapped_lane = spec_row_offset + lane; + sizes[static_cast(lane)] = V; + tree_slots[static_cast(mapped_lane)] = proposal.slot; + tree_state_slots[static_cast(mapped_lane)] = proposal.slot; + seq_lens_[static_cast(proposal.slot)] = + slots_.slot(proposal.slot).cur_pos; + for (int node = 0; node < V; ++node) { + const int tree_row = tree_base + node; + const int row = row_base + node; + flat_tokens[static_cast(row)] = + proposal.flat[static_cast(node)]; + parents[static_cast(tree_row)] = node == 0 + ? -1 : proposal.tree.parents[static_cast(node)]; + query_slots[static_cast(row)] = proposal.slot; + const int depth = node == 0 + ? 0 : proposal.tree.depths[static_cast(node) - 1]; + const int position = + slots_.slot(proposal.slot).cur_pos + depth; + tree_positions[static_cast(0) * total_tree + row] = + position; + tree_positions[static_cast(1) * total_tree + row] = + position; + tree_positions[static_cast(2) * total_tree + row] = + position; + for (int head = 0; head < n_head_kv; ++head) { + tree_rows[static_cast(head) * total_tree + row] = + static_cast(tree_scratch_base_) + + static_cast(proposal.slot) * + tree_scratch_stride_ + node; + } + } + } + + if (!b_.w_.embedder.embed( + flat_tokens.data(), total_tree, tree_embed.data())) { + result.error = "packed chain speculation embedding failed"; + return result; + } + ggml_backend_tensor_set(tree_sg.inp_embed, tree_embed.data(), 0, + sizeof(float) * tree_embed.size()); + ggml_backend_tensor_set(tree_sg.positions, tree_positions.data(), 0, + sizeof(int32_t) * tree_positions.size()); + ggml_backend_tensor_set(tree_sg.parent_ids, parents.data(), 0, + sizeof(int32_t) * parents.size()); + ggml_backend_tensor_set(tree_sg.tree_sizes, sizes.data(), 0, + sizeof(int32_t) * sizes.size()); + if (detail::target_paged_tree_active_slots_need_upload(tree_sg)) { + ggml_backend_tensor_set( + tree_sg.active_slot_ids, tree_slots.data(), 0, + sizeof(int32_t) * tree_slots.size()); + } + ggml_backend_tensor_set( + tree_sg.state_slot_ids, tree_state_slots.data(), 0, + sizeof(int32_t) * tree_state_slots.size()); + ggml_backend_tensor_set( + tree_sg.paged_query_seq_ids, query_slots.data(), 0, + sizeof(int32_t) * query_slots.size()); + if (tree_sg.paged_query_positions) { + ggml_backend_tensor_set( + tree_sg.paged_query_positions, query_positions.data(), 0, + sizeof(int32_t) * query_positions.size()); + } + ggml_backend_tensor_set(tree_sg.kv_write_rows, tree_rows.data(), 0, + sizeof(int64_t) * tree_rows.size()); + ggml_backend_tensor_set( + b_.cache_.paged_kv_seq_lens, seq_lens_.data(), 0, + sizeof(int32_t) * seq_lens_.size()); + if (ggml_backend_graph_compute(b_.target_backend_, tree_sg.gf) != + GGML_STATUS_SUCCESS) { + result.error = "packed chain speculation verify compute failed"; + return result; + } + t_verify_exec_end = timing_clock::now(); + + posterior.assign(static_cast(total_tree), -1); + ggml_backend_tensor_get( + tree_sg.argmax_tokens, posterior.data(), 0, + sizeof(int32_t) * posterior.size()); + t_posterior_end = timing_clock::now(); + + for (int lane = 0; lane < spec_count; ++lane) { + Proposal & proposal = proposals[static_cast(lane)]; + const int32_t * lane_posterior = + posterior.data() + static_cast(spec_row_offset + lane * V); + proposal.accepted = follow_verified_tree( + proposal.tree, lane_posterior, proposal.verify_bonus); + const int room = + slots_.max_context() - slots_.slot(proposal.slot).cur_pos; + truncate_verified_path( + proposal.accepted, static_cast(std::max(0, room)), + lane_posterior, proposal.verify_bonus); + if (proposal.accepted.empty()) { + result.error = "chain accepted path has no context headroom"; + return result; + } + proposal.path.reserve(proposal.accepted.size()); + for (int flat_index : proposal.accepted) { + proposal.path.push_back(flat_index == 0 + ? proposal.root + : proposal.tree.token_ids[ + static_cast(flat_index) - 1]); + } + const size_t safe_prefix = chain_min_tokens_safe_prefix( + proposal.path, + slots_.slot(proposal.slot).generated_tokens(), + min_tokens, + [&](int32_t token) { return token_is_eos(token); }); + proposal.path.resize(safe_prefix); + proposal.accepted.resize(safe_prefix); + if (!proposal.debug_depth_fields.empty()) { + static const bool selector_log_enabled = []() { + const char * value = + std::getenv("DFLASH_DFLASH2_SELECTOR_LOG"); + return value && std::atoi(value) != 0; + }(); + if (selector_log_enabled) { + const Qwen35Slot & sequence = slots_.slot(proposal.slot); + const size_t accepted_depth = + proposal.path.empty() ? 0 : proposal.path.size() - 1; + std::fprintf(stderr, + "[spec-selector] {\"request_id\":%llu,\"slot\":%d," + "\"score_kind\":\"%s\",\"generated\":%d," + "\"accepted_depth\":%zu,\"depths\":[", + static_cast(sequence.request_id), + proposal.slot, chain_activation_score_kind().c_str(), + sequence.generated_tokens(), accepted_depth); + for (size_t depth = 0; + depth < proposal.debug_depth_fields.size(); ++depth) { + std::fprintf(stderr, + "%s{\"depth\":%zu,\"accepted\":%s,%s}", + depth == 0 ? "" : ",", depth + 1, + depth < accepted_depth ? "true" : "false", + proposal.debug_depth_fields[depth].c_str()); + } + std::fprintf(stderr, "]}\n"); + } + } + replay_total += static_cast(proposal.path.size()); + } + } else { + const auto no_verify = timing_clock::now(); + t_verify_build_start = no_verify; + t_verify_build_end = no_verify; + t_verify_exec_end = no_verify; + t_posterior_end = no_verify; + } + + // Stage accepted path segments and all non-admitted AR peers. Nothing is + // published to slot history until the combined durable graph succeeds. + std::vector replay_segments; + std::vector replay_tokens; + std::vector replay_slots; + std::vector replay_positions; + std::vector replay_physical; + replay_segments.reserve(static_cast(spec_count)); + replay_tokens.reserve(static_cast(replay_total)); + replay_slots.reserve(static_cast(replay_total)); + replay_positions.reserve(static_cast(replay_total)); + replay_physical.reserve(static_cast(replay_total)); + seq_lens_.assign(static_cast(n_slots), 0); + int max_kv_len = 1; + int replay_offset = 0; + + for (Proposal & proposal : proposals) { + const Qwen35SlotManager::StepAppend app = slots_.append_tokens( + proposal.slot, proposal.path.data(), + static_cast(proposal.path.size())); + const bool table_ok = slots_.residency_active() || + upload_block_table_delta( + proposal.slot, app.first_new_block, + app.new_blocks.data(), app.new_blocks.size()); + if (!app.ok || app.physical_rows.size() != proposal.path.size() || + !table_ok) { + result.error = app.busy + ? "paged KV pool exhausted during chain speculation commit" + : "chain accepted-path K/V append failed"; + return result; + } + replay_segments.push_back({ + replay_offset, static_cast(proposal.path.size()), + proposal.slot, + }); + for (size_t row = 0; row < proposal.path.size(); ++row) { + replay_tokens.push_back(proposal.path[row]); + replay_slots.push_back(proposal.slot); + replay_positions.push_back(app.position + static_cast(row)); + replay_physical.push_back(app.physical_rows[row]); + } + replay_offset += static_cast(proposal.path.size()); + const int seq_len = + app.position + static_cast(proposal.path.size()); + seq_lens_[static_cast(proposal.slot)] = seq_len; + max_kv_len = std::max(max_kv_len, seq_len); + } + + for (int ar_index = 0; ar_index < ar_count; ++ar_index) { + ArLane & lane = ar_lanes[static_cast(ar_index)]; + const StepInput & in = inputs[lane.input_index]; + if (!direct_commit) { + const Qwen35SlotManager::StepAppend app = + slots_.append_token(in.slot, in.token); + if (!app.ok) { + result.error = app.busy + ? "paged KV pool exhausted during mixed AR commit" + : "mixed AR K/V append failed"; + return result; + } + const bool table_ok = slots_.residency_active() || + app.new_block < 0 || + upload_block_table_delta( + in.slot, app.new_block_index, &app.new_block, 1); + if (!table_ok) { + result.error = "mixed AR block-table update failed"; + return result; + } + lane.position = app.position; + lane.physical_row = app.physical_row; + } + seq_lens_[static_cast(in.slot)] = lane.position + 1; + max_kv_len = std::max(max_kv_len, lane.position + 1); + } + if (spec_count == 0 && ar_count == 0) { + result.decode.reserve(inputs.size()); + for (size_t i = 0; i < inputs.size(); ++i) { + DecodeOutput out; + out.slot = inputs[i].slot; + out.failed = true; + out.error = proposal_errors[i].empty() + ? "chain proposal lane made no progress" + : proposal_errors[i]; + result.decode.push_back(std::move(out)); + } + return result; + } + if (!upload_all_active_block_tables()) { + result.error = "chain mixed-step block-table refresh failed"; + return result; + } + t_commit_end = timing_clock::now(); + + if (direct_commit) { + const int spec_tree_rows = V * tree_bucket; + const int spec_row_offset = ar_count; + const int total_tree = spec_row_offset + spec_tree_rows; + std::vector accepted_prefixes( + static_cast(tree_bucket), 0); + std::vector commit_rows( + static_cast(spec_tree_rows), -1); + std::vector feature_commit_rows( + static_cast(total_tree), -1); + const int feature_cap = b_.cache_.target_feat_cap; + int replay_cursor = 0; + for (int lane = 0; lane < spec_count; ++lane) { + const Proposal & proposal = proposals[static_cast(lane)]; + if (proposal.path.size() != proposal.accepted.size()) { + result.error = "direct commit path/acceptance size mismatch"; + return result; + } + accepted_prefixes[static_cast(lane)] = + static_cast(proposal.path.size()); + for (size_t depth = 0; depth < proposal.accepted.size(); ++depth) { + const int node = proposal.accepted[depth]; + // DFlash2 proposals are chains. A branching tree needs an + // indexed journal-commit kernel rather than prefix commit. + if (node != static_cast(depth)) { + result.error = + "direct commit requires a contiguous chain acceptance"; + return result; + } + const int flat = lane * V + node; + const int source_row = spec_row_offset + flat; + if (replay_cursor >= static_cast(replay_physical.size())) { + result.error = "direct commit replay cursor overflow"; + return result; + } + commit_rows[static_cast(flat)] = + replay_physical[static_cast(replay_cursor)]; + feature_commit_rows[static_cast(source_row)] = + proposal.slot * feature_cap + + replay_positions[static_cast(replay_cursor)] % + feature_cap; + ++replay_cursor; + } + } + if (replay_cursor != replay_total) { + result.error = "direct commit replay cursor mismatch"; + return result; + } + for (int ar_index = 0; ar_index < ar_count; ++ar_index) { + const ArLane & ar = ar_lanes[static_cast(ar_index)]; + feature_commit_rows[static_cast(ar_index)] = + ar.slot * feature_cap + ar.position % feature_cap; + } + std::vector commit_slots( + static_cast(tree_bucket), -1); + for (int lane = 0; lane < spec_count; ++lane) { + commit_slots[static_cast(lane)] = + proposals[static_cast(lane)].slot; + } + ggml_backend_tensor_set( + tree_sg.accepted_prefixes, accepted_prefixes.data(), 0, + accepted_prefixes.size() * sizeof(int32_t)); + ggml_backend_tensor_set( + tree_sg.commit_slot_ids, commit_slots.data(), 0, + commit_slots.size() * sizeof(int32_t)); + ggml_backend_tensor_set( + tree_sg.commit_rows, commit_rows.data(), 0, + commit_rows.size() * sizeof(int64_t)); + ggml_backend_tensor_set( + tree_sg.feature_commit_rows, feature_commit_rows.data(), 0, + feature_commit_rows.size() * sizeof(int32_t)); + + const size_t n_delta = b_.cache_.ssm_state.size(); + if (tree_sg.delta_captures.size() != n_delta || + b_.cache_.conv_state.size() != n_delta || + !tree_sg.tree_features || !b_.cache_.target_feat) { + result.error = "direct commit capture set incomplete"; + return result; + } + std::vector journals; + std::vector states; + std::vector conv_inputs; + std::vector conv_states; + journals.reserve(n_delta); + states.reserve(n_delta); + conv_inputs.reserve(n_delta); + conv_states.reserve(n_delta); + for (size_t layer = 0; layer < n_delta; ++layer) { + const DeltaNetCapture & capture = + tree_sg.delta_captures[layer]; + if (!capture.transition_journal || !capture.conv_input || + !b_.cache_.ssm_state[layer] || + !b_.cache_.conv_state[layer]) { + result.error = "direct commit layer capture incomplete"; + return result; + } + journals.push_back(capture.transition_journal); + states.push_back(b_.cache_.ssm_state[layer]); + conv_inputs.push_back(capture.conv_input); + conv_states.push_back(b_.cache_.conv_state[layer]); + } + std::vector cache_tensors; + cache_tensors.reserve( + b_.cache_.attn_k.size() + b_.cache_.attn_v.size()); + for (ggml_tensor * tensor : b_.cache_.attn_k) { + if (tensor) cache_tensors.push_back(tensor); + } + for (ggml_tensor * tensor : b_.cache_.attn_v) { + if (tensor) cache_tensors.push_back(tensor); + } + if (cache_tensors.empty()) { + result.error = "direct commit K/V cache set empty"; + return result; + } + + t_replay_build_end = timing_clock::now(); + if (!ggml_backend_cuda_tree_cache_commit_many( + cache_tensors.data(), + static_cast(cache_tensors.size()), + tree_sg.commit_rows, tree_sg.commit_slot_ids, + tree_scratch_base_, tree_scratch_stride_)) { + result.error = "direct commit K/V promotion failed"; + return result; + } + if (!ggml_backend_cuda_tree_feature_commit( + tree_sg.tree_features, b_.cache_.target_feat, + tree_sg.feature_commit_rows)) { + result.error = "direct commit target-feature promotion failed"; + return result; + } + if (!ggml_backend_cuda_gdn_transition_journal_commit_many( + journals.data(), states.data(), conv_inputs.data(), + conv_states.data(), static_cast(n_delta), + tree_sg.accepted_prefixes, tree_sg.commit_slot_ids)) { + result.error = "direct commit recurrent-state promotion failed"; + return result; + } + ggml_backend_tensor_set( + b_.cache_.paged_kv_seq_lens, seq_lens_.data(), 0, + sizeof(int32_t) * seq_lens_.size()); + t_replay_exec_end = timing_clock::now(); + + std::vector write_slots; + write_slots.reserve(inputs.size()); + for (size_t i = 0; i < inputs.size(); ++i) { + if (chain_lane_executes(lane_disposition(i))) { + write_slots.push_back(inputs[i].slot); + } + } + if (!commit_residency_writes(write_slots)) { + result.error = "direct commit residency write failed"; + return result; + } + for (size_t i = 0; i < inputs.size(); ++i) { + if (chain_lane_executes(lane_disposition(i))) { + slots_.commit_step(inputs[i].slot); + } + } + for (int lane = 0; lane < spec_count; ++lane) { + Proposal & proposal = proposals[static_cast(lane)]; + const int graph_row = spec_row_offset + lane * V + + proposal.accepted.back(); + proposal.pending = sample_graph_row( + proposal.slot, graph_row, + &posterior[static_cast(graph_row)], &logits_buf_); + if (proposal.pending < 0) { + result.error = "direct commit speculative sampling failed"; + return result; + } + } + for (int ar_index = 0; ar_index < ar_count; ++ar_index) { + ArLane & ar = ar_lanes[static_cast(ar_index)]; + const int graph_row = ar_index; + ar.pending = sample_graph_row( + ar.slot, graph_row, + &posterior[static_cast(graph_row)], &logits_buf_); + if (ar.pending < 0) { + result.error = "direct commit AR sampling failed"; + return result; + } + } + t_sample_end = timing_clock::now(); + + for (size_t i = 0; i < inputs.size(); ++i) { + if (!chain_lane_executes(lane_disposition(i))) continue; + std::string reselect_error; + if (!maybe_reselect_residency(inputs[i].slot, reselect_error)) { + result.error = reselect_error.empty() + ? "KVFlash reselect failed" : reselect_error; + return result; + } + } + result.decode.reserve(inputs.size()); + for (size_t i = 0; i < inputs.size(); ++i) { + DecodeOutput out; + out.slot = inputs[i].slot; + const ChainLaneDisposition disposition = lane_disposition(i); + if (disposition == ChainLaneDisposition::Failed) { + out.failed = true; + out.error = proposal_errors[i]; + result.decode.push_back(std::move(out)); + continue; + } + if (disposition == ChainLaneDisposition::Speculation) { + Proposal & proposal = + proposals[static_cast(proposal_for_input[i])]; + out.token = proposal.pending; + out.spec_steps = 1; + out.spec_accepted_tokens = proposal.path.size() > 1 + ? static_cast(proposal.path.size() - 1) : 0; + out.target_forwards = 1; + out.committed_tokens.assign( + proposal.path.begin() + 1, proposal.path.end()); + } else { + ArLane & ar = + ar_lanes[static_cast(ar_for_input[i])]; + out.token = ar.pending; + out.target_forwards = 1; + } + attach_residency_telemetry(out); + result.decode.push_back(std::move(out)); + } + if (timing) { + const auto t_round_end = timing_clock::now(); + std::fprintf(stderr, + "[step-timing] {\"path\":\"spec-direct\",\"live\":%d," + "\"k\":%d,\"tree_bucket\":%d,\"tree_rows\":%d," + "\"replay_rows\":0,\"ar_lanes\":%d,\"ar_bucket\":0," + "\"max_kv_len\":%d,\"draft_us\":%.1f," + "\"draft_lanes\":%d,\"pre_us\":%.1f," + "\"verify_build_us\":%.1f,\"verify_exec_us\":%.1f," + "\"posterior_read_us\":%.1f,\"commit_cpu_us\":%.1f," + "\"replay_build_us\":%.1f,\"replay_exec_us\":%.1f," + "\"sample_read_us\":%.1f,\"finish_us\":%.1f," + "\"total_us\":%.1f,\"accepted_tokens\":%d," + "\"emitted_tokens\":%d,\"target_forwards\":%d}\n", + spec_count + ar_count, spec_count, tree_bucket, + total_tree, ar_count, max_kv_len, + round_draft_us_, round_draft_lanes_, + span_us(t_round_start, t_verify_build_start), + span_us(t_verify_build_start, t_verify_build_end), + span_us(t_verify_build_end, t_verify_exec_end), + span_us(t_verify_exec_end, t_posterior_end), + span_us(t_posterior_end, t_commit_end), + span_us(t_commit_end, t_replay_build_end), + span_us(t_replay_build_end, t_replay_exec_end), + span_us(t_replay_exec_end, t_sample_end), + span_us(t_sample_end, t_round_end), + span_us(t_round_start, t_round_end), + replay_total - spec_count, replay_total + ar_count, + spec_count + ar_count); + } + return result; + } + + // Launch 2: accepted path segments + compact AR rows in the same builder + // combination already used by mixed prefill/decode. + const int ar_bucket = chain_decode_bucket_width(ar_count); + const int n_total = replay_total + ar_bucket; + const int gather_rows = spec_count + ar_bucket; + const bool has_replay = replay_total > 0; + StepGraph & durable_sg = b_.sg_; + if (!build_target_step( + durable_sg, b_.w_, b_.cache_, b_.target_backend_, + 0, n_total, false, true, false, 0, 0, + b_.cfg_.kq_stride_pad, false, false, false, true, + ar_bucket > 0 ? ar_bucket : 1, 0, max_kv_len, + replay_total, replay_segments.data(), + static_cast(replay_segments.size()), + gather_rows, ar_bucket > 0) || + !durable_sg.kv_write_rows || !durable_sg.target_feat_rows || + (has_replay && + (!durable_sg.paged_query_seq_ids || + !durable_sg.paged_query_positions)) || + (ar_bucket > 0 && + (!durable_sg.active_slot_ids || !durable_sg.state_slot_ids)) || + !durable_sg.logits_row_indices || !durable_sg.argmax_tokens) { + result.error = "chain mixed commit/AR graph build failed"; + return result; + } + t_replay_build_end = timing_clock::now(); + + std::vector durable_tokens(static_cast(n_total), 0); + std::copy(replay_tokens.begin(), replay_tokens.end(), + durable_tokens.begin()); + for (int lane = 0; lane < ar_count; ++lane) { + durable_tokens[static_cast(replay_total + lane)] = + ar_lanes[static_cast(lane)].token; + } + embed_buf_.resize(static_cast(hidden) * n_total); + if (!b_.w_.embedder.embed( + durable_tokens.data(), n_total, embed_buf_.data())) { + result.error = "chain mixed commit/AR embedding failed"; + return result; + } + ggml_backend_tensor_set( + durable_sg.inp_embed, embed_buf_.data(), 0, + sizeof(float) * embed_buf_.size()); + + pos_buf_.assign(static_cast(4) * n_total, 0); + for (int row = 0; row < replay_total; ++row) { + const int position = replay_positions[static_cast(row)]; + pos_buf_[static_cast(0) * n_total + row] = position; + pos_buf_[static_cast(1) * n_total + row] = position; + pos_buf_[static_cast(2) * n_total + row] = position; + } + for (int lane = 0; lane < ar_count; ++lane) { + const int row = replay_total + lane; + const int position = ar_lanes[static_cast(lane)].position; + pos_buf_[static_cast(0) * n_total + row] = position; + pos_buf_[static_cast(1) * n_total + row] = position; + pos_buf_[static_cast(2) * n_total + row] = position; + } + ggml_backend_tensor_set(durable_sg.positions, pos_buf_.data(), 0, + sizeof(int32_t) * pos_buf_.size()); + + rows_buf_.assign( + static_cast(n_total) * n_head_kv, scratch_row_); + for (int head = 0; head < n_head_kv; ++head) { + for (int row = 0; row < replay_total; ++row) { + rows_buf_[static_cast(head) * n_total + row] = + replay_physical[static_cast(row)]; + } + for (int lane = 0; lane < ar_count; ++lane) { + rows_buf_[static_cast(head) * n_total + + replay_total + lane] = + ar_lanes[static_cast(lane)].physical_row; + } + } + ggml_backend_tensor_set( + durable_sg.kv_write_rows, rows_buf_.data(), 0, + sizeof(int64_t) * rows_buf_.size()); + + query_slot_ids_.assign(static_cast(n_total), -1); + query_positions_.assign(static_cast(n_total), -1); + for (int row = 0; row < replay_total; ++row) { + query_slot_ids_[static_cast(row)] = + replay_slots[static_cast(row)]; + query_positions_[static_cast(row)] = + replay_positions[static_cast(row)]; + } + for (int lane = 0; lane < ar_count; ++lane) { + const int row = replay_total + lane; + query_slot_ids_[static_cast(row)] = + ar_lanes[static_cast(lane)].slot; + query_positions_[static_cast(row)] = + ar_lanes[static_cast(lane)].position; + } + logits_rows_.clear(); + logits_rows_.reserve(static_cast(gather_rows)); + int path_end = 0; + for (const Proposal & proposal : proposals) { + path_end += static_cast(proposal.path.size()); + logits_rows_.push_back(path_end - 1); + } + for (int lane = 0; lane < ar_bucket; ++lane) { + logits_rows_.push_back(replay_total + lane); + } + if (has_replay) { + ggml_backend_tensor_set( + durable_sg.paged_query_seq_ids, query_slot_ids_.data(), 0, + sizeof(int32_t) * query_slot_ids_.size()); + ggml_backend_tensor_set( + durable_sg.paged_query_positions, query_positions_.data(), 0, + sizeof(int32_t) * query_positions_.size()); + } + ggml_backend_tensor_set( + durable_sg.logits_row_indices, logits_rows_.data(), 0, + sizeof(int32_t) * logits_rows_.size()); + + active_slot_ids_.assign(static_cast(ar_bucket), -1); + state_slot_ids_.assign(static_cast(ar_bucket), 0); + for (int lane = 0; lane < ar_count; ++lane) { + active_slot_ids_[static_cast(lane)] = + ar_lanes[static_cast(lane)].slot; + state_slot_ids_[static_cast(lane)] = + ar_lanes[static_cast(lane)].slot; + } + if (ar_bucket > 0) { + ggml_backend_tensor_set( + durable_sg.active_slot_ids, active_slot_ids_.data(), 0, + sizeof(int32_t) * active_slot_ids_.size()); + ggml_backend_tensor_set( + durable_sg.state_slot_ids, state_slot_ids_.data(), 0, + sizeof(int32_t) * state_slot_ids_.size()); + } + + const int feature_cap = b_.cache_.target_feat_cap; + const int dead_feature_row = feature_cap * n_slots; + feature_rows_.assign( + static_cast(n_total), dead_feature_row); + for (int row = 0; row < replay_total; ++row) { + feature_rows_[static_cast(row)] = + replay_slots[static_cast(row)] * feature_cap + + replay_positions[static_cast(row)] % feature_cap; + } + for (int lane = 0; lane < ar_count; ++lane) { + const ArLane & ar = ar_lanes[static_cast(lane)]; + feature_rows_[static_cast(replay_total + lane)] = + ar.slot * feature_cap + ar.position % feature_cap; + } + ggml_backend_tensor_set( + durable_sg.target_feat_rows, feature_rows_.data(), 0, + sizeof(int32_t) * feature_rows_.size()); + ggml_backend_tensor_set( + b_.cache_.paged_kv_seq_lens, seq_lens_.data(), 0, + sizeof(int32_t) * seq_lens_.size()); + + if (ggml_backend_graph_compute(b_.target_backend_, durable_sg.gf) != + GGML_STATUS_SUCCESS) { + result.error = "chain mixed commit/AR compute failed"; + return result; + } + t_replay_exec_end = timing_clock::now(); + + argmax_buf_.assign(static_cast(gather_rows), -1); + ggml_backend_tensor_get_async( + b_.target_backend_, durable_sg.argmax_tokens, + argmax_buf_.data(), 0, + sizeof(int32_t) * argmax_buf_.size()); + ggml_backend_synchronize(b_.target_backend_); + for (int lane = 0; lane < spec_count; ++lane) { + if (argmax_buf_[static_cast(lane)] < 0) { + result.error = "chain durable replay produced invalid token"; + return result; + } + } + for (int lane = 0; lane < ar_count; ++lane) { + if (argmax_buf_[static_cast(spec_count + lane)] < 0) { + result.error = "mixed AR durable step produced invalid token"; + return result; + } + } + + std::vector write_slots; + write_slots.reserve(inputs.size()); + for (size_t i = 0; i < inputs.size(); ++i) { + if (chain_lane_executes(lane_disposition(i))) { + write_slots.push_back(inputs[i].slot); + } + } + if (!commit_residency_writes(write_slots)) { + result.error = "chain mixed-step KV write commit failed"; + return result; + } + + // Publish the fed root/path before sampling the next token, matching the + // ordinary AR path's penalty history, RNG, and min-token-floor semantics. + for (size_t i = 0; i < inputs.size(); ++i) { + if (chain_lane_executes(lane_disposition(i))) { + slots_.commit_step(inputs[i].slot); + } + } + for (int lane = 0; lane < spec_count; ++lane) { + Proposal & proposal = proposals[static_cast(lane)]; + proposal.pending = sample_graph_row( + proposal.slot, lane, + &argmax_buf_[static_cast(lane)], &logits_buf_); + if (proposal.pending < 0) { + result.error = "chain durable replay sampling failed"; + return result; + } + } + for (int lane = 0; lane < ar_count; ++lane) { + ArLane & ar = ar_lanes[static_cast(lane)]; + const int gathered_row = spec_count + lane; + ar.pending = sample_graph_row( + ar.slot, gathered_row, + &argmax_buf_[static_cast(gathered_row)], &logits_buf_); + if (ar.pending < 0) { + result.error = "mixed AR durable sampling failed"; + return result; + } + } + t_sample_end = timing_clock::now(); + + for (size_t i = 0; i < inputs.size(); ++i) { + if (!chain_lane_executes(lane_disposition(i))) continue; + std::string reselect_error; + if (!maybe_reselect_residency(inputs[i].slot, reselect_error)) { + result.error = reselect_error.empty() + ? "KVFlash reselect failed" : reselect_error; + return result; + } + } + + result.decode.reserve(inputs.size()); + for (size_t i = 0; i < inputs.size(); ++i) { + DecodeOutput out; + out.slot = inputs[i].slot; + const ChainLaneDisposition disposition = lane_disposition(i); + if (disposition == ChainLaneDisposition::Failed) { + out.failed = true; + out.error = proposal_errors[i]; + result.decode.push_back(std::move(out)); + continue; + } + if (disposition == ChainLaneDisposition::Speculation) { + Proposal & proposal = + proposals[static_cast(proposal_for_input[i])]; + out.token = proposal.pending; + out.spec_steps = 1; + out.spec_accepted_tokens = + proposal.path.size() > 1 + ? static_cast(proposal.path.size() - 1) + : 0; + out.target_forwards = 2; + out.committed_tokens.assign( + proposal.path.begin() + 1, proposal.path.end()); + } else { + ArLane & ar = + ar_lanes[static_cast(ar_for_input[i])]; + out.token = ar.pending; + out.target_forwards = 1; + } + attach_residency_telemetry(out); + result.decode.push_back(std::move(out)); + } + if (timing) { + const auto t_round_end = timing_clock::now(); + std::fprintf(stderr, + "[step-timing] {\"path\":\"spec\",\"live\":%d,\"k\":%d," + "\"tree_bucket\":%d,\"tree_rows\":%d,\"replay_rows\":%d," + "\"ar_lanes\":%d,\"ar_bucket\":%d,\"max_kv_len\":%d," + "\"draft_us\":%.1f,\"draft_lanes\":%d," + "\"pre_us\":%.1f,\"verify_build_us\":%.1f," + "\"verify_exec_us\":%.1f,\"posterior_read_us\":%.1f," + "\"commit_cpu_us\":%.1f,\"replay_build_us\":%.1f," + "\"replay_exec_us\":%.1f,\"sample_read_us\":%.1f," + "\"finish_us\":%.1f,\"total_us\":%.1f," + "\"accepted_tokens\":%d,\"emitted_tokens\":%d," + "\"target_forwards\":%d}\n", + spec_count + ar_count, spec_count, + tree_bucket, V * tree_bucket, replay_total, + ar_count, ar_bucket, max_kv_len, + round_draft_us_, round_draft_lanes_, + span_us(t_round_start, t_verify_build_start), + span_us(t_verify_build_start, t_verify_build_end), + span_us(t_verify_build_end, t_verify_exec_end), + span_us(t_verify_exec_end, t_posterior_end), + span_us(t_posterior_end, t_commit_end), + span_us(t_commit_end, t_replay_build_end), + span_us(t_replay_build_end, t_replay_exec_end), + span_us(t_replay_exec_end, t_sample_end), + span_us(t_sample_end, t_round_end), + span_us(t_round_start, t_round_end), + replay_total - spec_count, replay_total + ar_count, + 2 * spec_count + ar_count); + } + return result; +} + + +std::optional Qwen35SeqEngine::step_ddtree( + const StepPlan & plan) { + StepResult result; + const int active = (int)plan.decode.size(); + const int bucket = decode_bucket_width(active); + const int T = tree_width_; + const int q_len = b_.dw_.block_size; + const int hidden = b_.w_.n_embd; + const int n_head_kv = b_.w_.n_head_kv; + const int n_slots = slots_.slot_count(); + const int K = b_.cfg_.ddtree_budget > q_len - 1 ? 8 : 1; + + struct Proposal { + int slot = -1; + int32_t root = -1; + DDTree tree; + std::vector flat; + std::vector accepted; + int32_t bonus = -1; + }; + std::vector proposals; + proposals.reserve((size_t)active); + std::vector noise((size_t)q_len, b_.w_.mask_token_id); + std::vector noise_embed((size_t)hidden * q_len); + std::vector logits((size_t)b_.w_.n_vocab * q_len); + std::vector top_lp((size_t)q_len * K); + std::vector top_ids((size_t)q_len * K); + + auto proposal_fallback = [&]() -> std::optional { + // begin_step updates persistent drafter bookkeeping before compute. + // A failed proposal graph may therefore leave only part of that + // cache valid. Reset all participating draft rings so a later + // speculative round rebuilds them from committed target features. + for (const StepInput & in : plan.decode) { + if (in.slot >= 0 && in.slot < (int)slot_draft_kv_.size() && + slot_draft_kv_[(size_t)in.slot]) { + draft_kv_reset(*slot_draft_kv_[(size_t)in.slot]); + } + } + return std::nullopt; + }; + + if (!build_lm_head_projection_step( + b_.proj_sg_, b_.w_, b_.target_backend_, q_len)) { + return proposal_fallback(); + } + + // Proposal is sequential by slot: immutable draft weights are shared, + // while each slot owns an independent persistent context-KV ring. + for (const StepInput & in : plan.decode) { + DraftKvState * draft = ensure_slot_draft_kv(in.slot); + DraftFeatureMirror * mirror = slot_feature_mirror(in.slot); + if (!draft || !mirror) return proposal_fallback(); + noise[0] = in.token; + std::fill(noise.begin() + 1, noise.end(), b_.w_.mask_token_id); + if (!b_.w_.embedder.embed( + noise.data(), q_len, noise_embed.data()) || + !draft_kv_begin_step(*draft, b_.dw_, b_.draft_backend_, + *mirror, slots_.slot(in.slot).cur_pos)) { + return proposal_fallback(); + } + ggml_backend_tensor_set( + draft->inp_embed, noise_embed.data(), 0, + sizeof(float) * noise_embed.size()); + if (ggml_backend_graph_compute(b_.draft_backend_, draft->gf) != + GGML_STATUS_SUCCESS) { + return proposal_fallback(); + } + // The draft and target backends own separate HIP streams even when + // both are placed on hip:0. Projection consumes the draft hidden + // state on the target stream, so establish the producer/consumer + // ordering explicitly before the cross-backend tensor copy. Without + // this barrier the first tree proposal can race stale hidden rows and + // collapse acceptance to the one-token fallback. + ggml_backend_synchronize(b_.draft_backend_); + ggml_backend_tensor_copy( + draft->hidden_states, b_.proj_sg_.hidden_input); + if (ggml_backend_graph_compute( + b_.target_backend_, b_.proj_sg_.gf) != GGML_STATUS_SUCCESS) { + return proposal_fallback(); + } + bool topk_ready = false; +#ifdef DFLASH27B_HAVE_DRAFT_TOPK + topk_ready = geometric_extract_draft_topk_cuda( + b_.proj_sg_.logits->data, q_len, b_.w_.n_vocab, K, + top_lp.data(), top_ids.data(), b_.cfg_.ddtree_temp); +#endif + if (!topk_ready) { + ggml_backend_tensor_get( + b_.proj_sg_.logits, logits.data(), 0, + sizeof(float) * logits.size()); + extract_draft_topk( + logits.data(), q_len, b_.w_.n_vocab, K, + top_lp.data(), top_ids.data(), b_.cfg_.ddtree_temp); + } + + Proposal p; + p.slot = in.slot; + p.root = in.token; + p.tree = build_ddtree( + top_lp.data() + K, top_ids.data() + K, + q_len - 1, K, b_.cfg_.ddtree_budget, + b_.cfg_.ddtree_chain_seed); + p.flat.assign((size_t)T, 0); + p.flat[0] = in.token; + for (int node = 0; node < p.tree.n_nodes; ++node) { + p.flat[(size_t)node + 1] = p.tree.token_ids[(size_t)node]; + } + proposals.push_back(std::move(p)); + } + + StepGraph & tree_sg = b_.sg_; + int max_prefix = 1; + for (const Proposal & p : proposals) { + max_prefix = std::max(max_prefix, slots_.slot(p.slot).cur_pos); + } + if (!build_target_step_paged_tree( + tree_sg, b_.w_, b_.cache_, b_.target_backend_, T, bucket, + max_prefix, tree_scratch_base_, tree_scratch_stride_, + b_.cfg_.kq_stride_pad)) { + result.error = "packed DDTree verify graph build failed"; + return result; + } + + const int total_tree = T * bucket; + std::vector flat_tokens((size_t)total_tree, 0); + std::vector parents((size_t)total_tree, -1); + std::vector sizes((size_t)bucket, 0); + // Negative active IDs identify bucket padding. Recurrent gathers cannot + // index a negative slab, so padded trees use slot 0 only for their + // read-only base-state gather; tree_size=0/query_slot=-1 keeps all of + // their attention/output rows inactive and tree mode never persists it. + std::vector tree_slots((size_t)bucket, -1); + std::vector tree_state_slots((size_t)bucket, 0); + std::vector query_slots((size_t)total_tree, -1); + std::vector tree_rows( + (size_t)total_tree * n_head_kv, scratch_row_); + std::vector tree_pos((size_t)4 * total_tree, 0); + std::vector tree_embed((size_t)hidden * total_tree, 0.0f); + seq_lens_.assign((size_t)n_slots, 0); + + for (int s = 0; s < active; ++s) { + const Proposal & p = proposals[(size_t)s]; + const int base = s * T; + sizes[(size_t)s] = p.tree.n_nodes + 1; + tree_slots[(size_t)s] = p.slot; + tree_state_slots[(size_t)s] = p.slot; + seq_lens_[(size_t)p.slot] = slots_.slot(p.slot).cur_pos; + for (int node = 0; node < sizes[(size_t)s]; ++node) { + const int row = base + node; + flat_tokens[(size_t)row] = p.flat[(size_t)node]; + parents[(size_t)row] = node == 0 ? -1 : + p.tree.parents[(size_t)node]; + query_slots[(size_t)row] = p.slot; + const int depth = node == 0 ? 0 : + p.tree.depths[(size_t)node - 1]; + const int pos = slots_.slot(p.slot).cur_pos + depth; + tree_pos[(size_t)0 * total_tree + row] = pos; + tree_pos[(size_t)1 * total_tree + row] = pos; + tree_pos[(size_t)2 * total_tree + row] = pos; + for (int h = 0; h < n_head_kv; ++h) { + tree_rows[(size_t)h * total_tree + row] = + (int64_t)tree_scratch_base_ + + (int64_t)p.slot * tree_scratch_stride_ + node; + } + } + } + if (!b_.w_.embedder.embed( + flat_tokens.data(), total_tree, tree_embed.data())) { + result.error = "packed DDTree embedding failed"; + return result; + } + ggml_backend_tensor_set(tree_sg.inp_embed, tree_embed.data(), 0, + sizeof(float) * tree_embed.size()); + ggml_backend_tensor_set(tree_sg.positions, tree_pos.data(), 0, + sizeof(int32_t) * tree_pos.size()); + ggml_backend_tensor_set(tree_sg.parent_ids, parents.data(), 0, + sizeof(int32_t) * parents.size()); + ggml_backend_tensor_set(tree_sg.tree_sizes, sizes.data(), 0, + sizeof(int32_t) * sizes.size()); + // Mapped-tree DeltaNet uses active_slot_ids only as a topology marker; + // gallocr may therefore optimize away its backing buffer. The actual + // state/attention mappings below are live graph inputs and remain + // mandatory. Upload the marker only if a future topology consumes it. + if (detail::target_paged_tree_active_slots_need_upload(tree_sg)) { + ggml_backend_tensor_set(tree_sg.active_slot_ids, tree_slots.data(), 0, + sizeof(int32_t) * tree_slots.size()); + } + ggml_backend_tensor_set(tree_sg.state_slot_ids, tree_state_slots.data(), 0, + sizeof(int32_t) * tree_state_slots.size()); + ggml_backend_tensor_set(tree_sg.paged_query_seq_ids, query_slots.data(), 0, + sizeof(int32_t) * query_slots.size()); + ggml_backend_tensor_set(tree_sg.kv_write_rows, tree_rows.data(), 0, + sizeof(int64_t) * tree_rows.size()); + ggml_backend_tensor_set(b_.cache_.paged_kv_seq_lens, seq_lens_.data(), 0, + sizeof(int32_t) * seq_lens_.size()); + if (ggml_backend_graph_compute(b_.target_backend_, tree_sg.gf) != + GGML_STATUS_SUCCESS) { + result.error = "packed DDTree verify compute failed"; + return result; + } + std::vector posterior((size_t)total_tree, -1); + ggml_backend_tensor_get(tree_sg.argmax_tokens, posterior.data(), 0, + sizeof(int32_t) * posterior.size()); + + int replay_total = 0; + for (int s = 0; s < active; ++s) { + Proposal & p = proposals[(size_t)s]; + p.accepted = follow_verified_tree( + p.tree, posterior.data() + (size_t)s * T, p.bonus); + const int room = slots_.max_context() - slots_.slot(p.slot).cur_pos; + truncate_verified_path( + p.accepted, (size_t)std::max(0, room), + posterior.data() + (size_t)s * T, p.bonus); + if (p.accepted.empty()) { + result.error = "DDTree accepted path has no context headroom"; + return result; + } + replay_total += (int)p.accepted.size(); + } + + std::vector replay_segments; + std::vector replay_tokens; + std::vector replay_slots; + std::vector replay_positions; + std::vector replay_rows; + std::vector replay_logits_rows; + replay_segments.reserve((size_t)active); + replay_tokens.reserve((size_t)replay_total); + replay_slots.reserve((size_t)replay_total); + replay_positions.reserve((size_t)replay_total); + replay_rows.assign((size_t)replay_total * n_head_kv, scratch_row_); + replay_logits_rows.reserve((size_t)active); + seq_lens_.assign((size_t)n_slots, 0); + + int replay_offset = 0; + for (Proposal & p : proposals) { + std::vector path; + path.reserve(p.accepted.size()); + for (int dfs : p.accepted) { + path.push_back(dfs == 0 ? p.root : + p.tree.token_ids[(size_t)dfs - 1]); + } + const Qwen35SlotManager::StepAppend app = slots_.append_tokens( + p.slot, path.data(), (int)path.size()); + const bool table_ok = slots_.residency_active() || + upload_block_table_delta(p.slot, app.first_new_block, + app.new_blocks.data(), app.new_blocks.size()); + if (!app.ok || app.physical_rows.size() != path.size() || + !table_ok) { + result.error = app.busy + ? "paged KV pool exhausted during DDTree replay" + : "DDTree replay K/V append failed"; + return result; + } + replay_segments.push_back( + {replay_offset, (int)path.size(), p.slot}); + for (size_t i = 0; i < path.size(); ++i) { + replay_tokens.push_back(path[i]); + replay_slots.push_back(p.slot); + replay_positions.push_back(app.position + (int)i); + for (int h = 0; h < n_head_kv; ++h) { + replay_rows[(size_t)h * replay_total + replay_offset + i] = + app.physical_rows[i]; + } + } + replay_offset += (int)path.size(); + replay_logits_rows.push_back(replay_offset - 1); + seq_lens_[(size_t)p.slot] = app.position + (int)path.size(); + } + + if (!upload_all_active_block_tables()) { + result.error = "DDTree replay block-table refresh failed"; + return result; + } + + StepGraph & replay_sg = b_.sg_; + if (!build_target_step( + replay_sg, b_.w_, b_.cache_, b_.target_backend_, + 0, replay_total, false, true, false, 0, 0, + b_.cfg_.kq_stride_pad, false, false, false, true, + 1, 0, *std::max_element(seq_lens_.begin(), seq_lens_.end()), + replay_total, replay_segments.data(), + (int)replay_segments.size(), active, false) || + !replay_sg.target_feat_rows || !replay_sg.paged_query_seq_ids || + !replay_sg.paged_query_positions || !replay_sg.logits_row_indices || + !replay_sg.argmax_tokens) { + result.error = "DDTree accepted-path replay graph build failed"; + return result; + } + embed_buf_.resize((size_t)hidden * replay_total); + if (!b_.w_.embedder.embed( + replay_tokens.data(), replay_total, embed_buf_.data())) { + result.error = "DDTree replay embedding failed"; + return result; + } + pos_buf_.assign((size_t)4 * replay_total, 0); + feature_rows_.resize((size_t)replay_total); + const int cap = b_.cache_.target_feat_cap; + for (int row = 0; row < replay_total; ++row) { + const int pos = replay_positions[(size_t)row]; + pos_buf_[(size_t)0 * replay_total + row] = pos; + pos_buf_[(size_t)1 * replay_total + row] = pos; + pos_buf_[(size_t)2 * replay_total + row] = pos; + feature_rows_[(size_t)row] = + replay_slots[(size_t)row] * cap + pos % cap; + } + ggml_backend_tensor_set(replay_sg.inp_embed, embed_buf_.data(), 0, + sizeof(float) * embed_buf_.size()); + ggml_backend_tensor_set(replay_sg.positions, pos_buf_.data(), 0, + sizeof(int32_t) * pos_buf_.size()); + ggml_backend_tensor_set(replay_sg.kv_write_rows, replay_rows.data(), 0, + sizeof(int64_t) * replay_rows.size()); + ggml_backend_tensor_set(replay_sg.paged_query_seq_ids, + replay_slots.data(), 0, + sizeof(int32_t) * replay_slots.size()); + ggml_backend_tensor_set(replay_sg.paged_query_positions, + replay_positions.data(), 0, + sizeof(int32_t) * replay_positions.size()); + ggml_backend_tensor_set(replay_sg.logits_row_indices, + replay_logits_rows.data(), 0, + sizeof(int32_t) * replay_logits_rows.size()); + ggml_backend_tensor_set(replay_sg.target_feat_rows, + feature_rows_.data(), 0, + sizeof(int32_t) * feature_rows_.size()); + ggml_backend_tensor_set(b_.cache_.paged_kv_seq_lens, seq_lens_.data(), 0, + sizeof(int32_t) * seq_lens_.size()); + if (ggml_backend_graph_compute(b_.target_backend_, replay_sg.gf) != + GGML_STATUS_SUCCESS) { + result.error = "DDTree accepted-path replay compute failed"; + return result; + } + std::vector replay_write_slots; + replay_write_slots.reserve(proposals.size()); + for (const Proposal & p : proposals) replay_write_slots.push_back(p.slot); + if (!commit_residency_writes(replay_write_slots)) { + result.error = "DDTree replay KV write commit failed"; + return result; + } + + // The replay is the durable target forward: its recurrent/KV/feature + // state is what the next step consumes. Use its posterior rather than + // the tree-verify posterior so the pending scalar remains exact even if + // the two graph shapes differ numerically. + std::vector replay_next((size_t)active, -1); + ggml_backend_tensor_get( + replay_sg.argmax_tokens, replay_next.data(), 0, + sizeof(int32_t) * replay_next.size()); + for (int s = 0; s < active; ++s) { + if (replay_next[(size_t)s] < 0) { + result.error = "DDTree replay produced an invalid pending token"; + return result; + } + proposals[(size_t)s].bonus = replay_next[(size_t)s]; + } + + for (Proposal & p : proposals) { + slots_.commit_step(p.slot); + std::string reselect_error; + if (!maybe_reselect_residency(p.slot, reselect_error)) { + result.error = reselect_error.empty() + ? "KVFlash reselect failed" : reselect_error; + return result; + } + } + + uint64_t cohort_emitted = 0; + for (const Proposal & p : proposals) { + // accepted contains the replay root plus accepted children. The + // output emits those children plus one separately computed pending + // scalar, so accepted.size() is this request's emitted yield. + cohort_emitted += (uint64_t)p.accepted.size(); + } + const bool suspend_cohort = + Qwen35SlotManager::ddtree_cohort_should_suspend( + cohort_emitted, active); + std::vector newly_suspended((size_t)slots_.slot_count(), false); + for (const Proposal & p : proposals) { + newly_suspended[(size_t)p.slot] = + slots_.record_ddtree_sample(p.slot, suspend_cohort); + } + + result.decode.reserve((size_t)active); + for (Proposal & p : proposals) { + DecodeOutput out; + out.slot = p.slot; + out.token = p.bonus; + out.ddtree_steps = 1; + const int accepted_children = (int)p.accepted.size() - 1; + out.ddtree_accepted_tokens = (uint64_t)accepted_children; + out.target_forwards = 2; + for (size_t i = 1; i < p.accepted.size(); ++i) { + const int dfs = p.accepted[i]; + out.committed_tokens.push_back( + p.tree.token_ids[(size_t)dfs - 1]); + } + if (newly_suspended[(size_t)p.slot]) { + out.ddtree_suspensions = 1; + const Qwen35Slot & seq = slots_.slot(p.slot); + std::fprintf(stderr, + "[parallel-ddtree] adaptive suspend request=%llu slot=%d " + "sample=%llu emitted=%d accepted_children=%d " + "cohort_emitted=%llu cohort_size=%d target_forwards=2 " + "floor=%d\n", + (unsigned long long)seq.request_id, p.slot, + (unsigned long long)seq.ddtree_sampled_steps, + accepted_children + 1, accepted_children, + (unsigned long long)cohort_emitted, active, + Qwen35SlotManager::kDdtreeMinEmittedTokens); + } + attach_residency_telemetry(out); + result.decode.push_back(std::move(out)); + } + return result; +} + bool Qwen35SeqEngine::token_is_eos(int32_t token) const { return b_.token_is_eos(token); } @@ -46,6 +2775,26 @@ SeqEngine::AdmitResult Qwen35SeqEngine::admit( AdmitResult result = slots_.admit(request_id, prompt, sampler); if (result.status == AdmitResult::Status::admitted) { reset_recurrent_slot(b_.cache_, result.slot); + if (slots_.residency_active()) { + slots_.slot(result.slot).kvflash_last_reselect_generated = + -std::max(1, b_.kvflash_tau_); + } + if (result.slot >= 0 && result.slot < (int)slot_draft_kv_.size() && + slot_draft_kv_[(size_t)result.slot]) { + draft_kv_reset(*slot_draft_kv_[(size_t)result.slot]); + } + if (result.slot >= 0 && + result.slot < (int)last_activation_estimate_.size()) { + last_activation_estimate_[(size_t)result.slot] = {}; + if (result.slot < + (int)prepared_chain_drafts_.size()) { + prepared_chain_drafts_[(size_t)result.slot].valid = false; + } + } + if (result.slot >= 0 && + result.slot < (int)adaptive_fallback_ar_.size()) { + adaptive_fallback_ar_[(size_t)result.slot] = 0; + } } return result; } @@ -102,6 +2851,77 @@ bool Qwen35SeqEngine::upload_block_table_delta( return true; } +bool Qwen35SeqEngine::upload_all_active_block_tables() { + if (!slots_.residency_active()) return true; + ggml_tensor * table = b_.cache_.paged_block_table; + if (!table) return false; + std::vector column((size_t)table->ne[0], -1); + std::vector snapshot; + for (int slot = 0; slot < slots_.slot_count(); ++slot) { + if (!slots_.is_active(slot)) continue; + std::fill(column.begin(), column.end(), -1); + if (!slots_.block_table_snapshot(slot, snapshot) || + snapshot.size() > column.size()) { + return false; + } + std::copy(snapshot.begin(), snapshot.end(), column.begin()); + ggml_backend_tensor_set( + table, column.data(), (size_t)slot * table->nb[1], + sizeof(int32_t) * column.size()); + } + return true; +} + +bool Qwen35SeqEngine::commit_residency_writes( + const std::vector & slots) { + if (!slots_.residency_active()) return true; + // Graph completion is not a host barrier for every backend. Pending pages + // become evictable only after all target writes are device-complete. + ggml_backend_synchronize(b_.target_backend_); + for (int slot : slots) { + if (!slots_.commit_residency_writes(slot)) return false; + } + return true; +} + +bool Qwen35SeqEngine::maybe_reselect_residency( + int slot, std::string & error) { + if (!slots_.residency_active()) return true; + Qwen35Slot & seq = slots_.slot(slot); + const int generated = seq.generated_tokens(); + const int tau = std::max( + b_.kvflash_tau_, (int)(seq.sample_history.size() / 45)); + if (generated - seq.kvflash_last_reselect_generated < tau) return true; + + b_.kvflash_ensure_scorer(); + std::vector scores; + const std::vector * score_ptr = nullptr; + if (b_.kvflash_scorer_) { + if (!b_.kvflash_scorer_->score_chunks( + seq.sample_history, PAGED_BLOCK_SIZE, scores)) { + // Short histories and recoverable drafter failures are expected + // scorer outcomes. Preserve service with the pager's explicit + // recency/LRU policy; only residency or transfer errors below are + // fatal to the request. + std::fprintf(stderr, + "[parallel-kvflash] scorer unavailable for slot %d; using LRU\n", + slot); + } else { + const size_t blocks = (seq.sample_history.size() + + PAGED_BLOCK_SIZE - 1) / PAGED_BLOCK_SIZE; + scores.resize(blocks, scores.empty() ? 0.0f : scores.back()); + score_ptr = &scores; + } + } + if (!slots_.reselect_residency(slot, score_ptr, &error)) return false; + seq.kvflash_last_reselect_generated = generated; + return upload_all_active_block_tables(); +} + +void Qwen35SeqEngine::attach_residency_telemetry(DecodeOutput & out) { + slots_.take_residency_telemetry(out.slot, out); +} + void Qwen35SeqEngine::fail_prefill( int slot, std::vector & prefill_outputs, const char * log_message, const char * client_message) { @@ -135,9 +2955,10 @@ Qwen35SeqEngine::PrefillStage Qwen35SeqEngine::stage_prefill_chunk( "prefill K/V allocation failed"); return PrefillStage{}; } - if (!upload_block_table_delta( - slot, chunk.first_new_block, chunk.new_blocks.data(), - chunk.new_blocks.size())) { + const bool table_ok = slots_.residency_active() || upload_block_table_delta( + slot, chunk.first_new_block, chunk.new_blocks.data(), + chunk.new_blocks.size()); + if (!table_ok) { fail_prefill( slot, prefill_outputs, "prefill block-table delta exceeds device capacity", "prefill block-table update failed"); @@ -207,10 +3028,403 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { } if (inputs.empty() && plan.prefills.empty()) return result; + if (ddtree_eligible(plan)) { + std::optional speculative = step_ddtree(plan); + if (speculative) return std::move(*speculative); + // Proposal setup failed before target/cache mutation. Preserve service + // by taking the existing packed AR path for this iteration. + } + + // One common wall-clock origin makes pure AR, adaptive k=0, and admitted + // speculative rounds directly comparable and feeds the online cost model + // even when diagnostic phase logging is disabled. + const bool timing = step_timing_enabled(); + using timing_clock = std::chrono::steady_clock; + const bool gate_cost_timing = + spec_mode_ == SpecMode::chain && + speculation_gate_ != nullptr; + const auto decode_round_started = timing || gate_cost_timing + ? timing_clock::now() : timing_clock::time_point{}; + std::optional pending_ar_gate_plan; + std::vector spec_service_ar(inputs.size(), 0); + + if (spec_mode_ == SpecMode::chain && !inputs.empty()) { + // New chain round: restart the [step-timing] draft attribution. + round_draft_us_ = 0.0; + round_draft_lanes_ = 0; + const auto chain_started = decode_round_started; + std::vector admitted(inputs.size(), 0); + SpecPlan gate_plan; + bool have_gate_plan = false; + + const char * force_value = std::getenv("DFLASH_SPEC_GATE_FORCE"); + const std::string force = force_value ? force_value : ""; + + if (speculation_gate_) { + std::vector candidates; + candidates.reserve(inputs.size()); + const bool use_activation_score = activation_scoring_enabled(); + for (const StepInput & in : inputs) { + const Qwen35Slot & seq = slots_.slot(in.slot); + SpeculationPolicy policy = in.speculation_policy; + if (policy == SpeculationPolicy::Adaptive) { + if (force == "all") policy = SpeculationPolicy::Always; + if (force == "none") policy = SpeculationPolicy::Never; + } + // The activation-score-off arm is an explicit AR ablation. It must + // not pay the one-time activation draft. + if (!use_activation_score && + policy == SpeculationPolicy::Adaptive) { + policy = SpeculationPolicy::Never; + } + const bool scoreable = + chain_activation_input_scoreable(in); + const bool can_speculate = + chain_spec_request_capable(in); + double activation_score = + std::numeric_limits::quiet_NaN(); + std::vector activation_hazards; + if (use_activation_score && scoreable && in.slot >= 0 && + in.slot < (int)last_activation_estimate_.size() && + std::isfinite(last_activation_estimate_[(size_t)in.slot].expected_yield)) { + activation_score = + last_activation_estimate_[(size_t)in.slot].expected_yield; + activation_hazards = last_activation_estimate_[ + (size_t)in.slot].conditional_hazards; + } + candidates.push_back({ + seq.request_id, in.slot, policy, + scoreable, can_speculate, activation_score, + std::move(activation_hazards), + chain_activation_score_kind(), + }); + } + const bool cohort_changed = + !spec_cohort_epoch_.has_value() || + !spec_cohort_epoch_->matches(candidates); + bool published_epoch = false; + if (cohort_changed) { + gate_plan = speculation_gate_->plan( + (int)inputs.size(), candidates, (int)inputs.size()); + } else { + gate_plan = spec_cohort_epoch_->plan; + } + have_gate_plan = true; + if (!gate_plan.valid) { + return fail_step(gate_plan.error.empty() + ? "adaptive speculation gate failed" : gate_plan.error); + } + auto replan_with_published_score = [&]() { + for (SpecCandidate & candidate : candidates) { + if (!candidate.scoreable || + candidate.policy == SpeculationPolicy::Never) { + continue; + } + const int slot = candidate.slot; + if (slot >= 0 && + slot < (int)last_activation_estimate_.size() && + std::isfinite(last_activation_estimate_[(size_t)slot].expected_yield)) { + candidate.activation_yield = + last_activation_estimate_[(size_t)slot].expected_yield; + candidate.conditional_hazards = + last_activation_estimate_[(size_t)slot].conditional_hazards; + } + } + gate_plan = speculation_gate_->plan( + (int)inputs.size(), candidates, (int)inputs.size()); + return gate_plan.valid; + }; + + auto reset_evaluation_lane = [&](int slot) { + if (slot >= 0 && + slot < (int)prepared_chain_drafts_.size()) { + prepared_chain_drafts_[(size_t)slot] = {}; + } + if (slot >= 0 && slot < (int)slot_draft_kv_.size() && + slot_draft_kv_[(size_t)slot]) { + draft_kv_reset(*slot_draft_kv_[(size_t)slot]); + } + if (slot >= 0 && + slot < (int)last_activation_estimate_.size()) { + last_activation_estimate_[(size_t)slot] = {}; + } + }; + auto commit_evaluation_fallback = + [&](const SpecPendingEvaluation & evaluation) { + reset_evaluation_lane(evaluation.slot); + if (speculation_gate_->record_evaluation_failure( + evaluation.request_id)) { + const std::string kind = + chain_activation_score_kind(); + const char * reason = + "activation_evaluation_failed"; + log_spec_evaluation_fallback( + evaluation.request_id, evaluation.slot, + kind, reason); + } + }; + + // Resolve every one-time evaluation action before one immediate + // replan. A packed bootstrap failure is retried once per lane so + // one broken request is marked evaluation-failed without poisoning a + // healthy scored peer or the cohort. + if (!gate_plan.pending_evaluations.empty()) { + std::vector score_evaluations; + std::vector bootstrap(inputs.size(), 0); + for (const SpecPendingEvaluation & evaluation : + gate_plan.pending_evaluations) { + if (evaluation.action == + SpecEvaluationAction::FallbackAR) { + commit_evaluation_fallback(evaluation); + continue; + } + score_evaluations.push_back(evaluation); + for (size_t i = 0; i < inputs.size(); ++i) { + if (inputs[i].slot == evaluation.slot) { + bootstrap[i] = 1; + } + } + } + + if (use_activation_score && !score_evaluations.empty()) { + const bool batch_scored = prepare_chain_drafts( + inputs, bootstrap, /*force_serial=*/false, + /*fail_fast_batch=*/true); + if (!batch_scored) { + for (const SpecPendingEvaluation & evaluation : + score_evaluations) { + reset_evaluation_lane(evaluation.slot); + } + for (const SpecPendingEvaluation & evaluation : + score_evaluations) { + std::vector one(inputs.size(), 0); + for (size_t i = 0; i < inputs.size(); ++i) { + if (inputs[i].slot == evaluation.slot) { + one[i] = 1; + } + } + const bool lane_scored = prepare_chain_drafts( + inputs, one, /*force_serial=*/true); + if (!lane_scored || evaluation.slot < 0 || + evaluation.slot >= + (int)last_activation_estimate_.size() || + !std::isfinite(last_activation_estimate_[ + (size_t)evaluation.slot].expected_yield)) { + commit_evaluation_fallback(evaluation); + } + } + } else { + for (const SpecPendingEvaluation & evaluation : + score_evaluations) { + if (evaluation.slot < 0 || + evaluation.slot >= + (int)last_activation_estimate_.size() || + !std::isfinite(last_activation_estimate_[ + (size_t)evaluation.slot].expected_yield)) { + commit_evaluation_fallback(evaluation); + } + } + } + } else { + for (const SpecPendingEvaluation & evaluation : + score_evaluations) { + commit_evaluation_fallback(evaluation); + } + } + + if (!replan_with_published_score()) { + return fail_step(gate_plan.error.empty() + ? "adaptive speculation gate failed" : gate_plan.error); + } + // This guard converts any unexpectedly unpublished score into + // the same request-local fallback rather than repeating the + // cold evaluation forever or failing a mixed cohort. + if (!gate_plan.pending_evaluations.empty()) { + const std::vector unresolved = + gate_plan.pending_evaluations; + for (const SpecPendingEvaluation & evaluation : unresolved) { + commit_evaluation_fallback(evaluation); + } + if (!replan_with_published_score()) { + return fail_step(gate_plan.error.empty() + ? "adaptive speculation gate failed" + : gate_plan.error); + } + } + if (!gate_plan.pending_evaluations.empty()) { + return fail_step( + "adaptive activation score did not resolve"); + } + } + if (cohort_changed) { + SpecCohortEpoch epoch; + epoch.id = next_spec_cohort_epoch_id_++; + epoch.request_ids = SpecCohortEpoch::ids(candidates); + epoch.plan = gate_plan; + spec_cohort_epoch_ = std::move(epoch); + log_spec_epoch(*spec_cohort_epoch_, *speculation_gate_); + published_epoch = true; + } + + // Discard bootstrap proposals for lanes routed to AR in the new + // epoch. Selected lanes keep that first proposal as useful work. + for (const SpecPlanScore & score : gate_plan.ordered) { + if (!published_epoch || score.admitted) continue; + const int slot = score.slot; + if (slot >= 0 && + slot < (int)prepared_chain_drafts_.size()) { + prepared_chain_drafts_[(size_t)slot] = {}; + } + if (slot >= 0 && slot < (int)slot_draft_kv_.size() && + slot_draft_kv_[(size_t)slot]) { + draft_kv_reset(*slot_draft_kv_[(size_t)slot]); + } + } + + // Apply the cached epoch subset. Planned prefills use a + // telemetered AR service round without changing the epoch plan. + for (const SpecPlanScore & score : gate_plan.ordered) { + if (!score.admitted) continue; + const int slot = score.slot; + bool found = false; + for (size_t i = 0; i < inputs.size(); ++i) { + if (inputs[i].slot == slot) { + admitted[i] = 1; + found = true; + } + } + if (!found) { + return fail_step( + "speculation gate admitted a missing decode lane"); + } + } + } else { + // Forced modes remain available without a cost profile. A normal + // Adaptive request fails closed once, request-locally, and then + // remains AR for its entire slot lifetime without failing UX. + for (size_t i = 0; i < inputs.size(); ++i) { + SpeculationPolicy policy = inputs[i].speculation_policy; + if (policy == SpeculationPolicy::Adaptive && + force == "all") { + policy = SpeculationPolicy::Always; + } else if (policy == SpeculationPolicy::Adaptive && + force == "none") { + policy = SpeculationPolicy::Never; + } else if (policy == SpeculationPolicy::Adaptive) { + policy = SpeculationPolicy::Never; + const int slot = inputs[i].slot; + if (slot >= 0 && + slot < (int)adaptive_fallback_ar_.size() && + !adaptive_fallback_ar_[(size_t)slot]) { + adaptive_fallback_ar_[(size_t)slot] = 1; + const uint64_t request_id = + slots_.slot(slot).request_id; + const char * reason = + adaptive_fallback_reason_.empty() + ? "cost_profile_unavailable" + : adaptive_fallback_reason_.c_str(); + log_spec_evaluation_fallback( + request_id, slot, + chain_activation_score_kind(), reason); + } + } + admitted[i] = policy == SpeculationPolicy::Always; + } + } + + const bool any_admitted = std::any_of( + admitted.begin(), admitted.end(), + [](uint8_t value) { return value != 0; }); + if (any_admitted && plan.prefills.empty()) { + StepResult speculative = + step_chain_spec(plan, admitted, decode_round_started); + const double measured_us = + std::chrono::duration( + std::chrono::steady_clock::now() - chain_started).count(); + const bool spec_completed = speculative.error.empty(); + bool proposal_failed = false; + std::vector accepted_lengths(inputs.size(), 0); + if (spec_completed) { + for (size_t i = 0; i < inputs.size(); ++i) { + if (!admitted[i]) continue; + const auto output = std::find_if( + speculative.decode.begin(), speculative.decode.end(), + [&](const DecodeOutput & item) { + return item.slot == inputs[i].slot; + }); + if (output == speculative.decode.end() || output->failed || + output->spec_steps == 0) { + proposal_failed = true; + break; + } + accepted_lengths[i] = + 1 + static_cast(output->spec_accepted_tokens); + } + } + const bool cost_sample_valid = + have_gate_plan && spec_completed && !proposal_failed; + if (cost_sample_valid) { + const ChainLaunchShape executed = chain_launch_shape( + admitted, accepted_lengths, + chain_verify_depth_for_round()); + const bool direct_commit = chain_direct_commit_enabled(); + const int priced_tree_rows = direct_commit + ? executed.tree_rows + + static_cast(inputs.size()) - + executed.spec_lanes + : executed.tree_rows; + speculation_gate_->observe_cost( + {static_cast(inputs.size()), executed.spec_lanes, + priced_tree_rows, + direct_commit ? 0 : executed.commit_rows, + round_draft_lanes_}, + measured_us); + } + if (have_gate_plan && spec_gate_debug_enabled() && + spec_completed && !proposal_failed) { + const double realized_tokens = spec_completed + ? initial_prediction_realized_tokens(gate_plan, speculative) + : std::numeric_limits::quiet_NaN(); + log_spec_gate_plan( + gate_plan, realized_tokens, + cost_sample_valid + ? measured_us + : std::numeric_limits::quiet_NaN()); + } + return speculative; + } + if (any_admitted) { + // Chain verification cannot share a target graph with prompt work. + // Use the existing packed AR+prefill graph for this service round + // so selected prompts make immediate progress. Routing remains + // fixed for the current epoch; the per-request metric distinguishes this bounded + // scheduling suspension from speculative execution. + spec_service_ar = admitted; + for (size_t i = 0; i < inputs.size(); ++i) { + if (!admitted[i]) continue; + const int slot = inputs[i].slot; + if (slot >= 0 && + slot < (int)prepared_chain_drafts_.size()) { + prepared_chain_drafts_[(size_t)slot] = {}; + } + if (slot >= 0 && slot < (int)slot_draft_kv_.size() && + slot_draft_kv_[(size_t)slot]) { + draft_kv_reset(*slot_draft_kv_[(size_t)slot]); + } + } + } + if (!any_admitted && have_gate_plan) { + if (gate_plan.admitted_count == 0 && plan.prefills.empty()) { + pending_ar_gate_plan = gate_plan; + } + } + } const TargetWeights & w = b_.w_; StepGraph & sg = b_.sg_; const int hidden = w.n_embd; const int n_head_kv = w.n_head_kv; + timing_clock::time_point t_ar_build_start, t_ar_build_end, + t_ar_exec_start, t_ar_exec_end, t_ar_read_end; decode_outputs.reserve(inputs.size()); prefill_outputs.reserve(plan.prefills.size()); @@ -242,9 +3456,10 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { output_rows_.push_back(compact_row); continue; } - if (app.new_block >= 0 && - !upload_block_table_delta( - in.slot, app.new_block_index, &app.new_block, 1)) { + const bool table_ok = slots_.residency_active() || app.new_block < 0 || + upload_block_table_delta(in.slot, app.new_block_index, + &app.new_block, 1); + if (!table_ok) { out.error = "decode block-table entry exceeds device capacity"; decode_outputs.push_back(std::move(out)); output_rows_.push_back(compact_row); @@ -278,6 +3493,9 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { } prefills.push_back(std::move(prefill)); } + if (!upload_all_active_block_tables()) { + return fail_step("active KVFlash block-table refresh failed"); + } const int live_count = (int)live_tokens_.size(); const bool with_decode = live_count > 0; @@ -324,12 +3542,13 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { : std::max(1, n_commits)) : 0; + if (timing) t_ar_build_start = timing_clock::now(); bool built = false; if (with_prefill) { built = build_target_step( sg, w, b_.cache_, b_.target_backend_, /*kv_start=*/0, /*n_tokens=*/n_total, - /*with_mask=*/false, /*capture=*/false, + /*with_mask=*/false, /*capture=*/capture_features_, /*capture_delta_intermediate=*/false, /*fa_window=*/0, /*logits_tail_rows=*/0, b_.cfg_.kq_stride_pad, @@ -347,7 +3566,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { built = build_target_step( sg, w, b_.cache_, b_.target_backend_, /*kv_start=*/0, /*n_tokens=*/decode_bucket, - /*with_mask=*/false, /*capture=*/false, + /*with_mask=*/false, /*capture=*/capture_features_, /*capture_delta_intermediate=*/false, /*fa_window=*/0, /*logits_tail_rows=*/0, b_.cfg_.kq_stride_pad, @@ -365,11 +3584,13 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { /*compact_slots=*/true); } if (!built || !sg.kv_write_rows || + (capture_features_ && !sg.target_feat_rows) || (with_prefill && (!sg.paged_query_seq_ids || !sg.paged_query_positions || !sg.logits_row_indices))) { return fail_step("packed prefill/decode graph build failed"); } + if (timing) t_ar_build_end = timing_clock::now(); embed_buf_.resize((size_t)hidden * n_total); int token_offset = 0; @@ -428,6 +3649,30 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { b_.target_backend_, sg.kv_write_rows, rows_buf_.data(), 0, sizeof(int64_t) * rows_buf_.size()); + if (capture_features_) { + const int cap = b_.cache_.target_feat_cap; + const int dead_row = cap * n_slots; + feature_rows_.assign((size_t)n_total, dead_row); + int feature_offset = 0; + for (size_t i = 0; i < prefills.size(); ++i) { + const PrefillStage & prefill = prefills[i]; + const int slot = plan.prefills[i].slot; + for (int row = 0; row < prefill.chunk; ++row) { + feature_rows_[(size_t)(feature_offset + row)] = + slot * cap + (prefill.kv_pos + row) % cap; + } + feature_offset += prefill.chunk; + } + for (int row = 0; row < live_count; ++row) { + feature_rows_[(size_t)(n_prefill + row)] = + live_slot_ids_[(size_t)row] * cap + + live_positions_[(size_t)row] % cap; + } + ggml_backend_tensor_set_async( + b_.target_backend_, sg.target_feat_rows, feature_rows_.data(), 0, + sizeof(int32_t) * feature_rows_.size()); + } + if (with_prefill) { query_slot_ids_.assign((size_t)n_total, -1); query_positions_.assign((size_t)n_total, -1); @@ -486,6 +3731,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { b_.target_backend_, b_.cache_.paged_kv_seq_lens, seq_lens_.data(), 0, sizeof(int32_t) * seq_lens_.size()); + if (timing) t_ar_exec_start = timing_clock::now(); ggml_status st = GGML_STATUS_FAILED; { const Qwen35RoctxRange roctx_compute( @@ -495,6 +3741,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { if (st != GGML_STATUS_SUCCESS) { return fail_step("packed prefill/decode compute failed"); } + if (timing) t_ar_exec_end = timing_clock::now(); const int decode_row0 = with_prefill ? n_commits : 0; const int argmax_rows = with_prefill ? gather_rows : decode_bucket; @@ -507,16 +3754,40 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { "qwen35.argmax_readback", roctx_metadata); ggml_backend_synchronize(b_.target_backend_); } + if (timing) t_ar_read_end = timing_clock::now(); + std::vector write_slots; + write_slots.reserve(live_slot_ids_.size() + prefills.size()); + write_slots.insert(write_slots.end(), + live_slot_ids_.begin(), live_slot_ids_.end()); + for (size_t i = 0; i < prefills.size(); ++i) { + write_slots.push_back(plan.prefills[i].slot); + } + if (!commit_residency_writes(write_slots)) { + return fail_step("KVFlash pending write commit failed"); + } for (size_t oi = 0; oi < inputs.size(); ++oi) { DecodeOutput & out = decode_outputs[oi]; if (out.failed) continue; slots_.commit_step(out.slot); const int row = decode_row0 + output_rows_[oi]; + out.spec_service_ar_steps = + spec_service_ar[oi] ? 1 : 0; out.token = sample_graph_row( out.slot, row, &argmax_buf_[(size_t)row], &logits_buf_); } + for (DecodeOutput & out : decode_outputs) { + if (out.failed) continue; + std::string reselect_error; + if (!maybe_reselect_residency(out.slot, reselect_error)) { + return fail_step(reselect_error.empty() + ? "KVFlash reselect failed" : reselect_error); + } + out.target_forwards = 1; + attach_residency_telemetry(out); + } + int commit_row = 0; for (size_t i = 0; i < prefills.size(); ++i) { const int slot = plan.prefills[i].slot; @@ -532,11 +3803,64 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { } prefill_outputs.push_back(std::move(out)); } + if (plan.prefills.empty() && live_count > 0 && + (timing || pending_ar_gate_plan.has_value())) { + const auto t_ar_end = timing_clock::now(); + auto span_us = [](timing_clock::time_point from, + timing_clock::time_point to) { + return std::chrono::duration( + to - from).count(); + }; + if (pending_ar_gate_plan) { + const double measured_us = + span_us(decode_round_started, t_ar_end); + speculation_gate_->observe_cost( + {live_count, 0, 0, decode_bucket, round_draft_lanes_}, + measured_us); + if (spec_gate_debug_enabled()) { + log_spec_gate_plan( + *pending_ar_gate_plan, 0.0, + measured_us); + } + } + if (timing) { + std::fprintf(stderr, + "[step-timing] {\"path\":\"ar\",\"live\":%d,\"k\":0," + "\"decode_bucket\":%d,\"n_prefill\":0,\"max_kv_len\":%d," + "\"draft_us\":%.1f,\"draft_lanes\":%d," + "\"pre_us\":%.1f,\"graph_build_us\":%.1f," + "\"graph_prepare_us\":%.1f,\"graph_exec_us\":%.1f," + "\"sample_read_us\":%.1f,\"finish_us\":%.1f," + "\"total_us\":%.1f,\"accepted_tokens\":0," + "\"emitted_tokens\":%d,\"target_forwards\":%d}\n", + live_count, decode_bucket, max_kv_len, + round_draft_us_, round_draft_lanes_, + span_us(decode_round_started, t_ar_build_start), + span_us(t_ar_build_start, t_ar_build_end), + span_us(t_ar_build_end, t_ar_exec_start), + span_us(t_ar_exec_start, t_ar_exec_end), + span_us(t_ar_exec_end, t_ar_read_end), + span_us(t_ar_read_end, t_ar_end), + span_us(decode_round_started, t_ar_end), + live_count, live_count); + } + } return result; } void Qwen35SeqEngine::retire(int slot) { if (!slots_.is_active(slot)) return; + const uint64_t request_id = slots_.slot(slot).request_id; + if (speculation_gate_) speculation_gate_->forget(request_id); + if (slot >= 0 && slot < (int)last_activation_estimate_.size()) { + last_activation_estimate_[(size_t)slot] = {}; + if (slot < (int)prepared_chain_drafts_.size()) { + prepared_chain_drafts_[(size_t)slot].valid = false; + } + } + if (slot >= 0 && slot < (int)adaptive_fallback_ar_.size()) { + adaptive_fallback_ar_[(size_t)slot] = 0; + } slots_.retire(slot); } diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.h b/server/src/qwen35/concurrency/qwen35_seq_engine.h index d9391784c..793d38e75 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.h +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.h @@ -22,10 +22,20 @@ #pragma once #include "common/concurrency/seq_engine.h" +#include "common/speculation/speculator.h" +#include "common/speculation/speculation_gate.h" +#include "common/dflash_draft_kv.h" +#include "common/dflash_feature_ring.h" +#include "common/ddtree.h" #include "qwen35_slot_manager.h" #include +#include #include +#include +#include +#include +#include #include namespace dflash::common { @@ -34,25 +44,28 @@ class Qwen35Backend; class Qwen35SeqEngine final : public SeqEngine { public: + enum class SpecMode { + none, + ddtree, + chain, + }; + // `pool` and `backend` must outlive the engine. `scratch_row` is the // first row of the block appended past the pool's index space, used as // the K/V write destination of graph-bucket padding rows. // `max_prefills` bounds scheduler-selected prompt slices per traversal. Qwen35SeqEngine(Qwen35Backend & backend, PagedKvPool & pool, int max_ctx, int64_t scratch_row, + int tree_width = 0, int tree_scratch_base = 0, + int tree_scratch_stride = 0, + SpecMode spec_mode = SpecMode::none, int max_prefills = 8, int mixed_prefill_tokens = 2048, int long_mixed_prefill_tokens = 4096, int long_prefill_threshold = 768, int idle_prefill_tokens = 4096, - int prefill_quantum = 512) - : max_prefills_(std::max(1, max_prefills)), - mixed_prefill_tokens_(std::max(1, mixed_prefill_tokens)), - long_mixed_prefill_tokens_(std::max(1, long_mixed_prefill_tokens)), - long_prefill_threshold_(std::max(1, long_prefill_threshold)), - idle_prefill_tokens_(std::max(1, idle_prefill_tokens)), - prefill_quantum_(std::max(1, prefill_quantum)), b_(backend), - slots_(pool, max_ctx), scratch_row_(scratch_row) {} + int prefill_quantum = 512); + ~Qwen35SeqEngine() override; int slot_count() const override { return slots_.slot_count(); } int max_context() const override { return slots_.max_context(); } @@ -62,6 +75,13 @@ class Qwen35SeqEngine final : public SeqEngine { const SamplerCfg & sampler) override; StepResult step(const StepPlan & plan) override; + // Fabricate a steady-state paged context and profile the three launch + // series used by the activation gate. Called once from backend init. + bool profile_spec_costs(int context_tokens); + // True only when the registered speculator adapter can produce a + // first-request activation score. A configured chain + // may still accept Adaptive requests and serve AR fallback when this is false. + bool activation_scoring_available() const; StepPlanLimits step_plan_limits(int decode_rows) const override { const bool mixed = decode_rows > 0; const int per_sequence = mixed ? 512 : 2048; @@ -104,6 +124,10 @@ class Qwen35SeqEngine final : public SeqEngine { bool upload_block_table_delta(int slot, int first_block, const int32_t * blocks, size_t count); + bool upload_all_active_block_tables(); + bool commit_residency_writes(const std::vector & slots); + bool maybe_reselect_residency(int slot, std::string & error); + void attach_residency_telemetry(DecodeOutput & out); void fail_prefill(int slot, std::vector & outputs, const char * log_message, const char * client_message); @@ -112,11 +136,79 @@ class Qwen35SeqEngine final : public SeqEngine { int32_t sample_graph_row(int slot, int logits_row, const int32_t * cached_argmax = nullptr, std::vector * logits_scratch = nullptr); + DraftFeatureMirror * slot_feature_mirror(int slot); + DraftKvState * ensure_slot_draft_kv(int slot); + bool ddtree_eligible(const StepPlan & plan) const; + bool chain_proposal_input_capable(const StepInput & input) const; + bool chain_activation_input_scoreable(const StepInput & input) const; + bool chain_spec_request_capable(const StepInput & input) const; + bool chain_spec_input_eligible(const StepInput & input) const; + bool spec_gate_debug_enabled() const; + struct PreparedChainDraft { + bool valid = false; + int generated = -1; + int32_t root = -1; + std::vector tokens; + ActivationEstimate estimate; + std::vector debug_depth_fields; + }; + bool prepare_chain_drafts( + const std::vector & inputs, + const std::vector & selected, + bool force_serial = false, + bool fail_fast_batch = false); + bool batched_drafting_enabled() const; + bool activation_scoring_enabled() const; + std::string chain_activation_score_kind() const; + // Verification depth may vary between rounds while the cohort route stays + // fixed; every returned depth must stay in [2, tree_width_]. + int chain_verify_depth_for_round() const { + return chain_verify_depth_; + } + // DFLASH_STEP_TIMING=1 emits one [step-timing] JSON line per decode + // round attributing wall time to draft, verify, readback, CPU commit, + // replay, and packed-AR phases. Diagnostic only; off by default. + static bool step_timing_enabled(); + // DDTree preserves its legacy best-effort AR fallback. Chain + // proposal failures are instead returned as lane-local DecodeOutput + // failures so an epoch-selected speculation lane never silently executes as AR. + std::optional step_ddtree(const StepPlan & plan); + StepResult step_chain_spec( + const StepPlan & plan, const std::vector & admitted, + std::chrono::steady_clock::time_point round_started); Qwen35Backend & b_; Qwen35SlotManager slots_; int64_t scratch_row_ = 0; - + int tree_width_ = 0; + // Root-inclusive verification depth. The drafter still produces + // tree_width_ tokens; this common cohort depth may vary between 2 and + // tree_width_ without changing the current cohort plan. + int chain_verify_depth_ = 0; + int tree_scratch_base_ = 0; + int tree_scratch_stride_ = 0; + bool capture_features_ = false; + SpecMode spec_mode_ = SpecMode::none; + ggml_context * feature_view_ctx_ = nullptr; + std::vector slot_feature_mirrors_; + std::vector> slot_draft_kv_; + + DraftKvBatchGraph batch_draft_graph_; + std::vector> dummy_draft_kv_; + std::vector prepared_chain_drafts_; + std::unique_ptr speculation_gate_; + std::optional spec_cohort_epoch_; + uint64_t next_spec_cohort_epoch_id_ = 1; + std::unique_ptr speculator_; + // Startup profile/adapter failure is a request-local AR outcome, never an + // admission or step error for a configured chain. + std::string adaptive_fallback_reason_ = "cost_profile_unavailable"; + std::vector adaptive_fallback_ar_; + std::vector last_activation_estimate_; + // Per-round draft cost accumulator for [step-timing]; reset at the top + // of each chain-speculation round, accumulated by prepare_chain_drafts. + double round_draft_us_ = 0.0; + int round_draft_lanes_ = 0; // Hoisted per-step buffers (reused across step() calls). std::vector output_rows_; std::vector live_tokens_; @@ -131,6 +223,7 @@ class Qwen35SeqEngine final : public SeqEngine { std::vector query_slot_ids_; std::vector query_positions_; std::vector logits_rows_; + std::vector feature_rows_; std::vector embed_buf_; std::vector pos_buf_; std::vector rows_buf_; diff --git a/server/src/qwen35/concurrency/qwen35_slot_manager.cpp b/server/src/qwen35/concurrency/qwen35_slot_manager.cpp index b8dcc36ec..e559a224e 100644 --- a/server/src/qwen35/concurrency/qwen35_slot_manager.cpp +++ b/server/src/qwen35/concurrency/qwen35_slot_manager.cpp @@ -2,11 +2,16 @@ #include #include +#include namespace dflash::common { -Qwen35SlotManager::Qwen35SlotManager(PagedKvPool & pool, int max_ctx) - : pool_(pool), max_ctx_(max_ctx) { +Qwen35SlotManager::Qwen35SlotManager( + PagedKvPool & pool, int max_ctx, int speculative_headroom, + PagedKvResidencyManager * residency) + : pool_(pool), max_ctx_(max_ctx), + headroom_tokens_(std::max(pool.block_size(), speculative_headroom)), + residency_(residency) { slots_.assign(pool.max_sequences(), Qwen35Slot{}); } @@ -21,7 +26,7 @@ int Qwen35SlotManager::decoding_count() const { uint32_t Qwen35SlotManager::decode_headroom_capacity(int logical_tokens) const { const uint64_t extended = static_cast(std::max(0, logical_tokens)) + - pool_.block_size(); + static_cast(headroom_tokens_); return static_cast(std::min( static_cast(max_ctx_), extended)); } @@ -84,6 +89,21 @@ bool Qwen35SlotManager::is_prefilling(int slot) const { return is_active(slot) && slots_[(size_t)slot].prefilling(); } +void Qwen35SlotManager::accumulate_residency_delta( + Qwen35Slot & slot, const PagedKvResidencyStats & before) { + if (!residency_) return; + const PagedKvResidencyStats after = residency_->stats(); + if (after.page_ins >= before.page_ins) { + slot.kvflash_page_ins_pending += after.page_ins - before.page_ins; + } + if (after.page_outs >= before.page_outs) { + slot.kvflash_page_outs_pending += after.page_outs - before.page_outs; + } + if (after.reselects >= before.reselects) { + slot.kvflash_reselects_pending += after.reselects - before.reselects; + } +} + bool Qwen35SlotManager::has_prefill_prompt_at_least(int tokens) const { if (tokens <= 0) return true; return std::any_of(slots_.begin(), slots_.end(), @@ -112,7 +132,7 @@ SeqEngine::AdmitResult Qwen35SlotManager::admit( // fail anyway. Hard-fail it up front instead of reporting busy. const uint64_t pool_capacity = (uint64_t)pool_.physical_block_count() * pool_.block_size(); - if ((uint64_t)prompt_len > pool_capacity) { + if (!residency_ && (uint64_t)prompt_len > pool_capacity) { r.error = "prompt needs " + std::to_string(prompt_len) + " KV tokens but the pool holds " + std::to_string(pool_capacity) + @@ -120,9 +140,15 @@ SeqEngine::AdmitResult Qwen35SlotManager::admit( return r; } + // Retirement keeps failed copy-stream ownership quarantined. Retry those + // barriers before deciding whether a sequence slot is actually available. + for (int i = 0; i < (int)slots_.size(); ++i) { + if (slots_[(size_t)i].retiring()) retire(i); + } + int slot = -1; for (int i = 0; i < (int)slots_.size(); i++) { - if (!slots_[(size_t)i].active()) { slot = i; break; } + if (slots_[(size_t)i].phase == Qwen35SlotPhase::free) { slot = i; break; } } if (slot < 0) { r.status = AdmitStatus::busy; @@ -132,7 +158,8 @@ SeqEngine::AdmitResult Qwen35SlotManager::admit( // A newly freed block belongs to any older decoder missing its rolling // next-page reserve before it can belong to this admission. - const PagedKvStatus headroom_status = protect_decode_headroom(); + const PagedKvStatus headroom_status = residency_ + ? PagedKvStatus::Ok : protect_decode_headroom(); if (headroom_status != PagedKvStatus::Ok) { r.status = headroom_status == PagedKvStatus::BlocksExhausted ? AdmitStatus::busy : AdmitStatus::failed; @@ -143,7 +170,7 @@ SeqEngine::AdmitResult Qwen35SlotManager::admit( } PagedKvSequenceHandle handle; - uint32_t reservation_capacity = + uint32_t reservation_capacity = residency_ ? 0 : decode_headroom_capacity(prompt_len); if (!capacity_fits_pool(reservation_capacity)) { // The prompt itself fits, but this physical pool can never hold its @@ -151,8 +178,9 @@ SeqEngine::AdmitResult Qwen35SlotManager::admit( // decode exhaustion later if the sequence reaches that boundary. reservation_capacity = static_cast(prompt_len); } - const PagedKvStatus status = pool_.acquire_reserved( - request_id, reservation_capacity, handle); + const PagedKvStatus status = residency_ + ? pool_.acquire(request_id, handle) + : pool_.acquire_reserved(request_id, reservation_capacity, handle); if (status != PagedKvStatus::Ok) { r.status = status == PagedKvStatus::SequenceSlotsExhausted || status == PagedKvStatus::BlocksExhausted @@ -163,8 +191,22 @@ SeqEngine::AdmitResult Qwen35SlotManager::admit( return r; } + if (residency_) { + const PagedKvResidencyStatus resident_status = + residency_->register_sequence(handle); + if (resident_status != PagedKvResidencyStatus::Ok) { + (void)pool_.release(handle); + r.error = std::string("KV residency registration failed: ") + + paged_kv_residency_status_string(resident_status); + return r; + } + } + Qwen35Slot & s = slots_[(size_t)slot]; s.phase = Qwen35SlotPhase::prefill; + s.request_id = request_id; + s.ddtree_suspended = false; + s.ddtree_sampled_steps = 0; s.handle = handle; s.cur_pos = 0; s.prompt_len = prompt_len; @@ -183,6 +225,33 @@ SeqEngine::AdmitResult Qwen35SlotManager::admit( return r; } +bool Qwen35SlotManager::ddtree_speculation_allowed(int slot) const { + return is_active(slot) && !slots_[(size_t)slot].ddtree_suspended; +} + +bool Qwen35SlotManager::ddtree_cohort_should_suspend( + uint64_t total_emitted, int active) { + const char * adaptive = std::getenv("DFLASH_DDTREE_ADAPTIVE"); + if (adaptive && std::atoi(adaptive) == 0) { + return false; + } + if (active <= 0) return false; + return total_emitted < + (uint64_t)active * (uint64_t)kDdtreeMinEmittedTokens; +} + +bool Qwen35SlotManager::record_ddtree_sample( + int slot, bool suspend_cohort) { + if (!is_active(slot)) return false; + Qwen35Slot & s = slots_[(size_t)slot]; + // A suspended request must never pay for another probe. + if (s.ddtree_suspended) return false; + ++s.ddtree_sampled_steps; + if (!suspend_cohort) return false; + s.ddtree_suspended = true; + return true; +} + Qwen35SlotManager::PrefillChunk Qwen35SlotManager::append_prefill( int slot, int n_tokens) { PrefillChunk out; @@ -194,7 +263,22 @@ Qwen35SlotManager::PrefillChunk Qwen35SlotManager::append_prefill( return out; } - PagedKvAppendResult app = pool_.append(s.handle, (uint32_t)n_tokens); + PagedKvAppendResult app; + if (residency_) { + const PagedKvResidencyStats before = residency_->stats(); + PagedKvResidentAppendResult resident = residency_->append( + s.handle, (uint32_t)n_tokens); + accumulate_residency_delta(s, before); + if (!resident) { + std::fprintf(stderr, + "[parallel-kvflash] prefill append failed for slot %d: %s\n", + slot, paged_kv_residency_status_string(resident.status)); + return out; + } + app = std::move(resident.pool_result); + } else { + app = pool_.append(s.handle, (uint32_t)n_tokens); + } if (!app) { // Admission reserved the whole prompt. Treat exhaustion here as a // broken invariant, not a retryable condition: retrying a batch of @@ -218,6 +302,17 @@ Qwen35SlotManager::PrefillChunk Qwen35SlotManager::append_prefill( out.new_blocks.push_back((int32_t)write.physical_block); } } + if (residency_) { + PagedKvSequenceSnapshot snapshot; + if (pool_.sequence(s.handle, snapshot) != PagedKvStatus::Ok) { + return out; + } + out.full_block_table.reserve(snapshot.block_table.size()); + for (uint32_t block : snapshot.block_table) { + out.full_block_table.push_back( + block == PAGED_KV_COLD_BLOCK ? -1 : (int32_t)block); + } + } s.cur_pos += n_tokens; out.ok = true; return out; @@ -230,43 +325,190 @@ void Qwen35SlotManager::commit_prefill(int slot) { s.phase = Qwen35SlotPhase::decode; } -Qwen35SlotManager::StepAppend Qwen35SlotManager::append_token(int slot, - int32_t fed_token) { +Qwen35SlotManager::StepAppend Qwen35SlotManager::append_tokens( + int slot, const int32_t * fed_tokens, int n_tokens) { StepAppend out; - if (!is_active(slot) || !slots_[(size_t)slot].decoding()) return out; + if (!is_active(slot) || !slots_[(size_t)slot].decoding() || + !fed_tokens || n_tokens < 1) { + return out; + } Qwen35Slot & s = slots_[(size_t)slot]; - if (s.cur_pos >= max_ctx_) { - // No context left; the scheduler should have stopped this slot. + if (!s.staged_tokens.empty() || s.cur_pos > max_ctx_ || + n_tokens > max_ctx_ - s.cur_pos) { return out; } - PagedKvAppendResult app = pool_.append( - s.handle, 1, /*only_first_last_slots=*/true); - if (!app || app.token_count != 1 || - app.last.logical_position != (uint32_t)s.cur_pos) { + + PagedKvAppendResult app; + if (residency_) { + const PagedKvResidencyStats before = residency_->stats(); + PagedKvResidentAppendResult resident = residency_->append( + s.handle, static_cast(n_tokens)); + accumulate_residency_delta(s, before); + if (!resident) { + out.busy = resident.status == + PagedKvResidencyStatus::PoolExhausted || + resident.status == PagedKvResidencyStatus::NoEvictableBlock; + return out; + } + app = std::move(resident.pool_result); + } else { + app = pool_.append(s.handle, static_cast(n_tokens)); + } + if (!app || app.token_count != static_cast(n_tokens)) { out.busy = app.status == PagedKvStatus::BlocksExhausted; return out; } - s.sample_history.push_back(fed_token); + if (app.write_slots.size() != static_cast(n_tokens) || + app.write_slots.front().logical_position != + static_cast(s.cur_pos) || + app.write_slots.back().logical_position != + static_cast(s.cur_pos + n_tokens - 1)) { + // Pool success guarantees this shape. Retain staged ownership so a + // fatal caller retires the sequence rather than double-appending. + s.staged_tokens.assign(fed_tokens, fed_tokens + n_tokens); + return out; + } + + out.physical_rows.reserve(app.write_slots.size()); + for (const PagedKvWriteSlot & write : app.write_slots) { + out.physical_rows.push_back( + static_cast(write.physical_token_index)); + if (write.block_offset == 0) { + if (out.first_new_block < 0) { + out.first_new_block = static_cast( + write.logical_position / pool_.block_size()); + } + out.new_blocks.push_back( + static_cast(write.physical_block)); + } + } + s.staged_tokens.assign(fed_tokens, fed_tokens + n_tokens); + if (residency_) { + PagedKvSequenceSnapshot snapshot; + if (pool_.sequence(s.handle, snapshot) != PagedKvStatus::Ok) { + return out; + } + out.full_block_table.reserve(snapshot.block_table.size()); + for (uint32_t block : snapshot.block_table) { + out.full_block_table.push_back( + block == PAGED_KV_COLD_BLOCK ? -1 : (int32_t)block); + } + } out.ok = true; - out.physical_row = (int64_t)app.last.physical_token_index; + out.count = n_tokens; out.position = s.cur_pos; - if ((uint32_t)s.cur_pos % pool_.block_size() == 0) { - out.new_block = (int32_t)app.last.physical_block; - out.new_block_index = s.cur_pos / (int)pool_.block_size(); + if (n_tokens == 1) { + out.physical_row = out.physical_rows.front(); + if (!out.new_blocks.empty()) { + out.new_block = out.new_blocks.front(); + out.new_block_index = out.first_new_block; + } } return out; } +Qwen35SlotManager::StepAppend Qwen35SlotManager::append_token( + int slot, int32_t fed_token) { + return append_tokens(slot, &fed_token, 1); +} + void Qwen35SlotManager::commit_step(int slot) { if (!is_active(slot)) return; - slots_[(size_t)slot].cur_pos += 1; + Qwen35Slot & s = slots_[(size_t)slot]; + if (s.staged_tokens.empty()) return; + s.sample_history.insert( + s.sample_history.end(), s.staged_tokens.begin(), + s.staged_tokens.end()); + s.cur_pos += static_cast(s.staged_tokens.size()); + s.staged_tokens.clear(); +} + +bool Qwen35SlotManager::block_table_snapshot( + int slot, std::vector & out) const { + out.clear(); + if (!is_active(slot)) return false; + PagedKvSequenceSnapshot snapshot; + if (pool_.sequence(slots_[(size_t)slot].handle, snapshot) != + PagedKvStatus::Ok) { + return false; + } + out.reserve(snapshot.block_table.size()); + for (uint32_t block : snapshot.block_table) { + out.push_back(block == PAGED_KV_COLD_BLOCK ? -1 : (int32_t)block); + } + return true; +} + +bool Qwen35SlotManager::commit_residency_writes(int slot) { + if (!residency_) return true; + if (!is_active(slot)) return false; + return residency_->commit_pending_writes(slots_[(size_t)slot].handle) == + PagedKvResidencyStatus::Ok; +} + +bool Qwen35SlotManager::reselect_residency( + int slot, const std::vector * scores, std::string * error) { + if (!residency_) return true; + if (!is_active(slot)) { + if (error) *error = "inactive KVFlash slot"; + return false; + } + Qwen35Slot & s = slots_[(size_t)slot]; + const PagedKvResidencyStats before = residency_->stats(); + const std::vector no_scores; + PagedKvResidencyStatus status = residency_->set_scores( + s.handle, scores ? *scores : no_scores); + if (status == PagedKvResidencyStatus::Ok) { + status = residency_->reselect(s.handle); + } + accumulate_residency_delta(s, before); + if (status != PagedKvResidencyStatus::Ok) { + if (error) { + *error = std::string("KVFlash reselect failed: ") + + paged_kv_residency_status_string(status); + } + return false; + } + return true; +} + +void Qwen35SlotManager::take_residency_telemetry( + int slot, SeqEngine::DecodeOutput & out) { + if (!residency_ || !is_active(slot)) return; + Qwen35Slot & s = slots_[(size_t)slot]; + out.kvflash_page_ins = s.kvflash_page_ins_pending; + out.kvflash_page_outs = s.kvflash_page_outs_pending; + out.kvflash_reselects = s.kvflash_reselects_pending; + uint32_t resident = 0; + if (pool_.resident_block_count(s.handle, resident) == PagedKvStatus::Ok) { + out.kvflash_resident_blocks = resident; + } + s.kvflash_page_ins_pending = 0; + s.kvflash_page_outs_pending = 0; + s.kvflash_reselects_pending = 0; } void Qwen35SlotManager::retire(int slot) { if (slot < 0 || slot >= (int)slots_.size()) return; Qwen35Slot & s = slots_[(size_t)slot]; - if (!s.active()) return; + if (s.phase == Qwen35SlotPhase::free) return; + if (residency_) { + const PagedKvResidencyStatus resident_status = + residency_->forget_sequence(s.handle); + if (resident_status != PagedKvResidencyStatus::Ok && + resident_status != PagedKvResidencyStatus::StaleHandle && + resident_status != PagedKvResidencyStatus::SequenceNotRegistered) { + std::fprintf(stderr, + "[parallel-kvflash] slot %d residency release failed: %s\n", + slot, paged_kv_residency_status_string(resident_status)); + // A failed copy-stream barrier leaves physical pages in flight. + // Keep the slot and its pool handle intact so a later retirement + // can retry forget_sequence without recycling those pages. + s.phase = Qwen35SlotPhase::retiring; + return; + } + } const PagedKvStatus status = pool_.release(s.handle); if (status != PagedKvStatus::Ok && status != PagedKvStatus::StaleHandle) { std::fprintf(stderr, "[parallel] slot %d release failed: %s\n", diff --git a/server/src/qwen35/concurrency/qwen35_slot_manager.h b/server/src/qwen35/concurrency/qwen35_slot_manager.h index 1da009f69..9a84fb3be 100644 --- a/server/src/qwen35/concurrency/qwen35_slot_manager.h +++ b/server/src/qwen35/concurrency/qwen35_slot_manager.h @@ -17,6 +17,7 @@ #pragma once #include "common/concurrency/paged_kv_pool.h" +#include "common/concurrency/paged_kv_residency.h" #include "common/sampler.h" #include "common/concurrency/seq_engine.h" @@ -31,10 +32,12 @@ enum class Qwen35SlotPhase { free, prefill, decode, + retiring, }; struct Qwen35Slot { Qwen35SlotPhase phase = Qwen35SlotPhase::free; + uint64_t request_id = 0; PagedKvSequenceHandle handle; // Prompt tokens are the immutable prefix of sample_history. Decode tokens // append to the same allocation, avoiding a second full prompt copy. @@ -45,6 +48,17 @@ struct Qwen35Slot { // Penalty history is recorded as fed rather than sampled: the scheduler // may override a sample before the model consumes it. std::vector sample_history; + // Decode rows allocated from the pool but not yet made durable by a + // successful target forward. A step stages one contiguous token range + // per slot, then commit_step() publishes all of it atomically. + std::vector staged_tokens; + + // Residency operation deltas are attributed to the requesting slot and + // held across prefill until a DecodeOutput can carry them upstream. + uint64_t kvflash_page_ins_pending = 0; + uint64_t kvflash_page_outs_pending = 0; + uint64_t kvflash_reselects_pending = 0; + int kvflash_last_reselect_generated = 0; int generated_tokens() const { return sample_history.size() > (size_t)prompt_len @@ -52,16 +66,37 @@ struct Qwen35Slot { : 0; } - bool active() const { return phase != Qwen35SlotPhase::free; } + // Real packed-tree samples are counted per request. A low aggregate-yield + // cohort sample suspends every participating request; ordinary AR keeps + // every target cache and feature-ring row current. + bool ddtree_suspended = false; + uint64_t ddtree_sampled_steps = 0; + + bool active() const { + return phase == Qwen35SlotPhase::prefill || + phase == Qwen35SlotPhase::decode; + } bool prefilling() const { return phase == Qwen35SlotPhase::prefill; } bool decoding() const { return phase == Qwen35SlotPhase::decode; } + bool retiring() const { return phase == Qwen35SlotPhase::retiring; } }; class Qwen35SlotManager { public: + // Packed DDTree pays for verify + accepted-path replay. Requiring six + // emitted tokens makes continuation earn at least three tokens per target + // forward before accounting for its additional draft/tree work. + static constexpr int kDdtreeMinEmittedTokens = 6; + // `max_ctx` is the per-sequence logical bound; slot count comes from the // pool's max_sequences. The pool must outlive the manager. - Qwen35SlotManager(PagedKvPool & pool, int max_ctx); + Qwen35SlotManager(PagedKvPool & pool, int max_ctx, + int speculative_headroom = 1, + PagedKvResidencyManager * residency = nullptr); + Qwen35SlotManager(const Qwen35SlotManager &) = delete; + Qwen35SlotManager & operator=(const Qwen35SlotManager &) = delete; + Qwen35SlotManager(Qwen35SlotManager &&) = delete; + Qwen35SlotManager & operator=(Qwen35SlotManager &&) = delete; // Claim a free slot and atomically reserve all K/V blocks needed by the // known prompt plus its next logical decode page when that page can exist @@ -80,6 +115,7 @@ class Qwen35SlotManager { // Delta to patch into the slot's device block-table column. std::vector new_blocks; int first_new_block = -1; + std::vector full_block_table; }; // Append `n_tokens` more prompt rows for a prefilling slot. Physical block @@ -93,17 +129,44 @@ class Qwen35SlotManager { struct StepAppend { bool ok = false; bool busy = false; // no physical block available right now + std::vector physical_rows; + std::vector new_blocks; + int first_new_block = -1; + int count = 0; + // Compatibility fields for the common one-token append. int64_t physical_row = -1; - int position = -1; // logical position the fed token is written at + int position = -1; // logical position of the first staged token int32_t new_block = -1; int new_block_index = -1; + std::vector full_block_table; }; - // Allocate the next decode token's cache row, report any new block-table - // entry, and log it to sample_history. cur_pos waits for commit_step(). + // Atomically allocate and stage a contiguous accepted path. The pool + // append is all-or-nothing; sample_history and cur_pos remain unchanged + // until commit_step(). A slot may have only one staged range at a time. + StepAppend append_tokens(int slot, const int32_t * fed_tokens, + int n_tokens); + + // Complete residency snapshots encode cold logical pages as -1. The + // engine pads each device column before upload. + bool residency_active() const { return residency_ != nullptr; } + bool block_table_snapshot(int slot, std::vector & out) const; + bool commit_residency_writes(int slot); + bool reselect_residency(int slot, const std::vector * scores, + std::string * error = nullptr); + void take_residency_telemetry(int slot, SeqEngine::DecodeOutput & out); + + bool ddtree_speculation_allowed(int slot) const; + // Compare aggregate emitted yield (accepted children plus one replay + // bonus per request) against the cohort continuation floor. + static bool ddtree_cohort_should_suspend(uint64_t total_emitted, int active); + // Returns true exactly once, when this sample newly suspends the request. + bool record_ddtree_sample(int slot, bool suspend_cohort); + + // One-token compatibility wrapper used by ordinary autoregressive decode. StepAppend append_token(int slot, int32_t fed_token); - // The batched step's compute succeeded: cur_pos++. + // Publish every staged token after a successful target forward. void commit_step(int slot); // Release the slot's blocks and clear its state. Safe on inactive slots @@ -111,6 +174,8 @@ class Qwen35SlotManager { void retire(int slot); int slot_count() const { return (int)slots_.size(); } + void accumulate_residency_delta(Qwen35Slot & slot, + const PagedKvResidencyStats & before); int max_context() const { return max_ctx_; } int decoding_count() const; bool is_active(int slot) const; @@ -131,6 +196,8 @@ class Qwen35SlotManager { PagedKvPool & pool_; int max_ctx_ = 0; + int headroom_tokens_; + PagedKvResidencyManager * residency_ = nullptr; std::vector slots_; }; diff --git a/server/src/qwen35/delta_transition_journal.cpp b/server/src/qwen35/delta_transition_journal.cpp new file mode 100644 index 000000000..8a1ff1638 --- /dev/null +++ b/server/src/qwen35/delta_transition_journal.cpp @@ -0,0 +1,142 @@ +#include "qwen35/delta_transition_journal.h" + +#include +#include +#include + +namespace dflash::qwen35 { +namespace { + +bool matrix_size(size_t rows, size_t cols, size_t & elements) { + if (rows == 0 || cols == 0 || + rows > std::numeric_limits::max() / cols) { + return false; + } + elements = rows * cols; + return true; +} + +bool transition_shape_valid( + const DeltaTransition & transition, + size_t rows, + size_t cols) { + const size_t expected_gate = + transition.gate_mode == DeltaTransitionGateMode::Scalar ? 1 : rows; + return transition.gate.size() == expected_gate && + transition.key.size() == rows && + transition.delta.size() == cols; +} + +float gate_at(const DeltaTransition & transition, size_t row) { + return transition.gate_mode == DeltaTransitionGateMode::Scalar + ? transition.gate[0] + : transition.gate[row]; +} + +} // namespace + +bool capture_delta_transition( + const std::vector & state, + size_t rows, + size_t cols, + const std::vector & key, + const std::vector & value, + const std::vector & gate, + float beta, + DeltaTransitionGateMode gate_mode, + DeltaTransition & output) { + size_t state_elements = 0; + const size_t expected_gate = + gate_mode == DeltaTransitionGateMode::Scalar ? 1 : rows; + if (!matrix_size(rows, cols, state_elements) || + state.size() != state_elements || key.size() != rows || + value.size() != cols || gate.size() != expected_gate) { + return false; + } + + DeltaTransition next; + next.gate_mode = gate_mode; + next.gate = gate; + next.key = key; + next.delta.resize(cols); + + for (size_t col = 0; col < cols; ++col) { + float projection = 0.0f; + for (size_t row = 0; row < rows; ++row) { + const float row_gate = gate_mode == DeltaTransitionGateMode::Scalar + ? 1.0f + : gate[row]; + projection += row_gate * state[col * rows + row] * key[row]; + } + const float scalar_gate = gate_mode == DeltaTransitionGateMode::Scalar + ? gate[0] + : 1.0f; + next.delta[col] = + (value[col] - scalar_gate * projection) * beta; + } + + output = std::move(next); + return true; +} + +bool apply_delta_transition( + const DeltaTransition & transition, + size_t rows, + size_t cols, + std::vector & state) { + size_t state_elements = 0; + if (!matrix_size(rows, cols, state_elements) || + state.size() != state_elements || + !transition_shape_valid(transition, rows, cols)) { + return false; + } + + for (size_t col = 0; col < cols; ++col) { + for (size_t row = 0; row < rows; ++row) { + const size_t index = col * rows + row; + state[index] = std::fma( + transition.key[row], transition.delta[col], + gate_at(transition, row) * state[index]); + } + } + return true; +} + +bool commit_delta_transition_prefix( + const DeltaTransitionJournal & journal, + size_t accepted, + std::vector & state) { + size_t state_elements = 0; + if (!matrix_size(journal.rows, journal.cols, state_elements) || + state.size() != state_elements || + accepted > journal.transitions.size()) { + return false; + } + for (size_t i = 0; i < accepted; ++i) { + if (!transition_shape_valid( + journal.transitions[i], journal.rows, journal.cols)) { + return false; + } + } + for (size_t i = 0; i < accepted; ++i) { + // Already validated, so this cannot partially fail. + apply_delta_transition( + journal.transitions[i], journal.rows, journal.cols, state); + } + return true; +} + +size_t delta_transition_float_count( + size_t rows, + size_t cols, + DeltaTransitionGateMode gate_mode) { + const size_t gate_values = + gate_mode == DeltaTransitionGateMode::Scalar ? 1 : rows; + if (rows > std::numeric_limits::max() - cols || + rows + cols > std::numeric_limits::max() - gate_values) { + return 0; + } + return rows + cols + gate_values; +} + +} // namespace dflash::qwen35 diff --git a/server/src/qwen35/delta_transition_journal.h b/server/src/qwen35/delta_transition_journal.h new file mode 100644 index 000000000..5dbe887ba --- /dev/null +++ b/server/src/qwen35/delta_transition_journal.h @@ -0,0 +1,74 @@ +#pragma once + +#include +#include + +namespace dflash::qwen35 { + +// Host-side contract for the compact journal emitted by a future GDN verify +// kernel. The persistent state uses the kernel's transposed layout: +// state[col * rows + row]. A transition contains exactly the values needed to +// repeat the state update, but none of the model projections: +// +// state' = gate * state + key (outer-product) delta +// +// In scalar mode gate has one value. In row-wise mode it has one value per +// state row. key and delta are the normalized/resolved values produced by +// verification; in particular, delta is captured after the state-dependent +// k^T S reduction. Therefore this journal is valid only for a chain replayed +// from the same base recurrent state. +enum class DeltaTransitionGateMode { + Scalar, + RowWise, +}; + +struct DeltaTransition { + DeltaTransitionGateMode gate_mode = DeltaTransitionGateMode::Scalar; + std::vector gate; + std::vector key; + std::vector delta; +}; + +struct DeltaTransitionJournal { + size_t rows = 0; + size_t cols = 0; + std::vector transitions; +}; + +// Resolve the state-dependent delta exactly once, as the verification kernel +// would. gate contains already-exponentiated multipliers (not raw/log gates). +// The output is unchanged when validation fails. +bool capture_delta_transition( + const std::vector & state, + size_t rows, + size_t cols, + const std::vector & key, + const std::vector & value, + const std::vector & gate, + float beta, + DeltaTransitionGateMode gate_mode, + DeltaTransition & output); + +// Apply one captured transition without evaluating projections or recomputing +// delta. The state is unchanged when validation fails. +bool apply_delta_transition( + const DeltaTransition & transition, + size_t rows, + size_t cols, + std::vector & state); + +// Commit transitions [0, accepted) to a persistent state. accepted == 0 is a +// no-op. Oversized or malformed prefixes fail before modifying state. +bool commit_delta_transition_prefix( + const DeltaTransitionJournal & journal, + size_t accepted, + std::vector & state); + +// Compact recurrent journal footprint per head/token, excluding allocator +// alignment. Runtime storage may further share keys across grouped V heads. +size_t delta_transition_float_count( + size_t rows, + size_t cols, + DeltaTransitionGateMode gate_mode); + +} // namespace dflash::qwen35 diff --git a/server/src/qwen35/gguf_target_loader.cpp b/server/src/qwen35/gguf_target_loader.cpp index 1f917ee69..4c41acb75 100644 --- a/server/src/qwen35/gguf_target_loader.cpp +++ b/server/src/qwen35/gguf_target_loader.cpp @@ -680,16 +680,71 @@ bool load_target_gguf_partial(const std::string & path, if (!t || !should_load_target_tensor(tname, plan.layer_begin, plan.layer_end, plan.load_output, plan.skip_expert_tensors)) { continue; } - alloc_total = align_up_size(alloc_total, alignment); TargetTensorAlloc a; a.tensor = t; a.file_offset = gguf_get_data_offset(gctx) + gguf_get_tensor_offset(gctx, tid); a.file_size = gguf_get_tensor_size(gctx, tid); - a.buffer_offset = alloc_total; - alloc_total += ggml_backend_buft_get_alloc_size(buft, t); allocs.push_back(a); } + // Stacked projections: place each (first, second) pair back to back in the + // weight buffer so one alias tensor spanning both rows serves a single + // GEMV. Only for the plain single-buffer path (the TP meta allocator + // places tensors itself) and only when the pair shares type/ne0 and the + // first tensor's byte size keeps the second one aligned. + const bool can_stack = !plan.metadata_only && !ggml_backend_buft_is_meta(buft) && + std::getenv("DFLASH_QWEN35_NO_STACK") == nullptr; + if (can_stack) { + auto find_alloc = [&](const std::string & name) -> int { + for (size_t i = 0; i < allocs.size(); i++) { + if (name == allocs[i].tensor->name) return (int)i; + } + return -1; + }; + // (first, second) suffix pairs; the alias tensor stacks first's rows + // then second's, so they are emitted in that order whichever member + // the file lists first. + static const char * const kPairs[][2] = { + { ".attn_gate.weight", ".attn_qkv.weight" }, + { ".ssm_beta.weight", ".ssm_alpha.weight" }, + }; + std::vector ordered; + ordered.reserve(allocs.size()); + std::vector taken(allocs.size(), false); + for (size_t i = 0; i < allocs.size(); i++) { + if (taken[i]) continue; + const std::string name = allocs[i].tensor->name; + int first = -1, second = -1; + if (name.rfind("blk.", 0) == 0) { + for (const auto & pr : kPairs) { + for (int m = 0; m < 2; m++) { + const size_t pos = name.find(pr[m]); + if (pos == std::string::npos) continue; + const std::string prefix = name.substr(0, pos); + first = find_alloc(prefix + pr[0]); + second = find_alloc(prefix + pr[1]); + break; + } + if (first >= 0 || second >= 0) break; + } + } + if (first >= 0 && second >= 0 && !taken[(size_t)first] && !taken[(size_t)second]) { + taken[(size_t)first] = taken[(size_t)second] = true; + ordered.push_back(allocs[(size_t)first]); + ordered.push_back(allocs[(size_t)second]); + continue; + } + taken[i] = true; + ordered.push_back(allocs[i]); + } + allocs.swap(ordered); + } + for (TargetTensorAlloc & a : allocs) { + alloc_total = align_up_size(alloc_total, alignment); + a.buffer_offset = alloc_total; + alloc_total += ggml_backend_buft_get_alloc_size(buft, a.tensor); + } + // The generic meta buffer allocator must see all tensors together so it // can allocate each device from its actual slices. The legacy loader's // monolithic backing buffer would reserve alloc_total on every rank. @@ -793,6 +848,47 @@ bool load_target_gguf_partial(const std::string & path, return false; } } + if (can_stack) { + // Alias tensors over adjacent pairs. They read the same bytes as + // the two source tensors (no copy, no extra VRAM). + ggml_init_params sip{}; + sip.mem_size = (2 * n_layer + 8) * ggml_tensor_overhead(); + sip.mem_buffer = nullptr; + sip.no_alloc = true; + out.stack_ctx = ggml_init(sip); + int n_stacked = 0; + auto make_stack = [&](ggml_tensor * first, ggml_tensor * second, + const char * name) -> ggml_tensor * { + if (!first || !second || !out.stack_ctx) return nullptr; + if (first->type != second->type || first->ne[0] != second->ne[0]) return nullptr; + if (!ggml_is_contiguous(first) || !ggml_is_contiguous(second)) return nullptr; + const char * f = (const char *)first->data; + const char * sd = (const char *)second->data; + if (!f || !sd || sd != f + ggml_nbytes(first)) return nullptr; + ggml_tensor * st = ggml_new_tensor_2d(out.stack_ctx, first->type, + first->ne[0], first->ne[1] + second->ne[1]); + // The alias must not need padding the backend would want to + // clear past its end (that would scribble on the next tensor). + if (ggml_backend_buft_get_alloc_size(buft, st) != ggml_nbytes(st)) return nullptr; + ggml_set_name(st, name); + if (ggml_backend_tensor_alloc(out.buf, st, first->data) != GGML_STATUS_SUCCESS) { + return nullptr; + } + n_stacked++; + return st; + }; + for (int il = 0; il < (int)n_layer; il++) { + TargetLayer & L = out.layers[il]; + char nm[96]; + std::snprintf(nm, sizeof(nm), "blk.%d.attn_gate_qkv.stacked", il); + L.wqkv_z = make_stack(L.wqkv_gate, L.wqkv, nm); + std::snprintf(nm, sizeof(nm), "blk.%d.ssm_beta_alpha.stacked", il); + L.ssm_ba = make_stack(L.ssm_beta, L.ssm_alpha, nm); + } + if (n_stacked > 0) { + std::fprintf(stderr, "[loader] stacked %d projection pairs (zero-copy aliases)\n", n_stacked); + } + } } const size_t data_start = gguf_get_data_offset(gctx); @@ -958,6 +1054,7 @@ bool load_target_gguf_partial(const std::string & path, void free_target_weights(TargetWeights & w) { if (w.buf) { ggml_backend_buffer_free(w.buf); w.buf = nullptr; } if (w.ctx) { ggml_free(w.ctx); w.ctx = nullptr; } + if (w.stack_ctx) { ggml_free(w.stack_ctx); w.stack_ctx = nullptr; } // CpuEmbedder destructor handles the mmap automatically. w.moe_hybrid.reset(); w.layers.clear(); diff --git a/server/src/qwen35/graph_builders.cpp b/server/src/qwen35/graph_builders.cpp index ad3d56b58..db929b4be 100644 --- a/server/src/qwen35/graph_builders.cpp +++ b/server/src/qwen35/graph_builders.cpp @@ -5,10 +5,90 @@ #include #include #include +#include #include namespace dflash::common { +bool detail::target_graph_capacity_for_parallel_segments( + int n_parallel_segments, + size_t & capacity) { + static constexpr int k_max_parallel_segments = 64; + static constexpr size_t k_base_capacity = 16384; + static constexpr int k_segments_per_capacity = 8; + static constexpr size_t k_max_capacity = + k_base_capacity * + (k_max_parallel_segments / k_segments_per_capacity); + + if (n_parallel_segments < 0 || + n_parallel_segments > k_max_parallel_segments) { + return false; + } + const int64_t scale = std::max( + 1, ((int64_t)n_parallel_segments + + k_segments_per_capacity - 1) / + k_segments_per_capacity); + if ((uint64_t)scale > + std::numeric_limits::max() / k_base_capacity) { + return false; + } + const size_t computed = k_base_capacity * (size_t)scale; + if (computed > k_max_capacity) return false; + capacity = computed; + return true; +} + +bool detail::target_paged_tree_graph_capacity( + int tree_width, + int n_tree_seqs, + size_t & capacity) { + static constexpr int tree_buckets[] = { + 1, 2, 3, 4, 6, 8, 12, 16, 24, 32, 48, 64, + }; + if (tree_width < 1 || tree_width > 256 || + std::find(std::begin(tree_buckets), std::end(tree_buckets), + n_tree_seqs) == std::end(tree_buckets) || + (int64_t)tree_width * n_tree_seqs > INT32_MAX) { + return false; + } + return target_graph_capacity_for_parallel_segments( + n_tree_seqs, capacity); +} + +bool detail::validate_target_paged_tree_layout( + const TargetCache & cache, + int tree_width, + int n_tree_seqs, + int paged_max_kv_len, + int tree_scratch_base, + int tree_scratch_stride) { + size_t graph_capacity = 0; + if (!target_paged_tree_graph_capacity( + tree_width, n_tree_seqs, graph_capacity) || + cache.n_seq_slots <= 1 || !cache.paged_block_table || + !cache.paged_kv_seq_lens || paged_max_kv_len < 1 || + tree_scratch_base <= 0 || + tree_scratch_base % PAGED_BLOCK_SIZE != 0 || + tree_scratch_stride < tree_width) { + return false; + } + + int physical_kv_rows = 0; + for (ggml_tensor * tensor : cache.attn_k) { + if (tensor) { + physical_kv_rows = (int)tensor->ne[1]; + break; + } + } + if (physical_kv_rows < 1) return false; + + const int64_t scratch_end = + (int64_t)tree_scratch_base + + (int64_t)(cache.n_seq_slots - 1) * tree_scratch_stride + + tree_width; + return scratch_end <= physical_kv_rows; +} + // ── build_layer_step ──────────────────────────────────────────── bool build_layer_step( @@ -332,6 +412,17 @@ bool build_target_step( } if (segment_total != n_prefill_tokens) return false; if (n_logits_rows > 0 && n_prefill_tokens == 0) return false; + size_t graph_capacity = 0; + if (!detail::target_graph_capacity_for_parallel_segments( + n_prefill_segments, graph_capacity)) { + return false; + } + // Experimental adaptive prefill: node count depends on aggregate width, + // ragged layout, and fused decode rows. Use the already-supported maximum + // capacity so the benchmark does not rely on an incomplete shape proxy. + if (n_prefill_tokens > 0) { + graph_capacity = std::max(graph_capacity, 131072); + } // Persistent thread_local arena: rebuilt step graphs land at identical // addresses, keeping the ggml-cuda CUDA-graph cache key (nodes[0]) and @@ -490,8 +581,14 @@ bool build_target_step( ggml_set_name(sg.logits_row_indices, "logits_row_indices"); ggml_set_input(sg.logits_row_indices); } + if (capture && paged_attention && cache.target_feat) { + sg.target_feat_rows = + ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, n_tokens); + ggml_set_name(sg.target_feat_rows, "target_feat_rows"); + ggml_set_input(sg.target_feat_rows); + } - sg.gf = ggml_new_graph_custom(sg.ctx, 16384, false); + sg.gf = ggml_new_graph_custom(sg.ctx, graph_capacity, false); // Step-invariant KV write: only when topology can't vary per step. // DFLASH_QWEN35_NO_KVPAD=1 restores the legacy cpy append + exact-length @@ -539,6 +636,7 @@ bool build_target_step( gi.paged_query_seq_ids = sg.paged_query_seq_ids; gi.paged_query_positions = sg.paged_query_positions; gi.logits_row_indices = sg.logits_row_indices; + gi.target_feat_rows = sg.target_feat_rows; gi.prefill_segments = prefill_segments; gi.n_prefill_segments = n_prefill_segments; @@ -630,6 +728,172 @@ bool build_target_step_tree( return ggml_gallocr_alloc_graph(sg.alloc, sg.gf); } +// ── build_target_step_paged_tree ──────────────────────────────── + +bool build_target_step_paged_tree( + StepGraph & sg, + const TargetWeights & w, + TargetCache & cache, + ggml_backend_t backend, + int tree_width, + int n_tree_seqs, + int paged_max_kv_len, + int tree_scratch_base, + int tree_scratch_stride, + int kq_stride_pad, + int mapped_ar_seqs, + bool capture_tree_commit) { + (void)kq_stride_pad; + step_graph_free(sg); + + if (mapped_ar_seqs < 0 || + (mapped_ar_seqs > 0 && !capture_tree_commit) || + mapped_ar_seqs + n_tree_seqs > cache.n_seq_slots || + mapped_ar_seqs + n_tree_seqs > 64) { + return false; + } + if (!detail::validate_target_paged_tree_layout( + cache, tree_width, n_tree_seqs, paged_max_kv_len, + tree_scratch_base, tree_scratch_stride)) { + return false; + } + size_t graph_capacity = 0; + if (!detail::target_paged_tree_graph_capacity( + tree_width, n_tree_seqs, graph_capacity)) { + return false; + } + const int n_tokens = mapped_ar_seqs + tree_width * n_tree_seqs; + const int n_mapped_seqs = mapped_ar_seqs + n_tree_seqs; + + ggml_init_params ip{}; + ip.mem_size = 512 * 1024 * 1024; + static thread_local std::vector g_tree_arena; + if (g_tree_arena.size() < ip.mem_size) g_tree_arena.resize(ip.mem_size); + ip.mem_buffer = g_tree_arena.data(); + ip.no_alloc = true; + sg.ctx = ggml_init(ip); + if (!sg.ctx) return false; + + // Salt graph addresses by the stable bucket shape so captured graphs for + // different T/S buckets never alias in ggml-cuda's topology cache. + for (int i = 0; i < tree_width + n_tree_seqs + + mapped_ar_seqs + n_tokens; ++i) { + (void)ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, 1); + } + + sg.inp_embed = ggml_new_tensor_3d( + sg.ctx, GGML_TYPE_F32, w.n_embd, n_tokens, 1); + sg.positions = + ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, 4 * n_tokens); + sg.parent_ids = ggml_new_tensor_2d( + sg.ctx, GGML_TYPE_I32, tree_width, n_tree_seqs); + sg.tree_sizes = + ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, n_tree_seqs); + sg.active_slot_ids = + ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, n_mapped_seqs); + sg.state_slot_ids = + ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, n_mapped_seqs); + sg.paged_query_seq_ids = + ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, n_tokens); + if (mapped_ar_seqs > 0) { + sg.paged_query_positions = + ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, n_tokens); + } + sg.kv_write_rows = ggml_new_tensor_2d( + sg.ctx, GGML_TYPE_I64, n_tokens, w.n_head_kv); + + const struct NamedInput { + ggml_tensor * tensor; + const char * name; + } inputs[] = { + {sg.inp_embed, "inp_embed"}, + {sg.positions, "positions"}, + {sg.parent_ids, "parent_ids"}, + {sg.tree_sizes, "tree_sizes"}, + {sg.active_slot_ids, "active_slot_ids"}, + {sg.state_slot_ids, "state_slot_ids"}, + {sg.paged_query_seq_ids, "paged_query_seq_ids"}, + {sg.paged_query_positions, "paged_query_positions"}, + {sg.kv_write_rows, "kv_write_rows"}, + }; + for (const NamedInput & input : inputs) { + if (!input.tensor) continue; + ggml_set_name(input.tensor, input.name); + ggml_set_input(input.tensor); + } + + sg.gf = ggml_new_graph_custom(sg.ctx, graph_capacity, false); + QwenGraphInputs gi{}; + gi.inp_embed = sg.inp_embed; + gi.positions = sg.positions; + gi.n_tokens = n_tokens; + gi.kv_start = 0; + gi.capture_layers = capture_tree_commit; + gi.capture_delta_intermediate = false; + gi.capture_tree_commit = capture_tree_commit; + gi.parent_ids = sg.parent_ids; + gi.tree_sizes = sg.tree_sizes; + gi.kv_write_rows = sg.kv_write_rows; + gi.paged_block_table = cache.paged_block_table; + gi.paged_kv_seq_lens = cache.paged_kv_seq_lens; + gi.active_slot_ids = sg.active_slot_ids; + gi.state_slot_ids = sg.state_slot_ids; + gi.paged_query_seq_ids = sg.paged_query_seq_ids; + gi.paged_query_positions = sg.paged_query_positions; + gi.n_seqs = n_tree_seqs; + gi.mapped_ar_seqs = mapped_ar_seqs; + gi.paged_max_kv_len = paged_max_kv_len; + gi.tree_width = tree_width; + gi.tree_scratch_base = tree_scratch_base; + gi.tree_scratch_stride = tree_scratch_stride; + + QwenGraphOutputs go = build_qwen35_graph(sg.ctx, sg.gf, w, cache, gi); + if (!go.logits) return false; + sg.logits = go.logits; + sg.delta_captures = std::move(go.delta_captures); + sg.tree_features = go.tree_features; + if (capture_tree_commit && + (!sg.tree_features || sg.delta_captures.empty())) { + return false; + } + ggml_set_output(sg.logits); + sg.argmax_tokens = ggml_argmax(sg.ctx, sg.logits); + ggml_set_name(sg.argmax_tokens, "paged_tree_verify_argmax"); + ggml_set_output(sg.argmax_tokens); + ggml_build_forward_expand(sg.gf, sg.argmax_tokens); + + if (!sg.alloc) { + sg.alloc = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(backend)); + } + if (!ggml_gallocr_alloc_graph(sg.alloc, sg.gf) || + !detail::target_paged_tree_uploads_ready(sg)) { + return false; + } + if (!capture_tree_commit) return true; + + ggml_init_params commit_params{}; + commit_params.mem_size = 16 * ggml_tensor_overhead(); + commit_params.no_alloc = true; + sg.commit_ctx = ggml_init(commit_params); + if (!sg.commit_ctx) return false; + sg.accepted_prefixes = ggml_new_tensor_1d( + sg.commit_ctx, GGML_TYPE_I32, n_tree_seqs); + sg.commit_slot_ids = ggml_new_tensor_1d( + sg.commit_ctx, GGML_TYPE_I32, n_tree_seqs); + sg.commit_rows = ggml_new_tensor_2d( + sg.commit_ctx, GGML_TYPE_I64, tree_width, n_tree_seqs); + sg.feature_commit_rows = ggml_new_tensor_1d( + sg.commit_ctx, GGML_TYPE_I32, n_tokens); + ggml_set_name(sg.accepted_prefixes, "accepted_prefixes"); + ggml_set_name(sg.commit_slot_ids, "commit_slot_ids"); + ggml_set_name(sg.commit_rows, "commit_rows"); + ggml_set_name(sg.feature_commit_rows, "feature_commit_rows"); + sg.commit_buffer = ggml_backend_alloc_ctx_tensors( + sg.commit_ctx, backend); + return sg.commit_buffer != nullptr; +} + // ── build_lm_head_projection_step ─────────────────────────────── diff --git a/server/src/qwen35/graph_builders.h b/server/src/qwen35/graph_builders.h index cdbaf75ed..b8da0fd7d 100644 --- a/server/src/qwen35/graph_builders.h +++ b/server/src/qwen35/graph_builders.h @@ -23,6 +23,60 @@ namespace dflash::common { +namespace detail { + +// Qwen's recurrent graph duplicates one small subgraph per ragged sequence. +// Return a graph capacity that covers every supported concurrent bucket while +// keeping the legacy allocation for the common <= 8-sequence case. +bool target_graph_capacity_for_parallel_segments( + int n_parallel_segments, + size_t & capacity); + +// Checked packed-tree shape/capacity contract. The public DDTree budget allows +// at most 255 children plus the root, and concurrent serving at most 64 slots. +bool target_paged_tree_graph_capacity( + int tree_width, + int n_tree_seqs, + size_t & capacity); + +// Model-free validation shared by the packed-tree builder and its shape +// tests. paged_max_kv_len is a logical launch bound and may exceed the +// bounded physical K/V pool; only the per-slot scratch slabs must fit in the +// physical tensor rows. +bool validate_target_paged_tree_layout( + const TargetCache & cache, + int tree_width, + int n_tree_seqs, + int paged_max_kv_len, + int tree_scratch_base, + int tree_scratch_stride); + +// `active_slot_ids` is a topology marker in mapped-tree graphs. It may be +// optimized out by gallocr because the actual recurrent and attention row +// mappings are carried by state_slot_ids and paged_query_seq_ids. Every other +// tensor listed here is read by a graph node and must have backend storage +// before the engine uploads metadata. +inline bool target_paged_tree_uploads_ready(const StepGraph & sg) { + const auto allocated = [](const ggml_tensor * tensor) { + return tensor && tensor->buffer; + }; + return sg.active_slot_ids && + allocated(sg.inp_embed) && allocated(sg.positions) && + allocated(sg.parent_ids) && allocated(sg.tree_sizes) && + allocated(sg.state_slot_ids) && + allocated(sg.paged_query_seq_ids) && + (!sg.paged_query_positions || + allocated(sg.paged_query_positions)) && + allocated(sg.kv_write_rows); +} + +inline bool target_paged_tree_active_slots_need_upload( + const StepGraph & sg) { + return sg.active_slot_ids && sg.active_slot_ids->buffer; +} + +} // namespace detail + // Layer-segmented prefill: process one target layer for chunk_start..chunk_start+n_tokens. bool build_layer_step( StepGraph & sg, @@ -109,6 +163,10 @@ bool build_hybrid_full_layer_step( // overrides logits_tail_rows. Multi-prompt steps need it because // committing rows are scattered. 0 keeps the tail-view behavior. // `logits_tail_rows` — logits/argmax only for the last n rows (0 = all). +// When `capture && paged_attention`, sg.target_feat_rows is an I32 graph +// input mapping every token to its slot-local feature-ring destination. This +// keeps accepted-path replay graph-stable and leaves legacy offset capture +// unchanged for callers that do not use paged serving. bool build_target_step( StepGraph & sg, const TargetWeights & w, @@ -146,6 +204,28 @@ bool build_target_step_tree( int fa_window = 0, int kq_stride_pad = KQ_MASK_PAD); +// Packed concurrent DDTree verify over a paged multi-slot cache. Tokens are +// flattened after an optional compact one-token AR prefix as +// [mapped_ar_seqs + tree_width*n_tree_seqs]. n_tree_seqs is a stable graph- +// bucket width; inactive trees use tree_size=0 and dead/safe row mappings. In +// particular, state_slot_ids padding must map to a valid harmless slot +// (normally 0), while active/paged sequence IDs may use -1. Speculative K/V is +// written into per-slot scratch slabs. With capture_tree_commit, recurrent +// transitions and target features are exposed for post-verification promotion. +bool build_target_step_paged_tree( + StepGraph & sg, + const TargetWeights & w, + TargetCache & cache, + ggml_backend_t backend, + int tree_width, + int n_tree_seqs, + int paged_max_kv_len, + int tree_scratch_base, + int tree_scratch_stride, + int kq_stride_pad = KQ_MASK_PAD, + int mapped_ar_seqs = 0, + bool capture_tree_commit = false); + // LM-head projection: project draft hidden states through the target output matrix. bool build_lm_head_projection_step( StepGraph & sg, diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index 490f86d3b..3213c1cdf 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -15,6 +15,8 @@ #include "common/geometric_sampler_cuda.h" #include #endif +#include "common/dspark_head.h" +#include "common/dflash2_head.h" #include "common/io_utils.h" #include "common/restore_delta.h" #include "qwen35_tensor_parallel.h" @@ -27,6 +29,7 @@ #include "flashprefill.h" #include +#include #include #include #include @@ -125,6 +128,7 @@ static int dflash_min_tokens_floor() { return value; } + static FILE * open_dflash_floor_log() { #if defined(_WIN32) // Simple append-mode log on Windows (no file size check). @@ -173,7 +177,7 @@ static FILE * open_dflash_floor_log() { // staging K/V or staging recurrent slab to reserve. static int64_t concurrent_fixed_cache_bytes( const TargetWeights & w, int max_ctx, int n_slots, - int64_t kv_bytes_per_token) { + int64_t kv_bytes_per_token, int64_t scratch_tokens) { const int64_t n_full_attn = w.n_layer / w.full_attention_interval; const int64_t n_delta = w.n_layer - n_full_attn; @@ -188,7 +192,8 @@ static int64_t concurrent_fixed_cache_bytes( state_per_layer * n_delta * (int64_t)n_slots; const int64_t target_feat = (int64_t)w.n_capture_layers * w.n_embd * - std::min(max_ctx, 4096) * (int64_t)sizeof(uint16_t); + ((int64_t)std::min(max_ctx, 4096) * n_slots + 1) * + (int64_t)sizeof(uint16_t); const int64_t q_capture = (int64_t)w.n_embd_head_k * w.n_head * n_full_attn * (int64_t)sizeof(float); @@ -196,12 +201,14 @@ static int64_t concurrent_fixed_cache_bytes( ((int64_t)paged_block_count(max_ctx) * n_slots + n_slots) * (int64_t)sizeof(int32_t); const int64_t scratch = - kv_bytes_per_token * PAGED_BLOCK_SIZE; + kv_bytes_per_token * scratch_tokens; return recurrent + target_feat + q_capture + paged_metadata + scratch; } } // namespace +static bool qwen35_dspark_enabled(); + #define IS_EOS_TOK(tok, w) \ ( ((w).eos_chat_id >= 0 && (tok) == (w).eos_chat_id) \ || ((w).eos_id >= 0 && (tok) == (w).eos_id ) ) @@ -215,6 +222,33 @@ static bool qwen35_empty_visible_output(const std::vector & tokens, return true; } +// Drafters trained on explicit target layers (GGUF dflash.target_layer_ids) +// override the evenly-spaced derivation: capturing different layers than the +// drafter was trained on silently destroys acceptance. +static void apply_drafter_capture_layer_ids(const DraftWeights & dw, TargetWeights & w) { + if (dw.capture_layer_ids.empty()) return; + const int n = (int)dw.capture_layer_ids.size(); + bool ok = (n == w.n_capture_layers); + for (int k = 0; ok && k < n; k++) + ok = dw.capture_layer_ids[k] >= 0 && dw.capture_layer_ids[k] < w.n_layer; + if (!ok) { + std::fprintf(stderr, + "[draft] drafter target_layer_ids invalid (n=%d, slots=%d); " + "keeping derived capture layers\n", n, w.n_capture_layers); + return; + } + bool changed = false; + for (int k = 0; k < n; k++) { + changed |= w.capture_layer_ids[k] != dw.capture_layer_ids[k]; + w.capture_layer_ids[k] = dw.capture_layer_ids[k]; + } + if (changed) { + std::printf("[draft] target capture layers from drafter GGUF:"); + for (int k = 0; k < n; k++) std::printf(" %d", w.capture_layer_ids[k]); + std::printf("\n"); + } +} + // ── Construction / destruction ────────────────────────────────────────── Qwen35Backend::Qwen35Backend(const Qwen35Config & cfg) : cfg_(cfg) {} @@ -243,6 +277,7 @@ KvFlashAutoBudget Qwen35Backend::make_kvflash_budget(const TargetWeights & w, bool Qwen35Backend::init() { configure_concurrent_hipblaslt_default(cfg_); + concurrent_decode_capabilities_ = {}; const bool use_remote_draft = cfg_.remote_draft.enabled(); const bool tensor_parallel = cfg_.device.is_tensor_parallel(); @@ -321,6 +356,7 @@ bool Qwen35Backend::init() { return false; } std::printf("[draft] loaded\n"); + apply_drafter_capture_layer_ids(dw_, w_); if (cfg_.draft_swa_window > 0) { dw_.swa_window = cfg_.draft_swa_window; @@ -331,10 +367,20 @@ bool Qwen35Backend::init() { } } + // Feature-gate validation normally catches this before construction, but + // keep backend arithmetic safe for direct/test callers too. + if (cfg_.ddtree_mode && + (cfg_.ddtree_budget < 1 || cfg_.ddtree_budget > 255)) { + set_last_error("--ddtree-budget must be in [1, 255]"); + return false; + } + // Create KV cache const int max_verify_tokens = cfg_.ddtree_mode ? std::max(dw_.block_size, cfg_.ddtree_budget + 1) : dw_.block_size; + const int n_slots = concurrent_slots(); + // kvflash (bounded residency): pool size from the env, rounded/floored/ // clamped by the shared reader (256-stride keeps FA vec-kernel // eligibility; the floor keeps eviction from deadlocking). @@ -366,20 +412,20 @@ bool Qwen35Backend::init() { if (!post_kvflash_init_gate()) return false; // KVFlash is resolved from env at init; this is the authoritative // paged×KVFlash compatibility check. - if (cfg_.paged_attention && kvflash_active()) { + if (cfg_.paged_attention && kvflash_active() && n_slots <= 1) { std::fprintf(stderr, - "[paged-attention] cannot be combined with KVFlash " - "(resident pool %d tokens)\n", kvflash_tokens_); - set_last_error("paged attention cannot be combined with KVFlash"); + "[paged-attention] single-sequence paged decode cannot be combined " + "with KVFlash (resident pool %d tokens)\n", kvflash_tokens_); + set_last_error( + "paged KVFlash requires --max-concurrency greater than 1"); return false; } // Paged mode sizes the KV cache to whole blocks; otherwise KVFlash // decides the allocation (0 = full max_ctx). - const int n_slots = concurrent_slots(); const int max_concurrent_prefills = n_slots > 1 ? std::clamp( env_int_or_default("DFLASH_MAX_CONCURRENT_PREFILLS", 8), - 1, std::min(n_slots, 8)) + 1, n_slots) : 1; const int mixed_prefill_tokens = std::max( 1, env_int_or_default("DFLASH_MIXED_PREFILL_TOKENS", 2048)); @@ -395,6 +441,32 @@ bool Qwen35Backend::init() { set_last_error("--max-concurrency requires --paged-attention"); return false; } + const bool concurrent_local_draft = + n_slots > 1 && cfg_.draft_path && + !use_remote_draft && !tensor_parallel && !split_gpus_ && + target_backend_ == draft_backend_; + const bool concurrent_local_ddtree = + concurrent_local_draft && cfg_.ddtree_mode; + const bool concurrent_local_chain = + concurrent_local_draft && !cfg_.ddtree_mode && + (dw_.selector.enabled || + (dw_.dspark.enabled && qwen35_dspark_enabled())) && + cfg_.speculation_policy != SpeculationPolicy::Never; + const bool concurrent_spec_tree = + concurrent_local_ddtree || concurrent_local_chain; + const Qwen35SeqEngine::SpecMode spec_mode = concurrent_local_ddtree + ? Qwen35SeqEngine::SpecMode::ddtree + : concurrent_local_chain + ? Qwen35SeqEngine::SpecMode::chain + : Qwen35SeqEngine::SpecMode::none; + const int tree_width = concurrent_local_ddtree + ? cfg_.ddtree_budget + 1 + : concurrent_local_chain ? dw_.block_size : 0; + const int tree_stride = concurrent_spec_tree + ? paged_token_capacity(tree_width) : 0; + const int64_t concurrent_scratch_tokens = + (int64_t)n_slots * tree_stride + PAGED_BLOCK_SIZE; + // Concurrent slots share one physical pool. An explicit // --kv-pool-tokens is rounded up to a whole block; otherwise capacity is // derived from device-free memory after subtracting fixed concurrent cache @@ -404,10 +476,17 @@ bool Qwen35Backend::init() { // pool's index space) as the write target of dead decode-batch rows. int64_t pool_tokens = 0; if (n_slots > 1) { - if (cfg_.kv_pool_tokens > 0) { + if (kvflash_active()) { + pool_tokens = paged_token_capacity(kvflash_tokens_); + std::fprintf(stderr, + "[parallel-kvflash] physical resident pool %lld tokens; " + "logical per-slot cap %d across %d slots " + "(--kv-pool-tokens does not expand resident VRAM)\n", + (long long)pool_tokens, cfg_.device.max_ctx, n_slots); + } else if (cfg_.kv_pool_tokens > 0) { pool_tokens = (int64_t)paged_token_capacity( (int)std::min( - cfg_.kv_pool_tokens, INT32_MAX - PAGED_BLOCK_SIZE)); + cfg_.kv_pool_tokens, INT32_MAX - concurrent_scratch_tokens)); } else { PagedKvAutoBudget budget; // TODO: Size tensor-parallel pools from each device's free memory @@ -416,7 +495,8 @@ bool Qwen35Backend::init() { budget.bytes_per_token = kvf_budget.bytes_per_token; budget.reserve_bytes = kvf_budget.reserve_bytes; budget.fixed_cache_bytes = concurrent_fixed_cache_bytes( - w_, cfg_.device.max_ctx, n_slots, budget.bytes_per_token); + w_, cfg_.device.max_ctx, n_slots, budget.bytes_per_token, + concurrent_scratch_tokens); pool_tokens = paged_kv_auto_pool_tokens( cfg_.device.max_ctx, n_slots, budget); const int64_t one_context = @@ -438,19 +518,19 @@ bool Qwen35Backend::init() { return false; } } - if (pool_tokens + PAGED_BLOCK_SIZE > INT32_MAX) { + if (pool_tokens + concurrent_scratch_tokens > INT32_MAX) { set_last_error("paged KV pool exceeds INT32_MAX tokens"); return false; } } const int ctx_alloc = n_slots > 1 - ? (int)(pool_tokens + PAGED_BLOCK_SIZE) + ? (int)(pool_tokens + concurrent_scratch_tokens) : (cfg_.paged_attention ? paged_token_capacity(cfg_.device.max_ctx) : kvflash_tokens_); if (!create_target_cache(w_, cfg_.device.max_ctx, max_verify_tokens, target_backend_, cache_, /*prefill_only=*/true, ctx_alloc, - cfg_.paged_attention, n_slots)) { + cfg_.paged_attention, n_slots, concurrent_spec_tree)) { std::fprintf(stderr, "cache: %s\n", dflash27b_last_error()); return false; } @@ -481,13 +561,107 @@ bool Qwen35Backend::init() { e.what()); return false; } + if (n_slots > 1 && kvflash_active()) { + std::string transfer_error; + paged_kv_transfer_ = QwenPagedKvResidencyTransfer::create( + cache_, target_backend_, cfg_.device.gpu, PAGED_BLOCK_SIZE, + &transfer_error); + if (!paged_kv_transfer_) { + set_last_error("concurrent KVFlash transfer init failed: " + + transfer_error); + return false; + } + PagedKvResidencyConfig residency_cfg; + residency_cfg.block_bytes = paged_kv_transfer_->block_bytes(); + residency_cfg.resident_budget_blocks = + paged_kv_pool_->physical_block_count(); + residency_cfg.sink_blocks = 1; + residency_cfg.tail_blocks = 4; + try { + paged_kv_residency_ = + std::make_unique( + *paged_kv_pool_, residency_cfg, + paged_kv_transfer_->callbacks()); + } catch (const std::exception & e) { + set_last_error(std::string( + "concurrent KVFlash residency init failed: ") + e.what()); + return false; + } + } if (n_slots > 1) { + const int tree_scratch_base = (int)pool_tokens; + const int64_t dead_scratch_row = + pool_tokens + (int64_t)n_slots * tree_stride; seq_engine_ = std::make_unique( *this, *paged_kv_pool_, cfg_.device.max_ctx, - /*scratch_row=*/pool_tokens, + dead_scratch_row, tree_width, tree_scratch_base, tree_stride, + spec_mode, max_concurrent_prefills, mixed_prefill_tokens, long_mixed_prefill_tokens, long_prefill_threshold, idle_prefill_tokens, prefill_quantum); + concurrent_decode_capabilities_.forced_speculation = + concurrent_local_ddtree || concurrent_local_chain; + concurrent_decode_capabilities_.adaptive = + concurrent_local_ddtree || concurrent_local_chain; + // Per-request decode_mode may select Adaptive even when the + // server default is forced speculation. A configured chain always + // accepts Adaptive: if activation scoring or startup profiling is + // unavailable, the engine uses an AR fallback instead of rejecting + // the request or failing its peers. + bool adaptive_scored = false; + if (concurrent_local_chain && + seq_engine_->activation_scoring_available()) { + int profile_ctx = 4096; + if (const char * value = + std::getenv("DFLASH_SPEC_PROFILE_CONTEXT")) { + profile_ctx = std::max(1, std::atoi(value)); + } + const bool profile_ready = + seq_engine_->profile_spec_costs(profile_ctx); + adaptive_scored = profile_ready; + if (!profile_ready) { + std::fprintf(stderr, + "[parallel-chain] adaptive scoring unavailable: " + "cost profile failed; adaptive requests will use " + "AR fallback (forced speculation remains " + "available)\n"); + } + } else if (concurrent_local_chain) { + std::fprintf(stderr, + "[parallel-chain] adaptive activation unavailable: " + "drafter has no compatible request-benefit adapter; " + "adaptive requests will use AR fallback " + "(forced speculation remains available)\n"); + } + if (concurrent_local_ddtree) { + const char * adaptive = std::getenv("DFLASH_DDTREE_ADAPTIVE"); + std::fprintf(stderr, + "[parallel-ddtree] enabled budget=%d width=%d mode=packed-verify-replay adaptive=%s\n", + cfg_.ddtree_budget, tree_width, + adaptive && std::atoi(adaptive) == 0 ? "off" : "on"); + } + if (concurrent_local_chain) { + std::fprintf(stderr, + "[parallel-chain] enabled producer=%s width=%d " + "mode=packed-chain-verify decode_mode=%s adaptive=%s\n", + dw_.selector.enabled ? "dflash2" : "dspark", + tree_width, + speculation_policy_name(cfg_.speculation_policy), + adaptive_scored ? "scored" : "fallback-ar"); + } else if (!cfg_.ddtree_mode && + cfg_.speculation_policy != SpeculationPolicy::Never) { + std::fprintf(stderr, + "[parallel-chain] unavailable for this concurrent " + "configuration; speculation/adaptive requests will be " + "rejected at admission\n"); + } else if (!cfg_.ddtree_mode && concurrent_local_draft && + (dw_.selector.enabled || + (dw_.dspark.enabled && qwen35_dspark_enabled()))) { + std::fprintf(stderr, + "[parallel-chain] disabled by decode_mode=ar; " + "per-request speculation/adaptive overrides will be " + "rejected at admission\n"); + } std::printf("[parallel] %d decode slots, up to %d packed prefills " "(mixed short/long %d/%d at >=%d tokens, " "idle %d, quantum %d), " @@ -509,7 +683,7 @@ bool Qwen35Backend::init() { cfg_.device.max_ctx); std::fflush(stdout); } - if (kvflash_active()) { + if (kvflash_active() && !(cfg_.paged_attention && n_slots > 1)) { KvFlashConfig pc; pc.pool_tokens = kvflash_tokens_; if (!kvflash_pager_.attach(pc, cache_.attn_k, cache_.attn_v)) { @@ -540,7 +714,8 @@ bool Qwen35Backend::init() { // Init feature mirror when draft model is available (needed for spec decode). // On single-GPU, this is an F32 conversion buffer; on split-GPU, a cross-device mirror. - if (cfg_.draft_path && !use_remote_draft) { + if (cfg_.draft_path && !use_remote_draft && + !(n_slots > 1 && concurrent_local_ddtree)) { const int mirror_cap = std::min({cfg_.draft_ctx_max, cfg_.device.max_ctx, cache_.target_feat_cap > 0 ? cache_.target_feat_cap : cfg_.device.max_ctx}); if (!draft_feature_mirror_init(feature_mirror_, draft_backend_, @@ -813,6 +988,7 @@ bool Qwen35Backend::unpark(ParkTarget target) { std::fprintf(stderr, "[unpark] draft: %s\n", dflash27b_last_error()); return false; } + apply_drafter_capture_layer_ids(dw_, w_); // Re-apply rope overrides after reload. if (dw_.rope_theta != w_.rope_theta && w_.rope_theta > 0.0f) dw_.rope_theta = w_.rope_theta; @@ -1184,7 +1360,12 @@ DFlashTarget * Qwen35Backend::dflash_target() { void Qwen35Backend::shutdown() { const bool use_remote_draft = cfg_.remote_draft.enabled(); + concurrent_decode_capabilities_ = {}; + seq_engine_.reset(); end_paged_sequence(); + paged_kv_residency_.reset(); + paged_kv_transfer_.reset(); + paged_kv_pool_.reset(); free_drafter(); step_graph_destroy(sg_); step_graph_destroy(draft_sg_); @@ -2318,6 +2499,64 @@ bool Qwen35Backend::sync_local_draft_features(int start_pos, int n_tokens) { // ── DFlash speculative decode loop ───────────────────────────────────── +static bool qwen35_dspark_enabled() { + static const bool kEnabled = []() { + const char * e = std::getenv("DFLASH_QWEN35_DSPARK"); + return e == nullptr || std::string(e) != "0"; + }(); + return kEnabled; +} + +// Confidence-gate threshold for adaptive block length (0 = gate off, verify +// the full drafted block). The drafter's AcceptRatePredictor scores each +// draft position; the chain is truncated at the first position below the +// threshold and only the confident prefix is verified. +static float qwen35_dspark_confidence_threshold() { + static const float kThreshold = []() { + const char * e = std::getenv("DFLASH_QWEN35_DSPARK_CONFIDENCE_THRESHOLD"); + if (!e) return 0.0f; + float threshold = (float)std::atof(e); + if (threshold < 0.0f) threshold = 0.0f; + if (threshold > 1.0f) threshold = 1.0f; + return threshold; + }(); + return kThreshold; +} + +// Adaptive speculation policy: a spec step (draft + heads + width-q verify) +// costs about DFLASH_QWEN35_SPEC_STEP_RATIO plain-decode steps, so it only +// pays off while the drafter gets more than (ratio - 1) of its tokens +// accepted per step. Below that (low-acceptance prose) the loop runs +// DFLASH_QWEN35_AR_BURST plain-decode steps inside the spec loop (target +// forward on the seed token only, features still captured for the drafter), +// then probes with one spec step. Set DFLASH_QWEN35_SPEC_STEP_RATIO=0 to +// disable the policy. +struct Qwen35AdaptiveSpecPolicy { + float step_ratio = 1.9f; // spec/plain step cost used until both step kinds have been timed + // (measured 54-55 vs 28.6 ms on gfx1201 for width-8 and width-16 verify) + int burst = 40; // plain-decode steps per burst (each burst ends with one spec probe step) + float ema_alpha = 0.1f; // slow EMA: ~10-step memory so bursty acceptance does not flap + bool enabled() const { return step_ratio > 1.0f && burst > 0; } + // Enter a burst only clearly below break-even (hysteresis against noise). + // `ratio` is the live spec/plain step-time ratio once measured. + float accept_threshold(float ratio) const { return 0.8f * (ratio - 1.0f); } + float accept_threshold() const { return accept_threshold(step_ratio); } +}; + +static Qwen35AdaptiveSpecPolicy qwen35_adaptive_spec_policy() { + static const Qwen35AdaptiveSpecPolicy kPolicy = []() { + Qwen35AdaptiveSpecPolicy p; + if (const char * e = std::getenv("DFLASH_QWEN35_SPEC_STEP_RATIO")) { + p.step_ratio = (float)std::atof(e); + } + if (const char * e = std::getenv("DFLASH_QWEN35_AR_BURST")) { + p.burst = std::atoi(e); + } + return p; + }(); + return kPolicy; +} + bool Qwen35Backend::do_spec_decode(int committed, int n_gen, std::vector & out_tokens, const DaemonIO & io, @@ -2508,8 +2747,36 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, auto t_dec0 = std::chrono::steady_clock::now(); + // Adaptive speculation state (see Qwen35AdaptiveSpecPolicy). + const Qwen35AdaptiveSpecPolicy adaptive = qwen35_adaptive_spec_policy(); + // Start well above the burst threshold so an unlucky opening does not + // park a predictable stream in plain decode. The probe step that ends a + // burst updates the EMA with a fast alpha (see below) so a stream that + // turned predictable leaves plain decode quickly. + float accepted_ema = 2.0f * adaptive.accept_threshold(); + int ar_burst_left = 0; + int n_ar_burst_steps = 0; + bool probe_step = false; // first spec step after a burst + // Live step-time EMAs (seconds) for the break-even ratio; 0 = not yet measured. + double t_spec_step_ema = 0.0; + double t_ar_step_ema = 0.0; + auto live_step_ratio = [&]() { + return (t_spec_step_ema > 0.0 && t_ar_step_ema > 0.0) + ? (float)(t_spec_step_ema / t_ar_step_ema) : adaptive.step_ratio; + }; + while (n_generated < n_gen) { const int need_commit_budget = n_gen - n_generated; + // Plain-decode step inside the spec loop: no drafter forward, verify + // the seed token only. Features are still captured, so the drafter + // resumes cleanly on the next probe step. + const bool ar_step = adaptive.enabled() && ar_burst_left > 0; + if (ar_step) { + ar_burst_left--; + n_ar_burst_steps++; + probe_step = (ar_burst_left == 0); + } + const auto t_step_start = std::chrono::steady_clock::now(); // Budget hook: no tail-off here. The close-token injection fires // during the emit phase (step 8) after acceptance+replay, mirroring @@ -2548,105 +2815,107 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, return false; } - // 2. Draft compute + // 2. Draft compute (skipped on plain-decode burst steps) constexpr int DRAFT_CTX_MAX_DEFAULT = 2048; - const int ring_cap = use_remote_draft ? remote_draft_.ring_cap() : feature_mirror_.cap; - const int draft_ctx = std::min(committed, - std::min(ring_cap, std::max(DRAFT_CTX_MAX_DEFAULT, cfg_.draft_ctx_max))); - const int draft_start = committed - draft_ctx; - int mirror_slot0 = 0; - const bool use_mirror_view = - !use_remote_draft && - draft_feature_mirror_can_view(feature_mirror_, committed, draft_ctx, mirror_slot0); - - const auto profile_draft_start = profile_start(); - if (use_remote_draft) { - local_hidden.clear(); - if (!remote_draft_.propose(committed, draft_ctx, noise_embed, local_hidden)) { - std::fprintf(stderr, "spec-decode: remote draft propose failed\n"); - step_graph_destroy(draft_sg); - return false; - } - } else { - // [TAG_DRAFT_KV] ring-cached drafter context KV: append newly - // committed rows instead of re-encoding the whole feature window. - static const bool draft_kv_on = []() { - const char * e = std::getenv("DFLASH_DRAFT_KV"); - return !(e && e[0] == '0' && e[1] == '\0'); - }(); - bool use_draft_kv = draft_kv_on && feature_mirror_.target_feat != nullptr; - if (use_draft_kv && draft_kv_.gf && - draft_kv_.built_for != (const void *)&dw_) { - draft_kv_free(draft_kv_); - } - if (use_draft_kv && !draft_kv_.gf) { - const int kv_cap = std::min(ring_cap, - std::max(DRAFT_CTX_MAX_DEFAULT, cfg_.draft_ctx_max)); - if (!draft_kv_init(draft_kv_, dw_, draft_backend_, kv_cap, nullptr)) { - draft_kv_free(draft_kv_); - use_draft_kv = false; - std::fprintf(stderr, - "spec-decode: draft-kv init failed; using legacy draft path\n"); - } - } - if (use_draft_kv) { - if (!draft_kv_begin_step(draft_kv_, dw_, draft_backend_, - feature_mirror_, committed)) { - std::fprintf(stderr, "spec-decode: draft-kv step prep failed\n"); + if (!ar_step) { + const int ring_cap = use_remote_draft ? remote_draft_.ring_cap() : feature_mirror_.cap; + const int draft_ctx = std::min(committed, + std::min(ring_cap, std::max(DRAFT_CTX_MAX_DEFAULT, cfg_.draft_ctx_max))); + const int draft_start = committed - draft_ctx; + int mirror_slot0 = 0; + const bool use_mirror_view = + !use_remote_draft && + draft_feature_mirror_can_view(feature_mirror_, committed, draft_ctx, mirror_slot0); + + const auto profile_draft_start = profile_start(); + if (use_remote_draft) { + local_hidden.clear(); + if (!remote_draft_.propose(committed, draft_ctx, noise_embed, local_hidden)) { + std::fprintf(stderr, "spec-decode: remote draft propose failed\n"); step_graph_destroy(draft_sg); return false; } - ggml_backend_tensor_set(draft_kv_.inp_embed, noise_embed.data(), 0, - sizeof(float) * noise_embed.size()); - if (ggml_backend_graph_compute(draft_backend_, draft_kv_.gf) != - GGML_STATUS_SUCCESS) { - std::fprintf(stderr, "spec-decode: draft-kv compute failed\n"); - step_graph_destroy(draft_sg); - return false; - } - local_hidden.resize((size_t)hidden * q_len); - ggml_backend_tensor_get(draft_kv_.hidden_states, local_hidden.data(), 0, - sizeof(float) * local_hidden.size()); } else { - if (!build_draft_step(draft_sg, dw_, /*lm_head=*/nullptr, draft_backend_, - draft_ctx, use_mirror_view ? &feature_mirror_ : nullptr, - committed, - /*ctx_len_max=*/std::min(ring_cap, std::max(DRAFT_CTX_MAX_DEFAULT, cfg_.draft_ctx_max)))) { - std::fprintf(stderr, "spec-decode: draft build failed\n"); - step_graph_destroy(draft_sg); - return false; - } - if (!use_mirror_view && - !copy_feature_ring_range_to_tensor(feature_mirror_, draft_sg.target_hidden_cat, - draft_start, draft_ctx)) { - std::fprintf(stderr, "spec-decode: feature copy failed\n"); - step_graph_destroy(draft_sg); - return false; + // [TAG_DRAFT_KV] ring-cached drafter context KV: append newly + // committed rows instead of re-encoding the whole feature window. + static const bool draft_kv_on = []() { + const char * e = std::getenv("DFLASH_DRAFT_KV"); + return !(e && e[0] == '0' && e[1] == '\0'); + }(); + bool use_draft_kv = draft_kv_on && feature_mirror_.target_feat != nullptr; + if (use_draft_kv && draft_kv_.gf && + draft_kv_.built_for != (const void *)&dw_) { + draft_kv_free(draft_kv_); } - ggml_backend_tensor_set(draft_sg.inp_embed, noise_embed.data(), 0, - sizeof(float) * noise_embed.size()); - pos_k.resize((size_t)draft_ctx + q_len); - for (int i = 0; i < q_len; i++) pos_q[i] = draft_ctx + i; - for (int i = 0; i < draft_ctx + q_len; i++) pos_k[i] = i; - ggml_backend_tensor_set(draft_sg.positions, pos_q.data(), 0, - sizeof(int32_t) * pos_q.size()); - ggml_backend_tensor_set(draft_sg.positions_k, pos_k.data(), 0, - sizeof(int32_t) * pos_k.size()); - - auto st = ggml_backend_graph_compute(draft_backend_, draft_sg.gf); - if (st != GGML_STATUS_SUCCESS) { - std::fprintf(stderr, "spec-decode: draft compute failed\n"); - step_graph_destroy(draft_sg); - return false; + if (use_draft_kv && !draft_kv_.gf) { + const int kv_cap = std::min(ring_cap, + std::max(DRAFT_CTX_MAX_DEFAULT, cfg_.draft_ctx_max)); + if (!draft_kv_init(draft_kv_, dw_, draft_backend_, kv_cap, nullptr)) { + draft_kv_free(draft_kv_); + use_draft_kv = false; + std::fprintf(stderr, + "spec-decode: draft-kv init failed; using legacy draft path\n"); + } } + if (use_draft_kv) { + if (!draft_kv_begin_step(draft_kv_, dw_, draft_backend_, + feature_mirror_, committed)) { + std::fprintf(stderr, "spec-decode: draft-kv step prep failed\n"); + step_graph_destroy(draft_sg); + return false; + } + ggml_backend_tensor_set(draft_kv_.inp_embed, noise_embed.data(), 0, + sizeof(float) * noise_embed.size()); + if (ggml_backend_graph_compute(draft_backend_, draft_kv_.gf) != + GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "spec-decode: draft-kv compute failed\n"); + step_graph_destroy(draft_sg); + return false; + } + local_hidden.resize((size_t)hidden * q_len); + ggml_backend_tensor_get(draft_kv_.hidden_states, local_hidden.data(), 0, + sizeof(float) * local_hidden.size()); + } else { + if (!build_draft_step(draft_sg, dw_, /*lm_head=*/nullptr, draft_backend_, + draft_ctx, use_mirror_view ? &feature_mirror_ : nullptr, + committed, + /*ctx_len_max=*/std::min(ring_cap, std::max(DRAFT_CTX_MAX_DEFAULT, cfg_.draft_ctx_max)))) { + std::fprintf(stderr, "spec-decode: draft build failed\n"); + step_graph_destroy(draft_sg); + return false; + } + if (!use_mirror_view && + !copy_feature_ring_range_to_tensor(feature_mirror_, draft_sg.target_hidden_cat, + draft_start, draft_ctx)) { + std::fprintf(stderr, "spec-decode: feature copy failed\n"); + step_graph_destroy(draft_sg); + return false; + } + ggml_backend_tensor_set(draft_sg.inp_embed, noise_embed.data(), 0, + sizeof(float) * noise_embed.size()); + pos_k.resize((size_t)draft_ctx + q_len); + for (int i = 0; i < q_len; i++) pos_q[i] = draft_ctx + i; + for (int i = 0; i < draft_ctx + q_len; i++) pos_k[i] = i; + ggml_backend_tensor_set(draft_sg.positions, pos_q.data(), 0, + sizeof(int32_t) * pos_q.size()); + ggml_backend_tensor_set(draft_sg.positions_k, pos_k.data(), 0, + sizeof(int32_t) * pos_k.size()); + + auto st = ggml_backend_graph_compute(draft_backend_, draft_sg.gf); + if (st != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "spec-decode: draft compute failed\n"); + step_graph_destroy(draft_sg); + return false; + } - // Read draft hidden states to host for LM-head projection. - local_hidden.resize((size_t)hidden * q_len); - ggml_backend_tensor_get(draft_sg.hidden_states, local_hidden.data(), 0, - sizeof(float) * local_hidden.size()); + // Read draft hidden states to host for LM-head projection. + local_hidden.resize((size_t)hidden * q_len); + ggml_backend_tensor_get(draft_sg.hidden_states, local_hidden.data(), 0, + sizeof(float) * local_hidden.size()); + } } - } - profile_add(profile_draft_s, profile_draft_start); + profile_add(profile_draft_s, profile_draft_start); + } // !ar_step // ── DDTree tree-structured verify ──────────────────────────────── // When --ddtree is on and the target supports tree verify, build a @@ -2678,17 +2947,119 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, kvflash_pager_.identity_prefix_covers(committed)); const bool use_tree_verify = cfg_.ddtree_mode && target->supports_tree_verify() && kvflash_tree_ok && - !use_remote_draft && q_len > 1 && tree_special_inactive; + !use_remote_draft && q_len > 1 && tree_special_inactive && !ar_step; + // Chain-verify length for this step. The DSpark confidence gate may + // truncate the drafted block (adaptive block length); q_len stays the + // buffer-sizing upper bound. + int v_len = q_len; // DDTree consumes top-K rows directly. Avoid projecting the same // hidden block once for argmax and again for top-K on every step. - if (!use_tree_verify) { - if (!target->project_hidden_to_tokens(local_hidden.data(), q_len, draft_tok)) { - std::fprintf(stderr, "spec-decode: projection failed\n"); - step_graph_destroy(draft_sg); - return false; + if (ar_step) { + draft_tok.assign(1, last_tok); + v_len = 1; + } else if (!use_tree_verify) { + const auto profile_project_start = profile_start(); + // DFlash 2 selector (top-k candidates + low-rank path score) when + // the drafter ships it. + bool used_dspark = false; + if (dw_.selector.enabled && q_len > 1 && !sampled_verify && !use_remote_draft) { + static std::atomic s_sel_logged{false}; + if (!s_sel_logged.exchange(true)) { + std::fprintf(stderr, + "[qwen35-spec] DFlash 2 selector active for greedy chain decode " + "(rank=%d top_k=%d)\n", dw_.selector.rank, dw_.selector.top_k); + } + if (dflash2_select_chain(dw_, draft_backend_, *target, + local_hidden.data(), q_len, last_tok, draft_tok)) { + used_dspark = true; + v_len = std::max(1, (int)draft_tok.size()); + } else { + static std::atomic s_sel_warned{false}; + if (!s_sel_warned.exchange(true)) { + std::fprintf(stderr, + "[qwen35-spec] DFlash 2 selector failed; falling back to " + "base DFlash projection\n"); + } + } + } + // DSpark heads (markov bigram correction + optional confidence + // gate) when the drafter ships them; mirrors the laguna hook. + if (!used_dspark && qwen35_dspark_enabled() && dw_.dspark.enabled && + q_len > 1 && !sampled_verify && !use_remote_draft) { + static std::atomic s_dspark_logged{false}; + if (!s_dspark_logged.exchange(true)) { + std::fprintf(stderr, + "[qwen35-spec] DSpark Markov head active for greedy chain decode " + "(rank=%d vocab=%d confidence_dim=%d)\n", + dw_.dspark.markov_rank, dw_.dspark.vocab_size, + dw_.dspark.confidence_dim); + } + static const bool fused_dspark = []() { + const char * e = std::getenv("DFLASH_QWEN35_FUSED_DSPARK"); + return !(e && e[0] == '0' && e[1] == '\0'); + }(); + bool ds_ok = false; + const float conf_threshold = qwen35_dspark_confidence_threshold(); + if (fused_dspark) { + // One graph for every candidate: markov-corrected tokens + // plus (when gated) the confidence score per position. + std::vector conf_scores; + ds_ok = dspark_markov_correct_greedy_chain_fused( + dw_, draft_backend_, target->lm_head_tensor(), + local_hidden.data(), q_len, last_tok, draft_tok, + conf_threshold > 0.0f ? &conf_scores : nullptr); + if (ds_ok && conf_threshold > 0.0f) { + // Truncate the chain at the first low-confidence + // position: draft_tok[0] is the seed, candidate i + // scores conf_scores[i-1]. + size_t keep = 1; + while (keep < draft_tok.size() && + keep - 1 < conf_scores.size() && + conf_scores[keep - 1] >= conf_threshold) { + ++keep; + } + static const bool conf_debug = []() { + const char * e = std::getenv("DFLASH_QWEN35_DSPARK_CONF_DEBUG"); + return e && e[0] == '1'; + }(); + if (conf_debug) { + std::fprintf(stderr, "[dspark-conf] keep=%zu/%zu:", keep, draft_tok.size()); + for (float c : conf_scores) std::fprintf(stderr, " %.3f", c); + std::fprintf(stderr, "\n"); + } + draft_tok.resize(keep); + } + } + if (!ds_ok) { + ds_ok = dspark_markov_correct_greedy_chain(dw_, draft_backend_, *target, + local_hidden.data(), q_len, + last_tok, conf_threshold, + draft_tok); + } + if (ds_ok) { + used_dspark = true; + // Confidence gate truncates the drafted chain: verify + // only the confident prefix this step. + v_len = std::max(1, (int)draft_tok.size()); + } else { + static std::atomic s_dspark_warned{false}; + if (!s_dspark_warned.exchange(true)) { + std::fprintf(stderr, + "[qwen35-spec] DSpark Markov head failed; falling back to " + "base DFlash projection\n"); + } + } } - draft_tok[0] = last_tok; + if (!used_dspark) { + if (!target->project_hidden_to_tokens(local_hidden.data(), q_len, draft_tok)) { + std::fprintf(stderr, "spec-decode: projection failed\n"); + step_graph_destroy(draft_sg); + return false; + } + draft_tok[0] = last_tok; + } + profile_add(profile_project_s, profile_project_start); } if (use_tree_verify) { @@ -2697,7 +3068,25 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, std::vector top_lp; std::vector top_ids; const auto profile_project_start = profile_start(); - if (!target->project_hidden_to_topk(local_hidden.data(), q_len, K, + static const bool dspark_tree = []() { + const char * e = std::getenv("DFLASH_QWEN35_DSPARK_TREE"); + return !(e && e[0] == '0' && e[1] == '\0'); + }(); + bool topk_ok = false; + if (dspark_tree && qwen35_dspark_enabled() && dw_.dspark.enabled) { + static std::atomic s_dstree_logged{false}; + if (!s_dstree_logged.exchange(true)) { + std::fprintf(stderr, + "[qwen35-spec] DSpark Markov head active for DDTree candidates\n"); + } + topk_ok = dspark_markov_project_topk(dw_, draft_backend_, + target->lm_head_tensor(), + local_hidden.data(), q_len, K, + cfg_.ddtree_temp, last_tok, + top_lp, top_ids); + } + if (!topk_ok && + !target->project_hidden_to_topk(local_hidden.data(), q_len, K, cfg_.ddtree_temp, top_lp, top_ids)) { std::fprintf(stderr, "spec-decode: ddtree topk projection failed\n"); step_graph_destroy(draft_sg); @@ -2961,7 +3350,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, int hint_fill = 0; if (hint_tokens && n_generated < (int)hint_tokens->size()) { const int hint_avail = (int)hint_tokens->size() - n_generated; - hint_fill = std::min(hint_avail, q_len - 1); + hint_fill = std::min(hint_avail, v_len - 1); for (int i = 0; i < hint_fill; i++) { draft_tok[1 + i] = (*hint_tokens)[n_generated + i]; } @@ -2972,13 +3361,18 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, io.observer("draft", draft_tok); } - // 4. Verify: snapshot KV, run target forward over draft tokens - if (!target->snapshot_kv()) { + // 4. Verify: snapshot KV, run target forward over draft tokens. + // A plain-decode step verifies only the (always accepted) seed, so + // it never rolls back: skip the snapshot copy. + const auto profile_snapshot_start = profile_start(); + if (!ar_step && !target->snapshot_kv()) { step_graph_destroy(draft_sg); return false; } + profile_add(profile_snapshot_s, profile_snapshot_start); int verify_last_tok = -1; + const auto profile_verify_start = profile_start(); if (!target->verify_batch(draft_tok, committed, verify_last_tok, &target_tok, /*capture_ssm_intermediates=*/true)) { std::fprintf(stderr, "spec-decode: verify failed\n"); @@ -2986,6 +3380,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, step_graph_destroy(draft_sg); return false; } + profile_add(profile_verify_s, profile_verify_start); target_forwards++; // 5. Acceptance. Greedy: longest matching prefix between draft and @@ -2996,13 +3391,13 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, int accept_n = 1; int bonus_tok = -1; if (sampled_verify) { - if (!target->read_verify_logits(q_len, verify_logits)) { + if (!target->read_verify_logits(v_len, verify_logits)) { std::fprintf(stderr, "spec-decode: verify logits read failed\n"); target->restore_kv(); step_graph_destroy(draft_sg); return false; } - const int vocab_v = (int)(verify_logits.size() / (size_t)q_len); + const int vocab_v = (int)(verify_logits.size() / (size_t)v_len); static const bool kSvDebug = []() { const char * e = std::getenv("DFLASH_SV_DEBUG"); return e != nullptr && std::string(e) == "1"; @@ -3011,7 +3406,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, // Row-alignment check: CPU argmax over each bulk-read row must // equal the GPU argmax (target_tok). Divergence = misaligned // or stale bulk read. - for (int i = 0; i < q_len; i++) { + for (int i = 0; i < v_len; i++) { const float * row = verify_logits.data() + (size_t)i * vocab_v; int am = 0; float best = row[0]; for (int v = 1; v < vocab_v; v++) @@ -3036,7 +3431,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, verify_history = out_tokens; verify_history.push_back(draft_tok[0]); bool mismatched = false; - for (int i = 0; i < q_len - 1; i++) { + for (int i = 0; i < v_len - 1; i++) { const int s = sample_logits( verify_logits.data() + (size_t)i * vocab_v, vocab_v, sampler_, verify_history, sampler_rng_); @@ -3057,11 +3452,11 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, } (void)mismatched; } else { - for (int i = 0; i < q_len - 1; i++) { + for (int i = 0; i < v_len - 1; i++) { if (draft_tok[i + 1] == target_tok[i]) accept_n++; else break; } - bonus_tok = (accept_n < q_len) ? target_tok[accept_n - 1] : -1; + bonus_tok = (accept_n < v_len) ? target_tok[accept_n - 1] : -1; } // Track hint acceptance telemetry. if (hint_fill > 0) { @@ -3086,7 +3481,14 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, int replay_last_tok = -1; bool fast_rolled_back = false; - if (use_fast_rollback) { + if (ar_step) { + // Seed-only verify: the recurrent state already sits after the + // one committed token; nothing to restore. + bonus_tok = -1; + commit_n = std::min(accept_n, need_commit_budget); + replay_last_tok = target_tok[commit_n - 1]; + fast_rolled_back = true; + } else if (use_fast_rollback) { // Fast rollback: restore SSM from captured intermediates, skip replay. // Implicit bonus: target_tok[commit_n-1] seeds next draft as draft_tok[0], // always accepted on next step. @@ -3095,7 +3497,10 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, // budget (need_commit_budget), so committing accept_n would emit // more tokens than requested. commit_n was already clamped above. commit_n = std::min(accept_n, need_commit_budget); - if (target->rollback_to(committed, commit_n)) { + const auto profile_rollback_start = profile_start(); + const bool rolled = target->rollback_to(committed, commit_n); + profile_add(profile_rollback_s, profile_rollback_start); + if (rolled) { replay_last_tok = target_tok[commit_n - 1]; fast_rolled_back = true; rollback_diag.record_fast_rollback(accept_n); @@ -3121,11 +3526,13 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, for (int i = 0; i < commit_n; i++) { replay_batch[i] = (i < accept_n) ? draft_tok[i] : bonus_tok; } + const auto profile_replay_start = profile_start(); if (!target->verify_batch(replay_batch, committed, replay_last_tok, nullptr)) { std::fprintf(stderr, "spec-decode: replay failed\n"); step_graph_destroy(draft_sg); return false; } + profile_add(profile_replay_s, profile_replay_start); target_forwards++; } @@ -3142,10 +3549,12 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, return false; } } else if (feature_mirror_.target_feat && cache_.target_feat) { + const auto profile_feature_start = profile_start(); if (!sync_local_draft_features(committed, commit_n)) { step_graph_destroy(draft_sg); return false; } + profile_add(profile_feature_s, profile_feature_start); } // 8. Emit committed tokens (stop at EOS) @@ -3289,6 +3698,27 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, n_accept_sum += std::min(accept_n, emitted); n_draft_steps++; + // Adaptive policy update on real spec steps: EMA of accepted draft + // tokens (the seed is always accepted); a low EMA schedules a burst + // of plain-decode steps, the step after the burst is a spec probe. + if (adaptive.enabled()) { + const double t_step = std::chrono::duration( + std::chrono::steady_clock::now() - t_step_start).count(); + double & t_ema = ar_step ? t_ar_step_ema : t_spec_step_ema; + t_ema = (t_ema > 0.0) ? 0.9 * t_ema + 0.1 * t_step : t_step; + } + if (adaptive.enabled() && !ar_step) { + const float accepted_drafts = (float)std::max(0, accept_n - 1); + // A probe (first spec step after a burst) weighs its result + // heavily: it is the only evidence about the current text. + const float alpha = probe_step ? 0.5f : adaptive.ema_alpha; + accepted_ema = (1.0f - alpha) * accepted_ema + alpha * accepted_drafts; + probe_step = false; + if (accepted_ema < adaptive.accept_threshold(live_step_ratio())) { + ar_burst_left = adaptive.burst; + } + } + // Notify observer with accepted tokens for this step. if (io.observer) { io.observer("verify", replay_tok); @@ -3375,6 +3805,12 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, n_generated > 0 ? n_generated / decode_s : 0.0, n_draft_steps, n_accept_sum, total_draft_pos, accept_pct, n_draft_steps > 0 ? (double)n_generated / (double)n_draft_steps : 0.0); + if (n_ar_burst_steps > 0) { + std::fprintf(stderr, "[spec-decode] adaptive: %d of %d steps ran as plain decode " + "(step ratio %.2f, accept threshold %.2f drafts/step, burst %d)\n", + n_ar_burst_steps, n_draft_steps, live_step_ratio(), + adaptive.accept_threshold(live_step_ratio()), adaptive.burst); + } if (tp_profile) { std::fprintf(stderr, "[spec-profile] draft=%.3fs project=%.3fs snapshot=%.3fs " diff --git a/server/src/qwen35/qwen35_backend.h b/server/src/qwen35/qwen35_backend.h index a6a508f73..6f9545b5d 100644 --- a/server/src/qwen35/qwen35_backend.h +++ b/server/src/qwen35/qwen35_backend.h @@ -18,13 +18,16 @@ #include "placement/remote_draft_config.h" #include "step_graph.h" #include "ddtree.h" +#include "common/speculation_policy.h" #include "dflash_feature_ring.h" #include "common/dflash_draft_kv.h" #include "common/concurrency/paged_kv_pool.h" #include "concurrency/qwen35_seq_engine.h" #include "internal.h" // TargetWeights, TargetCache, DraftWeights, PrefixSnapshot #include "qwen3/qwen3_drafter.h" // DrafterContext, load_drafter, free_drafter, drafter_score_and_compress -#include "kvflash_pager.h" // bounded KV residency pool +#include "kvflash_pager.h" +#include "common/concurrency/paged_kv_residency.h" +#include "common/concurrency/qwen_paged_kv_transfer.h" #include "kvflash_scorer.h" // chunk-relevance policy interface #include "kvflash_qk.h" // target-QK scorer (pooled keys + query) @@ -84,6 +87,7 @@ struct Qwen35Config { float ddtree_temp = 1.0f; bool ddtree_chain_seed = true; bool use_feature_mirror = false; + SpeculationPolicy speculation_policy = SpeculationPolicy::Adaptive; }; // ── Backend class ─────────────────────────────────────────────────────── @@ -145,6 +149,10 @@ class Qwen35Backend : public ModelBackend { // attention); null otherwise, which is what tells the server to serve // one request at a time through generate(). SeqEngine * seq_engine() override; + ConcurrentDecodeCapabilities concurrent_decode_capabilities() + const override { + return concurrent_decode_capabilities_; + } // EOS identity of the loaded weights. Model-level, so it stays on the // backend and is shared by the AR decode path and the engine. @@ -295,6 +303,8 @@ class Qwen35Backend : public ModelBackend { // Page size comes from PAGED_BLOCK_SIZE (paged_attention_config.h), // shared with the graph builder and the cache's block-aligned sizing. std::unique_ptr paged_kv_pool_; + std::unique_ptr paged_kv_transfer_; + std::unique_ptr paged_kv_residency_; std::optional paged_sequence_; PagedKvRequestId paged_request_id_ = 0; @@ -312,6 +322,7 @@ class Qwen35Backend : public ModelBackend { // hence the friendship — and owns everything else concurrent serving // needs (Qwen35SlotManager, slot prefill, the batched decode step). std::unique_ptr seq_engine_; + ConcurrentDecodeCapabilities concurrent_decode_capabilities_; friend class Qwen35SeqEngine; // DFLASH_MIN_TOKENS floor for the slot paths (mirrors do_ar_decode's diff --git a/server/src/qwen35/qwen35_dflash_target.h b/server/src/qwen35/qwen35_dflash_target.h index 3c8864b6b..cc8a37c3d 100644 --- a/server/src/qwen35/qwen35_dflash_target.h +++ b/server/src/qwen35/qwen35_dflash_target.h @@ -75,6 +75,7 @@ class Qwen35DFlashTarget : public DFlashTarget { int hidden_size() const override { return w_.n_embd; } int mask_token_id() const override; + ggml_tensor * lm_head_tensor() override { return w_.output; } const std::vector & capture_layer_ids() const override; // kvflash mode: verify writes are slot-mapped via the pager and the diff --git a/server/src/qwen35/qwen35_target_graph.cpp b/server/src/qwen35/qwen35_target_graph.cpp index 1082df6f5..fb1a9a779 100644 --- a/server/src/qwen35/qwen35_target_graph.cpp +++ b/server/src/qwen35/qwen35_target_graph.cpp @@ -80,12 +80,14 @@ bool create_target_cache(const TargetWeights & w, bool prefill_only, int ctx_alloc, bool paged_attention, - int n_seq_slots) { + int n_seq_slots, + bool concurrent_tree) { return create_target_cache_partial(w, max_ctx, max_verify_tokens, backend, out, prefill_only, 0, w.n_layer, true, ctx_alloc, /*f32_ssm_intermediates=*/false, - paged_attention, n_seq_slots); + paged_attention, n_seq_slots, + concurrent_tree); } // concurrent_fixed_cache_bytes() in qwen35_backend.cpp mirrors this @@ -103,7 +105,8 @@ bool create_target_cache_partial(const TargetWeights & w, int ctx_alloc, bool f32_ssm_intermediates, bool paged_attention, - int n_seq_slots) { + int n_seq_slots, + bool concurrent_tree) { if (layer_begin < 0) layer_begin = 0; if (layer_end < 0 || layer_end > w.n_layer) layer_end = w.n_layer; if (layer_begin > layer_end) { @@ -115,6 +118,11 @@ bool create_target_cache_partial(const TargetWeights & w, set_last_error("multi-slot target cache requires paged attention"); return false; } + if (concurrent_tree && (!paged_attention || n_seq_slots <= 1)) { + set_last_error( + "concurrent tree cache requires paged multi-slot serving"); + return false; + } out.backend = backend; out.max_ctx = max_ctx; out.cur_pos = 0; @@ -148,7 +156,13 @@ bool create_target_cache_partial(const TargetWeights & w, // Graph-level FWHT K-rotation (TurboQuant-style outlier spreading with // standard quant types that keep fast FA kernel paths on all arches). // Skip for TQ3_0 K cache — that type already applies WHT during quantization. - out.kv_k_rotated = (kv_k_type != GGML_TYPE_TQ3_0); + // DFLASH_KV_ROTATE=0 turns it off (two fewer launches per attention layer; + // with q8_0/f16 caches the rotation is precision-neutral). + static const bool kv_rotate_env = []() { + const char * e = std::getenv("DFLASH_KV_ROTATE"); + return !(e && e[0] == '0' && e[1] == '\0'); + }(); + out.kv_k_rotated = (kv_k_type != GGML_TYPE_TQ3_0) && kv_rotate_env; const bool needs_256_stride = kv_k_type == GGML_TYPE_TQ3_0 || kv_v_type == GGML_TYPE_TQ3_0; @@ -227,7 +241,14 @@ bool create_target_cache_partial(const TargetWeights & w, out.target_feat_cap = std::min(max_ctx, TARGET_FEAT_CAP_DEFAULT); if (allocate_target_feat) { const int fc_in = w.n_capture_layers * w.n_embd; - out.target_feat = ggml_new_tensor_2d(out.base_ctx, GGML_TYPE_BF16, fc_in, out.target_feat_cap); + // Concurrent slots own disjoint feature rings. The final row is + // dead scratch for padded bucket rows because set_rows does not + // accept negative destination indices. + const int feat_rows = multi_slot + ? out.target_feat_cap * n_seq_slots + 1 + : out.target_feat_cap; + out.target_feat = ggml_new_tensor_2d( + out.base_ctx, GGML_TYPE_BF16, fc_in, feat_rows); ggml_set_name(out.target_feat, "target_feat"); } else { out.target_feat = nullptr; @@ -266,9 +287,10 @@ bool create_target_cache_partial(const TargetWeights & w, } // ── Rollback context: snapshots + intermediates ─────────────────── - // Multi-slot caches skip these entirely: concurrent serving is paged and - // therefore AR-only (no spec-decode rollback), and the tensors are the - // single largest optional allocation (~0.8 GB at 48 delta layers). + // Multi-slot caches skip these entirely. Packed tree verification gathers + // the selected slots' base state without mutating it, then a bounded + // replay commits accepted paths. T*S recurrent captures would be tens of + // GiB at useful Strix concurrency and are intentionally not allocated. if (!prefill_only && !multi_slot) { const int rb_tensors = 4 * n_delta; ggml_init_params ip{}; @@ -673,10 +695,19 @@ bool ensure_ssm_snapshot(TargetCache & c, ggml_backend_t backend) { static ggml_tensor * build_swiglu_ffn(ggml_context * ctx, ggml_tensor * cur, const TargetLayer & L) { - ggml_tensor * gate = apply_scale2(ctx, ggml_mul_mat(ctx, L.w_gate, cur), L.w_gate_s); // [inter, n_tokens] - gate = ggml_silu(ctx, gate); - ggml_tensor * up = apply_scale2(ctx, ggml_mul_mat(ctx, L.w_up, cur), L.w_up_s); - ggml_tensor * gu = ggml_mul(ctx, gate, up); + ggml_tensor * gate = ggml_mul_mat(ctx, L.w_gate, cur); // [inter, n_tokens] + ggml_tensor * up = ggml_mul_mat(ctx, L.w_up, cur); + ggml_tensor * gu; + if (L.w_gate_s == 1.0f && L.w_up_s == 1.0f) { + // GLU node right after the two matmuls: the CUDA/HIP backend fuses + // mul_mat(gate) + mul_mat(up) + swiglu into a single vector kernel + // for single-token decode. + gu = ggml_swiglu_split(ctx, gate, up); + } else { + gate = ggml_silu(ctx, apply_scale2(ctx, gate, L.w_gate_s)); + up = apply_scale2(ctx, up, L.w_up_s); + gu = ggml_mul(ctx, gate, up); + } return apply_scale2(ctx, ggml_mul_mat(ctx, L.w_down, gu), L.w_down_s); // [hidden, n_tokens] } @@ -722,7 +753,16 @@ static ggml_tensor * build_full_attn_block( int paged_max_kv_len = 0, // Compact decode row -> physical block-table column. Negative ids are // graph-bucket padding rows. - ggml_tensor * active_slot_ids = nullptr + ggml_tensor * active_slot_ids = nullptr, + // Packed paged-tree verification. Query rows are flattened + // sequence-major; row mappings are supplied through + // paged_query_seq_ids, while parent/tree metadata describes each tree. + ggml_tensor * paged_tree_parent_ids = nullptr, + ggml_tensor * paged_tree_sizes = nullptr, + int tree_width = 0, + int tree_scratch_base = 0, + int tree_scratch_stride = 0, + int paged_logical_max_ctx = 0 ) { const int head_dim = w.n_embd_head_k; const int n_head = w.n_head; @@ -810,9 +850,13 @@ static ggml_tensor * build_full_attn_block( Kcur_T = ggml_turbo_wht(ctx, Kcur_T, 0); } + const bool paged_tree = paged_tree_parent_ids || paged_tree_sizes; + GGML_ASSERT((paged_tree_parent_ids == nullptr) == + (paged_tree_sizes == nullptr)); const bool ragged = paged_query_seq_ids != nullptr; - GGML_ASSERT(!ragged || (paged_block_table && paged_query_positions && - kv_write_rows)); + GGML_ASSERT(!ragged || (paged_block_table && kv_write_rows)); + GGML_ASSERT(!ragged || paged_tree || paged_query_positions); + GGML_ASSERT(!paged_tree || (ragged && tree_width > 0)); if (kv_write_rows) { // Step-invariant: the destination tensor stays fixed while the input // indices carry contiguous, KVFlash, or paged physical rows. @@ -878,12 +922,31 @@ static ggml_tensor * build_full_attn_block( ggml_tensor * row_seq_ids, ggml_tensor * row_positions, bool dense_token_layout) { - const int padded = ((std::max(1, launch_kv_len) + 255) / 256) * 256; - const int launch_len = std::min(padded, (int)cache_k->ne[1]); + // max_kv_seq_len sizes the logical partition grid. In bounded + // KVFlash mode the block table maps that logical range onto a much + // smaller physical K/V pool, so cache_k->ne[1] is not a valid clamp. + // Bound against both sources of logical capacity instead, doing the + // 256-window rounding in i64 to avoid signed overflow at large + // configured contexts. Per-row kv_seq_lens remains the exact runtime + // bound, and the paged kernel bounds every resolved physical row. + GGML_ASSERT(paged_block_table && cache_k && cache_v); + const int64_t table_capacity = + (int64_t)paged_block_table->ne[0] * PAGED_BLOCK_SIZE; + const int64_t logical_capacity = + std::min(paged_logical_max_ctx, table_capacity); + GGML_ASSERT(logical_capacity > 0 && logical_capacity <= INT32_MAX); + const int64_t requested = + std::min(std::max(1, launch_kv_len), + logical_capacity); + const int64_t padded = ((requested + 255) / 256) * 256; + const int launch_len = + (int)std::min(padded, logical_capacity); ggml_tensor * out = ggml_paged_attn_ext( ctx, q, cache_k, cache_v, paged_block_table, paged_kv_seq_lens, row_seq_ids, row_positions, kq_scale, - PAGED_BLOCK_SIZE, launch_len); + PAGED_BLOCK_SIZE, launch_len, + paged_tree_parent_ids, paged_tree_sizes, + tree_width, tree_scratch_base, tree_scratch_stride); if (dense_token_layout) { out = ggml_cont(ctx, ggml_permute(ctx, out, 0, 2, 1, 3)); } @@ -891,7 +954,20 @@ static ggml_tensor * build_full_attn_block( }; ggml_tensor * attn = nullptr; - if (ragged) { + if (paged_tree) { + // ── Packed concurrent tree verify. Every query row selects its + // physical sequence/scratch slab. The paged kernel combines the + // committed block-table prefix with only this node.s ancestor chain. + // A mixed graph uses causal positions for the compact AR prefix and + // -1 for the tree tail; a pure tree keeps positions absent. + ggml_tensor * Qfa = q_segment(0, n_tokens); + if (q_fa_out) *q_fa_out = Qfa; + const int launch_kv_len = paged_max_kv_len > 0 + ? paged_max_kv_len : kv_start + n_tokens; + attn = paged_read(Qfa, launch_kv_len, + paged_query_seq_ids, paged_query_positions, + /*dense_token_layout=*/n_tokens > 1); + } else if (ragged) { // ── Ragged concurrent step: prefill chunk rows and decode rows all // read the pool through one call, each row clamped to its own // inclusive position. This step's chunk rows are visible to their @@ -899,6 +975,7 @@ static ggml_tensor * build_full_attn_block( // attention in the graph; cross-sequence isolation is structural // (each row's seq id selects its own block-table column). ggml_tensor * Qfa = q_segment(0, n_tokens); + if (q_fa_out) *q_fa_out = Qfa; const int launch_kv_len = paged_max_kv_len > 0 ? paged_max_kv_len : kv_start + n_tokens; attn = paged_read(Qfa, launch_kv_len, @@ -920,8 +997,8 @@ static ggml_tensor * build_full_attn_block( // bound only over-sizes the partition grid, and partitions past the // real length exit with a zero-weight sentinel. // Batched decode: kv_len (kv_start + n_tokens) describes one sequence; - // the launch bound must cover the longest live slot instead. Clamped - // because ggml_paged_attn asserts max_kv_seq_len <= k->ne[1]. + // the launch bound must cover the longest live slot instead. Bounded + // paged pools may be physically smaller than this logical span. const int launch_kv_len = paged_max_kv_len > 0 ? paged_max_kv_len : kv_len; attn = paged_read( Qfa, launch_kv_len, active_slot_ids, /*row_positions=*/nullptr, @@ -1015,6 +1092,7 @@ static ggml_tensor * build_delta_net_block( int n_prefill_segments = 0, ggml_tensor * active_slot_ids = nullptr, ggml_tensor * state_slot_ids = nullptr, + int mapped_ar_seqs = 0, bool allow_inplace_state = false ) { const int head_k_dim = w.ssm_d_state; @@ -1032,32 +1110,75 @@ static ggml_tensor * build_delta_net_block( prefill_total += prefill_segments[i].n_tokens; } GGML_ASSERT((active_slot_ids == nullptr) == (state_slot_ids == nullptr)); + const bool mapped_tree = active_slot_ids && parent_ids; + GGML_ASSERT(mapped_ar_seqs >= 0); + GGML_ASSERT(mapped_ar_seqs == 0 || mapped_tree); + GGML_ASSERT(!active_slot_ids || !cap || + (!cap->ssm_intermediate_states && !cap->conv_input)); GGML_ASSERT(!active_slot_ids || - (!cap && !parent_ids && prefill_total + n_seqs == n_tokens)); + (mapped_tree + ? (!ragged && prefill_total == 0 && + n_tokens >= mapped_ar_seqs && + (n_tokens - mapped_ar_seqs) % n_seqs == 0 && + active_slot_ids->ne[0] == + mapped_ar_seqs + n_seqs && + state_slot_ids->ne[0] == + mapped_ar_seqs + n_seqs) + : (mapped_ar_seqs == 0 && + prefill_total + n_seqs == n_tokens))); if (!active_slot_ids) { GGML_ASSERT(n_seqs == 1); GGML_ASSERT(prefill_total == 0 || prefill_total == n_tokens); } GGML_ASSERT(!ragged || (!cap && !parent_ids)); - const bool can_skip_gdn_intermediate = skip_gdn_intermediate && !parent_ids && !cap; - - // ── Whole-batch projections ───────────────────────────────────── - // qkv_mixed = wqkv @ cur [10240, n_tokens] - ggml_tensor * qkv_2d = apply_scale2(ctx, ggml_mul_mat(ctx, L.wqkv, cur), L.wqkv_s); - // z = wqkv_gate @ cur [inner, n_tokens] - ggml_tensor * z = apply_scale2(ctx, ggml_mul_mat(ctx, L.wqkv_gate, cur), L.wqkv_gate_s); + // Row slices of stacked projections are strided for multi-token inputs. + // Materialize only the small beta/alpha slices; qkv keeps its explicit + // column stride and z is made contiguous at the final per-segment gate. + auto contig = [&](ggml_tensor * t) { + return ggml_is_contiguous(t) ? t : ggml_cont(ctx, t); + }; - // beta = sigmoid(ssm_beta @ cur) [dt_rank, n_tokens] - ggml_tensor * beta_2d = apply_scale2(ctx, ggml_mul_mat(ctx, L.ssm_beta, cur), L.ssm_beta_s); - beta_2d = ggml_sigmoid(ctx, beta_2d); + // ── Whole-batch projections ───────────────────────────────────── + // One GEMM over the zero-copy stacked (z | qkv) alias when possible. + ggml_tensor * qkv_2d = nullptr; + ggml_tensor * z = nullptr; + const bool stacked_qkv_z = + L.wqkv_z && L.wqkv_s == 1.0f && L.wqkv_gate_s == 1.0f; + if (stacked_qkv_z) { + const int64_t n_z = L.wqkv_gate->ne[1]; + ggml_tensor * qkvz = ggml_mul_mat(ctx, L.wqkv_z, cur); + const size_t e = ggml_element_size(qkvz); + z = ggml_view_2d(ctx, qkvz, n_z, n_tokens, qkvz->nb[1], 0); + qkv_2d = ggml_view_2d(ctx, qkvz, conv_channels, n_tokens, + qkvz->nb[1], (size_t)n_z * e); + } else { + qkv_2d = apply_scale2(ctx, ggml_mul_mat(ctx, L.wqkv, cur), L.wqkv_s); + z = apply_scale2(ctx, ggml_mul_mat(ctx, L.wqkv_gate, cur), L.wqkv_gate_s); + } + + // One GEMM over the zero-copy stacked (beta | alpha) alias when possible. + ggml_tensor * beta_2d = nullptr; + ggml_tensor * alpha_2d = nullptr; + const bool stacked_ba = + L.ssm_ba && L.ssm_beta_s == 1.0f && L.ssm_alpha_s == 1.0f; + if (stacked_ba) { + ggml_tensor * ba = ggml_mul_mat(ctx, L.ssm_ba, cur); + const size_t e = ggml_element_size(ba); + beta_2d = contig(ggml_view_2d( + ctx, ba, num_v_heads, n_tokens, ba->nb[1], 0)); + alpha_2d = contig(ggml_view_2d( + ctx, ba, num_v_heads, n_tokens, ba->nb[1], + (size_t)num_v_heads * e)); + } else { + beta_2d = apply_scale2( + ctx, ggml_mul_mat(ctx, L.ssm_beta, cur), L.ssm_beta_s); + alpha_2d = apply_scale2( + ctx, ggml_mul_mat(ctx, L.ssm_alpha, cur), L.ssm_alpha_s); + } - // alpha = ssm_alpha @ cur [dt_rank, n_tokens] - // g = softplus(alpha + ssm_dt_bias) * ssm_a (-A_log.exp() * softplus) - ggml_tensor * alpha = apply_scale2(ctx, ggml_mul_mat(ctx, L.ssm_alpha, cur), L.ssm_alpha_s); - alpha = ggml_add(ctx, alpha, L.ssm_dt_bias); - alpha = ggml_softplus(ctx, alpha); - ggml_tensor * g_2d = ggml_mul(ctx, alpha, L.ssm_a); + static const bool fused_kernels_env = + std::getenv("DFLASH_QWEN35_NO_FUSED_KERNELS") == nullptr; // ── Token-axis segments: prompt chunks first, then the decode batch ── struct DeltaSeg { @@ -1065,8 +1186,11 @@ static ggml_tensor * build_delta_net_block( int T; // timesteps per sequence int S; // sequences bool active; // compact decode segment (slot-mapped) + bool tree; // mapped tree: gather-only, no persistence ggml_tensor * conv_st; ggml_tensor * ssm_st; + ggml_tensor * active_ids; + ggml_tensor * state_ids; }; std::vector segs; segs.reserve((size_t)n_prefill_segments + 1); @@ -1082,14 +1206,35 @@ static ggml_tensor * build_delta_net_block( ssm_state->ne[0], ssm_state->ne[1], ssm_state->ne[2], 1, ssm_state->nb[1], ssm_state->nb[2], ssm_state->nb[3], (size_t)pf.seq_slot * ssm_state->nb[3]); - segs.push_back({pf.token_offset, pf.n_tokens, 1, false, c, s}); + segs.push_back({pf.token_offset, pf.n_tokens, 1, + false, false, c, s, nullptr, nullptr}); } if (active_slot_ids) { - segs.push_back({prefill_total, 1, n_seqs, true, - conv_state, ssm_state}); + if (mapped_tree && mapped_ar_seqs > 0) { + ggml_tensor * ar_active = ggml_view_1d( + ctx, active_slot_ids, mapped_ar_seqs, 0); + ggml_tensor * ar_state = ggml_view_1d( + ctx, state_slot_ids, mapped_ar_seqs, 0); + segs.push_back({0, 1, mapped_ar_seqs, true, false, + conv_state, ssm_state, ar_active, ar_state}); + } + const int tree_tokens = mapped_tree + ? (n_tokens - mapped_ar_seqs) / n_seqs : 1; + const size_t slot_offset = + (size_t)mapped_ar_seqs * active_slot_ids->nb[0]; + ggml_tensor * segment_active = mapped_ar_seqs > 0 + ? ggml_view_1d(ctx, active_slot_ids, n_seqs, slot_offset) + : active_slot_ids; + ggml_tensor * segment_state = mapped_ar_seqs > 0 + ? ggml_view_1d(ctx, state_slot_ids, n_seqs, slot_offset) + : state_slot_ids; + segs.push_back({prefill_total + mapped_ar_seqs, tree_tokens, + n_seqs, true, mapped_tree, conv_state, ssm_state, + segment_active, segment_state}); } else if (segs.empty()) { // No general [timesteps x sequences] mode: one multi-token sequence. - segs.push_back({0, n_tokens, n_seqs, false, conv_state, ssm_state}); + segs.push_back({0, n_tokens, n_seqs, false, false, + conv_state, ssm_state, nullptr, nullptr}); } const int n_segs = (int)segs.size(); @@ -1109,25 +1254,60 @@ static ggml_tensor * build_delta_net_block( const int seg_seqs = seg.S; const int seg_tokens = seg.T * seg.S; const bool seg_active = seg.active; + const bool seg_tree = seg.tree; + DeltaNetCapture * seg_cap = mapped_tree + ? (seg_tree ? cap : nullptr) : cap; + ggml_tensor * seg_parent_ids = seg_tree ? parent_ids : nullptr; + const bool can_skip_gdn_intermediate = + skip_gdn_intermediate && !seg_parent_ids && !seg_cap; // Plain one-token decode has no in-graph consumer of the updated state: // the next graph evaluation is the first read. Write the final state // directly into its persistent slab and avoid materializing/copying a // second S_v x S_v x H_v state. The active-aware path also updates each // mapped physical slab directly; only its negative bucket-padding rows // use the result tensor's retained scratch state region. - const bool inplace_state = seg_active || + const bool dense_chain = !ragged && !active_slot_ids && !seg_tree; + const bool inplace_state = dense_chain || (seg_active && !seg_tree) || (allow_inplace_state && can_skip_gdn_intermediate && !ragged && n_seq_tokens == 1); - ggml_tensor * qkv_mixed = ggml_reshape_3d(ctx, - seg_cols(qkv_2d, seg.off, seg_tokens), - conv_channels, n_seq_tokens, seg_seqs); + ggml_tensor * qkv_seg = seg_cols(qkv_2d, seg.off, seg_tokens); + ggml_tensor * qkv_mixed = stacked_qkv_z + ? ggml_view_3d(ctx, qkv_seg, + conv_channels, n_seq_tokens, seg_seqs, + qkv_seg->nb[1], qkv_seg->nb[1] * n_seq_tokens, 0) + : ggml_reshape_3d(ctx, qkv_seg, + conv_channels, n_seq_tokens, seg_seqs); + + // Chunked delta-net path is opt-in and chain-only. + bool use_chunked = false; + if (can_skip_gdn_intermediate && n_seq_tokens > 1) { + if (const char * s_env = std::getenv("DFLASH27B_CHUNKED")) { + use_chunked = (std::atoi(s_env) != 0); + } + } + const bool fused_conv = fused_kernels_env && !seg_active && !seg_tree; + const bool raw_gates = fused_kernels_env && !seg_tree && !use_chunked; + ggml_tensor * beta = ggml_reshape_4d(ctx, seg_cols(beta_2d, seg.off, seg_tokens), 1, num_v_heads, n_seq_tokens, seg_seqs); - ggml_tensor * g_tensor = ggml_reshape_4d(ctx, - seg_cols(g_2d, seg.off, seg_tokens), - 1, num_v_heads, n_seq_tokens, seg_seqs); + ggml_tensor * alpha = ggml_reshape_3d(ctx, + seg_cols(alpha_2d, seg.off, seg_tokens), + num_v_heads, n_seq_tokens, seg_seqs); + ggml_tensor * g_tensor = nullptr; + if (raw_gates) { + // The kernel applies sigmoid(beta) and softplus(alpha + dt_bias) * A. + g_tensor = ggml_reshape_4d( + ctx, alpha, 1, num_v_heads, n_seq_tokens, seg_seqs); + } else { + beta = ggml_sigmoid(ctx, beta); + alpha = ggml_add(ctx, alpha, L.ssm_dt_bias); + alpha = ggml_softplus(ctx, alpha); + g_tensor = ggml_mul(ctx, alpha, L.ssm_a); + g_tensor = ggml_reshape_4d( + ctx, g_tensor, 1, num_v_heads, n_seq_tokens, seg_seqs); + } // ── Fetch conv state [kernel-1, conv_channels] and prepend to qkv_mixed // along the token axis to form the convolution input. @@ -1138,7 +1318,7 @@ static ggml_tensor * build_delta_net_block( ggml_tensor * all_conv = ggml_reshape_2d( ctx, seg.conv_st, slab, seg.conv_st->ne[2]); ggml_tensor * gathered = - ggml_get_rows(ctx, all_conv, state_slot_ids); + ggml_get_rows(ctx, all_conv, seg.state_ids); conv_states_r = ggml_reshape_3d( ctx, gathered, w.ssm_d_conv - 1, conv_channels, seg_seqs); } else { @@ -1146,66 +1326,88 @@ static ggml_tensor * build_delta_net_block( w.ssm_d_conv - 1, conv_channels, seg_seqs); } - // qkv_mixed currently is [conv_channels, n_tokens, n_seqs]; we need - // [n_tokens, conv_channels, n_seqs] to concat on dim 0. - ggml_tensor * qkv_T = ggml_transpose(ctx, qkv_mixed); - - ggml_tensor * conv_input = ggml_concat(ctx, conv_states_r, qkv_T, 0); - // I0 domain: [0,K_conv-2] are prefix-history rows; tree token flat slot t - // (root-inclusive, including synthetic root t=0) is stored at - // conv_input row (K_conv-1)+t. - // conv_input: [kernel-1 + n_tokens, conv_channels, n_seqs] - - // For spec-decode rollback: copy the full conv_input into the persistent - // cache buffer via an in-graph ggml_cpy. This avoids marking conv_input as - // a graph output (which would force the gallocr to preserve its memory - // past graph_compute). After graph_compute, the cache buffer's data is - // always valid; the rollback code slices it at commit_n. - if (cap && cap->conv_input) { - // conv_input may be shorter than the pre-allocated cache - // (e.g. during prefill when n_tokens < max_verify_tokens). - // Copy into a matching-sized view of the cache destination. - const int64_t ci_len = conv_input->ne[0]; - ggml_tensor * dst; - if (ci_len == cap->conv_input->ne[0]) { - dst = cap->conv_input; - } else { - dst = ggml_view_3d(ctx, cap->conv_input, - ci_len, cap->conv_input->ne[1], cap->conv_input->ne[2], - cap->conv_input->nb[1], cap->conv_input->nb[2], 0); + ggml_tensor * conv_out = nullptr; + if (fused_conv) { + // One kernel: window = [conv_state | x], silu(conv), history + // write-back, and (when capturing) the rollback window copy. + ggml_tensor * ci_dst = nullptr; + if (seg_cap && seg_cap->conv_input) { + const int64_t ci_len = (w.ssm_d_conv - 1) + n_tokens; + ci_dst = (ci_len == seg_cap->conv_input->ne[0]) + ? seg_cap->conv_input + : ggml_view_3d(ctx, seg_cap->conv_input, + ci_len, seg_cap->conv_input->ne[1], seg_cap->conv_input->ne[2], + seg_cap->conv_input->nb[1], seg_cap->conv_input->nb[2], 0); } - GGML_ASSERT(ggml_nelements(conv_input) == ggml_nelements(dst)); - ggml_build_forward_expand(gf, ggml_cpy(ctx, conv_input, dst)); - } - - // ── Save the last (kernel-1) steps back to the conv state - ggml_tensor * last_conv = ggml_view_3d(ctx, conv_input, - w.ssm_d_conv - 1, conv_channels, seg_seqs, - conv_input->nb[1], conv_input->nb[2], - (conv_input->ne[0] - (w.ssm_d_conv - 1)) * ggml_element_size(conv_input)); - if (seg_active) { - const int64_t slab = - (int64_t)(w.ssm_d_conv - 1) * conv_channels; - ggml_tensor * compact_last = ggml_reshape_2d( - ctx, ggml_cont(ctx, last_conv), slab, seg_seqs); - ggml_tensor * all_conv = ggml_reshape_2d( - ctx, seg.conv_st, slab, seg.conv_st->ne[2]); - ggml_build_forward_expand( - gf, ggml_set_rows_masked( - ctx, all_conv, compact_last, active_slot_ids)); + conv_out = ggml_ssm_conv_step(ctx, qkv_mixed, L.ssm_conv1d, conv_states_r, ci_dst); } else { - ggml_build_forward_expand(gf, ggml_cpy(ctx, last_conv, seg.conv_st)); - } + // qkv_mixed currently is [conv_channels, n_tokens, n_seqs]; we need + // [n_tokens, conv_channels, n_seqs] to concat on dim 0. + ggml_tensor * qkv_T = ggml_transpose(ctx, qkv_mixed); + + ggml_tensor * conv_input = ggml_concat(ctx, conv_states_r, qkv_T, 0); + // I0 domain: [0,K_conv-2] are prefix-history rows; tree token flat slot t + // (root-inclusive, including synthetic root t=0) is stored at + // conv_input row (K_conv-1)+t. + // conv_input: [kernel-1 + n_tokens, conv_channels, n_seqs] + + // For spec-decode rollback: copy the full conv_input into the persistent + // cache buffer via an in-graph ggml_cpy. This avoids marking conv_input as + // a graph output (which would force the gallocr to preserve its memory + // past graph_compute). After graph_compute, the cache buffer's data is + // always valid; the rollback code slices it at commit_n. + if (seg_cap && seg_cap->conv_input) { + // conv_input may be shorter than the pre-allocated cache + // (e.g. during prefill when n_tokens < max_verify_tokens). + // Copy into a matching-sized view of the cache destination. + const int64_t ci_len = conv_input->ne[0]; + ggml_tensor * dst; + if (ci_len == seg_cap->conv_input->ne[0]) { + dst = seg_cap->conv_input; + } else { + dst = ggml_view_3d(ctx, seg_cap->conv_input, + ci_len, seg_cap->conv_input->ne[1], seg_cap->conv_input->ne[2], + seg_cap->conv_input->nb[1], seg_cap->conv_input->nb[2], 0); + } + GGML_ASSERT(ggml_nelements(conv_input) == ggml_nelements(dst)); + ggml_build_forward_expand(gf, ggml_cpy(ctx, conv_input, dst)); + } - // ── 1D conv + silu - // Tree mode: use the parent-chain-aware variant so sibling nodes gather - // their conv window from their actual tree parent instead of the DFS - // predecessor. Without this, siblings get garbage logits (the conv - // output would mix unrelated branches). - ggml_tensor * conv_out = parent_ids - ? ggml_ssm_conv_tree(ctx, conv_input, L.ssm_conv1d, parent_ids) - : ggml_ssm_conv (ctx, conv_input, L.ssm_conv1d); - conv_out = ggml_silu(ctx, conv_out); + if (seg_cap && seg_tree && !seg_cap->conv_input) { + seg_cap->conv_input = conv_input; + ggml_set_output(seg_cap->conv_input); + } + + // ── Save the last (kernel-1) steps back to conv_state + ggml_tensor * last_conv = ggml_view_3d(ctx, conv_input, + w.ssm_d_conv - 1, conv_channels, seg_seqs, + conv_input->nb[1], conv_input->nb[2], + (conv_input->ne[0] - (w.ssm_d_conv - 1)) * ggml_element_size(conv_input)); + if (seg_active && !seg_tree) { + const int64_t slab = + (int64_t)(w.ssm_d_conv - 1) * conv_channels; + ggml_tensor * compact_last = ggml_reshape_2d( + ctx, ggml_cont(ctx, last_conv), slab, seg_seqs); + ggml_tensor * all_conv = ggml_reshape_2d( + ctx, seg.conv_st, slab, seg.conv_st->ne[2]); + ggml_build_forward_expand( + gf, ggml_set_rows_masked( + ctx, all_conv, compact_last, seg.active_ids)); + } else if (!seg_tree) { + ggml_build_forward_expand( + gf, ggml_cpy(ctx, last_conv, seg.conv_st)); + } + + // ── 1D conv + silu + // Tree mode: use the parent-chain-aware variant so sibling nodes gather + // their conv window from their actual tree parent instead of the DFS + // predecessor. Without this, siblings get garbage logits (the conv + // output would mix unrelated branches). + conv_out = seg_parent_ids + ? ggml_ssm_conv_tree(ctx, conv_input, L.ssm_conv1d, seg_parent_ids) + : ggml_ssm_conv (ctx, conv_input, L.ssm_conv1d); + conv_out = ggml_silu(ctx, conv_out); + } // conv_out: [conv_channels, n_tokens, n_seqs] const int64_t q_offset = 0; @@ -1234,27 +1436,58 @@ static ggml_tensor * build_delta_net_block( row_size * n_seq_tokens, v_offset * elt); - // L2 norm on Q and K - q_c = ggml_l2_norm(ctx, q_c, w.rms_eps); - k_c = ggml_l2_norm(ctx, k_c, w.rms_eps); - - // Repeat Q and K from num_k_heads to num_v_heads so they match V's layout - // (only needed if not using the fused op's broadcast support). - if (num_k_heads != num_v_heads) { + // L2 norm on Q and K: q and k heads are adjacent in conv_out, so one + // launch over the [head_k_dim, 2*num_k_heads] slab normalizes both. + { + ggml_tensor * qk_c = ggml_view_4d(ctx, conv_out, + head_k_dim, 2 * num_k_heads, n_seq_tokens, seg_seqs, + head_k_dim * elt, + row_size, + row_size * n_seq_tokens, + q_offset * elt); + ggml_tensor * qk_n = ggml_l2_norm(ctx, qk_c, w.rms_eps); // contiguous [hd, 2*Hk, T, S] + const size_t ne_ = ggml_element_size(qk_n); + q_c = ggml_view_4d(ctx, qk_n, head_k_dim, num_k_heads, n_seq_tokens, seg_seqs, + qk_n->nb[1], qk_n->nb[2], qk_n->nb[3], 0); + k_c = ggml_view_4d(ctx, qk_n, head_k_dim, num_k_heads, n_seq_tokens, seg_seqs, + qk_n->nb[1], qk_n->nb[2], qk_n->nb[3], + (size_t)num_k_heads * head_k_dim * ne_); + } + + // Repeat Q and K from num_k_heads to num_v_heads so they match V's layout. + // The fused gated_delta_net kernels broadcast heads themselves (v head h + // reads q/k head h % num_k_heads, the same tiling ggml_repeat produces), + // so only the chunked path needs the materialized copies. + if (num_k_heads != num_v_heads && use_chunked) { q_c = ggml_repeat_4d(ctx, q_c, head_k_dim, num_v_heads, n_seq_tokens, seg_seqs); k_c = ggml_repeat_4d(ctx, k_c, head_k_dim, num_v_heads, n_seq_tokens, seg_seqs); } // ── SSM state (recurrent): reshape to [S_v, S_v, H_v, n_seqs] - ggml_tensor * s = seg_active - ? seg.ssm_st - : ggml_reshape_4d(ctx, seg.ssm_st, + ggml_tensor * s = nullptr; + if (seg_tree) { + // Packed tree verification starts each tree from the owning slot's + // base state. Gather compact slabs, then leave the persistent tensor + // untouched; accepted paths are committed by later direct promotion. + const int64_t slab = + (int64_t)head_v_dim * head_v_dim * num_v_heads; + ggml_tensor * all_ssm = ggml_reshape_2d( + ctx, seg.ssm_st, slab, seg.ssm_st->ne[3]); + ggml_tensor * gathered = + ggml_get_rows(ctx, all_ssm, seg.state_ids); + s = ggml_reshape_4d(ctx, gathered, head_v_dim, head_v_dim, num_v_heads, seg_seqs); + } else { + s = seg_active + ? seg.ssm_st + : ggml_reshape_4d(ctx, seg.ssm_st, + head_v_dim, head_v_dim, num_v_heads, seg_seqs); + } // ── Fused Gated DeltaNet op — returns packed (output | new_state [| intermediates]). // In tree mode, the kernel uses parent_ids to reload state at DFS // branch transitions (ported from sglang's retrieve_parent_token path). - // When `cap->ssm_intermediate_states` is present AND we are in tree + // When `seg_cap->ssm_intermediate_states` is present AND we are in tree // mode, use the _tree_persist variant: the kernel writes per-token // intermediate states DIRECTLY into the persistent cache buffer, // eliminating the downstream ggml_cpy that would otherwise copy them. @@ -1270,10 +1503,10 @@ static ggml_tensor * build_delta_net_block( // path is never quantized. In tree mode, n_seq_tokens is root-inclusive and // flat slot t is persisted directly at ne[3] slot t. // Q8_0 intermediates fall through to the guarded legacy copy path below. - ggml_tensor * persist_inter = (cap && cap->ssm_intermediate_states - && (cap->ssm_intermediate_states->type == GGML_TYPE_F32 - || cap->ssm_intermediate_states->type == GGML_TYPE_F16)) - ? cap->ssm_intermediate_states + ggml_tensor * persist_inter = (seg_cap && seg_cap->ssm_intermediate_states + && (seg_cap->ssm_intermediate_states->type == GGML_TYPE_F32 + || seg_cap->ssm_intermediate_states->type == GGML_TYPE_F16)) + ? seg_cap->ssm_intermediate_states : nullptr; // Chunked delta-net path: chain-only (no parent_ids), no per-token @@ -1284,13 +1517,6 @@ static ggml_tensor * build_delta_net_block( // default — port produces correct shape but slightly wrong final state, // causing AL degradation and loopy output. Set DFLASH27B_CHUNKED=1 to // opt in for A/B testing while debugging. - bool use_chunked = false; - if (can_skip_gdn_intermediate && n_seq_tokens > 1) { - if (const char * s_env = std::getenv("DFLASH27B_CHUNKED")) { - use_chunked = (std::atoi(s_env) != 0); - } - } - ggml_tensor * output = nullptr; if (use_chunked) { @@ -1302,14 +1528,14 @@ static ggml_tensor * build_delta_net_block( ggml_build_forward_expand(gf, ggml_cpy(ctx, r.new_state, s)); } else { ggml_tensor * result; - if (seg_active) { + if (seg_active && !seg_tree) { result = ggml_gated_delta_net_active_inplace( - ctx, q_c, k_c, v_c, g_tensor, beta, s, active_slot_ids); - } else if (parent_ids) { + ctx, q_c, k_c, v_c, g_tensor, beta, s, seg.active_ids); + } else if (seg_parent_ids) { // Tree verify: _tree_persist wires src[7] internally. result = persist_inter - ? ggml_gated_delta_net_tree_persist(ctx, q_c, k_c, v_c, g_tensor, beta, s, parent_ids, persist_inter) - : ggml_gated_delta_net_tree(ctx, q_c, k_c, v_c, g_tensor, beta, s, parent_ids); + ? ggml_gated_delta_net_tree_persist(ctx, q_c, k_c, v_c, g_tensor, beta, s, seg_parent_ids, persist_inter) + : ggml_gated_delta_net_tree(ctx, q_c, k_c, v_c, g_tensor, beta, s, seg_parent_ids); } else { // Non-tree (chain/prefill). When capture is requested, set src[7] so // the kernel writes per-token intermediates directly to the persistent @@ -1323,6 +1549,19 @@ static ggml_tensor * build_delta_net_block( result->src[7] = persist_inter; } } + if (seg_cap && seg_tree) { + const int64_t journal_width = + g_tensor->ne[0] == head_v_dim ? 3*head_v_dim : 2*head_v_dim + 1; + seg_cap->transition_journal = ggml_new_tensor_4d( + ctx, GGML_TYPE_F32, journal_width, num_v_heads, + n_seq_tokens, seg_seqs); + ggml_set_output(seg_cap->transition_journal); + ggml_gated_delta_net_set_transition_journal( + result, seg_cap->transition_journal); + } + if (raw_gates) { + ggml_gated_delta_net_set_raw_gates(result, L.ssm_dt_bias, L.ssm_a); + } if (can_skip_gdn_intermediate) { ggml_gated_delta_net_set_skip_intermediate(result, true); } @@ -1337,7 +1576,7 @@ static ggml_tensor * build_delta_net_block( S_v * H_v * r_elt, S_v * H_v * n_seq_tokens * r_elt, 0); - if (!inplace_state) { + if (!inplace_state && !seg_tree) { ggml_tensor * new_state = ggml_view_4d(ctx, result, S_v, S_v, H_v, seg_seqs, S_v * r_elt, @@ -1345,8 +1584,8 @@ static ggml_tensor * build_delta_net_block( S_v * S_v * H_v * r_elt, S_v * H_v * n_seq_tokens * seg_seqs * r_elt); - // Persist new_state back to cache. Both compact active decode and the - // plain in-place AR path write state from the GDN kernel directly. + // Persist new_state back to cache. Mapped trees deliberately skip + // this branch: their gathered base state is read-only. ggml_build_forward_expand(gf, ggml_cpy(ctx, new_state, seg.ssm_st)); } @@ -1361,10 +1600,10 @@ static ggml_tensor * build_delta_net_block( // forces gallocr to preserve ~50 MB per layer × 48 layers of otherwise // transient memory and inflates graph_build by ~35 ms), we create a VIEW // into the intermediate region and ggml_cpy it into the persistent cache - // buffer cap->ssm_intermediate_states. The gallocr is unaware of the + // buffer seg_cap->ssm_intermediate_states. The gallocr is unaware of the // persistent cache, so verify_build stays cheap. Matches SGLang's // mamba_caches.intermediate_ssm pattern. - if (cap && cap->ssm_intermediate_states && !persist_inter) { + if (seg_cap && seg_cap->ssm_intermediate_states && !persist_inter) { // This path is only reachable when the intermediate buffer is a type // persist routing can't handle (persist requires F32/F16; the cache // allocates F16, so this is normally dead). If the result tensor has no @@ -1373,13 +1612,13 @@ static ggml_tensor * build_delta_net_block( GGML_ABORT( "non-tree GDN intermediate capture requires an F32/F16 persist buffer " "(got type %d); use F16 intermediates (the default) or the tree-verify path.", - (int)cap->ssm_intermediate_states->type); + (int)seg_cap->ssm_intermediate_states->type); } } // ── Gated output norm: rms_norm(output) * silu(z_4d) ggml_tensor * z_4d = ggml_reshape_4d(ctx, - seg_cols(z, seg.off, seg_tokens), + contig(seg_cols(z, seg.off, seg_tokens)), head_v_dim, num_v_heads, n_seq_tokens, seg_seqs); ggml_tensor * output_n = ggml_rms_norm(ctx, rms_norm_input_f32(ctx, output), w.rms_eps); output_n = ggml_mul(ctx, output_n, L.ssm_norm); @@ -1540,7 +1779,7 @@ QwenGraphOutputs build_qwen35_graph( // If the caller requested capture, size the output list to the total delta- // net layer count so we can index by dn_idx as we iterate the layers. QwenGraphOutputs og_early{}; - if (in.capture_delta_intermediate) { + if (in.capture_delta_intermediate || in.capture_tree_commit) { const int n_full_attn = w.n_layer / w.full_attention_interval; const int n_delta = w.n_layer - n_full_attn; og_early.delta_captures.resize(n_delta); @@ -1555,6 +1794,14 @@ QwenGraphOutputs build_qwen35_graph( const int hidden = w.n_embd; const float eps = w.rms_eps; + const bool capture_with_rows = + in.capture_layers && cache.target_feat && in.target_feat_rows; + const bool capture_tree_features = + in.capture_layers && in.capture_tree_commit && cache.target_feat; + std::vector capture_slices; + if (capture_with_rows || capture_tree_features) { + capture_slices.assign((size_t)N_CAPTURE, nullptr); + } for (int il = 0; il < w.n_layer; il++) { const TargetLayer & L = w.layers[il]; @@ -1584,7 +1831,13 @@ QwenGraphOutputs build_qwen35_graph( in.paged_query_seq_ids, in.paged_query_positions, in.paged_max_kv_len, - in.active_slot_ids); + in.active_slot_ids, + in.parent_ids, + in.tree_sizes, + in.tree_width, + in.tree_scratch_base, + in.tree_scratch_stride, + cache.max_ctx); if (want_q_cap && q_fa) { // Last token's Q, all heads: src [head_dim, 1, n_head] view of // [head_dim, n_tokens, n_head]; dst = q_cap plane fa_idx @@ -1603,15 +1856,17 @@ QwenGraphOutputs build_qwen35_graph( fa_idx++; } else { DeltaNetCapture * cap_ptr = nullptr; - if (in.capture_delta_intermediate) { + if (in.capture_delta_intermediate || in.capture_tree_commit) { cap_ptr = &og_early.delta_captures[dn_idx]; // Point at the persistent per-layer cache buffers so // build_delta_net_block can ggml_cpy into them during graph // execution. The caller (test_dflash.cpp spec loop) reads from // these tensors post-compute; their ->data pointers are always // valid because they're cache-resident, not gallocr-managed. + if (in.capture_delta_intermediate) { cap_ptr->ssm_intermediate_states = cache.ssm_intermediate[dn_idx]; cap_ptr->conv_input = cache.conv_input_cache[dn_idx]; + } } ggml_tensor * conv_st = cache.conv_state[dn_idx]; ggml_tensor * ssm_st = cache.ssm_state[dn_idx]; @@ -1642,6 +1897,7 @@ QwenGraphOutputs build_qwen35_graph( in.n_prefill_segments, in.active_slot_ids, in.state_slot_ids, + in.mapped_ar_seqs, /*allow_inplace_state=*/ in.n_prefill_tokens == 0); dn_idx++; @@ -1676,6 +1932,13 @@ QwenGraphOutputs build_qwen35_graph( if (CAPTURE_LAYERS[k] == il) { capture_idx = k; break; } } if (capture_idx >= 0) { + ggml_tensor * cur_2d = + ggml_reshape_2d(ctx, cur, hidden, n_tokens); + if (capture_with_rows || capture_tree_features) { + capture_slices[(size_t)capture_idx] = cur_2d; + inpL = cur; + continue; + } const size_t elt = ggml_element_size(cache.target_feat); const size_t col_stride = cache.target_feat->nb[1]; const int cap = cache.target_feat_cap; @@ -1683,8 +1946,6 @@ QwenGraphOutputs build_qwen35_graph( const int pre_n = std::min(n_tokens, cap - slot_start); const int post_n = n_tokens - pre_n; - ggml_tensor * cur_2d = ggml_reshape_2d(ctx, cur, hidden, n_tokens); - // First slice: [slot_start..slot_start+pre_n) in the ring. { const size_t offset = @@ -1714,6 +1975,27 @@ QwenGraphOutputs build_qwen35_graph( inpL = cur; } + if (capture_with_rows || capture_tree_features) { + GGML_ASSERT(!capture_slices.empty()); + ggml_tensor * feat_cat = capture_slices[0]; + GGML_ASSERT(feat_cat); + for (int k = 1; k < (int)capture_slices.size(); ++k) { + GGML_ASSERT(capture_slices[(size_t)k]); + feat_cat = ggml_concat( + ctx, feat_cat, capture_slices[(size_t)k], 0); + } + feat_cat = ggml_cont(ctx, feat_cat); + if (capture_tree_features) { + og_early.tree_features = ggml_cast(ctx, feat_cat, GGML_TYPE_BF16); + ggml_set_output(og_early.tree_features); + ggml_build_forward_expand(gf, og_early.tree_features); + } else { + ggml_build_forward_expand( + gf, ggml_set_rows( + ctx, cache.target_feat, feat_cat, in.target_feat_rows)); + } + } + // 2. Final norm ggml_tensor * out = rms_norm_mul(ctx, inpL, w.out_norm, w.rms_eps); diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index e92f495bf..a66370003 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -764,6 +764,7 @@ json build_props_body(const ServerConfig & config, {"build_info", std::string(kServerName) + " v" DFLASH_SERVER_VERSION " props_schema=" + std::to_string(kPropsSchema)}, {"speculative_mode", speculative_mode}, + {"decode_mode", speculation_policy_name(config.decode_mode)}, {"server", server}, {"model", { {"arch", config.arch}, @@ -777,6 +778,7 @@ json build_props_body(const ServerConfig & config, {"kv_cache_v", config.kv_cache_v}, {"lazy_draft", config.lazy_draft}, {"draft_residency", draft_residency_policy_name(config.draft_residency)}, + {"decode_mode", speculation_policy_name(config.decode_mode)}, {"target_sharding", config.target_sharding}, // Prefill chunk size (bargs.chunk). Surfaced so snapshot // tooling captures the full config — bench consumers @@ -1379,11 +1381,27 @@ int HttpServer::run() { std::fprintf(stderr, "[server] listening on http://%s:%d\n", config_.host.c_str(), config_.port); - // A backend-provided sequence engine replaces the one-request worker - // with the concurrent scheduler. Upstream forwarding stays on the - // classic path even when the local backend exposes an engine. - if (SeqEngine * engine = backend_.seq_engine(); - engine && config_.pflash_upstream_base.empty()) { + // A backend-provided sequence engine replaces the one-request worker. + // Local PFlash stays on this path too: scheduler admission prepares each + // prompt exactly once, then admits the effective tokens. Persistent + // residency is the explicit safety contract that lets compression run + // without parking model state owned by other live slots. + SeqEngine * engine = backend_.seq_engine(); + if (engine && config_.pflash_upstream_base.empty()) { + const ConcurrentPflashPlan pflash_plan = + resolve_concurrent_pflash_plan(config_, drafter_tokenizer_ != nullptr); + if (!pflash_plan.ok()) { + std::fprintf(stderr, "[server] %s\n", pflash_plan.error.c_str()); + socket_close(listen_fd_); + listen_fd_ = kInvalidSocket; + return 2; + } + if (pflash_plan.force_skip_park && !config_.pflash_skip_park) { + config_.pflash_skip_park = true; + std::fprintf(stderr, + "[server] concurrent PFlash: persistent residency enables " + "skip-park for live sequence safety\n"); + } worker_thread_ = std::thread([this, engine]() { scheduler_loop(*engine); }); } else { @@ -1645,6 +1663,21 @@ bool HttpServer::parse_common_request_fields( req.stream = body.value("stream", false); req.model = body.value("model", config_.model_name); req.disk_cache_policy = config_.disk_cache_policy; + if (body.contains("decode_mode")) { + if (!body["decode_mode"].is_string()) { + send_error(fd, 400, + "decode_mode must be ar, speculation, or adaptive"); + return false; + } + SpeculationPolicy policy; + if (!parse_speculation_policy( + body["decode_mode"].get(), policy)) { + send_error(fd, 400, + "decode_mode must be ar, speculation, or adaptive"); + return false; + } + req.decode_mode = policy; + } // Accept the output-token names used by each supported API dialect. // Default when the client omits all three: --default-max-tokens, so @@ -4231,6 +4264,7 @@ std::string HttpServer::format_http_response( case 400: reason = "Bad Request"; break; case 404: reason = "Not Found"; break; case 405: reason = "Method Not Allowed"; break; + case 409: reason = "Conflict"; break; case 413: reason = "Payload Too Large"; break; case 500: reason = "Internal Server Error"; break; case 503: reason = "Service Unavailable"; break; diff --git a/server/src/server/http_server.h b/server/src/server/http_server.h index 90e02412e..3644dea01 100644 --- a/server/src/server/http_server.h +++ b/server/src/server/http_server.h @@ -19,6 +19,7 @@ #include "tokenizer.h" #include "chat_template.h" #include "tool_memory.h" +#include "common/speculation_policy.h" #include "prefix_cache.h" #include "disk_prefix_cache.h" #include "freeze_history.h" @@ -33,6 +34,7 @@ #include #include +#include #include #include #include @@ -212,6 +214,8 @@ struct ServerConfig { bool lazy_draft = false; // legacy alias for request-scoped draft residency DraftResidencyPolicy draft_residency = DraftResidencyPolicy::Auto; + // Default speculative-decode policy; individual requests may override it. + SpeculationPolicy decode_mode = SpeculationPolicy::Adaptive; // Disk prefix cache std::string disk_cache_dir; // empty = disabled size_t disk_cache_budget_mb = 4096; // max disk usage in MB @@ -250,6 +254,49 @@ bool should_clamp_flowkv_disk_cache( bool flowkv, const DiskPrefixCachePolicy & policy); } // namespace http_detail +// Resolve the prompt-compression contract for a backend-provided sequence +// engine. Concurrent PFlash must keep both models resident: parking either +// model while another slot owns live device state invalidates that slot. The +// explicit persistent policy is therefore the opt-in that also implies the +// effective skip-park behavior; callers do not need a second, redundant CLI +// flag on large-memory concurrent hosts. +struct ConcurrentPflashPlan { + bool enabled = false; + bool force_skip_park = false; + std::string error; + + bool ok() const { return error.empty(); } +}; + +inline ConcurrentPflashPlan resolve_concurrent_pflash_plan( + const ServerConfig & config, bool drafter_tokenizer_available) { + ConcurrentPflashPlan plan; + if (config.pflash_mode == ServerConfig::PflashMode::OFF || + !config.pflash_upstream_base.empty()) { + return plan; + } + plan.enabled = true; + if (!drafter_tokenizer_available) { + plan.error = + "concurrent PFlash requires a loaded --prefill-drafter tokenizer"; + return plan; + } + if (config.draft_residency != DraftResidencyPolicy::Persistent) { + plan.error = + "concurrent PFlash requires --draft-residency persistent so " + "prompt compression cannot park live target/draft state"; + return plan; + } + if (config.prefix_cache_cap > 0 || config.prefill_cache_cap > 0 || + !config.disk_cache_dir.empty()) { + plan.error = + "concurrent paged PFlash does not support prefix/prefill " + "snapshots; disable the snapshot caches"; + return plan; + } + plan.force_skip_park = true; + return plan; +} // ─── Parsed request ───────────────────────────────────────────────────── @@ -292,6 +339,7 @@ struct ParsedRequest { DiskPrefixCachePolicy disk_cache_policy; // PPP: stable pin cut for tool-heavy requests (0 = use default boundary). int pin_end_token = 0; + std::optional decode_mode; }; // Parse request sampler fields, applying model-card defaults where present. @@ -605,6 +653,13 @@ struct ServerJob { // server-side prefill/elapsed telemetry does not erase queueing delay. std::chrono::steady_clock::time_point parallel_started_at{}; std::unique_ptr emitter; + // Prompt preparation (FlowKV/PFlash) is expensive and may load a resident + // drafter. Cache its result on the job so a pool-full retry never runs it + // twice. The original request tokens remain untouched for API accounting; + // this vector is the effective prompt admitted to the sequence engine. + bool parallel_prompt_prepared = false; + bool parallel_prompt_compressed = false; + std::vector parallel_prompt_tokens; }; // ─── Parse session_id from a chat-completion JSON body ────────────────── diff --git a/server/src/server/scheduler.cpp b/server/src/server/scheduler.cpp index 3c0b8f7c3..7c6b85614 100644 --- a/server/src/server/scheduler.cpp +++ b/server/src/server/scheduler.cpp @@ -36,6 +36,20 @@ struct SchedSlot { double prefill_s = 0.0; int n_gen_cap = 0; int completion_tokens = 0; + int effective_prompt_tokens = 0; + bool prompt_compressed = false; + uint64_t engine_request_id = 0; + uint64_t ddtree_steps = 0; + uint64_t ddtree_accepted_tokens = 0; + uint64_t ddtree_suspensions = 0; + uint64_t spec_steps = 0; + uint64_t spec_accepted_tokens = 0; + uint64_t spec_service_ar_steps = 0; + uint64_t target_forwards = 0; + uint64_t kvflash_page_ins = 0; + uint64_t kvflash_page_outs = 0; + uint64_t kvflash_resident_blocks = 0; + uint64_t kvflash_reselects = 0; bool client_disconnected = false; bool failed = false; std::string error; @@ -238,7 +252,7 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { stop_job_stream(s.job, &s.send_buffer); const double decode_s = std::chrono::duration( std::chrono::steady_clock::now() - s.decode_started_at).count(); - const int prompt_tokens = (int)req.prompt_tokens.size(); + const int prompt_tokens = s.effective_prompt_tokens; GenTimings gen_timings{ s.prefill_s, decode_s, @@ -253,9 +267,10 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { perf.prompt_tokens = (int)req.prompt_tokens.size(); perf.completion_tokens = s.completion_tokens; perf.prefill_tok_s = s.prefill_s > 0.0 - ? (double)req.prompt_tokens.size() / s.prefill_s : 0.0; + ? (double)prompt_tokens / s.prefill_s : 0.0; perf.decode_tok_s = decode_s > 0.0 ? (double)s.completion_tokens / decode_s : 0.0; + perf.pflash = s.prompt_compressed; status_.record_perf(perf); } @@ -292,17 +307,42 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { std::chrono::steady_clock::now() - s.started_at).count(); const int out_tokens = (int)s.gen_tokens.size(); std::fprintf(stderr, - "[server] chat DONE %s ok=%s in=%zu out=%d %.1fs %.1f tok/s " + "[server] chat DONE %s ok=%s in=%zu effective_in=%d out=%d %.1fs %.1f tok/s " "finish=%s slot=%d prefill=%.1fs decode=%.1fs(%.1ftok/s) parallel\n", req.response_id.c_str(), (!s.failed && backend_ok) ? "true" : "false", - req.prompt_tokens.size(), out_tokens, elapsed_s, + req.prompt_tokens.size(), prompt_tokens, out_tokens, elapsed_s, elapsed_s > 0.0 ? out_tokens / elapsed_s : 0.0, s.client_disconnected ? "client_disconnect" : s.emitter->finish_reason().c_str(), idx, s.prefill_s, decode_s, decode_s > 0.0 ? out_tokens / decode_s : 0.0); + const json concurrency_metrics = { + {"request_id", req.response_id}, + {"response_id", req.response_id}, + {"engine_request_id", s.engine_request_id}, + {"raw_prompt_tokens", req.prompt_tokens.size()}, + {"effective_prompt_tokens", prompt_tokens}, + {"output_tokens", out_tokens}, + {"pflash_applied", s.prompt_compressed}, + {"pflash_input_tokens", req.prompt_tokens.size()}, + {"pflash_output_tokens", prompt_tokens}, + {"ddtree_steps", s.ddtree_steps}, + {"ddtree_accepted_tokens", s.ddtree_accepted_tokens}, + {"ddtree_suspensions", s.ddtree_suspensions}, + {"spec_steps", s.spec_steps}, + {"spec_accepted_tokens", s.spec_accepted_tokens}, + {"spec_service_ar_steps", s.spec_service_ar_steps}, + {"target_forwards", s.target_forwards}, + {"kvflash_page_ins", s.kvflash_page_ins}, + {"kvflash_page_outs", s.kvflash_page_outs}, + {"kvflash_resident_blocks", s.kvflash_resident_blocks}, + {"kvflash_reselects", s.kvflash_reselects}, + }; + std::fprintf(stderr, "[concurrency-metrics] %s\n", + concurrency_metrics.dump().c_str()); + engine.retire(idx); // A retirement may have released the blocks the head job needs. deferred_retry_at = {}; @@ -392,6 +432,24 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { return AdmissionDisposition::Retired; } + const SpeculationPolicy decode_mode = resolve_speculation_policy( + config_.decode_mode, req.decode_mode); + const ConcurrentDecodeCapabilities decode_capabilities = + backend_.concurrent_decode_capabilities(); + if (!decode_capabilities.supports(decode_mode)) { + const std::string message = + std::string("decode_mode=") + + speculation_policy_name(decode_mode) + + " is unavailable for this server's concurrent decode " + "configuration"; + std::fprintf(stderr, + "[server] concurrent admission rejected %s: %s\n", + req.response_id.c_str(), message.c_str()); + send_error(job->fd, 409, message); + finish_job(job); + return AdmissionDisposition::Retired; + } + if (!job->announced) { job->announced = true; std::fprintf(stderr, @@ -431,9 +489,65 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { start_job_stream(job); } + // Apply the same FlowKV/PFlash precedence and overflow checks as the + // classic worker. Keep the prepared prompt on the job because an + // atomically-busy admission is retried at the FIFO head later. + if (!job->parallel_prompt_prepared) { + PreparedPrompt prepared = prepare_prompt(req); + if (prepared.error_status != 0) { + const std::string message = prepared.error.empty() + ? "prompt preparation failed" + : prepared.error; + std::fprintf(stderr, + "[server] concurrent prompt preparation failed: %s\n", + message.c_str()); + if (req.stream && job->sse_started) { + stop_job_stream(job); + for (const std::string & chunk : + sse_error_close_chunks(message)) { + send_job_bytes(job, chunk.data(), chunk.size()); + } + } else { + send_error(job->fd, prepared.error_status, message); + } + finish_job(job); + return AdmissionDisposition::Retired; + } + // Paged sequence engines cannot restore the classic snapshot + // format. Startup normally disables those caches; keep this + // check as a hard guard for embedded/non-CLI callers. + if (prepared.full_cache_hit_slot >= 0 || + prepared.full_cache_served_tokens >= 0) { + const std::string message = + "concurrent paged serving cannot restore a prefix snapshot"; + if (req.stream && job->sse_started) { + stop_job_stream(job); + for (const std::string & chunk : + sse_error_close_chunks(message)) { + send_job_bytes(job, chunk.data(), chunk.size()); + } + } else { + send_error(job->fd, 409, message); + } + finish_job(job); + return AdmissionDisposition::Retired; + } + job->parallel_prompt_tokens = std::move(prepared.tokens); + job->parallel_prompt_compressed = prepared.compressed; + job->parallel_prompt_prepared = true; + std::fprintf(stderr, + "[server] concurrent prompt READY %s raw=%zu effective=%zu " + "pflash=%s\n", + req.response_id.c_str(), req.prompt_tokens.size(), + job->parallel_prompt_tokens.size(), + prepared.compressed ? "true" : "false"); + } + const auto & effective_prompt = job->parallel_prompt_tokens; + // Admission only claims the slot and queues the prompt. Prefill // advances one chunk per engine step alongside live decode. - auto ar = engine.admit(next_request_id, req.prompt_tokens, + const uint64_t engine_request_id = next_request_id; + auto ar = engine.admit(engine_request_id, effective_prompt, req.sampler); if (ar.status == SeqEngine::AdmitResult::Status::busy) return AdmissionDisposition::Deferred; @@ -464,9 +578,12 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { s.admission_order = next_admission_order++; s.started_at = started_at; s.decode_started_at = started_at; // sane on prefill failure + s.effective_prompt_tokens = (int)effective_prompt.size(); + s.prompt_compressed = job->parallel_prompt_compressed; + s.engine_request_id = engine_request_id; s.n_gen_cap = std::min( n_gen_cap, - engine.max_context() - (int)req.prompt_tokens.size() + 1); + engine.max_context() - s.effective_prompt_tokens + 1); s.emitter = std::move(job->emitter); s.send_buffer.mark_progress(std::chrono::steady_clock::now()); if (budget_active && !config_.think_close_token_ids.empty() && @@ -606,8 +723,15 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { prefill_candidates.clear(); for (int i = 0; i < n_slots; i++) { if (slots[(size_t)i].job && !slots[(size_t)i].prefilling) { - step_plan.decode.push_back( - {i, slots[(size_t)i].pending_tok}); + SeqEngine::StepInput input; + input.slot = i; + input.token = slots[(size_t)i].pending_tok; + input.allow_speculation = + slots[(size_t)i].hook.close_token_ids.empty(); + input.speculation_policy = resolve_speculation_policy( + config_.decode_mode, + slots[(size_t)i].job->req.decode_mode); + step_plan.decode.push_back(input); } else if (slots[(size_t)i].job) { prefill_candidates.push_back( {i, slots[(size_t)i].admission_order}); @@ -617,9 +741,6 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { engine.step_plan_limits((int)step_plan.decode.size()); step_plan.prefills = plan_prefill_slices( prefill_candidates, step_limits, prefill_round_robin_start); - if (!prefill_candidates.empty()) { - ++prefill_round_robin_start; - } SeqEngine::StepResult step_result = engine.step(step_plan); const std::string protocol_error = @@ -655,9 +776,26 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { s.finished = true; continue; } - advance_slot(s, out.token); + s.ddtree_steps += out.ddtree_steps; + s.ddtree_accepted_tokens += out.ddtree_accepted_tokens; + s.ddtree_suspensions += out.ddtree_suspensions; + s.spec_steps += out.spec_steps; + s.spec_accepted_tokens += out.spec_accepted_tokens; + s.spec_service_ar_steps += out.spec_service_ar_steps; + s.target_forwards += out.target_forwards; + s.kvflash_page_ins += out.kvflash_page_ins; + s.kvflash_page_outs += out.kvflash_page_outs; + s.kvflash_resident_blocks = std::max( + s.kvflash_resident_blocks, out.kvflash_resident_blocks); + s.kvflash_reselects += out.kvflash_reselects; + consume_decode_output_tokens(out, [&](int32_t token) { + advance_slot(s, token); + return !s.finished; + }); } using PrefillStatus = SeqEngine::PrefillOutput::Status; + const bool prefill_progressed = + prefill_result_made_progress(step_result); for (const auto & out : step_result.prefills) { if (out.slot < 0 || out.slot >= n_slots) continue; SchedSlot & s = slots[(size_t)out.slot]; @@ -668,6 +806,12 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { s.finished = true; continue; } + if (out.status == PrefillStatus::deferred) { + // The exclusive decode graph made no prompt progress. Leave + // the slot untouched so normal FIFO planning retries it after + // the current decode wave drains. + continue; + } if (out.status == PrefillStatus::completed) { s.prefilling = false; publish_live_count(); @@ -681,6 +825,9 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { continue; } } + // Deferred means the engine intentionally left the selected slices + // untouched. Keep allocation fairness stable across such rounds. + if (prefill_progressed) ++prefill_round_robin_start; // Phase 4 — Non-blocking flush of every live slot's chunks. Progress // resets the stall clock; a reader that makes no progress for 30 s // or lets the buffer hit the cap is dropped (its slot retires). diff --git a/server/src/server/server_main.cpp b/server/src/server/server_main.cpp index 7373aace4..67f8058fd 100644 --- a/server/src/server/server_main.cpp +++ b/server/src/server/server_main.cpp @@ -122,6 +122,8 @@ static void print_usage(const char * prog) { " --no-fast-rollback Disable speculative fast rollback, even with --ddtree\n" " --ddtree Enable DDTree speculative decode\n" " --ddtree-budget DDTree budget (default: 22)\n" + " --decode-mode Speculative decode policy: ar, speculation,\n" + " or adaptive (default: adaptive)\n" " --verify-width laguna chain spec verify width (default: base 8,\n" " trimmed per step by drafter confidence; N = fixed base)\n" " --adaptive-experts [tau] MoE expert-count gating on verify batches\n" @@ -412,6 +414,15 @@ int main(int argc, char ** argv) { bargs.fast_rollback = true; } else if (std::strcmp(argv[i], "--ddtree-budget") == 0 && i + 1 < argc) { bargs.ddtree_budget = std::atoi(argv[++i]); + } else if (std::strcmp(argv[i], "--decode-mode") == 0 && i + 1 < argc) { + SpeculationPolicy policy; + if (!parse_speculation_policy(argv[++i], policy)) { + std::fprintf(stderr, + "[server] --decode-mode expects ar, speculation, or adaptive\n"); + return 2; + } + bargs.speculation_policy = policy; + sconfig.decode_mode = policy; } else if (std::strcmp(argv[i], "--adaptive-experts") == 0) { const char * tau = "0.80"; if (i + 1 < argc && argv[i + 1][0] != '-') { @@ -1115,6 +1126,8 @@ int main(int argc, char ** argv) { } std::fprintf(stderr, "[server] │ ddtree_budget = %d\n", bargs.ddtree_budget); std::fprintf(stderr, "[server] │ prefix_cache = %d slots\n", sconfig.prefix_cache_cap); + std::fprintf(stderr, "[server] │ decode_mode = %s\n", + speculation_policy_name(bargs.speculation_policy)); std::fprintf(stderr, "[server] │ prefill_cache = %d slots\n", sconfig.prefill_cache_cap); std::fprintf(stderr, "[server] │ cors = %s\n", sconfig.enable_cors ? "ON" : "off"); std::fprintf(stderr, "[server] │ cache_type_k = %s\n", @@ -1158,7 +1171,11 @@ int main(int argc, char ** argv) { sconfig.draft_path = bargs.draft_path ? bargs.draft_path : ""; sconfig.fa_window = bargs.fa_window; sconfig.ddtree_budget = bargs.ddtree_budget; - sconfig.speculative_enabled = bargs.ddtree_mode; + sconfig.speculative_enabled = + bargs.speculation_policy != SpeculationPolicy::Never && + (bargs.ddtree_mode || + (bargs.paged_attention && bargs.max_concurrency > 1 && + bargs.draft_path != nullptr)); sconfig.target_sharding = bargs.device.is_layer_split(); // KV type: report the operator's choice if set, else the family default // the backend resolves (the tq3_0 auto policy was removed; laguna uses diff --git a/server/test/bench_paged_attention.cpp b/server/test/bench_paged_attention.cpp index 49fc00960..1e5afe9c8 100644 --- a/server/test/bench_paged_attention.cpp +++ b/server/test/bench_paged_attention.cpp @@ -725,7 +725,8 @@ bool run_case( ggml_tensor * paged_output = ggml_paged_attn_ext( ctx.value, q_paged, k_paged, v_paged, table, kv_seq_lens_tensor, nullptr, nullptr, 1.0f / std::sqrt(static_cast(D)), - BLOCK_SIZE, max_context); + BLOCK_SIZE, max_context, + nullptr, nullptr, 0, 0, 0); ggml_tensor * contiguous_output = ggml_flash_attn_ext( ctx.value, q_contiguous, k_contiguous, v_contiguous, padding_mask, 1.0f / std::sqrt(static_cast(D)), diff --git a/server/test/seq_engine_contract.h b/server/test/seq_engine_contract.h index 63ca035ab..17073646c 100644 --- a/server/test/seq_engine_contract.h +++ b/server/test/seq_engine_contract.h @@ -151,7 +151,7 @@ inline std::vector check_seq_engine_contract(SeqEngine & engine) { require(remaining[(size_t)output.slot] > 1, "prefill reported advanced for its final token"); --remaining[(size_t)output.slot]; - } else { + } else if (output.status == PrefillStatus::completed) { require(remaining[(size_t)output.slot] == 1, "prefill reported completion before its final token"); remaining[(size_t)output.slot] = 0; @@ -258,6 +258,19 @@ inline std::vector check_seq_engine_contract(SeqEngine & engine) { return violations; } + // Scheduler policy can retain commit authority for selected slots. A + // conforming engine may still use its ordinary AR implementation, but it + // must not return already-committed children for a disabled input. + SeqEngine::StepPlan no_speculation; + no_speculation.decode = decode_inputs(); + for (SeqEngine::StepInput & input : no_speculation.decode) { + input.allow_speculation = false; + } + if (!execute(no_speculation)) { + retire_all(); + return violations; + } + // A full engine is retryable admission pressure, not a request error. if (n_slots == 2) { const SeqEngine::AdmitResult full = engine.admit(3, {31}, greedy); diff --git a/server/test/smoke_draft_graph.cpp b/server/test/smoke_draft_graph.cpp index 544f8b51f..57e9b69dd 100644 --- a/server/test/smoke_draft_graph.cpp +++ b/server/test/smoke_draft_graph.cpp @@ -20,6 +20,7 @@ #include "ggml-alloc.h" #include "ggml-backend.h" #include "ggml-cuda.h" +#include #include #include @@ -103,11 +104,22 @@ int main(int argc, char ** argv) { gi.positions_k = pos_k; DraftGraphOutputs go = build_draft_graph(gctx, w, gi); - if (!go.hidden_states) { std::fprintf(stderr, "build_draft_graph returned null\n"); return 1; } + if (!go.hidden_prenorm || !go.hidden_states) { + std::fprintf(stderr, "build_draft_graph returned null output\n"); + return 1; + } + ggml_tensor * rebuilt_hidden = ggml_rms_norm( + gctx, go.hidden_prenorm, DFLASH27B_RMS_EPS); + rebuilt_hidden = ggml_mul(gctx, rebuilt_hidden, w.out_norm); + ggml_set_name(rebuilt_hidden, "rebuilt_draft_hidden_out"); + ggml_set_output(go.hidden_prenorm); ggml_set_output(go.hidden_states); + ggml_set_output(rebuilt_hidden); ggml_cgraph * gf = ggml_new_graph(gctx); + ggml_build_forward_expand(gf, go.hidden_prenorm); ggml_build_forward_expand(gf, go.hidden_states); + ggml_build_forward_expand(gf, rebuilt_hidden); std::printf("graph built: n_nodes=%d\n", ggml_graph_n_nodes(gf)); // ── 5. Allocate graph + all input tensors on the backend @@ -158,6 +170,29 @@ int main(int argc, char ** argv) { } std::vector out(n_out_elems); ggml_backend_tensor_get(go.hidden_states, out.data(), 0, sizeof(float) * out.size()); + std::vector prenorm(n_out_elems); + std::vector rebuilt(n_out_elems); + ggml_backend_tensor_get(go.hidden_prenorm, prenorm.data(), 0, + sizeof(float) * prenorm.size()); + ggml_backend_tensor_get(rebuilt_hidden, rebuilt.data(), 0, + sizeof(float) * rebuilt.size()); + + double max_prenorm_delta = 0.0; + double max_rebuild_error = 0.0; + for (size_t i = 0; i < out.size(); ++i) { + max_prenorm_delta = std::max( + max_prenorm_delta, std::fabs((double)prenorm[i] - out[i])); + max_rebuild_error = std::max( + max_rebuild_error, std::fabs((double)rebuilt[i] - out[i])); + } + if (max_prenorm_delta < 1e-5 || max_rebuild_error > 1e-5) { + std::fprintf(stderr, + "FAIL: pre-norm calibration output delta=%.8g rebuild_error=%.8g\n", + max_prenorm_delta, max_rebuild_error); + return 1; + } + std::printf("pre-norm export OK: delta=%.6g rebuild_error=%.6g\n", + max_prenorm_delta, max_rebuild_error); int n_nan = 0, n_inf = 0; double sum = 0.0, sumsq = 0.0; diff --git a/server/test/test_chain_spec_shapes.cpp b/server/test/test_chain_spec_shapes.cpp new file mode 100644 index 000000000..6814dd05e --- /dev/null +++ b/server/test/test_chain_spec_shapes.cpp @@ -0,0 +1,120 @@ +#include "common/concurrency/chain_spec_shapes.h" +#include "host_check.h" + +#include +#include + +using namespace dflash::common; + +static int g_checks = 0; + +int main() { + const std::vector draft = {10, 11, 12, 13}; + const DDTree tree = make_chain_verify_tree(draft); + CHECK(tree.n_nodes == 3); + CHECK((tree.token_ids == std::vector{11, 12, 13})); + CHECK((tree.depths == std::vector{1, 2, 3})); + CHECK((tree.parents == std::vector{-1, 0, 1, 2})); + + CHECK(resolve_chain_verify_depth(0, 4) == 4); + CHECK(resolve_chain_verify_depth(2, 4) == 2); + CHECK(resolve_chain_verify_depth(4, 4) == 4); + CHECK(resolve_chain_verify_depth(1, 4) == 0); + CHECK(resolve_chain_verify_depth(5, 4) == 0); + CHECK(resolve_chain_verify_depth(0, 1) == 0); + std::vector short_draft = draft; + CHECK(truncate_chain_proposal(short_draft, 2)); + CHECK((short_draft == std::vector{10, 11})); + const DDTree short_tree = make_chain_verify_tree(short_draft); + CHECK(short_tree.n_nodes == 1); + CHECK((short_tree.parents == std::vector{-1, 0})); + const std::vector before_invalid = short_draft; + CHECK(!truncate_chain_proposal(short_draft, 1)); + CHECK(short_draft == before_invalid); + CHECK(!truncate_chain_proposal(short_draft, 3)); + CHECK(short_draft == before_invalid); + const std::vector bucket_inputs = { + 1, 2, 3, 4, 5, 6, 7, 8, 9, 12, 13, 16, 17, + }; + const std::vector bucket_expected = { + 1, 2, 3, 4, 6, 6, 8, 8, 12, 12, 16, 16, 24, + }; + for (size_t i = 0; i < bucket_inputs.size(); ++i) { + CHECK(chain_decode_bucket_width(bucket_inputs[i]) == + bucket_expected[i]); + } + + int pending = -1; + const int32_t full_posterior[] = {11, 12, 13, 14}; + std::vector accepted = + follow_verified_tree(tree, full_posterior, pending); + CHECK((accepted == std::vector{0, 1, 2, 3})); + CHECK(pending == 14); + + const int32_t rejected_posterior[] = {11, 99, 13, 14}; + accepted = follow_verified_tree(tree, rejected_posterior, pending); + CHECK((accepted == std::vector{0, 1})); + CHECK(pending == 99); + + CHECK(truncate_verified_path( + accepted, 1, rejected_posterior, pending)); + CHECK((accepted == std::vector{0})); + CHECK(pending == 11); + + const ChainLaunchShape mixed = chain_launch_shape( + {1, 0, 1, 0, 0, 0}, {4, 0, 2, 0, 0, 0}, 16); + CHECK(mixed.spec_lanes == 2); + CHECK(mixed.tree_bucket == 2); + CHECK(mixed.tree_rows == 32); + CHECK(mixed.ar_lanes == 4); + CHECK(mixed.ar_bucket == 4); + CHECK(mixed.accepted_rows == 6); + CHECK(mixed.commit_rows == 10); + + const ChainLaunchShape all_spec = chain_launch_shape( + {1, 1, 1}, {1, 2, 3}, 16); + CHECK(all_spec.tree_bucket == 3); + CHECK(all_spec.ar_bucket == 0); + CHECK(all_spec.commit_rows == 6); + + const ChainLaunchShape ar_after_spec_failures = chain_launch_shape( + {0, 0}, {0, 0}, 16); + CHECK(ar_after_spec_failures.spec_lanes == 0); + CHECK(ar_after_spec_failures.tree_bucket == 0); + CHECK(ar_after_spec_failures.tree_rows == 0); + CHECK(ar_after_spec_failures.ar_lanes == 2); + CHECK(ar_after_spec_failures.ar_bucket == 2); + CHECK(ar_after_spec_failures.commit_rows == 2); + + CHECK(chain_lane_disposition(false, false) == + ChainLaneDisposition::AR); + CHECK(chain_lane_disposition(true, false) == + ChainLaneDisposition::Speculation); + CHECK(chain_lane_disposition(true, true) == + ChainLaneDisposition::Failed); + CHECK(chain_lane_disposition(false, true) == + ChainLaneDisposition::Failed); + CHECK(chain_lane_executes(ChainLaneDisposition::AR)); + CHECK(chain_lane_executes(ChainLaneDisposition::Speculation)); + CHECK(!chain_lane_executes(ChainLaneDisposition::Failed)); + + const auto eos = [](int32_t token) { return token == 2; }; + const std::vector eos_first_child = {10, 2, 11}; + CHECK(chain_min_tokens_safe_prefix( + eos_first_child, 0, 3, eos) == 1); + CHECK(chain_min_tokens_safe_prefix( + eos_first_child, 2, 3, eos) == 2); + const std::vector eos_second_child = {10, 11, 2, 12}; + CHECK(chain_min_tokens_safe_prefix( + eos_second_child, 0, 3, eos) == 2); + CHECK(chain_min_tokens_safe_prefix( + eos_second_child, 1, 3, eos) == 3); + const std::vector eos_root = {2, 11, 12}; + CHECK(chain_min_tokens_safe_prefix( + eos_root, 0, 3, eos) == eos_root.size()); + CHECK(chain_min_tokens_safe_prefix( + eos_first_child, 0, 0, eos) == 2); + + std::printf("chain spec shape tests passed: %d checks\n", g_checks); + return 0; +} diff --git a/server/test/test_ddtree_path.cpp b/server/test/test_ddtree_path.cpp new file mode 100644 index 000000000..a4118ce1e --- /dev/null +++ b/server/test/test_ddtree_path.cpp @@ -0,0 +1,48 @@ +#include "common/ddtree.h" +#include "host_check.h" + +#include +#include +#include + +using dflash::common::DDTree; +using dflash::common::follow_verified_tree; +using dflash::common::truncate_verified_path; + +static int g_checks = 0; + +int main() { + DDTree tree; + tree.n_nodes = 2; + tree.token_ids = {11, 22}; + tree.depths = {1, 2}; + tree.parents = {-1, 0, 1}; + tree.child_maps.resize(3); + tree.child_maps[0][11] = 1; + tree.child_maps[1][22] = 2; + + const int32_t posterior[] = {11, 22, 33}; + int pending = -1; + std::vector accepted = + follow_verified_tree(tree, posterior, pending); + CHECK((accepted == std::vector{0, 1, 2})); + CHECK(pending == 33); + + // Truncating after node 1 means node 2's token becomes pending. Keeping + // the old value (33) would skip token 22 and describe uncommitted state. + CHECK(truncate_verified_path(accepted, 2, posterior, pending)); + CHECK((accepted == std::vector{0, 1})); + CHECK(pending == 22); + + // An unchanged path preserves the already-computed pending token. + CHECK(!truncate_verified_path(accepted, 2, posterior, pending)); + CHECK(pending == 22); + + // No headroom is represented explicitly and never dereferences a tip. + CHECK(truncate_verified_path(accepted, 0, posterior, pending)); + CHECK(accepted.empty()); + CHECK(pending == -1); + + std::puts("ddtree path tests passed"); + return 0; +} diff --git a/server/test/test_delta_transition_journal.cpp b/server/test/test_delta_transition_journal.cpp new file mode 100644 index 000000000..a5859b2e6 --- /dev/null +++ b/server/test/test_delta_transition_journal.cpp @@ -0,0 +1,186 @@ +#include "qwen35/delta_transition_journal.h" +#include "host_check.h" + +#include +#include +#include +#include + +using dflash::qwen35::DeltaTransition; +using dflash::qwen35::DeltaTransitionGateMode; +using dflash::qwen35::DeltaTransitionJournal; +using dflash::qwen35::apply_delta_transition; +using dflash::qwen35::capture_delta_transition; +using dflash::qwen35::commit_delta_transition_prefix; +using dflash::qwen35::delta_transition_float_count; + +static int g_checks = 0; + +namespace { + +struct RawStep { + std::vector key; + std::vector value; + std::vector gate; + float beta = 0.0f; +}; + +bool near(const std::vector & lhs, const std::vector & rhs) { + if (lhs.size() != rhs.size()) return false; + for (size_t i = 0; i < lhs.size(); ++i) { + const float scale = std::max( + 1.0f, std::max(std::fabs(lhs[i]), std::fabs(rhs[i]))); + if (std::fabs(lhs[i] - rhs[i]) > 4e-6f * scale) return false; + } + return true; +} + +// Independent spelling of the existing gated_delta_net_cuda recurrence. It +// deliberately consumes raw step inputs rather than a captured transition. +void replay_reference( + std::vector & state, + size_t rows, + size_t cols, + const RawStep & step, + DeltaTransitionGateMode gate_mode) { + for (size_t col = 0; col < cols; ++col) { + float projection = 0.0f; + for (size_t row = 0; row < rows; ++row) { + const float projection_gate = + gate_mode == DeltaTransitionGateMode::RowWise + ? step.gate[row] + : 1.0f; + projection += projection_gate * + state[col * rows + row] * step.key[row]; + } + const float scalar_gate = + gate_mode == DeltaTransitionGateMode::Scalar + ? step.gate[0] + : 1.0f; + const float delta = + (step.value[col] - scalar_gate * projection) * step.beta; + for (size_t row = 0; row < rows; ++row) { + const size_t index = col * rows + row; + const float update_gate = + gate_mode == DeltaTransitionGateMode::Scalar + ? step.gate[0] + : step.gate[row]; + state[index] = std::fma( + step.key[row], delta, update_gate * state[index]); + } + } +} + +std::vector initial_state(size_t rows, size_t cols, int layer) { + std::vector state(rows * cols); + for (size_t i = 0; i < state.size(); ++i) { + state[i] = 0.013f * static_cast(i + 1) - + 0.07f * static_cast(layer + 1); + } + return state; +} + +std::vector make_steps( + size_t rows, + size_t cols, + size_t count, + int layer, + DeltaTransitionGateMode gate_mode) { + std::vector steps(count); + for (size_t t = 0; t < count; ++t) { + RawStep & step = steps[t]; + step.key.resize(rows); + step.value.resize(cols); + step.gate.resize( + gate_mode == DeltaTransitionGateMode::Scalar ? 1 : rows); + for (size_t row = 0; row < rows; ++row) { + step.key[row] = 0.021f * static_cast(row + 1) - + 0.009f * static_cast(t + layer); + } + for (size_t col = 0; col < cols; ++col) { + step.value[col] = 0.031f * static_cast(col + 1) + + 0.017f * static_cast(t + 2 * layer); + } + for (size_t row = 0; row < step.gate.size(); ++row) { + step.gate[row] = 0.78f + + 0.011f * static_cast((row + t + layer) % 7); + } + step.beta = 0.42f + 0.03f * static_cast(t % 4); + } + return steps; +} + +void prove_all_prefixes(DeltaTransitionGateMode gate_mode) { + constexpr size_t rows = 8; + constexpr size_t cols = 7; + constexpr size_t tokens = 6; + constexpr int layers = 2; + + for (int layer = 0; layer < layers; ++layer) { + const std::vector base = initial_state(rows, cols, layer); + const std::vector steps = + make_steps(rows, cols, tokens, layer, gate_mode); + + DeltaTransitionJournal journal; + journal.rows = rows; + journal.cols = cols; + std::vector verify_state = base; + for (const RawStep & step : steps) { + DeltaTransition transition; + CHECK(capture_delta_transition( + verify_state, rows, cols, step.key, step.value, step.gate, + step.beta, gate_mode, transition)); + CHECK(apply_delta_transition( + transition, rows, cols, verify_state)); + journal.transitions.push_back(transition); + } + + for (size_t accepted = 1; accepted <= tokens; ++accepted) { + std::vector replayed = base; + for (size_t t = 0; t < accepted; ++t) { + replay_reference( + replayed, rows, cols, steps[t], gate_mode); + } + + std::vector committed = base; + CHECK(commit_delta_transition_prefix( + journal, accepted, committed)); + CHECK(near(committed, replayed)); + } + + std::vector full_replay = base; + for (const RawStep & step : steps) { + replay_reference(full_replay, rows, cols, step, gate_mode); + } + CHECK(near(verify_state, full_replay)); + } +} + +} // namespace + +int main() { + prove_all_prefixes(DeltaTransitionGateMode::Scalar); + prove_all_prefixes(DeltaTransitionGateMode::RowWise); + + CHECK(delta_transition_float_count( + 128, 128, DeltaTransitionGateMode::Scalar) == 257); + CHECK(delta_transition_float_count( + 128, 128, DeltaTransitionGateMode::RowWise) == 384); + + // Contract guards are fail-closed and transactional. + DeltaTransitionJournal malformed; + malformed.rows = 2; + malformed.cols = 2; + malformed.transitions.push_back(DeltaTransition{}); + std::vector state = {1.0f, 2.0f, 3.0f, 4.0f}; + const std::vector original = state; + CHECK(!commit_delta_transition_prefix(malformed, 1, state)); + CHECK(state == original); + CHECK(!commit_delta_transition_prefix(malformed, 2, state)); + CHECK(state == original); + CHECK(commit_delta_transition_prefix(malformed, 0, state)); + CHECK(state == original); + + std::printf("delta transition journal tests passed: %d checks\n", g_checks); + return 0; +} diff --git a/server/test/test_dflash2_benefit.cpp b/server/test/test_dflash2_benefit.cpp new file mode 100644 index 000000000..c6bbbaf09 --- /dev/null +++ b/server/test/test_dflash2_benefit.cpp @@ -0,0 +1,225 @@ +#include "common/dflash2_benefit.h" +#include "common/speculation/speculation_gate.h" +#include "host_check.h" + +#include +#include +#include +#include + +using namespace dflash::common; + +static int g_checks = 0; + +static DFlash2BenefitModelSignature seeded_signature() { + DFlash2BenefitModelSignature value; + value.target_layers = 64; + value.target_hidden = 5120; + value.target_vocab = 248320; + value.draft_layers = 5; + value.draft_hidden = 5120; + value.draft_block_size = 8; + value.selector_rank = 256; + value.selector_top_k = 16; + value.selector_vocab = 248320; + value.conv_kernel_size = 2; + value.conv_group_size = 16; + value.target_file_size = 15195272800ULL; + value.draft_file_size = 2045471776ULL; + return value; +} + +static DFlash2SelectorTrace trace_from( + std::initializer_list> values) { + DFlash2SelectorTrace trace; + for (const auto & value : values) { + DFlash2DepthSignal signal; + signal.selected_log_prob = value.first; + signal.selector_winner_mass = value.second; + trace.depths.push_back(signal); + } + return trace; +} + +int main() { + DFlash2BenefitProvider provider(seeded_signature()); + CHECK(provider.ready()); + CHECK(std::string(provider.score_kind()) == + "qwen38-dflash2-selector-benefit-v1"); + CHECK(provider.config().lm_log_weight == 0.10); + CHECK(provider.config().hazard_scale == 1.0); + CHECK(provider.config().yield_scale == 1.0); + + // Retained C1 first-block traces for he08 code and prose. The adapter is + // continuous and content-agnostic: code has the higher expected yield, while + // prose remains lower and cohort-dependent. Values are (selected LM + // log-prob, selector winner mass). + const DFlash2SelectorTrace code_like = trace_from({ + {0.0f, 0.99999988f}, {-7.6293945e-06f, 1.0f}, + {-0.44184685f, 0.64686394f}, {-2.0253334f, 1.0f}, + {-1.2170925f, 1.0f}, {-0.84755516f, 0.99942774f}, + {-3.7030106f, 0.86279351f}, + }); + const DFlash2SelectorTrace prose_like = trace_from({ + {-0.046934128f, 0.95176214f}, {-1.527895f, 0.37046611f}, + {-3.398098f, 0.98939508f}, {-3.0824165f, 0.41014573f}, + {-0.30518532f, 0.47964928f}, {-4.0057392f, 0.99999905f}, + {-1.676815f, 0.89187294f}, + }); + + DFlash2BenefitEstimate code; + DFlash2BenefitEstimate prose; + std::string error; + CHECK(provider.estimate(code_like, 8, code, &error)); + CHECK(error.empty()); + CHECK(provider.estimate(prose_like, 8, prose, &error)); + CHECK(code.conditional_hazards.size() == 7); + CHECK(prose.conditional_hazards.size() == 7); + CHECK(std::abs(code.expected_yield - 5.330600824) < 1e-6); + CHECK(std::abs(prose.expected_yield - 2.684506084) < 1e-6); + CHECK(code.expected_yield > prose.expected_yield + 2.5); + + // Feed the retained request benefits into representative C2 profile + // costs. These are joint economics, not a topic classifier: code wins + // alone, prose loses alone and as an all-prose C2 cohort, while a mixed + // cohort must rank code first and may admit prose only as a prefix-k + // amortization result. + const SpecCostTables observed_costs{ + {{8, 16}, {42912.6, 47572.8}}, + {{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}, + {29322.0, 34968.2, 38639.9, 38639.9, + 40077.6, 40077.6, 40077.6, 40077.6, + 42360.3, 42360.3, 42360.3, 42360.3, + 42360.3, 42360.3, 42360.3, 44412.2}}, + {{1, 2}, {7934.4, 13967.3}}, + }; + SpecStepGeometry observed_geometry; + observed_geometry.tree_width = 8; + observed_geometry.bucket = [](int lanes) { return lanes; }; + auto benefit_candidate = [](uint64_t id, int slot, double score, + SpeculationPolicy policy = + SpeculationPolicy::Adaptive) { + return SpecCandidate{ + id, slot, policy, true, true, score, {}, + kDFlash2BenefitAdapterVersion}; + }; + + SpeculationGate code_c1(observed_costs, observed_geometry, 8); + SpecPlan gate_plan = code_c1.plan( + 1, {benefit_candidate(100, 5, code.expected_yield)}, 1); + CHECK(gate_plan.admitted_count == 1); + CHECK(code_c1.initial_score_kind(100) == + kDFlash2BenefitAdapterVersion); + + SpeculationGate prose_c1(observed_costs, observed_geometry, 8); + gate_plan = prose_c1.plan( + 1, {benefit_candidate(200, 3, prose.expected_yield)}, 1); + CHECK(gate_plan.admitted_count == 0); + + SpeculationGate prose_c2(observed_costs, observed_geometry, 8); + gate_plan = prose_c2.plan(2, { + benefit_candidate(300, 4, prose.expected_yield), + benefit_candidate(301, 1, prose.expected_yield)}, 2); + CHECK(gate_plan.admitted_count == 0); + + SpeculationGate mixed_c2(observed_costs, observed_geometry, 8); + gate_plan = mixed_c2.plan(2, { + benefit_candidate(401, 3, prose.expected_yield), + benefit_candidate(400, 5, code.expected_yield)}, 2); + CHECK(gate_plan.admitted_count >= 1); + CHECK(!gate_plan.admitted_request_ids.empty()); + CHECK(gate_plan.admitted_request_ids[0] == 400); + CHECK(gate_plan.ordered.size() == 2); + CHECK(gate_plan.ordered[0].request_id == 400); + CHECK(gate_plan.ordered[0].slot == 5); + CHECK(gate_plan.ordered[1].request_id == 401); + CHECK(gate_plan.ordered[1].slot == 3); + CHECK(gate_plan.ordered[0].admitted); + CHECK(gate_plan.ordered[1].admitted == + (gate_plan.admitted_count == 2)); + + SpeculationGate code_with_ar_peer( + observed_costs, observed_geometry, 8); + gate_plan = code_with_ar_peer.plan(2, { + benefit_candidate(500, 2, code.expected_yield), + benefit_candidate(501, 7, prose.expected_yield, + SpeculationPolicy::Never)}, 2); + CHECK(gate_plan.admitted_count == 1); + CHECK(gate_plan.admitted_request_ids.size() == 1); + CHECK(gate_plan.admitted_request_ids[0] == 500); + + // Maximum depth consumes exactly block_size-1 signals; diagnostic tail + // values beyond that depth cannot affect the score. + DFlash2SelectorTrace with_tail = code_like; + with_tail.depths.push_back(trace_from({{-100.0f, 0.001f}}).depths[0]); + DFlash2BenefitEstimate tail; + CHECK(provider.estimate(with_tail, 8, tail, &error)); + CHECK(tail.conditional_hazards.size() == 7); + CHECK(std::abs(tail.expected_yield - code.expected_yield) < 1e-12); + CHECK(!provider.estimate(code_like, 9, tail, &error)); + CHECK(error.find("outside") != std::string::npos); + + // A partial, malformed, or nonfinite first trace fails closed and never + // publishes a synthetic request score. + DFlash2SelectorTrace missing = code_like; + missing.depths.pop_back(); + CHECK(!provider.estimate(missing, 8, tail, &error)); + CHECK(error.find("missing") != std::string::npos); + DFlash2SelectorTrace malformed = code_like; + malformed.depths[2].selector_winner_mass = 0.0f; + CHECK(!provider.estimate(malformed, 8, tail, &error)); + malformed = code_like; + malformed.depths[2].selected_log_prob = 0.1f; + CHECK(!provider.estimate(malformed, 8, tail, &error)); + malformed = code_like; + malformed.depths[2].selected_log_prob = + std::numeric_limits::quiet_NaN(); + CHECK(!provider.estimate(malformed, 8, tail, &error)); + double unpublished = std::numeric_limits::quiet_NaN(); + CHECK(!provider.publish_once(malformed, 8, unpublished, &error)); + CHECK(std::isnan(unpublished)); + + // Request lifetime is external and explicit: once the first valid score + // is published, neither a lower later trace nor a malformed trace can + // overwrite it. + double first_score = std::numeric_limits::quiet_NaN(); + CHECK(provider.publish_once(code_like, 8, first_score, &error)); + CHECK(std::abs(first_score - code.expected_yield) < 1e-12); + CHECK(provider.publish_once(prose_like, 8, first_score, &error)); + CHECK(std::abs(first_score - code.expected_yield) < 1e-12); + CHECK(provider.publish_once(malformed, 8, first_score, &error)); + CHECK(std::abs(first_score - code.expected_yield) < 1e-12); + + // Unknown models and invalid adapter versions/coefficients are not + // silently generalized. A conservative scale may only reduce yield. + DFlash2BenefitModelSignature unknown = seeded_signature(); + unknown.selector_rank = 128; + DFlash2BenefitProvider unknown_provider(unknown); + CHECK(!unknown_provider.ready()); + CHECK(unknown_provider.error().find("unsupported") != std::string::npos); + CHECK(!unknown_provider.estimate(code_like, 8, tail, &error)); + + DFlash2BenefitConfig bad_version; + bad_version.adapter_version = "future-unfitted-adapter"; + DFlash2BenefitProvider version_provider( + seeded_signature(), bad_version); + CHECK(!version_provider.ready()); + + DFlash2BenefitConfig conservative; + conservative.hazard_scale = 0.9; + DFlash2BenefitProvider conservative_provider( + seeded_signature(), conservative); + CHECK(conservative_provider.ready()); + CHECK(conservative_provider.estimate(code_like, 8, tail, &error)); + CHECK(tail.expected_yield < code.expected_yield); + + DFlash2BenefitConfig invalid_coefficients; + invalid_coefficients.lm_log_weight = + std::numeric_limits::infinity(); + DFlash2BenefitProvider invalid_provider( + seeded_signature(), invalid_coefficients); + CHECK(!invalid_provider.ready()); + + std::printf("DFlash2 benefit adapter: %d checks passed\n", g_checks); + return 0; +} diff --git a/server/test/test_dflash2_selector_validation.cpp b/server/test/test_dflash2_selector_validation.cpp new file mode 100644 index 000000000..bbc46acc8 --- /dev/null +++ b/server/test/test_dflash2_selector_validation.cpp @@ -0,0 +1,86 @@ +#include "common/dflash2_selector_validation.h" +#include "host_check.h" + +#include +#include + +using namespace dflash::common; + +static int g_checks = 0; + +static DFlash2SelectorLayout valid_layout() { + DFlash2SelectorLayout layout; + layout.rank = 32; + layout.top_k = 16; + layout.hproj_rank = 32; + layout.pred_rank = 32; + layout.pred_vocab = 151936; + layout.succ_rank = 32; + layout.succ_vocab = 151936; + layout.target_output_vocab = 151936; + layout.target_declared_vocab = 151936; + return layout; +} + +int main() { + std::string error; + DFlash2SelectorLayout layout = valid_layout(); + CHECK(validate_dflash2_selector_layout(layout, error)); + CHECK(error.empty()); + + for (int K = 1; K <= 8; ++K) { + layout = valid_layout(); + layout.top_k = K; + CHECK(validate_dflash2_selector_layout(layout, error)); + } + for (int K : {12, 16}) { + layout = valid_layout(); + layout.top_k = K; + CHECK(validate_dflash2_selector_layout(layout, error)); + } + for (int K : {0, 9, 10, 11, 13, 14, 15, 17}) { + layout = valid_layout(); + layout.top_k = K; + CHECK(!validate_dflash2_selector_layout(layout, error)); + CHECK(error.find("top_k=") != std::string::npos); + CHECK(error.find("unsupported") != std::string::npos); + } + + layout = valid_layout(); + layout.succ_vocab--; + CHECK(!validate_dflash2_selector_layout(layout, error)); + CHECK(error.find("codebook vocab mismatch") != std::string::npos); + + layout = valid_layout(); + layout.target_output_vocab--; + CHECK(!validate_dflash2_selector_layout(layout, error)); + CHECK(error.find("target vocab mismatch") != std::string::npos); + + layout = valid_layout(); + layout.target_declared_vocab = 0; + layout.target_output_vocab--; + CHECK(!validate_dflash2_selector_layout(layout, error)); + CHECK(error.find("target output/lm_head") != std::string::npos); + + layout = valid_layout(); + layout.target_output_vocab = 0; + layout.target_declared_vocab--; + CHECK(!validate_dflash2_selector_layout(layout, error)); + CHECK(error.find("target.n_vocab") != std::string::npos); + + layout = valid_layout(); + layout.pred_vocab = 8; + layout.succ_vocab = 8; + layout.target_output_vocab = 0; + layout.target_declared_vocab = 0; + CHECK(!validate_dflash2_selector_layout(layout, error)); + CHECK(error.find("exceeds codebook vocab") != std::string::npos); + + layout = valid_layout(); + layout.succ_rank = 31; + CHECK(!validate_dflash2_selector_layout(layout, error)); + CHECK(error.find("rank mismatch") != std::string::npos); + + std::printf("dflash2 selector validation: %d checks passed\n", g_checks); + return 0; +} diff --git a/server/test/test_draft_topk_cuda.cpp b/server/test/test_draft_topk_cuda.cpp index a1defe7c7..dd507e6ee 100644 --- a/server/test/test_draft_topk_cuda.cpp +++ b/server/test/test_draft_topk_cuda.cpp @@ -30,6 +30,7 @@ using dflash::common::extract_draft_topk; using dflash::common::geometric_extract_draft_topk_cuda; +using dflash::common::geometric_draft_topk_cuda_supports_k; namespace { @@ -124,6 +125,30 @@ namespace { struct DraftTopkCudaFixture {}; } +TEST_CASE(DraftTopkCudaFixture, draft_topk_cuda_dispatch_contract_host_only) { + for (int K = -1; K <= 18; ++K) { + const bool expected = (K >= 1 && K <= 8) || K == 12 || K == 16; + CHECK(geometric_draft_topk_cuda_supports_k(K) == expected); + } + + // Unsupported K must be rejected before pointer inspection or any CUDA + // call. This makes the fallback contract testable on a host with no GPU. + const void * invalid_device_pointer = + reinterpret_cast(uintptr_t{1}); + std::vector log_probs(64, 123.0f); + std::vector token_ids(64, 456); + for (int K : {0, 9, 10, 11, 13, 14, 15, 17, 64}) { + CHECK(!geometric_extract_draft_topk_cuda( + invalid_device_pointer, 1, 128, K, + log_probs.data(), token_ids.data(), 1.0f)); + CHECK(log_probs[0] == 123.0f); + CHECK(token_ids[0] == 456); + } + CHECK(!geometric_extract_draft_topk_cuda( + invalid_device_pointer, 1, 8, 16, + log_probs.data(), token_ids.data(), 1.0f)); +} + TEST_CASE(DraftTopkCudaFixture, draft_topk_cuda_suite) { int dev_count = 0; if (cudaGetDeviceCount(&dev_count) != cudaSuccess || dev_count == 0) { @@ -131,8 +156,7 @@ TEST_CASE(DraftTopkCudaFixture, draft_topk_cuda_suite) { return; } - // The kernel supports K up to kMaxK (=8 in geometric_draft_topk_cuda.cu); larger K is - // handled by a documented CPU fallback (returns false), checked separately. + // Exercise every instantiated dispatch family, including DFlash 2's K=16. const Case cases[] = { // Realistic decode shape: Qwen3.5 vocab, small position batch. {15, 151936, 8, 1.0f}, @@ -145,6 +169,8 @@ TEST_CASE(DraftTopkCudaFixture, draft_topk_cuda_suite) { {3, 257, 8, 1.0f}, // vocab barely above K, non-power-of-two {1, 151936, 1, 1.0f}, // K=1 (argmax + log_z) {15, 151936, 4, 1.0f}, + {3, 4096, 12, 1.0f}, + {3, 4096, 16, 1.0f}, }; int failures = 0; @@ -154,24 +180,27 @@ TEST_CASE(DraftTopkCudaFixture, draft_topk_cuda_suite) { idx++; } - // Fallback contract: K beyond the kernel's supported range must return false - // (not silently produce wrong output) so the caller can use the CPU path. + // Fallback contract: both in-range dispatch holes and K beyond the maximum + // must return false so the caller can use the CPU path. { - const int n = 4, vocab = 4096, big_K = 64; + const int n = 4, vocab = 4096; std::vector h(n * vocab, 0.f); float * d = nullptr; if (cudaMalloc(&d, h.size() * sizeof(float)) == cudaSuccess) { cudaMemcpy(d, h.data(), h.size() * sizeof(float), cudaMemcpyHostToDevice); - std::vector lp(n * big_K); - std::vector ids(n * big_K); - bool ret = geometric_extract_draft_topk_cuda(d, n, vocab, big_K, - lp.data(), ids.data(), 1.0f); + for (int K : {9, 10, 11, 13, 14, 15, 64}) { + std::vector lp((size_t)n * K); + std::vector ids((size_t)n * K); + bool ret = geometric_extract_draft_topk_cuda( + d, n, vocab, K, lp.data(), ids.data(), 1.0f); + const bool pass = !ret; + printf(" [%s] fallback contract: K=%d returned %s\n", + pass ? "PASS" : "FAIL", K, + ret ? "true" : "false"); + if (!pass) failures++; + idx++; + } cudaFree(d); - const bool pass = !ret; // expect false - printf(" [%s] fallback contract: K=%d (>kMaxK) returned %s\n", - pass ? "PASS" : "FAIL", big_K, ret ? "true" : "false"); - if (!pass) failures++; - idx++; } } diff --git a/server/test/test_dspark_batched_head.cpp b/server/test/test_dspark_batched_head.cpp new file mode 100644 index 000000000..64ee5634f --- /dev/null +++ b/server/test/test_dspark_batched_head.cpp @@ -0,0 +1,224 @@ +#include "common/dspark_head.h" +#include "host_check.h" + +#include "ggml-alloc.h" +#include "ggml-backend.h" +#include "ggml-cpu.h" + +#include +#include +#include + +using namespace dflash::common; + +static int g_checks = 0; + +int main() { + constexpr int hidden = 3; + constexpr int vocab = 5; + constexpr int rank = 2; + constexpr int q_len = 4; + constexpr int lanes = 2; + constexpr int confidence_dim = hidden + rank; + + ggml_backend_t backend = ggml_backend_cpu_init(); + CHECK(backend != nullptr); + if (!backend) return 1; + + ggml_init_params weights_params{}; + weights_params.mem_size = ggml_tensor_overhead() * 8; + weights_params.no_alloc = true; + ggml_context * weights_ctx = ggml_init(weights_params); + CHECK(weights_ctx != nullptr); + if (!weights_ctx) { + ggml_backend_free(backend); + return 1; + } + + ggml_tensor * lm_head = ggml_new_tensor_2d( + weights_ctx, GGML_TYPE_F32, hidden, vocab); + ggml_tensor * markov_w1 = ggml_new_tensor_2d( + weights_ctx, GGML_TYPE_F32, rank, vocab); + ggml_tensor * markov_w2 = ggml_new_tensor_2d( + weights_ctx, GGML_TYPE_F32, rank, vocab); + ggml_tensor * confidence_w = ggml_new_tensor_2d( + weights_ctx, GGML_TYPE_F32, confidence_dim, 1); + ggml_tensor * confidence_b = ggml_new_tensor_1d( + weights_ctx, GGML_TYPE_F32, 1); + ggml_backend_buffer_t weights_buf = + ggml_backend_alloc_ctx_tensors(weights_ctx, backend); + CHECK(weights_buf != nullptr); + if (!weights_buf) { + ggml_free(weights_ctx); + ggml_backend_free(backend); + return 1; + } + + std::vector lm((size_t)hidden * vocab); + std::vector w1((size_t)rank * vocab); + std::vector w2((size_t)rank * vocab); + std::vector cw((size_t)confidence_dim); + for (int token = 0; token < vocab; ++token) { + for (int h = 0; h < hidden; ++h) { + lm[(size_t)token * hidden + h] = + 0.031f * (float)(token + 1) * (float)(h + 1); + } + for (int r = 0; r < rank; ++r) { + w1[(size_t)token * rank + r] = + 0.017f * (float)(token + 1 + r); + w2[(size_t)token * rank + r] = + 0.013f * (float)(token + 1) * (float)(r + 1); + } + } + for (int i = 0; i < confidence_dim; ++i) { + cw[(size_t)i] = 0.021f * (float)(i + 1); + } + const float cb = -0.11f; + ggml_backend_tensor_set(lm_head, lm.data(), 0, sizeof(float) * lm.size()); + ggml_backend_tensor_set( + markov_w1, w1.data(), 0, sizeof(float) * w1.size()); + ggml_backend_tensor_set( + markov_w2, w2.data(), 0, sizeof(float) * w2.size()); + ggml_backend_tensor_set( + confidence_w, cw.data(), 0, sizeof(float) * cw.size()); + ggml_backend_tensor_set(confidence_b, &cb, 0, sizeof(cb)); + + DraftWeights dw; + dw.n_embd = hidden; + dw.block_size = q_len; + dw.dspark.enabled = true; + dw.dspark.markov_rank = rank; + dw.dspark.vocab_size = vocab; + dw.dspark.confidence_dim = confidence_dim; + dw.dspark.markov_w1 = markov_w1; + dw.dspark.markov_w2 = markov_w2; + dw.dspark.confidence_w = confidence_w; + dw.dspark.confidence_b = confidence_b; + + std::vector> hidden_host( + lanes, std::vector((size_t)hidden * q_len)); + std::vector> prenorm_host( + lanes, std::vector((size_t)hidden * q_len)); + const int32_t seeds[lanes] = {1, 3}; + for (int lane = 0; lane < lanes; ++lane) { + for (int position = 0; position < q_len; ++position) { + for (int h = 0; h < hidden; ++h) { + const size_t index = + (size_t)position * hidden + h; + hidden_host[(size_t)lane][index] = + 0.07f * (float)(1 + lane + 2 * position + h); + prenorm_host[(size_t)lane][index] = + hidden_host[(size_t)lane][index] + + 0.019f * (float)(h + 1); + } + } + } + + std::vector> serial_tokens(lanes); + std::vector> serial_confidence(lanes); + for (int lane = 0; lane < lanes; ++lane) { + CHECK(dspark_markov_correct_greedy_chain_fused( + dw, backend, lm_head, hidden_host[(size_t)lane].data(), + q_len, seeds[lane], serial_tokens[(size_t)lane], + &serial_confidence[(size_t)lane], + prenorm_host[(size_t)lane].data())); + } + + std::vector arena(4u * 1024u * 1024u); + ggml_init_params graph_params{}; + graph_params.mem_size = arena.size(); + graph_params.mem_buffer = arena.data(); + graph_params.no_alloc = true; + ggml_context * graph_ctx = ggml_init(graph_params); + CHECK(graph_ctx != nullptr); + ggml_cgraph * graph = + ggml_new_graph_custom(graph_ctx, 2048, false); + + std::vector hidden_inputs(lanes); + std::vector prenorm_inputs(lanes); + for (int lane = 0; lane < lanes; ++lane) { + hidden_inputs[(size_t)lane] = ggml_new_tensor_2d( + graph_ctx, GGML_TYPE_F32, hidden, q_len); + prenorm_inputs[(size_t)lane] = ggml_new_tensor_2d( + graph_ctx, GGML_TYPE_F32, hidden, q_len); + ggml_set_input(hidden_inputs[(size_t)lane]); + ggml_set_input(prenorm_inputs[(size_t)lane]); + } + ggml_tensor * seed_input = + ggml_new_tensor_1d(graph_ctx, GGML_TYPE_I32, lanes); + ggml_set_input(seed_input); + + DSparkBatchedChainOutputs outputs; + CHECK(build_dspark_markov_batched_chain( + graph_ctx, graph, dw, lm_head, + hidden_inputs, prenorm_inputs, seed_input, + q_len, true, outputs)); + CHECK(outputs.n_lanes == lanes); + CHECK(outputs.q_len == q_len); + CHECK(outputs.tokens.size() == (size_t)(q_len - 1)); + CHECK(outputs.confidence.size() == (size_t)(q_len - 1)); + + ggml_gallocr_t allocator = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(backend)); + CHECK(allocator != nullptr); + const bool graph_allocated = + allocator && ggml_gallocr_alloc_graph(allocator, graph); + CHECK(graph_allocated); + if (graph_allocated) { + for (int lane = 0; lane < lanes; ++lane) { + ggml_backend_tensor_set( + hidden_inputs[(size_t)lane], + hidden_host[(size_t)lane].data(), 0, + sizeof(float) * hidden_host[(size_t)lane].size()); + ggml_backend_tensor_set( + prenorm_inputs[(size_t)lane], + prenorm_host[(size_t)lane].data(), 0, + sizeof(float) * prenorm_host[(size_t)lane].size()); + } + ggml_backend_tensor_set( + seed_input, seeds, 0, sizeof(seeds)); + CHECK(ggml_backend_graph_compute(backend, graph) == + GGML_STATUS_SUCCESS); + + std::vector depth_tokens( + (size_t)(q_len - 1) * lanes); + std::vector depth_confidence( + (size_t)(q_len - 1) * lanes); + for (int depth = 0; depth < q_len - 1; ++depth) { + ggml_backend_tensor_get_async( + backend, outputs.tokens[(size_t)depth], + depth_tokens.data() + (size_t)depth * lanes, + 0, sizeof(int32_t) * lanes); + ggml_backend_tensor_get_async( + backend, outputs.confidence[(size_t)depth], + depth_confidence.data() + (size_t)depth * lanes, + 0, sizeof(float) * lanes); + } + ggml_backend_synchronize(backend); + + for (int lane = 0; lane < lanes; ++lane) { + CHECK(serial_tokens[(size_t)lane].size() == (size_t)q_len); + CHECK(serial_confidence[(size_t)lane].size() == + (size_t)(q_len - 1)); + CHECK(serial_tokens[(size_t)lane][0] == seeds[lane]); + for (int depth = 0; depth < q_len - 1; ++depth) { + CHECK(serial_tokens[(size_t)lane][(size_t)depth + 1] == + depth_tokens[(size_t)depth * lanes + lane]); + CHECK(std::fabs( + serial_confidence[(size_t)lane][(size_t)depth] - + depth_confidence[ + (size_t)depth * lanes + lane]) < 1e-6f); + } + } + } + if (allocator) ggml_gallocr_free(allocator); + + ggml_free(graph_ctx); + ggml_backend_buffer_free(weights_buf); + ggml_free(weights_ctx); + ggml_backend_free(backend); + std::printf( + "DSpark batched-head parity tests passed: %d checks\n", + g_checks); + return 0; +} diff --git a/server/test/test_feature_gate.cpp b/server/test/test_feature_gate.cpp index d01b9da04..64cbc7f3b 100644 --- a/server/test/test_feature_gate.cpp +++ b/server/test/test_feature_gate.cpp @@ -376,11 +376,72 @@ void test_feature_gate_paged_attention_requires_plain_ar_decode() { BackendArgs draft = base; draft.draft_path = "/nonexistent/draft.gguf"; CHECK(!gate_result(draft, "qwen35", PlacementBackend::Cuda).empty()); + // A local DSpark chain cluster is admitted under concurrent paged + // serving on either GPU backend. Forced AR may carry an unused drafter + // even without concurrent slots. + BackendArgs concurrent_chain = draft; + concurrent_chain.max_concurrency = 16; + CHECK(gate_result( + concurrent_chain, "qwen35", PlacementBackend::Cuda).empty()); + CHECK(gate_result( + concurrent_chain, "qwen35", PlacementBackend::Hip).empty()); + + BackendArgs forced_ar = draft; + forced_ar.speculation_policy = SpeculationPolicy::Never; + CHECK(gate_result( + forced_ar, "qwen35", PlacementBackend::Cuda).empty()); BackendArgs ddtree = base; ddtree.ddtree_mode = true; CHECK(!gate_result(ddtree, "qwen35", PlacementBackend::Cuda).empty()); + BackendArgs concurrent_ddtree = base; + concurrent_ddtree.max_concurrency = 16; + concurrent_ddtree.draft_path = "/nonexistent/draft.gguf"; + concurrent_ddtree.ddtree_mode = true; + concurrent_ddtree.ddtree_budget = 22; + CHECK(gate_result( + concurrent_ddtree, "qwen35", PlacementBackend::Cuda).empty()); + CHECK(gate_result( + concurrent_ddtree, "qwen35", PlacementBackend::Hip).empty()); + + BackendArgs tensor_ddtree = concurrent_ddtree; + CHECK(parse_placement_device_list( + "cuda:0,cuda:1", tensor_ddtree.device)); + tensor_ddtree.device.split_mode = TargetSplitMode::Tensor; + CHECK(!gate_result( + tensor_ddtree, "qwen35", PlacementBackend::Cuda).empty()); + + + BackendFeatureConfig concurrent_pflash; + concurrent_pflash.pflash_enabled = true; + concurrent_pflash.pflash_drafter_configured = true; + CHECK(gate_result(concurrent_ddtree, "qwen35", + PlacementBackend::Hip, concurrent_pflash).empty()); + + BackendArgs concurrent_plain = base; + concurrent_plain.max_concurrency = 16; + CHECK(gate_result(concurrent_plain, "qwen35", + PlacementBackend::Hip, concurrent_pflash).empty()); + BackendFeatureConfig concurrent_kvflash; + concurrent_kvflash.kvflash_enabled = true; + CHECK(gate_result(concurrent_plain, "qwen35", + PlacementBackend::Hip, concurrent_kvflash).empty()); + + BackendArgs bad_budget = concurrent_ddtree; + for (int value : {0, -1, 256, INT_MAX}) { + bad_budget.ddtree_budget = value; + CHECK(!gate_result( + bad_budget, "qwen35", PlacementBackend::Hip).empty()); + } + + BackendArgs remote_ddtree = concurrent_ddtree; + remote_ddtree.remote_draft.ipc_bin = "/usr/bin/draft-ipc"; + remote_ddtree.draft_device.backend = PlacementBackend::Cuda; + remote_ddtree.device.backend = PlacementBackend::Hip; + CHECK(!gate_result( + remote_ddtree, "qwen35", PlacementBackend::Hip).empty()); + BackendArgs windowed = base; windowed.fa_window = 4096; CHECK(!gate_result( @@ -481,6 +542,38 @@ void test_feature_gate_parallel_and_kv_pool_rules() { pool.kv_pool_tokens = max_pool_tokens; CHECK(gate_result(pool, "qwen35", PlacementBackend::Cuda).empty()); + BackendArgs tree_pool = paged; + tree_pool.max_concurrency = 16; + tree_pool.draft_path = "/nonexistent/draft.gguf"; + tree_pool.ddtree_mode = true; + tree_pool.ddtree_budget = 22; + const long long tree_scratch = + (long long)tree_pool.max_concurrency * + paged_token_capacity(tree_pool.ddtree_budget + 1); + const long long max_tree_pool_tokens = + ((long long)INT_MAX - PAGED_BLOCK_SIZE - tree_scratch) / + PAGED_BLOCK_SIZE * PAGED_BLOCK_SIZE; + tree_pool.kv_pool_tokens = max_tree_pool_tokens; + CHECK(gate_result( + tree_pool, "qwen35", PlacementBackend::Cuda).empty()); + tree_pool.kv_pool_tokens = max_tree_pool_tokens + PAGED_BLOCK_SIZE; + CHECK(!gate_result( + tree_pool, "qwen35", PlacementBackend::Cuda).empty()); + BackendArgs chain_pool = paged; + chain_pool.max_concurrency = 16; + chain_pool.draft_path = "/nonexistent/draft.gguf"; + const long long chain_scratch = + (long long)chain_pool.max_concurrency * paged_token_capacity(16); + const long long max_chain_pool_tokens = + ((long long)INT_MAX - PAGED_BLOCK_SIZE - chain_scratch) / + PAGED_BLOCK_SIZE * PAGED_BLOCK_SIZE; + chain_pool.kv_pool_tokens = max_chain_pool_tokens; + CHECK(gate_result( + chain_pool, "qwen35", PlacementBackend::Hip).empty()); + chain_pool.kv_pool_tokens = max_chain_pool_tokens + PAGED_BLOCK_SIZE; + CHECK(!gate_result( + chain_pool, "qwen35", PlacementBackend::Hip).empty()); + // The automatic pool is memory-derived, so a logical slot/context product // larger than the physical tensor address space is legal. BackendArgs overflow = paged; diff --git a/server/test/test_gdn_transition_journal.cpp b/server/test/test_gdn_transition_journal.cpp new file mode 100644 index 000000000..fd25d7fec --- /dev/null +++ b/server/test/test_gdn_transition_journal.cpp @@ -0,0 +1,672 @@ +// Standalone GPU proof for the compact linear-chain GDN transition journal. +// The captured [gate, key, state-dependent delta] tuples must reconstruct +// every accepted prefix without rerunning the target recurrence. +#include "ggml-backend.h" +#include "ggml-cuda.h" +#include "ggml.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +constexpr int S = 128; +constexpr int H = 48; +constexpr int KEY_HEADS = 16; +constexpr int T = 6; +constexpr int B = 4; +constexpr int PHYSICAL_SLOTS = 3; +constexpr float FIELD_TOLERANCE = 5.0e-5f; +constexpr float STATE_TOLERANCE = 2.0e-4f; + +size_t qkv_index(int sequence, int token, int head, int value) { + return (((size_t) sequence*T + token)*H + head)*S + value; +} + +size_t key_index(int sequence, int token, int head, int value) { + return (((size_t) sequence*T + token)*KEY_HEADS + + head%KEY_HEADS)*S + value; +} + +size_t scalar_index(int sequence, int token, int head) { + return ((size_t) sequence*T + token)*H + head; +} + +size_t state_index(int slot, int head, int col, int row) { + return (((size_t) slot*H + head)*S + col)*S + row; +} + +size_t journal_index( + int sequence, int token, int head, int width, int value) { + return ((((size_t) sequence*T + token)*H + head)*width) + value; +} + +float sigmoid(float x) { + return 1.0f/(1.0f + std::exp(-x)); +} + +float softplus(float x) { + return x > 20.0f ? x : std::log1p(std::exp(x)); +} + +bool compare_vectors( + const char * label, + const std::vector & actual, + const std::vector & expected, + float tolerance) { + if (actual.size() != expected.size()) { + std::fprintf(stderr, "%s: size mismatch %zu != %zu\n", label, + actual.size(), expected.size()); + return false; + } + float max_error = 0.0f; + size_t worst = 0; + for (size_t i = 0; i < actual.size(); ++i) { + if (!std::isfinite(actual[i]) || !std::isfinite(expected[i])) { + std::fprintf(stderr, + "%s: non-finite value at %zu (actual %.9g expected %.9g)\n", + label, i, actual[i], expected[i]); + return false; + } + const float error = std::fabs(actual[i] - expected[i]); + if (error > max_error) { + max_error = error; + worst = i; + } + } + if (max_error > tolerance || !std::isfinite(max_error)) { + std::fprintf(stderr, + "%s: max error %.9g at %zu (actual %.9g expected %.9g, tolerance %.9g)\n", + label, max_error, worst, actual[worst], expected[worst], + tolerance); + return false; + } + return true; +} + +struct Inputs { + std::vector q; + std::vector k; + std::vector v; + std::vector g; + std::vector beta; + std::vector state; + std::vector dt_bias; + std::vector gate_A; +}; + +Inputs make_inputs(bool kda, bool raw_gates) { + std::mt19937 rng(20260819 + 17*kda + 31*raw_gates); + std::uniform_real_distribution small(-0.25f, 0.25f); + std::uniform_real_distribution state_dist(-0.06f, 0.06f); + std::uniform_real_distribution gate_dist(0.82f, 0.98f); + std::uniform_real_distribution beta_dist(0.15f, 0.85f); + std::uniform_real_distribution raw_dist(-1.5f, 1.5f); + + Inputs in; + const size_t qkv_elements = (size_t) S*H*T*B; + const size_t qk_elements = (size_t) S*KEY_HEADS*T*B; + in.q.resize(qk_elements); + in.k.resize(qk_elements); + in.v.resize(qkv_elements); + in.g.resize((size_t) (kda ? S : 1)*H*T*B); + in.beta.resize((size_t) H*T*B); + in.state.resize((size_t) S*S*H*B); + in.dt_bias.resize(H); + in.gate_A.resize(H); + + for (float & value : in.q) value = small(rng); + for (float & value : in.v) value = small(rng); + for (float & value : in.state) value = state_dist(rng); + + // The target feeds normalized/shared keys to GDN. Normalize every + // sequence/token/head vector before capture so the proof exercises that + // exact resolved input rather than an arbitrary projection. + for (int sequence = 0; sequence < B; ++sequence) { + for (int token = 0; token < T; ++token) { + for (int head = 0; head < KEY_HEADS; ++head) { + float norm2 = 0.0f; + for (int row = 0; row < S; ++row) { + const float value = small(rng); + in.k[key_index(sequence, token, head, row)] = value; + norm2 += value*value; + } + const float inverse_norm = 1.0f/std::sqrt(norm2); + for (int row = 0; row < S; ++row) { + in.k[key_index(sequence, token, head, row)] *= inverse_norm; + } + } + } + } + + if (raw_gates) { + for (float & value : in.g) value = raw_dist(rng); + for (float & value : in.beta) value = raw_dist(rng); + for (int head = 0; head < H; ++head) { + in.dt_bias[head] = -0.35f + 0.12f*head; + in.gate_A[head] = -0.12f - 0.07f*head; + } + } else { + for (float & value : in.g) value = std::log(gate_dist(rng)); + for (float & value : in.beta) value = beta_dist(rng); + } + return in; +} + +float resolved_gate( + const Inputs & in, bool kda, bool raw_gates, + int sequence, int token, int head, int row) { + if (kda) { + return std::exp(in.g[qkv_index(sequence, token, head, row)]); + } + const float raw_or_log = in.g[scalar_index(sequence, token, head)]; + if (!raw_gates) return std::exp(raw_or_log); + return std::exp( + softplus(raw_or_log + in.dt_bias[head])*in.gate_A[head]); +} + +float resolved_beta( + const Inputs & in, bool raw_gates, + int sequence, int token, int head) { + const float value = in.beta[scalar_index(sequence, token, head)]; + return raw_gates ? sigmoid(value) : value; +} + +std::vector ordinary_recurrence( + const Inputs & in, bool kda, bool raw_gates, + int accepted_prefix) { + std::vector state = in.state; + for (int sequence = 0; sequence < B; ++sequence) { + for (int token = 0; token < accepted_prefix; ++token) { + for (int head = 0; head < H; ++head) { + const float beta = resolved_beta( + in, raw_gates, sequence, token, head); + for (int col = 0; col < S; ++col) { + float projection = 0.0f; + for (int row = 0; row < S; ++row) { + const float gate = resolved_gate( + in, kda, raw_gates, + sequence, token, head, row); + const float state_value = + state[state_index(sequence, head, col, row)]; + const float key = + in.k[key_index(sequence, token, head, row)]; + projection += (kda ? gate : 1.0f)*state_value*key; + } + const float scalar_gate = resolved_gate( + in, kda, raw_gates, + sequence, token, head, 0); + const float delta = + (in.v[qkv_index(sequence, token, head, col)] - + (kda ? projection : scalar_gate*projection))*beta; + for (int row = 0; row < S; ++row) { + const float gate = resolved_gate( + in, kda, raw_gates, + sequence, token, head, row); + const float key = + in.k[key_index(sequence, token, head, row)]; + float & state_value = + state[state_index(sequence, head, col, row)]; + state_value = std::fma(key, delta, gate*state_value); + } + } + } + } + } + return state; +} + +std::vector expected_journal( + const Inputs & in, bool kda, bool raw_gates) { + const int gate_values = kda ? S : 1; + const int width = gate_values + 2*S; + std::vector journal((size_t) width*H*T*B); + std::vector state = in.state; + for (int sequence = 0; sequence < B; ++sequence) { + for (int token = 0; token < T; ++token) { + for (int head = 0; head < H; ++head) { + for (int row = 0; row < gate_values; ++row) { + journal[journal_index(sequence, token, head, width, row)] = + resolved_gate(in, kda, raw_gates, + sequence, token, head, row); + } + for (int row = 0; row < S; ++row) { + journal[journal_index( + sequence, token, head, width, + gate_values + row)] = + in.k[key_index(sequence, token, head, row)]; + } + const float beta = resolved_beta( + in, raw_gates, sequence, token, head); + for (int col = 0; col < S; ++col) { + float projection = 0.0f; + for (int row = 0; row < S; ++row) { + const float gate = resolved_gate( + in, kda, raw_gates, + sequence, token, head, row); + projection += (kda ? gate : 1.0f)* + state[state_index(sequence, head, col, row)]* + in.k[key_index(sequence, token, head, row)]; + } + const float scalar_gate = resolved_gate( + in, kda, raw_gates, + sequence, token, head, 0); + const float delta = + (in.v[qkv_index(sequence, token, head, col)] - + (kda ? projection : scalar_gate*projection))*beta; + journal[journal_index( + sequence, token, head, width, + gate_values + S + col)] = delta; + for (int row = 0; row < S; ++row) { + const float gate = resolved_gate( + in, kda, raw_gates, + sequence, token, head, row); + float & value = + state[state_index(sequence, head, col, row)]; + value = std::fma( + in.k[key_index(sequence, token, head, row)], + delta, gate*value); + } + } + } + } + } + return journal; +} + +struct CaseTensors { + ggml_context * ctx = nullptr; + ggml_backend_buffer_t buffer = nullptr; + ggml_tensor * journal = nullptr; + ggml_tensor * identity_state = nullptr; + ggml_tensor * mapped_state = nullptr; + ggml_tensor * accepted = nullptr; + ggml_tensor * slots = nullptr; +}; + +void destroy(CaseTensors & tensors) { + if (tensors.buffer) ggml_backend_buffer_free(tensors.buffer); + if (tensors.ctx) ggml_free(tensors.ctx); + tensors = {}; +} + +bool run_case( + ggml_backend_t backend, bool kda, bool raw_gates, + bool report_timing) { + const char * name = raw_gates ? "scalar-raw" : kda ? "kda" : "scalar"; + const int gate_values = kda ? S : 1; + const int width = gate_values + 2*S; + const Inputs inputs = make_inputs(kda, raw_gates); + + ggml_init_params params{}; + params.mem_size = 8*1024*1024; + params.no_alloc = true; + CaseTensors tensors; + tensors.ctx = ggml_init(params); + if (!tensors.ctx) return false; + + ggml_tensor * q = ggml_new_tensor_4d( + tensors.ctx, GGML_TYPE_F32, S, KEY_HEADS, T, B); + ggml_tensor * k = ggml_new_tensor_4d( + tensors.ctx, GGML_TYPE_F32, S, KEY_HEADS, T, B); + ggml_tensor * v = ggml_new_tensor_4d( + tensors.ctx, GGML_TYPE_F32, S, H, T, B); + ggml_tensor * g = ggml_new_tensor_4d( + tensors.ctx, GGML_TYPE_F32, kda ? S : 1, H, T, B); + ggml_tensor * beta = ggml_new_tensor_4d( + tensors.ctx, GGML_TYPE_F32, 1, H, T, B); + ggml_tensor * capture_state = ggml_new_tensor_4d( + tensors.ctx, GGML_TYPE_F32, S, S, H, B); + tensors.journal = ggml_new_tensor_4d( + tensors.ctx, GGML_TYPE_F32, width, H, T, B); + tensors.identity_state = ggml_new_tensor_4d( + tensors.ctx, GGML_TYPE_F32, S, S, H, B); + tensors.mapped_state = ggml_new_tensor_4d( + tensors.ctx, GGML_TYPE_F32, S, S, H, PHYSICAL_SLOTS); + tensors.accepted = ggml_new_tensor_1d( + tensors.ctx, GGML_TYPE_I32, B); + tensors.slots = ggml_new_tensor_1d( + tensors.ctx, GGML_TYPE_I32, B); + ggml_tensor * dt_bias = nullptr; + ggml_tensor * gate_A = nullptr; + if (raw_gates) { + dt_bias = ggml_new_tensor_1d(tensors.ctx, GGML_TYPE_F32, H); + gate_A = ggml_new_tensor_1d(tensors.ctx, GGML_TYPE_F32, H); + } + + ggml_tensor * result = ggml_gated_delta_net( + tensors.ctx, q, k, v, g, beta, capture_state); + ggml_gated_delta_net_set_skip_intermediate(result, true); + if (raw_gates) { + ggml_gated_delta_net_set_raw_gates(result, dt_bias, gate_A); + } + ggml_gated_delta_net_set_transition_journal(result, tensors.journal); + ggml_set_output(result); + ggml_cgraph * graph = ggml_new_graph(tensors.ctx); + ggml_build_forward_expand(graph, result); + + tensors.buffer = ggml_backend_alloc_ctx_tensors(tensors.ctx, backend); + if (!tensors.buffer) { + std::fprintf(stderr, "%s: GPU tensor allocation failed\n", name); + destroy(tensors); + return false; + } + auto upload_f32 = [](ggml_tensor * tensor, + const std::vector & values) { + ggml_backend_tensor_set(tensor, values.data(), 0, + values.size()*sizeof(float)); + }; + upload_f32(q, inputs.q); + upload_f32(k, inputs.k); + upload_f32(v, inputs.v); + upload_f32(g, inputs.g); + upload_f32(beta, inputs.beta); + upload_f32(capture_state, inputs.state); + if (raw_gates) { + upload_f32(dt_bias, inputs.dt_bias); + upload_f32(gate_A, inputs.gate_A); + } + + bool ok = ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS; + std::vector actual_journal((size_t) width*H*T*B); + if (ok) { + ggml_backend_tensor_get( + tensors.journal, actual_journal.data(), 0, + actual_journal.size()*sizeof(float)); + ok = compare_vectors( + name, actual_journal, + expected_journal(inputs, kda, raw_gates), FIELD_TOLERANCE); + } + + const std::vector identity_slots{0, 1, 2, 3}; + std::vector accepted(B); + std::vector actual_state(inputs.state.size()); + for (int prefix = 0; ok && prefix <= T; ++prefix) { + std::fill(accepted.begin(), accepted.end(), prefix); + upload_f32(tensors.identity_state, inputs.state); + ggml_backend_tensor_set(tensors.accepted, accepted.data(), 0, + accepted.size()*sizeof(accepted[0])); + ggml_backend_tensor_set(tensors.slots, identity_slots.data(), 0, + identity_slots.size()*sizeof(identity_slots[0])); + ok = ggml_backend_cuda_gdn_transition_journal_commit( + tensors.journal, tensors.identity_state, + tensors.accepted, tensors.slots); + if (ok) { + ggml_backend_tensor_get( + tensors.identity_state, actual_state.data(), 0, + actual_state.size()*sizeof(float)); + char label[64]; + std::snprintf(label, sizeof(label), "%s prefix %d", name, prefix); + ok = compare_vectors( + label, actual_state, + ordinary_recurrence(inputs, kda, raw_gates, prefix), + STATE_TOLERANCE); + } + } + + // Compact lanes {0,2,3} map to physical slots {2,0,1}; lane 1 is bucket + // padding. Each physical base must match the state used to capture that + // lane's state-dependent deltas. + const std::vector mapped_slots{2, -1, 0, 1}; + const std::vector mapped_prefixes{T, T, 2, 4}; + std::vector mapped_base((size_t) S*S*H*PHYSICAL_SLOTS); + for (int sequence : {0, 2, 3}) { + const int slot = mapped_slots[(size_t) sequence]; + for (int head = 0; head < H; ++head) { + for (int col = 0; col < S; ++col) { + for (int row = 0; row < S; ++row) { + mapped_base[state_index(slot, head, col, row)] = + inputs.state[state_index(sequence, head, col, row)]; + } + } + } + } + std::vector mapped_expected = mapped_base; + for (int sequence : {0, 2, 3}) { + const int slot = mapped_slots[(size_t) sequence]; + const std::vector lane_state = ordinary_recurrence( + inputs, kda, raw_gates, mapped_prefixes[(size_t) sequence]); + for (int head = 0; head < H; ++head) { + for (int col = 0; col < S; ++col) { + for (int row = 0; row < S; ++row) { + mapped_expected[state_index(slot, head, col, row)] = + lane_state[state_index(sequence, head, col, row)]; + } + } + } + } + std::vector mapped_actual(mapped_base.size()); + if (ok) { + upload_f32(tensors.mapped_state, mapped_base); + ggml_backend_tensor_set(tensors.accepted, mapped_prefixes.data(), 0, + mapped_prefixes.size()*sizeof(mapped_prefixes[0])); + ggml_backend_tensor_set(tensors.slots, mapped_slots.data(), 0, + mapped_slots.size()*sizeof(mapped_slots[0])); + ok = ggml_backend_cuda_gdn_transition_journal_commit( + tensors.journal, tensors.mapped_state, + tensors.accepted, tensors.slots); + if (ok) { + ggml_backend_tensor_get( + tensors.mapped_state, mapped_actual.data(), 0, + mapped_actual.size()*sizeof(float)); + ok = compare_vectors( + "permuted/padded slots", mapped_actual, mapped_expected, + STATE_TOLERANCE); + } + } + // Out-of-range ids are padding too. + if (ok) { + const std::vector out_of_range_slots{2, 99, 0, 1}; + upload_f32(tensors.mapped_state, mapped_base); + ggml_backend_tensor_set(tensors.slots, out_of_range_slots.data(), 0, + out_of_range_slots.size()*sizeof(out_of_range_slots[0])); + ok = ggml_backend_cuda_gdn_transition_journal_commit( + tensors.journal, tensors.mapped_state, + tensors.accepted, tensors.slots); + if (ok) { + ggml_backend_tensor_get( + tensors.mapped_state, mapped_actual.data(), 0, + mapped_actual.size()*sizeof(float)); + ok = compare_vectors( + "out-of-range padded slot", mapped_actual, mapped_expected, + STATE_TOLERANCE); + } + } + + // Host validation is transactional: malformed prefixes and duplicate live + // slots are rejected before the state kernel can launch. + if (ok) { + const std::vector unchanged = inputs.state; + std::vector invalid_prefix(B, 1); + invalid_prefix[0] = T + 1; + upload_f32(tensors.identity_state, unchanged); + ggml_backend_tensor_set(tensors.accepted, invalid_prefix.data(), 0, + invalid_prefix.size()*sizeof(invalid_prefix[0])); + ggml_backend_tensor_set(tensors.slots, identity_slots.data(), 0, + identity_slots.size()*sizeof(identity_slots[0])); + ok = !ggml_backend_cuda_gdn_transition_journal_commit( + tensors.journal, tensors.identity_state, + tensors.accepted, tensors.slots); + const std::vector duplicate_slots{0, 0, 2, 3}; + accepted.assign(B, 1); + ggml_backend_tensor_set(tensors.accepted, accepted.data(), 0, + accepted.size()*sizeof(accepted[0])); + ggml_backend_tensor_set(tensors.slots, duplicate_slots.data(), 0, + duplicate_slots.size()*sizeof(duplicate_slots[0])); + ok = ok && !ggml_backend_cuda_gdn_transition_journal_commit( + tensors.journal, tensors.identity_state, + tensors.accepted, tensors.slots); + ggml_backend_tensor_get( + tensors.identity_state, actual_state.data(), 0, + actual_state.size()*sizeof(float)); + ok = ok && compare_vectors( + "transactional validation", actual_state, unchanged, 0.0f); + } + + if (ok && report_timing) { + constexpr int repetitions = 25; + std::vector elapsed_us; + elapsed_us.reserve(repetitions); + accepted.assign(B, T); + ggml_backend_tensor_set(tensors.accepted, accepted.data(), 0, + accepted.size()*sizeof(accepted[0])); + ggml_backend_tensor_set(tensors.slots, identity_slots.data(), 0, + identity_slots.size()*sizeof(identity_slots[0])); + for (int repetition = 0; repetition < repetitions; ++repetition) { + upload_f32(tensors.identity_state, inputs.state); + const auto start = std::chrono::steady_clock::now(); + const bool committed = + ggml_backend_cuda_gdn_transition_journal_commit( + tensors.journal, tensors.identity_state, + tensors.accepted, tensors.slots); + const auto stop = std::chrono::steady_clock::now(); + if (!committed) { + ok = false; + break; + } + elapsed_us.push_back( + std::chrono::duration(stop - start).count()); + } + if (ok) { + std::sort(elapsed_us.begin(), elapsed_us.end()); + std::printf( + "gdn journal commit S=%d H=%d T=%d B=%d median %.1f us (synchronous Phase 1)\n", + S, H, T, B, elapsed_us[elapsed_us.size()/2]); + } + } + + std::printf("gdn transition journal %-10s: %s\n", name, + ok ? "PASS" : "FAIL"); + destroy(tensors); + return ok; +} + + +bool run_grouped_tree_case(ggml_backend_t backend) { + const Inputs inputs = make_inputs(/*kda=*/false, /*raw_gates=*/false); + constexpr int width = 2*S + 1; + + ggml_init_params params{}; + params.mem_size = 8*1024*1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + if (!ctx) return false; + + ggml_tensor * q = ggml_new_tensor_4d( + ctx, GGML_TYPE_F32, S, KEY_HEADS, T, B); + ggml_tensor * k = ggml_new_tensor_4d( + ctx, GGML_TYPE_F32, S, KEY_HEADS, T, B); + ggml_tensor * v = ggml_new_tensor_4d( + ctx, GGML_TYPE_F32, S, H, T, B); + ggml_tensor * g = ggml_new_tensor_4d( + ctx, GGML_TYPE_F32, 1, H, T, B); + ggml_tensor * beta = ggml_new_tensor_4d( + ctx, GGML_TYPE_F32, 1, H, T, B); + ggml_tensor * base_state = ggml_new_tensor_4d( + ctx, GGML_TYPE_F32, S, S, H, B); + ggml_tensor * parents = ggml_new_tensor_2d( + ctx, GGML_TYPE_I32, T, B); + ggml_tensor * journal = ggml_new_tensor_4d( + ctx, GGML_TYPE_F32, width, H, T, B); + ggml_tensor * committed_state = ggml_new_tensor_4d( + ctx, GGML_TYPE_F32, S, S, H, B); + ggml_tensor * accepted = ggml_new_tensor_1d( + ctx, GGML_TYPE_I32, B); + ggml_tensor * slots = ggml_new_tensor_1d( + ctx, GGML_TYPE_I32, B); + + ggml_tensor * result = ggml_gated_delta_net_tree( + ctx, q, k, v, g, beta, base_state, parents); + ggml_gated_delta_net_set_transition_journal(result, journal); + ggml_set_output(result); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, result); + + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + if (!buffer) { + ggml_free(ctx); + return false; + } + auto upload_f32 = [](ggml_tensor * tensor, + const std::vector & values) { + ggml_backend_tensor_set( + tensor, values.data(), 0, values.size()*sizeof(float)); + }; + upload_f32(q, inputs.q); + upload_f32(k, inputs.k); + upload_f32(v, inputs.v); + upload_f32(g, inputs.g); + upload_f32(beta, inputs.beta); + upload_f32(base_state, inputs.state); + + std::vector parent_ids((size_t) T*B, -1); + for (int sequence : {0, 2}) { + for (int token = 1; token < T; ++token) { + parent_ids[(size_t) sequence*T + token] = token - 1; + } + } + ggml_backend_tensor_set( + parents, parent_ids.data(), 0, + parent_ids.size()*sizeof(parent_ids[0])); + + bool ok = ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS; + const std::vector prefixes{T, 1, T, 1}; + const std::vector identity_slots{0, 1, 2, 3}; + upload_f32(committed_state, inputs.state); + ggml_backend_tensor_set( + accepted, prefixes.data(), 0, prefixes.size()*sizeof(prefixes[0])); + ggml_backend_tensor_set( + slots, identity_slots.data(), 0, + identity_slots.size()*sizeof(identity_slots[0])); + ok = ok && ggml_backend_cuda_gdn_transition_journal_commit( + journal, committed_state, accepted, slots); + + std::vector expected = inputs.state; + const size_t slot_elements = (size_t) S*S*H; + for (int sequence = 0; sequence < B; ++sequence) { + const std::vector lane = ordinary_recurrence( + inputs, /*kda=*/false, /*raw_gates=*/false, + prefixes[(size_t) sequence]); + const size_t offset = (size_t) sequence*slot_elements; + std::copy_n(lane.begin() + offset, slot_elements, + expected.begin() + offset); + } + std::vector actual(expected.size()); + if (ok) { + ggml_backend_tensor_get( + committed_state, actual.data(), 0, + actual.size()*sizeof(float)); + ok = compare_vectors( + "grouped tree chain/root commit", actual, expected, + STATE_TOLERANCE); + } + + std::printf("gdn grouped tree journal : %s\n", + ok ? "PASS" : "FAIL"); + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + return ok; +} +} // namespace + +int main() { + setenv("DFLASH_GDN_FORCE_GROUPED_COLS", "1", 1); + ggml_backend_t backend = ggml_backend_cuda_init(0); + if (!backend) { + std::fprintf(stderr, "GPU backend unavailable\n"); + return 1; + } + bool ok = run_case(backend, false, false, true); + ok = run_case(backend, true, false, false) && ok; + ok = run_case(backend, false, true, false) && ok; + ok = run_grouped_tree_case(backend) && ok; + ggml_backend_free(backend); + return ok ? 0 : 1; +} diff --git a/server/test/test_paged_attention.cpp b/server/test/test_paged_attention.cpp index 78338f717..3e550c1de 100644 --- a/server/test/test_paged_attention.cpp +++ b/server/test/test_paged_attention.cpp @@ -30,6 +30,13 @@ struct TestCase { bool corrupt_blocks; }; +struct TreeMetadata { + int width; + int scratch_stride; + std::vector parent_ids; + std::vector tree_sizes; +}; + int clamped_seq_len(const TestCase & test_case, int seq) { return std::max( 0, std::min( @@ -51,6 +58,36 @@ bool block_is_valid(int32_t block, int physical_blocks) { return block >= 0 && block < physical_blocks; } +bool tree_visible( + const TreeMetadata & tree, + int tree_seq, + int query_node, + int candidate) { + const int tree_size = tree.tree_sizes[tree_seq]; + if (tree_size < 0 || tree_size > tree.width || + query_node < 0 || query_node >= tree_size || + candidate < 0 || candidate >= tree_size) { + return false; + } + + int current = query_node; + for (int depth = 0; depth < tree_size; ++depth) { + if (current == candidate) { + return true; + } + if (current < 0 || current >= tree_size) { + return false; + } + const int parent = + tree.parent_ids[tree_seq * tree.width + current]; + if (parent == current) { + return false; + } + current = parent; + } + return false; +} + std::vector make_block_table( const TestCase & test_case, int physical_blocks) { @@ -127,7 +164,9 @@ std::vector reference_attention( const std::vector & k, const std::vector & v, const std::vector * active_slot_ids = nullptr, - const std::vector * query_positions = nullptr) { + const std::vector * query_positions = nullptr, + const TreeMetadata * tree = nullptr, + int tree_scratch_base = 0) { std::vector output(q.size(), 0.0f); const float scale = 1.0f / std::sqrt(static_cast(D)); const int q_per_kv = N_HEAD / N_HEAD_KV; @@ -147,50 +186,74 @@ std::vector reference_attention( // cached tokens [0, position]. kv_seq_len = (*query_positions)[seq] + 1; } + const int tree_seq = tree ? seq / tree->width : 0; + const int query_node = tree ? seq % tree->width : -1; + const int tree_size = tree ? tree->tree_sizes[tree_seq] : 0; + if (tree && + (tree_size < 0 || tree_size > tree->width || + query_node >= tree_size)) { + continue; + } + + std::vector physical_rows; + physical_rows.reserve(kv_seq_len + (tree ? tree->width : 0)); + for (int token = 0; token < kv_seq_len; ++token) { + const int block = + block_table[ + physical_seq * test_case.max_blocks + + token / BLOCK_SIZE]; + physical_rows.push_back( + block_is_valid(block, physical_blocks) + ? block * BLOCK_SIZE + token % BLOCK_SIZE + : -1); + } + if (tree) { + for (int candidate = 0; candidate < tree->width; ++candidate) { + physical_rows.push_back( + tree_visible(*tree, tree_seq, query_node, candidate) + ? tree_scratch_base + + physical_seq * tree->scratch_stride + candidate + : -1); + } + } for (int head = 0; head < N_HEAD; ++head) { const int kv_head = head / q_per_kv; const float * q_row = q.data() + (static_cast(head) * n_seq + seq) * D; - std::vector scores(kv_seq_len); + std::vector scores(physical_rows.size(), -INFINITY); float max_score = -INFINITY; - for (int token = 0; token < kv_seq_len; ++token) { - const int block = - block_table[ - physical_seq * test_case.max_blocks + token / BLOCK_SIZE]; - if (!block_is_valid(block, physical_blocks)) { - // Mirrors the kernel: invalid blocks contribute nothing. - scores[token] = -INFINITY; - continue; - } - const int physical = block * BLOCK_SIZE + token % BLOCK_SIZE; + for (size_t row = 0; row < physical_rows.size(); ++row) { + const int physical = physical_rows[row]; + if (physical < 0) continue; const float * k_row = k.data() + (static_cast(kv_head) * pool_tokens + physical) * D; float dot = 0.0f; for (int d = 0; d < D; ++d) dot += q_row[d] * k_row[d]; - scores[token] = dot * scale; - max_score = std::max(max_score, scores[token]); + scores[row] = dot * scale; + max_score = std::max(max_score, scores[row]); } float denominator = 0.0f; for (float & score : scores) { + if (!std::isfinite(score)) { + score = 0.0f; + continue; + } score = std::exp(score - max_score); denominator += score; } float * out_row = output.data() + (static_cast(head) * n_seq + seq) * D; - for (int token = 0; token < kv_seq_len; ++token) { - const int block = - block_table[ - physical_seq * test_case.max_blocks + token / BLOCK_SIZE]; - if (!block_is_valid(block, physical_blocks)) continue; - const int physical = block * BLOCK_SIZE + token % BLOCK_SIZE; + for (size_t row = 0; row < physical_rows.size(); ++row) { + const int physical = physical_rows[row]; + if (physical < 0 || denominator == 0.0f) continue; const float * v_row = v.data() + (static_cast(kv_head) * pool_tokens + physical) * D; - const float probability = scores[token] / denominator; + const float probability = scores[row] / denominator; for (int d = 0; d < D; ++d) { out_row[d] += probability * v_row[d]; } @@ -205,7 +268,8 @@ bool run_case(ggml_backend_t backend, ggml_type k_type, ggml_type v_type, const std::vector * active_slot_ids = nullptr, - const std::vector * query_positions = nullptr) { + const std::vector * query_positions = nullptr, + const TreeMetadata * tree = nullptr) { const int physical_n_seq = static_cast(test_case.kv_seq_lens.size()); const int n_seq = active_slot_ids ? static_cast(active_slot_ids->size()) @@ -214,8 +278,23 @@ bool run_case(ggml_backend_t backend, GGML_ASSERT(!query_positions || (active_slot_ids && query_positions->size() == active_slot_ids->size())); + GGML_ASSERT(!tree || (active_slot_ids && !query_positions)); + if (tree) { + GGML_ASSERT(tree->width > 0); + GGML_ASSERT(tree->scratch_stride >= tree->width); + GGML_ASSERT( + tree->parent_ids.size() == + static_cast(tree->width) * tree->tree_sizes.size()); + GGML_ASSERT( + n_seq == + tree->width * static_cast(tree->tree_sizes.size())); + } const int physical_blocks = count_physical_blocks(test_case); - const int pool_tokens = physical_blocks * BLOCK_SIZE; + const int tree_scratch_base = physical_blocks * BLOCK_SIZE; + const int pool_tokens = tree + ? tree_scratch_base + physical_n_seq * tree->scratch_stride + : tree_scratch_base; + GGML_ASSERT(pool_tokens % BLOCK_SIZE == 0); const std::vector block_table = make_block_table(test_case, physical_blocks); @@ -250,13 +329,26 @@ bool run_case(ggml_backend_t backend, positions = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_seq); ggml_set_input(positions); } + ggml_tensor * parents = nullptr; + ggml_tensor * sizes = nullptr; + if (tree) { + parents = ggml_new_tensor_2d( + ctx, GGML_TYPE_I32, tree->width, tree->tree_sizes.size()); + sizes = ggml_new_tensor_1d( + ctx, GGML_TYPE_I32, tree->tree_sizes.size()); + ggml_set_input(parents); + ggml_set_input(sizes); + } const float scale = 1.0f / std::sqrt(static_cast(D)); const int max_kv_seq_len = *std::max_element( test_case.kv_seq_lens.begin(), test_case.kv_seq_lens.end()); ggml_tensor * output = ggml_paged_attn_ext( ctx, q, k, v, table, kv_seq_lens, active, positions, - scale, BLOCK_SIZE, max_kv_seq_len); + scale, BLOCK_SIZE, max_kv_seq_len, parents, sizes, + tree ? tree->width : 0, + tree ? tree_scratch_base : 0, + tree ? tree->scratch_stride : 0); ggml_set_output(output); ggml_cgraph * graph = ggml_new_graph(ctx); ggml_build_forward_expand(graph, output); @@ -315,6 +407,14 @@ bool run_case(ggml_backend_t backend, positions, query_positions->data(), 0, query_positions->size() * sizeof((*query_positions)[0])); } + if (tree) { + ggml_backend_tensor_set( + parents, tree->parent_ids.data(), 0, + tree->parent_ids.size() * sizeof(tree->parent_ids[0])); + ggml_backend_tensor_set( + sizes, tree->tree_sizes.data(), 0, + tree->tree_sizes.size() * sizeof(tree->tree_sizes[0])); + } ok = ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS; } @@ -327,7 +427,7 @@ bool run_case(ggml_backend_t backend, reference_attention( test_case, block_table, pool_tokens, physical_blocks, q_data, k_reference, v_reference, active_slot_ids, - query_positions); + query_positions, tree, tree ? tree_scratch_base : 0); max_abs_error = 0.0f; for (size_t i = 0; i < actual.size(); ++i) { if (!std::isfinite(actual[i])) { @@ -340,10 +440,11 @@ bool run_case(ggml_backend_t backend, ok = ok && max_abs_error < MAX_ABS_ERROR; } - std::printf("paged attention %-11s K=%-4s V=%-4s active=%s pos=%s max_abs=%.6g %s\n", + std::printf("paged attention %-11s K=%-4s V=%-4s active=%s pos=%s tree=%s max_abs=%.6g %s\n", test_case.name, ggml_type_name(k_type), ggml_type_name(v_type), active_slot_ids ? "yes" : "no", query_positions ? "yes" : "no", + tree ? "yes" : "no", max_abs_error, ok ? "PASS" : "FAIL"); ggml_gallocr_free(allocator); ggml_free(ctx); @@ -373,7 +474,8 @@ bool rejects_unlaunchable_gqa(ggml_backend_t backend) { ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); ggml_tensor * output = ggml_paged_attn_ext( ctx, q, k, v, table, kv_seq_lens, nullptr, nullptr, - 1.0f / std::sqrt(static_cast(D)), BLOCK_SIZE, 1); + 1.0f / std::sqrt(static_cast(D)), BLOCK_SIZE, 1, + nullptr, nullptr, 0, 0, 0); const bool rejected = !ggml_backend_supports_op(backend, output); std::printf("paged attention unlaunchable GQA support %s\n", @@ -404,6 +506,38 @@ void run_paged_attention_case(const TestCase & test_case) { ggml_backend_free(backend); } +void run_tree_case() { + ggml_backend_t backend = ggml_backend_cuda_init(0); + REQUIRE_NOT_NULL(backend); + const TestCase tree_case{"tree", 65, {1025, 17, 257}, false}; + const TreeMetadata tree_metadata{ + 22, 32, + { + -1, 0, 0, 1, 1, 2, 2, 3, 3, 4, 4, + 5, 5, 6, 6, 7, 8, 9, 10, 11, 12, 13, + -1, 0, 0, 1, 1, 2, 2, 3, 4, + -1, -1, -1, -1, -1, -1, -1, + -1, -1, -1, -1, -1, -1, + }, + {22, 9}, + }; + const std::vector tree_slot_ids{ + 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, + 2, 2, 2, 2, 2, 2, 2, 2, + 2, 2, 2, 2, 2, 2, 2, 2, + 2, 2, 2, 2, 2, 2, + }; + CHECK(run_case(backend, tree_case, GGML_TYPE_F16, GGML_TYPE_F16, + &tree_slot_ids, nullptr, &tree_metadata)); + CHECK(run_case(backend, tree_case, GGML_TYPE_Q4_0, GGML_TYPE_Q8_0, + &tree_slot_ids, nullptr, &tree_metadata)); + CHECK(run_case(backend, tree_case, GGML_TYPE_Q8_0, GGML_TYPE_Q4_0, + &tree_slot_ids, nullptr, &tree_metadata)); + ggml_backend_free(backend); +} + void run_active_slot_case(const TestCase & test_case, const std::vector & active_slot_ids) { ggml_backend_t backend = ggml_backend_cuda_init(0); @@ -466,6 +600,10 @@ TEST_CASE(PagedAttention, CompactThreeSlotBucketMatchesReference) { }, {2, 1, 0, -1}); } +TEST_CASE(PagedAttention, PackedTreesMatchReference) { + run_tree_case(); +} + TEST_CASE(PagedAttention, RaggedCausalPositionsMatchReference) { // Interleaved query rows from two sequences attend the paged pool causally // through per-row positions. Sequence 1 spans 65 logical blocks, so its diff --git a/server/test/test_paged_kv_pool.cpp b/server/test/test_paged_kv_pool.cpp index e7fd60c07..a91d5775e 100644 --- a/server/test/test_paged_kv_pool.cpp +++ b/server/test/test_paged_kv_pool.cpp @@ -401,6 +401,83 @@ TEST_CASE(PagedKvPoolFixture, invalid_arguments) { })); } + +TEST_CASE(PagedKvPoolFixture, cold_block_roundtrip_and_release) { + PagedKvPool pool(/*physical_block_count=*/4, + /*max_sequences=*/2, /*block_size=*/16); + const auto first = acquire(pool, 1); + CHECK(pool.append(first, 33).status == PagedKvStatus::Ok); + CHECK(equals(sequence(pool, first).block_table, {0, 1, 2})); + CHECK(pool.free_block_count() == 1); + + uint32_t released = 99; + CHECK(pool.page_out_block(first, 1, released) == PagedKvStatus::Ok); + CHECK(released == 1); + CHECK(pool.free_block_count() == 2); + CHECK(sequence(pool, first).block_table[1] == PAGED_KV_COLD_BLOCK); + uint32_t resident = 0; + uint32_t owned = 0; + CHECK(pool.resident_block_count(first, resident) == PagedKvStatus::Ok); + CHECK(resident == 2); + CHECK(pool.owned_block_count(first, owned) == PagedKvStatus::Ok); + CHECK(owned == 3); // logical appended capacity includes the cold block + + uint32_t unchanged = 123; + CHECK(pool.page_out_block(first, 1, unchanged) == + PagedKvStatus::BlockNotResident); + CHECK(unchanged == 123); + CHECK(pool.page_in_block(first, 99, unchanged) == + PagedKvStatus::LogicalBlockOutOfRange); + CHECK(unchanged == 123); + + uint32_t restored = 99; + CHECK(pool.page_in_block(first, 1, restored) == PagedKvStatus::Ok); + CHECK(restored == 1); + CHECK(equals(sequence(pool, first).block_table, {0, 1, 2})); + CHECK(pool.page_in_block(first, 1, unchanged) == + PagedKvStatus::BlockAlreadyResident); + + CHECK(pool.page_out_block(first, 1, released) == PagedKvStatus::Ok); + // release() skips the sentinel and returns only resident blocks. + CHECK(pool.release(first) == PagedKvStatus::Ok); + CHECK(pool.free_block_count() == 4); +} + +TEST_CASE(PagedKvPoolFixture, append_remaps_cold_partial_head_atomically) { + PagedKvPool pool(/*physical_block_count=*/2, + /*max_sequences=*/2, /*block_size=*/16); + const auto first = acquire(pool, 1); + const auto second = acquire(pool, 2); + CHECK(pool.append(first, 8).status == PagedKvStatus::Ok); + CHECK(pool.append(second, 16).status == PagedKvStatus::Ok); + CHECK(pool.free_block_count() == 0); + + uint32_t released = 99; + CHECK(pool.page_out_block(first, 0, released) == PagedKvStatus::Ok); + CHECK(released == 0); + auto appended = pool.append(first, 4); + CHECK(appended.status == PagedKvStatus::Ok); + CHECK(appended.remapped_cold_blocks.size() == 1); + CHECK(appended.remapped_cold_blocks[0].logical_block == 0); + CHECK(appended.remapped_cold_blocks[0].physical_block == 0); + CHECK(appended.write_slots.front().physical_block == 0); + CHECK(appended.write_slots.front().logical_position == 8); + + // A cold head plus a newly-opened block needs two physical allocations. + CHECK(pool.page_out_block(first, 0, released) == PagedKvStatus::Ok); + const auto before = sequence(pool, first); + appended = pool.append(first, 8); // positions 12..19 + CHECK(appended.status == PagedKvStatus::BlocksExhausted); + const auto after = sequence(pool, first); + CHECK(after.kv_seq_len == before.kv_seq_len); + CHECK(after.block_table == before.block_table); + CHECK(after.block_table[0] == PAGED_KV_COLD_BLOCK); + CHECK(appended.remapped_cold_blocks.empty()); + + pool.reset(); + CHECK(pool.free_block_count() == 2); +} + TEST_CASE(PagedKvPoolFixture, auto_pool_sizing) { PagedKvAutoBudget budget; budget.free_bytes = 10'000; diff --git a/server/test/test_paged_kv_residency.cpp b/server/test/test_paged_kv_residency.cpp new file mode 100644 index 000000000..274c7f1e2 --- /dev/null +++ b/server/test/test_paged_kv_residency.cpp @@ -0,0 +1,660 @@ +// Pure-host tests for multi-sequence paged K/V residency. Transfers use a +// deterministic mock device array but preserve async queue/barrier semantics. + +#define GENERATE_UNIT_TEST_MAIN +#include "CppUnitTestFramework.hpp" +#include "common/concurrency/paged_kv_residency.h" + +#include +#include +#include +#include +#include +#include +#include + +using namespace dflash::common; + +namespace { +struct PagedKvResidencyFixture {}; + +PagedKvSequenceHandle acquire(PagedKvPool & pool, uint64_t request) { + PagedKvSequenceHandle handle; + if (pool.acquire(request, handle) != PagedKvStatus::Ok) { + throw std::runtime_error("acquire failed"); + } + return handle; +} + +PagedKvSequenceSnapshot snapshot(PagedKvPool & pool, + PagedKvSequenceHandle handle) { + PagedKvSequenceSnapshot out; + if (pool.sequence(handle, out) != PagedKvStatus::Ok) { + throw std::runtime_error("sequence failed"); + } + return out; +} + +struct MockTransfers { + struct Pending { + std::function apply; + }; + + explicit MockTransfers(uint32_t blocks, size_t bytes) + : block_bytes(bytes), device((size_t)blocks * bytes, 0) {} + + PagedKvResidencyTransferOps callbacks() { + return { + [this](size_t bytes) -> void * { + if (fail_alloc || bytes != block_bytes) return nullptr; + allocations++; + return new uint8_t[bytes]; + }, + [this](void * ptr) { + frees++; + delete[] static_cast(ptr); + }, + [this](PagedKvSequenceHandle, uint32_t, uint32_t physical, + void * host, size_t bytes) { + if (fail_copy_out || bytes != block_bytes) return false; + pending.push_back({[this, physical, host, bytes] { + std::memcpy(host, &device[(size_t)physical * block_bytes], bytes); + }}); + return !fail_copy_out_after_queue; + }, + [this](PagedKvSequenceHandle, uint32_t, uint32_t physical, + const void * host, size_t bytes) { + if (fail_copy_in || bytes != block_bytes) return false; + pending.push_back({[this, physical, host, bytes] { + std::memcpy(&device[(size_t)physical * block_bytes], host, bytes); + }}); + return !fail_copy_in_after_queue; + }, + [this] { + syncs++; + if (fail_sync) return false; + for (Pending & op : pending) op.apply(); + pending.clear(); + return true; + }, + }; + } + + void fill(uint32_t physical, uint8_t value) { + std::fill_n(&device[(size_t)physical * block_bytes], block_bytes, value); + } + + bool block_is(uint32_t physical, uint8_t value) const { + const auto begin = device.begin() + (size_t)physical * block_bytes; + return std::all_of(begin, begin + block_bytes, + [value](uint8_t byte) { return byte == value; }); + } + + size_t block_bytes; + std::vector device; + std::vector pending; + int allocations = 0; + int frees = 0; + int syncs = 0; + bool fail_alloc = false; + bool fail_copy_out = false; + bool fail_copy_in = false; + bool fail_copy_out_after_queue = false; + bool fail_copy_in_after_queue = false; + bool fail_sync = false; +}; + +PagedKvResidencyConfig config(size_t block_bytes, uint32_t budget, + uint32_t sink = 0, uint32_t tail = 0) { + return {block_bytes, budget, sink, tail}; +} + +} // namespace + +TEST_CASE(PagedKvResidencyFixture, page_roundtrip_is_bit_exact_and_barriered) { + PagedKvPool pool(/*physical blocks=*/3, /*sequences=*/1, /*block size=*/4); + MockTransfers io(3, 32); + PagedKvResidencyManager pager(pool, config(32, 3), io.callbacks()); + const auto handle = acquire(pool, 1); + CHECK(pager.register_sequence(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(handle, 12)); + CHECK(pager.commit_pending_writes(handle) == PagedKvResidencyStatus::Ok); + const auto table = snapshot(pool, handle).block_table; + io.fill(table[0], 0x11); + io.fill(table[1], 0x22); + io.fill(table[2], 0x33); + + CHECK(pager.evict_block(handle, 1) == PagedKvResidencyStatus::Ok); + CHECK(io.syncs == 1); + CHECK(snapshot(pool, handle).block_table[1] == PAGED_KV_COLD_BLOCK); + io.fill(table[1], 0xEE); // recycled device bytes must not affect backing + + CHECK(pager.ensure_resident(handle, {1}) == PagedKvResidencyStatus::Ok); + CHECK(io.syncs == 2); + const uint32_t restored = snapshot(pool, handle).block_table[1]; + CHECK(io.block_is(restored, 0x22)); + CHECK(pager.stats().page_outs == 1); + CHECK(pager.stats().page_ins == 1); + CHECK(pager.stats().resident_blocks == 3); + CHECK(pager.stats().host_bytes == 32); + CHECK(pager.stats().moved_bytes == 64); +} + +TEST_CASE(PagedKvResidencyFixture, fair_share_reclaims_borrowed_pages) { + PagedKvPool pool(/*physical blocks=*/6, /*sequences=*/2, /*block size=*/4); + MockTransfers io(6, 16); + PagedKvResidencyManager pager(pool, config(16, 6), io.callbacks()); + const auto first = acquire(pool, 1); + CHECK(pager.register_sequence(first) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(first, 20)); // borrows five of six pages while alone + CHECK(pager.commit_pending_writes(first) == PagedKvResidencyStatus::Ok); + CHECK(pager.fair_quota(first) == 6); + + const auto second = acquire(pool, 2); + CHECK(pager.register_sequence(second) == PagedKvResidencyStatus::Ok); + CHECK(pager.fair_quota(first) == 3); + CHECK(pager.fair_quota(second) == 3); + CHECK(pager.append(second, 12)); + CHECK(pager.commit_pending_writes(second) == PagedKvResidencyStatus::Ok); + CHECK(pager.stats().page_outs == 2); + CHECK(pager.stats().resident_blocks == 6); + + uint32_t first_resident = 0; + uint32_t second_resident = 0; + CHECK(pool.resident_block_count(first, first_resident) == PagedKvStatus::Ok); + CHECK(pool.resident_block_count(second, second_resident) == PagedKvStatus::Ok); + CHECK(first_resident == 3); + CHECK(second_resident == 3); +} + +TEST_CASE(PagedKvResidencyFixture, sink_and_tail_are_never_auto_evicted) { + PagedKvPool pool(/*physical blocks=*/5, /*sequences=*/2, /*block size=*/4); + MockTransfers io(5, 16); + PagedKvResidencyManager pager( + pool, config(16, 5, /*sink=*/1, /*tail=*/1), io.callbacks()); + const auto first = acquire(pool, 1); + CHECK(pager.register_sequence(first) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(first, 16)); // logical blocks 0..3 + CHECK(pager.commit_pending_writes(first) == PagedKvResidencyStatus::Ok); + const auto second = acquire(pool, 2); + CHECK(pager.register_sequence(second) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(second, 8)); // needs two pages; only one was free + CHECK(pager.commit_pending_writes(second) == PagedKvResidencyStatus::Ok); + + const auto first_table = snapshot(pool, first).block_table; + CHECK(first_table[0] != PAGED_KV_COLD_BLOCK); // sink + CHECK(first_table[3] != PAGED_KV_COLD_BLOCK); // tail + CHECK(first_table[1] == PAGED_KV_COLD_BLOCK || + first_table[2] == PAGED_KV_COLD_BLOCK); + CHECK(pager.evict_block(first, 0) == + PagedKvResidencyStatus::NoEvictableBlock); +} + +TEST_CASE(PagedKvResidencyFixture, append_restores_cold_partial_head) { + PagedKvPool pool(/*physical blocks=*/2, /*sequences=*/1, /*block size=*/4); + MockTransfers io(2, 16); + PagedKvResidencyManager pager(pool, config(16, 2), io.callbacks()); + const auto handle = acquire(pool, 1); + CHECK(pager.register_sequence(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(handle, 2)); + CHECK(pager.commit_pending_writes(handle) == PagedKvResidencyStatus::Ok); + const uint32_t head = snapshot(pool, handle).block_table[0]; + io.fill(head, 0xA5); + CHECK(pager.evict_block(handle, 0) == PagedKvResidencyStatus::Ok); + io.fill(head, 0x00); + + const auto append = pager.append(handle, 1); + CHECK(append); + CHECK(append.pool_result.remapped_cold_blocks.empty()); + CHECK(append.pool_result.write_slots[0].logical_position == 2); + CHECK(io.block_is(append.pool_result.write_slots[0].physical_block, 0xA5)); + CHECK(pager.stats().page_ins == 1); + CHECK(pager.commit_pending_writes(handle) == PagedKvResidencyStatus::Ok); +} + +TEST_CASE(PagedKvResidencyFixture, scores_drive_reselection) { + PagedKvPool pool(/*physical blocks=*/3, /*sequences=*/2, /*block size=*/4); + MockTransfers io(3, 16); + PagedKvResidencyManager pager(pool, config(16, 3), io.callbacks()); + const auto handle = acquire(pool, 1); + CHECK(pager.register_sequence(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(handle, 12)); + CHECK(pager.commit_pending_writes(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.evict_block(handle, 2) == PagedKvResidencyStatus::Ok); + CHECK(pager.evict_block(handle, 0) == PagedKvResidencyStatus::Ok); + CHECK(pager.stats().resident_blocks == 1); + const auto peer = acquire(pool, 2); + CHECK(pager.register_sequence(peer) == PagedKvResidencyStatus::Ok); + CHECK(pager.fair_quota(handle) == 2); + CHECK(pager.set_scores(handle, {1.0f, 2.0f, 9.0f}) == + PagedKvResidencyStatus::Ok); + CHECK(pager.reselect(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.stats().reselects == 1); + CHECK(pager.is_resident(handle, 2)); + CHECK(pager.is_resident(handle, 1)); + CHECK(!pager.is_resident(handle, 0)); +} + +TEST_CASE(PagedKvResidencyFixture, allocation_and_copy_failures_are_explicit) { + PagedKvPool pool(/*physical blocks=*/2, /*sequences=*/1, /*block size=*/4); + MockTransfers io(2, 16); + PagedKvResidencyManager pager(pool, config(16, 2), io.callbacks()); + const auto handle = acquire(pool, 1); + CHECK(pager.register_sequence(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(handle, 8)); + CHECK(pager.commit_pending_writes(handle) == PagedKvResidencyStatus::Ok); + + io.fail_alloc = true; + CHECK(pager.evict_block(handle, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::HostAllocationFailed); + CHECK(pager.is_resident(handle, 0)); + io.fail_alloc = false; + io.fail_copy_out = true; + CHECK(pager.evict_block(handle, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::TransferFailed); + CHECK(pager.is_resident(handle, 0)); +} + +TEST_CASE(PagedKvResidencyFixture, + failed_copy_out_prefix_is_barriered_before_retry) { + PagedKvPool pool(/*physical blocks=*/1, /*sequences=*/1, /*block size=*/4); + MockTransfers io(1, 16); + PagedKvResidencyManager pager(pool, config(16, 1), io.callbacks()); + const auto handle = acquire(pool, 1); + CHECK(pager.register_sequence(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(handle, 4)); + CHECK(pager.commit_pending_writes(handle) == PagedKvResidencyStatus::Ok); + + io.fail_copy_out_after_queue = true; + CHECK(pager.evict_block(handle, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::TransferFailed); + CHECK(io.pending.empty()); + CHECK(pager.is_resident(handle, 0)); + CHECK(pager.stats().page_outs == 0); + + io.fail_copy_out_after_queue = false; + CHECK(pager.evict_block(handle, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::Ok); + CHECK(!pager.is_resident(handle, 0)); +} + +TEST_CASE(PagedKvResidencyFixture, + rejected_append_remap_is_rolled_back_after_barrier) { + PagedKvPool pool(/*physical blocks=*/1, /*sequences=*/1, /*block size=*/4); + MockTransfers io(1, 16); + PagedKvResidencyManager pager(pool, config(16, 1), io.callbacks()); + const auto handle = acquire(pool, 1); + CHECK(pager.register_sequence(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(handle, 2)); + CHECK(pager.commit_pending_writes(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.evict_block(handle, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::Ok); + + io.fail_copy_in = true; + const auto append = pool.append(handle, 1); + CHECK(append.status == PagedKvStatus::Ok); + CHECK(append.remapped_cold_blocks.size() == 1); + CHECK(pager.observe_append(handle, append) == + PagedKvResidencyStatus::TransferFailed); + CHECK(snapshot(pool, handle).block_table[0] == PAGED_KV_COLD_BLOCK); + CHECK(pool.free_block_count() == 1); + CHECK(pager.stats().page_ins == 0); +} + +TEST_CASE(PagedKvResidencyFixture, + failed_copy_in_prefix_quarantines_mapping_until_barrier) { + PagedKvPool pool(/*physical blocks=*/1, /*sequences=*/1, /*block size=*/4); + MockTransfers io(1, 16); + PagedKvResidencyManager pager(pool, config(16, 1), io.callbacks()); + const auto handle = acquire(pool, 1); + CHECK(pager.register_sequence(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(handle, 4)); + CHECK(pager.commit_pending_writes(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.evict_block(handle, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::Ok); + + io.fail_copy_in_after_queue = true; + io.fail_sync = true; + CHECK(pager.ensure_resident(handle, {0}) == + PagedKvResidencyStatus::TransferFailed); + CHECK(snapshot(pool, handle).block_table[0] != PAGED_KV_COLD_BLOCK); + CHECK(pool.free_block_count() == 0); + CHECK(io.pending.size() == 1); + + io.fail_sync = false; + CHECK(pager.synchronize_before_read() == PagedKvResidencyStatus::Ok); + CHECK(snapshot(pool, handle).block_table[0] == PAGED_KV_COLD_BLOCK); + CHECK(pool.free_block_count() == 1); + CHECK(pager.stats().page_ins == 0); +} + +TEST_CASE(PagedKvResidencyFixture, forget_frees_host_backing_and_stale_is_rejected) { + PagedKvPool pool(/*physical blocks=*/2, /*sequences=*/1, /*block size=*/4); + MockTransfers io(2, 16); + PagedKvResidencyManager pager(pool, config(16, 2), io.callbacks()); + const auto old_handle = acquire(pool, 1); + CHECK(pager.register_sequence(old_handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(old_handle, 8)); + CHECK(pager.commit_pending_writes(old_handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.evict_block(old_handle, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::Ok); + CHECK(io.allocations == 1); + CHECK(pager.forget_sequence(old_handle) == PagedKvResidencyStatus::Ok); + CHECK(io.frees == 1); + CHECK(pool.release(old_handle) == PagedKvStatus::Ok); + + const auto replacement = acquire(pool, 2); + CHECK(replacement.generation != old_handle.generation); + CHECK(pager.register_sequence(replacement) == PagedKvResidencyStatus::Ok); + CHECK(pager.touch(old_handle, 0) == PagedKvResidencyStatus::StaleHandle); +} + +TEST_CASE(PagedKvResidencyFixture, + forget_retries_failed_barrier_before_releasing_backing) { + PagedKvPool pool(/*physical blocks=*/1, /*sequences=*/1, /*block size=*/4); + MockTransfers io(1, 16); + PagedKvResidencyManager pager(pool, config(16, 1), io.callbacks()); + const auto handle = acquire(pool, 1); + CHECK(pager.register_sequence(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(handle, 4)); + CHECK(pager.commit_pending_writes(handle) == PagedKvResidencyStatus::Ok); + + io.fail_sync = true; + CHECK(pager.evict_block(handle, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::TransferFailed); + CHECK(pool.free_block_count() == 0); + CHECK(io.frees == 0); + CHECK(pager.forget_sequence(handle) == + PagedKvResidencyStatus::TransferFailed); + CHECK(pool.free_block_count() == 0); + CHECK(io.frees == 0); + + // Teardown retries the quarantined barrier instead of being rejected by + // the manager-wide failure latch. + io.fail_sync = false; + CHECK(pager.forget_sequence(handle) == PagedKvResidencyStatus::Ok); + CHECK(io.frees == 1); + CHECK(pool.release(handle) == PagedKvStatus::Ok); +} + +TEST_CASE(PagedKvResidencyFixture, + invalid_restore_request_does_not_leak_reservations) { + PagedKvPool pool(/*physical blocks=*/1, /*sequences=*/2, /*block size=*/4); + MockTransfers io(1, 16); + PagedKvResidencyManager pager(pool, config(16, 1), io.callbacks()); + const auto first = acquire(pool, 1); + const auto second = acquire(pool, 2); + CHECK(pager.register_sequence(first) == PagedKvResidencyStatus::Ok); + CHECK(pager.register_sequence(second) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(first, 4)); + CHECK(pager.commit_pending_writes(first) == PagedKvResidencyStatus::Ok); + + CHECK(pager.ensure_resident(first, {0, 1}) == + PagedKvResidencyStatus::InvalidArgument); + + // The valid prefix remains evictable, allowing the peer to make progress. + CHECK(pager.append(second, 4)); + CHECK(snapshot(pool, first).block_table[0] == PAGED_KV_COLD_BLOCK); +} + +TEST_CASE(PagedKvResidencyFixture, + staged_writes_are_never_recycled_across_slots) { + PagedKvPool pool(/*physical blocks=*/3, /*sequences=*/2, /*block size=*/4); + MockTransfers io(3, 16); + PagedKvResidencyManager pager(pool, config(16, 3), io.callbacks()); + const auto first = acquire(pool, 1); + const auto second = acquire(pool, 2); + CHECK(pager.register_sequence(first) == PagedKvResidencyStatus::Ok); + CHECK(pager.register_sequence(second) == PagedKvResidencyStatus::Ok); + + CHECK(pager.append(first, 8)); + CHECK(pager.append(second, 4)); + const auto first_staged = snapshot(pool, first).block_table; + const auto second_staged = snapshot(pool, second).block_table; + CHECK(first_staged.size() == 2); + CHECK(second_staged.size() == 1); + + const auto blocked = pager.append(first, 4); + CHECK(blocked.status == PagedKvResidencyStatus::NoEvictableBlock); + CHECK(pager.stats().page_outs == 0); + CHECK(pager.evict_block(second, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::NoEvictableBlock); + CHECK(pager.reselect(first) == PagedKvResidencyStatus::Ok); + CHECK(pager.stats().page_outs == 0); + CHECK(snapshot(pool, first).block_table == first_staged); + CHECK(snapshot(pool, second).block_table == second_staged); + + // Simulate one synchronized packed target graph, then make room. The still + // pending peer page cannot be the victim even though another slot grows. + CHECK(pager.commit_pending_writes(first) == PagedKvResidencyStatus::Ok); + const auto grown = pager.append(first, 4); + CHECK(grown); + CHECK(pager.stats().page_outs == 1); + CHECK(snapshot(pool, second).block_table == second_staged); + CHECK(pager.commit_pending_writes(second) == PagedKvResidencyStatus::Ok); + CHECK(pager.commit_pending_writes(first) == PagedKvResidencyStatus::Ok); +} + +TEST_CASE(PagedKvResidencyFixture, + cross_slot_eviction_is_visible_in_full_victim_snapshot) { + PagedKvPool pool(/*physical blocks=*/4, /*sequences=*/2, /*block size=*/4); + MockTransfers io(4, 16); + PagedKvResidencyManager pager(pool, config(16, 4), io.callbacks()); + const auto first = acquire(pool, 1); + CHECK(pager.register_sequence(first) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(first, 12)); + CHECK(pager.commit_pending_writes(first) == PagedKvResidencyStatus::Ok); + const auto before = snapshot(pool, first).block_table; + + const auto second = acquire(pool, 2); + CHECK(pager.register_sequence(second) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(second, 8)); + const auto victim = snapshot(pool, first).block_table; + const auto requester = snapshot(pool, second).block_table; + CHECK(victim.size() == before.size()); + CHECK(requester.size() == 2); + + size_t cold = victim.size(); + for (size_t logical = 0; logical < victim.size(); ++logical) { + if (victim[logical] == PAGED_KV_COLD_BLOCK) { + CHECK(cold == victim.size()); + cold = logical; + } + } + CHECK(cold < victim.size()); + CHECK(std::find(requester.begin(), requester.end(), before[cold]) != + requester.end()); + CHECK(pager.stats().page_outs == 1); + CHECK(pager.commit_pending_writes(second) == PagedKvResidencyStatus::Ok); +} + +TEST_CASE(PagedKvResidencyFixture, + sixteen_slots_progress_with_adaptive_fair_quotas) { + constexpr uint32_t kSlots = 16; + constexpr uint32_t kPoolBlocks = 96; + PagedKvPool pool(kPoolBlocks, kSlots, /*block size=*/4); + MockTransfers io(kPoolBlocks, 16); + PagedKvResidencyManager pager( + pool, config(16, kPoolBlocks, /*sink=*/1, /*tail=*/4), + io.callbacks()); + + std::vector handles; + handles.reserve(kSlots); + for (uint32_t slot = 0; slot < kSlots; ++slot) { + handles.push_back(acquire(pool, slot + 1)); + CHECK(pager.register_sequence(handles.back()) == + PagedKvResidencyStatus::Ok); + } + for (const auto handle : handles) { + CHECK(pager.append(handle, 4)); + CHECK(pager.commit_pending_writes(handle) == + PagedKvResidencyStatus::Ok); + } + for (const auto handle : handles) { + CHECK(pager.append(handle, 28)); // eight logical blocks total + CHECK(pager.commit_pending_writes(handle) == + PagedKvResidencyStatus::Ok); + } + + CHECK(pager.stats().resident_blocks == kPoolBlocks); + CHECK(pager.stats().page_outs == kSlots * 2); + for (const auto handle : handles) { + CHECK(pager.fair_quota(handle) == kPoolBlocks / kSlots); + const auto table = snapshot(pool, handle).block_table; + CHECK(table.size() == 8); + CHECK(table[0] != PAGED_KV_COLD_BLOCK); + for (size_t logical = 4; logical < 8; ++logical) { + CHECK(table[logical] != PAGED_KV_COLD_BLOCK); + } + } +} + +TEST_CASE(PagedKvResidencyFixture, + failed_copy_out_barrier_retains_the_device_mapping) { + PagedKvPool pool(/*physical blocks=*/2, /*sequences=*/1, /*block size=*/4); + MockTransfers io(2, 16); + PagedKvResidencyManager pager(pool, config(16, 2), io.callbacks()); + const auto handle = acquire(pool, 1); + CHECK(pager.register_sequence(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(handle, 8)); + CHECK(pager.commit_pending_writes(handle) == PagedKvResidencyStatus::Ok); + const auto before = snapshot(pool, handle).block_table; + + io.fail_sync = true; + CHECK(pager.evict_block(handle, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::TransferFailed); + CHECK(snapshot(pool, handle).block_table == before); + CHECK(pool.free_block_count() == 0); + CHECK(pager.stats().page_outs == 0); +} + +TEST_CASE(PagedKvResidencyFixture, + failed_copy_in_barrier_quarantines_mapping_until_stream_drains) { + PagedKvPool pool(/*physical blocks=*/1, /*sequences=*/1, /*block size=*/4); + MockTransfers io(1, 16); + PagedKvResidencyManager pager(pool, config(16, 1), io.callbacks()); + const auto handle = acquire(pool, 1); + CHECK(pager.register_sequence(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(handle, 4)); + CHECK(pager.commit_pending_writes(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.evict_block(handle, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::Ok); + + io.fail_sync = true; + CHECK(pager.ensure_resident(handle, {0}) == + PagedKvResidencyStatus::TransferFailed); + const uint32_t quarantined = snapshot(pool, handle).block_table[0]; + CHECK(quarantined != PAGED_KV_COLD_BLOCK); + CHECK(pool.free_block_count() == 0); + CHECK(pager.stats().page_ins == 0); + + // No operation may reuse or mutate the quarantined destination while the + // failed stream is not known to be drained. + CHECK(pager.evict_block(handle, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::TransferFailed); + CHECK(io.pending.size() == 1); + + io.fail_sync = false; + CHECK(pager.synchronize_before_read() == PagedKvResidencyStatus::Ok); + CHECK(snapshot(pool, handle).block_table[0] == quarantined); + CHECK(pool.free_block_count() == 0); + CHECK(pager.stats().page_ins == 1); +} + +TEST_CASE(PagedKvResidencyFixture, + rebalance_evicts_only_blocks_above_fair_quota) { + PagedKvPool pool(/*physical blocks=*/6, /*sequences=*/2, /*block size=*/4); + MockTransfers io(6, 16); + PagedKvResidencyManager pager(pool, config(16, 6), io.callbacks()); + const auto borrower = acquire(pool, 1); + CHECK(pager.register_sequence(borrower) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(borrower, 20)); + CHECK(pager.commit_pending_writes(borrower) == + PagedKvResidencyStatus::Ok); + + const auto peer = acquire(pool, 2); + CHECK(pager.register_sequence(peer) == PagedKvResidencyStatus::Ok); + CHECK(pager.fair_quota(borrower) == 3); + CHECK(pager.rebalance() == PagedKvResidencyStatus::Ok); + + uint32_t borrower_resident = 0; + CHECK(pool.resident_block_count(borrower, borrower_resident) == + PagedKvStatus::Ok); + CHECK(borrower_resident == 3); + CHECK(pager.stats().page_outs == 2); +} + +TEST_CASE(PagedKvResidencyFixture, + append_reserves_partial_head_and_new_block_together) { + PagedKvPool pool(/*physical blocks=*/2, /*sequences=*/2, /*block size=*/4); + MockTransfers io(2, 16); + PagedKvResidencyManager pager(pool, config(16, 2), io.callbacks()); + const auto first = acquire(pool, 1); + const auto second = acquire(pool, 2); + CHECK(pager.register_sequence(first) == PagedKvResidencyStatus::Ok); + CHECK(pager.register_sequence(second) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(first, 2)); + CHECK(pager.commit_pending_writes(first) == PagedKvResidencyStatus::Ok); + CHECK(pager.evict_block(first, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::Ok); + CHECK(pager.append(second, 4)); + CHECK(pager.commit_pending_writes(second) == PagedKvResidencyStatus::Ok); + + const auto grown = pager.append(first, 3); + CHECK(grown); + const auto table = snapshot(pool, first).block_table; + CHECK(table.size() == 2); + CHECK(table[0] != PAGED_KV_COLD_BLOCK); + CHECK(table[1] != PAGED_KV_COLD_BLOCK); + CHECK(snapshot(pool, second).block_table[0] == PAGED_KV_COLD_BLOCK); +} + +TEST_CASE(PagedKvResidencyFixture, + resident_partial_head_stays_within_budget_during_growth) { + PagedKvPool pool(/*physical blocks=*/2, /*sequences=*/2, /*block size=*/4); + MockTransfers io(2, 16); + PagedKvResidencyManager pager( + pool, config(16, /*budget=*/2, /*sink=*/0, /*tail=*/0), + io.callbacks()); + const auto first = acquire(pool, 1); + const auto second = acquire(pool, 2); + CHECK(pager.register_sequence(first) == PagedKvResidencyStatus::Ok); + CHECK(pager.register_sequence(second) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(first, 2)); + CHECK(pager.commit_pending_writes(first) == PagedKvResidencyStatus::Ok); + const uint32_t append_head = snapshot(pool, first).block_table[0]; + CHECK(pager.append(second, 4)); + CHECK(pager.commit_pending_writes(second) == PagedKvResidencyStatus::Ok); + + const auto grown = pager.append(first, 3); + CHECK(grown); + const auto table = snapshot(pool, first).block_table; + CHECK(table.size() == 2); + CHECK(table[0] == append_head); + CHECK(table[1] != PAGED_KV_COLD_BLOCK); + CHECK(snapshot(pool, second).block_table[0] == PAGED_KV_COLD_BLOCK); + CHECK(pager.stats().resident_blocks == 2); +} +TEST_CASE(PagedKvResidencyFixture, + requested_restore_set_is_protected_as_one_batch) { + PagedKvPool pool(/*physical blocks=*/3, /*sequences=*/2, /*block size=*/4); + MockTransfers io(3, 16); + PagedKvResidencyManager pager(pool, config(16, 3), io.callbacks()); + const auto first = acquire(pool, 1); + const auto second = acquire(pool, 2); + CHECK(pager.register_sequence(first) == PagedKvResidencyStatus::Ok); + CHECK(pager.register_sequence(second) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(first, 8)); + CHECK(pager.commit_pending_writes(first) == PagedKvResidencyStatus::Ok); + CHECK(pager.evict_block(first, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::Ok); + CHECK(pager.evict_block(first, 1, /*allow_protected=*/true) == + PagedKvResidencyStatus::Ok); + CHECK(pager.append(second, 8)); + CHECK(pager.commit_pending_writes(second) == PagedKvResidencyStatus::Ok); + + CHECK(pager.ensure_resident(first, {0, 1}) == + PagedKvResidencyStatus::Ok); + CHECK(pager.is_resident(first, 0)); + CHECK(pager.is_resident(first, 1)); +} diff --git a/server/test/test_qwen_paged_kv_transfer_layout.cpp b/server/test/test_qwen_paged_kv_transfer_layout.cpp new file mode 100644 index 000000000..59c9d0e7a --- /dev/null +++ b/server/test/test_qwen_paged_kv_transfer_layout.cpp @@ -0,0 +1,86 @@ +#define GENERATE_UNIT_TEST_MAIN +#include "CppUnitTestFramework.hpp" +#include "common/concurrency/qwen_paged_kv_transfer.h" + +#include +#include +#include + +using namespace dflash::common; + +namespace { +struct QwenPagedKvTransferLayoutFixture {}; + +QwenPagedKvTensorLayout packed(size_t row_bytes, uint64_t rows, + uint64_t heads) { + return { + row_bytes, + row_bytes, + row_bytes * static_cast(rows), + row_bytes * static_cast(rows) * + static_cast(heads), + rows, + heads, + }; +} +} // namespace + +TEST_CASE(QwenPagedKvTransferLayoutFixture, + packs_mixed_kv_types_and_all_heads) { + const std::vector tensors = { + packed(8, 64, 2), + packed(4, 64, 2), + packed(6, 64, 2), + packed(2, 64, 2), + }; + QwenPagedKvBlockLayout plan; + std::string error; + CHECK(plan_qwen_paged_kv_block_layout( + tensors, /*block_size=*/16, plan, &error)); + CHECK(error.empty()); + CHECK(plan.physical_block_count == 4); + CHECK(plan.tensor_offsets == + std::vector({0, 256, 384, 576})); + CHECK(plan.tensor_head_bytes == + std::vector({128, 64, 96, 32})); + CHECK(plan.block_bytes == 640); +} + +TEST_CASE(QwenPagedKvTransferLayoutFixture, + accepts_row_and_head_padding_without_storing_padding) { + QwenPagedKvTensorLayout tensor; + tensor.row_bytes = 6; + tensor.row_stride = 8; + tensor.head_stride = 520; + tensor.physical_rows = 64; + tensor.heads = 2; + tensor.storage_bytes = 520 + 63 * 8 + 6; + + QwenPagedKvBlockLayout plan; + CHECK(plan_qwen_paged_kv_block_layout( + {tensor}, /*block_size=*/16, plan)); + CHECK(plan.tensor_head_bytes == std::vector({96})); + CHECK(plan.block_bytes == 192); +} + +TEST_CASE(QwenPagedKvTransferLayoutFixture, + rejects_mismatched_rows_short_storage_and_overflow) { + QwenPagedKvBlockLayout plan; + std::string error; + CHECK(!plan_qwen_paged_kv_block_layout( + {packed(8, 64, 2), packed(8, 32, 2)}, 16, plan, &error)); + CHECK(!error.empty()); + + QwenPagedKvTensorLayout short_tensor = packed(8, 64, 2); + short_tensor.storage_bytes--; + CHECK(!plan_qwen_paged_kv_block_layout( + {short_tensor}, 16, plan, &error)); + + QwenPagedKvTensorLayout overflow = packed(8, 64, 2); + overflow.row_bytes = std::numeric_limits::max(); + overflow.row_stride = overflow.row_bytes; + overflow.head_stride = overflow.row_bytes; + overflow.storage_bytes = overflow.row_bytes; + CHECK(!plan_qwen_paged_kv_block_layout( + {overflow}, 16, plan, &error)); +} diff --git a/server/test/test_recurrent_snapshot.cpp b/server/test/test_recurrent_snapshot.cpp index 6fb5de4d6..41b38aea8 100644 --- a/server/test/test_recurrent_snapshot.cpp +++ b/server/test/test_recurrent_snapshot.cpp @@ -1,5 +1,6 @@ #include "CppUnitTestFramework.hpp" #include "internal.h" +#include "qwen35/graph_builders.h" #include "ggml-backend.h" #include "ggml-cpu.h" @@ -11,6 +12,7 @@ using namespace CppUnitTestFramework; using dflash::common::TargetCache; +using dflash::common::StepGraph; using dflash::common::restore_ssm_state; using dflash::common::snapshot_ssm_state; @@ -32,11 +34,159 @@ static std::vector get_tensor(const ggml_tensor * tensor) { return values; } +TEST_CASE(RecurrentSnapshotFixture, hardens_feature_smoke_paths) { + size_t graph_capacity = 0; + CHECK(dflash::common::detail:: + target_graph_capacity_for_parallel_segments( + 0, graph_capacity) && graph_capacity == 16384); + CHECK(dflash::common::detail:: + target_graph_capacity_for_parallel_segments( + 8, graph_capacity) && graph_capacity == 16384); + CHECK(dflash::common::detail:: + target_graph_capacity_for_parallel_segments( + 16, graph_capacity) && graph_capacity == 32768); + CHECK(dflash::common::detail:: + target_graph_capacity_for_parallel_segments( + 64, graph_capacity) && graph_capacity == 131072); + CHECK(!dflash::common::detail:: + target_graph_capacity_for_parallel_segments( + 65, graph_capacity)); + CHECK(dflash::common::detail::target_paged_tree_graph_capacity( + 23, 16, graph_capacity) && graph_capacity == 32768); + CHECK(!dflash::common::detail::target_paged_tree_graph_capacity( + 257, 16, graph_capacity)); + + // Mapped-tree active_slot_ids is a topology marker and can legitimately + // be left without gallocr storage. Required state/query/write metadata + // remains allocated and uploadable. + { + ggml_backend_t tree_backend = ggml_backend_cpu_init(); + CHECK(tree_backend != nullptr); + ggml_init_params marker_params{}; + marker_params.mem_size = 4 * ggml_tensor_overhead(); + marker_params.no_alloc = true; + ggml_context * marker_ctx = ggml_init(marker_params); + ggml_init_params live_params{}; + live_params.mem_size = 16 * ggml_tensor_overhead(); + live_params.no_alloc = true; + ggml_context * live_ctx = ggml_init(live_params); + CHECK(marker_ctx != nullptr); + CHECK(live_ctx != nullptr); + if (tree_backend && marker_ctx && live_ctx) { + StepGraph tree; + tree.active_slot_ids = + ggml_new_tensor_1d(marker_ctx, GGML_TYPE_I32, 2); + ggml_tensor * unallocated_state_ids = + ggml_new_tensor_1d(marker_ctx, GGML_TYPE_I32, 2); + tree.inp_embed = ggml_new_tensor_2d( + live_ctx, GGML_TYPE_F32, 4, 4); + tree.positions = + ggml_new_tensor_1d(live_ctx, GGML_TYPE_I32, 16); + tree.parent_ids = + ggml_new_tensor_2d(live_ctx, GGML_TYPE_I32, 2, 2); + tree.tree_sizes = + ggml_new_tensor_1d(live_ctx, GGML_TYPE_I32, 2); + tree.state_slot_ids = + ggml_new_tensor_1d(live_ctx, GGML_TYPE_I32, 2); + tree.paged_query_seq_ids = + ggml_new_tensor_1d(live_ctx, GGML_TYPE_I32, 4); + tree.kv_write_rows = + ggml_new_tensor_2d(live_ctx, GGML_TYPE_I64, 4, 1); + ggml_backend_buffer_t live_buffer = + ggml_backend_alloc_ctx_tensors(live_ctx, tree_backend); + CHECK(live_buffer != nullptr); + if (live_buffer) { + CHECK(tree.active_slot_ids->buffer == nullptr); + CHECK(dflash::common::detail:: + target_paged_tree_uploads_ready(tree)); + CHECK(!dflash::common::detail:: + target_paged_tree_active_slots_need_upload(tree)); + + const int32_t state_ids[] = {0, 1}; + ggml_backend_tensor_set(tree.state_slot_ids, state_ids, 0, + sizeof(state_ids)); + tree.state_slot_ids = unallocated_state_ids; + CHECK(!dflash::common::detail:: + target_paged_tree_uploads_ready(tree)); + ggml_backend_buffer_free(live_buffer); + } + } + if (live_ctx) ggml_free(live_ctx); + if (marker_ctx) ggml_free(marker_ctx); + if (tree_backend) ggml_backend_free(tree_backend); + } + +} + +TEST_CASE(RecurrentSnapshotFixture, validates_paged_tree_layout) { + // The packed-tree launch length is logical. KVFlash may keep a much + // smaller physical resident pool, provided every tree scratch slab still + // fits within that pool. + { + ggml_init_params shape_params{}; + shape_params.mem_size = 8 * ggml_tensor_overhead(); + shape_params.no_alloc = true; + ggml_context * shape_ctx = ggml_init(shape_params); + CHECK(shape_ctx != nullptr); + if (shape_ctx) { + TargetCache shape_cache; + shape_cache.n_seq_slots = 2; + shape_cache.paged_block_table = + ggml_new_tensor_2d(shape_ctx, GGML_TYPE_I32, 4, 2); + shape_cache.paged_kv_seq_lens = + ggml_new_tensor_1d(shape_ctx, GGML_TYPE_I32, 2); + shape_cache.attn_k = { + ggml_new_tensor_4d(shape_ctx, GGML_TYPE_F16, 4, 64, 1, 1), + }; + CHECK(dflash::common::detail::validate_target_paged_tree_layout( + shape_cache, 8, 2, 4096, 32, 16)); + CHECK(!dflash::common::detail::validate_target_paged_tree_layout( + shape_cache, 8, 2, 4096, 48, 16)); + CHECK(!dflash::common::detail::validate_target_paged_tree_layout( + shape_cache, 8, 5, 4096, 32, 16)); + ggml_free(shape_ctx); + } + } + +} + TEST_CASE(RecurrentSnapshotFixture, snapshot_and_restore_recurrent_state) { ggml_backend_t backend = ggml_backend_cpu_init(); CHECK(backend != nullptr); if (!backend) SKIP("CPU backend is unavailable"); + // C16 x width-23 selects a 32K graph. Prove that graph traversal and + // gallocr can cross the old 16K hard ceiling without asserting. + { + size_t graph_capacity = 0; + CHECK(dflash::common::detail::target_paged_tree_graph_capacity( + 23, 16, graph_capacity)); + ggml_init_params graph_params{}; + graph_params.mem_size = 32 * 1024 * 1024; + graph_params.no_alloc = true; + ggml_context * graph_ctx = ggml_init(graph_params); + CHECK(graph_ctx != nullptr); + if (graph_ctx) { + ggml_tensor * input = + ggml_new_tensor_1d(graph_ctx, GGML_TYPE_F32, 1); + ggml_set_input(input); + ggml_cgraph * graph = ggml_new_graph_custom( + graph_ctx, graph_capacity, false); + for (int i = 0; i < 16385; ++i) { + ggml_build_forward_expand( + graph, ggml_dup(graph_ctx, input)); + } + CHECK(ggml_graph_n_nodes(graph) == 16385); + ggml_gallocr_t graph_alloc = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(backend)); + CHECK(graph_alloc != nullptr); + CHECK(graph_alloc && + ggml_gallocr_alloc_graph(graph_alloc, graph)); + if (graph_alloc) ggml_gallocr_free(graph_alloc); + ggml_free(graph_ctx); + } + } + ggml_init_params params{}; params.mem_size = 8 * ggml_tensor_overhead(); params.no_alloc = true; diff --git a/server/test/test_seq_batch_plan.cpp b/server/test/test_seq_batch_plan.cpp index 8f3b87178..3755dc741 100644 --- a/server/test/test_seq_batch_plan.cpp +++ b/server/test/test_seq_batch_plan.cpp @@ -119,17 +119,107 @@ int main() { SeqEngine::StepPlan work; work.decode = {{0, 7}}; work.prefills = {{1, 4}}; + CHECK(work.decode[0].allow_speculation); SeqEngine::StepResult good; good.decode.push_back({0, 11, false, {}}); good.prefills.push_back({ 1, SeqEngine::PrefillOutput::Status::advanced, -1, {}}); CHECK(validate_step_result(work, good, 2).empty()); + CHECK(prefill_result_made_progress(good)); + + SeqEngine::StepResult burst = good; + burst.decode[0].committed_tokens = {8, 9, 10}; + burst.decode[0].ddtree_steps = 1; + burst.decode[0].ddtree_accepted_tokens = 3; + burst.decode[0].ddtree_suspensions = 1; + burst.decode[0].target_forwards = 1; + CHECK(validate_step_result(work, burst, 2).empty()); + + SeqEngine::StepResult chain_burst = good; + chain_burst.decode[0].committed_tokens = {8, 9}; + chain_burst.decode[0].spec_steps = 1; + chain_burst.decode[0].spec_accepted_tokens = 2; + chain_burst.decode[0].target_forwards = 2; + CHECK(validate_step_result(work, chain_burst, 2).empty()); + + SeqEngine::StepResult chain_service = good; + chain_service.decode[0].spec_service_ar_steps = 1; + chain_service.decode[0].target_forwards = 1; + CHECK(validate_step_result(work, chain_service, 2).empty()); + + SeqEngine::StepResult mixed_chain_paths = chain_burst; + mixed_chain_paths.decode[0].spec_service_ar_steps = 1; + CHECK(!validate_step_result(work, mixed_chain_paths, 2).empty()); + + SeqEngine::StepResult orphan_chain_service = good; + orphan_chain_service.decode[0].spec_service_ar_steps = 1; + CHECK(!validate_step_result(work, orphan_chain_service, 2).empty()); + + SeqEngine::StepResult orphan_chain_acceptance = good; + orphan_chain_acceptance.decode[0].spec_accepted_tokens = 1; + CHECK(!validate_step_result(work, orphan_chain_acceptance, 2).empty()); + + SeqEngine::StepResult failed_chain = good; + failed_chain.decode[0] = {0, -1, true, "decode failed"}; + failed_chain.decode[0].spec_steps = 1; + CHECK(!validate_step_result(work, failed_chain, 2).empty()); + + SeqEngine::StepResult orphan_suspension = good; + orphan_suspension.decode[0].ddtree_suspensions = 1; + CHECK(!validate_step_result(work, orphan_suspension, 2).empty()); + + SeqEngine::StepResult failed_suspension = good; + failed_suspension.decode[0].failed = true; + failed_suspension.decode[0].token = -1; + failed_suspension.decode[0].error = "decode failed"; + failed_suspension.decode[0].ddtree_steps = 1; + failed_suspension.decode[0].ddtree_suspensions = 1; + CHECK(!validate_step_result(work, failed_suspension, 2).empty()); + + // A scheduler stop in the committed prefix must hide the remaining burst + // and final pending token. The backend state is discarded at retirement. + std::vector delivered; + const bool delivered_all = consume_decode_output_tokens( + burst.decode[0], [&](int32_t token) { + delivered.push_back(token); + return token != 9; + }); + CHECK(!delivered_all); + CHECK((delivered == std::vector{8, 9})); + + SeqEngine::StepResult malformed_burst = burst; + malformed_burst.decode[0].committed_tokens = {8, -1}; + CHECK(!validate_step_result(work, malformed_burst, 2).empty()); + + SeqEngine::StepPlan speculation_disabled = work; + speculation_disabled.decode[0].allow_speculation = false; + CHECK(!validate_step_result( + speculation_disabled, burst, 2).empty()); + CHECK(validate_step_result( + speculation_disabled, good, 2).empty()); SeqEngine::StepResult complete = good; complete.prefills[0] = { 1, SeqEngine::PrefillOutput::Status::completed, 12, {}}; CHECK(validate_step_result(work, complete, 2).empty()); + CHECK(prefill_result_made_progress(complete)); + + SeqEngine::StepResult deferred = good; + deferred.prefills[0] = { + 1, SeqEngine::PrefillOutput::Status::deferred, -1, {}}; + CHECK(validate_step_result(work, deferred, 2).empty()); + CHECK(!prefill_result_made_progress(deferred)); + deferred.prefills[0].token = 12; + CHECK(!validate_step_result(work, deferred, 2).empty()); + deferred.prefills[0].token = -1; + deferred.prefills[0].error = "not an error"; + CHECK(!validate_step_result(work, deferred, 2).empty()); + SeqEngine::StepPlan idle_prefill = work; + idle_prefill.decode.clear(); + deferred.prefills[0].error.clear(); + deferred.decode.clear(); + CHECK(!validate_step_result(idle_prefill, deferred, 2).empty()); SeqEngine::StepResult missing_decode = good; missing_decode.decode.clear(); @@ -149,6 +239,7 @@ int main() { prefill_failure.prefills[0] = { 1, SeqEngine::PrefillOutput::Status::failed, -1, "prefill failed"}; CHECK(validate_step_result(work, prefill_failure, 2).empty()); + CHECK(!prefill_result_made_progress(prefill_failure)); SeqEngine::StepResult bad_row_failure = prefill_failure; bad_row_failure.prefills.back().error.clear(); @@ -171,6 +262,13 @@ int main() { success_with_error.decode[0].error = "contradictory diagnostic"; CHECK(!validate_step_result(work, success_with_error, 2).empty()); + SeqEngine::StepResult failed_burst = good; + failed_burst.decode[0].failed = true; + failed_burst.decode[0].token = -1; + failed_burst.decode[0].error = "decode failed"; + failed_burst.decode[0].committed_tokens = {8}; + CHECK(!validate_step_result(work, failed_burst, 2).empty()); + SeqEngine::StepResult failed; failed.error = "device compute failed"; CHECK(validate_step_result(work, failed, 2).empty()); diff --git a/server/test/test_seq_engine_contract.cpp b/server/test/test_seq_engine_contract.cpp index be3f031b4..373aac441 100644 --- a/server/test/test_seq_engine_contract.cpp +++ b/server/test/test_seq_engine_contract.cpp @@ -22,7 +22,9 @@ struct Faults { bool lose_other_pending = false; bool overconsume_prefill = false; bool drop_second_completion = false; + bool defer_first_mixed_prefill = false; bool retire_leaks = false; + bool burst_when_speculation_disabled = false; }; struct FakeCapabilities { @@ -111,6 +113,10 @@ class FakeSeqEngine final : public SeqEngine { 100 + input.slot + (int32_t)slot.fed.size(), false, {}, }); + if (faults_.burst_when_speculation_disabled && + !input.allow_speculation) { + result.decode.back().committed_tokens.push_back(91); + } } std::vector completed_this_step; @@ -122,6 +128,14 @@ class FakeSeqEngine final : public SeqEngine { if (!slot.active || !slot.prefilling || slot.remaining <= 0) { continue; } + if (faults_.defer_first_mixed_prefill && + !plan.decode.empty() && !deferred_mixed_prefill_) { + deferred_mixed_prefill_ = true; + result.prefills.push_back({ + slice.slot, PrefillOutput::Status::deferred, -1, {}, + }); + continue; + } int consumed = std::min(slice.max_tokens, slot.remaining); if (faults_.overconsume_prefill) consumed = slice.max_tokens + 1; slot.remaining -= consumed; @@ -231,6 +245,7 @@ class FakeSeqEngine final : public SeqEngine { std::vector slots_; Faults faults_; FakeCapabilities capabilities_; + bool deferred_mixed_prefill_ = false; }; static void print_violations(const char * label, @@ -279,6 +294,17 @@ int main() { CHECK(violations.empty()); } + { + Faults faults; + faults.defer_first_mixed_prefill = true; + FakeSeqEngine engine(2, faults); + const auto violations = check_seq_engine_contract(engine); + if (!violations.empty()) { + print_violations("conforming-deferred-prefill", violations); + } + CHECK(violations.empty()); + } + struct Case { const char * label; bool Faults::*fault; @@ -302,6 +328,9 @@ int main() { "omitted an output"}, {"retire-leak", &Faults::retire_leaks, "succeed while a slot is free"}, + {"ignore-speculation-gate", + &Faults::burst_when_speculation_disabled, + "disabled speculation"}, }; for (const Case & test : cases) { diff --git a/server/test/test_seq_slot_manager.cpp b/server/test/test_seq_slot_manager.cpp index a64329e5f..7a2d7a26c 100644 --- a/server/test/test_seq_slot_manager.cpp +++ b/server/test/test_seq_slot_manager.cpp @@ -7,9 +7,13 @@ #include "qwen35/concurrency/qwen35_slot_manager.h" #include "host_check.h" +#include "scoped_env.h" +#include #include #include +#include +#include using namespace dflash::common; @@ -85,18 +89,22 @@ int main() { mgr.commit_prefill(0); CHECK(mgr.slot(0).cur_pos == 20); - // Decode appends: row allocation + sample_history; cur_pos advances - // separately after the step's compute. + // Decode append stages row allocation and fed-token history; both + // history and cur_pos publish only after the target compute succeeds. auto st = mgr.append_token(0, /*fed_token=*/42); CHECK(st.ok); CHECK(st.position == 20); CHECK(st.physical_row == 20); // tail of the prompt's last block CHECK(st.new_block < 0 && st.new_block_index < 0); CHECK(mgr.slot(0).cur_pos == 20); - CHECK(mgr.slot(0).sample_history.size() == 21 && - mgr.slot(0).sample_history.back() == 42); + CHECK(mgr.slot(0).sample_history.size() == 20); + CHECK(mgr.slot(0).staged_tokens.size() == 1 && + mgr.slot(0).staged_tokens.back() == 42); mgr.commit_step(0); CHECK(mgr.slot(0).cur_pos == 21); + CHECK(mgr.slot(0).sample_history.size() == 21 && + mgr.slot(0).sample_history.back() == 42); + CHECK(mgr.slot(0).staged_tokens.empty()); // Second admission lands in slot 1 with non-identity rows. auto b = admit(mgr, 2, prompt_tokens(20), greedy_sampler()); @@ -333,6 +341,51 @@ int main() { CHECK(pool.free_block_count() == 0); } + // Accepted-path replay stages multiple rows and publishes history and + // logical position only after the target compute succeeds. + { + PagedKvPool pool(8, 1, /*block_size=*/4); + Qwen35SlotManager mgr(pool, /*max_ctx=*/32, + /*speculative_headroom=*/7); + auto a = admit(mgr, 1, prompt_tokens(3), greedy_sampler()); + CHECK(is_admitted(a)); + CHECK(mgr.append_prefill(a.slot, 3).ok); + mgr.commit_prefill(a.slot); + + const int32_t accepted[] = {41, 42, 43, 44, 45, 46}; + auto staged = mgr.append_tokens(a.slot, accepted, 6); + CHECK(staged.ok && !staged.busy && staged.count == 6); + CHECK(staged.position == 3 && staged.physical_rows.size() == 6); + CHECK(staged.first_new_block == 1); + CHECK(staged.new_blocks.size() == 2); + CHECK(mgr.slot(a.slot).cur_pos == 3); + CHECK(mgr.slot(a.slot).sample_history.size() == 3); + CHECK(mgr.slot(a.slot).staged_tokens == + std::vector(accepted, accepted + 6)); + CHECK(!mgr.append_token(a.slot, 99).ok); + + // A failed target step never publishes staged history. Scheduler + // retirement releases both materialized rows and the staged host range. + const uint32_t free_while_staged = pool.free_block_count(); + mgr.retire(a.slot); + CHECK(!mgr.is_active(a.slot)); + CHECK(pool.free_block_count() > free_while_staged); + + a = admit(mgr, 2, prompt_tokens(3), greedy_sampler()); + CHECK(is_admitted(a)); + CHECK(mgr.append_prefill(a.slot, 3).ok); + mgr.commit_prefill(a.slot); + staged = mgr.append_tokens(a.slot, accepted, 6); + CHECK(staged.ok); + + mgr.commit_step(a.slot); + CHECK(mgr.slot(a.slot).cur_pos == 9); + CHECK(mgr.slot(a.slot).staged_tokens.empty()); + CHECK(mgr.slot(a.slot).sample_history.size() == 9); + CHECK(std::equal(accepted, accepted + 6, + mgr.slot(a.slot).sample_history.end() - 6)); + } + // Context exhaustion: append_token refuses past max_ctx. { PagedKvPool pool(4, 1, /*block_size=*/16); @@ -439,6 +492,134 @@ int main() { CHECK(!mgr.has_prefill_prompt_at_least(768)); } + // Adaptive DDTree judges aggregate cohort yield rather than allowing one + // unlucky request to suppress higher-yield peers. Equality at the + // six-token average continues; a low average suspends all participants. + { + const luce_test::ScopedEnvVar adaptive("DFLASH_DDTREE_ADAPTIVE", nullptr); + CHECK(Qwen35SlotManager::ddtree_cohort_should_suspend(5, 1)); + CHECK(!Qwen35SlotManager::ddtree_cohort_should_suspend(6, 1)); + // Mixed synthetic cohort: one child-less row plus an eleven-token row + // exactly meets the continuation floor; one fewer token does not. + CHECK(!Qwen35SlotManager::ddtree_cohort_should_suspend(12, 2)); + CHECK(Qwen35SlotManager::ddtree_cohort_should_suspend(11, 2)); + + PagedKvPool pool(8, 2, /*block_size=*/16); + Qwen35SlotManager mgr(pool, 64); + auto a = admit(mgr, 101, prompt_tokens(4), greedy_sampler()); + auto b = admit(mgr, 102, prompt_tokens(4), greedy_sampler()); + CHECK(is_admitted(a)); + CHECK(is_admitted(b)); + CHECK(mgr.append_prefill(a.slot, 4).ok); + CHECK(mgr.append_prefill(b.slot, 4).ok); + mgr.commit_prefill(a.slot); + mgr.commit_prefill(b.slot); + CHECK(mgr.slot(a.slot).request_id == 101); + CHECK(mgr.slot(b.slot).request_id == 102); + CHECK(mgr.ddtree_speculation_allowed(a.slot)); + CHECK(mgr.ddtree_speculation_allowed(b.slot)); + + // The keep decision records a real sample without latching either + // participant, including the low-yield member. + CHECK(!mgr.record_ddtree_sample(a.slot, false)); + CHECK(!mgr.record_ddtree_sample(b.slot, false)); + CHECK(mgr.ddtree_speculation_allowed(a.slot)); + CHECK(mgr.ddtree_speculation_allowed(b.slot)); + CHECK(mgr.slot(a.slot).ddtree_sampled_steps == 1); + CHECK(mgr.slot(b.slot).ddtree_sampled_steps == 1); + + // A low cohort decision is applied identically and atomically to both. + CHECK(mgr.record_ddtree_sample(a.slot, true)); + CHECK(mgr.record_ddtree_sample(b.slot, true)); + CHECK(!mgr.ddtree_speculation_allowed(a.slot)); + CHECK(!mgr.ddtree_speculation_allowed(b.slot)); + CHECK(mgr.slot(a.slot).ddtree_sampled_steps == 2); + CHECK(mgr.slot(b.slot).ddtree_sampled_steps == 2); + + // Suspended requests cannot pay for another bad probe. + CHECK(!mgr.record_ddtree_sample(a.slot, true)); + CHECK(mgr.slot(a.slot).ddtree_sampled_steps == 2); + + mgr.retire(a.slot); + auto reused = admit(mgr, 103, prompt_tokens(4), greedy_sampler()); + CHECK(is_admitted(reused) && reused.slot == a.slot); + CHECK(mgr.slot(reused.slot).request_id == 103); + CHECK(mgr.ddtree_speculation_allowed(reused.slot)); + CHECK(mgr.slot(reused.slot).ddtree_sampled_steps == 0); + } + + // A failed residency barrier quarantines retirement ownership, and a + // later admission retries it before considering the slot reusable. + { + PagedKvPool pool(/*physical_block_count=*/1, + /*max_sequences=*/1, /*block_size=*/4); + bool fail_sync = false; + PagedKvResidencyTransferOps transfers{ + [](size_t bytes) -> void * { return new uint8_t[bytes]; }, + [](void * ptr) { delete[] static_cast(ptr); }, + [](PagedKvSequenceHandle, uint32_t, uint32_t, + void *, size_t) { return true; }, + [](PagedKvSequenceHandle, uint32_t, uint32_t, + const void *, size_t) { return true; }, + [&fail_sync] { return !fail_sync; }, + }; + PagedKvResidencyConfig config; + config.block_bytes = 16; + config.resident_budget_blocks = 1; + config.sink_blocks = 0; + config.tail_blocks = 0; + PagedKvResidencyManager residency(pool, config, std::move(transfers)); + Qwen35SlotManager mgr(pool, /*max_ctx=*/16, + /*speculative_headroom=*/1, &residency); + + auto first = admit(mgr, 301, prompt_tokens(4), greedy_sampler()); + CHECK(is_admitted(first)); + CHECK(mgr.append_prefill(first.slot, 4).ok); + CHECK(residency.commit_pending_writes(mgr.slot(first.slot).handle) == + PagedKvResidencyStatus::Ok); + mgr.commit_prefill(first.slot); + + fail_sync = true; + CHECK(residency.evict_block( + mgr.slot(first.slot).handle, 0, + /*allow_protected=*/true) == + PagedKvResidencyStatus::TransferFailed); + mgr.retire(first.slot); + CHECK(!mgr.is_active(first.slot)); + CHECK(mgr.slot(first.slot).retiring()); + CHECK(pool.active_sequence_count() == 1); + + auto blocked = admit(mgr, 302, prompt_tokens(4), greedy_sampler()); + CHECK(!is_admitted(blocked)); + CHECK(is_busy(blocked)); + CHECK(mgr.slot(first.slot).retiring()); + CHECK(pool.active_sequence_count() == 1); + + fail_sync = false; + auto recovered = admit(mgr, 303, prompt_tokens(4), greedy_sampler()); + CHECK(is_admitted(recovered)); + CHECK(recovered.slot == first.slot); + CHECK(mgr.is_prefilling(recovered.slot)); + CHECK(pool.active_sequence_count() == 1); + } + + // Long-prefill policy follows active request length and clears on retire. + { + PagedKvPool pool(128, 2, /*block_size=*/16); + Qwen35SlotManager mgr(pool, 2048); + CHECK(!mgr.has_prefill_prompt_at_least(768)); + auto short_req = + admit(mgr, 201, prompt_tokens(512), greedy_sampler()); + CHECK(is_admitted(short_req)); + CHECK(!mgr.has_prefill_prompt_at_least(768)); + auto long_req = + admit(mgr, 202, prompt_tokens(800), greedy_sampler()); + CHECK(is_admitted(long_req)); + CHECK(mgr.has_prefill_prompt_at_least(768)); + mgr.retire(long_req.slot); + CHECK(!mgr.has_prefill_prompt_at_least(768)); + } + std::printf("OK test_seq_slot_manager (%d checks)\n", g_checks); return 0; } diff --git a/server/test/test_server_unit.cpp b/server/test/test_server_unit.cpp index 8373859dc..829bdbaae 100644 --- a/server/test/test_server_unit.cpp +++ b/server/test/test_server_unit.cpp @@ -21,6 +21,7 @@ #include "server/http_server.h" #include "server/chat_template.h" #include "common/sampler.h" +#include "common/concurrency/seq_engine.h" #include "common/backend_precision.h" #include "common/backend_ipc.h" #include "common/moe_hybrid_ffn_eval.h" @@ -1947,6 +1948,33 @@ TEST_CASE(ServerUnitFixture, test_stop_sequence_holdback_extends) { TEST_ASSERT(em.accumulated_text().find("suffix") == std::string::npos); } +TEST_CASE(ServerUnitFixture, + test_concurrent_scheduler_burst_stops_at_eos) { + SeqEngine::DecodeOutput burst; + burst.slot = 0; + burst.committed_tokens = {101, 2, 103}; + burst.token = 104; + + std::vector emitted; + int completion_tokens = 0; + const bool consumed_all = consume_decode_output_tokens( + burst, [&](int32_t token) { + emitted.push_back(token); + ++completion_tokens; + // Mirrors scheduler.cpp: advance_slot marks the slot finished on + // EOS and its callback immediately stops the rest of the burst. + return token != 2; + }); + + TEST_ASSERT(!consumed_all); + TEST_ASSERT((emitted == std::vector{101, 2})); + TEST_ASSERT(completion_tokens == 2); + TEST_ASSERT(std::find(emitted.begin(), emitted.end(), 103) == + emitted.end()); + TEST_ASSERT(std::find(emitted.begin(), emitted.end(), 104) == + emitted.end()); +} + // ═══════════════════════════════════════════════════════════════════════ // Prefix cache hash tests (model-free) // ═══════════════════════════════════════════════════════════════════════ @@ -2369,6 +2397,49 @@ TEST_CASE(ServerUnitFixture, test_pflash_config_modes) { TEST_ASSERT(cfg.pflash_mode != ServerConfig::PflashMode::AUTO); } +TEST_CASE(ServerUnitFixture, test_concurrent_pflash_persistent_forces_skip_park) { + ServerConfig cfg; + cfg.pflash_mode = ServerConfig::PflashMode::AUTO; + cfg.draft_residency = DraftResidencyPolicy::Persistent; + cfg.prefix_cache_cap = 0; + + const auto plan = resolve_concurrent_pflash_plan( + cfg, /*drafter_tokenizer_available=*/true); + TEST_ASSERT(plan.ok()); + TEST_ASSERT(plan.enabled); + TEST_ASSERT(plan.force_skip_park); +} + +TEST_CASE(ServerUnitFixture, test_concurrent_pflash_rejects_unsafe_residency) { + ServerConfig cfg; + cfg.pflash_mode = ServerConfig::PflashMode::ALWAYS; + cfg.draft_residency = DraftResidencyPolicy::Auto; + cfg.prefix_cache_cap = 0; + + auto plan = resolve_concurrent_pflash_plan( + cfg, /*drafter_tokenizer_available=*/true); + TEST_ASSERT(!plan.ok()); + TEST_ASSERT(plan.error.find("--draft-residency persistent") != + std::string::npos); + + cfg.draft_residency = DraftResidencyPolicy::Persistent; + plan = resolve_concurrent_pflash_plan( + cfg, /*drafter_tokenizer_available=*/false); + TEST_ASSERT(!plan.ok()); + TEST_ASSERT(plan.error.find("--prefill-drafter") != std::string::npos); +} + +TEST_CASE(ServerUnitFixture, test_concurrent_pflash_rejects_snapshot_caches) { + ServerConfig cfg; + cfg.pflash_mode = ServerConfig::PflashMode::AUTO; + cfg.draft_residency = DraftResidencyPolicy::Persistent; + // Default prefix_cache_cap is intentionally non-zero. + const auto plan = resolve_concurrent_pflash_plan( + cfg, /*drafter_tokenizer_available=*/true); + TEST_ASSERT(!plan.ok()); + TEST_ASSERT(plan.error.find("snapshots") != std::string::npos); +} + TEST_CASE(ServerUnitFixture, test_pflash_compress_request_struct) { ModelBackend::CompressRequest req; req.input_ids = {1, 2, 3, 4, 5}; @@ -2509,6 +2580,46 @@ TEST_CASE(ServerUnitFixture, test_parse_request_sampler_applies_defaults_and_ove TEST_ASSERT(std::fabs(sampler.pres_pen - 0.3f) < 0.001f); TEST_ASSERT(std::fabs(sampler.rep_pen - 1.1f) < 0.001f); } +TEST_CASE(ServerUnitFixture, test_speculation_policy_parse_name_and_fold) { + SpeculationPolicy policy = SpeculationPolicy::Adaptive; + TEST_ASSERT(parse_speculation_policy("ar", policy)); + TEST_ASSERT(policy == SpeculationPolicy::Never); + TEST_ASSERT(std::string(speculation_policy_name(policy)) == "ar"); + + TEST_ASSERT(parse_speculation_policy("speculation", policy)); + TEST_ASSERT(policy == SpeculationPolicy::Always); + TEST_ASSERT(std::string(speculation_policy_name(policy)) == "speculation"); + + TEST_ASSERT(parse_speculation_policy("adaptive", policy)); + TEST_ASSERT(policy == SpeculationPolicy::Adaptive); + TEST_ASSERT(!parse_speculation_policy("always", policy)); + + TEST_ASSERT(resolve_speculation_policy( + SpeculationPolicy::Always, std::nullopt) == SpeculationPolicy::Always); + TEST_ASSERT(resolve_speculation_policy( + SpeculationPolicy::Always, SpeculationPolicy::Never) == + SpeculationPolicy::Never); + + const ConcurrentDecodeCapabilities ar_only{}; + TEST_ASSERT(ar_only.supports(SpeculationPolicy::Never)); + TEST_ASSERT(!ar_only.supports(SpeculationPolicy::Always)); + TEST_ASSERT(!ar_only.supports(SpeculationPolicy::Adaptive)); + + const ConcurrentDecodeCapabilities forced_only{ + /*forced_speculation=*/true, + /*adaptive=*/false, + }; + TEST_ASSERT(forced_only.supports(SpeculationPolicy::Never)); + TEST_ASSERT(forced_only.supports(SpeculationPolicy::Always)); + TEST_ASSERT(!forced_only.supports(SpeculationPolicy::Adaptive)); + + const ConcurrentDecodeCapabilities adaptive{ + /*forced_speculation=*/true, + /*adaptive=*/true, + }; + TEST_ASSERT(adaptive.supports(SpeculationPolicy::Always)); + TEST_ASSERT(adaptive.supports(SpeculationPolicy::Adaptive)); +} TEST_CASE(ServerUnitFixture, test_require_messages_array_rejects_invalid) { const json valid = {{"messages", json::array({ @@ -4857,6 +4968,7 @@ TEST_CASE(ServerUnitFixture, test_props_runtime_shape) { cfg.chunk = 512; cfg.target_device = "auto:0"; cfg.draft_device = "auto:0"; + cfg.decode_mode = SpeculationPolicy::Always; TEST_ASSERT(cfg.admission_coalesce_ms == 20); Tokenizer tok; @@ -4872,6 +4984,8 @@ TEST_CASE(ServerUnitFixture, test_props_runtime_shape) { TEST_ASSERT(rt["kv_cache_v"].get() == "tq3_0"); TEST_ASSERT(rt["lazy_draft"].get() == false); TEST_ASSERT(rt["draft_residency"].get() == "persistent"); + TEST_ASSERT(rt["decode_mode"].get() == "speculation"); + TEST_ASSERT(body["decode_mode"].get() == "speculation"); TEST_ASSERT(rt["target_sharding"].get() == false); TEST_ASSERT(rt["chunk"].get() == 512); TEST_ASSERT(rt["target_device"].get() == "auto:0"); diff --git a/server/test/test_spec_cost_profile.cpp b/server/test/test_spec_cost_profile.cpp new file mode 100644 index 000000000..990c4603c --- /dev/null +++ b/server/test/test_spec_cost_profile.cpp @@ -0,0 +1,126 @@ +#include "common/speculation/spec_cost_profile.h" +#include "host_check.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +using namespace dflash::common; + +static int g_checks = 0; + +int main() { + const SpecProfileGrid grid = build_spec_profile_grid( + 3, 16, 4, [](int lanes) { return lanes; }); + CHECK((grid.tree_rows == std::vector{16, 32, 48})); + CHECK((grid.draft_lanes == std::vector{1, 2, 3})); + CHECK(std::binary_search(grid.step_rows.begin(), grid.step_rows.end(), 1)); + CHECK(std::binary_search(grid.step_rows.begin(), grid.step_rows.end(), 12)); + CHECK(std::is_sorted(grid.step_rows.begin(), grid.step_rows.end())); + CHECK(std::adjacent_find(grid.step_rows.begin(), grid.step_rows.end()) == + grid.step_rows.end()); + + const SpecProfileGrid bucketed = build_spec_profile_grid( + 5, 7, 7, [](int lanes) { + if (lanes <= 1) return 1; + if (lanes <= 2) return 2; + if (lanes <= 4) return 4; + return 6; + }); + CHECK((bucketed.tree_rows == std::vector{7, 14, 28, 42})); + CHECK(bucketed.step_rows.back() == 35); + CHECK(build_spec_profile_grid(0, 16, 4, {}).tree_rows.empty()); + + SpecProfileGrid small; + small.tree_rows = {2, 1, 2}; + small.step_rows = {4}; + small.draft_lanes = {3}; + std::unordered_map tree_calls; + const double noise[] = {-2.0, 1.0, 0.0, 2.0, -1.0}; + auto tree_runner = [&](int index) { + const int call = tree_calls[index]++; + if (call == 0) return 10000.0; + const double intended = index == 1 ? 10.0 : 5.0; + return intended + noise[(call - 1) % 5]; + }; + int step_calls = 0; + int draft_calls = 0; + const SpecCostProfileResult profiled = SpecCostProfiler{}.profile( + small, + tree_runner, + [&](int index) { + ++step_calls; + return index == 4 ? 20.0 : 0.0; + }, + [&](int index) { + ++draft_calls; + return index == 3 ? 30.0 : 0.0; + }, + "adapter-score-v1"); + CHECK(profiled.ok()); + CHECK(profiled.tables.speculator_id == "adapter-score-v1"); + CHECK((profiled.tables.tree_cost.indices == std::vector{1, 2})); + CHECK(profiled.tables.tree_cost.costs[0] == 10.0); + CHECK(profiled.tables.tree_cost.costs[1] == 10.0); + CHECK(tree_calls[1] == 6 && tree_calls[2] == 6); + CHECK(step_calls == 6); + CHECK(draft_calls == 6); + + SpecCostProfileResult bad = SpecCostProfiler{}.profile( + {}, [](int) { return 1.0; }, [](int) { return 1.0; }, + [](int) { return 1.0; }, "adapter-score-v1"); + CHECK(!bad.ok() && !bad.error.empty()); + bad = SpecCostProfiler{}.profile( + small, [](int) { return 1.0; }, [](int) { return 1.0; }, + [](int) { return 1.0; }, "", 5); + CHECK(!bad.ok()); + int invalid_calls = 0; + bad = SpecCostProfiler{}.profile( + {{1}, {1}, {1}}, + [&](int) { + ++invalid_calls; + return invalid_calls == 2 + ? std::numeric_limits::quiet_NaN() : 1.0; + }, + [](int) { return 1.0; }, + [](int) { return 1.0; }, + "adapter-score-v1"); + CHECK(!bad.ok()); + CHECK(bad.tables.tree_cost.indices.empty()); + + const std::filesystem::path cache_dir = + std::filesystem::temp_directory_path() / + ("dflash-spec-profile-test-" + std::to_string(getpid())); + const std::filesystem::path cache_path = cache_dir / "profile"; + std::string cache_error; + CHECK(save_spec_cost_profile( + cache_path.string(), "identity-a", profiled.tables, cache_error)); + SpecCostTables loaded; + CHECK(load_spec_cost_profile( + cache_path.string(), "identity-a", loaded, cache_error)); + CHECK(loaded.speculator_id == profiled.tables.speculator_id); + CHECK(loaded.tree_cost.indices == profiled.tables.tree_cost.indices); + CHECK(loaded.tree_cost.costs == profiled.tables.tree_cost.costs); + CHECK(!load_spec_cost_profile( + cache_path.string(), "identity-b", loaded, cache_error)); + CHECK(loaded.tree_cost.indices.empty()); + + { + std::ofstream corrupt(cache_path, std::ios::trunc); + corrupt << "not a profile\n"; + } + CHECK(!load_spec_cost_profile( + cache_path.string(), "identity-a", loaded, cache_error)); + std::error_code remove_error; + std::filesystem::remove_all(cache_dir, remove_error); + + std::printf("spec cost profile tests passed: %d checks\n", g_checks); + return 0; +} diff --git a/server/test/test_speculation_gate.cpp b/server/test/test_speculation_gate.cpp new file mode 100644 index 000000000..de67ffc57 --- /dev/null +++ b/server/test/test_speculation_gate.cpp @@ -0,0 +1,511 @@ +#include "common/speculation/speculation_gate.h" +#include "common/speculation/speculator.h" +#include "common/speculation/survival_score.h" +#include "host_check.h" + +#include +#include +#include +#include + +using namespace dflash::common; + +static int g_checks = 0; + +static SpecCostSeries series(int max_index, double cost) { + SpecCostSeries out; + for (int i = 1; i <= max_index; ++i) { + out.indices.push_back(i); + out.costs.push_back(cost); + } + return out; +} + +static SpecCostTables constant_costs(double tree, double step, double draft) { + return {series(128, tree), series(128, step), series(16, draft)}; +} + +static SpecStepGeometry geometry() { + SpecStepGeometry out; + out.tree_width = 4; + out.bucket = [](int lanes) { return lanes; }; + return out; +} + +static SpecCandidate candidate( + uint64_t id, int slot, double activation_score, + SpeculationPolicy policy = SpeculationPolicy::Adaptive, + bool scoreable = true, bool can_speculate = true, + std::vector hazards = {}, + std::string score_kind = "test-score-v1") { + return {id, slot, policy, scoreable, can_speculate, activation_score, + std::move(hazards), std::move(score_kind)}; +} + +int main() { + // A drafter without a registered adapter (including a DSpark-only GGUF) + // is not scoreable and reports the generic fallback reason. + CHECK(!speculator_is_ready(nullptr)); + CHECK(std::string(speculator_fallback_reason(nullptr)) == + "no_speculator_adapter"); + const ConfidenceVectorScorer confidence_scorer; + CHECK(std::abs(confidence_scorer.score({0.5f, 0.5f}, 4) + .expected_yield - 1.75) < 1e-9); + CHECK(confidence_scorer.score({2.0f, -1.0f}, 4).expected_yield == 2.0); + CHECK(confidence_scorer.score({}, 4).expected_yield == 1.0); + + // A first finite score is a complete cached evaluation. Expensive + // speculation routes both requests to AR and prices k=0 as the + // exact pure-AR candidate (no draft tax). + SpeculationGate costly(constant_costs(100.0, 10.0, 100.0), + geometry(), 4); + CHECK(costly.valid()); + SpecPlan plan = costly.plan(2, { + candidate(1, 0, 4.0), candidate(2, 1, 4.0)}, 2); + CHECK(plan.valid); + CHECK(plan.admitted_count == 0); + CHECK(std::abs(plan.goodput - plan.ar_goodput) < 1e-12); + CHECK(plan.draft_lanes == 0); + CHECK(plan.profiled_cost == 10.0); + CHECK(plan.cost_scale == 1.0); + CHECK(plan.predicted_cost == plan.profiled_cost); + CHECK(costly.initial_score(1) == 4.0); + + plan = costly.plan(2, { + candidate(1, 0, NAN), candidate(2, 1, NAN)}, 2); + CHECK(plan.pending_evaluations.empty()); + CHECK(plan.ordered.size() == 2); + CHECK(plan.admitted_count == 0); + CHECK(std::abs(plan.goodput - plan.ar_goodput) < 1e-12); + + // The initial ranking is a prefix argmax. The first score remains immutable when later inputs present reversed scores. + SpeculationGate prefix(constant_costs(1.0, 10.0, 1.0), + geometry(), 4); + plan = prefix.plan(3, { + candidate(10, 0, 4.0), candidate(11, 1, 1.0), + candidate(12, 2, 4.0, SpeculationPolicy::Never)}, 3); + CHECK(plan.admitted_count == 1); + CHECK((plan.admitted_request_ids == std::vector{10})); + CHECK(plan.ordered.size() == 2); + CHECK(plan.ordered[0].admitted); + CHECK(plan.ordered[0].source == SpecScoreSource::Fresh); + CHECK(std::string(spec_score_source_name(plan.ordered[0].source)) == + "fresh"); + plan = prefix.plan(3, { + candidate(10, 0, 1.0), candidate(11, 1, 4.0), + candidate(12, 2, 4.0, SpeculationPolicy::Never)}, 3); + CHECK((plan.admitted_request_ids == std::vector{10})); + CHECK(plan.ordered.size() == 2); + CHECK(!plan.ordered[0].forced); + CHECK(plan.ordered[0].source == SpecScoreSource::Initial); + CHECK(prefix.initial_score(10) == 4.0); + CHECK(prefix.initial_score(11) == 1.0); + + // A cold cohort is evaluated before planning: every scoreable undecided lane without a score + // is returned for bootstrap, and no scored-but-undecided peer commits until + // the immediate replan. The immediate replan ranks every lane from cached scores. + SpecCostTables crossover = constant_costs(1.0, 4.0, 1.0); + SpeculationGate bootstrap(crossover, geometry(), 4); + plan = bootstrap.plan(2, { + candidate(20, 0, NAN), candidate(21, 1, 4.0)}, 2); + CHECK(plan.valid); + CHECK(plan.admitted_count == 0); + CHECK(plan.ordered.empty()); + CHECK(plan.unavailable_count == 1); + CHECK(plan.pending_evaluations.size() == 1); + CHECK(plan.pending_evaluations[0].request_id == 20); + CHECK(plan.pending_evaluations[0].slot == 0); + CHECK(plan.pending_evaluations[0].action == + SpecEvaluationAction::Score); + CHECK(bootstrap.initial_score(21) == 4.0); + + plan = bootstrap.plan(2, { + candidate(20, 0, 1.0), candidate(21, 1, NAN)}, 2); + CHECK(plan.valid); + CHECK(plan.pending_evaluations.empty()); + CHECK(plan.admitted_count == 1); + CHECK((plan.admitted_request_ids == std::vector{21})); + CHECK(bootstrap.initial_score(20) == 1.0); + CHECK(bootstrap.initial_score(21) == 4.0); + + plan = bootstrap.plan(2, { + candidate(20, 0, 4.0), candidate(21, 1, 1.0)}, 2); + CHECK(plan.pending_evaluations.empty()); + CHECK(plan.admitted_count == 1); + CHECK((plan.admitted_request_ids == std::vector{21})); + CHECK(plan.ordered.size() == 2); + CHECK(!plan.ordered[0].forced); + CHECK(plan.ordered[0].source == SpecScoreSource::Initial); + CHECK(bootstrap.initial_score(20) == 1.0); + CHECK(bootstrap.initial_score(21) == 4.0); + + // The gate side of the no-adapter contract emits FallbackAR and records it once; later finite values cannot invent a score. + // Request-lifetime scoreability is separate from permanent executor + // support. Unsupported requests still bootstrap and retain an activation score, + // then commit directly to AR. A request that cannot evaluate an activation score + // receives an explicit failed-evaluation action and AR with no + // synthetic score. + SpeculationGate support(crossover, geometry(), 4); + plan = support.plan(1, { + candidate(30, 0, NAN, SpeculationPolicy::Adaptive, true, false)}, 1); + CHECK(plan.valid); + CHECK(plan.pending_evaluations.size() == 1); + CHECK(plan.pending_evaluations[0].request_id == 30); + CHECK(plan.pending_evaluations[0].slot == 0); + CHECK(plan.pending_evaluations[0].action == + SpecEvaluationAction::Score); + plan = support.plan(1, { + candidate(30, 0, 4.0, SpeculationPolicy::Adaptive, true, false)}, 1); + CHECK(plan.valid); + CHECK(plan.admitted_count == 0); + CHECK(plan.ordered.size() == 1); + CHECK(plan.ordered[0].execution_unsupported); + CHECK(plan.ordered[0].source == SpecScoreSource::Fresh); + CHECK(support.initial_score(30) == 4.0); + plan = support.plan(1, { + candidate(31, 0, NAN, SpeculationPolicy::Adaptive, false, false)}, 1); + CHECK(plan.valid); + CHECK(plan.pending_evaluations.size() == 1); + CHECK(plan.pending_evaluations[0].request_id == 31); + CHECK(plan.pending_evaluations[0].slot == 0); + CHECK(plan.pending_evaluations[0].action == + SpecEvaluationAction::FallbackAR); + CHECK(support.record_evaluation_failure(31)); + CHECK(!support.record_evaluation_failure(31)); + CHECK(support.evaluation_failed(31)); + CHECK(!support.has_score(31)); + CHECK(std::isnan(support.initial_score(31))); + plan = support.plan(1, {candidate(31, 0, 4.0)}, 1); + CHECK(plan.valid); + CHECK(plan.pending_evaluations.empty()); + CHECK(plan.ordered.size() == 1); + CHECK(plan.ordered[0].execution_unsupported); + CHECK(plan.ordered[0].source == SpecScoreSource::Unavailable); + + // A failed lane and an already-scored healthy lane activate atomically: + // the former becomes AR while the latter still receives its + // score-based mode on the immediate replan. + SpeculationGate mixed_activation(crossover, geometry(), 4); + plan = mixed_activation.plan(2, { + candidate(32, 0, NAN, SpeculationPolicy::Adaptive, false, false), + candidate(33, 1, 4.0)}, 2); + CHECK(plan.valid); + CHECK(plan.pending_evaluations.size() == 1); + CHECK(plan.pending_evaluations[0].request_id == 32); + CHECK(plan.pending_evaluations[0].action == + SpecEvaluationAction::FallbackAR); + CHECK(mixed_activation.record_evaluation_failure(32)); + plan = mixed_activation.plan(2, { + candidate(32, 0, NAN, SpeculationPolicy::Adaptive, false, false), + candidate(33, 1, NAN)}, 2); + CHECK(plan.valid); + CHECK(plan.pending_evaluations.empty()); + CHECK(mixed_activation.evaluation_failed(32)); + CHECK(mixed_activation.initial_score(33) == 4.0); + CHECK((plan.admitted_request_ids == std::vector{33})); + + // Explicit Always/Never are configured execution policies, not adaptive + // activation decisions. Always remains a non-negotiable baseline. + SpeculationGate policies(constant_costs(100.0, 10.0, 100.0), + geometry(), 4); + plan = policies.plan(3, { + candidate(40, 0, 1.0, SpeculationPolicy::Never), + candidate(41, 1, NAN, SpeculationPolicy::Always), + candidate(42, 2, 4.0)}, 2); + CHECK(plan.valid); + CHECK(plan.admitted_count >= 1); + CHECK(plan.admitted_request_ids.front() == 41); + CHECK(plan.ordered.front().forced); + CHECK(plan.ordered.front().source == SpecScoreSource::Unavailable); + plan = policies.plan(2, { + candidate(43, 0, 4.0, SpeculationPolicy::Always), + candidate(44, 1, 4.0, SpeculationPolicy::Always)}, 1); + CHECK(!plan.valid); + CHECK(!plan.error.empty()); + plan = policies.plan(1, { + candidate(46, 0, NAN, SpeculationPolicy::Always, + /*scoreable=*/false, /*can_speculate=*/false)}, 1); + CHECK(plan.valid); + CHECK(plan.admitted_count == 1); + CHECK(plan.ordered.front().forced); + CHECK(plan.admitted_request_ids.front() == 46); + + SpeculationGate never(constant_costs(1.0, 2.0, 1.0), geometry(), 4); + plan = never.plan(1, { + candidate(45, 0, NAN, SpeculationPolicy::Never)}, 1); + CHECK(plan.admitted_count == 0); + CHECK(plan.ordered.empty()); + + // Capacity, malformed shapes, always-draft pricing, and lookup clamps. + SpeculationGate capacity(constant_costs(1.0, 10.0, 1.0), geometry(), 4); + plan = capacity.plan(1, {candidate(50, 0, 4.0)}, 0); + CHECK(plan.admitted_count == 0); + plan = capacity.plan(2, {candidate(51, 0, 4.0)}, 1); + CHECK(!plan.valid); + SpeculationGate always_draft( + constant_costs(1.0, 10.0, 1.0), geometry(), 4); + plan = always_draft.plan(1, {candidate(52, 0, 4.0)}, 1, 4); + CHECK(plan.draft_lanes == 4); + + SpecCostSeries sparse{{2, 4}, {1.0, 2.0}}; + CHECK(sparse.valid()); + SpecCostLookup lookup = sparse.lookup(1); + CHECK(lookup.clamped && lookup.profiled_index == 2); + lookup = sparse.lookup(3); + CHECK(!lookup.clamped && lookup.rounded_up && lookup.profiled_index == 4); + lookup = sparse.lookup(9); + CHECK(lookup.clamped && lookup.profiled_index == 4); + SpecCostSeries invalid{{2, 1}, {1.0, 2.0}}; + CHECK(!invalid.valid()); + + int clamp_logs = 0; + SpecCostTables tiny{series(1, 1.0), series(1, 2.0), series(1, 1.0)}; + SpeculationGate clamped(tiny, geometry(), 4, + [&](const char *, int, int) { ++clamp_logs; }); + plan = clamped.plan(2, { + candidate(60, 0, 4.0), candidate(61, 1, 4.0)}, 2); + CHECK(plan.cost_lookup_clamped); + CHECK(clamp_logs > 0); + + // Every request is ranked only from its own immutable, already-calibrated + // adapter score. Re-presenting a different score cannot change the request, + // and one request never supplies a prior for another. + SpeculationGate fitted( + constant_costs(1.0, 10.0, 1.0), geometry(), 4); + plan = fitted.plan(1, { + candidate( + 700, 0, 2.0, SpeculationPolicy::Adaptive, true, true, + {0.5, 0.25}, "adapter-score-v1")}, 1); + CHECK(fitted.initial_score(700) == 2.0); + CHECK(fitted.initial_score_kind(700) == "adapter-score-v1"); + CHECK((fitted.initial_hazards(700) == + std::vector{0.5, 0.25})); + CHECK(plan.ordered[0].expected_yield == 2.0); + CHECK(plan.initial_predicted_tokens == 2.0); + + plan = fitted.plan(1, {candidate(700, 0, 1.0)}, 1); + CHECK(!plan.ordered[0].forced); + CHECK(plan.ordered[0].source == SpecScoreSource::Initial); + CHECK(plan.ordered[0].expected_yield == 2.0); + CHECK(fitted.initial_score(700) == 2.0); + + plan = fitted.plan(1, {candidate(701, 0, 4.0)}, 1); + CHECK(plan.ordered[0].expected_yield == 4.0); + CHECK(fitted.initial_score(701) == 4.0); + CHECK(fitted.initial_score(700) == 2.0); + + // forget() removes only this request's activation state. A repeated ID is + // cold again; fixed configuration and shape-cost feedback are independent. + fitted.forget(700); + CHECK(!fitted.has_state(700)); + CHECK(!fitted.has_score(700)); + CHECK(std::isnan(fitted.initial_score(700))); + plan = fitted.plan(1, {candidate(700, 0, NAN)}, 1); + CHECK(plan.pending_evaluations.size() == 1); + CHECK(plan.pending_evaluations[0].slot == 0); + CHECK(plan.pending_evaluations[0].action == + SpecEvaluationAction::Score); + + // Interpolate the nonlinear profile across fractional expected rows. A + // tiny score increase across 16.5 rows must preserve the 16 -> 17 cost + // cliff without inheriting lround()'s full 100-us decision discontinuity. + SpecCostSeries cliff_step = series(32, 100.0); + for (size_t i = 16; i < cliff_step.costs.size(); ++i) { + cliff_step.costs[i] = 200.0; + } + const SpecCostTables cliff_costs{ + series(128, 1.0), cliff_step, series(16, 1.0)}; + SpeculationGate cliff_gate(cliff_costs, geometry(), 8); + const SpecPlan below_cliff = cliff_gate.plan(4, { + candidate(800, 0, 4.12475, SpeculationPolicy::Always), + candidate(801, 1, 4.12475, SpeculationPolicy::Always), + candidate(802, 2, 4.12475, SpeculationPolicy::Always), + candidate(803, 3, 4.12475, SpeculationPolicy::Always)}, 4); + const SpecPlan above_cliff = cliff_gate.plan(4, { + candidate(804, 0, 4.12525, SpeculationPolicy::Always), + candidate(805, 1, 4.12525, SpeculationPolicy::Always), + candidate(806, 2, 4.12525, SpeculationPolicy::Always), + candidate(807, 3, 4.12525, SpeculationPolicy::Always)}, 4); + CHECK(std::abs(below_cliff.expected_step_rows - 16.499) < 1e-9); + CHECK(std::abs(above_cliff.expected_step_rows - 16.501) < 1e-9); + CHECK(std::abs(below_cliff.profiled_cost - 151.9) < 1e-9); + CHECK(std::abs(above_cliff.profiled_cost - 152.1) < 1e-9); + CHECK(above_cliff.profiled_cost > below_cliff.profiled_cost); + CHECK(above_cliff.profiled_cost - below_cliff.profiled_cost < 1.0); + + // Direct promotion prices one full-live verification graph and no replay. + // With an expensive replay-row cliff, the legacy executor therefore picks + // AR while the one-launch executor correctly picks all four Spec lanes. + SpecCostSeries replay_step = series(128, 200.0); + for (size_t i = 0; i < 4; ++i) replay_step.costs[i] = 10.0; + const SpecCostTables replay_priced{ + series(128, 20.0), replay_step, series(16, 1.0)}; + SpeculationGate legacy_replay_gate( + replay_priced, geometry(), 4); + SpeculationGate direct_commit_gate( + replay_priced, geometry(), 4, {}, {}, true); + std::vector direct_candidates; + for (int i = 0; i < 4; ++i) + direct_candidates.push_back(candidate(850 + i, i, 4.0)); + const SpecPlan legacy_replay_plan = legacy_replay_gate.plan( + 4, direct_candidates, 4); + const SpecPlan direct_commit_plan = direct_commit_gate.plan( + 4, direct_candidates, 4); + CHECK(legacy_replay_plan.admitted_count == 0); + CHECK(direct_commit_plan.admitted_count == 4); + CHECK(direct_commit_plan.tree_rows == geometry().tree_rows(4)); + CHECK(direct_commit_plan.expected_step_rows == 0.0); + CHECK(direct_commit_plan.profiled_cost == 21.0); + + // A compact mixed graph carries one row per AR peer before the + // bucketed speculative tree. It must not price every AR peer at the + // full verification depth. + SpeculationGate direct_mixed_gate( + replay_priced, geometry(), 4, {}, {}, true); + const SpecPlan direct_mixed_plan = direct_mixed_gate.plan(4, { + candidate(870, 0, 4.0, SpeculationPolicy::Always), + candidate(871, 1, 4.0, SpeculationPolicy::Always), + candidate(872, 2, 0.0, SpeculationPolicy::Never), + candidate(873, 3, 0.0, SpeculationPolicy::Never)}, 4); + CHECK(direct_mixed_plan.admitted_count == 2); + CHECK(direct_mixed_plan.tree_rows == geometry().tree_rows(2) + 2); + CHECK(direct_mixed_plan.expected_step_rows == 0.0); + CHECK(direct_mixed_plan.profiled_cost == 21.0); + + // A measured direct shape with zero step rows is valid online feedback and + // remains isolated under that exact execution-shape key. + direct_commit_gate.observe_cost({4, 4, 16, 0, 4}, 84.0); + std::vector forced_direct_candidates; + for (int i = 0; i < 4; ++i) { + forced_direct_candidates.push_back(candidate(860 + i, i, 4.0, + SpeculationPolicy::Always)); + } + const SpecPlan observed_direct_plan = direct_commit_gate.plan( + 4, forced_direct_candidates, 4); + CHECK(observed_direct_plan.admitted_count == 4); + CHECK(observed_direct_plan.cost_scale == 4.0); + CHECK(observed_direct_plan.predicted_cost == 84.0); + + // A realized row-4 sample belongs only to the row-4 executable shape. It + // must not be written under the gate's earlier row-1 expectation. + SpeculationGate shape_feedback( + constant_costs(1.0, 10.0, 1.0), geometry(), 4); + SpecPlan expected_row_one = shape_feedback.plan(1, { + candidate(900, 0, 1.0, SpeculationPolicy::Always)}, 1); + CHECK(expected_row_one.expected_step_rows == 1.0); + CHECK(expected_row_one.cost_scale == 1.0); + shape_feedback.observe_cost({1, 1, 4, 4, 1}, 48.0); + expected_row_one = shape_feedback.plan(1, { + candidate(901, 0, 1.0, SpeculationPolicy::Always)}, 1); + const SpecPlan observed_row_four = shape_feedback.plan(1, { + candidate(902, 0, 4.0, SpeculationPolicy::Always)}, 1); + CHECK(expected_row_one.cost_scale == 1.0); + CHECK(expected_row_one.predicted_cost == 12.0); + CHECK(observed_row_four.cost_scale == 4.0); + CHECK(observed_row_four.predicted_cost == 48.0); + + // One-time activation drafting on a k=0 round is likewise isolated from + // the pure-AR shape used by the gate's future counterfactual. + SpeculationGate draft_shape_feedback( + constant_costs(100.0, 10.0, 100.0), geometry(), 4); + draft_shape_feedback.observe_cost({1, 0, 0, 1, 1}, 440.0); + const SpecPlan pure_ar_after_bootstrap = draft_shape_feedback.plan( + 1, {candidate(903, 0, 4.0)}, 1); + CHECK(pure_ar_after_bootstrap.admitted_count == 0); + CHECK(pure_ar_after_bootstrap.cost_scale == 1.0); + CHECK(pure_ar_after_bootstrap.predicted_cost == 10.0); + + // Shape-local total-cost feedback can change the next cohort epoch's + // route for an existing request. The engine, not request state in the + // gate, keeps the current epoch stable between membership changes. + SpeculationGate cost_feedback( + constant_costs(1.0, 10.0, 1.0), geometry(), 4); + plan = cost_feedback.plan(1, {candidate(950, 0, 4.0)}, 1); + CHECK(plan.admitted_count == 1); + CHECK(plan.profiled_cost == 12.0); + CHECK(plan.cost_scale == 1.0); + cost_feedback.observe_cost({1, 1, 4, 4, 1}, 48.0); + plan = cost_feedback.plan(1, {candidate(950, 0, NAN)}, 1); + CHECK(plan.admitted_count == 0); + CHECK(std::abs(plan.goodput - plan.ar_goodput) < 1e-12); + CHECK(plan.draft_lanes == 0); + + cost_feedback.forget(950); + plan = cost_feedback.plan(1, {candidate(951, 0, 4.0)}, 1); + CHECK(plan.admitted_count == 0); + CHECK(std::abs(plan.goodput - plan.ar_goodput) < 1e-12); + plan = cost_feedback.plan(2, { + candidate(952, 0, 4.0), candidate(953, 1, 4.0)}, 2); + CHECK(plan.admitted_count == 2); + CHECK(plan.cost_scale == 1.0); + + SpeculationGate ar_feedback( + constant_costs(100.0, 10.0, 100.0), geometry(), 4); + SpecPlan ar_plan = ar_feedback.plan(1, {candidate(960, 0, 4.0)}, 1); + CHECK(ar_plan.admitted_count == 0); + ar_feedback.observe_cost({1, 0, 0, 1, 0}, 20.0); + ar_plan = ar_feedback.plan(1, {candidate(960, 0, NAN)}, 1); + CHECK(ar_plan.admitted_count == 0); + CHECK(ar_plan.profiled_cost == 10.0); + CHECK(ar_plan.cost_scale == 2.0); + CHECK(ar_plan.predicted_cost == 20.0); + CHECK(std::abs(ar_plan.goodput - ar_plan.ar_goodput) < 1e-12); + + // Adaptive gains below the default 1% safety margin commit AR. A zero + // margin admits the same new request, while explicit Always is unchanged. + const SpecCostTables near_break_even = + constant_costs(1.0, 100.0, 1.0); + SpeculationGate margin_gate(near_break_even, geometry(), 4); + plan = margin_gate.plan(1, {candidate(970, 0, 1.03)}, 1); + CHECK(plan.admitted_count == 0); + plan = margin_gate.plan(1, {candidate(970, 0, 4.0)}, 1); + CHECK(plan.admitted_count == 0); + CHECK(plan.ordered.size() == 1); + CHECK(!plan.ordered[0].admitted); + CHECK(margin_gate.initial_score(970) == 1.03); + + SpecGateConfig zero_margin; + zero_margin.adaptive_gain_margin = 0.0; + SpeculationGate no_margin( + zero_margin, near_break_even, geometry(), 4); + plan = no_margin.plan(1, {candidate(971, 0, 1.03)}, 1); + CHECK(plan.admitted_count == 1); + plan = margin_gate.plan(1, { + candidate(972, 0, 1.0, SpeculationPolicy::Always)}, 1); + CHECK(plan.admitted_count == 1); + CHECK(plan.ordered[0].forced); + + // Cost-aware endpoints remain deterministic for each new request. + SpeculationGate pays(constant_costs(1.0, 10.0, 1.0), geometry(), 4); + plan = pays.plan(2, { + candidate(980, 0, 4.0), candidate(981, 1, 4.0)}, 2); + CHECK(plan.admitted_count == 2); + SpeculationGate cannot(constant_costs(100.0, 10.0, 100.0), + geometry(), 4); + std::vector eight; + for (int i = 0; i < 8; ++i) + eight.push_back(candidate(990 + i, i, 4.0)); + plan = cannot.plan(8, eight, 8); + CHECK(plan.admitted_count == 0); + + // Epoch identity follows the live decode request set, not slot + // occupancy or changing scores. Refill in the same slot is a new epoch. + const std::vector cohort = { + candidate(1000, 3, 4.0), candidate(1001, 7, 2.0)}; + SpecCohortEpoch epoch; + epoch.id = 7; + epoch.request_ids = SpecCohortEpoch::ids(cohort); + epoch.plan = plan; + CHECK(epoch.matches(cohort)); + CHECK(epoch.matches({ + candidate(1000, 9, NAN), candidate(1001, 2, NAN)})); + CHECK(!epoch.matches({candidate(1000, 3, NAN)})); + CHECK(epoch.matches({ + candidate(1001, 7, NAN), candidate(1000, 3, NAN)})); + CHECK(!epoch.matches({ + candidate(1000, 3, NAN), candidate(1002, 7, NAN)})); + CHECK((epoch.request_ids == std::vector{1000, 1001})); + + std::printf("speculation gate tests passed: %d checks\n", g_checks); + return 0; +} diff --git a/server/tests/test_quantize_draft_q8.py b/server/tests/test_quantize_draft_q8.py index 4fd201eda..01158c0ea 100644 --- a/server/tests/test_quantize_draft_q8.py +++ b/server/tests/test_quantize_draft_q8.py @@ -29,6 +29,15 @@ def add_uint32(self, key, value): def add_array(self, key, value): self.calls.append(("array", key, value)) + def add_string(self, key, value): + self.calls.append(("string", key, value)) + + def add_float32(self, key, value): + self.calls.append(("float32", key, value)) + + def add_quantization_version(self, value): + self.calls.append(("quantization_version", value)) + class Qwen36SwaMetadataTest(unittest.TestCase): @staticmethod @@ -108,6 +117,51 @@ def test_converter_cli_writes_profile_only_when_requested(self): self.assertIsNone(window) self.assertIsNone(pattern) + def test_dflash2_tensor_mapping_is_complete(self): + expected = { + "layers.0.attention_conv.base_kernel": "blk.0.attn_conv.base", + "layers.0.attention_conv.kernel_projection.weight": "blk.0.attn_conv.proj.weight", + "layers.0.mlp_conv.base_kernel": "blk.0.ffn_conv.base", + "layers.0.mlp_conv.kernel_projection.weight": "blk.0.ffn_conv.proj.weight", + "candidate_selector.hidden_projection.weight": "dflash.selector.hproj.weight", + "candidate_selector.predecessor_codebook": "dflash.selector.pred_cb", + "candidate_selector.successor_codebook": "dflash.selector.succ_cb", + } + for source, output in expected.items(): + self.assertEqual(MODULE.map_name(source), output) + self.assertTrue(MODULE.is_norm_tensor("blk.0.attn_conv.base")) + + def test_dflash2_metadata_is_emitted_from_resolved_profile(self): + profile = dict( + hidden=32, + n_layer=1, + n_head=2, + n_head_kv=1, + head_dim=16, + intermediate=64, + vocab=64, + n_target_layers=1, + rope_theta=10_000_000.0, + rms_eps=1e-6, + mask_token_id=63, + block_size=8, + ctx_len=4096, + capture_layer_ids=[7], + conv_kernel_size=2, + conv_group_size=16, + selector_rank=32, + selector_top_k=16, + ) + writer = RecordingWriter() + MODULE.add_arch_metadata(writer, profile) + prefix = "qwen35-dflash-draft.dflash." + self.assertIn(("array", prefix + "target_layer_ids", [7]), writer.calls) + self.assertIn(("uint32", prefix + "block_size", 8), writer.calls) + self.assertIn(("uint32", prefix + "dflash2.conv_kernel_size", 2), writer.calls) + self.assertIn(("uint32", prefix + "dflash2.conv_group_size", 16), writer.calls) + self.assertIn(("uint32", prefix + "dflash2.selector_rank", 32), writer.calls) + self.assertIn(("uint32", prefix + "dflash2.selector_top_k", 16), writer.calls) + def test_gguf_round_trip_preserves_types_and_values(self): with tempfile.TemporaryDirectory() as tmp: path = Path(tmp) / "metadata.gguf"