diff --git a/.agents/specs/dots3-note.md b/.agents/specs/dots3-note.md index ae8787ced..8552b5e65 100644 --- a/.agents/specs/dots3-note.md +++ b/.agents/specs/dots3-note.md @@ -16,18 +16,19 @@ parallelism for Dots3 NOTE"). **NOT present at our parity pin.** fleet device **`thor:gpu0` through an `rc` lease and never by `ssh`** — the host address is recorded in `environment.md` to identify the box, not as a way into it. §6.3 records what that host can and cannot carry for this model, measured. -**Status:** W4b-1 — the SLIDING arm's maths and the whole §2.3 machinery are -ported as host code (§7 W4b-1, evidence §4.7), on top of W4a's full-attention -layer on the decode path (§4.6), W3's host reference (§4.5), W2's whole weight -map (§4.4) and W1's config + registry (§4.1). The arch RESOLVES, parses, accounts -for 38006/38006 of the released checkpoint's tensors, and DECODES one config -shape — every layer `full_attention` with a dense MLP — through -`ModelRegistry::Forward`, over an `mla::ForwardMlaAttentionBlock` that now -carries dots3-note's two LoRA rescales, its `k_rope_only_layernorm` and its -headwise gate. The RELEASED checkpoint still REFUSES BY NAME, at its first MoE -layer and its first sliding layer, and so do GGUF and both towers. No GPU has -been used at any brick and no tensor byte of the checkpoint has been -downloaded: the committed fixtures are the released `config.json` and a +**Status:** W4b-2 — **BOTH attention geometries are on the decode path** +(§7 W4b-2, evidence §4.8), on top of W4b-1's host maths (§4.7), W4a's +full-attention layer (§4.6), W3's host reference (§4.5), W2's whole weight map +(§4.4) and W1's config + registry (§4.1). The arch RESOLVES, parses, accounts +for 38006/38006 of the released checkpoint's tensors, and DECODES a config whose +layers are any mix of `full_attention` and `sliding_attention` with dense MLPs — +through `ModelRegistry::Forward`, over an `mla::ForwardMlaAttentionBlock` that +carries dots3-note's two LoRA rescales, its `k_rope_only_layernorm`, its +headwise gate and now its 513-wide window, reading a PADDED 1088-wide MLA cache +row narrowed to each layer's own logical width. The RELEASED checkpoint still +REFUSES BY NAME, now at its first MoE layer (W5), and so do GGUF and both +towers. No GPU has been used at any brick and no tensor byte of the checkpoint +has been downloaded: the committed fixtures are the released `config.json` and a headers-only projection of the complete shard index. The row stays `SPIKE`. --- @@ -1793,6 +1794,612 @@ Re-run at this head, all four unmoved from the numbers §4.6 recorded: file was touched, so no byte-identity probe was owed and none was run — W4a's six-arm fingerprint stands as the seam's evidence and W4b-1 adds nothing to it. +### 4.8 W4b-2 put the SLIDING arm on the decode path, over a PADDED cache + +**LANDED at W4b-2** (`row/MODEL-MM-dots3-note-W4b-2`; `include/vt/ops.h` + +`src/vt/ops.cpp` + `src/vt/cpu/cpu_mla_attn.cpp` + `src/vt/cpu/cpu_mla_prefill.cpp` ++ `src/vt/cuda/cuda_mla_attn.cu` + `src/vt/cuda/cuda_flash_attn_fa2.cu` (the two +windowed kernels), `include/vllm/model_executor/models/mla_attention.h` + +`src/vllm/model_executor/layers/attention/mla_attention.cpp` + +`include/vllm/model_executor/layers/attention/mla_chunked_context.h` + +`include/vllm/v1/attention/backend.h` + `src/vllm/v1/attention/backend.cpp` +(the seam), `src/vllm/model_executor/models/dots3_note.h` + +`dots3_note_device.cpp` (the wiring), and four test files). CPU-only. No GPU +lease was taken; §6.3's Thor lease is owed for the CUDA half and is named below +rather than implied. + +**Upstream re-derived at vLLM `origin/main` = `bc2d63e650`.** +`git diff --stat d9fbe526c0 origin/main -- vllm/models/dots3_note/` is EMPTY, so +`vllm/models/dots3_note/` is byte-identical to the tree W4b-1 read and every +§2.3 anchor holds unchanged at the newer head. The anchors this brick leans on, +each re-derived rather than copied: `_gather_swa_kv_kernel:49`, +`_apply_swa_score_mask_kernel:119` (the two mask predicates at `:151` and +`:152`), `_build_sliding_window_metadata:192`, `Dots3NoteFlashAttnPrefillBackend +:258` with `run_sliding_window:279` and its `window_size=(sliding_window - 1, 0)` +at **`:300`**, `Dots3NoteTritonMLAImpl:439` (the subclass that keeps +`self.sliding_window`, `:454-468`) with `_forward_swa_mqa:470` and +`forward_mha:565`, `Dots3NotePaddedSparseImpl:697` with `_logical_cache:700-702` +and `do_kv_cache_update:704-720`, `Dots3NoteSlidingAttention:329` (its scale +`qk_head_dim**-0.5` at `:446`, its rope at `:401-409`, `sliding_window= +config.sliding_window_size` at `:457`, `self.indexer = None` / `is_sparse = +False` at `:432-434`), `Dots3NotePaddedMLAAttention:204-216` and the +`physical_head_size` it is fed at `:283`, and the layer dispatch at `:501-505`. + +#### The window is a KERNEL BOUND here and a GATHER upstream, and that is the same function + +Upstream's decode does not window the paged kernel. `_forward_swa_mqa` gathers +`[max(seq_len - GATHER_LEN, 0), ...)` into a workspace, where `GATHER_LEN` is +`sliding_window + query_len - 1` **rounded up to 8** (`:484`), and then masks the +scores with `kv_positions <= query_position` and +`kv_positions >= query_position - WINDOW_SIZE + 1` (`:151-152`). The gather is a +SUPERSET and the mask is what makes it exact; the round-up exists because Triton +needs a power-of-two tile. + +Walking the paged block table directly over `[seq_len - W, seq_len)` reaches the +identical key set with no gather, no workspace and no mask. So the port is a +`window_size` on the two MLA ops rather than two new ops, and the shared seam is +extended rather than bypassed — which is also why W4b-1's +`GatherSwaKv` / `ApplySwaScoreMask` / `BuildSlidingWindowMetadata` stay HOST +reference code driving the gate's oracle, exactly as W3's `ForwardFullAttention` +did after W4a. That is recorded here rather than left for a reader to infer, and +it is listed under `## Owed`. + +`std::nullopt` is the ABSENT state on both args structs, and it is a NOT-TAKEN +branch rather than a wide window: + +| where | what absence does | +|---|---| +| `cpu_mla_attn.cpp` | `j_start` stays 0 — the same `for (j = 0; j < seq_len; ++j)` loop the op had | +| `cuda_mla_attn.cu` | `kv_start` stays 0 in BOTH split stages, so the partition over `[0, seq_len)` is the one that was there | +| `cuda_flash_attn_fa2.cu` (MLA prefill) | `is_local` is false, so `is_causal` / `window_size_left` / `window_size_right` take their previous assignments and the template dispatch is unchanged | +| `cpu_mla_prefill.cpp` | the `first` lower bound stays 0 | +| `mla::MlaBlockDims::sliding_window` | 0, so no `window_size` is ever constructed | +| `v1::TritonMLAImpl::sliding_window` | 0, same | + +**The op gates prove that rather than assert it.** On both ops a window at least +as wide as the longest sequence is compared against no window **bit-for-bit**, +not to a tolerance. A mask applied after the fact could not pass that. + +#### The padded row needed ZERO `vt` changes, and the narrowing is one line + +`Tensor::Slice(2, 0, logical)` shrinks `shape[2]` and keeps both leading strides +(`tensor.cpp:70-84`), every MLA cache op sources its strides from the tensor, and +that IS upstream's `kv_cache[..., : self.head_size]`. The narrowing is one line +in `Dots3NoteModel::ForwardDevice`. This is the correction §4.7 already recorded, +executed: the physical row is the 1088 both classes share, a FULL layer reads its +logical 576 out of the head of it, and a SLIDING layer's logical row IS the +physical one by construction, so the slice is the identity there and is written +unconditionally rather than branched. + +**The evidence is the RAW cache bytes after a real forward**, not an argument: +the gate reads the pool back and asserts that lanes `[6, 10)` of every slot a +FULL layer wrote are still ZERO, with the CONTROL that the same lanes on the +sliding layers carry **28** non-zero values. Without the control the assertion +would pass on a fixture that produced zeros anyway. + +#### Two of W4a's three refusals are LIFTED; the third is NARROWED; one is NEW + +| refusal | at W4a | at W4b-2 | +|---|---|---| +| any `sliding_attention` layer | config level | **LIFTED** — runs through the same seam over `Dots3NoteSlidingAttnMlaDims` | +| a PADDED physical latent row | config level | **LIFTED** — `Slice(2, 0, logical)`, no `vt` change | +| a KV cache row disagreeing with the config | per step | **KEPT**, and now compared against the PHYSICAL row, which is what `MakeDots3NoteKVCache` tells the allocator | +| `seq_len > index_topk` | per step, always | **KEPT and NARROWED** — asked only of a config that HAS a full layer, because `Dots3NoteSlidingAttention` sets `self.indexer = None` / `is_sparse = False` (`model.py:432-434`) | +| a windowed prefill with chunked CONTEXT | — | **NEW**, in the seam | +| any MoE layer | config level | unchanged — W5 | +| a nextn tail | config level | unchanged — W10 | + +**The new refusal is a scope statement with an upstream reason.** A sliding +layer's prefill gathers only `min(seq_len, query_len + W - 1)` keys and runs ONE +varlen call per request group (`attention.py:206`, `:594-654`); upstream never +merges context chunks under a window, so `forward_mha`'s LSE merge has no +windowed form to mirror. The seam throws rather than merging an UNwindowed +context into a windowed suffix, which is a silently wrong answer. Owed to W4b-3. + +**The released `dots-studio/dots3-note-prev` config still refuses**, now at layer +1's MoE rather than at layer 2's sliding attention, so nothing a user can run +changed. `test_dots3_note_scaffold`'s forward-refusal case was updated to name +the piece the released config ACTUALLY trips on — a string that outlives the +refusal it describes is the failure this row keeps recording. + +#### The gate, met + +`test_dots3_note_attn` — **36 cases / 3028 assertions**, CPU-only, no GPU, no +checkpoint, no speed claim (30/2418 at W4b-1, 18/638 at W4a, 12/198 at W3). +Six new cases. + +The bench is a MIXED config — `layer_types` `{full, sliding, full}`, every layer +a dense MLP, physical row 10 against the full arm's logical 6 — loaded through +`ModelRegistry::Resolve` → `reg.factory->load_weights` and run through +`ModelRegistry::Forward` **twice against one cache pool**: a six-token PREFILL, +then a DECODE of the seventh, over a SHUFFLED block table `{1, 0}`. Both halves +are compared against a whole-model double reference that dispatches per layer +kind into W3's `ref::Forward` and W4b-1's `sref::Forward` — a materialized MHA +with no cache, no paging, no gather and the window as the direct positional +predicate `s <= t && t - s < W`. + +Running two steps against one pool is the point. The decode step reads K/V the +PREFILL step wrote, through the padded physical row and the shuffled table, so +"what the decode read out of the cache" has to equal "what a fresh full-sequence +forward computes". + +**What the instrument measured, printed by the gate rather than assumed:** the +prefill's window cuts **3 of 6** queries and drops **6** keys; the decode query +at position 6 keeps **3** of its 7. At `window >= tokens` the windowed answer IS +the causal answer and every assertion here would pass on a port with no window +at all, which is why both counts are asserted BY NUMBER. + +**The bound is `6e-2` and the three ratios are kept SEPARATE**, because §4.6's +review finding F1 is that merging them overstates the headroom: + +| ratio | value | what it says | +|---|---:|---| +| bound / residue | 0.06 / 0.0254 = **2.36x** | headroom above the bf16 floor | +| nearest mechanism / bound | 0.158 / 0.06 = **2.63x** | headroom below the nearest defect | +| nearest mechanism / residue | 0.158 / 0.0254 = **6.22x** | separation of the whole instrument — a statement about the FIXTURE | + +**The fixture was RETUNED twice, and both times a measurement forced it.** The +first draft ran four layers `{full, sliding, full, sliding}` with +`swa_rope_theta` 41 against 137 and amplified the sliding arm's k_pe rows 6x. It +measured a residue of **0.119** with the nearest mechanism — the sliding arm +inheriting the model-level rope theta — at **0.106**, i.e. the nearest defect +sat UNDER the quantisation floor and the instrument could not see it. Two +changes fixed it, and neither was widening the bound: the thetas became 3 +against 1300, orders apart the way the released 5e4 against 8e7 is (W4b-1's +0.0300 relative on ONE layer is simply too small to survive a bf16 model), and +the schedule dropped to three layers, which is still full/sliding/full so a +per-layer field leaking in EITHER direction is wrong. Residue 0.119 → 0.0254, +nearest mechanism 0.106 → 0.158. + +**Each sliding-only mechanism is shown EXERCISED, not merely compiled**, by +neutralising it in the REFERENCE and measuring the device arm drifting away. +Both ratios are given and labelled, for the reason above: + +| the reference with … | device-vs-reference | / the 6e-2 BOUND | / the 0.0254 RESIDUE | +|---|---:|---:|---:| +| **no window at all** — plain causal attention | 0.818662 | 13.6x | 32.2x | +| the sliding arm inheriting the MODEL-level rope theta | 0.182502 | 3.04x | 7.18x | +| the sliding arm's **q** LoRA rescale dropped | 0.158023 | 2.63x | 6.22x | +| the sliding arm's **kv** LoRA rescale dropped | 1.10842 | 18.5x | 43.6x | +| the sliding arm's `k_rope_only_layernorm` dropped | 0.672635 | 11.2x | 26.5x | +| the sliding arm's headwise gate made lane-wise | 0.200210 | 3.34x | 7.88x | + +The two LoRA scales are neutralised SEPARATELY as well as being different +numbers on this fixture (`sqrt(16/3)` and `sqrt(16/6)`), so an arm that dropped +both at once could not distinguish a port carrying both from one carrying only +the q. Upstream's released ranks make the two SLIDING scales EQUAL at +`sqrt(5120/1024)`; the fixture deliberately does not copy them, and the +released-config case pins the released values separately. + +#### The op gates, and why their oracle is the op itself + +`vt::MlaDecodeAttention` and `vt::MlaPrefillAttention` are gated WITHOUT writing +a windowed reference, deliberately: a reference that recomputed +`seq_len - 1 - left` or `iq + (lk - lq) - left` a second time would share the +arithmetic it is supposed to check, which is the shared-helper trap this project +keeps naming. + +- **Decode.** The windowed call over a length-`n` PAGED sequence is compared + against an UNWINDOWED call over a freshly built single-page cache holding + exactly that request's last `min(W, n)` keys — a path already gated against + the ported `ref_mla` oracle. The window is **13** against pages of **16**, so + its start lands INSIDE a page and a port that rounded to a page boundary is + caught. The boundary is pinned from BOTH sides: a window one key WIDER is a + different answer on every request the window cut. +- **Prefill.** The windowed multi-query call is compared against an EXPANDED + batch in which every query becomes its own single-query request carrying only + the keys its window admits, run UNWINDOWED. With `lq == 1` the bottom-right + causal bound admits every key handed in, so the expansion needs no mask of its + own. **475** (query, key) pairs are dropped across 57 queries. +- Both ops refuse a `right != 0` window BY NAME, and the prefill additionally + refuses a NON-causal one: FlashAttention's local mask REPLACES the causal + specialization (`is_causal = causal && !is_local`), so "everything forward, + windowed backward" has no finite spelling. Upstream never asks for one. + +`test_ops_mla_attn` **15 cases / 246290 assertions** (11 / 197113 at the base +SHA `925a4a587`); `test_ops_mla_prefill` **6 cases / 329772 assertions** +(4 / 242156 at the base). Both base numbers are MEASURED — the two test files +were checked out at the base SHA, rebuilt and run, then restored and verified +byte-for-byte — rather than counted off `TEST_CASE` lines. + +#### The DeepSeek-V2 path is byte-identical, MEASURED again on six arms + +`mla::ForwardMlaAttentionBlock` still has FOUR callers — `deepseek_v2`, +`minicpm3`, `kimi_linear` and `dots3_note` — and DeepSeek-V2-Lite carries a +SACRED token-exact gate that cannot be run on a box with no GPU and no +V2-Lite checkpoint. W4a's standard applies unchanged, and the probe was rebuilt +rather than reused. + +The BEFORE arm is a separate `git archive` tree at the base SHA `925a4a587`, +with a byte-identical probe appended (`md5sum` equal on both files), its own +`cmake` configure and its own build — so there is no previous binary a failed +compile could fall back to. Both arms print the compiler exit and refuse a binary +older than its source. + +| arm | geometry | bytes | BASE `925a4a587` | HEAD | +|---|---|---:|---|---| +| 0 | V2-Lite, f32, MIXED (2 decode + 2 prefill, one with context) | 106496 | `2071435139082975929` | identical | +| 1 | V2-Lite, bf16, same batch | 53248 | `15607516550467795365` | identical | +| 2 | V3 q_lora branch, f32 | 86016 | `5937425064452249605` | identical | +| 3 | V3 q_lora branch, bf16 | 43008 | `4610065661939359460` | identical | +| 4 | MiniCPM3 (`is_neox_style=true`), f32 | 30720 | `7108812291202172077` | identical | +| 5 | MiniCPM3 (`is_neox_style=true`), bf16 | 15360 | `16826999257951116139` | identical | + +Six for six. + +**Arms 0 and 1 reproduce W4a's recorded fingerprints exactly — and arms 2 to 5 +do NOT. The first draft of this paragraph generalised from the first two, and +the #1969 review caught it.** What is actually measured: + +| arm | §4.6, base `d7d1ee914` | §4.8, base `925a4a587` | | +|---|---|---|---| +| 0 | `2071435139082975929` | `2071435139082975929` | same | +| 1 | `15607516550467795365` | `15607516550467795365` | same | +| 2 | `4982522374592074643` | `5937425064452249605` | DIFFERENT | +| 3 | `3757253798370478450` | `4610065661939359460` | DIFFERENT | +| 4 | `9024916185557934982` | `7108812291202172077` | DIFFERENT | +| 5 | `16077001697345918067` | `16826999257951116139` | DIFFERENT | + +**It is a PROBE difference and not a behaviour change, and that conclusion is +measured rather than argued.** Three checks, each cheap and each re-run at this +head: + +1. `git log --oneline d7d1ee914..925a4a587` over `mla_attention.{h,cpp}`, + `cpu_mla_attn.cpp`, `cpu_mla_prefill.cpp`, `deepseek_v2.cpp` and + `minicpm3.cpp` returns EXACTLY ONE commit — `446ac1806`, which is W4a itself + — and W4a's own gate is the §4.6 table asserting byte-identity across it. + Widening the sweep to all of `src/vt` and `include/vt` over the same range + adds only `Conv3d` and `Exl3Gemm`, neither of which the MLA block reaches. + So base-to-base the executing chain is unchanged, and by transitivity all + four measurements describe one function. +2. `grep -r 2071435139082975929` over the tree hits ONLY this file. The probe + was never committed — §4.6 says so in its own words — so the two tables were + produced by two independently hand-written instruments in two sessions. +3. Arms 0 and 1 are V2-Lite, whose dims both authors would write the same way + and which take no `q_lora_scale` and no `is_neox_style`. Arms 2 to 5 carry + free parameters that each author chose, and `q_lora_scale` in particular did + not EXIST at §4.6's base: `git show d7d1ee914:…/mla_attention.h | grep -c + q_lora_scale` is **0** and the same grep at `925a4a587` is **2**, because + 446ac1806 introduced the field. §4.6's probe had to be byte-identical across + its own two trees, so it could not mention the field at all; §4.8's probe + could. The byte COUNTS agree on all six arms (86016, 43008, 30720, 15360), + so the two probes agreed on shapes and differed on values — which is exactly + the signature of a differing scalar or weight fill. + +**The transferable rule, which is why this stays in the record rather than +being quietly corrected: a fingerprint from an uncommitted, hand-written probe +is not a cross-session reproducible quantity.** Two probes sharing a prose +label — "the V3 q_lora arm" — are two different instruments, and comparing +their outputs measures the authors, not the code. What each table legitimately +says is base-vs-head identity WITHIN its own session, which is the claim each +was built to make. The cross-table agreement on arms 0 and 1 is a pleasant +coincidence of a parameter-free fixture, not evidence of reproducibility. + +Committing the probe is what would fix this, and it is NOT done here: neither +scratch tree survives, and writing a third probe would produce a third set of +numbers and no more reproducibility than two. It is listed under `## Owed`. + +**The probe's own false-green, caught by the harness rather than by luck.** The +first BASE run used doctest's `-ts=` (the test-SUITE filter) instead of +`--test-case=`. It matched ZERO cases, printed `[doctest] test cases: 0 | 0 +passed | 0 failed | 13 skipped` and `Status: SUCCESS!`, and exited 0. Read +without checking the case count that is a clean pass with no fingerprints — the +third of the four failure modes `scripts/mutation-harness.py`'s own docstring +enumerates, met in the one place that was hand-driven rather than run through the +harness. The rule generalises: **a filter that matches nothing is not a result, +and only a NON-ZERO case count says so.** + +**Both DeepSeek gates were re-run at this head.** `test_mla_attention_block` +**12 cases / 2247715 assertions** and `test_deepseek_v2_forward` **11 / 1052**, +both unmoved from the numbers §4.6 recorded; +`test_deepseek_v2_decode_graph_seam` **3 / 230** and `test_ops_mla_cache` +**9 / 2947** likewise. `test_dots3_note_scaffold` reads **26 / 110819** — one +assertion more than §4.6's 110818, because its forward-refusal case's single +`Contains("sliding-window MLA")` became two, `Contains("MoE layer")` and +`Contains("W5")`, naming the piece the released config now actually trips on. + +**NOT run, and named rather than implied:** the SACRED DeepSeek-V2-Lite e2e token +gate. It needs a ~29.26 GiB checkpoint on a CUDA host; this brick ran CPU-only on +a box with neither. + +#### The mutation table + +**The driver is the one this tree ships.** Every row below was produced by +`scripts/mutation-harness.py` (row LTX25-RES2S-LOOP, +[#921](https://github.com/mudler/vllm.cpp/issues/921)) rather than by a scratch +script — which is what W4b-1's own section says it should have done. The harness +implements the four guards this project has paid for: it REFUSES an anchor that +does not occur exactly once, prints the COMPILER EXIT beside every row, runs the +whole binary and asserts a NON-ZERO case count, and re-stamps every restore so a +stale object cannot carry a previous binary forward. It also refuses to start on +a dirty tree, which is how a restore failure stays distinguishable from an edit. +A second guard runs before the first compile: `check_anchors.py` asserts every +anchor in the plan occurs exactly once, so a stale plan costs nothing. + +```sh +python3 scripts/mutation-harness.py --build build-w4b2 \ + --test test_dots3_note_attn --plan $SCRATCH/w4b2_plan.jsonl +``` + +**Every count below was RE-MEASURED at the final baseline** (36 cases / 3028 +assertions), not carried over from the first pass. The first pass ran at +35/3025, before the case M16 exposed as missing existed; a table that mixed the +two would be an instrument reporting on a state it was not given, which is the +failure this row keeps naming. The verdicts agreed across both passes. + +| id | mutation | compiler exit | result | cases | assertions | +|---|---|---:|---|---:|---:| +| M1 | the sliding window is never resolved from the config | 0 | RED | 2 | 3 | +| M2 | the window reaches the ops ONE WIDER (`sliding_window`, not `- 1`) | 0 | RED | 1 | 1 | +| M3 | the CPU decode's window START is off by one (`seq_len - left`) | 0 | RED | 1 | 1 | +| M4 | the CPU decode ignores the window | **1** | **NOT A RESULT** | — | — | +| M4b | the same, with `(void)j_start` so it compiles | 0 | RED | 1 | 1 | +| M5 | the CPU prefill's PASS-1 loop ignores the window | 0 | **GREEN** | 0 | 0 | +| M5b | the CPU prefill's `first` bound is forced to 0 in ALL THREE passes | 0 | RED | 1 | 1 | +| M6 | the padded cache is read at the LOGICAL stride | 0 | RED | 1 | 24 | +| M7 | the `_logical_cache` narrowing is dropped — a full layer reads the physical row | 0 | RED | 3 | 0 (threw) | +| M8 | a SLIDING layer runs the FULL arm's `MlaBlockDims` | 0 | RED | 4 | 1 | +| M9 | one shared rope cache for both geometries | **1** | **NOT A RESULT** | — | — | +| M9b | the same, compiling | 0 | RED | 3 | 3 | +| M10 | the materializer uses the FULL dims for every layer | **1** | **NOT A RESULT** | — | — | +| M10b | the same, compiling | 0 | RED | 4 | 3 | +| M11 | REACHABILITY: the DECODE production call site is deleted | 0 | RED | 1 | 1 | +| M12 | REACHABILITY: the PREFILL production call site is deleted | 0 | RED | 1 | 1 | +| M13 | the per-step cache-row check is deleted | 0 | RED | 2 | 2 | +| M14 | the `index_topk` refusal is asked of EVERY config | **1** | **NOT A RESULT** | — | — | +| M14b | the same, compiling | 0 | RED | 1 | 1 | +| M15 | the `index_topk` refusal is deleted outright | **1** | **NOT A RESULT** | — | — | +| M15b | the same, compiling | 0 | RED | 2 | 2 | +| M16 | the windowed-prefill-with-CONTEXT refusal is deleted | 0 | **GREEN** | 0 | 0 | +| M16b | the same, after the missing case landed | 0 | RED | 1 | 1 | +| M17 | the sliding softmax scale uses the LATENT row, not `qk_head_dim` | 0 | RED | 1 | 1 | +| M18 | the window is ONE WIDER at the impl (`dims.sliding_window + 1`) | 0 | RED | 1 | 1 | +| M19 | M18, measured on `test_deepseek_v2_forward` | 0 | **GREEN** | 0 | 0 | +| M20 | M18, measured on `test_mla_attention_block` | 0 | RED | 3 | 4 | + +**FIVE of twenty-seven rows FAILED TO BUILD on their first attempt, and every one +of them is printed rather than tidied away.** M4, M9, M10, M14 and M15 each leave +a variable unread (`j_start`, `rope_swa`, `sliding`, `has_full_layer`) and +`-Werror=unused-variable` stops the build — compiler exit 1, `NOT A RESULT`. +That is the failure mode a hand-driven pass reads as a passing test, and the +proportion is the point: **nearly one row in five** would have been a false +green here. W4b-1 hit it once in twenty-six; the harness caught it five times in +one plan without anyone looking. The `b` rows re-run the identical defect behind +`((void)x, …)`. + +**M11 and M12 are the reachability rows, and they are the ones that say the +window is REACHED.** M11 stops `impl.sliding_window` being assigned from +`dims.sliding_window`, so no window ever reaches `vt::MlaDecodeAttention`; M12 +passes a literal 0 to `ForwardMlaPrefillMha`, so none reaches +`vt::MlaPrefillAttention`. Both go red, so the window on the decode path comes +from the config through the real loader and the shared seam, not from a struct +the test typed. + +**THREE GREEN ROWS, and each says something different.** + +**M5 is a green the DRIVER earned, not the code — and it is the reason M5b +exists.** Forcing the PASS-1 loop to start at 0 makes the kernel compute logits +for out-of-window keys and take the running MAX over them. It changes nothing: +softmax is invariant to the constant subtracted before `exp`, and passes 2 and 3 +still sum only `[first, visible)`. The only observable is the LSE, which this +config never merges because it has no chunked context. So the row measures the +mutation, not the guard. M5b moves the `first` DEFINITION instead, which reaches +all three passes with one substitution, and reds. Kept in the table rather than +deleted, because a reader who sees "prefill window ignored, GREEN" and stops has +learnt the wrong thing. + +**M16 is a green that found a REAL GATE GAP, and it is the most valuable row +here.** Deleting the windowed-prefill-with-context refusal left the gate +completely green — because the case asserting that refusal never made it out of +the draft into the committed file. A refusal whose test does not exist is +indistinguishable from a refusal that works. The case is in the gate now, with +two controls (no window ⇒ the call proceeds and fails with a DIFFERENT exception +type; a window with no chunk list ⇒ it does not fire either), and M16b reds. + +**M19 is a green that MAPS THE GATES rather than exposing a hole**, and its pair +M20 is why. Both apply the same defect — a window leaking onto the DeepSeek path +at `impl.sliding_window`. `test_deepseek_v2_forward` does not see it; +`test_mla_attention_block`, whose cases include decode-only and MIXED batches, +reds on 3 cases / 4 assertions. Together they say the field IS on the DeepSeek +path and its 0 IS load-bearing — measured, which is what the byte-identity table +one section up needs and cannot supply on its own, since identical output could +also mean nothing ever read the field. + +**The MECHANISM this section first gave for M19's green was wrong, and the +#1969 review measured it wrong.** The claim was that `test_deepseek_v2_forward` +misses the defect "because its CPU synthetic forward drives the PREFILL half and +`impl.sliding_window` only reaches the decode MQA" — which implies the prefill +half of the same leak WOULD be caught. It is not. Leaking the window into the +PREFILL call instead (`suffix_lse_t, dims.sliding_window + 1`) was run against +that binary at this head through `scripts/mutation-harness.py`: compiler exit +**0**, EXIT 0, **11 cases / 1052 assertions, zero failures — SURVIVED**. The +same mutation on `test_mla_attention_block` is compiler exit 0, EXIT 1, 4 cases +and 2 assertions failing — DETECTED. So the defect is real and detectable, and +the deepseek binary's blindness is a property of the binary. + +**The real reason is that `test_deepseek_v2_forward`'s CPU cases have NO VALUE +ORACLE for the attention output at all.** `:443` asserts finiteness and +run-to-run determinism, `:464` asserts fusion-catalog ADOPT equals the hand-call +fallback, `:482` asserts a zero-routed MoE layer equals a dense one, `:511` +asserts the shared expert changes the logits, and `:537` repeats `:443` at the +real V2-Lite head dims. Every one of them compares the model against ITSELF +under another configuration, or against nothing. A defect that moves both sides +of such a comparison by the same amount is invisible whichever half of attention +it lands in. That both halves are also prefill-only (`RunTiny` builds +`PrefillMeta`, `:299`) is true and is a SECOND reason the decode leak +specifically is unreached — but it is not the reason the binary is blind, and +stating it as the reason draws the map wrong. The M19 conclusion survives; its +explanation did not. + +#### The #1969 REVIEW-REPAIR rows, measured after the repair + +Five more rows, all through `scripts/mutation-harness.py` on this head, all with +the compiler exit printed. They close the review's F5 and ground its F3 and F4. + +| id | mutation | binary | compiler exit | result | cases | assertions | +|---|---|---|---:|---|---:|---:| +| R1 | the `MlaBlockDims::sliding_window < 0` refusal is DELETED | `test_mla_attention_block` | 0 | **RED** | 1 | 1 | +| R1-control | the same, on the ROW gate | `test_dots3_note_attn` | 0 | GREEN | 0 | 0 | +| R2 | the per-step `ld.head_size() <= physical_row` check is DELETED | `test_dots3_note_attn` | 0 | GREEN | 0 | 0 | +| R3 | the CONFIG-level `physical >= logical` refusal is DELETED | `test_dots3_note_scaffold` | 0 | **RED** | 1 | 1 | +| R4 | the PREFILL window leaks onto the DeepSeek path | `test_deepseek_v2_forward` | 0 | GREEN | 0 | 0 | +| R4b | the same | `test_mla_attention_block` | 0 | **RED** | 4 | 2 | + +**R1 and R1-control are the before/after pair.** The refusal shipped with no +test, and deleting it leaves the ROW gate completely green — which is what the +review measured and what R1-control reproduces here on an unchanged binary. The +new case in `test_mla_attention_block` is what turns it red, and it carries +controls on both sides of the boundary (0 is ABSENT and legal, 513 is legal) so +it cannot pass on an implementation that refused every window. + +**R2 stays GREEN and is not a hole, and R3 is the measurement that says so.** +The per-step `ld.head_size() <= physical_row` check is UNREACHABLE, not merely +untested: `physical_latent_row()` IS `swa.latent_row()` (`dots3_note.h:192`), so +on a sliding layer it is an identity, and on a full layer +`ParseDots3NoteParams` has already refused the violating config at load +(`dots3_note.cpp:389`). R3 deletes THAT check and `test_dots3_note_scaffold` +reds, so the closure is gated rather than assumed. The review's diagnosis — that +`Tensor::Slice` backstops it, so deleting it only downgrades the message — +understates the position: the backstop is not reached either, because no input +the loader accepts can violate the invariant. The check is kept as the +executable spelling of upstream's `assert physical_head_size >= self.head_size` +(`model.py:210`), the site says all of this in a comment, and `## Owed` carries +it as an untested assertion rather than as a gated refusal. + +**R4 and R4b are what correct M19's mechanism**, and the correction is written +into the M19 paragraph above rather than only here. + +#### The merge that built clean and threw, and what caught it + +`origin/main` moved twice while this brick was in flight, and the second +integration is worth recording. It merged with no conflict, it COMPILED with no +error, and `test_dots3_note_attn` then went from 36/3028 green to **5 cases +throwing**: + +``` +vt: resident weight: EMPTY tensor has no host bytes to alias + (host-alias arm, dtype f32, rank 0) dense_attn_block.h:206 +``` + +The incoming commit was [#1952](https://github.com/mudler/vllm.cpp/pull/1952)'s +review finding [#1953](https://github.com/mudler/vllm.cpp/issues/1953): +`ResidentWeight` now REFUSES an empty weight, because an empty one aliases a null +host pointer that no downstream op can detect — every op validates rank, shape, +dtype and device, and the shape comes from the CALLER rather than from the bytes. +That refusal is right, and W4b-2 was on the wrong side of it: it made BOTH rope +caches resident up front while `MaterializeDots3NoteDevice` deliberately leaves +the unused one empty (each is 64 MiB at the released 524288 positions). The fix +is one guard per cache, and the comment now says the guard is the CONTRACT rather +than an optimization. + +**Nothing about this was visible to the merge.** The two branches touched +different files, `git merge` reported clean, and `cmake --build` exited 0 — the +"merge-tree CLEAN is not merge-tree BUILDS" note one step further along, where it +builds too and only the gate can see it. The only thing that caught it was +re-running the focused gate set AFTER the merge and BEFORE the push, which is the +sequence AGENTS.md asks for and the reason it asks. + +#### The 33 sliding layers get a FULL-LENGTH KV spec, and upstream gives them a windowed one + +A divergence this brick creates and does not close, recorded here because a +token gate structurally cannot see it and `porting.md` asks for the memory +format to be compared with the oracle explicitly. + +Upstream's `MLAAttention.get_kv_cache_spec` branches on the window +(`vllm/model_executor/layers/attention/mla_attention.py:1215-1219` @ +`bc2d63e650`): + +```python +if self.sliding_window is not None: + return SlidingWindowMLASpec(**common_kwargs, sliding_window=self.sliding_window) +return MLAAttentionSpec(**common_kwargs, non_causal_multi_token_decode=...) +``` + +and every sliding layer takes that branch, because +`Dots3NoteSlidingAttention.__init__` passes +`sliding_window=config.sliding_window_size` into `MLAAttention` +(`vllm/models/dots3_note/nvidia/model.py:457`). Ours emits ONE uniform +`v1::MLAAttentionSpec` for all 46 layers +(`src/vllm/model_executor/models/dots3_note.cpp`, +`MakeDots3NoteKVCache`). + +**There is no correctness consequence** — the window is applied on READ, by the +two ops, and the gate above proves that it is. The consequence is ALLOCATION. +`SlidingWindowSpec.max_admission_blocks_per_request` caps a windowed layer at +`cdiv(min(sliding_window - 1 + extra_retained + in_flight, max_model_len), +block_size) + 1` blocks (`vllm/v1/kv_cache_interface.py:696-722`), against a +full-length spec's whole `max_model_len`. On the released config that is 513 +against 524288, on 33 of 46 layers — 72% of the tower holding a full-length +latent cache where upstream holds roughly a window's worth. It is the single +largest memory property of this architecture, and no token gate can report it, +because the tokens are right either way. + +**THREE pieces are missing, not one**, and that is why this is scoped to W4b-3 +rather than fixed in place. `SlidingWindowMLASpec` is one of the specs +`include/vllm/v1/kv_cache_interface.h` records as deliberately omitted at MLA +campaign T1. `max_memory_usage_bytes` is omitted from that header too, so the +tree cannot yet express the saving even if the spec existed. And a second spec +kind forces the heterogeneous per-layer GROUP SPLIT in `kv_cache_utils`, which +`MakeDots3NoteKVCache`'s own comment has been deferring since W2. Owed, under +`## Owed`, with this row's W4b-3. + +#### The CUDA half is WRITTEN and NOT GATED, and that is the largest debt here + +Both CUDA changes are small and local — `kv_start` in the two MLA-decode split +stages, and the `is_local` normalization the paged FA-2 launcher already performs +one function above the MLA one. Neither has been RUN: this box has no GPU. Under +§6.3 the row's designated CUDA host is `thor:gpu0` through an `rc` lease. + +**The two halves did NOT have the same amount of gate, and the first draft of +this paragraph read as though they did.** It named the two CUDA files together +and then said "the CUDA-vs-CPU window parity case is present in +`test_ops_mla_attn` and SKIPS without a device" — literally true of the DECODE +half and silent about the other. `test_ops_mla_prefill`'s two windowed cases +were CPU-only, and its only `HasCuda()` cases are pre-existing and unwindowed. +So the FA-2 MLA-prefill launcher's `is_local` block had NO case on ANY device, +and a later lease would have discharged the decode half against a record that +read as covering both. The #1969 review found this. The repair adds +`CUDA MLA prefill: the sliding window matches the CPU reference` to +`test_ops_mla_prefill`, mirroring the decode sibling, and comparing the windowed +device call against the windowed CPU op, against the unwindowed device call +(the window must BITE on the device, or the launcher could be dropping +`window_size` on the floor), and against the file's expanded single-query +oracle. + +**Even with that case, NEITHER half has executed, and the assertion count is +what says so.** Both skip without a device, and doctest counts a case that +returns before its first assertion as PASSED with ZERO assertions. So +`test_ops_mla_prefill` reads **7 cases / 329772 assertions** at this head +against 6 / 329772 before — one more case and not one more assertion. A skip is +not a pass, the number is printed here rather than described, and `## Owed` +names the prefill half separately from the decode half so a lease cannot close +one and be read as closing both. + +**Neither has been COMPILED here either, and that is a separate statement.** The +only compile verification these two files can get on this change is CI's +`cuda-fat-build`; at the time of writing it had not reported on `fa96f9557`. Until +it does, the CUDA half is neither compiled nor executed, and this section says so +rather than letting "written" imply "builds". + +**The lease that would gate it CANNOT BE TAKEN, and the blocker is the FLEET +rather than scheduling.** Measured here with `rc devices` on 2026-08-26, not +taken from a report: + +``` +DEVICE STATE HOLDER ELAPSED COMMAND +dgx:gpu0 unhealthy (no contact 1h27m50s) - - - +orin:gpu0 unknown (no contact 1m17s) - - - +thor:gpu0 unhealthy (no contact 1h16m32s) - - - +``` + +Both CUDA hosts this row could use are QUARANTINED — `thor:gpu0`, the designated +one, and `dgx:gpu0` — and the third device is `unknown` rather than healthy, so +even it is not currently reporting; it is an `orin` (sm_87) and not this row's +host in any case. Clearing a quarantined device needs an admin token, which is a +human's decision and not an agent's. So the debt below is blocked on HARDWARE +RECOVERY, and a reader should not infer that a lease was available and nobody +took it. Owed, and listed under `## Owed`. + ## 5. Gates **Correctness first, and the gate form is chosen by measurement, not in advance** @@ -2183,15 +2790,35 @@ dispatchable in order, under the constraints that answer imposes. §4.7 carries the evidence, the mutation table and the two fixture defects a green mutation found. **No device path changed** and none of W4a's three refusals is lifted. -- **W4b-2 — the sliding arm ON the decode path.** The padded/heterogeneous KV - cache in `vt` (a physical row wider than the row a layer reads), a windowed - decode and prefill through the shared MLA seam, and the three refusals W4a - handed on: the DSA lightning indexer's SELECTION, so a request whose `seq_len` - exceeds `index_topk` stops being refused; the PADDED physical latent row, - refused at config level today; and the per-step check for a KV cache row that - disagrees with the config it was built from, which stays because an engine - allocates the cache separately. Owes the SAME byte-identity evidence W4a - produced for the seam's four callers. All are in `## Owed`. +- **W4b-2 — the sliding arm ON the decode path. DONE** + (`row/MODEL-MM-dots3-note-W4b-2`, evidence §4.8, upstream re-derived at + `bc2d63e650`). Both attention geometries run through + `mla::ForwardMlaAttentionBlock`, reached from `ModelRegistry::Forward`, over a + PADDED physical KV row narrowed on read with `Tensor::Slice(2, 0, logical)` — + upstream's `_logical_cache`, and ZERO `vt` cache ops changed. + `vt::MlaDecodeAttention` and `vt::MlaPrefillAttention` each grew an optional + `AttentionWindow`, whose absent state is a not-taken branch proven + bit-identical on both ops. Two of W4a's three refusals are LIFTED (the sliding + layer, the padded row), the `index_topk` one is KEPT and NARROWED to configs + that have a full layer, and the per-step cache-row check is KEPT against the + PHYSICAL row. The seam's byte-identity was re-measured on six arms in a + separate `git archive` tree and arms 0-1 reproduce W4a's fingerprints exactly. + **The CUDA half is written, NOT compiled here and NOT run** — CI's + `cuda-fat-build` is its only compile check, and the `rc` lease that would + execute it is blocked on fleet recovery rather than on scheduling (§4.8) — and + a windowed prefill with chunked CONTEXT is refused by name; both are `## Owed` + against W4b-3. +- **W4b-3 — the DSA lightning indexer's SELECTION on the device path, and the + two debts W4b-2 named.** The split line is that the indexer shares nothing + with the sliding window: the sliding layers carry no indexer at all + (`self.indexer = None` / `is_sparse = False`, model.py:432-434), so lifting + `seq_len > index_topk` is about the FULL layers and needs the indexer weights + on device, its logits, its top-k and a SPARSE MLA attention kernel on both + backends — none of which the window touches. It also carries the windowed + prefill with chunked CONTEXT, and the `rc` lease on §6.3's `thor:gpu0` that + gates W4b-2's CUDA half — which cannot be scheduled at all until the fleet is + back: both CUDA hosts read `unhealthy` on 2026-08-26 and clearing a + quarantined device is an admin-token decision. All three are in `## Owed`. - **W5 — MoE.** Ungrouped `noaux_tc` at 256/8 + the shared expert. Mostly routing our existing path at new dims. - **W6 — vision tower.** Dense ViT half first, then the pyramid MoE and the @@ -2254,19 +2881,114 @@ Carried openly under option B (§6.4), not waived: carries the deltas.** Both halves of W3's entry are discharged — §4.6 is the evidence — and the entry is kept here rather than deleted so a reader who followed W3's `## Owed` link lands on the answer instead of a gap. -- **The SLIDING half of everything W4a did — HALF CLOSED at W4b-1.** 33 of the - 46 layers are `sliding_attention`. Their MATHS now exists and is gated: the - sliding geometry, the windowed metadata, the KV gather, the score mask and the - padded/heterogeneous KV spec are `dots3_note_attn.{h,cpp}`, §4.7 is the - evidence. **What is still owed is the DECODE PATH**, which is W4b-2: - `Dots3NoteModel::ForwardDevice` still refuses a sliding layer by name, and - lifting that needs a `vt` MLA cache whose PHYSICAL row is wider than the row a - layer reads — a change inside `vt::ConcatAndCacheMla`, - `vt::MlaDecodeAttention` and the MLA prefill gather on both backends, with the - CUDA half unverifiable on a CPU-only box and the byte-identity obligation W4a - recorded for the seam's four callers. Owner: row +- **CLOSED at W4b-2: the SLIDING half of everything W4a did is on the decode + path.** 33 of the 46 layers are `sliding_attention`; both attention geometries + now run through `mla::ForwardMlaAttentionBlock` over a PADDED physical KV row, + reached from `ModelRegistry::Forward`. §4.8 is the evidence. The entry is kept + here rather than deleted so a reader who followed W4b-1's `## Owed` link lands + on the answer instead of a gap. The paragraph W4b-1 wrote here — that lifting + the refusal needed changes inside `vt::ConcatAndCacheMla`, + `vt::MlaDecodeAttention` and the MLA prefill gather — was the FALSE constraint + §4.7 already corrects: the cache ops are stride-driven and ZERO of them + changed. What the window needed was a `window_size` on two of them, which is + the additive shape this tree uses everywhere else. Owner: row `MODEL-MM-dots3-note-dots3-note-for-causal-lm`. Issue [#699](https://github.com/mudler/vllm.cpp/issues/699). +- **The CUDA half of the windowed decode and prefill is WRITTEN and NOT RUN.** + `cuda_mla_attn.cu` moves `kv_start` in both split stages and + `cuda_flash_attn_fa2.cu`'s MLA prefill launcher performs the `is_local` + normalization its paged sibling already performs; neither has executed, + because W4b-2 ran on a box with no GPU, and neither has been COMPILED here — + CI's `cuda-fat-build` is the only compile verification this change can give + them. **The two halves are named separately here on purpose, because W4b-2's + first record merged them and the #1969 review caught it.** DECODE: the + CUDA-vs-CPU window parity case is `test_ops_mla_attn`'s "CUDA mla_decode: the + sliding window matches the CPU reference". PREFILL: the matching case in + `test_ops_mla_prefill` did not exist until the #1969 repair added it, so the + FA-2 MLA-prefill launcher's `is_local` block had no case on any device. BOTH + cases SKIP without a device and have therefore never executed — + `test_ops_mla_prefill` reads 7 cases / 329772 assertions, one more case than + before the repair and not one more assertion, which is what a skip looks like + in a count. A lease closes ONE half at a time and this entry stays open until + both are run. §6.3's designated host `thor:gpu0` through an `rc` + lease is what discharges it, and **that lease cannot currently be taken**: + `rc devices` on 2026-08-26 reads `thor:gpu0` and `dgx:gpu0` both `unhealthy` + with no contact for over an hour and `orin:gpu0` `unknown`, and clearing a + quarantined device needs an admin token, which is a human's call. The blocker + is FLEET RECOVERY, not scheduling — §4.8 carries the measurement. Owner: this + row, **W4b-3**. Issue [#699](https://github.com/mudler/vllm.cpp/issues/699). +- **A windowed PREFILL that also carries chunked CONTEXT is refused by name.** + Upstream caps a sliding layer's gather at `min(seq_len, query_len + W - 1)` and + runs one varlen call per request group (`attention.py:206`, `:594-654`), so + `forward_mha`'s LSE merge has no windowed form to mirror and the seam throws + rather than merging an unwindowed context into a windowed suffix. Reachable in + production by a chunked prefill of a long prompt; not reachable on the + RELEASED checkpoint, which refuses at its first MoE layer. Owner: this row, + **W4b-3**. Issue [#699](https://github.com/mudler/vllm.cpp/issues/699). +- **W4b-1's `dots3_note_attn.{h,cpp}` sliding functions stay HOST REFERENCE + code.** `ForwardSlidingAttention`, `GatherSwaKv`, `ApplySwaScoreMask`, + `BuildSlidingWindowMetadata`, `WritePaddedMlaCache` and + `NarrowLogicalCacheRows` have no production call site and did not gain one at + W4b-2, because the device path reaches the same key set through the paged + block table instead of upstream's Triton gather-plus-mask (§4.8). They are the + gate's oracle, which is the status W3's `ForwardFullAttention` has had since + W4a. Stated rather than left to be inferred, per `## Nothing lands dead`. + Owner: this row. Issue [#699](https://github.com/mudler/vllm.cpp/issues/699). +- **The 33 sliding layers report a FULL-LENGTH `MLAAttentionSpec`; upstream + gives them a `SlidingWindowMLASpec`.** `MLAAttention.get_kv_cache_spec` + branches on the window and returns `SlidingWindowMLASpec(..., sliding_window= + self.sliding_window)` when one is set + (`vllm/model_executor/layers/attention/mla_attention.py:1215-1219` @ + `bc2d63e650`), and every sliding layer sets one, because + `Dots3NoteSlidingAttention` passes `sliding_window=config.sliding_window_size` + into `MLAAttention` (`vllm/models/dots3_note/nvidia/model.py:457`). + `MakeDots3NoteKVCache` emits one uniform `v1::MLAAttentionSpec` for all 46. + **No correctness consequence** — the window is applied on READ and the W4b-2 + gate proves it — but 72% of the tower then holds a full-length latent cache + where upstream caps a windowed layer at + `cdiv(min(sliding_window - 1 + extra_retained + in_flight, max_model_len), + block_size) + 1` blocks + (`SlidingWindowSpec.max_admission_blocks_per_request`, + `vllm/v1/kv_cache_interface.py:696-722`), i.e. 513 against 524288 on the + released config. It is the largest memory property of this architecture and + a token gate structurally cannot see it, which is the class `porting.md` + names. THREE pieces are missing, not one: `SlidingWindowMLASpec` is on + `include/vllm/v1/kv_cache_interface.h`'s deliberately-omitted list from MLA + campaign T1, `max_memory_usage_bytes` is omitted from the same header so the + saving is not yet expressible, and a second spec kind forces the + heterogeneous per-layer GROUP SPLIT in `kv_cache_utils`. §4.8 carries the + derivation. Owner: this row, **W4b-3**. Issue + [#699](https://github.com/mudler/vllm.cpp/issues/699). +- **One refusal in the device forward is UNREACHABLE and therefore untested.** + `Dots3NoteModel::ForwardDevice`'s `VT_CHECK(ld.head_size() <= physical_row)` + cannot fire through any production entry point, and this is recorded rather + than dressed up as a gated refusal. Both sides come from the same parsed + config: `physical_latent_row()` IS `swa.latent_row()` + (`dots3_note.h:192`), so on a SLIDING layer the comparison is an identity, and + on a FULL layer `ParseDots3NoteParams` has already refused + `physical_latent_row() < full.latent_row()` at load + (`dots3_note.cpp:389`, gated at + `tests/vllm/models/test_dots3_note_scaffold.cpp:720-722`). Deleting it therefore + leaves the gate green — MEASURED, `scripts/mutation-harness.py` at this head, + compiler exit 0, `test_dots3_note_attn` 36 cases / 3028 assertions, SURVIVED — + and the #1969 review's finding that it is "backstopped by `Tensor::Slice`" + understates it, because the backstop is not reached either. It is kept as the + executable spelling of upstream's `assert physical_head_size >= self.head_size` + (`model.py:210`); making it load-bearing means giving the forward an input the + loader cannot produce, which is not a shape this row wants. Owner: this row. + Issue [#699](https://github.com/mudler/vllm.cpp/issues/699). +- **The six-arm DeepSeek byte-identity probe is not committed, so neither §4.6's + nor §4.8's fingerprints can be reproduced.** Both tables are valid + base-vs-head statements within their own session and neither is reproducible + across sessions; §4.8 records the measurement that proves the four differing + arms are a probe difference and not a behaviour change, and the general rule + that a fingerprint from an uncommitted hand-written probe is not a + cross-session quantity. Not fixed in the #1969 repair: neither scratch tree + survives, and a third hand-written probe would produce a third set of numbers + and no more reproducibility than two. What discharges it is committing the + probe — its arm definitions, dims, seeds and scalar values — beside the + DeepSeek gates, once. Owner: this row. Issue + [#699](https://github.com/mudler/vllm.cpp/issues/699). - **The DSA lightning indexer's SELECTION is not on the device path.** W3 ported the selection maths as a host reference and W4a did not wire it: the shared MLA seam computes DENSE attention, which is upstream's answer only while @@ -2279,7 +3001,11 @@ Carried openly under option B (§6.4), not waived: indexer at all (`self.indexer = None` / `is_sparse = False`, model.py:432-434), so nothing W4b-1 wrote touches the FULL arm's selection, and `Dots3NoteSlidingAttnDimsFrom` REFUSES a params object whose sliding arm claims - one. Owner: this row, **W4b-2**. Issue + one. **W4b-2 did not lift it either, and it could not have, for the same + reason** — but it NARROWED who is asked: the per-step bound is now checked only + for a config that HAS a full-attention layer, so a pure-SWA schedule is no + longer refused for a mechanism it does not carry. §4.8 records the + measurement. Owner: this row, **W4b-3**. Issue [#699](https://github.com/mudler/vllm.cpp/issues/699). - **The PADDED physical latent row.** `MakeDots3NoteKVCache` already reports the 1088-wide row both classes share, and W4a refuses any config whose physical row @@ -2296,10 +3022,13 @@ Carried openly under option B (§6.4), not waived: records. `vt`'s MLA cache ops are STRIDE-DRIVEN, `Tensor::Slice(2, 0, logical)` is upstream's `kv_cache[..., : head_size]`, and a probe wrote, gathered and decoded through a physical-7 / logical-5 view at 30/30 with no `vt` change at - all. What W4b-2 is actually blocked on is the WINDOW: - `vt::MlaDecodeAttention` attends the whole `seq_len` (`cpu_mla_attn.cpp:94`) - with no window and no per-slot `valid`. **W4b-2**. - Owner: this row. Issue + all. **CLOSED at W4b-2**: the config-level refusal is gone, the narrowing is + one `Tensor::Slice(2, 0, logical)` in `Dots3NoteModel::ForwardDevice`, no `vt` + op changed, and the gate reads the RAW cache bytes after a real forward to + assert the pad lanes of every full-layer slot are untouched (§4.8). The + PER-STEP refusal stays and is not the same check: an engine allocates the + cache separately from the config it was built from, so a row that disagrees is + an input only the forward can see. Owner: this row. Issue [#699](https://github.com/mudler/vllm.cpp/issues/699). - **The nextn tail on the device path.** W4a refuses a config with `num_nextn_predict_layers > 0` rather than enumerating, loading and never diff --git a/docs/FEATURES.md b/docs/FEATURES.md index 7ed0ba57a..7f06d1978 100644 --- a/docs/FEATURES.md +++ b/docs/FEATURES.md @@ -150,7 +150,7 @@ speed-pending, which [BENCHMARKS.md](BENCHMARKS.md) tracks. | `LagunaForCausalLM` | poolside/Laguna-S-2.1-NVFP4, GGUF-Q4_K, Laguna-XS | byte-exact near-tie (distributional vs vLLM) | vLLM parity+ 1.03x, default on, via the `laguna-gen` CLI; the registered engine forward VT_CHECKs non-bf16 (`ARCH-ONE-SURFACE` fold) | | `KimiLinearForCausalLM` | Kimi-Linear-48B-A3B (KDA + NoPE-MLA + MoE) | **Folded onto the shared paged runner (ROW 7 §21, #122): engine==CLI 128/128 byte-identical; vs golden 122/128 (the intrinsic near-tie profile); FA2 paged MLA default-ON; SACRED post-fold green** | Served via `vllm_engine_load` + `vllm_complete_tokens` (ABI v13); server 19.0 tok/s wall vs vLLM ~21 (~0.90×), speed residual open | | `KimiK3ForConditionalGeneration` | Kimi-K3 (2.8T MoE) | scaffold: registry+config+enumeration gated, forward refuses | HW-infeasible (~1.56 TB); no run | -| `Dots3NoteForCausalLM` | `dots-studio/dots3-note-prev` @`1e1e7b0c` (280B-A16B multimodal MoE, ~576 GB bf16; the `-fp8` sibling is ~290 GB). Headers only — no tensor byte downloaded | W1+W2 scaffold: registry + config gated off the REAL released `config.json`, with one assertion per §4 config trap (ungrouped 1/1 router, GPT-J indexer RoPE, one nextn layer, the two LoRA rescales, the sliding theta); name map accounted **38006/38006** over the WHOLE released index — 35381 language, 2195 vision, 430 audio, with the two tower files carried as named W6/W7 deferrals rather than dropped; load, GGUF and forward all REFUSE BY NAME | **No oracle, on any host we own** (~290 GB fp8 against a 122 GiB ceiling), so NO number is claimable on any axis and the e2e gate is an open gap by construction ([spec](../.agents/specs/dots3-note.md) §6.4, #699) | +| `Dots3NoteForCausalLM` | `dots-studio/dots3-note-prev` @`1e1e7b0c` (280B-A16B multimodal MoE, ~576 GB bf16; the `-fp8` sibling is ~290 GB). Headers only — no tensor byte downloaded | W1+W2 scaffold: registry + config gated off the REAL released `config.json`, with one assertion per §4 config trap (ungrouped 1/1 router, GPT-J indexer RoPE, one nextn layer, the two LoRA rescales, the sliding theta); name map accounted **38006/38006** over the WHOLE released index — 35381 language, 2195 vision, 430 audio, with the two tower files carried as named W6/W7 deferrals rather than dropped; W4a+W4b-2 put BOTH attention geometries on the DECODE PATH, reached through `ModelRegistry::Forward`: the 13 full-attention layers with the two LoRA rescales, `k_rope_only_layernorm` and the headwise gate, and the 33 sliding-window layers over a PADDED 1088-wide MLA cache row that each layer narrows to its own logical width on read; `vt::MlaDecodeAttention` and `vt::MlaPrefillAttention` grew an optional window whose absent state is bit-identical to no window. The RELEASED checkpoint still REFUSES BY NAME at its first MoE layer (W5), and so do GGUF, the nextn tail (W10) and both towers (W6/W7) | **No oracle, on any host we own** (~290 GB fp8 against a 122 GiB ceiling), so NO number is claimable on any axis and the e2e gate is an open gap by construction ([spec](../.agents/specs/dots3-note.md) §6.4, #699) | | `NemotronHForCausalLM` | Nemotron-3.5-Lightning-30B-A3B-NVFP4 (`nvidia` @`29f2d174`) | config+enumeration+KV-shape gated; hybrid forward COMPUTES; loader materializes 18487/18487 as SHIPPED; A3 e2e gate 96/96 `STRICT PASS` on GB10 at `0ea5d249f` (#1221); NO run against current `main` | **PAGED (#810 A2-P): K/V go to the runner's pages; conv+SSM rows carry at the metadata's state indices.** G-SAFE: `num_reqs <= 1`. Device `lm_head` (A2-Q2b), UNMEASURED. Owed: FP8 mamba (A2-Q1), MTP, GGUF | | `MuseGlimmerForCausalLM` | real tensors, **bf16 depth 4/52 only**: 5 prefill argmax positions match a torch transcription of vllm#51655 and HF. GGUF full depth generates coherently (#347, #359) but is **NOT token-exact** | text forward + loader vs an fp32 reference, per-mechanism property tests, scaffold 11/11, GGUF gate 17/17. An ABSENT config key now takes the architecture's constant (#412): GGUF post-norms ran at 1e-5, not 1e-8 | no vLLM denominator (pin cannot load it); SECONDARY llama.cpp, same GGUF, GB10 CPU: prefill tie **0.997x**, decode 0.232x, RSS 1.92x (#333) | | `MuseGlimmerForConditionalGeneration` | vision: **no reference run of any kind**; enumeration gated vs the released 30B index (1436/1436). Image/video need bf16 safetensors: `mmproj-kquant.gguf` is refused by name | perception encoder loaded and wired, so an image or video prompt runs; `perception_emb_norm` now armed by default (#405). Reachability plus placeholder scatter only, no image or video correctness | not measurable; anchored to open vllm#51655 | diff --git a/include/vllm/model_executor/layers/attention/mla_chunked_context.h b/include/vllm/model_executor/layers/attention/mla_chunked_context.h index 42512ac64..b696db1bb 100644 --- a/include/vllm/model_executor/layers/attention/mla_chunked_context.h +++ b/include/vllm/model_executor/layers/attention/mla_chunked_context.h @@ -320,6 +320,13 @@ inline void ComputeMlaPrefillContext(vt::Queue& q, const vt::Tensor& query, // The final merge is PREFIX = context, SUFFIX = new tokens (:2413-2420) — the // order matters, and it carries `prefill_tokens_with_context` so query rows // belonging to context-free requests take the suffix verbatim. +// +// `sliding_window` (dots3-note W4b-2, #699) is 0 for every DeepSeek / MiniCPM3 / +// Kimi-Linear caller, which leaves `window_size` at `std::nullopt` and this +// function byte-identical. > 0 is `sliding_window_size`, and it reaches the +// new-tokens call as the `AttentionWindow{W - 1, 0}` pair upstream hands +// FlashAttention (`vllm/models/dots3_note/nvidia/attention.py:300` @ +// `bc2d63e650`). inline void ForwardMlaPrefillMha(vt::Queue& q, vt::Tensor& output, const vt::Tensor& query, const vt::Tensor& key, const vt::Tensor& value, const vt::Tensor& kv_cache, const vt::Tensor& block_table, @@ -328,8 +335,22 @@ inline void ForwardMlaPrefillMha(vt::Queue& q, vt::Tensor& output, const vt::Ten const MlaUpProjectFn& up_project, float scale, int32_t max_query_len, int32_t prefill_tokens_with_context, MlaPrefillContextBuffers& bufs, vt::Tensor& suffix_output, - vt::Tensor& suffix_lse) { + vt::Tensor& suffix_lse, int64_t sliding_window = 0) { const bool has_context = !chunks.empty(); + // A windowed prefill that ALSO has chunked context has no upstream form to + // mirror: a sliding layer gathers only `min(seq_len, query_len + W - 1)` keys + // and runs one varlen call per request group (attention.py:206, :594-654), so + // the LSE merge below never runs windowed upstream. Refuse rather than merge + // an unwindowed context into a windowed suffix, which is a silently wrong + // answer and exactly the class this row keeps naming. + if (sliding_window > 0 && has_context) { + throw std::invalid_argument( + "MLA prefill: a SLIDING-WINDOW layer with chunked CONTEXT is not ported. " + "Upstream's windowed prefill caps the gather at the window instead of " + "merging context chunks (dots3-note attention.py:206, :594-654), so there " + "is no windowed form of this merge to mirror. See " + ".agents/specs/dots3-note.md `## Owed` and issue #699."); + } // ":2381-2392" — the causal pass over the new tokens. `return_softmax_lse` is // True exactly when there is context to merge with (:2385). @@ -338,6 +359,9 @@ inline void ForwardMlaPrefillMha(vt::Queue& q, vt::Tensor& output, const vt::Ten args.causal = true; args.max_seqlen_q = max_query_len; args.max_seqlen_k = max_query_len; + if (sliding_window > 0) { + args.window_size = vt::AttentionWindow{static_cast(sliding_window - 1), 0}; + } vt::Tensor& new_out = has_context ? suffix_output : output; vt::MlaPrefillAttention(q, new_out, has_context ? &suffix_lse : nullptr, query, key, value, cu_seqlens_q, cu_seqlens_q, args); diff --git a/include/vllm/model_executor/models/mla_attention.h b/include/vllm/model_executor/models/mla_attention.h index ae91c7581..d206bbb14 100644 --- a/include/vllm/model_executor/models/mla_attention.h +++ b/include/vllm/model_executor/models/mla_attention.h @@ -155,6 +155,39 @@ struct MlaBlockDims { // when `has_q_lora()` is false rather than dropping it silently. double kv_lora_scale = 1.0; + // ─── dots3-note's SLIDING WINDOW (W4b-2, #699) ──────────────────────────── + // 0 is the ABSENT state and it is a NOT-TAKEN branch, not a wide window: at + // 0 no `window_size` reaches `vt::MlaDecodeAttention` or + // `vt::MlaPrefillAttention` at all, both keep `std::nullopt`, and their + // loops keep the full-context bounds they had. Every DeepSeek / MiniCPM3 / + // Kimi-Linear registration leaves it 0. + // + // > 0 is `sliding_window_size` — 513 on dots3-note's 33 `sliding_attention` + // layers (`vllm/models/dots3_note/nvidia/model.py:456` passes it to + // `MLAAttention`; `attention.py:439-468` @ `bc2d63e650` is the impl subclass + // that keeps it). It reaches the two ops as `AttentionWindow{W - 1, 0}`, + // which is literally the pair upstream hands FlashAttention on the prefill + // half (`attention.py:300`) and exactly the key set its decode mask keeps + // (`:151-152`). + // + // WHAT IT DOES NOT COVER, and the seam refuses it BY NAME rather than + // serving a wrong answer: a windowed prefill that also has CHUNKED CONTEXT. + // Upstream never builds one — a sliding layer's prefill gathers only + // `min(seq_len, query_len + W - 1)` keys and runs ONE varlen call per chunk + // of requests (`attention.py:206, :594-654`), so the chunked-context merge + // this seam inherits from `DeepseekV2` has no windowed counterpart upstream + // to mirror. Owed to the row; see `.agents/specs/dots3-note.md` `## Owed`. + // + // There is deliberately NO `has_sliding_window()` accessor beside this field, + // unlike `has_q_lora()` below. Every consumer wants the VALUE, not the + // predicate: `ForwardMlaAttentionBlock` passes `dims.sliding_window` to + // `ForwardMlaPrefillMha` and assigns it to `impl.sliding_window` + // unconditionally, precisely so a 0 cannot be skipped and leave a previous + // layer's 513 in place. W4b-2 shipped the accessor with no caller in `src`, + // `include` or `tests`, and its review removed it under `## Nothing lands + // dead` rather than inventing a call site for it. + int64_t sliding_window = 0; + // `self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim` (:969) — 192. int64_t qk_head_dim() const { return qk_nope_head_dim + qk_rope_head_dim; } // The MLA cache head_size `kv_lora_rank + qk_rope_head_dim` diff --git a/include/vllm/v1/attention/backend.h b/include/vllm/v1/attention/backend.h index b49cc57b7..846dcbe9f 100644 --- a/include/vllm/v1/attention/backend.h +++ b/include/vllm/v1/attention/backend.h @@ -644,6 +644,24 @@ class TritonMLAImpl final : public AttentionImpl { // runner's queue in when the DeepSeek-V2 forward lands. vt::Queue* queue = nullptr; + // ─── the SLIDING-WINDOW arm (dots3-note W4b-2, #699) ────────────────────── + // Upstream expresses this as a SUBCLASS: `Dots3NoteTritonMLAImpl( + // TritonMLAImpl)` passes `sliding_window=None` up to the base — so the base's + // own rejection at `triton_mla.py:165-171` still fires for everyone else — + // and keeps the value on itself as `self.sliding_window` + // (`vllm/models/dots3_note/nvidia/attention.py:439-468` @ `bc2d63e650`), + // which `_forward_swa_mqa` (`:470-563`) then reads. This tree has one MLA + // impl and a registry that hands it out, so the subclass is expressed as a + // FIELD whose absent state is 0 — the additive shape used for every other + // family-specific MLA extension here. + // + // 0 leaves `MlaDecodeAttentionArgs::window_size` at `std::nullopt`, i.e. the + // full-context decode every DeepSeek / MiniCPM3 / Kimi-Linear caller gets. + // `layer.window_size` — the per-LAYER window on `AttentionLayer` — stays + // REFUSED in `forward_mqa`, because that is the base class's rejection and + // dots3-note does not set it either. + int64_t sliding_window = 0; + // The DECODE entry point — the 1:1 counterpart of `forward_mqa` // (triton_mla.py:189-260). `q` is the already-concatenated // [num_reqs, num_heads, kv_lora_rank + qk_rope_head_dim] query (upstream diff --git a/include/vt/ops.h b/include/vt/ops.h index 331f14f45..2a24f4d62 100644 --- a/include/vt/ops.h +++ b/include/vt/ops.h @@ -1401,6 +1401,34 @@ struct MlaDecodeAttentionArgs { // used to derive `num_kv_splits` when that is 0; an upper bound is safe. When // both are 0 the impl falls back to 1 split. int32_t max_seq_len = 0; + // ─── the SLIDING-WINDOW decode arm (dots3-note W4b-2, #699) ─────────────── + // OPTIONAL local-attention bounds, the same `AttentionWindow` convention + // `PagedAttentionArgs::window_size` already uses: for the bottom-right + // aligned absolute query position `p = seq_len - 1` (MLA decode is ONE query + // per row), visible keys are `[p - left, p + right]` intersected with + // `[0, seq_len)`. `std::nullopt` — every DeepSeek / MiniCPM3 / Kimi-Linear + // registration — leaves the full-context loop byte-identical: the window is + // not a mask applied afterwards, it is the loop's START BOUND, so an absent + // window is a NOT-TAKEN branch rather than a no-op. + // + // UPSTREAM. `TritonMLAImpl` itself REJECTS a sliding window + // (triton_mla.py:165-171); dots3-note SUBCLASSES it — + // `Dots3NoteTritonMLAImpl.__init__` passes `sliding_window=None` to super and + // keeps the value on itself (`vllm/models/dots3_note/nvidia/attention.py` + // :439-468 @ `bc2d63e650`), then `_forward_swa_mqa` (`:470-563`) gathers a + // window-bounded slice of the paged latent and masks the scores with + // `kv_positions >= query_position - WINDOW_SIZE + 1` (`:152`) and + // `kv_positions <= query_position` (`:151`). `WINDOW_SIZE` is + // `sliding_window_size` = 513, so `left == sliding_window - 1`, matching the + // `window_size=(sliding_window - 1, 0)` upstream hands FlashAttention on the + // PREFILL half (`:300`). The gather is upstream's Triton WORKSPACE strategy; + // the paged kernels here read the block table directly over the same key + // range, which is the same function with no gather and no mask. + // + // `right` must be 0: an MLA decode query IS the last position of its own + // sequence, so a positive right bound could only admit keys that do not + // exist. Anything else is refused BY NAME in ops.cpp rather than ignored. + std::optional window_size = std::nullopt; }; // Arguments for vt::MlaPrefillAttention (MLA campaign W5). Mirrors the scalar @@ -1423,6 +1451,26 @@ struct MlaPrefillAttentionArgs { // same fallback the FA-2 paged prefill launcher uses. int32_t max_seqlen_q = 0; int32_t max_seqlen_k = 0; + // ─── the SLIDING-WINDOW prefill arm (dots3-note W4b-2, #699) ────────────── + // OPTIONAL local-attention bounds, the `AttentionWindow` convention. Query + // `i` of a request whose query length is `Lq` and key length is `Lk` sits at + // the bottom-right aligned position `p = i + (Lk - Lq)`; visible keys are + // `[p - left, p + right]` intersected with `[0, Lk)`. + // + // UPSTREAM is literally this pair: `Dots3NoteFlashAttnPrefillBackend. + // run_sliding_window` calls `_flash_attn_varlen_diff_headdims(..., causal= + // True, window_size=(sliding_window - 1, 0))` + // (`vllm/models/dots3_note/nvidia/attention.py:279-305` @ `bc2d63e650`, the + // window at `:300`), so `left == sliding_window - 1` and `right == 0`. + // + // `causal` must be TRUE whenever this is set, and `right` must be 0. Both are + // refused BY NAME in ops.cpp rather than approximated: FlashAttention's local + // mask REPLACES the causal specialization (the adapter normalizes + // `is_causal = causal && !is_local`, cuda_flash_attn_fa2.cu:472-476), so a + // non-causal window would have to be spelled with an infinite right bound, + // which this struct cannot say. Upstream never asks for one — every windowed + // call it makes is the causal `(W-1, 0)` pair above. + std::optional window_size = std::nullopt; }; // Router SCORING function. softmax over all E is the Qwen3.6 / DeepSeek-V2 diff --git a/src/vllm/model_executor/layers/attention/mla_attention.cpp b/src/vllm/model_executor/layers/attention/mla_attention.cpp index 65a6cbaab..b0154d708 100644 --- a/src/vllm/model_executor/layers/attention/mla_attention.cpp +++ b/src/vllm/model_executor/layers/attention/mla_attention.cpp @@ -124,6 +124,15 @@ void MlaBlockDims::Validate() const { "q_a_layernorm output to rescale on the DIRECT q_proj branch " "(deepseek_v2.py:1028-1034)"); } + // dots3-note's sliding window (#699 W4b-2). 0 is ABSENT; a negative value is + // a caller that computed `sliding_window - 1` one layer too early, which + // would otherwise reach the ops as a window that admits nothing. + if (sliding_window < 0) { + throw std::invalid_argument( + "MlaBlockDims: sliding_window must be >= 0 (0 means ABSENT — the full " + "context; dots3-note's sliding layers set `sliding_window_size` 513, " + "model.py:456)"); + } } // mla_attention.py:880-900 + :959-962. Upstream's chain is @@ -546,11 +555,16 @@ void ForwardMlaAttentionBlock(Dev d, const MlaBlockDims& dims, const MlaBlockWei MlaUpProjectFn up = MakeMlaUpProjectFn(d, dims, w, up_scratch); Tensor suffix_out_t = suffix_out.t(), suffix_lse_t = suffix_lse.t(); Tensor kv_cache_ro = kv_cache; + // dots3-note's sliding layers (#699 W4b-2). 0 — every DeepSeek / MiniCPM3 / + // Kimi-Linear registration — leaves the call byte-identical; > 0 becomes + // the `(W - 1, 0)` FlashAttention pair upstream's `run_sliding_window` + // passes (attention.py:300 @ bc2d63e650), and refuses a windowed prefill + // that also has chunked context BY NAME. ForwardMlaPrefillMha(d.q, prefill_out, q_prefill, key_t, value, kv_cache_ro, meta.prefill_block_table, meta.prefill_cu_seqlens_q, meta.chunks, up, dims.scale, meta.max_query_len, meta.prefill_tokens_with_context, bufs, suffix_out_t, - suffix_lse_t); + suffix_lse_t, dims.sliding_window); } // ─── 5b. DECODE — the ABSORBED MQA form (mla_attention.py:739-830) ─────── @@ -582,6 +596,14 @@ void ForwardMlaAttentionBlock(Dev d, const MlaBlockDims& dims, const MlaBlockWei impl.head_size = static_cast(dims.head_size()); impl.scale = dims.scale; impl.queue = &d.q; // W4 deviation (i), wired here. + // dots3-note's windowed decode (#699 W4b-2). Upstream expresses it as the + // `Dots3NoteTritonMLAImpl` subclass keeping `self.sliding_window` + // (attention.py:439-468 @ bc2d63e650); here it is the impl's field, and 0 + // is every DeepSeek / MiniCPM3 / Kimi-Linear caller's value. It is assigned + // UNCONDITIONALLY rather than under a guard because `impl` is the caller's + // object and may be reused across layers of DIFFERENT kinds — a guard would + // let a sliding layer's 513 leak into the next full layer. + impl.sliding_window = dims.sliding_window; v1::AttentionLayer layer{}; impl.forward_mqa(layer, mqa_q_t, kv_cache, meta.decode, mqa_out_t, nullptr); // `self._v_up_proj(attn_out, out=mqa_output_slice)` (:830, :1024-1034): diff --git a/src/vllm/model_executor/models/dots3_note.cpp b/src/vllm/model_executor/models/dots3_note.cpp index fe0fd9e86..4393f8732 100644 --- a/src/vllm/model_executor/models/dots3_note.cpp +++ b/src/vllm/model_executor/models/dots3_note.cpp @@ -687,8 +687,31 @@ v1::KVCacheConfig MakeDots3NoteKVCache(const HfConfig& config, int block_size, // read (`attention.py`::Dots3NotePaddedSparseImpl._logical_cache). Reporting // 576 here would under-allocate every sliding layer by 512 rows. // - // The heterogeneous per-layer group split, the DSA index cache and the - // windowed metadata are W4's, and are NOT represented here. + // Three things are still NOT represented here, and W4b-2 changed which brick + // owes the first of them. The comment used to say all three were "W4's", + // which stopped being true the moment the sliding layers ran. + // + // the heterogeneous per-layer W4b-3. Upstream gives a sliding layer a + // KV-cache GROUP SPLIT `SlidingWindowMLASpec` + // (`mla_attention.py:1215-1219` @ + // `bc2d63e650`, fed + // `sliding_window=config.sliding_window_size` + // at `model.py:457`), which is a SECOND spec + // kind and therefore a second group. We emit + // one uniform `MLAAttentionSpec` for all 46 + // layers, so 33 of them hold a full-length + // latent cache where upstream caps them near + // the window. No correctness consequence — the + // window is applied on READ, and W4b-2's gate + // proves it — but it is the largest memory + // property of this architecture and a token + // gate cannot see it. `## Owed` carries it. + // the DSA index cache W4b-3, with the indexer's SELECTION + // the windowed metadata never, here: `_build_sliding_window_metadata` + // is upstream's Triton gather bound, and the + // port walks the paged block table instead + // (spec §4.8), so there is no metadata to + // allocate. const Dots3NoteParams p = ParseDots3NoteParams(config); v1::KVCacheConfig kv; kv.num_blocks = num_blocks; diff --git a/src/vllm/model_executor/models/dots3_note.h b/src/vllm/model_executor/models/dots3_note.h index 7838b0268..e8a08a148 100644 --- a/src/vllm/model_executor/models/dots3_note.h +++ b/src/vllm/model_executor/models/dots3_note.h @@ -323,6 +323,12 @@ struct Dots3NoteDenseMlp { }; struct Dots3NoteLayerDeviceWeights { + // WHICH of the two attention geometries this layer runs — `config.layer_types + // [layer_idx] == "sliding_attention"` selects `Dots3NoteSlidingAttention` + // over `Dots3NoteFullAttention` upstream (`model.py:501-505` @ + // `bc2d63e650`). Every tensor in `attn` is shaped by this, so it is stored + // beside them rather than re-read from the params at each use (W4b-2, #699). + Dots3NoteLayerKind kind = Dots3NoteLayerKind::kFullAttention; OwnedTensor input_layernorm; // [hidden] OwnedTensor post_attention_layernorm; // [hidden] Dots3NoteMlaLayerWeights attn; @@ -335,11 +341,24 @@ struct Dots3NoteLayerDeviceWeights { // accounting still runs, and the forward refuses BY NAME. struct Dots3NoteDeviceWeights { bool present = false; + // The FULL-attention geometry (13 of 46 layers). mla::MlaBlockDims mla{}; + // The SLIDING geometry (33 of 46 layers) — a DIFFERENT head count, latent + // rank, NoPE width, rope theta and softmax scale, plus the 513 window. Not a + // parameterisation of the one above; see the table in + // `.agents/specs/dots3-note.md` §4.7. `sliding_window == 0` here means the + // config has no sliding layer and the struct is unused. (W4b-2, #699.) + mla::MlaBlockDims swa_mla{}; OwnedTensor embed_tokens; // [vocab, hidden] (embed lookup; NOT transposed) OwnedTensor final_norm; // [hidden] OwnedTensor lm_head; // [hidden, vocab] Matmul-B; EMPTY when tied - OwnedTensor rope_cos_sin_cache; // [max_position_embeddings, qk_rope_head_dim] + // TWO rope caches, because the two geometries carry DIFFERENT thetas — 8e7 + // on the full layers and `swa_rope_theta` 5e4 on the sliding ones + // (`model.py:401-409` @ `bc2d63e650`). Sharing one would be numerically + // silent, which is why they are separate tensors and not one with a flag. + // Each is EMPTY when no layer of that kind exists. + OwnedTensor rope_cos_sin_cache; // [max_position_embeddings, qk_rope_head_dim] + OwnedTensor swa_rope_cos_sin_cache; // [max_position_embeddings, swa_qk_rope_head_dim] std::vector layers; }; @@ -393,6 +412,15 @@ Dots3NoteDeviceWeights MaterializeDots3NoteDevice( // the forward builds instead of typing one by hand. mla::MlaBlockDims Dots3NoteFullAttnMlaDims(const Dots3NoteParams& params); +// The `mla::MlaBlockDims` a dots3-note SLIDING layer runs +// (`model.py`::Dots3NoteSlidingAttention.__init__ :341-460 @ `bc2d63e650`): +// `swa_*` geometry throughout, the softmax scale is a PLAIN `qk_head_dim**-0.5` +// with no YaRN and no mscale (`:446`), the rope is `rope_type="default"` at +// `swa_rope_theta` (`:401-409`), and `sliding_window` carries +// `config.sliding_window_size` (`:457`). Exported for the same reason the full +// one is: the gate drives the struct the forward builds, never one it typed. +mla::MlaBlockDims Dots3NoteSlidingAttnMlaDims(const Dots3NoteParams& params); + // W1 loader: resolves the config, accounts for 100% of the checkpoint's tensors // (refusing BY NAME on the first unclaimed or missing one), and returns an // UNMATERIALIZED model. It never returns a half-built tower and never silently diff --git a/src/vllm/model_executor/models/dots3_note_device.cpp b/src/vllm/model_executor/models/dots3_note_device.cpp index 27e3f5b41..bd610118a 100644 --- a/src/vllm/model_executor/models/dots3_note_device.cpp +++ b/src/vllm/model_executor/models/dots3_note_device.cpp @@ -19,24 +19,54 @@ // `ModelRegistry::Forward`. Everything else — the released checkpoint // included — still refuses BY NAME, naming the brick that owes it. // +// ─── W4b-2 ADDED THE SLIDING ARM, AND THE PADDED ROW ───────────────────────── +// The 33 `sliding_attention` layers now run, through the SAME +// `mla::ForwardMlaAttentionBlock` over a SECOND `mla::MlaBlockDims` +// (`Dots3NoteSlidingAttnMlaDims`) whose `sliding_window` becomes the +// `AttentionWindow{W - 1, 0}` pair `vt::MlaDecodeAttention` and +// `vt::MlaPrefillAttention` learned at W4b-2. The physical MLA cache row is the +// PADDED 1088 both classes share, and the full layers read their logical 576 +// out of the head of it with `Tensor::Slice(2, 0, ...)` — upstream's +// `Dots3NotePaddedSparseImpl._logical_cache` (attention.py:700-702), and no +// `vt` op changed to make it work. +// // ─── WHAT IS STILL REFUSED, AND BY WHICH BRICK ─────────────────────────────── -// sliding_attention layers W4b — the windowed metadata, the gather, the -// score mask, the padded/heterogeneous KV -// spec (spec §2.3) // MoE layers W5 — the ungrouped noaux_tc router at 256/8 -// seq_len > index_topk W4b — the DSA lightning indexer's SELECTION is +// seq_len > index_topk W4b-3 — the DSA lightning indexer's SELECTION is // not on the device path, so dense // attention is only the same answer while -// the top-k selects every causal candidate +// the top-k selects every causal candidate. +// Only asked of a config that HAS a full +// layer: the sliding layers carry no indexer +// (`self.indexer = None` / `is_sparse = +// False`, model.py:432-434) +// a windowed prefill with +// chunked CONTEXT W4b-3 — refused inside the seam; upstream caps a +// sliding layer's gather at the window +// instead of merging context chunks +// (attention.py:206, :594-654), so there is +// no windowed form of that merge to mirror +// a KV cache row that +// disagrees with the config kept — an engine allocates the cache separately // the vision / audio towers W6 / W7 — never part of the language forward // the nextn tail W10 // // ─── WHAT THIS IS A PORT OF (file:line on BOTH sides) ──────────────────────── // BEYOND-PIN. `dots3_note` does NOT exist at our parity pin `555967922` -// (0.26.0.dev0). Every anchor below was RE-DERIVED at upstream `origin/main` = -// `06ecec7a84`; `git log 06ecec7a84..origin/main -- vllm/models/dots3_note/` -// and `.. -- vllm/model_executor/models/deepseek_v2.py` are both EMPTY at the -// time of writing, so the same numbers hold at the newer head. +// (0.26.0.dev0). Every anchor below was first derived at upstream `06ecec7a84` +// and RE-DERIVED at `bc2d63e650`, which is the revision the row's spec and +// W4b-2 read; naming one number on both sides is the point of an anchor. +// MEASURED rather than asserted: `git diff 06ecec7a84 bc2d63e650 -- +// vllm/models/dots3_note/` is EMPTY, and the ONLY delta in +// `vllm/model_executor/models/deepseek_v2.py` is two lines inside +// `DeepseekV2Attention` — `q_lora_rank: int | None` at `:457` and +// `self.q_lora_rank` at `:495` — a class dots3-note does not subclass, and an +// equal-length edit, so every line number below is unmoved. +// +// The previous form of this comment said that deepseek_v2.py diff was EMPTY. +// It is not, and it was not silently wrong for free: a claim of emptiness goes +// stale without a symptom, while naming the two lines makes the next reader's +// check a one-command `git diff` instead of a re-derivation. // // OURS <- UPSTREAM // Dots3NoteFullAttnMlaDims <- `vllm/models/dots3_note/nvidia/model.py` @@ -131,6 +161,23 @@ mla::DeepseekYarnRopeParams FullAttnRope(const Dots3NoteParams& p) { return r; } +// The rope dots3-note's SLIDING layers run. `Dots3NoteSlidingAttention.__init__` +// builds `get_rope(..., rope_parameters={"rope_type": "default", "rope_theta": +// config.swa_rope_theta}, is_neox_style=False)` (model.py:401-409 @ +// `bc2d63e650`) — the SAME plain form as the full layers at a DIFFERENT theta, +// 5e4 against 8e7. Sharing the full layers' cache here is numerically silent +// and is spec §4 trap 6; it is a separate function so a reader sees the two +// side by side (W4b-2, #699). +mla::DeepseekYarnRopeParams SwaRope(const Dots3NoteParams& p) { + mla::DeepseekYarnRopeParams r; + r.yarn = false; + r.scaling_factor = 1.0; + r.base = p.swa.rope_theta; + r.rotary_dim = p.swa.qk_rope_head_dim; + r.original_max_position_embeddings = p.max_position_embeddings; + return r; +} + // The absorbed decode forms, exactly `MLAAttention.process_weights_after_ // loading` (mla_attention.py:1066-1196 @ 06ecec7a84; the two permutes are // :1178 and :1180). Both forms are kept: `kv_b_proj` feeds the @@ -247,18 +294,39 @@ mla::MlaBlockDims Dots3NoteFullAttnMlaDims(const Dots3NoteParams& p) { return d; } +mla::MlaBlockDims Dots3NoteSlidingAttnMlaDims(const Dots3NoteParams& p) { + const Dots3NoteAttnParams& w = p.swa; + mla::MlaBlockDims d; + d.hidden_size = p.hidden_size; + d.num_heads = w.num_attention_heads; + d.qk_nope_head_dim = w.qk_nope_head_dim; + d.qk_rope_head_dim = w.qk_rope_head_dim; + d.v_head_dim = w.v_head_dim; + d.kv_lora_rank = w.kv_lora_rank; + d.q_lora_rank = w.q_lora_rank; + d.rms_norm_eps = static_cast(p.rms_norm_eps); + d.is_neox_style = w.rope_is_neox_style; + d.q_lora_scale = w.q_lora_scale; + d.kv_lora_scale = w.kv_lora_scale; + // `scale=qk_head_dim**-0.5` (model.py:446) — the SLIDING arm builds + // `MLAAttention` with the bare inverse square root and NO rope_scaling block + // at all, so there is no YaRN ramp and no mscale^2. Passing the rope through + // `MlaAttentionScale` yields the same number because `SwaRope().yarn` is + // false, and it is routed that way rather than written by hand so the full + // and sliding arms cannot drift apart on the one factor that is silent. + d.scale = mla::MlaAttentionScale(d, SwaRope(p)); + // `sliding_window=config.sliding_window_size` (model.py:457) — 513. + d.sliding_window = w.sliding_window; + d.Validate(); + return d; +} + int64_t Dots3NoteDenseEquivalentMaxSeqLen(const Dots3NoteParams& params) { return params.index_topk; } std::string Dots3NoteDeviceRefusal(const Dots3NoteParams& p) { for (size_t l = 0; l < p.layer_types.size(); ++l) { - if (p.layer_types[l] == Dots3NoteLayerKind::kSlidingAttention) { - return "layer " + std::to_string(l) + - " is `sliding_attention` — the sliding-window MLA (the windowed " - "metadata, the KV gather, the score mask and the padded/" - "heterogeneous KV spec of `nvidia/attention.py`) is W4b"; - } if (p.is_moe_layer(static_cast(l))) { return "layer " + std::to_string(l) + " is a MoE layer — the ungrouped noaux_tc router at " + @@ -266,24 +334,25 @@ std::string Dots3NoteDeviceRefusal(const Dots3NoteParams& p) { std::to_string(p.num_experts_per_tok) + " plus the shared expert is W5"; } } - // The PADDED physical latent row. `MakeDots3NoteKVCache` reports the row both - // attention classes share — `swa_kv_lora_rank + swa_qk_rope_head_dim` - // (model.py:204-217) — and the full layers read their own logical width out - // of the head of it. Narrowing on read is - // `Dots3NotePaddedSparseImpl._logical_cache`, and it is W4b. + // ─── LIFTED at W4b-2 (#699): the sliding layer and the PADDED row ───────── + // W4a refused both here. Both now run. // - // This is checked HERE, at config level, and not only at the forward — review - // finding F5. The forward's own cache-row assertion stays (an engine can hand - // a cache that disagrees with the config it was built from, and a test does - // exactly that), but leaving the config case to it meant the LOADER - // materialized a whole tower for a config the very next call refuses. - if (p.physical_latent_row() != p.full.latent_row()) { - return "the physical MLA cache row is " + - std::to_string(p.physical_latent_row()) + " but the full layers read " + - std::to_string(p.full.latent_row()) + - " — narrowing a PADDED row back to the logical one " - "(`Dots3NotePaddedSparseImpl._logical_cache`) is W4b"; - } + // The SLIDING layer runs through the same `mla::ForwardMlaAttentionBlock` + // over `Dots3NoteSlidingAttnMlaDims`, whose `sliding_window` reaches + // `vt::MlaDecodeAttention` and `vt::MlaPrefillAttention` as the + // `AttentionWindow{W - 1, 0}` pair upstream hands FlashAttention + // (attention.py:300 @ bc2d63e650). + // + // The PADDED physical row runs because the MLA cache ops are STRIDE-DRIVEN: + // `Tensor::Slice(2, 0, logical)` shrinks `shape[2]` and KEEPS both leading + // strides (tensor.cpp), which IS upstream's `kv_cache[..., : self.head_size]` + // (`Dots3NotePaddedSparseImpl._logical_cache`, attention.py:700-702). The + // narrowing is ONE line in the forward and no `vt` op changed. W4b-1 first + // claimed the ops address the cache contiguously; that was false, and the + // correction is spec §4.7. + // + // WHAT STILL REFUSES here is only what has no upstream form to mirror yet. + // The nextn tail. `Dots3NoteMTPModel` is deliberately not registered and the // backbone forward has no place to put an extra block, so a checkpoint that // ships one is refused rather than silently having it enumerated, loaded and @@ -308,10 +377,19 @@ Dots3NoteDeviceWeights MaterializeDots3NoteDevice( Dots3NoteDeviceWeights w; w.mla = Dots3NoteFullAttnMlaDims(p); - const mla::MlaBlockDims& d = w.mla; + w.swa_mla = Dots3NoteSlidingAttnMlaDims(p); const int64_t H = p.hidden_size, V = p.vocab_size, I = p.intermediate_size; - const int64_t N = d.num_heads, R = d.qk_rope_head_dim, L = d.kv_lora_rank; - const int64_t QL = d.q_lora_rank; + // Which geometries this config actually uses. A rope cache is 2 * 64 bytes + // per position over `max_position_embeddings` — 64 MiB each at the released + // 524288 — so the unused one is never built (W4b-2, #699). + bool any_full = false, any_sliding = false; + for (const Dots3NoteLayerKind k : p.layer_types) { + if (k == Dots3NoteLayerKind::kSlidingAttention) { + any_sliding = true; + } else { + any_full = true; + } + } w.embed_tokens = LoadBf16Direct(get, "model.embed_tokens.weight"); RequireShape(w.embed_tokens, "model.embed_tokens.weight", {V, H}); @@ -322,20 +400,35 @@ Dots3NoteDeviceWeights MaterializeDots3NoteDevice( w.lm_head = LoadBf16Transposed(get, "lm_head.weight"); } - // `_compute_cos_sin_cache` over the whole positional range, once per model. - { + // `_compute_cos_sin_cache` over the whole positional range, once per model — + // and ONCE PER GEOMETRY, because the two thetas differ (8e7 against + // `swa_rope_theta` 5e4, model.py:230-238 / :401-409). + const auto build_rope = [&](const mla::DeepseekYarnRopeParams& rp) { const std::vector cache = - mla::BuildDeepseekRopeCosSinCache(FullAttnRope(p), p.max_position_embeddings); - w.rope_cos_sin_cache = MakeOwned(DType::kBF16, {p.max_position_embeddings, R}); - auto* dst = reinterpret_cast(w.rope_cos_sin_cache.bytes.data()); + mla::BuildDeepseekRopeCosSinCache(rp, p.max_position_embeddings); + OwnedTensor t = + MakeOwned(DType::kBF16, {p.max_position_embeddings, rp.rotary_dim}); + auto* dst = reinterpret_cast(t.bytes.data()); for (size_t i = 0; i < cache.size(); ++i) dst[i] = vt::F32ToBF16(cache[i]); - } + return t; + }; + if (any_full) w.rope_cos_sin_cache = build_rope(FullAttnRope(p)); + if (any_sliding) w.swa_rope_cos_sin_cache = build_rope(SwaRope(p)); w.layers.resize(static_cast(p.num_hidden_layers)); for (int64_t l = 0; l < p.num_hidden_layers; ++l) { const std::string pre = "model.layers." + std::to_string(l) + "."; const std::string sa = pre + "self_attn."; Dots3NoteLayerDeviceWeights& lw = w.layers[static_cast(l)]; + // `attention_cls = Dots3NoteSlidingAttention if config.layer_types[ + // layer_idx] == "sliding_attention" else Dots3NoteFullAttention` + // (model.py:501-505 @ bc2d63e650). EVERY shape below is this choice's, + // which is why the kind is resolved before the first tensor is read. + lw.kind = p.kind_of(l); + const bool sliding = lw.kind == Dots3NoteLayerKind::kSlidingAttention; + const mla::MlaBlockDims& d = sliding ? w.swa_mla : w.mla; + const int64_t N = d.num_heads, R = d.qk_rope_head_dim, L = d.kv_lora_rank; + const int64_t QL = d.q_lora_rank; lw.input_layernorm = LoadBf16Direct(get, pre + "input_layernorm.weight"); RequireShape(lw.input_layernorm, pre + "input_layernorm.weight", {H}); lw.post_attention_layernorm = @@ -395,9 +488,10 @@ ForwardLogits Dots3NoteModel::ForwardDevice( const std::string why = Dots3NoteDeviceRefusal(p); VT_CHECK(why.empty(), "Dots3NoteForCausalLM forward: not ported — " + why + - ". W4a covers the full-attention layer with a dense MLP only; the " - "sliding-window MLA is W4, the MoE is W5, the vision/audio towers " - "are W6/W7. See .agents/specs/dots3-note.md and issue #699."); + ". W4a/W4b-2 cover BOTH attention geometries — full and " + "sliding-window — with a dense MLP; the MoE is W5, the " + "vision/audio towers are W6/W7, the nextn tail is W10. See " + ".agents/specs/dots3-note.md and issue #699."); VT_CHECK(weights.materialized && weights.device.present, "Dots3NoteForCausalLM forward: the language tower was not " "materialized — the loader only materializes a config the device " @@ -407,15 +501,28 @@ ForwardLogits Dots3NoteModel::ForwardDevice( // candidate and dense attention IS upstream's answer; past that it is a // different answer, and W3 measured that difference at 0.392 on the layer // output. Refuse rather than serve dense attention on a sparse model. + // + // W4b-2 narrows WHO this is asked of, and the narrowing is upstream's own + // statement rather than a convenience: `Dots3NoteSlidingAttention` sets + // `self.indexer = None` and `is_sparse = False` (model.py:432-434), so a + // sliding layer has no selection to get wrong and a config with no FULL layer + // has no DSA anywhere. The released checkpoint has 13 full layers and is + // unaffected. + const bool has_full_layer = + std::any_of(p.layer_types.begin(), p.layer_types.end(), [](Dots3NoteLayerKind k) { + return k == Dots3NoteLayerKind::kFullAttention; + }); const int64_t topk = Dots3NoteDenseEquivalentMaxSeqLen(p); - for (int32_t sl : attn_meta.seq_lens) { - VT_CHECK(static_cast(sl) <= topk, - "Dots3NoteForCausalLM forward: a request needs " + std::to_string(sl) + - " keys but `index_topk` is " + std::to_string(topk) + - " — past that the DSA lightning indexer PRUNES " - "(model.py:171), and the sparse selection is not on the " - "device path yet (W4b). Refusing rather than serving dense " - "attention on a sparse model. See issue #699."); + if (has_full_layer) { + for (int32_t sl : attn_meta.seq_lens) { + VT_CHECK(static_cast(sl) <= topk, + "Dots3NoteForCausalLM forward: a request needs " + std::to_string(sl) + + " keys but `index_topk` is " + std::to_string(topk) + + " — past that the DSA lightning indexer PRUNES " + "(model.py:171), and the sparse selection is not on the " + "device path yet (W4b-3). Refusing rather than serving dense " + "attention on a sparse model. See issue #699."); + } } const Dots3NoteDeviceWeights& dw = weights.device; @@ -447,28 +554,85 @@ ForwardLogits Dots3NoteModel::ForwardDevice( const int64_t block_size = attn_kv[0].block_size; MlaStep step = BuildMlaStep(d, positions, attn_meta, block_size, p.max_position_embeddings); - const Tensor rope = ResidentWeight(d, dw.rope_cos_sin_cache); - step.rope_cache = &rope; + // One resident rope cache per GEOMETRY, and each is made resident ONLY when + // the config has a layer of that kind. `MaterializeDots3NoteDevice` leaves + // the other one EMPTY to avoid building a 64 MiB table nothing reads, and + // since #1953 `ResidentWeight` REFUSES an empty weight BY NAME rather than + // aliasing a null host pointer — so the guard is the contract, not an + // optimization. A layer only ever reads its own kind's cache, so the one + // that stays an empty `Tensor` here is never dereferenced. + Tensor rope_full{}; + Tensor rope_swa{}; + if (!dw.rope_cos_sin_cache.Empty()) rope_full = ResidentWeight(d, dw.rope_cos_sin_cache); + if (!dw.swa_rope_cos_sin_cache.Empty()) { + rope_swa = ResidentWeight(d, dw.swa_rope_cos_sin_cache); + } + // `MlaStep::rope_cache` is deliberately LEFT NULL here, and that is a + // statement rather than an omission. It exists for the one-geometry models — + // `deepseek_v2.cpp:500`, `minicpm3.cpp:200`, `kimi_linear_device.cpp:2210` + // each dereference it — because one per-MODEL table serves every layer there. + // dots3-note has TWO tables, and which one a layer reads is a property of the + // LAYER (`layer_rope` below), so no single per-model value is correct. + // Assigning `&rope_full` here would hand a future `*step.rope_cache` reader + // the FULL arm's rope on a sliding layer — a silently wrong answer — or, on a + // pure-SWA config, an EMPTY `Tensor` whose failure surfaces somewhere else + // entirely. A null pointer fails at the first read, loudly, in this function. v1::TritonMLAImpl impl; const float eps = static_cast(p.rms_norm_eps); + // The PHYSICAL MLA cache row both attention classes share: + // `physical_head_size = swa_kv_lora_rank + swa_qk_rope_head_dim` + // (`Dots3NotePaddedMLAAttention.get_kv_cache_spec`, model.py:204-216, fed at + // :283). 1088 on the released config, against the full layers' logical 576. + const int64_t physical_row = p.physical_latent_row(); for (int64_t l = 0; l < p.num_hidden_layers; ++l) { const Dots3NoteLayerDeviceWeights& lw = dw.layers[static_cast(l)]; const PagedKvCache& kv = attn_kv[static_cast(l)]; - // The PHYSICAL row is the padded 1088 both classes share - // (`physical_latent_row()`, model.py:204-217); the full layers read their - // logical 576 out of the head of it. W4a runs a schedule with no sliding - // layer, so the two coincide unless the config pads deliberately — and a - // padded row is exactly what `_logical_cache` narrows upstream, which is - // W4b. Refuse the padded case by name rather than reading a wrong stride. - VT_CHECK(kv.num_kv_heads == 1 && kv.head_size == dw.mla.head_size(), + const bool sliding = lw.kind == Dots3NoteLayerKind::kSlidingAttention; + const mla::MlaBlockDims& ld = sliding ? dw.swa_mla : dw.mla; + const Tensor& layer_rope = sliding ? rope_swa : rope_full; + // The PER-STEP cache-row check STAYS, and it is not the config-level one + // W4b-2 lifted: an engine allocates the KV cache separately from the config + // it was built from, so a row that disagrees is an input this forward can + // see and the config cannot. It now compares against the PHYSICAL row, + // because that is what the allocator is told to give + // (`MakeDots3NoteKVCache`). + VT_CHECK(kv.num_kv_heads == 1 && kv.head_size == physical_row, "dots3-note forward: the MLA cache row is " + - std::to_string(kv.head_size) + " but this layer reads " + - std::to_string(dw.mla.head_size()) + - " — narrowing a PADDED physical row back to the logical one " - "(`Dots3NotePaddedSparseImpl._logical_cache`) is W4b"); + std::to_string(kv.head_size) + " but this model's PHYSICAL row is " + + std::to_string(physical_row) + + " (`physical_head_size = swa_kv_lora_rank + swa_qk_rope_head_dim`, " + "model.py:204-216) — refusing rather than reading a wrong stride"); Tensor kv_cache = MakeTensor(kv.data, kv.dtype, d.q.device, - {kv.num_blocks, kv.block_size, kv.head_size}); + {kv.num_blocks, kv.block_size, physical_row}); + // `Dots3NotePaddedSparseImpl._logical_cache` (attention.py:700-702): + // `kv_cache[..., : self.head_size]`. `Tensor::Slice` shrinks `shape[2]` and + // KEEPS both leading strides, and every MLA cache op reads those strides + // from the tensor, so the narrowed view addresses the padded rows correctly + // with no op change. A SLIDING layer's logical row IS the physical one by + // construction (the padding exists for the full layers), so the slice is + // the identity there and is written unconditionally rather than branched. + // + // THIS ONE IS UNREACHABLE, and saying so is cheaper than leaving the next + // reader to discover it with a mutation. Unlike the cache-row check above, + // both sides come from the SAME parsed config: `physical_latent_row()` IS + // `swa.latent_row()` (`dots3_note.h:192`), so on a SLIDING layer the + // comparison is an identity; and on a FULL layer + // `ParseDots3NoteParams` has already refused + // `physical_latent_row() < full.latent_row()` at load + // (`dots3_note.cpp:389`, gated at + // `tests/vllm/models/test_dots3_note_scaffold.cpp:720-722`). No input the + // loader accepts can make it fire, which is why the W4b-2 review's + // mutation of it SURVIVED. It is kept as the executable spelling of + // upstream's own `assert` and is listed under `## Owed` as an untested + // assertion rather than presented as a gated refusal. + VT_CHECK(ld.head_size() <= physical_row, + "dots3-note forward: a layer reads " + std::to_string(ld.head_size()) + + " latent lanes but the physical row is only " + + std::to_string(physical_row) + + " — upstream asserts `physical_head_size >= self.head_size` " + "(model.py:210)"); + kv_cache = kv_cache.Slice(2, 0, ld.head_size()); // The residual add + RMSNorm goes through the SHARED `vt::FusedChain` // catalog (AGENTS.md: route model fusion through vt::FusedChain), with the @@ -486,8 +650,8 @@ ForwardLogits Dots3NoteModel::ForwardDevice( DBuf attn(d, DType::kBF16, {T, H}); Tensor attn_t = attn.t(); - const mla::MlaBlockWeights mw = ResidentMla(d, lw.attn, rope); - mla::ForwardMlaAttentionBlock(d, dw.mla, mw, dhn.t(), step.positions, kv_cache, + const mla::MlaBlockWeights mw = ResidentMla(d, lw.attn, layer_rope); + mla::ForwardMlaAttentionBlock(d, ld, mw, dhn.t(), step.positions, kv_cache, step.slot_mapping, step.meta, impl, attn_t); DBuf dh2(d, DType::kBF16, {T, H}); diff --git a/src/vllm/v1/attention/backend.cpp b/src/vllm/v1/attention/backend.cpp index 427b38288..98d129684 100644 --- a/src/vllm/v1/attention/backend.cpp +++ b/src/vllm/v1/attention/backend.cpp @@ -312,6 +312,14 @@ void TritonMLAImpl::forward_mqa(const AttentionLayer& layer, const vt::Tensor& q args.scale = scale; // `:253` self.scale args.num_kv_splits = metadata.num_kv_splits; args.max_seq_len = metadata.max_seq_len; + // dots3-note's windowed decode (#699 W4b-2): `_forward_swa_mqa` keeps keys + // `kv_pos >= query_pos - WINDOW_SIZE + 1` (attention.py:152 @ bc2d63e650), + // i.e. the inclusive left distance is `sliding_window - 1` — the same pair + // upstream hands FlashAttention on the prefill half (`:300`). 0 leaves this + // `std::nullopt`, which is the full-context loop the op already had. + if (sliding_window > 0) { + args.window_size = vt::AttentionWindow{static_cast(sliding_window - 1), 0}; + } // `:242-259` decode_attention_fwd(q, kv_c_and_k_pe_cache, kv_c_cache, o, lse, // block_table, seq_lens, attn_logits, num_kv_splits, scale, PAGE_SIZE, ...) // — the two "K" and "V" arguments are the SAME buffer, which our single diff --git a/src/vt/cpu/cpu_mla_attn.cpp b/src/vt/cpu/cpu_mla_attn.cpp index 0fd60ab86..53b7a413b 100644 --- a/src/vt/cpu/cpu_mla_attn.cpp +++ b/src/vt/cpu/cpu_mla_attn.cpp @@ -80,8 +80,28 @@ void MlaDecodeAttentionKernel(Queue&, Tensor& out, Tensor* lse, const Tensor& qu std::vector acc(static_cast(v_head_dim)); std::vector q_row(static_cast(head_size)); + // The SLIDING-WINDOW start bound (dots3-note W4b-2, #699). An MLA decode + // query is the LAST position of its own sequence, `p = seq_len - 1`, so + // `AttentionWindow{left, 0}` admits keys `[p - left, p]` — i.e. the loop + // START moves and nothing else does. Absent (`std::nullopt`, which is every + // DeepSeek / MiniCPM3 / Kimi-Linear caller) the start stays 0 and this + // kernel is the byte-identical full-context loop it was. + // + // UPSTREAM computes the same set the other way round: `_forward_swa_mqa` + // gathers `[max(seq_len - GATHER_LEN, 0), ...)` into a workspace + // (`vllm/models/dots3_note/nvidia/attention.py:76-79` @ `bc2d63e650`) and + // then masks with `kv_positions >= query_position - WINDOW_SIZE + 1` (`:152`) + // and `kv_positions <= query_position` (`:151`). GATHER_LEN is rounded up to + // the Triton tile (`:484`), so the gather is a SUPERSET and the mask is what + // makes it exact; walking the block table directly needs no superset and no + // mask, and reaches the identical key set. + const int64_t win_left = args.window_size.has_value() ? args.window_size->left : -1; + for (int64_t b = 0; b < batch; ++b) { const int64_t seq_len = seq[b]; + // `p - left` with `p = seq_len - 1`, clamped at 0. + const int64_t j_start = + win_left < 0 ? int64_t{0} : std::max(0, seq_len - 1 - win_left); for (int64_t h = 0; h < heads; ++h) { const int64_t q_off = b * query.stride[0] + h * query.stride[1]; for (int64_t d = 0; d < head_size; ++d) q_row[static_cast(d)] = LoadF(query.data, query.dtype, q_off + d); @@ -91,7 +111,7 @@ void MlaDecodeAttentionKernel(Queue&, Tensor& out, Tensor* lse, const Tensor& qu float l = 0.0f; std::fill(acc.begin(), acc.end(), 0.0f); - for (int64_t j = 0; j < seq_len; ++j) { + for (int64_t j = j_start; j < seq_len; ++j) { const int64_t blk_slot = j / block_size; VT_CHECK(blk_slot < max_blocks, "cpu mla_decode_attention: seq_len exceeds the block_table row"); diff --git a/src/vt/cpu/cpu_mla_prefill.cpp b/src/vt/cpu/cpu_mla_prefill.cpp index 6f87cc1ef..365c12466 100644 --- a/src/vt/cpu/cpu_mla_prefill.cpp +++ b/src/vt/cpu/cpu_mla_prefill.cpp @@ -81,6 +81,17 @@ void MlaPrefillAttentionKernel(Queue&, Tensor& out, Tensor* lse, const Tensor& q std::vector logits; std::vector acc(static_cast(v_head_dim)); + // The SLIDING-WINDOW lower bound (dots3-note W4b-2, #699). `left` is the + // inclusive distance behind the bottom-right aligned query position, exactly + // FlashAttention's `window_size=(left, right)` — which is what upstream hands + // it: `run_sliding_window(..., causal=True, window_size=(sliding_window - 1, + // 0))` (`vllm/models/dots3_note/nvidia/attention.py:279-305` @ `bc2d63e650`, + // the pair at `:300`). Absent (`std::nullopt`, every DeepSeek / MiniCPM3 / + // Kimi-Linear caller) the lower bound stays 0 and the loop is the + // byte-identical one it was; ops.cpp refuses a window with `causal=false`, so + // the upper bound below is always the causal one when this is set. + const int64_t win_left = args.window_size.has_value() ? args.window_size->left : -1; + for (int64_t b = 0; b < num_reqs; ++b) { const int64_t q_begin = qsl[b]; const int64_t q_end = qsl[b + 1]; @@ -100,13 +111,20 @@ void MlaPrefillAttentionKernel(Queue&, Tensor& out, Tensor* lse, const Tensor& q const int64_t visible = args.causal ? std::min(len_k, std::max(0, iq + causal_shift + 1)) : len_k; + // `p - left`, clamped at 0. `p = iq + causal_shift` is the same + // bottom-right position the causal bound above uses, so the two agree by + // construction rather than by two independent derivations. + const int64_t first = + win_left < 0 ? int64_t{0} + : std::min(visible, + std::max(0, iq + causal_shift - win_left)); for (int64_t h = 0; h < num_heads; ++h) { const int64_t q_off = t * query.stride[0] + h * query.stride[1]; // PASS 1 — the raw logits and their max. logits.assign(static_cast(visible), 0.0f); float m = -std::numeric_limits::infinity(); - for (int64_t j = 0; j < visible; ++j) { + for (int64_t j = first; j < visible; ++j) { const int64_t k_off = (k_begin + j) * key.stride[0] + h * key.stride[1]; float dot = 0.0f; for (int64_t d = 0; d < qk_head_dim; ++d) { @@ -120,13 +138,13 @@ void MlaPrefillAttentionKernel(Queue&, Tensor& out, Tensor* lse, const Tensor& q // PASS 2 — the exp-sum, then the weighted sum. No running rescale. float l = 0.0f; - for (int64_t j = 0; j < visible; ++j) { + for (int64_t j = first; j < visible; ++j) { const float p = std::exp(logits[static_cast(j)] - m); logits[static_cast(j)] = p; l += p; } std::fill(acc.begin(), acc.end(), 0.0f); - for (int64_t j = 0; j < visible; ++j) { + for (int64_t j = first; j < visible; ++j) { const int64_t v_off = (k_begin + j) * value.stride[0] + h * value.stride[1]; const float p = logits[static_cast(j)]; for (int64_t d = 0; d < v_head_dim; ++d) { diff --git a/src/vt/cuda/cuda_flash_attn_fa2.cu b/src/vt/cuda/cuda_flash_attn_fa2.cu index 0f5d5c851..d895cc8ee 100644 --- a/src/vt/cuda/cuda_flash_attn_fa2.cu +++ b/src/vt/cuda/cuda_flash_attn_fa2.cu @@ -882,11 +882,26 @@ void LaunchMlaPrefillFA2Bf16(cudaStream_t s, Tensor& out, float* lse_out, p.philox_args = at::PhiloxCudaState(0, 0); // `causal=True` for new tokens (flash_attn.py:223), `causal=False` for a - // context chunk (`:246`). No local window on any MLA path — TritonMLAImpl - // rejects `sliding_window` outright (triton_mla.py:165-171). - p.is_causal = args.causal; - p.window_size_left = -1; - p.window_size_right = args.causal ? 0 : -1; + // context chunk (`:246`). `TritonMLAImpl` itself rejects `sliding_window` + // (triton_mla.py:165-171), but dots3-note's SWA prefill backend asks FA for + // exactly one: `run_sliding_window` calls + // `_flash_attn_varlen_diff_headdims(..., causal=True, + // window_size=(sliding_window - 1, 0))` + // (`vllm/models/dots3_note/nvidia/attention.py:279-305` @ `bc2d63e650`, the + // pair at `:300`) — dots3-note W4b-2, #699. + // + // The NORMALIZATION is the paged launcher's, verbatim (`:466-476` above), and + // it is not optional: the pinned FA-2 API routes a finite window through the + // LOCAL specialization, whose compile-time `Is_causal` sibling deliberately + // IGNORES `window_size_left` (flash_fwd_launch_template.h LOCAL_SWITCH). A + // windowed call that kept `is_causal` true would therefore silently drop the + // left bound and attend the whole context. `std::nullopt` — every DeepSeek / + // MiniCPM3 / Kimi-Linear caller — takes the identical assignments this + // launcher had, so their dispatch is unchanged by construction. + const bool is_local = args.window_size.has_value(); + p.is_causal = args.causal && !is_local; + p.window_size_left = is_local ? args.window_size->left : -1; + p.window_size_right = is_local ? args.window_size->right : (args.causal ? 0 : -1); p.is_seqlens_k_cumulative = true; p.is_rotary_interleaved = false; p.rotary_dim = 0; @@ -904,19 +919,23 @@ void LaunchMlaPrefillFA2Bf16(cudaStream_t s, Tensor& out, float* lse_out, p.num_splits = 1; p.o_batch_stride = static_cast(max_seqlen_q) * p.o_row_stride; + // `run_causal` is `p.is_causal`, i.e. a finite window dispatches the + // NON-causal template so its runtime LOCAL_SWITCH picks the exact local mask + // (the paged launcher's `:500` rule, same reason). + const bool run_causal = args.causal && !is_local; if (d == 256) { - if (args.causal) { + if (run_causal) { FLASH_NAMESPACE::run_mha_fwd_splitkv_dispatch(p, s); } else { FLASH_NAMESPACE::run_mha_fwd_splitkv_dispatch(p, s); } } else if (d == 128) { - if (args.causal) { + if (run_causal) { FLASH_NAMESPACE::run_mha_fwd_splitkv_dispatch(p, s); } else { FLASH_NAMESPACE::run_mha_fwd_splitkv_dispatch(p, s); } - } else if (args.causal) { + } else if (run_causal) { FLASH_NAMESPACE::run_mha_fwd_splitkv_dispatch(p, s); } else { FLASH_NAMESPACE::run_mha_fwd_splitkv_dispatch(p, s); diff --git a/src/vt/cuda/cuda_mla_attn.cu b/src/vt/cuda/cuda_mla_attn.cu index 648380e6b..6ba6428d9 100644 --- a/src/vt/cuda/cuda_mla_attn.cu +++ b/src/vt/cuda/cuda_mla_attn.cu @@ -154,18 +154,28 @@ __global__ __launch_bounds__(kThreads) void MlaDecodeStage1( const int32_t* __restrict__ block_table, const int32_t* __restrict__ seq_lens, int64_t q_s0, int64_t q_s1, int64_t c_s0, int64_t c_s1, int64_t bt_s0, int head_size, int v_head_dim, int block_size, int num_heads, int valid_h, int num_splits, - int64_t mid_s0, int64_t mid_s1, int64_t mid_s2, float scale) { + int64_t mid_s0, int64_t mid_s1, int64_t mid_s2, float scale, int win_left) { const int b = static_cast(blockIdx.x); const int split = static_cast(blockIdx.z); const int warp = static_cast(threadIdx.x) >> 5; const int lane = static_cast(threadIdx.x) & 31; const int seq_len = seq_lens[b]; - // `:352-354`: kv_len_per_split = cdiv(seq_len, NUM_KV_SPLITS); - // start = kv_len_per_split * split_kv_id; + // The SLIDING-WINDOW start (dots3-note W4b-2, #699): the decode query is the + // LAST position of its sequence, so `AttentionWindow{left, 0}` keeps keys + // `[seq_len - 1 - left, seq_len)`. `win_left < 0` is the ABSENT state, which + // is every DeepSeek / MiniCPM3 / Kimi-Linear caller, and it puts `kv_start` + // back at 0 — the identical partition over `[0, seq_len)` this kernel had. + // The SPLIT GRID is partitioned over the WINDOWED range, not over the whole + // sequence, so the splits stay balanced instead of leaving all but the last + // one empty. Stage 2 recomputes the same partition from the same two inputs. + const int kv_start = win_left < 0 ? 0 : max(0, seq_len - 1 - win_left); + const int kv_len = seq_len - kv_start; + // `:352-354`: kv_len_per_split = cdiv(kv_len, NUM_KV_SPLITS); + // start = kv_start + kv_len_per_split * split_kv_id; // end = min(start + kv_len_per_split, seq_len). - const int per_split = (seq_len + num_splits - 1) / num_splits; - const int split_start = per_split * split; + const int per_split = (kv_len + num_splits - 1) / num_splits; + const int split_start = kv_start + per_split * split; const int split_end = min(split_start + per_split, seq_len); if (split_end <= split_start) return; // `:361` — block-uniform, so safe here @@ -287,11 +297,16 @@ __global__ void MlaDecodeStage2(T* __restrict__ out, float* __restrict__ lse, const float* __restrict__ mid, const int32_t* __restrict__ seq_lens, int64_t o_s0, int64_t o_s1, int64_t lse_s0, int64_t mid_s0, int64_t mid_s1, - int64_t mid_s2, int v_head_dim, int num_splits) { + int64_t mid_s2, int v_head_dim, int num_splits, + int win_left) { const int b = static_cast(blockIdx.x); const int h = static_cast(blockIdx.y); const int seq_len = seq_lens[b]; - const int per_split = (seq_len + num_splits - 1) / num_splits; + // The SAME partition stage 1 wrote, recomputed from the same two inputs + // (#699 W4b-2). Deriving it twice from `seq_lens` and `win_left` is what + // keeps the two kernels' notion of "which splits are empty" identical. + const int kv_start = win_left < 0 ? 0 : max(0, seq_len - 1 - win_left); + const int per_split = (seq_len - kv_start + num_splits - 1) / num_splits; const float* base = mid + b * mid_s0 + static_cast(h) * mid_s1; for (int d0 = 0; d0 < v_head_dim; d0 += static_cast(blockDim.x)) { @@ -301,7 +316,7 @@ __global__ void MlaDecodeStage2(T* __restrict__ out, float* __restrict__ lse, float l = 0.0f; float acc = 0.0f; for (int s = 0; s < num_splits; ++s) { - const int split_start = per_split * s; + const int split_start = kv_start + per_split * s; const int split_end = min(split_start + per_split, seq_len); if (split_end <= split_start) continue; // `:610` const float* p = base + static_cast(s) * mid_s2; @@ -451,7 +466,7 @@ void LaunchStage1(cudaStream_t s, dim3 grid, size_t smem, int n_tile, float* mid const int32_t* seq_lens, int64_t q_s0, int64_t q_s1, int64_t c_s0, int64_t c_s1, int64_t bt_s0, int head_size, int v_head_dim, int block_size, int num_heads, int valid_h, int num_splits, int64_t mid_s0, int64_t mid_s1, - int64_t mid_s2, float scale) { + int64_t mid_s2, float scale, int win_left) { if (n_tile == kNTile) { const void* fn = reinterpret_cast(&MlaDecodeStage1); if (smem > 48u * 1024u) { @@ -461,7 +476,8 @@ void LaunchStage1(cudaStream_t s, dim3 grid, size_t smem, int n_tile, float* mid } MlaDecodeStage1<<>>( mid, query, kv_cache, block_table, seq_lens, q_s0, q_s1, c_s0, c_s1, bt_s0, head_size, - v_head_dim, block_size, num_heads, valid_h, num_splits, mid_s0, mid_s1, mid_s2, scale); + v_head_dim, block_size, num_heads, valid_h, num_splits, mid_s0, mid_s1, mid_s2, scale, + win_left); } else { const void* fn = reinterpret_cast(&MlaDecodeStage1); if (smem > 48u * 1024u) { @@ -471,7 +487,8 @@ void LaunchStage1(cudaStream_t s, dim3 grid, size_t smem, int n_tile, float* mid } MlaDecodeStage1<<>>( mid, query, kv_cache, block_table, seq_lens, q_s0, q_s1, c_s0, c_s1, bt_s0, head_size, - v_head_dim, block_size, num_heads, valid_h, num_splits, mid_s0, mid_s1, mid_s2, scale); + v_head_dim, block_size, num_heads, valid_h, num_splits, mid_s0, mid_s1, mid_s2, scale, + win_left); } } @@ -549,12 +566,18 @@ void LaunchMlaDecode(Queue& q, Tensor& out, Tensor* lse, const Tensor& query, const auto* cp = kv_cache.Ptr(); const int32_t* btp = block_table.Ptr(); const int32_t* slp = seq_lens.Ptr(); + // dots3-note W4b-2 (#699). -1 is the ABSENT window and restores the full + // `[0, seq_len)` partition exactly; ops.cpp has already refused any window + // whose `right` is non-zero, so `left` is the whole contract here. + const int win_left = + args.window_size.has_value() ? static_cast(args.window_size->left) : -1; #define VT_MLA_STAGE1(DVREGS) \ LaunchStage1(s, grid1, smem, n_tile, mid, qp, cp, btp, slp, query.stride[0], \ query.stride[1], kv_cache.stride[0], kv_cache.stride[1], \ block_table.stride[0], head_size, v_head_dim, block_size, \ - num_heads, valid_h, num_splits, mid_s0, mid_s1, mid_s2, args.scale) + num_heads, valid_h, num_splits, mid_s0, mid_s1, mid_s2, args.scale, \ + win_left) if (v_head_dim <= 64) { VT_MLA_STAGE1(2); } else if (v_head_dim <= 128) { @@ -580,7 +603,7 @@ void LaunchMlaDecode(Queue& q, Tensor& out, Tensor* lse, const Tensor& query, MlaDecodeStage2<<(threads2), 0, s>>>( out.Ptr(), lse != nullptr ? lse->Ptr() : nullptr, mid, slp, out.stride[0], out.stride[1], lse != nullptr ? lse->stride[0] : 0, mid_s0, mid_s1, mid_s2, v_head_dim, - num_splits); + num_splits, win_left); Check(cudaGetLastError(), "stage2 launch"); } diff --git a/src/vt/ops.cpp b/src/vt/ops.cpp index 8cc2588b3..fb435ea6d 100644 --- a/src/vt/ops.cpp +++ b/src/vt/ops.cpp @@ -3600,6 +3600,22 @@ void MlaDecodeAttention(Queue& q, Tensor& out, Tensor* lse, const Tensor& query, "(the auto cache path; the fp8 KV-cache branch is out of scope)"); VT_CHECK(args.scale > 0.0f, "mla_decode_attention: args.scale must be > 0"); VT_CHECK(args.num_kv_splits >= 0, "mla_decode_attention: args.num_kv_splits must be >= 0"); + // The sliding-window arm (dots3-note W4b-2, #699). `left` is the inclusive + // distance behind the query, so the window WIDTH is `left + 1` and upstream's + // `sliding_window_size` 513 arrives as `left == 512`. A zero-width window + // would leave a decode row with no keys at all, which upstream cannot + // produce (`WINDOW_SIZE` is a positive config field), so it is refused rather + // than silently emitting zeros. + if (args.window_size.has_value()) { + VT_CHECK(args.window_size->left >= 0, + "mla_decode_attention: window_size.left must be >= 0 (it is the INCLUSIVE " + "distance behind the query; sliding_window 513 is left == 512)"); + VT_CHECK(args.window_size->right == 0, + "mla_decode_attention: window_size.right must be 0 — an MLA decode query IS " + "the last position of its own sequence, so a positive right bound could only " + "admit keys that do not exist. Upstream's dots3-note window is " + "(sliding_window - 1, 0) (attention.py:300 @ bc2d63e650)."); + } // Indexing is stride-driven on the leading dims (a cross-layer cache view has // gaps — cf. upstream `_page_stride`, triton_decode_attention.py:59-65), so we // require only unit innermost strides. @@ -3669,6 +3685,25 @@ void MlaPrefillAttention(Queue& q, Tensor& out, Tensor* lse, const Tensor& query VT_CHECK(args.scale > 0.0f, "mla_prefill_attention: args.scale must be > 0"); VT_CHECK(args.max_seqlen_q >= 0 && args.max_seqlen_k >= 0, "mla_prefill_attention: args.max_seqlen_q/max_seqlen_k must be >= 0"); + // The sliding-window arm (dots3-note W4b-2, #699). Upstream's only windowed + // prefill call is `causal=True, window_size=(sliding_window - 1, 0)` + // (attention.py:279-305 @ bc2d63e650). A NON-causal window is refused rather + // than approximated: FlashAttention's local mask replaces the causal + // specialization entirely, so "all keys forward, windowed backward" would need + // an infinite right bound this struct cannot express. + if (args.window_size.has_value()) { + VT_CHECK(args.window_size->left >= 0, + "mla_prefill_attention: window_size.left must be >= 0 (the INCLUSIVE distance " + "behind the bottom-right aligned query position)"); + VT_CHECK(args.window_size->right == 0, + "mla_prefill_attention: window_size.right must be 0 — upstream's windowed MLA " + "prefill is the causal (sliding_window - 1, 0) pair (attention.py:300)"); + VT_CHECK(args.causal, + "mla_prefill_attention: a window requires causal=true. FlashAttention's local " + "mask REPLACES the causal specialization (is_causal = causal && !is_local), so " + "a non-causal window cannot be spelled with a finite right bound. Upstream " + "never asks for one (attention.py:299-301)."); + } // Stride-driven on the token/head axes (a workspace slice is a strided view); // the innermost head_dim must be packed. VT_CHECK(query.stride[2] == 1 && key.stride[2] == 1 && value.stride[2] == 1 && diff --git a/tests/vllm/model_executor/layers/attention/test_mla_attention_block.cpp b/tests/vllm/model_executor/layers/attention/test_mla_attention_block.cpp index d25be7df6..793548a21 100644 --- a/tests/vllm/model_executor/layers/attention/test_mla_attention_block.cpp +++ b/tests/vllm/model_executor/layers/attention/test_mla_attention_block.cpp @@ -1189,6 +1189,30 @@ TEST_CASE("MLA block: the dots3-note fields REFUSE what they cannot represent, b CHECK_THROWS_WITH_AS(zero.Validate(), doctest::Contains("kv_lora_scale"), std::invalid_argument); + // (b2) a NEGATIVE sliding window (dots3-note W4b-2, #699). This refusal + // shipped with no test, and its review's mutation SURVIVED. It is not + // decoration: the two ops read the window as `> 0`, so a negative value + // does not throw anywhere downstream — it makes every `sliding_window > + // 0` test false and silently degrades a windowed layer to FULL + // attention, which is a wrong answer that no token differs loudly + // enough to name. The value a caller most plausibly arrives at is + // `sliding_window - 1` computed one layer too early, so 0 stays the + // ABSENT state and -1 is refused rather than folded into it. + MlaBlockDims neg = d; + neg.sliding_window = -1; + CHECK_THROWS_WITH_AS(neg.Validate(), doctest::Contains("sliding_window"), + std::invalid_argument); + // The CONTROLS, both sides of the boundary: 0 is ABSENT and legal (it is + // what every DeepSeek / MiniCPM3 / Kimi-Linear registration carries), and + // a positive window is legal. Without these the case above would also + // pass on an implementation that refused EVERY window. + MlaBlockDims absent = d; + absent.sliding_window = 0; + CHECK_NOTHROW(absent.Validate()); + MlaBlockDims windowed = d; + windowed.sliding_window = 513; // `sliding_window_size`, model.py:456 + CHECK_NOTHROW(windowed.Validate()); + HostWeights hw = MakeWeights(d, rp, 512, 5252u); const std::vector reqs = {{0, 3}}; const auto hidden = RoundBf16(RandF32(static_cast(3 * d.hidden_size), 5253u, 0.8f)); diff --git a/tests/vllm/models/test_dots3_note_attn.cpp b/tests/vllm/models/test_dots3_note_attn.cpp index cd4be65f5..35216aa67 100644 --- a/tests/vllm/models/test_dots3_note_attn.cpp +++ b/tests/vllm/models/test_dots3_note_attn.cpp @@ -2023,16 +2023,17 @@ TEST_CASE( CHECK(worst_prod <= scale * std::pow(2.0, -7)); } TEST_CASE("dots3-note W4a: what the device path still REFUSES, by name") { - // (a) a sliding layer — W4b. + // (a) a sliding layer — LIFTED at W4b-2. It is kept here as an ACCEPTANCE + // rather than deleted, because a reader following W4a's evidence lands + // on the answer instead of a gap: the same config W4a refused by name is + // now empty, and the W4b-2 cases below run it end to end. { w4a::DeviceSpec s; nlohmann::json doc = w4a::DeviceConfigDoc(s); doc["layer_types"] = nlohmann::json::array({"full_attention", "sliding_attention"}); TempConfig cfg(doc); const Dots3NoteParams p = ParseDots3NoteParams(LoadHfConfig(cfg.path())); - const std::string why = w4a::Dots3NoteDeviceRefusal(p); - CHECK(why.find("sliding_attention") != std::string::npos); - CHECK(why.find("W4b") != std::string::npos); + CHECK(w4a::Dots3NoteDeviceRefusal(p).empty()); } // (b) a MoE layer — W5. { @@ -2062,9 +2063,9 @@ TEST_CASE("dots3-note W4a: what the device path still REFUSES, by name") { CHECK_THROWS_WITH_AS(b.RunDevice(), doctest::Contains("index_topk"), std::runtime_error); } - // (e) a PADDED physical latent row — W4b. The refusal is at CONFIG level, so - // the loader never materializes a tower the forward then refuses (review - // finding F5). + // (e) a PADDED physical latent row — LIFTED at W4b-2, and kept here for the + // same reason (a): the config W4a refused at CONFIG level now passes, + // and `Tensor::Slice(2, 0, logical)` is what narrows it on read. { w4a::DeviceSpec s; nlohmann::json doc = w4a::DeviceConfigDoc(s); @@ -2072,9 +2073,7 @@ TEST_CASE("dots3-note W4a: what the device path still REFUSES, by name") { TempConfig cfg(doc); const Dots3NoteParams p = ParseDots3NoteParams(LoadHfConfig(cfg.path())); REQUIRE(p.physical_latent_row() > p.full.latent_row()); - const std::string why = w4a::Dots3NoteDeviceRefusal(p); - CHECK(why.find("PADDED") != std::string::npos); - CHECK(why.find("W4b") != std::string::npos); + CHECK(w4a::Dots3NoteDeviceRefusal(p).empty()); } // (f) a nextn tail — W10. `Dots3NoteMTPModel` is deliberately unregistered // and the backbone forward has nowhere to put an extra block, so a @@ -2097,8 +2096,11 @@ TEST_CASE("dots3-note W4a: what the device path still REFUSES, by name") { { const w4a::DeviceBench b; CHECK(w4a::Dots3NoteDeviceRefusal(b.params).empty()); + // W4b-2 changed what this check COMPARES AGAINST — the physical row rather + // than the full arm's logical one — because a padded row is now legal and + // the allocator is told the physical number. The message moved with it. CHECK_THROWS_WITH_AS(b.RunDeviceWithCacheRow(b.params.physical_latent_row() + 4), - doctest::Contains("_logical_cache"), std::runtime_error); + doctest::Contains("PHYSICAL row"), std::runtime_error); } } @@ -3256,16 +3258,17 @@ TEST_CASE( } TEST_CASE( - "dots3-note W4b-1: the DEVICE path still refuses the sliding arm and the " - "padded row, by name") { - // W4b-1 is HOST code. `Dots3NoteModel::ForwardDevice` is unchanged, and the - // boundary is asserted rather than described: the same config this file - // computes a sliding layer for is still refused by the device predicate, and - // so is the padded physical row. W4b-2 lifts both. + "dots3-note W4b-1: the MoE layer is what the DEVICE path refuses on this " + "bench, and W4b-2 lifted the other two") { + // W4b-1 wrote HOST code and lifted nothing. This case recorded the boundary + // it left: the sliding layer and the PADDED physical row were both refused + // by `Dots3NoteDeviceRefusal`. **W4b-2 lifted both**, so the assertions + // below are the ACCEPTANCES that replaced them — kept in place rather than + // deleted, so a reader following W4b-1's evidence lands on the answer. const w4b::SwaBench b; - // The bench's own config, which is the released schedule's shape: MoE from - // layer 1 and sliding from layer 2. `Dots3NoteDeviceRefusal` walks the layer - // list in order, so this one is refused at the MoE layer — and that is worth + // The bench's own config is the released schedule's shape: MoE from layer 1 + // and sliding from layer 2. `Dots3NoteDeviceRefusal` walks the layer list in + // order, so this one is refused at the MoE layer — and that is worth // asserting rather than working around, because it is what the RELEASED // checkpoint does too (spec §4.6). const std::string why = vllm::Dots3NoteDeviceRefusal(b.params); @@ -3273,29 +3276,924 @@ TEST_CASE( CHECK_FALSE(why.empty()); CHECK(why.find("MoE") != std::string::npos); - // With the MoE layers out of the way the SLIDING refusal is what fires. This - // is the branch W4b-1 does not lift: the sliding maths now exists as host - // code and the decode path still will not run it. + // With the MoE layers out of the way NOTHING is refused any more: the + // sliding layers run, and so does the PADDED physical row (6 + 4 = 10 + // against the full arm's 4 + 4 = 8). W4b-2's own cases run this exact + // geometry through `ModelRegistry::Forward`. nlohmann::json d = w4b::SwaConfigDoc(b.spec); d["first_k_dense_replace"] = 4; // every layer dense + TempConfig cfg_nextn(d); + const Dots3NoteParams p_nextn = ParseDots3NoteParams(LoadHfConfig(cfg_nextn.path())); + const std::string why_nextn = vllm::Dots3NoteDeviceRefusal(p_nextn); + MESSAGE("with no MoE layer, the refusal is: " << why_nextn); + // NOT the sliding layer any more — the NEXTN tail, which this fixture + // inherits from the released config's §4 trap 3 default of 1 and which W10 + // owns. That is worth pinning: it says the sliding refusal is gone rather + // than merely reordered behind something else. + CHECK(why_nextn.find("nextn") != std::string::npos); + CHECK(why_nextn.find("W10") != std::string::npos); + d["num_nextn_predict_layers"] = 0; TempConfig cfg_swa(d); const Dots3NoteParams p_swa = ParseDots3NoteParams(LoadHfConfig(cfg_swa.path())); const std::string why_swa = vllm::Dots3NoteDeviceRefusal(p_swa); - MESSAGE("with no MoE layer, the refusal is: " << why_swa); - CHECK(why_swa.find("sliding_attention") != std::string::npos); - CHECK(why_swa.find("W4b") != std::string::npos); + MESSAGE("with no MoE and no nextn tail, the refusal is now: '" << why_swa + << "' (empty)"); + CHECK(why_swa.empty()); + CHECK(p_swa.physical_latent_row() == 10); + CHECK(p_swa.full.latent_row() == 8); + CHECK(p_swa.swa.latent_row() == 10); + // The padded row is real on this config: two spare lanes on every physical + // row a full layer writes into. + CHECK(p_swa.physical_latent_row() - p_swa.full.latent_row() == 2); +} - // And with the sliding layers gone too, the PADDED physical row is still - // refused — 6 + 4 = 10 against the full arm's 4 + 4 = 8. That refusal is - // W4b-2's to lift and it is NOT lifted here, which is why this case exists. - d["layer_types"] = nlohmann::json::array({"full_attention", "full_attention", - "full_attention", "full_attention"}); - TempConfig cfg_pad(d); - const Dots3NoteParams p = ParseDots3NoteParams(LoadHfConfig(cfg_pad.path())); - const std::string why2 = vllm::Dots3NoteDeviceRefusal(p); - MESSAGE("with no sliding layer either, the refusal is: " << why2); - CHECK_FALSE(why2.empty()); - CHECK(why2.find("_logical_cache") != std::string::npos); - CHECK(p.physical_latent_row() == 10); - CHECK(p.full.latent_row() == 8); +// ═════════════════════════════════════════════════════════════════════════════ +// W4b-2 — the SLIDING arm ON THE DECODE PATH, over a PADDED KV cache. +// +// ─── WHAT THIS ESTABLISHES, AND WHAT IT CANNOT ─────────────────────────────── +// A MIXED config — layers `{full, sliding, full}`, every one with a +// dense MLP — is loaded through the REAL registry and run through +// `ModelRegistry::Forward` TWICE against one KV cache pool: a PREFILL of six +// tokens, then a DECODE of the seventh. Its logits are compared against a +// whole-model double reference that dispatches per layer kind — W3's +// `ref::Forward` for the full layers, W4b-1's `sref::Forward` for the sliding +// ones — with no vt op, no cache and no window arithmetic anywhere inside it. +// +// It CANNOT say the model matches vLLM. No oracle for `dots3_note` runs on any +// host this project owns (spec §6.2, §6.4 option B), so this is a consistency +// gate between two independent implementations of the same formula, one of +// them the reference W3/W4b-1 transcribed straight from the python. +// +// ─── WHAT THE PADDED ROW MEANS HERE, MEASURED ──────────────────────────────── +// The physical cache row is 10 (`swa_kv_lora_rank` 6 + `swa_qk_rope_head_dim` +// 4); a FULL layer reads 6 (`kv_lora_rank` 2 + 4). Those four lanes are not +// decoration: the case reads the RAW cache bytes after the forward and asserts +// that every full layer's written slots still carry ZERO in lanes [6, 10), +// which is exactly upstream's `_logical_cache` narrowing +// (`vllm/models/dots3_note/nvidia/attention.py:700-702` @ `bc2d63e650`) and is +// FALSE for any port that writes at the physical stride. +// +// ─── WHAT IS NOT REACHED ───────────────────────────────────────────────────── +// The MoE layers (W5), the vision and audio towers (W6/W7), the nextn tail +// (W10), the DSA lightning indexer's SELECTION and a windowed prefill that +// also carries chunked CONTEXT (both W4b-3). Each is refused BY NAME and the +// refusals have their own case below. +// ═════════════════════════════════════════════════════════════════════════════ +namespace { +namespace w4b2 { + +using vllm::Dots3NoteDeviceRefusal; +using vllm::Dots3NoteFullAttnMlaDims; +using vllm::Dots3NoteLayerKind; +using vllm::Dots3NoteSlidingAttnMlaDims; +using vllm::PagedKvCache; +using vllm::dots3_note::Dots3NoteSlidingAttnDimsFrom; +using w4a::Bf16All; +using w4a::StOut; + +// ───────────────────────────────────────────────────────────────────────────── +// The MIXED bench. Every number is chosen so a mechanism this brick added is +// OBSERVABLE, and the reasons are the ones W4a's finding F1 and W4b-1's two +// green mutations already paid for: +// +// * `layer_types` ALTERNATES full/sliding/full. A block of full +// layers followed by a block of sliding ones would let a per-layer field +// that is never RESET — the impl's `sliding_window`, the rope cache, the +// dims — still produce the right answer; alternating makes a leak in +// either direction wrong. +// * `window` 3 against a 6-token prompt, so THREE of the six prefill queries +// really lose a key and the decode query at position 6 keeps 3 of its 7. +// At `window >= tokens` the windowed answer IS the causal answer and every +// assertion here would pass on a port with no window at all. Both counts +// are printed by the case. +// * `swa_kv_lora` 6 against the full arm's 2, so the physical row (10) is +// genuinely WIDER than a full layer's logical row (6) and the padding is +// four real lanes rather than zero. +// * `swa_heads` 3 against the full arm's 2 and `swa_qk_nope` 8 against 4, so +// the two geometries disagree in head count, latent rank, NoPE width AND +// softmax scale — a layer that ran the wrong `MlaBlockDims` cannot produce +// the right answer by coincidence. +// * `swa_rope_theta` 3 against the full arm's 1300 — two thetas ORDERS +// apart, mirroring the released 5e4 against 8e7. W4b-1's fixture used 41 +// against 137 and measured a shared rope cache at only 0.0300 relative on +// one LAYER; over a whole bf16 model that sits UNDER the quantisation +// residue, so this fixture separates the two thetas the way upstream does +// rather than the way a tidy pair of close numbers would. +// * `page_size` 4 with a SHUFFLED block table {1, 0}: logical page 0 maps to +// physical page 1. A contiguous table makes a paged read and a flat read +// the same answer, which would leave the block lookup unproven — and the +// windowed decode's key range starts INSIDE logical page 1, so the lookup +// is on the windowed path and not only the prefill one. +// * `q_lora` 3 over `hidden` 16 gives rescales sqrt(16/3)=2.309 (q) and +// sqrt(16/2)=2.828 / sqrt(16/6)=1.633 (kv, full / sliding) — all far from +// 1.0 and different from each other, so neither a missing nor a swapped +// scale can hide on either arm. +// ───────────────────────────────────────────────────────────────────────────── +struct Spec { + int64_t hidden = 16; + int64_t vocab = 12; + int64_t inter = 10; + int64_t max_pos = 32; + double rms_eps = 1e-3; + // the FULL arm — W4a's geometry, unchanged, so the two bricks' fixtures agree + int64_t full_heads = 2; + int64_t full_qk_nope = 4; + int64_t qk_rope = 4; + int64_t v_head = 8; + int64_t q_lora = 3; + int64_t full_kv_lora = 2; + double rope_theta = 1300.0; + // the SLIDING arm + int64_t swa_heads = 3; + int64_t swa_qk_nope = 8; + int64_t swa_kv_lora = 6; + double swa_rope_theta = 3.0; + int64_t window = 3; + // >= prompt + 1, so the DSA top-k selects every causal candidate on the full + // layers and dense attention IS upstream's answer. The refusal past this + // bound has its own case. + int64_t index_topk = 32; + int64_t index_n_heads = 2; + int64_t index_head_dim = 6; + int64_t prompt = 6; + int64_t page_size = 4; + bool tie_word_embeddings = false; + // {full, sliding, full} + std::vector kinds{Dots3NoteLayerKind::kFullAttention, + Dots3NoteLayerKind::kSlidingAttention, + Dots3NoteLayerKind::kFullAttention}; + int64_t layers() const { return static_cast(kinds.size()); } +}; + +nlohmann::json ConfigDoc(const Spec& s) { + nlohmann::json d = FixtureConfigDoc(); + d["hidden_size"] = s.hidden; + d["num_hidden_layers"] = s.layers(); + nlohmann::json lt = nlohmann::json::array(); + for (Dots3NoteLayerKind k : s.kinds) { + lt.push_back(k == Dots3NoteLayerKind::kSlidingAttention ? "sliding_attention" + : "full_attention"); + } + d["layer_types"] = lt; + d["num_attention_heads"] = s.full_heads; + d["num_key_value_heads"] = s.full_heads; + d["qk_nope_head_dim"] = s.full_qk_nope; + d["qk_rope_head_dim"] = s.qk_rope; + d["v_head_dim"] = s.v_head; + d["q_lora_rank"] = s.q_lora; + d["kv_lora_rank"] = s.full_kv_lora; + d["rope_theta"] = s.rope_theta; + d["rms_norm_eps"] = s.rms_eps; + d["max_position_embeddings"] = s.max_pos; + d["index_n_heads"] = s.index_n_heads; + d["index_head_dim"] = s.index_head_dim; + d["index_topk"] = s.index_topk; + d["swa_num_attention_heads"] = s.swa_heads; + d["swa_num_key_value_heads"] = s.swa_heads; + d["swa_q_lora_rank"] = s.q_lora; + d["swa_kv_lora_rank"] = s.swa_kv_lora; + d["swa_qk_nope_head_dim"] = s.swa_qk_nope; + d["swa_qk_rope_head_dim"] = s.qk_rope; + d["swa_v_head_dim"] = s.v_head; + d["swa_rope_theta"] = s.swa_rope_theta; + d["sliding_window_size"] = s.window; + d["vocab_size"] = s.vocab; + d["intermediate_size"] = s.inter; + d["moe_intermediate_size"] = 6; + d["n_routed_experts"] = 4; + d["num_experts_per_tok"] = 2; + // Every layer DENSE: W5 owns the MoE. + d["first_k_dense_replace"] = s.layers(); + d["num_nextn_predict_layers"] = 0; + d["tie_word_embeddings"] = s.tie_word_embeddings; + return d; +} + +// The whole tiny mixed model in double, every weight ALREADY bf16-rounded so +// the comparison measures the FORWARD and not the weights' storage width. +struct Weights { + std::vector embed; + std::vector final_norm; + std::vector lm_head; + struct Layer { + Dots3NoteLayerKind kind = Dots3NoteLayerKind::kFullAttention; + std::vector input_ln; + std::vector post_ln; + FullAttnWeights full; // populated iff kind == full + SlidingAttnWeights swa; // populated iff kind == sliding + std::vector gate_proj; + std::vector up_proj; + std::vector down_proj; + }; + std::vector layers; +}; + +std::vector Ln(Rng& r, int64_t n) { + std::vector v = r.fill(n, 0.3); + for (double& x : v) x += 1.0; + return Bf16All(v); +} + +Weights MakeWeights(const Spec& s, const FullAttnDims& fd, const SlidingAttnDims& sd, + uint64_t seed) { + Rng r(seed); + Weights w; + w.embed = Bf16All(r.fill(s.vocab * s.hidden, 0.7)); + w.final_norm = Ln(r, s.hidden); + w.lm_head = Bf16All(r.fill(s.vocab * s.hidden, 0.5)); + for (int64_t l = 0; l < s.layers(); ++l) { + Weights::Layer lw; + lw.kind = s.kinds[static_cast(l)]; + lw.input_ln = Ln(r, s.hidden); + lw.post_ln = Ln(r, s.hidden); + const uint64_t lseed = seed + 0x1000ULL * static_cast(l + 1); + if (lw.kind == Dots3NoteLayerKind::kFullAttention) { + lw.full = TinyWeights(fd, lseed); + // `k_rope_only_layernorm` is made OBSERVABLE the way W4a's bench makes it + // observable, and for the reason spec §4.6 records: RoPE preserves the L2 + // norm of every rotated pair exactly, so the only part of the norm that + // does not commute with the rotation is the PER-LANE weight. Weights + // hugging 1.0 let "norm AFTER the rope" slip under any sane bound. + { + const int64_t R = fd.qk_rope_head_dim, H = fd.hidden_size; + const int64_t off = fd.kv_lora_rank * H; + for (int64_t i = 0; i < R * H; ++i) + lw.full.kv_a_proj_with_mqa[static_cast(off + i)] *= 6.0; + for (int64_t i = 0; i < R; ++i) + lw.full.k_rope_only_layernorm[static_cast(i)] = (i % 2 == 0) ? 2.5 : 0.3; + } + lw.full.q_a_proj = Bf16All(lw.full.q_a_proj); + lw.full.q_a_layernorm = Bf16All(lw.full.q_a_layernorm); + lw.full.kv_a_layernorm = Bf16All(lw.full.kv_a_layernorm); + lw.full.kv_a_proj_with_mqa = Bf16All(lw.full.kv_a_proj_with_mqa); + lw.full.k_rope_only_layernorm = Bf16All(lw.full.k_rope_only_layernorm); + lw.full.q_b_proj = Bf16All(lw.full.q_b_proj); + lw.full.kv_b_proj = Bf16All(lw.full.kv_b_proj); + lw.full.o_proj = Bf16All(lw.full.o_proj); + lw.full.g_proj = Bf16All(lw.full.g_proj); + lw.full.indexer_wq_b = Bf16All(lw.full.indexer_wq_b); + lw.full.indexer_wk = Bf16All(lw.full.indexer_wk); + lw.full.indexer_weights_proj = Bf16All(lw.full.indexer_weights_proj); + lw.full.indexer_k_norm_weight = Bf16All(lw.full.indexer_k_norm_weight); + lw.full.indexer_k_norm_bias = Bf16All(lw.full.indexer_k_norm_bias); + } else { + // `w4b::SwaWeights` already alternates the k_pe norm 2.5/0.3 within each + // rotated pair, for the same reason, and it needs no extra amplification + // of the rope rows: dropping the norm on this arm moves the logits by + // more than 1.5 relative, which the case prints. + lw.swa = w4b::SwaWeights(sd, lseed); + lw.swa.q_a_proj = Bf16All(lw.swa.q_a_proj); + lw.swa.q_a_layernorm = Bf16All(lw.swa.q_a_layernorm); + lw.swa.kv_a_layernorm = Bf16All(lw.swa.kv_a_layernorm); + lw.swa.kv_a_proj_with_mqa = Bf16All(lw.swa.kv_a_proj_with_mqa); + lw.swa.k_rope_only_layernorm = Bf16All(lw.swa.k_rope_only_layernorm); + lw.swa.q_b_proj = Bf16All(lw.swa.q_b_proj); + lw.swa.kv_b_proj = Bf16All(lw.swa.kv_b_proj); + lw.swa.o_proj = Bf16All(lw.swa.o_proj); + lw.swa.g_proj = Bf16All(lw.swa.g_proj); + } + lw.gate_proj = Bf16All(r.fill(s.inter * s.hidden, 0.5)); + lw.up_proj = Bf16All(r.fill(s.inter * s.hidden, 0.5)); + lw.down_proj = Bf16All(r.fill(s.hidden * s.inter, 0.5)); + w.layers.push_back(std::move(lw)); + } + return w; +} + +// The on-disk entries in the names `EnumerateDots3NoteTensors` claims. The five +// indexer tensors exist ONLY on the full layers, which is upstream's own shape +// (`Dots3NoteSlidingAttention` sets `self.indexer = None`, model.py:432) and +// what the W1 enumerator already encodes. +std::vector CheckpointOf(const Spec& s, const FullAttnDims& fd, + const SlidingAttnDims& sd, const Weights& w) { + const int64_t H = s.hidden; + std::vector e; + e.push_back({"model.embed_tokens.weight", {s.vocab, H}, w.embed}); + e.push_back({"model.norm.weight", {H}, w.final_norm}); + if (!s.tie_word_embeddings) e.push_back({"lm_head.weight", {s.vocab, H}, w.lm_head}); + for (int64_t l = 0; l < s.layers(); ++l) { + const Weights::Layer& lw = w.layers[static_cast(l)]; + const bool sliding = lw.kind == Dots3NoteLayerKind::kSlidingAttention; + const std::string p = "model.layers." + std::to_string(l) + "."; + const std::string sa = p + "self_attn."; + const int64_t N = sliding ? sd.num_heads : fd.num_heads; + const int64_t QK = sliding ? sd.qk_head_dim() : fd.qk_head_dim(); + const int64_t P = sliding ? sd.qk_nope_head_dim : fd.qk_nope_head_dim; + const int64_t V = sliding ? sd.v_head_dim : fd.v_head_dim; + const int64_t L = sliding ? sd.kv_lora_rank : fd.kv_lora_rank; + const int64_t R = sliding ? sd.qk_rope_head_dim : fd.qk_rope_head_dim; + const int64_t QL = sliding ? sd.q_lora_rank : fd.q_lora_rank; + const SlidingAttnWeights& sw = lw.swa; + const FullAttnWeights& fw = lw.full; + e.push_back({p + "input_layernorm.weight", {H}, lw.input_ln}); + e.push_back({p + "post_attention_layernorm.weight", {H}, lw.post_ln}); + e.push_back({sa + "q_a_proj.weight", {QL, H}, sliding ? sw.q_a_proj : fw.q_a_proj}); + e.push_back({sa + "q_a_layernorm.weight", + {QL}, + sliding ? sw.q_a_layernorm : fw.q_a_layernorm}); + e.push_back({sa + "q_b_proj.weight", + {N * QK, QL}, + sliding ? sw.q_b_proj : fw.q_b_proj}); + e.push_back({sa + "kv_a_proj_with_mqa.weight", + {L + R, H}, + sliding ? sw.kv_a_proj_with_mqa : fw.kv_a_proj_with_mqa}); + e.push_back({sa + "kv_a_layernorm.weight", + {L}, + sliding ? sw.kv_a_layernorm : fw.kv_a_layernorm}); + e.push_back({sa + "kv_b_proj.weight", + {N * (P + V), L}, + sliding ? sw.kv_b_proj : fw.kv_b_proj}); + e.push_back({sa + "o_proj.weight", {H, N * V}, sliding ? sw.o_proj : fw.o_proj}); + e.push_back({sa + "g_proj.weight", {N, H}, sliding ? sw.g_proj : fw.g_proj}); + e.push_back({sa + "k_rope_only_layernorm.weight", + {R}, + sliding ? sw.k_rope_only_layernorm : fw.k_rope_only_layernorm}); + if (!sliding) { + e.push_back({sa + "indexer.wq_b.weight", + {fd.index_n_heads * fd.index_head_dim, fd.q_lora_rank}, + fw.indexer_wq_b}); + e.push_back({sa + "indexer.wk.weight", {fd.index_head_dim, H}, fw.indexer_wk}); + e.push_back({sa + "indexer.k_norm.weight", + {fd.index_head_dim}, + fw.indexer_k_norm_weight}); + e.push_back({sa + "indexer.k_norm.bias", + {fd.index_head_dim}, + fw.indexer_k_norm_bias}); + e.push_back({sa + "indexer.weights_proj.weight", + {fd.index_n_heads, H}, + fw.indexer_weights_proj}); + } + e.push_back({p + "mlp.gate_proj.weight", {s.inter, H}, lw.gate_proj}); + e.push_back({p + "mlp.up_proj.weight", {s.inter, H}, lw.up_proj}); + e.push_back({p + "mlp.down_proj.weight", {H, s.inter}, lw.down_proj}); + } + return e; +} + +// ───────────────────────────────────────────────────────────────────────────── +// The MIXED whole-model reference. The residual stream, the MLP and the lm_head +// are W4a's; the attention DISPATCHES on the layer kind, into W3's +// `ref::Forward` or W4b-1's `sref::Forward`. Neither of those knows about a +// cache, a block table, a slot mapping, a gather or a physical row: the sliding +// one takes the window as the direct positional predicate `s <= t && t - s < W` +// over a materialized MHA. That is the whole point — the device arm's answer +// has to survive being computed a completely different way. +std::vector RefModel(const Spec& s, const FullAttnDims& fd, + const SlidingAttnDims& sd, const Weights& w, + const std::vector& tokens, + const std::vector& positions, const ref::Opts& fo, + const w4b::sref::Opts& so) { + const int64_t T = static_cast(tokens.size()), H = s.hidden; + std::vector hidden(static_cast(T * H)); + for (int64_t t = 0; t < T; ++t) { + for (int64_t c = 0; c < H; ++c) { + hidden[static_cast(t * H + c)] = + w.embed[static_cast(tokens[static_cast(t)] * H + c)]; + } + } + std::vector res(static_cast(T * H), 0.0); + for (int64_t l = 0; l < s.layers(); ++l) { + const Weights::Layer& lw = w.layers[static_cast(l)]; + for (size_t i = 0; i < res.size(); ++i) res[i] += hidden[i]; + const std::vector x = ref::Rms(res, lw.input_ln, T, H, s.rms_eps); + const std::vector a = + lw.kind == Dots3NoteLayerKind::kSlidingAttention + ? w4b::sref::Forward(sd, lw.swa, x, positions, T, so).out + : ref::Forward(fd, lw.full, x, positions, T, fo).out; + for (size_t i = 0; i < res.size(); ++i) res[i] += a[i]; + const std::vector y = ref::Rms(res, lw.post_ln, T, H, s.rms_eps); + const std::vector g = ref::Dense(y, lw.gate_proj, T, H, s.inter); + const std::vector u = ref::Dense(y, lw.up_proj, T, H, s.inter); + std::vector act(g.size()); + for (size_t i = 0; i < g.size(); ++i) act[i] = (g[i] / (1.0 + std::exp(-g[i]))) * u[i]; + hidden = ref::Dense(act, lw.down_proj, T, s.inter, H); + } + for (size_t i = 0; i < res.size(); ++i) res[i] += hidden[i]; + const std::vector z = ref::Rms(res, w.final_norm, T, H, s.rms_eps); + return ref::Dense(z, w.lm_head, T, H, s.vocab); +} + +// The PHYSICAL slot a logical position lands in, through the SHUFFLED block +// table. Derived here by hand rather than taken from the forward, so the test +// and the code cannot share one arithmetic bug. +int64_t SlotOf(const Spec& s, const std::vector& block_table, int64_t pos) { + const int64_t page = pos / s.page_size; + REQUIRE(page < static_cast(block_table.size())); + return static_cast(block_table[static_cast(page)]) * s.page_size + + pos % s.page_size; +} + +// `Dots3NoteFullAttnDimsFrom` REFUSES a schedule with no `full_attention` +// layer, by design (spec §1.1) — and one of the cases below drives exactly that +// schedule, to show that a pure-SWA config is not refused for a DSA indexer it +// does not have. An all-sliding bench never reads these dims, so it gets a +// default-constructed struct rather than the refusal. +FullAttnDims FullDimsOrDefault(const Dots3NoteParams& p) { + for (Dots3NoteLayerKind k : p.layer_types) { + if (k == Dots3NoteLayerKind::kFullAttention) return Dots3NoteFullAttnDimsFrom(p); + } + return FullAttnDims{}; +} + +// The mixed bench. The registration, the config and the loaded model all come +// from the REAL registry over the REAL loader; the geometry comes from the +// released `config.json` with the fields above overridden, so every W1 +// validation still applies to it. +struct Bench { + Spec spec; + TempConfig cfg; + HfConfig config; + Dots3NoteParams params; + // The HOST reference geometries (W3 / W4b-1), which drive the weights and the + // double reference... + FullAttnDims fdims; + SlidingAttnDims sdims; + // ...and the DEVICE seam geometries the forward itself builds, which is where + // the softmax scale and the window live. + vllm::mla::MlaBlockDims mfull; + vllm::mla::MlaBlockDims mswa; + Weights w; + std::vector entries; + std::vector tokens; // the prompt PLUS the token the decode step runs + std::vector positions; + std::vector block_table{1, 0}; // SHUFFLED + + explicit Bench(Spec s = Spec{}) + : spec(s), + cfg(ConfigDoc(s)), + config(LoadHfConfig(cfg.path())), + params(ParseDots3NoteParams(config)), + fdims(FullDimsOrDefault(params)), + sdims(Dots3NoteSlidingAttnDimsFrom(params)), + mfull(Dots3NoteFullAttnMlaDims(params)), + mswa(Dots3NoteSlidingAttnMlaDims(params)), + w(MakeWeights(s, fdims, sdims, 0x243F6A8885A308D3ULL)), + entries(CheckpointOf(s, fdims, sdims, w)) { + for (int64_t t = 0; t <= spec.prompt; ++t) { + tokens.push_back(static_cast((t * 5 + 1) % spec.vocab)); + positions.push_back(static_cast(t)); + } + } + + // The whole point of this brick: PREFILL then DECODE, both through + // `ModelRegistry::Forward`, against ONE cache pool. Returns the decode step's + // [1, vocab] logits, and optionally the raw cache bytes afterwards. + std::vector RunPrefillThenDecode( + std::vector>* cache_out = nullptr, + int64_t cache_row_override = 0) const { + const vllm::ModelRegistration& reg = ModelRegistry::Resolve(config); + w4a::TempCheckpoint ckpt(entries); + std::vector shards; + shards.push_back(vllm::SafetensorsFile::Open(ckpt.file())); + const vllm::ModelSource source = vllm::ModelSource::FromSafetensors(shards); + std::unique_ptr model = reg.factory->load_weights(reg, config, source); + REQUIRE(model != nullptr); + + const int64_t row = + cache_row_override > 0 ? cache_row_override : params.physical_latent_row(); + w4a::MlaCachePool pool(spec.layers(), row, /*num_blocks=*/2, spec.page_size); + vt::Queue queue{vt::Device{vt::DeviceType::kCPU, 0}, nullptr}; + const std::vector no_gather; + std::vector gdn_state; + vllm::v1::GDNAttentionMetadata gdn_meta{}; + + // ── step 1: PREFILL the prompt ───────────────────────────────────────── + { + vllm::v1::CommonAttentionMetadata m; + m.num_reqs = 1; + m.num_actual_tokens = static_cast(spec.prompt); + m.query_start_loc = {0, static_cast(spec.prompt)}; + m.query_start_loc_cpu = m.query_start_loc; + m.seq_lens = {static_cast(spec.prompt)}; + m.seq_lens_cpu = m.seq_lens; + m.max_query_len = static_cast(spec.prompt); + m.max_seq_len = static_cast(spec.prompt); + m.block_table_num_cols = static_cast(block_table.size()); + m.block_table_tensor = block_table; + for (int64_t t = 0; t < spec.prompt; ++t) { + m.slot_mapping.push_back(SlotOf(spec, block_table, t)); + } + m.causal = true; + const std::vector ids(tokens.begin(), tokens.begin() + spec.prompt); + const std::vector pos(positions.begin(), positions.begin() + spec.prompt); + const vllm::ModelForwardInput in{.token_ids = ids, + .positions = pos, + .attn_meta = m, + .gdn_meta = gdn_meta, + .attn_kv = pool.attn_kv, + .gdn_state = gdn_state, + .config = config, + .queue = queue, + .logits_indices = no_gather, + .num_reqs = 1}; + const vllm::ForwardLogits fl = ModelRegistry::Forward(*model, in); + REQUIRE(fl.on_device()); + REQUIRE(fl.rows == spec.prompt); + } + + // ── step 2: DECODE the next token, against the cache step 1 wrote ────── + std::vector out; + { + vllm::v1::CommonAttentionMetadata m; + m.num_reqs = 1; + m.num_actual_tokens = 1; + m.query_start_loc = {0, 1}; + m.query_start_loc_cpu = m.query_start_loc; + m.seq_lens = {static_cast(spec.prompt + 1)}; + m.seq_lens_cpu = m.seq_lens; + m.max_query_len = 1; + m.max_seq_len = static_cast(spec.prompt + 1); + m.block_table_num_cols = static_cast(block_table.size()); + m.block_table_tensor = block_table; + m.slot_mapping = {SlotOf(spec, block_table, spec.prompt)}; + m.causal = true; + const std::vector ids{tokens.back()}; + const std::vector pos{positions.back()}; + const vllm::ModelForwardInput in{.token_ids = ids, + .positions = pos, + .attn_meta = m, + .gdn_meta = gdn_meta, + .attn_kv = pool.attn_kv, + .gdn_state = gdn_state, + .config = config, + .queue = queue, + .logits_indices = no_gather, + .num_reqs = 1}; + const vllm::ForwardLogits fl = ModelRegistry::Forward(*model, in); + REQUIRE(fl.on_device()); + REQUIRE(fl.rows == 1); + REQUIRE(fl.vocab == spec.vocab); + const auto* src = static_cast(fl.device_tensor.data); + out.assign(static_cast(fl.vocab), 0.0); + for (size_t i = 0; i < out.size(); ++i) out[i] = static_cast(src[i]); + } + if (cache_out != nullptr) *cache_out = pool.buf; + return out; + } + + // The reference's logits for the LAST position, i.e. the row the decode step + // produces. The reference recomputes the whole 7-token sequence from scratch + // in double and has no cache at all, so "what the decode step read out of the + // cache" has to equal "what a fresh full-sequence forward computes". + std::vector RefLastRow(const ref::Opts& fo = ref::Opts{}, + const w4b::sref::Opts& so = w4b::sref::Opts{}) const { + const std::vector all = + RefModel(spec, fdims, sdims, w, tokens, positions, fo, so); + const int64_t T = static_cast(tokens.size()); + return std::vector(all.begin() + static_cast((T - 1) * spec.vocab), + all.end()); + } +}; + +// The bf16 agreement bound, chosen for SEPARATION and not to hug the residue — +// W4a's review finding F1 applied to a THREE-layer model whose activation +// stream is bf16 end to end while the reference is double throughout. +// +// THREE, not four. The first draft of this fixture ran `{full, sliding, full, +// sliding}`; the retune §4.8 records dropped it to `{full, sliding, full}`, +// which is still full/sliding/full so a per-layer field leaking in EITHER +// direction is wrong. `Spec::kinds` is the committed truth and the case +// `REQUIRE`s `num_hidden_layers == 3`; two comments kept saying four until the +// W4b-2 review read them against the code. +// +// THREE ratios, kept SEPARATE, because merging any two of them overstates the +// headroom: spec §4.6 records a draft that did exactly that and was wrong by +// 2.8x. The cases PRINT all three from the numbers they just measured, so this +// constant is the only thing written down in advance. +constexpr double kMixedRel = 6e-2; + +} // namespace w4b2 +} // namespace + +// ───────────────────────────────────────────────────────────────────────────── +TEST_CASE( + "dots3-note W4b-2: the SLIDING layer is REACHED through " + "ModelRegistry::Forward, over a PADDED cache, and agrees with the " + "independent reference") { + const w4b2::Bench b; + // The scope boundary first: this config is one the device path ACCEPTS, and + // it is genuinely mixed and genuinely padded. + CHECK(w4b2::Dots3NoteDeviceRefusal(b.params).empty()); + REQUIRE(b.params.num_hidden_layers == 3); + int64_t n_full = 0, n_sliding = 0; + for (int64_t l = 0; l < b.params.num_hidden_layers; ++l) { + if (b.params.kind_of(l) == vllm::Dots3NoteLayerKind::kSlidingAttention) { + ++n_sliding; + } else { + ++n_full; + } + CHECK_FALSE(b.params.is_moe_layer(l)); + } + CHECK(n_full == 2); + CHECK(n_sliding == 1); + // The PADDING is real: the physical row is wider than a full layer's logical + // one. Without this the narrowing is the identity and proves nothing. + REQUIRE(b.params.physical_latent_row() == 10); + REQUIRE(b.params.full.latent_row() == 6); + REQUIRE(b.params.swa.latent_row() == b.params.physical_latent_row()); + // The two geometries really differ, on every axis that could hide a mistake. + CHECK(b.mswa.num_heads != b.mfull.num_heads); + CHECK(b.mswa.kv_lora_rank != b.mfull.kv_lora_rank); + CHECK(b.mswa.qk_nope_head_dim != b.mfull.qk_nope_head_dim); + CHECK(b.mswa.scale != b.mfull.scale); + CHECK(b.mswa.sliding_window == 3); + CHECK(b.mfull.sliding_window == 0); + CHECK(b.mswa.head_size() == b.params.physical_latent_row()); + CHECK(b.mfull.head_size() < b.params.physical_latent_row()); + + // The WINDOW bites, COUNTED rather than assumed. Prefill: query t sees keys + // [t - 2, t], so queries 3, 4 and 5 lose 1, 2 and 3 keys. Decode: the query + // at position 6 keeps keys 4..6 of 0..6, i.e. it loses four. + int64_t prefill_queries_that_lose = 0, prefill_keys_dropped = 0; + for (int64_t t = 0; t < b.spec.prompt; ++t) { + const int64_t lost = std::max(0, t - (b.spec.window - 1)); + if (lost > 0) ++prefill_queries_that_lose; + prefill_keys_dropped += lost; + } + const int64_t decode_keys_dropped = b.spec.prompt + 1 - b.spec.window; + MESSAGE("W4b-2 window bite: prefill " + << prefill_queries_that_lose << " of " << b.spec.prompt + << " queries lose a key (" << prefill_keys_dropped + << " keys dropped); the decode query at position " << b.spec.prompt + << " keeps " << b.spec.window << " of " << (b.spec.prompt + 1) << " keys"); + REQUIRE(prefill_queries_that_lose == 3); + REQUIRE(prefill_keys_dropped == 6); + REQUIRE(decode_keys_dropped == 4); + + std::vector> cache; + const std::vector got = b.RunPrefillThenDecode(&cache); + const std::vector want = b.RefLastRow(); + const Diff d = Compare(got, want); + MESSAGE("W4b-2 mixed prefill+decode vs the independent reference: max|diff| " + << d.max_abs << " over a scale of " << d.max_mag << " = " << d.max_rel + << " relative; bound " << w4b2::kMixedRel << " = " + << (w4b2::kMixedRel / d.max_rel) << "x the residue"); + CHECK(d.max_rel < w4b2::kMixedRel); + + // ── the PADDED row's spare lanes, read out of the RAW cache ────────────── + // A FULL layer writes `full.latent_row()` lanes into a `physical_row`-wide + // slot. Upstream narrows the cache view on read AND on write + // (`_logical_cache`, attention.py:700-720), so the tail of every physical row + // it touches stays exactly as the allocator left it — zero. A port that wrote + // at the physical stride, or read at the logical one, breaks this. + const int64_t phys = b.params.physical_latent_row(); + const int64_t logical = b.params.full.latent_row(); + int64_t full_slots_checked = 0, sliding_nonzero_pad_lanes = 0; + for (int64_t l = 0; l < b.params.num_hidden_layers; ++l) { + const bool sliding = + b.params.kind_of(l) == vllm::Dots3NoteLayerKind::kSlidingAttention; + const std::vector& buf = cache[static_cast(l)]; + for (int64_t t = 0; t <= b.spec.prompt; ++t) { + const int64_t slot = w4b2::SlotOf(b.spec, b.block_table, t); + for (int64_t c = logical; c < phys; ++c) { + const uint16_t v = buf[static_cast(slot * phys + c)]; + if (sliding) { + if (v != 0) ++sliding_nonzero_pad_lanes; + } else { + CHECK(v == 0); + } + } + if (!sliding) ++full_slots_checked; + } + } + MESSAGE("W4b-2 padded row: " << full_slots_checked + << " full-layer slots checked, lanes [" << logical + << ", " << phys + << ") all ZERO; the same lanes on the sliding layers " + "carry " + << sliding_nonzero_pad_lanes << " non-zero values"); + CHECK(full_slots_checked == 2 * (b.spec.prompt + 1)); + // The CONTROL: those lanes are not zero everywhere. A sliding layer's logical + // row IS the physical one, so it writes them — which is what makes the + // all-zero assertion above a statement about the NARROWING rather than about + // the fixture happening to produce zeros. + CHECK(sliding_nonzero_pad_lanes > 0); +} + +TEST_CASE("dots3-note W4b-2: the mixed device forward is DETERMINISTIC run to run") { + const w4b2::Bench b; + const std::vector a = b.RunPrefillThenDecode(); + const std::vector c = b.RunPrefillThenDecode(); + REQUIRE(a.size() == c.size()); + for (size_t i = 0; i < a.size(); ++i) CHECK(a[i] == c[i]); +} + +TEST_CASE( + "dots3-note W4b-2: the WINDOW is what makes the answer different, and each " + "sliding-only mechanism is EXERCISED on the device path") { + const w4b2::Bench b; + const std::vector got = b.RunPrefillThenDecode(); + const Diff base = Compare(got, b.RefLastRow()); + + // The reference with `windowed = false` models a port that never noticed + // `sliding_window_size` and ran plain causal attention on the 33 sliding + // layers. That is THE defect this brick exists to prevent, and no shape check + // can see it. + w4b::sref::Opts unwindowed; + unwindowed.windowed = false; + const Diff no_win = Compare(got, b.RefLastRow(ref::Opts{}, unwindowed)); + MESSAGE("W4b-2 separation: residue " << base.max_rel << "; the bound " + << w4b2::kMixedRel << " is " + << (w4b2::kMixedRel / base.max_rel) + << "x the residue"); + MESSAGE("W4b-2 with NO WINDOW at all: " << no_win.max_rel << " relative = " + << (no_win.max_rel / w4b2::kMixedRel) + << "x the BOUND, " + << (no_win.max_rel / base.max_rel) + << "x the RESIDUE"); + CHECK(no_win.max_rel > w4b2::kMixedRel); + + // Each sliding-only mechanism, neutralised in the REFERENCE, with the device + // arm drifting AWAY. Ratios are given against both the bound and the residue, + // LABELLED, because the two are different statements (spec §4.6 F1). + struct Arm { + std::string what; + w4b::sref::Opts so; + }; + std::vector arms; + { + Arm a{"the sliding arm inheriting the MODEL-level rope theta", {}}; + a.so.rope_theta_override = b.spec.rope_theta; + arms.push_back(a); + } + { + Arm a{"the sliding arm's q LoRA rescale dropped", {}}; + a.so.apply_q_lora_rescale = false; + arms.push_back(a); + } + { + Arm a{"the sliding arm's kv LoRA rescale dropped", {}}; + a.so.apply_kv_lora_rescale = false; + arms.push_back(a); + } + { + Arm a{"the sliding arm's k_rope_only_layernorm dropped", {}}; + a.so.k_rope_only_norm = false; + arms.push_back(a); + } + { + Arm a{"the sliding arm's headwise gate made lane-wise", {}}; + a.so.headwise_gate = false; + arms.push_back(a); + } + for (const Arm& a : arms) { + const Diff dd = Compare(got, b.RefLastRow(ref::Opts{}, a.so)); + MESSAGE("W4b-2 with " << a.what << ": " << dd.max_rel << " relative = " + << (dd.max_rel / w4b2::kMixedRel) << "x the BOUND, " + << (dd.max_rel / base.max_rel) << "x the RESIDUE"); + CHECK(dd.max_rel > w4b2::kMixedRel); + } +} + +TEST_CASE( + "dots3-note W4b-2: the SLIDING geometry the DEVICE forward runs comes off " + "the RELEASED config") { + // Not a hand-typed struct: the released `config.json`, through + // `ParseDots3NoteParams` -> `Dots3NoteSlidingAttnMlaDims`, which is the + // function `MaterializeDots3NoteDevice` and `ForwardDevice` both call. + TempConfig cfg(FixtureConfigDoc()); + const HfConfig config = LoadHfConfig(cfg.path()); + const Dots3NoteParams p = ParseDots3NoteParams(config); + const vllm::mla::MlaBlockDims sd = vllm::Dots3NoteSlidingAttnMlaDims(p); + const vllm::mla::MlaBlockDims fd = vllm::Dots3NoteFullAttnMlaDims(p); + CHECK(sd.num_heads == 64); + CHECK(sd.kv_lora_rank == 1024); + CHECK(sd.qk_nope_head_dim == 192); + CHECK(sd.qk_rope_head_dim == 64); + CHECK(sd.v_head_dim == 128); + CHECK(sd.q_lora_rank == 1024); + CHECK(sd.head_size() == 1088); + CHECK(sd.head_size() == p.physical_latent_row()); + // `sliding_window=config.sliding_window_size` (model.py:457). + CHECK(sd.sliding_window == 513); + CHECK(fd.sliding_window == 0); + // `scale = qk_head_dim ** -0.5` (model.py:446) — 256^-0.5, NOT the full arm's + // 192^-0.5. No YaRN and no mscale on either arm. + CHECK(sd.scale == doctest::Approx(1.0 / std::sqrt(256.0)).epsilon(1e-6)); + CHECK(fd.scale == doctest::Approx(1.0 / std::sqrt(192.0)).epsilon(1e-6)); + // Both geometries are GPT-J (spec §4 item 6, #1804). + CHECK_FALSE(sd.is_neox_style); + CHECK_FALSE(fd.is_neox_style); + // The released ranks make the two SLIDING rescales EQUAL at sqrt(5120/1024); + // the full arm's kv rescale is sqrt(5120/512) and differs. + CHECK(sd.q_lora_scale == doctest::Approx(std::sqrt(5120.0 / 1024.0)).epsilon(1e-12)); + CHECK(sd.kv_lora_scale == doctest::Approx(std::sqrt(5120.0 / 1024.0)).epsilon(1e-12)); + CHECK(fd.kv_lora_scale == doctest::Approx(std::sqrt(5120.0 / 512.0)).epsilon(1e-12)); + // The full layers read 576 out of the 1088-wide physical row: the padding is + // 512 real lanes on 13 of the 46 layers. + CHECK(fd.head_size() == 576); + CHECK(p.physical_latent_row() - fd.head_size() == 512); +} + +TEST_CASE("dots3-note W4b-2: what the device path STILL refuses, by name") { + // (a) a MoE layer — W5. The RELEASED checkpoint trips this at layer 1, so + // nothing a user can run changed at W4b-2. + { + nlohmann::json doc = w4b2::ConfigDoc(w4b2::Spec{}); + doc["first_k_dense_replace"] = 1; + TempConfig cfg(doc); + const Dots3NoteParams p = ParseDots3NoteParams(LoadHfConfig(cfg.path())); + const std::string why = w4b2::Dots3NoteDeviceRefusal(p); + CHECK(why.find("MoE") != std::string::npos); + CHECK(why.find("W5") != std::string::npos); + } + // (b) the RELEASED config still refuses, at its MoE layer. + { + TempConfig cfg(FixtureConfigDoc()); + const Dots3NoteParams p = ParseDots3NoteParams(LoadHfConfig(cfg.path())); + const std::string why = w4b2::Dots3NoteDeviceRefusal(p); + MESSAGE("W4b-2 released-config refusal: " << why); + CHECK(why.find("MoE") != std::string::npos); + } + // (c) a sequence past `index_topk`, on a config that HAS a full layer. The + // DSA selection is still not on the device path (W4b-3). + { + w4b2::Spec s; + s.index_topk = 2; // < the prompt + const w4b2::Bench b(s); + CHECK(w4b2::Dots3NoteDeviceRefusal(b.params).empty()); // the CONFIG is fine + CHECK_THROWS_WITH_AS(b.RunPrefillThenDecode(), doctest::Contains("index_topk"), + std::runtime_error); + } + // (d) the SAME `index_topk` on a config with NO full layer RUNS, because a + // sliding layer carries no indexer at all (`self.indexer = None`, + // model.py:432). This is the narrowing W4b-2 made, asserted rather than + // described — it is what stops a pure-SWA config being refused for a + // mechanism it does not have. + { + w4b2::Spec s; + s.index_topk = 2; + s.kinds = {vllm::Dots3NoteLayerKind::kSlidingAttention, + vllm::Dots3NoteLayerKind::kSlidingAttention}; + const w4b2::Bench b(s); + CHECK(w4b2::Dots3NoteDeviceRefusal(b.params).empty()); + CHECK_NOTHROW((void)b.RunPrefillThenDecode()); + } + // (e) a nextn tail — W10, unchanged. + { + nlohmann::json doc = w4b2::ConfigDoc(w4b2::Spec{}); + doc["num_nextn_predict_layers"] = 1; + TempConfig cfg(doc); + const Dots3NoteParams p = ParseDots3NoteParams(LoadHfConfig(cfg.path())); + const std::string why = w4b2::Dots3NoteDeviceRefusal(p); + CHECK(why.find("nextn") != std::string::npos); + CHECK(why.find("W10") != std::string::npos); + } + // (f) a KV cache whose row disagrees with the config it was built from. The + // config-level checks cannot see this — an engine allocates the cache + // separately — so the PER-STEP assertion stays, and this is what makes it + // reached rather than defensive decoration. + { + const w4b2::Bench b; + CHECK(w4b2::Dots3NoteDeviceRefusal(b.params).empty()); + CHECK_THROWS_WITH_AS( + b.RunPrefillThenDecode(nullptr, b.params.physical_latent_row() + 3), + doctest::Contains("PHYSICAL row"), std::runtime_error); + } +} + +TEST_CASE( + "dots3-note W4b-2: a windowed PREFILL that also has chunked CONTEXT is " + "refused BY NAME inside the shared seam") { + // Upstream's windowed prefill caps its gather at + // `min(seq_len, query_len + W - 1)` and runs one varlen call per request group + // (attention.py:206, :594-654); it never merges context chunks under a window, + // so there is no windowed form of `forward_mha`'s LSE merge to mirror. The + // seam refuses rather than merging an UNwindowed context into a windowed + // suffix, which would be silently wrong. W4b-3 owes it. + // + // The operands are deliberately EMPTY: the refusal is at the top of the + // function, before anything is read, which is the same discipline W4a's + // finding F6 applied to the gate's dtype check. The control below is what + // makes that a statement about the WINDOW rather than about the empty + // tensors — with `sliding_window == 0` the same call reaches the op and fails + // there instead, with a different type and a different message. + vt::Queue q{vt::Device{vt::DeviceType::kCPU, 0}, nullptr}; + vllm::mla::MlaPrefillContextBuffers bufs{}; + std::vector chunks(1); + vt::Tensor empty{}; + vllm::mla::MlaUpProjectFn up; + CHECK_THROWS_WITH_AS( + vllm::mla::ForwardMlaPrefillMha(q, empty, empty, empty, empty, empty, empty, empty, + chunks, up, 1.0f, 1, 0, bufs, empty, empty, + /*sliding_window=*/513), + doctest::Contains("SLIDING-WINDOW layer with chunked CONTEXT"), + std::invalid_argument); + // The CONTROL. Without a window the guard is a NOT-TAKEN branch, so the call + // proceeds into `vt::MlaPrefillAttention` and fails on the empty operands — + // a `std::runtime_error` from VT_CHECK, not the `std::invalid_argument` + // above. Asserting the TYPE is what stops this control passing vacuously. + CHECK_THROWS_AS( + vllm::mla::ForwardMlaPrefillMha(q, empty, empty, empty, empty, empty, empty, empty, + chunks, up, 1.0f, 1, 0, bufs, empty, empty, + /*sliding_window=*/0), + std::runtime_error); + // And with a window but NO context the guard does not fire either: the chunk + // list is what pairs with it, so an ordinary windowed prefill is unaffected. + const std::vector no_chunks; + CHECK_THROWS_AS( + vllm::mla::ForwardMlaPrefillMha(q, empty, empty, empty, empty, empty, empty, empty, + no_chunks, up, 1.0f, 1, 0, bufs, empty, empty, + /*sliding_window=*/513), + std::runtime_error); } diff --git a/tests/vllm/models/test_dots3_note_scaffold.cpp b/tests/vllm/models/test_dots3_note_scaffold.cpp index df9bbff78..5e06c798f 100644 --- a/tests/vllm/models/test_dots3_note_scaffold.cpp +++ b/tests/vllm/models/test_dots3_note_scaffold.cpp @@ -1737,9 +1737,16 @@ TEST_CASE("dots3-note: the forward REFUSES BY NAME through the REAL loaded model CHECK_THROWS_WITH_AS(reg.factory->forward(*model, input), doctest::Contains("Dots3NoteForCausalLM forward"), std::runtime_error); - // ...name the missing piece rather than only failing... - CHECK_THROWS_WITH_AS(reg.factory->forward(*model, input), - doctest::Contains("sliding-window MLA"), + // ...name the missing piece rather than only failing. The RELEASED config's + // first unrepresentable layer is layer 1's MoE (W5): W4b-2 put both attention + // geometries — full AND sliding-window — on the decode path, so the sliding + // layer at index 2 is no longer what stops this checkpoint. Naming the piece + // the released config ACTUALLY trips on is the point of the assertion; a + // string that outlives the refusal it describes is the failure this row keeps + // recording. + CHECK_THROWS_WITH_AS(reg.factory->forward(*model, input), doctest::Contains("MoE layer"), + std::runtime_error); + CHECK_THROWS_WITH_AS(reg.factory->forward(*model, input), doctest::Contains("W5"), std::runtime_error); // ...and point at the record that owns the brick. CHECK_THROWS_WITH_AS(reg.factory->forward(*model, input), diff --git a/tests/vt/test_ops_mla_attn.cpp b/tests/vt/test_ops_mla_attn.cpp index c20fb8819..3df9c2430 100644 --- a/tests/vt/test_ops_mla_attn.cpp +++ b/tests/vt/test_ops_mla_attn.cpp @@ -37,6 +37,7 @@ #include #include #include +#include #include #include @@ -620,3 +621,174 @@ TEST_CASE("mla_decode rejects malformed operands") { std::runtime_error); } } + +// ─────────────────────────────────────────────────────────────────────────── +// THE SLIDING WINDOW (dots3-note W4b-2, #699). +// +// `MlaDecodeAttentionArgs::window_size` is FlashAttention's `(left, right)` +// pair. An MLA decode query IS the last position of its own sequence, so +// `{left, 0}` keeps keys `[seq_len - 1 - left, seq_len)`. +// +// THE ORACLE IS THE OP ITSELF, ON A DIFFERENT INPUT, and that is deliberate: +// no windowed reference is written here, because a reference that computed +// `seq_len - 1 - left` a second time would share the arithmetic it is supposed +// to check (the shared-helper trap this project keeps naming). Instead the +// windowed call over a length-n paged sequence is compared against an +// UNWINDOWED call over a freshly built cache holding exactly that request's +// last `min(W, n)` keys — a path already gated against the ported `ref_mla` +// oracle above. The two agree only if the window keeps precisely the right key +// set, and the case also pins the boundary from the other side by showing that +// a window one key WIDER is a different answer. +// +// UPSTREAM: `_forward_swa_mqa` gathers `[max(seq_len - GATHER_LEN, 0), ...)` +// (`vllm/models/dots3_note/nvidia/attention.py:76-79` @ `bc2d63e650`) and masks +// with `kv_positions >= query_position - WINDOW_SIZE + 1` (`:152`) and +// `kv_positions <= query_position` (`:151`). +namespace { + +// One request's last `keep` keys, copied into a fresh single-page cache. The +// page is `keep` wide, the block table is {0} and `seq_lens` is {keep}, so the +// unwindowed op reads exactly those keys and nothing else. +void RunTruncatedCpu(const Case& c, int b, int keep, std::vector& out, + const MlaDecodeAttentionArgs& base) { + const int n = c.seq_lens[static_cast(b)]; + const int start = std::max(0, n - keep); + const int len = n - start; + REQUIRE(len > 0); + std::vector page(static_cast(len) * c.head_size, 0.0f); + for (int j = 0; j < len; ++j) { + const int src = start + j; + const int blk = c.block_table[static_cast(b) * c.max_blocks + src / c.block_size]; + const float* row = + c.cache.data() + + (static_cast(blk) * c.block_size + src % c.block_size) * c.head_size; + std::copy(row, row + c.head_size, page.begin() + static_cast(j) * c.head_size); + } + std::vector q(c.q.begin() + static_cast(b) * c.heads * c.head_size, + c.q.begin() + static_cast(b + 1) * c.heads * c.head_size); + std::vector bt{0}; + std::vector sl{len}; + out.assign(static_cast(c.heads) * c.v_head_dim, 0.0f); + Tensor t_out = Contig(out.data(), DType::kF32, Cpu(), {1, c.heads, c.v_head_dim}); + Tensor t_q = Contig(q.data(), DType::kF32, Cpu(), {1, c.heads, c.head_size}); + Tensor t_c = Contig(page.data(), DType::kF32, Cpu(), {1, len, c.head_size}); + Tensor t_bt = Contig(bt.data(), DType::kI32, Cpu(), {1, 1}); + Tensor t_sl = Contig(sl.data(), DType::kI32, Cpu(), {1}); + MlaDecodeAttentionArgs args = base; + args.window_size = std::nullopt; // the UNWINDOWED path, over the sliced keys + args.max_seq_len = len; + Queue qq = CpuQ(); + vt::MlaDecodeAttention(qq, t_out, nullptr, t_q, t_c, t_bt, t_sl, args); +} + +} // namespace + +TEST_CASE("mla_decode CPU: a sliding window attends the LAST W keys and nothing else") { + // Ragged, and every length straddles a different relation to the window: + // 5 < W (the window is the whole sequence), 13 == W exactly, 64 and 100 are + // several pages past it. W is 13 — NOT a multiple of the 16-wide page — so + // the window start lands INSIDE a page and a port that rounded to a page + // boundary is caught. + const Case c = MakeCase({5, 13, 64, 100}, 4, 101u); + constexpr int kWindow = 13; + MlaDecodeAttentionArgs args; + args.scale = static_cast(LiteScale()); + args.window_size = vt::AttentionWindow{kWindow - 1, 0}; + std::vector got; + RunCpu(c, got, nullptr, args); + + int64_t requests_the_window_cut = 0, keys_dropped = 0; + for (int b = 0; b < c.bs; ++b) { + const int n = c.seq_lens[static_cast(b)]; + if (n > kWindow) { + ++requests_the_window_cut; + keys_dropped += n - kWindow; + } + std::vector want; + RunTruncatedCpu(c, b, kWindow, want, args); + const std::vector slice( + got.begin() + static_cast(b) * c.heads * c.v_head_dim, + got.begin() + static_cast(b + 1) * c.heads * c.v_head_dim); + CHECK(MaxAbsDiff(slice, want) < 1e-5); + } + MESSAGE("mla_decode window " << kWindow << ": " << requests_the_window_cut << " of " + << c.bs << " requests really lose keys (" + << keys_dropped << " dropped in total)"); + // The fixture has to BITE, or every assertion above is vacuous. + REQUIRE(requests_the_window_cut == 2); + REQUIRE(keys_dropped == (64 - kWindow) + (100 - kWindow)); + + // The other side of the boundary: a window one key WIDER is a different + // answer on every request the window actually cut. Without this the case + // would pass for a port that kept W + 1 keys. + for (int b = 0; b < c.bs; ++b) { + if (c.seq_lens[static_cast(b)] <= kWindow) continue; + std::vector wider; + RunTruncatedCpu(c, b, kWindow + 1, wider, args); + const std::vector slice( + got.begin() + static_cast(b) * c.heads * c.v_head_dim, + got.begin() + static_cast(b + 1) * c.heads * c.v_head_dim); + CHECK(MaxAbsDiff(slice, wider) > 1e-3); + } + + // And the CONTROL that says the window did anything at all: the same call + // with no window is a different answer. + MlaDecodeAttentionArgs full = args; + full.window_size = std::nullopt; + std::vector unwindowed; + RunCpu(c, unwindowed, nullptr, full); + CHECK(MaxAbsDiff(got, unwindowed) > 1e-3); +} + +TEST_CASE("mla_decode CPU: a window at least as wide as the sequence is BIT-IDENTICAL") { + // The ABSENT state is a not-taken branch, and the widest possible window is + // the same key set as no window at all — so the two must agree to the LAST + // BIT, not merely to a tolerance. This is what says the window is a loop + // BOUND and not a mask applied afterwards. + const Case c = MakeCase({1, 16, 33, 64}, 4, 202u); + MlaDecodeAttentionArgs args; + args.scale = static_cast(LiteScale()); + std::vector unwindowed; + RunCpu(c, unwindowed, nullptr, args); + args.window_size = vt::AttentionWindow{63, 0}; // == the longest sequence + std::vector windowed; + RunCpu(c, windowed, nullptr, args); + REQUIRE(unwindowed.size() == windowed.size()); + for (size_t i = 0; i < unwindowed.size(); ++i) CHECK(unwindowed[i] == windowed[i]); +} + +TEST_CASE("mla_decode rejects a window shape upstream never produces") { + const Case c = MakeCase({8}, 4, 303u); + std::vector out; + MlaDecodeAttentionArgs args; + args.scale = static_cast(LiteScale()); + // A positive RIGHT bound could only admit keys past the decode query, which + // IS the last position of its sequence. + args.window_size = vt::AttentionWindow{3, 1}; + CHECK_THROWS_WITH_AS(RunCpu(c, out, nullptr, args), doctest::Contains("window_size.right"), + std::runtime_error); + args.window_size = vt::AttentionWindow{-1, 0}; + CHECK_THROWS_WITH_AS(RunCpu(c, out, nullptr, args), doctest::Contains("window_size.left"), + std::runtime_error); +} + +TEST_CASE("CUDA mla_decode: the sliding window matches the CPU reference") { + if (!HasCuda()) return; + // The split-KV schedule partitions the WINDOWED range, so this also covers + // the case where a split is empty because the window is shorter than the + // sequence. Both `num_kv_splits` arms are driven: the derived one and the + // forced single split. + const Case c = MakeCase({5, 13, 64, 100}, 4, 101u); + for (int splits : {0, 1, 4}) { + MlaDecodeAttentionArgs args; + args.scale = static_cast(LiteScale()); + args.window_size = vt::AttentionWindow{12, 0}; + args.num_kv_splits = splits; + args.max_seq_len = 100; + std::vector cpu; + RunCpu(c, cpu, nullptr, args); + std::vector gpu; + RunCuda(c, gpu, nullptr, args, /*bf16=*/false); + CHECK(MaxAbsDiff(gpu, cpu) < 1e-3); + } +} diff --git a/tests/vt/test_ops_mla_prefill.cpp b/tests/vt/test_ops_mla_prefill.cpp index fbc0f7f06..3cf734c6a 100644 --- a/tests/vt/test_ops_mla_prefill.cpp +++ b/tests/vt/test_ops_mla_prefill.cpp @@ -32,6 +32,7 @@ #include #include #include +#include #include #include @@ -475,3 +476,298 @@ TEST_CASE("CUDA MLA prefill is run-to-run bit-exact") { } } } + +// ─────────────────────────────────────────────────────────────────────────── +// THE SLIDING WINDOW (dots3-note W4b-2, #699). +// +// `MlaPrefillAttentionArgs::window_size` is FlashAttention's `(left, right)` +// pair, exactly as upstream hands it: `run_sliding_window` calls +// `_flash_attn_varlen_diff_headdims(..., causal=True, +// window_size=(sliding_window - 1, 0))` +// (`vllm/models/dots3_note/nvidia/attention.py:279-305` @ `bc2d63e650`, the +// pair at `:300`). +// +// THE ORACLE IS THE OP ITSELF, ON A DIFFERENT INPUT. `RefPrefill` is NOT given +// a window, deliberately: a reference that recomputed `iq + (lk - lq) - left` +// would share the bottom-right arithmetic it is supposed to check. Instead the +// windowed multi-query call is compared against an EXPANDED batch in which +// every query becomes its own single-query request carrying only the keys its +// window admits, run through the UNWINDOWED path that `RefPrefill` already +// gates. With `lq == 1` the bottom-right causal bound admits every key handed +// in, so the expansion needs no mask of its own. +namespace { + +// Expand `c` into one request per query, each carrying only that query's +// windowed key range, and run the UNWINDOWED op over it. Returns the [total_q, +// heads, dv] output in the original query order. +std::vector RunExpandedWindowCpu(const std::vector& q_lens, + const std::vector& k_lens, int heads, + int win_left, const std::vector& q, + const std::vector& k, + const std::vector& v, double scale, + int64_t* keys_dropped) { + const std::vector cu_q = Cumsum(q_lens); + const std::vector cu_k = Cumsum(k_lens); + const int total_q = cu_q.back(); + std::vector sub_q(static_cast(total_q), 1); + std::vector sub_k; + std::vector kk, vv; + *keys_dropped = 0; + for (size_t b = 0; b + 1 < cu_q.size(); ++b) { + const int lq = cu_q[b + 1] - cu_q[b]; + const int lk = cu_k[b + 1] - cu_k[b]; + for (int iq = 0; iq < lq; ++iq) { + const int p = iq + (lk - lq); // the bottom-right query position + const int hi = std::min(lk, p + 1); // the causal bound, inclusive of p + const int lo = std::max(0, p - win_left); + REQUIRE(hi >= lo); + sub_k.push_back(hi - lo); + *keys_dropped += lo; + for (int j = lo; j < hi; ++j) { + const size_t kr = (static_cast(cu_k[b] + j) * heads) * kQkHeadDim; + kk.insert(kk.end(), k.begin() + static_cast(kr), + k.begin() + static_cast(kr + static_cast(heads) * kQkHeadDim)); + const size_t vr = (static_cast(cu_k[b] + j) * heads) * kVHeadDim; + vv.insert(vv.end(), v.begin() + static_cast(vr), + v.begin() + static_cast(vr + static_cast(heads) * kVHeadDim)); + } + } + } + const std::vector cu_sq = Cumsum(sub_q); + const std::vector cu_sk = Cumsum(sub_k); + const int total_k = cu_sk.back(); + std::vector out(static_cast(total_q) * heads * kVHeadDim, + std::numeric_limits::quiet_NaN()); + std::vector qq = q; + std::vector a = cu_sq, bb = cu_sk; + Tensor tq = Contig(qq.data(), DType::kF32, Cpu(), {total_q, heads, kQkHeadDim}); + Tensor tk = Contig(kk.data(), DType::kF32, Cpu(), {std::max(total_k, 1), heads, kQkHeadDim}); + Tensor tv = Contig(vv.data(), DType::kF32, Cpu(), {std::max(total_k, 1), heads, kVHeadDim}); + Tensor to = Contig(out.data(), DType::kF32, Cpu(), {total_q, heads, kVHeadDim}); + Tensor tcq = Contig(a.data(), DType::kI32, Cpu(), {static_cast(a.size())}); + Tensor tck = Contig(bb.data(), DType::kI32, Cpu(), {static_cast(bb.size())}); + MlaPrefillAttentionArgs args; + args.scale = static_cast(scale); + args.causal = true; // lq == 1 everywhere, so this admits every key handed in + Queue q0 = CpuQ(); + vt::MlaPrefillAttention(q0, to, nullptr, tq, tk, tv, tcq, tck, args); + return out; +} + +} // namespace + +TEST_CASE("MLA prefill CPU: a sliding window keeps exactly the last W keys per query") { + const std::vector q_lens{7, 1, 33, 16}; + const std::vector k_lens = q_lens; // a fresh prompt: seq_len == query_len + constexpr int kWindow = 5; + const int h = 4; + const double scale = LiteScale(); + const std::vector cu_q = Cumsum(q_lens); + const int total_q = cu_q.back(); + const auto q = RandF32(static_cast(total_q) * h * kQkHeadDim, 909u); + const auto k = RandF32(static_cast(total_q) * h * kQkHeadDim, 911u); + const auto v = RandF32(static_cast(total_q) * h * kVHeadDim, 913u); + + std::vector out(static_cast(total_q) * h * kVHeadDim, + std::numeric_limits::quiet_NaN()); + std::vector cu_a = cu_q, cu_b = cu_q; + Tensor tq = Contig(const_cast(q.data()), DType::kF32, Cpu(), + {total_q, h, kQkHeadDim}); + Tensor tk = Contig(const_cast(k.data()), DType::kF32, Cpu(), + {total_q, h, kQkHeadDim}); + Tensor tv = Contig(const_cast(v.data()), DType::kF32, Cpu(), + {total_q, h, kVHeadDim}); + Tensor to = Contig(out.data(), DType::kF32, Cpu(), {total_q, h, kVHeadDim}); + Tensor tcq = Contig(cu_a.data(), DType::kI32, Cpu(), {static_cast(cu_a.size())}); + Tensor tck = Contig(cu_b.data(), DType::kI32, Cpu(), {static_cast(cu_b.size())}); + MlaPrefillAttentionArgs args; + args.scale = static_cast(scale); + args.causal = true; + args.window_size = vt::AttentionWindow{kWindow - 1, 0}; + Queue q0 = CpuQ(); + vt::MlaPrefillAttention(q0, to, nullptr, tq, tk, tv, tcq, tck, args); + + int64_t dropped = 0; + const std::vector want = + RunExpandedWindowCpu(q_lens, k_lens, h, kWindow - 1, q, k, v, scale, &dropped); + MESSAGE("MLA prefill window " << kWindow << ": " << dropped + << " (query, key) pairs dropped across " << total_q + << " queries"); + // The fixture has to BITE: 7+1+33+16 = 57 queries, and every query past + // position 4 in its own request loses at least one key. + REQUIRE(dropped > 0); + CHECK(MaxAbsDiff(out, want) < 2e-4); + + // The CONTROL: without the window the same call is a different answer. + std::vector full(static_cast(total_q) * h * kVHeadDim, + std::numeric_limits::quiet_NaN()); + Tensor tof = Contig(full.data(), DType::kF32, Cpu(), {total_q, h, kVHeadDim}); + MlaPrefillAttentionArgs no_win = args; + no_win.window_size = std::nullopt; + vt::MlaPrefillAttention(q0, tof, nullptr, tq, tk, tv, tcq, tck, no_win); + CHECK(MaxAbsDiff(out, full) > 1e-3); + + // A window at least as wide as the longest request is BIT-IDENTICAL to no + // window: the absent state is a not-taken branch, not a mask. + std::vector wide(static_cast(total_q) * h * kVHeadDim, + std::numeric_limits::quiet_NaN()); + Tensor tow = Contig(wide.data(), DType::kF32, Cpu(), {total_q, h, kVHeadDim}); + MlaPrefillAttentionArgs wide_args = args; + wide_args.window_size = vt::AttentionWindow{32, 0}; // == the longest request + vt::MlaPrefillAttention(q0, tow, nullptr, tq, tk, tv, tcq, tck, wide_args); + for (size_t i = 0; i < full.size(); ++i) CHECK(full[i] == wide[i]); +} + +// The CUDA half of the windowed prefill (dots3-note W4b-2, #699). +// +// WHY IT EXISTS SEPARATELY. The CPU case above gates `cpu_mla_prefill.cpp`, and +// `test_ops_mla_attn`'s "CUDA mla_decode: the sliding window matches the CPU +// reference" gates `cuda_mla_attn.cu`. Neither reaches the FA-2 MLA-prefill +// launcher's `is_local` normalization in `cuda_flash_attn_fa2.cu`, which is the +// OTHER CUDA file W4b-2 changed. Without this case that normalization has no +// gate on any device, and a later GPU lease would discharge the decode half +// while the record read as though it had closed both. +// +// IT SKIPS ON A BOX WITH NO CUDA DEVICE, which is where W4b-2 and this repair +// ran. A skip is not a pass: until the row's designated host is reachable this +// case has never EXECUTED, and the spec's §4.8 and `## Owed` say so. +// +// WHAT IT DOES NOT ASSERT, and why that is not a weaker bar. The CPU case +// asserts that a window at least as wide as the longest request is +// BIT-IDENTICAL to no window, because on CPU the absent state is a not-taken +// branch. That claim is FALSE on the GPU by construction and asserting it would +// be a defect: a finite `window_size` sets `p.is_causal = false` and dispatches +// FA-2's LOCAL specialization, a different compile-time template from the +// causal one, so agreement there is numerical rather than byte-wise. +TEST_CASE("CUDA MLA prefill: the sliding window matches the CPU reference") { + if (!HasCuda()) return; + Backend& b = vt::GetBackend(DeviceType::kCUDA); + QueueGuard g(b); + + // The CPU case's fixture, unchanged, so the two halves are the same + // experiment: 7+1+33+16 = 57 queries and a window of 5 that cuts every query + // past position 4 in its own request. + const std::vector q_lens{7, 1, 33, 16}; + const std::vector k_lens = q_lens; // a fresh prompt: seq_len == query_len + constexpr int kWindow = 5; + const int h = kHeadsLite; + const double scale = LiteScale(); + const std::vector cu_q = Cumsum(q_lens); + const int total_q = cu_q.back(); + + // The CUDA MLA prefill is bf16-only — the launcher instantiates + // `cutlass::bfloat16_t` at every head-dim arm. Round-trip the inputs so the + // CPU arm reads EXACTLY the values the kernel does; otherwise the comparison + // charges bf16 input quantisation to the window. + const auto qf = Bf16Round(RandF32(static_cast(total_q) * h * kQkHeadDim, 909u)); + const auto kf = Bf16Round(RandF32(static_cast(total_q) * h * kQkHeadDim, 911u)); + const auto vf = Bf16Round(RandF32(static_cast(total_q) * h * kVHeadDim, 913u)); + + // ── the CPU arm, windowed and unwindowed ───────────────────────────────── + auto run_cpu = [&](std::optional win) { + std::vector q = qf, k = kf, v = vf; + std::vector ca = cu_q, cb = cu_q; + std::vector out(static_cast(total_q) * h * kVHeadDim, + std::numeric_limits::quiet_NaN()); + Tensor tq = Contig(q.data(), DType::kF32, Cpu(), {total_q, h, kQkHeadDim}); + Tensor tk = Contig(k.data(), DType::kF32, Cpu(), {total_q, h, kQkHeadDim}); + Tensor tv = Contig(v.data(), DType::kF32, Cpu(), {total_q, h, kVHeadDim}); + Tensor to = Contig(out.data(), DType::kF32, Cpu(), {total_q, h, kVHeadDim}); + Tensor tcq = Contig(ca.data(), DType::kI32, Cpu(), {static_cast(ca.size())}); + Tensor tck = Contig(cb.data(), DType::kI32, Cpu(), {static_cast(cb.size())}); + MlaPrefillAttentionArgs args; + args.scale = static_cast(scale); + args.causal = true; + args.window_size = win; + Queue q0 = CpuQ(); + vt::MlaPrefillAttention(q0, to, nullptr, tq, tk, tv, tcq, tck, args); + return out; + }; + const std::vector cpu_win = run_cpu(vt::AttentionWindow{kWindow - 1, 0}); + const std::vector cpu_none = run_cpu(std::nullopt); + + // ── the CUDA arm, windowed and unwindowed ──────────────────────────────── + const auto qb = ToBf16(qf); + const auto kb = ToBf16(kf); + const auto vb = ToBf16(vf); + auto cu_v = cu_q; + DeviceTensor dq(b, g.q, DType::kBF16, {total_q, h, kQkHeadDim}, qb.data()); + DeviceTensor dk(b, g.q, DType::kBF16, {total_q, h, kQkHeadDim}, kb.data()); + DeviceTensor dv(b, g.q, DType::kBF16, {total_q, h, kVHeadDim}, vb.data()); + DeviceTensor dcu(b, g.q, DType::kI32, {static_cast(cu_v.size())}, cu_v.data()); + + auto run_cuda = [&](std::optional win) { + // NaN-poison the output: a kernel that fails to write FAILS the gate + // rather than reading back whatever the allocator handed us. + const std::vector poison(static_cast(total_q) * h * kVHeadDim, + vt::F32ToBF16(std::numeric_limits::quiet_NaN())); + DeviceTensor dout(b, g.q, DType::kBF16, {total_q, h, kVHeadDim}, poison.data()); + MlaPrefillAttentionArgs args; + args.scale = static_cast(scale); + args.causal = true; + args.window_size = win; + vt::MlaPrefillAttention(g.q, dout.tensor(), nullptr, dq.tensor(), dk.tensor(), + dv.tensor(), dcu.tensor(), dcu.tensor(), args); + b.Synchronize(g.q); + std::vector got(poison.size()); + dout.Download(g.q, got.data()); + return FromBf16(got); + }; + const std::vector gpu_win = run_cuda(vt::AttentionWindow{kWindow - 1, 0}); + const std::vector gpu_none = run_cuda(std::nullopt); + + // The bar is the file's own CUDA-vs-reference bar: FA-2 accumulates in bf16 + // against a two-pass f32 CPU arm at these dims. + CHECK(MaxAbsDiff(gpu_win, cpu_win) < 3e-2); + + // THE WINDOW HAS TO BITE ON THE DEVICE, not only on the CPU. Without this the + // case above would pass on a launcher that dropped `window_size` on the floor + // — which is the exact defect the `is_local` normalization exists to prevent, + // because FA-2's `Is_causal` template IGNORES `window_size_left`. + CHECK(MaxAbsDiff(gpu_win, gpu_none) > 1e-2); + // …and the unwindowed device call still agrees with the unwindowed CPU one, + // so the difference above is the window and not a broken device arm. + CHECK(MaxAbsDiff(gpu_none, cpu_none) < 3e-2); + + // The independent semantic oracle, the same one the CPU case uses: every + // query re-expressed as its own single-query request carrying only the keys + // its window admits, run UNWINDOWED. This is what makes the comparison more + // than "two spellings of one arithmetic". + int64_t dropped = 0; + const std::vector expanded = + RunExpandedWindowCpu(q_lens, k_lens, h, kWindow - 1, qf, kf, vf, scale, &dropped); + MESSAGE("CUDA MLA prefill window " << kWindow << ": " << dropped + << " (query, key) pairs dropped across " << total_q + << " queries"); + REQUIRE(dropped > 0); + CHECK(MaxAbsDiff(gpu_win, expanded) < 3e-2); +} + +TEST_CASE("MLA prefill rejects a window shape upstream never produces") { + const std::vector lens{4}; + const int h = 2; + std::vector cu = Cumsum(lens); + auto q = RandF32(4 * static_cast(h) * kQkHeadDim, 1u); + auto k = q; + auto v = RandF32(4 * static_cast(h) * kVHeadDim, 2u); + std::vector out(4 * static_cast(h) * kVHeadDim, 0.0f); + Tensor tq = Contig(q.data(), DType::kF32, Cpu(), {4, h, kQkHeadDim}); + Tensor tk = Contig(k.data(), DType::kF32, Cpu(), {4, h, kQkHeadDim}); + Tensor tv = Contig(v.data(), DType::kF32, Cpu(), {4, h, kVHeadDim}); + Tensor to = Contig(out.data(), DType::kF32, Cpu(), {4, h, kVHeadDim}); + Tensor tcq = Contig(cu.data(), DType::kI32, Cpu(), {static_cast(cu.size())}); + Queue q0 = CpuQ(); + MlaPrefillAttentionArgs args; + args.scale = 0.1f; + args.causal = true; + args.window_size = vt::AttentionWindow{2, 1}; + CHECK_THROWS_WITH_AS(vt::MlaPrefillAttention(q0, to, nullptr, tq, tk, tv, tcq, tcq, args), + doctest::Contains("window_size.right"), std::runtime_error); + // A NON-causal window has no FlashAttention spelling — the local mask + // REPLACES the causal specialization, so "everything forward, windowed + // backward" would need an infinite right bound. Upstream never asks for one. + args.window_size = vt::AttentionWindow{2, 0}; + args.causal = false; + CHECK_THROWS_WITH_AS(vt::MlaPrefillAttention(q0, to, nullptr, tq, tk, tv, tcq, tcq, args), + doctest::Contains("causal=true"), std::runtime_error); +}