Repository navigation
[draft] feat(rdna2): LeapQuant config flag, validator, numerics tests and opt-in page packing (follow-up to #168) - #170
Draft
BlivionIaG wants to merge 13 commits into
Draft
BlivionIaG wants to merge 13 commits into
BlivionIaG wants to merge 13 commits into
Conversation
…t-in) VLLM_RDNA_GDN_LEAPQUANT_PACK=1 (needs VLLM_RDNA_GDN_LEAPQUANT=1, default off) keeps GDN and PLE pages at their natural size and lets the KV-cache planner use the fewest groups whose layers fit one block, instead of 3 groups of 12. In this allocator the unit of capacity is the KV-cache group, not the page. Layers of a group already sit side by side in one block and every group draws block ids from one pool. Padding GDN pages to the attention page caps a block at 13 pages, and the repeat-pattern split makes 3 groups of 12 for any page size, so a smaller LeapQuant page frees nothing. The flag removes both limits. - leapquant_gdn_pack_enabled(): the flag; raises if LeapQuant is off. - Platform._align_hybrid_block_size: skip the Mamba page padding under the flag. - _get_packed_kv_cache_groups: Mamba buckets use cdiv(layers, layers per block) groups under the flag and log the split. Flag off is unchanged. - serve_rdna.sh: list the flag in the PRINT=1 environment plan. Flash-Next TP4 MTP2 from the real planner (CPU): the flag is a no-op at the 512-token recipe (15 layers fit a block, 3 groups stay). From 598 tokens up (608, 640, 1024) 36 GDN layers fit 2 groups of 18, so one request holds 10 GDN block ids instead of 15. Capacity at 262,144 tokens stays within 0.15 % of the 512-token recipe (608: 273,846, 640: 273,220, against 273,456); 1k contexts gain +7.6 % (608) and +2.0 % (640). Tables and byte math: docs/rdna2/leapquant-page-packing-2026-11-14.md. Tests: 29 CPU tests (tests/v1/core/test_leapquant_page_packing.py) run the real planner and allocate_kv_cache without a GPU: flag-off layout identical to the base branch, group counts, capacity-log math, per-layer view offsets, strides and 16-byte alignment, grouped normalize offsets. Two GPU kernel tests for packed pages are added and have not run. Nothing here ran on a GPU. Signed-off-by: BlivionIaG <kev29lt@gmail.com>
The one-layer PLE bucket logged "1 Mamba layers". Use "layer(s)". Signed-off-by: BlivionIaG <kev29lt@gmail.com>
Chosen option (fit GDN layers into a block by bytes), rejected options, byte math, the capacity table from the real planner, a what-if for merging the PLE layer into a GDN group, risks and the pending GPU checklist. Credits mirror the first LeapQuant note. Signed-off-by: BlivionIaG <kev29lt@gmail.com>
…rate states The LeapQuant RFC (vllm-project#64725, lihaoyang-amd, sections 2.2 and 5) reports NaN in the fold for exactly low-rank states and asks for a max-abs scaled, zero-guarded normalization plus regression tests. The RFC has no code. CPU evidence on the oracle of PR #168 (7b4aff7): it never gives NaN for a low-rank state. It loses accuracy instead: a rounding-noise compensator column above the absolute 1e-12 limit is normalized, stays half parallel to the earlier columns, and the INT8 residual carries a state-sized error (relative error 1.6e-3 to 2.4e-3 against 1.2e-6 after this change for rank 1; the outcome also depended on the state scale). It gives NaN only for states of about 1e37 and more, of any rank: the sum of |R| for the smoothing factor overflows, and z * bscale * csmooth becomes 0 * inf. Oracle (leapquant_gdn.py): - _orthonormalize scales each column by its max |entry|, sweeps twice (Kahan's "twice is enough", Parlett 1980) and zeroes a column when the second sweep removed more than half of the largest entry (dependence test of Daniel, Gragg, Kaufman and Stewart, Math. Comp. 30, 1976). Every column is zero or unit. - A dead column of the previous comp_k restarts from its fixed integer column. - The smoothing factor divides |R| by V before the sum (bit-identical to mean() for V a power of two). - Header words outside the lq_header_ok range read as an empty ring, like the kernels (the old oracle raised IndexError or AssertionError there). HIP fold (gdn_leapquant_rdna2.cu): the same logic. NOT BUILT, NOT RUN, no HIP compiler was available; it is marked as pending GPU validation. Tests: tests/kernels/mamba/test_gdn_leapquant_oracle.py, 122 CPU tests (92 for the fold, 30 for stale state of recycled blocks). Against the old oracle 50 of them fail (4 NaN, 35 accuracy or orthogonality, 11 invalid headers). Signed-off-by: BlivionIaG <kev29lt@gmail.com>
…ing GPU validation) Tests for the safety list of the LeapQuant RFC (vllm-project#64725, lihaoyang-amd, sections 2.2, 4 item 11 and 5). The RFC has no code; the tests are original work for this fork. They skip without a gfx103x GPU. STATUS: written without GPU access, never run on a GPU. 46 tests: - fold of zero, tiny, huge, exactly low-rank and float-limit states: finite, accurate, orthonormal compensators; a degenerate head leaves its neighbours bitwise unchanged; a rank-1 trajectory through window flushes - determinism: store, normalize, replay and decode (and the MonoKernel path), 200 repeats under a second-stream matmul load, bitwise equal outputs, states and blobs; the output buffers carry a different poison in each repeat - stale state of recycled blocks: random bytes, a previous tenant, the same with NaN, a valid-looking header; with and without the normalize and materialize pass of the serving prefill path; MonoKernel on a zero block and after a store; ten headers outside the range check; rows past the cursor; the align pre-copy overwriting every byte of a dirty destination - flush race: full windows on many (block, head) pairs with MTP accepted counts, one grouped normalize over six layers against per-layer launches, bitwise, repeated, with and without load; batched decode against singleton launches Only the first two header words [cursor, pending] are compared. Signed-off-by: BlivionIaG <kev29lt@gmail.com>
…enchmark Item 13 of section 4 of the LeapQuant RFC (vllm-project#64725, lihaoyang-amd) asks to extend benchmarks/replayssm/e2e_decode_speedup.py. The script is from vllm-project#48018 (Johnny-Liou); its header keeps that credit and records this change. - new arm `leapquant`: sets VLLM_RDNA_GDN_LEAPQUANT=1 in its own worker process, imports leapquant_gdn only there (gfx10x, fails loudly elsewhere); every other arm refuses to start when the flag is set in the environment - `--modes` (default `standard,replayssm`, so the old behaviour is unchanged; the first mode is the baseline of the speedup lines), `--num-speculative-tokens` (MTP), `--monokernel`, `--tensor-parallel-size`, repeatable `--llm-kwarg KEY=JSON` and `--env KEY=VALUE`, `--dry-run` (imports neither torch nor vllm) - the FP32 baseline against LeapQuant on the RDNA path is `--modes standard,leapquant --mamba-ssm-cache-dtype float32` - NOT done: the 16-bit per-window baseline of the RFC. It needs the Gated DeltaNet ReplaySSM decode kernel (vllm-project#62004), which this tree does not have; use_replayssm is Mamba2 only. The docstring says so. tests/benchmarks/test_replayssm_e2e_decode_speedup.py: 17 CPU tests pin the arguments, the arm and environment rules, the unchanged default arms, the worker commands, --help and --dry-run without torch. No benchmark was run. Signed-off-by: BlivionIaG <kev29lt@gmail.com>
…s [skip ci] One section per item of the safety list of the LeapQuant RFC (vllm-project/vllm vllm-project#64725, lihaoyang-amd): what changed, what was verified on the CPU, what is pending GPU validation. - NaN question answered with CPU evidence: the oracle of PR #168 never gives NaN for low-rank states (it loses accuracy: 1.6e-3 to 2.4e-3 against 1.2e-6 for rank 1) and gives NaN only for states of about 1e37 and more (the smoothing sum overflows), for any rank. Old-vs-new tables, the design that failed, the host emulation of lq_orthonormalize (g++, not the HIP compiler). - Determinism: status "not yet run". - Stale state of recycled blocks: what the serving prefill path does with the leftovers, findings (no stale read at the contract level; two header cases that the kernels cannot check), tests. - Flush race: review by construction with line numbers for normalize, decode, the MonoKernel and the preconditions; the one disagreement case. - Benchmark arm, commands, the 16-bit baseline NOT done and why. - Pending GPU checklist and not-done list. Credits the RFC for the item list, Kahan/Parlett and Daniel-Gragg-Kaufman-Stewart for the two textbook results, and vllm-project#48018 (Johnny-Liou) for the benchmark. No external code was used. Signed-off-by: BlivionIaG <kev29lt@gmail.com>
The new GPU tests never ran on a GPU, so a failure is either a kernel bug or a test bug. Signed-off-by: BlivionIaG <kev29lt@gmail.com>
…te it in one place
Replace the VLLM_RDNA_GDN_LEAPQUANT environment variable with
--replayssm-state-dtype {auto,leapquant} (CacheConfig.replayssm_state_dtype).
There is no alias: the engine does not read the variable any more.
- verify_leapquant_gdn is the one place that refuses an unsupported setup. It
runs at config time (VllmConfig) and for every LeapQuant GDN layer. The
refusals that were spread over the layer code move into it with the same
meaning. New config-time refusals: --use-replayssm, stochastic rounding, KV
connectors, KV offloading, head dimensions other than 128, and models
without GDN layers. MTP and --monokernel stay supported.
- MambaSpec carries a state format identity (format, layout version, head
dimensions, rank, window, ring capacity, record dtype) from
LeapQuantLayout.state_format. FP32 and LeapQuant state never share a KV
cache group, and the V2 align copy marker follows the spec.
- KVConnectorFactory.create_connector refuses a KV cache config that holds
LeapQuant state.
- The scheduler growth gate, the max-memory estimate, and the GDN state
calculators read the flag. Flag-off hashes, dtypes, shapes, page sizes and
launcher plans equal the first PR head (CPU runs, see the receipts).
- The launcher takes REPLAYSSM_STATE_DTYPE and passes the flag after
--block-size. It stops when VLLM_RDNA_GDN_LEAPQUANT=1 is still set, because
the engine would ignore it and serve FP32 state. The
flashnext-leapquant-mtp2 plan equals the old one except for the flag
spelling.
Credits: the list of checks comes from RFC vllm-project#64725
(lihaoyang-amd); the flag name and its place next to --replayssm-buffer-len
come from draft PR vllm-project#52962 (askliar). No code was taken from
either. Nothing in this change ran on a GPU.
Signed-off-by: BlivionIaG <kev29lt@gmail.com>
…eipts [skip ci] New leapquant-config-guards doc: the flag table, the rejected setups, why --use-replayssm is rejected, why the window stays 32, the state format identity, the connector refusal, flag-off identity, the launcher dry run, the tests, the checks that still need a GPU, and credits (RFC vllm#64725, vllm#52962, vllm#48018, vllm#62004, arXiv 2609.38166). The first LeapQuant doc names the new flag and says that its receipts used the environment variable. bench_results/2026-11-14_leapquant-config-guards holds the dry-run plans before and after and the CPU receipts. Signed-off-by: BlivionIaG <kev29lt@gmail.com>
Signed-off-by: BlivionIaG <kev29lt@gmail.com>
…-packing branch # Conflicts: # tools/rdna/serve_rdna.sh # vllm/model_executor/layers/mamba/ops/leapquant_gdn.py Signed-off-by: BlivionIaG <kev29lt@gmail.com>
BlivionIaG
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Draft. Follow-up to #168 (merged). Nothing here has run on a GPU, and the HIP change is not built.
Four parts, each its own commits:
--replayssm-state-dtype {auto,leapquant}(CacheConfig.replayssm_state_dtype) replaces the env variableVLLM_RDNA_GDN_LEAPQUANT. There is no alias. The launcher key isREPLAYSSM_STATE_DTYPE, and the launcher exits if the old variable is still set.verify_leapquant_gdnrejects every unsupported setup at config time, per layer and at connector creation.MambaSpec.state_formatidentifies the format.VLLM_RDNA_GDN_LEAPQUANT_PACK=1). It gains nothing at the shipped 512-token recipe.Flag off, the KV-cache config and the config hashes are unchanged (observed on CPU).
flowchart LR F["--replayssm-state-dtype leapquant"] --> V["verify_leapquant_gdn<br/>config / layer / connector"] V -->|ok| S["MambaSpec.state_format<br/>leapquant-int8/v1/..."] V -->|unsupported| X["stop at start-up<br/>names flag and cause"] S --> P["planner: pack GDN layers<br/>(only with the PACK switch)"]1-2. Flag, validator, format id
--use-replayssm, any KV connector or offloading, V1 runner with align mode, AITER GDN, FlashInfer GDN, fused CUDA spec decode, models without GDN.--monokernel.--use-replayssmis rejected, because Qwen GDN models do not declare ReplaySSM support. The window is fixed at 32, and--replayssm-buffer-lenis not read. All connectors are rejected, stricter than the RFC.MambaSpec.state_format(for exampleleapquant-int8/v1/k128v128/r4/p32/cap34/rec-float16) is part of the spec key. There is no marker in the blob, because the kernels are not rebuilt here.--helpoption, one newvllm:cache_config_infolabel (replayssm_state_dtype),MambaSpecgainsstate_format(defaultNone).docs/rdna2/leapquant-config-guards-2026-11-14.md.3. Numerics
gdn_leapquant_rdna2.cu.benchmarks/replayssm/e2e_decode_speedup.pygains a LeapQuant arm. The 16-bit per-window baseline arm is not done, because it needs the GDN ReplaySSM kernel from vllm#62004.docs/rdna2/leapquant-gdn-numerics-and-isolation-2026-11-14.md.4. Page packing
Layers of one KV-cache group already share a block, and all groups draw from one block-id pool. Capacity depends on bytes per block and on how many groups a request holds. Today 3 GDN groups of 12 come from padding the GDN page to the attention page. With the switch, the planner skips the padding and uses
cdiv(layers, layers per block)groups.CPU planner, Flash-Next TP4, MTP2 (capacity tokens at
max_model_len262,144):At max context packing moves capacity by 0.14% or less. Merging the PLE group into a GDN group would add more (+21% to +41% at 1k-4k tokens in the planner what-if, not built). Details:
docs/rdna2/leapquant-page-packing-2026-11-14.md.Verification
test_leapquant_config.py::test_cli_flag_parses_and_restricts_its_values: it callsEngineArgs.from_cli_argswith--model dummy, which needs an HF download and the box runs offline. It was not run on the config-guards branch alone.test_gdn_leapquant_rdna2.py,test_gdn_leapquant_safety_rdna2.py) skip without a GPU.test_mamba_align_chunk_split.py::test_flash_next_tp4_materializes_replay_state[3200|3602].Pending GPU validation
flashnext-leapquant-mtp2with the flag: KV capacity 273,456 tokens, group sizes(512,512,512,512,8,512), greedy, prefix-cache and concurrency probes; compare with feat(rdna2): add LeapQuant GDN state with 512-token pages #168's numbers_rocm_Cwith the fold change; run the kernel tests (22 existing, 46 new) on gfx1030--monokernelunchanged with the flagCredits
LeapQuant (arXiv 2609.38166), ReplaySSM (vllm#48018, vllm-project#62004) and the RFC vllm#64725 (lihaoyang-amd) are credited in file headers and docs. The flag name and its place next to
--replayssm-buffer-lenfollow draft vllm#52962 (askliar). No outside code was copied; the validator, layout id, fold guard, tests and planner change are original to this fork.