diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml index 143fc427a49e..ab3d25c81c82 100644 --- a/.github/workflows/pre-commit.yml +++ b/.github/workflows/pre-commit.yml @@ -47,9 +47,13 @@ jobs: pre-commit: needs: pre-run-check if: always() && (needs.pre-run-check.result == 'success' || needs.pre-run-check.result == 'skipped') - runs-on: [self-hosted, linux, x64, vllm-runners] + # Fork integration branches do not have access to the upstream runner pool. + runs-on: ubuntu-latest steps: - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + with: + # PR-diff linting needs the fork base commit as well as the head. + fetch-depth: 0 - uses: actions/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548 # v6.1.0 with: python-version: "3.12" @@ -60,6 +64,32 @@ jobs: - run: echo "::add-matcher::.github/workflows/matchers/actionlint.json" - run: echo "::add-matcher::.github/workflows/matchers/markdownlint.json" - run: echo "::add-matcher::.github/workflows/matchers/mypy.json" - - uses: pre-commit/action@2c7b3805fd2a0fd8c1884dcaebf91fc102a13ecd # v3.0.1 - with: - extra_args: --all-files --hook-stage manual + - run: python -m pip install pre-commit + - name: Check changed shell scripts + if: github.event_name == 'pull_request' + shell: bash + env: + BASE_SHA: ${{ github.event.pull_request.base.sha }} + HEAD_SHA: ${{ github.event.pull_request.head.sha }} + run: | + mapfile -d '' shell_files < <( + git diff --name-only --diff-filter=ACMR -z \ + "$BASE_SHA" "$HEAD_SHA" -- '*.sh' + ) + if ((${#shell_files[@]})); then + shellcheck -s bash "${shell_files[@]}" + fi + - name: Check pull request changes + if: github.event_name == 'pull_request' + env: + # The repository hook scans every shell script and ignores the PR file + # list. The previous step preserves changed-file coverage on the fork. + SKIP: shellcheck + run: >- + pre-commit run --show-diff-on-failure --color=always + --from-ref ${{ github.event.pull_request.base.sha }} + --to-ref ${{ github.event.pull_request.head.sha }} + --hook-stage manual + - name: Check all files on push + if: github.event_name == 'push' + run: pre-commit run --show-diff-on-failure --color=always --all-files --hook-stage manual diff --git a/docs/design/glm52_sparse_ckv.md b/docs/design/glm52_sparse_ckv.md new file mode 100644 index 000000000000..e4790a394e47 --- /dev/null +++ b/docs/design/glm52_sparse_ckv.md @@ -0,0 +1,346 @@ +# GLM-5.2 Sparse CKV under Decode Context Parallelism + +## Status + +This document describes an experimental, feature-gated implementation for +GLM-5.2 sparse MLA on PCIe-connected Blackwell GPUs. It combines mixed +replicated/sharded KV allocation, full-CKV prefill lookahead, and selected-record +CKV exchange during decode. + +The complete stack is validated on one 4-GPU host with TP4/DCP4/MTP3 and the +native 432-byte `nvfp4_ds_mla` + BF16-RoPE cache. Generic transport supports +other widths and DCP sizes, but those are not end-to-end model claims. +All features default off; unsupported batches retain stock DCP. + +## Goals and Non-goals + +Goals: + +1. Preserve DCP KV capacity while recovering prefill and decode overhead. +2. Keep model math unchanged; move storage, communication, and scheduling. +3. Deduplicate MTP candidates per sequence. +4. Support bounded C1-C8 decode with explicit stock-DCP fallbacks. +5. Keep CUDA graph replay allocation-free and isolate target/draft state. +6. Separate vLLM policy from record-format-agnostic Sparkinfer transport. + +Non-goals: + +- Change checkpoint quantization or resident KV representation. +- Claim validated FP8-RoPE, TP6/TP8, or DCP6/DCP8 model support. +- Replace stock DCP for unsupported or overflowing batches. +- Provide host/network fallback inside CUDA-IPC transports. + +## Problem and Model Structure + +The stock sparse-MLA path exposes query/head communication, indexer work, KV +writes, and CKV retrieval on a mostly serial critical path. CKV traffic grows +with context, while a DCP-sharded indexer adds another distributed lookup. + +GLM-5.2 starts with three Full layers and then repeats +`Full -> Shared -> Shared -> Shared`. A Full layer computes top-K; the next +three Shared layers reuse it. This permits layer-specific CKV preparation ahead +of use. + +The design combines: + +1. Replicated indexer KV with DCP-sharded main CKV. +2. Full native CKV gather and depth-N prefetch for prefill. +3. Selected-record exchange deduplicated across each sequence's MTP rows. + +## Architecture + +```mermaid +flowchart LR + F["Full layer: finalized top-K"] --> U["Per-sequence union + remap"] + R["Replicated indexer KV"] --> F + U --> O["Owner/local-slot conversion"] + O --> X["Selected-record exchange"] + X --> W["Native sparse workspace"] + W --> A["Existing B12X sparse attention"] + F --> S["S1/S2/S3 lookahead"] + S --> X + X --> D["Direct CUDA-IPC"] + X --> C["CE compact/DMA/unpack"] +``` + +vLLM owns model structure, union/dedup, remapping, eligibility, fallback, and +attention metadata. Sparkinfer receives fixed-width records plus a finalized +destination map; it does not interpret CKV or compute top-K. + +## Cache ABI + +Validated: `nvfp4_ds_mla`, `KV_FP8_ROPE=0`. + +| Region | Bytes | Meaning | +| --- | ---: | --- | +| Latent payload | 256 | Packed NVFP4 latent | +| Latent scales | 32 | Group scales | +| Alignment | 16 | Resident-cache padding | +| RoPE payload | 128 | BF16 RoPE | + +The sparse workspace preserves this native 432-byte geometry so existing B12X +attention consumes it unchanged. Sparkinfer is byte-record generic; vLLM +currently gates sparse decode to this ABI. + +## Mixed KV Allocation and Replicated Indexer + +Replicated indexer KV and DCP-sharded main CKV need different page geometry and +block counts. Padding the indexer page is invalid because its backend requires +identity stride. + +The allocator groups them separately but allocates in lockstep. Both groups +cover the same global token span per logical block, use identical block IDs, +and advance together through scheduling, prefix promotion, block tables, and +GPU tensors. External KV loads are rejected in lockstep mode because external +ownership cannot guarantee matching cross-group IDs. + +`VLLM_DCP_REPLICATE_INDEXER_CACHE=1` applies only to real target indexer +layers. It requires the V2 runner, B12X sparse indexer, DCP2-8, and PCP1. +Effective indexer DCP becomes one, removing indexer-K sharding and its DCP +merge. Main CKV remains sharded; draft/MTP cache policy remains separate. + +## Full-CKV Prefill + +`VLLM_B12X_MLA_CKV_GATHER=1`: + +1. Packs rank-local native paged CKV. +2. Gathers it on a dedicated CKV communicator. +3. Runs local query heads over global CKV. +4. Avoids normal DCP query gather and LSE merge. + +It requires pure prefill, complete DCP metadata, non-DBO workspace, aligned +heads, a token count inside configured bounds, and no active CUDA capture. +Ineligible steps use stock DCP. + +`VLLM_B12X_MLA_CKV_PREFETCH_DEPTH=N` allocates a `depth + 1` ring and queues +future layers on a side stream. At the target layer, current-chunk KV written +after prefetch began is patched into gathered history before attention. +Lookahead supports `fp8_ds_mla` and `nvfp4_ds_mla`; unsupported formats use +synchronous full gather. + +## Sparse Selected-Record Decode + +`VLLM_B12X_MLA_SPARSE_DECODE_CKV_GATHER=1`: + +1. Finalizes top-K for every query row. +2. Groups rows by request. +3. Builds a stable dense union per destination/request across MTP rows. +4. Produces exact row-to-union remaps. +5. Converts logical tokens to owner-local paged-cache slots. +6. Exchanges only owned native records. +7. Patches the current token when lookahead preceded its KV write. +8. Runs existing local-head B12X attention over the sparse workspace. + +MTP3 has up to four rows/request and a worst-case `4 * topk` union. Dedup makes +wire traffic follow unique records per sequence instead of speculative depth. +Dedup never crosses unrelated requests. + +The workspace is pooled. `POOL_RECORDS=0` covers the configured fast-path +request count. A smaller pool saves VRAM and lowers the fast-path concurrency +bound. Overflow and unsupported shapes use stock DCP. + +## Shared-Layer Bulk Lookahead + +After a Full layer, depth three discovers S1/S2/S3. They share indices and +remap geometry but retain distinct CKV values. + +With CE and `VLLM_B12X_MLA_SPARSE_DECODE_BULK_PREFETCH=1`, exactly three +compatible layers use one pack, transfer, and synchronization sequence. S1 may +wait for the combined packet; S2/S3 consume ready workspaces without new +inter-rank barriers. Missing API support, incompatible layers, or layered +capacity overflow falls back to per-layer exchange. + +The pooled sparse path supports C1-C8. The bulk CE speedup was measured only +within its layered capacity; larger batches retain the per-layer path. + +## Transport Backends + +### Direct CUDA-IPC + +An SM kernel scans each destination map and peer-writes owned records into a +destination CUDA-IPC slab at final selected positions. Barrier phase 0 makes +writes visible; the destination copies its slab into caller output; phase 1 +prevents reuse. This avoids NCCL gather/recombine but retains one local D2D copy. + +### Copy Engine + +CE compacts owned records per destination, preserves final positions, rotates +destination order by rank, and transfers a contiguous primary packet with +`cudaMemcpyAsync`. Owner skew uses an IPC overflow region. One arrival barrier +precedes unpack; release is deferred until slab reuse. CE costs staging VRAM but +uses fewer SM resources and implements the three-layer bulk API. + +### Shared Contract + +Both require contiguous CUDA records/output, a contiguous integer destination +map, exactly one owner per destination slot, rank-identical call order and +active count, CUDA IPC/P2P, one stream-affine channel per stream, eager binding +before capture, and cleanup before process-group destruction. + +NCCL is used for initialization and rank-consistency handshakes. Record data +moves through CUDA IPC peer stores or peer DMA. No host/network fallback exists +inside the transport. + +`transport=auto` prefers CE and falls back to direct if CE initialization is +unavailable. `ce` is strict. `direct` never constructs CE state. + +## CUDA Graph and Correctness + +Sparse state is keyed by workspace identity and execution-lane key, preventing +target and draft passes from aliasing rings, events, workspaces, or channels. +Graph capacity must cover `max_fast_requests * rows_per_request`. + +Initialization rank-consistently preloads the dense-union extension and binds +transport channels. Replay changes data/maps without allocations or address +changes. Full-CKV prefill is disabled during capture; selected-record decode is +graph-safe after eager setup. + +Correctness invariants: + +1. Exactly one source owns every destination position. +2. Stable union/remap is rank-identical. +3. MTP dedup is per request. +4. Replicated-indexer and sharded-CKV blocks advance in lockstep. +5. Current-layer writes patch prefetched history before use. +6. S1/S2/S3 share indices, not values. +7. Every rank issues communication in identical order. +8. Target/draft and caller-stream state cannot alias. +9. Graph replay allocates nothing. +10. Unsupported shapes use stock DCP. + +## Configuration + +All features default off. + +| Variable | Default | Purpose | +| --- | --- | --- | +| `VLLM_DCP_REPLICATE_INDEXER_CACHE` | `0` | Replicate target indexer KV | +| `VLLM_DCP_GLOBAL_TOPK` | `1` | Preserve global top-K semantics | +| `VLLM_DCP_QUERY_SPLIT` | `0` | Remove duplicate query work where applicable | +| `VLLM_B12X_MLA_CKV_GATHER` | `0` | Enable full-CKV prefill | +| `...CKV_GATHER_MIN_TOKENS` | `16` | Minimum prefill size | +| `...CKV_GATHER_MAX_TOKENS` | `524288` | Workspace ceiling | +| `VLLM_B12X_MLA_CKV_PREFETCH_DEPTH` | `1` | Lookahead; 0 is synchronous | +| `...SPARSE_DECODE_CKV_GATHER` | `0` | Enable selected-record decode | +| `...SPARSE_DECODE_TRANSPORT` | `direct` | `auto`, `ce`, or `direct` | +| `...SPARSE_DECODE_BULK_PREFETCH` | `0` | One CE exchange for S1/S2/S3 | +| `...SPARSE_DECODE_MAX_SEQS` | `8` | Fast-path request ceiling | +| `...SPARSE_DECODE_POOL_RECORDS` | `0` | Pool records; 0 is automatic | + +Query splitting is independent and gives no duplicate-work reduction at +TP4/DCP4 because TP/DCP is one. + +## Eligibility and Fallback + +| Condition | Full-CKV prefill | Sparse decode | +| --- | --- | --- | +| DCP | >1 | 2-8 | +| Workspace | non-DBO | non-DBO | +| Heads | aligned | aligned | +| Cache | gather generic; lookahead gated | 432-byte NVFP4 + BF16 RoPE | +| Shape | pure prefill | bounded uniform decode rows | +| Capture | stock fallback | allowed after eager setup | +| Missing metadata | stock fallback | stock fallback | +| Request/pool overflow | N/A | stock fallback | +| CE layered overflow | synchronous/per-layer | per-layer | + +## Validated Profile and Results + +Hardware: ASUS WRX90E-SAGE SE, Threadripper PRO 9965WX, 128 GiB RAM, +4x RTX PRO 6000 Blackwell 96 GB at 400 W, full peer access, ACS/IOMMU off. + +| Setting | Value | +| --- | --- | +| Parallelism | TP4 / DCP4 / MTP3 | +| Model limit | 300,000 | +| Max sequences / batch / graph | 8 / 2,048 / 32 | +| GPU utilization | 0.98 | +| KV allocation | 3,426,112,942 bytes/GPU | +| Reported KV | 302,047 tokens | +| Cache | `nvfp4_ds_mla`, BF16 RoPE | +| Features | replicated indexer, depth-3 full CKV, CE sparse, bulk S1/S2/S3 | +| TP all-reduce | NCCL | + +Daily five-run temperature-zero coding: 105.2 tok/s median, 104.5 mean, +101.6-108.0 range. + +| Cold prefill | 8K | 32K | 64K | 128K | 256K | +| --- | ---: | ---: | ---: | ---: | ---: | +| tok/s | 3,326 | 3,139 | 2,966 | 2,719 | 2,356 | + +Matched 90K coding: DCP4/MTP3 CE sparse improved 82.9 -> 99.2 tok/s (+19.6%) +at unchanged 302,080-token KV. At MTP0: 44.4 -> 45.4 (+2.4%). Bulk S1/S2/S3 +improved the 64K+ prefill geometric mean by 8.5% and coding +99.63 -> 103.79 with unchanged KV. + +Deterministic output matched byte-for-byte twice at about 8K, 64K, and 128K. +Results are configuration-sensitive and are not TP8/DCP8 projections. + +## Verification + +- 160 vLLM policy/workspace tests passed; 18 skipped. +- 3 mixed-DCP allocator tests passed. +- 57 Sparkinfer transport tests passed; 4 GPU-only tests skipped in CI. +- Ruff, formatting, mypy 3.10-3.13, Clang-format, import/API checks passed. +- Four-GPU startup, target/MTP graph capture, API/transport smoke, long-prefill + correctness, and C2/C4 decode gates passed. +- Transport GPU tests cover eager/graph replay, changing maps, empty/full + capacity, odd/production widths, rank-consistent failure, overflow, layered + exchange, DCP4, and offsets above 2 GiB. + +## Failure Behavior and Observability + +Startup logs report active gather, prefetch, sparse decode, transport, and bulk +exchange. Eligibility logs report format/topology/request fallback. Transport +initialization validates rank agreement, extension load, allocation, IPC +handles, and peer access. Barriers time out and device-trap instead of spinning +forever. Strict CE surfaces failure; automatic mode may choose direct. + +JIT compilation can contaminate latency. Warm every measured shape or enable +compile-progress diagnostics before recording results. + +## Remaining Work + +Required validation: + +- TP6/TP8 and DCP6/DCP8 full-model correctness and scaling. +- Matched final-image DCP1/DCP4/DCP8 comparisons. +- C1-C8, 8K-256K, MTP0/MTP3 stress and prolonged agent/tool-call tests. +- CE/direct comparison on switched and non-switched fabrics. +- Missing-P2P, pool-overflow, graph, and allocation-failure fallbacks. +- FP8-RoPE, FP8/NF3 KV, calibrated outer scales, and >128K exact output. +- Rebuild/GPU-test exact clean vLLM head; packaged validation used the + runtime-equivalent pre-cleanup head. + +Candidate optimizations: + +- Peer-write directly into final attention storage, removing local copy/unpack. +- Finish indexer/Q/KV stream overlap. +- Generalize vLLM record geometry beyond 432 bytes. +- Evaluate sparse Full-layer transfer against query gather + LSE merge. +- Skip the 16-byte alignment field on wire after padding-poison validation. +- Tune split-K and launch geometry by topology. + +## Upstream Decomposition + +1. Sparkinfer direct fixed-record CUDA-IPC exchange. +2. Sparkinfer CE exchange and three-layer API. +3. vLLM mixed MLA lockstep allocation. +4. vLLM target indexer replication. +5. vLLM native full-CKV depth-N prefetch. +6. vLLM selected-record policy, union/remap, and fallback. +7. vLLM CE three-layer integration. +8. Keep deployment recipes and host tuning outside runtime PRs. + +Each PR should remain feature-gated, focused, tested, and linked to this +document. Transport PRs stay record-format agnostic; the 432-byte restriction +belongs in vLLM until more formats are validated. + +## Credits + +The design originates from Koush's CKV-gather, shared-layer lookahead, sparse +decode, and per-sequence MTP-deduplication work. It builds on Luke Alonso's B12X +and Sparkinfer infrastructure and local-inference-lab's GLM-5.2 integration. + +Implementation, testing, benchmarking, and documentation were developed with +OpenAI Codex assistance and manually reviewed and validated on the target host. diff --git a/kv-scales/README.md b/kv-scales/README.md new file mode 100644 index 000000000000..6ed40a94def9 --- /dev/null +++ b/kv-scales/README.md @@ -0,0 +1,49 @@ +# GLM-5.2 NVFP4 MLA KV outer scales + +Per-layer calibration for the `nvfp4_ds_mla` KV cache writer +(`VLLM_NVFP4_MLA_SCALES_FILE`, format `nvfp4_ds_mla_outer_scale_v1`). + +GLM-5.2's post-RMSNorm 512-dim `kv_c` latent spans a ~240x amplitude range +across layers. With the default outer scale of 1.0, shallow layers quantize +with E4M3 block scales at or below the subnormal floor, and shallow-layer KV +error is strongly amplified downstream. `s_l = max_abs(kv_c_normed) / (6*448)` +re-centers every layer so its largest block scale lands at the top of the +E4M3 range. + +## Files + +- `glm52-nvfp4-nf3-hybrid_mla_outer_scales_v1.json` — calibrated on + `madeby561/GLM-5.2-MXFP8-NVFP4-NF3-Hybrid` (K64R16), Salesforce/wikitext + (wikitext-2-raw-v1, test), 2048-token context, TP4; per-layer envelope with + an independent community capture of the same base model for shallow-layer + headroom. `max_abs` per layer is included for auditability. + +## Usage + +```bash +VLLM_NVFP4_MLA_SCALES_FILE=/path/to/glm52-nvfp4-nf3-hybrid_mla_outer_scales_v1.json \ + ./serve-glm52.sh # nvfp4_ds_mla + B12X_MLA_SPARSE only; inert otherwise +``` + +## Results (teacher-forced prefill KLD vs BF16 reference, 5 fresh boots each) + +| KV config | mean +/- sd | max ctx (4x96GB) | +| --- | --- | --- | +| fp8_ds_mla | 0.1263 +/- 0.0030 | 373k | +| nvfp4_ds_mla + scales, bf16 rope | 0.1345 +/- 0.0035 | 550k | +| nvfp4_ds_mla + scales, fp8 rope (`KV_FP8_ROPE=1`) | 0.1356 +/- 0.0054 | 600k+ | +| nvfp4_ds_mla, no scales, bf16 rope | 0.158 | 550k | +| nvfp4_ds_mla, no scales, fp8 rope | 0.168 | 600k+ | + +Protocol: local-inference-lab/rtx6kpro `benchmarks/glm52-kld-evaluation.md` +(festr2 2026-07-08 reference logits, one fixed 2048-token window, 2047 +positions, full 154,880 vocab, `KL(ref || candidate)`). + +## Cache invalidation (required) + +CuTeDSL folds the outer-scale multiply out of the kernel when `latent_scale` +traces at exactly 1.0. b12x builds without the identity/dynamic compile-spec +fact (see lukealonso/b12x PR "mla: split latent_scale identity/dynamic +compile-cache entries") replay stale identity cubins from persistent compile +caches, silently dropping the restore. Clear mounted b12x compile caches once +when enabling scales on such builds. diff --git a/kv-scales/glm52-nvfp4-nf3-hybrid_mla_outer_scales_v1.json b/kv-scales/glm52-nvfp4-nf3-hybrid_mla_outer_scales_v1.json new file mode 100644 index 000000000000..7ae9d9351c4e --- /dev/null +++ b/kv-scales/glm52-nvfp4-nf3-hybrid_mla_outer_scales_v1.json @@ -0,0 +1,178 @@ +{ + "format": "nvfp4_ds_mla_outer_scale_v1", + "num_layers": 78, + "latent_dim": 512, + "denominator": 2688.0, + "formula": "s_l = max_abs(kv_c_normed) / (6 * 448)", + "model": "madeby561/GLM-5.2-MXFP8-NVFP4-NF3-Hybrid (local K64R16 build; own capture)", + "dataset": { + "name": "Salesforce/wikitext", + "config": "wikitext-2-raw-v1", + "split": "test", + "context_length": 2048, + "windows": 1, + "note": "the canonical festr2-0708 census window" + }, + "created_utc": "2026-07-20T16:40:43.312260+00:00", + "hook": "mla.py kv_a_layernorm output (bind-mount capture patch), per-TP-rank agreement checked", + "max_abs": [ + 0.04833984375, + 0.021728515625, + 0.03662109375, + 0.1005859375, + 0.023681640625, + 0.03515625, + 0.046875, + 0.033203125, + 0.259765625, + 0.1474609375, + 0.7734375, + 1.0, + 1.1171875, + 0.150390625, + 0.76171875, + 0.392578125, + 0.25390625, + 0.455078125, + 0.80859375, + 0.80078125, + 0.96875, + 0.78515625, + 0.58203125, + 0.330078125, + 0.89453125, + 0.73046875, + 1.0390625, + 1.578125, + 1.0078125, + 0.92578125, + 1.25, + 1.5859375, + 1.234375, + 1.9140625, + 1.9453125, + 1.7578125, + 1.7890625, + 2.546875, + 2.296875, + 2.4375, + 1.921875, + 2.171875, + 2.875, + 2.546875, + 2.765625, + 3.234375, + 2.5625, + 3.390625, + 2.40625, + 2.546875, + 3.09375, + 2.90625, + 3.671875, + 2.46875, + 2.609375, + 2.359375, + 3.078125, + 3.171875, + 2.578125, + 2.359375, + 3.890625, + 2.953125, + 4.09375, + 4.09375, + 5.0625, + 5.1875, + 2.984375, + 2.875, + 2.984375, + 3.296875, + 3.828125, + 2.859375, + 3.59375, + 3.734375, + 4.25, + 4.1875, + 4.84375, + 3.75 + ], + "scales": [ + 1.7983572823660715e-05, + 8.083525158110119e-06, + 1.3623918805803572e-05, + 3.742036365327381e-05, + 8.810134161086309e-06, + 1.3078962053571428e-05, + 1.743861607142857e-05, + 1.2352353050595238e-05, + 9.663899739583333e-05, + 5.485897972470238e-05, + 0.00028773716517857144, + 0.0003720238095238095, + 0.00041562034970238094, + 5.5948893229166664e-05, + 0.0002833775111607143, + 0.00014604840959821428, + 9.445917038690477e-05, + 0.00016929989769345238, + 0.00030081612723214287, + 0.0002979096912202381, + 0.0003603980654761905, + 0.00029209681919642856, + 0.00021652948288690475, + 0.0001227969215029762, + 0.00033278692336309525, + 0.00027175176711309525, + 0.0003865559895833333, + 0.0005871000744047619, + 0.0003749302455357143, + 0.0003444126674107143, + 0.0004650297619047619, + 0.0005900065104166666, + 0.0004592168898809524, + 0.0007120768229166666, + 0.0007237025669642857, + 0.0006539481026785714, + 0.0006655738467261905, + 0.0009474981398809524, + 0.0008544921875, + 0.0009068080357142857, + 0.0007149832589285714, + 0.0008079892113095238, + 0.0010695684523809525, + 0.0009474981398809524, + 0.0010288783482142857, + 0.0012032645089285715, + 0.0009533110119047619, + 0.0012613932291666667, + 0.0008951822916666666, + 0.0009474981398809524, + 0.0011509486607142857, + 0.0010811941964285715, + 0.001366024925595238, + 0.0009184337797619048, + 0.0009707496279761905, + 0.0008777436755952381, + 0.0011451357886904762, + 0.0011800130208333333, + 0.0009591238839285714, + 0.0008777436755952381, + 0.0014474051339285715, + 0.0010986328125, + 0.0015229724702380952, + 0.0015229724702380952, + 0.0018833705357142857, + 0.001929873511904762, + 0.001110258556547619, + 0.0010695684523809525, + 0.001110258556547619, + 0.0012265159970238095, + 0.0014241536458333333, + 0.0010637555803571428, + 0.0013369605654761905, + 0.0013892764136904762, + 0.0015811011904761905, + 0.0015578497023809525, + 0.0018019903273809525, + 0.0013950892857142857 + ] +} \ No newline at end of file diff --git a/serve-glm52.sh b/serve-glm52.sh index 3af05f4daa9d..e2577952468f 100755 --- a/serve-glm52.sh +++ b/serve-glm52.sh @@ -108,6 +108,8 @@ MOE_BACKEND="${MOE_BACKEND:-b12x}" MOE_SPEC_BACKEND="${MOE_SPEC_BACKEND:-b12x}" ATTENTION_BACKEND="${ATTENTION_BACKEND:-B12X_MLA_SPARSE}" KV_CACHE_DTYPE="${KV_CACHE_DTYPE:-fp8}" +# Calibrated per-layer outer scales for nvfp4_ds_mla KV (empty = disabled). +export VLLM_NVFP4_MLA_SCALES_FILE="${VLLM_NVFP4_MLA_SCALES_FILE:-}" GLM51_PROFILE="${GLM51_PROFILE:-0}" GLM52_CAUSAL_CASCADE="${GLM52_CAUSAL_CASCADE:-0}" GLM52_DSPARK="${GLM52_DSPARK:-0}" diff --git a/tests/model_executor/layers/test_sparse_attn_indexer_b12x.py b/tests/model_executor/layers/test_sparse_attn_indexer_b12x.py index 7f98f8979647..069b10b5af0c 100644 --- a/tests/model_executor/layers/test_sparse_attn_indexer_b12x.py +++ b/tests/model_executor/layers/test_sparse_attn_indexer_b12x.py @@ -3,6 +3,7 @@ import sys import types +from typing import Any, cast import pytest import torch @@ -36,6 +37,44 @@ def get_simultaneous( return tensors +def test_replicated_constructor_uses_single_rank_without_dcp_group(monkeypatch): + def init_module(self): + torch.nn.Module.__init__(self) + + def fail_dcp_group(): + pytest.fail("replicated indexer construction must not query the DCP group") + + monkeypatch.setattr(indexer_mod.CustomOp, "__init__", init_module) + monkeypatch.setattr( + indexer_mod, + "get_current_vllm_config", + lambda: types.SimpleNamespace( + parallel_config=types.SimpleNamespace( + decode_context_parallel_size=4, + cp_kv_cache_interleave_size=1, + ) + ), + ) + monkeypatch.setattr(indexer_mod, "get_dcp_group", fail_dcp_group) + monkeypatch.setattr(indexer_mod, "use_b12x_sparse_indexer", lambda: True) + + indexer = indexer_mod.SparseAttnIndexer( + k_cache=torch.nn.Identity(), + quant_block_size=128, + scale_fmt="ue8m0", + topk_tokens=4, + head_dim=128, + max_model_len=4096, + max_total_seq_len=4096, + topk_indices_buffer=torch.empty((2, 4), dtype=torch.int32), + dcp_replicated=True, + ) + + assert indexer.dcp_replicated is True + assert indexer.dcp_world_size == 1 + assert indexer.dcp_rank == 0 + + def _install_fake_b12x_indexer( monkeypatch, calls: list[tuple], @@ -138,16 +177,17 @@ def build_paged_mqa_schedule_metadata(seq_lens, block_size, num_sms, *, out): out.fill_(7) return out - indexer_mod.PAGED_INDEX_PAGE_SIZE = 64 - indexer_mod.Caps = _Caps - indexer_mod.SOURCE_LAYOUT_PAGED = "paged" - indexer_mod.plan = plan_indexer_scratch - indexer_mod.index_topk_fp8 = index_topk_fp8 - indexer_mod.uses_paged_schedule = uses_paged_mqa_schedule - indexer_mod.plan_paged_schedule = build_paged_mqa_schedule_metadata - - sparkinfer_mod.attention = attention_mod - attention_mod.nsa_indexer = indexer_mod + fake_indexer_mod = cast(Any, indexer_mod) + fake_indexer_mod.PAGED_INDEX_PAGE_SIZE = 64 + fake_indexer_mod.Caps = _Caps + fake_indexer_mod.SOURCE_LAYOUT_PAGED = "paged" + fake_indexer_mod.plan = plan_indexer_scratch + fake_indexer_mod.index_topk_fp8 = index_topk_fp8 + fake_indexer_mod.uses_paged_schedule = uses_paged_mqa_schedule + fake_indexer_mod.plan_paged_schedule = build_paged_mqa_schedule_metadata + + cast(Any, sparkinfer_mod).attention = attention_mod + cast(Any, attention_mod).nsa_indexer = indexer_mod monkeypatch.setitem(sys.modules, "sparkinfer", sparkinfer_mod) monkeypatch.setitem(sys.modules, "sparkinfer.attention", attention_mod) monkeypatch.setitem( @@ -159,7 +199,7 @@ def build_paged_mqa_schedule_metadata(seq_lens, block_size, num_sms, *, out): def _install_fake_b12x_dcp_merge(monkeypatch, run_row_topk, *, world_size: int): tiled_topk_mod = types.ModuleType("sparkinfer.attention.nsa_indexer.tiled_topk") - tiled_topk_mod.run_row_topk = run_row_topk + cast(Any, tiled_topk_mod).run_row_topk = run_row_topk monkeypatch.setitem( sys.modules, "sparkinfer.attention.nsa_indexer.tiled_topk", @@ -411,7 +451,7 @@ def test_b12x_prefill_indexer_requires_packed_contiguous_route(monkeypatch): 64 * 576, ], ) -def test_sparse_attn_indexer_decode_uses_non_shared_b12x_binding( +def test_replicated_decode_skips_dcp_merge_and_keeps_global_topk_ids( monkeypatch, page_stride0, ): @@ -437,6 +477,12 @@ def test_sparse_attn_indexer_decode_uses_non_shared_b12x_binding( lambda: torch.uint8, raising=False, ) + monkeypatch.setattr(indexer_mod, "_dcp_global_topk_requested", lambda: True) + + def fail_dcp_merge(**kwargs): + pytest.fail("replicated indexer must not all-gather top-k candidates") + + monkeypatch.setattr(indexer_mod, "_merge_b12x_dcp_topk", fail_dcp_merge) q_rows = 2 num_heads = 1 @@ -505,6 +551,7 @@ def test_sparse_attn_indexer_decode_uses_non_shared_b12x_binding( ) assert result is topk_indices_buffer + # B12X indexes the full replicated cache, so its logical IDs are already global. assert topk_indices_buffer.tolist() == [[123] * topk, [123] * topk] assert calls == [ ( diff --git a/tests/models/test_dcp_shard_draft_defaults.py b/tests/models/test_dcp_shard_draft_defaults.py index e6319599046c..143ae10ba1aa 100644 --- a/tests/models/test_dcp_shard_draft_defaults.py +++ b/tests/models/test_dcp_shard_draft_defaults.py @@ -3,16 +3,33 @@ from types import SimpleNamespace +import pytest import torch -from vllm.model_executor.models.deepseek_v2 import DeepseekV32IndexerCache +import vllm.model_executor.models.deepseek_v2 as deepseek_v2 +from vllm.model_executor.models.deepseek_v2 import ( + DeepseekV32IndexerCache, + _replicate_indexer_cache_under_dcp, +) -def _vllm_config(num_hidden_layers: int = 78): +def _vllm_config( + num_hidden_layers: int = 78, + dcp_size: int = 4, + pcp_size: int = 1, +): return SimpleNamespace( + use_v2_model_runner=True, cache_config=SimpleNamespace(block_size=256), + compilation_config=SimpleNamespace(static_forward_context={}), + attention_config=SimpleNamespace(backend="B12X_MLA_SPARSE"), model_config=SimpleNamespace( - hf_config=SimpleNamespace(num_hidden_layers=num_hidden_layers) + hf_config=SimpleNamespace(num_hidden_layers=num_hidden_layers), + max_model_len=4096, + ), + parallel_config=SimpleNamespace( + decode_context_parallel_size=dcp_size, + prefill_context_parallel_size=pcp_size, ), ) @@ -26,17 +43,157 @@ def _indexer_cache(layer_id: int = 78): return cache +def _get_indexer_spec(layer_id: int, config): + cache = _indexer_cache(layer_id) + cache.dcp_replicated = _replicate_indexer_cache_under_dcp(cache.prefix, config) + return cache, cache.get_kv_cache_spec(config) + + +def test_target_indexer_replication_defaults_off(monkeypatch): + monkeypatch.delenv("VLLM_DCP_REPLICATE_INDEXER_CACHE", raising=False) + + cache, spec = _get_indexer_spec(12, _vllm_config()) + + assert spec.dcp_replicated is False + assert spec.block_size == cache.cache_config.block_size + + def test_dcp_shard_draft_defaults_to_sharded(monkeypatch): monkeypatch.delenv("VLLM_DCP_SHARD_DRAFT", raising=False) + monkeypatch.setenv("VLLM_DCP_REPLICATE_INDEXER_CACHE", "1") - spec = _indexer_cache().get_kv_cache_spec(_vllm_config()) + cache, spec = _get_indexer_spec(78, _vllm_config()) assert spec.dcp_replicated is False + assert spec.block_size == cache.cache_config.block_size def test_dcp_shard_draft_can_restore_replicated_legacy_mode(monkeypatch): monkeypatch.setenv("VLLM_DCP_SHARD_DRAFT", "0") + monkeypatch.setenv("VLLM_DCP_REPLICATE_INDEXER_CACHE", "1") + + cache, spec = _get_indexer_spec(78, _vllm_config()) + + assert spec.dcp_replicated is True + assert spec.block_size == cache.cache_config.block_size + + +@pytest.mark.parametrize("dcp_size", [2, 4, 6, 8]) +def test_target_indexer_replication_equalizes_global_block_coverage( + monkeypatch, dcp_size: int +): + monkeypatch.setenv("VLLM_DCP_REPLICATE_INDEXER_CACHE", "1") + monkeypatch.setattr(deepseek_v2, "use_b12x_sparse_indexer", lambda: True) - spec = _indexer_cache().get_kv_cache_spec(_vllm_config()) + config = _vllm_config(dcp_size=dcp_size) + cache, spec = _get_indexer_spec(12, config) assert spec.dcp_replicated is True + assert spec.block_size == dcp_size * cache.cache_config.block_size + + +def test_target_indexer_replication_rejects_pcp(monkeypatch): + monkeypatch.setenv("VLLM_DCP_REPLICATE_INDEXER_CACHE", "1") + monkeypatch.setattr(deepseek_v2, "use_b12x_sparse_indexer", lambda: True) + + config = _vllm_config(dcp_size=4, pcp_size=2) + with pytest.raises(NotImplementedError, match="DCP2 through DCP8 with PCP1"): + _replicate_indexer_cache_under_dcp("model.layers.12.indexer.k_cache", config) + + +def test_target_indexer_replication_requires_v2_model_runner(monkeypatch): + monkeypatch.setenv("VLLM_DCP_REPLICATE_INDEXER_CACHE", "1") + + config = _vllm_config() + config.use_v2_model_runner = False + with pytest.raises(NotImplementedError, match="V2 model runner"): + _replicate_indexer_cache_under_dcp("model.layers.12.indexer.k_cache", config) + + +def test_real_cache_construction_replicates_target_but_not_draft(monkeypatch): + monkeypatch.setenv("VLLM_DCP_REPLICATE_INDEXER_CACHE", "1") + monkeypatch.delenv("VLLM_DCP_SHARD_DRAFT", raising=False) + monkeypatch.setattr(deepseek_v2, "use_b12x_sparse_indexer", lambda: True) + + config = _vllm_config() + monkeypatch.setattr(deepseek_v2, "get_current_vllm_config", lambda: config) + + target = DeepseekV32IndexerCache( + head_dim=132, + dtype=torch.uint8, + prefix="model.layers.12.self_attn.indexer.k_cache", + cache_config=config.cache_config, + ) + draft = DeepseekV32IndexerCache( + head_dim=132, + dtype=torch.uint8, + prefix="model.layers.78.self_attn.indexer.k_cache", + cache_config=config.cache_config, + ) + + assert target.dcp_replicated is True + assert target.get_kv_cache_spec(config).dcp_replicated is True + assert draft.dcp_replicated is False + assert draft.get_kv_cache_spec(config).dcp_replicated is False + assert config.compilation_config.static_forward_context == { + target.prefix: target, + draft.prefix: draft, + } + + +def test_indexer_constructor_forwards_cache_replication(monkeypatch): + captured: dict[str, object] = {} + + class FakeCache(torch.nn.Module): + def __init__(self, *args, **kwargs): + super().__init__() + self.dcp_replicated = True + + def fake_sparse_attn_indexer(*args, **kwargs): + captured["args"] = args + captured["kwargs"] = kwargs + return torch.nn.Identity() + + monkeypatch.setattr(deepseek_v2, "DeepseekV32IndexerCache", FakeCache) + monkeypatch.setattr(deepseek_v2, "SparseAttnIndexer", fake_sparse_attn_indexer) + monkeypatch.setattr( + deepseek_v2, + "ReplicatedLinear", + lambda *args, **kwargs: torch.nn.Identity(), + ) + monkeypatch.setattr( + deepseek_v2, + "MergedColumnParallelLinear", + lambda *args, **kwargs: torch.nn.Identity(), + ) + monkeypatch.setattr( + deepseek_v2, + "LayerNorm", + lambda *args, **kwargs: torch.nn.Identity(), + ) + monkeypatch.setattr(deepseek_v2.current_platform, "is_cuda", lambda: False) + + vllm_config = _vllm_config() + model_config = SimpleNamespace( + index_topk=4, + index_n_heads=1, + index_head_dim=128, + qk_rope_head_dim=64, + ) + indexer = deepseek_v2.Indexer( + vllm_config=vllm_config, + config=model_config, + hidden_size=256, + q_lora_rank=128, + quant_config=None, + cache_config=vllm_config.cache_config, + topk_indices_buffer=torch.empty((2, 4), dtype=torch.int32), + prefix="model.layers.12.self_attn.indexer", + ) + + sparse_args = captured["args"] + sparse_kwargs = captured["kwargs"] + assert isinstance(sparse_args, tuple) + assert isinstance(sparse_kwargs, dict) + assert sparse_args[0] is indexer.k_cache + assert sparse_kwargs["dcp_replicated"] is True diff --git a/tests/v1/attention/test_b12x_ckv_prefetch_policy.py b/tests/v1/attention/test_b12x_ckv_prefetch_policy.py new file mode 100644 index 000000000000..a0548ae05128 --- /dev/null +++ b/tests/v1/attention/test_b12x_ckv_prefetch_policy.py @@ -0,0 +1,372 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import gc +from types import SimpleNamespace + +import pytest +import torch + +import vllm.v1.worker.workspace as workspace +from vllm.v1.attention.backends.mla import b12x_sparse_ckv_decode +from vllm.v1.attention.backends.mla.b12x_mla_sparse import ( + B12xMLASparseImpl, + _ckv_prefetch_ring_slots, + _ckv_prefetch_supports_format, + _ckv_prefetch_target_indices, + _ckv_workspace_identity, + _CKVPrefetchStateRegistry, +) + + +@pytest.mark.parametrize( + ("depth", "expected_slots", "expected_targets"), + [ + (0, 1, []), + (1, 2, [2]), + (3, 4, [2, 3, 4]), + ], +) +def test_ckv_prefetch_depth_controls_ring_and_targets( + depth, expected_slots, expected_targets +): + caches = [torch.empty(0) for _ in range(6)] + + assert _ckv_prefetch_ring_slots(depth) == expected_slots + assert _ckv_prefetch_target_indices(1, depth, caches, {}) == expected_targets + + +def test_ckv_prefetch_supports_native_full_record_formats(): + assert _ckv_prefetch_supports_format("nvfp4_ds_mla") + assert _ckv_prefetch_supports_format("fp8_ds_mla") + assert not _ckv_prefetch_supports_format("auto") + + +def test_ckv_prefetch_targets_stop_at_first_unregistered_layer(): + caches = [torch.empty(0), torch.empty(0), torch.empty(0), None, torch.empty(0)] + pending = {2: (object(), 0)} + + assert _ckv_prefetch_target_indices(1, 3, caches, pending) == [] + + +@pytest.mark.parametrize("record_bytes", [368, 432]) +def test_ckv_workspace_reuses_local_staging_across_ring_slots(record_bytes): + impl = object.__new__(B12xMLASparseImpl) + impl._ckv_gather_enabled = True + impl._ckv_workspace_slots = 4 + impl._ckv_local_capacity = 8 + impl._kv_record_bytes = record_bytes + impl.dcp_world_size = 4 + impl.device = torch.device("cpu") + impl._ckv_workspace_nbytes = ( + (1 + impl._ckv_workspace_slots * impl.dcp_world_size) + * impl._ckv_local_capacity + * impl._kv_record_bytes + ) + workspace = torch.empty(impl._ckv_workspace_nbytes, dtype=torch.uint8) + + local_0, gathered_0 = impl._ckv_workspace_views(workspace, 0) + local_3, gathered_3 = impl._ckv_workspace_views(workspace, 3) + + assert local_0.data_ptr() == local_3.data_ptr() + assert gathered_0.shape == gathered_3.shape == (32, record_bytes) + assert gathered_0.data_ptr() != gathered_3.data_ptr() + + +def test_ckv_workspace_rejects_ring_slot_outside_depth(): + impl = object.__new__(B12xMLASparseImpl) + impl._ckv_gather_enabled = True + impl._ckv_workspace_slots = 2 + impl._ckv_local_capacity = 1 + impl._kv_record_bytes = 432 + impl.dcp_world_size = 2 + impl.device = torch.device("cpu") + impl._ckv_workspace_nbytes = 5 * impl._kv_record_bytes + workspace = torch.empty(impl._ckv_workspace_nbytes, dtype=torch.uint8) + + with pytest.raises(ValueError, match="outside"): + impl._ckv_workspace_views(workspace, 2) + + +class _FakeEvent: + def __init__(self): + self.wait_calls = 0 + + def wait(self): + self.wait_calls += 1 + + +class _FakeStream: + def __init__(self): + self.waited_events = [] + + def wait_event(self, event): + self.waited_events.append(event) + + +def test_ckv_bulk_prefetch_ticket_waits_once_for_three_shared_layers(): + registry = _CKVPrefetchStateRegistry() + state = registry.for_workspace(torch.empty(16, dtype=torch.uint8)) + event = _FakeEvent() + stream = _FakeStream() + + ticket = state.register_pending_group({3: 0, 4: 1, 5: 2}, event) + pending = [state.pop_pending_layer(layer) for layer in (3, 4, 5)] + + assert all(item is not None for item in pending) + assert all(item.ticket is ticket for item in pending if item is not None) + for item in pending: + assert item is not None + item.ticket.wait_on_stream_once(stream) + + assert stream.waited_events == [event] + assert state.pending_layers == {} + + +def test_ckv_prefetch_incomplete_step_recovers_with_stream_wait(): + registry = _CKVPrefetchStateRegistry() + state = registry.for_workspace(torch.empty(16, dtype=torch.uint8)) + cache = torch.empty(0) + event = _FakeEvent() + state.register_cache(3, cache) + state.register_pending_group({3: 0}, event) + state.last_layer_idx = 2 + + state.enter_layer(0) + + assert event.wait_calls == 1 + assert state.pending_layers == {} + assert state.layer_caches[3] is cache + assert state.last_layer_idx == 0 + + +def test_ckv_prefetch_first_request_discovers_caches_without_lookahead(): + registry = _CKVPrefetchStateRegistry() + state = registry.for_workspace(torch.empty(16, dtype=torch.uint8)) + caches = [torch.empty(0) for _ in range(4)] + + for layer_idx, cache in enumerate(caches): + state.enter_layer(layer_idx) + state.register_cache(layer_idx, cache) + + assert ( + _ckv_prefetch_target_indices( + layer_idx, 3, state.layer_caches, state.pending_layers + ) + == [] + ) + assert state.gather_stream is None + + state.enter_layer(0) + + assert _ckv_prefetch_target_indices(0, 3, state.layer_caches, {}) == [1, 2, 3] + + +def test_ckv_prefetch_target_and_draft_lifecycles_are_isolated(monkeypatch): + shared_workspace = torch.empty(16, dtype=torch.uint8) + target_lane = ("target", 1) + draft_lane = ("draft", 2) + registry = _CKVPrefetchStateRegistry() + target_state = registry.for_workspace( + shared_workspace, execution_lane_key=target_lane + ) + draft_state = registry.for_workspace( + shared_workspace, execution_lane_key=draft_lane + ) + target_cache = torch.empty(0) + draft_cache = torch.empty(0) + target_event = _FakeEvent() + target_ring = target_state.get_ckv_workspace(64) + draft_ring = draft_state.get_ckv_workspace(64) + + target_state.register_cache(1, target_cache) + target_pending = target_state.register_pending_group({1: 1}, target_event) + draft_state.register_cache(1, draft_cache) + + class SparseState: + def __init__(self, *, layout, device, exchange): + self.layout = layout + self.device = device + self.exchange = exchange + + monkeypatch.setattr(b12x_sparse_ckv_decode, "SparseCKVDecodeState", SparseState) + target_exchange = object() + draft_exchange = object() + layout = object() + target_sparse = target_state.get_sparse_decode_state(layout, target_exchange) + draft_sparse = draft_state.get_sparse_decode_state(layout, draft_exchange) + + assert target_state is not draft_state + assert target_ring.untyped_storage().data_ptr() != ( + draft_ring.untyped_storage().data_ptr() + ) + assert target_state.layer_caches[1] is target_cache + assert target_state.pending_layers[1].ticket is target_pending + assert target_state.pending_layers[1].buf_idx == 1 + assert draft_state.layer_caches[1] is draft_cache + assert draft_state.pending_layers == {} + assert target_sparse.exchange is target_exchange + assert draft_sparse.exchange is draft_exchange + + +def test_ckv_prefetch_lazily_owns_one_stream_per_workspace_lane(monkeypatch): + current_ubatch = [0] + monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: current_ubatch[0]) + monkeypatch.setattr(torch.accelerator, "empty_cache", lambda: None) + manager = workspace.WorkspaceManager(torch.device("cpu"), num_ubatches=2) + (target_workspace,) = manager.get_simultaneous(((16,), torch.uint8)) + (target_workspace_reused,) = manager.get_simultaneous(((16,), torch.uint8)) + current_ubatch[0] = 1 + (draft_workspace,) = manager.get_simultaneous(((16,), torch.uint8)) + + created_streams = [] + + def create_stream(*, device): + stream = SimpleNamespace(device=device) + created_streams.append(stream) + return stream + + monkeypatch.setattr(torch.cuda, "Stream", create_stream) + registry = _CKVPrefetchStateRegistry() + target_state = registry.for_workspace(target_workspace) + target_state_reused = registry.for_workspace(target_workspace_reused) + draft_state = registry.for_workspace(draft_workspace) + + assert created_streams == [] + assert target_state_reused is target_state + assert target_state.get_gather_stream() is target_state.get_gather_stream() + assert draft_state.get_gather_stream() is draft_state.get_gather_stream() + assert target_state.gather_stream is not draft_state.gather_stream + assert len(created_streams) == 2 + + +def test_ckv_prefetch_ring_survives_intervening_workspace_borrow(monkeypatch): + monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: 0) + monkeypatch.setattr(torch.accelerator, "empty_cache", lambda: None) + manager = workspace.WorkspaceManager(torch.device("cpu"), num_ubatches=1) + (lane_workspace,) = manager.get_simultaneous(((256,), torch.uint8)) + registry = _CKVPrefetchStateRegistry() + state = registry.for_workspace(lane_workspace) + + assert state.ckv_workspace is None + ring = state.get_ckv_workspace(64) + ring.fill_(0xA5) + + # WorkspaceManager callers all borrow from offset zero. An intervening + # indexer/MoE scratch allocation must not alias cross-layer CKV state. + (intervening_workspace,) = manager.get_simultaneous(((128,), torch.uint8)) + intervening_workspace.zero_() + + assert ring.untyped_storage().data_ptr() != ( + intervening_workspace.untyped_storage().data_ptr() + ) + assert torch.all(ring == 0xA5) + + +def test_ckv_prefetch_ring_resize_drains_pending_generation(): + registry = _CKVPrefetchStateRegistry() + state = registry.for_workspace(torch.empty(16, dtype=torch.uint8)) + + first_ring = state.get_ckv_workspace(64) + assert state.get_ckv_workspace(64) is first_ring + assert state.ckv_workspace_generation == 1 + + event = _FakeEvent() + state.register_pending_group({1: 0}, event) + resized_ring = state.get_ckv_workspace(128) + + assert resized_ring is not first_ring + assert resized_ring.numel() == 128 + assert state.ckv_workspace_generation == 2 + assert event.wait_calls == 1 + assert state.pending_layers == {} + + +def test_ckv_prefetch_workspace_identity_invalidates_changed_geometry(): + registry = _CKVPrefetchStateRegistry() + workspace_buffer = torch.empty(16, dtype=torch.uint8) + same_geometry = workspace_buffer.view_as(workspace_buffer) + changed_geometry = workspace_buffer[:8] + old_state = registry.for_workspace(workspace_buffer) + event = _FakeEvent() + old_state.register_pending_group({1: 0}, event) + + assert registry.for_workspace(same_geometry) is old_state + + resized_state = registry.for_workspace(changed_geometry) + + assert resized_state is not old_state + assert event.wait_calls == 1 + assert len(registry.states) == 1 + + +def test_ckv_prefetch_resize_retires_only_matching_execution_lane(): + workspace_buffer = torch.empty(257, dtype=torch.uint8) + first_workspace = workspace_buffer[:16] + draft_workspace = first_workspace + resized_workspace = workspace_buffer + target_lane = ("target", 1) + draft_lane = ("draft", 2) + registry = _CKVPrefetchStateRegistry() + cache = torch.empty(0) + first_state = registry.for_workspace( + first_workspace, + 0, + cache, + execution_lane_key=target_lane, + ) + first_state.register_cache(0, cache) + draft_cache = torch.empty(0) + draft_state = registry.for_workspace( + draft_workspace, + 0, + draft_cache, + execution_lane_key=draft_lane, + ) + draft_state.register_cache(0, draft_cache) + event = _FakeEvent() + first_state.register_pending_group({1: 0}, event) + + resized_state = registry.for_workspace( + resized_workspace, + 0, + cache, + execution_lane_key=target_lane, + ) + + assert _ckv_workspace_identity(first_workspace) != _ckv_workspace_identity( + resized_workspace + ) + assert resized_state is not first_state + assert resized_state.layer_caches == [] + assert event.wait_calls == 1 + assert ( + registry.for_workspace(draft_workspace, execution_lane_key=draft_lane) + is draft_state + ) + assert len(registry.states) == 2 + + +def test_ckv_prefetch_registry_retires_released_profile_workspace(): + registry = _CKVPrefetchStateRegistry() + profile_workspace = torch.empty(16, dtype=torch.uint8) + state = registry.for_workspace(profile_workspace) + event = _FakeEvent() + state.register_pending_group({1: 0}, event) + + del profile_workspace + gc.collect() + registry.begin_step() + + assert event.wait_calls == 1 + assert registry.states == {} + + +def test_ckv_gather_uses_capture_fallback_without_reading_prefetch_state( + monkeypatch, +): + impl = object.__new__(B12xMLASparseImpl) + impl._ckv_gather_enabled = True + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + + assert not impl.dcp_prefill_ckv_gather_eligible(SimpleNamespace(), 128) diff --git a/tests/v1/attention/test_b12x_mla_fp8_rope_writer.py b/tests/v1/attention/test_b12x_mla_fp8_rope_writer.py index 0063050b8fa6..ad82540fb910 100644 --- a/tests/v1/attention/test_b12x_mla_fp8_rope_writer.py +++ b/tests/v1/attention/test_b12x_mla_fp8_rope_writer.py @@ -24,6 +24,14 @@ class _WriterInitializationFailure(RuntimeError): pass +class _FakePackageModule(types.ModuleType): + __path__: list[str] + attention: types.ModuleType + _shared: types.ModuleType + mla: types.ModuleType + kv_cache: types.ModuleType + + def _initialize_writer_seam( impl: B12xMLASparseImpl, *, @@ -75,13 +83,13 @@ def _install_fake_writer_package( monkeypatch: pytest.MonkeyPatch, writer, ) -> None: - sparkinfer_module = types.ModuleType("sparkinfer") + sparkinfer_module = _FakePackageModule("sparkinfer") sparkinfer_module.__path__ = [] - attention_module = types.ModuleType("sparkinfer.attention") + attention_module = _FakePackageModule("sparkinfer.attention") attention_module.__path__ = [] - shared_module = types.ModuleType("sparkinfer.attention._shared") + shared_module = _FakePackageModule("sparkinfer.attention._shared") shared_module.__path__ = [] - mla_module = types.ModuleType("sparkinfer.attention._shared.mla") + mla_module = _FakePackageModule("sparkinfer.attention._shared.mla") mla_module.__path__ = [] kv_cache_module = types.ModuleType(_WRITER_MODULE) @@ -236,6 +244,106 @@ def reject_stock_writer(*args, **kwargs): assert actual_scale is k_scale +def test_sparse_decode_current_token_patch_uses_compact_writer( + monkeypatch: pytest.MonkeyPatch, +): + writer_calls = [] + + def public_writer(kv_c_arg, k_pe_arg, kv_cache_arg, slot_mapping_arg, scale_arg): + writer_calls.append( + (kv_c_arg, k_pe_arg, kv_cache_arg, slot_mapping_arg, scale_arg) + ) + + def set_patch_slots(union_indices, request_ids, causal_lens, output): + del union_indices, request_ids, causal_lens + output.copy_(torch.tensor([9, 3], dtype=output.dtype)) + + def reject_stock_writer(*args, **kwargs): + pytest.fail("compact sparse patch fell back to the stock MLA writer") + + monkeypatch.setattr( + b12x_mla_sparse, + "_find_sparse_decode_current_token_slot", + set_patch_slots, + ) + monkeypatch.setattr(ops, "concat_and_cache_mla", reject_stock_writer) + impl = _enabled_impl(public_writer) + impl._ckv_current_chunk_kv_c = torch.zeros((2, 512), dtype=torch.bfloat16) + impl._ckv_current_chunk_kpe = torch.zeros((2, 1, 64), dtype=torch.bfloat16) + + gathered_cache = torch.empty((1, 2, 368), dtype=torch.uint8) + metadata = types.SimpleNamespace(req_id_per_token=torch.tensor([0, 0])) + union_indices = torch.empty((2, 1), dtype=torch.int32) + causal_lens = torch.tensor([1, 2], dtype=torch.int32) + patch_slots = torch.empty(2, dtype=torch.int64) + k_scale = torch.tensor(0.125) + + impl._append_current_token_to_sparse_decode_gathered( + gathered_cache, + metadata, + union_indices, + causal_lens, + patch_slots, + types.SimpleNamespace(_k_scale=k_scale), + ) + + assert len(writer_calls) == 1 + actual_kv_c, actual_k_pe, actual_cache, actual_slots, actual_scale = writer_calls[0] + assert actual_kv_c.data_ptr() == impl._ckv_current_chunk_kv_c.data_ptr() + assert actual_k_pe.shape == (2, 64) + assert actual_cache is gathered_cache + assert torch.equal(actual_slots, torch.tensor([9, 3])) + assert actual_scale is k_scale + + +def test_sparse_decode_current_token_patch_preserves_stock_432_writer( + monkeypatch: pytest.MonkeyPatch, +): + stock_calls = [] + + def stock_writer(*args): + stock_calls.append(args) + + def set_patch_slots(union_indices, request_ids, causal_lens, output): + del union_indices, request_ids, causal_lens + output.fill_(7) + + monkeypatch.setattr( + b12x_mla_sparse, + "_find_sparse_decode_current_token_slot", + set_patch_slots, + ) + monkeypatch.setattr(ops, "concat_and_cache_mla", stock_writer) + impl = object.__new__(B12xMLASparseImpl) + impl._kv_fp8_rope = False + impl.kv_cache_dtype = "nvfp4_ds_mla" + impl._ckv_current_chunk_kv_c = torch.zeros((1, 512), dtype=torch.bfloat16) + impl._ckv_current_chunk_kpe = torch.zeros((1, 1, 64), dtype=torch.bfloat16) + + gathered_cache = torch.empty((1, 1, 432), dtype=torch.uint8) + patch_slots = torch.empty(1, dtype=torch.int64) + k_scale = torch.tensor(0.25) + impl._append_current_token_to_sparse_decode_gathered( + gathered_cache, + types.SimpleNamespace(req_id_per_token=torch.tensor([0])), + torch.empty((1, 1), dtype=torch.int32), + torch.tensor([1], dtype=torch.int32), + patch_slots, + types.SimpleNamespace(_k_scale=k_scale), + ) + + assert len(stock_calls) == 1 + actual_kv_c, actual_k_pe, actual_cache, actual_slots, actual_dtype, actual_scale = ( + stock_calls[0] + ) + assert actual_kv_c.data_ptr() == impl._ckv_current_chunk_kv_c.data_ptr() + assert actual_k_pe.shape == (1, 64) + assert actual_cache is gathered_cache + assert torch.equal(actual_slots, torch.tensor([7])) + assert actual_dtype == "nvfp4_ds_mla" + assert actual_scale is k_scale + + def test_enabled_mode_rejects_wrong_cache_dtype_before_calling_writer(): writer_calls = [] diff --git a/tests/v1/attention/test_b12x_sparse_ckv_decode_policy.py b/tests/v1/attention/test_b12x_sparse_ckv_decode_policy.py new file mode 100644 index 000000000000..7a0dc72f0923 --- /dev/null +++ b/tests/v1/attention/test_b12x_sparse_ckv_decode_policy.py @@ -0,0 +1,929 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import sys +from contextlib import nullcontext +from types import ModuleType +from typing import Any, ClassVar, cast + +import pytest +import torch + +from vllm.distributed import parallel_state +from vllm.v1.attention.backends.mla import b12x_sparse_ckv_decode +from vllm.v1.attention.backends.mla.b12x_mla_sparse import ( + B12xMLASparseImpl, + _sparse_decode_supports_format, +) +from vllm.v1.attention.backends.mla.b12x_sparse_ckv_decode import ( + build_dense_union_remap, + dense_union_remap_reference, + owner_and_local_ordinal, + plan_sparse_ckv_decode, + sparse_decode_batch_eligible, + sparse_decode_prefetch_targets, + try_exchange_sparse_decode_shared_layers, +) + + +def _layout( + *, + dcp: int = 4, + requests: int = 8, + pool: int = 0, + record_bytes: int = 432, +): + return plan_sparse_ckv_decode( + dcp_world_size=dcp, + topk=2048, + rows_per_request=4, + max_requests=requests, + pool_records=pool, + record_bytes=record_bytes, + prefetch_depth=3, + ) + + +@pytest.mark.parametrize(("record_bytes", "kv_fp8_rope"), [(368, True), (432, False)]) +def test_sparse_decode_supports_native_nvfp4_record_formats(record_bytes, kv_fp8_rope): + assert _sparse_decode_supports_format("nvfp4_ds_mla", record_bytes, kv_fp8_rope) + + +@pytest.mark.parametrize( + ("kv_cache_dtype", "record_bytes", "kv_fp8_rope"), + [ + ("nvfp4_ds_mla", 368, False), + ("nvfp4_ds_mla", 432, True), + ("nvfp4_ds_mla", 304, False), + ("nvfp4_ds_mla", 656, False), + ("fp8_ds_mla", 656, False), + ], +) +def test_sparse_decode_rejects_unsupported_record_formats( + kv_cache_dtype, record_bytes, kv_fp8_rope +): + assert not _sparse_decode_supports_format(kv_cache_dtype, record_bytes, kv_fp8_rope) + + +def test_bulk_shared_prefetch_exchanges_three_layers_once_without_new_workspace(): + class Exchange: + def __init__(self): + self.calls = [] + + def exchange_layers(self, **kwargs): + self.calls.append(kwargs) + + layout = _layout() + exchange = Exchange() + records = (object(), object(), object()) + outputs = (object(), object(), object()) + transport_indices = object() + workspace_geometry = (layout.workspace_slots, layout.workspace_bytes) + + used_bulk = try_exchange_sparse_decode_shared_layers( + enabled=True, + transport="ce", + exchange=exchange, + records_by_layer=records, + local_indices_by_destination=transport_indices, + outputs_by_layer=outputs, + active_records=layout.active_records(1), + pool_records=layout.pool_records, + ) + + assert used_bulk + assert exchange.calls == [ + { + "records_by_layer": records, + "local_indices_by_destination": transport_indices, + "outputs_by_layer": outputs, + } + ] + assert (layout.workspace_slots, layout.workspace_bytes) == workspace_geometry + + +def test_bulk_shared_prefetch_is_default_off(monkeypatch): + name = "VLLM_B12X_MLA_SPARSE_DECODE_BULK_PREFETCH" + monkeypatch.delenv(name, raising=False) + + assert b12x_sparse_ckv_decode.envs.environment_variables[name]() is False + + monkeypatch.setenv(name, "1") + assert b12x_sparse_ckv_decode.envs.environment_variables[name]() is True + + +@pytest.mark.parametrize("value", ["enabled", "1 ", ""]) +def test_bulk_shared_prefetch_rejects_invalid_value(monkeypatch, value): + name = "VLLM_B12X_MLA_SPARSE_DECODE_BULK_PREFETCH" + monkeypatch.setenv(name, value) + + with pytest.raises(ValueError, match="Valid options"): + b12x_sparse_ckv_decode.envs.environment_variables[name]() + + +def test_sparse_decode_transport_rejects_unknown_value(monkeypatch): + name = "VLLM_B12X_MLA_SPARSE_DECODE_TRANSPORT" + monkeypatch.setenv(name, "unknown") + + with pytest.raises(ValueError, match="Valid options"): + b12x_sparse_ckv_decode.envs.environment_variables[name]() + + +@pytest.mark.parametrize( + ("enabled", "transport"), + [(False, "ce"), (True, "direct")], +) +def test_bulk_shared_prefetch_preserves_per_layer_fallback(enabled, transport): + class Exchange: + pass + + used_bulk = try_exchange_sparse_decode_shared_layers( + enabled=enabled, + transport=transport, + exchange=Exchange(), + records_by_layer=(object(), object(), object()), + local_indices_by_destination=object(), + outputs_by_layer=(object(), object(), object()), + active_records=8192, + pool_records=65536, + ) + + assert not used_bulk + + +def test_bulk_shared_prefetch_old_b12x_falls_back_cleanly(): + class LegacyCopyExchange: + pass + + used_bulk = try_exchange_sparse_decode_shared_layers( + enabled=True, + transport="ce", + exchange=LegacyCopyExchange(), + records_by_layer=(object(), object(), object()), + local_indices_by_destination=object(), + outputs_by_layer=(object(), object(), object()), + active_records=8192, + pool_records=65536, + ) + + assert not used_bulk + + +def test_bulk_shared_prefetch_requires_exact_depth_three(): + class Exchange: + def exchange_layers(self, **kwargs): + raise AssertionError("partial groups must use the per-layer path") + + used_bulk = try_exchange_sparse_decode_shared_layers( + enabled=True, + transport="ce", + exchange=Exchange(), + records_by_layer=(object(), object()), + local_indices_by_destination=object(), + outputs_by_layer=(object(), object()), + active_records=8192, + pool_records=65536, + ) + + assert not used_bulk + + +def test_bulk_shared_prefetch_falls_back_above_layered_capacity(): + class Exchange: + layered_max_records = 8191 + + def exchange_layers(self, **kwargs): + raise AssertionError("oversized batches must use the per-layer path") + + used_bulk = try_exchange_sparse_decode_shared_layers( + enabled=True, + transport="ce", + exchange=Exchange(), + records_by_layer=(object(), object(), object()), + local_indices_by_destination=object(), + outputs_by_layer=(object(), object(), object()), + active_records=8192, + pool_records=65536, + ) + + assert not used_bulk + + +def test_bulk_shared_prefetch_rejects_pool_overflow(): + class Exchange: + def exchange_layers(self, **kwargs): + raise AssertionError("overflow must fail before exchange") + + with pytest.raises(ValueError, match="exceed pool"): + try_exchange_sparse_decode_shared_layers( + enabled=True, + transport="ce", + exchange=Exchange(), + records_by_layer=(object(), object(), object()), + local_indices_by_destination=object(), + outputs_by_layer=(object(), object(), object()), + active_records=65537, + pool_records=65536, + ) + + +def test_bulk_shared_prefetch_reuses_existing_sparse_state(monkeypatch): + layout = plan_sparse_ckv_decode( + dcp_world_size=4, + topk=2, + rows_per_request=2, + max_requests=1, + pool_records=4, + record_bytes=8, + prefetch_depth=3, + ) + + class Output: + def __init__(self, slot): + self.slot = slot + self.streams = [] + + def record_stream(self, stream): + self.streams.append(stream) + + outputs = [Output(slot) for slot in range(layout.workspace_slots)] + + class Workspace: + def __getitem__(self, key): + slot, records = key + assert records == slice(None, layout.active_records(1), None) + return outputs[slot] + + class Event: + def __init__(self): + self.streams = [] + + def record(self, stream): + self.streams.append(stream) + + class Stream: + def __init__(self): + self.waited_for = [] + + def wait_stream(self, stream): + self.waited_for.append(stream) + + class Exchange: + layered_max_records = layout.pool_records + + def __init__(self): + self.calls = [] + + def exchange_layers(self, **kwargs): + self.calls.append(kwargs) + + class State: + payload_workspace = Workspace() + transport_local_slots = {1: object()} + complete_events = [Event() for _ in range(layout.workspace_slots)] + + class PrefetchState: + def get_sparse_decode_state(self, actual_layout, actual_exchange): + assert actual_layout is layout + assert actual_exchange is exchange + return state + + class Metadata: + num_reqs = 1 + + exchange = Exchange() + state = State() + stream = Stream() + current_stream = object() + monkeypatch.setitem( + b12x_sparse_ckv_decode._SELECTED_RECORD_STREAMS, + id(exchange), + stream, + ) + monkeypatch.setattr( + b12x_sparse_ckv_decode.envs, + "VLLM_B12X_MLA_SPARSE_DECODE_TRANSPORT", + "ce", + ) + monkeypatch.setattr(torch.cuda, "current_stream", lambda: current_stream) + monkeypatch.setattr(torch.cuda, "stream", lambda _stream: nullcontext()) + + impl = object.__new__(B12xMLASparseImpl) + impl._sparse_decode_bulk_prefetch = True + impl._sparse_decode_layout = layout + impl._sparse_decode_exchange = exchange + impl.block_size = 1 + impl.cp_kv_cache_interleave_size = 1 + targets = [] + target_kv_caches = [] + for target_idx in (1, 2, 3): + target_impl = object.__new__(B12xMLASparseImpl) + target_impl._sparse_decode_layout = layout + target_impl._sparse_decode_exchange = exchange + target_impl.block_size = 1 + target_impl.cp_kv_cache_interleave_size = 1 + target_kv = torch.zeros((1, 1, layout.record_bytes), dtype=torch.uint8) + target_kv_caches.append(target_kv) + targets.append((target_idx, target_impl, target_kv)) + + def fail_allocation(*args, **kwargs): + raise AssertionError("bulk prefetch must reuse allocated sparse state") + + for allocator in ("empty", "empty_like", "zeros", "zeros_like"): + monkeypatch.setattr(torch, allocator, fail_allocation) + result = impl._dcp_gather_sparse_decode_shared_layers( + targets, + Metadata(), + PrefetchState(), + ) + + assert result is not None + layer_slots, complete = result + assert layer_slots == {1: 1, 2: 2, 3: 3} + assert stream.waited_for == [current_stream] + assert all(output.streams == [stream] for output in outputs[1:4]) + assert complete is state.complete_events[1] + assert complete.streams == [stream] + assert len(exchange.calls) == 1 + call = exchange.calls[0] + assert all( + actual.data_ptr() == expected.data_ptr() + and actual.dtype == expected.dtype + and actual.shape[-1] == layout.record_bytes + for actual, expected in zip( + call["records_by_layer"], target_kv_caches, strict=True + ) + ) + assert call["outputs_by_layer"] == tuple(outputs[1:4]) + assert call["local_indices_by_destination"] is state.transport_local_slots[1] + + +def _install_fake_b12x_transport(monkeypatch, *, ce_error=None, include_ce=True): + calls = [] + + class InitializationError(RuntimeError): + pass + + class DirectExchange: + init_kwargs: ClassVar[list[dict[str, Any]]] = [] + + @classmethod + def from_process_group(cls, **kwargs): + calls.append("direct") + cls.init_kwargs.append(kwargs) + return cls() + + def close(self): + pass + + class CopyExchange: + init_kwargs: ClassVar[list[dict[str, Any]]] = [] + + @classmethod + def from_process_group(cls, **kwargs): + calls.append("ce") + cls.init_kwargs.append(kwargs) + if ce_error is not None: + raise InitializationError(ce_error) + return cls() + + def close(self): + pass + + package = ModuleType("sparkinfer") + package.__path__ = [] + comm = ModuleType("sparkinfer.comm") + comm.__path__ = [] + pcie = ModuleType("sparkinfer.comm.pcie") + fake_pcie = cast(Any, pcie) + fake_pcie.SelectedRecordExchange = DirectExchange + fake_pcie.SelectedRecordExchangeInitializationError = InitializationError + if include_ce: + fake_pcie.SelectedRecordCopyExchange = CopyExchange + cast(Any, package).comm = comm + cast(Any, comm).pcie = pcie + monkeypatch.setitem(sys.modules, "sparkinfer", package) + monkeypatch.setitem(sys.modules, "sparkinfer.comm", comm) + monkeypatch.setitem(sys.modules, "sparkinfer.comm.pcie", pcie) + return calls, DirectExchange, CopyExchange + + +def _prepare_transport_factory_test(monkeypatch, transport): + monkeypatch.setattr( + b12x_sparse_ckv_decode.envs, + "VLLM_B12X_MLA_SPARSE_DECODE_TRANSPORT", + transport, + ) + monkeypatch.setattr(b12x_sparse_ckv_decode, "_SELECTED_RECORD_EXCHANGES", {}) + monkeypatch.setattr(b12x_sparse_ckv_decode, "_SELECTED_RECORD_STREAMS", {}) + stream = object() + monkeypatch.setattr(torch.cuda, "Stream", lambda **kwargs: stream) + return stream + + +def test_selected_record_transport_auto_prefers_copy_engine(monkeypatch): + stream = _prepare_transport_factory_test(monkeypatch, "auto") + calls, _, CopyExchange = _install_fake_b12x_transport(monkeypatch) + + exchange = b12x_sparse_ckv_decode.get_selected_record_exchange( + process_group=object(), + device=torch.device("cuda", 0), + layout=_layout(), + lane_key=("target", 1), + ) + + assert isinstance(exchange, CopyExchange) + assert calls == ["ce"] + assert b12x_sparse_ckv_decode.get_selected_record_stream(exchange) is stream + + +@pytest.mark.parametrize("transport", ["ce", "direct"]) +def test_selected_record_transport_preserves_368_byte_record_size( + monkeypatch, transport +): + _prepare_transport_factory_test(monkeypatch, transport) + calls, DirectExchange, CopyExchange = _install_fake_b12x_transport(monkeypatch) + layout = _layout(record_bytes=368) + + exchange = b12x_sparse_ckv_decode.get_selected_record_exchange( + process_group=object(), + device=torch.device("cuda", 0), + layout=layout, + lane_key=("target", 1), + ) + + expected_class = CopyExchange if transport == "ce" else DirectExchange + assert isinstance(exchange, expected_class) + assert calls == [transport] + init_kwargs = expected_class.init_kwargs[-1] + assert init_kwargs["record_bytes"] == 368 + assert init_kwargs["max_records"] == layout.pool_records + + +def test_selected_record_transport_isolates_368_and_432_byte_abis(monkeypatch): + _prepare_transport_factory_test(monkeypatch, "direct") + calls, _, _ = _install_fake_b12x_transport(monkeypatch) + process_group = object() + + record_432 = b12x_sparse_ckv_decode.get_selected_record_exchange( + process_group=process_group, + device=torch.device("cuda", 0), + layout=_layout(record_bytes=432), + lane_key=("target", 1), + ) + record_368 = b12x_sparse_ckv_decode.get_selected_record_exchange( + process_group=process_group, + device=torch.device("cuda", 0), + layout=_layout(record_bytes=368), + lane_key=("target", 1), + ) + + assert record_432 is not record_368 + assert calls == ["direct", "direct"] + + +def test_selected_record_transport_auto_falls_back_to_direct(monkeypatch): + _prepare_transport_factory_test(monkeypatch, "auto") + calls, DirectExchange, _ = _install_fake_b12x_transport( + monkeypatch, ce_error="CE unavailable" + ) + + exchange = b12x_sparse_ckv_decode.get_selected_record_exchange( + process_group=object(), + device=torch.device("cuda", 0), + layout=_layout(), + lane_key=("target", 1), + ) + + assert isinstance(exchange, DirectExchange) + assert calls == ["ce", "direct"] + + +def test_selected_record_transport_auto_supports_older_b12x(monkeypatch): + _prepare_transport_factory_test(monkeypatch, "auto") + calls, DirectExchange, _ = _install_fake_b12x_transport( + monkeypatch, include_ce=False + ) + + exchange = b12x_sparse_ckv_decode.get_selected_record_exchange( + process_group=object(), + device=torch.device("cuda", 0), + layout=_layout(), + lane_key=("target", 1), + ) + + assert isinstance(exchange, DirectExchange) + assert calls == ["direct"] + + +def test_selected_record_transport_direct_never_constructs_copy_engine(monkeypatch): + _prepare_transport_factory_test(monkeypatch, "direct") + calls, DirectExchange, _ = _install_fake_b12x_transport(monkeypatch) + + exchange = b12x_sparse_ckv_decode.get_selected_record_exchange( + process_group=object(), + device=torch.device("cuda", 0), + layout=_layout(), + lane_key=("target", 1), + ) + + assert isinstance(exchange, DirectExchange) + assert calls == ["direct"] + + +def test_selected_record_transport_isolates_target_and_draft_lanes(monkeypatch): + _prepare_transport_factory_test(monkeypatch, "direct") + calls, _, _ = _install_fake_b12x_transport(monkeypatch) + monkeypatch.setattr(torch.cuda, "Stream", lambda **kwargs: object()) + process_group = object() + layout = _layout() + + target = b12x_sparse_ckv_decode.get_selected_record_exchange( + process_group=process_group, + device=torch.device("cuda", 0), + layout=layout, + lane_key=("target", 1), + ) + draft = b12x_sparse_ckv_decode.get_selected_record_exchange( + process_group=process_group, + device=torch.device("cuda", 0), + layout=layout, + lane_key=("draft", 1), + ) + + assert target is not draft + assert b12x_sparse_ckv_decode.get_selected_record_stream(target) is not ( + b12x_sparse_ckv_decode.get_selected_record_stream(draft) + ) + assert calls == ["direct", "direct"] + + +def test_selected_record_transport_strict_ce_does_not_fallback(monkeypatch): + _prepare_transport_factory_test(monkeypatch, "ce") + calls, _, _ = _install_fake_b12x_transport(monkeypatch, ce_error="CE unavailable") + + with pytest.raises(RuntimeError, match="strict copy-engine"): + b12x_sparse_ckv_decode.get_selected_record_exchange( + process_group=object(), + device=torch.device("cuda", 0), + layout=_layout(), + lane_key=("target", 1), + ) + + assert calls == ["ce"] + + +def test_selected_record_transport_rejects_unknown_mode(monkeypatch): + _prepare_transport_factory_test(monkeypatch, "mystery") + + with pytest.raises(ValueError, match="must be auto, ce, or direct"): + b12x_sparse_ckv_decode.get_selected_record_exchange( + process_group=object(), + device=torch.device("cuda", 0), + layout=_layout(), + lane_key=("target", 1), + ) + + +@pytest.mark.parametrize("record_bytes", [368, 432]) +@pytest.mark.parametrize("dcp", [2, 3, 4, 5, 6, 7, 8]) +def test_sparse_decode_layout_supports_dcp_two_through_eight(dcp, record_bytes): + layout = _layout(dcp=dcp, record_bytes=record_bytes) + + assert layout.dcp_world_size == dcp + assert layout.per_request_capacity == 8192 + assert layout.pool_records == 65536 + assert layout.max_fast_requests == 8 + assert layout.record_bytes == record_bytes + assert layout.workspace_slots == 4 + + +@pytest.mark.parametrize("dcp", [1, 9]) +def test_sparse_decode_layout_rejects_unsupported_dcp(dcp): + with pytest.raises(ValueError, match="DCP world sizes 2-8"): + _layout(dcp=dcp) + + +@pytest.mark.parametrize("concurrency", [1, 2, 4, 8]) +def test_sparse_decode_batch_policy_accepts_c1_through_c8(concurrency): + layout = _layout() + rows = concurrency * layout.rows_per_request + + assert sparse_decode_batch_eligible( + layout, + num_requests=concurrency, + num_rows=rows, + max_query_len=layout.rows_per_request, + num_actual_tokens=rows, + has_required_metadata=True, + ) + assert layout.active_records(concurrency) == concurrency * 8192 + + +def test_sparse_decode_pool_overflow_falls_back_before_transport(): + layout = _layout(pool=3 * 8192) + + assert layout.max_fast_requests == 3 + assert sparse_decode_batch_eligible( + layout, + num_requests=3, + num_rows=12, + max_query_len=4, + num_actual_tokens=12, + has_required_metadata=True, + ) + assert not sparse_decode_batch_eligible( + layout, + num_requests=4, + num_rows=16, + max_query_len=4, + num_actual_tokens=16, + has_required_metadata=True, + ) + with pytest.raises(ValueError, match="outside sparse capacity"): + layout.active_records(4) + + +def test_mtp3_union_is_per_sequence_dense_stable_and_exact(): + indices = torch.tensor( + [ + [ + [10, 11, 10, -1], + [12, 11, 13, 10], + [13, 14, 12, -1], + [10, 15, 14, 15], + ], + [ + [10, 20, 10, -1], + [21, 20, 22, 10], + [22, 23, 21, -1], + [10, 24, 23, 24], + ], + ], + dtype=torch.int32, + ) + + union, remap, counts = dense_union_remap_reference(indices) + + assert union[0, :6].tolist() == [10, 11, 12, 13, 14, 15] + assert union[1, :6].tolist() == [10, 20, 21, 22, 23, 24] + assert counts.tolist() == [6, 6] + for request in range(indices.shape[0]): + for row in range(indices.shape[1]): + for column in range(indices.shape[2]): + token = int(indices[request, row, column]) + slot = int(remap[request, row, column]) + if token < 0: + assert slot == -1 + else: + assert int(union[request, slot]) == token + + +def test_destination_unions_and_remaps_are_independent_and_exact(): + indices = torch.tensor( + [ + [[[1, 2], [2, 3]], [[10, 11], [11, 12]]], + [[[4, 5], [5, 6]], [[10, 13], [13, 14]]], + ], + dtype=torch.int32, + ) + + union, remap, counts = dense_union_remap_reference(indices) + + assert counts.tolist() == [[3, 3], [3, 3]] + assert union[0, 0, :3].tolist() == [1, 2, 3] + assert union[1, 0, :3].tolist() == [4, 5, 6] + assert union[0, 1, :3].tolist() == [10, 11, 12] + assert union[1, 1, :3].tolist() == [10, 13, 14] + reconstructed = torch.gather( + union.unsqueeze(-2).expand(*indices.shape[:-2], indices.shape[-2], -1), + -1, + remap.to(torch.int64), + ) + assert torch.equal(reconstructed, indices) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_cuda_union_accepts_runtime_rows_below_reserved_mtp_rows(): + device = torch.device("cuda") + indices = torch.tensor([[[[1, 2]]], [[[3, 4]]]], dtype=torch.int32, device=device) + union = torch.empty((2, 2, 8), dtype=torch.int32, device=device) + remap = torch.empty((2, 2, 4, 2), dtype=torch.int32, device=device) + counts = torch.empty((2, 2), dtype=torch.int32, device=device) + hash_keys = torch.empty((2, 2, 16), dtype=torch.int32, device=device) + hash_first = torch.empty_like(hash_keys) + first_to_dense = torch.empty_like(union) + + build_dense_union_remap( + indices, + union, + remap, + counts, + hash_keys, + hash_first, + first_to_dense, + num_requests=1, + ) + torch.accelerator.synchronize() + + assert counts[:, 0].cpu().tolist() == [2, 2] + assert union[0, 0, :2].cpu().tolist() == [1, 2] + assert union[1, 0, :2].cpu().tolist() == [3, 4] + assert remap[:, 0, 0, :].cpu().tolist() == [[0, 1], [0, 1]] + + +def test_union_preflight_caches_rank_consistent_success(monkeypatch): + process_group = object() + loads = 0 + reductions = 0 + monkeypatch.setattr(b12x_sparse_ckv_decode, "_UNION_PREFLIGHT_RESULTS", {}) + + def load_extension(): + nonlocal loads + loads += 1 + + def all_reduce(status, *, op, group): + nonlocal reductions + reductions += 1 + assert int(status.item()) == 1 + assert op is torch.distributed.ReduceOp.MIN + assert group is process_group + + monkeypatch.setattr( + b12x_sparse_ckv_decode, "preload_dense_union_extension", load_extension + ) + monkeypatch.setattr(b12x_sparse_ckv_decode.dist, "all_reduce", all_reduce) + + for _ in range(2): + b12x_sparse_ckv_decode.preload_dense_union_extension_consistently( + process_group, torch.device("cpu") + ) + + assert loads == 1 + assert reductions == 1 + + +def test_union_preflight_propagates_remote_rank_failure(monkeypatch): + process_group = object() + reductions = 0 + monkeypatch.setattr(b12x_sparse_ckv_decode, "_UNION_PREFLIGHT_RESULTS", {}) + + def all_reduce(status, *, op, group): + nonlocal reductions + reductions += 1 + status.zero_() + + monkeypatch.setattr( + b12x_sparse_ckv_decode, "preload_dense_union_extension", lambda: None + ) + monkeypatch.setattr(b12x_sparse_ckv_decode.dist, "all_reduce", all_reduce) + + for _ in range(2): + with pytest.raises(RuntimeError, match="another DCP rank"): + b12x_sparse_ckv_decode.preload_dense_union_extension_consistently( + process_group, torch.device("cpu") + ) + + assert reductions == 1 + + +def test_selected_record_cleanup_closes_once_and_clears_registries(monkeypatch): + class Exchange: + def __init__(self): + self.close_calls = 0 + + def close(self): + self.close_calls += 1 + + first = Exchange() + second = Exchange() + monkeypatch.setattr( + b12x_sparse_ckv_decode, + "_SELECTED_RECORD_EXCHANGES", + {("target",): first, ("target-alias",): first, ("draft",): second}, + ) + monkeypatch.setattr( + b12x_sparse_ckv_decode, + "_SELECTED_RECORD_STREAMS", + {id(first): object(), id(second): object()}, + ) + monkeypatch.setattr( + b12x_sparse_ckv_decode, + "_UNION_PREFLIGHT_RESULTS", + {(1, "cuda", 0): (True, "")}, + ) + + b12x_sparse_ckv_decode.close_selected_record_exchanges() + + assert first.close_calls == 1 + assert second.close_calls == 1 + assert b12x_sparse_ckv_decode._SELECTED_RECORD_EXCHANGES == {} + assert b12x_sparse_ckv_decode._SELECTED_RECORD_STREAMS == {} + assert b12x_sparse_ckv_decode._UNION_PREFLIGHT_RESULTS == {} + + +def test_model_parallel_cleanup_runs_before_group_destroy(monkeypatch): + events = [] + + class Group: + def destroy(self): + events.append("group") + + monkeypatch.setattr( + parallel_state, + "_MODEL_PARALLEL_CLEANUP_HOOKS", + [lambda: events.append("cleanup")], + ) + monkeypatch.setattr(parallel_state, "_TP", Group()) + for name in ( + "_DCP", + "_QUERY_SPLIT", + "_DCP_CKV_PREFETCH", + "_PCP", + "_PP", + "_DP", + "_EP", + "_EPLB", + ): + monkeypatch.setattr(parallel_state, name, None) + + parallel_state.destroy_model_parallel() + + assert events == ["cleanup", "group"] + + +@pytest.mark.parametrize("concurrency", [2, 4, 8]) +def test_multi_sequence_union_never_aliases_request_local_positions(concurrency): + rows = [] + for request in range(concurrency): + base = request * 1000 + rows.append( + [ + [0, base + 1, base + 2], + [base + 2, base + 3, 0], + [base + 3, base + 4, base + 1], + [base + 4, base + 5, 0], + ] + ) + indices = torch.tensor(rows, dtype=torch.int32) + + union, remap, counts = dense_union_remap_reference(indices) + + assert counts.tolist() == [6] * concurrency + pooled_union = union.reshape(-1) + for request in range(concurrency): + reconstructed = torch.gather( + union[request].expand(indices.shape[1], -1), + 1, + remap[request].to(torch.int64), + ) + assert torch.equal(reconstructed, indices[request]) + pooled_slots = remap[request].to(torch.int64) + request * union.shape[-1] + assert torch.equal(pooled_union[pooled_slots], indices[request]) + + +def test_shared_layer_prefetch_is_exactly_full_to_s1_s2_s3(): + # GLM: Full 0-2, Shared 3-5, then Full 6. + emits_topk = [True, True, True, False, False, False, True, False] + + assert sparse_decode_prefetch_targets(2, 3, emits_topk) == [3, 4, 5] + assert sparse_decode_prefetch_targets(1, 3, emits_topk) == [] + assert sparse_decode_prefetch_targets(6, 3, emits_topk) == [7] + + +@pytest.mark.parametrize("dcp", [2, 3, 4, 5, 6, 7, 8]) +@pytest.mark.parametrize("interleave", [1, 2, 4]) +def test_owner_arithmetic_round_trips_global_token(dcp, interleave): + for logical_token in range(1024): + owner, local = owner_and_local_ordinal( + logical_token, + dcp_world_size=dcp, + interleave=interleave, + ) + group = local // interleave + lane = local % interleave + reconstructed = (group * dcp + owner) * interleave + lane + + assert 0 <= owner < dcp + assert reconstructed == logical_token + + +def test_sparse_decode_requires_uniform_rows_and_complete_metadata(): + layout = _layout() + + assert not sparse_decode_batch_eligible( + layout, + num_requests=2, + num_rows=7, + max_query_len=4, + num_actual_tokens=7, + has_required_metadata=True, + ) + assert not sparse_decode_batch_eligible( + layout, + num_requests=2, + num_rows=8, + max_query_len=4, + num_actual_tokens=8, + has_required_metadata=False, + ) diff --git a/tests/v1/attention/test_indexer_dcp_localize.py b/tests/v1/attention/test_indexer_dcp_localize.py index 3fafe2c4fdef..ab01ad01434d 100644 --- a/tests/v1/attention/test_indexer_dcp_localize.py +++ b/tests/v1/attention/test_indexer_dcp_localize.py @@ -1,24 +1,130 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + import pytest import torch import vllm.model_executor.layers.sparse_attn_indexer as sparse_indexer +import vllm.v1.attention.backends.mla.indexer as indexer_backend from vllm.platforms import current_platform from vllm.utils.import_utils import has_cutedsl -from vllm.v1.attention.backends.mla.indexer import build_prefill_chunk_metadata +from vllm.v1.attention.backend import CommonAttentionMetadata +from vllm.v1.attention.backends.mla.indexer import ( + DeepseekV32IndexerMetadataBuilder, + build_prefill_chunk_metadata, + get_indexer_max_num_blocks_per_req, +) from vllm.v1.attention.backends.mla.sparse_utils import ( triton_filter_and_convert_dcp_index, ) from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens from vllm.v1.attention.ops.common import CPTritonContext, correct_attn_out +from vllm.v1.kv_cache_interface import MLAAttentionSpec def _local_count(length: int, rank: int, world: int, interleave: int) -> int: return sum(1 for pos in range(length) if (pos // interleave) % world == rank) +@pytest.mark.parametrize("dcp_size", [2, 4, 6, 8]) +def test_replicated_indexer_metadata_covers_full_context(dcp_size: int): + max_model_len = 131_071 + storage_block_size = 64 + manager_block_size = storage_block_size * dcp_size + + sharded_blocks = get_indexer_max_num_blocks_per_req( + max_model_len=max_model_len, + block_size=storage_block_size, + configured_cp_world_size=dcp_size, + dcp_replicated=False, + ) + replicated_blocks = get_indexer_max_num_blocks_per_req( + max_model_len=max_model_len, + block_size=manager_block_size, + configured_cp_world_size=dcp_size, + dcp_replicated=True, + ) + expected_blocks = (max_model_len + manager_block_size - 1) // manager_block_size + + assert sharded_blocks == expected_blocks + assert replicated_blocks == expected_blocks + + +def test_replicated_builder_uses_global_lengths_without_dcp_localization(monkeypatch): + def fail_dcp_path(*args, **kwargs): + pytest.fail("replicated indexer metadata must not enter DCP localization") + + monkeypatch.setattr(indexer_backend, "get_dcp_group", fail_dcp_path) + monkeypatch.setattr(indexer_backend, "num_compute_units", lambda device: 4) + monkeypatch.setattr(indexer_backend.envs, "VLLM_USE_B12X_SPARSE_INDEXER", True) + + spec = MLAAttentionSpec( + block_size=1024, + num_kv_heads=1, + head_size=132, + dtype=torch.uint8, + dcp_replicated=True, + ) + config = SimpleNamespace( + scheduler_config=SimpleNamespace( + max_num_batched_tokens=8, + max_num_seqs=4, + ), + parallel_config=SimpleNamespace( + decode_context_parallel_size=4, + prefill_context_parallel_size=1, + cp_kv_cache_interleave_size=1, + ), + speculative_config=None, + attention_config=SimpleNamespace(use_fp4_indexer_cache=False), + model_config=SimpleNamespace(max_model_len=4096), + ) + builder = DeepseekV32IndexerMetadataBuilder( + kv_cache_spec=spec, + layer_names=["model.layers.12.self_attn.indexer.k_cache"], + vllm_config=config, + device=torch.device("cpu"), + ) + monkeypatch.setattr(builder, "_dcp_localize_decode_seq_lens", fail_dcp_path) + monkeypatch.setattr( + builder, + "_maybe_build_b12x_schedule_metadata", + lambda *args, **kwargs: None, + ) + + global_seq_lens = torch.tensor([9], dtype=torch.int32) + local_seq_lens = torch.tensor([3], dtype=torch.int32) + common_metadata = CommonAttentionMetadata( + query_start_loc=torch.tensor([0, 1], dtype=torch.int32), + query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.int32), + seq_lens=global_seq_lens, + num_reqs=1, + num_actual_tokens=1, + max_query_len=1, + max_seq_len=9, + block_table_tensor=torch.zeros((1, 1), dtype=torch.int32), + slot_mapping=torch.tensor([8], dtype=torch.int64), + dcp_local_seq_lens=local_seq_lens, + dcp_local_seq_lens_cpu=local_seq_lens.clone(), + seq_lens_cpu_upper_bound=global_seq_lens.clone(), + _seq_lens_cpu=global_seq_lens.clone(), + ) + + metadata = builder.build(0, common_metadata) + + assert builder.dcp_world_size == 1 + assert builder.dcp_rank == 0 + assert metadata.seq_lens.tolist() == [9] + assert metadata.slot_mapping.tolist() == [8] + assert metadata.decode is not None + assert metadata.decode.seq_lens.tolist() == [[9]] + assert metadata.decode.max_seq_len == 9 + assert metadata.decode.active_width.tolist() == [9] + assert metadata.decode.global_seq_lens is None + + def _global_to_local_indices( global_indices: torch.Tensor, rank: int, @@ -89,6 +195,16 @@ def _attention_from_indices( v: torch.Tensor, indices: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: + if k.shape[0] == 0: + return ( + torch.zeros_like(q), + torch.full( + (q.shape[0],), + float("-inf"), + dtype=q.dtype, + device=q.device, + ), + ) valid = indices >= 0 safe_indices = indices.clamp_min(0) selected_k = k[safe_indices] @@ -253,7 +369,7 @@ def _merge_local_topks_global_with_fake_dcp( sparse_indexer.get_dcp_group = original_get_dcp_group -@pytest.mark.parametrize("world", [1, 2, 4]) +@pytest.mark.parametrize("world", [1, 2, 4, 6, 8]) @pytest.mark.parametrize("interleave", [1, 2, 4]) def test_get_dcp_local_seq_lens_matches_naive(world: int, interleave: int): seq_lens = torch.arange(0, 33, dtype=torch.int32) @@ -341,9 +457,11 @@ def test_get_dcp_local_seq_lens_must_run_after_decode_expansion(): @pytest.mark.parametrize("interleave", [1, 2]) -def test_sparse_dcp_attention_matches_global_topk_attention(interleave: int): +@pytest.mark.parametrize("world", [2, 4, 6, 8]) +def test_sparse_dcp_attention_matches_global_topk_attention( + interleave: int, world: int +): torch.manual_seed(0) - world = 2 topk = 3 num_queries = 4 max_seq_len = 13 diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 0e791d57b728..2a7297462c04 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -1683,6 +1683,170 @@ def test_resolve_kv_cache_block_sizes_mixed_dcp_replicated_groups(): assert hash_block_size == 64 +@pytest.mark.parametrize( + ("model_version", "record_bytes"), [(None, 432), ("glm_fp8_rope", 368)] +) +def test_replicated_mla_uses_lockstep_pool_capacity_and_contiguous_tensors( + model_version, record_bytes +): + vllm_config = SimpleNamespace( + model_config=SimpleNamespace(max_model_len=262144), + parallel_config=SimpleNamespace( + decode_context_parallel_size=4, + prefill_context_parallel_size=1, + ), + cache_config=SimpleNamespace(num_gpu_blocks_override=None), + kv_transfer_config=None, + ) + specs: dict[str, KVCacheSpec] = { + f"target.{i}": MLAAttentionSpec( + block_size=64, + num_kv_heads=1, + head_size=576, + dtype=torch.uint8, + cache_dtype_str="nvfp4_ds_mla", + model_version=model_version, + ) + for i in range(3) + } + specs.update( + { + f"indexer.{i}": MLAAttentionSpec( + block_size=256, + num_kv_heads=1, + head_size=132, + dtype=torch.uint8, + dcp_replicated=True, + ) + for i in range(2) + } + ) + + grouped_specs = kv_cache_utils.group_and_unify_kv_cache_specs(specs, 4, 1) + assert grouped_specs is not None + groups = kv_cache_utils._get_kv_cache_groups_uniform_groups(grouped_specs) + assert all( + isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) for group in groups + ) + assert kv_cache_utils._use_lockstep_mla_allocation(groups, 4, 1) + + bytes_per_pool_block = 3 * (64 * record_bytes) + 2 * (256 * 132) + request_blocks = 262144 // 256 + required_memory = bytes_per_pool_block * request_blocks + kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups( + vllm_config, + groups, + available_memory=required_memory * 2, + ) + + assert kv_cache_config.num_blocks == request_blocks * 2 + assert all(tensor.block_stride == 0 for tensor in kv_cache_config.kv_cache_tensors) + assert sum(tensor.size for tensor in kv_cache_config.kv_cache_tensors) == ( + required_memory * 2 + ) + assert ( + kv_cache_utils._max_memory_usage_bytes_from_groups(vllm_config, groups) + == required_memory + ) + assert get_max_concurrency_for_kv_cache_config( + vllm_config, kv_cache_config + ) == pytest.approx(2.0) + + +def test_lockstep_mla_predicate_rejects_nonmatching_layouts(): + sharded = MLAAttentionSpec( + block_size=64, + num_kv_heads=1, + head_size=128, + dtype=torch.float32, + ) + replicated = MLAAttentionSpec( + block_size=256, + num_kv_heads=1, + head_size=128, + dtype=torch.float32, + dcp_replicated=True, + ) + groups = [ + KVCacheGroupSpec(["target"], sharded), + KVCacheGroupSpec(["indexer"], replicated), + ] + assert not kv_cache_utils._use_lockstep_mla_allocation(groups, 1, 1) + + mismatched = MLAAttentionSpec( + block_size=64, + num_kv_heads=1, + head_size=128, + dtype=torch.float32, + dcp_replicated=True, + ) + mismatched_groups = [groups[0], KVCacheGroupSpec(["indexer"], mismatched)] + assert not kv_cache_utils._use_lockstep_mla_allocation(mismatched_groups, 4, 1) + assert ( + kv_cache_utils.group_and_unify_kv_cache_specs( + {"target": sharded, "indexer": mismatched}, 4, 1 + ) + is None + ) + + non_mla = FullAttentionSpec( + block_size=256, + num_kv_heads=1, + head_size=128, + dtype=torch.float32, + dcp_replicated=True, + ) + assert not kv_cache_utils._use_lockstep_mla_allocation( + [groups[0], KVCacheGroupSpec(["draft"], non_mla)], 4, 1 + ) + + +def test_lockstep_mla_equal_page_sizes_use_distinct_tensors(): + sharded = MLAAttentionSpec( + block_size=64, + num_kv_heads=1, + head_size=128, + dtype=torch.bfloat16, + ) + replicated = MLAAttentionSpec( + block_size=256, + num_kv_heads=1, + head_size=32, + dtype=torch.bfloat16, + dcp_replicated=True, + ) + assert sharded.page_size_bytes == replicated.page_size_bytes + groups = [ + KVCacheGroupSpec(["target"], sharded), + KVCacheGroupSpec(["indexer"], replicated), + ] + vllm_config = SimpleNamespace( + cache_config=SimpleNamespace(num_gpu_blocks_override=None), + parallel_config=SimpleNamespace( + decode_context_parallel_size=4, + prefill_context_parallel_size=1, + ), + kv_transfer_config=None, + ) + bytes_per_pool_block = sharded.page_size_bytes + replicated.page_size_bytes + + kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups( + vllm_config, + groups, + available_memory=3 * bytes_per_pool_block, + ) + + assert kv_cache_config.num_blocks == 3 + assert [tensor.shared_by for tensor in kv_cache_config.kv_cache_tensors] == [ + ["target"], + ["indexer"], + ] + assert [tensor.size for tensor in kv_cache_config.kv_cache_tensors] == [ + 3 * sharded.page_size_bytes, + 3 * replicated.page_size_bytes, + ] + + def test_dsv4_engine_capacity_uses_worker_kv_cache_config(): from vllm.v1.engine.core import EngineCore diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index dd89769315b9..74fdcfccb56f 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -177,6 +177,183 @@ def make_kv_cache_config_hybrid_model( ) +def make_lockstep_mla_manager(num_blocks: int = 5) -> KVCacheManager: + global_block_size = 256 + kv_cache_config = KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["target"], + MLAAttentionSpec( + block_size=global_block_size // 4, + num_kv_heads=1, + head_size=432, + dtype=torch.uint8, + ), + ), + KVCacheGroupSpec( + ["indexer"], + MLAAttentionSpec( + block_size=global_block_size, + num_kv_heads=1, + head_size=132, + dtype=torch.uint8, + dcp_replicated=True, + ), + ), + ], + ) + return KVCacheManager( + kv_cache_config=kv_cache_config, + max_model_len=4 * global_block_size, + scheduler_block_size=global_block_size, + hash_block_size=global_block_size, + enable_caching=True, + dcp_world_size=4, + ) + + +def test_mixed_mla_groups_share_block_ids_hashes_and_eviction_order(): + block_size = 256 + manager = make_lockstep_mla_manager() + assert manager.coordinator.lockstep_mla_allocations + request = make_request( + "lockstep", + list(range(4 * block_size)), + block_size, + sha256, + ) + + new_blocks = manager.allocate_slots( + request, + num_new_tokens=4 * block_size, + full_sequence_must_fit=True, + ) + assert new_blocks is not None + target_ids, indexer_ids = manager.get_block_ids(request.request_id) + assert target_ids == indexer_ids + assert new_blocks.get_block_ids() == (target_ids, target_ids) + assert manager.block_pool.get_num_free_blocks() == 0 + assert all( + manager.block_pool.blocks[block_id].ref_cnt == 2 for block_id in target_ids + ) + + for block_hash, block_id in zip(request.block_hashes, target_ids): + group_keys = [ + make_block_hash_with_group_id(block_hash, group_id) for group_id in range(2) + ] + assert group_keys[0] != group_keys[1] + cached = manager.block_pool.get_cached_block(block_hash, [0, 1]) + assert cached is not None + assert [block.block_id for block in cached] == [block_id, block_id] + + manager.free(request) + assert manager.block_pool.get_num_free_blocks() == 4 + assert all( + manager.block_pool.blocks[block_id].ref_cnt == 0 for block_id in target_ids + ) + free_ids = [ + block.block_id + for block in manager.block_pool.free_block_queue.get_all_free_blocks() + ] + assert free_ids == target_ids[::-1] + + replay = make_request( + "lockstep-replay", + list(range(4 * block_size)), + block_size, + sha256, + ) + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(replay) + assert num_computed_tokens == 3 * block_size + assert computed_blocks.get_block_ids() == ( + target_ids[:3], + target_ids[:3], + ) + + replay_new_blocks = manager.allocate_slots( + replay, + num_new_tokens=block_size, + num_new_computed_tokens=num_computed_tokens, + new_computed_blocks=computed_blocks, + ) + assert replay_new_blocks is not None + assert manager.get_block_ids(replay.request_id) == (target_ids, target_ids) + assert all( + manager.block_pool.blocks[block_id].ref_cnt == 2 for block_id in target_ids + ) + manager.free(replay) + assert [ + block.block_id + for block in manager.block_pool.free_block_queue.get_all_free_blocks() + ] == target_ids[::-1] + + evicted = manager.block_pool.get_new_blocks(1)[0] + assert evicted.block_id == target_ids[-1] + assert manager.block_pool.get_cached_block(request.block_hashes[-1], [0]) is None + assert manager.block_pool.get_cached_block(request.block_hashes[-1], [1]) is None + + +def test_lockstep_mla_rejects_external_computed_blocks(): + manager = make_lockstep_mla_manager() + + with pytest.raises(NotImplementedError, match="External KV loads"): + manager.coordinator.allocate_new_computed_blocks( + "external", + ([], []), + num_local_computed_tokens=0, + num_external_computed_tokens=256, + ) + + assert manager.get_block_ids("external") == ([], []) + + +def test_lockstep_group_hashes_promote_partial_block_together(): + hash_block_size = 64 + block_size = 4 * hash_block_size + pool = BlockPool( + num_gpu_blocks=2, + enable_caching=True, + hash_block_size=hash_block_size, + ) + block = pool.get_new_blocks(1)[0] + request = make_request( + "promotion", + list(range(block_size)), + hash_block_size, + sha256, + ) + + for group_id in range(2): + pool.cache_partial_block( + request=request, + block=block, + num_tokens=2 * hash_block_size, + kv_cache_group_id=group_id, + block_size=block_size, + ) + partial_hash = request.block_hashes[1] + partial_cached = pool.get_cached_block(partial_hash, [0, 1]) + assert partial_cached == [block, block] + + for group_id in range(2): + pool.cache_full_blocks( + request=request, + blocks=[block], + num_cached_blocks=0, + num_full_blocks=1, + block_size=block_size, + kv_cache_group_id=group_id, + ) + + assert pool.get_cached_block(partial_hash, [0]) is None + assert pool.get_cached_block(partial_hash, [1]) is None + full_hash = request.block_hashes[-1] + assert pool.get_cached_block(full_hash, [0, 1]) == [block, block] + assert block.block_hash_num_tokens == block_size + + def make_kv_cache_config_three_types( block_size: int, num_blocks: int, third_spec_type: str = "mamba" ) -> KVCacheConfig: diff --git a/tests/v1/worker/test_attn_utils.py b/tests/v1/worker/test_attn_utils.py index 7e65d650f7ed..2591a426913d 100644 --- a/tests/v1/worker/test_attn_utils.py +++ b/tests/v1/worker/test_attn_utils.py @@ -1,10 +1,20 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + import torch -from vllm.v1.kv_cache_interface import FullAttentionSpec, KVQuantMode -from vllm.v1.worker.gpu.attn_utils import _reshape_kv_cache +from vllm.v1.core.kv_cache_utils import _get_kv_cache_config_packed +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, + KVQuantMode, + MLAAttentionSpec, + UniformTypeKVCacheSpecs, +) +from vllm.v1.worker.gpu.attn_utils import _allocate_kv_cache, _reshape_kv_cache from vllm.v1.worker.utils import AttentionGroup @@ -149,6 +159,104 @@ def get_kv_cache_stride_order( return (0, 1, 2, 3) +class FakeSingleKVBackend: + @staticmethod + def get_kv_cache_shape( + num_blocks: int, + block_size: int, + num_kv_heads: int, + head_size: int, + cache_dtype_str: str = "auto", + ) -> tuple[int, ...]: + assert num_kv_heads == 1 + return (num_blocks, block_size, head_size) + + @staticmethod + def get_kv_cache_stride_order( + include_num_layers_dimension: bool = False, + ) -> tuple[int, ...]: + assert not include_num_layers_dimension + return (0, 1, 2) + + @staticmethod + def get_kv_cache_block_dim( + block_size: int, + num_kv_heads: int, + head_size: int, + cache_dtype_str: str = "auto", + ) -> int: + return 0 + + +def test_lockstep_manager_pages_reshape_to_contiguous_kernel_pages(): + manager_blocks = 3 + kernel_block_size = 64 + main_layer = MLAAttentionSpec( + block_size=kernel_block_size, + num_kv_heads=1, + head_size=328, + dtype=torch.bfloat16, + ) + indexer_layer = MLAAttentionSpec( + block_size=kernel_block_size * 4, + num_kv_heads=1, + head_size=132, + dtype=torch.bfloat16, + dcp_replicated=True, + ) + main_spec = UniformTypeKVCacheSpecs( + block_size=main_layer.block_size, + kv_cache_specs={"main": main_layer}, + ) + indexer_spec = UniformTypeKVCacheSpecs( + block_size=indexer_layer.block_size, + kv_cache_specs={"indexer": indexer_layer}, + ) + groups = [ + KVCacheGroupSpec(["main"], main_spec), + KVCacheGroupSpec(["indexer"], indexer_spec), + ] + config = SimpleNamespace( + cache_config=SimpleNamespace(num_gpu_blocks_override=None), + parallel_config=SimpleNamespace( + decode_context_parallel_size=4, + prefill_context_parallel_size=1, + ), + kv_transfer_config=None, + ) + bytes_per_manager_block = main_layer.page_size_bytes + indexer_layer.page_size_bytes + num_blocks, tensors = _get_kv_cache_config_packed( + config, + groups, + available_memory=bytes_per_manager_block * manager_blocks, + ) + kv_cache_config = KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=tensors, + kv_cache_groups=groups, + ) + raw_tensors = _allocate_kv_cache(kv_cache_config, {}, torch.device("cpu")) + attn_groups = [ + AttentionGroup(FakeSingleKVBackend, ["main"], main_layer, 0), + AttentionGroup(FakeSingleKVBackend, ["indexer"], indexer_layer, 1), + ] + + kv_caches = _reshape_kv_cache( + attn_groups, + raw_tensors, + "auto", + [kernel_block_size, kernel_block_size], + {}, + kv_cache_config, + ) + + assert kv_caches["main"].shape == (manager_blocks, 64, 328) + assert kv_caches["indexer"].shape == (manager_blocks * 4, 64, 132) + assert kv_caches["main"].is_contiguous() + assert kv_caches["indexer"].is_contiguous() + assert all(tensor.block_stride == 0 for tensor in tensors) + + def test_reshape_padded_diff_kv_cache_does_not_infer_kv_dim(): num_blocks = 3 spec = FullAttentionSpec( diff --git a/tests/v1/worker/test_gpu_block_table.py b/tests/v1/worker/test_gpu_block_table.py index 4c15915912ab..954f788781a5 100644 --- a/tests/v1/worker/test_gpu_block_table.py +++ b/tests/v1/worker/test_gpu_block_table.py @@ -177,3 +177,49 @@ def test_compute_slot_mappings_applies_padding_mask(): assert slot_mappings.cpu().tolist() == [ [32, PAD_SLOT_ID, 34, 48, PAD_SLOT_ID, PAD_SLOT_ID, PAD_SLOT_ID, PAD_SLOT_ID] ] + + +def test_compute_slot_mappings_mixed_sharded_and_replicated_groups(): + device = torch.device("cuda") + block_tables = BlockTables( + block_sizes=[64, 256], + max_num_reqs=1, + max_num_batched_tokens=8, + max_num_blocks_per_group=[1, 1], + device=device, + kernel_block_sizes=[64, 64], + cp_size=4, + cp_rank=2, + group_cp_sizes=[4, 1], + ) + block_tables.append_block_ids( + req_index=0, + new_block_ids=([5], [5]), + overwrite=True, + ) + block_tables.apply_staged_writes() + + idx_mapping = torch.tensor([0], dtype=torch.int32, device=device) + query_start_loc = torch.tensor([0, 8], dtype=torch.int32, device=device) + positions = torch.arange(8, dtype=torch.int64, device=device) + slot_mappings = block_tables.compute_slot_mappings( + idx_mapping, + query_start_loc, + positions, + num_tokens_padded=8, + ) + torch.accelerator.synchronize() + + assert slot_mappings.cpu().tolist() == [ + [ + PAD_SLOT_ID, + PAD_SLOT_ID, + 320, + PAD_SLOT_ID, + PAD_SLOT_ID, + PAD_SLOT_ID, + 321, + PAD_SLOT_ID, + ], + list(range(1280, 1288)), + ] diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index d8c28899c1f5..561690dfacae 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -33,7 +33,7 @@ from dataclasses import dataclass from datetime import timedelta from multiprocessing import shared_memory -from typing import TYPE_CHECKING, Any, Protocol +from typing import TYPE_CHECKING, Any, Protocol, cast from unittest.mock import patch import torch @@ -41,6 +41,7 @@ import torch.distributed._functional_collectives as funcol import torch.distributed._symmetric_memory from torch.distributed import Backend, ProcessGroup, Store +from torch.distributed.distributed_c10d import GroupName import vllm.envs as envs from vllm.distributed.device_communicators.base_device_communicator import ( @@ -75,7 +76,20 @@ class Handle(Protocol): def is_completed(self) -> bool: ... - def wait(self) -> None: ... + def wait(self, timeout: timedelta = ...) -> bool: ... + + +class _SplitGroupWithBackend(Protocol): + """Callable shape for the split_group backend extension used by vLLM.""" + + def __call__( + self, + *, + split_ranks: list[list[int]], + group_desc: str, + backend: str, + timeout: timedelta | None, + ) -> ProcessGroup | None: ... def _split_tensor_dict( @@ -221,7 +235,7 @@ def patched_fused_scaled_matmul_reduce_scatter_fake( C, reduce_op, orig_scatter_dim, # need original scatter dim for 3D+ output tensor here - group_name, + GroupName(group_name), ) res = funcol.wait_tensor(res) return res @@ -275,7 +289,8 @@ def _create_subgroups_split_group( ) device_backend_str = _device_backend_str(torch_distributed_backend) - self_device_group = torch.distributed.split_group( + split_group = cast(_SplitGroupWithBackend, torch.distributed.split_group) + self_device_group = split_group( split_ranks=group_ranks, group_desc=f"{group_name}:device", backend=device_backend_str, @@ -286,12 +301,14 @@ def _create_subgroups_split_group( # was bound to via ``device_id``), so a cpu-only filter is rejected. # Include the device backend in the filter; only the gloo backend is # actually used for CPU collectives on this group. - self_cpu_group = torch.distributed.split_group( + self_cpu_group = split_group( split_ranks=group_ranks, group_desc=f"{group_name}:cpu", backend=f"cpu:gloo,{device_backend_str}", timeout=get_cpu_distributed_timeout_or_none(), ) + assert self_device_group is not None + assert self_cpu_group is not None return self_device_group, self_cpu_group @@ -758,9 +775,7 @@ def reduce_scatter_head_major( ) -> torch.Tensor: """Reduce-scatter and preserve a physically head-major output view.""" if self.world_size <= 1 or dim != 1: - raise RuntimeError( - "reduce_scatter_head_major requires DCP heads on dim 1" - ) + raise RuntimeError("reduce_scatter_head_major requires DCP heads on dim 1") if self.device_communicator is None: raise RuntimeError( "reduce_scatter_head_major requires a device communicator" @@ -1107,6 +1122,7 @@ def isend_tensor_dict( handle = torch.distributed.isend( tensor, dst=self.ranks[dst], group=comm_group ) + assert handle is not None if tensor.is_cuda: tensor.record_stream(torch.cuda.current_stream(tensor.device)) handles.append(handle) @@ -1210,6 +1226,7 @@ def irecv_tensor_dict( handle = torch.distributed.irecv( slice_tensor, src=self.ranks[src], group=comm_group ) + assert handle is not None handles.append(handle) def _postprocess( @@ -1230,6 +1247,7 @@ def _postprocess( handle = torch.distributed.irecv( full_tensor, src=self.ranks[src], group=comm_group ) + assert handle is not None handles.append(handle) tensor_dict[key] = full_tensor else: @@ -1836,6 +1854,7 @@ def initialize_model_parallel( """ # Get world size and rank. Ensure some consistencies. assert torch.distributed.is_initialized() + assert decode_context_model_parallel_size is not None from vllm.config import get_current_vllm_config @@ -1890,11 +1909,14 @@ def initialize_model_parallel( # Build the tensor model-parallel groups. global _TP assert _TP is None, "tensor model parallel group is already initialized" - group_ranks = all_ranks.view(-1, tensor_model_parallel_size).unbind(0) - group_ranks = [x.tolist() for x in group_ranks] + group_ranks: list[list[int]] = [ + x.tolist() for x in all_ranks.view(-1, tensor_model_parallel_size).unbind(0) + ] if enable_elastic_ep: - group_ranks = local_all_ranks.view(-1, tensor_model_parallel_size).unbind(0) - group_ranks = [x.tolist() for x in group_ranks] + group_ranks = [ + x.tolist() + for x in local_all_ranks.view(-1, tensor_model_parallel_size).unbind(0) + ] # message queue broadcaster is only used in tensor model parallel group _TP = init_model_parallel_group( group_ranks, @@ -1911,13 +1933,17 @@ def initialize_model_parallel( # dcp_size must not exceed tp_size, because the world size does not # change by DCP, it simply reuses the GPUs of TP group, and split one # TP group into tp_size//dcp_size DCP groups. - group_ranks = all_ranks.reshape(-1, decode_context_model_parallel_size).unbind(0) - group_ranks = [x.tolist() for x in group_ranks] + group_ranks = [ + x.tolist() + for x in all_ranks.reshape(-1, decode_context_model_parallel_size).unbind(0) + ] if enable_elastic_ep: - group_ranks = local_all_ranks.reshape( - -1, decode_context_model_parallel_size - ).unbind(0) - group_ranks = [x.tolist() for x in group_ranks] + group_ranks = [ + x.tolist() + for x in local_all_ranks.reshape( + -1, decode_context_model_parallel_size + ).unbind(0) + ] _DCP = init_model_parallel_group( group_ranks, get_world_group().local_rank, @@ -1944,14 +1970,15 @@ def initialize_model_parallel( group_name="query_split", ) - # A dedicated communicator over the DCP ranks for the transient ckv - # prefetch gather (Fix B). The prefetch runs on a side stream and would - # otherwise share the DCP communicator with the indexer's DCP top-k - # merge on the default stream; concurrent collectives on one NCCL - # communicator from two streams is unsupported. Same ranks as ``_DCP``. + # A dedicated communicator over the DCP ranks for full-CKV prefetch and + # selected-record decode exchange. These paths run on a side stream and + # must not share the DCP communicator with the indexer's top-k merge on + # the default stream. Same ranks as ``_DCP``. global _DCP_CKV_PREFETCH assert _DCP_CKV_PREFETCH is None, "DCP ckv prefetch group is already initialized" - if decode_context_model_parallel_size > 1 and envs.VLLM_B12X_MLA_CKV_GATHER: + if decode_context_model_parallel_size > 1 and ( + envs.VLLM_B12X_MLA_CKV_GATHER or envs.VLLM_B12X_MLA_SPARSE_DECODE_CKV_GATHER + ): _DCP_CKV_PREFETCH = init_model_parallel_group( group_ranks, get_world_group().local_rank, @@ -1961,19 +1988,19 @@ def initialize_model_parallel( global _PCP assert _PCP is None, "prefill context parallel group is already initialized" - group_ranks = ( - all_ranks.transpose(3, 4) + group_ranks = [ + x.tolist() + for x in all_ranks.transpose(3, 4) .reshape(-1, prefill_context_model_parallel_size) .unbind(0) - ) - group_ranks = [x.tolist() for x in group_ranks] + ] if enable_elastic_ep: - group_ranks = ( - local_all_ranks.transpose(1, 2) + group_ranks = [ + x.tolist() + for x in local_all_ranks.transpose(1, 2) .reshape(-1, prefill_context_model_parallel_size) .unbind(0) - ) - group_ranks = [x.tolist() for x in group_ranks] + ] _PCP = init_model_parallel_group( group_ranks, get_world_group().local_rank, backend, group_name="pcp" ) @@ -1981,26 +2008,31 @@ def initialize_model_parallel( # Build the pipeline model-parallel groups. global _PP assert _PP is None, "pipeline model parallel group is already initialized" - group_ranks = ( - all_ranks.transpose(2, 4).reshape(-1, pipeline_model_parallel_size).unbind(0) - ) - group_ranks = [x.tolist() for x in group_ranks] + group_ranks = [ + x.tolist() + for x in all_ranks.transpose(2, 4) + .reshape(-1, pipeline_model_parallel_size) + .unbind(0) + ] if enable_elastic_ep: - group_ranks = ( - local_all_ranks.transpose(0, 2) + group_ranks = [ + x.tolist() + for x in local_all_ranks.transpose(0, 2) .reshape(-1, pipeline_model_parallel_size) .unbind(0) - ) - group_ranks = [x.tolist() for x in group_ranks] + ] _PP = init_model_parallel_group( group_ranks, get_world_group().local_rank, backend, group_name="pp" ) global _DP assert _DP is None, "data parallel group is already initialized" - group_ranks = all_ranks.transpose(1, 4).reshape(-1, data_parallel_size).unbind(0) - group_ranks = [x.tolist() for x in group_ranks] + group_ranks = [ + x.tolist() + for x in all_ranks.transpose(1, 4).reshape(-1, data_parallel_size).unbind(0) + ] if enable_elastic_ep: + assert coord_store is not None _DP = _init_stateless_group( group_ranks, "dp", @@ -2017,8 +2049,9 @@ def initialize_model_parallel( assert _EP is None, "expert parallel group is already initialized" # Don't create EP group for dense models. if config.model_config is None or config.model_config.is_moe: - group_ranks = ( - all_ranks.transpose(1, 2) + group_ranks = [ + x.tolist() + for x in all_ranks.transpose(1, 2) .reshape( -1, data_parallel_size @@ -2026,9 +2059,9 @@ def initialize_model_parallel( * tensor_model_parallel_size, ) .unbind(0) - ) - group_ranks = [x.tolist() for x in group_ranks] + ] if enable_elastic_ep: + assert coord_store is not None _EP = _init_stateless_group( group_ranks, "ep", @@ -2049,6 +2082,7 @@ def initialize_model_parallel( assert _EPLB is None, "EPLB group is already initialized" if config.parallel_config.enable_eplb: if enable_elastic_ep: + assert coord_store is not None _EPLB = _init_stateless_group( group_ranks, "eplb", @@ -2155,6 +2189,15 @@ def model_parallel_is_initialized(): _TP_STATE_PATCHED = False +_MODEL_PARALLEL_CLEANUP_HOOKS: list[Callable[[], None]] = [] + + +def register_model_parallel_cleanup_hook(callback: Callable[[], None]) -> None: + """Register idempotent cleanup that must run before groups are destroyed.""" + if callback not in _MODEL_PARALLEL_CLEANUP_HOOKS: + _MODEL_PARALLEL_CLEANUP_HOOKS.append(callback) + + def get_tensor_model_parallel_world_size() -> int: """Return world size for the tensor model parallel group.""" return get_tp_group().world_size @@ -2173,6 +2216,12 @@ def get_node_count() -> int: def destroy_model_parallel(): """Set the groups to none and destroy them.""" + for cleanup in _MODEL_PARALLEL_CLEANUP_HOOKS: + try: + cleanup() + except Exception: + logger.exception("Model-parallel cleanup hook failed") + global _TP if _TP: diff --git a/vllm/envs.py b/vllm/envs.py index 2cd6ed0ff4f4..a01707890995 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -74,11 +74,18 @@ VLLM_DCP_A2A_MAX_TOKENS: int = 0 VLLM_DCP_A2A_LARGE_BACKEND: Literal["ag_rs", "a2a"] = "ag_rs" VLLM_DCP_SHARD_DRAFT: str | None = None + VLLM_DCP_REPLICATE_INDEXER_CACHE: bool = False VLLM_DCP_GLOBAL_TOPK: bool = True VLLM_DCP_QUERY_SPLIT: bool = False VLLM_B12X_MLA_CKV_GATHER: bool = False + VLLM_B12X_MLA_SPARSE_DECODE_CKV_GATHER: bool = False + VLLM_B12X_MLA_SPARSE_DECODE_TRANSPORT: str = "direct" + VLLM_B12X_MLA_SPARSE_DECODE_BULK_PREFETCH: bool = False + VLLM_B12X_MLA_SPARSE_DECODE_MAX_SEQS: int = 8 + VLLM_B12X_MLA_SPARSE_DECODE_POOL_RECORDS: int = 0 VLLM_B12X_MLA_CKV_GATHER_MIN_TOKENS: int = 16 VLLM_B12X_MLA_CKV_GATHER_MAX_TOKENS: int = 524288 + VLLM_B12X_MLA_CKV_PREFETCH_DEPTH: int = 1 VLLM_MINIMAX_M3_ENABLE_TORCH_COMPILE: bool = False VLLM_B12X_CUDAGRAPH_PIECEWISE_PREWARM: bool = False VLLM_B12X_MOE_FORCE_MODELOPT_PREP: bool = False @@ -438,6 +445,26 @@ def _get_validated_env() -> str | None: return _get_validated_env +def env_bool_with_choices( + env_name: str, + default: bool = False, +) -> Callable[[], bool]: + """Create a case-insensitive, strictly validated boolean env getter.""" + getter = env_with_choices( + env_name, + "1" if default else "0", + ["0", "1", "false", "true", "no", "yes", "off", "on"], + case_sensitive=False, + ) + + def _get_validated_bool() -> bool: + value = getter() + assert value is not None + return value.lower() in ("1", "true", "yes", "on") + + return _get_validated_bool + + def env_list_with_choices( env_name: str, default: list[str], @@ -1159,6 +1186,11 @@ def _resolve_rust_frontend_path() -> str | None: # target indexer cache and native MTP drafts, replicated for external # (Eagle-style) drafts. "VLLM_DCP_SHARD_DRAFT": lambda: os.getenv("VLLM_DCP_SHARD_DRAFT", None), + # Replicate the target model's sparse-indexer K cache on every DCP rank. + "VLLM_DCP_REPLICATE_INDEXER_CACHE": lambda: ( + os.getenv("VLLM_DCP_REPLICATE_INDEXER_CACHE", "0").lower() + in ("1", "true", "yes", "on") + ), # Under DCP, gather sparse-indexer logits across ranks and select a global # top-k instead of a per-rank local top-k. "VLLM_DCP_GLOBAL_TOPK": lambda: ( @@ -1170,12 +1202,44 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_B12X_MLA_CKV_GATHER": lambda: ( os.getenv("VLLM_B12X_MLA_CKV_GATHER", "0").lower() in ("1", "true", "yes", "on") ), + # Exchange only sparse-MLA-selected native CKV records during DCP decode. + "VLLM_B12X_MLA_SPARSE_DECODE_CKV_GATHER": lambda: ( + os.getenv("VLLM_B12X_MLA_SPARSE_DECODE_CKV_GATHER", "0").lower() + in ("1", "true", "yes", "on") + ), + # Transport for sparse selected-record decode. Keep the existing direct + # path as the default because copy-engine staging has additional VRAM cost. + "VLLM_B12X_MLA_SPARSE_DECODE_TRANSPORT": env_with_choices( + "VLLM_B12X_MLA_SPARSE_DECODE_TRANSPORT", + "direct", + ["auto", "ce", "direct"], + ), + # Combine exactly three eligible Shared-layer selected-record prefetches + # into one CE exchange. Sparse decode must also be enabled. This is + # effective only with transport=ce|auto and sufficient layered capacity; + # direct transport, partial groups, and oversized groups retain the + # existing per-layer path. + "VLLM_B12X_MLA_SPARSE_DECODE_BULK_PREFETCH": env_bool_with_choices( + "VLLM_B12X_MLA_SPARSE_DECODE_BULK_PREFETCH" + ), + "VLLM_B12X_MLA_SPARSE_DECODE_MAX_SEQS": lambda: int( + os.getenv("VLLM_B12X_MLA_SPARSE_DECODE_MAX_SEQS", "8") + ), + # Zero sizes the pool for the full configured request count. + "VLLM_B12X_MLA_SPARSE_DECODE_POOL_RECORDS": lambda: int( + os.getenv("VLLM_B12X_MLA_SPARSE_DECODE_POOL_RECORDS", "0") + ), "VLLM_B12X_MLA_CKV_GATHER_MIN_TOKENS": lambda: int( os.getenv("VLLM_B12X_MLA_CKV_GATHER_MIN_TOKENS", "16") ), "VLLM_B12X_MLA_CKV_GATHER_MAX_TOKENS": lambda: int( os.getenv("VLLM_B12X_MLA_CKV_GATHER_MAX_TOKENS", "524288") ), + # Number of future full-CKV layer gathers to queue. Zero keeps the + # synchronous gather path without allocating lookahead ring slots. + "VLLM_B12X_MLA_CKV_PREFETCH_DEPTH": lambda: int( + os.getenv("VLLM_B12X_MLA_CKV_PREFETCH_DEPTH", "1") + ), # Diagnostic flag retained for local experiments. MiniMax M3 compile is # fail-closed in the model until the no-break path is validated. "VLLM_MINIMAX_M3_ENABLE_TORCH_COMPILE": lambda: bool( diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 67e71c6291d1..8790932f3bb6 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -1018,12 +1018,22 @@ def forward_impl( workspace_gather_used = False ckv_gather_used = False if self.impl.dcp_world_size > 1: - ckv_gather_selector = getattr( + prefill_ckv_gather_selector = getattr( self.impl, "dcp_prefill_ckv_gather_eligible", None ) - ckv_gather_used = bool( - callable(ckv_gather_selector) - and ckv_gather_selector(attn_metadata, num_mqa_tokens) + sparse_decode_ckv_gather_selector = getattr( + self.impl, "dcp_sparse_decode_ckv_gather_eligible", None + ) + prefill_ckv_gather_used = bool( + callable(prefill_ckv_gather_selector) + and prefill_ckv_gather_selector(attn_metadata, num_mqa_tokens) + ) + sparse_decode_ckv_gather_used = bool( + callable(sparse_decode_ckv_gather_selector) + and sparse_decode_ckv_gather_selector(attn_metadata, num_mqa_tokens) + ) + ckv_gather_used = ( + prefill_ckv_gather_used or sparse_decode_ckv_gather_used ) if not ckv_gather_used: if not self.impl.can_return_lse_for_decode: @@ -1080,8 +1090,7 @@ def forward_impl( mqa_q = torch.cat(mqa_q, dim=-1) if ckv_gather_used: logger.info_once( - "Keeping local query heads for transient full-CKV " - "B12X sparse MLA prefill" + "Keeping local query heads for B12X sparse MLA CKV gather" ) elif dcp_use_b12x: mqa_q = dcp_b12x_all_gather_heads( diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index b4fe13d67ad4..c4ad9034f594 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -3,6 +3,7 @@ """Custom Sparse Attention Indexer layers.""" import os +from typing import cast import torch @@ -10,7 +11,7 @@ from vllm import _custom_ops as ops from vllm._aiter_ops import rocm_aiter_ops from vllm.config import CUDAGraphMode, get_current_vllm_config -from vllm.distributed import get_dcp_group, get_query_split_group +from vllm.distributed import GroupCoordinator, get_dcp_group, get_query_split_group from vllm.forward_context import get_forward_context from vllm.model_executor.custom_op import CustomOp from vllm.model_executor.layers.quantization.utils.quant_utils import ( @@ -1732,17 +1733,18 @@ def sparse_attn_indexer( cp_kv_cache_interleave_size=cp_kv_cache_interleave_size, ) if qs_active: - gathered_indices = qs_group.all_gather( + active_qs_group = cast(GroupCoordinator, qs_group) + gathered_indices = active_qs_group.all_gather( topk_indices.contiguous(), dim=0 ) topk_indices_buffer[ chunk.token_start : chunk.token_end, :topk_tokens ].copy_(gathered_indices) if topk_scores is not None: - gathered_scores = qs_group.all_gather( + gathered_scores = active_qs_group.all_gather( topk_scores.contiguous(), dim=0 ) - topk_scores_buffer[ + cast(torch.Tensor, topk_scores_buffer)[ chunk.token_start : chunk.token_end, :topk_tokens ].copy_(gathered_scores) continue @@ -2155,6 +2157,7 @@ def __init__( topk_scores_buffer: torch.Tensor | None = None, output_physical_slots: bool = False, num_q_heads: int | None = None, + dcp_replicated: bool = False, ): super().__init__() self.k_cache = k_cache @@ -2174,7 +2177,9 @@ def __init__( # during model construction) and pass them into the custom op, rather # than threading them through per-step metadata. parallel_config = get_current_vllm_config().parallel_config - self.dcp_world_size = parallel_config.decode_context_parallel_size + self.dcp_replicated = bool(dcp_replicated) + configured_dcp_world_size = parallel_config.decode_context_parallel_size + self.dcp_world_size = 1 if self.dcp_replicated else configured_dcp_world_size self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_world_size > 1 else 0 self.cp_kv_cache_interleave_size = parallel_config.cp_kv_cache_interleave_size self.use_b12x_sparse_indexer = use_b12x_sparse_indexer() diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 134500d172a6..a18f3ca1709a 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -617,6 +617,41 @@ def forward( return output +def _replicate_indexer_cache_under_dcp(prefix: str, vllm_config: VllmConfig) -> bool: + layer_id = extract_layer_index(prefix) + num_hidden_layers = getattr( + vllm_config.model_config.hf_config, "num_hidden_layers", None + ) + if layer_id is None or num_hidden_layers is None: + return False + + parallel_config = vllm_config.parallel_config + dcp_size = parallel_config.decode_context_parallel_size + pcp_size = parallel_config.prefill_context_parallel_size + if dcp_size * pcp_size <= 1: + return False + + if int(layer_id) >= int(num_hidden_layers): + return False + + requested = envs.VLLM_DCP_REPLICATE_INDEXER_CACHE + if requested and not vllm_config.use_v2_model_runner: + raise NotImplementedError( + "Replicated sparse-indexer KV requires the V2 model runner's " + "per-group context-parallel block tables." + ) + if requested and (not 2 <= dcp_size <= 8 or pcp_size != 1): + raise NotImplementedError( + "Replicated sparse-indexer KV currently supports DCP2 through " + "DCP8 with PCP1." + ) + if requested and not use_b12x_sparse_indexer(): + raise RuntimeError( + "VLLM_DCP_REPLICATE_INDEXER_CACHE requires the B12X sparse indexer." + ) + return requested + + class DeepseekV32IndexerCache(torch.nn.Module, AttentionLayerBase): def __init__( self, head_dim: int, dtype: torch.dtype, prefix: str, cache_config: CacheConfig @@ -627,30 +662,44 @@ def __init__( self.prefix = prefix self.cache_config = cache_config self.dtype = dtype - compilation_config = get_current_vllm_config().compilation_config + vllm_config = get_current_vllm_config() + self.dcp_replicated = _replicate_indexer_cache_under_dcp(prefix, vllm_config) + compilation_config = vllm_config.compilation_config if prefix in compilation_config.static_forward_context: raise ValueError(f"Duplicate layer name: {prefix}") compilation_config.static_forward_context[prefix] = self def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: + block_size = self.cache_config.block_size + if self.dcp_replicated: + parallel_config = vllm_config.parallel_config + block_size *= ( + parallel_config.decode_context_parallel_size + * parallel_config.prefill_context_parallel_size + ) + logger.info_once( + "Using a DCP-replicated sparse-indexer K cache with %d-token " + "manager pages.", + block_size, + ) layer_id = extract_layer_index(self.prefix) num_hidden_layers = getattr( vllm_config.model_config.hf_config, "num_hidden_layers", None ) raw = envs.VLLM_DCP_SHARD_DRAFT shard_draft = ("1" if raw is None else raw).lower() in ("1", "true", "yes") - dcp_replicated = ( + draft_replicated = ( not shard_draft and layer_id is not None and num_hidden_layers is not None and int(layer_id) >= int(num_hidden_layers) ) return MLAAttentionSpec( # Only has one vector instead of K + V - block_size=self.cache_config.block_size, + block_size=block_size, num_kv_heads=1, head_size=self.head_dim, dtype=self.dtype, - dcp_replicated=dcp_replicated, + dcp_replicated=self.dcp_replicated or draft_replicated, ) def forward(self): ... @@ -751,6 +800,7 @@ def __init__( topk_scores_buffer=self.topk_scores_buffer, output_physical_slots=self.output_physical_slots, num_q_heads=self.n_head, + dcp_replicated=self.k_cache.dcp_replicated, ) self.is_inplace_rope = is_inplace_rope diff --git a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py index 0ed33b889860..fbc9235006c9 100644 --- a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py @@ -27,6 +27,8 @@ import inspect import os +import weakref +from collections.abc import Hashable from dataclasses import dataclass from typing import TYPE_CHECKING, Any, ClassVar, cast @@ -77,7 +79,6 @@ _EXTEND_PREWARM_DONE: set[ tuple[int | None, int, int, int, int, int, bool, str, bool] ] = set() -_CKV_GATHER_WORKSPACES: dict[tuple[str, int | None], torch.Tensor] = {} _KV_FP8_ROPE_REQUESTED = os.getenv("KV_FP8_ROPE", "0") == "1" @@ -146,18 +147,272 @@ def _env_int(name: str, default: int) -> int: return parsed -def _get_ckv_gather_workspace(device: torch.device, nbytes: int) -> torch.Tensor: - key = (device.type, device.index) - workspace = _CKV_GATHER_WORKSPACES.get(key) - if workspace is None: - workspace = torch.empty((nbytes,), dtype=torch.uint8, device=device) - _CKV_GATHER_WORKSPACES[key] = workspace - elif workspace.numel() < nbytes: - raise RuntimeError( - "CKV gather workspace cannot grow after attention layers retain " - f"aliases: existing={workspace.numel()} requested={nbytes}" +def _ckv_prefetch_supports_format(kv_cache_dtype: str) -> bool: + return kv_cache_dtype in ("fp8_ds_mla", "nvfp4_ds_mla") + + +def _sparse_decode_supports_format( + kv_cache_dtype: str, record_bytes: int, kv_fp8_rope: bool +) -> bool: + if kv_cache_dtype != "nvfp4_ds_mla": + return False + expected_record_bytes = 368 if kv_fp8_rope else 432 + return int(record_bytes) == expected_record_bytes + + +def _ckv_prefetch_ring_slots(depth: int) -> int: + return max(0, int(depth)) + 1 + + +def _ckv_prefetch_target_indices( + layer_idx: int, + depth: int, + layer_caches: list[torch.Tensor | None], + pending_layers: dict[int, "_CKVPrefetchPendingLayer"], +) -> list[int]: + targets: list[int] = [] + for distance in range(1, max(0, int(depth)) + 1): + target_idx = layer_idx + distance + if target_idx in pending_layers: + continue + if target_idx >= len(layer_caches) or layer_caches[target_idx] is None: + break + targets.append(target_idx) + return targets + + +@dataclass(frozen=True) +class _CKVWorkspaceIdentity: + device: torch.device + storage_generation: int + data_ptr: int + storage_offset: int + shape: tuple[int, ...] + stride: tuple[int, ...] + dtype: torch.dtype + + +@dataclass +class _CKVPrefetchTicket: + event: Any + wait_scheduled: bool = False + + def wait_once(self) -> None: + if not self.wait_scheduled: + self.event.wait() + self.wait_scheduled = True + + def wait_on_stream_once(self, stream: Any) -> None: + if not self.wait_scheduled: + stream.wait_event(self.event) + self.wait_scheduled = True + + +@dataclass(frozen=True) +class _CKVPrefetchPendingLayer: + ticket: _CKVPrefetchTicket + buf_idx: int + + +def _ckv_workspace_identity(workspace: torch.Tensor) -> _CKVWorkspaceIdentity: + storage = workspace.untyped_storage() + return _CKVWorkspaceIdentity( + device=workspace.device, + # WorkspaceManager does not expose its allocation generation. PyTorch's + # storage handle changes when MRV2 replaces the backing allocation. + storage_generation=int(storage._cdata), + data_ptr=workspace.data_ptr(), + storage_offset=workspace.storage_offset(), + shape=tuple(workspace.shape), + stride=tuple(workspace.stride()), + dtype=workspace.dtype, + ) + + +class _CKVPrefetchState: + """Cross-layer state for one workspace allocation and execution lane.""" + + def __init__( + self, + workspace_identity: _CKVWorkspaceIdentity, + workspace: torch.Tensor, + ) -> None: + self.workspace_identity = workspace_identity + self.workspace_storage_ref = weakref.ref(workspace.untyped_storage()) + self.layer_caches: list[torch.Tensor | None] = [] + self.layer_impls: list[B12xMLASparseImpl | None] = [] + self.pending_layers: dict[int, _CKVPrefetchPendingLayer] = {} + self.sparse_decode_state: Any | None = None + self.gather_stream: torch.cuda.Stream | None = None + self.ckv_workspace: torch.Tensor | None = None + self.ckv_workspace_generation = 0 + self.last_layer_idx: int | None = None + + def begin_step(self) -> None: + tickets = { + id(pending.ticket): pending.ticket + for pending in self.pending_layers.values() + } + for ticket in tickets.values(): + # Preserve ring ordering without blocking the host indefinitely. + # The next main-stream gather is enqueued after these dependencies. + ticket.wait_once() + self.pending_layers.clear() + self.last_layer_idx = None + + def register_pending_group( + self, + layer_slots: dict[int, int], + event: Any, + ) -> _CKVPrefetchTicket: + if not layer_slots: + raise ValueError("prefetch group must contain at least one layer") + duplicate_layers = self.pending_layers.keys() & layer_slots.keys() + if duplicate_layers: + raise RuntimeError( + f"prefetch group overlaps pending layers {sorted(duplicate_layers)}" + ) + ticket = _CKVPrefetchTicket(event) + self.pending_layers.update( + { + layer_idx: _CKVPrefetchPendingLayer(ticket, int(buf_idx)) + for layer_idx, buf_idx in layer_slots.items() + } + ) + return ticket + + def pop_pending_layer( + self, + layer_idx: int | None, + ) -> _CKVPrefetchPendingLayer | None: + if layer_idx is None: + return None + return self.pending_layers.pop(layer_idx, None) + + def enter_layer(self, layer_idx: int) -> None: + if self.last_layer_idx is not None and layer_idx <= self.last_layer_idx: + self.begin_step() + self.last_layer_idx = layer_idx + + def register_cache(self, layer_idx: int, kv_cache: torch.Tensor) -> None: + while len(self.layer_caches) <= layer_idx: + self.layer_caches.append(None) + self.layer_caches[layer_idx] = kv_cache + + def register_impl(self, layer_idx: int, impl: "B12xMLASparseImpl") -> None: + while len(self.layer_impls) <= layer_idx: + self.layer_impls.append(None) + self.layer_impls[layer_idx] = impl + + def get_sparse_decode_state(self, layout, exchange): + if self.sparse_decode_state is None: + from vllm.v1.attention.backends.mla.b12x_sparse_ckv_decode import ( + SparseCKVDecodeState, + ) + + self.sparse_decode_state = SparseCKVDecodeState( + layout=layout, + device=self.workspace_identity.device, + exchange=exchange, + ) + elif self.sparse_decode_state.layout != layout: + raise RuntimeError("sparse CKV layout changed within one execution lane") + elif self.sparse_decode_state.exchange is not exchange: + raise RuntimeError("sparse CKV exchange changed within one execution lane") + return self.sparse_decode_state + + def get_gather_stream(self) -> torch.cuda.Stream: + if self.gather_stream is None: + self.gather_stream = torch.cuda.Stream( + device=self.workspace_identity.device + ) + return self.gather_stream + + def get_ckv_workspace(self, nbytes: int) -> torch.Tensor: + if nbytes <= 0: + raise ValueError(f"CKV workspace size must be positive, got {nbytes}") + workspace = self.ckv_workspace + if workspace is None or workspace.numel() != nbytes: + # Pending side-stream writes must be ordered before replacing their + # backing allocation. record_stream() in _dcp_gather_ckv keeps the + # old storage alive until those writes complete. + self.begin_step() + workspace = torch.empty( + (nbytes,), + dtype=torch.uint8, + device=self.workspace_identity.device, + ) + self.ckv_workspace = workspace + self.ckv_workspace_generation += 1 + return workspace + + +class _CKVPrefetchStateRegistry: + """Builder-owned states partitioned by lane-scoped CKV workspace.""" + + def __init__(self) -> None: + self.states: dict[ + tuple[_CKVWorkspaceIdentity, Hashable | None], _CKVPrefetchState + ] = {} + + def _retire( + self, + keys: list[tuple[_CKVWorkspaceIdentity, Hashable | None]], + ) -> None: + for key in keys: + self.states.pop(key).begin_step() + + def _prune_released_workspaces(self) -> None: + self._retire( + [ + key + for key, state in self.states.items() + if state.workspace_storage_ref() is None + ] ) - return workspace[:nbytes] + + def begin_step(self) -> None: + self._prune_released_workspaces() + for state in self.states.values(): + state.begin_step() + + def for_workspace( + self, + workspace: torch.Tensor, + layer_idx: int | None = None, + kv_cache: torch.Tensor | None = None, + *, + execution_lane_key: Hashable | None = None, + ) -> _CKVPrefetchState: + self._prune_released_workspaces() + identity = _ckv_workspace_identity(workspace) + registry_key = (identity, execution_lane_key) + state = self.states.get(registry_key) + if state is None: + # A resized view may retain its address, while an allocation + # replacement changes it. A known layer cache identifies the same + # execution lane across the latter without merging target/draft. + stale_keys = [ + existing_key + for existing_key, existing_state in self.states.items() + if existing_key[1] == execution_lane_key + and ( + ( + existing_key[0].device == identity.device + and existing_key[0].data_ptr == identity.data_ptr + ) + or ( + layer_idx is not None + and kv_cache is not None + and layer_idx < len(existing_state.layer_caches) + and existing_state.layer_caches[layer_idx] is kv_cache + ) + ) + ] + self._retire(stale_keys) + state = _CKVPrefetchState(identity, workspace) + self.states[registry_key] = state + return state def _dcp_all_gather_current_stream( @@ -195,6 +450,94 @@ def _dcp_all_gather_current_stream( output_tensor.copy_(gathered) +@triton.jit +def _find_sparse_decode_current_token_slot_kernel( + union_indices_ptr, + req_id_ptr, + global_seq_len_ptr, + output_slot_ptr, + union_stride0, + capacity, + BLOCK_N: tl.constexpr, +): + row = tl.program_id(0) + offsets = tl.arange(0, BLOCK_N) + req = tl.load(req_id_ptr + row) + selected = tl.load( + union_indices_ptr + req * union_stride0 + offsets, + mask=offsets < capacity, + other=-1, + ) + current_token = tl.load(global_seq_len_ptr + row) - 1 + matched = tl.max(tl.where(selected == current_token, offsets, -1), axis=0) + flattened_slot = tl.where(matched >= 0, req * capacity + matched, -1) + tl.store(output_slot_ptr + row, flattened_slot) + + +def _find_sparse_decode_current_token_slot( + union_indices: torch.Tensor, + req_id_per_token: torch.Tensor, + global_seq_len: torch.Tensor, + output_slot: torch.Tensor, +) -> None: + if union_indices.ndim != 2 or union_indices.dtype != torch.int32: + raise TypeError("sparse CKV unions must be rank-2 int32") + if req_id_per_token.shape != global_seq_len.shape: + raise ValueError("sparse CKV request ids must match causal lengths") + if output_slot.shape != global_seq_len.shape or output_slot.dtype != torch.int64: + raise TypeError("sparse CKV patch slots must be matching int64 rows") + capacity = int(union_indices.shape[1]) + _find_sparse_decode_current_token_slot_kernel[(global_seq_len.numel(),)]( + union_indices, + req_id_per_token, + global_seq_len, + output_slot, + union_indices.stride(0), + capacity, + BLOCK_N=triton.next_power_of_2(capacity), + ) + + +@triton.jit +def _offset_sparse_decode_remap_kernel( + remap_ptr, + numel, + request_stride, + rows_per_request, + topk, + BLOCK_SIZE: tl.constexpr, +): + positions = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = positions < numel + values = tl.load(remap_ptr + positions, mask=mask, other=-1) + request = positions // (rows_per_request * topk) + values = tl.where(values >= 0, values + request * request_stride, values) + tl.store(remap_ptr + positions, values, mask=mask) + + +def _offset_sparse_decode_remap( + remap: torch.Tensor, + *, + request_stride: int, + rows_per_request: int, + topk: int, +) -> None: + if remap.dtype != torch.int32 or not remap.is_contiguous(): + raise TypeError("sparse CKV remap must be contiguous int32") + numel = int(remap.numel()) + if numel == 0: + return + block_size = 256 + _offset_sparse_decode_remap_kernel[(_cdiv(numel, block_size),)]( + remap, + numel, + int(request_stride), + int(rows_per_request), + int(topk), + BLOCK_SIZE=block_size, + ) + + @triton.jit def _mask_page_table_after_nsa_len_kernel( page_table_ptr, @@ -531,6 +874,7 @@ class B12xMLASparseMetadata(AttentionMetadata): dcp_local_total_tokens: int = 0 dcp_padded_total_tokens: int = 0 dcp_ckv_gather_eligible: bool = False + ckv_prefetch_registry: _CKVPrefetchStateRegistry | None = None block_size: int = 64 topk_tokens: int = 2048 @@ -571,6 +915,11 @@ def __init__( from vllm import envs as envs_mod ckv_gather_requested = envs_mod.VLLM_B12X_MLA_CKV_GATHER + sparse_decode_requested = envs_mod.VLLM_B12X_MLA_SPARSE_DECODE_CKV_GATHER + ckv_metadata_requested = ckv_gather_requested or sparse_decode_requested + self.ckv_prefetch_registry = ( + _CKVPrefetchStateRegistry() if ckv_metadata_requested else None + ) # Max-batched-token scratch buffers so cudagraph capture sees stable # allocations (sliced per build()). self.cache_seq_lens_per_token_buffer = torch.empty( @@ -672,16 +1021,8 @@ def build( from vllm import envs as envs_mod - if envs_mod.VLLM_B12X_MLA_CKV_GATHER: - # Reset the cross-layer prefetch pipeline once per step so the - # first layer always sync-gathers. The event/buf-idx are class - # state that would otherwise leak across chunks (the last layer - # of a chunk consumes but never re-arms), scheduling layer 0 - # of subsequent chunks onto a stale gathered buffer. The - # layer->cache registry is intentionally left intact (stable - # cache pointers across chunks). - B12xMLASparseImpl._shared_gather_event = None - B12xMLASparseImpl._shared_gather_buf_idx = 0 + if self.ckv_prefetch_registry is not None: + self.ckv_prefetch_registry.begin_step() if ( use_dcp @@ -886,12 +1227,17 @@ def build( dcp_local_cu_seq_lens=dcp_local_cu_seq_lens, global_cache_seq_lens_per_req=( cm.seq_lens[: cm.num_reqs] - if use_dcp and envs_mod.VLLM_B12X_MLA_CKV_GATHER + if use_dcp + and ( + envs_mod.VLLM_B12X_MLA_CKV_GATHER + or envs_mod.VLLM_B12X_MLA_SPARSE_DECODE_CKV_GATHER + ) else None ), dcp_local_total_tokens=dcp_local_total_tokens, dcp_padded_total_tokens=dcp_padded_total_tokens, dcp_ckv_gather_eligible=dcp_ckv_gather_eligible, + ckv_prefetch_registry=self.ckv_prefetch_registry, block_size=self.kv_cache_spec.block_size, topk_tokens=self.topk_tokens, ) @@ -900,7 +1246,7 @@ def build( class B12xMLASparseImpl(MLAAttentionImpl[B12xMLASparseMetadata]): """b12x unified sparse-MLA implementation (decode + extend/prefill).""" - is_sparse: bool = True + is_sparse: ClassVar[bool] = True can_return_lse_for_decode: bool = True # B12X handles decode and extend inside its own top-k MQA kernels; the # generic dense-MHA prefill path assumes cache layouts it never validated. @@ -1067,6 +1413,8 @@ def __init__( spec = getattr(vllm_config, "speculative_config", None) if spec is not None and getattr(spec, "num_speculative_tokens", None): q_per_req = 1 + int(spec.num_speculative_tokens) + self._sparse_decode_rows_per_request = q_per_req + self._sparse_decode_emits_topk = indexer is not None if self.spec_extend_as_decode: q_per_req = max(q_per_req, self.spec_decode_max_q) self._decode_max_rows = min(max_num_seqs * q_per_req, max_batched) @@ -1148,40 +1496,109 @@ def _make_plan( self._extend_plan = _make_plan( "extend", max_batched, self._kernel_num_heads, max_num_seqs ) + + from vllm import envs as envs_mod + from vllm.v1.attention.backends.mla.b12x_sparse_ckv_decode import ( + get_selected_record_exchange, + plan_sparse_ckv_decode, + preload_dense_union_extension_consistently, + ) + + model_type = str( + getattr(vllm_config.model_config.hf_config, "model_type", "target") + ) + self._ckv_execution_lane_key = (model_type, id(vllm_config)) + sparse_decode_requested = envs_mod.VLLM_B12X_MLA_SPARSE_DECODE_CKV_GATHER + self._sparse_decode_enabled = sparse_decode_requested and ( + 2 <= self.dcp_world_size <= 8 + and self.num_heads % _HEAD_ALIGNMENT == 0 + and self.dcp_workspace_non_dbo + and _sparse_decode_supports_format( + self.kv_cache_dtype, + self._kv_record_bytes, + self._kv_fp8_rope, + ) + ) + sparse_max_requests = min( + max_num_seqs, + max(1, int(envs_mod.VLLM_B12X_MLA_SPARSE_DECODE_MAX_SEQS)), + ) + self._sparse_decode_layout = None + self._sparse_decode_exchange = None + self._sparse_decode_bulk_prefetch = bool( + envs_mod.VLLM_B12X_MLA_SPARSE_DECODE_BULK_PREFETCH + ) + if self._sparse_decode_enabled: + self._sparse_decode_layout = plan_sparse_ckv_decode( + dcp_world_size=self.dcp_world_size, + topk=self.topk_tokens, + rows_per_request=self._sparse_decode_rows_per_request, + max_requests=sparse_max_requests, + pool_records=int(envs_mod.VLLM_B12X_MLA_SPARSE_DECODE_POOL_RECORDS), + record_bytes=self._kv_record_bytes, + prefetch_depth=max(0, int(envs_mod.VLLM_B12X_MLA_CKV_PREFETCH_DEPTH)), + ) + graph_rows = int( + vllm_config.compilation_config.max_cudagraph_capture_size or 0 + ) + required_rows = ( + self._sparse_decode_layout.max_fast_requests + * self._sparse_decode_layout.rows_per_request + ) + if graph_rows and graph_rows < required_rows: + raise ValueError( + "selected-record CKV decode requires max cudagraph " + f"capture size >= {required_rows}; configured {graph_rows}" + ) + + from vllm.distributed.parallel_state import get_dcp_ckv_prefetch_group + + process_group = get_dcp_ckv_prefetch_group().device_group + if process_group is None: + raise RuntimeError("DCP CKV group has no device process group") + preload_dense_union_extension_consistently(process_group, self.device) + self._sparse_decode_exchange = get_selected_record_exchange( + process_group=process_group, + device=self.device, + layout=self._sparse_decode_layout, + lane_key=self._ckv_execution_lane_key, + ) + elif sparse_decode_requested: + logger.warning_once( + "Ignoring selected-record CKV decode for unsupported " + "topology/format: DCP=%d local_heads=%d dtype=%s bytes=%d " + "KV_FP8_ROPE=%d DBO=%d", + self.dcp_world_size, + self.num_heads, + self.kv_cache_dtype, + self._kv_record_bytes, + int(self._kv_fp8_rope), + int(not self.dcp_workspace_non_dbo), + ) + + self._sparse_decode_plan = ( + _make_plan( + "decode", + self._sparse_decode_layout.max_fast_requests + * self._sparse_decode_layout.rows_per_request, + self.num_heads, + self._sparse_decode_layout.max_fast_requests + * self._sparse_decode_layout.rows_per_request, + ) + if self._sparse_decode_layout is not None + else None + ) # One caller-owned uint8 scratch tensor covers either path (the larger # layout); the per-mode materializer carves its views from the prefix. self._scratch_nbytes = max( int(self._decode_plan.layout.nbytes), int(self._extend_plan.layout.nbytes), + int(self._sparse_decode_plan.layout.nbytes) + if self._sparse_decode_plan is not None + else 0, ) - # Pre-touch q-concat + the attention scratch TOGETHER so the workspace - # manager grows during warmup (before lock_workspace() runs - # post-cudagraph-capture) and so the two always come from ONE - # get_simultaneous call -> distinct, non-overlapping offsets. The manager - # packs every call from offset 0, so borrowing q and the scratch the kernel - # writes in separate calls would alias them. - workspace_specs: list[tuple[tuple[int, ...], torch.dtype]] = [ - ( - (max_batched, self._kernel_num_heads, self.q_head_dim), - torch.bfloat16, - ) - ] - if self._pad_heads: - workspace_specs.append( - ( - (max_batched, self._input_num_heads, self.kv_lora_rank), - torch.bfloat16, - ) - ) - workspace_specs.append(((self._scratch_nbytes,), torch.uint8)) - self._workspace_specs = tuple(workspace_specs) - self._borrow_workspaces() - self._prewarm_extend_kernels_once(max_batched) - # CKV gather setup (Fix B). - from vllm import envs as envs_mod - ckv_gather_requested = envs_mod.VLLM_B12X_MLA_CKV_GATHER self._ckv_gather_enabled = ckv_gather_requested and ( self.dcp_world_size > 1 @@ -1199,6 +1616,20 @@ def _make_plan( self._ckv_kernel_num_heads = self.num_heads self._ckv_gather_max_tokens = envs_mod.VLLM_B12X_MLA_CKV_GATHER_MAX_TOKENS self._ckv_gather_min_tokens = envs_mod.VLLM_B12X_MLA_CKV_GATHER_MIN_TOKENS + configured_prefetch_depth = int(envs_mod.VLLM_B12X_MLA_CKV_PREFETCH_DEPTH) + if configured_prefetch_depth < 0: + logger.warning_once( + "Ignoring negative VLLM_B12X_MLA_CKV_PREFETCH_DEPTH=%d; using 0", + configured_prefetch_depth, + ) + self._ckv_prefetch_supported = ( + self._ckv_gather_enabled + and _ckv_prefetch_supports_format(self.kv_cache_dtype) + ) + self._ckv_prefetch_depth = ( + max(0, configured_prefetch_depth) if self._ckv_prefetch_supported else 0 + ) + self._ckv_workspace_slots = _ckv_prefetch_ring_slots(self._ckv_prefetch_depth) self._ckv_local_capacity = ( _cdiv( _cdiv(self._ckv_gather_max_tokens, max(1, self.dcp_world_size)) @@ -1208,18 +1639,12 @@ def _make_plan( * self.block_size ) self._ckv_workspace_nbytes = ( - 2 - * (self.dcp_world_size + 1) + (1 + self._ckv_workspace_slots * self.dcp_world_size) * self._ckv_local_capacity * self._kv_record_bytes if self._ckv_gather_enabled else 0 ) - self._ckv_workspace = ( - _get_ckv_gather_workspace(self.device, self._ckv_workspace_nbytes) - if self._ckv_gather_enabled - else None - ) # Separate extend plan for the gathered-cache path: full local heads # (no head all-gather), global seq lens. @@ -1234,28 +1659,49 @@ def _make_plan( else: self._ckv_extend_plan = None - # Layer prefetch (side stream + events + ping-pong). - # _shared_* are class-level: layer L kicks off the prefetch for - # layer L+1, and layer L+1 (a different impl instance) consumes it. - self._ckv_prefetch_supported = self._ckv_gather_enabled and ( - self.kv_cache_dtype == "fp8_ds_mla" or self._kv_fp8_rope - ) - if self._ckv_gather_enabled: - self._ckv_gather_stream = torch.cuda.Stream(device=self.device) + # Pre-touch q-concat and attention scratch together. Cross-layer CKV + # data cannot live in WorkspaceManager: every caller borrows from + # offset zero, so intervening indexer/MoE scratch would alias it. The + # builder-owned state allocates a dedicated ring per workspace lane. + workspace_specs: list[tuple[tuple[int, ...], torch.dtype]] = [ + ( + (max_batched, self._kernel_num_heads, self.q_head_dim), + torch.bfloat16, + ) + ] + if self._pad_heads: + workspace_specs.append( + ( + (max_batched, self._input_num_heads, self.kv_lora_rank), + torch.bfloat16, + ) + ) + workspace_specs.append(((self._scratch_nbytes,), torch.uint8)) + self._workspace_specs = tuple(workspace_specs) + self._borrow_workspaces() + self._prewarm_extend_kernels_once(max_batched) + + # The registry isolates target/draft and graph workspaces. Full-CKV + # prefetch creates its lane stream lazily; the stream-affine selected- + # record exchange owns one dedicated stream created during init. + if self._ckv_gather_enabled or self._sparse_decode_enabled: self._ckv_current_chunk_kv_c: torch.Tensor | None = None self._ckv_current_chunk_kpe: torch.Tensor | None = None - B12xMLASparseImpl._all_layer_kv_caches: list[torch.Tensor | None] = [] - B12xMLASparseImpl._shared_gather_event: torch.cuda.Event | None = None - B12xMLASparseImpl._shared_gather_buf_idx = 0 - if not self._ckv_prefetch_supported: + if self._ckv_gather_enabled and not self._ckv_prefetch_supported: logger.warning_once( "CKV gather prefetch disabled for kv_cache_dtype=%s " "(KV_FP8_ROPE=%s); falling back to synchronous gather.", self.kv_cache_dtype, int(self._kv_fp8_rope), ) + elif self._ckv_gather_enabled and self._ckv_prefetch_depth > 0: + logger.info_once( + "Using native CKV layer prefetch with depth=%d and " + "%d workspace slots.", + self._ckv_prefetch_depth, + self._ckv_workspace_slots, + ) else: - self._ckv_gather_stream = None self._ckv_current_chunk_kv_c = None self._ckv_current_chunk_kpe = None @@ -1601,6 +2047,27 @@ def _validate_ckv_workspace(self, ckv_workspace: torch.Tensor) -> None: ): raise RuntimeError("B12X CKV gather borrowed an invalid workspace") + def _ckv_workspace_views( + self, ckv_workspace: torch.Tensor, buf_idx: int + ) -> tuple[torch.Tensor, torch.Tensor]: + self._validate_ckv_workspace(ckv_workspace) + if not 0 <= int(buf_idx) < self._ckv_workspace_slots: + raise ValueError( + f"CKV gather buffer index {buf_idx} is outside " + f"[0, {self._ckv_workspace_slots})" + ) + records = ckv_workspace.view(-1, self._kv_record_bytes) + local_buffer = records[: self._ckv_local_capacity] + gathered_base = ( + self._ckv_local_capacity + + buf_idx * self.dcp_world_size * self._ckv_local_capacity + ) + gathered_buffer = records[ + gathered_base : gathered_base + + self.dcp_world_size * self._ckv_local_capacity + ] + return local_buffer, gathered_buffer + def dcp_prefill_ckv_gather_eligible( self, attn_metadata: B12xMLASparseMetadata, @@ -1633,6 +2100,326 @@ def dcp_prefill_ckv_gather_eligible( ) ) + def dcp_sparse_decode_ckv_gather_eligible( + self, + attn_metadata: B12xMLASparseMetadata, + num_tokens: int, + ) -> bool: + """Select bounded, uniform-row decode batches for sparse CKV exchange.""" + if self._sparse_decode_layout is None: + return False + from vllm.v1.attention.backends.mla.b12x_sparse_ckv_decode import ( + sparse_decode_batch_eligible, + ) + + has_required_metadata = all( + tensor is not None + for tensor in ( + attn_metadata.req_id_per_token, + attn_metadata.page_table_1, + attn_metadata.nsa_cache_seqlens, + attn_metadata.global_cache_seq_lens_per_req, + attn_metadata.ckv_prefetch_registry, + ) + ) + eligible = sparse_decode_batch_eligible( + self._sparse_decode_layout, + num_requests=int(attn_metadata.num_reqs), + num_rows=int(num_tokens), + max_query_len=int(attn_metadata.max_query_len), + num_actual_tokens=int(attn_metadata.num_actual_tokens), + has_required_metadata=has_required_metadata, + ) + if ( + self._sparse_decode_enabled + and int(attn_metadata.num_reqs) + > self._sparse_decode_layout.max_fast_requests + ): + logger.info_once( + "Selected-record CKV decode falling back to stock DCP: " + "active requests=%d exceeds pooled maximum=%d", + int(attn_metadata.num_reqs), + self._sparse_decode_layout.max_fast_requests, + ) + return bool(self._sparse_decode_enabled and eligible) + + def _dcp_gather_sparse_decode_ckv( + self, + kv_cache: torch.Tensor, + attn_metadata: B12xMLASparseMetadata, + topk_indices: torch.Tensor, + prefetch_state: _CKVPrefetchState, + *, + buf_idx: int, + build_union: bool, + wait_for_completion: bool, + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.cuda.Event, + ]: + """Exchange one dense per-sequence MTP union through generic B12X P2P.""" + layout = self._sparse_decode_layout + exchange = self._sparse_decode_exchange + if layout is None or exchange is None: + raise RuntimeError("selected-record CKV decode was not initialized") + num_rows = int(topk_indices.shape[0]) + num_requests = int(attn_metadata.num_reqs) + rows_per_request = num_rows // num_requests + active_records = layout.active_records(num_requests) + if not self.dcp_sparse_decode_ckv_gather_eligible(attn_metadata, num_rows): + raise RuntimeError( + "selected-record CKV gather received an ineligible batch" + ) + if not 0 <= int(buf_idx) < layout.workspace_slots: + raise ValueError( + f"selected-record workspace slot {buf_idx} is outside " + f"[0, {layout.workspace_slots})" + ) + if ( + kv_cache.dtype != torch.uint8 + or kv_cache.ndim != 3 + or tuple(kv_cache.shape[1:]) != (self.block_size, layout.record_bytes) + or not kv_cache.is_contiguous() + ): + raise ValueError( + "selected-record CKV decode requires contiguous native paged cache" + ) + + assert attn_metadata.req_id_per_token is not None + assert attn_metadata.page_table_1 is not None + assert attn_metadata.nsa_cache_seqlens is not None + assert attn_metadata.global_cache_seq_lens_per_req is not None + state = prefetch_state.get_sparse_decode_state(layout, exchange) + output = state.payload_workspace[buf_idx, :active_records] + caller_stream = torch.cuda.current_stream() + + from vllm.distributed.parallel_state import get_dcp_ckv_prefetch_group + from vllm.v1.attention.backends.mla.b12x_sparse_ckv_decode import ( + build_dense_union_remap, + get_selected_record_stream, + ) + + gather_stream = get_selected_record_stream(exchange) + gather_stream.wait_stream(caller_stream) + output.record_stream(gather_stream) + + with torch.cuda.stream(gather_stream): + if build_union: + local_topk = topk_indices[:num_rows, : layout.topk] + all_rank_topk = state.all_rank_topk[: layout.dcp_world_size * num_rows] + _dcp_all_gather_current_stream( + get_dcp_ckv_prefetch_group(), + local_topk, + all_rank_topk, + ) + destination_topk = all_rank_topk.view( + layout.dcp_world_size, + num_requests, + rows_per_request, + layout.topk, + ) + build_dense_union_remap( + destination_topk, + state.union_indices, + state.remap, + state.union_counts, + state.hash_keys, + state.hash_first_positions, + state.first_to_dense, + num_requests=num_requests, + ) + for request in range(num_requests): + row_start = request * rows_per_request + state.selected_indices[ + row_start : row_start + rows_per_request + ].copy_( + state.remap[ + self.dcp_rank, + request, + :rows_per_request, + ], + non_blocking=True, + ) + selected = state.selected_indices[:num_rows, : layout.topk] + _offset_sparse_decode_remap( + selected, + request_stride=layout.per_request_capacity, + rows_per_request=rows_per_request, + topk=layout.topk, + ) + + transport_indices = state.transport_local_slots[num_requests] + for destination in range(layout.dcp_world_size): + destination_slots = state.local_slots[destination, :num_requests] + triton_filter_and_convert_dcp_index( + state.request_ids[:num_requests], + attn_metadata.block_table, + state.union_indices[destination, :num_requests], + dcp_size=layout.dcp_world_size, + dcp_rank=self.dcp_rank, + cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size, + BLOCK_SIZE=attn_metadata.block_size, + NUM_TOPK_TOKENS=layout.per_request_capacity, + compact_valid_to_front=False, + out=destination_slots, + ) + transport_indices[destination].copy_( + destination_slots.reshape(-1), non_blocking=True + ) + exchange.exchange( + kv_cache.reshape(-1, layout.record_bytes), + transport_indices, + output, + ) + complete = state.complete_events[buf_idx] + complete.record(gather_stream) + + if wait_for_completion: + caller_stream.wait_event(complete) + selected_indices = state.selected_indices[:num_rows, : layout.topk] + nsa_cache_seqlens = attn_metadata.nsa_cache_seqlens[:num_rows] + if wait_for_completion: + global_causal_lens = _global_causal_lens_for_ckv_gather( + attn_metadata.global_cache_seq_lens_per_req, + attn_metadata.query_start_loc, + attn_metadata.req_id_per_token, + num_rows, + ) + nsa_cache_seqlens.copy_(global_causal_lens, non_blocking=True) + nsa_cache_seqlens.clamp_(max=layout.topk) + _mask_page_table_after_nsa_len(selected_indices, nsa_cache_seqlens) + gathered_cache = output.view(-1, self.block_size, layout.record_bytes) + return gathered_cache, selected_indices, nsa_cache_seqlens, complete + + def _dcp_gather_sparse_decode_shared_layers( + self, + targets: list[tuple[int, "B12xMLASparseImpl", torch.Tensor]], + attn_metadata: B12xMLASparseMetadata, + prefetch_state: _CKVPrefetchState, + ) -> tuple[dict[int, int], torch.cuda.Event] | None: + """Exchange one Full layer's S1/S2/S3 CKV payloads together.""" + layout = self._sparse_decode_layout + exchange = self._sparse_decode_exchange + if ( + not self._sparse_decode_bulk_prefetch + or layout is None + or exchange is None + or len(targets) != 3 + ): + return None + + num_requests = int(attn_metadata.num_reqs) + active_records = layout.active_records(num_requests) + state = prefetch_state.get_sparse_decode_state(layout, exchange) + layer_slots: dict[int, int] = {} + records_by_layer: list[torch.Tensor] = [] + outputs_by_layer: list[torch.Tensor] = [] + for target_idx, target_impl, target_kv in targets: + if ( + target_impl._sparse_decode_layout != layout + or target_impl._sparse_decode_exchange is not exchange + or target_impl.block_size != self.block_size + or target_impl.cp_kv_cache_interleave_size + != self.cp_kv_cache_interleave_size + or target_kv.dtype != torch.uint8 + or target_kv.ndim != 3 + or tuple(target_kv.shape[1:]) != (self.block_size, layout.record_bytes) + or not target_kv.is_contiguous() + ): + return None + buf_idx = target_idx % layout.workspace_slots + if buf_idx in layer_slots.values(): + return None + layer_slots[target_idx] = buf_idx + records_by_layer.append(target_kv.reshape(-1, layout.record_bytes)) + outputs_by_layer.append(state.payload_workspace[buf_idx, :active_records]) + + from vllm import envs as envs_mod + from vllm.v1.attention.backends.mla.b12x_sparse_ckv_decode import ( + get_selected_record_stream, + try_exchange_sparse_decode_shared_layers, + ) + + transport = envs_mod.VLLM_B12X_MLA_SPARSE_DECODE_TRANSPORT + if transport not in {"auto", "ce"}: + return None + if not callable(getattr(exchange, "exchange_layers", None)): + return None + layered_capacity = getattr(exchange, "layered_max_records", None) + if layered_capacity is not None and active_records > int(layered_capacity): + return None + + gather_stream = get_selected_record_stream(exchange) + gather_stream.wait_stream(torch.cuda.current_stream()) + for output in outputs_by_layer: + output.record_stream(gather_stream) + + with torch.cuda.stream(gather_stream): + used_bulk = try_exchange_sparse_decode_shared_layers( + enabled=True, + transport=transport, + exchange=exchange, + records_by_layer=tuple(records_by_layer), + local_indices_by_destination=( + state.transport_local_slots[num_requests] + ), + outputs_by_layer=tuple(outputs_by_layer), + active_records=active_records, + pool_records=layout.pool_records, + ) + if not used_bulk: + return None + complete = state.complete_events[next(iter(layer_slots.values()))] + complete.record(gather_stream) + return layer_slots, complete + + def _append_current_token_to_sparse_decode_gathered( + self, + gathered_cache: torch.Tensor, + attn_metadata: B12xMLASparseMetadata, + union_indices: torch.Tensor, + global_causal_lens: torch.Tensor, + patch_slots: torch.Tensor, + layer: AttentionLayer, + ) -> None: + """Patch records written after a Shared layer's history prefetch.""" + if self._ckv_current_chunk_kv_c is None: + raise RuntimeError("selected-record prefetch is missing current CKV") + if self._ckv_current_chunk_kpe is None: + raise RuntimeError("selected-record prefetch is missing current RoPE") + if attn_metadata.req_id_per_token is None: + raise RuntimeError("selected-record prefetch is missing request ids") + num_rows = int(global_causal_lens.numel()) + _find_sparse_decode_current_token_slot( + union_indices, + attn_metadata.req_id_per_token[:num_rows], + global_causal_lens, + patch_slots[:num_rows], + ) + k_pe = self._ckv_current_chunk_kpe[:num_rows] + if k_pe.ndim == 3: + k_pe = k_pe.squeeze(1) + k_scale = getattr(layer, "_k_scale", None) + if self._kv_fp8_rope: + self._concat_and_cache_nvfp4_mla_fp8_rope( + self._ckv_current_chunk_kv_c[:num_rows], + k_pe, + gathered_cache, + patch_slots[:num_rows], + k_scale, + ) + else: + ops.concat_and_cache_mla( + self._ckv_current_chunk_kv_c[:num_rows], + k_pe, + gathered_cache, + patch_slots[:num_rows], + self.kv_cache_dtype, + k_scale, + ) + def _dcp_gather_ckv( self, kv_cache: torch.Tensor, @@ -1658,18 +2445,11 @@ def _dcp_gather_ckv( assert attn_metadata.dcp_local_cu_seq_lens is not None padded_tokens = attn_metadata.dcp_padded_total_tokens local_tokens = attn_metadata.dcp_local_total_tokens - self._validate_ckv_workspace(ckv_workspace) - half_nbytes = ( - (self.dcp_world_size + 1) * self._ckv_local_capacity * self._kv_record_bytes + local_buffer, gathered_buffer = self._ckv_workspace_views( + ckv_workspace, buf_idx ) - ws_half = ckv_workspace.view(-1, self._kv_record_bytes) - base = buf_idx * (half_nbytes // self._kv_record_bytes) - local_buffer = ws_half[base : base + self._ckv_local_capacity] - gathered_buffer = ws_half[ - base + self._ckv_local_capacity : base - + self._ckv_local_capacity * (self.dcp_world_size + 1) - ] if stream is not None: + ckv_workspace.record_stream(stream) # The side stream must observe the default stream's prior writes # to the paged KV cache (this and earlier steps' do_kv_cache_update) # before gathering history off it; there is otherwise no ordering @@ -1767,7 +2547,7 @@ def _append_current_chunk_to_gathered( if self._ckv_current_chunk_kv_c is None or num_actual_toks == 0: return kv_c = self._ckv_current_chunk_kv_c[:num_actual_toks] - k_pe_flat = self._ckv_current_chunk_kpe[:num_actual_toks] + k_pe_flat = cast(torch.Tensor, self._ckv_current_chunk_kpe)[:num_actual_toks] if k_pe_flat.ndim == 3: k_pe_flat = k_pe_flat.squeeze(1) @@ -1777,7 +2557,8 @@ def _append_current_chunk_to_gathered( if global_seq_lens is None: return global_seq_lens = global_seq_lens[:num_reqs] - req_ids = attn_metadata.req_id_per_token[:num_actual_toks].to(torch.int64) + req_id_per_token = cast(torch.Tensor, attn_metadata.req_id_per_token) + req_ids = req_id_per_token[:num_actual_toks].to(torch.int64) global_seq_per_token = global_seq_lens[req_ids].to(torch.int32) # Map each current-chunk token to its absolute position within its @@ -1798,7 +2579,7 @@ def _append_current_chunk_to_gathered( + global_pos % interleave ).to(torch.int64) - rank_req_starts = attn_metadata.dcp_rank_req_starts + rank_req_starts = cast(torch.Tensor, attn_metadata.dcp_rank_req_starts) flat_idx = owner * num_reqs + req_ids # ``dcp_rank_req_starts`` is a ``[dcp, :num_reqs]`` slice of a # ``[dcp, max_seqs]`` buffer, so it is non-contiguous whenever @@ -1818,7 +2599,7 @@ def _append_current_chunk_to_gathered( slots, k_scale, ) - elif self.kv_cache_dtype == "fp8_ds_mla": + elif self.kv_cache_dtype in ("fp8_ds_mla", "nvfp4_ds_mla"): ops.concat_and_cache_mla( kv_c, k_pe_flat, @@ -1831,7 +2612,7 @@ def _append_current_chunk_to_gathered( raise RuntimeError( "CKV gather prefetch append is not yet supported for " f"kv_cache_dtype={self.kv_cache_dtype!r}; disable prefetch " - "or use fp8_ds_mla / KV_FP8_ROPE." + "or use fp8_ds_mla / nvfp4_ds_mla." ) def _sync_warmup(self) -> None: @@ -1958,15 +2739,18 @@ def forward_mqa( use_ckv_gather = self.dcp_prefill_ckv_gather_eligible( attn_metadata, int(query_rows) ) + use_sparse_decode_ckv_gather = self.dcp_sparse_decode_ckv_gather_eligible( + attn_metadata, int(query_rows) + ) + use_local_query_heads = use_ckv_gather or use_sparse_decode_ckv_gather workspace_tensors = self._borrow_workspaces() q_workspace = workspace_tensors[0] dense_out_workspace = workspace_tensors[1] if self._pad_heads else None - ckv_workspace = self._ckv_workspace scratch_storage = workspace_tensors[-1] expected_input_heads = ( - self.num_heads if use_ckv_gather else self._input_num_heads + self.num_heads if use_local_query_heads else self._input_num_heads ) - if use_ckv_gather: + if use_local_query_heads: local_q_numel = ( self._max_batched * self._ckv_kernel_num_heads * self.q_head_dim ) @@ -2013,6 +2797,16 @@ def forward_mqa( assert self.topk_indices_buffer is not None topk_indices = self.topk_indices_buffer[:num_actual_toks] per_token_cache = attn_metadata.cache_seq_lens_per_token[:num_actual_toks] + sparse_decode_global_causal_lens = None + if use_sparse_decode_ckv_gather: + assert attn_metadata.global_cache_seq_lens_per_req is not None + assert attn_metadata.req_id_per_token is not None + sparse_decode_global_causal_lens = _global_causal_lens_for_ckv_gather( + attn_metadata.global_cache_seq_lens_per_req, + attn_metadata.query_start_loc, + attn_metadata.req_id_per_token, + num_actual_toks, + ) if use_ckv_gather: assert attn_metadata.req_id_per_token is not None assert attn_metadata.ckv_page_table_1 is not None @@ -2055,6 +2849,12 @@ def forward_mqa( out=nsa_cache_seqlens, ) _mask_page_table_after_nsa_len(selected_indices, nsa_cache_seqlens) + elif use_sparse_decode_ckv_gather: + # The dense remap is materialized after the current layer's native + # cache is available below. Keep placeholders here so the stock + # rank-local conversion is not entered. + selected_indices = topk_indices + nsa_cache_seqlens = per_token_cache elif self.dcp_world_size > 1: # The indexer globally merges logical top-k ids across DCP ranks. # Compact just this rank's winners into local physical cache slots; @@ -2120,29 +2920,174 @@ def forward_mqa( "B12X_MLA_SPARSE requires a contiguous native paged KV cache; " f"got stride={tuple(kv_cache.stride())}" ) - if use_ckv_gather: - if ckv_workspace is None: - raise RuntimeError("CKV gather workspace was not borrowed") + if use_sparse_decode_ckv_gather: + layout = self._sparse_decode_layout + exchange = self._sparse_decode_exchange + if layout is None or exchange is None: + raise RuntimeError("selected-record CKV decode was not initialized") + if sparse_decode_global_causal_lens is None: + raise RuntimeError( + "selected-record CKV decode is missing global causal lengths" + ) layer_idx = self._resolve_layer_index(layer) + registry = attn_metadata.ckv_prefetch_registry + if registry is None: + raise RuntimeError( + "selected-record CKV decode is missing its lane registry" + ) + prefetch_state = registry.for_workspace( + q_workspace, + layer_idx, + kv_cache, + execution_lane_key=self._ckv_execution_lane_key, + ) if layer_idx is not None: - while len(B12xMLASparseImpl._all_layer_kv_caches) <= layer_idx: - B12xMLASparseImpl._all_layer_kv_caches.append(None) - B12xMLASparseImpl._all_layer_kv_caches[layer_idx] = kv_cache - if B12xMLASparseImpl._shared_gather_event is not None: - B12xMLASparseImpl._shared_gather_event.wait() - half_nbytes = ( - (self.dcp_world_size + 1) - * self._ckv_local_capacity - * self._kv_record_bytes + prefetch_state.enter_layer(layer_idx) + prefetch_state.register_cache(layer_idx, kv_cache) + prefetch_state.register_impl(layer_idx, self) + sparse_state = prefetch_state.get_sparse_decode_state(layout, exchange) + pending = prefetch_state.pop_pending_layer(layer_idx) + if pending is not None: + pending.ticket.wait_on_stream_once(torch.cuda.current_stream()) + current_buf_idx = pending.buf_idx + active_records = layout.active_records(int(attn_metadata.num_reqs)) + kv_cache = sparse_state.payload_workspace[ + current_buf_idx, :active_records + ].view(-1, self.block_size, layout.record_bytes) + self._append_current_token_to_sparse_decode_gathered( + kv_cache, + attn_metadata, + sparse_state.union_indices[ + self.dcp_rank, : int(attn_metadata.num_reqs) + ], + sparse_decode_global_causal_lens, + sparse_state.patch_slots, + layer, + ) + selected_indices = sparse_state.selected_indices[ + :num_actual_toks, : layout.topk + ] + nsa_cache_seqlens = cast(torch.Tensor, attn_metadata.nsa_cache_seqlens)[ + :num_actual_toks + ] + nsa_cache_seqlens.copy_( + sparse_decode_global_causal_lens, non_blocking=True ) - ws_half = ckv_workspace.view(-1, self._kv_record_bytes) - base = B12xMLASparseImpl._shared_gather_buf_idx * ( - half_nbytes // self._kv_record_bytes + nsa_cache_seqlens.clamp_(max=layout.topk) + _mask_page_table_after_nsa_len(selected_indices, nsa_cache_seqlens) + else: + current_buf_idx = ( + layer_idx % layout.workspace_slots if layer_idx is not None else 0 ) - gathered_buffer = ws_half[ - base + self._ckv_local_capacity : base - + self._ckv_local_capacity * (self.dcp_world_size + 1) + ( + kv_cache, + selected_indices, + nsa_cache_seqlens, + _, + ) = self._dcp_gather_sparse_decode_ckv( + kv_cache, + attn_metadata, + topk_indices, + prefetch_state, + buf_idx=current_buf_idx, + build_union=True, + wait_for_completion=True, + ) + + logger.info_once( + "Using selected-record CKV decode for C<=%d/MTP%d " + "(DCP%d topk=%d pool=%d record_bytes=%d depth=%d)", + layout.max_fast_requests, + layout.rows_per_request - 1, + layout.dcp_world_size, + layout.topk, + layout.pool_records, + layout.record_bytes, + layout.prefetch_depth, + ) + if ( + self._sparse_decode_emits_topk + and layer_idx is not None + and layout.prefetch_depth > 0 + ): + from vllm.v1.attention.backends.mla.b12x_sparse_ckv_decode import ( + sparse_decode_prefetch_targets, + ) + + emits_topk_by_layer = [ + None if impl is None else impl._sparse_decode_emits_topk + for impl in prefetch_state.layer_impls ] + targets = sparse_decode_prefetch_targets( + layer_idx, + layout.prefetch_depth, + emits_topk_by_layer, + ) + target_entries: list[tuple[int, B12xMLASparseImpl, torch.Tensor]] = [] + for target_idx in targets: + if target_idx in prefetch_state.pending_layers: + continue + target_impl = prefetch_state.layer_impls[target_idx] + target_kv = prefetch_state.layer_caches[target_idx] + if target_impl is None or target_kv is None: + break + target_entries.append((target_idx, target_impl, target_kv)) + + bulk_pending = self._dcp_gather_sparse_decode_shared_layers( + target_entries, + attn_metadata, + prefetch_state, + ) + if bulk_pending is not None: + layer_slots, target_event = bulk_pending + prefetch_state.register_pending_group( + layer_slots, + target_event, + ) + else: + for target_idx, target_impl, target_kv in target_entries: + target_buf_idx = target_idx % layout.workspace_slots + _, _, _, target_event = ( + target_impl._dcp_gather_sparse_decode_ckv( + target_kv, + attn_metadata, + topk_indices, + prefetch_state, + buf_idx=target_buf_idx, + build_union=False, + wait_for_completion=False, + ) + ) + prefetch_state.register_pending_group( + {target_idx: target_buf_idx}, + target_event, + ) + if use_ckv_gather: + layer_idx = self._resolve_layer_index(layer) + prefetch_registry = attn_metadata.ckv_prefetch_registry + if prefetch_registry is None: + raise RuntimeError("CKV gather requires a prefetch state registry") + prefetch_state = prefetch_registry.for_workspace( + q_workspace, + layer_idx, + kv_cache, + execution_lane_key=self._ckv_execution_lane_key, + ) + ckv_workspace = prefetch_state.get_ckv_workspace(self._ckv_workspace_nbytes) + if layer_idx is not None and self._ckv_prefetch_depth > 0: + prefetch_state.enter_layer(layer_idx) + prefetch_state.register_cache(layer_idx, kv_cache) + pending = ( + prefetch_state.pop_pending_layer(layer_idx) + if self._ckv_prefetch_depth > 0 + else None + ) + if pending is not None: + pending.ticket.wait_once() + current_buf_idx = pending.buf_idx + _, gathered_buffer = self._ckv_workspace_views( + ckv_workspace, current_buf_idx + ) kv_cache = gathered_buffer[ : self.dcp_world_size * self._ckv_local_capacity ].view(-1, self.block_size, self._kv_record_bytes) @@ -2150,7 +3095,17 @@ def forward_mqa( kv_cache, attn_metadata, layer, num_actual_toks ) else: - kv_cache = self._dcp_gather_ckv(kv_cache, attn_metadata, ckv_workspace) + current_buf_idx = ( + layer_idx % self._ckv_workspace_slots + if layer_idx is not None + else 0 + ) + kv_cache = self._dcp_gather_ckv( + kv_cache, + attn_metadata, + ckv_workspace, + buf_idx=current_buf_idx, + ) logger.info_once( "Using transient full-CKV gather for B12X sparse MLA prefill " "(capacity=%d logical tokens)", @@ -2158,52 +3113,86 @@ def forward_mqa( ) if ( self._ckv_prefetch_supported + and self._ckv_prefetch_depth > 0 and layer_idx is not None - and layer_idx + 1 < len(B12xMLASparseImpl._all_layer_kv_caches) - and B12xMLASparseImpl._all_layer_kv_caches[layer_idx + 1] is not None ): - next_kv = B12xMLASparseImpl._all_layer_kv_caches[layer_idx + 1] - next_buf_idx = 1 - B12xMLASparseImpl._shared_gather_buf_idx - self._dcp_gather_ckv( - next_kv, - attn_metadata, - ckv_workspace, - buf_idx=next_buf_idx, - stream=self._ckv_gather_stream, + # The first eligible request only discovers one layer cache at + # a time. It intentionally has no lookahead and primes the + # cache registry for subsequent requests. + targets = _ckv_prefetch_target_indices( + layer_idx, + self._ckv_prefetch_depth, + prefetch_state.layer_caches, + prefetch_state.pending_layers, ) - B12xMLASparseImpl._shared_gather_event = torch.cuda.Event( - blocking=False + prefetch_stream = ( + prefetch_state.get_gather_stream() if targets else None ) - B12xMLASparseImpl._shared_gather_event.record(self._ckv_gather_stream) - B12xMLASparseImpl._shared_gather_buf_idx = next_buf_idx - - use_decode_kernel = attn_metadata.max_query_len <= 1 or ( - self.spec_extend_as_decode - and attn_metadata.max_query_len <= self.spec_decode_max_q - and num_actual_toks <= attn_metadata.num_reqs * self.spec_decode_max_q - and num_actual_toks <= self._decode_max_rows + for target_idx in targets: + assert prefetch_stream is not None + target_kv = prefetch_state.layer_caches[target_idx] + assert target_kv is not None + target_buf_idx = target_idx % self._ckv_workspace_slots + self._dcp_gather_ckv( + target_kv, + attn_metadata, + ckv_workspace, + buf_idx=target_buf_idx, + stream=prefetch_stream, + ) + target_event = torch.cuda.Event(blocking=False) + target_event.record(prefetch_stream) + prefetch_state.register_pending_group( + {target_idx: target_buf_idx}, + target_event, + ) + + use_decode_kernel = ( + use_sparse_decode_ckv_gather + or (attn_metadata.max_query_len <= 1) + or ( + self.spec_extend_as_decode + and attn_metadata.max_query_len <= self.spec_decode_max_q + and num_actual_toks <= attn_metadata.num_reqs * self.spec_decode_max_q + and num_actual_toks <= self._decode_max_rows + ) ) if use_decode_kernel: cache_seqlens = ( - attn_metadata.cache_seq_lens_per_req + sparse_decode_global_causal_lens + if use_sparse_decode_ckv_gather + else attn_metadata.cache_seq_lens_per_req if attn_metadata.max_query_len <= 1 else attn_metadata.cache_seq_lens_per_token[:num_actual_toks] ) + if cache_seqlens is None: + raise RuntimeError( + "selected-record CKV decode is missing global causal lengths" + ) decode_q = q_all - if self._pad_heads: + if self._pad_heads and not use_sparse_decode_ckv_gather: decode_q = q_buffer[:, : self._kernel_num_heads] decode_q[:, self._input_num_heads :, :].zero_() # Eager bind maps caller-owned scratch into views. forced_num_splits # pins the planner choice for this captured graph; the merge kernel is # specialized on that count and needs no device-side control fill. - binding = self._decode_plan.bind( + decode_plan = ( + self._sparse_decode_plan + if use_sparse_decode_ckv_gather + else self._decode_plan + ) + if decode_plan is None: + raise RuntimeError( + "selected-record CKV decode plan was not initialized" + ) + binding = decode_plan.bind( scratch=scratch_storage, q=decode_q, selected_indices=selected_indices, cache_seqlens_int32=cache_seqlens, nsa_cache_seqlens_int32=nsa_cache_seqlens, ) - if self.need_to_return_lse_for_decode: + if self.need_to_return_lse_for_decode and not use_sparse_decode_ckv_gather: out, lse = cast( tuple[torch.Tensor, torch.Tensor], self._sparse_mla_decode_forward( @@ -2235,7 +3224,7 @@ def forward_mqa( **kernel_format_kwargs, ), ) - if self._pad_heads: + if self._pad_heads and not use_sparse_decode_ckv_gather: assert dense_out_workspace is not None dense_out = dense_out_workspace[:num_actual_toks] dense_out.copy_(out[:, : self._input_num_heads, :]) diff --git a/vllm/v1/attention/backends/mla/b12x_sparse_ckv_decode.py b/vllm/v1/attention/backends/mla/b12x_sparse_ckv_decode.py new file mode 100644 index 000000000000..44609ba7df7d --- /dev/null +++ b/vllm/v1/attention/backends/mla/b12x_sparse_ckv_decode.py @@ -0,0 +1,527 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Planning and deterministic union construction for sparse DCP CKV decode. + +The B12X transport only exchanges fixed-width records. This module owns the +vLLM policy which decides which records each request needs and how MTP rows map +into one per-request, densely packed sparse workspace. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from functools import lru_cache +from pathlib import Path +from typing import Protocol + +import torch +import torch.distributed as dist +from torch.distributed import ProcessGroup +from torch.utils.cpp_extension import load + +from vllm import envs +from vllm.distributed.parallel_state import register_model_parallel_cleanup_hook +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +@dataclass(frozen=True) +class SparseCKVDecodeLayout: + dcp_world_size: int + topk: int + rows_per_request: int + max_requests: int + per_request_capacity: int + pool_records: int + max_fast_requests: int + record_bytes: int + prefetch_depth: int + workspace_slots: int + + @property + def workspace_bytes(self) -> int: + return self.workspace_slots * self.pool_records * self.record_bytes + + def active_records(self, num_requests: int) -> int: + if not 1 <= num_requests <= self.max_fast_requests: + raise ValueError( + f"request count {num_requests} is outside sparse capacity " + f"[1, {self.max_fast_requests}]" + ) + return num_requests * self.per_request_capacity + + +def plan_sparse_ckv_decode( + *, + dcp_world_size: int, + topk: int, + rows_per_request: int, + max_requests: int, + pool_records: int, + record_bytes: int, + prefetch_depth: int, +) -> SparseCKVDecodeLayout: + """Return the static pooled-workspace layout for selected-record decode.""" + if not 2 <= dcp_world_size <= 8: + raise ValueError("selected-record CKV decode supports DCP world sizes 2-8") + if min(topk, rows_per_request, max_requests, record_bytes) <= 0: + raise ValueError("topk, rows, requests, and record width must be positive") + if prefetch_depth < 0: + raise ValueError("prefetch depth cannot be negative") + + per_request_capacity = topk * rows_per_request + requested_pool = int(pool_records) + if requested_pool <= 0: + requested_pool = max_requests * per_request_capacity + max_fast_requests = min(max_requests, requested_pool // per_request_capacity) + if max_fast_requests <= 0: + raise ValueError( + "selected-record pool cannot hold one request's worst-case MTP union" + ) + # Capacity beyond the largest schedulable batch only wastes VRAM. + effective_pool = min( + requested_pool, + max_requests * per_request_capacity, + ) + return SparseCKVDecodeLayout( + dcp_world_size=dcp_world_size, + topk=topk, + rows_per_request=rows_per_request, + max_requests=max_requests, + per_request_capacity=per_request_capacity, + pool_records=effective_pool, + max_fast_requests=max_fast_requests, + record_bytes=record_bytes, + prefetch_depth=prefetch_depth, + workspace_slots=prefetch_depth + 1, + ) + + +def sparse_decode_batch_eligible( + layout: SparseCKVDecodeLayout, + *, + num_requests: int, + num_rows: int, + max_query_len: int, + num_actual_tokens: int, + has_required_metadata: bool, +) -> bool: + """Pure policy gate used before entering the selected-record fast path.""" + if not has_required_metadata or not 1 <= num_requests <= layout.max_fast_requests: + return False + if not 1 <= num_rows <= num_requests * layout.rows_per_request: + return False + if num_rows % num_requests: + return False + rows_per_request = num_rows // num_requests + return ( + rows_per_request <= layout.rows_per_request + and max_query_len == rows_per_request + and num_actual_tokens == num_rows + ) + + +def sparse_decode_prefetch_targets( + layer_index: int, + depth: int, + emits_topk_by_layer: list[bool | None], +) -> list[int]: + """Return consecutive Shared layers following one Full/indexer layer.""" + targets: list[int] = [] + for distance in range(1, max(0, depth) + 1): + target = layer_index + distance + if target >= len(emits_topk_by_layer): + break + emits_topk = emits_topk_by_layer[target] + if emits_topk is None or emits_topk: + break + targets.append(target) + return targets + + +def try_exchange_sparse_decode_shared_layers( + *, + enabled: bool, + transport: str, + exchange: object, + records_by_layer: tuple[torch.Tensor, ...], + local_indices_by_destination: torch.Tensor, + outputs_by_layer: tuple[torch.Tensor, ...], + active_records: int, + pool_records: int, +) -> bool: + """Use one copy-engine exchange for a Full layer's S1/S2/S3 payloads.""" + if not enabled: + return False + if transport not in {"auto", "ce"}: + logger.info_once( + "Bulk S1/S2/S3 CKV prefetch requires copy-engine transport; " + "using the per-layer exchange path" + ) + return False + if len(records_by_layer) != 3 or len(outputs_by_layer) != 3: + return False + if not 0 < active_records <= pool_records: + raise ValueError( + f"active sparse records {active_records} exceed pool {pool_records}" + ) + exchange_layers = getattr(exchange, "exchange_layers", None) + if exchange_layers is None: + logger.warning_once( + "Bulk S1/S2/S3 CKV prefetch requires B12X " + "PCIeSelectedRecordCopyExchange.exchange_layers; using the " + "per-layer exchange path" + ) + return False + layered_capacity = getattr(exchange, "layered_max_records", None) + if layered_capacity is not None and active_records > int(layered_capacity): + logger.info_once( + "Bulk S1/S2/S3 CKV prefetch capacity %d is below this batch's " + "%d records; using the per-layer exchange path", + int(layered_capacity), + active_records, + ) + return False + exchange_layers( + records_by_layer=records_by_layer, + local_indices_by_destination=local_indices_by_destination, + outputs_by_layer=outputs_by_layer, + ) + logger.info_once("Using one bulk CKV exchange for S1/S2/S3 prefetch") + return True + + +def owner_and_local_ordinal( + logical_token: int, + *, + dcp_world_size: int, + interleave: int, +) -> tuple[int, int]: + """Map one global logical token to its DCP owner and owner-local ordinal.""" + if logical_token < 0: + return -1, -1 + if dcp_world_size <= 0 or interleave <= 0: + raise ValueError("DCP world size and interleave must be positive") + owner = (logical_token // interleave) % dcp_world_size + local_ordinal = ( + logical_token // (dcp_world_size * interleave) * interleave + + logical_token % interleave + ) + return owner, local_ordinal + + +def dense_union_remap_reference( + indices: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """CPU oracle: stable dense union and exact remap for every group. + + ``indices`` has shape ``[..., rows, topk]``. Every leading group receives + an independent stable union in flattened row order. + """ + if indices.device.type != "cpu" or indices.dtype != torch.int32: + raise ValueError("reference union requires a CPU int32 tensor") + if indices.ndim < 2: + raise ValueError("indices must include row and top-k dimensions") + rows, topk = indices.shape[-2:] + capacity = rows * topk + groups = indices.reshape(-1, capacity) + union = torch.full_like(groups, -1) + remap = torch.full_like(groups, -1) + counts = torch.zeros(groups.shape[0], dtype=torch.int32) + for group_index, row in enumerate(groups.tolist()): + slots: dict[int, int] = {} + for input_index, token in enumerate(row): + if token < 0: + continue + slot = slots.get(token) + if slot is None: + slot = len(slots) + slots[token] = slot + union[group_index, slot] = token + remap[group_index, input_index] = slot + counts[group_index] = len(slots) + leading = indices.shape[:-2] + return ( + union.reshape(*leading, capacity), + remap.reshape(indices.shape), + counts.reshape(leading), + ) + + +class SparseCKVDecodeState: + """Persistent graph-safe union/remap state for one execution lane.""" + + def __init__( + self, + *, + layout: SparseCKVDecodeLayout, + device: torch.device, + exchange, + ) -> None: + self.layout = layout + self.exchange = exchange + capacity = layout.per_request_capacity + hash_capacity = 1 << (2 * capacity - 1).bit_length() + max_rows = layout.max_requests * layout.rows_per_request + state_shape = (layout.dcp_world_size, layout.max_requests) + self.all_rank_topk = torch.empty( + (layout.dcp_world_size * max_rows, layout.topk), + dtype=torch.int32, + device=device, + ) + self.union_indices = torch.empty( + (*state_shape, capacity), dtype=torch.int32, device=device + ) + self.remap = torch.empty( + ( + *state_shape, + layout.rows_per_request, + layout.topk, + ), + dtype=torch.int32, + device=device, + ) + self.union_counts = torch.empty(state_shape, dtype=torch.int32, device=device) + self.hash_keys = torch.empty( + (*state_shape, hash_capacity), dtype=torch.int32, device=device + ) + self.hash_first_positions = torch.empty_like(self.hash_keys) + self.first_to_dense = torch.empty_like(self.union_indices) + self.local_slots = torch.empty_like(self.union_indices) + # PCIeSelectedRecordExchange requires one contiguous destination-major + # index plane. A [:num_requests] view of local_slots retains the + # max-request row stride and is therefore non-contiguous. Preallocate + # the eight small exact-shape planes once so decode never allocates or + # copies through a temporary contiguous() tensor during graph replay. + self.transport_local_slots = { + request_count: torch.empty( + ( + layout.dcp_world_size, + request_count * layout.per_request_capacity, + ), + dtype=torch.int32, + device=device, + ) + for request_count in range(1, layout.max_fast_requests + 1) + } + self.selected_indices = torch.empty( + (max_rows, layout.topk), dtype=torch.int32, device=device + ) + # Cross-layer prefetch outlives a WorkspaceManager borrow. Keep the + # payload in lane-owned storage so unrelated scratch users cannot + # overwrite S1/S2/S3 records before attention consumes them. + self.payload_workspace = torch.empty( + ( + layout.workspace_slots, + layout.pool_records, + layout.record_bytes, + ), + dtype=torch.uint8, + device=device, + ) + self.request_ids = torch.arange( + layout.max_requests, dtype=torch.int32, device=device + ) + self.patch_slots = torch.empty(max_rows, dtype=torch.int64, device=device) + self.complete_events = [ + torch.cuda.Event(blocking=False) for _ in range(layout.workspace_slots) + ] + + +class _ClosableSelectedRecordExchange(Protocol): + def close(self) -> None: ... + + +_SELECTED_RECORD_EXCHANGES: dict[ + tuple[object, ...], _ClosableSelectedRecordExchange +] = {} +_SELECTED_RECORD_STREAMS: dict[int, torch.cuda.Stream] = {} +_UNION_PREFLIGHT_RESULTS: dict[tuple[int, str, int | None], tuple[bool, str]] = {} + + +def close_selected_record_exchanges() -> None: + """Release IPC slabs before their DCP process groups are destroyed.""" + exchanges = list( + {id(value): value for value in _SELECTED_RECORD_EXCHANGES.values()}.values() + ) + _SELECTED_RECORD_EXCHANGES.clear() + _SELECTED_RECORD_STREAMS.clear() + _UNION_PREFLIGHT_RESULTS.clear() + + first_error: Exception | None = None + for exchange in exchanges: + try: + exchange.close() + except Exception as exc: # pragma: no cover - defensive shutdown path + if first_error is None: + first_error = exc + if first_error is not None: + raise RuntimeError( + "failed to close a selected-record exchange" + ) from first_error + + +register_model_parallel_cleanup_hook(close_selected_record_exchanges) + + +def get_selected_record_exchange( + *, + process_group: ProcessGroup, + device: torch.device, + layout: SparseCKVDecodeLayout, + lane_key: tuple[str, int], +): + """Return one generic B12X exchange per target/draft execution lane.""" + transport = envs.VLLM_B12X_MLA_SPARSE_DECODE_TRANSPORT + if transport not in {"auto", "ce", "direct"}: + raise ValueError( + "VLLM_B12X_MLA_SPARSE_DECODE_TRANSPORT must be auto, ce, or " + f"direct; got {transport!r}" + ) + key = ( + id(process_group), + device.type, + device.index, + layout.dcp_world_size, + layout.pool_records, + layout.record_bytes, + transport, + lane_key, + ) + exchange = _SELECTED_RECORD_EXCHANGES.get(key) + if exchange is None: + from sparkinfer.comm.pcie import ( + SelectedRecordExchange, + SelectedRecordExchangeInitializationError, + ) + + ce_error: Exception | None = None + if transport in {"auto", "ce"}: + try: + from sparkinfer.comm.pcie import SelectedRecordCopyExchange + + exchange = SelectedRecordCopyExchange.from_process_group( + process_group=process_group, + device=device, + max_records=layout.pool_records, + record_bytes=layout.record_bytes, + ) + except ( + ImportError, + SelectedRecordExchangeInitializationError, + ) as exc: + ce_error = exc + if transport == "ce": + raise RuntimeError( + "strict copy-engine selected-record transport failed " + "to initialize" + ) from exc + if exchange is None: + if ce_error is not None: + logger.warning_once( + "Copy-engine selected-record transport is unavailable; " + "falling back to direct peer writes: %s", + ce_error, + ) + exchange = SelectedRecordExchange.from_process_group( + process_group=process_group, + device=device, + max_records=layout.pool_records, + record_bytes=layout.record_bytes, + ) + logger.info_once( + "Using %s for sparse selected-record CKV exchange", + type(exchange).__name__, + ) + _SELECTED_RECORD_EXCHANGES[key] = exchange + if id(exchange) not in _SELECTED_RECORD_STREAMS: + _SELECTED_RECORD_STREAMS[id(exchange)] = torch.cuda.Stream(device=device) + return exchange + + +def get_selected_record_stream(exchange) -> torch.cuda.Stream: + """Return the dedicated stream created with one stream-affine exchange.""" + stream = _SELECTED_RECORD_STREAMS.get(id(exchange)) + if stream is None: + raise RuntimeError("selected-record exchange has no execution stream") + return stream + + +@lru_cache(maxsize=1) +def _load_union_extension(): + source = Path(__file__).with_name("b12x_sparse_ckv_union.cu") + return load( + name="vllm_b12x_sparse_ckv_union_ext", + sources=[str(source)], + extra_cuda_cflags=["-O3"], + verbose=False, + ) + + +def preload_dense_union_extension() -> None: + """Compile the union helper before CUDA graph warmup or capture.""" + _load_union_extension() + + +def preload_dense_union_extension_consistently( + process_group: ProcessGroup, + device: torch.device, +) -> None: + """Preload on every DCP rank and fail or proceed as one group.""" + key = (id(process_group), device.type, device.index) + cached = _UNION_PREFLIGHT_RESULTS.get(key) + if cached is not None: + ok, message = cached + if not ok: + raise RuntimeError(message) + return + + local_error: Exception | None = None + try: + preload_dense_union_extension() + except Exception as exc: # pragma: no cover - exercised by fault injection + local_error = exc + + status = torch.tensor( + [0 if local_error is not None else 1], + dtype=torch.int32, + device=device, + ) + dist.all_reduce(status, op=dist.ReduceOp.MIN, group=process_group) + globally_ok = bool(status.item()) + if globally_ok: + _UNION_PREFLIGHT_RESULTS[key] = (True, "") + return + + message = ( + f"selected-record CKV union extension failed on this DCP rank: {local_error}" + if local_error is not None + else "selected-record CKV union extension failed on another DCP rank" + ) + _UNION_PREFLIGHT_RESULTS[key] = (False, message) + raise RuntimeError(message) + + +def build_dense_union_remap( + indices: torch.Tensor, + union_indices: torch.Tensor, + remap: torch.Tensor, + union_counts: torch.Tensor, + hash_keys: torch.Tensor, + hash_first_positions: torch.Tensor, + first_to_dense: torch.Tensor, + *, + num_requests: int, +) -> None: + """Build stable dense per-destination/per-request unions on CUDA.""" + _load_union_extension().build_dense_union_remap( + indices, + union_indices, + remap, + union_counts, + hash_keys, + hash_first_positions, + first_to_dense, + int(num_requests), + ) diff --git a/vllm/v1/attention/backends/mla/b12x_sparse_ckv_union.cu b/vllm/v1/attention/backends/mla/b12x_sparse_ckv_union.cu new file mode 100644 index 000000000000..d875b18b5557 --- /dev/null +++ b/vllm/v1/attention/backends/mla/b12x_sparse_ckv_union.cu @@ -0,0 +1,246 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +#include +#include +#include +#include +#include + +#include +#include + +namespace { + +__device__ __forceinline__ uint32_t hash_token_index(uint32_t value) { + value ^= value >> 16; + value *= 0x7feb352dU; + value ^= value >> 15; + value *= 0x846ca68bU; + value ^= value >> 16; + return value; +} + +__device__ __forceinline__ int32_t find_first_position( + int32_t key, const int32_t* hash_keys, const int32_t* hash_first_positions, + int64_t hash_capacity) { + const uint32_t mask = static_cast(hash_capacity - 1); + uint32_t slot = hash_token_index(static_cast(key)) & mask; + for (int64_t probe = 0; probe < hash_capacity; ++probe) { + const int32_t stored = hash_keys[slot]; + if (stored == key) { + return hash_first_positions[slot]; + } + if (stored == -1) { + return -1; + } + slot = (slot + 1) & mask; + } + return -1; +} + +__global__ void build_dense_union_remap_kernel( + const int32_t* __restrict__ indices, int32_t* __restrict__ union_indices, + int32_t* __restrict__ remap, int32_t* __restrict__ union_counts, + int32_t* __restrict__ hash_keys, int32_t* __restrict__ hash_first_positions, + int32_t* __restrict__ first_to_dense, int destinations, int num_requests, + int max_requests, int64_t input_entries, int64_t capacity, + int64_t hash_capacity) { + const int group = static_cast(blockIdx.x); + const int destination = group / num_requests; + const int request = group % num_requests; + if (destination >= destinations) { + return; + } + const int state_group = destination * max_requests + request; + const int64_t input_base = static_cast(group) * input_entries; + const int64_t state_base = static_cast(state_group) * capacity; + const int64_t hash_base = static_cast(state_group) * hash_capacity; + + for (int64_t i = threadIdx.x; i < hash_capacity; i += blockDim.x) { + hash_keys[hash_base + i] = -1; + hash_first_positions[hash_base + i] = std::numeric_limits::max(); + } + for (int64_t i = threadIdx.x; i < capacity; i += blockDim.x) { + union_indices[state_base + i] = -1; + first_to_dense[state_base + i] = -1; + } + for (int64_t i = threadIdx.x; i < input_entries; i += blockDim.x) { + remap[state_base + i] = -1; + } + if (threadIdx.x == 0) { + union_counts[state_group] = 0; + } + __syncthreads(); + + const uint32_t mask = static_cast(hash_capacity - 1); + for (int64_t i = threadIdx.x; i < input_entries; i += blockDim.x) { + const int32_t key = indices[input_base + i]; + if (key < 0) { + continue; + } + uint32_t slot = hash_token_index(static_cast(key)) & mask; + for (int64_t probe = 0; probe < hash_capacity; ++probe) { + const int32_t previous = atomicCAS(hash_keys + hash_base + slot, -1, key); + if (previous == -1 || previous == key) { + atomicMin(hash_first_positions + hash_base + slot, + static_cast(i)); + break; + } + slot = (slot + 1) & mask; + } + } + __syncthreads(); + + __shared__ int32_t first_flags[256]; + __shared__ int32_t running_count; + if (threadIdx.x == 0) { + running_count = 0; + } + __syncthreads(); + + for (int64_t tile = 0; tile < input_entries; tile += blockDim.x) { + const int64_t i = tile + threadIdx.x; + int32_t is_first = 0; + int32_t key = -1; + if (i < input_entries) { + key = indices[input_base + i]; + if (key >= 0) { + is_first = find_first_position(key, hash_keys + hash_base, + hash_first_positions + hash_base, + hash_capacity) == i; + } + } + first_flags[threadIdx.x] = is_first; + __syncthreads(); + if (threadIdx.x == 0) { + int32_t prefix = running_count; + const int64_t remaining = input_entries - tile; + const int32_t valid = + static_cast(remaining < static_cast(blockDim.x) + ? remaining + : static_cast(blockDim.x)); + for (int32_t lane = 0; lane < valid; ++lane) { + const int32_t flag = first_flags[lane]; + first_flags[lane] = prefix; + prefix += flag; + } + running_count = prefix; + } + __syncthreads(); + if (i < input_entries && is_first) { + const int32_t dense_slot = first_flags[threadIdx.x]; + first_to_dense[state_base + i] = dense_slot; + union_indices[state_base + dense_slot] = key; + } + __syncthreads(); + } + + if (threadIdx.x == 0) { + union_counts[state_group] = running_count; + } + __syncthreads(); + + for (int64_t i = threadIdx.x; i < input_entries; i += blockDim.x) { + const int32_t key = indices[input_base + i]; + if (key < 0) { + continue; + } + const int32_t first = + find_first_position(key, hash_keys + hash_base, + hash_first_positions + hash_base, hash_capacity); + if (first >= 0) { + remap[state_base + i] = first_to_dense[state_base + first]; + } + } +} + +void validate_cuda_int32_contiguous(const torch::Tensor& value, + const char* name) { + TORCH_CHECK(value.is_cuda(), name, " must be CUDA"); + TORCH_CHECK(value.scalar_type() == torch::kInt32, name, " must be int32"); + TORCH_CHECK(value.is_contiguous(), name, " must be contiguous"); +} + +void build_dense_union_remap(torch::Tensor indices, torch::Tensor union_indices, + torch::Tensor remap, torch::Tensor union_counts, + torch::Tensor hash_keys, + torch::Tensor hash_first_positions, + torch::Tensor first_to_dense, + int64_t num_requests) { + validate_cuda_int32_contiguous(indices, "indices"); + validate_cuda_int32_contiguous(union_indices, "union_indices"); + validate_cuda_int32_contiguous(remap, "remap"); + validate_cuda_int32_contiguous(union_counts, "union_counts"); + validate_cuda_int32_contiguous(hash_keys, "hash_keys"); + validate_cuda_int32_contiguous(hash_first_positions, "hash_first_positions"); + validate_cuda_int32_contiguous(first_to_dense, "first_to_dense"); + TORCH_CHECK(indices.dim() == 4, + "indices must be [destination, request, row, topk]"); + TORCH_CHECK(num_requests > 0 && num_requests == indices.size(1), + "active request count must match indices"); + TORCH_CHECK(union_indices.dim() == 3, + "union_indices must be [destination, request, capacity]"); + TORCH_CHECK(hash_keys.dim() == 3, + "hash_keys must be [destination, request, hash capacity]"); + + const auto device = indices.device(); + TORCH_CHECK(union_indices.device() == device, + "union_indices device mismatch"); + TORCH_CHECK(remap.device() == device, "remap device mismatch"); + TORCH_CHECK(union_counts.device() == device, "union_counts device mismatch"); + TORCH_CHECK(hash_keys.device() == device, "hash_keys device mismatch"); + TORCH_CHECK(hash_first_positions.device() == device, + "hash_first_positions device mismatch"); + TORCH_CHECK(first_to_dense.device() == device, + "first_to_dense device mismatch"); + + const int destinations = static_cast(indices.size(0)); + const int max_requests = static_cast(union_indices.size(1)); + const int64_t input_entries = indices.size(2) * indices.size(3); + const int64_t capacity = union_indices.size(2); + const int64_t hash_capacity = hash_keys.size(2); + TORCH_CHECK(capacity >= input_entries, "union capacity is too small"); + TORCH_CHECK((hash_capacity & (hash_capacity - 1)) == 0, + "hash capacity must be a power of two"); + TORCH_CHECK(hash_capacity >= 2 * input_entries, + "hash capacity must be at least 2x input entries"); + TORCH_CHECK(union_indices.size(0) == destinations, + "union destination mismatch"); + TORCH_CHECK(num_requests <= max_requests, + "active requests exceed union state capacity"); + TORCH_CHECK(remap.dim() == 4 && remap.size(0) == destinations && + remap.size(1) == max_requests && + remap.size(2) >= indices.size(2) && + remap.size(3) == indices.size(3) && + remap.size(2) * remap.size(3) == capacity, + "remap shape mismatch"); + TORCH_CHECK(union_counts.dim() == 2 && union_counts.size(0) == destinations && + union_counts.size(1) == max_requests, + "count shape mismatch"); + TORCH_CHECK(hash_keys.sizes() == hash_first_positions.sizes(), + "hash shape mismatch"); + TORCH_CHECK( + hash_keys.size(0) == destinations && hash_keys.size(1) == max_requests, + "hash state shape mismatch"); + TORCH_CHECK(first_to_dense.sizes() == union_indices.sizes(), + "dense-slot shape mismatch"); + + constexpr int threads = 256; + const int blocks = destinations * static_cast(num_requests); + const auto stream = c10::cuda::getCurrentCUDAStream(indices.get_device()); + build_dense_union_remap_kernel<<>>( + indices.data_ptr(), union_indices.data_ptr(), + remap.data_ptr(), union_counts.data_ptr(), + hash_keys.data_ptr(), hash_first_positions.data_ptr(), + first_to_dense.data_ptr(), destinations, + static_cast(num_requests), max_requests, input_entries, capacity, + hash_capacity); + AT_CUDA_CHECK(cudaGetLastError()); +} + +} // namespace + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { + module.def("build_dense_union_remap", &build_dense_union_remap); +} diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index f0172a10b95d..5fdcd378a1e5 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -31,7 +31,6 @@ split_decodes_and_prefills, ) from vllm.v1.kv_cache_interface import AttentionSpec, MLAAttentionSpec -from vllm.v1.worker.cp_utils import get_total_cp_world_size logger = init_logger(__name__) @@ -272,6 +271,17 @@ def get_max_prefill_buffer_size(vllm_config: VllmConfig): return max_model_len * 40 +def get_indexer_max_num_blocks_per_req( + max_model_len: int, + block_size: int, + configured_cp_world_size: int, + dcp_replicated: bool, +) -> int: + """Size indexer metadata for the cache group's effective DCP layout.""" + effective_cp_world_size = 1 if dcp_replicated else configured_cp_world_size + return cdiv(max_model_len, block_size * effective_cp_world_size) + + def _supports_varlen_paged_mqa_logits() -> bool: if ( envs.VLLM_USE_B12X_SPARSE_INDEXER @@ -347,7 +357,12 @@ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) scheduler_config = self.vllm_config.scheduler_config parallel_config = self.vllm_config.parallel_config - self.dcp_world_size = parallel_config.decode_context_parallel_size + self.dcp_replicated = bool(getattr(self.kv_cache_spec, "dcp_replicated", False)) + configured_dcp_world_size = parallel_config.decode_context_parallel_size + configured_cp_world_size = ( + configured_dcp_world_size * parallel_config.prefill_context_parallel_size + ) + self.dcp_world_size = 1 if self.dcp_replicated else configured_dcp_world_size self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_world_size > 1 else 0 self.cp_kv_cache_interleave_size = parallel_config.cp_kv_cache_interleave_size # The DCP sparse-indexer code is parameterized by interleave size, but @@ -444,9 +459,11 @@ def __init__(self, *args, **kwargs): dtype=torch.int32, device=self.device, ) - max_num_blocks_per_req = cdiv( + max_num_blocks_per_req = get_indexer_max_num_blocks_per_req( self.vllm_config.model_config.max_model_len, - self.kv_cache_spec.block_size * get_total_cp_world_size(), + self.kv_cache_spec.block_size, + configured_cp_world_size, + self.dcp_replicated, ) self.expanded_block_table_buffer = torch.zeros( ( @@ -986,7 +1003,7 @@ def build( # range and miss valid tokens. Keep the global seq_lens here and # localize the expanded bounds further down. global_seq_lens_for_decode: torch.Tensor | None = None - if dcp_local_seq_lens is not None: + if use_dcp_local_kv: global_seq_lens_for_decode = common_attn_metadata.seq_lens[:num_decodes] seq_lens = common_attn_metadata.seq_lens[:num_decodes] block_table = common_attn_metadata.block_table_tensor[:num_decodes, ...] @@ -1050,7 +1067,7 @@ def build( # Uncompressed DCP localizes after per-token expansion. Compressed # DCP must first convert global logical lengths to compressed-token # lengths and only then shard those retained tokens across ranks. - if dcp_local_seq_lens is not None and self.compress_ratio == 1: + if use_dcp_local_kv and self.compress_ratio == 1: seq_lens = self._dcp_localize_decode_seq_lens( seq_lens, num_decodes, seq_lens_is_buffer_view ) @@ -1060,7 +1077,7 @@ def build( if self.compress_ratio > 1: if seq_lens_is_buffer_view: seq_lens //= self.compress_ratio - if dcp_local_seq_lens is not None: + if use_dcp_local_kv: seq_lens.copy_( get_dcp_local_seq_lens( seq_lens, @@ -1072,7 +1089,7 @@ def build( else: # Copy to avoid mutating shared state; keeps CG address stable. compressed_decode_seq_lens = seq_lens // self.compress_ratio - if dcp_local_seq_lens is not None: + if use_dcp_local_kv: compressed_decode_seq_lens = get_dcp_local_seq_lens( compressed_decode_seq_lens, self.dcp_world_size, @@ -1123,7 +1140,7 @@ def build( use_native, ( common_attn_metadata.dcp_local_seq_lens_cpu[:num_decodes] - if dcp_local_seq_lens is not None + if use_dcp_local_kv and common_attn_metadata.dcp_local_seq_lens_cpu is not None else None ), diff --git a/vllm/v1/core/block_pool.py b/vllm/v1/core/block_pool.py index b42c7662d2ba..7e241ec14035 100644 --- a/vllm/v1/core/block_pool.py +++ b/vllm/v1/core/block_pool.py @@ -282,14 +282,16 @@ def cache_full_blocks( block_hash, kv_cache_group_id ) if blk.block_hash is not None: - # The only valid case where a "new full block" already has a - # hash is partial->full promotion of the same cache block. - assert ( - blk.block_hash_num_tokens is not None - and blk.block_hash_num_tokens < num_hash_tokens - ) - removed_hashes = self._remove_cached_block_hashes(blk) - self._emit_block_removed_events(removed_hashes) + assert blk.block_hash_num_tokens is not None + if blk.block_hash_num_tokens < num_hash_tokens: + removed_hashes = self._remove_cached_block_hashes(blk) + self._emit_block_removed_events(removed_hashes) + else: + # Lockstep groups attach group-specific keys to one block + # at the same token boundary. + assert blk.block_hash_num_tokens == num_hash_tokens + assert get_block_hash(blk.block_hash) == block_hash + assert get_group_id(blk.block_hash) != kv_cache_group_id self._insert_block_hash( block_hash_with_group_id, blk, diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index aebd694bb854..f3727a5a3311 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -11,6 +11,7 @@ from vllm.v1.core.kv_cache_utils import ( BlockHash, KVCacheBlock, + _use_lockstep_mla_allocation, is_deepseek_v4_hybrid_kv_cache_config, ) from vllm.v1.core.single_type_kv_cache_manager import ( @@ -119,6 +120,11 @@ def __init__( ) for i, kv_cache_group in enumerate(self.kv_cache_config.kv_cache_groups) ) + self.lockstep_mla_allocations = _use_lockstep_mla_allocation( + kv_cache_config.kv_cache_groups, + dcp_world_size, + pcp_world_size, + ) # A positive retention interval must be a multiple of the base hit granularity # (``scheduler_block_size``) to land on real cache-hit boundaries. @@ -164,31 +170,37 @@ def get_num_blocks_to_allocate( Returns: The number of blocks to allocate. """ - num_blocks_to_allocate = 0 + blocks_by_group: list[int] = [] for i, manager in enumerate(self.single_type_managers): if isinstance(manager, CrossAttentionManager): # For cross-attention, we issue a single static allocation # of blocks based on the number of encoder input tokens. - num_blocks_to_allocate += manager.get_num_blocks_to_allocate( - request_id, - num_encoder_tokens, - [], - 0, - 0, - num_encoder_tokens, - apply_admission_cap=apply_admission_cap, + blocks_by_group.append( + manager.get_num_blocks_to_allocate( + request_id, + num_encoder_tokens, + [], + 0, + 0, + num_encoder_tokens, + apply_admission_cap=apply_admission_cap, + ) ) else: - num_blocks_to_allocate += manager.get_num_blocks_to_allocate( - request_id, - num_tokens, - new_computed_blocks[i], - total_computed_tokens, - num_local_computed_tokens, - num_tokens_main_model, - apply_admission_cap=apply_admission_cap, + blocks_by_group.append( + manager.get_num_blocks_to_allocate( + request_id, + num_tokens, + new_computed_blocks[i], + total_computed_tokens, + num_local_computed_tokens, + num_tokens_main_model, + apply_admission_cap=apply_admission_cap, + ) ) - return num_blocks_to_allocate + if self.lockstep_mla_allocations: + return max(blocks_by_group, default=0) + return sum(blocks_by_group) def allocate_new_computed_blocks( self, @@ -208,6 +220,11 @@ def allocate_new_computed_blocks( num_local_computed_tokens: The number of local computed tokens. num_external_computed_tokens: The number of external computed tokens. """ + if self.lockstep_mla_allocations and num_external_computed_tokens: + raise NotImplementedError( + "External KV loads are not supported with lockstep MLA allocations." + ) + # A running request is already tracked in num_cached_block and won't # have new prefix-cache hits, so this is a no-op for it. if any( @@ -236,6 +253,17 @@ def allocate_new_computed_blocks( num_external_computed_tokens, ) + if self.lockstep_mla_allocations: + group_blocks = [ + manager.req_to_blocks[request_id] + for manager in self.single_type_managers + ] + block_ids = [block.block_id for block in group_blocks[0]] + assert all( + [block.block_id for block in blocks] == block_ids + for blocks in group_blocks[1:] + ) + def allocate_new_blocks( self, request_id: str, @@ -260,16 +288,38 @@ def allocate_new_blocks( Returns: The new allocated blocks. """ - return tuple( - manager.allocate_new_blocks( - request_id, - num_encoder_tokens - if isinstance(manager, CrossAttentionManager) - else num_tokens, - num_tokens_main_model, + if not self.lockstep_mla_allocations: + return tuple( + manager.allocate_new_blocks( + request_id, + num_encoder_tokens + if isinstance(manager, CrossAttentionManager) + else num_tokens, + num_tokens_main_model, + ) + for manager in self.single_type_managers ) - for manager in self.single_type_managers + + managers = self.single_type_managers + primary_ids = [ + block.block_id for block in managers[0].req_to_blocks[request_id] + ] + assert all( + [block.block_id for block in manager.req_to_blocks[request_id]] + == primary_ids + for manager in managers[1:] + ) + + new_blocks = managers[0].allocate_new_blocks( + request_id, + num_tokens, + num_tokens_main_model, ) + for manager in managers[1:]: + if new_blocks: + self.block_pool.touch(new_blocks) + manager.req_to_blocks[request_id].extend(new_blocks) + return tuple(list(new_blocks) for _ in managers) def cache_blocks(self, request: Request, num_computed_tokens: int) -> None: """ @@ -749,10 +799,10 @@ def find_longest_cache_hit( sparse-retention group has not cached yet (0 unless hybrid). """ if self.disable_prefix_cache_for_dcp_hybrid: - blocks: tuple[list[KVCacheBlock], ...] = tuple( + empty_blocks: tuple[list[KVCacheBlock], ...] = tuple( [] for _ in range(len(self.kv_cache_config.kv_cache_groups)) ) - return blocks, 0, 0 + return empty_blocks, 0, 0 num_groups = len(self.kv_cache_config.kv_cache_groups) hit_length = max_cache_hit_length @@ -830,8 +880,8 @@ def find_longest_cache_hit( # length shrunk; invalidate previous eagle verifications eagle_verified.clear() curr_hit_length = _new_hit_length - for group_id, blocks in zip(group_ids, hit_blocks): - hit_blocks_by_group[group_id] = blocks + for group_id, group_blocks in zip(group_ids, hit_blocks): + hit_blocks_by_group[group_id] = group_blocks hit_length_by_group[group_id] = _new_hit_length longest_hit_length = max(longest_hit_length, curr_hit_length) diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 01aa1196d714..fbd71d5bb232 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -1002,6 +1002,16 @@ def _pool_bytes_per_block( `available_memory` into `num_blocks`. Used to compute the effective KV cache capacity once `num_gpu_blocks_override` is applied. """ + parallel_config = vllm_config.parallel_config + if _use_lockstep_mla_allocation( + kv_cache_groups, + parallel_config.decode_context_parallel_size, + parallel_config.prefill_context_parallel_size, + ): + return sum( + page_size + for page_size, _ in _get_lockstep_mla_tensor_slots(kv_cache_groups) + ) if len(kv_cache_groups) == 1 and isinstance( kv_cache_groups[0].kv_cache_spec, UniformTypeKVCacheSpecs ): @@ -1313,10 +1323,91 @@ def _bucket_layers_by_page_size( return buckets +def _is_lockstep_mla_spec_layout( + kv_cache_specs: Iterable[KVCacheSpec], + dcp_world_size: int, + pcp_world_size: int, +) -> bool: + specs = list(kv_cache_specs) + if ( + len(specs) < 2 + or not isinstance(dcp_world_size, int) + or not isinstance(pcp_world_size, int) + ): + return False + cp_world_size = dcp_world_size * pcp_world_size + if cp_world_size <= 1: + return False + if not all(isinstance(spec, MLAAttentionSpec) for spec in specs): + return False + + replication_modes = {bool(getattr(spec, "dcp_replicated", False)) for spec in specs} + global_block_sizes = { + spec.block_size + if getattr(spec, "dcp_replicated", False) + else spec.block_size * cp_world_size + for spec in specs + } + return replication_modes == {False, True} and len(global_block_sizes) == 1 + + +def _use_lockstep_mla_allocation( + kv_cache_groups: list[KVCacheGroupSpec], + dcp_world_size: int, + pcp_world_size: int, +) -> bool: + """Return whether mixed MLA groups can share one block-ID namespace.""" + if len(kv_cache_groups) < 2: + return False + + layer_specs: list[KVCacheSpec] = [] + for group in kv_cache_groups: + group_spec = group.kv_cache_spec + group_layer_specs = ( + list(group_spec.kv_cache_specs.values()) + if isinstance(group_spec, UniformTypeKVCacheSpecs) + else [group_spec] + ) + group_modes = { + bool(getattr(spec, "dcp_replicated", False)) for spec in group_layer_specs + } + if len(group_modes) != 1: + return False + layer_specs.extend(group_layer_specs) + + return _is_lockstep_mla_spec_layout(layer_specs, dcp_world_size, pcp_world_size) + + +def _get_lockstep_mla_tensor_slots( + kv_cache_groups: list[KVCacheGroupSpec], +) -> list[tuple[int, list[str]]]: + slots: list[tuple[int, list[str]]] = [] + for group in kv_cache_groups: + group_spec = group.kv_cache_spec + if isinstance(group_spec, UniformTypeKVCacheSpecs): + slots.extend( + (group_spec.kv_cache_specs[layer_name].page_size_bytes, [layer_name]) + for layer_name in group.layer_names + ) + else: + slots.extend( + (group_spec.page_size_bytes, [layer_name]) + for layer_name in group.layer_names + ) + return slots + + def _use_packed_kv_cache_config( vllm_config: VllmConfig, kv_cache_groups: list[KVCacheGroupSpec], ) -> bool: + parallel_config = vllm_config.parallel_config + if _use_lockstep_mla_allocation( + kv_cache_groups, + parallel_config.decode_context_parallel_size, + parallel_config.prefill_context_parallel_size, + ): + return True is_dsv4 = all( isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) for group in kv_cache_groups @@ -1346,11 +1437,27 @@ def _get_kv_cache_config_packed( groups at the same slot share a tensor (they have independent block tables so block-id namespaces never collide). Each emitted tensor aliases one physical backing allocation, with per-block data laid out contiguously. - """ + Lockstep MLA groups instead get distinct contiguous tensors because their + block IDs intentionally coincide. + """ + parallel_config = vllm_config.parallel_config + if _use_lockstep_mla_allocation( + kv_cache_groups, + parallel_config.decode_context_parallel_size, + parallel_config.prefill_context_parallel_size, + ): + tensor_slots = _get_lockstep_mla_tensor_slots(kv_cache_groups) + total_num_bytes_per_block = sum(page_size for page_size, _ in tensor_slots) + num_blocks = available_memory // total_num_bytes_per_block + num_blocks = may_override_num_blocks(vllm_config, num_blocks) + return num_blocks, [ + KVCacheTensor(size=page_size * num_blocks, shared_by=slot) + for page_size, slot in tensor_slots + ] + # buckets = {page_size: [[layer_names], [layer_names], ...]} buckets = _bucket_layers_by_page_size(kv_cache_groups) total_num_bytes_per_block = sum(ps * len(slots) for ps, slots in buckets.items()) - num_blocks = available_memory // total_num_bytes_per_block num_blocks = may_override_num_blocks(vllm_config, num_blocks) @@ -1557,6 +1664,8 @@ def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]): def group_and_unify_kv_cache_specs( kv_cache_spec: dict[str, KVCacheSpec], + dcp_world_size: int = 1, + pcp_world_size: int = 1, ) -> list[UniformTypeKVCacheSpecs] | None: """ Group the KV cache specs and unify each group into one UniformTypeKVCacheSpecs. @@ -1565,29 +1674,33 @@ def group_and_unify_kv_cache_specs( has_swa = any( isinstance(spec, SlidingWindowMLASpec) for spec in kv_cache_spec.values() ) - # DFlash-under-DCP draft: full-attention layers replicated on every DCP - # rank. They have a different page size than the MLA target and need their - # own group, but the DeepseekV4 multi-group allocator (with page-size - # padding) handles exactly that, so route them through here too. + lockstep_mla_layout = _is_lockstep_mla_spec_layout( + kv_cache_spec.values(), dcp_world_size, pcp_world_size + ) + # Replicated layers need a group separate from sharded layers because they + # have different token ownership and block-table semantics. has_repl = any( getattr(spec, "dcp_replicated", False) - and not isinstance(spec, MLAAttentionSpec) + and (not isinstance(spec, MLAAttentionSpec) or lockstep_mla_layout) for spec in kv_cache_spec.values() ) if not (has_swa or has_repl): return None - # SlidingWindowMLASpec models with uniform page sizes don't need tuple packing. + # Other uniform page layouts do not need tuple packing. page_sizes = {spec.page_size_bytes for spec in kv_cache_spec.values()} - if len(page_sizes) <= 1: + if len(page_sizes) <= 1 and not lockstep_mla_layout: return None mla_specs: dict[str, KVCacheSpec] = {} grouped_swa_mla_specs: dict[tuple[int, int, bool, bool], dict[str, KVCacheSpec]] = ( defaultdict(dict) ) - # dcp_replicated non-MLA groups (e.g. the DFlash draft), keyed by block_size. - grouped_repl_specs: dict[tuple[int], dict[str, KVCacheSpec]] = defaultdict(dict) + # Keep incompatible replicated types and page layouts separate. This + # includes DFlash drafts and replicated MLA indexer caches. + grouped_repl_specs: dict[ + tuple[int] | tuple[str, int, int], dict[str, KVCacheSpec] + ] = defaultdict(dict) # NOTE: Here we group SWA layers by (block_size, sliding_window, # dcp_replicated, dcp_sharded), which separates SWA layers, C4I+C4A # layers, C128A layers, and replicated compressor-state groups. @@ -1601,10 +1714,17 @@ def group_and_unify_kv_cache_specs( spec.dcp_sharded, ) ][name] = spec + elif getattr(spec, "dcp_replicated", False) and ( + not isinstance(spec, MLAAttentionSpec) or lockstep_mla_layout + ): + key = ( + (type(spec).__name__, spec.block_size, spec.page_size_bytes) + if isinstance(spec, MLAAttentionSpec) + else (spec.block_size,) + ) + grouped_repl_specs[key][name] = spec elif isinstance(spec, MLAAttentionSpec): mla_specs[name] = spec - elif getattr(spec, "dcp_replicated", False): - grouped_repl_specs[(spec.block_size,)][name] = spec if len(mla_specs) == 0: # No full-MLA group to anchor the DeepseekV4 layout; let the generic @@ -1698,11 +1818,17 @@ def _get_kv_cache_groups_uniform_groups( ] swa_mla_specs = grouped_specs[1:] - # Non-first groups are SWA-MLA, full-attention dcp_replicated drafts, or - # sliding-window dcp_replicated drafts. All are padded to MLA buckets - # identically. + # Non-first groups are SWA-MLA or DCP-replicated caches. assert all( - isinstance(spec, (SlidingWindowMLASpec, FullAttentionSpec, SlidingWindowSpec)) + isinstance( + spec, + ( + MLAAttentionSpec, + SlidingWindowMLASpec, + FullAttentionSpec, + SlidingWindowSpec, + ), + ) for group in swa_mla_specs for spec in group.kv_cache_specs.values() ) @@ -1716,19 +1842,30 @@ def _get_kv_cache_groups_uniform_groups( for sm_spec in swa_mla_specs: sm_page_sizes = sm_spec.get_page_sizes() layers_per_size: dict[int, list[str]] = defaultdict(list) - assert max(sm_page_sizes) <= max(all_page_sizes) + is_replicated_mla = all( + isinstance(spec, MLAAttentionSpec) + and getattr(spec, "dcp_replicated", False) + for spec in sm_spec.kv_cache_specs.values() + ) + if not is_replicated_mla: + assert max(sm_page_sizes) <= max(all_page_sizes) # Unify page size by padding layers' page_size to the nearest larger page_size. # Compute candidate (nearest larger page_size) for each unique page size. size_to_candidate: dict[int, int] = {} for ps in sm_page_sizes: - size_to_candidate[ps] = min(x for x in all_page_sizes if x >= ps) + if ps <= max(all_page_sizes): + size_to_candidate[ps] = min(x for x in all_page_sizes if x >= ps) # Pad and collect layer names per page size. for layer_name, layer_spec in sm_spec.kv_cache_specs.items(): current_size = layer_spec.page_size_bytes - candidate = size_to_candidate[current_size] - if current_size < candidate: - object.__setattr__(layer_spec, "page_size_padded", candidate) + if is_replicated_mla: + candidate = current_size + else: + assert current_size <= max(all_page_sizes) + candidate = size_to_candidate[current_size] + if current_size < candidate: + object.__setattr__(layer_spec, "page_size_padded", candidate) layers_per_size[candidate].append(layer_name) # NOTE(yifan): for now, inside a UniformKV group, each page_size should # have the same number of layers. This also means we don't need to pad layers @@ -1816,7 +1953,11 @@ def get_kv_cache_groups( # full attention, or all layers are sliding window attention with the # same window size). Put all layers into one group. return _get_kv_cache_groups_uniform_type(uniform_spec) - elif grouped_specs := group_and_unify_kv_cache_specs(kv_cache_spec): + elif grouped_specs := group_and_unify_kv_cache_specs( + kv_cache_spec, + vllm_config.parallel_config.decode_context_parallel_size, + vllm_config.parallel_config.prefill_context_parallel_size, + ): # DeepseekV4 case: All layers need the same number of token slots, # yet some layers are full attention while others are sliding window # attention in different sizes. Need to group layers into multiple @@ -1903,6 +2044,21 @@ def _max_memory_usage_bytes_from_groups( if not kv_cache_groups: return 0 + parallel_config = vllm_config.parallel_config + if _use_lockstep_mla_allocation( + kv_cache_groups, + parallel_config.decode_context_parallel_size, + parallel_config.prefill_context_parallel_size, + ): + request_blocks = max( + cdiv( + group.kv_cache_spec.max_memory_usage_bytes(vllm_config), + group.kv_cache_spec.page_size_bytes, + ) + for group in kv_cache_groups + ) + return request_blocks * _pool_bytes_per_block(vllm_config, kv_cache_groups) + if len(kv_cache_groups) == 1 and isinstance( kv_cache_groups[0].kv_cache_spec, UniformTypeKVCacheSpecs ):