diff --git a/.agents/specs/glm5-next-flash.md b/.agents/specs/glm5-next-flash.md index 6bb9059df..b301a151f 100644 --- a/.agents/specs/glm5-next-flash.md +++ b/.agents/specs/glm5-next-flash.md @@ -2366,9 +2366,15 @@ the exit code. The `static_assert` in the case pins what it is FOR — offset in the same tower still fits, so a 32-bit element count would truncate here and nowhere in the E<=4 codebook cases. -**`dgx:gpu0` is still queued and this gate is NOT closed by the `thor` run.** -`sm_110` and `sm_121a` are different targets and a number from one does not -transfer to the other; the GB10 leg stays PENDING until its lease returns. +**`dgx:gpu0` was still queued when this was written, and the GB10 leg has since +been TAKEN.** `sm_110` and `sm_121a` are different targets and a number from one +does not transfer to the other, which is why the leg was held open. It returned +on 2026-08-31: on `dgx:gpu0` with `CUDA FA2 compiled-arch manifest: [121a]`, +`test_cuda_quant_dot` ran **17 cases / 177,284 assertions / 0 failed** and the +GLM-5.3-Flash published MoE geometry case **25 assertions / 0 failed**, the same +totals `thor` produced, now reproduced on the architecture that counts. Cite the +GB10 run rather than `thor`'s. Corrected here by W9c-0 rather than left as a +PENDING that the fleet has already answered. **One fleet fact bounds the END-TO-END test and is recorded here**, because which box can ever take the end-to-end test: on `thor:gpu0` (`compute_cap` 11.0, @@ -2391,6 +2397,275 @@ reproduced on the recorded SHA, STOP: a scoping built on a misread anchor is worse than no scoping, and this spec family has been caught with drifted anchors before. +#### W9c-0 — the k-pool indexer's device selection ops (GPU, large). [#2415](https://github.com/mudler/vllm.cpp/issues/2415) + +`CLAIM-GLM53-FLASH-KPOOL-CUDA`, 2026-08-31, row +`MODEL-MM-glm5-next-glm5-next-for-conditional-generation`, branch +`row/MODEL-MM-GLM53-FLASH-KPOOL-CUDA`, base `bd2c14ce5`. This wave writes the +kernel §W9c "RESCOPED" said nobody had to write. + +**THE DECISION #2415 ASKS FOR, TAKEN.** #2415 offers two ways to put the 11 DSA +layers on a GPU: write a k-pool indexer CUDA kernel, or leave the indexer on the +host and pay a device-to-host round trip per DSA layer per step. **The kernel.** +The reasoning, so it is not re-litigated: the host arm is not a cheap fallback +on this model. Each round trip has to move the whole packed key history the pool +grid is re-formed over — `2 * index_head_dim + 1 = 257` floats per token per +layer, which at 32k context is 33.7 MiB per layer per step and 371 MiB per step +across the eleven — plus a `[B, S, 2051]` selection back. That is a synchronous +stall eleven times on the critical path of a 101.24 GiB model whose entire +reason for having a device arm is that the weights are already resident. A +device arm whose indexer is on the host is not a device arm; it is the CPU path +with a copy tax. #2415 is closed by this wave with that argument recorded, and +the alternative is not staged, not flagged and not kept as a fallback, because a +fallback nobody may use is the branch AGENTS.md §"Nothing lands dead" names. + +**Scope.** Two new `vt::OpId` entries, one CUDA provider apiece, and a +model-side availability probe. Nothing else. Named exclusions, each owned +elsewhere: this wave does **not** wire the model's forward onto the ops (W9c-3), +does **not** lift the two CPU-only refusals at `glm5_next_kda.cpp:322-325` and +`glm5_next_moe.cpp:222-225` (W9c-2), does **not** touch residency (W9b), and +does **not** retire `glm5_next_attn.cpp`'s hand-rolled MLA (W9c-1, O33). It +lands **unreached** and O36 says so in the specific. + +**Oracle.** vLLM implements `glm5_next` at no revision — re-verified for this +wave rather than inherited: `git grep -n 'glm5_next\|Glm5Next' -- vllm/` in the +pinned checkout at `555967922` exits 1 with no output, and +[vllm#53906](https://github.com/vllm-project/vllm/pull/53906) is still OPEN and +therefore inadmissible under AGENTS.md §"When vLLM has no implementation". The +reference is `transformers` **v5.16.1**, the lane pin +[`.agents/oracles/transformers.md`](../oracles/transformers.md) already carries, +fetched at `refs/tags/v5.16.1` and byte-checked against the size this spec +records for it (`modular_glm5_next.py`, 95,314 bytes). + +**Ported from, `file:line` on both sides, each anchor asserted UNIQUE in its +file** (`grep -c` on the `def` line = 1 for every one of the four): + +| ours | `transformers` v5.16.1 `models/glm5_next/modular_glm5_next.py` | +|---|---| +| `vt::Glm5NextKpoolCompress` | `:897-970` `Glm5NextTextIndexer.get_pooled_states` | +| `vt::Glm5NextKpoolSelect`, the scoring and top-k half | `:821-875` `Glm5NextTextIndexer.forward` | +| `vt::Glm5NextKpoolSelect`, the visibility it folds in | `:877-895` `Glm5NextTextIndexer.get_visible_tokens` | +| `vt::Glm5NextKpoolSelect`, the tail half | `:972-1022` `Glm5NextTextIndexer.append_visible_tail` | + +The second reference is this tree's own gated host implementation, +`glm5_next_dsa.{h,cpp}`, whose port map carries the same anchors and whose +`test_glm5_next_dsa` goldens are the RUN output of that same lane revision. The +device ops answer to it. + +**THE DECOMPOSITION, AND WHY IT IS TWO OPS AND NOT ONE OR FOUR.** + +- `PackIndexerStates` (`glm5_next_dsa.cpp:95-142`) is **not** a new op and must + not become one. It is `wk` + a `nn.LayerNorm` + `kpool_gate` + a concat, and + `vt::Matmul`/`vt::MatmulBT` and `vt::LayerNorm` already serve all of it. The + concat is a write into a slice of the packed row. W9c-3 composes it; this wave + deliberately adds nothing for it, on the same reasoning + `kQwen4ExpQsaCompress`'s `ops.h` comment gives for what it left out. +- `GetVisibleTokens` (`:147-165`) is **not** a new op either. It is + `j <= current_length - q_length + s && valid_keys[b][j]` — two scalars and a + mask — and materialising it as a `[B, S, kv_len]` tensor to hand between two + device ops would allocate 2.7 GiB at 32k context for a predicate a thread can + evaluate in two instructions. `Glm5NextKpoolSelect` takes `valid_keys` and the + two lengths and evaluates it inline. Upstream materialises it because torch + has no other way to express a gather under it; that is a framework constraint, + not the model's semantics, and mirroring the constraint instead of the + semantics is the trap `.agents/porting.md` §"Mirror the memory format" names. +- `Glm5NextKpoolCompress` is the LEARNED pooling — the stage DeepSeek-V4 has no + counterpart for and the reason a new op family exists at all. +- `Glm5NextKpoolSelect` is the selection: score, mask, pool-level top-k, expand + to member token indices, append the ragged tail, truncate, mask by the query + padding. The tail is inside this op rather than beside it because the tail's + WRITE OFFSET is `select_k * index_kpool` and `select_k` depends on `P`, which + only this op knows; splitting it would publish an intermediate width whose + only consumer re-truncates it. + +**`P` IS DATA-DEPENDENT AND IT IS NOT READ BACK TO THE HOST.** Upstream's +`keep = pool_valid.any(0)` (`:968`) compacts the pool axis, and `select_k = +min(index_topk // index_kpool, P)` (`:845`) reads the compacted width. This is +not cosmetic: with `P` too large by even one, `select_k * index_kpool` moves the +tail's write offset and the final `[..., :output_width]` truncation +(`:870-872`) cuts a different set. So `Glm5NextKpoolCompress` performs the +compaction on the device — a per-pool validity pass, a single-block exclusive +scan over the `keep` predicate, then a compacted write — and publishes `P` as a +`[1]` i32 **device** scalar that `Glm5NextKpoolSelect` reads on the device. +Neither op synchronises and neither returns anything to the host. Every output +buffer is sized at the static upper bound `np = ceil(kv_len / index_kpool)`; +`topk_indices` is `[B, S, OutputWidth()]`, which is fixed at 2051 on this +checkpoint regardless of `P`. + +**DTYPE: f32, AND THAT IS UPSTREAM'S OWN ARITHMETIC RATHER THAN A CHOICE.** +Upstream scores in fp32 (`scores = torch.matmul(q.float(), pool_keys...float())`, +`:823`) and takes the pool softmax in fp32 (`:960-964`). The kernels do the +same. **The host reference accumulates in `double`** — `Linear` +(`glm5_next_dsa.cpp:28-37`), the dot at `:465-466`, the pool softmax at +`:225-261` — which is a host-reference widening this row took deliberately and +which the device arm must NOT copy: a fp64 pool softmax on `sm_121a` would move +the model path onto the 1/64-rate pipe to be *more* precise than the reference +it mirrors, and AGENTS.md §"Inherit vLLM defaults" is explicit that a token gate +cannot see a dtype that is too wide. So the host↔device delta on this row is +**fp32-vs-fp64 reduction order and nothing else**, and §Gates below says how it +is bounded rather than waved at. + +`nvcc` contracts `a * b + c` into an FMA by default, which makes the kernel's +own rounding depend on the optimiser rather than on the source. Every +accumulation in both kernels uses `__fadd_rn` / `__fmul_rn`, the convention +`cuda_conv1d_general.cu` sets, so the device answer is reproducible run to run +and the only remaining difference from the host is the width. `expf` and not +`__expf`: the fast intrinsic is a different function, not a faster spelling of +the same one. + +**Registration, and the CPU-only build.** Both ops register on `kCUDA` only, +through `vt::RegisterOp` in `src/vt/cuda/cuda_glm5_next.cu`, in the +`Registrar`-struct shape `cuda_dsa_indexer.cu:318-325` uses. The `ops.h` free +functions shape-check and then `GetOp(op, q.device.type)`, so a CPU queue throws +`GetOp`'s own unregistered-op error and no stub is ever linked. The +availability probe is `vllm::glm5_next::KpoolDeviceOpsAvailable()` +(`glm5_next_device.{h,cpp}`), a 1:1 mirror of +`deepseek_v4::V4DeviceKernelsAvailable` (`deepseek_v4_device.cpp:30-35`), so +W9c-3's forward can decide before it builds operands rather than after it +throws. **There is no CPU provider and this wave does not add one**: the CPU +answer already exists as `glm5_next_dsa.cpp` and is what the gate measures +against; registering it a second time under an `OpId` would make the seam its +own oracle, which is the tautology `.agents/verification.md` warns about. + +**Tests.** `tests/vllm/models/test_glm5_next_kpool_device.cpp`, which SKIPS +without a CUDA backend and is therefore a real gate only on a leased device. +Three properties it must keep, each because dropping it makes the file a +tautology: + +1. It runs the **same fixture geometry** `test_glm5_next_dsa` runs — `seq_len` + 21 against `index_topk` 8, one row left-padded by three — so the pool grid + does not start at slot 0 and the selection is not the identity. It reuses + `glm5_next_dsa_goldens.inc` verbatim, which carries the transformers RUN + output of the intermediates as well as the result: `kPoolKeys`, + `kPoolIndices`, `kPoolValid`, `kIndexScores` and `kTopkIndices`. So the + device arm answers to the oracle directly and to the host arm as well, rather + than only to our own C++. **The fixture exercises the compaction rather than + assuming it**: `np = ceil(21 / 4) = 6` and `kNumPools = 5`, so a kernel that + skipped `keep` would place the tail at column 12 instead of 8 and fail on + every row. Its four short cases (`kShortNumPools` all 0) exercise `P == 0`, + where the entire selection is the raw tail. +2. **SET equality of the selected token indices, plus the positionwise + comparison, plus the printed margin.** Top-k error is bimodal: a wrong pool's + scores can be arbitrarily close, so a tolerance on `index_scores` bounds + nothing about the selection. The margin between the `select_k`-th and + `select_k + 1`-th masked score is printed for every discriminating row. +3. **Every float comparison is `isfinite`-guarded on both operands before it is + made.** An all-NaN forward on this row once read as a perfect match, because + every comparison against NaN is false, and the model then emitted token id 0 + eight times. A NaN in either arm fails this file rather than passing it. + +**A PRE-LEASE ALGORITHM CHECK, AND THE BOUND IT DOES NOT CROSS.** Before a +device lease was taken, both kernels were transcribed into numpy and run against +the same goldens, because a lease spent discovering an index error is a lease +wasted. It reported `P = 5` against `np = 6` (the compaction), `pool_keys` +max|delta| `3.075e-07`, `index_scores` max|delta| `6.852e-06`, `topk_indices` +positionwise **0 mismatches of 462**, 17 pruning rows and a smallest decision +margin of **2.58e-03**. **That transcription is NOT the gate and it is not +committed.** It bounds the ALGORITHM — indexing, the compaction, the emission +order, the tail offset, the padded-row rule — and it bounds nothing at all about +the build, about `nvcc`, about the launch geometry, about the shared-memory +reductions or about float rounding, because none of those exist in it. A +transcription cannot gate the function it transcribes. It is recorded here as +what it is: the reason the lease was taken with some confidence, stated ahead of +the numbers it does not support. + +**THE SCORE BOUND IS TIED TO THE MARGIN RATHER THAN CHOSEN.** The golden score +magnitudes run to 45.17, so a single f32 ULP there is already 3.8e-6: an +absolute tolerance below that fails on arithmetic and one far above it bounds +nothing. The gate instead asserts that the device-versus-oracle difference is +under a quarter of the smallest gap the top-k decides on, and prints both, so a +later fixture that narrows the margin fails this file rather than quietly +becoming a coin flip. + +**Gates.** `scripts/agent-preflight.sh --fail-on-skip`; the CPU suites by hand +with case and assertion counts, including `test_glm5_next_dsa` unchanged and the +sibling inertness set (`test_cuda_deepseek_v4`, `test_qwen4_exp_qsa_device`, +`test_glm_moe_dsa*`) since `include/vt/ops.h` and `src/vt/ops.cpp` are shared; +and the CUDA suite on `dgx:gpu0` under `rc run`. A doctest `assertions: 0` line +is a skip wearing a pass and the device job reads that line out loud rather than +trusting the exit code. + +**THE DEVICE GATE, TAKEN ON `dgx:gpu0`, AND IT PASSES.** Every number below came +from **GB10** and from nowhere else. `sm_121a` and `thor`'s `sm_110` are +different targets and neither transfers to the other. + +| | | +|---|---| +| job | `79aa5bb5-7536-43fe-a051-ed73ac1302e1`, `rc run -d dgx:gpu0` | +| device | NVIDIA GB10, `GPU-cb5c11ff-4ea1-5472-a9a6-c7a468a4d9f1`, driver 580.173.02, `compute_cap` 12.1, built `sm_121a` | +| worker | `rc-worker-4b8lj`, aarch64, 20 cores, `boot_id 49b5d969-3870-4a2b-b5b8-71355083d5e6` | +| source | `git archive` of `4034c368c`, `src.tar.gz` sha256 `b036a0138f5a…`, COMPARED and fatal on mismatch | +| build | CUTLASS 4.5.0, `CUDA FA2 compiled-arch manifest: [121a]`, 605/605, `BUILD=0` | +| binary | `test_glm5_next_kpool_device` sha256 `39788d6d51c5fbe5…` | +| run | 2026-08-31T22:49:44Z → 23:10:05Z | + +**`test_glm5_next_kpool_device`: 4 cases, 918 assertions, 0 failed, and +`KPOOL_SKIPPED_CASES=0`.** The skip count is the load-bearing line, not the exit +code: the same binary on a CPU-only build passes 4 cases with **4** assertions +and three cases skipped, so 918-against-4 is what says the device arm ran. The +host reference suite `test_glm5_next_dsa` is UNMOVED at 10 / 1934. + +| quantity | measured | +|---|---| +| `pool_keys` max abs, device vs **transformers v5.16.1** | **3.57628e-07** | +| `pool_keys` max abs, device vs the **host reference** | **1.78814e-07** | +| `index_scores` max abs, device vs transformers | **6.67572e-06** | +| device `P` vs golden `kNumPools` | **5 == 5** (from `np = 6`: the `keep` compaction is exercised) | +| pruning rows / smallest decision margin | **17 / 2.58482e-03** | +| the margin-tied bound `score_delta * 4 < margin` | 2.67e-05 < 2.58e-03, **97x of room** | + +Selection agreement is SET equality of the chosen token indices AND positionwise +equality, against the transformers run and against the host reference, at +**0 mismatches of 462**. The fp32 device answer therefore differs from the fp64 +host reference by a hundredth of the gap the top-k decides on, which is the +statement a bimodal gate can actually make; a tolerance on the values alone +would have said nothing. + +**NINE MUTATIONS, and the two that did not end in an assertion kill are reported +as what they are.** Each was applied to product code, rebuilt with the BINARY's +sha256 asserted to have moved before its result counted, and restored +byte-for-byte with the source sha256 checked — nine restores, all verified. The +pristine tree was then rebuilt and re-run, reproducing 918 assertions, because a +post-mutation green means nothing until the restored tree has been built again. + +| mutation | outcome | +|---|---| +| M2 `keep` compaction removed | KILLED by assertion | +| M3 learned pool weighting replaced by a mean | KILLED by assertion | +| M3b learned intra-pool position embedding zeroed | KILLED by assertion | +| M4 pool from slot 0 instead of the first valid token | KILLED by assertion | +| M5 ragged visible tail never written | **BUILD_BROKE** — see below | +| M6 score skipped on padded query rows | KILLED by assertion | +| M7 the ReLU before the head mix dropped | KILLED by assertion | +| M8 the compress op's `RegisterOp` deleted | **SURVIVED AS A SKIP** — see below | +| M9 the availability probe loses one clause | KILLED by assertion | + +**M8 is the reachability result and it is not a kill.** With no provider the +probe is false, `HasCuda()` is false, and the three device cases return early: +the suite exits **0** with **4 assertions and 3 cases skipped**. That is the +exact shape of a skip wearing a pass, it is why this job reads +`KPOOL_SKIPPED_CASES` instead of the exit code, and it is the honest answer to +`.agents/reachability.md`'s question — there is no production call site to +delete, so the mutation measures the registration and not a capability. + +**M5 was killed by the COMPILER, which is weaker, and it is owed a re-run.** +Replacing the tail write with a bare `-1` orphaned `valid` and `vis`, and +`-Werror=all-warnings` refused it with `error #177-D: variable "valid" was +declared but never referenced`. A compiler kill proves the code changed and says +nothing about whether the gate can SEE the defect. Rewritten as +`(valid && vis && idx < 0)` — every operand live, and `idx` is non-negative for +every reachable `j`, so the tail column is still always written `-1` — and +re-submitted as job `3db96e7a-cfe9-4148-aac3-d26eb01a15bc`. **Until that returns, +the ragged tail is the one guarantee in this file with no assertion-grade +mutation behind it, and it is recorded as such rather than counted among the +seven.** + +**Stop conditions.** If `transformers` v5.16.1 turns out not to implement the +k-pool indexer, or if the published artifact's tensors do not match what it +expects, STOP and return the evidence — inventing the semantics of a learned +pooling is worse than leaving the device arm refused. If the device gate reds, +that is the RESULT and it lands as one. + ## Tests to port `tests/models/` in transformers `v5.16.1` is the upstream suite. What is @@ -4404,8 +4679,125 @@ Debts this row carries, each visible rather than waived: [#2099](https://github.com/mudler/vllm.cpp/issues/2099)'s lane-pin finding is the same shape one document over. +- **O36 — THE K-POOL DEVICE OPS LAND UNREACHED, and this entry names what is + unreached and who owns the wiring.** W9c-0 registers + `vt::OpId::kGlm5NextKpoolCompress` and `vt::OpId::kGlm5NextKpoolSelect` on + `kCUDA` and gates them on `dgx:gpu0` against the transformers v5.16.1 goldens + and the host reference. **Nothing on this model's production path calls + either.** `ModelRegistry::Forward` reaches `glm5_next_dsa.cpp`'s host + `SelectIndexerTopkFromPacked` and continues to; `glm5_next_forward.cpp:231-238` + still refuses a non-CPU queue by name; and the only callers of the two ops are + the device gate and the availability probe. This is the shape + `.agents/reachability.md` calls "the test-only driver", and it is staged + deliberately rather than disclosed after the fact. + + Three things, in the specific, as AGENTS.md §"Nothing lands dead" requires. + **What is unreached:** the `DeviceType::kCUDA` registration of both new OpIds + as reached FROM THIS MODEL, and `vllm::glm5_next::KpoolDeviceOpsAvailable()`, + which no production predicate consults yet. **The row that owns the wiring:** + `MODEL-MM-glm5-next-glm5-next-for-conditional-generation`, wave **W9c-3**, the + compose that constructs a CUDA queue for this forward and deletes the refusal + at `glm5_next_forward.cpp:231-238`. W9c-3 also owns the two operands these ops + need and the forward does not yet produce on a device: the packed indexer row + (`PackIndexerStates`, which W9c-0 deliberately did not make an op because + `vt::Matmul` and `vt::LayerNorm` already serve it) and a `k_pass` that is not + de-paged into host memory at `glm5_next_kv.cpp:450-456`. **The issue that + tracks it:** [#2410](https://github.com/mudler/vllm.cpp/issues/2410). + + The reachability mutation `.agents/reachability.md` asks for is reported as + what it is. There is no production call site to delete, so the question is + already answered, and the mutation that IS run instead deletes the + `RegisterOp` line and shows the device gate reds — which proves the gate + measures the registered provider and not a directly-called kernel, and proves + nothing about a capability. Read it that way. +- **O37 — `--device cuda` ON THE REAL ARTIFACT DOES NOT REACH THE REFUSAL THIS + ROW KEEPS CITING, and the row had never measured it.** W9c-0's lease drove + `vllm-cli --model --device cuda` on + `GLM-5.3-Flash-UD-Q2_K_XL` on `dgx:gpu0` (job + `79aa5bb5-7536-43fe-a051-ed73ac1302e1`, 2026-08-31T23:30Z, on `4034c368c`). + It loads, auto-fits the KV cache, enters the engine and dies in the KV + BINDING: `'model.layers.3.self_attn.indexer.k_cache' resolved to attn_kv index + 45 but only 22 cache(s) arrived` (`glm5_next_kv.cpp:127`), `CLI_DEVICE_CUDA=1`. + + **The guard at `glm5_next_forward.cpp:231-238` is still in the tree and never + fires**, because `ResolveAttnCache` throws before the forward is entered. So + the sentence this row repeats — "`--device cuda` still refuses by name at + `glm5_next_forward.cpp:231-238`" — describes code that exists but not any + behaviour a user can reach, and W9c-3, scoped as "the compose that deletes + that refusal", will meet the KV binding first and earlier. The 22 is W5b-2c's + own expected count (11 latents + 11 indexer caches, O29), so the name index is + producing indices from a larger space than the vector it indexes. + + **The prior probe measured nothing and is not a baseline.** It passed the + checkpoint DIRECTORY instead of a shard, died in `hf_config` in five seconds + over a GGUF artifact that carries no `config.json`, and its capture came back + empty. This is the first run that reached the engine. + + **What is NOT established: whether `--device cpu` reproduces it.** That was not + measured, this lease drove the CUDA arm only, and the auto-fit above (256 + blocks, `max_model_len` 8192, `max_num_seqs` 1) is memory-dependent and may + differ per device. O30's ` Paris.` predates several waves and a different + resolved config. A general KV-binding defect and a device-only one have + different owners, so the polarity is left open rather than guessed, and one + `--device cpu` run at the same shard and config settles it. Tracked by + [#2445](https://github.com/mudler/vllm.cpp/issues/2445); not repaired in flow + because the fix lives in the multi-KV index mapping W5b-2c owns and the + measurement above has to come first. + ## Now +`ACTIVE`, 2026-08-31. **The k-pool indexer has a device implementation, and it +is the first kernel this campaign has had to write.** W9c-0 +(`CLAIM-GLM53-FLASH-KPOOL-CUDA`, issue +[#2415](https://github.com/mudler/vllm.cpp/issues/2415)) answers the question +#2415 posed — a k-pool CUDA kernel, or eleven device-to-host round trips per +step — by building the kernel, and closes #2415 with the argument recorded in +§W9c-0. The row's lifecycle state does not move, because O1 does not, and +`--device cuda` on this model still refuses by name at +`glm5_next_forward.cpp:231-238`. + +**Two ops, CUDA only, and no CPU provider on purpose.** +`vt::Glm5NextKpoolCompress` is upstream's `get_pooled_states` +(`modular_glm5_next.py:897-970`) — the learned per-channel `index_kpool`-way +softmax over the pool's members, the `keep` compaction, and `P` published as a +device scalar so nothing synchronises. `vt::Glm5NextKpoolSelect` is the +selection half of `forward` (`:821-875`) with `get_visible_tokens` (`:877-895`) +folded in as a predicate and `append_visible_tail` (`:972-1022`) folded in as +the tail write. The CPU answer already exists as `glm5_next_dsa.cpp` and is what +the gate measures against; registering it a second time under an `OpId` would +make the seam its own oracle. + +**W9c's "no kernel needs writing for correctness" is now formally retired.** It +was falsified by W9c-1's measurement and it is discharged by this wave rather +than left as a correction. What survives it unchanged is the rest of the +rescoping: every OTHER primitive family this model needs does have a registered +CUDA provider, and the remaining device work is still a compose. + +**GATED ON `dgx:gpu0` (GB10, `sm_121a`): 4 cases, 918 assertions, 0 failed, and +ZERO skipped cases.** The device arm agrees with the transformers v5.16.1 run to +3.58e-07 on the pooled keys and 6.68e-06 on the scores, and its selection is SET- +and position-identical to both the oracle and the host reference at 0 mismatches +of 462 — against a decision margin of 2.58e-03, so the numerical difference is a +hundredth of the gap the top-k turns on. Seven of nine mutations were killed by +assertions. M8 (the registration deleted) SURVIVED AS A SKIP, which is the +reachability finding rather than a pass, and M5 was killed by the compiler rather +than an assertion and is owed a re-run; §W9c-0 records both instead of counting +them. + +**The ops are UNREACHED and O36 says so in the specific.** Nothing on the +production path calls either one; W9c-3 owns the wiring and +[#2410](https://github.com/mudler/vllm.cpp/issues/2410) tracks it. Do not read +the device gate as a capability claim — it measures two registered providers at a +small real geometry, and that is all it measures. No throughput number and no +generation from this artifact: O6 is unchanged and there is still no +denominator. + +The next actions are W9c-3 (the compose that reaches these ops), W9b +(keep-quant residency), W6 (the vision tower) and W7b (the first fitting +artifact). + +### Before W9c-0 + `ACTIVE`, 2026-08-31. **The device arm is a PORT, not a kernel campaign, and §W9c said otherwise on a premise this wave falsified.** W9c's rescoping (`CLAIM-GLM53-FLASH-CUDA`, row `MODEL-MM-GLM53-FLASH-CUDA`, issue diff --git a/CMakeLists.txt b/CMakeLists.txt index 8eeae6d9f..10dba9a69 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -823,6 +823,7 @@ add_library(vllm STATIC src/vllm/model_executor/models/glm5_next_moe.cpp src/vllm/model_executor/models/glm5_next_kv.cpp src/vllm/model_executor/models/glm5_next_registry.cpp + src/vllm/model_executor/models/glm5_next_device.cpp src/vllm/model_executor/models/glm_moe_dsa.cpp src/vllm/model_executor/models/glm_moe_dsa_forward.cpp src/vllm/model_executor/models/glm_moe_dsa_loader.cpp @@ -1823,6 +1824,7 @@ if(VLLM_CPP_CUDA) src/vt/cuda/cuda_combine_tokens.cu src/vt/cuda/cuda_sample.cu src/vt/cuda/cuda_deepseek_v4.cu + src/vt/cuda/cuda_glm5_next.cu src/vt/cuda/cuda_exl3.cu src/vt/cuda/cuda_laguna.cu src/vt/cuda/cuda_minimax_h3.cu diff --git a/include/vllm/model_executor/models/glm5_next_device.h b/include/vllm/model_executor/models/glm5_next_device.h new file mode 100644 index 000000000..042fd3f9c --- /dev/null +++ b/include/vllm/model_executor/models/glm5_next_device.h @@ -0,0 +1,29 @@ +// GLM-5.3-Flash W9c-0 — the availability probe for this model's device-only op +// family. Issue #2415, `.agents/specs/glm5-next-flash.md` section W9c-0. +// +// The k-pool DSA indexer's two ops (`vt::OpId::kGlm5NextKpoolCompress`, +// `kGlm5NextKpoolSelect`) are registered on `kCUDA` only, by +// `src/vt/cuda/cuda_glm5_next.cu`. On a CPU-only build nothing is registered for +// them, so `vt::GetOp` throws and `vt::Glm5NextKpoolCompress` surfaces a clean +// device-only error rather than linking a stub that returns a plausible wrong +// selection. That is the arrangement `deepseek_v4_device.h` sets out for the +// four V4 families, and this header is its one-family mirror. +// +// The probe exists so a forward can decide BEFORE it builds operands rather +// than after it throws — the shape `deepseek_v4_device.cpp:30-35` uses. **No +// production path consults it yet**: W9c-3 owns the compose that constructs a +// CUDA queue for this model and deletes the refusal at +// `glm5_next_forward.cpp:231-238`, and the row's spec records the debt as O36. +#ifndef VLLM_MODEL_EXECUTOR_MODELS_GLM5_NEXT_DEVICE_H_ +#define VLLM_MODEL_EXECUTOR_MODELS_GLM5_NEXT_DEVICE_H_ + +namespace vllm::glm5_next { + +// True iff BOTH k-pool ops have a CUDA provider. Both, because the family is +// only useful as a pair: the compress op publishes the compacted pool count the +// select op reads, so half a family is not a usable half of the capability. +bool KpoolDeviceOpsAvailable(); + +} // namespace vllm::glm5_next + +#endif // VLLM_MODEL_EXECUTOR_MODELS_GLM5_NEXT_DEVICE_H_ diff --git a/include/vt/ops.h b/include/vt/ops.h index ce20e66a6..aa20a3b4d 100644 --- a/include/vt/ops.h +++ b/include/vt/ops.h @@ -726,6 +726,68 @@ enum class OpId : uint8_t { // those BY NAME. // Appended before kCount so no existing op's id shifts. kQwen4ExpPleGate, + // MODEL-MM-GLM53-FLASH W9c-0 ([#2415]) — GLM-5.3-Flash's K-POOL DSA indexer, + // the one primitive family on this model's critical path that nothing in this + // tree implemented on a device. + // + // THIS IS NOT THE LIGHTNING INDEXER AND THE PAIR ABOVE CANNOT SERVE IT. + // `kDsaIndexerLogits` / `kDsaTopkSelect` score RAW TOKENS and pick the top + // `index_topk` of them. `Glm5NextTextIndexer` (transformers v5.16.1, + // `models/glm5_next/modular_glm5_next.py`) compresses `index_kpool = 4` + // consecutive VALID tokens into one candidate under a LEARNED per-channel + // 4-way softmax, picks the top `index_topk / index_kpool = 512` POOLS, expands + // each back into its member token indices, and then appends the ragged visible + // tail raw and UNSCORED. Feeding raw candidates into a pooled top-k, or pooled + // candidates into a consumer expecting raw ones, yields plausible indices + // either way — which is why this is a separate op family and not a mode. + // + // Composition out of what already exists was checked and is not available: + // this `OpId` inventory has no general softmax, no general reduction and no + // pooled gather, and its only top-k entries are `kMoeRouterTopK` and the + // sampler's `kTopKValuesIndices`, neither of which selects over pools or + // expands a pool back to its members. + // + // kGlm5NextKpoolCompress — `get_pooled_states` (`:897-970`). The pool grid + // starts at the first VALID token and not at slot 0, so a left-padded row + // groups differently; a pool is a candidate only when ALL of its members are + // valid; and `keep = pool_valid.any(0)` (`:968`) COMPACTS the pool axis. That + // compaction is not cosmetic: `select_k = min(index_topk // index_kpool, P)` + // reads the compacted width, and a `P` too large by one moves the ragged + // tail's write offset and changes what the final truncation keeps. So the op + // compacts ON THE DEVICE and publishes `P` as a `[1]` i32 DEVICE scalar. It + // never synchronises and never returns anything to the host. + // + // kGlm5NextKpoolSelect — the selection half of `forward` (`:821-875`), with + // `get_visible_tokens` (`:877-895`) folded in as a PREDICATE rather than + // materialised (a `[B, S, kv_len]` visibility tensor is 2.7 GiB at 32k + // context for something a thread evaluates in two instructions; upstream + // materialises it because torch has no other way to gather under it) and + // `append_visible_tail` (`:972-1022`) folded in as the tail write, because the + // tail's write offset is `select_k * index_kpool` and only this op knows `P`. + // Output width is `index_topk + index_kpool - 1` = **2051** on the published + // checkpoint, not 2048, and it carries `-1` sentinels and duplicates, which + // upstream absorbs downstream with `scatter_add_` + `ne(0)`. + // + // f32 THROUGHOUT, which is upstream's own arithmetic and not a choice: + // `:823` scores in fp32 and `:960-964` takes the pool softmax in fp32. The + // HOST reference `glm5_next_dsa.cpp` accumulates in `double`, a + // host-reference widening the device deliberately does not copy — a fp64 pool + // softmax would put the model path on the 1/64-rate pipe to be more precise + // than the thing it mirrors. + // + // Registered on kCUDA ONLY (src/vt/cuda/cuda_glm5_next.cu). There is + // deliberately NO CPU provider: the CPU answer is `glm5_next_dsa.cpp`, which + // is this family's ORACLE, and registering it a second time under these ids + // would make the seam its own oracle. A CPU queue is therefore refused BY NAME + // by the dispatcher. + // + // Additive, and NOT REACHED from any production entry point yet: only + // `tests/vllm/models/test_glm5_next_kpool_device.cpp` and the availability + // probe `vllm::glm5_next::KpoolDeviceOpsAvailable()` call them. W9c-3 owns the + // wiring and #2410 tracks it; the spec's O36 says so in the specific. + // Appended before kCount so no existing op's id shifts. + kGlm5NextKpoolCompress, + kGlm5NextKpoolSelect, kCount }; @@ -2294,6 +2356,37 @@ using Qwen4ExpQsaGatherAttentionFn = void (*)(Queue&, Tensor& /*out*/, const Tensor& /*block_ids*/, const Tensor& /*kv_lens*/, const Qwen4ExpQsaAttnArgs&); +// GLM-5.3-Flash k-pool DSA indexer (MODEL-MM-GLM53-FLASH W9c-0, #2415). The +// pooled-candidate selection `vt::DsaIndexerLogits` / `vt::DsaTopkSelect` +// cannot express, because those two score RAW TOKENS. See the OpId comments. +struct Glm5NextKpoolSelectArgs { + // `config.index_topk` — 2048 on the published checkpoint. The pool budget is + // `index_topk // index_kpool` and `validate_architecture` enforces that the + // division is exact (`configuration_glm5_next.py:219-220`). + int64_t index_topk = 0; + // `cache_layer.get_seq_length()` (`modular_glm5_next.py:811`). It differs from + // `kv_len` only on a STATIC cache padded to a maximum length; the query at + // step `s` sits at `current_length - seq_len + s`, which is what lets a decode + // step's single query see the whole cached prefix instead of only slot 0. + int64_t current_length = 0; + // `self.softmax_scale = self.head_dim ** -0.5` (`modeling_glm5_next.py:765`). + // The INDEXER head dim, not the MLA one. + float softmax_scale = 0.0f; + // `config.index_kpool_always_select_tail`. True on the published checkpoint, + // and it is what widens the output to `index_topk + index_kpool - 1`. + bool always_select_tail = true; +}; +using Glm5NextKpoolCompressFn = void (*)(Queue&, Tensor& /*pool_keys*/, + Tensor& /*pool_indices*/, Tensor& /*pool_valid*/, + Tensor& /*num_pools*/, const Tensor& /*packed*/, + const Tensor& /*ape*/); +using Glm5NextKpoolSelectFn = + void (*)(Queue&, Tensor& /*topk_indices*/, Tensor& /*index_scores*/, + const Tensor& /*q_states*/, const Tensor& /*head_weights*/, + const Tensor& /*pool_keys*/, const Tensor& /*pool_indices*/, + const Tensor& /*pool_valid*/, const Tensor& /*num_pools*/, + const Tensor& /*valid_keys*/, const Tensor& /*q_mask*/, + const Glm5NextKpoolSelectArgs&); using GdnStateGatherFn = void (*)(Queue&, Tensor&, const Tensor&, const Tensor&, const Tensor*); using GdnStateScatterFn = @@ -4859,6 +4952,56 @@ void DsaIndexerLogits(Queue& q, Tensor& logits, const Tensor& q_states, const Te void DsaTopkSelect(Queue& q, Tensor& indices, Tensor& counts, const Tensor& logits, const Tensor& win_start, const Tensor& win_end); +// GLM-5.3-Flash's k-pool DSA indexer — the LEARNED pooled compression +// (`Glm5NextTextIndexer.get_pooled_states`, `modular_glm5_next.py:897-970` @ +// transformers v5.16.1) and the pooled selection that consumes it +// (`forward` `:821-875`, with `get_visible_tokens` `:877-895` folded in as a +// predicate and `append_visible_tail` `:972-1022` as the tail write). vLLM +// implements `glm5_next` at no revision, so transformers is the reference here +// under AGENTS.md "When vLLM has no implementation". CUDA only. +// +// packed f32 [batch, kv_len, 2 * head_dim + 1] — `concat[k, gate, valid]` +// (`:798-801`), the FULL key history the pool grid is re-formed +// over on every call +// ape f32 [index_kpool, head_dim] — `index_kpool_compress_ape` +// pool_keys f32 [batch, np, head_dim] np = ceil(kv_len / index_kpool) +// pool_indices i32 [batch, np, index_kpool] -1 for an invalid member +// pool_valid i32 [batch, np] 1 iff ALL members are valid +// num_pools i32 [1] `P` AFTER `keep` (`:968-970`) +// +// `np` is the STATIC upper bound and `P <= np` is the live width. Only +// `[0, P)` carries meaning; the slack is zeroed (and `pool_indices` filled with +// the -1 sentinel) so a downstream read of it is empty rather than undefined. +// `num_pools` stays on the DEVICE: the whole point of the family is that the +// eleven DSA layers stop paying a device-to-host round trip per step. +void Glm5NextKpoolCompress(Queue& q, Tensor& pool_keys, Tensor& pool_indices, + Tensor& pool_valid, Tensor& num_pools, const Tensor& packed, + const Tensor& ape); + +// The selection over those pooled candidates. +// +// q_states f32 [batch, seq_len, index_n_heads, head_dim] — `wq_b(q_resid)` (`:795`) +// head_weights f32 [batch, seq_len, index_n_heads] — `weights_proj(hidden)` BEFORE +// the `n_heads ** -0.5` scale, which this op applies (`:827`) +// valid_keys i32 [batch, kv_len] — the packed row's validity channel (`:814`) +// q_mask i32 [batch, seq_len] — the query-side padding mask (`:873`) +// topk_indices i32 [batch, seq_len, index_topk (+ index_kpool - 1)] — **2051** +// wide on the published checkpoint, NOT 2048, carrying -1 +// sentinels and possible duplicates, which upstream absorbs with +// `scatter_add_` + `ne(0)` (`:1119-1129`) +// index_scores f32 [batch, seq_len, np] — the per-pool score BEFORE the +// validity mask (`:828`). It is an output rather than an internal +// because it is the only way a gate can show the selection is a +// strict separation and not a coin flip: top-k error is BIMODAL, +// so a tolerance on the selected values passes a wrong set whose +// values happen to be close. +void Glm5NextKpoolSelect(Queue& q, Tensor& topk_indices, Tensor& index_scores, + const Tensor& q_states, const Tensor& head_weights, + const Tensor& pool_keys, const Tensor& pool_indices, + const Tensor& pool_valid, const Tensor& num_pools, + const Tensor& valid_keys, const Tensor& q_mask, + const Glm5NextKpoolSelectArgs& args); + // --- MLA prefill attention (MLA campaign W5) -------------------------------- // The MHA prefill half of Multi-head Latent Attention — upstream's // "Compute Friendly Approach" (mla_attention.py:66-89): `kv_b_proj` has already diff --git a/scripts/check-device-leakage.py b/scripts/check-device-leakage.py index 45ccb0afc..d08add244 100644 --- a/scripts/check-device-leakage.py +++ b/scripts/check-device-leakage.py @@ -416,6 +416,28 @@ "resolvers so they become device-parameterized lookups " "(`GetOp(op, runner.device.type)`) instead of hardcoding kCUDA."), }, + "src/vllm/model_executor/models/glm5_next_device.cpp": { + "kcuda": (2, "the GLM-5.3-Flash k-pool DSA indexer's device PROBE TU " + "(W9c-0, #2415) — the SAME shape as laguna_device.cpp below: " + "2 `OpRegistered` lookups asking whether cuda_glm5_next.cu " + "registered `kGlm5NextKpool{Compress,Select}` on kCUDA. The TU " + "is always compiled and holds NO CUDA code. There is " + "deliberately no CPU provider for either op — the CPU answer " + "is `glm5_next_dsa.cpp`, which is the ORACLE these kernels are " + "gated against, and registering it again under the same ids " + "would make the seam its own golden — so a CPU build finds " + "nothing on `(op, kCUDA)`, the probe returns false, and " + "`vt::Glm5NextKpoolCompress` refuses a CPU queue BY NAME " + "through `GetOp`. `tests/vllm/models/" + "test_glm5_next_kpool_device.cpp` asserts that refusal and " + "asserts both ops are absent on kCPU. FOLLOW-UP (deferred, " + "shares the DeepSeek-V4 and Laguna row): thread the runner " + "`DeviceType` through this resolver so it becomes " + "`OpRegistered(op, runner.device.type)` instead of hardcoding " + "kCUDA. W9c-3 (#2410) owns the forward that would supply that " + "device, and until it lands there is no runner to read one " + "from."), + }, "src/vllm/model_executor/models/laguna_device.cpp": { "kcuda": (2, "the Laguna CUDA device-forward RESOLVER TU — the SAME shape as " "deepseek_v4_device.cpp above: 1 `GetOp` + 1 `OpRegistered` " diff --git a/src/vllm/model_executor/models/glm5_next_device.cpp b/src/vllm/model_executor/models/glm5_next_device.cpp new file mode 100644 index 000000000..e4c879ec4 --- /dev/null +++ b/src/vllm/model_executor/models/glm5_next_device.cpp @@ -0,0 +1,16 @@ +// GLM-5.3-Flash W9c-0 — the OpProvider-seam probe for the k-pool indexer's +// device ops. Always compiled (CPU + CUDA); it holds NO CUDA code. It only asks +// the op table whether `src/vt/cuda/cuda_glm5_next.cu` registered the pair under +// `kCUDA`. See glm5_next_device.h. +#include "vllm/model_executor/models/glm5_next_device.h" + +#include "vt/ops.h" // OpId, OpRegistered + +namespace vllm::glm5_next { + +bool KpoolDeviceOpsAvailable() { + return vt::OpRegistered(vt::OpId::kGlm5NextKpoolCompress, vt::DeviceType::kCUDA) && + vt::OpRegistered(vt::OpId::kGlm5NextKpoolSelect, vt::DeviceType::kCUDA); +} + +} // namespace vllm::glm5_next diff --git a/src/vt/cuda/cuda_glm5_next.cu b/src/vt/cuda/cuda_glm5_next.cu new file mode 100644 index 000000000..4e1019417 --- /dev/null +++ b/src/vt/cuda/cuda_glm5_next.cu @@ -0,0 +1,571 @@ +// GLM-5.3-Flash's k-pool DSA indexer (CUDA) — W9c-0, #2415. +// +// The device sibling of the host reference +// `src/vllm/model_executor/models/glm5_next_dsa.cpp`, which carries the full +// port map and stays the ORACLE for this file. There is no CPU provider for +// these two ops on purpose: registering the oracle under the same OpId would +// make the seam its own golden. +// +// PORTED FROM, transformers **v5.16.1** (`refs/tags/v5.16.1`, +// `93c8b7b485963a10800c91f55304db6be211c2bd`), the lane pin +// `.agents/oracles/transformers.md` records for `model_type: glm5_next`, file +// `src/transformers/models/glm5_next/modular_glm5_next.py` (95,314 bytes, +// sha256 `666faa54d8ff84d1642f55192b9e9df67a4d1b3a56cd789e4b7b2fb3e7b7a815`). +// vLLM implements this architecture at NO revision — `git grep -n +// 'glm5_next\|Glm5Next' -- vllm/` exits 1 with no output at the parity pin +// `5559679229` — so under AGENTS.md "When vLLM has no implementation" +// transformers is the reference here. +// +// Glm5NextKpoolCompressCuda <- `Glm5NextTextIndexer.get_pooled_states` :897-970 +// Glm5NextKpoolSelectCuda <- `Glm5NextTextIndexer.forward` :821-875 +// `Glm5NextTextIndexer.get_visible_tokens` :877-895 +// `Glm5NextTextIndexer.append_visible_tail`:972-1022 +// +// ─── DECOMPOSITION, and why `P` never reaches the host ────────────────────── +// +// `keep = pool_valid.any(0)` (`:968`) compacts the pool axis, and +// `select_k = min(index_topk // index_kpool, P)` (`:845`) reads the COMPACTED +// width. A `P` one too large moves the ragged tail's write offset by +// `index_kpool` columns and changes what `[..., :output_width]` (`:872`) keeps, +// so the compaction has to happen and its result has to be exact. It is done in +// three launches — per-pool validity, a single-block exclusive scan over the +// `keep` predicate, then a compacted write — and `P` is published as a `[1]` +// i32 DEVICE scalar that the select kernel reads on the device. Neither op +// synchronises, and every output buffer is sized at the STATIC upper bound +// `np = ceil(kv_len / index_kpool)`. The whole point of the family is that the +// eleven DSA layers stop paying a device-to-host round trip per step. +// +// ─── DETERMINISM AND ROUNDING ─────────────────────────────────────────────── +// +// `nvcc` contracts `a * b + c` into an FMA by default, which makes the answer +// depend on the optimiser rather than on the source. Every accumulation here +// spells the rounding out with `__fadd_rn` / `__fmul_rn` / `__fdiv_rn` — the +// convention `cuda_conv1d_general.cu` sets — in the SAME order the host +// reference accumulates, so the arithmetic is bit-reproducible run to run and +// the only remaining difference from the host is that the host widens to +// `double` and upstream (`:823`, `:960-964`) and this file do not. `expf` and +// not `__expf`: the fast intrinsic is a different function, not a faster +// spelling of the same one. +// +// ─── HONEST SCOPE ─────────────────────────────────────────────────────────── +// +// This is a REFERENCE-GRADE selection, exactly as `cuda_dsa_indexer.cu`'s top-k +// says of itself: the pool selection is one block per query with a bounded +// O(select_k * P) pass rather than a sort, because its job is to give the same +// answer as the host arm on any input. A threshold-based selection is a later +// brick and it will be gated against exactly this. +#include +#include + +#include + +#include +#include +#include +#include + +#include "vt/ops.h" + +namespace vt::cuda { +namespace { + +void CheckKpool(cudaError_t err, const char* what) { + if (err != cudaSuccess) { + throw std::runtime_error(std::string("vt cuda: glm5_next kpool: ") + what + ": " + + cudaGetErrorString(err)); + } +} + +cudaStream_t AsStream(Queue& q) { return static_cast(q.handle); } + +constexpr int kMetaThreads = 256; +constexpr int kScanThreads = 1024; +constexpr int kPoolThreads = 128; +constexpr int kSelectThreads = 256; + +__device__ __forceinline__ int64_t ClampIdx(int64_t v, int64_t hi) { + return v < 0 ? 0 : (v > hi ? hi : v); +} + +__global__ void FillI32Kernel(int32_t* p, int32_t v, int64_t n) { + const int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i < n) p[i] = v; +} + +// ── (A) `first_key` and per-pool validity ──────────────────────────────────── +// `first_key` is `valid_keys.long().argmax(-1)` where any key is valid and +// `seq_len` where none is (`:938-942`). It is what makes a LEFT-PADDED row +// group differently from an unpadded one; pooling from slot 0 instead passes +// every unpadded fixture. A pool is valid iff every member is in range and +// valid (`:955-956`). +__global__ __launch_bounds__(kMetaThreads) void KpoolMetaKernel( + int32_t* __restrict__ first_key, int32_t* __restrict__ raw_valid, + const float* __restrict__ packed, int64_t kv_len, int64_t np, int kpool, + int64_t row_stride, int64_t valid_off) { + const int64_t b = blockIdx.x; + const float* prow = packed + b * kv_len * row_stride; + const int tid = static_cast(threadIdx.x); + + __shared__ int red[kMetaThreads]; + int local = static_cast(kv_len); + for (int64_t j = tid; j < kv_len; j += kMetaThreads) { + if (prow[j * row_stride + valid_off] != 0.0f) { + local = static_cast(j); + break; + } + } + red[tid] = local; + __syncthreads(); + for (int stride = kMetaThreads >> 1; stride > 0; stride >>= 1) { + if (tid < stride) red[tid] = min(red[tid], red[tid + stride]); + __syncthreads(); + } + const int fk = red[0]; + if (tid == 0) first_key[b] = fk; + __syncthreads(); + + for (int64_t p = tid; p < np; p += kMetaThreads) { + int all_valid = 1; + for (int j = 0; j < kpool; ++j) { + const int64_t idx = static_cast(fk) + p * kpool + j; + // `safe_indices = pool_indices.clamp(0, seq_len - 1)` (`:948`) makes the + // gather legal; `& (pool_indices < seq_len)` (`:955`) throws the + // out-of-range members away again. + const int64_t safe = ClampIdx(idx, kv_len - 1); + const bool ok = idx < kv_len && prow[safe * row_stride + valid_off] != 0.0f; + all_valid = all_valid && (ok ? 1 : 0); + } + raw_valid[b * np + p] = all_valid; + } +} + +// ── (B) the `keep` compaction ──────────────────────────────────────────────── +// `keep = pool_valid.any(0)` (`:968`), then `[:, keep]` (`:970`) renumbers every +// pool after a dropped one. One block, a chunked exclusive scan with a running +// base, so the destination slot of every kept pool and `P` itself are produced +// without a host round trip. +__global__ __launch_bounds__(kScanThreads) void KpoolKeepScanKernel( + int32_t* __restrict__ dst, int32_t* __restrict__ num_pools, + const int32_t* __restrict__ raw_valid, int64_t batch, int64_t np) { + __shared__ int sh[kScanThreads]; + __shared__ int base; + const int tid = static_cast(threadIdx.x); + if (tid == 0) base = 0; + __syncthreads(); + + for (int64_t chunk = 0; chunk < np; chunk += kScanThreads) { + const int64_t p = chunk + tid; + int keep = 0; + if (p < np) { + for (int64_t b = 0; b < batch; ++b) { + if (raw_valid[b * np + p] != 0) { + keep = 1; + break; + } + } + } + sh[tid] = keep; + __syncthreads(); + // Hillis-Steele inclusive scan, fixed shape, same order every launch. + for (int off = 1; off < kScanThreads; off <<= 1) { + int add = 0; + if (tid >= off) add = sh[tid - off]; + __syncthreads(); + sh[tid] += add; + __syncthreads(); + } + if (p < np) dst[p] = keep != 0 ? (base + sh[tid] - 1) : -1; + __syncthreads(); + if (tid == kScanThreads - 1) base += sh[kScanThreads - 1]; + __syncthreads(); + } + if (tid == 0) num_pools[0] = base; +} + +// ── (C) the LEARNED pool weighting ─────────────────────────────────────────── +// `:959-965`: `index_head_dim` INDEPENDENT softmaxes, one per channel, over the +// pool's `index_kpool` members, with the learned intra-pool absolute-position +// embedding added to each member's gate score. It is NOT a mean, and a mean +// passes every shape check this file could carry. +// +// `torch.nan_to_num(logits.softmax(dim=2))` (`:962-964`) is the fully-invalid +// pool: every logit is `-inf`, the softmax is NaN, and the NaN is zeroed so the +// pool contributes nothing instead of poisoning the row. Here that is the +// `isfinite(mx)` guard, and the output buffer's pre-zero supplies the zero. +__device__ __forceinline__ bool MemberValid(const float* prow, int64_t idx, int64_t kv_len, + int64_t row_stride, int64_t valid_off) { + if (idx >= kv_len) return false; + const int64_t safe = ClampIdx(idx, kv_len - 1); + return prow[safe * row_stride + valid_off] != 0.0f; +} + +__global__ __launch_bounds__(kPoolThreads) void KpoolPoolKernel( + float* __restrict__ pool_keys, int32_t* __restrict__ pool_indices, + int32_t* __restrict__ pool_valid, const float* __restrict__ packed, + const float* __restrict__ ape, const int32_t* __restrict__ first_key, + const int32_t* __restrict__ raw_valid, const int32_t* __restrict__ dst, + int64_t kv_len, int64_t np, int head_dim, int kpool, int64_t row_stride, + int64_t valid_off) { + const int64_t p = blockIdx.x; + const int64_t b = blockIdx.y; + const int slot = dst[p]; + if (slot < 0) return; // dropped by `keep` + + const float* prow = packed + b * kv_len * row_stride; + const int64_t base = static_cast(first_key[b]) + p * kpool; + + if (threadIdx.x == 0) { + pool_valid[b * np + slot] = raw_valid[b * np + p]; + // `pool_indices.masked_fill(~grouped_valid_keys, -1)` (`:957`). The -1 is + // written BEFORE `pool_end` is read, so an invalid LAST member makes the + // pool's visibility probe read slot 0 after the clamp at `:831`. + for (int j = 0; j < kpool; ++j) { + const int64_t idx = base + j; + pool_indices[(b * np + slot) * kpool + j] = + MemberValid(prow, idx, kv_len, row_stride, valid_off) ? static_cast(idx) + : -1; + } + } + + // Three passes over the pool's members rather than three register arrays: the + // member count is a RUNTIME value, so an array indexed by the loop variable + // spills to local memory, and re-reading `index_kpool <= 16` gate scores is + // cheaper than that spill. Every pass recomputes the same logit in the same + // order, so the three agree by construction. + for (int c = static_cast(threadIdx.x); c < head_dim; c += kPoolThreads) { + float mx = -CUDART_INF_F; + for (int j = 0; j < kpool; ++j) { + const int64_t idx = base + j; + if (!MemberValid(prow, idx, kv_len, row_stride, valid_off)) continue; + const int64_t safe = ClampIdx(idx, kv_len - 1); + const float lg = + __fadd_rn(prow[safe * row_stride + head_dim + c], ape[j * head_dim + c]); + mx = fmaxf(mx, lg); + } + // `torch.nan_to_num(logits.softmax(dim=2))` (`:962-964`): a pool with NO + // valid member softmaxes to NaN and is then zeroed, so it contributes + // nothing rather than poisoning the row. The pre-zeroed buffer is the zero. + if (!isfinite(mx)) continue; + + float denom = 0.0f; + for (int j = 0; j < kpool; ++j) { + const int64_t idx = base + j; + // An invalid member's probability is exactly 0 and adding it changes no + // f32 sum, which is why the host reference can add it unconditionally. + if (!MemberValid(prow, idx, kv_len, row_stride, valid_off)) continue; + const int64_t safe = ClampIdx(idx, kv_len - 1); + const float lg = + __fadd_rn(prow[safe * row_stride + head_dim + c], ape[j * head_dim + c]); + denom = __fadd_rn(denom, expf(__fadd_rn(lg, -mx))); + } + + float acc = 0.0f; + for (int j = 0; j < kpool; ++j) { + const int64_t idx = base + j; + if (!MemberValid(prow, idx, kv_len, row_stride, valid_off)) continue; + const int64_t safe = ClampIdx(idx, kv_len - 1); + const float lg = + __fadd_rn(prow[safe * row_stride + head_dim + c], ape[j * head_dim + c]); + const float e = expf(__fadd_rn(lg, -mx)); + if (e == 0.0f) continue; // underflowed: the host reference skips it too + acc = __fadd_rn(acc, __fmul_rn(__fdiv_rn(e, denom), prow[safe * row_stride + c])); + } + pool_keys[(b * np + slot) * head_dim + c] = acc; + } +} + +// ── (D) score, select, expand, append the tail ─────────────────────────────── +// One block per (batch row, query). `:823-828` is the score: a per-head dot, +// then the ReLU, THEN the head mix — the ReLU is before the mix, so a head that +// dislikes a pool contributes zero rather than a negative vote. `:831-837` is +// the candidacy: a pool is selectable only if its LAST member is visible to the +// query and all its members are valid. `:839-842` masks a non-candidate with +// `torch.finfo(dtype).min` and NOT with `-inf`, which is observable rather than +// cosmetic: a query whose candidates are all invalid still gets finite scores +// and a well-defined (then discarded) top-k. +// +// The selection reproduces `torch.topk`'s CPU rule — larger score wins, an exact +// tie breaks toward the SMALLER pool index — and emits in the DESCENDING order +// `topk` returns, because the emission order decides which member indices land +// in which columns and the tail is written at `select_k * index_kpool`. +__global__ __launch_bounds__(kSelectThreads) void KpoolSelectKernel( + int32_t* __restrict__ topk_indices, float* __restrict__ index_scores, + const float* __restrict__ q_states, const float* __restrict__ head_weights, + const float* __restrict__ pool_keys, const int32_t* __restrict__ pool_indices, + const int32_t* __restrict__ pool_valid, const int32_t* __restrict__ num_pools, + const int32_t* __restrict__ valid_keys, const int32_t* __restrict__ q_mask, + int64_t np, int64_t kv_len, int64_t seq_len, int n_heads, int head_dim, int kpool, + int64_t index_topk, int64_t out_w, int64_t current_length, float softmax_scale, + float head_scale, int always_tail) { + const int64_t s = blockIdx.x; + const int64_t b = blockIdx.y; + const int tid = static_cast(threadIdx.x); + const int64_t r = b * seq_len + s; + int32_t* dst = topk_indices + r * out_w; + + // `F.pad(..., value=-1)` (`:870`) as the initial state, so every column this + // kernel does not write is already the invalid sentinel. + for (int64_t i = tid; i < out_w; i += kSelectThreads) dst[i] = -1; + __syncthreads(); + + const int P = num_pools[0]; + const int64_t q_pos = current_length - seq_len + s; + + __shared__ float best_val[kSelectThreads]; + __shared__ int best_idx[kSelectThreads]; + __shared__ int red[kSelectThreads]; + __shared__ float pick_val; + __shared__ int pick_idx; + + const float* qrow = q_states + r * n_heads * head_dim; + const float* wrow = head_weights + r * n_heads; + + // The score is computed for EVERY query row, including a padded one. Upstream + // masks `topk_indices` at `:873` and never masks `index_scores` at `:828`, so + // returning early on the padding mask here would leave a padded row's scores + // zero and disagree with the reference on exactly the rows a left-padded + // fixture adds. + for (int p = tid; p < P; p += kSelectThreads) { + float acc = 0.0f; + const float* pk = pool_keys + (b * np + p) * head_dim; + for (int h = 0; h < n_heads; ++h) { + const float* qh = qrow + static_cast(h) * head_dim; + float dot = 0.0f; + for (int c = 0; c < head_dim; ++c) dot = __fadd_rn(dot, __fmul_rn(qh[c], pk[c])); + const float relu = fmaxf(0.0f, __fmul_rn(dot, softmax_scale)); + acc = __fadd_rn(acc, __fmul_rn(__fmul_rn(wrow[h], head_scale), relu)); + } + index_scores[r * np + p] = acc; + } + __syncthreads(); + + // `topk_indices.masked_fill(~attention_mask[..., None], -1)` (`:873`): a padded + // query selects nothing and keeps the -1 prefill. The whole block shares one + // query, so this return is uniform. + if (q_mask[r] == 0) return; + + if (tid == 0) { + pick_val = CUDART_INF_F; + pick_idx = -1; + } + __syncthreads(); + + const int select_k = min(static_cast(index_topk / kpool), P); + for (int slot = 0; slot < select_k; ++slot) { + const float bar_v = pick_val; + const int bar_i = pick_idx; + float bv = -CUDART_INF_F; + int bi = -1; + for (int p = tid; p < P; p += kSelectThreads) { + // Candidacy is recomputed rather than cached: it is two loads, and a + // `[B, S, P]` candidate buffer is the allocation this family exists to + // avoid. + const int last = pool_indices[(b * np + p) * kpool + kpool - 1]; + const int64_t safe = ClampIdx(last, kv_len - 1); + const bool vis = safe <= q_pos && valid_keys[b * kv_len + safe] != 0; + const bool cand = vis && pool_valid[b * np + p] != 0; + const float v = cand ? index_scores[r * np + p] : -FLT_MAX; + // Strictly after the previous pick in (score desc, index asc) order. + const bool avail = bar_i < 0 || v < bar_v || (v == bar_v && p > bar_i); + if (!avail) continue; + if (bi < 0 || v > bv || (v == bv && p < bi)) { + bv = v; + bi = p; + } + } + best_val[tid] = bv; + best_idx[tid] = bi; + __syncthreads(); + for (int stride = kSelectThreads >> 1; stride > 0; stride >>= 1) { + if (tid < stride) { + const float ov = best_val[tid + stride]; + const int oi = best_idx[tid + stride]; + const bool better = oi >= 0 && (best_idx[tid] < 0 || ov > best_val[tid] || + (ov == best_val[tid] && oi < best_idx[tid])); + if (better) { + best_val[tid] = ov; + best_idx[tid] = oi; + } + } + __syncthreads(); + } + if (tid == 0) { + pick_val = best_val[0]; + pick_idx = best_idx[0]; + const int p = pick_idx; + bool keep = false; + if (p >= 0) { + const int last = pool_indices[(b * np + p) * kpool + kpool - 1]; + const int64_t safe = ClampIdx(last, kv_len - 1); + const bool vis = safe <= q_pos && valid_keys[b * kv_len + safe] != 0; + keep = vis && pool_valid[b * np + p] != 0; + } + // `selected_valid` masks the WHOLE expanded pool to -1 (`:853,:859-862`): + // a pool the top-k had to pick because nothing better existed does not + // become a real selection. + for (int m = 0; m < kpool; ++m) { + const int64_t col = static_cast(slot) * kpool + m; + if (col >= out_w) break; + dst[col] = keep ? pool_indices[(b * np + p) * kpool + m] : -1; + } + } + __syncthreads(); + } + + // `append_visible_tail` (`:972-1022`). `max_tail_width == 0` at + // `index_kpool == 1` returns unchanged (`:985-986`). + if (always_tail == 0 || kpool <= 1) return; + + int local_count = 0; + int local_first = static_cast(kv_len); + for (int64_t j = tid; j < kv_len; j += kSelectThreads) { + const bool valid = valid_keys[b * kv_len + j] != 0; + if (valid && j <= q_pos) ++local_count; + if (valid && j < local_first) local_first = static_cast(j); + } + red[tid] = local_count; + __syncthreads(); + for (int stride = kSelectThreads >> 1; stride > 0; stride >>= 1) { + if (tid < stride) red[tid] += red[tid + stride]; + __syncthreads(); + } + const int visible_count = red[0]; + __syncthreads(); + red[tid] = local_first; + __syncthreads(); + for (int stride = kSelectThreads >> 1; stride > 0; stride >>= 1) { + if (tid < stride) red[tid] = min(red[tid], red[tid + stride]); + __syncthreads(); + } + const int first = red[0]; + + if (tid == 0) { + const int tail_count = visible_count % kpool; + const int64_t tail_start = static_cast(first) + visible_count - tail_count; + for (int j = 0; j < kpool - 1; ++j) { + const int64_t col = static_cast(select_k) * kpool + j; + if (col >= out_w) break; + const int64_t idx = tail_start + j; + // `tail_valid` drops the fill positions and anything past the cache + // (`:1013`); `tail_visible` re-checks the padding mask (`:1016-1017`). + const bool valid = j < tail_count && idx < kv_len; + const int64_t safe = ClampIdx(idx, kv_len - 1); + const bool vis = safe <= q_pos && valid_keys[b * kv_len + safe] != 0; + dst[col] = (valid && vis) ? static_cast(idx) : -1; + } + } +} + +// ── the launchers ──────────────────────────────────────────────────────────── + +void Glm5NextKpoolCompressCuda(Queue& q, Tensor& pool_keys, Tensor& pool_indices, + Tensor& pool_valid, Tensor& num_pools, const Tensor& packed, + const Tensor& ape) { + cudaStream_t s = AsStream(q); + const int64_t batch = packed.shape[0]; + const int64_t kv_len = packed.shape[1]; + const int64_t row_stride = packed.shape[2]; + const int kpool = static_cast(ape.shape[0]); + const int head_dim = static_cast(ape.shape[1]); + const int64_t np = pool_keys.shape[1]; + + // Deterministic slack. Only `[0, P)` carries meaning, and leaving the rest + // uninitialised would make a downstream read of it undefined instead of + // obviously empty. + CheckKpool(cudaMemsetAsync(pool_keys.data, 0, + static_cast(batch * np * head_dim) * sizeof(float), s), + "memset pool_keys"); + CheckKpool(cudaMemsetAsync(pool_valid.data, 0, + static_cast(batch * np) * sizeof(int32_t), s), + "memset pool_valid"); + { + const int64_t n = batch * np * kpool; + if (n > 0) { + const unsigned grid = static_cast((n + 255) / 256); + FillI32Kernel<<>>(pool_indices.Ptr(), -1, n); + CheckKpool(cudaGetLastError(), "fill pool_indices launch"); + } + } + + // Per-call scratch on the stream, the shape `cuda_mla_prefill.cu:125` uses. A + // pooled workspace is a later brick; this family is not on a hot path yet + // because nothing calls it (spec O36). + void* scratch = nullptr; + const size_t scratch_bytes = + (static_cast(batch) + static_cast(batch * np) + + static_cast(np)) * + sizeof(int32_t); + CheckKpool(cudaMallocAsync(&scratch, scratch_bytes, s), "cudaMallocAsync kpool scratch"); + int32_t* first_key = static_cast(scratch); + int32_t* raw_valid = first_key + batch; + int32_t* dst = raw_valid + batch * np; + + KpoolMetaKernel<<(batch), kMetaThreads, 0, s>>>( + first_key, raw_valid, packed.Ptr(), kv_len, np, kpool, row_stride, + static_cast(2 * head_dim)); + CheckKpool(cudaGetLastError(), "kpool meta launch"); + + KpoolKeepScanKernel<<<1, kScanThreads, 0, s>>>(dst, num_pools.Ptr(), raw_valid, + batch, np); + CheckKpool(cudaGetLastError(), "kpool keep-scan launch"); + + const dim3 pool_grid(static_cast(np), static_cast(batch)); + KpoolPoolKernel<<>>( + pool_keys.Ptr(), pool_indices.Ptr(), pool_valid.Ptr(), + packed.Ptr(), ape.Ptr(), first_key, raw_valid, dst, kv_len, np, + head_dim, kpool, row_stride, static_cast(2 * head_dim)); + CheckKpool(cudaGetLastError(), "kpool pool launch"); + + CheckKpool(cudaFreeAsync(scratch, s), "cudaFreeAsync kpool scratch"); +} + +void Glm5NextKpoolSelectCuda(Queue& q, Tensor& topk_indices, Tensor& index_scores, + const Tensor& q_states, const Tensor& head_weights, + const Tensor& pool_keys, const Tensor& pool_indices, + const Tensor& pool_valid, const Tensor& num_pools, + const Tensor& valid_keys, const Tensor& q_mask, + const Glm5NextKpoolSelectArgs& args) { + cudaStream_t s = AsStream(q); + const int64_t batch = q_states.shape[0]; + const int64_t seq_len = q_states.shape[1]; + const int n_heads = static_cast(q_states.shape[2]); + const int head_dim = static_cast(q_states.shape[3]); + const int64_t np = pool_keys.shape[1]; + const int kpool = static_cast(pool_indices.shape[2]); + const int64_t kv_len = valid_keys.shape[1]; + const int64_t out_w = topk_indices.shape[2]; + + CheckKpool(cudaMemsetAsync(index_scores.data, 0, + static_cast(batch * seq_len * np) * sizeof(float), s), + "memset index_scores"); + + // `weights = self.weights_proj(hidden) * (self.n_heads ** -0.5)` (`:827`). + // Computed on the HOST in double and narrowed once, which is exactly what the + // reference does (`glm5_next_dsa.cpp:438`). `rsqrtf` is an approximation and + // would put a device-only rounding into a constant both arms share. + const float head_scale = + static_cast(1.0 / std::sqrt(static_cast(n_heads))); + const dim3 grid(static_cast(seq_len), static_cast(batch)); + KpoolSelectKernel<<>>( + topk_indices.Ptr(), index_scores.Ptr(), q_states.Ptr(), + head_weights.Ptr(), pool_keys.Ptr(), pool_indices.Ptr(), + pool_valid.Ptr(), num_pools.Ptr(), valid_keys.Ptr(), + q_mask.Ptr(), np, kv_len, seq_len, n_heads, head_dim, kpool, + args.index_topk, out_w, args.current_length, args.softmax_scale, head_scale, + args.always_select_tail ? 1 : 0); + CheckKpool(cudaGetLastError(), "kpool select launch"); +} + +struct Registrar { + Registrar() { + RegisterOp(OpId::kGlm5NextKpoolCompress, DeviceType::kCUDA, + reinterpret_cast( + static_cast(&Glm5NextKpoolCompressCuda))); + RegisterOp(OpId::kGlm5NextKpoolSelect, DeviceType::kCUDA, + reinterpret_cast( + static_cast(&Glm5NextKpoolSelectCuda))); + } +} registrar; + +} // namespace +} // namespace vt::cuda diff --git a/src/vt/op_provider.cpp b/src/vt/op_provider.cpp index 0b156edc8..77a0e6f01 100644 --- a/src/vt/op_provider.cpp +++ b/src/vt/op_provider.cpp @@ -539,6 +539,10 @@ const char* OpNameImpl(OpId op) { return "Qwen4ExpGatedResidualWriteBack"; case OpId::kQwen4ExpQsaCompress: return "Qwen4ExpQsaCompress"; + case OpId::kGlm5NextKpoolCompress: + return "Glm5NextKpoolCompress"; + case OpId::kGlm5NextKpoolSelect: + return "Glm5NextKpoolSelect"; case OpId::kQwen4ExpQsaGatherAttention: return "Qwen4ExpQsaGatherAttention"; case OpId::kConv2d: diff --git a/src/vt/ops.cpp b/src/vt/ops.cpp index acf7047a5..d26a653b5 100644 --- a/src/vt/ops.cpp +++ b/src/vt/ops.cpp @@ -4338,6 +4338,166 @@ void DsaTopkSelect(Queue& q, Tensor& indices, Tensor& counts, const Tensor& logi q, indices, counts, logits, win_start, win_end); } +void Glm5NextKpoolCompress(Queue& q, Tensor& pool_keys, Tensor& pool_indices, + Tensor& pool_valid, Tensor& num_pools, const Tensor& packed, + const Tensor& ape) { + VT_CHECK(packed.rank == 3, + "glm5_next_kpool_compress: packed must be rank-3 [batch, kv_len, " + "2 * index_head_dim + 1] — the indexer cache row is " + "`concat[k, gate_scores, valid]` (modular_glm5_next.py:798-801 @ " + "transformers v5.16.1), not the 128-wide key DeepSeek-V4 caches"); + VT_CHECK(ape.rank == 2, + "glm5_next_kpool_compress: ape must be rank-2 [index_kpool, index_head_dim]"); + VT_CHECK(pool_keys.rank == 3 && pool_indices.rank == 3 && pool_valid.rank == 2 && + num_pools.rank == 1, + "glm5_next_kpool_compress: pool_keys/pool_indices must be rank-3, pool_valid " + "rank-2, num_pools rank-1"); + const int64_t batch = packed.shape[0]; + const int64_t kv_len = packed.shape[1]; + const int64_t kpool = ape.shape[0]; + const int64_t head_dim = ape.shape[1]; + VT_CHECK(batch > 0 && kv_len > 0 && head_dim > 0, + "glm5_next_kpool_compress: batch/kv_len/index_head_dim must be > 0"); + VT_CHECK(kpool >= 1, + "glm5_next_kpool_compress: index_kpool must be >= 1 " + "(configuration_glm5_next.py:216-217); it is 4 on the published checkpoint " + "and 16 in the config class, so a defaulted value is wrong by a factor of " + "four"); + VT_CHECK(packed.shape[2] == 2 * head_dim + 1, + "glm5_next_kpool_compress: packed's row must be 2 * index_head_dim + 1 wide"); + const int64_t np = (kv_len + kpool - 1) / kpool; + VT_CHECK(pool_keys.shape[0] == batch && pool_keys.shape[1] == np && + pool_keys.shape[2] == head_dim, + "glm5_next_kpool_compress: pool_keys must be [batch, ceil(kv_len / " + "index_kpool), index_head_dim] — the STATIC upper bound, because the live " + "width P is only known after the `keep` compaction (:968-970)"); + VT_CHECK(pool_indices.shape[0] == batch && pool_indices.shape[1] == np && + pool_indices.shape[2] == kpool, + "glm5_next_kpool_compress: pool_indices must be [batch, np, index_kpool]"); + VT_CHECK(pool_valid.shape[0] == batch && pool_valid.shape[1] == np, + "glm5_next_kpool_compress: pool_valid must be [batch, np]"); + VT_CHECK(num_pools.shape[0] == 1, + "glm5_next_kpool_compress: num_pools must be a [1] DEVICE scalar — reading P " + "back to the host is the round trip this op family exists to remove"); + VT_CHECK(packed.dtype == DType::kF32 && ape.dtype == DType::kF32 && + pool_keys.dtype == DType::kF32, + "glm5_next_kpool_compress: packed/ape/pool_keys must be f32 — upstream scores " + "and pools in fp32 (:823, :960-964)"); + VT_CHECK(pool_indices.dtype == DType::kI32 && pool_valid.dtype == DType::kI32 && + num_pools.dtype == DType::kI32, + "glm5_next_kpool_compress: pool_indices/pool_valid/num_pools must be i32"); + VT_CHECK(packed.IsContiguous() && ape.IsContiguous() && pool_keys.IsContiguous() && + pool_indices.IsContiguous() && pool_valid.IsContiguous() && + num_pools.IsContiguous(), + "glm5_next_kpool_compress: every operand must be contiguous"); + VT_CHECK(packed.device == q.device && ape.device == q.device && + pool_keys.device == q.device && pool_indices.device == q.device && + pool_valid.device == q.device && num_pools.device == q.device, + "glm5_next_kpool_compress: device mismatch (operands/queue)"); + reinterpret_cast( + GetOp(OpId::kGlm5NextKpoolCompress, q.device.type))(q, pool_keys, pool_indices, + pool_valid, num_pools, packed, + ape); +} + +void Glm5NextKpoolSelect(Queue& q, Tensor& topk_indices, Tensor& index_scores, + const Tensor& q_states, const Tensor& head_weights, + const Tensor& pool_keys, const Tensor& pool_indices, + const Tensor& pool_valid, const Tensor& num_pools, + const Tensor& valid_keys, const Tensor& q_mask, + const Glm5NextKpoolSelectArgs& args) { + VT_CHECK(q_states.rank == 4, + "glm5_next_kpool_select: q_states must be rank-4 [batch, seq_len, " + "index_n_heads, index_head_dim]"); + VT_CHECK(head_weights.rank == 3, + "glm5_next_kpool_select: head_weights must be rank-3 [batch, seq_len, " + "index_n_heads] — the weights_proj output BEFORE the n_heads ** -0.5 scale, " + "which this op applies (:827)"); + VT_CHECK(pool_keys.rank == 3 && pool_indices.rank == 3 && pool_valid.rank == 2, + "glm5_next_kpool_select: the pooled candidate set must come from " + "vt::Glm5NextKpoolCompress unchanged"); + VT_CHECK(num_pools.rank == 1 && num_pools.shape[0] == 1, + "glm5_next_kpool_select: num_pools must be the [1] DEVICE scalar the compress " + "op published"); + VT_CHECK(valid_keys.rank == 2 && q_mask.rank == 2, + "glm5_next_kpool_select: valid_keys must be [batch, kv_len] and q_mask " + "[batch, seq_len]"); + VT_CHECK(topk_indices.rank == 3 && index_scores.rank == 3, + "glm5_next_kpool_select: topk_indices/index_scores must be rank-3"); + const int64_t batch = q_states.shape[0]; + const int64_t seq_len = q_states.shape[1]; + const int64_t n_heads = q_states.shape[2]; + const int64_t head_dim = q_states.shape[3]; + const int64_t np = pool_keys.shape[1]; + const int64_t kpool = pool_indices.shape[2]; + const int64_t kv_len = valid_keys.shape[1]; + VT_CHECK(batch > 0 && seq_len > 0 && n_heads > 0 && head_dim > 0 && kv_len > 0, + "glm5_next_kpool_select: batch/seq_len/index_n_heads/index_head_dim/kv_len " + "must be > 0"); + VT_CHECK(kpool >= 1, "glm5_next_kpool_select: index_kpool must be >= 1"); + VT_CHECK(args.index_topk > 0, "glm5_next_kpool_select: index_topk must be > 0"); + VT_CHECK(args.index_topk % kpool == 0, + "glm5_next_kpool_select: index_topk must be divisible by index_kpool — the " + "pool budget `index_topk // index_kpool` is exact upstream " + "(configuration_glm5_next.py:219-220)"); + VT_CHECK(kv_len >= seq_len, + "glm5_next_kpool_select: kv_len must be at least seq_len — the current window " + "is always part of the key history it selects over"); + VT_CHECK(args.current_length >= seq_len, + "glm5_next_kpool_select: current_length must be at least seq_len; the query at " + "step s sits at current_length - seq_len + s (:892)"); + VT_CHECK(pool_keys.shape[0] == batch && pool_keys.shape[2] == head_dim, + "glm5_next_kpool_select: pool_keys must be [batch, np, index_head_dim]"); + VT_CHECK(pool_indices.shape[0] == batch && pool_indices.shape[1] == np, + "glm5_next_kpool_select: pool_indices must be [batch, np, index_kpool]"); + VT_CHECK(pool_valid.shape[0] == batch && pool_valid.shape[1] == np, + "glm5_next_kpool_select: pool_valid must be [batch, np]"); + VT_CHECK(np == (kv_len + kpool - 1) / kpool, + "glm5_next_kpool_select: np must be ceil(kv_len / index_kpool), the same " + "static bound the compress op allocated against"); + VT_CHECK(head_weights.shape[0] == batch && head_weights.shape[1] == seq_len && + head_weights.shape[2] == n_heads, + "glm5_next_kpool_select: head_weights must be [batch, seq_len, index_n_heads]"); + VT_CHECK(q_mask.shape[0] == batch && q_mask.shape[1] == seq_len, + "glm5_next_kpool_select: q_mask must be [batch, seq_len]"); + VT_CHECK(valid_keys.shape[0] == batch, + "glm5_next_kpool_select: valid_keys must be [batch, kv_len]"); + const int64_t out_w = args.index_topk + (args.always_select_tail ? kpool - 1 : 0); + VT_CHECK(topk_indices.shape[0] == batch && topk_indices.shape[1] == seq_len && + topk_indices.shape[2] == out_w, + "glm5_next_kpool_select: topk_indices must be [batch, seq_len, index_topk + " + "index_kpool - 1] with the always-kept tail (:864-867) — 2051 and not 2048 on " + "the published checkpoint; sizing it 2048 truncates the tail the model always " + "keeps"); + VT_CHECK(index_scores.shape[0] == batch && index_scores.shape[1] == seq_len && + index_scores.shape[2] == np, + "glm5_next_kpool_select: index_scores must be [batch, seq_len, np]"); + VT_CHECK(q_states.dtype == DType::kF32 && head_weights.dtype == DType::kF32 && + pool_keys.dtype == DType::kF32 && index_scores.dtype == DType::kF32, + "glm5_next_kpool_select: q_states/head_weights/pool_keys/index_scores must be " + "f32"); + VT_CHECK(pool_indices.dtype == DType::kI32 && pool_valid.dtype == DType::kI32 && + num_pools.dtype == DType::kI32 && valid_keys.dtype == DType::kI32 && + q_mask.dtype == DType::kI32 && topk_indices.dtype == DType::kI32, + "glm5_next_kpool_select: every integer operand must be i32"); + VT_CHECK(q_states.IsContiguous() && head_weights.IsContiguous() && + pool_keys.IsContiguous() && pool_indices.IsContiguous() && + pool_valid.IsContiguous() && num_pools.IsContiguous() && + valid_keys.IsContiguous() && q_mask.IsContiguous() && + topk_indices.IsContiguous() && index_scores.IsContiguous(), + "glm5_next_kpool_select: every operand must be contiguous"); + VT_CHECK(q_states.device == q.device && head_weights.device == q.device && + pool_keys.device == q.device && pool_indices.device == q.device && + pool_valid.device == q.device && num_pools.device == q.device && + valid_keys.device == q.device && q_mask.device == q.device && + topk_indices.device == q.device && index_scores.device == q.device, + "glm5_next_kpool_select: device mismatch (operands/queue)"); + reinterpret_cast(GetOp(OpId::kGlm5NextKpoolSelect, + q.device.type))( + q, topk_indices, index_scores, q_states, head_weights, pool_keys, pool_indices, + pool_valid, num_pools, valid_keys, q_mask, args); +} + void MlaPrefillAttention(Queue& q, Tensor& out, Tensor* lse, const Tensor& query, const Tensor& key, const Tensor& value, const Tensor& cu_seqlens_q, const Tensor& cu_seqlens_k, const MlaPrefillAttentionArgs& args) { diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 2af116cc9..133a49b28 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -1106,6 +1106,19 @@ target_include_directories(test_glm5_next_dsa PRIVATE ${CMAKE_SOURCE_DIR}/src) target_include_directories(test_glm5_next_dsa PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/vllm/models/fixtures) +# GLM-5.3-Flash W9c-0 -- the k-pool indexer's two DEVICE ops (#2415). The CUDA +# arm of what `test_glm5_next_dsa` gates on the host, run at the SAME fixture +# geometry and against the SAME transformers v5.16.1 goldens, so the device +# answer is compared with the oracle directly and not only with our own C++. +# `np = ceil(21/4) = 6` against `kNumPools = 5` makes the `keep` compaction +# observable; a kernel that skipped it puts the ragged tail three columns late. +# The device cases SKIP without a CUDA backend, so this is a real gate only on a +# leased device -- an `assertions: 0` line here is a skip wearing a pass. +vllm_cpp_add_test(test_glm5_next_kpool_device vllm/models/test_glm5_next_kpool_device.cpp) +target_include_directories(test_glm5_next_kpool_device PRIVATE ${CMAKE_SOURCE_DIR}/src) +target_include_directories(test_glm5_next_kpool_device PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/vllm/models/fixtures) + # GLM-5.3-Flash W5 -- the 288 routed + 1 shared expert MoE, its grouped # `noaux_tc` router and the clamped-SwiGLU epilogue (#2223). The router is gated # at the PUBLISHED 288/top-8 and the composed block at 8 experts, because the diff --git a/tests/vllm/models/test_glm5_next_kpool_device.cpp b/tests/vllm/models/test_glm5_next_kpool_device.cpp new file mode 100644 index 000000000..3b0c70093 --- /dev/null +++ b/tests/vllm/models/test_glm5_next_kpool_device.cpp @@ -0,0 +1,626 @@ +// GLM-5.3-Flash W9c-0 gate — the k-pool DSA indexer's two DEVICE ops, against +// the transformers v5.16.1 RUN goldens and the landed host reference. +// +// Row MODEL-MM-glm5-next-glm5-next-for-conditional-generation, issue #2415, +// `.agents/specs/glm5-next-flash.md` section W9c-0. +// +// ─── WHAT IS UNDER TEST, AND AGAINST WHAT ─────────────────────────────────── +// +// `vt::Glm5NextKpoolCompress` and `vt::Glm5NextKpoolSelect` are the CUDA arm of +// the k-pool indexer. There is no CPU provider and there deliberately is not +// one: the CPU answer already exists as `glm5_next_dsa.cpp`, gated by +// `test_glm5_next_dsa` against the same goldens, and registering it a second +// time under an `OpId` would make the seam its own oracle. +// +// So this file compares the device arm against TWO independent things: +// * `glm5_next_dsa_goldens.inc` — the RUN output of the unmodified +// `Glm5NextTextIndexer` at transformers v5.16.1, captured by +// `fixtures/gen_glm5_next_dsa_goldens.py`. Nothing in it is transcribed +// from our C++, and it carries the INTERMEDIATES (`kPoolKeys`, +// `kPoolIndices`, `kPoolValid`, `kIndexScores`) as well as the result. +// * `vllm::glm5_next::GetPooledStates` / `SelectIndexerTopkFromPacked` — the +// host reference, which accumulates in `double` where the device (and +// upstream, `modular_glm5_next.py:823,960-964`) accumulate in fp32. +// +// ─── THREE PROPERTIES, EACH LOAD-BEARING ──────────────────────────────────── +// +// 1. `seq_len` 21 against `index_topk` 8. At or below `index_topk` the +// selection is the identity and the pooling is unobservable. Row 1 is +// LEFT-PADDED by three, so its pool grid starts at token 3 and a kernel +// that pools from slot 0 passes row 0 and fails row 1. And +// `np = ceil(21/4) = 6` against `kNumPools = 5`: the fixture EXERCISES the +// `keep` compaction (`modular_glm5_next.py:968-970`) rather than assuming +// it, because a kernel that skipped it would put the ragged tail at column +// 12 instead of 8 and miss the `[..., :output_width]` truncation entirely. +// +// 2. Selection error is BIMODAL, not continuous. A wrong-but-adjacent pool's +// scores can be arbitrarily close, so a tolerance on `index_scores` bounds +// nothing about which pools were chosen. Every selection assertion here is +// SET equality of the real token indices PLUS the positionwise comparison, +// and the minimum decision margin between the `select_k`-th and +// `select_k + 1`-th masked score is COMPUTED AND PRINTED rather than +// assumed adequate. +// +// 3. Every float comparison is `isfinite`-guarded on BOTH operands before it +// is made. An all-NaN forward on this row once read as a perfect match, +// because every comparison against NaN is false, and the model then emitted +// token id 0 eight times. A NaN in either arm reds this file. +// +// The device cases SKIP on a CPU-only build (no CUDA backend registered). A +// doctest `assertions: 0` line is a skip wearing a pass; the device job reads +// that line rather than the exit code. +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "glm5_next_dsa_goldens.inc" +#include "vllm/model_executor/models/glm5_next_device.h" +#include "vllm/model_executor/models/glm5_next_dsa.h" +#include "vt/backend.h" +#include "vt/ops.h" + +namespace g = glm5_next_dsa_goldens; + +using vllm::glm5_next::GetPooledStates; +using vllm::glm5_next::IndexerDims; +using vllm::glm5_next::IndexerSelection; +using vllm::glm5_next::IndexerWeights; +using vllm::glm5_next::PackIndexerStates; +using vllm::glm5_next::PooledStates; +using vllm::glm5_next::SelectIndexerTopkFromPacked; +using vt::Backend; +using vt::Device; +using vt::DeviceType; +using vt::DType; +using vt::Queue; +using vt::Tensor; + +namespace { + +Device Cpu() { return Device{DeviceType::kCPU, 0}; } +Device Gpu() { return Device{DeviceType::kCUDA, 0}; } + +Tensor Contig(void* data, DType dt, Device dev, const std::vector& shape) { + Tensor t; + t.data = data; + t.dtype = dt; + t.device = dev; + t.rank = static_cast(shape.size()); + int64_t stride = 1; + for (int i = t.rank - 1; i >= 0; --i) { + t.shape[i] = shape[static_cast(i)]; + t.stride[i] = stride; + stride *= shape[static_cast(i)]; + } + return t; +} + +bool HasCuda() { + try { + vt::GetBackend(DeviceType::kCUDA); + return vllm::glm5_next::KpoolDeviceOpsAvailable(); + } catch (const std::runtime_error&) { + return false; + } +} + +struct QueueGuard { + Backend& b; + Queue q; + explicit QueueGuard(Backend& backend) : b(backend), q(backend.CreateQueue()) {} + ~QueueGuard() { b.DestroyQueue(q); } + QueueGuard(const QueueGuard&) = delete; + QueueGuard& operator=(const QueueGuard&) = delete; +}; + +class DeviceTensor { + public: + DeviceTensor(Backend& b, Queue& q, DType dt, const std::vector& shape, + const void* host = nullptr) + : b_(b) { + int64_t numel = 1; + for (auto s : shape) numel *= s; + bytes_ = static_cast(numel) * vt::SizeOf(dt); + p_ = b_.Alloc(bytes_ == 0 ? 1 : bytes_); + if (host != nullptr) b_.Copy(q, p_, host, bytes_); + t_ = Contig(p_, dt, Gpu(), shape); + } + ~DeviceTensor() { b_.Free(p_); } + DeviceTensor(const DeviceTensor&) = delete; + DeviceTensor& operator=(const DeviceTensor&) = delete; + Tensor& tensor() { return t_; } + void Download(Queue& q, void* dst) { + b_.Copy(q, dst, p_, bytes_); + b_.Synchronize(q); + } + + private: + Backend& b_; + void* p_ = nullptr; + size_t bytes_ = 0; + Tensor t_; +}; + +// ─── the fixture, read off the goldens rather than restated ───────────────── + +IndexerDims FixtureDims() { + IndexerDims d; + d.hidden_size = g::kHidden; + d.q_lora_rank = g::kQLora; + d.n_heads = g::kNHeads; + d.head_dim = g::kHeadDim; + d.index_topk = g::kIndexTopk; + d.index_kpool = g::kIndexKpool; + d.always_select_tail = true; + return d; +} + +IndexerWeights FixtureWeights() { + IndexerWeights w; + w.wq_b = g::kWqB; + w.wk = g::kWk; + w.k_norm_weight = g::kKNormWeight; + w.k_norm_bias = g::kKNormBias; + w.weights_proj = g::kWeightsProj; + w.kpool_ape = g::kKpoolApe; + w.kpool_gate = g::kKpoolGate; + return w; +} + +std::vector MaskBytes(const int32_t* a, size_t n) { + std::vector v(n); + for (size_t i = 0; i < n; ++i) v[i] = a[i] != 0 ? 1 : 0; + return v; +} + +// `nn.Linear` on ONE row, fp32 — the same arithmetic W9c-3 will get from +// `vt::MatmulBT`. It builds the two operands the select op takes and the host +// reference builds internally, so both arms see byte-identical inputs. +std::vector LinearRows(const float* w, const std::vector& x, int64_t rows, + int64_t out_features, int64_t in_features) { + std::vector out(static_cast(rows * out_features), 0.0f); + for (int64_t r = 0; r < rows; ++r) { + for (int64_t o = 0; o < out_features; ++o) { + const float* wr = w + o * in_features; + const float* xr = x.data() + r * in_features; + double acc = 0.0; + for (int64_t i = 0; i < in_features; ++i) + acc += static_cast(wr[i]) * static_cast(xr[i]); + out[static_cast(r * out_features + o)] = static_cast(acc); + } + } + return out; +} + +// ─── the guarded comparisons ──────────────────────────────────────────────── +// +// `Finite` is checked on BOTH operands FIRST and asserted, never folded into the +// tolerance: `std::abs(nan - nan) <= tol` is false, so an unguarded CHECK on two +// NaNs reports a failure, but an unguarded `CHECK(a == b || close)` on one NaN +// silently reports nothing at all. The pattern this file uses reds on a NaN. +void RequireFinite(const std::vector& v, const char* what) { + size_t bad = 0; + for (float x : v) + if (!std::isfinite(x)) ++bad; + INFO("non-finite entries in " << what << ": " << bad << " of " << v.size()); + REQUIRE(bad == 0); +} + +double MaxAbsDiff(const std::vector& a, const std::vector& b) { + REQUIRE(a.size() == b.size()); + double worst = 0.0; + for (size_t i = 0; i < a.size(); ++i) { + REQUIRE(std::isfinite(a[i])); + REQUIRE(std::isfinite(b[i])); + worst = std::max(worst, std::abs(static_cast(a[i]) - static_cast(b[i]))); + } + return worst; +} + +// The set of REAL token indices a row selected. `-1` is the invalid sentinel and +// duplicates are possible — upstream absorbs them with `scatter_add_` + `ne(0)` +// (`modular_glm5_next.py:1119-1129`) — so a set is the model's own semantics. +// The positionwise comparison is asserted separately and is strictly stronger. +std::set SelectedSet(const int32_t* row, int64_t width) { + std::set s; + for (int64_t i = 0; i < width; ++i) + if (row[i] >= 0) s.insert(row[i]); + return s; +} + +// The decision margin of one query row: the gap between the score of the +// `select_k`-th pool and the best pool that did NOT make the cut. A top-k error +// either flips this gap or it does not, so the gate reports the smallest gap the +// fixture offers instead of asserting a tolerance that bounds nothing. +double RowMargin(const std::vector& scores, const std::vector& candidate, + int64_t num_pools, int64_t select_k) { + std::vector masked; + masked.reserve(static_cast(num_pools)); + for (int64_t p = 0; p < num_pools; ++p) { + if (candidate[static_cast(p)] == 0) continue; + masked.push_back(static_cast(scores[static_cast(p)])); + } + if (static_cast(masked.size()) <= select_k) return -1.0; // no boundary to cross + std::sort(masked.begin(), masked.end(), std::greater()); + return masked[static_cast(select_k) - 1] - masked[static_cast(select_k)]; +} + +} // namespace + +TEST_CASE("glm5_next k-pool: the device ops are registered on CUDA and absent on CPU") { + // The registration is the seam. `KpoolDeviceOpsAvailable` is the probe + // `deepseek_v4_device.cpp:30-35` sets the shape of, and W9c-3's forward is + // meant to consult it BEFORE it builds operands rather than after it throws. + // + // The first CHECK compares the probe against the two lookups it is made of, + // so it measures CONSISTENCY and not correctness — said plainly rather than + // dressed up. What it does buy is the `&&`: a probe that lost one of its two + // clauses would report the family available with half of it registered, and + // the reviewer's mutation of this file deletes a clause to show that. + const bool cuda_present = vt::OpRegistered(vt::OpId::kGlm5NextKpoolCompress, + DeviceType::kCUDA) && + vt::OpRegistered(vt::OpId::kGlm5NextKpoolSelect, + DeviceType::kCUDA); + CHECK(vllm::glm5_next::KpoolDeviceOpsAvailable() == cuda_present); + + // No CPU provider, on purpose: the CPU answer is `glm5_next_dsa.cpp` and it is + // this file's oracle, not a second registration of itself. + CHECK_FALSE(vt::OpRegistered(vt::OpId::kGlm5NextKpoolCompress, DeviceType::kCPU)); + CHECK_FALSE(vt::OpRegistered(vt::OpId::kGlm5NextKpoolSelect, DeviceType::kCPU)); + + // So a CPU queue is refused by name rather than served a wrong answer. + const IndexerDims d = FixtureDims(); + std::vector packed_dummy( + static_cast(g::kBatch * g::kSeqLen * (2 * d.head_dim + 1)), 0.0f); + std::vector ape_dummy(static_cast(d.index_kpool * d.head_dim), 0.0f); + const int64_t np = (g::kSeqLen + d.index_kpool - 1) / d.index_kpool; + std::vector keys_dummy(static_cast(g::kBatch * np * d.head_dim), 0.0f); + std::vector idx_dummy(static_cast(g::kBatch * np * d.index_kpool), -1); + std::vector valid_dummy(static_cast(g::kBatch * np), 0); + std::vector np_dummy(1, 0); + + Queue cpu_q{Cpu(), nullptr}; + Tensor t_packed = Contig(packed_dummy.data(), DType::kF32, Cpu(), + {g::kBatch, g::kSeqLen, 2 * d.head_dim + 1}); + Tensor t_ape = Contig(ape_dummy.data(), DType::kF32, Cpu(), {d.index_kpool, d.head_dim}); + Tensor t_keys = Contig(keys_dummy.data(), DType::kF32, Cpu(), {g::kBatch, np, d.head_dim}); + Tensor t_idx = + Contig(idx_dummy.data(), DType::kI32, Cpu(), {g::kBatch, np, d.index_kpool}); + Tensor t_valid = Contig(valid_dummy.data(), DType::kI32, Cpu(), {g::kBatch, np}); + Tensor t_np = Contig(np_dummy.data(), DType::kI32, Cpu(), {1}); + CHECK_THROWS( + vt::Glm5NextKpoolCompress(cpu_q, t_keys, t_idx, t_valid, t_np, t_packed, t_ape)); +} + +TEST_CASE("glm5_next k-pool device: compress agrees with the transformers goldens") { + if (!HasCuda()) { + MESSAGE("no CUDA backend: the k-pool compress device gate is SKIPPED"); + return; + } + Backend& backend = vt::GetBackend(DeviceType::kCUDA); + QueueGuard qg(backend); + + const IndexerDims d = FixtureDims(); + const IndexerWeights w = FixtureWeights(); + const std::vector hidden(g::kHiddenStates, + g::kHiddenStates + g::kBatch * g::kSeqLen * g::kHidden); + const std::vector mask = + MaskBytes(g::kMask, static_cast(g::kBatch * g::kSeqLen)); + const std::vector packed = + PackIndexerStates(d, w, hidden, mask, g::kBatch, g::kSeqLen); + RequireFinite(packed, "packed indexer states"); + + const int64_t np = (g::kSeqLen + d.index_kpool - 1) / d.index_kpool; + REQUIRE(np == 6); // and kNumPools is 5, so the `keep` compaction is exercised + + const std::vector ape(g::kKpoolApe, + g::kKpoolApe + d.index_kpool * d.head_dim); + DeviceTensor dev_packed(backend, qg.q, DType::kF32, + {g::kBatch, g::kSeqLen, 2 * d.head_dim + 1}, packed.data()); + DeviceTensor dev_ape(backend, qg.q, DType::kF32, {d.index_kpool, d.head_dim}, ape.data()); + DeviceTensor dev_keys(backend, qg.q, DType::kF32, {g::kBatch, np, d.head_dim}); + DeviceTensor dev_idx(backend, qg.q, DType::kI32, {g::kBatch, np, d.index_kpool}); + DeviceTensor dev_valid(backend, qg.q, DType::kI32, {g::kBatch, np}); + DeviceTensor dev_np(backend, qg.q, DType::kI32, {1}); + + vt::Glm5NextKpoolCompress(qg.q, dev_keys.tensor(), dev_idx.tensor(), dev_valid.tensor(), + dev_np.tensor(), dev_packed.tensor(), dev_ape.tensor()); + backend.Synchronize(qg.q); + + std::vector got_np(1, -1); + dev_np.Download(qg.q, got_np.data()); + INFO("device P = " << got_np[0] << ", golden kNumPools = " << g::kNumPools); + REQUIRE(got_np[0] == static_cast(g::kNumPools)); + + const int64_t P = g::kNumPools; + std::vector got_keys(static_cast(g::kBatch * np * d.head_dim)); + std::vector got_idx(static_cast(g::kBatch * np * d.index_kpool)); + std::vector got_valid(static_cast(g::kBatch * np)); + dev_keys.Download(qg.q, got_keys.data()); + dev_idx.Download(qg.q, got_idx.data()); + dev_valid.Download(qg.q, got_valid.data()); + + // Only the COMPACTED prefix [0, P) carries meaning; the slack up to `np` is + // allocation, and the ops contract says nothing about it. + std::vector keys_pref, golden_keys(g::kPoolKeys, g::kPoolKeys + g::kBatch * P * d.head_dim); + for (int64_t b = 0; b < g::kBatch; ++b) + for (int64_t p = 0; p < P; ++p) + for (int64_t c = 0; c < d.head_dim; ++c) + keys_pref.push_back(got_keys[static_cast((b * np + p) * d.head_dim + c)]); + RequireFinite(keys_pref, "device pool_keys"); + const double keys_delta = MaxAbsDiff(keys_pref, golden_keys); + MESSAGE("pool_keys max|device - transformers| = " << keys_delta); + CHECK(keys_delta < 2e-5); + + for (int64_t b = 0; b < g::kBatch; ++b) { + for (int64_t p = 0; p < P; ++p) { + CHECK(got_valid[static_cast(b * np + p)] == + g::kPoolValid[static_cast(b * P + p)]); + for (int64_t m = 0; m < d.index_kpool; ++m) { + CHECK(got_idx[static_cast((b * np + p) * d.index_kpool + m)] == + g::kPoolIndices[static_cast((b * P + p) * d.index_kpool + m)]); + } + } + } + + // And against the host reference, which is the arm W9c-3 replaces. Its + // `double` accumulation is a host-reference widening the device deliberately + // does not copy, so this bound is fp32-vs-fp64 reduction and nothing else. + const PooledStates host = GetPooledStates(d, w, packed, g::kBatch, g::kSeqLen); + REQUIRE(host.num_pools == P); + const double host_delta = MaxAbsDiff(keys_pref, host.pool_keys); + MESSAGE("pool_keys max|device - host reference| = " << host_delta); + CHECK(host_delta < 2e-5); +} + +TEST_CASE("glm5_next k-pool device: the selection is SET-equal and positionwise equal") { + if (!HasCuda()) { + MESSAGE("no CUDA backend: the k-pool select device gate is SKIPPED"); + return; + } + Backend& backend = vt::GetBackend(DeviceType::kCUDA); + QueueGuard qg(backend); + + const IndexerDims d = FixtureDims(); + const IndexerWeights w = FixtureWeights(); + const std::vector hidden(g::kHiddenStates, + g::kHiddenStates + g::kBatch * g::kSeqLen * g::kHidden); + const std::vector q_resid(g::kQResid, + g::kQResid + g::kBatch * g::kSeqLen * g::kQLora); + const std::vector mask = + MaskBytes(g::kMask, static_cast(g::kBatch * g::kSeqLen)); + const std::vector packed = + PackIndexerStates(d, w, hidden, mask, g::kBatch, g::kSeqLen); + const int64_t rows = g::kBatch * g::kSeqLen; + const int64_t np = (g::kSeqLen + d.index_kpool - 1) / d.index_kpool; + const int64_t P = g::kNumPools; + const int64_t out_w = d.OutputWidth(); + REQUIRE(out_w == g::kOutputWidth); + + // The two operands the op takes rather than computes: `wq_b(q_resid)` and + // `weights_proj(hidden)` (`modular_glm5_next.py:795,827`). W9c-3 gets both + // from `vt::MatmulBT`; here they are built once so both arms see the same + // bytes. + const std::vector q_states = + LinearRows(w.wq_b, q_resid, rows, d.n_heads * d.head_dim, d.q_lora_rank); + const std::vector head_w = LinearRows(w.weights_proj, hidden, rows, d.n_heads, + d.hidden_size); + RequireFinite(q_states, "wq_b(q_resid)"); + RequireFinite(head_w, "weights_proj(hidden)"); + + std::vector valid_keys(static_cast(g::kBatch * g::kSeqLen), 0); + for (int64_t i = 0; i < g::kBatch * g::kSeqLen; ++i) + valid_keys[static_cast(i)] = + packed[static_cast(i * (2 * d.head_dim + 1) + 2 * d.head_dim)] != 0.0f ? 1 : 0; + std::vector q_mask(g::kMask, g::kMask + g::kBatch * g::kSeqLen); + const std::vector ape(g::kKpoolApe, g::kKpoolApe + d.index_kpool * d.head_dim); + + DeviceTensor dev_packed(backend, qg.q, DType::kF32, + {g::kBatch, g::kSeqLen, 2 * d.head_dim + 1}, packed.data()); + DeviceTensor dev_ape(backend, qg.q, DType::kF32, {d.index_kpool, d.head_dim}, ape.data()); + DeviceTensor dev_keys(backend, qg.q, DType::kF32, {g::kBatch, np, d.head_dim}); + DeviceTensor dev_idx(backend, qg.q, DType::kI32, {g::kBatch, np, d.index_kpool}); + DeviceTensor dev_valid(backend, qg.q, DType::kI32, {g::kBatch, np}); + DeviceTensor dev_np(backend, qg.q, DType::kI32, {1}); + vt::Glm5NextKpoolCompress(qg.q, dev_keys.tensor(), dev_idx.tensor(), dev_valid.tensor(), + dev_np.tensor(), dev_packed.tensor(), dev_ape.tensor()); + + DeviceTensor dev_q(backend, qg.q, DType::kF32, + {g::kBatch, g::kSeqLen, d.n_heads, d.head_dim}, q_states.data()); + DeviceTensor dev_hw(backend, qg.q, DType::kF32, {g::kBatch, g::kSeqLen, d.n_heads}, + head_w.data()); + DeviceTensor dev_vk(backend, qg.q, DType::kI32, {g::kBatch, g::kSeqLen}, valid_keys.data()); + DeviceTensor dev_qm(backend, qg.q, DType::kI32, {g::kBatch, g::kSeqLen}, q_mask.data()); + DeviceTensor dev_topk(backend, qg.q, DType::kI32, {g::kBatch, g::kSeqLen, out_w}); + DeviceTensor dev_scores(backend, qg.q, DType::kF32, {g::kBatch, g::kSeqLen, np}); + + vt::Glm5NextKpoolSelectArgs args; + args.index_topk = d.index_topk; + args.current_length = g::kSeqLen; + args.softmax_scale = d.softmax_scale(); + args.always_select_tail = true; + vt::Glm5NextKpoolSelect(qg.q, dev_topk.tensor(), dev_scores.tensor(), dev_q.tensor(), + dev_hw.tensor(), dev_keys.tensor(), dev_idx.tensor(), + dev_valid.tensor(), dev_np.tensor(), dev_vk.tensor(), + dev_qm.tensor(), args); + backend.Synchronize(qg.q); + + std::vector got_topk(static_cast(rows * out_w)); + std::vector got_scores(static_cast(rows * np)); + dev_topk.Download(qg.q, got_topk.data()); + dev_scores.Download(qg.q, got_scores.data()); + + // The scores over the LIVE pools only. `[P, np)` is allocation slack. + std::vector scores_pref, golden_scores(g::kIndexScores, + g::kIndexScores + rows * P); + for (int64_t r = 0; r < rows; ++r) + for (int64_t p = 0; p < P; ++p) + scores_pref.push_back(got_scores[static_cast(r * np + p)]); + RequireFinite(scores_pref, "device index_scores"); + const double score_delta = MaxAbsDiff(scores_pref, golden_scores); + + // POSITIONWISE against the transformers run, which is strictly stronger than + // the set comparison and pins the emission ORDER as well as the membership. + int64_t mismatched = 0; + for (int64_t r = 0; r < rows; ++r) + for (int64_t i = 0; i < out_w; ++i) + if (got_topk[static_cast(r * out_w + i)] != + g::kTopkIndices[static_cast(r * out_w + i)]) + ++mismatched; + INFO("positionwise mismatches against transformers: " << mismatched << " of " + << rows * out_w); + CHECK(mismatched == 0); + + // SET equality, which is the model's own semantics (duplicates and -1 both + // absorbed downstream), asserted separately so a positionwise regression and + // a selection regression are distinguishable in the output. + int64_t set_mismatch = 0; + for (int64_t r = 0; r < rows; ++r) { + const std::set got = SelectedSet(got_topk.data() + r * out_w, out_w); + const std::set want = SelectedSet(g::kTopkIndices + r * out_w, out_w); + if (got != want) ++set_mismatch; + } + CHECK(set_mismatch == 0); + + // And against the host reference over the same packed history. + const IndexerSelection host = SelectIndexerTopkFromPacked( + d, w, hidden, q_resid, mask, packed, g::kBatch, g::kSeqLen, g::kSeqLen); + int64_t host_mismatch = 0; + for (size_t i = 0; i < host.topk_indices.size(); ++i) + if (host.topk_indices[i] != got_topk[i]) ++host_mismatch; + INFO("positionwise mismatches against the host reference: " << host_mismatch); + CHECK(host_mismatch == 0); + + // THE MARGIN. Computed from the host reference's own candidate mask, printed, + // and required to be strictly positive on the rows that actually prune — a + // fixture whose boundary gap is zero decides the top-k by the tie rule alone + // and cannot detect a scoring defect at all. + const int64_t select_k = d.SelectK(P); + double worst = std::numeric_limits::infinity(); + int64_t pruning_rows = 0; + for (int64_t b = 0; b < g::kBatch; ++b) { + for (int64_t s = 0; s < g::kSeqLen; ++s) { + const int64_t r = b * g::kSeqLen + s; + std::vector row(host.index_scores.begin() + r * P, + host.index_scores.begin() + (r + 1) * P); + std::vector cand(static_cast(P), 0); + for (int64_t p = 0; p < P; ++p) { + const int32_t last = + host.pooled.pool_indices[static_cast((b * P + p) * d.index_kpool + + d.index_kpool - 1)]; + const int64_t safe = last < 0 ? 0 : (last >= g::kSeqLen ? g::kSeqLen - 1 : last); + const bool vis = safe <= s && valid_keys[static_cast(b * g::kSeqLen + safe)] != 0; + cand[static_cast(p)] = + (vis && host.pooled.pool_valid[static_cast(b * P + p)] != 0) ? 1 : 0; + } + const double m = RowMargin(row, cand, P, select_k); + if (m < 0.0) continue; // fewer candidates than the budget: nothing pruned + ++pruning_rows; + worst = std::min(worst, m); + } + } + MESSAGE("pruning rows = " << pruning_rows << ", smallest decision margin = " << worst); + MESSAGE("index_scores max|device - transformers| = " << score_delta); + REQUIRE(pruning_rows > 0); + CHECK(std::isfinite(worst)); + CHECK(worst > 0.0); + + // THE SCORE BOUND IS TIED TO THE MARGIN, not chosen. An absolute tolerance on + // a score would be a number with no meaning: the golden magnitudes here run to + // 45, so one f32 ULP is already 3.8e-6 and a tolerance below that fails on + // arithmetic while a tolerance far above it stops bounding anything. The + // question a selection gate can actually answer is whether the numerical + // difference is small against the gap the top-k decides on, so that is what is + // asserted — with a factor of four of room, and both numbers printed so a + // future fixture that narrows the margin fails here instead of silently + // becoming a coin flip. + CHECK(score_delta * 4.0 < worst); +} + +TEST_CASE("glm5_next k-pool device: P == 0 serves the raw visible tail alone") { + if (!HasCuda()) { + MESSAGE("no CUDA backend: the k-pool short-context device gate is SKIPPED"); + return; + } + Backend& backend = vt::GetBackend(DeviceType::kCUDA); + QueueGuard qg(backend); + + const IndexerDims d = FixtureDims(); + const IndexerWeights w = FixtureWeights(); + const int64_t out_w = d.OutputWidth(); + const std::vector ape(g::kKpoolApe, g::kKpoolApe + d.index_kpool * d.head_dim); + + int64_t h_off = 0, q_off = 0, m_off = 0, t_off = 0; + for (int64_t c = 0; c < g::kShortCases; ++c) { + const int64_t S = g::kShortSeqLen[c]; + const std::vector hidden(g::kShortHidden + h_off, + g::kShortHidden + h_off + S * g::kHidden); + const std::vector q_resid(g::kShortQResid + q_off, + g::kShortQResid + q_off + S * g::kQLora); + const std::vector mask = MaskBytes(g::kShortMask + m_off, static_cast(S)); + const std::vector packed = PackIndexerStates(d, w, hidden, mask, 1, S); + const int64_t np = (S + d.index_kpool - 1) / d.index_kpool; + + DeviceTensor dev_packed(backend, qg.q, DType::kF32, {1, S, 2 * d.head_dim + 1}, + packed.data()); + DeviceTensor dev_ape(backend, qg.q, DType::kF32, {d.index_kpool, d.head_dim}, ape.data()); + DeviceTensor dev_keys(backend, qg.q, DType::kF32, {1, np, d.head_dim}); + DeviceTensor dev_idx(backend, qg.q, DType::kI32, {1, np, d.index_kpool}); + DeviceTensor dev_valid(backend, qg.q, DType::kI32, {1, np}); + DeviceTensor dev_np(backend, qg.q, DType::kI32, {1}); + vt::Glm5NextKpoolCompress(qg.q, dev_keys.tensor(), dev_idx.tensor(), dev_valid.tensor(), + dev_np.tensor(), dev_packed.tensor(), dev_ape.tensor()); + backend.Synchronize(qg.q); + std::vector got_np(1, -1); + dev_np.Download(qg.q, got_np.data()); + INFO("short case " << c << " (S = " << S << ")"); + CHECK(got_np[0] == static_cast(g::kShortNumPools[c])); + + const std::vector q_states = + LinearRows(w.wq_b, q_resid, S, d.n_heads * d.head_dim, d.q_lora_rank); + const std::vector head_w = + LinearRows(w.weights_proj, hidden, S, d.n_heads, d.hidden_size); + std::vector valid_keys(static_cast(S), 0); + for (int64_t i = 0; i < S; ++i) + valid_keys[static_cast(i)] = + packed[static_cast(i * (2 * d.head_dim + 1) + 2 * d.head_dim)] != 0.0f ? 1 : 0; + std::vector q_mask(g::kShortMask + m_off, g::kShortMask + m_off + S); + + DeviceTensor dev_q(backend, qg.q, DType::kF32, {1, S, d.n_heads, d.head_dim}, + q_states.data()); + DeviceTensor dev_hw(backend, qg.q, DType::kF32, {1, S, d.n_heads}, head_w.data()); + DeviceTensor dev_vk(backend, qg.q, DType::kI32, {1, S}, valid_keys.data()); + DeviceTensor dev_qm(backend, qg.q, DType::kI32, {1, S}, q_mask.data()); + DeviceTensor dev_topk(backend, qg.q, DType::kI32, {1, S, out_w}); + DeviceTensor dev_scores(backend, qg.q, DType::kF32, {1, S, np}); + + vt::Glm5NextKpoolSelectArgs args; + args.index_topk = d.index_topk; + args.current_length = S; + args.softmax_scale = d.softmax_scale(); + args.always_select_tail = true; + vt::Glm5NextKpoolSelect(qg.q, dev_topk.tensor(), dev_scores.tensor(), dev_q.tensor(), + dev_hw.tensor(), dev_keys.tensor(), dev_idx.tensor(), + dev_valid.tensor(), dev_np.tensor(), dev_vk.tensor(), + dev_qm.tensor(), args); + backend.Synchronize(qg.q); + + std::vector got(static_cast(S * out_w)); + dev_topk.Download(qg.q, got.data()); + for (int64_t i = 0; i < S * out_w; ++i) { + CHECK(got[static_cast(i)] == + g::kShortTopk[static_cast(t_off + i)]); + } + + h_off += S * g::kHidden; + q_off += S * g::kQLora; + m_off += S; + t_off += S * out_w; + } +}