diff --git a/.gitattributes b/.gitattributes index b8ef1173..4e6ce988 100644 --- a/.gitattributes +++ b/.gitattributes @@ -2,6 +2,8 @@ *.patch text eol=lf runtime/** text eol=lf runtime/sparkring/*.gz -text +runtime/glm53-spark-mtp3-mesh/compute/*.tar.gz -text +runtime/glm53-spark-mtp3-mesh/compute/*.patch whitespace=-trailing-space runtime/exl3/patches/*.patch whitespace=-trailing-space # Runtime receipts hash these source files byte-for-byte before copying them, diff --git a/README.md b/README.md index 2cdc519e..4265b9e8 100644 --- a/README.md +++ b/README.md @@ -29,7 +29,7 @@ reproducible benchmarks, and [test results](performance/). | Model / predictor | Layout | Context | Sequences | Batch | Guide | |---|---|---:|---:|---:|---| -| **GLM-5.3 Flash NVFP4-Spark · MTP3 mesh** | TP4/DCP4 | 1M | 16 | 8,192 | [Quickstart](docs/GLM53_SPARK_MTP3_MESH_QUICKSTART.md) | +| **GLM-5.3 Flash NVFP4-Spark · MTP3 cache/checkpoint mesh** | TP4/DCP4 | 1M | 16 | 8,192 | [Quickstart](docs/GLM53_MTP3_CACHE_CHECKPOINTS_QUICKSTART.md) | | GLM-5.3 Flash NVFP4 · BF16 DFlash2 | TP4/DCP4; DCP1/2 | 1M | 16 | 8,192 | [Quickstart](docs/GLM53_JJ_R8_GB10_SPARKCACHE_TP4_QUICKSTART.md) | | GLM-5.2 EXL3 3.5-bpw | TP4/DCP4 | 1M | 16 | 4,096 | [Quickstart](docs/GLM52_35BPW_QUICKSTART.md) | | DeepSeek-V4-Flash-0731 | TP4/DCP1 | 1M | 32 | 4,096 | [Quickstart](docs/DEEPSEEK_V4_FLASH_QUICKSTART.md) | diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index c3d03400..1fbe44b8 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -55,7 +55,17 @@ two compatibility patches. The public GLM-5.3 builder does not consume them; its `NCCL_SWITCHLESS_RING_ONLY` parameter and diagnostics are an independent SparkRing implementation. -## 3. vLLM (referenced and patched) +## 3. vLLM (referenced, patched, and selected source included) + +`runtime/glm53-spark-mtp3-mesh/compute/vllm-compute-files.tar.gz` includes +24 vLLM Python files for GLM metadata reuse, dense-kernel integration, +the NVFP4 proposal head, deferred-weight ownership, and independent draft +and rejection-sampling randomness. The adjacent patch provides a readable diff; +`source-lock.json` records the base, donor revisions, and exact file hashes. +These files derive from Local Inference Lab's vLLM fork at `3512b066` and +`a8c796f3`, loader correction `17e341b9`, and independent-RNG backport +`44e6766e` from [PR 653](https://github.com/local-inference-lab/vllm/pull/653), +under Apache-2.0 with their contributor notices retained. The unified diffs under `runtime/deepseek0731-gb10/patches/` contain context and removed lines from vLLM, pinned to the source revision recorded by that @@ -245,8 +255,10 @@ artifacts that are not distributed in this repository: repository notices. SparkRing records these identities and validates compatible image content; it -does not redistribute the model weights, vLLM source, or B12X model-kernel -package. The selected B12X communication source in Section 11 is included. Operators +does not include model weights or the complete vLLM/B12X source trees in this +Git repository. The selected vLLM files in Section 3 and B12X communication +source in Section 11 are included. Published runtime images contain the +pinned vLLM and B12X packages under their respective licenses. Operators must obtain each artifact under its own terms. The exact operator-image composition is in `runtime/glm53-flash-jj-r8-gb10/pins.json` and `runtime/glm53-flash-jj-r8-gb10/glm53-dcp4-sircl-public-image-receipt.json`. @@ -287,3 +299,16 @@ The native-MTP3 profile references Operators must obtain and use that checkpoint under its own license and notices. Native MTP uses its included prediction layer and does not require an external DFlash checkpoint. + +## 12. B12X selector source and MoE scale sharing + +The native-MTP3 compute package downloads B12X revision +`ef308bac0f3b3eb8fea63e4013afc0c2ea1c6301`, including its shared native NVFP4 +scales for A4/A16 MoE paths. Three selector Python files are included in +`runtime/glm53-spark-mtp3-mesh/compute/b12x-selector-files.tar.gz` from +[B12X PR 316](https://github.com/local-inference-lab/b12x/pull/316), revision +`9ac142824b4edb750892a0fb63d914230086495d`. They implement the top-k-512 +candidate buffer, exact overflow handling, and omission of unused terminal +scores. These files are licensed under Apache-2.0; source notices are retained +and the downloaded B12X archive supplies the license. The compute source lock +binds the source archive and each base/result file hash. diff --git a/docs/GLM53_JJ_R8_GB10_SPARKCACHE_TP4_QUICKSTART.md b/docs/GLM53_JJ_R8_GB10_SPARKCACHE_TP4_QUICKSTART.md index 66e93de2..8892c6e3 100644 --- a/docs/GLM53_JJ_R8_GB10_SPARKCACHE_TP4_QUICKSTART.md +++ b/docs/GLM53_JJ_R8_GB10_SPARKCACHE_TP4_QUICKSTART.md @@ -305,6 +305,11 @@ SPARKCACHE_ACCESS_MODE=read-write # restore existing entries and publish new o SPARKCACHE_ACCESS_MODE=restore-only # restore existing entries; never capture new prompts ``` +Keep `SPARKCACHE_ASYNC_PAGE_CAPTURE=auto` from the environment template to +disable capture automatically in restore-only mode or when SparkCache is off. +An explicit `1` with either setting is rejected. Explicit `0` in read-write +mode selects synchronous publication. + The GLM-5.3 profile retains a verified shared GPU prefix for up to five minutes so one restore can serve an extended request queue: @@ -360,7 +365,7 @@ and `IMAGE_ID`; keep the page-tail settings unchanged. The recommended DCP4 profile enables bounded asynchronous page capture: ```bash -SPARKCACHE_ASYNC_PAGE_CAPTURE=1 +SPARKCACHE_ASYNC_PAGE_CAPTURE=auto SPARKCACHE_ASYNC_CAPTURE_SLOT_BYTES=auto SPARKCACHE_ASYNC_CAPTURE_SLOT_COUNT=2 ``` @@ -368,12 +373,26 @@ SPARKCACHE_ASYNC_CAPTURE_SLOT_COUNT=2 The `auto` slot policy selects 8 GiB for DCP1, 5 GiB for DCP2, or 3 GiB for DCP4. Two capture slots let the background publisher consume one completed capture while a later capture uses the other. Restore separately overlaps -bounded NVMe reads and CUDA placement through two 256 MiB mapped arenas. More -restore arenas are not part of this profile because measured arena waits did -not justify the additional unified-memory pressure. DCP1 and DCP2 page-tail +bounded NVMe reads and CUDA placement through two 256 MiB mapped arenas per +load lane. Eight lanes reserve 4 GiB per rank; adding the DCP4 capture slots +gives 10 GiB per rank (40 GiB across TP4) of configured payload capacity, +before KV, models, metadata, retained objects and transport. Restore-only +omits the 6 GiB capture ring. DCP1 and DCP2 page-tail capture have no matching live record; use complete snapshots or test those layouts separately. +Print the resolved memory plan without host or GPU access before launching: + +```bash +SPARKRING_PRINT_MEMORY_PLAN=1 bash runtime/glm53-flash-jj-r8-gb10/launch-rank.sh \ + 0 /path/to/rank.env +``` + +Set `SPARKCACHE_BUFFER_BUDGET_BYTES` in that configuration to enforce a +per-rank restore-plus-capture payload ceiling; zero disables it. The report +shows KV separately and names excluded memory. It does not predict total +resident memory or replace a hardware capacity test. + The environment template enables `DFLASH_WARMUP=1`. Rank 0 waits for the API, then exercises every concurrency from C1 through C16 and scheduled prompt spans covering DFlash's Triton block-size specializations. DFlash depth seven diff --git a/docs/GLM53_MTP3_CACHE_CHECKPOINTS_QUICKSTART.md b/docs/GLM53_MTP3_CACHE_CHECKPOINTS_QUICKSTART.md new file mode 100644 index 00000000..0b1fb27f --- /dev/null +++ b/docs/GLM53_MTP3_CACHE_CHECKPOINTS_QUICKSTART.md @@ -0,0 +1,98 @@ +# GLM-5.3 native MTP3 with verified caching and recurrent checkpoints + +Status: **research-only**. The image build and source checks are implemented. +The published image's vLLM, B12X, transport, and warmup files match the deployed +composition in 5,308 byte comparisons. It incorporates SparkCache main at the +revision below. Source equivalence is not a serving soak of this rebuilt image. + +## Prerequisites + +Use four NVIDIA Sparks with the [managed mesh host setup](GLM53_SPARK_MESH_HOST_SETUP.md). +The model uses TP4/DCP4 and native MTP depth three. Keep the reviewed fabric, +driver, GID and interface settings; this guide does not reconfigure NICs. +Follow [host prerequisites](PREREQUISITES.md#four-spark-managed-hardware-forwarded-mesh), +including optional reboot preparation after large GPU workloads. Stop model +workloads through the managed lifecycle before replacing containers or services. + +Download the target `local-inference-lab/GLM-5.3-Flash-NVFP4-Spark` at revision +`df116c4fb16b1d37ae43d2cfd624de26ffbc832e`. No external draft is required. + +## Pull and verify the image + +Use the same immutable image on all four ranks: + +```bash +image='ghcr.io/fujitsupolycom/sparkring-glm53-sparkcache@sha256:11a556a54041fd823d152a7f051ac4f7c617dc539030df26e93008392fee0746' +image_id='sha256:6921a6c163ea40b603e19a0332330efe3dbccbf4dce9f6cbbf6b756c9231835a' +docker pull "$image" +test "$(docker image inspect "$image" --format '{{.Id}}')" = "$image_id" +docker run --rm --network none --entrypoint python3 "$image" \ + -S -B /opt/sparkring/bin/verify-performance.py +``` + +The [public image contract](../runtime/glm53-spark-mtp3-mesh/performance/public-image.json) +is the renderer and installer input for this composition. Do not substitute the +compute parent's `image-receipt.json`. Source builds use +[the performance build instructions](../runtime/glm53-spark-mtp3-mesh/performance/README.md). + +## Extract and render + +Choose unused artifact and export-container names. The following commands do +not start a model. Create `/srv/sparkring/artifacts` with operator permissions: + +```bash +docker create --name mtp3-artifact-export "$image" +docker cp mtp3-artifact-export:/opt/spark-sircl /srv/sparkring/artifacts/cache-checkpoints-bundle +docker cp mtp3-artifact-export:/opt/sparkring/bin/mlx5-rdma-tx-marker /srv/sparkring/artifacts/mlx5-rdma-tx-marker +docker rm mtp3-artifact-export +cp runtime/glm53-spark-mtp3-mesh/performance/public-image.json /srv/sparkring/verified-image-receipt.json +``` + +Edit the private site and fabric inputs described in the +[managed deployment guide](../runtime/glm53-spark-mtp3-mesh/MANAGED_MESH.md). +Set the site's bundle root to the extracted bundle and its marker path to the +extracted marker. The marker SHA-256 is +`2828c07e4255c4962c77425be2c88969e7eb7dd4b1bf9e36485bc705bb5d6d64`. +Use a distinct container prefix and cache directory; do not rename entries from +another cache namespace into this one. + +```bash +python runtime/glm53-spark-mtp3-mesh/profile.py render \ + --site /srv/sparkring/site/mtp3-mesh.json \ + --bundle /srv/sparkring/artifacts/cache-checkpoints-bundle \ + --image-receipt /srv/sparkring/verified-image-receipt.json \ + --output build/mtp3-cache-checkpoints-launch +``` + +The renderer selects the matching image, transport manifest, optimized native +placement digest, and cache namespace. Review the generated rank environments +and plan, then follow the managed guide's create-only container, installation, +and coordinated startup steps with this launch directory and verified receipt. +Do not use direct `docker start` to bypass the memory and four-rank gates. + +## Runtime behavior and limits + +SparkCache revision `48bbd2be4a7b972e56632a2d7b934bac5460f272` provides bounded +restoration, tiled CUDA placement, authenticated history reconstruction, +publication dependency protection, and publication-backlog gauges. Periodic +full captures remain off by default. They trade more writes for shorter +history processing and require a separately selected workload configuration. + +The profile retains 24 GiB KV capacity per rank, an 8,192-token batch budget, +40 GiB persistent-cache maximum and 32 GiB low watermark. Those capacity +defaults are not evidence of a 40 GiB multimodal stress test. The recorded +cache-pressure soak used 2 GiB per rank and an opt-in full-capture cadence. + +The image includes explicit two-checkpoint prefill, complete convolution-state +exports, stream-ordered transport, CTA publication synchronization, delayed +proxy-thread startup, and linked payload/doorbell posts. Checkpoint source and +ownership hashes are preserved by the build rather than inferred from labels. + +GLM-5.3 supports `reasoning_effort=low`, `high`, or `max`. Explicit thinking-off +chat requests return HTTP 400; warmup uses supported reasoning settings. This +does not implement reasoning-free generation. Responses API behavior is not +qualified by the chat-completion checks. + +Run bounded semantic, repeated-prefix, multimodal, and idle-probe checks after +deployment. Neither source equivalence nor the parent's measurements qualify +unattended availability, every cache boundary, or a universal speedup. diff --git a/docs/GLM53_SPARK_MTP3_MESH_QUICKSTART.md b/docs/GLM53_SPARK_MTP3_MESH_QUICKSTART.md index fa259e9a..257c1a93 100644 --- a/docs/GLM53_SPARK_MTP3_MESH_QUICKSTART.md +++ b/docs/GLM53_SPARK_MTP3_MESH_QUICKSTART.md @@ -1,17 +1,11 @@ # Run GLM-5.3 Flash Spark with native MTP3 and hardware-forwarded mesh -Status: **research-only**. The profile's composition, managed host service, -and CPU checks are **implemented**. The -[managed functional record](../performance/records/glm53-flash/spark-mtp3-managed-mesh-functional-20260905.md) -qualifies bounded installer, policy-scoped fault/recovery, post-recovery -readiness, and one persistent-cache recall case for the published managed -image identified below. -Broader cache/failure coverage and unattended serving remain unqualified. - -The [public application-install record](../performance/records/glm53-flash/spark-mtp3-public-application-install-20260905.md) -covers fresh public checkouts, extracted image artifacts, empty application -caches, installation, native correctness, and model-restart cache restoration -on four prepared hosts. It does not qualify a factory-reset OS/network setup. +Status: **research-only** profile with **implemented** source packaging and +managed host services. The published image passed the native, GPU stream, +serving, idle rank-loss, restart, and persistent-recall checks in its +[exact-image validation record](../performance/records/glm53-flash/spark-mtp3-compute-stream-safety-20260906.md). +That record defines the qualified conditions; in-flight collective failure +containment and unattended availability are not established. **Starting with four stock Sparks and no image?** Follow [the managed-mesh prerequisite section](PREREQUISITES.md#four-spark-managed-hardware-forwarded-mesh) @@ -19,16 +13,30 @@ first. It reuses the shared blank-cluster bootstrap and adds the secondary data interfaces, GID/MTU checks, and driver configuration required below. Return here to pull the published image and deploy the model. +**Preparing hosts that have already run GPU workloads?** Consider a reboot +before installation and the first model start, after stopping active workloads. +It can reduce memory fragmentation even when plenty of RAM appears free. +The [startup memory gate](../runtime/glm53-spark-mtp3-mesh/MANAGED_MESH.md#automatic-startup-memory-preparation) +checks available memory and large contiguous free blocks; rebooting does not +replace that check. A reboot interrupts every workload on that host. + The profile uses the `GLM-5.3-Flash-NVFP4-Spark` target's built-in multi-token -predictor with three speculative tokens. Graph-native SIRCL handles most +predictor with three speculative tokens. Its separate proposal head is packed +to NVFP4 at model load and uses BF16 activations. The target/verifier head +retains its BF16 checkpoint representation. Graph-native SIRCL handles most captured target verification, fused SIRCL handles large eager prefill, and RoCEnante handles selected small all-reduces. Patched NCCL retains the other collectives. The host fabric supplies hardware-forwarded paths between -opposite ranks without extra diagonal cables. +opposite ranks without extra diagonal cables. No external draft checkpoint or +DFlash model is used. The [profile contract](../runtime/glm53-spark-mtp3-mesh/README.md) and [pins](../runtime/glm53-spark-mtp3-mesh/pins.json) are the canonical inputs. -The [throughput record](../performance/records/glm53-flash/spark-mtp3-mesh-20260905.md) +The packaged RoCEnante runtime orders shared staging buffers across streams +and preserves its one-stream-per-CUDA-capture guard. CPU regressions and +four-rank GPU tests cover alternating streams, misaligned buffers, +changed-input graph replay, and second-stream capture rejection. +The [proposal-head throughput record](../performance/records/glm53-flash/spark-mtp3-nvfp4-proposal-head-20260905.md) reports observations, not a general performance guarantee. ## Attribution and design origins @@ -46,16 +54,38 @@ managed deployment. It does not claim to originate RoCEnante or install both complete PRs unchanged. The [vendored-source provenance](../third_party/b12x_roce/README.md) identifies the included code and retained license. +The proposal-head and metadata implementation is derived from +[Local Inference Lab vLLM revision `3512b066`](https://github.com/local-inference-lab/vllm/commit/3512b066e7796128c0c380ccc558182960f2f0ea), +as retained in +[revision `a8c796f3`](https://github.com/local-inference-lab/vllm/commit/a8c796f3af74106b2d8d441e9ec54588936a5388). +The compute source uses Local Inference Lab B12X +[revision `ef308bac`](https://github.com/local-inference-lab/b12x/commit/ef308bac0f3b3eb8fea63e4013afc0c2ea1c6301) +for shared native MoE scale storage, with three source-checked selector files +from [PR 316](https://github.com/local-inference-lab/b12x/pull/316). The vLLM +composition also includes deferred-weight ownership from `17e341b9` and +independent draft/rejection randomness from +[PR 653](https://github.com/local-inference-lab/vllm/pull/653). Exact source and +file identities are in the compute source lock. + ## Recorded benchmark observations +The [compute matrices](../performance/records/glm53-flash/spark-mtp3-compute-matrices-20260905.md) +compare proposal-head, loader/RNG, scale-sharing, and selector configurations +at C1/C2/C4/C8/C12/C16. They include averages across 8K/32K/64K context rows +and source-hashed individual cells. Their compute images and transport bundle +are identified separately from the combined image requiring qualification. + See the [consolidated validation report](../performance/records/glm53-flash/spark-mtp3-validation-summary-20260905.md) for completed checks, repeat counts, and the remaining test plan. The [profile results table](../runtime/glm53-spark-mtp3-mesh/README.md#operator-benchmark-observations) -shows the full concurrency matrix: aggregate decode reached **231.3 tok/s at -8K/C16**, and concurrency-one prefill scouts measured **2,703–2,787 prompt -tokens/s** across 8K–128K contexts. The linked record provides the measured -configuration, sampling settings, and single-run measurement conditions. +shows the completed three-run C1/C2/C4/C8 screen. At 8K, aggregate decode means +were **51.6, 76.9, 120.8, and 168.8 tok/s**. Against two shared-BF16-head +controls using the same CUDA version, B12X kernels, metadata reuse, and dense-kernel +integration, C1 improved **8.22% raw** and +**4.90% in normalized sequence steps/s**. Higher concurrency was mixed and +prefill means were flat within 0.36% over 8K–128K. The linked record provides +the receipt hashes, exact settings, and limitations. A separate [Estonia long-context accuracy benchmark](../performance/records/glm53-flash/spark-mtp3-country-recall-20260905.md) completed **30/30 correct answers at C8** on one repeated 133,208-token prompt, @@ -70,7 +100,7 @@ assume a Linux Bash shell and a checkout containing this guide: ```bash set -euo pipefail -git clone --branch codex/glm53-spark-mtp3-mesh https://github.com/FujitsuPolycom/sparkring.git +git clone --branch main https://github.com/FujitsuPolycom/sparkring.git cd sparkring git rev-parse HEAD test -f runtime/glm53-spark-mtp3-mesh/managed_install.py @@ -293,6 +323,15 @@ This addresses the greedy-only warmup gap tracked in Completed warmup establishes that its requests ran, not comprehensive sampling correctness or thinking-enabled generation coverage. +The pinned image supplies the defaults +`VLLM_MTP_NVFP4_LM_HEAD=1`, `VLLM_LM_HEAD_A16=1`, and +`VLLM_MXFP8_LM_HEAD=0`. The target checkpoint's unquantized `lm_head.weight` +initializes a distinct proposal-head copy on each tensor-parallel rank. For +154,880 vocabulary rows, width 4,096, and TP4, the packed NVFP4 values and +scales add approximately 85.08 MiB per rank. The retained BF16 target head +remains allocated, so 85.08 MiB is an added proposal allocation, not a net +model-memory reduction. + For a full MTP3 verification batch, target rows are approximately `Q = 4 × active requests`. Draft execution and partial batches can use different shapes. A capture list is not proof that every live step replays a @@ -333,26 +372,38 @@ uses eight I/O workers, eight load threads, eight pending operations, and two separate from SIRCL's two 64 MiB transport arenas. Native MTP's cache draft identity is the target checkpoint. The profile uses -the dedicated namespace in `pins.json`; external-draft-tagged entries must -not be renamed into it. The `draft_policy=separate` field describes cache -registration layout, not an external draft model. The linked functional record -includes an uncached publication and stopped-container restoration under this -identity. It covers one recall prompt and does not qualify other checkpoints, -all context lengths, or concurrent cache workloads. +the dedicated namespace +`glm53-spark-df116c4f-mtp3-nvfp4-a16-c139f3670-mesh69313e19-tail-cow-v2`; +shared-BF16-head and external-DFlash entries must not be renamed into it. The +`draft_policy=separate` field describes cache registration layout, not an +external draft model. The +[exact-image validation record](../performance/records/glm53-flash/spark-mtp3-compute-stream-safety-20260906.md) +includes publication and restoration of a 26,624-token prefix after all model +containers stopped and restarted under this namespace. The answer, external-hit +counters, and all four restore logs agree. This qualifies the recorded recall +case, not every context length or concurrent cache workload. ## Obtain the image and target Pull the published Linux/ARM64 managed image on every Spark. No local build is required. The [registry receipt](../runtime/glm53-spark-mtp3-mesh/public-image.json) -records anonymous access and its match to the tested image. The separate +records anonymous access and the published image's manifest/config identities. The separate [content receipt](../runtime/glm53-spark-mtp3-mesh/image-receipt.json) is the input accepted by the renderer, installer, and native qualification runner. Keep both with the checkout; do not substitute `public-image.json` for the content receipt. +The content and registry receipts identify the same public image. The +[compute-image equivalence record](../runtime/glm53-spark-mtp3-mesh/compute-image-equivalence.json) +verifies all 4,891 vLLM, 385 B12X, and 150 SparkCache package files plus the +selected environment against tested private image +`sha256:3b4768e5ba31cadcc882dffa06d7b667af44abdf157d5c11b7ac7fe962e80c43`. +The mounted transport is checked separately. The published config-image ID +below is also the exact image used for the linked runtime validation. + ```bash -mtp_image='ghcr.io/fujitsupolycom/sparkring-glm53-sparkcache@sha256:23f00af873ccc784cfb742b7be2a29c6d3c20ebec9741843c025320bb9c04685' -mtp_image_id='sha256:26273b8e358df139ae913610a5d43084ff0fd08aafe282ef633a3bc74afefe47' +mtp_image='ghcr.io/fujitsupolycom/sparkring-glm53-sparkcache@sha256:67dc0ae453baaae6831ccec1d259b4ef8b236a8b0dc9f747d901b95c66ec1987' +mtp_image_id='sha256:2e41b1e934a85ff7c21b780532db2f0a0e978df081e52f4ae2bf11f8992fb24f' docker pull "${mtp_image}" test "$(docker image inspect "${mtp_image}" --format '{{.Id}}')" = "${mtp_image_id}" @@ -403,7 +454,7 @@ docker rm sparkring-mtp3-extract cp runtime/glm53-spark-mtp3-mesh/image-receipt.json /srv/sparkring/verified-image-receipt.json printf '%s %s\n' \ - '4204fabc93303226b9a120b094ef3c82ed4aadd1d7f97cfbe291204c027ed45f' \ + '69313e19e881ec93e9ed3bd150d2f24fc6b444488ac729a69f45d038e2243500' \ '/srv/sparkring/artifacts/mtp3-mesh-bundle/sparkring-overlay-manifest.json' \ '2828c07e4255c4962c77425be2c88969e7eb7dd4b1bf9e36485bc705bb5d6d64' \ '/srv/sparkring/artifacts/mlx5-rdma-tx-marker' | sha256sum --check @@ -601,6 +652,22 @@ not performance claims. Managed markers remain under supervision throughout model loading and qualification; the test runner does not own or stop the mesh service. +### GPU stream safety + +For the selected stream-safety cases, use the same stopped-model test window +and rendered image receipt as the native check: + +```bash +python3 runtime/glm53-spark-mtp3-mesh/qualification/run_native.py \ + --launch /srv/sparkring/mtp3-mesh-launch \ + --image-receipt /srv/sparkring/verified-image-receipt.json \ + --mode streams --rows 4 64 --port 30140 \ + --output /path/to/private-receipts/stream-checks --execute-authorized +``` + +This checks alternating caller streams with misaligned buffers, changed-input +graph replay, and rejection of a second stream in one CUDA capture. It does +not inject in-flight link or GPU failures. ### Model output and persistent-cache restoration The managed startup command checks available and contiguous memory on all four diff --git a/docs/PREREQUISITES.md b/docs/PREREQUISITES.md index 11e7148a..fefa2cf7 100644 --- a/docs/PREREQUISITES.md +++ b/docs/PREREQUISITES.md @@ -86,6 +86,15 @@ cluster installer. The managed profile adds these requirements: +Before installation and the first model start, consider rebooting hosts that +have run large models or repeated GPU workloads. Stop active workloads first. +A reboot can reduce unified-memory fragmentation and restore large contiguous +free blocks; a high total-free-memory reading alone does not establish loader +readiness. This is a recommendation, not a requirement for every installation. +The [managed startup memory check](../runtime/glm53-spark-mtp3-mesh/MANAGED_MESH.md#automatic-startup-memory-preparation) +still applies after reboot and attempts compaction before requiring further +recovery. Reboots interrupt all workloads on the affected host. + - Four physical cables in the cycle `0-1-2-3-0`, with each rank's port 0/f0 connected to the next rank's port 1/f1; management stays on a separate LAN. - Four configured RDMA functions per host: primary and Socket Direct diff --git a/docs/profiles/README.md b/docs/profiles/README.md index dff6e672..b995c38e 100644 --- a/docs/profiles/README.md +++ b/docs/profiles/README.md @@ -28,6 +28,13 @@ published configuration hash. ## Research-only profiles +The [GLM-5.3 native-MTP3 cache/checkpoint mesh](../GLM53_MTP3_CACHE_CHECKPOINTS_QUICKSTART.md) +publishes the integrated compute, recurrent checkpoint, optimized SparkCache, +and transport composition for TP4/DCP4. Its +[recipe](../../recipes/glm53-mtp3-cache-checkpoints-tp4.json) pins the image and +source inputs. File equivalence with the deployed runtime is verified; that +does not constitute a serving soak of the rebuilt image. + The [GLM-5.3 Flash NVFP4-Spark TP2 profile](../GLM53_FLASH_SPARK_TP2_EXPERIMENTAL_QUICKSTART.md) is research-only. Model-neutral image and matching SparkCache source pins are public; child-image serving qualification remains pending. diff --git a/integrations/lil/glm53-mtp3.json b/integrations/lil/glm53-mtp3.json index 72abe75d..a79b1aa7 100644 --- a/integrations/lil/glm53-mtp3.json +++ b/integrations/lil/glm53-mtp3.json @@ -1,14 +1,54 @@ { "schema": "sparkring-lil-descriptor/v1", "profile": "glm53-spark-mtp3-mesh", - "transport": {"provider": "sircl", "rail_mode": "dual"}, + "transport": { + "provider": "sircl", + "rail_mode": "dual" + }, "lil_revision": "329cde801b847294005cb16765692032a6cdf206", "sources": { - "runtime": {"path": "runtime/glm53-flash-jj-r8-gb10/pins.json", "sha256": "dc49370d911ecc32f1c7ea7656f41837334a8e5ed800db263174781bee2b2313"}, - "models": {"path": "runtime/glm53-flash/pins.json", "sha256": "d8e48252bad6dfebe152f14c91e6e1e4eb01a811231e1142b65ce7b25f7f75dd"}, - "mesh": {"path": "runtime/glm53-spark-mtp3-mesh/pins.json", "sha256": "400b190cc66b32a6c3d0707d763c55d9814f1d5dc8be9adabf2e437168258a6c"}, - "public_image": {"path": "runtime/glm53-spark-mtp3-mesh/public-image.json", "sha256": "1399d2db292791475420775132dde8978019c7a31bffc411a458686385a2340e"} + "runtime": { + "path": "runtime/glm53-flash-jj-r8-gb10/pins.json", + "sha256": "dc49370d911ecc32f1c7ea7656f41837334a8e5ed800db263174781bee2b2313" + }, + "models": { + "path": "runtime/glm53-flash/pins.json", + "sha256": "d8e48252bad6dfebe152f14c91e6e1e4eb01a811231e1142b65ce7b25f7f75dd" + }, + "mesh": { + "path": "runtime/glm53-spark-mtp3-mesh/pins.json", + "sha256": "05fddc68d2e74e00ed29126c09258766aff8a852b67681836023c097234d728b" + }, + "public_image": { + "path": "runtime/glm53-spark-mtp3-mesh/public-image.json", + "sha256": "e4e518cd17f9bdccd12484be43c7e83ff70751992e62cca292c5a18281b38124" + } + }, + "defaults": { + "tp": 4, + "dcp": 4, + "max_model_len": 1048576, + "max_num_seqs": 16, + "max_num_batched_tokens": 8192, + "kv_cache_memory_bytes": 25769803776, + "speculator": "mtp", + "speculative_tokens": 3, + "port": 8015 }, - "defaults": {"tp": 4, "dcp": 4, "max_model_len": 1048576, "max_num_seqs": 16, "max_num_batched_tokens": 8192, "kv_cache_memory_bytes": 25769803776, "speculator": "mtp", "speculative_tokens": 3, "port": 8015}, - "supported": {"tp": [4], "dcp": [4], "speculator": ["mtp"], "speculative_tokens": [3], "max_num_seqs": 16, "max_model_len": 1048576} + "supported": { + "tp": [ + 4 + ], + "dcp": [ + 4 + ], + "speculator": [ + "mtp" + ], + "speculative_tokens": [ + 3 + ], + "max_num_seqs": 16, + "max_model_len": 1048576 + } } diff --git a/integrations/lil/test_export.py b/integrations/lil/test_export.py index 2ee726ce..0fbeac16 100644 --- a/integrations/lil/test_export.py +++ b/integrations/lil/test_export.py @@ -212,9 +212,12 @@ def test_mtp3_uses_published_mesh_image_and_target_identity(): descriptor = module.read_json(HERE / "glm53-mtp3.json") site = module.read_json(HERE / "site-mtp3.example.json") bundle = module.export(descriptor, site, fabric(), "mtp-fixture") + public_image = module.read_json( + HERE.parents[1] / descriptor["sources"]["public_image"]["path"] + ) for r in bundle["ranks"]: argv = r["argv"] - assert any("sha256:23f00af8" in a for a in argv) + assert public_image["public_reference"] in argv assert not any("/dflash-draft" in a for a in argv) speculation = json.loads(argv[argv.index("--speculative-config") + 1]) assert ( diff --git a/performance/harnesses/validation/README.md b/performance/harnesses/validation/README.md index 5bef1c6b..d80b2af6 100644 --- a/performance/harnesses/validation/README.md +++ b/performance/harnesses/validation/README.md @@ -81,3 +81,116 @@ unique-prefix samples per context. It calibrates prompts using `/tokenize`, records actual usage and cache evidence, and writes incremental JSONL receipts. It requests one output token; its metric is prompt tokens divided by TTFT. See the shared runbook for commands and the complete measurement protocol. + +## Growing-conversation soak and idle probes + +Status: **implemented**. `conversation_soak.py` reuses the prefill probe's +HTTP and streaming helpers to measure concurrent conversations that grow by a +calibrated number of prompt tokens each turn. It runs sequential, unique-prefix +5,300-token probes before and immediately after the conversation workers finish. +It never changes the serving stack, speculation method, cache access mode, or +namespace. Select those externally and record their identities with `--metadata`. +The `--arm` option labels a receipt; it does not configure the server. + +Print the request and output bounds without any network activity first. This +example targets the GLM-5.3 native-MTP3 profile: + +```bash +python3 performance/harnesses/validation/conversation_soak.py \ + --plan --model glm-5.3-flash-spark --arm baseline-mtp3 --context-limit 1m \ + --concurrency 2 --start-tokens 32768 --max-turns-per-agent 4 \ + --tail-tokens 2048 --max-tokens 64 --probe-output-tokens 64 \ + --duration-seconds 300 --max-soak-prompt-tokens 1000000 \ + --reset-tokens 160000 --seed 2026090601 +``` + +Replace `--plan` with `--endpoint http://192.0.2.10:8015` and +`--output /existing-results-directory/baseline-mtp3.jsonl` to execute. The +short shape admits at most eight conversation requests and six idle probes. +Every request is guarded by successful server tokenization and the declared +context limit. Calibration preserves prior messages and includes the same +chat-template options used for generation. + +For a longer MTP3 run, use `--concurrency 4 --start-tokens 100000 +--max-turns-per-agent 200 --duration-seconds 3600 --max-soak-prompt-tokens +50000000 --reset-tokens 160000 --max-tokens 512 --probe-output-tokens 300`. +Keep the remaining options explicit and run `--plan` first. The time limit +stops admission of new soak requests; calibration, probes, and requests already +in flight can extend wall time. The HTTP timeout is an inactivity timeout. +The global token ceiling counts full prompt tokens admitted across all workers, +including cached tokens, and excludes the separately bounded idle probes. +Each worker stops when its next prompt cannot fit the remaining shared token +budget; other workers may still admit smaller prompts. Token or turn limits +can end the run before the configured duration. Workers rotate to a fresh +conversation before its next prompt and output would exceed `--reset-tokens`. +Request errors stop further soak admissions; other in-flight requests can finish. + +Use identical seeds, shapes, template settings, MTP3 settings and model +identities for paired runs. Provide isolated persistent namespaces externally +when cold independence matters, or explicitly record retained cache state. +The same seed produces the same initial synthetic prompts, while actual +assistant replies grow subsequent histories; receipts preserve those replies +and prompt hashes. Distinct seeds create fresh contexts but do not establish a +byte-identical paired comparison. Probe prefixes differ between the before and +after phases; reported cached-token counts remain available to check whether +either population was already cached. A run name does not isolate GPU prefix +caching or disk state. + +`--image /path/to/fixture.png --image-every 10` adds the same local PNG or JPEG +to every tenth continuation. The receipt records its SHA-256. Keep the fixture +alongside the receipt when reproducing a multimodal run. Image bytes are sent +as a data URL and never fetched from another server. The image must be at most +4 MiB. Tokenization must account for it; an endpoint that cannot tokenize this +shape fails calibration instead of using an estimated count. + +Each turn records the client request ID sent in `X-Request-ID`, response and +server request IDs when returned, start time, actual usage, tokenized prompt +count, prompt hash, assistant text and reasoning, finish reason, elapsed time, +TTFT, and timestamp offsets for every content or reasoning delta. TTFT runs +from client submission to the first nonempty content or reasoning delta. +The decode estimate is `(completion_tokens - 1) / (last_delta - first_delta)`; +streaming chunks may contain several tokens, so this is not a token-level +inter-token latency measurement. The full delta timestamps and usage allow +recalculation. `finish_reason=length` is accepted for the bounded load and is +recorded as output-budget exhaustion, not a correctness pass. + +The summary compares continuation turns with reported cached tokens below or +above half the prompt, using the condition of a positive prompt increase below +10,000 tokens in the same conversation. Missing cached-token usage is excluded +from that classification. These are server-reported numbers: the harness does +not infer local reuse, external restore, or recomputation from latency. Correlate +request IDs with connector and vLLM logs to establish those causes. Tokenization +uses the serving host's CPU and is excluded from each chat latency; its overhead +and in-flight work affect the interval before the after probes. + +Analyze a saved receipt offline with: + +```bash +python3 performance/harnesses/validation/conversation_soak.py \ + --model glm-5.3-flash-spark --arm analysis --context-limit 1m \ + --analyze /existing-results-directory/baseline-mtp3.jsonl +``` + +JSONL records are flushed after each completion and existing output files are +never overwritten. API keys come from the environment and are omitted from +receipts; error bodies are omitted as well. Preserve model/image/source IDs, +topology, actual KV allocation, access mode, namespace, concurrent traffic and +initial store occupancy in the supplied metadata or an accompanying evidence +record. One short successful run does not qualify a long soak or a performance +improvement. Offline tests use in-memory HTTP fixtures and contact no hosts. + +`analyze_conversation_reuse.py` joins a saved receipt to the compact +`sparkcache-reuse-trace/v1` records in saved rank logs. Supply the expected +physical ranks explicitly. A verified-restore classification requires a +successful completion for the same engine request, digest and token span on +every expected rank; an offer alone is insufficient. Scheduler lease-attachment +events are reported separately. API-only evidence is labeled `reported_cached`: +positive counts do not identify local versus external reuse, and zero counts +do not exclude a GPU lease. + +The join accepts recorded request/response IDs and the runtime's exact +eight-hex-character engine suffix. Ambiguous joins cannot establish a restore +quorum. The output preserves per-rank queue, service and phase timings without +summing token spans across ranks. Log captures can contain requests outside the +receipt window; inspect unmatched IDs and timestamps before assigning them to a +missing-rank or request failure. diff --git a/performance/harnesses/validation/analyze_conversation_reuse.py b/performance/harnesses/validation/analyze_conversation_reuse.py new file mode 100644 index 00000000..ff4dd5de --- /dev/null +++ b/performance/harnesses/validation/analyze_conversation_reuse.py @@ -0,0 +1,233 @@ +"""Join saved conversation receipts to rank-local SparkCache reuse traces; no network access.""" +import argparse +from collections import Counter, defaultdict +import hashlib +import json +import math +from pathlib import Path +import re +import statistics + +TRACE_SCHEMA = "sparkcache-reuse-trace/v1" +MARKER = "spark-context-cache-reuse:" + + +def parse_trace(line): + """Accept Docker log prefixes and JSON log envelopes without interpreting prose.""" + raw = line.split(MARKER, 1)[1].lstrip() if MARKER in line and not line.lstrip().startswith("{") else line.lstrip() + try: + value, _ = json.JSONDecoder().raw_decode(raw) + except (ValueError, TypeError): + return None + if not isinstance(value, dict): + return None + if value.get("schema") == TRACE_SCHEMA: + return value + for field in ("log", "message"): + if isinstance(value.get(field), str) and MARKER in value[field]: + return parse_trace(value[field]) + return None + + +def aliases(record): + values = set() + for name in ("request_id", "response_id", "server_request_id"): + value = record.get(name) + if isinstance(value, str) and value: + values.add(value) + values.add(value.removeprefix("chatcmpl-") if value.startswith("chatcmpl-") else "chatcmpl-" + value) + return values + + +def finite_number(value): + return type(value) in (int, float) and math.isfinite(value) and value >= 0 + + +def median(values): + valid = [value for value in values if finite_number(value)] + return statistics.median(valid) if valid else None + + +def latest_completions(traces, expected): + candidates = defaultdict(list) + for trace in traces: + rank = trace.get("rank") + if (trace.get("event") == "worker_restore_completed" and trace.get("role") == "worker" + and type(rank) is int and rank in expected and type(trace.get("time_ns")) is int): + candidates[rank].append(trace) + latest, conflicts = {}, [] + for rank, rows in candidates.items(): + newest = max(row["time_ns"] for row in rows) + current = [row for row in rows if row["time_ns"] == newest] + identities = {json.dumps(row, sort_keys=True) for row in current} + if len(identities) > 1: + conflicts.append(rank) + else: + latest[rank] = current[0] + return latest, conflicts + + +def classify(turn, traces, expected, *, ambiguous=False): + latest, conflicts = latest_completions(traces, expected) + verified = len(latest) == len(expected) and not conflicts + identities = set() + for row in latest.values(): + span = row.get("requested_span_tokens") + if type(span) is not int or not isinstance(row.get("digest"), str): + verified = False + continue + verified &= (row.get("outcome") == "verified" and span > 0 + and type(row.get("verified_span_tokens")) is int + and row.get("verified_span_tokens") == span + and bool(row["digest"])) + identities.add((row.get("digest"), span, row.get("request_id"))) + verified = bool(verified and len(identities) == 1 and turn.get("valid") and not ambiguous) + attached = [row for row in traces if row.get("event") == "gpu_lease_attached" + and row.get("role") == "scheduler" and type(row.get("lease_span_tokens")) is int + and row["lease_span_tokens"] > 0] + cached = turn.get("cached_tokens_reported") + usage = turn.get("usage") or {} + prompt = usage.get("prompt_tokens") + positive_report = type(cached) is int and type(prompt) is int and 0 < cached <= prompt + if ambiguous or not turn.get("valid"): + source = "unknown" + elif verified: + source = "verified_all_rank_external_restore" + elif attached: + source = "gpu_lease_attached" + elif positive_report: + source = "reported_cached" + else: + source = "unknown" + return {"type": "turn_analysis", "request_id": turn.get("request_id"), + "response_id": turn.get("response_id"), "identity": turn.get("identity"), + "phase": turn.get("phase"), "continuation": turn.get("continuation", False), + "elapsed_seconds": turn.get("elapsed_seconds"), "ttft_seconds": turn.get("ttft_seconds"), + "source": source, "trace_join_ambiguous": ambiguous, + "cached_tokens_reported": cached, "prompt_tokens_reported": prompt, + "all_rank_worker_verification": verified, + "verified_span_tokens": next(iter(identities))[1] if verified else None, + "gpu_lease_attachment_observed": bool(attached), + "lease_span_tokens_observed": max((row["lease_span_tokens"] for row in attached), default=None), + "offer_count": sum(row.get("event") == "external_restore_offer" for row in traces), + "expected_ranks": sorted(expected), "completed_ranks": sorted(latest), + "missing_ranks": sorted(expected - latest.keys()), "conflicting_ranks": sorted(conflicts), + "worker_completions_by_rank": {str(rank): row for rank, row in sorted(latest.items())}, + "trace_events": traces, + "reported_cache_interpretation": "API cached tokens only; positive counts do not distinguish local reuse from external restore, and zero counts do not exclude GPU lease reuse"} + + +def analyze(receipts, traces, expected_ranks): + expected = set(expected_ranks) + if not expected or any(type(rank) is not int or rank < 0 for rank in expected): + raise ValueError("Expected physical ranks must be explicit nonnegative integers") + turns = [row for row in receipts if row.get("type") == "turn"] + owners, response_owners, joined, ambiguous = defaultdict(set), defaultdict(set), defaultdict(list), set() + for index, turn in enumerate(turns): + for alias in aliases(turn): + owners[alias].add(index) + response_id = turn.get("response_id") + if isinstance(response_id, str) and response_id: + response_owners[response_id].add(index) + seen, unmatched, duplicated = set(), 0, 0 + for trace in traces: + encoded = json.dumps(trace, sort_keys=True) + if encoded in seen: + duplicated += 1 + continue + seen.add(encoded) + request_id = trace.get("request_id") + matches = set(owners.get(request_id, set())) if isinstance(request_id, str) else set() + if isinstance(request_id, str): + # vLLM appends an eight-hex engine nonce to the observed API response + # ID. Only that exact suffix grammar is accepted; never prefix-match. + match = re.fullmatch(r"(.+)-[0-9a-fA-F]{8}", request_id) + if match: + matches.update(response_owners.get(match.group(1), set())) + if len(matches) == 1: + joined[next(iter(matches))].append(trace) + elif matches: + ambiguous.update(matches) + else: + unmatched += 1 + rows = [classify(turn, joined[index], expected, ambiguous=index in ambiguous) + for index, turn in enumerate(turns)] + by_source = {} + by_phase_and_source = {} + for source in sorted({row["source"] for row in rows}): + population = [row for row in rows if row["source"] == source] + by_source[source] = {"turns": len(population), + "median_latency_seconds": median(row["elapsed_seconds"] for row in population), + "median_ttft_seconds": median(row["ttft_seconds"] for row in population)} + for phase in sorted({row.get("phase") or "unknown" for row in rows}): + by_phase_and_source[phase] = {} + for source in by_source: + population = [row for row in rows if (row.get("phase") or "unknown") == phase and row["source"] == source] + if population: + by_phase_and_source[phase][source] = { + "turns": len(population), "continuations": sum(row["continuation"] for row in population), + "median_latency_seconds": median(row["elapsed_seconds"] for row in population), + "median_ttft_seconds": median(row["ttft_seconds"] for row in population)} + timings = {} + for rank in sorted(expected): + population = [trace for row in rows for trace in row["trace_events"] + if trace.get("role") == "worker" and trace.get("event") == "worker_restore_completed" + and type(trace.get("rank")) is int and trace["rank"] == rank] + phases = sorted({phase for row in population for phase in (row.get("phase_ms") or {})}) + timings[str(rank)] = {"completions": len(population), + "outcomes": dict(Counter(row.get("outcome", "unknown") for row in population)), + **{f"median_{field}": median(row.get(field) for row in population) + for field in ("queue_wait_ms", "service_ms", "end_to_end_ms")}, + "median_phase_ms": {phase: median((row.get("phase_ms") or {}).get(phase) for row in population) + for phase in phases}} + return {"schema": "sparkring-conversation-reuse-analysis/v1", "trace_schema": TRACE_SCHEMA, + "expected_ranks": sorted(expected), "turns": rows, + "summary": {"turns": len(rows), "by_source": by_source, "by_phase_and_source": by_phase_and_source, + "per_rank_restore_timings": timings, + "unmatched_trace_events": unmatched, "duplicate_trace_events": duplicated, + "ambiguous_turns": len(ambiguous)}, + "limits": ["No recomputation or local-miss inference from zero or missing API cache counts", + "Offers do not prove an executed restore; every expected physical rank must verify the same span and digest", + "Worker verification plus successful client completion does not expose vLLM's internal receive-aggregation decision", + "An observed lease attachment proves attachment, not its full subsequent lifetime", + "Rank token counts are never summed; source classes may coexist and the strongest observed evidence is selected"]} + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--receipt", type=Path, required=True) + parser.add_argument("--log", type=Path, action="append", required=True, + help="Saved rank log; repeat for every physical rank") + parser.add_argument("--expected-ranks", required=True, help="Physical ranks, for example 0,1,2,3") + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + try: + ranks = [int(value) for value in args.expected_ranks.split(",")] + if len(set(ranks)) != len(ranks): + raise ValueError("Expected ranks must be unique") + receipts = [json.loads(line) for line in args.receipt.read_text(encoding="utf-8").splitlines() if line.strip()] + traces, malformed, sources = [], 0, [] + for path in args.log: + sources.append({"path": str(path), "sha256": hashlib.sha256(path.read_bytes()).hexdigest()}) + with path.open(encoding="utf-8", errors="replace") as stream: + for line in stream: + record = parse_trace(line) + if record: + traces.append(record) + elif MARKER in line: + malformed += 1 + result = analyze(receipts, traces, ranks) + result["inputs"] = {"receipt": str(args.receipt), + "receipt_sha256": hashlib.sha256(args.receipt.read_bytes()).hexdigest(), "logs": sources} + result["summary"]["malformed_trace_lines"] = malformed + with args.output.open("x", encoding="utf-8") as stream: + json.dump(result, stream, indent=2) + stream.write("\n") + except (ValueError, OSError) as error: + parser.error(str(error)) + print(json.dumps(result["summary"], sort_keys=True)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/performance/harnesses/validation/conversation_soak.py b/performance/harnesses/validation/conversation_soak.py new file mode 100644 index 00000000..266cadc5 --- /dev/null +++ b/performance/harnesses/validation/conversation_soak.py @@ -0,0 +1,323 @@ +"""Measure bounded growing conversations and idle probes; never change server cache state.""" +import argparse +import base64 +from concurrent.futures import ThreadPoolExecutor +import hashlib +import json +import math +import os +from pathlib import Path +import statistics +import sys +import threading +import time +import uuid + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from needle_hunt import base_url # noqa: E402 +from prefill_probe import build_text, events, has_token, request, tokens # noqa: E402 + +SCHEMA = "sparkring-conversation-soak/v1" + + +def digest(value): + return hashlib.sha256(json.dumps(value, sort_keys=True, separators=(",", ":")).encode()).hexdigest() + + +def fixture(identity, chars): + text = build_text(identity, chars) + return text.replace("Read the notes and reply OK.", + "Analyze these maintenance notes in detail and recommend next actions.", 1) + + +def count_tokens(config, messages): + payload = {"model": config["model"], "messages": messages, "add_generation_prompt": True, + "chat_template_kwargs": config["chat_template_kwargs"]} + with request(config["endpoint"] + "/tokenize", payload, config["api_key"], config["timeout"]) as response: + count = json.load(response).get("count") + if type(count) is not int or count <= 0: + raise ValueError("Tokenizer must return a positive integer count") + return count + + +def calibrated_user(config, history, identity, target_total, image=None): + """Calibrate the added user message while preserving every preceding message byte.""" + previous = count_tokens(config, history) if history else 0 + chars = max(256, (target_total - previous) * 5) + for _ in range(10): + text = fixture(identity, chars) + content = text if image is None else [ + {"type": "text", "text": text}, {"type": "image_url", "image_url": {"url": image}}] + messages = history + [{"role": "user", "content": content}] + count = count_tokens(config, messages) + if abs(count - target_total) <= max(8, (target_total - previous) * 0.01): + return messages, count + chars = max(128, min(64 * 1024 * 1024, round(chars * max(1, target_total - previous) + / max(1, count - previous)))) + raise ValueError("Token calibration did not converge") + + +def stream_turn(config, messages, count, max_tokens, identity, phase, *, request_id=None, clock=time.perf_counter): + if count + max_tokens > config["context_limit"]: + raise ValueError("Tokenized prompt plus output exceeds the context limit") + request_id = request_id or "soak-" + uuid.uuid4().hex + payload = {"model": config["model"], "messages": messages, "max_tokens": max_tokens, + "stream": True, "stream_options": {"include_usage": True}, + "temperature": config["temperature"], "seed": config["seed"], "top_p": 1, + "chat_template_kwargs": config["chat_template_kwargs"]} + effort = config["probe_reasoning_effort"] if phase != "soak" else config["reasoning_effort"] + if effort: + payload["reasoning_effort"] = effort + chunks, content, reasoning, usage, response_id, finish = [], [], [], None, None, None + started_unix, started = time.time(), clock() + with request(config["endpoint"] + "/v1/chat/completions", payload, config["api_key"], + config["timeout"], request_id=request_id) as response: + header_id = response.headers.get("X-Request-ID") if hasattr(response, "headers") else None + for event in events(response): + now = clock() + if event.get("id"): + response_id = event["id"] + if has_token(event): + chunks.append(now - started) + if event.get("usage"): + usage = event["usage"] + for choice in event.get("choices", []): + delta = choice.get("delta", {}) + if delta.get("content"): + content.append(delta["content"]) + if delta.get("reasoning_content") or delta.get("reasoning"): + reasoning.append(delta.get("reasoning_content") or delta["reasoning"]) + finish = choice.get("finish_reason") or finish + elapsed = clock() - started + if not chunks or not isinstance(usage, dict) or type(usage.get("prompt_tokens")) is not int or usage["prompt_tokens"] <= 0: + raise ValueError("Missing first delta or authoritative prompt usage") + completion = usage.get("completion_tokens") + if type(completion) is not int or completion < 1 or finish not in ("stop", "length"): + raise ValueError("Missing output usage or normal finish reason") + cached = (usage.get("prompt_tokens_details") or {}).get("cached_tokens") + if cached is not None and (type(cached) is not int or not 0 <= cached <= usage["prompt_tokens"]): + raise ValueError("Invalid reported cached-token count") + assistant = {"role": "assistant", "content": "".join(content)} + if reasoning: + assistant["reasoning_content"] = "".join(reasoning) + span = chunks[-1] - chunks[0] + record = {"type": "turn", "phase": phase, "identity": identity, "valid": True, + "request_id": request_id, "server_request_id": header_id, "response_id": response_id, + "started_unix": started_unix, "prompt_sha256": digest(messages), + "assistant": assistant, "tokenized_prompt_tokens": count, "usage": usage, + "cached_tokens_reported": cached, "cached_fraction_reported": cached / usage["prompt_tokens"] + if cached is not None and usage["prompt_tokens"] > 0 else None, + "ttft_seconds": chunks[0], "elapsed_seconds": elapsed, + "content_delta_offsets_seconds": chunks, "decode_span_seconds": span, + "decode_tokens_per_second_estimate": (completion - 1) / span if span > 0 else None, + "finish_reason": finish, "output_budget_exhausted": finish == "length"} + return record, assistant + + +def median(values): + available = [value for value in values if value is not None] + return statistics.median(available) if available else None + + +def summarize(records): + good = [r for r in records if r.get("type") == "turn" and r.get("valid")] + continuation = [r for r in good if r.get("continuation")] + known = [r for r in continuation if r["cached_fraction_reported"] is not None] + low = [r for r in known if r["cached_fraction_reported"] < 0.5] + high = [r for r in known if r["cached_fraction_reported"] >= 0.5] + probes = {phase: {"samples": len(rows), + "cached_token_counts_reported": [r["cached_tokens_reported"] for r in rows], + "median_ttft_seconds": median(r["ttft_seconds"] for r in rows), + "median_decode_tokens_per_second_estimate": median(r["decode_tokens_per_second_estimate"] for r in rows)} + for phase in ("before", "after") + for rows in [[r for r in good if r["phase"] == phase]]} + return {"type": "summary", "valid_turns": len(good), + "soak_turns": sum(r["phase"] == "soak" for r in good), + "errors": sum(r.get("type") == "error" for r in records), + "continuations": len(continuation), "continuations_with_cache_usage": len(known), + "continuations_below_half_cached": len(low), + "fraction_below_half_cached": len(low) / len(known) if known else None, + "low_cached_median_latency_seconds": median(r["elapsed_seconds"] for r in low), + "high_cached_median_latency_seconds": median(r["elapsed_seconds"] for r in high), + "probes": probes, + "cache_interpretation": "server-reported cached tokens; no inference of local hit, external restore, or recompute"} + + +def execute(config, output): + lock, stop = threading.Lock(), threading.Event() + records, consumed = [], 0 + image = config.get("image_data") + safe_config = {key: value for key, value in config.items() if key not in ("api_key", "image_data")} + with Path(output).open("x", encoding="utf-8") as stream: + def emit(record): + with lock: + records.append(record) + line = json.dumps(record) + stream.write(line + "\n") + stream.flush() + print(line, flush=True) + + emit({"type": "start", "schema": SCHEMA, "config": safe_config, + "harness_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest(), + "http_helper_sha256": hashlib.sha256(Path(__file__).with_name("prefill_probe.py").read_bytes()).hexdigest(), + "started_unix": time.time()}) + + def probe(phase, index): + identity = f"{config['seed']}-probe-{phase}-{index}" + request_id = "soak-" + uuid.uuid4().hex + try: + messages, count = calibrated_user(config, [], identity, config["probe_tokens"]) + record, _ = stream_turn(config, messages, count, config["probe_output_tokens"], identity, + phase, request_id=request_id) + emit(record) + except Exception as error: + emit({"type": "error", "phase": phase, "identity": identity, + "request_id": request_id, "error": type(error).__name__}) + stop.set() + + def worker(agent, deadline): + nonlocal consumed + history, prior_count, conversation, turn = [], None, 0, 0 + for _ in range(config["max_turns_per_agent"]): + if stop.is_set() or time.monotonic() >= deadline: + break + if prior_count is not None and prior_count + config["tail_tokens"] + config["max_tokens"] > config["reset_tokens"]: + history, prior_count, turn = [], None, 0 + conversation += 1 + identity = f"{config['seed']}-agent-{agent}-conversation-{conversation}-turn-{turn}" + request_id = "soak-" + uuid.uuid4().hex + try: + initial = config["start_tokens"][agent % len(config["start_tokens"])] + target = initial if prior_count is None else prior_count + config["tail_tokens"] + picture = image if image and turn > 0 and turn % config["image_every"] == 0 else None + messages, count = calibrated_user(config, history, identity, target, picture) + if count + config["max_tokens"] > config["context_limit"]: + raise ValueError("Context limit exceeded") + with lock: + if stop.is_set() or time.monotonic() >= deadline or consumed + count > config["max_soak_prompt_tokens"]: + break + consumed += count + record, assistant = stream_turn(config, messages, count, config["max_tokens"], identity, + "soak", request_id=request_id) + record.update(agent=agent, conversation=conversation, turn=turn, + prompt_token_growth=count - prior_count if prior_count is not None else None, + continuation=prior_count is not None and 0 < count - prior_count < 10000, + image_added=picture is not None) + emit(record) + history, prior_count, turn = messages + [assistant], count, turn + 1 + except Exception as error: + emit({"type": "error", "phase": "soak", "identity": identity, + "request_id": request_id, "error": type(error).__name__}) + stop.set() + break + + try: + for index in range(config["probe_repeats"]): + probe("before", index) + if stop.is_set(): + break + deadline = time.monotonic() + config["duration_seconds"] + with ThreadPoolExecutor(max_workers=config["concurrency"]) as pool: + list(pool.map(lambda agent: worker(agent, deadline), range(config["concurrency"]))) + if not stop.is_set(): + for index in range(config["probe_repeats"]): + probe("after", index) + if stop.is_set(): + break + except Exception as error: + emit({"type": "error", "phase": "probe", "error": type(error).__name__}) + summary = summarize(records) + summary["soak_prompt_tokens_admitted"] = consumed + summary["success"] = not summary["errors"] and summary["soak_turns"] > 0 + summary["completed_unix"] = time.time() + emit(summary) + return 0 if summary["success"] else 2 + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--endpoint") + parser.add_argument("--model", required=True) + parser.add_argument("--arm", required=True, help="Evidence label only; never changes server configuration") + parser.add_argument("--seed", type=int, default=20260906) + parser.add_argument("--concurrency", type=int, default=4) + parser.add_argument("--start-tokens", default="75000,106000,128000,150000") + parser.add_argument("--tail-tokens", type=tokens, default=2048) + parser.add_argument("--reset-tokens", type=tokens, default=300000) + parser.add_argument("--context-limit", type=tokens, required=True) + parser.add_argument("--max-tokens", type=tokens, default=512) + parser.add_argument("--duration-seconds", type=float, default=300) + parser.add_argument("--max-turns-per-agent", type=int, default=20) + parser.add_argument("--max-soak-prompt-tokens", type=int, default=10000000) + parser.add_argument("--temperature", type=float, default=1) + parser.add_argument("--reasoning-effort", default="low") + parser.add_argument("--probe-reasoning-effort", default="low") + parser.add_argument("--probe-tokens", type=tokens, default=5300) + parser.add_argument("--probe-output-tokens", type=tokens, default=300) + parser.add_argument("--probe-repeats", type=int, default=3) + parser.add_argument("--timeout", type=float, default=900, help="HTTP inactivity timeout, not a wall-clock deadline") + parser.add_argument("--api-key-env", default="OPENAI_API_KEY") + parser.add_argument("--chat-template-kwargs", default="{}") + parser.add_argument("--metadata", default="{}", help="JSON evidence identities; do not include credentials") + parser.add_argument("--image", type=Path, help="Optional local PNG or JPEG, at most 4 MiB") + parser.add_argument("--image-every", type=int, default=10) + parser.add_argument("--output", type=Path) + parser.add_argument("--plan", action="store_true", help="Print bounds without tokenization or HTTP requests") + parser.add_argument("--analyze", type=Path, help="Summarize an existing JSONL receipt without HTTP") + args = parser.parse_args() + if args.analyze: + print(json.dumps(summarize([json.loads(line) for line in args.analyze.read_text(encoding="utf-8").splitlines()]))) + return 0 + config = vars(args).copy() + for name in ("output", "plan", "analyze", "image"): + config.pop(name) + try: + config["start_tokens"] = [tokens(value) for value in args.start_tokens.split(",")] + config["chat_template_kwargs"] = json.loads(args.chat_template_kwargs) + config["metadata"] = json.loads(args.metadata) + if not isinstance(config["chat_template_kwargs"], dict) or not isinstance(config["metadata"], dict): + raise ValueError("Template options and metadata must be JSON objects") + if not 1 <= args.concurrency <= 32 or not 1 <= args.max_turns_per_agent <= 10000 or not 0 <= args.probe_repeats <= 20 or args.image_every < 1: + raise ValueError("Invalid concurrency, repetition, or image frequency") + if not math.isfinite(args.duration_seconds) or not 0 < args.duration_seconds <= 14400: + raise ValueError("Duration must be positive and at most 4 hours") + if not 1 <= args.max_soak_prompt_tokens <= 1000000000: + raise ValueError("Soak token budget must be between 1 and 1000000000") + if not math.isfinite(args.timeout) or not 0 < args.timeout <= 3600 or not 0 <= args.temperature <= 2: + raise ValueError("Invalid timeout or temperature") + if max(config["start_tokens"]) + args.max_tokens > args.reset_tokens or args.reset_tokens > args.context_limit: + raise ValueError("Starting prompts plus output must fit reset threshold and context limit") + if args.probe_tokens + args.probe_output_tokens > args.context_limit: + raise ValueError("Probe prompt plus output must fit context limit") + if args.endpoint: + config["endpoint"] = base_url(args.endpoint) + if args.image: + if args.image.stat().st_size > 4 * 1024 * 1024: + raise ValueError("Image exceeds 4 MiB") + raw = args.image.read_bytes() + mime = "image/png" if raw.startswith(b"\x89PNG\r\n\x1a\n") else "image/jpeg" if raw.startswith(b"\xff\xd8\xff") else None + if mime is None: + raise ValueError("Image must be PNG or JPEG") + config["image_sha256"] = hashlib.sha256(raw).hexdigest() + image_data = "data:" + mime + ";base64," + base64.b64encode(raw).decode() + else: + image_data = None + except (ValueError, OSError, argparse.ArgumentTypeError) as error: + parser.error(str(error)) + if args.plan: + print(json.dumps({"schema": SCHEMA, "config": config, + "max_chat_requests": args.concurrency * args.max_turns_per_agent + 2 * args.probe_repeats, + "max_output_tokens": args.concurrency * args.max_turns_per_agent * args.max_tokens + + 2 * args.probe_repeats * args.probe_output_tokens, + "duration_scope": "soak admission window; probes, calibration and in-flight requests may extend elapsed time", + "status": "implemented; no hardware qualification"})) + return 0 + if not args.endpoint or not args.output: + parser.error("Execution requires --endpoint and --output") + config["api_key"] = os.environ.get(args.api_key_env, "") + config["image_data"] = image_data + return execute(config, args.output) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/performance/harnesses/validation/prefill_probe.py b/performance/harnesses/validation/prefill_probe.py index d89134f3..a47b9c62 100644 --- a/performance/harnesses/validation/prefill_probe.py +++ b/performance/harnesses/validation/prefill_probe.py @@ -20,8 +20,10 @@ def tokens(value): return result -def request(url, payload, key, timeout): +def request(url, payload, key, timeout, request_id=None): headers = {'Content-Type': 'application/json'} + if request_id: + headers['X-Request-ID'] = request_id if key: headers['Authorization'] = 'Bearer ' + key class NoRedirect(urllib.request.HTTPRedirectHandler): diff --git a/performance/harnesses/validation/test_analyze_conversation_reuse.py b/performance/harnesses/validation/test_analyze_conversation_reuse.py new file mode 100644 index 00000000..2f688aa5 --- /dev/null +++ b/performance/harnesses/validation/test_analyze_conversation_reuse.py @@ -0,0 +1,149 @@ +"""Reuse evidence tests are entirely offline; incomplete quorum never verifies.""" +import importlib.util +import json +from pathlib import Path + +import pytest + +spec = importlib.util.spec_from_file_location("reuse_analysis_test", Path(__file__).with_name("analyze_conversation_reuse.py")) +analysis = importlib.util.module_from_spec(spec) +spec.loader.exec_module(analysis) + + +def turn(**overrides): + return {"type": "turn", "request_id": "soak-one", "response_id": "chatcmpl-soak-one", "valid": True, + "phase": "soak", "cached_tokens_reported": 0, "usage": {"prompt_tokens": 12000}, + "elapsed_seconds": 5, "ttft_seconds": 1, **overrides} + + +def completion(rank, **overrides): + return {"schema": analysis.TRACE_SCHEMA, "event": "worker_restore_completed", + "request_id": "chatcmpl-soak-one", "role": "worker", "rank": rank, + "dcp_rank": rank % 2, "time_ns": 100, "digest": "0123456789ab", + "requested_span_tokens": 10000, "verified_span_tokens": 10000, "outcome": "verified", + "queue_wait_ms": 1 + rank, "service_ms": 20, "end_to_end_ms": 21 + rank, + "phase_ms": {"read": 12, "place": 8}, **overrides} + + +def test_log_prefix_and_docker_envelope(): + record = completion(0) + line = "2026-09-06 (Worker_TP0) INFO " + analysis.MARKER + json.dumps(record) + assert analysis.parse_trace(line) == record + assert analysis.parse_trace(json.dumps({"log": line, "stream": "stderr"})) == record + assert analysis.parse_trace("not a trace") is None + assert analysis.parse_trace(analysis.MARKER + "{partial") is None + + +def test_every_physical_rank_required_and_tokens_not_summed(): + report = analysis.analyze([turn()], [completion(rank) for rank in range(4)], [0, 1, 2, 3]) + row = report["turns"][0] + assert row["source"] == "verified_all_rank_external_restore" + assert row["verified_span_tokens"] == 10000 + assert row["missing_ranks"] == [] + assert report["summary"]["per_rank_restore_timings"]["3"]["median_queue_wait_ms"] == 4 + assert report["summary"]["per_rank_restore_timings"]["2"]["median_phase_ms"] == {"read": 12, "place": 8} + + +@pytest.mark.parametrize("traces", [ + [completion(0), completion(1)], + [completion(rank, outcome="recompute", verified_span_tokens=0) for rank in range(4)], + [completion(rank, digest="different" if rank == 3 else "0123456789ab") for rank in range(4)], + [completion(rank, requested_span_tokens=11000 if rank == 3 else 10000) for rank in range(4)], + [completion(rank, role="scheduler") for rank in range(4)], +]) +def test_partial_failed_or_mismatched_completions_do_not_verify(traces): + row = analysis.analyze([turn()], traces, [0, 1, 2, 3])["turns"][0] + assert row["source"] == "unknown" + assert row["verified_span_tokens"] is None + + +def test_offer_is_not_restore_and_zero_api_not_local_miss(): + offer = {"schema": analysis.TRACE_SCHEMA, "request_id": "chatcmpl-soak-one", "role": "scheduler", + "event": "external_restore_offer", "selected_span_tokens": 10000} + row = analysis.analyze([turn()], [offer], [0, 1, 2, 3])["turns"][0] + assert row["source"] == "unknown" + assert row["offer_count"] == 1 + assert not row["all_rank_worker_verification"] + missing = analysis.analyze([turn(cached_tokens_reported=None)], [], [0])["turns"][0] + assert missing["source"] == "unknown" + + +def test_gpu_attachment_and_positive_report_are_separate_evidence(): + attached = {"schema": analysis.TRACE_SCHEMA, "request_id": "soak-one", "event": "gpu_lease_attached", + "role": "scheduler", "lease_span_tokens": 9000} + row = analysis.analyze([turn(cached_tokens_reported=8000)], [attached], [0, 1, 2, 3])["turns"][0] + assert row["source"] == "gpu_lease_attached" + assert not row["all_rank_worker_verification"] + assert row["lease_span_tokens_observed"] == 9000 + assert analysis.analyze([turn(cached_tokens_reported=8000)], [], [0])["turns"][0]["source"] == "reported_cached" + + +def test_latest_failure_overrides_prior_success_and_duplicates_do_not_make_quorum(): + records = [completion(rank) for rank in range(4)] + records += [completion(3, time_ns=200, outcome="recompute", verified_span_tokens=0)] + report = analysis.analyze([turn()], records + records, [0, 1, 2, 3]) + assert report["turns"][0]["source"] == "unknown" + assert report["summary"]["duplicate_trace_events"] == 5 + assert analysis.analyze([turn()], [completion(0)] * 4, [0, 1, 2, 3])["turns"][0]["source"] == "unknown" + + +def test_conflicting_same_time_and_ambiguous_ids_do_not_verify(): + records = [completion(rank) for rank in range(4)] + [completion(0, outcome="recompute", verified_span_tokens=0)] + row = analysis.analyze([turn()], records, [0, 1, 2, 3])["turns"][0] + assert row["source"] == "unknown" and row["conflicting_ranks"] == [0] + report = analysis.analyze([turn(), turn()], [completion(rank) for rank in range(4)], [0, 1, 2, 3]) + assert report["summary"]["ambiguous_turns"] == 2 + assert all(row["source"] == "unknown" for row in report["turns"]) + + +def test_unmatched_id_not_fuzzily_associated_and_rank_set_explicit(): + report = analysis.analyze([turn()], [completion(0, request_id="chatcmpl-soak-one-other")], [0]) + assert report["summary"]["unmatched_trace_events"] == 1 + with pytest.raises(ValueError, match="explicit"): + analysis.analyze([turn()], [], []) + + +def test_malformed_completion_never_verifies(): + row = analysis.analyze([turn()], [completion(0, requested_span_tokens=[10000])], [0])["turns"][0] + assert row["source"] == "unknown" + row = analysis.analyze([turn()], [completion(0, digest=[])], [0])["turns"][0] + assert row["source"] == "unknown" + + +def test_phase_populations_remain_separate_and_retry_timings_retained(): + turns = [turn(), turn(phase="before", request_id="probe-id", response_id="chatcmpl-probe-id", elapsed_seconds=1)] + report = analysis.analyze(turns, [completion(0, time_ns=50, outcome="recompute", verified_span_tokens=0), completion(0)], [0]) + phases = report["summary"]["by_phase_and_source"] + assert phases["soak"]["verified_all_rank_external_restore"]["median_latency_seconds"] == 5 + assert phases["before"]["unknown"]["median_latency_seconds"] == 1 + assert report["summary"]["per_rank_restore_timings"]["0"]["outcomes"] == {"recompute": 1, "verified": 1} + + +def test_observed_response_id_engine_nonce_joins_strictly(): + response_id = "chatcmpl-soak-cc0bdb28896248beab09f145c9418947" + rows = [completion(rank, request_id=response_id + "-b4884e31") for rank in range(4)] + report = analysis.analyze([turn(response_id=response_id)], rows, [0, 1, 2, 3]) + assert report["turns"][0]["source"] == "verified_all_rank_external_restore" + assert report["summary"]["unmatched_trace_events"] == 0 + for suffix in ("-b4884e3", "-b4884e311", "-b4884e3z", "-b4884e31-more"): + report = analysis.analyze([turn(response_id=response_id)], [completion(0, request_id=response_id + suffix)], [0]) + assert report["turns"][0]["source"] == "unknown" + assert report["summary"]["unmatched_trace_events"] == 1 + + +def test_engine_nonce_requires_observed_response_id_and_refuses_collisions(): + engine_id = "chatcmpl-soak-one-1234abcd" + row = analysis.analyze([turn(response_id=None)], [completion(0, request_id=engine_id)], [0])["turns"][0] + assert row["source"] == "unknown" + turns = [turn(), turn(request_id=engine_id, response_id="other-api-id")] + report = analysis.analyze(turns, [completion(0, request_id=engine_id)], [0]) + assert report["summary"]["ambiguous_turns"] == 2 + assert all(row["source"] == "unknown" for row in report["turns"]) + + +def test_different_engine_requests_cannot_be_combined_into_quorum(): + rows = [completion(0, request_id="chatcmpl-soak-one-1234abcd"), + completion(1, request_id="chatcmpl-soak-one-5678abcd")] + row = analysis.analyze([turn()], rows, [0, 1])["turns"][0] + assert row["source"] == "unknown" + assert not row["all_rank_worker_verification"] diff --git a/performance/harnesses/validation/test_conversation_soak.py b/performance/harnesses/validation/test_conversation_soak.py new file mode 100644 index 00000000..0bb8bb6d --- /dev/null +++ b/performance/harnesses/validation/test_conversation_soak.py @@ -0,0 +1,168 @@ +"""Bounded soak tests use in-memory HTTP fixtures and never contact a host.""" +import importlib.util +import io +import json +from pathlib import Path +import sys + +import pytest + +spec = importlib.util.spec_from_file_location("soak_under_test", Path(__file__).with_name("conversation_soak.py")) +soak = importlib.util.module_from_spec(spec) +spec.loader.exec_module(soak) + + +def config(**overrides): + return {"endpoint": "http://192.0.2.1:8015", "model": "glm-5.3-flash-spark", + "api_key": "test-key", "timeout": 1, "context_limit": 8192, + "temperature": 1, "seed": 20260906, "arm": "mtp3-test", "reasoning_effort": "low", + "probe_reasoning_effort": "low", "chat_template_kwargs": {"enable_thinking": True}, + "probe_tokens": 300, "probe_output_tokens": 64, "probe_repeats": 1, + "start_tokens": [1000], "max_tokens": 64, "tail_tokens": 200, + "reset_tokens": 3000, "concurrency": 2, "duration_seconds": 30, + "max_turns_per_agent": 3, "max_soak_prompt_tokens": 20000, + "image_every": 10, **overrides} + + +def fake_count(messages): + length = 0 + for message in messages: + content = message["content"] + if isinstance(content, list): + length += sum(len(part.get("text", "")) + (500 if part["type"] == "image_url" else 0) + for part in content) + else: + length += len(content) + length += len(message.get("reasoning_content", "")) + 20 + return max(1, length // 5) + + +def fake_http(calls, *, cached=0, missing_usage=False, error=False): + def send(url, payload, key, timeout, request_id=None): + calls.append((url, payload, request_id)) + count = fake_count(payload["messages"]) + if url.endswith("/tokenize"): + return io.BytesIO(json.dumps({"count": count}).encode()) + if error: + raise ValueError("credential echoed by server: test-key") + records = [ + {"id": "chatcmpl-fixture", "choices": [{"delta": {"reasoning_content": "think"}}]}, + {"choices": [{"delta": {"content": "answer"}, "finish_reason": "length"}]}, + ] + if not missing_usage: + records.append({"usage": {"prompt_tokens": count, "completion_tokens": 3, + "prompt_tokens_details": {"cached_tokens": cached}}}) + return io.BytesIO(("".join("data: " + json.dumps(record) + "\n" for record in records) + + "data: [DONE]\n").encode()) + return send + + +def test_calibration_preserves_conversation_and_template(monkeypatch): + calls = [] + monkeypatch.setattr(soak, "request", fake_http(calls)) + history = [{"role": "user", "content": "original"}, {"role": "assistant", "content": "reply"}] + original = json.dumps(history) + messages, count = soak.calibrated_user(config(), history, "seed-agent-1", 1000) + assert abs(count - 1000) <= 10 + assert messages[:2] == history + assert json.dumps(history) == original + assert all(payload["chat_template_kwargs"] == {"enable_thinking": True} for _, payload, _ in calls) + + +def test_stream_records_ids_usage_reasoning_and_delta_clock(monkeypatch): + calls = [] + monkeypatch.setattr(soak, "request", fake_http(calls, cached=1)) + ticks = iter([100, 101, 103, 104, 105]) + record, assistant = soak.stream_turn(config(), [{"role": "user", "content": "fixture"}], + 100, 64, "identity", "soak", request_id="fixture-id", clock=lambda: next(ticks)) + assert record["request_id"] == calls[0][2] == "fixture-id" + assert record["response_id"] == "chatcmpl-fixture" + assert record["cached_tokens_reported"] == 1 + assert record["ttft_seconds"] == 1 + assert record["elapsed_seconds"] == 5 + assert record["content_delta_offsets_seconds"] == [1, 3] + assert record["decode_tokens_per_second_estimate"] == 1 + assert assistant == {"role": "assistant", "content": "answer", "reasoning_content": "think"} + assert record["output_budget_exhausted"] + assert calls[0][1]["reasoning_effort"] == "low" + + +def test_guard_and_missing_usage_fail_closed(monkeypatch): + calls = [] + monkeypatch.setattr(soak, "request", fake_http(calls, missing_usage=True)) + with pytest.raises(ValueError, match="context limit"): + soak.stream_turn(config(), [], 8192, 64, "id", "soak") + assert not calls + with pytest.raises(ValueError, match="authoritative"): + soak.stream_turn(config(), [{"role": "user", "content": "text"}], 50, 64, "id", "soak") + + +def test_bounded_conversations_and_before_after_probes(monkeypatch, tmp_path): + calls = [] + monkeypatch.setattr(soak, "request", fake_http(calls)) + path = tmp_path / "receipt.jsonl" + assert soak.execute(config(), path) == 0 + records = [json.loads(line) for line in path.read_text().splitlines()] + turns = [record for record in records if record["type"] == "turn"] + assert len(turns) == 8 + assert [record["phase"] for record in turns][0] == "before" + assert [record["phase"] for record in turns][-1] == "after" + assert sum(record.get("continuation", False) for record in turns) == 4 + assert len({record["request_id"] for record in turns}) == 8 + assert len({record["prompt_sha256"] for record in turns}) == 8 + assert records[-1]["soak_turns"] == 6 + assert records[-1]["fraction_below_half_cached"] == 1 + assert "test-key" not in path.read_text() + with pytest.raises(FileExistsError): + soak.execute(config(), path) + + +def test_global_admission_budget_counts_all_agents(monkeypatch, tmp_path): + monkeypatch.setattr(soak, "request", fake_http([])) + path = tmp_path / "budget.jsonl" + assert soak.execute(config(max_soak_prompt_tokens=2100), path) == 0 + summary = json.loads(path.read_text().splitlines()[-1]) + assert summary["soak_prompt_tokens_admitted"] <= 2100 + assert summary["soak_turns"] == 2 + assert summary["probes"]["after"]["samples"] == 1 + + +def test_failure_stops_before_soak_and_redacts_error(monkeypatch, tmp_path): + calls = [] + monkeypatch.setattr(soak, "request", fake_http(calls, error=True)) + path = tmp_path / "failed.jsonl" + assert soak.execute(config(), path) == 2 + records = [json.loads(line) for line in path.read_text().splitlines()] + errors = [record for record in records if record["type"] == "error"] + assert len(errors) == 1 and errors[0]["request_id"] + assert records[-1]["soak_turns"] == 0 + assert "test-key" not in path.read_text() + assert sum(url.endswith("/v1/chat/completions") for url, _, _ in calls) == 1 + + +def test_unknown_cache_usage_is_not_a_miss(): + rows = [{"type": "turn", "phase": "soak", "valid": True, "continuation": True, + "cached_fraction_reported": None, "elapsed_seconds": 10}] + summary = soak.summarize(rows) + assert summary["continuations"] == 1 + assert summary["continuations_with_cache_usage"] == 0 + assert summary["fraction_below_half_cached"] is None + + +def test_fixture_seed_reproducible_and_image_added(monkeypatch): + assert soak.fixture("seed-one", 1000) == soak.fixture("seed-one", 1000) + assert soak.fixture("seed-one", 1000) != soak.fixture("seed-two", 1000) + monkeypatch.setattr(soak, "request", fake_http([])) + messages, count = soak.calibrated_user(config(), [], "image-seed", 1000, "data:image/png;base64,fixture") + assert messages[0]["content"][1]["type"] == "image_url" + assert abs(count - 1000) <= 10 + + +def test_plan_requires_no_endpoint_or_key(monkeypatch, capsys): + monkeypatch.setattr(soak, "request", lambda *args, **kwargs: pytest.fail("unexpected HTTP")) + monkeypatch.setattr(sys, "argv", ["conversation_soak.py", "--plan", "--model", "glm-5.3-flash-spark", + "--arm", "baseline-mtp3", "--context-limit", "1m"]) + assert soak.main() == 0 + plan = json.loads(capsys.readouterr().out) + assert plan["max_chat_requests"] == 86 + assert "api_key" not in plan["config"] diff --git a/performance/records/glm53-flash/mtp3-cache-history-observations.json.gz b/performance/records/glm53-flash/mtp3-cache-history-observations.json.gz new file mode 100644 index 00000000..17c3d990 Binary files /dev/null and b/performance/records/glm53-flash/mtp3-cache-history-observations.json.gz differ diff --git a/performance/records/glm53-flash/mtp3-cache-history-validation.md b/performance/records/glm53-flash/mtp3-cache-history-validation.md new file mode 100644 index 00000000..4a007783 --- /dev/null +++ b/performance/records/glm53-flash/mtp3-cache-history-validation.md @@ -0,0 +1,77 @@ +# Persistent-cache pressure with native MTP3 decoding + +Status: **research-only**. This record qualifies only the functional and +measurement conditions below; it is not a production throughput guarantee. + +## Conditions + +GLM-5.3-Flash-NVFP4-Spark revision +`df116c4fb16b1d37ae43d2cfd624de26ffbc832e` ran on four NVIDIA GB10 hosts with +TP4/DCP4, native MTP depth three, 24 GiB KV capacity per rank, and dual-rail +hardware-forwarded RoCE transport. The serving image was +`sha256:75050f7b4dd7287f1ecb3e7e34226d24aa6c5b3e012a8bd6c8399e411cfbd908`, +containing SparkCache `607ccef061d0f511f45a2a8a93f74514c955d3a3` and the source +transforms in `runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/`. + +The transport library SHA-256 was +`056243fad27d224b82e437925ffa2aed42037e6bd29f239f56076a832f6ca5cb`; +the SparkCache placement library SHA-256 was +`2657cdd2e54a097c9544e4c79ae62c0646db6db123ff24e4f0c384238c3a1e8d`. +The persisted cache used a 2 GiB maximum and 1.5 GiB low watermark per rank, +eight restore lanes, and a 16,384-token full-capture interval. Periodic full +capture is opt-in and increases write traffic. + +Four text-only conversations started around 100K tokens and grew by about +2K tokens per turn, with a 512-token response limit and rotation at 160K. +Admission lasted 3,600 seconds; a 100M prompt-token ceiling was not reached. +Thinking was enabled. Three approximately 5.3K-token, 300-output-token probes +ran before and after load. Startup and warmup completed before measurement. + +## Measurement + +The harness `performance/harnesses/validation/conversation_soak.py` records +client monotonic-clock TTFT and stream-delta offsets. Decode rate is estimated +as `(completion_tokens - 1) / (last_delta - first_delta)`; it is not exact +token-level inter-token latency. Reported cached tokens come from API usage. +They do not independently distinguish local hits from external restores. + +The numeric observations are in +[`mtp3-cache-history-observations.json.gz`](mtp3-cache-history-observations.json.gz). +The artifact identifies the raw receipt and harness by SHA-256 and preserves +all 551 response timings and stream offsets. Endpoint addresses, request IDs, +and prompt/generated text are omitted; numeric observations are unchanged. +The artifact verifier recomputes response counts and probe medians. + +## Result + +The run completed 551 responses without request errors: 545 conversation +requests and six probes. It admitted 70,074,391 prompt tokens, including +cached tokens, in 60m58.46s including probes and drain. All 525 continuations +reported at least half their prompt tokens cached. + +| Probe median | Before load | After load | +|---|---:|---:| +| Client TTFT | 2.9436 s | 2.9798 s | +| Estimated decode rate | 51.20 tokens/s | 47.80 tokens/s | + +The decode estimate decreased 6.65%; individual probe ranges overlap. Three +samples per side do not support a precise confidence interval or a claim of +zero performance drift. + +## Conclusion + +This four-rank composition completed the specified sustained cache workload +without request errors or a population of low-reported-cache continuations. +The probes did not reproduce the approximately 35% post-load slowdown +described in SparkCache issue #60 under these conditions. + +## Limitations + +The image combines SparkCache and runtime changes, so this is not an isolated +comparison of any individual patch. No matched original-image or restore-only +control was completed. The 2 GiB test policy is not the original issues' +40 GiB cache policy; multimodal traffic, C8/C16, and near-1M serving are not +qualified. Some resumed schedules still miss safely. One startup attempt +exited for an undetermined reason; the same image passed startup on retry. +This record does not qualify unattended startup availability or close +SparkCache issues #60 or #61. diff --git a/performance/records/glm53-flash/mtp3-integrated-image-source-equivalence.md b/performance/records/glm53-flash/mtp3-integrated-image-source-equivalence.md new file mode 100644 index 00000000..58b77d9f --- /dev/null +++ b/performance/records/glm53-flash/mtp3-integrated-image-source-equivalence.md @@ -0,0 +1,50 @@ +# Source equivalence for the native-MTP3 cache/checkpoint image + +Status: **research-only**. Image composition and file verification are implemented. + +## Conditions + +The published image has registry digest +`sha256:11a556a54041fd823d152a7f051ac4f7c617dc539030df26e93008392fee0746` +and config ID +`sha256:6921a6c163ea40b603e19a0332330efe3dbccbf4dce9f6cbbf6b756c9231835a`. +It derives from the published native-MTP3 compute image identified by +`runtime/glm53-spark-mtp3-mesh/performance/prepare.py`, with SparkCache +`48bbd2be4a7b972e56632a2d7b934bac5460f272` and the packaged runtime transforms. + +The comparison image, observed running on a four-GB10 TP4/DCP4 cluster, was +`sha256:84c4546d37e8504fc98a359435b262a5d8e7bfe9c892af7429e876db5bc98422`. +It includes explicit recurrent checkpoints and the GLM reasoning contract. +The cluster remained running during the build and read-only comparison. + +## Measurement + +SHA-256 hashes from the rebuilt image's installed-file receipt were compared +with files in the running container. The population contains every recorded +vLLM and B12X file, the transport bundle, and the warmup helper: 5,308 files. +SparkCache is excluded from this equality claim because the rebuild explicitly +selects its merged main revision. The runtime's checkpoint and source preimage +checks ran before installation; no ownership digest was accepted merely by +recomputing it from an unknown runtime. + +## Result + +All 5,308 compared files matched. The image verifier checked 5,472 files in the +rebuilt container without loading a model or initializing CUDA. An independent +registry pull and file-verification run on another Spark also passed. Anonymous +registry manifest access verified the config ID recorded in the public receipt. + +## Conclusion + +The published build preserves the compared deployed runtime components and +incorporates the selected SparkCache source. Repository-relative build inputs, +source inventories, strict patches, and the published-image contract reproduce +the composition without private workspace paths. + +## Limitations + +No serving soak was run on this exact rebuilt image. Source equivalence does not +transfer every measurement from the comparison image. Earlier cache-pressure +tests used a 2 GiB/rank namespace; 40 GiB multimodal stress behavior remains +unqualified. Native binary rebuilds can differ with toolchains and require +separate verification. This record does not establish unattended availability. diff --git a/performance/records/glm53-flash/spark-mtp3-compute-matrices-20260905.json b/performance/records/glm53-flash/spark-mtp3-compute-matrices-20260905.json new file mode 100644 index 00000000..ed512880 --- /dev/null +++ b/performance/records/glm53-flash/spark-mtp3-compute-matrices-20260905.json @@ -0,0 +1,1465 @@ +{ + "schema": "sparkring-glm53-compute-matrices/v1", + "status": "research-only", + "aggregation": "Arithmetic mean over three context rows per concurrency; each context has equal weight.", + "contexts": [ + 8192, + 32768, + 65536 + ], + "concurrencies": [ + 1, + 2, + 4, + 8, + 12, + 16 + ], + "normalized_definition": "Aggregate sequence steps/s, not batched engine iterations/s.", + "runs": [ + { + "configuration": "Public mesh MTP3", + "image": null, + "source_receipt_name": "glm-5.3-flash-spark-dcp4-dflash7-bf16-SIRCL-and-MESH-20260904-233405.json", + "source_receipt_sha256": "f0916f6b72cb8256225169b44c4f11e3ca764a5dd854977b8963686197b843fa", + "metadata": { + "version": "0.4.32", + "model": "glm-5.3-flash-spark", + "duration_per_test": 20.0, + "temperature": 1.0, + "max_tokens": 2048, + "dcp_size": 4 + }, + "cells": [ + { + "context_tokens": 8192, + "concurrency": 1, + "aggregate_tps": 48.186139604737214, + "server_steps_per_s": 17.813360020027464, + "server_spec_accept_length": 2.705056179775281, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 32768, + "concurrency": 1, + "aggregate_tps": 49.9, + "server_steps_per_s": 17.75, + "server_spec_accept_length": 2.8112676056338026, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 65536, + "concurrency": 1, + "aggregate_tps": 43.01558752984392, + "server_steps_per_s": 17.635891286916262, + "server_spec_accept_length": 2.4390934844192635, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 8192, + "concurrency": 2, + "aggregate_tps": 75.84311766170075, + "server_steps_per_s": 27.479390457137953, + "server_spec_accept_length": 2.76, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 8192, + "concurrency": 4, + "aggregate_tps": 112.22394711780265, + "server_steps_per_s": 43.066753467786825, + "server_spec_accept_length": 2.605813953488372, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 8192, + "concurrency": 8, + "aggregate_tps": 168.80835257585548, + "server_steps_per_s": 59.83335006554259, + "server_spec_accept_length": 2.821308724832215, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 8192, + "concurrency": 12, + "aggregate_tps": 193.35819391325268, + "server_steps_per_s": 69.54243096176405, + "server_spec_accept_length": 2.7804347826086957, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 8192, + "concurrency": 16, + "aggregate_tps": 231.3372852411072, + "server_steps_per_s": 81.9853310561318, + "server_spec_accept_length": 2.8216911764705883, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 32768, + "concurrency": 2, + "aggregate_tps": 76.8423682237785, + "server_steps_per_s": 27.17961528851463, + "server_spec_accept_length": 2.827205882352941, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 32768, + "concurrency": 4, + "aggregate_tps": 118.97960206497459, + "server_steps_per_s": 41.69799027719413, + "server_spec_accept_length": 2.8533653846153846, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 32768, + "concurrency": 8, + "aggregate_tps": 164.56330128210502, + "server_steps_per_s": 59.69551282053231, + "server_spec_accept_length": 2.7567114093959733, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 32768, + "concurrency": 12, + "aggregate_tps": 197.31954623009523, + "server_steps_per_s": 68.66780443723486, + "server_spec_accept_length": 2.8735380116959064, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 32768, + "concurrency": 16, + "aggregate_tps": 222.73572110305662, + "server_steps_per_s": 78.76626312349859, + "server_spec_accept_length": 2.8278061224489797, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 65536, + "concurrency": 2, + "aggregate_tps": 76.37566402715933, + "server_steps_per_s": 28.064548461423374, + "server_spec_accept_length": 2.7214285714285715, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 65536, + "concurrency": 4, + "aggregate_tps": 119.02971984175005, + "server_steps_per_s": 42.499874705601705, + "server_spec_accept_length": 2.8007075471698113, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 65536, + "concurrency": 8, + "aggregate_tps": 165.85, + "server_steps_per_s": 59.78371883448483, + "server_spec_accept_length": 2.774166666666667, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 65536, + "concurrency": 12, + "aggregate_tps": 192.30188679245282, + "server_steps_per_s": 69.43396226415094, + "server_spec_accept_length": 2.769565217391304, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 65536, + "concurrency": 16, + "aggregate_tps": 220.88050314465409, + "server_steps_per_s": 78.8930817610063, + "server_spec_accept_length": 2.799744897959184, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + } + ], + "prefill": { + "8192": { + "ttft_seconds": 3.031, + "tok_per_sec": 2703.0, + "samples": 1, + "method": "integrated_scout" + }, + "32768": { + "ttft_seconds": 11.797, + "tok_per_sec": 2778.0, + "samples": 1, + "method": "integrated_scout" + }, + "65536": { + "ttft_seconds": 23.516, + "tok_per_sec": 2787.0, + "samples": 1, + "method": "integrated_scout" + }, + "131072": { + "ttft_seconds": 47.39, + "tok_per_sec": 2766.0, + "samples": 1, + "method": "scout_only" + } + }, + "column_averages": { + "1": { + "aggregate_tps": 47.03390904486038, + "server_steps_per_s": 17.733083768981242, + "aggregate_tps_vs_public_percent": 0.0, + "server_steps_per_s_vs_public_percent": 0.0 + }, + "2": { + "aggregate_tps": 76.35371663754619, + "server_steps_per_s": 27.57451806902532, + "aggregate_tps_vs_public_percent": 0.0, + "server_steps_per_s_vs_public_percent": 0.0 + }, + "4": { + "aggregate_tps": 116.74442300817576, + "server_steps_per_s": 42.42153948352755, + "aggregate_tps_vs_public_percent": 0.0, + "server_steps_per_s_vs_public_percent": 0.0 + }, + "8": { + "aggregate_tps": 166.4072179526535, + "server_steps_per_s": 59.77086057351991, + "aggregate_tps_vs_public_percent": 0.0, + "server_steps_per_s_vs_public_percent": 0.0 + }, + "12": { + "aggregate_tps": 194.32654231193357, + "server_steps_per_s": 69.21473255438329, + "aggregate_tps_vs_public_percent": 0.0, + "server_steps_per_s_vs_public_percent": 0.0 + }, + "16": { + "aggregate_tps": 224.9845031629393, + "server_steps_per_s": 79.8815586468789, + "aggregate_tps_vs_public_percent": 0.0, + "server_steps_per_s_vs_public_percent": 0.0 + } + } + }, + { + "configuration": "NVFP4 MTP proposal head", + "image": "sha256:04d5a35b03e99f68c37a05514d221988a3eb70a5b8fdcfa859025ca1cbc25e74", + "source_receipt_name": "glm-5.3-flash-spark-dcp4-MTP3-SIRCL-and-MESH-r2420260905-124723.json", + "source_receipt_sha256": "29e86dbbd0522cd9c42151b828826fbc8b48a938ef86ba86d0cd0f428af30bf3", + "metadata": { + "version": "0.4.32", + "model": "glm-5.3-flash-spark", + "duration_per_test": 20.0, + "temperature": 1.0, + "max_tokens": 2048, + "dcp_size": 4 + }, + "cells": [ + { + "context_tokens": 8192, + "concurrency": 1, + "aggregate_tps": 50.0, + "server_steps_per_s": 18.9, + "server_spec_accept_length": 2.6455026455026456, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 32768, + "concurrency": 1, + "aggregate_tps": 42.05, + "server_steps_per_s": 18.4, + "server_spec_accept_length": 2.2853260869565215, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 65536, + "concurrency": 1, + "aggregate_tps": 49.877309830134514, + "server_steps_per_s": 18.829185237078892, + "server_spec_accept_length": 2.648936170212766, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 8192, + "concurrency": 2, + "aggregate_tps": 77.62031148263905, + "server_steps_per_s": 28.8447092993549, + "server_spec_accept_length": 2.6909722222222223, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 8192, + "concurrency": 4, + "aggregate_tps": 122.38147739784979, + "server_steps_per_s": 43.299589054767495, + "server_spec_accept_length": 2.826388888888889, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 8192, + "concurrency": 8, + "aggregate_tps": 163.18348118094238, + "server_steps_per_s": 60.14133213056844, + "server_spec_accept_length": 2.7133333333333334, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 8192, + "concurrency": 12, + "aggregate_tps": 200.8827807590449, + "server_steps_per_s": 70.42182876047416, + "server_spec_accept_length": 2.8525641025641026, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 8192, + "concurrency": 16, + "aggregate_tps": 225.35779076752496, + "server_steps_per_s": 82.18403168508102, + "server_spec_accept_length": 2.7421116504854366, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 32768, + "concurrency": 2, + "aggregate_tps": 75.15, + "server_steps_per_s": 29.5, + "server_spec_accept_length": 2.5474576271186438, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 32768, + "concurrency": 4, + "aggregate_tps": 122.28737533215585, + "server_steps_per_s": 43.10128802690738, + "server_spec_accept_length": 2.837209302325581, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 32768, + "concurrency": 8, + "aggregate_tps": 167.0177550399394, + "server_steps_per_s": 61.390309960626375, + "server_spec_accept_length": 2.7205882352941178, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 32768, + "concurrency": 12, + "aggregate_tps": 204.43374460863902, + "server_steps_per_s": 72.22389407174686, + "server_spec_accept_length": 2.8305555555555557, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 32768, + "concurrency": 16, + "aggregate_tps": 223.45, + "server_steps_per_s": 80.29690905047809, + "server_spec_accept_length": 2.7827970297029703, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 65536, + "concurrency": 2, + "aggregate_tps": 79.75, + "server_steps_per_s": 29.2, + "server_spec_accept_length": 2.731164383561644, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 65536, + "concurrency": 4, + "aggregate_tps": 120.53683208949174, + "server_steps_per_s": 43.06675346778683, + "server_spec_accept_length": 2.7988372093023255, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 65536, + "concurrency": 8, + "aggregate_tps": 175.2551405152696, + "server_steps_per_s": 61.53536775407056, + "server_spec_accept_length": 2.8480392156862746, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 65536, + "concurrency": 12, + "aggregate_tps": 194.6, + "server_steps_per_s": 70.43803680981594, + "server_spec_accept_length": 2.7627118644067794, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 65536, + "concurrency": 16, + "aggregate_tps": 225.81132075471697, + "server_steps_per_s": 80.50314465408805, + "server_spec_accept_length": 2.8049999999999997, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + } + ], + "prefill": { + "8192": { + "ttft_seconds": 3.062, + "tok_per_sec": 2675.0, + "samples": 1, + "method": "integrated_scout" + }, + "32768": { + "ttft_seconds": 11.875, + "tok_per_sec": 2759.0, + "samples": 1, + "method": "integrated_scout" + }, + "65536": { + "ttft_seconds": 23.609, + "tok_per_sec": 2776.0, + "samples": 1, + "method": "integrated_scout" + }, + "131072": { + "ttft_seconds": 47.625, + "tok_per_sec": 2752.0, + "samples": 1, + "method": "scout_only" + } + }, + "column_averages": { + "1": { + "aggregate_tps": 47.309103276711504, + "server_steps_per_s": 18.70972841235963, + "aggregate_tps_vs_public_percent": 0.5850975124960245, + "aggregate_tps_vs_preceding_percent": 0.5850975124960245, + "server_steps_per_s_vs_public_percent": 5.507472113151235, + "server_steps_per_s_vs_preceding_percent": 5.507472113151235 + }, + "2": { + "aggregate_tps": 77.50677049421301, + "server_steps_per_s": 29.181569766451634, + "aggregate_tps_vs_public_percent": 1.5101476489224774, + "aggregate_tps_vs_preceding_percent": 1.5101476489224774, + "server_steps_per_s_vs_public_percent": 5.828031856816129, + "server_steps_per_s_vs_preceding_percent": 5.828031856816129 + }, + "4": { + "aggregate_tps": 121.73522827316579, + "server_steps_per_s": 43.15587684982057, + "aggregate_tps_vs_public_percent": 4.2749838805066664, + "aggregate_tps_vs_preceding_percent": 4.2749838805066664, + "server_steps_per_s_vs_public_percent": 1.7310483665454068, + "server_steps_per_s_vs_preceding_percent": 1.7310483665454068 + }, + "8": { + "aggregate_tps": 168.48545891205046, + "server_steps_per_s": 61.02233661508846, + "aggregate_tps_vs_public_percent": 1.248888711058349, + "aggregate_tps_vs_preceding_percent": 1.248888711058349, + "server_steps_per_s_vs_public_percent": 2.093789564948967, + "server_steps_per_s_vs_preceding_percent": 2.093789564948967 + }, + "12": { + "aggregate_tps": 199.9721751225613, + "server_steps_per_s": 71.02791988067898, + "aggregate_tps_vs_public_percent": 2.905229899868922, + "aggregate_tps_vs_preceding_percent": 2.905229899868922, + "server_steps_per_s_vs_public_percent": 2.619655179438918, + "server_steps_per_s_vs_preceding_percent": 2.619655179438918 + }, + "16": { + "aggregate_tps": 224.87303717408065, + "server_steps_per_s": 80.99469512988239, + "aggregate_tps_vs_public_percent": -0.049543851817168694, + "aggregate_tps_vs_preceding_percent": -0.049543851817168694, + "server_steps_per_s_vs_public_percent": 1.3934836799118777, + "server_steps_per_s_vs_preceding_percent": 1.3934836799118777 + } + } + }, + { + "configuration": "Proposal head + loader/RNG fixes", + "image": "sha256:e3e83b5ef49c8787ae1a1b4945c22bb55f3432d43d74fdc1b1c3564d702bc2d5", + "source_receipt_name": "glm-5.3-flash-spark-dcp4-MTP3-SIRCL-and-MESH-r2420260905-145256.json", + "source_receipt_sha256": "bab411558cc190e3d19bc33d9ee45444f11975c81c32283e1c4dc1ce33ab4037", + "metadata": { + "version": "0.4.32", + "model": "glm-5.3-flash-spark", + "duration_per_test": 20.0, + "temperature": 1.0, + "max_tokens": 2048, + "dcp_size": 4 + }, + "cells": [ + { + "context_tokens": 8192, + "concurrency": 1, + "aggregate_tps": 47.43557668254504, + "server_steps_per_s": 18.764073054804207, + "server_spec_accept_length": 2.528, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 32768, + "concurrency": 1, + "aggregate_tps": 51.541232986247394, + "server_steps_per_s": 18.76501200955609, + "server_spec_accept_length": 2.7466666666666666, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 65536, + "concurrency": 1, + "aggregate_tps": 49.27637838639795, + "server_steps_per_s": 18.691040077599222, + "server_spec_accept_length": 2.6363636363636367, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 8192, + "concurrency": 2, + "aggregate_tps": 78.73502731426919, + "server_steps_per_s": 29.569488297529485, + "server_spec_accept_length": 2.6627118644067798, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 8192, + "concurrency": 4, + "aggregate_tps": 123.08001204685665, + "server_steps_per_s": 43.77070575238948, + "server_spec_accept_length": 2.811926605504587, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 8192, + "concurrency": 8, + "aggregate_tps": 178.71103538113528, + "server_steps_per_s": 63.34569509864133, + "server_spec_accept_length": 2.8212025316455698, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 8192, + "concurrency": 12, + "aggregate_tps": 200.3524672710372, + "server_steps_per_s": 72.50755287014164, + "server_spec_accept_length": 2.7631944444444443, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 8192, + "concurrency": 16, + "aggregate_tps": 228.58294698602262, + "server_steps_per_s": 82.24148357169068, + "server_spec_accept_length": 2.7794117647058822, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 32768, + "concurrency": 2, + "aggregate_tps": 80.0600450338313, + "server_steps_per_s": 30.022516887686738, + "server_spec_accept_length": 2.666666666666667, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 32768, + "concurrency": 4, + "aggregate_tps": 120.52721258879146, + "server_steps_per_s": 43.90097223608371, + "server_spec_accept_length": 2.7454337899543377, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 32768, + "concurrency": 8, + "aggregate_tps": 172.28783026373765, + "server_steps_per_s": 60.83101091774244, + "server_spec_accept_length": 2.832236842105263, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 32768, + "concurrency": 12, + "aggregate_tps": 201.13442425438606, + "server_steps_per_s": 71.67955024588552, + "server_spec_accept_length": 2.8060224089635852, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 32768, + "concurrency": 16, + "aggregate_tps": 233.04766200459304, + "server_steps_per_s": 81.79221169709588, + "server_spec_accept_length": 2.8492647058823533, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 65536, + "concurrency": 2, + "aggregate_tps": 79.22279533306506, + "server_steps_per_s": 28.544243577653024, + "server_spec_accept_length": 2.7754385964912283, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 65536, + "concurrency": 4, + "aggregate_tps": 112.23417563245603, + "server_steps_per_s": 42.43182386817776, + "server_spec_accept_length": 2.6450471698113205, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 65536, + "concurrency": 8, + "aggregate_tps": 164.1588925662828, + "server_steps_per_s": 60.186578392770954, + "server_spec_accept_length": 2.7275, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 65536, + "concurrency": 12, + "aggregate_tps": 208.2, + "server_steps_per_s": 73.8, + "server_spec_accept_length": 2.821138211382114, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 65536, + "concurrency": 16, + "aggregate_tps": 225.31798909844872, + "server_steps_per_s": 79.95154451880438, + "server_spec_accept_length": 2.8181818181818183, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + } + ], + "prefill": { + "8192": { + "ttft_seconds": 3.047, + "tok_per_sec": 2689.0, + "samples": 1, + "method": "integrated_scout" + }, + "32768": { + "ttft_seconds": 11.813, + "tok_per_sec": 2774.0, + "samples": 1, + "method": "integrated_scout" + }, + "65536": { + "ttft_seconds": 23.547, + "tok_per_sec": 2783.0, + "samples": 1, + "method": "integrated_scout" + }, + "131072": { + "ttft_seconds": 47.438, + "tok_per_sec": 2763.0, + "samples": 1, + "method": "scout_only" + } + }, + "column_averages": { + "1": { + "aggregate_tps": 49.41772935173013, + "server_steps_per_s": 18.740041713986507, + "aggregate_tps_vs_public_percent": 5.068301477123849, + "aggregate_tps_vs_preceding_percent": 4.457125434581255, + "server_steps_per_s_vs_public_percent": 5.678414189677716, + "server_steps_per_s_vs_preceding_percent": 0.16201892918366045 + }, + "2": { + "aggregate_tps": 79.33928922705519, + "server_steps_per_s": 29.378749587623084, + "aggregate_tps_vs_public_percent": 3.9101863288222516, + "aggregate_tps_vs_preceding_percent": 2.3643337493709726, + "server_steps_per_s_vs_public_percent": 6.543111702193172, + "server_steps_per_s_vs_preceding_percent": 0.6756998432556394 + }, + "4": { + "aggregate_tps": 118.61380008936804, + "server_steps_per_s": 43.367833952216984, + "aggregate_tps_vs_public_percent": 1.6012560026626543, + "aggregate_tps_vs_preceding_percent": -2.564112482537484, + "server_steps_per_s_vs_public_percent": 2.230693369949188, + "server_steps_per_s_vs_preceding_percent": 0.4911430791546767 + }, + "8": { + "aggregate_tps": 171.7192527370519, + "server_steps_per_s": 61.454428136384905, + "aggregate_tps_vs_public_percent": 3.1921901283812115, + "aggregate_tps_vs_preceding_percent": 1.9193311077898345, + "server_steps_per_s_vs_public_percent": 2.8167028995578214, + "server_steps_per_s_vs_preceding_percent": 0.7080874729887165 + }, + "12": { + "aggregate_tps": 203.22896384180774, + "server_steps_per_s": 72.66236770534239, + "aggregate_tps_vs_public_percent": 4.581166022901795, + "aggregate_tps_vs_preceding_percent": 1.6286209405135255, + "server_steps_per_s_vs_public_percent": 4.981071260009906, + "server_steps_per_s_vs_preceding_percent": 2.301134296779561 + }, + "16": { + "aggregate_tps": 228.98286602968813, + "server_steps_per_s": 81.32841326253032, + "aggregate_tps_vs_public_percent": 1.777172565460261, + "aggregate_tps_vs_preceding_percent": 1.8276218915591613, + "server_steps_per_s_vs_public_percent": 1.811249855611008, + "server_steps_per_s_vs_preceding_percent": 0.41202467903951323 + } + } + }, + { + "configuration": "Loader/RNG + MoE scale sharing", + "image": "sha256:aa221d22f83d90ac497c7e58b2382c5c4d428fb697f0191d6afe971faf67d439", + "source_receipt_name": "moe-sharing-matrix.json", + "source_receipt_sha256": "1d3b93c057fc3bbb7ab2557e78858dc283a06cd9ff3e62ddcedf13f5942f37b7", + "metadata": { + "version": "0.4.32", + "model": "glm-5.3-flash-spark", + "duration_per_test": 20.0, + "temperature": 1.0, + "max_tokens": 2048, + "dcp_size": 4 + }, + "cells": [ + { + "context_tokens": 8192, + "concurrency": 1, + "aggregate_tps": 49.877309830716065, + "server_steps_per_s": 18.979418098234326, + "server_spec_accept_length": 2.627968337730871, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 32768, + "concurrency": 1, + "aggregate_tps": 48.925835044218296, + "server_steps_per_s": 18.829185237078896, + "server_spec_accept_length": 2.598404255319149, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 65536, + "concurrency": 1, + "aggregate_tps": 49.076067905152435, + "server_steps_per_s": 18.879262857390273, + "server_spec_accept_length": 2.59946949602122, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 8192, + "concurrency": 2, + "aggregate_tps": 78.15, + "server_steps_per_s": 29.3, + "server_spec_accept_length": 2.667235494880546, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 8192, + "concurrency": 4, + "aggregate_tps": 119.69575660565239, + "server_steps_per_s": 43.43474779837219, + "server_spec_accept_length": 2.7557603686635943, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 8192, + "concurrency": 8, + "aggregate_tps": 173.8059326283227, + "server_steps_per_s": 61.94067370497355, + "server_spec_accept_length": 2.8060064935064934, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 8192, + "concurrency": 12, + "aggregate_tps": 199.55687597390622, + "server_steps_per_s": 72.51120398749053, + "server_spec_accept_length": 2.752083333333333, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 8192, + "concurrency": 16, + "aggregate_tps": 226.29462041078403, + "server_steps_per_s": 82.0330565327409, + "server_spec_accept_length": 2.758578431372549, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 32768, + "concurrency": 2, + "aggregate_tps": 81.17582252569352, + "server_steps_per_s": 29.145175021563006, + "server_spec_accept_length": 2.7852233676975944, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 32768, + "concurrency": 4, + "aggregate_tps": 114.37728479252559, + "server_steps_per_s": 42.66613250579326, + "server_spec_accept_length": 2.68075117370892, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 32768, + "concurrency": 8, + "aggregate_tps": 170.02664521590413, + "server_steps_per_s": 62.33975164628063, + "server_spec_accept_length": 2.72741935483871, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 32768, + "concurrency": 12, + "aggregate_tps": 200.82131410145885, + "server_steps_per_s": 70.91346153807126, + "server_spec_accept_length": 2.8319209039548023, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 32768, + "concurrency": 16, + "aggregate_tps": 229.90475230432568, + "server_steps_per_s": 82.2456281807671, + "server_spec_accept_length": 2.795343137254902, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 65536, + "concurrency": 2, + "aggregate_tps": 80.3786057687884, + "server_steps_per_s": 29.146634615224208, + "server_spec_accept_length": 2.7577319587628866, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 65536, + "concurrency": 4, + "aggregate_tps": 118.62877762754627, + "server_steps_per_s": 43.702701348213076, + "server_spec_accept_length": 2.7144495412844036, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 65536, + "concurrency": 8, + "aggregate_tps": 170.4363490778019, + "server_steps_per_s": 61.64931945503581, + "server_spec_accept_length": 2.7646103896103895, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 65536, + "concurrency": 12, + "aggregate_tps": 201.6112890318451, + "server_steps_per_s": 70.85668534849657, + "server_spec_accept_length": 2.8453389830508478, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 65536, + "concurrency": 16, + "aggregate_tps": 229.17713252064863, + "server_steps_per_s": 82.78910881061573, + "server_spec_accept_length": 2.7682038834951457, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + } + ], + "prefill": { + "8192": { + "ttft_seconds": 3.063, + "tok_per_sec": 2675.0, + "samples": 1, + "method": "integrated_scout" + }, + "32768": { + "ttft_seconds": 11.86, + "tok_per_sec": 2763.0, + "samples": 1, + "method": "integrated_scout" + }, + "65536": { + "ttft_seconds": 23.609, + "tok_per_sec": 2776.0, + "samples": 1, + "method": "integrated_scout" + }, + "131072": { + "ttft_seconds": 47.532, + "tok_per_sec": 2758.0, + "samples": 1, + "method": "scout_only" + } + }, + "column_averages": { + "1": { + "aggregate_tps": 49.2930709266956, + "server_steps_per_s": 18.895955397567832, + "aggregate_tps_vs_public_percent": 4.803262003335629, + "aggregate_tps_vs_preceding_percent": -0.2522544573978269, + "server_steps_per_s_vs_public_percent": 6.5576390645641025, + "server_steps_per_s_vs_preceding_percent": 0.8319815182959767 + }, + "2": { + "aggregate_tps": 79.90147609816064, + "server_steps_per_s": 29.19726987892907, + "aggregate_tps_vs_public_percent": 4.646479067228393, + "aggregate_tps_vs_preceding_percent": 0.7085857165881171, + "server_steps_per_s_vs_public_percent": 5.884968889906372, + "server_steps_per_s_vs_preceding_percent": -0.6177244138752203 + }, + "4": { + "aggregate_tps": 117.56727300857474, + "server_steps_per_s": 43.26786055079284, + "aggregate_tps_vs_public_percent": 0.7048302430184172, + "aggregate_tps_vs_preceding_percent": -0.882297911376928, + "server_steps_per_s_vs_public_percent": 1.9950267660463261, + "server_steps_per_s_vs_preceding_percent": -0.23052431332930556 + }, + "8": { + "aggregate_tps": 171.42297564067624, + "server_steps_per_s": 61.976581602096665, + "aggregate_tps_vs_public_percent": 3.0141467117428933, + "aggregate_tps_vs_preceding_percent": -0.17253574753748335, + "server_steps_per_s_vs_public_percent": 3.6902949153018394, + "server_steps_per_s_vs_preceding_percent": 0.8496596283557478 + }, + "12": { + "aggregate_tps": 200.6631597024034, + "server_steps_per_s": 71.42711695801945, + "aggregate_tps_vs_public_percent": 3.2608090048235727, + "aggregate_tps_vs_preceding_percent": -1.262518929832046, + "server_steps_per_s_vs_public_percent": 3.196406779290606, + "server_steps_per_s_vs_preceding_percent": -1.6999869207841911 + }, + "16": { + "aggregate_tps": 228.4588350785861, + "server_steps_per_s": 82.3559311747079, + "aggregate_tps_vs_public_percent": 1.5442538782907356, + "aggregate_tps_vs_preceding_percent": -0.22885159933061505, + "server_steps_per_s_vs_public_percent": 3.097551637377416, + "server_steps_per_s_vs_preceding_percent": 1.2634181228406982 + } + } + }, + { + "configuration": "MoE scale sharing + top-k selector", + "image": "sha256:3b4768e5ba31cadcc882dffa06d7b667af44abdf157d5c11b7ac7fe962e80c43", + "source_receipt_name": "topk512-matrix.json", + "source_receipt_sha256": "fe5c17bd0121b92522fa862ed2ef33b59f90efd8675864c1438dcf31f571a90b", + "metadata": { + "version": "0.4.32", + "model": "glm-5.3-flash-spark", + "duration_per_test": 20.0, + "temperature": 1.0, + "max_tokens": 2048, + "dcp_size": 4 + }, + "cells": [ + { + "context_tokens": 8192, + "concurrency": 1, + "aggregate_tps": 52.6, + "server_steps_per_s": 18.85, + "server_spec_accept_length": 2.790450928381963, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 32768, + "concurrency": 1, + "aggregate_tps": 46.2, + "server_steps_per_s": 17.958466453674124, + "server_spec_accept_length": 2.5769230769230766, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 65536, + "concurrency": 1, + "aggregate_tps": 50.05, + "server_steps_per_s": 18.75, + "server_spec_accept_length": 2.6693333333333333, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 8192, + "concurrency": 2, + "aggregate_tps": 82.47784065284291, + "server_steps_per_s": 29.145175021223178, + "server_spec_accept_length": 2.8298969072164946, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 8192, + "concurrency": 4, + "aggregate_tps": 121.63853973634211, + "server_steps_per_s": 43.4852769732887, + "server_spec_accept_length": 2.7972350230414746, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 8192, + "concurrency": 8, + "aggregate_tps": 170.11367619974143, + "server_steps_per_s": 60.49376533685241, + "server_spec_accept_length": 2.812086092715232, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 8192, + "concurrency": 12, + "aggregate_tps": 198.77522337094305, + "server_steps_per_s": 71.07720108415539, + "server_spec_accept_length": 2.7966101694915255, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 8192, + "concurrency": 16, + "aggregate_tps": 233.02236934401157, + "server_steps_per_s": 82.65623432607212, + "server_spec_accept_length": 2.8191747572815533, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 32768, + "concurrency": 2, + "aggregate_tps": 81.02558966397989, + "server_steps_per_s": 29.34548550252918, + "server_spec_accept_length": 2.7610921501706485, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 32768, + "concurrency": 4, + "aggregate_tps": 121.54115586690078, + "server_steps_per_s": 43.832874655992214, + "server_spec_accept_length": 2.7728310502283104, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 32768, + "concurrency": 8, + "aggregate_tps": 170.027083960276, + "server_steps_per_s": 63.396529240645684, + "server_spec_accept_length": 2.681962025316456, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 32768, + "concurrency": 12, + "aggregate_tps": 199.21874999999926, + "server_steps_per_s": 70.91346153846128, + "server_spec_accept_length": 2.8093220338983054, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 32768, + "concurrency": 16, + "aggregate_tps": 241.1358173076914, + "server_steps_per_s": 84.13461538461507, + "server_spec_accept_length": 2.866071428571429, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 65536, + "concurrency": 2, + "aggregate_tps": 78.37147578747127, + "server_steps_per_s": 29.145175021283244, + "server_spec_accept_length": 2.689003436426117, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 65536, + "concurrency": 4, + "aggregate_tps": 119.18473634133011, + "server_steps_per_s": 44.06830587410525, + "server_spec_accept_length": 2.7045454545454546, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 65536, + "concurrency": 8, + "aggregate_tps": 181.63152886974765, + "server_steps_per_s": 64.49997496118968, + "server_spec_accept_length": 2.815993788819876, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 65536, + "concurrency": 12, + "aggregate_tps": 188.9440184433421, + "server_steps_per_s": 69.16253195008278, + "server_spec_accept_length": 2.7318840579710146, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + }, + { + "context_tokens": 65536, + "concurrency": 16, + "aggregate_tps": 230.29816283505642, + "server_steps_per_s": 82.72261821102288, + "server_spec_accept_length": 2.783980582524272, + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false + } + ], + "prefill": { + "8192": { + "ttft_seconds": 3.094, + "tok_per_sec": 2648.0, + "samples": 1, + "method": "integrated_scout" + }, + "32768": { + "ttft_seconds": 11.859, + "tok_per_sec": 2763.0, + "samples": 1, + "method": "integrated_scout" + }, + "65536": { + "ttft_seconds": 23.532, + "tok_per_sec": 2785.0, + "samples": 1, + "method": "integrated_scout" + }, + "131072": { + "ttft_seconds": 47.422, + "tok_per_sec": 2764.0, + "samples": 1, + "method": "scout_only" + } + }, + "column_averages": { + "1": { + "aggregate_tps": 49.61666666666667, + "server_steps_per_s": 18.519488817891375, + "aggregate_tps_vs_public_percent": 5.491267203291317, + "aggregate_tps_vs_preceding_percent": 0.656473078036246, + "server_steps_per_s_vs_public_percent": 4.434677347465721, + "server_steps_per_s_vs_preceding_percent": -1.9923130201975003 + }, + "2": { + "aggregate_tps": 80.62496870143136, + "server_steps_per_s": 29.211945181678534, + "aggregate_tps_vs_public_percent": 5.594032945587912, + "aggregate_tps_vs_preceding_percent": 0.9054808979772622, + "server_steps_per_s_vs_public_percent": 5.938189412973083, + "server_steps_per_s_vs_preceding_percent": 0.05026258554419183 + }, + "4": { + "aggregate_tps": 120.78814398152433, + "server_steps_per_s": 43.795485834462056, + "aggregate_tps_vs_public_percent": 3.46373802632558, + "aggregate_tps_vs_preceding_percent": 2.7395982661898355, + "server_steps_per_s_vs_public_percent": 3.238794177820936, + "server_steps_per_s_vs_preceding_percent": 1.2194392719044433 + }, + "8": { + "aggregate_tps": 173.92409634325503, + "server_steps_per_s": 62.796756512895925, + "aggregate_tps_vs_public_percent": 4.517158860705339, + "aggregate_tps_vs_preceding_percent": 1.4590347024552086, + "server_steps_per_s_vs_public_percent": 5.062493513296618, + "server_steps_per_s_vs_preceding_percent": 1.3233626147130195 + }, + "12": { + "aggregate_tps": 195.64599727142814, + "server_steps_per_s": 70.38439819089982, + "aggregate_tps_vs_public_percent": 0.6789885436115917, + "aggregate_tps_vs_preceding_percent": -2.500290755122181, + "server_steps_per_s_vs_public_percent": 1.6899084824137711, + "server_steps_per_s_vs_preceding_percent": -1.4598360000061095 + }, + "16": { + "aggregate_tps": 234.81878316225314, + "server_steps_per_s": 83.17115597390335, + "aggregate_tps_vs_public_percent": 4.371092169042257, + "aggregate_tps_vs_preceding_percent": 2.7838486007684127, + "server_steps_per_s_vs_public_percent": 4.118093565958447, + "server_steps_per_s_vs_preceding_percent": 0.989879887905154 + } + } + } + ], + "conditions": "Four DGX Sparks, GLM-5.3-Flash-NVFP4-Spark, native MTP3, TP4/DCP4, SIRCL/mesh and SparkCache.", + "interpretation": "Descriptive matrices from separate runs, not a matched isolated transport comparison or repeated-run confidence interval.", + "topk_run_continuity": "Eight cells preceded benchmark-client power loss; ten resumed with the exact same 0.4.32 harness. Four server containers stayed running.", + "topk_host_state": "Top-k startup followed a host reboot and the memory startup gate; reboot effects are not isolated.", + "transport": "Measured compute images retain transport bundle4204fabc. These runs do not qualify the combined stream-safety bundle69313e19 image.", + "image_identity_source": "Deployment receipts associate the four compute images with runs; the historical reference is identified by its hashed benchmark receipt only." +} diff --git a/performance/records/glm53-flash/spark-mtp3-compute-matrices-20260905.md b/performance/records/glm53-flash/spark-mtp3-compute-matrices-20260905.md new file mode 100644 index 00000000..428831a5 --- /dev/null +++ b/performance/records/glm53-flash/spark-mtp3-compute-matrices-20260905.md @@ -0,0 +1,62 @@ +# GLM native-MTP3 compute comparison + +Status: research-only measurements. All five matrices contain 18 completed +decode cells with no request errors, underfilled cells, capacity-limited cells, +or warm-up timeouts. + +## Conditions + +Four DGX Sparks serve GLM-5.3-Flash-NVFP4-Spark with TP4/DCP4, native MTP3, +SIRCL/hardware-forwarded mesh, and SparkCache. The harness version is 0.4.32; +temperature is 1, the maximum output is 2,048 tokens, and each decode cell has +a 20-second measurement window. Contexts are 8K, 32K, and 64K tokens. Each +column below is their equally weighted arithmetic mean, not per-request +throughput and not a mean of repeated benchmark runs. + +The [sanitized matrices](spark-mtp3-compute-matrices-20260905.json) retain all +cell values, acceptance lengths, prefill scouts, configuration labels, full +compute-image identities, and hashes of the original receipts. The historical +mesh reference is identified by its benchmark receipt; its filename contains +a DFlash label but the associated serving configuration used native MTP3. + +## Aggregate decode tokens per second + +The compute additions are cumulative within the four compute-image rows. + +| Configuration | C1 | C2 | C4 | C8 | C12 | C16 | +|---|---:|---:|---:|---:|---:|---:| +| Mesh MTP3 reference | 47.0 | 76.4 | 116.7 | 166.4 | 194.3 | 225.0 | +| NVFP4 MTP proposal head | 47.3 | 77.5 | 121.7 | 168.5 | 200.0 | 224.9 | +| Proposal head + loader/RNG fixes | 49.4 | 79.3 | 118.6 | 171.7 | 203.2 | 229.0 | +| Loader/RNG + MoE scale sharing | 49.3 | 79.9 | 117.6 | 171.4 | 200.7 | 228.5 | +| MoE scale sharing + top-k selector | 49.6 | 80.6 | 120.8 | 173.9 | 195.6 | 234.8 | + +## MTP-normalized aggregate sequence steps per second + +Normalization divides output throughput by observed accepted length. These +are aggregate sequence-step rates, not batched engine iterations per second. + +| Configuration | C1 | C2 | C4 | C8 | C12 | C16 | +|---|---:|---:|---:|---:|---:|---:| +| Mesh MTP3 reference | 17.7 | 27.6 | 42.4 | 59.8 | 69.2 | 79.9 | +| NVFP4 MTP proposal head | 18.7 | 29.2 | 43.2 | 61.0 | 71.0 | 81.0 | +| Proposal head + loader/RNG fixes | 18.7 | 29.4 | 43.4 | 61.5 | 72.7 | 81.3 | +| Loader/RNG + MoE scale sharing | 18.9 | 29.2 | 43.3 | 62.0 | 71.4 | 82.4 | +| MoE scale sharing + top-k selector | 18.5 | 29.2 | 43.8 | 62.8 | 70.4 | 83.2 | + +The selector image's C4/C8/C16 normalized means exceed the scale-sharing +image's means, while C1/C12 are lower. This single matrix does not establish +repeatability of those differences. The independent-RNG correction changes +sampling behavior; raw token rates alone must not be interpreted as compute +speed when acceptance differs. + +## Scope + +The top-k matrix resumed ten cells after the benchmark client lost power; +the serving containers stayed up and the resumed harness matched the recorded +source hash. Its initial deployment followed a host reboot to restore memory +contiguity. The comparison does not isolate reboot effects. + +These compute images use transport bundle `4204fabc`. They are not performance +qualification of the combined image containing stream-safety bundle +`69313e19`; that image requires its own deployment and validation receipts. diff --git a/performance/records/glm53-flash/spark-mtp3-compute-stream-safety-20260906.json b/performance/records/glm53-flash/spark-mtp3-compute-stream-safety-20260906.json new file mode 100644 index 00000000..a2fab68d --- /dev/null +++ b/performance/records/glm53-flash/spark-mtp3-compute-stream-safety-20260906.json @@ -0,0 +1,4103 @@ +{ + "schema": "sparkring-glm53-compute-stream-validation/v1", + "status": "qualified", + "scope": "Exact-image native all-reduce, selected GPU stream invariants, bounded serving matrix, idle rank-loss shutdown, restart and one persistent-prefix recall.", + "image": "sha256:2e41b1e934a85ff7c21b780532db2f0a0e978df081e52f4ae2bf11f8992fb24f", + "registry_manifest": "sha256:67dc0ae453baaae6831ccec1d259b4ef8b236a8b0dc9f747d901b95c66ec1987", + "compute_source_lock_sha256": "139f36701e0e47f45bf99fba2cc2fa59b417f2ee801dad3a064455d5b464a459", + "transport_manifest_sha256": "69313e19e881ec93e9ed3bd150d2f24fc6b444488ac729a69f45d038e2243500", + "startup_memory_gate": "Host source from draft PR222; no watchdog or memory thresholds relaxed.", + "host_source_manifest": { + "composition": "PR219 compute/profile contracts and PR222 memory-gated host startup", + "files": { + "runtime/glm53-spark-mtp3-mesh/managed_memory.py": "a338586470ebcae67df13bb30d4f6707d2b10da0f67eecaf6c3b5c1fcc22a367", + "runtime/glm53-spark-mtp3-mesh/managed_service.py": "e75dcd5394951a5801f6c5f742d1aef69f91d9fbfbd324be012667b893c4ed08", + "runtime/glm53-spark-mtp3-mesh/managed_network.py": "f2c9538e356537f6e9c63fcbd6ee086d4641491b034b07b8b86860be7c678877", + "runtime/glm53-spark-mtp3-mesh/managed_units.py": "4a8f47de02e2acba28a4a6a0847f95ed9877ee8cf9ad15e040bd82ef7540e786", + "runtime/glm53-spark-mtp3-mesh/managed_cluster.py": "c772836f208f24e8bedb5c6b57c8818146dade2a46b88bc18f974fd1f16b023f", + "runtime/glm53-spark-mtp3-mesh/managed_install.py": "cbb9f89f2ff3d2b2d0957fc50522d48c5813c79815c2b9acc16a82ae3b4b57b0", + "runtime/glm53-spark-mtp3-mesh/profile.py": "9e5e2b8e8bb1a0f74982f816723ebbfed361b2d94c494a50983057167bf367ff", + "runtime/glm53-spark-mtp3-mesh/inspect_fabric.py": "def378fea0534c4302c14b0b22be5d40f7bfe456880748e3ee6c7b7237203cc6", + "runtime/glm53-spark-mtp3-mesh/pins.json": "05fddc68d2e74e00ed29126c09258766aff8a852b67681836023c097234d728b", + "runtime/glm53-flash-jj-r8-gb10/pins.json": "dc49370d911ecc32f1c7ea7656f41837334a8e5ed800db263174781bee2b2313", + "runtime/glm53-flash-jj-r8-gb10/warmup_dflash.py": "f41c38eef41d15d63dcfc49cd6643357ca1a3ae18200ddbe4f8692d0b767ee79", + "runtime/glm53-flash-jj-r8-gb10/launch-rank.sh": "8a4511eb80e3d12d9daccfc9c1d2e6305c67926e24017f4e90528cc3be6c3ba7", + "runtime/glm53-flash-jj-r8-gb10/runtime.env.example": "707e0568d0410e8911efbb15f2d4e97e779ef1ac81712de70265f4e7769cc911", + "runtime/glm53-flash-jj-r8-gb10/sircl-fused.env.example": "f47619ad2aaece7425e1958ddf19fa9b6508fbe0964fb66fe88aee0f864c3038", + "spark_transport/experiments/cx7_hairpin_diagonal/__init__.py": "67fd1b48093286b47f2cf32c7fcd69d3e6d29bf6484abc7216a82da0e10e7b41", + "spark_transport/experiments/cx7_hairpin_diagonal/fabric.py": "cbf2f596e744ab50caac66db41e91c1cd3a5625f3cc6ba77ffbd1d667147d544", + "spark_transport/experiments/glm53_rocenante_overlay/build_bundle.py": "ca5f5502b55fbe095bdcdf1283b1f5df85b0f9727873970df727475d5f305a3b", + "runtime/glm53-spark-mtp3-mesh/compute/source-lock.json": "139f36701e0e47f45bf99fba2cc2fa59b417f2ee801dad3a064455d5b464a459" + } + }, + "native": [ + { + "schema": "b12x.rocenante-virtual-diagonal-evidence/v1", + "status": "research-only", + "payload_bytes": 32768, + "warmups": 2, + "samples": 3, + "graph_operations_per_replay": 3, + "ranks": [ + { + "rank": 0, + "eager_samples_us": [ + 444.41598653793335, + 108.99200290441513, + 66.39999896287918 + ], + "graph_samples_us": [ + 240.48000574111938, + 21.5786670645078, + 21.290667355060577 + ], + "eager_median_us": 108.99200290441513, + "graph_median_us": 21.5786670645078, + "correctness_cases": [ + { + "name": "rank-specific-index-pattern-a", + "mode": "eager", + "input_formula": "((rank + 1) * ((index % 17) - 8)) + rank", + "expected_formula": "10 * ((index % 17) - 8) + 6", + "input_sha256": "48d35fdc0ee602ef417b66e84d5b8161066b7c09efef2aee301b94678645fa55", + "expected_sha256": "202a06d669dc01f7432eda3e514cf5740ea9652f609c8e0dfec40716b2651522", + "output_sha256": "202a06d669dc01f7432eda3e514cf5740ea9652f609c8e0dfec40716b2651522", + "passed": true + }, + { + "name": "constant-timing-reference", + "mode": "eager", + "input_formula": "rank + 1", + "expected_formula": "10", + "input_sha256": "98a9a14f351766d2743a1b9d2967cbe2f5993ef7d56889f9d78bb2974269267d", + "expected_sha256": "d6de16041970b40bf4c7653ad60f95ae46976f16550b1d54988463980a19fbc2", + "output_sha256": "d6de16041970b40bf4c7653ad60f95ae46976f16550b1d54988463980a19fbc2", + "passed": true + }, + { + "name": "graph-input-mutation-a", + "mode": "graph-replay", + "input_formula": "((rank + 1) * ((index % 17) - 8)) + rank", + "expected_formula": "10 * ((index % 17) - 8) + 6", + "input_sha256": "48d35fdc0ee602ef417b66e84d5b8161066b7c09efef2aee301b94678645fa55", + "expected_sha256": "202a06d669dc01f7432eda3e514cf5740ea9652f609c8e0dfec40716b2651522", + "output_sha256": "202a06d669dc01f7432eda3e514cf5740ea9652f609c8e0dfec40716b2651522", + "passed": true + }, + { + "name": "graph-input-mutation-b", + "mode": "graph-replay", + "input_formula": "((4 - rank) * (((index * 5 + 3) % 19) - 9)) - rank", + "expected_formula": "10 * (((index * 5 + 3) % 19) - 9) - 6", + "input_sha256": "f71e11bb4c4c1043db1c26bcb9aa8ff4d2523663e9ff05385e17ef534205e60b", + "expected_sha256": "b282b975648deeef7ecb7491a997047c1f9561901f148450b08ee5b80712328a", + "output_sha256": "b282b975648deeef7ecb7491a997047c1f9561901f148450b08ee5b80712328a", + "passed": true + }, + { + "name": "constant-timing-output", + "mode": "graph-replay", + "input_formula": "rank + 1", + "expected_formula": "10", + "input_sha256": "98a9a14f351766d2743a1b9d2967cbe2f5993ef7d56889f9d78bb2974269267d", + "expected_sha256": "d6de16041970b40bf4c7653ad60f95ae46976f16550b1d54988463980a19fbc2", + "output_sha256": "d6de16041970b40bf4c7653ad60f95ae46976f16550b1d54988463980a19fbc2", + "passed": true + } + ], + "stats": { + "world_size": 4, + "rank": 0, + "hcas": [ + "rocep1s0f0", + "rocep1s0f1", + "roceP2p1s0f0", + "roceP2p1s0f1" + ], + "max_size": 2097152, + "max_gather_bytes": 2097152, + "slot_bytes": 2097152, + "epoch": 33, + "error_seq": 0, + "error_peer": 0, + "ctrl_seq": 33, + "spin_limit": 20000000, + "opposite_paths": 2, + "ops_posted": 33, + "writes_completed": 198, + "last_seq": 33, + "two_wave_activations": 0, + "two_wave_threshold_bytes": 196608, + "wave_mode": "two", + "peer_hca": { + "1": [ + 0, + 2 + ], + "2": [ + 0, + 3 + ], + "3": [ + 1, + 3 + ] + } + }, + "path_counters": [ + { + "peer_rank": 1, + "path_index": 0, + "device": "rocep1s0f0", + "payload_writes": 33, + "payload_bytes": 540672, + "physical_hop_payload_bytes": 540672, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 464, + "remote_qp_number": 847, + "local_hca_index": 0, + "remote_hca_index": 1, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 1, + "path_index": 1, + "device": "roceP2p1s0f0", + "payload_writes": 33, + "payload_bytes": 540672, + "physical_hop_payload_bytes": 540672, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 457, + "remote_qp_number": 713, + "local_hca_index": 2, + "remote_hca_index": 3, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 2, + "path_index": 0, + "device": "rocep1s0f0", + "payload_writes": 33, + "payload_bytes": 540672, + "physical_hop_payload_bytes": 1081344, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 465, + "remote_qp_number": 847, + "local_hca_index": 0, + "remote_hca_index": 1, + "physical_hops": 2, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 2, + "path_index": 1, + "device": "roceP2p1s0f1", + "payload_writes": 33, + "payload_bytes": 540672, + "physical_hop_payload_bytes": 1081344, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 714, + "remote_qp_number": 457, + "local_hca_index": 3, + "remote_hca_index": 2, + "physical_hops": 2, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 3, + "path_index": 0, + "device": "rocep1s0f1", + "payload_writes": 33, + "payload_bytes": 540672, + "physical_hop_payload_bytes": 540672, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 849, + "remote_qp_number": 464, + "local_hca_index": 1, + "remote_hca_index": 0, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 3, + "path_index": 1, + "device": "roceP2p1s0f1", + "payload_writes": 33, + "payload_bytes": 540672, + "physical_hop_payload_bytes": 540672, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 715, + "remote_qp_number": 457, + "local_hca_index": 3, + "remote_hca_index": 2, + "physical_hops": 1, + "retries": null, + "retry_events": null + } + ], + "expected_counters": { + "operations_per_rank": 33, + "paths_per_rank": 6, + "flags_per_path": 33, + "payload_bytes_path_0": 540672, + "payload_bytes_path_1": 540672 + } + }, + { + "rank": 1, + "eager_samples_us": [ + 255.71200251579285, + 119.87199634313583, + 55.80800026655197 + ], + "graph_samples_us": [ + 191.4880077044169, + 20.67199970285098, + 22.133332987626392 + ], + "eager_median_us": 119.87199634313583, + "graph_median_us": 22.133332987626392, + "correctness_cases": [ + { + "name": "rank-specific-index-pattern-a", + "mode": "eager", + "input_formula": "((rank + 1) * ((index % 17) - 8)) + rank", + "expected_formula": "10 * ((index % 17) - 8) + 6", + "input_sha256": "080d97a7d56f4c6201ed6913d36927d48fdee761dbaba585f2eeee5c93f0d0a8", + "expected_sha256": "202a06d669dc01f7432eda3e514cf5740ea9652f609c8e0dfec40716b2651522", + "output_sha256": "202a06d669dc01f7432eda3e514cf5740ea9652f609c8e0dfec40716b2651522", + "passed": true + }, + { + "name": "constant-timing-reference", + "mode": "eager", + "input_formula": "rank + 1", + "expected_formula": "10", + "input_sha256": "3ee6d0365f9d1e5297cff3c44da13ac69e147d518f411bd33cc9f5e40166c94d", + "expected_sha256": "d6de16041970b40bf4c7653ad60f95ae46976f16550b1d54988463980a19fbc2", + "output_sha256": "d6de16041970b40bf4c7653ad60f95ae46976f16550b1d54988463980a19fbc2", + "passed": true + }, + { + "name": "graph-input-mutation-a", + "mode": "graph-replay", + "input_formula": "((rank + 1) * ((index % 17) - 8)) + rank", + "expected_formula": "10 * ((index % 17) - 8) + 6", + "input_sha256": "080d97a7d56f4c6201ed6913d36927d48fdee761dbaba585f2eeee5c93f0d0a8", + "expected_sha256": "202a06d669dc01f7432eda3e514cf5740ea9652f609c8e0dfec40716b2651522", + "output_sha256": "202a06d669dc01f7432eda3e514cf5740ea9652f609c8e0dfec40716b2651522", + "passed": true + }, + { + "name": "graph-input-mutation-b", + "mode": "graph-replay", + "input_formula": "((4 - rank) * (((index * 5 + 3) % 19) - 9)) - rank", + "expected_formula": "10 * (((index * 5 + 3) % 19) - 9) - 6", + "input_sha256": "f64ed4bf2210447de8c95ffa74f5e5839d722b49c3eb63b552dd3c948528dec2", + "expected_sha256": "b282b975648deeef7ecb7491a997047c1f9561901f148450b08ee5b80712328a", + "output_sha256": "b282b975648deeef7ecb7491a997047c1f9561901f148450b08ee5b80712328a", + "passed": true + }, + { + "name": "constant-timing-output", + "mode": "graph-replay", + "input_formula": "rank + 1", + "expected_formula": "10", + "input_sha256": "3ee6d0365f9d1e5297cff3c44da13ac69e147d518f411bd33cc9f5e40166c94d", + "expected_sha256": "d6de16041970b40bf4c7653ad60f95ae46976f16550b1d54988463980a19fbc2", + "output_sha256": "d6de16041970b40bf4c7653ad60f95ae46976f16550b1d54988463980a19fbc2", + "passed": true + } + ], + "stats": { + "world_size": 4, + "rank": 1, + "hcas": [ + "rocep1s0f0", + "rocep1s0f1", + "roceP2p1s0f0", + "roceP2p1s0f1" + ], + "max_size": 2097152, + "max_gather_bytes": 2097152, + "slot_bytes": 2097152, + "epoch": 33, + "error_seq": 0, + "error_peer": 0, + "ctrl_seq": 33, + "spin_limit": 20000000, + "opposite_paths": 2, + "ops_posted": 33, + "writes_completed": 198, + "last_seq": 33, + "two_wave_activations": 0, + "two_wave_threshold_bytes": 196608, + "wave_mode": "two", + "peer_hca": { + "0": [ + 1, + 3 + ], + "2": [ + 0, + 2 + ], + "3": [ + 0, + 3 + ] + } + }, + "path_counters": [ + { + "peer_rank": 0, + "path_index": 0, + "device": "rocep1s0f1", + "payload_writes": 33, + "payload_bytes": 540672, + "physical_hop_payload_bytes": 540672, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 847, + "remote_qp_number": 464, + "local_hca_index": 1, + "remote_hca_index": 0, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 0, + "path_index": 1, + "device": "roceP2p1s0f1", + "payload_writes": 33, + "payload_bytes": 540672, + "physical_hop_payload_bytes": 540672, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 713, + "remote_qp_number": 457, + "local_hca_index": 3, + "remote_hca_index": 2, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 2, + "path_index": 0, + "device": "rocep1s0f0", + "payload_writes": 33, + "payload_bytes": 540672, + "physical_hop_payload_bytes": 540672, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 465, + "remote_qp_number": 848, + "local_hca_index": 0, + "remote_hca_index": 1, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 2, + "path_index": 1, + "device": "roceP2p1s0f0", + "payload_writes": 33, + "payload_bytes": 540672, + "physical_hop_payload_bytes": 540672, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 458, + "remote_qp_number": 714, + "local_hca_index": 2, + "remote_hca_index": 3, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 3, + "path_index": 0, + "device": "rocep1s0f0", + "payload_writes": 33, + "payload_bytes": 540672, + "physical_hop_payload_bytes": 1081344, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 466, + "remote_qp_number": 848, + "local_hca_index": 0, + "remote_hca_index": 1, + "physical_hops": 2, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 3, + "path_index": 1, + "device": "roceP2p1s0f1", + "payload_writes": 33, + "payload_bytes": 540672, + "physical_hop_payload_bytes": 1081344, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 715, + "remote_qp_number": 458, + "local_hca_index": 3, + "remote_hca_index": 2, + "physical_hops": 2, + "retries": null, + "retry_events": null + } + ], + "expected_counters": { + "operations_per_rank": 33, + "paths_per_rank": 6, + "flags_per_path": 33, + "payload_bytes_path_0": 540672, + "payload_bytes_path_1": 540672 + } + }, + { + "rank": 2, + "eager_samples_us": [ + 207.20000565052032, + 125.2799928188324, + 49.15200173854828 + ], + "graph_samples_us": [ + 42.87999868392944, + 20.810666183630627, + 21.877333521842957 + ], + "eager_median_us": 125.2799928188324, + "graph_median_us": 21.877333521842957, + "correctness_cases": [ + { + "name": "rank-specific-index-pattern-a", + "mode": "eager", + "input_formula": "((rank + 1) * ((index % 17) - 8)) + rank", + "expected_formula": "10 * ((index % 17) - 8) + 6", + "input_sha256": "52ed8ae4094366d3b404eda20e086af8c2ea7673c41993998ab566c3c751df7b", + "expected_sha256": "202a06d669dc01f7432eda3e514cf5740ea9652f609c8e0dfec40716b2651522", + "output_sha256": "202a06d669dc01f7432eda3e514cf5740ea9652f609c8e0dfec40716b2651522", + "passed": true + }, + { + "name": "constant-timing-reference", + "mode": "eager", + "input_formula": "rank + 1", + "expected_formula": "10", + "input_sha256": "70e5ef9e0c8aa7e0ba8e4276d397164259c3cfa25b0de39f76b4d3e6a37c8329", + "expected_sha256": "d6de16041970b40bf4c7653ad60f95ae46976f16550b1d54988463980a19fbc2", + "output_sha256": "d6de16041970b40bf4c7653ad60f95ae46976f16550b1d54988463980a19fbc2", + "passed": true + }, + { + "name": "graph-input-mutation-a", + "mode": "graph-replay", + "input_formula": "((rank + 1) * ((index % 17) - 8)) + rank", + "expected_formula": "10 * ((index % 17) - 8) + 6", + "input_sha256": "52ed8ae4094366d3b404eda20e086af8c2ea7673c41993998ab566c3c751df7b", + "expected_sha256": "202a06d669dc01f7432eda3e514cf5740ea9652f609c8e0dfec40716b2651522", + "output_sha256": "202a06d669dc01f7432eda3e514cf5740ea9652f609c8e0dfec40716b2651522", + "passed": true + }, + { + "name": "graph-input-mutation-b", + "mode": "graph-replay", + "input_formula": "((4 - rank) * (((index * 5 + 3) % 19) - 9)) - rank", + "expected_formula": "10 * (((index * 5 + 3) % 19) - 9) - 6", + "input_sha256": "6ba9e176e052357b0bc90b26f9d5625780e57bcfcf091ccc891d4265273baf08", + "expected_sha256": "b282b975648deeef7ecb7491a997047c1f9561901f148450b08ee5b80712328a", + "output_sha256": "b282b975648deeef7ecb7491a997047c1f9561901f148450b08ee5b80712328a", + "passed": true + }, + { + "name": "constant-timing-output", + "mode": "graph-replay", + "input_formula": "rank + 1", + "expected_formula": "10", + "input_sha256": "70e5ef9e0c8aa7e0ba8e4276d397164259c3cfa25b0de39f76b4d3e6a37c8329", + "expected_sha256": "d6de16041970b40bf4c7653ad60f95ae46976f16550b1d54988463980a19fbc2", + "output_sha256": "d6de16041970b40bf4c7653ad60f95ae46976f16550b1d54988463980a19fbc2", + "passed": true + } + ], + "stats": { + "world_size": 4, + "rank": 2, + "hcas": [ + "rocep1s0f0", + "rocep1s0f1", + "roceP2p1s0f0", + "roceP2p1s0f1" + ], + "max_size": 2097152, + "max_gather_bytes": 2097152, + "slot_bytes": 2097152, + "epoch": 33, + "error_seq": 0, + "error_peer": 0, + "ctrl_seq": 33, + "spin_limit": 20000000, + "opposite_paths": 2, + "ops_posted": 33, + "writes_completed": 198, + "last_seq": 33, + "two_wave_activations": 0, + "two_wave_threshold_bytes": 196608, + "wave_mode": "two", + "peer_hca": { + "0": [ + 1, + 2 + ], + "1": [ + 1, + 3 + ], + "3": [ + 0, + 2 + ] + } + }, + "path_counters": [ + { + "peer_rank": 0, + "path_index": 0, + "device": "rocep1s0f1", + "payload_writes": 33, + "payload_bytes": 540672, + "physical_hop_payload_bytes": 1081344, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 847, + "remote_qp_number": 465, + "local_hca_index": 1, + "remote_hca_index": 0, + "physical_hops": 2, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 0, + "path_index": 1, + "device": "roceP2p1s0f0", + "payload_writes": 33, + "payload_bytes": 540672, + "physical_hop_payload_bytes": 1081344, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 457, + "remote_qp_number": 714, + "local_hca_index": 2, + "remote_hca_index": 3, + "physical_hops": 2, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 1, + "path_index": 0, + "device": "rocep1s0f1", + "payload_writes": 33, + "payload_bytes": 540672, + "physical_hop_payload_bytes": 540672, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 848, + "remote_qp_number": 465, + "local_hca_index": 1, + "remote_hca_index": 0, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 1, + "path_index": 1, + "device": "roceP2p1s0f1", + "payload_writes": 33, + "payload_bytes": 540672, + "physical_hop_payload_bytes": 540672, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 714, + "remote_qp_number": 458, + "local_hca_index": 3, + "remote_hca_index": 2, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 3, + "path_index": 0, + "device": "rocep1s0f0", + "payload_writes": 33, + "payload_bytes": 540672, + "physical_hop_payload_bytes": 540672, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 466, + "remote_qp_number": 849, + "local_hca_index": 0, + "remote_hca_index": 1, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 3, + "path_index": 1, + "device": "roceP2p1s0f0", + "payload_writes": 33, + "payload_bytes": 540672, + "physical_hop_payload_bytes": 540672, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 459, + "remote_qp_number": 715, + "local_hca_index": 2, + "remote_hca_index": 3, + "physical_hops": 1, + "retries": null, + "retry_events": null + } + ], + "expected_counters": { + "operations_per_rank": 33, + "paths_per_rank": 6, + "flags_per_path": 33, + "payload_bytes_path_0": 540672, + "payload_bytes_path_1": 540672 + } + }, + { + "rank": 3, + "eager_samples_us": [ + 50.464000552892685, + 138.84800672531128, + 54.1439987719059 + ], + "graph_samples_us": [ + 113.7600044409434, + 22.154666483402252, + 21.205333371957142 + ], + "eager_median_us": 54.1439987719059, + "graph_median_us": 22.154666483402252, + "correctness_cases": [ + { + "name": "rank-specific-index-pattern-a", + "mode": "eager", + "input_formula": "((rank + 1) * ((index % 17) - 8)) + rank", + "expected_formula": "10 * ((index % 17) - 8) + 6", + "input_sha256": "f2fd2757da7269e75a3e951341d0aca81d50ffeb15cb777beabb60cdb57c23a6", + "expected_sha256": "202a06d669dc01f7432eda3e514cf5740ea9652f609c8e0dfec40716b2651522", + "output_sha256": "202a06d669dc01f7432eda3e514cf5740ea9652f609c8e0dfec40716b2651522", + "passed": true + }, + { + "name": "constant-timing-reference", + "mode": "eager", + "input_formula": "rank + 1", + "expected_formula": "10", + "input_sha256": "6bb75bbeb7bf2338fbfa910d2cc216be460410287c94fa595007b348b0ed5a18", + "expected_sha256": "d6de16041970b40bf4c7653ad60f95ae46976f16550b1d54988463980a19fbc2", + "output_sha256": "d6de16041970b40bf4c7653ad60f95ae46976f16550b1d54988463980a19fbc2", + "passed": true + }, + { + "name": "graph-input-mutation-a", + "mode": "graph-replay", + "input_formula": "((rank + 1) * ((index % 17) - 8)) + rank", + "expected_formula": "10 * ((index % 17) - 8) + 6", + "input_sha256": "f2fd2757da7269e75a3e951341d0aca81d50ffeb15cb777beabb60cdb57c23a6", + "expected_sha256": "202a06d669dc01f7432eda3e514cf5740ea9652f609c8e0dfec40716b2651522", + "output_sha256": "202a06d669dc01f7432eda3e514cf5740ea9652f609c8e0dfec40716b2651522", + "passed": true + }, + { + "name": "graph-input-mutation-b", + "mode": "graph-replay", + "input_formula": "((4 - rank) * (((index * 5 + 3) % 19) - 9)) - rank", + "expected_formula": "10 * (((index * 5 + 3) % 19) - 9) - 6", + "input_sha256": "b206433a2115b49319dc85427a12a65c4651bf36545ec41cdb522d6c2516a1bb", + "expected_sha256": "b282b975648deeef7ecb7491a997047c1f9561901f148450b08ee5b80712328a", + "output_sha256": "b282b975648deeef7ecb7491a997047c1f9561901f148450b08ee5b80712328a", + "passed": true + }, + { + "name": "constant-timing-output", + "mode": "graph-replay", + "input_formula": "rank + 1", + "expected_formula": "10", + "input_sha256": "6bb75bbeb7bf2338fbfa910d2cc216be460410287c94fa595007b348b0ed5a18", + "expected_sha256": "d6de16041970b40bf4c7653ad60f95ae46976f16550b1d54988463980a19fbc2", + "output_sha256": "d6de16041970b40bf4c7653ad60f95ae46976f16550b1d54988463980a19fbc2", + "passed": true + } + ], + "stats": { + "world_size": 4, + "rank": 3, + "hcas": [ + "rocep1s0f0", + "rocep1s0f1", + "roceP2p1s0f0", + "roceP2p1s0f1" + ], + "max_size": 2097152, + "max_gather_bytes": 2097152, + "slot_bytes": 2097152, + "epoch": 33, + "error_seq": 0, + "error_peer": 0, + "ctrl_seq": 33, + "spin_limit": 20000000, + "opposite_paths": 2, + "ops_posted": 33, + "writes_completed": 198, + "last_seq": 33, + "two_wave_activations": 0, + "two_wave_threshold_bytes": 196608, + "wave_mode": "two", + "peer_hca": { + "0": [ + 0, + 2 + ], + "1": [ + 1, + 2 + ], + "2": [ + 1, + 3 + ] + } + }, + "path_counters": [ + { + "peer_rank": 0, + "path_index": 0, + "device": "rocep1s0f0", + "payload_writes": 33, + "payload_bytes": 540672, + "physical_hop_payload_bytes": 540672, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 464, + "remote_qp_number": 849, + "local_hca_index": 0, + "remote_hca_index": 1, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 0, + "path_index": 1, + "device": "roceP2p1s0f0", + "payload_writes": 33, + "payload_bytes": 540672, + "physical_hop_payload_bytes": 540672, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 457, + "remote_qp_number": 715, + "local_hca_index": 2, + "remote_hca_index": 3, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 1, + "path_index": 0, + "device": "rocep1s0f1", + "payload_writes": 33, + "payload_bytes": 540672, + "physical_hop_payload_bytes": 1081344, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 848, + "remote_qp_number": 466, + "local_hca_index": 1, + "remote_hca_index": 0, + "physical_hops": 2, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 1, + "path_index": 1, + "device": "roceP2p1s0f0", + "payload_writes": 33, + "payload_bytes": 540672, + "physical_hop_payload_bytes": 1081344, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 458, + "remote_qp_number": 715, + "local_hca_index": 2, + "remote_hca_index": 3, + "physical_hops": 2, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 2, + "path_index": 0, + "device": "rocep1s0f1", + "payload_writes": 33, + "payload_bytes": 540672, + "physical_hop_payload_bytes": 540672, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 849, + "remote_qp_number": 466, + "local_hca_index": 1, + "remote_hca_index": 0, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 2, + "path_index": 1, + "device": "roceP2p1s0f1", + "payload_writes": 33, + "payload_bytes": 540672, + "physical_hop_payload_bytes": 540672, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 715, + "remote_qp_number": 459, + "local_hca_index": 3, + "remote_hca_index": 2, + "physical_hops": 1, + "retries": null, + "retry_events": null + } + ], + "expected_counters": { + "operations_per_rank": 33, + "paths_per_rank": 6, + "flags_per_path": 33, + "payload_bytes_path_0": 540672, + "payload_bytes_path_1": 540672 + } + } + ] + }, + { + "schema": "b12x.rocenante-virtual-diagonal-evidence/v1", + "status": "research-only", + "payload_bytes": 163840, + "warmups": 2, + "samples": 3, + "graph_operations_per_replay": 3, + "ranks": [ + { + "rank": 0, + "eager_samples_us": [ + 1161.5999937057495, + 145.31199634075165, + 79.13599908351898 + ], + "graph_samples_us": [ + 122.1440037091573, + 61.86666587988535, + 77.60000228881836 + ], + "eager_median_us": 145.31199634075165, + "graph_median_us": 77.60000228881836, + "correctness_cases": [ + { + "name": "rank-specific-index-pattern-a", + "mode": "eager", + "input_formula": "((rank + 1) * ((index % 17) - 8)) + rank", + "expected_formula": "10 * ((index % 17) - 8) + 6", + "input_sha256": "c5353ec97297a7430ac63087c23328c4a6c7870de568a57ec33cb3f5e24b1423", + "expected_sha256": "f8085672fe1ec140c1b4b6630e0a375ccc703cfb9a8ecc98e7c55d2f182bd954", + "output_sha256": "f8085672fe1ec140c1b4b6630e0a375ccc703cfb9a8ecc98e7c55d2f182bd954", + "passed": true + }, + { + "name": "constant-timing-reference", + "mode": "eager", + "input_formula": "rank + 1", + "expected_formula": "10", + "input_sha256": "12c37dce350901b16b44060cfc088f8347c7e165e1e6b4ee37461ddc2b0ce427", + "expected_sha256": "38ef226f69c13e18514f88f1836409267aa4cfb380b806a52a198756e24adc26", + "output_sha256": "38ef226f69c13e18514f88f1836409267aa4cfb380b806a52a198756e24adc26", + "passed": true + }, + { + "name": "graph-input-mutation-a", + "mode": "graph-replay", + "input_formula": "((rank + 1) * ((index % 17) - 8)) + rank", + "expected_formula": "10 * ((index % 17) - 8) + 6", + "input_sha256": "c5353ec97297a7430ac63087c23328c4a6c7870de568a57ec33cb3f5e24b1423", + "expected_sha256": "f8085672fe1ec140c1b4b6630e0a375ccc703cfb9a8ecc98e7c55d2f182bd954", + "output_sha256": "f8085672fe1ec140c1b4b6630e0a375ccc703cfb9a8ecc98e7c55d2f182bd954", + "passed": true + }, + { + "name": "graph-input-mutation-b", + "mode": "graph-replay", + "input_formula": "((4 - rank) * (((index * 5 + 3) % 19) - 9)) - rank", + "expected_formula": "10 * (((index * 5 + 3) % 19) - 9) - 6", + "input_sha256": "0400d22ed348280b33c32240f6dad0dbdc4e01736f5dd9daca54126ea27f4888", + "expected_sha256": "a773986c394cf2f8b67267812f58785f84fd7be5ceace0f32ce3d8730abbef1a", + "output_sha256": "a773986c394cf2f8b67267812f58785f84fd7be5ceace0f32ce3d8730abbef1a", + "passed": true + }, + { + "name": "constant-timing-output", + "mode": "graph-replay", + "input_formula": "rank + 1", + "expected_formula": "10", + "input_sha256": "12c37dce350901b16b44060cfc088f8347c7e165e1e6b4ee37461ddc2b0ce427", + "expected_sha256": "38ef226f69c13e18514f88f1836409267aa4cfb380b806a52a198756e24adc26", + "output_sha256": "38ef226f69c13e18514f88f1836409267aa4cfb380b806a52a198756e24adc26", + "passed": true + } + ], + "stats": { + "world_size": 4, + "rank": 0, + "hcas": [ + "rocep1s0f0", + "rocep1s0f1", + "roceP2p1s0f0", + "roceP2p1s0f1" + ], + "max_size": 2097152, + "max_gather_bytes": 2097152, + "slot_bytes": 2097152, + "epoch": 33, + "error_seq": 0, + "error_peer": 0, + "ctrl_seq": 33, + "spin_limit": 20000000, + "opposite_paths": 2, + "ops_posted": 33, + "writes_completed": 198, + "last_seq": 33, + "two_wave_activations": 0, + "two_wave_threshold_bytes": 196608, + "wave_mode": "two", + "peer_hca": { + "1": [ + 0, + 2 + ], + "2": [ + 0, + 3 + ], + "3": [ + 1, + 3 + ] + } + }, + "path_counters": [ + { + "peer_rank": 1, + "path_index": 0, + "device": "rocep1s0f0", + "payload_writes": 33, + "payload_bytes": 2703360, + "physical_hop_payload_bytes": 2703360, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 467, + "remote_qp_number": 850, + "local_hca_index": 0, + "remote_hca_index": 1, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 1, + "path_index": 1, + "device": "roceP2p1s0f0", + "payload_writes": 33, + "payload_bytes": 2703360, + "physical_hop_payload_bytes": 2703360, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 460, + "remote_qp_number": 716, + "local_hca_index": 2, + "remote_hca_index": 3, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 2, + "path_index": 0, + "device": "rocep1s0f0", + "payload_writes": 33, + "payload_bytes": 2703360, + "physical_hop_payload_bytes": 5406720, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 468, + "remote_qp_number": 850, + "local_hca_index": 0, + "remote_hca_index": 1, + "physical_hops": 2, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 2, + "path_index": 1, + "device": "roceP2p1s0f1", + "payload_writes": 33, + "payload_bytes": 2703360, + "physical_hop_payload_bytes": 5406720, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 717, + "remote_qp_number": 460, + "local_hca_index": 3, + "remote_hca_index": 2, + "physical_hops": 2, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 3, + "path_index": 0, + "device": "rocep1s0f1", + "payload_writes": 33, + "payload_bytes": 2703360, + "physical_hop_payload_bytes": 2703360, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 852, + "remote_qp_number": 467, + "local_hca_index": 1, + "remote_hca_index": 0, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 3, + "path_index": 1, + "device": "roceP2p1s0f1", + "payload_writes": 33, + "payload_bytes": 2703360, + "physical_hop_payload_bytes": 2703360, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 718, + "remote_qp_number": 460, + "local_hca_index": 3, + "remote_hca_index": 2, + "physical_hops": 1, + "retries": null, + "retry_events": null + } + ], + "expected_counters": { + "operations_per_rank": 33, + "paths_per_rank": 6, + "flags_per_path": 33, + "payload_bytes_path_0": 2703360, + "payload_bytes_path_1": 2703360 + } + }, + { + "rank": 1, + "eager_samples_us": [ + 255.16799092292786, + 152.0639955997467, + 5669.312000274658 + ], + "graph_samples_us": [ + 86.66666348775227, + 58.84799857934316, + 71.19999825954437 + ], + "eager_median_us": 255.16799092292786, + "graph_median_us": 71.19999825954437, + "correctness_cases": [ + { + "name": "rank-specific-index-pattern-a", + "mode": "eager", + "input_formula": "((rank + 1) * ((index % 17) - 8)) + rank", + "expected_formula": "10 * ((index % 17) - 8) + 6", + "input_sha256": "6664adc78c736e7e12488891b5cce2f17a973e4afbf5281c941827c0cb5d12f5", + "expected_sha256": "f8085672fe1ec140c1b4b6630e0a375ccc703cfb9a8ecc98e7c55d2f182bd954", + "output_sha256": "f8085672fe1ec140c1b4b6630e0a375ccc703cfb9a8ecc98e7c55d2f182bd954", + "passed": true + }, + { + "name": "constant-timing-reference", + "mode": "eager", + "input_formula": "rank + 1", + "expected_formula": "10", + "input_sha256": "75fbaf49ad48bcbe0f9410ce15291b79664a7fd94f9c5148ffaf5bf20a15fadf", + "expected_sha256": "38ef226f69c13e18514f88f1836409267aa4cfb380b806a52a198756e24adc26", + "output_sha256": "38ef226f69c13e18514f88f1836409267aa4cfb380b806a52a198756e24adc26", + "passed": true + }, + { + "name": "graph-input-mutation-a", + "mode": "graph-replay", + "input_formula": "((rank + 1) * ((index % 17) - 8)) + rank", + "expected_formula": "10 * ((index % 17) - 8) + 6", + "input_sha256": "6664adc78c736e7e12488891b5cce2f17a973e4afbf5281c941827c0cb5d12f5", + "expected_sha256": "f8085672fe1ec140c1b4b6630e0a375ccc703cfb9a8ecc98e7c55d2f182bd954", + "output_sha256": "f8085672fe1ec140c1b4b6630e0a375ccc703cfb9a8ecc98e7c55d2f182bd954", + "passed": true + }, + { + "name": "graph-input-mutation-b", + "mode": "graph-replay", + "input_formula": "((4 - rank) * (((index * 5 + 3) % 19) - 9)) - rank", + "expected_formula": "10 * (((index * 5 + 3) % 19) - 9) - 6", + "input_sha256": "4b38daba9747bc84a8bf035e50531c0b182ca834a2a28a00b55317f9c75db659", + "expected_sha256": "a773986c394cf2f8b67267812f58785f84fd7be5ceace0f32ce3d8730abbef1a", + "output_sha256": "a773986c394cf2f8b67267812f58785f84fd7be5ceace0f32ce3d8730abbef1a", + "passed": true + }, + { + "name": "constant-timing-output", + "mode": "graph-replay", + "input_formula": "rank + 1", + "expected_formula": "10", + "input_sha256": "75fbaf49ad48bcbe0f9410ce15291b79664a7fd94f9c5148ffaf5bf20a15fadf", + "expected_sha256": "38ef226f69c13e18514f88f1836409267aa4cfb380b806a52a198756e24adc26", + "output_sha256": "38ef226f69c13e18514f88f1836409267aa4cfb380b806a52a198756e24adc26", + "passed": true + } + ], + "stats": { + "world_size": 4, + "rank": 1, + "hcas": [ + "rocep1s0f0", + "rocep1s0f1", + "roceP2p1s0f0", + "roceP2p1s0f1" + ], + "max_size": 2097152, + "max_gather_bytes": 2097152, + "slot_bytes": 2097152, + "epoch": 33, + "error_seq": 0, + "error_peer": 0, + "ctrl_seq": 33, + "spin_limit": 20000000, + "opposite_paths": 2, + "ops_posted": 33, + "writes_completed": 198, + "last_seq": 33, + "two_wave_activations": 0, + "two_wave_threshold_bytes": 196608, + "wave_mode": "two", + "peer_hca": { + "0": [ + 1, + 3 + ], + "2": [ + 0, + 2 + ], + "3": [ + 0, + 3 + ] + } + }, + "path_counters": [ + { + "peer_rank": 0, + "path_index": 0, + "device": "rocep1s0f1", + "payload_writes": 33, + "payload_bytes": 2703360, + "physical_hop_payload_bytes": 2703360, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 850, + "remote_qp_number": 467, + "local_hca_index": 1, + "remote_hca_index": 0, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 0, + "path_index": 1, + "device": "roceP2p1s0f1", + "payload_writes": 33, + "payload_bytes": 2703360, + "physical_hop_payload_bytes": 2703360, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 716, + "remote_qp_number": 460, + "local_hca_index": 3, + "remote_hca_index": 2, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 2, + "path_index": 0, + "device": "rocep1s0f0", + "payload_writes": 33, + "payload_bytes": 2703360, + "physical_hop_payload_bytes": 2703360, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 468, + "remote_qp_number": 851, + "local_hca_index": 0, + "remote_hca_index": 1, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 2, + "path_index": 1, + "device": "roceP2p1s0f0", + "payload_writes": 33, + "payload_bytes": 2703360, + "physical_hop_payload_bytes": 2703360, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 461, + "remote_qp_number": 717, + "local_hca_index": 2, + "remote_hca_index": 3, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 3, + "path_index": 0, + "device": "rocep1s0f0", + "payload_writes": 33, + "payload_bytes": 2703360, + "physical_hop_payload_bytes": 5406720, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 469, + "remote_qp_number": 851, + "local_hca_index": 0, + "remote_hca_index": 1, + "physical_hops": 2, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 3, + "path_index": 1, + "device": "roceP2p1s0f1", + "payload_writes": 33, + "payload_bytes": 2703360, + "physical_hop_payload_bytes": 5406720, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 718, + "remote_qp_number": 461, + "local_hca_index": 3, + "remote_hca_index": 2, + "physical_hops": 2, + "retries": null, + "retry_events": null + } + ], + "expected_counters": { + "operations_per_rank": 33, + "paths_per_rank": 6, + "flags_per_path": 33, + "payload_bytes_path_0": 2703360, + "payload_bytes_path_1": 2703360 + } + }, + { + "rank": 2, + "eager_samples_us": [ + 707.647979259491, + 115.55200070142746, + 120.4800009727478 + ], + "graph_samples_us": [ + 271.8613346417745, + 59.14666752020518, + 79.83999947706859 + ], + "eager_median_us": 120.4800009727478, + "graph_median_us": 79.83999947706859, + "correctness_cases": [ + { + "name": "rank-specific-index-pattern-a", + "mode": "eager", + "input_formula": "((rank + 1) * ((index % 17) - 8)) + rank", + "expected_formula": "10 * ((index % 17) - 8) + 6", + "input_sha256": "b83703e6cd9c085f96af128d6f27dcc084689dbd778a4c319f836b871faa29ea", + "expected_sha256": "f8085672fe1ec140c1b4b6630e0a375ccc703cfb9a8ecc98e7c55d2f182bd954", + "output_sha256": "f8085672fe1ec140c1b4b6630e0a375ccc703cfb9a8ecc98e7c55d2f182bd954", + "passed": true + }, + { + "name": "constant-timing-reference", + "mode": "eager", + "input_formula": "rank + 1", + "expected_formula": "10", + "input_sha256": "3e8d4c9ba4b5404b015873cb17affd09878c8cc461e7c81c759b4ffe6e9f9792", + "expected_sha256": "38ef226f69c13e18514f88f1836409267aa4cfb380b806a52a198756e24adc26", + "output_sha256": "38ef226f69c13e18514f88f1836409267aa4cfb380b806a52a198756e24adc26", + "passed": true + }, + { + "name": "graph-input-mutation-a", + "mode": "graph-replay", + "input_formula": "((rank + 1) * ((index % 17) - 8)) + rank", + "expected_formula": "10 * ((index % 17) - 8) + 6", + "input_sha256": "b83703e6cd9c085f96af128d6f27dcc084689dbd778a4c319f836b871faa29ea", + "expected_sha256": "f8085672fe1ec140c1b4b6630e0a375ccc703cfb9a8ecc98e7c55d2f182bd954", + "output_sha256": "f8085672fe1ec140c1b4b6630e0a375ccc703cfb9a8ecc98e7c55d2f182bd954", + "passed": true + }, + { + "name": "graph-input-mutation-b", + "mode": "graph-replay", + "input_formula": "((4 - rank) * (((index * 5 + 3) % 19) - 9)) - rank", + "expected_formula": "10 * (((index * 5 + 3) % 19) - 9) - 6", + "input_sha256": "8aed95419a9fe77e6131f40d29c181ff843bedacd4d2a9687f38d211c4c256a1", + "expected_sha256": "a773986c394cf2f8b67267812f58785f84fd7be5ceace0f32ce3d8730abbef1a", + "output_sha256": "a773986c394cf2f8b67267812f58785f84fd7be5ceace0f32ce3d8730abbef1a", + "passed": true + }, + { + "name": "constant-timing-output", + "mode": "graph-replay", + "input_formula": "rank + 1", + "expected_formula": "10", + "input_sha256": "3e8d4c9ba4b5404b015873cb17affd09878c8cc461e7c81c759b4ffe6e9f9792", + "expected_sha256": "38ef226f69c13e18514f88f1836409267aa4cfb380b806a52a198756e24adc26", + "output_sha256": "38ef226f69c13e18514f88f1836409267aa4cfb380b806a52a198756e24adc26", + "passed": true + } + ], + "stats": { + "world_size": 4, + "rank": 2, + "hcas": [ + "rocep1s0f0", + "rocep1s0f1", + "roceP2p1s0f0", + "roceP2p1s0f1" + ], + "max_size": 2097152, + "max_gather_bytes": 2097152, + "slot_bytes": 2097152, + "epoch": 33, + "error_seq": 0, + "error_peer": 0, + "ctrl_seq": 33, + "spin_limit": 20000000, + "opposite_paths": 2, + "ops_posted": 33, + "writes_completed": 198, + "last_seq": 33, + "two_wave_activations": 0, + "two_wave_threshold_bytes": 196608, + "wave_mode": "two", + "peer_hca": { + "0": [ + 1, + 2 + ], + "1": [ + 1, + 3 + ], + "3": [ + 0, + 2 + ] + } + }, + "path_counters": [ + { + "peer_rank": 0, + "path_index": 0, + "device": "rocep1s0f1", + "payload_writes": 33, + "payload_bytes": 2703360, + "physical_hop_payload_bytes": 5406720, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 850, + "remote_qp_number": 468, + "local_hca_index": 1, + "remote_hca_index": 0, + "physical_hops": 2, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 0, + "path_index": 1, + "device": "roceP2p1s0f0", + "payload_writes": 33, + "payload_bytes": 2703360, + "physical_hop_payload_bytes": 5406720, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 460, + "remote_qp_number": 717, + "local_hca_index": 2, + "remote_hca_index": 3, + "physical_hops": 2, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 1, + "path_index": 0, + "device": "rocep1s0f1", + "payload_writes": 33, + "payload_bytes": 2703360, + "physical_hop_payload_bytes": 2703360, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 851, + "remote_qp_number": 468, + "local_hca_index": 1, + "remote_hca_index": 0, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 1, + "path_index": 1, + "device": "roceP2p1s0f1", + "payload_writes": 33, + "payload_bytes": 2703360, + "physical_hop_payload_bytes": 2703360, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 717, + "remote_qp_number": 461, + "local_hca_index": 3, + "remote_hca_index": 2, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 3, + "path_index": 0, + "device": "rocep1s0f0", + "payload_writes": 33, + "payload_bytes": 2703360, + "physical_hop_payload_bytes": 2703360, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 469, + "remote_qp_number": 852, + "local_hca_index": 0, + "remote_hca_index": 1, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 3, + "path_index": 1, + "device": "roceP2p1s0f0", + "payload_writes": 33, + "payload_bytes": 2703360, + "physical_hop_payload_bytes": 2703360, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 462, + "remote_qp_number": 718, + "local_hca_index": 2, + "remote_hca_index": 3, + "physical_hops": 1, + "retries": null, + "retry_events": null + } + ], + "expected_counters": { + "operations_per_rank": 33, + "paths_per_rank": 6, + "flags_per_path": 33, + "payload_bytes_path_0": 2703360, + "payload_bytes_path_1": 2703360 + } + }, + { + "rank": 3, + "eager_samples_us": [ + 205.08800446987152, + 178.9119988679886, + 87.55200356245041 + ], + "graph_samples_us": [ + 181.59999450047812, + 88.55467041333516, + 60.80000102519989 + ], + "eager_median_us": 178.9119988679886, + "graph_median_us": 88.55467041333516, + "correctness_cases": [ + { + "name": "rank-specific-index-pattern-a", + "mode": "eager", + "input_formula": "((rank + 1) * ((index % 17) - 8)) + rank", + "expected_formula": "10 * ((index % 17) - 8) + 6", + "input_sha256": "dcee6c3b8ab802537925c1ca8a6b9a61e8dfa599e321c78de57ab4e3492896c2", + "expected_sha256": "f8085672fe1ec140c1b4b6630e0a375ccc703cfb9a8ecc98e7c55d2f182bd954", + "output_sha256": "f8085672fe1ec140c1b4b6630e0a375ccc703cfb9a8ecc98e7c55d2f182bd954", + "passed": true + }, + { + "name": "constant-timing-reference", + "mode": "eager", + "input_formula": "rank + 1", + "expected_formula": "10", + "input_sha256": "5dd7020a56031ec31ca5632546126efc36de0f095e462f6a88593782f98d1c4f", + "expected_sha256": "38ef226f69c13e18514f88f1836409267aa4cfb380b806a52a198756e24adc26", + "output_sha256": "38ef226f69c13e18514f88f1836409267aa4cfb380b806a52a198756e24adc26", + "passed": true + }, + { + "name": "graph-input-mutation-a", + "mode": "graph-replay", + "input_formula": "((rank + 1) * ((index % 17) - 8)) + rank", + "expected_formula": "10 * ((index % 17) - 8) + 6", + "input_sha256": "dcee6c3b8ab802537925c1ca8a6b9a61e8dfa599e321c78de57ab4e3492896c2", + "expected_sha256": "f8085672fe1ec140c1b4b6630e0a375ccc703cfb9a8ecc98e7c55d2f182bd954", + "output_sha256": "f8085672fe1ec140c1b4b6630e0a375ccc703cfb9a8ecc98e7c55d2f182bd954", + "passed": true + }, + { + "name": "graph-input-mutation-b", + "mode": "graph-replay", + "input_formula": "((4 - rank) * (((index * 5 + 3) % 19) - 9)) - rank", + "expected_formula": "10 * (((index * 5 + 3) % 19) - 9) - 6", + "input_sha256": "510f45a03ff47ff11cb80c4ae7e33edd5fa19707bd7fd3c80e7998f71d34a173", + "expected_sha256": "a773986c394cf2f8b67267812f58785f84fd7be5ceace0f32ce3d8730abbef1a", + "output_sha256": "a773986c394cf2f8b67267812f58785f84fd7be5ceace0f32ce3d8730abbef1a", + "passed": true + }, + { + "name": "constant-timing-output", + "mode": "graph-replay", + "input_formula": "rank + 1", + "expected_formula": "10", + "input_sha256": "5dd7020a56031ec31ca5632546126efc36de0f095e462f6a88593782f98d1c4f", + "expected_sha256": "38ef226f69c13e18514f88f1836409267aa4cfb380b806a52a198756e24adc26", + "output_sha256": "38ef226f69c13e18514f88f1836409267aa4cfb380b806a52a198756e24adc26", + "passed": true + } + ], + "stats": { + "world_size": 4, + "rank": 3, + "hcas": [ + "rocep1s0f0", + "rocep1s0f1", + "roceP2p1s0f0", + "roceP2p1s0f1" + ], + "max_size": 2097152, + "max_gather_bytes": 2097152, + "slot_bytes": 2097152, + "epoch": 33, + "error_seq": 0, + "error_peer": 0, + "ctrl_seq": 33, + "spin_limit": 20000000, + "opposite_paths": 2, + "ops_posted": 33, + "writes_completed": 198, + "last_seq": 33, + "two_wave_activations": 0, + "two_wave_threshold_bytes": 196608, + "wave_mode": "two", + "peer_hca": { + "0": [ + 0, + 2 + ], + "1": [ + 1, + 2 + ], + "2": [ + 1, + 3 + ] + } + }, + "path_counters": [ + { + "peer_rank": 0, + "path_index": 0, + "device": "rocep1s0f0", + "payload_writes": 33, + "payload_bytes": 2703360, + "physical_hop_payload_bytes": 2703360, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 467, + "remote_qp_number": 852, + "local_hca_index": 0, + "remote_hca_index": 1, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 0, + "path_index": 1, + "device": "roceP2p1s0f0", + "payload_writes": 33, + "payload_bytes": 2703360, + "physical_hop_payload_bytes": 2703360, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 460, + "remote_qp_number": 718, + "local_hca_index": 2, + "remote_hca_index": 3, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 1, + "path_index": 0, + "device": "rocep1s0f1", + "payload_writes": 33, + "payload_bytes": 2703360, + "physical_hop_payload_bytes": 5406720, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 851, + "remote_qp_number": 469, + "local_hca_index": 1, + "remote_hca_index": 0, + "physical_hops": 2, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 1, + "path_index": 1, + "device": "roceP2p1s0f0", + "payload_writes": 33, + "payload_bytes": 2703360, + "physical_hop_payload_bytes": 5406720, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 461, + "remote_qp_number": 718, + "local_hca_index": 2, + "remote_hca_index": 3, + "physical_hops": 2, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 2, + "path_index": 0, + "device": "rocep1s0f1", + "payload_writes": 33, + "payload_bytes": 2703360, + "physical_hop_payload_bytes": 2703360, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 852, + "remote_qp_number": 469, + "local_hca_index": 1, + "remote_hca_index": 0, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 2, + "path_index": 1, + "device": "roceP2p1s0f1", + "payload_writes": 33, + "payload_bytes": 2703360, + "physical_hop_payload_bytes": 2703360, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 718, + "remote_qp_number": 462, + "local_hca_index": 3, + "remote_hca_index": 2, + "physical_hops": 1, + "retries": null, + "retry_events": null + } + ], + "expected_counters": { + "operations_per_rank": 33, + "paths_per_rank": 6, + "flags_per_path": 33, + "payload_bytes_path_0": 2703360, + "payload_bytes_path_1": 2703360 + } + } + ] + }, + { + "schema": "b12x.rocenante-virtual-diagonal-evidence/v1", + "status": "research-only", + "payload_bytes": 229376, + "warmups": 2, + "samples": 3, + "graph_operations_per_replay": 3, + "ranks": [ + { + "rank": 0, + "eager_samples_us": [ + 247.5840002298355, + 222.84799814224243, + 50.11200159788132 + ], + "graph_samples_us": [ + 63.97866706053416, + 89.32266632715861, + 61.237335205078125 + ], + "eager_median_us": 222.84799814224243, + "graph_median_us": 63.97866706053416, + "correctness_cases": [ + { + "name": "rank-specific-index-pattern-a", + "mode": "eager", + "input_formula": "((rank + 1) * ((index % 17) - 8)) + rank", + "expected_formula": "10 * ((index % 17) - 8) + 6", + "input_sha256": "cd404c52ade92a17da6209b4cf2ee77348b5412b717f9f263965947995c676a5", + "expected_sha256": "93a73648c332a1d4054183a6b8d177e63092ccb0f28cb04743c47cd32da4721c", + "output_sha256": "93a73648c332a1d4054183a6b8d177e63092ccb0f28cb04743c47cd32da4721c", + "passed": true + }, + { + "name": "constant-timing-reference", + "mode": "eager", + "input_formula": "rank + 1", + "expected_formula": "10", + "input_sha256": "7caecc2288cee7679c4bf5be7d8c5dee3c5229072b85f148afae1af4c0b0e4b5", + "expected_sha256": "1a64b11191799028236a93dc0ffb15330e07140967c672eb0475b5e63682ea87", + "output_sha256": "1a64b11191799028236a93dc0ffb15330e07140967c672eb0475b5e63682ea87", + "passed": true + }, + { + "name": "graph-input-mutation-a", + "mode": "graph-replay", + "input_formula": "((rank + 1) * ((index % 17) - 8)) + rank", + "expected_formula": "10 * ((index % 17) - 8) + 6", + "input_sha256": "cd404c52ade92a17da6209b4cf2ee77348b5412b717f9f263965947995c676a5", + "expected_sha256": "93a73648c332a1d4054183a6b8d177e63092ccb0f28cb04743c47cd32da4721c", + "output_sha256": "93a73648c332a1d4054183a6b8d177e63092ccb0f28cb04743c47cd32da4721c", + "passed": true + }, + { + "name": "graph-input-mutation-b", + "mode": "graph-replay", + "input_formula": "((4 - rank) * (((index * 5 + 3) % 19) - 9)) - rank", + "expected_formula": "10 * (((index * 5 + 3) % 19) - 9) - 6", + "input_sha256": "52450971e24b57afd045cbfa9af8ce817c1edf1a1007bb858c525a14cc5b31c8", + "expected_sha256": "f2020cf5aee4a14c69ba387ab5f7ecc167433ecc88f795b63cc205821f894c29", + "output_sha256": "f2020cf5aee4a14c69ba387ab5f7ecc167433ecc88f795b63cc205821f894c29", + "passed": true + }, + { + "name": "constant-timing-output", + "mode": "graph-replay", + "input_formula": "rank + 1", + "expected_formula": "10", + "input_sha256": "7caecc2288cee7679c4bf5be7d8c5dee3c5229072b85f148afae1af4c0b0e4b5", + "expected_sha256": "1a64b11191799028236a93dc0ffb15330e07140967c672eb0475b5e63682ea87", + "output_sha256": "1a64b11191799028236a93dc0ffb15330e07140967c672eb0475b5e63682ea87", + "passed": true + } + ], + "stats": { + "world_size": 4, + "rank": 0, + "hcas": [ + "rocep1s0f0", + "rocep1s0f1", + "roceP2p1s0f0", + "roceP2p1s0f1" + ], + "max_size": 2097152, + "max_gather_bytes": 2097152, + "slot_bytes": 2097152, + "epoch": 33, + "error_seq": 0, + "error_peer": 0, + "ctrl_seq": 33, + "spin_limit": 20000000, + "opposite_paths": 2, + "ops_posted": 33, + "writes_completed": 198, + "last_seq": 33, + "two_wave_activations": 33, + "two_wave_threshold_bytes": 196608, + "wave_mode": "two", + "peer_hca": { + "1": [ + 0, + 2 + ], + "2": [ + 0, + 3 + ], + "3": [ + 1, + 3 + ] + } + }, + "path_counters": [ + { + "peer_rank": 1, + "path_index": 0, + "device": "rocep1s0f0", + "payload_writes": 33, + "payload_bytes": 3784704, + "physical_hop_payload_bytes": 3784704, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 470, + "remote_qp_number": 853, + "local_hca_index": 0, + "remote_hca_index": 1, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 1, + "path_index": 1, + "device": "roceP2p1s0f0", + "payload_writes": 33, + "payload_bytes": 3784704, + "physical_hop_payload_bytes": 3784704, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 463, + "remote_qp_number": 719, + "local_hca_index": 2, + "remote_hca_index": 3, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 2, + "path_index": 0, + "device": "rocep1s0f0", + "payload_writes": 33, + "payload_bytes": 3784704, + "physical_hop_payload_bytes": 7569408, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 471, + "remote_qp_number": 853, + "local_hca_index": 0, + "remote_hca_index": 1, + "physical_hops": 2, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 2, + "path_index": 1, + "device": "roceP2p1s0f1", + "payload_writes": 33, + "payload_bytes": 3784704, + "physical_hop_payload_bytes": 7569408, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 720, + "remote_qp_number": 463, + "local_hca_index": 3, + "remote_hca_index": 2, + "physical_hops": 2, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 3, + "path_index": 0, + "device": "rocep1s0f1", + "payload_writes": 33, + "payload_bytes": 3784704, + "physical_hop_payload_bytes": 3784704, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 855, + "remote_qp_number": 470, + "local_hca_index": 1, + "remote_hca_index": 0, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 3, + "path_index": 1, + "device": "roceP2p1s0f1", + "payload_writes": 33, + "payload_bytes": 3784704, + "physical_hop_payload_bytes": 3784704, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 721, + "remote_qp_number": 463, + "local_hca_index": 3, + "remote_hca_index": 2, + "physical_hops": 1, + "retries": null, + "retry_events": null + } + ], + "expected_counters": { + "operations_per_rank": 33, + "paths_per_rank": 6, + "flags_per_path": 33, + "payload_bytes_path_0": 3784704, + "payload_bytes_path_1": 3784704 + } + }, + { + "rank": 1, + "eager_samples_us": [ + 598.8159775733948, + 125.2799928188324, + 44188.44985961914 + ], + "graph_samples_us": [ + 181.3973387082418, + 86.61333719889323, + 59.04000004132589 + ], + "eager_median_us": 598.8159775733948, + "graph_median_us": 86.61333719889323, + "correctness_cases": [ + { + "name": "rank-specific-index-pattern-a", + "mode": "eager", + "input_formula": "((rank + 1) * ((index % 17) - 8)) + rank", + "expected_formula": "10 * ((index % 17) - 8) + 6", + "input_sha256": "9d31d0c21c09ef82790ff09044224bab4dcf862cfce6b307f590980c6cd706cf", + "expected_sha256": "93a73648c332a1d4054183a6b8d177e63092ccb0f28cb04743c47cd32da4721c", + "output_sha256": "93a73648c332a1d4054183a6b8d177e63092ccb0f28cb04743c47cd32da4721c", + "passed": true + }, + { + "name": "constant-timing-reference", + "mode": "eager", + "input_formula": "rank + 1", + "expected_formula": "10", + "input_sha256": "2259bf8ca07ea75b30498c193744fcb7297784c10811ada2f668ead07c696006", + "expected_sha256": "1a64b11191799028236a93dc0ffb15330e07140967c672eb0475b5e63682ea87", + "output_sha256": "1a64b11191799028236a93dc0ffb15330e07140967c672eb0475b5e63682ea87", + "passed": true + }, + { + "name": "graph-input-mutation-a", + "mode": "graph-replay", + "input_formula": "((rank + 1) * ((index % 17) - 8)) + rank", + "expected_formula": "10 * ((index % 17) - 8) + 6", + "input_sha256": "9d31d0c21c09ef82790ff09044224bab4dcf862cfce6b307f590980c6cd706cf", + "expected_sha256": "93a73648c332a1d4054183a6b8d177e63092ccb0f28cb04743c47cd32da4721c", + "output_sha256": "93a73648c332a1d4054183a6b8d177e63092ccb0f28cb04743c47cd32da4721c", + "passed": true + }, + { + "name": "graph-input-mutation-b", + "mode": "graph-replay", + "input_formula": "((4 - rank) * (((index * 5 + 3) % 19) - 9)) - rank", + "expected_formula": "10 * (((index * 5 + 3) % 19) - 9) - 6", + "input_sha256": "42e0c790effb0da9acf2a617268417e1815d7b403065c34d189f84ec21058599", + "expected_sha256": "f2020cf5aee4a14c69ba387ab5f7ecc167433ecc88f795b63cc205821f894c29", + "output_sha256": "f2020cf5aee4a14c69ba387ab5f7ecc167433ecc88f795b63cc205821f894c29", + "passed": true + }, + { + "name": "constant-timing-output", + "mode": "graph-replay", + "input_formula": "rank + 1", + "expected_formula": "10", + "input_sha256": "2259bf8ca07ea75b30498c193744fcb7297784c10811ada2f668ead07c696006", + "expected_sha256": "1a64b11191799028236a93dc0ffb15330e07140967c672eb0475b5e63682ea87", + "output_sha256": "1a64b11191799028236a93dc0ffb15330e07140967c672eb0475b5e63682ea87", + "passed": true + } + ], + "stats": { + "world_size": 4, + "rank": 1, + "hcas": [ + "rocep1s0f0", + "rocep1s0f1", + "roceP2p1s0f0", + "roceP2p1s0f1" + ], + "max_size": 2097152, + "max_gather_bytes": 2097152, + "slot_bytes": 2097152, + "epoch": 33, + "error_seq": 0, + "error_peer": 0, + "ctrl_seq": 33, + "spin_limit": 20000000, + "opposite_paths": 2, + "ops_posted": 33, + "writes_completed": 198, + "last_seq": 33, + "two_wave_activations": 33, + "two_wave_threshold_bytes": 196608, + "wave_mode": "two", + "peer_hca": { + "0": [ + 1, + 3 + ], + "2": [ + 0, + 2 + ], + "3": [ + 0, + 3 + ] + } + }, + "path_counters": [ + { + "peer_rank": 0, + "path_index": 0, + "device": "rocep1s0f1", + "payload_writes": 33, + "payload_bytes": 3784704, + "physical_hop_payload_bytes": 3784704, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 853, + "remote_qp_number": 470, + "local_hca_index": 1, + "remote_hca_index": 0, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 0, + "path_index": 1, + "device": "roceP2p1s0f1", + "payload_writes": 33, + "payload_bytes": 3784704, + "physical_hop_payload_bytes": 3784704, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 719, + "remote_qp_number": 463, + "local_hca_index": 3, + "remote_hca_index": 2, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 2, + "path_index": 0, + "device": "rocep1s0f0", + "payload_writes": 33, + "payload_bytes": 3784704, + "physical_hop_payload_bytes": 3784704, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 471, + "remote_qp_number": 854, + "local_hca_index": 0, + "remote_hca_index": 1, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 2, + "path_index": 1, + "device": "roceP2p1s0f0", + "payload_writes": 33, + "payload_bytes": 3784704, + "physical_hop_payload_bytes": 3784704, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 464, + "remote_qp_number": 720, + "local_hca_index": 2, + "remote_hca_index": 3, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 3, + "path_index": 0, + "device": "rocep1s0f0", + "payload_writes": 33, + "payload_bytes": 3784704, + "physical_hop_payload_bytes": 7569408, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 472, + "remote_qp_number": 854, + "local_hca_index": 0, + "remote_hca_index": 1, + "physical_hops": 2, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 3, + "path_index": 1, + "device": "roceP2p1s0f1", + "payload_writes": 33, + "payload_bytes": 3784704, + "physical_hop_payload_bytes": 7569408, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 721, + "remote_qp_number": 464, + "local_hca_index": 3, + "remote_hca_index": 2, + "physical_hops": 2, + "retries": null, + "retry_events": null + } + ], + "expected_counters": { + "operations_per_rank": 33, + "paths_per_rank": 6, + "flags_per_path": 33, + "payload_bytes_path_0": 3784704, + "payload_bytes_path_1": 3784704 + } + }, + { + "rank": 2, + "eager_samples_us": [ + 743.6479926109314, + 56.57599866390228, + 161.43999993801117 + ], + "graph_samples_us": [ + 118.02666385968526, + 74.15466507275899, + 65.13066589832306 + ], + "eager_median_us": 161.43999993801117, + "graph_median_us": 74.15466507275899, + "correctness_cases": [ + { + "name": "rank-specific-index-pattern-a", + "mode": "eager", + "input_formula": "((rank + 1) * ((index % 17) - 8)) + rank", + "expected_formula": "10 * ((index % 17) - 8) + 6", + "input_sha256": "2c6e99a15b93496b40f8cc02f8177cdf5a550336950f15fa6ce0f585102f549a", + "expected_sha256": "93a73648c332a1d4054183a6b8d177e63092ccb0f28cb04743c47cd32da4721c", + "output_sha256": "93a73648c332a1d4054183a6b8d177e63092ccb0f28cb04743c47cd32da4721c", + "passed": true + }, + { + "name": "constant-timing-reference", + "mode": "eager", + "input_formula": "rank + 1", + "expected_formula": "10", + "input_sha256": "8ee6105aea1f6f4488d981245473db1326292aedfac99f7cef071db169f6807c", + "expected_sha256": "1a64b11191799028236a93dc0ffb15330e07140967c672eb0475b5e63682ea87", + "output_sha256": "1a64b11191799028236a93dc0ffb15330e07140967c672eb0475b5e63682ea87", + "passed": true + }, + { + "name": "graph-input-mutation-a", + "mode": "graph-replay", + "input_formula": "((rank + 1) * ((index % 17) - 8)) + rank", + "expected_formula": "10 * ((index % 17) - 8) + 6", + "input_sha256": "2c6e99a15b93496b40f8cc02f8177cdf5a550336950f15fa6ce0f585102f549a", + "expected_sha256": "93a73648c332a1d4054183a6b8d177e63092ccb0f28cb04743c47cd32da4721c", + "output_sha256": "93a73648c332a1d4054183a6b8d177e63092ccb0f28cb04743c47cd32da4721c", + "passed": true + }, + { + "name": "graph-input-mutation-b", + "mode": "graph-replay", + "input_formula": "((4 - rank) * (((index * 5 + 3) % 19) - 9)) - rank", + "expected_formula": "10 * (((index * 5 + 3) % 19) - 9) - 6", + "input_sha256": "855cc661f79e1119c78724e352f7d1ce459668530e9962b22d7b43d892865833", + "expected_sha256": "f2020cf5aee4a14c69ba387ab5f7ecc167433ecc88f795b63cc205821f894c29", + "output_sha256": "f2020cf5aee4a14c69ba387ab5f7ecc167433ecc88f795b63cc205821f894c29", + "passed": true + }, + { + "name": "constant-timing-output", + "mode": "graph-replay", + "input_formula": "rank + 1", + "expected_formula": "10", + "input_sha256": "8ee6105aea1f6f4488d981245473db1326292aedfac99f7cef071db169f6807c", + "expected_sha256": "1a64b11191799028236a93dc0ffb15330e07140967c672eb0475b5e63682ea87", + "output_sha256": "1a64b11191799028236a93dc0ffb15330e07140967c672eb0475b5e63682ea87", + "passed": true + } + ], + "stats": { + "world_size": 4, + "rank": 2, + "hcas": [ + "rocep1s0f0", + "rocep1s0f1", + "roceP2p1s0f0", + "roceP2p1s0f1" + ], + "max_size": 2097152, + "max_gather_bytes": 2097152, + "slot_bytes": 2097152, + "epoch": 33, + "error_seq": 0, + "error_peer": 0, + "ctrl_seq": 33, + "spin_limit": 20000000, + "opposite_paths": 2, + "ops_posted": 33, + "writes_completed": 198, + "last_seq": 33, + "two_wave_activations": 33, + "two_wave_threshold_bytes": 196608, + "wave_mode": "two", + "peer_hca": { + "0": [ + 1, + 2 + ], + "1": [ + 1, + 3 + ], + "3": [ + 0, + 2 + ] + } + }, + "path_counters": [ + { + "peer_rank": 0, + "path_index": 0, + "device": "rocep1s0f1", + "payload_writes": 33, + "payload_bytes": 3784704, + "physical_hop_payload_bytes": 7569408, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 853, + "remote_qp_number": 471, + "local_hca_index": 1, + "remote_hca_index": 0, + "physical_hops": 2, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 0, + "path_index": 1, + "device": "roceP2p1s0f0", + "payload_writes": 33, + "payload_bytes": 3784704, + "physical_hop_payload_bytes": 7569408, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 463, + "remote_qp_number": 720, + "local_hca_index": 2, + "remote_hca_index": 3, + "physical_hops": 2, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 1, + "path_index": 0, + "device": "rocep1s0f1", + "payload_writes": 33, + "payload_bytes": 3784704, + "physical_hop_payload_bytes": 3784704, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 854, + "remote_qp_number": 471, + "local_hca_index": 1, + "remote_hca_index": 0, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 1, + "path_index": 1, + "device": "roceP2p1s0f1", + "payload_writes": 33, + "payload_bytes": 3784704, + "physical_hop_payload_bytes": 3784704, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 720, + "remote_qp_number": 464, + "local_hca_index": 3, + "remote_hca_index": 2, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 3, + "path_index": 0, + "device": "rocep1s0f0", + "payload_writes": 33, + "payload_bytes": 3784704, + "physical_hop_payload_bytes": 3784704, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 472, + "remote_qp_number": 855, + "local_hca_index": 0, + "remote_hca_index": 1, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 3, + "path_index": 1, + "device": "roceP2p1s0f0", + "payload_writes": 33, + "payload_bytes": 3784704, + "physical_hop_payload_bytes": 3784704, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 465, + "remote_qp_number": 721, + "local_hca_index": 2, + "remote_hca_index": 3, + "physical_hops": 1, + "retries": null, + "retry_events": null + } + ], + "expected_counters": { + "operations_per_rank": 33, + "paths_per_rank": 6, + "flags_per_path": 33, + "payload_bytes_path_0": 3784704, + "payload_bytes_path_1": 3784704 + } + }, + { + "rank": 3, + "eager_samples_us": [ + 491.4880096912384, + 124.44800138473511, + 49696.83074951172 + ], + "graph_samples_us": [ + 174.76266622543335, + 74.81599847475688, + 70.66666583220164 + ], + "eager_median_us": 491.4880096912384, + "graph_median_us": 74.81599847475688, + "correctness_cases": [ + { + "name": "rank-specific-index-pattern-a", + "mode": "eager", + "input_formula": "((rank + 1) * ((index % 17) - 8)) + rank", + "expected_formula": "10 * ((index % 17) - 8) + 6", + "input_sha256": "54b6da6dab14daa1d0c60cda432b7128d18131c8850214039c8dda8ab3b47fc3", + "expected_sha256": "93a73648c332a1d4054183a6b8d177e63092ccb0f28cb04743c47cd32da4721c", + "output_sha256": "93a73648c332a1d4054183a6b8d177e63092ccb0f28cb04743c47cd32da4721c", + "passed": true + }, + { + "name": "constant-timing-reference", + "mode": "eager", + "input_formula": "rank + 1", + "expected_formula": "10", + "input_sha256": "6857eae090823b65a353ee31b96b2a5bc27174935c77dd87d8b5b8f40a8385ea", + "expected_sha256": "1a64b11191799028236a93dc0ffb15330e07140967c672eb0475b5e63682ea87", + "output_sha256": "1a64b11191799028236a93dc0ffb15330e07140967c672eb0475b5e63682ea87", + "passed": true + }, + { + "name": "graph-input-mutation-a", + "mode": "graph-replay", + "input_formula": "((rank + 1) * ((index % 17) - 8)) + rank", + "expected_formula": "10 * ((index % 17) - 8) + 6", + "input_sha256": "54b6da6dab14daa1d0c60cda432b7128d18131c8850214039c8dda8ab3b47fc3", + "expected_sha256": "93a73648c332a1d4054183a6b8d177e63092ccb0f28cb04743c47cd32da4721c", + "output_sha256": "93a73648c332a1d4054183a6b8d177e63092ccb0f28cb04743c47cd32da4721c", + "passed": true + }, + { + "name": "graph-input-mutation-b", + "mode": "graph-replay", + "input_formula": "((4 - rank) * (((index * 5 + 3) % 19) - 9)) - rank", + "expected_formula": "10 * (((index * 5 + 3) % 19) - 9) - 6", + "input_sha256": "e798f851c47d8673b2cab4cfa59c533604e990164d883b2393caee06d3f1eb5e", + "expected_sha256": "f2020cf5aee4a14c69ba387ab5f7ecc167433ecc88f795b63cc205821f894c29", + "output_sha256": "f2020cf5aee4a14c69ba387ab5f7ecc167433ecc88f795b63cc205821f894c29", + "passed": true + }, + { + "name": "constant-timing-output", + "mode": "graph-replay", + "input_formula": "rank + 1", + "expected_formula": "10", + "input_sha256": "6857eae090823b65a353ee31b96b2a5bc27174935c77dd87d8b5b8f40a8385ea", + "expected_sha256": "1a64b11191799028236a93dc0ffb15330e07140967c672eb0475b5e63682ea87", + "output_sha256": "1a64b11191799028236a93dc0ffb15330e07140967c672eb0475b5e63682ea87", + "passed": true + } + ], + "stats": { + "world_size": 4, + "rank": 3, + "hcas": [ + "rocep1s0f0", + "rocep1s0f1", + "roceP2p1s0f0", + "roceP2p1s0f1" + ], + "max_size": 2097152, + "max_gather_bytes": 2097152, + "slot_bytes": 2097152, + "epoch": 33, + "error_seq": 0, + "error_peer": 0, + "ctrl_seq": 33, + "spin_limit": 20000000, + "opposite_paths": 2, + "ops_posted": 33, + "writes_completed": 198, + "last_seq": 33, + "two_wave_activations": 33, + "two_wave_threshold_bytes": 196608, + "wave_mode": "two", + "peer_hca": { + "0": [ + 0, + 2 + ], + "1": [ + 1, + 2 + ], + "2": [ + 1, + 3 + ] + } + }, + "path_counters": [ + { + "peer_rank": 0, + "path_index": 0, + "device": "rocep1s0f0", + "payload_writes": 33, + "payload_bytes": 3784704, + "physical_hop_payload_bytes": 3784704, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 470, + "remote_qp_number": 855, + "local_hca_index": 0, + "remote_hca_index": 1, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 0, + "path_index": 1, + "device": "roceP2p1s0f0", + "payload_writes": 33, + "payload_bytes": 3784704, + "physical_hop_payload_bytes": 3784704, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 463, + "remote_qp_number": 721, + "local_hca_index": 2, + "remote_hca_index": 3, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 1, + "path_index": 0, + "device": "rocep1s0f1", + "payload_writes": 33, + "payload_bytes": 3784704, + "physical_hop_payload_bytes": 7569408, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 854, + "remote_qp_number": 472, + "local_hca_index": 1, + "remote_hca_index": 0, + "physical_hops": 2, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 1, + "path_index": 1, + "device": "roceP2p1s0f0", + "payload_writes": 33, + "payload_bytes": 3784704, + "physical_hop_payload_bytes": 7569408, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 464, + "remote_qp_number": 721, + "local_hca_index": 2, + "remote_hca_index": 3, + "physical_hops": 2, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 2, + "path_index": 0, + "device": "rocep1s0f1", + "payload_writes": 33, + "payload_bytes": 3784704, + "physical_hop_payload_bytes": 3784704, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 855, + "remote_qp_number": 472, + "local_hca_index": 1, + "remote_hca_index": 0, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 2, + "path_index": 1, + "device": "roceP2p1s0f1", + "payload_writes": 33, + "payload_bytes": 3784704, + "physical_hop_payload_bytes": 3784704, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 721, + "remote_qp_number": 465, + "local_hca_index": 3, + "remote_hca_index": 2, + "physical_hops": 1, + "retries": null, + "retry_events": null + } + ], + "expected_counters": { + "operations_per_rank": 33, + "paths_per_rank": 6, + "flags_per_path": 33, + "payload_bytes_path_0": 3784704, + "payload_bytes_path_1": 3784704 + } + } + ] + }, + { + "schema": "b12x.rocenante-virtual-diagonal-evidence/v1", + "status": "research-only", + "payload_bytes": 524288, + "warmups": 2, + "samples": 3, + "graph_operations_per_replay": 3, + "ranks": [ + { + "rank": 0, + "eager_samples_us": [ + 504.9920082092285, + 100.03200173377991, + 247.39199876785278 + ], + "graph_samples_us": [ + 196.58666849136353, + 198.27200969060263, + 114.54932888348897 + ], + "eager_median_us": 247.39199876785278, + "graph_median_us": 196.58666849136353, + "correctness_cases": [ + { + "name": "rank-specific-index-pattern-a", + "mode": "eager", + "input_formula": "((rank + 1) * ((index % 17) - 8)) + rank", + "expected_formula": "10 * ((index % 17) - 8) + 6", + "input_sha256": "5454745e573fa5778e29b37d413f7de0e88b9dc4e8fa8d57d2e94de4852ffbbb", + "expected_sha256": "510b1e3d9f4094f5df8c9eb19ab8c4c47e810dc028bd479ba4b64c86dadb3df8", + "output_sha256": "510b1e3d9f4094f5df8c9eb19ab8c4c47e810dc028bd479ba4b64c86dadb3df8", + "passed": true + }, + { + "name": "constant-timing-reference", + "mode": "eager", + "input_formula": "rank + 1", + "expected_formula": "10", + "input_sha256": "c5b0030d4b73d266c9c673643eab9fc1dc4e86d42add28efe442304139077aad", + "expected_sha256": "33fc9307e1ef3253c401b96f932255aab1a8d3d7a317168e9eded52d2dc7bc37", + "output_sha256": "33fc9307e1ef3253c401b96f932255aab1a8d3d7a317168e9eded52d2dc7bc37", + "passed": true + }, + { + "name": "graph-input-mutation-a", + "mode": "graph-replay", + "input_formula": "((rank + 1) * ((index % 17) - 8)) + rank", + "expected_formula": "10 * ((index % 17) - 8) + 6", + "input_sha256": "5454745e573fa5778e29b37d413f7de0e88b9dc4e8fa8d57d2e94de4852ffbbb", + "expected_sha256": "510b1e3d9f4094f5df8c9eb19ab8c4c47e810dc028bd479ba4b64c86dadb3df8", + "output_sha256": "510b1e3d9f4094f5df8c9eb19ab8c4c47e810dc028bd479ba4b64c86dadb3df8", + "passed": true + }, + { + "name": "graph-input-mutation-b", + "mode": "graph-replay", + "input_formula": "((4 - rank) * (((index * 5 + 3) % 19) - 9)) - rank", + "expected_formula": "10 * (((index * 5 + 3) % 19) - 9) - 6", + "input_sha256": "bedbe8f86cf6239d446f137a9bd98e58556d555e60472e926749847adc960b69", + "expected_sha256": "efc2693cd87c395d15cc8e5b6414ffd3dedad9d70217a88b46599ca34aa54395", + "output_sha256": "efc2693cd87c395d15cc8e5b6414ffd3dedad9d70217a88b46599ca34aa54395", + "passed": true + }, + { + "name": "constant-timing-output", + "mode": "graph-replay", + "input_formula": "rank + 1", + "expected_formula": "10", + "input_sha256": "c5b0030d4b73d266c9c673643eab9fc1dc4e86d42add28efe442304139077aad", + "expected_sha256": "33fc9307e1ef3253c401b96f932255aab1a8d3d7a317168e9eded52d2dc7bc37", + "output_sha256": "33fc9307e1ef3253c401b96f932255aab1a8d3d7a317168e9eded52d2dc7bc37", + "passed": true + } + ], + "stats": { + "world_size": 4, + "rank": 0, + "hcas": [ + "rocep1s0f0", + "rocep1s0f1", + "roceP2p1s0f0", + "roceP2p1s0f1" + ], + "max_size": 2097152, + "max_gather_bytes": 2097152, + "slot_bytes": 2097152, + "epoch": 33, + "error_seq": 0, + "error_peer": 0, + "ctrl_seq": 33, + "spin_limit": 20000000, + "opposite_paths": 2, + "ops_posted": 33, + "writes_completed": 198, + "last_seq": 33, + "two_wave_activations": 33, + "two_wave_threshold_bytes": 196608, + "wave_mode": "two", + "peer_hca": { + "1": [ + 0, + 2 + ], + "2": [ + 0, + 3 + ], + "3": [ + 1, + 3 + ] + } + }, + "path_counters": [ + { + "peer_rank": 1, + "path_index": 0, + "device": "rocep1s0f0", + "payload_writes": 33, + "payload_bytes": 8650752, + "physical_hop_payload_bytes": 8650752, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 473, + "remote_qp_number": 856, + "local_hca_index": 0, + "remote_hca_index": 1, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 1, + "path_index": 1, + "device": "roceP2p1s0f0", + "payload_writes": 33, + "payload_bytes": 8650752, + "physical_hop_payload_bytes": 8650752, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 466, + "remote_qp_number": 722, + "local_hca_index": 2, + "remote_hca_index": 3, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 2, + "path_index": 0, + "device": "rocep1s0f0", + "payload_writes": 33, + "payload_bytes": 8650752, + "physical_hop_payload_bytes": 17301504, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 474, + "remote_qp_number": 856, + "local_hca_index": 0, + "remote_hca_index": 1, + "physical_hops": 2, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 2, + "path_index": 1, + "device": "roceP2p1s0f1", + "payload_writes": 33, + "payload_bytes": 8650752, + "physical_hop_payload_bytes": 17301504, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 723, + "remote_qp_number": 466, + "local_hca_index": 3, + "remote_hca_index": 2, + "physical_hops": 2, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 3, + "path_index": 0, + "device": "rocep1s0f1", + "payload_writes": 33, + "payload_bytes": 8650752, + "physical_hop_payload_bytes": 8650752, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 858, + "remote_qp_number": 473, + "local_hca_index": 1, + "remote_hca_index": 0, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 3, + "path_index": 1, + "device": "roceP2p1s0f1", + "payload_writes": 33, + "payload_bytes": 8650752, + "physical_hop_payload_bytes": 8650752, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 724, + "remote_qp_number": 466, + "local_hca_index": 3, + "remote_hca_index": 2, + "physical_hops": 1, + "retries": null, + "retry_events": null + } + ], + "expected_counters": { + "operations_per_rank": 33, + "paths_per_rank": 6, + "flags_per_path": 33, + "payload_bytes_path_0": 8650752, + "payload_bytes_path_1": 8650752 + } + }, + { + "rank": 1, + "eager_samples_us": [ + 1352.4800539016724, + 206.496000289917, + 169.3120002746582 + ], + "graph_samples_us": [ + 112.61866490046184, + 200.91732343037924, + 23482.90252685547 + ], + "eager_median_us": 206.496000289917, + "graph_median_us": 200.91732343037924, + "correctness_cases": [ + { + "name": "rank-specific-index-pattern-a", + "mode": "eager", + "input_formula": "((rank + 1) * ((index % 17) - 8)) + rank", + "expected_formula": "10 * ((index % 17) - 8) + 6", + "input_sha256": "f31670bd652760a135fc801b266d512d7a9011d05c6f3716807777ff5b49a1dd", + "expected_sha256": "510b1e3d9f4094f5df8c9eb19ab8c4c47e810dc028bd479ba4b64c86dadb3df8", + "output_sha256": "510b1e3d9f4094f5df8c9eb19ab8c4c47e810dc028bd479ba4b64c86dadb3df8", + "passed": true + }, + { + "name": "constant-timing-reference", + "mode": "eager", + "input_formula": "rank + 1", + "expected_formula": "10", + "input_sha256": "fc536e1d3546dd8630ac053c00be0b69b3c1444b8087e7eb9503d1fc0fb13e0b", + "expected_sha256": "33fc9307e1ef3253c401b96f932255aab1a8d3d7a317168e9eded52d2dc7bc37", + "output_sha256": "33fc9307e1ef3253c401b96f932255aab1a8d3d7a317168e9eded52d2dc7bc37", + "passed": true + }, + { + "name": "graph-input-mutation-a", + "mode": "graph-replay", + "input_formula": "((rank + 1) * ((index % 17) - 8)) + rank", + "expected_formula": "10 * ((index % 17) - 8) + 6", + "input_sha256": "f31670bd652760a135fc801b266d512d7a9011d05c6f3716807777ff5b49a1dd", + "expected_sha256": "510b1e3d9f4094f5df8c9eb19ab8c4c47e810dc028bd479ba4b64c86dadb3df8", + "output_sha256": "510b1e3d9f4094f5df8c9eb19ab8c4c47e810dc028bd479ba4b64c86dadb3df8", + "passed": true + }, + { + "name": "graph-input-mutation-b", + "mode": "graph-replay", + "input_formula": "((4 - rank) * (((index * 5 + 3) % 19) - 9)) - rank", + "expected_formula": "10 * (((index * 5 + 3) % 19) - 9) - 6", + "input_sha256": "cda0cfab90e33789a8d5560893aa2a659cecc52259677456582e47ccdfcb32af", + "expected_sha256": "efc2693cd87c395d15cc8e5b6414ffd3dedad9d70217a88b46599ca34aa54395", + "output_sha256": "efc2693cd87c395d15cc8e5b6414ffd3dedad9d70217a88b46599ca34aa54395", + "passed": true + }, + { + "name": "constant-timing-output", + "mode": "graph-replay", + "input_formula": "rank + 1", + "expected_formula": "10", + "input_sha256": "fc536e1d3546dd8630ac053c00be0b69b3c1444b8087e7eb9503d1fc0fb13e0b", + "expected_sha256": "33fc9307e1ef3253c401b96f932255aab1a8d3d7a317168e9eded52d2dc7bc37", + "output_sha256": "33fc9307e1ef3253c401b96f932255aab1a8d3d7a317168e9eded52d2dc7bc37", + "passed": true + } + ], + "stats": { + "world_size": 4, + "rank": 1, + "hcas": [ + "rocep1s0f0", + "rocep1s0f1", + "roceP2p1s0f0", + "roceP2p1s0f1" + ], + "max_size": 2097152, + "max_gather_bytes": 2097152, + "slot_bytes": 2097152, + "epoch": 33, + "error_seq": 0, + "error_peer": 0, + "ctrl_seq": 33, + "spin_limit": 20000000, + "opposite_paths": 2, + "ops_posted": 33, + "writes_completed": 198, + "last_seq": 33, + "two_wave_activations": 33, + "two_wave_threshold_bytes": 196608, + "wave_mode": "two", + "peer_hca": { + "0": [ + 1, + 3 + ], + "2": [ + 0, + 2 + ], + "3": [ + 0, + 3 + ] + } + }, + "path_counters": [ + { + "peer_rank": 0, + "path_index": 0, + "device": "rocep1s0f1", + "payload_writes": 33, + "payload_bytes": 8650752, + "physical_hop_payload_bytes": 8650752, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 856, + "remote_qp_number": 473, + "local_hca_index": 1, + "remote_hca_index": 0, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 0, + "path_index": 1, + "device": "roceP2p1s0f1", + "payload_writes": 33, + "payload_bytes": 8650752, + "physical_hop_payload_bytes": 8650752, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 722, + "remote_qp_number": 466, + "local_hca_index": 3, + "remote_hca_index": 2, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 2, + "path_index": 0, + "device": "rocep1s0f0", + "payload_writes": 33, + "payload_bytes": 8650752, + "physical_hop_payload_bytes": 8650752, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 474, + "remote_qp_number": 857, + "local_hca_index": 0, + "remote_hca_index": 1, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 2, + "path_index": 1, + "device": "roceP2p1s0f0", + "payload_writes": 33, + "payload_bytes": 8650752, + "physical_hop_payload_bytes": 8650752, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 467, + "remote_qp_number": 723, + "local_hca_index": 2, + "remote_hca_index": 3, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 3, + "path_index": 0, + "device": "rocep1s0f0", + "payload_writes": 33, + "payload_bytes": 8650752, + "physical_hop_payload_bytes": 17301504, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 475, + "remote_qp_number": 857, + "local_hca_index": 0, + "remote_hca_index": 1, + "physical_hops": 2, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 3, + "path_index": 1, + "device": "roceP2p1s0f1", + "payload_writes": 33, + "payload_bytes": 8650752, + "physical_hop_payload_bytes": 17301504, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 724, + "remote_qp_number": 467, + "local_hca_index": 3, + "remote_hca_index": 2, + "physical_hops": 2, + "retries": null, + "retry_events": null + } + ], + "expected_counters": { + "operations_per_rank": 33, + "paths_per_rank": 6, + "flags_per_path": 33, + "payload_bytes_path_0": 8650752, + "payload_bytes_path_1": 8650752 + } + }, + { + "rank": 2, + "eager_samples_us": [ + 285.8560085296631, + 321.4080035686493, + 117.50400066375732 + ], + "graph_samples_us": [ + 360.21331946055096, + 129.9199958642324, + 162.8053287665049 + ], + "eager_median_us": 285.8560085296631, + "graph_median_us": 162.8053287665049, + "correctness_cases": [ + { + "name": "rank-specific-index-pattern-a", + "mode": "eager", + "input_formula": "((rank + 1) * ((index % 17) - 8)) + rank", + "expected_formula": "10 * ((index % 17) - 8) + 6", + "input_sha256": "f47b91b28cd1441d1cdece5388c65b3d168d1e908f6d7c1174dbbb45e707cb04", + "expected_sha256": "510b1e3d9f4094f5df8c9eb19ab8c4c47e810dc028bd479ba4b64c86dadb3df8", + "output_sha256": "510b1e3d9f4094f5df8c9eb19ab8c4c47e810dc028bd479ba4b64c86dadb3df8", + "passed": true + }, + { + "name": "constant-timing-reference", + "mode": "eager", + "input_formula": "rank + 1", + "expected_formula": "10", + "input_sha256": "36ecc3fa4521e727d5f750f398b915b33e833db4288e6cab028fd272d8d08180", + "expected_sha256": "33fc9307e1ef3253c401b96f932255aab1a8d3d7a317168e9eded52d2dc7bc37", + "output_sha256": "33fc9307e1ef3253c401b96f932255aab1a8d3d7a317168e9eded52d2dc7bc37", + "passed": true + }, + { + "name": "graph-input-mutation-a", + "mode": "graph-replay", + "input_formula": "((rank + 1) * ((index % 17) - 8)) + rank", + "expected_formula": "10 * ((index % 17) - 8) + 6", + "input_sha256": "f47b91b28cd1441d1cdece5388c65b3d168d1e908f6d7c1174dbbb45e707cb04", + "expected_sha256": "510b1e3d9f4094f5df8c9eb19ab8c4c47e810dc028bd479ba4b64c86dadb3df8", + "output_sha256": "510b1e3d9f4094f5df8c9eb19ab8c4c47e810dc028bd479ba4b64c86dadb3df8", + "passed": true + }, + { + "name": "graph-input-mutation-b", + "mode": "graph-replay", + "input_formula": "((4 - rank) * (((index * 5 + 3) % 19) - 9)) - rank", + "expected_formula": "10 * (((index * 5 + 3) % 19) - 9) - 6", + "input_sha256": "ce005080013f27240ee25d161bbbd04d3328b5988715ad979c9c1251c4256048", + "expected_sha256": "efc2693cd87c395d15cc8e5b6414ffd3dedad9d70217a88b46599ca34aa54395", + "output_sha256": "efc2693cd87c395d15cc8e5b6414ffd3dedad9d70217a88b46599ca34aa54395", + "passed": true + }, + { + "name": "constant-timing-output", + "mode": "graph-replay", + "input_formula": "rank + 1", + "expected_formula": "10", + "input_sha256": "36ecc3fa4521e727d5f750f398b915b33e833db4288e6cab028fd272d8d08180", + "expected_sha256": "33fc9307e1ef3253c401b96f932255aab1a8d3d7a317168e9eded52d2dc7bc37", + "output_sha256": "33fc9307e1ef3253c401b96f932255aab1a8d3d7a317168e9eded52d2dc7bc37", + "passed": true + } + ], + "stats": { + "world_size": 4, + "rank": 2, + "hcas": [ + "rocep1s0f0", + "rocep1s0f1", + "roceP2p1s0f0", + "roceP2p1s0f1" + ], + "max_size": 2097152, + "max_gather_bytes": 2097152, + "slot_bytes": 2097152, + "epoch": 33, + "error_seq": 0, + "error_peer": 0, + "ctrl_seq": 33, + "spin_limit": 20000000, + "opposite_paths": 2, + "ops_posted": 33, + "writes_completed": 198, + "last_seq": 33, + "two_wave_activations": 33, + "two_wave_threshold_bytes": 196608, + "wave_mode": "two", + "peer_hca": { + "0": [ + 1, + 2 + ], + "1": [ + 1, + 3 + ], + "3": [ + 0, + 2 + ] + } + }, + "path_counters": [ + { + "peer_rank": 0, + "path_index": 0, + "device": "rocep1s0f1", + "payload_writes": 33, + "payload_bytes": 8650752, + "physical_hop_payload_bytes": 17301504, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 856, + "remote_qp_number": 474, + "local_hca_index": 1, + "remote_hca_index": 0, + "physical_hops": 2, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 0, + "path_index": 1, + "device": "roceP2p1s0f0", + "payload_writes": 33, + "payload_bytes": 8650752, + "physical_hop_payload_bytes": 17301504, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 466, + "remote_qp_number": 723, + "local_hca_index": 2, + "remote_hca_index": 3, + "physical_hops": 2, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 1, + "path_index": 0, + "device": "rocep1s0f1", + "payload_writes": 33, + "payload_bytes": 8650752, + "physical_hop_payload_bytes": 8650752, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 857, + "remote_qp_number": 474, + "local_hca_index": 1, + "remote_hca_index": 0, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 1, + "path_index": 1, + "device": "roceP2p1s0f1", + "payload_writes": 33, + "payload_bytes": 8650752, + "physical_hop_payload_bytes": 8650752, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 723, + "remote_qp_number": 467, + "local_hca_index": 3, + "remote_hca_index": 2, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 3, + "path_index": 0, + "device": "rocep1s0f0", + "payload_writes": 33, + "payload_bytes": 8650752, + "physical_hop_payload_bytes": 8650752, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 475, + "remote_qp_number": 858, + "local_hca_index": 0, + "remote_hca_index": 1, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 3, + "path_index": 1, + "device": "roceP2p1s0f0", + "payload_writes": 33, + "payload_bytes": 8650752, + "physical_hop_payload_bytes": 8650752, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 468, + "remote_qp_number": 724, + "local_hca_index": 2, + "remote_hca_index": 3, + "physical_hops": 1, + "retries": null, + "retry_events": null + } + ], + "expected_counters": { + "operations_per_rank": 33, + "paths_per_rank": 6, + "flags_per_path": 33, + "payload_bytes_path_0": 8650752, + "payload_bytes_path_1": 8650752 + } + }, + { + "rank": 3, + "eager_samples_us": [ + 1071.903944015503, + 188.6720061302185, + 5829.472064971924 + ], + "graph_samples_us": [ + 237.06666628519693, + 124.56533312797546, + 170.66667477289835 + ], + "eager_median_us": 1071.903944015503, + "graph_median_us": 170.66667477289835, + "correctness_cases": [ + { + "name": "rank-specific-index-pattern-a", + "mode": "eager", + "input_formula": "((rank + 1) * ((index % 17) - 8)) + rank", + "expected_formula": "10 * ((index % 17) - 8) + 6", + "input_sha256": "db58c857abdd08c9f5f79b6d74e6f9250366a29761fb62b5a2e0593a0d1e76c7", + "expected_sha256": "510b1e3d9f4094f5df8c9eb19ab8c4c47e810dc028bd479ba4b64c86dadb3df8", + "output_sha256": "510b1e3d9f4094f5df8c9eb19ab8c4c47e810dc028bd479ba4b64c86dadb3df8", + "passed": true + }, + { + "name": "constant-timing-reference", + "mode": "eager", + "input_formula": "rank + 1", + "expected_formula": "10", + "input_sha256": "692d6abae8f5b17bb4fb00d1c72ac0b84dc21c09552d9c78834d730fc3a920c4", + "expected_sha256": "33fc9307e1ef3253c401b96f932255aab1a8d3d7a317168e9eded52d2dc7bc37", + "output_sha256": "33fc9307e1ef3253c401b96f932255aab1a8d3d7a317168e9eded52d2dc7bc37", + "passed": true + }, + { + "name": "graph-input-mutation-a", + "mode": "graph-replay", + "input_formula": "((rank + 1) * ((index % 17) - 8)) + rank", + "expected_formula": "10 * ((index % 17) - 8) + 6", + "input_sha256": "db58c857abdd08c9f5f79b6d74e6f9250366a29761fb62b5a2e0593a0d1e76c7", + "expected_sha256": "510b1e3d9f4094f5df8c9eb19ab8c4c47e810dc028bd479ba4b64c86dadb3df8", + "output_sha256": "510b1e3d9f4094f5df8c9eb19ab8c4c47e810dc028bd479ba4b64c86dadb3df8", + "passed": true + }, + { + "name": "graph-input-mutation-b", + "mode": "graph-replay", + "input_formula": "((4 - rank) * (((index * 5 + 3) % 19) - 9)) - rank", + "expected_formula": "10 * (((index * 5 + 3) % 19) - 9) - 6", + "input_sha256": "1808a13e28838bc253caf156b3c19c85889232790076f5f4a33eea2f2b40b9b8", + "expected_sha256": "efc2693cd87c395d15cc8e5b6414ffd3dedad9d70217a88b46599ca34aa54395", + "output_sha256": "efc2693cd87c395d15cc8e5b6414ffd3dedad9d70217a88b46599ca34aa54395", + "passed": true + }, + { + "name": "constant-timing-output", + "mode": "graph-replay", + "input_formula": "rank + 1", + "expected_formula": "10", + "input_sha256": "692d6abae8f5b17bb4fb00d1c72ac0b84dc21c09552d9c78834d730fc3a920c4", + "expected_sha256": "33fc9307e1ef3253c401b96f932255aab1a8d3d7a317168e9eded52d2dc7bc37", + "output_sha256": "33fc9307e1ef3253c401b96f932255aab1a8d3d7a317168e9eded52d2dc7bc37", + "passed": true + } + ], + "stats": { + "world_size": 4, + "rank": 3, + "hcas": [ + "rocep1s0f0", + "rocep1s0f1", + "roceP2p1s0f0", + "roceP2p1s0f1" + ], + "max_size": 2097152, + "max_gather_bytes": 2097152, + "slot_bytes": 2097152, + "epoch": 33, + "error_seq": 0, + "error_peer": 0, + "ctrl_seq": 33, + "spin_limit": 20000000, + "opposite_paths": 2, + "ops_posted": 33, + "writes_completed": 198, + "last_seq": 33, + "two_wave_activations": 33, + "two_wave_threshold_bytes": 196608, + "wave_mode": "two", + "peer_hca": { + "0": [ + 0, + 2 + ], + "1": [ + 1, + 2 + ], + "2": [ + 1, + 3 + ] + } + }, + "path_counters": [ + { + "peer_rank": 0, + "path_index": 0, + "device": "rocep1s0f0", + "payload_writes": 33, + "payload_bytes": 8650752, + "physical_hop_payload_bytes": 8650752, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 473, + "remote_qp_number": 858, + "local_hca_index": 0, + "remote_hca_index": 1, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 0, + "path_index": 1, + "device": "roceP2p1s0f0", + "payload_writes": 33, + "payload_bytes": 8650752, + "physical_hop_payload_bytes": 8650752, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 466, + "remote_qp_number": 724, + "local_hca_index": 2, + "remote_hca_index": 3, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 1, + "path_index": 0, + "device": "rocep1s0f1", + "payload_writes": 33, + "payload_bytes": 8650752, + "physical_hop_payload_bytes": 17301504, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 857, + "remote_qp_number": 475, + "local_hca_index": 1, + "remote_hca_index": 0, + "physical_hops": 2, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 1, + "path_index": 1, + "device": "roceP2p1s0f0", + "payload_writes": 33, + "payload_bytes": 8650752, + "physical_hop_payload_bytes": 17301504, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 467, + "remote_qp_number": 724, + "local_hca_index": 2, + "remote_hca_index": 3, + "physical_hops": 2, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 2, + "path_index": 0, + "device": "rocep1s0f1", + "payload_writes": 33, + "payload_bytes": 8650752, + "physical_hop_payload_bytes": 8650752, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 858, + "remote_qp_number": 475, + "local_hca_index": 1, + "remote_hca_index": 0, + "physical_hops": 1, + "retries": null, + "retry_events": null + }, + { + "peer_rank": 2, + "path_index": 1, + "device": "roceP2p1s0f1", + "payload_writes": 33, + "payload_bytes": 8650752, + "physical_hop_payload_bytes": 8650752, + "flag_writes": 33, + "send_completions": 33, + "completion_errors": 0, + "qp_number": 724, + "remote_qp_number": 468, + "local_hca_index": 3, + "remote_hca_index": 2, + "physical_hops": 1, + "retries": null, + "retry_events": null + } + ], + "expected_counters": { + "operations_per_rank": 33, + "paths_per_rank": 6, + "flags_per_path": 33, + "payload_bytes_path_0": 8650752, + "payload_bytes_path_1": 8650752 + } + } + ] + } + ], + "streams": [ + { + "checks": [ + { + "rank": 0, + "payload_bytes": 32768, + "alternating_stream_calls": 16, + "misaligned_input_output": true, + "changed_input_graph_replays": 2, + "capture_stream_rejected": true, + "passed": true + }, + { + "rank": 1, + "payload_bytes": 32768, + "alternating_stream_calls": 16, + "misaligned_input_output": true, + "changed_input_graph_replays": 2, + "capture_stream_rejected": true, + "passed": true + }, + { + "rank": 2, + "payload_bytes": 32768, + "alternating_stream_calls": 16, + "misaligned_input_output": true, + "changed_input_graph_replays": 2, + "capture_stream_rejected": true, + "passed": true + }, + { + "rank": 3, + "payload_bytes": 32768, + "alternating_stream_calls": 16, + "misaligned_input_output": true, + "changed_input_graph_replays": 2, + "capture_stream_rejected": true, + "passed": true + } + ], + "passed": true + }, + { + "checks": [ + { + "rank": 0, + "payload_bytes": 524288, + "alternating_stream_calls": 16, + "misaligned_input_output": true, + "changed_input_graph_replays": 2, + "capture_stream_rejected": true, + "passed": true + }, + { + "rank": 1, + "payload_bytes": 524288, + "alternating_stream_calls": 16, + "misaligned_input_output": true, + "changed_input_graph_replays": 2, + "capture_stream_rejected": true, + "passed": true + }, + { + "rank": 2, + "payload_bytes": 524288, + "alternating_stream_calls": 16, + "misaligned_input_output": true, + "changed_input_graph_replays": 2, + "capture_stream_rejected": true, + "passed": true + }, + { + "rank": 3, + "payload_bytes": 524288, + "alternating_stream_calls": 16, + "misaligned_input_output": true, + "changed_input_graph_replays": 2, + "capture_stream_rejected": true, + "passed": true + } + ], + "passed": true + } + ], + "startup": { + "ready": true, + "seconds": 437.3279999999995 + }, + "restart": { + "ready": true, + "seconds": 289.6409999999996 + }, + "idle_rank_loss": { + "injected_rank": 3, + "passed": true, + "all_model_ranks_stopped_seconds": 10.063000000000102, + "scope": "No in-flight request, link loss, or stalled GPU collective was injected." + }, + "persistent_cache": { + "prompt_sha256": "3d2bc5228895566b1497e6f35f6c5aa051685f99438f27226134dfcfab15c277", + "prompt_tokens": 27274, + "external_hit_tokens": 26624, + "all_rank_restore_tokens": [ + 26624, + 26624, + 26624, + 26624 + ], + "semantic_passed": true, + "expected_answer": "cobalt orchard lantern", + "cold_request_seconds": 11.85899999999856, + "restored_request_seconds": 2.9219999999986612, + "persistent_restore_proven": true + }, + "matrices": [ + { + "receipt_sha256": "49bcec98b5e39d10a4493a1056c9ac7589a9caeba47c0f64fada18f586f7f1fa", + "receipt_name": "combined-matrix.json", + "metadata": { + "version": "0.4.32", + "model": "glm-5.3-flash-spark", + "temperature": 1.0, + "max_tokens": 2048, + "duration_per_test": 20.0, + "context_lengths": [ + 8192, + 32768, + 65536 + ], + "concurrency_levels": [ + 1, + 2, + 4, + 8, + 12, + 16 + ] + }, + "prefill": { + "8192": { + "tok_per_sec": 2648.0, + "ttft_seconds": 3.094, + "method": "integrated_scout", + "samples": 1 + }, + "32768": { + "tok_per_sec": 2745.0, + "ttft_seconds": 11.938, + "method": "integrated_scout", + "samples": 1 + }, + "65536": { + "tok_per_sec": 2687.0, + "ttft_seconds": 24.391, + "method": "integrated_scout", + "samples": 1 + }, + "131072": { + "tok_per_sec": 2750.0, + "ttft_seconds": 47.656, + "method": "scout_only", + "samples": 1 + } + }, + "cells": [ + { + "context_tokens": 8192, + "concurrency": 1, + "aggregate_tps": 50.728629375534254, + "server_steps_per_s": 18.578797135560915, + "server_spec_accept_length": 2.730458221024259, + "num_errors": 0 + }, + { + "context_tokens": 32768, + "concurrency": 1, + "aggregate_tps": 45.68654923939065, + "server_steps_per_s": 17.163730984787506, + "server_spec_accept_length": 2.6618075801749272, + "num_errors": 0 + }, + { + "context_tokens": 65536, + "concurrency": 1, + "aggregate_tps": 49.53715286464704, + "server_steps_per_s": 18.864148111082763, + "server_spec_accept_length": 2.6259946949602124, + "num_errors": 0 + }, + { + "context_tokens": 8192, + "concurrency": 2, + "aggregate_tps": 74.5752518418301, + "server_steps_per_s": 28.29061162364332, + "server_spec_accept_length": 2.636042402826855, + "num_errors": 0 + }, + { + "context_tokens": 8192, + "concurrency": 4, + "aggregate_tps": 117.88841631223075, + "server_steps_per_s": 42.03152364273083, + "server_spec_accept_length": 2.8047619047619046, + "num_errors": 0 + }, + { + "context_tokens": 8192, + "concurrency": 8, + "aggregate_tps": 181.0396634615316, + "server_steps_per_s": 64.50320512820268, + "server_spec_accept_length": 2.8066770186335406, + "num_errors": 0 + }, + { + "context_tokens": 8192, + "concurrency": 12, + "aggregate_tps": 200.85, + "server_steps_per_s": 71.76774193548387, + "server_spec_accept_length": 2.798611111111111, + "num_errors": 0 + }, + { + "context_tokens": 8192, + "concurrency": 16, + "aggregate_tps": 230.49165539017244, + "server_steps_per_s": 83.39598055430461, + "server_spec_accept_length": 2.7638221153846154, + "num_errors": 0 + }, + { + "context_tokens": 32768, + "concurrency": 2, + "aggregate_tps": 78.83526286774108, + "server_steps_per_s": 28.968074976194753, + "server_spec_accept_length": 2.7214532871972317, + "num_errors": 0 + }, + { + "context_tokens": 32768, + "concurrency": 4, + "aggregate_tps": 120.3, + "server_steps_per_s": 44.0, + "server_spec_accept_length": 2.7340909090909093, + "num_errors": 0 + }, + { + "context_tokens": 32768, + "concurrency": 8, + "aggregate_tps": 178.8907866974868, + "server_steps_per_s": 63.90033155832721, + "server_spec_accept_length": 2.7995283018867925, + "num_errors": 0 + }, + { + "context_tokens": 32768, + "concurrency": 12, + "aggregate_tps": 204.42433383610452, + "server_steps_per_s": 73.60482654600517, + "server_spec_accept_length": 2.777322404371585, + "num_errors": 0 + }, + { + "context_tokens": 32768, + "concurrency": 16, + "aggregate_tps": 229.71952263456936, + "server_steps_per_s": 82.17936451986348, + "server_spec_accept_length": 2.795343137254902, + "num_errors": 0 + }, + { + "context_tokens": 65536, + "concurrency": 2, + "aggregate_tps": 72.30784627702025, + "server_steps_per_s": 28.82305844675686, + "server_spec_accept_length": 2.5086805555555554, + "num_errors": 0 + }, + { + "context_tokens": 65536, + "concurrency": 4, + "aggregate_tps": 122.33962642094659, + "server_steps_per_s": 43.267063949119056, + "server_spec_accept_length": 2.8275462962962963, + "num_errors": 0 + }, + { + "context_tokens": 65536, + "concurrency": 8, + "aggregate_tps": 170.77808356268258, + "server_steps_per_s": 62.046534901179726, + "server_spec_accept_length": 2.7524193548387097, + "num_errors": 0 + }, + { + "context_tokens": 65536, + "concurrency": 12, + "aggregate_tps": 205.38364779874215, + "server_steps_per_s": 73.05660377358491, + "server_spec_accept_length": 2.8112947658402203, + "num_errors": 0 + }, + { + "context_tokens": 65536, + "concurrency": 16, + "aggregate_tps": 235.47508573734683, + "server_steps_per_s": 82.30784748840223, + "server_spec_accept_length": 2.860906862745098, + "num_errors": 0 + } + ] + }, + { + "receipt_sha256": "a5cd0abe72a3bfd7035e9019fc4152ad9279bb6cd6edc4795c69d081c3b6882b", + "receipt_name": "combined-confirmation-1.json", + "metadata": { + "version": "0.4.32", + "model": "glm-5.3-flash-spark", + "temperature": 1.0, + "max_tokens": 2048, + "duration_per_test": 20.0, + "context_lengths": [ + 8192, + 65536 + ], + "concurrency_levels": [ + 1, + 2 + ] + }, + "prefill": { + "8192": { + "tok_per_sec": 2661.0, + "ttft_seconds": 3.078, + "method": "integrated_scout", + "samples": 1 + }, + "65536": { + "tok_per_sec": 2770.0, + "ttft_seconds": 23.656, + "method": "integrated_scout", + "samples": 1 + } + }, + "cells": [ + { + "context_tokens": 8192, + "concurrency": 1, + "aggregate_tps": 49.47668886774713, + "server_steps_per_s": 18.678952376183886, + "server_spec_accept_length": 2.648793565683646, + "num_errors": 0 + }, + { + "context_tokens": 65536, + "concurrency": 1, + "aggregate_tps": 53.23251039110365, + "server_steps_per_s": 18.85837762683897, + "server_spec_accept_length": 2.822751322751323, + "num_errors": 0 + }, + { + "context_tokens": 8192, + "concurrency": 2, + "aggregate_tps": 76.66132906324916, + "server_steps_per_s": 29.485126562788135, + "server_spec_accept_length": 2.6, + "num_errors": 0 + }, + { + "context_tokens": 65536, + "concurrency": 2, + "aggregate_tps": 79.2456615508071, + "server_steps_per_s": 29.692045340555573, + "server_spec_accept_length": 2.668918918918919, + "num_errors": 0 + } + ] + }, + { + "receipt_sha256": "1b07b7de5e927e1d755e8f3e33aa63c0765508e358f3bffc73a3af55e5f94006", + "receipt_name": "combined-confirmation-2.json", + "metadata": { + "version": "0.4.32", + "model": "glm-5.3-flash-spark", + "temperature": 1.0, + "max_tokens": 2048, + "duration_per_test": 20.0, + "context_lengths": [ + 8192, + 65536 + ], + "concurrency_levels": [ + 1, + 2 + ] + }, + "prefill": { + "8192": { + "tok_per_sec": 2661.0, + "ttft_seconds": 3.078, + "method": "integrated_scout", + "samples": 1 + }, + "65536": { + "tok_per_sec": 2767.0, + "ttft_seconds": 23.688, + "method": "integrated_scout", + "samples": 1 + } + }, + "cells": [ + { + "context_tokens": 8192, + "concurrency": 1, + "aggregate_tps": 48.25, + "server_steps_per_s": 18.950000000000003, + "server_spec_accept_length": 2.5461741424802113, + "num_errors": 0 + }, + { + "context_tokens": 65536, + "concurrency": 1, + "aggregate_tps": 50.89071257005508, + "server_steps_per_s": 18.81505204163295, + "server_spec_accept_length": 2.704787234042553, + "num_errors": 0 + }, + { + "context_tokens": 8192, + "concurrency": 2, + "aggregate_tps": 75.48400040124346, + "server_steps_per_s": 29.09017955662539, + "server_spec_accept_length": 2.594827586206897, + "num_errors": 0 + }, + { + "context_tokens": 65536, + "concurrency": 2, + "aggregate_tps": 80.048149262714, + "server_steps_per_s": 28.799334150732836, + "server_spec_accept_length": 2.779513888888889, + "num_errors": 0 + } + ] + } + ], + "raw_receipt_sha256": { + "image-c139f3670-receipt.json": "91e552e5274716d568f131b515a9f76a663b66e6cf3a72528675d3209ccd3857", + "model-ready.json": "3fa1942c036ea40fd28861eafd9a0fa96c1df3ee397dbbf4ec8b0094911d4171", + "model-ready-after-restart.json": "c72c18eeb15551cfb04f5b6555602efa9444e7ad144230e2665ad0cabcdb0622", + "idle-rank-loss.json": "af4d0cb0ab54b318b3fb58bdca45e3fba3bae5e6682f36d42e41e55cabde1d30", + "cache-publication.log": "630ee5413cc19bf4a7225a62c08356c8c5ace826e7e444ca89af097c75b3181f", + "cache-restore.log": "d18b34f62b45bfaf0424ccfcf1f25adb91a60ae16e3d5b48627efa69074ffd91", + "benchmark-command.json": "44f4e119f787e4a2086398dda5ce7164157040414c1b593248370f8d56af0a81" + }, + "limits": [ + "Idle rank-loss containment does not establish in-flight GPU failure containment or resolve upstream RoCEnante #313.", + "One persistent recall fixture does not establish all cache boundaries or model accuracy.", + "Serving measurements are observations, not a guaranteed speedup." + ] +} diff --git a/performance/records/glm53-flash/spark-mtp3-compute-stream-safety-20260906.md b/performance/records/glm53-flash/spark-mtp3-compute-stream-safety-20260906.md new file mode 100644 index 00000000..ad51d179 --- /dev/null +++ b/performance/records/glm53-flash/spark-mtp3-compute-stream-safety-20260906.md @@ -0,0 +1,63 @@ +# GLM native-MTP3 compute and stream-safety image validation + +Status: **qualified** for the bounded checks listed below. The serving profile +remains **research-only**. + +## Configuration + +Four NVIDIA DGX Sparks serve GLM-5.3-Flash-NVFP4-Spark with TP4/DCP4, native +MTP3, an NVFP4/BF16 proposal head, a BF16 verifier head, 24 GiB FP8 KV per +rank, an 8,192-token scheduling limit, and 16 sequences. SIRCL and RoCEnante +use the hardware-forwarded physical ring. SparkCache uses the profile's +compute-specific namespace. + +- Tested and published config-image ID: + `sha256:2e41b1e934a85ff7c21b780532db2f0a0e978df081e52f4ae2bf11f8992fb24f`. +- Registry manifest: + `sha256:67dc0ae453baaae6831ccec1d259b4ef8b236a8b0dc9f747d901b95c66ec1987`. +- Compute source lock: + `139f36701e0e47f45bf99fba2cc2fa59b417f2ee801dad3a064455d5b464a459`. +- Transport manifest: + `69313e19e881ec93e9ed3bd150d2f24fc6b444488ac729a69f45d038e2243500`. + +The [numeric record](spark-mtp3-compute-stream-safety-20260906.json) includes +per-rank collective evidence, stream checks, all serving cells, source hashes, +and cache measurements. Host startup used the memory gate proposed in +[PR 222](https://github.com/FujitsuPolycom/sparkring/pull/222). + +## Results + +| Check | Result | +|---|---| +| Installed compute package comparison | 4,891 vLLM, 385 B12X, and 150 SparkCache files match the compute-tested image; selected environment values match | +| Four-rank native all-reduce | Exact BF16 results at 4, 20, 28, and 64 token rows; graph input mutations and per-QP completion counters pass | +| GPU stream safety | At 4 and 64 rows, 16 alternating-stream calls with misaligned input/output, two changed-input graph replays, and second-stream capture rejection pass | +| Full serving matrix | 18 cells complete without errors, capacity limitation, underfilling, or warm-up timeout | +| Focused repetitions | Two four-cell 8K/64K C1/C2 runs complete cleanly; 64K prefill measures 2,770 and 2,767 tok/s | +| Idle model-rank loss | Deliberately stopping rank 3 causes all model ranks to stop within 10.063 seconds | +| Recovery and restart | Same image and container identities pass managed recovery, startup memory checks, and four-rank readiness | +| Persistent prefix restore | Correct phrase recall after restart, 26,624 external-hit tokens, matching restore logs on all four ranks | +| Public image access | Anonymous manifest/config verification and Docker pull pass; pull reused local image layers | + +The serving harness is version 0.4.32, pinned by source hash. It uses +temperature 1, 2,048 maximum output tokens, 20-second cells, 8K/32K/64K +contexts and C1/C2/C4/C8/C12/C16. The full matrix's initial 64K prefill scout +measured 2,687 tok/s; that lower result did not repeat in the two focused runs. +Individual decode cells vary. These results establish a functioning composed +image, not a guaranteed performance improvement over every tested configuration. + +The persistent-cache fixture contains 27,274 prompt tokens. Both requests +returned `cobalt orchard lantern` with a normal stop. Full request time was +11.859 seconds cold and 2.922 seconds after restart; restore logs reported +26,624 tokens on every rank, and the engine reported 26,624 external-hit +tokens. This proves persistent restoration for that fixture rather than merely +fast GPU-prefix reuse within one process. + +## Scope + +The rank-loss test used an idle model. It does not establish containment of +an in-flight stalled GPU collective or resolve upstream RoCEnante issue #313. +One recall fixture does not qualify every cache boundary, concurrent cache +workload, or general model accuracy. The recorded checks do not exercise +unattended high availability. Watchdog and memory-protection thresholds were +not relaxed. All four model containers were healthy after restoration. diff --git a/performance/records/glm53-flash/spark-mtp3-nvfp4-proposal-head-20260905.json b/performance/records/glm53-flash/spark-mtp3-nvfp4-proposal-head-20260905.json new file mode 100644 index 00000000..eadc6683 --- /dev/null +++ b/performance/records/glm53-flash/spark-mtp3-nvfp4-proposal-head-20260905.json @@ -0,0 +1,74 @@ +{ + "schema": "sparkring-saved-proposal-head-comparison/v1", + "status": "research-only", + "conditions": { + "hardware": "4 NVIDIA DGX Spark GB10 systems", + "topology": "TP4/DCP4/PP1 direct-cycle mesh", + "model": "local-inference-lab/GLM-5.3-Flash-NVFP4-Spark@df116c4fb16b1d37ae43d2cfd624de26ffbc832e", + "speculation": "native MTP3; no external DFlash model", + "sampling": {"temperature": 1.0, "ignore_eos": true}, + "measurement_seconds": 20.0, + "decode_context_tokens": 8192, + "concurrency": [1, 2, 4, 8], + "proposal_head": "runtime NVFP4 weights with BF16 activations", + "target_verifier_head": "retained BF16 checkpoint representation", + "cuda_toolchain": "13.3.33", + "b12x_revision": "b58f34eaf978277621efced6678e6713fd7122e4", + "vllm_head_donor_revision": "3512b066e7796128c0c380ccc558182960f2f0ea", + "vllm_complete_donor_revision": "a8c796f3af74106b2d8d441e9ec54588936a5388", + "metadata_dense_base_image_id": "sha256:f35ed3d1df1ee57f66ba571a491d6dd5575d1a2e619cbc63d065008894777ffe", + "tested_private_image_id": "sha256:04d5a35b03e99f68c37a05514d221988a3eb70a5b8fdcfa859025ca1cbc25e74", + "published_image": "ghcr.io/fujitsupolycom/sparkring-glm53-sparkcache@sha256:1b97e1dc9cb93c39f887f40bab24359a9b6ec998c28d2417b160f2103cd5fd86", + "published_image_id": "sha256:dd6c51efaf4127df863ac85c3be3fe46f260b34c7ab2deb384669fffdbe857df", + "published_source_lock_sha256": "2a444f7c1ad4319f64afbffb16403491a4863934372cf818f18c510d58c0d00e", + "published_compute_equivalence": "all 4891 vLLM, 385 B12X, and 150 SparkCache package files plus selected environment match the tested private image", + "published_exact_image_serving_repeated": false, + "tested_image_receipt_sha256": "aa8f1bf3d892af7b0a8059626ec91a83e2dbc5a39c8d58c0fb82c614e88dbaba", + "model_ready_receipt_sha256": "fe786299a6459cd57dcf5362591b56d94569dd473b0d1fefb47ca374399983f3", + "live_control_receipt_sha256": "b0ebcf65be43998791f66da819ebaa98aa4b4323b9b76cc1e363d977750e3341", + "proposal_head_manifest_sha256": "9408c82bcb5a47ce0cb580d6835b23924b1d40a29352ac97f72d2fd8a77b9fb1", + "final_private_image_identity_available_in_receipts": true + }, + "receipts": { + "proposal_head": [ + { + "basename": "glm-5.3-flash-spark-dcp4-MTP3-SIRCL-and-MESH-r2420260905-122857.json", + "sha256": "05d846465153230f4312a59aefef12c499d81634855f3c84cb6757baf653c0e4" + }, + { + "basename": "glm-5.3-flash-spark-dcp4-MTP3-SIRCL-and-MESH-r2420260905-123541.json", + "sha256": "9664bc0410db423d3edc759106f10bc24a0fe1d15952102cdd74913ff96c2149" + }, + { + "basename": "glm-5.3-flash-spark-dcp4-MTP3-SIRCL-and-MESH-r2420260905-124222.json", + "sha256": "4213df0fc47ba53308ea7dc076e3e1c77e9097e5eaef95db4e09ca4cd404e434" + } + ], + "shared_bf16_proposal_head_control": [ + { + "basename": "glm-5.3-flash-spark-dcp4-MTP3-SIRCL-and-MESH-r2420260905-114306.json", + "sha256": "f7ae32cef4c9ea606477a28ff0465403e78f9519e2d771e4af0f13017b671677" + }, + { + "basename": "glm-5.3-flash-spark-dcp4-MTP3-SIRCL-and-MESH-r2420260905-115629.json", + "sha256": "7c3ff3210e805b8a801d4e8ea49152d0545a982707cc9010a174f77ff3d4e673" + } + ], + "published_context_only": { + "basename": "glm-5.3-flash-spark-dcp4-dflash7-bf16-SIRCL-and-MESH-20260904-233405.json", + "sha256": "f0916f6b72cb8256225169b44c4f11e3ca764a5dd854977b8963686197b843fa" + } + }, + "decode_8k": [ + {"concurrency": 1, "control_raw_tps": [48.839071256871314, 46.61270983199115], "proposal_raw_tps": [49.35, 52.7922337868845, 52.8], "control_steps_per_s": [18.014411529173845, 17.93565147876187], "proposal_steps_per_s": [18.7, 18.8150520415816, 19.05], "raw_mean_delta_percent": 8.21675755680824, "steps_mean_delta_percent": 4.8956011177602}, + {"concurrency": 2, "control_raw_tps": [74.55964771796953, 77.1430004517391], "proposal_raw_tps": [74.6559919940477, 77.0811406806786, 78.9], "control_steps_per_s": [28.122497998321393, 28.93488271324153], "proposal_steps_per_s": [28.2211658744255, 29.2687816368766, 29.600000000000005], "raw_mean_delta_percent": 1.35491390882609, "steps_mean_delta_percent": 1.75715093028919}, + {"concurrency": 4, "control_raw_tps": [120.7303004461975, 118.67889540432174], "proposal_raw_tps": [123.991379741799, 115.25, 123.18949531411437], "control_steps_per_s": [43.336510006445636, 43.30175913400928], "proposal_steps_per_s": [43.9031724550588, 41.8, 44.10364356241686], "raw_mean_delta_percent": 0.923685287957832, "steps_mean_delta_percent": -0.115874655007697}, + {"concurrency": 8, "control_raw_tps": [163.50687130131578, 171.3898781157694], "proposal_raw_tps": [164.267442443414, 173.45, 168.77003205133715], "control_steps_per_s": [59.384090681214076, 60.189597231174496], "proposal_steps_per_s": [61.393389175798, 61.0767548906789, 60.89743589745578], "raw_mean_delta_percent": 0.824602483468539, "steps_mean_delta_percent": 2.23407460582565} + ], + "prefill_prompt_tps_means": [ + {"context_tokens": 8192, "control": 2668.0, "proposal_head": 2670.33333333333, "delta_percent": 0.0874562718640837}, + {"context_tokens": 16384, "control": 2709.0, "proposal_head": 2707.0, "delta_percent": -0.0738279808047304}, + {"context_tokens": 65536, "control": 2782.5, "proposal_head": 2778.33333333333, "delta_percent": -0.149745432764292}, + {"context_tokens": 131072, "control": 2749.0, "proposal_head": 2739.33333333333, "delta_percent": -0.351643021704862} + ] +} diff --git a/performance/records/glm53-flash/spark-mtp3-nvfp4-proposal-head-20260905.md b/performance/records/glm53-flash/spark-mtp3-nvfp4-proposal-head-20260905.md new file mode 100644 index 00000000..7efe5072 --- /dev/null +++ b/performance/records/glm53-flash/spark-mtp3-nvfp4-proposal-head-20260905.md @@ -0,0 +1,146 @@ +# Native-MTP3 NVFP4 proposal-head comparison + +Status: **research-only**. These measurements support the proposal-head default +inside the research-only GLM-5.3 native-MTP3 mesh profile. They are not a general +model-quality or production-serving claim. + +## Conditions + +The measured system used four NVIDIA DGX Spark GB10 nodes in TP4/DCP4/PP1, +`local-inference-lab/GLM-5.3-Flash-NVFP4-Spark` revision +`df116c4fb16b1d37ae43d2cfd624de26ffbc832e`, native MTP depth three, FP8 KV, +24 GiB KV allocation per rank, SparkCache, and the profile's hybrid +SIRCL/RoCEnante transport. No DFlash model was loaded. + +The measured proposal-head image is +`sha256:04d5a35b03e99f68c37a05514d221988a3eb70a5b8fdcfa859025ca1cbc25e74`. +Its image receipt has SHA-256 +`aa8f1bf3d892af7b0a8059626ec91a83e2dbc5a39c8d58c0fb82c614e88dbaba`; +the completed model-readiness and live-control receipts have SHA-256 +`fe786299a6459cd57dcf5362591b56d94569dd473b0d1fefb47ca374399983f3` +and `b0ebcf65be43998791f66da819ebaa98aa4b4323b9b76cc1e363d977750e3341`. +The image is layered on compute image +`sha256:f35ed3d1df1ee57f66ba571a491d6dd5575d1a2e619cbc63d065008894777ffe`, +which records CUDA 13.3.33, complete B12X revision +`b58f34eaf978277621efced6678e6713fd7122e4`, the uniform-speculation metadata +port from Local Inference Lab vLLM revision +`3512b066e7796128c0c380ccc558182960f2f0ea`, and dense wrapper changes from +revision `a8c796f3af74106b2d8d441e9ec54588936a5388`. The complete B12X tree also +contains MoE and other package changes in B12X revision `b58f34ea`, so these runs do not isolate a dense +kernel contribution. + +The changed variable was a separate runtime-NVFP4 proposal head with BF16 +activations. The target/verifier head kept its BF16 checkpoint representation. +The proposal head adds 85.08 MiB of persistent packed weight and scale storage +per rank while the retained BF16 target head remains allocated. It is not a net +85.08 MiB model-memory reduction. The shared-BF16-head control used the same +CUDA version, B12X kernels, metadata reuse, and dense-kernel integration without +the separate proposal allocation. Thus the proposal-head comparison preserves +the verifier implementation. Comparisons against the separate +[mesh matrix](spark-mtp3-mesh-20260905.md) also differ in CUDA, metadata, +dense-kernel integration, MoE, and B12X code; retaining a BF16 verifier head +does not imply that every target-side kernel matches between those configurations. + +The benchmark used harness version 0.4.32, temperature 1.0, ignored EOS, a +20-second sustained-decode window, 8,192-token decode context, and concurrency +1, 2, 4, and 8. The proposal-head configuration has three repetitions and the +control has two. All compared 8K cells had +zero request errors, no warmup timeout, and no underfill flag. + +The benchmark JSON does not embed the image ID. The image, launch, and completed +four-rank readiness receipts bind the benchmark endpoint to the tested private +image. The compute-equivalent public derivative is +`ghcr.io/fujitsupolycom/sparkring-glm53-sparkcache@sha256:1b97e1dc9cb93c39f887f40bab24359a9b6ec998c28d2417b160f2103cd5fd86` +with image ID +`sha256:dd6c51efaf4127df863ac85c3be3fe46f260b34c7ab2deb384669fffdbe857df`. +Its compute-equivalence record matches all vLLM, B12X, and SparkCache package +files and the selected environment to the tested image. Exact-image serving, +restart, and persistent-cache checks were not repeated on the public image ID. + +## Measurement + +[Sanitized per-run samples](spark-mtp3-nvfp4-proposal-samples-20260905.json) +include all decode windows, token counts, per-context prefill scouts, and +recorded settings for the two control and three proposal repetitions. +Each decode cell requests five seconds of warmup, allows 900 seconds to reach +the requested concurrency, and measures a 20-second client wall-clock window. +The harness version is 0.4.32; its exact historical source commit was not +captured. The recorded equivalent command is: + +```bash +python llm_decode_bench.py --host http://RANK0 --port 8015 \ + --model glm-5.3-flash-spark --temperature 1.0 --token-targeting exact \ + --display-mode live --no-hw-monitor --dcp-size 4 \ + --concurrency 1,2,4,8 --contexts 8k,16k --max-tokens 2048 \ + --duration 20 --decode-warmup-seconds 5 \ + --cell-warmup-timeout-seconds 900 --output result.json +``` + +Variability is reported as the minimum and maximum across repetitions, not +as a confidence interval. At 8K/C1, control output throughput spans +46.61–48.84 tok/s and proposal throughput spans 49.35–52.80 tok/s; +normalized throughput spans 17.94–18.01 and 18.70–19.05 respectively. +The sample file contains the same ranges for every measured 8K concurrency. + +Raw throughput is aggregate output tokens per second. Normalized throughput is +the harness's aggregate sequence steps per second, computed from drafted and +non-speculative request work. It is not a count of batched engine iterations. +Speculative acceptance can move raw tokens per second, so both metrics are +reported. The machine-readable values and receipt hashes are in +[`spark-mtp3-nvfp4-proposal-head-20260905.json`](spark-mtp3-nvfp4-proposal-head-20260905.json). + +Proposal-head receipt hashes: + +- `122857`: `05d846465153230f4312a59aefef12c499d81634855f3c84cb6757baf653c0e4` +- `123541`: `9664bc0410db423d3edc759106f10bc24a0fe1d15952102cdd74913ff96c2149` +- `124222`: `4213df0fc47ba53308ea7dc076e3e1c77e9097e5eaef95db4e09ca4cd404e434` + +Shared-BF16-head control hashes: + +- `114306`: `f7ae32cef4c9ea606477a28ff0465403e78f9519e2d771e4af0f13017b671677` +- `115629`: `7c3ff3210e805b8a801d4e8ea49152d0545a982707cc9010a174f77ff3d4e673` + +The previously published native-MTP3 `233405` receipt is retained as historical context +with SHA-256 `f0916f6b72cb8256225169b44c4f11e3ca764a5dd854977b8963686197b843fa`. +Its basename contains a stale DFlash7 label, but the public performance record +identifies the measured runtime as native MTP3. Its matrix and decode warmup +differ, and its JSON does not attest an image identity. It is not used to +calculate the proposal-head delta. + +## Result + +| 8K concurrency | Control raw tok/s | NVFP4-head raw tok/s | Raw change | Control steps/s | NVFP4-head steps/s | Normalized change | +|---:|---:|---:|---:|---:|---:|---:| +| 1 | 47.73 | 51.65 | **+8.22%** | 17.98 | 18.86 | **+4.90%** | +| 2 | 75.85 | 76.88 | +1.35% | 28.53 | 29.03 | +1.76% | +| 4 | 119.70 | 120.81 | +0.92% | 43.32 | 43.27 | -0.12% | +| 8 | 167.45 | 168.83 | +0.82% | 59.79 | 61.12 | +2.23% | + +The C1 repetitions show a small-batch benefit: approximately 8% +raw throughput and 5% normalized throughput. C2, C4, and C8 are mixed and do +not support a monotonic concurrency-wide speedup claim. + +Mean prefill scouts were unchanged within ±0.36%: 8K 2,668 versus 2,670.3 +prompt tok/s, 16K 2,709 versus 2,707, 64K 2,782.5 versus 2,778.3, and 128K +2,749 versus 2,739.3. The head change is decode-facing; these observations do +not establish a prefill improvement. + +## Conclusion + +The bounded evidence supports making runtime NVFP4/BF16 proposal-head execution +the default within this opt-in, research-only profile. It improves repeated C1 +decode while remaining approximately flat or mixed at the measured higher +concurrencies and prefill contexts. Relative to the compute-image control, the +verifier implementation and BF16 head are fixed; the proposal change can alter +speculative acceptance but not the target distribution under standard +rejection sampling. + +## Limitations + +The two control and three proposal-head repetitions are a small sample. They +cover only C1/C2/C4/C8 at 8K for +the direct comparison, not the full C1–C16 and 8K/32K/64K matrix. They do not +include isolated dense or MoE attribution, a general +accuracy evaluation, host reboot, long soak, failure containment, or unattended +high availability. A full C1–C16 and 8K/32K/64K matrix requires a separate +completed receipt before it can extend this record. diff --git a/performance/records/glm53-flash/spark-mtp3-nvfp4-proposal-samples-20260905.json b/performance/records/glm53-flash/spark-mtp3-nvfp4-proposal-samples-20260905.json new file mode 100644 index 00000000..ac9ca257 --- /dev/null +++ b/performance/records/glm53-flash/spark-mtp3-nvfp4-proposal-samples-20260905.json @@ -0,0 +1,1272 @@ +{ + "schema": "sparkring-sanitized-proposal-benchmark-samples/v1", + "harness_revision": "Recorded version 0.4.32; exact historical source commit was not captured", + "timing": "Client wall-clock measurement window; aggregate output uses OpenAI continuous-usage token counts", + "variability": "Minimum and maximum across repetitions at each 8K concurrency; not confidence intervals", + "ranges": { + "control": [ + { + "concurrency": 1, + "aggregate_tps_min": 46.61270983199115, + "aggregate_tps_max": 48.839071256871314, + "normalized_min": 17.93565147876187, + "normalized_max": 18.014411529173845 + }, + { + "concurrency": 2, + "aggregate_tps_min": 74.55964771796953, + "aggregate_tps_max": 77.1430004517391, + "normalized_min": 28.122497998321393, + "normalized_max": 28.93488271324153 + }, + { + "concurrency": 4, + "aggregate_tps_min": 118.67889540432174, + "aggregate_tps_max": 120.7303004461975, + "normalized_min": 43.30175913400928, + "normalized_max": 43.336510006445636 + }, + { + "concurrency": 8, + "aggregate_tps_min": 163.50687130131578, + "aggregate_tps_max": 171.3898781157694, + "normalized_min": 59.384090681214076, + "normalized_max": 60.189597231174496 + } + ], + "proposal": [ + { + "concurrency": 1, + "aggregate_tps_min": 49.35, + "aggregate_tps_max": 52.8, + "normalized_min": 18.7, + "normalized_max": 19.05 + }, + { + "concurrency": 2, + "aggregate_tps_min": 74.65599199404768, + "aggregate_tps_max": 78.9, + "normalized_min": 28.22116587442553, + "normalized_max": 29.600000000000005 + }, + { + "concurrency": 4, + "aggregate_tps_min": 115.25, + "aggregate_tps_max": 123.99137974179853, + "normalized_min": 41.8, + "normalized_max": 44.10364356241686 + }, + { + "concurrency": 8, + "aggregate_tps_min": 164.26744244341373, + "aggregate_tps_max": 173.45, + "normalized_min": 60.89743589745578, + "normalized_max": 61.39338917579799 + } + ] + }, + "runs": [ + { + "id": "114306", + "source_sha256": "f7ae32cef4c9ea606477a28ff0465403e78f9519e2d771e4af0f13017b671677", + "metadata": { + "version": "0.4.32", + "engine": "vllm", + "model": "glm-5.3-flash-spark", + "decode_mode": "duration", + "duration_per_test": 20, + "decode_warmup_seconds": 5, + "decode_warmup_context": 16384, + "decode_warmup_concurrency": 1, + "cell_warmup_timeout_seconds": 900, + "max_tokens": 2048, + "temperature": 1, + "ignore_eos": true, + "dcp_size": 4, + "concurrency_levels": [ + 1, + 2, + 4, + 8 + ], + "context_lengths": [ + 8192, + 16384 + ], + "prefill_mode": "integrated_decode_scout" + }, + "prefill": { + "8192": { + "ttft_seconds": 3.078, + "prefill_seconds": 3.078, + "tok_per_sec": 2661, + "client_ttft_seconds": 3.078, + "client_tok_per_sec": 2661, + "prompt_tokens": 8192, + "samples": 1, + "method": "integrated_scout", + "server_validation": { + "method": "prometheus:kv_computed", + "tok_per_sec": 2688, + "prefill_seconds": 3.047, + "prompt_tokens": 8192, + "request_prompt_tokens": 0, + "cached_tokens": 0, + "token_source": "", + "samples": 1, + "invalid_reason": "" + }, + "hardware_summary": {} + }, + "16384": { + "ttft_seconds": 6.047, + "prefill_seconds": 6.047, + "tok_per_sec": 2709, + "client_ttft_seconds": 6.047, + "client_tok_per_sec": 2709, + "prompt_tokens": 16384, + "samples": 1, + "method": "integrated_scout", + "server_validation": { + "method": "prometheus:kv_computed", + "tok_per_sec": 2722, + "prefill_seconds": 6.019, + "prompt_tokens": 16384, + "request_prompt_tokens": 0, + "cached_tokens": 0, + "token_source": "", + "samples": 1, + "invalid_reason": "" + }, + "hardware_summary": {} + }, + "65536": { + "ttft_seconds": 23.578, + "prefill_seconds": 23.578, + "tok_per_sec": 2780, + "client_ttft_seconds": 23.578, + "client_tok_per_sec": 2780, + "prompt_tokens": 65536, + "samples": 1, + "method": "scout_only", + "server_validation": { + "method": "", + "tok_per_sec": 0, + "prefill_seconds": 0, + "prompt_tokens": 0, + "request_prompt_tokens": 0, + "cached_tokens": 0, + "token_source": "", + "samples": 0, + "invalid_reason": "" + }, + "hardware_summary": {} + }, + "131072": { + "ttft_seconds": 47.531, + "prefill_seconds": 47.531, + "tok_per_sec": 2758, + "client_ttft_seconds": 47.531, + "client_tok_per_sec": 2758, + "prompt_tokens": 131072, + "samples": 1, + "method": "scout_only", + "server_validation": { + "method": "", + "tok_per_sec": 0, + "prefill_seconds": 0, + "prompt_tokens": 0, + "request_prompt_tokens": 0, + "cached_tokens": 0, + "token_source": "", + "samples": 0, + "invalid_reason": "" + }, + "hardware_summary": {} + } + }, + "decode": [ + { + "context_tokens": 8192, + "concurrency": 1, + "aggregate_tps": 48.839071256871314, + "server_steps_per_s": 18.014411529173845, + "server_spec_accept_length": 2.7111111111111112, + "measurement_wall_seconds": 20.015, + "client_output_tokens": 976, + "server_output_tokens": 976, + "aggregate_source": "openai_continuous_usage", + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false, + "warmup_duration": 12.907 + }, + { + "context_tokens": 16384, + "concurrency": 1, + "aggregate_tps": 49.93745308985227, + "server_steps_per_s": 17.813360020027464, + "server_spec_accept_length": 2.803370786516854, + "measurement_wall_seconds": 20, + "client_output_tokens": 998, + "server_output_tokens": 998, + "aggregate_source": "openai_continuous_usage", + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false, + "warmup_duration": 13.453 + }, + { + "context_tokens": 8192, + "concurrency": 2, + "aggregate_tps": 74.55964771796953, + "server_steps_per_s": 28.122497998321393, + "server_spec_accept_length": 2.6512455516014235, + "measurement_wall_seconds": 20.015, + "client_output_tokens": 1490, + "server_output_tokens": 1490, + "aggregate_source": "openai_continuous_usage", + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false, + "warmup_duration": 6.625 + }, + { + "context_tokens": 8192, + "concurrency": 4, + "aggregate_tps": 120.7303004461975, + "server_steps_per_s": 43.336510006445636, + "server_spec_accept_length": 2.7858796296296298, + "measurement_wall_seconds": 20, + "client_output_tokens": 2407, + "server_output_tokens": 2407, + "aggregate_source": "openai_continuous_usage", + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false, + "warmup_duration": 8.704 + }, + { + "context_tokens": 8192, + "concurrency": 8, + "aggregate_tps": 163.50687130131578, + "server_steps_per_s": 59.384090681214076, + "server_spec_accept_length": 2.753378378378378, + "measurement_wall_seconds": 20, + "client_output_tokens": 3260, + "server_output_tokens": 3260, + "aggregate_source": "openai_continuous_usage", + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false, + "warmup_duration": 12.234 + }, + { + "context_tokens": 16384, + "concurrency": 2, + "aggregate_tps": 75.0062546910707, + "server_steps_per_s": 28.42131598701011, + "server_spec_accept_length": 2.6390845070422535, + "measurement_wall_seconds": 20.016, + "client_output_tokens": 1499, + "server_output_tokens": 1499, + "aggregate_source": "openai_continuous_usage", + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false, + "warmup_duration": 7.156 + }, + { + "context_tokens": 16384, + "concurrency": 4, + "aggregate_tps": 113.46664661967247, + "server_steps_per_s": 41.71421133423315, + "server_spec_accept_length": 2.7200956937799043, + "measurement_wall_seconds": 20, + "client_output_tokens": 2264, + "server_output_tokens": 2274, + "aggregate_source": "openai_continuous_usage", + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false, + "warmup_duration": 9.187 + }, + { + "context_tokens": 16384, + "concurrency": 8, + "aggregate_tps": 163.28993011960046, + "server_steps_per_s": 59.524408023893756, + "server_spec_accept_length": 2.743243243243243, + "measurement_wall_seconds": 20.016, + "client_output_tokens": 3248, + "server_output_tokens": 3248, + "aggregate_source": "openai_continuous_usage", + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false, + "warmup_duration": 13.297 + } + ] + }, + { + "id": "115629", + "source_sha256": "7c3ff3210e805b8a801d4e8ea49152d0545a982707cc9010a174f77ff3d4e673", + "metadata": { + "version": "0.4.32", + "engine": "vllm", + "model": "glm-5.3-flash-spark", + "decode_mode": "duration", + "duration_per_test": 20, + "decode_warmup_seconds": 5, + "decode_warmup_context": 16384, + "decode_warmup_concurrency": 1, + "cell_warmup_timeout_seconds": 900, + "max_tokens": 2048, + "temperature": 1, + "ignore_eos": true, + "dcp_size": 4, + "concurrency_levels": [ + 1, + 2, + 4, + 8 + ], + "context_lengths": [ + 8192, + 16384 + ], + "prefill_mode": "integrated_decode_scout" + }, + "prefill": { + "8192": { + "ttft_seconds": 3.063, + "prefill_seconds": 3.063, + "tok_per_sec": 2675, + "client_ttft_seconds": 3.063, + "client_tok_per_sec": 2675, + "prompt_tokens": 8192, + "samples": 1, + "method": "integrated_scout", + "server_validation": { + "method": "prometheus:kv_computed", + "tok_per_sec": 2693, + "prefill_seconds": 3.042, + "prompt_tokens": 8192, + "request_prompt_tokens": 0, + "cached_tokens": 0, + "token_source": "", + "samples": 1, + "invalid_reason": "" + }, + "hardware_summary": {} + }, + "16384": { + "ttft_seconds": 6.047, + "prefill_seconds": 6.047, + "tok_per_sec": 2709, + "client_ttft_seconds": 6.047, + "client_tok_per_sec": 2709, + "prompt_tokens": 16384, + "samples": 1, + "method": "integrated_scout", + "server_validation": { + "method": "prometheus:kv_computed", + "tok_per_sec": 2724, + "prefill_seconds": 6.014, + "prompt_tokens": 16384, + "request_prompt_tokens": 0, + "cached_tokens": 0, + "token_source": "", + "samples": 1, + "invalid_reason": "" + }, + "hardware_summary": {} + }, + "65536": { + "ttft_seconds": 23.532, + "prefill_seconds": 23.532, + "tok_per_sec": 2785, + "client_ttft_seconds": 23.532, + "client_tok_per_sec": 2785, + "prompt_tokens": 65536, + "samples": 1, + "method": "scout_only", + "server_validation": { + "method": "", + "tok_per_sec": 0, + "prefill_seconds": 0, + "prompt_tokens": 0, + "request_prompt_tokens": 0, + "cached_tokens": 0, + "token_source": "", + "samples": 0, + "invalid_reason": "" + }, + "hardware_summary": {} + }, + "131072": { + "ttft_seconds": 47.828, + "prefill_seconds": 47.828, + "tok_per_sec": 2740, + "client_ttft_seconds": 47.828, + "client_tok_per_sec": 2740, + "prompt_tokens": 131072, + "samples": 1, + "method": "scout_only", + "server_validation": { + "method": "", + "tok_per_sec": 0, + "prefill_seconds": 0, + "prompt_tokens": 0, + "request_prompt_tokens": 0, + "cached_tokens": 0, + "token_source": "", + "samples": 0, + "invalid_reason": "" + }, + "hardware_summary": {} + } + }, + "decode": [ + { + "context_tokens": 8192, + "concurrency": 1, + "aggregate_tps": 46.61270983199115, + "server_steps_per_s": 17.93565147876187, + "server_spec_accept_length": 2.598885793871866, + "measurement_wall_seconds": 20.016, + "client_output_tokens": 933, + "server_output_tokens": 933, + "aggregate_source": "openai_continuous_usage", + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false, + "warmup_duration": 12.437 + }, + { + "context_tokens": 16384, + "concurrency": 1, + "aggregate_tps": 47.774049777056554, + "server_steps_per_s": 17.92150874604316, + "server_spec_accept_length": 2.6657381615598883, + "measurement_wall_seconds": 20.016, + "client_output_tokens": 954, + "server_output_tokens": 957, + "aggregate_source": "openai_continuous_usage", + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false, + "warmup_duration": 13.937 + }, + { + "context_tokens": 8192, + "concurrency": 2, + "aggregate_tps": 77.1430004517391, + "server_steps_per_s": 28.93488271324153, + "server_spec_accept_length": 2.6660899653979238, + "measurement_wall_seconds": 20.016, + "client_output_tokens": 1538, + "server_output_tokens": 1541, + "aggregate_source": "openai_continuous_usage", + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false, + "warmup_duration": 6.656 + }, + { + "context_tokens": 8192, + "concurrency": 4, + "aggregate_tps": 118.67889540432174, + "server_steps_per_s": 43.30175913400928, + "server_spec_accept_length": 2.7407407407407405, + "measurement_wall_seconds": 20.016, + "client_output_tokens": 2368, + "server_output_tokens": 2368, + "aggregate_source": "openai_continuous_usage", + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false, + "warmup_duration": 8.719 + }, + { + "context_tokens": 8192, + "concurrency": 8, + "aggregate_tps": 171.3898781157694, + "server_steps_per_s": 60.189597231174496, + "server_spec_accept_length": 2.8475, + "measurement_wall_seconds": 20, + "client_output_tokens": 3417, + "server_output_tokens": 3417, + "aggregate_source": "openai_continuous_usage", + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false, + "warmup_duration": 12.266 + }, + { + "context_tokens": 16384, + "concurrency": 2, + "aggregate_tps": 75.65, + "server_steps_per_s": 28.000000000000004, + "server_spec_accept_length": 2.7017857142857142, + "measurement_wall_seconds": 20, + "client_output_tokens": 1513, + "server_output_tokens": 1513, + "aggregate_source": "openai_continuous_usage", + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false, + "warmup_duration": 7.11 + }, + { + "context_tokens": 16384, + "concurrency": 4, + "aggregate_tps": 117.38194201055911, + "server_steps_per_s": 42.26551154305115, + "server_spec_accept_length": 2.7772511848341233, + "measurement_wall_seconds": 20, + "client_output_tokens": 2344, + "server_output_tokens": 2344, + "aggregate_source": "openai_continuous_usage", + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false, + "warmup_duration": 9.687 + }, + { + "context_tokens": 16384, + "concurrency": 8, + "aggregate_tps": 184.4835363105187, + "server_steps_per_s": 63.74981205840255, + "server_spec_accept_length": 2.893867924528302, + "measurement_wall_seconds": 20.015, + "client_output_tokens": 3681, + "server_output_tokens": 3681, + "aggregate_source": "openai_continuous_usage", + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false, + "warmup_duration": 11.797 + } + ] + }, + { + "id": "122857", + "source_sha256": "05d846465153230f4312a59aefef12c499d81634855f3c84cb6757baf653c0e4", + "metadata": { + "version": "0.4.32", + "engine": "vllm", + "model": "glm-5.3-flash-spark", + "decode_mode": "duration", + "duration_per_test": 20, + "decode_warmup_seconds": 5, + "decode_warmup_context": 16384, + "decode_warmup_concurrency": 1, + "cell_warmup_timeout_seconds": 900, + "max_tokens": 2048, + "temperature": 1, + "ignore_eos": true, + "dcp_size": 4, + "concurrency_levels": [ + 1, + 2, + 4, + 8 + ], + "context_lengths": [ + 8192, + 16384 + ], + "prefill_mode": "integrated_decode_scout" + }, + "prefill": { + "8192": { + "ttft_seconds": 3.079, + "prefill_seconds": 3.079, + "tok_per_sec": 2661, + "client_ttft_seconds": 3.079, + "client_tok_per_sec": 2661, + "prompt_tokens": 8192, + "samples": 1, + "method": "integrated_scout", + "server_validation": { + "method": "prometheus:kv_computed", + "tok_per_sec": 2694, + "prefill_seconds": 3.041, + "prompt_tokens": 8192, + "request_prompt_tokens": 0, + "cached_tokens": 0, + "token_source": "", + "samples": 1, + "invalid_reason": "" + }, + "hardware_summary": {} + }, + "16384": { + "ttft_seconds": 6.031, + "prefill_seconds": 6.031, + "tok_per_sec": 2717, + "client_ttft_seconds": 6.031, + "client_tok_per_sec": 2717, + "prompt_tokens": 16384, + "samples": 1, + "method": "integrated_scout", + "server_validation": { + "method": "prometheus:kv_computed", + "tok_per_sec": 2728, + "prefill_seconds": 6.006, + "prompt_tokens": 16384, + "request_prompt_tokens": 0, + "cached_tokens": 0, + "token_source": "", + "samples": 1, + "invalid_reason": "" + }, + "hardware_summary": {} + }, + "65536": { + "ttft_seconds": 23.579, + "prefill_seconds": 23.579, + "tok_per_sec": 2779, + "client_ttft_seconds": 23.579, + "client_tok_per_sec": 2779, + "prompt_tokens": 65536, + "samples": 1, + "method": "scout_only", + "server_validation": { + "method": "", + "tok_per_sec": 0, + "prefill_seconds": 0, + "prompt_tokens": 0, + "request_prompt_tokens": 0, + "cached_tokens": 0, + "token_source": "", + "samples": 0, + "invalid_reason": "" + }, + "hardware_summary": {} + }, + "131072": { + "ttft_seconds": 47.516, + "prefill_seconds": 47.516, + "tok_per_sec": 2758, + "client_ttft_seconds": 47.516, + "client_tok_per_sec": 2758, + "prompt_tokens": 131072, + "samples": 1, + "method": "scout_only", + "server_validation": { + "method": "", + "tok_per_sec": 0, + "prefill_seconds": 0, + "prompt_tokens": 0, + "request_prompt_tokens": 0, + "cached_tokens": 0, + "token_source": "", + "samples": 0, + "invalid_reason": "" + }, + "hardware_summary": {} + } + }, + "decode": [ + { + "context_tokens": 8192, + "concurrency": 1, + "aggregate_tps": 49.35, + "server_steps_per_s": 18.7, + "server_spec_accept_length": 2.63903743315508, + "measurement_wall_seconds": 20, + "client_output_tokens": 987, + "server_output_tokens": 987, + "aggregate_source": "openai_continuous_usage", + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false, + "warmup_duration": 12.39 + }, + { + "context_tokens": 16384, + "concurrency": 1, + "aggregate_tps": 51.830337022278336, + "server_steps_per_s": 18.52871951521061, + "server_spec_accept_length": 2.7972972972972974, + "measurement_wall_seconds": 20, + "client_output_tokens": 1035, + "server_output_tokens": 1035, + "aggregate_source": "openai_continuous_usage", + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false, + "warmup_duration": 13.953 + }, + { + "context_tokens": 8192, + "concurrency": 2, + "aggregate_tps": 74.65599199404768, + "server_steps_per_s": 28.22116587442553, + "server_spec_accept_length": 2.6453900709219855, + "measurement_wall_seconds": 20, + "client_output_tokens": 1492, + "server_output_tokens": 1492, + "aggregate_source": "openai_continuous_usage", + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false, + "warmup_duration": 6.578 + }, + { + "context_tokens": 8192, + "concurrency": 4, + "aggregate_tps": 123.99137974179853, + "server_steps_per_s": 43.903172455058815, + "server_spec_accept_length": 2.8242009132420094, + "measurement_wall_seconds": 20, + "client_output_tokens": 2474, + "server_output_tokens": 2474, + "aggregate_source": "openai_continuous_usage", + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false, + "warmup_duration": 10.125 + }, + { + "context_tokens": 8192, + "concurrency": 8, + "aggregate_tps": 164.26744244341373, + "server_steps_per_s": 61.39338917579799, + "server_spec_accept_length": 2.6756535947712417, + "measurement_wall_seconds": 20, + "client_output_tokens": 3275, + "server_output_tokens": 3275, + "aggregate_source": "openai_continuous_usage", + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false, + "warmup_duration": 13.36 + }, + { + "context_tokens": 16384, + "concurrency": 2, + "aggregate_tps": 78.62580128207695, + "server_steps_per_s": 28.645833333342686, + "server_spec_accept_length": 2.744755244755245, + "measurement_wall_seconds": 20, + "client_output_tokens": 1570, + "server_output_tokens": 1570, + "aggregate_source": "openai_continuous_usage", + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false, + "warmup_duration": 7.141 + }, + { + "context_tokens": 16384, + "concurrency": 4, + "aggregate_tps": 120.54040530406225, + "server_steps_per_s": 44.03302476860722, + "server_spec_accept_length": 2.7375, + "measurement_wall_seconds": 20, + "client_output_tokens": 2409, + "server_output_tokens": 2409, + "aggregate_source": "openai_continuous_usage", + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false, + "warmup_duration": 9.703 + }, + { + "context_tokens": 16384, + "concurrency": 8, + "aggregate_tps": 168.41787743246968, + "server_steps_per_s": 60.730983862215936, + "server_spec_accept_length": 2.77317880794702, + "measurement_wall_seconds": 20, + "client_output_tokens": 3350, + "server_output_tokens": 3350, + "aggregate_source": "openai_continuous_usage", + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false, + "warmup_duration": 11.89 + } + ] + }, + { + "id": "123541", + "source_sha256": "9664bc0410db423d3edc759106f10bc24a0fe1d15952102cdd74913ff96c2149", + "metadata": { + "version": "0.4.32", + "engine": "vllm", + "model": "glm-5.3-flash-spark", + "decode_mode": "duration", + "duration_per_test": 20, + "decode_warmup_seconds": 5, + "decode_warmup_context": 16384, + "decode_warmup_concurrency": 1, + "cell_warmup_timeout_seconds": 900, + "max_tokens": 2048, + "temperature": 1, + "ignore_eos": true, + "dcp_size": 4, + "concurrency_levels": [ + 1, + 2, + 4, + 8 + ], + "context_lengths": [ + 8192, + 16384 + ], + "prefill_mode": "integrated_decode_scout" + }, + "prefill": { + "8192": { + "ttft_seconds": 3.062, + "prefill_seconds": 3.062, + "tok_per_sec": 2675, + "client_ttft_seconds": 3.062, + "client_tok_per_sec": 2675, + "prompt_tokens": 8192, + "samples": 1, + "method": "integrated_scout", + "server_validation": { + "method": "prometheus:kv_computed", + "tok_per_sec": 2697, + "prefill_seconds": 3.038, + "prompt_tokens": 8192, + "request_prompt_tokens": 0, + "cached_tokens": 0, + "token_source": "", + "samples": 1, + "invalid_reason": "" + }, + "hardware_summary": {} + }, + "16384": { + "ttft_seconds": 6.063, + "prefill_seconds": 6.063, + "tok_per_sec": 2702, + "client_ttft_seconds": 6.063, + "client_tok_per_sec": 2702, + "prompt_tokens": 16384, + "samples": 1, + "method": "integrated_scout", + "server_validation": { + "method": "prometheus:kv_computed", + "tok_per_sec": 2721, + "prefill_seconds": 6.022, + "prompt_tokens": 16384, + "request_prompt_tokens": 0, + "cached_tokens": 0, + "token_source": "", + "samples": 1, + "invalid_reason": "" + }, + "hardware_summary": {} + }, + "65536": { + "ttft_seconds": 23.609, + "prefill_seconds": 23.609, + "tok_per_sec": 2776, + "client_ttft_seconds": 23.609, + "client_tok_per_sec": 2776, + "prompt_tokens": 65536, + "samples": 1, + "method": "scout_only", + "server_validation": { + "method": "", + "tok_per_sec": 0, + "prefill_seconds": 0, + "prompt_tokens": 0, + "request_prompt_tokens": 0, + "cached_tokens": 0, + "token_source": "", + "samples": 0, + "invalid_reason": "" + }, + "hardware_summary": {} + }, + "131072": { + "ttft_seconds": 47.672, + "prefill_seconds": 47.672, + "tok_per_sec": 2749, + "client_ttft_seconds": 47.672, + "client_tok_per_sec": 2749, + "prompt_tokens": 131072, + "samples": 1, + "method": "scout_only", + "server_validation": { + "method": "", + "tok_per_sec": 0, + "prefill_seconds": 0, + "prompt_tokens": 0, + "request_prompt_tokens": 0, + "cached_tokens": 0, + "token_source": "", + "samples": 0, + "invalid_reason": "" + }, + "hardware_summary": {} + } + }, + "decode": [ + { + "context_tokens": 8192, + "concurrency": 1, + "aggregate_tps": 52.792233786884466, + "server_steps_per_s": 18.815052041581573, + "server_spec_accept_length": 2.8058510638297873, + "measurement_wall_seconds": 20.015, + "client_output_tokens": 1055, + "server_output_tokens": 1055, + "aggregate_source": "openai_continuous_usage", + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false, + "warmup_duration": 11.938 + }, + { + "context_tokens": 16384, + "concurrency": 1, + "aggregate_tps": 52.93204466912869, + "server_steps_per_s": 18.97941809801303, + "server_spec_accept_length": 2.788918205804749, + "measurement_wall_seconds": 20, + "client_output_tokens": 1057, + "server_output_tokens": 1057, + "aggregate_source": "openai_continuous_usage", + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false, + "warmup_duration": 13.39 + }, + { + "context_tokens": 8192, + "concurrency": 2, + "aggregate_tps": 77.08114068067856, + "server_steps_per_s": 29.268781636876643, + "server_spec_accept_length": 2.633561643835616, + "measurement_wall_seconds": 20, + "client_output_tokens": 1538, + "server_output_tokens": 1538, + "aggregate_source": "openai_continuous_usage", + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false, + "warmup_duration": 7.187 + }, + { + "context_tokens": 8192, + "concurrency": 4, + "aggregate_tps": 115.25, + "server_steps_per_s": 41.8, + "server_spec_accept_length": 2.757177033492823, + "measurement_wall_seconds": 20, + "client_output_tokens": 2305, + "server_output_tokens": 2305, + "aggregate_source": "openai_continuous_usage", + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false, + "warmup_duration": 8.203 + }, + { + "context_tokens": 8192, + "concurrency": 8, + "aggregate_tps": 173.45, + "server_steps_per_s": 61.076754890678934, + "server_spec_accept_length": 2.8398692810457513, + "measurement_wall_seconds": 20, + "client_output_tokens": 3469, + "server_output_tokens": 3476, + "aggregate_source": "openai_continuous_usage", + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false, + "warmup_duration": 12.219 + }, + { + "context_tokens": 16384, + "concurrency": 2, + "aggregate_tps": 76.45, + "server_steps_per_s": 29.000000000000004, + "server_spec_accept_length": 2.636206896551724, + "measurement_wall_seconds": 20, + "client_output_tokens": 1529, + "server_output_tokens": 1529, + "aggregate_source": "openai_continuous_usage", + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false, + "warmup_duration": 7.578 + }, + { + "context_tokens": 16384, + "concurrency": 4, + "aggregate_tps": 118.96395944170077, + "server_steps_per_s": 43.16835659065935, + "server_spec_accept_length": 2.755813953488372, + "measurement_wall_seconds": 20, + "client_output_tokens": 2370, + "server_output_tokens": 2370, + "aggregate_source": "openai_continuous_usage", + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false, + "warmup_duration": 9.219 + }, + { + "context_tokens": 16384, + "concurrency": 8, + "aggregate_tps": 171.9245283018868, + "server_steps_per_s": 61.18238993710692, + "server_spec_accept_length": 2.810032894736842, + "measurement_wall_seconds": 20, + "client_output_tokens": 3417, + "server_output_tokens": 3417, + "aggregate_source": "openai_continuous_usage", + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false, + "warmup_duration": 13.766 + } + ] + }, + { + "id": "124222", + "source_sha256": "4213df0fc47ba53308ea7dc076e3e1c77e9097e5eaef95db4e09ca4cd404e434", + "metadata": { + "version": "0.4.32", + "engine": "", + "model": "glm-5.3-flash-spark", + "decode_mode": "duration", + "duration_per_test": 20, + "decode_warmup_seconds": 5, + "decode_warmup_context": 16384, + "decode_warmup_concurrency": 1, + "cell_warmup_timeout_seconds": 900, + "max_tokens": 2048, + "temperature": 1, + "ignore_eos": true, + "dcp_size": 4, + "concurrency_levels": [ + 1, + 2, + 4, + 8 + ], + "context_lengths": [ + 8192, + 16384 + ], + "prefill_mode": "integrated_decode_scout" + }, + "prefill": { + "8192": { + "ttft_seconds": 3.062, + "prefill_seconds": 3.062, + "tok_per_sec": 2675, + "client_ttft_seconds": 3.062, + "client_tok_per_sec": 2675, + "prompt_tokens": 8192, + "samples": 1, + "method": "integrated_scout", + "server_validation": { + "method": "prometheus:kv_computed", + "tok_per_sec": 2697, + "prefill_seconds": 3.038, + "prompt_tokens": 8192, + "request_prompt_tokens": 0, + "cached_tokens": 0, + "token_source": "", + "samples": 1, + "invalid_reason": "" + }, + "hardware_summary": {} + }, + "16384": { + "ttft_seconds": 6.063, + "prefill_seconds": 6.063, + "tok_per_sec": 2702, + "client_ttft_seconds": 6.063, + "client_tok_per_sec": 2702, + "prompt_tokens": 16384, + "samples": 1, + "method": "integrated_scout", + "server_validation": { + "method": "prometheus:kv_computed", + "tok_per_sec": 2725, + "prefill_seconds": 6.013, + "prompt_tokens": 16384, + "request_prompt_tokens": 0, + "cached_tokens": 0, + "token_source": "", + "samples": 1, + "invalid_reason": "" + }, + "hardware_summary": {} + }, + "65536": { + "ttft_seconds": 23.578, + "prefill_seconds": 23.578, + "tok_per_sec": 2780, + "client_ttft_seconds": 23.578, + "client_tok_per_sec": 2780, + "prompt_tokens": 65536, + "samples": 1, + "method": "scout_only", + "server_validation": { + "method": "", + "tok_per_sec": 0, + "prefill_seconds": 0, + "prompt_tokens": 0, + "request_prompt_tokens": 0, + "cached_tokens": 0, + "token_source": "", + "samples": 0, + "invalid_reason": "" + }, + "hardware_summary": {} + }, + "131072": { + "ttft_seconds": 48.344, + "prefill_seconds": 48.344, + "tok_per_sec": 2711, + "client_ttft_seconds": 48.344, + "client_tok_per_sec": 2711, + "prompt_tokens": 131071, + "samples": 1, + "method": "scout_only", + "server_validation": { + "method": "", + "tok_per_sec": 0, + "prefill_seconds": 0, + "prompt_tokens": 0, + "request_prompt_tokens": 0, + "cached_tokens": 0, + "token_source": "", + "samples": 0, + "invalid_reason": "" + }, + "hardware_summary": {} + } + }, + "decode": [ + { + "context_tokens": 8192, + "concurrency": 1, + "aggregate_tps": 52.8, + "server_steps_per_s": 19.05, + "server_spec_accept_length": 2.7716535433070866, + "measurement_wall_seconds": 20, + "client_output_tokens": 1056, + "server_output_tokens": 1056, + "aggregate_source": "openai_continuous_usage", + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false, + "warmup_duration": 11.844 + }, + { + "context_tokens": 16384, + "concurrency": 1, + "aggregate_tps": 47.012390088058716, + "server_steps_per_s": 18.984812150331894, + "server_spec_accept_length": 2.4763157894736842, + "measurement_wall_seconds": 20.016, + "client_output_tokens": 941, + "server_output_tokens": 941, + "aggregate_source": "openai_continuous_usage", + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false, + "warmup_duration": 13.984 + }, + { + "context_tokens": 8192, + "concurrency": 2, + "aggregate_tps": 78.9, + "server_steps_per_s": 29.600000000000005, + "server_spec_accept_length": 2.6655405405405403, + "measurement_wall_seconds": 20, + "client_output_tokens": 1578, + "server_output_tokens": 1578, + "aggregate_source": "openai_continuous_usage", + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false, + "warmup_duration": 7.141 + }, + { + "context_tokens": 8192, + "concurrency": 4, + "aggregate_tps": 123.18949531411437, + "server_steps_per_s": 44.10364356241686, + "server_spec_accept_length": 2.793181818181818, + "measurement_wall_seconds": 20.016, + "client_output_tokens": 2458, + "server_output_tokens": 2458, + "aggregate_source": "openai_continuous_usage", + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false, + "warmup_duration": 9.156 + }, + { + "context_tokens": 8192, + "concurrency": 8, + "aggregate_tps": 168.77003205133715, + "server_steps_per_s": 60.89743589745578, + "server_spec_accept_length": 2.771381578947368, + "measurement_wall_seconds": 20.015, + "client_output_tokens": 3370, + "server_output_tokens": 3370, + "aggregate_source": "openai_continuous_usage", + "num_errors": 0, + "underfilled": false, + "warmup_timed_out": false, + "capacity_limited": false, + "warmup_duration": 12.813 + } + ] + } + ] +} diff --git a/performance/records/glm53-flash/test_mtp3_cache_history_record.py b/performance/records/glm53-flash/test_mtp3_cache_history_record.py new file mode 100644 index 00000000..6487c09f --- /dev/null +++ b/performance/records/glm53-flash/test_mtp3_cache_history_record.py @@ -0,0 +1,22 @@ +import gzip +import json +import statistics +from pathlib import Path + + +def test_mtp3_cache_history_numeric_observations(): + path = Path(__file__).with_name('mtp3-cache-history-observations.json.gz') + record = json.loads(gzip.decompress(path.read_bytes())) + turns = record['turns'] + assert len(turns) == 551 and all(t['valid'] for t in turns) + assert record['summary']['errors'] == 0 + soak = [t for t in turns if t['phase'] == 'soak'] + assert len(soak) == 545 + assert sum(t['tokenized_prompt_tokens'] for t in soak) == 70074391 + assert sum(t['cached_tokens_reported'] > 0 for t in soak) == 525 + for phase in ('before', 'after'): + probes = [t for t in turns if t['phase'] == phase] + assert len(probes) == 3 + expected = record['summary']['probes'][phase] + assert statistics.median(t['ttft_seconds'] for t in probes) == expected['median_ttft_seconds'] + assert statistics.median(t['decode_tokens_per_second_estimate'] for t in probes) == expected['median_decode_tokens_per_second_estimate'] diff --git a/recipes/README.md b/recipes/README.md index 43fa9978..be7f7ff8 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -10,6 +10,7 @@ installation guides. | Model profile | Status | Topology | Recipe | Operator guide | |---|---|---|---|---| +| GLM-5.3 native MTP3 with recurrent checkpoints and optimized SparkCache | research-only | four Sparks, TP4/DCP4 mesh | [Cache/checkpoint recipe](glm53-mtp3-cache-checkpoints-tp4.json) | [Cache/checkpoint quickstart](../docs/GLM53_MTP3_CACHE_CHECKPOINTS_QUICKSTART.md) | | GLM-5.3 Flash NVFP4-Spark + native MTP3 + managed mesh + SparkCache | research-only | four Sparks, TP4/DCP4, hardware-forwarded opposite peers | [`glm53-spark-mtp3-managed-mesh-tp4.json`](glm53-spark-mtp3-managed-mesh-tp4.json) | [Managed-mesh quickstart](../docs/GLM53_SPARK_MTP3_MESH_QUICKSTART.md) | | GLM-5.3 Flash NVFP4 + BF16 DFlash2 | implemented; DCP4 preferred | four Sparks, TP4 with DCP1/DCP2/DCP4 | [`glm53-flash-nvfp4-dflash2-bf16-tp4.json`](glm53-flash-nvfp4-dflash2-bf16-tp4.json) | [GLM-5.3 quickstart](../docs/GLM53_JJ_R8_GB10_SPARKCACHE_TP4_QUICKSTART.md) | | GLM-5.2 EXL3 3.5-bpw | implemented | four Sparks, TP4/DCP4 | [`glm52-exl3-r7-3.5bpw.json`](glm52-exl3-r7-3.5bpw.json) | [GLM-5.2 quickstart](../docs/GLM52_35BPW_QUICKSTART.md) | diff --git a/recipes/glm53-mtp3-cache-checkpoints-tp4.json b/recipes/glm53-mtp3-cache-checkpoints-tp4.json new file mode 100644 index 00000000..2ed4bfa4 --- /dev/null +++ b/recipes/glm53-mtp3-cache-checkpoints-tp4.json @@ -0,0 +1,150 @@ +{ + "schema": "sparkring-recipe/v1", + "recipe_id": "glm53-mtp3-cache-checkpoints-tp4", + "status": "research-only", + "hardware": { + "platform": "linux/arm64", + "cuda_arch": "sm_121", + "ranks": 4, + "topology": "direct-cycle-4", + "physical_ports_per_rank": 2, + "socket_direct_functions_per_port": 2, + "ethernet_mtu": 9000, + "roce_mtu": 4096, + "gid_index": 3, + "tested_nic_configuration": { + "firmware": "28.45.4028", + "eswitch_mode": "legacy", + "flow_steering_mode": "hmfs", + "hairpin_num_queues_per_function": 4, + "hairpin_queue_size_packets": 1024, + "hw_tc_offload": true + }, + "requirements": "Verified direct-neighbor RoCE and the canonical HCA/peer map; hardware-only ConnectX-7 source marking and intermediate forwarding. No physical diagonal cables are required." + }, + "model": { + "repository": "local-inference-lab/GLM-5.3-Flash-NVFP4-Spark", + "revision": "df116c4fb16b1d37ae43d2cfd624de26ffbc832e", + "config_sha256": "e1c0246a44ebefb5fd6383fb57aebbf7ac69ff6e7b23e989c0571b279a0eca23", + "weight_index_sha256": "db30fc7c5a70ccfb3b1c46637bb4ddb04226b95a5dfc451dffccb96a4f0ff544", + "quantization": "ModelOpt mixed NVFP4 routed experts and MXFP8 projections", + "external_draft_required": false + }, + "runtime": { + "operator_contract": "runtime/glm53-spark-mtp3-mesh/pins.json", + "public_image_contract": "runtime/glm53-spark-mtp3-mesh/performance/public-image.json", + "image_receipt": "runtime/glm53-spark-mtp3-mesh/performance/public-image.json", + "compute_contract": "runtime/glm53-spark-mtp3-mesh/compute/source-lock.json", + "compute_contract_sha256": "139f36701e0e47f45bf99fba2cc2fa59b417f2ee801dad3a064455d5b464a459", + "image": "ghcr.io/fujitsupolycom/sparkring-glm53-sparkcache@sha256:11a556a54041fd823d152a7f051ac4f7c617dc539030df26e93008392fee0746", + "image_id": "sha256:6921a6c163ea40b603e19a0332330efe3dbccbf4dce9f6cbbf6b756c9231835a", + "site_template": "runtime/glm53-spark-mtp3-mesh/site.example.json", + "fabric_template": "runtime/glm53-spark-mtp3-mesh/fabric.example.json", + "renderer": "runtime/glm53-spark-mtp3-mesh/profile.py", + "installer": "runtime/glm53-spark-mtp3-mesh/managed_install.py", + "lifecycle_controller": "runtime/glm53-spark-mtp3-mesh/managed_cluster.py", + "quickstart": "docs/GLM53_MTP3_CACHE_CHECKPOINTS_QUICKSTART.md" + }, + "serving_common": { + "served_model_name": "glm-5.3-flash-spark", + "tensor_parallel_size": 4, + "decode_context_parallel_size": 4, + "pipeline_parallel_size": 1, + "dtype": "bfloat16", + "max_model_len": 1048576, + "max_num_seqs": 16, + "max_num_batched_tokens": 8192, + "prefill_schedule_interval": 2, + "kv_cache_dtype": "fp8", + "kv_cache_memory_bytes_per_rank": 25769803776, + "async_scheduling": true, + "chunked_prefill": true, + "native_prefix_caching": true, + "attention_backend": "B12X", + "moe_backend": "B12X", + "linear_backend": "B12X", + "kda_prefill_backend": "b12x", + "cudagraph_capture_sizes": [ + 4, + 8, + 12, + 16, + 20, + 24, + 28, + 32, + 36, + 40, + 44, + 48, + 52, + 56, + 60, + 64 + ], + "speculation": { + "method": "mtp", + "num_speculative_tokens": 3, + "draft_tensor_parallel_size": 4, + "proposal_head": { + "weight_quantization": "runtime_nvfp4", + "activation_dtype": "bfloat16", + "checkpoint_source": "target lm_head.weight", + "persistent_allocation_bytes_per_rank": 89210880, + "target_verifier_head": "retained BF16 checkpoint representation", + "environment": { + "VLLM_MTP_NVFP4_LM_HEAD": "1", + "VLLM_LM_HEAD_A16": "1", + "VLLM_MXFP8_LM_HEAD": "0" + } + } + }, + "readiness_warmup_temperature": 1.0, + "readiness_warmup_thinking": false + }, + "sparkcache": { + "enabled": true, + "commit": "66057174301a4759ca3a45207ea41016689449cb", + "access_mode": "read-write", + "publication": "tail-cow-v2", + "capacity_bytes_per_rank": 42949672960, + "low_watermark_bytes": 34359738368, + "minimum_span_tokens": 4096, + "maximum_span_tokens": 1048576, + "shared_gpu_prefix_lease_seconds": 300, + "ttl_seconds": 0, + "cuda_restore": true, + "async_page_capture": true, + "capture_slots_per_rank": 2, + "capture_slot_bytes": 3221225472, + "namespace": "glm53-spark-mtp3-cache-checkpoints-c0fd5567-tail-cow-v2", + "identity_contract": "runtime/glm53-spark-mtp3-mesh/pins.json#/cache_identity", + "source_commit": "48bbd2be4a7b972e56632a2d7b934bac5460f272", + "periodic_full_capture_interval_tokens": 0 + }, + "transport": { + "bundle_manifest_sha256": "c0fd5567442b08b908cc193f36d0864e262573c7e5d232509479a823cface742", + "routing_contract": "spark_transport/experiments/glm53_rocenante_overlay/overlay_contract.json", + "captured_sircl_query_rows": [ + 16, + 20, + 24, + 28, + 32 + ], + "large_eager_prefill": "dual-rail fused SIRCL", + "small_collectives": "RoCEnante on admitted rows; other captured rows use SIRCL", + "fallback": "patched NCCL", + "rocenante_all_gather": false, + "managed_marker_expiry": null, + "reserved_rdma_udp_source_port": 65535, + "management_health_port": 9975, + "automatic_failure_restart": false, + "management_guide": "runtime/glm53-spark-mtp3-mesh/MANAGED_MESH.md" + }, + "evidence": { + "record": "runtime/glm53-spark-mtp3-mesh/performance/public-image.json", + "status": "research-only", + "scope": "File-verified image and 5308 deployed-runtime byte comparisons; not a serving soak of this rebuilt image." + } +} diff --git a/recipes/glm53-spark-mtp3-managed-mesh-tp4.json b/recipes/glm53-spark-mtp3-managed-mesh-tp4.json index 1d60217f..9d9006eb 100644 --- a/recipes/glm53-spark-mtp3-managed-mesh-tp4.json +++ b/recipes/glm53-spark-mtp3-managed-mesh-tp4.json @@ -34,8 +34,10 @@ "operator_contract": "runtime/glm53-spark-mtp3-mesh/pins.json", "public_image_contract": "runtime/glm53-spark-mtp3-mesh/public-image.json", "image_receipt": "runtime/glm53-spark-mtp3-mesh/image-receipt.json", - "image": "ghcr.io/fujitsupolycom/sparkring-glm53-sparkcache@sha256:23f00af873ccc784cfb742b7be2a29c6d3c20ebec9741843c025320bb9c04685", - "image_id": "sha256:26273b8e358df139ae913610a5d43084ff0fd08aafe282ef633a3bc74afefe47", + "compute_contract": "runtime/glm53-spark-mtp3-mesh/compute/source-lock.json", + "compute_contract_sha256": "139f36701e0e47f45bf99fba2cc2fa59b417f2ee801dad3a064455d5b464a459", + "image": "ghcr.io/fujitsupolycom/sparkring-glm53-sparkcache@sha256:67dc0ae453baaae6831ccec1d259b4ef8b236a8b0dc9f747d901b95c66ec1987", + "image_id": "sha256:2e41b1e934a85ff7c21b780532db2f0a0e978df081e52f4ae2bf11f8992fb24f", "site_template": "runtime/glm53-spark-mtp3-mesh/site.example.json", "fabric_template": "runtime/glm53-spark-mtp3-mesh/fabric.example.json", "renderer": "runtime/glm53-spark-mtp3-mesh/profile.py", @@ -62,8 +64,41 @@ "moe_backend": "B12X", "linear_backend": "B12X", "kda_prefill_backend": "b12x", - "cudagraph_capture_sizes": [4, 8, 12, 16, 20, 24, 28, 32, 36, 40, 44, 48, 52, 56, 60, 64], - "speculation": {"method": "mtp", "num_speculative_tokens": 3, "draft_tensor_parallel_size": 4}, + "cudagraph_capture_sizes": [ + 4, + 8, + 12, + 16, + 20, + 24, + 28, + 32, + 36, + 40, + 44, + 48, + 52, + 56, + 60, + 64 + ], + "speculation": { + "method": "mtp", + "num_speculative_tokens": 3, + "draft_tensor_parallel_size": 4, + "proposal_head": { + "weight_quantization": "runtime_nvfp4", + "activation_dtype": "bfloat16", + "checkpoint_source": "target lm_head.weight", + "persistent_allocation_bytes_per_rank": 89210880, + "target_verifier_head": "retained BF16 checkpoint representation", + "environment": { + "VLLM_MTP_NVFP4_LM_HEAD": "1", + "VLLM_LM_HEAD_A16": "1", + "VLLM_MXFP8_LM_HEAD": "0" + } + } + }, "readiness_warmup_temperature": 1.0, "readiness_warmup_thinking": false }, @@ -82,12 +117,19 @@ "async_page_capture": true, "capture_slots_per_rank": 2, "capture_slot_bytes": 3221225472, + "namespace": "glm53-spark-df116c4f-mtp3-nvfp4-a16-c139f3670-mesh69313e19-tail-cow-v2", "identity_contract": "runtime/glm53-spark-mtp3-mesh/pins.json#/cache_identity" }, "transport": { - "bundle_manifest_sha256": "4204fabc93303226b9a120b094ef3c82ed4aadd1d7f97cfbe291204c027ed45f", + "bundle_manifest_sha256": "69313e19e881ec93e9ed3bd150d2f24fc6b444488ac729a69f45d038e2243500", "routing_contract": "spark_transport/experiments/glm53_rocenante_overlay/overlay_contract.json", - "captured_sircl_query_rows": [16, 20, 24, 28, 32], + "captured_sircl_query_rows": [ + 16, + 20, + 24, + 28, + 32 + ], "large_eager_prefill": "dual-rail fused SIRCL", "small_collectives": "RoCEnante on admitted rows; other captured rows use SIRCL", "fallback": "patched NCCL", @@ -99,14 +141,14 @@ "management_guide": "runtime/glm53-spark-mtp3-mesh/MANAGED_MESH.md" }, "evidence": { - "record": "performance/records/glm53-flash/spark-mtp3-managed-mesh-functional-20260905.md", - "status": "qualified", - "scope": "Bounded native correctness, managed lifecycle, model startup and persistent recall checks for the recorded image and conditions", + "record": "performance/records/glm53-flash/spark-mtp3-compute-stream-safety-20260906.md", + "proposal_head_performance_record": "performance/records/glm53-flash/spark-mtp3-nvfp4-proposal-head-20260905.md", + "status": "research-only", + "scope": "Exact-image native all-reduce, selected GPU stream invariants, bounded serving matrix, idle rank-loss shutdown, restart and one persistent-prefix recall.", "limitations": [ - "This opt-in profile does not replace the recommended DFlash/SIRCL profile.", - "No host-reboot, prolonged-soak, unattended high-availability or general cache/model-quality qualification is claimed.", - "Fault detection has a bounded observation window; partial streamed output can precede shutdown.", - "Readiness warmup covers temperature-one sampling with thinking disabled, not all mixed-prefill/decode shapes." + "Idle rank-loss containment does not establish in-flight GPU failure containment or resolve upstream RoCEnante #313.", + "One persistent recall fixture does not establish all cache boundaries or model accuracy.", + "Serving measurements are observations, not a guaranteed speedup." ] } } diff --git a/runtime/glm53-flash-jj-r8-gb10/README.md b/runtime/glm53-flash-jj-r8-gb10/README.md index 95bfd2c4..bee09a47 100644 --- a/runtime/glm53-flash-jj-r8-gb10/README.md +++ b/runtime/glm53-flash-jj-r8-gb10/README.md @@ -178,7 +178,8 @@ ports 19006/19007 and secondary ports 19106/19107. The derivation reserves two ports for each admitted capacity Q1024/Q2048/Q4096/Q8192. SIRCL's two transport slots are independent from SparkCache's two 3-GiB -asynchronous page-capture slots and two 256-MiB restore arenas. +asynchronous page-capture slots and sixteen 256-MiB restore arenas (two for +each of eight load lanes). #### Recorded functional evidence @@ -228,6 +229,11 @@ With the connector enabled, `SPARKCACHE_ACCESS_MODE=read-write` restores and publishes persistent entries. `restore-only` reuses compatible entries but does not capture or publish new prompt state. Missing entries are computed by vLLM normally. `store-only` and `disabled` are diagnostic modes. +The environment template sets `SPARKCACHE_ASYNC_PAGE_CAPTURE=auto`: capture +is enabled for `read-write` and `store-only`, and disabled for `restore-only`, +`disabled`, or `SPARKCACHE_ENABLED=0`. Explicit `1` still rejects a mode that +cannot publish. Explicit `0` uses synchronous publication in publishing modes. +The launcher defaults to `0` when no capture setting is supplied. ### Choose the persistent publication format @@ -277,14 +283,43 @@ not discard Triton, TorchInductor, B12X, or vLLM compilation caches. Each rank keeps its own persistent copy under `CACHE_HOST_ROOT`; the four ranks do not write to one network-shared compilation directory. -Set `SPARKCACHE_ASYNC_PAGE_CAPTURE=1` to capture manager pages through the +Set `SPARKCACHE_ASYNC_PAGE_CAPTURE=auto` to capture manager pages through the bounded CUDA ring. `SPARKCACHE_ASYNC_CAPTURE_SLOT_BYTES` defaults to 8 GiB for DCP1, 5 GiB for DCP2, and 3 GiB for DCP4. The DCP4 profile uses two 3 GiB capture slots, so the background publisher can consume one completed capture while a later capture uses the other. Restore separately pipelines bounded -NVMe reads and CUDA placement through two 256 MiB mapped arenas. A third arena -is not part of the profile because the two-stage pipeline has no measured -arena wait that would justify more unified-memory pressure. +NVMe reads and CUDA placement through two 256 MiB mapped arenas **per load +lane**. Eight lanes reserve 4 GiB per rank for restore payloads. With the 6 GiB +capture ring, the DCP4 profile configures 10 GiB per rank (40 GiB across TP4), +in addition to the 24 GiB per-rank KV allocation. Restore-only configures +4 GiB per rank and no capture slots. These figures describe payload capacities; +control arrays, Python objects, shared bases, transport, model weights and +allocator overhead require additional memory. `SPARKCACHE_LOAD_THREADS` defaults +to eight; throughput and memory-pressure effects need hardware measurements. + +### Inspect configured memory before launch + +Status: **implemented**. The launcher can print a JSON allocation plan without +Docker, GPUs, checkpoint files, or cache directories. It sources the same trusted +shell configuration as a launch and resolves its byte counts and capture mode: + +```bash +SPARKRING_PRINT_MEMORY_PLAN=1 bash runtime/glm53-flash-jj-r8-gb10/launch-rank.sh \ + 0 runtime/glm53-flash-jj-r8-gb10/runtime.env.example +``` + +Set `SPARKCACHE_BUFFER_BUDGET_BYTES` in the configuration to reject restore +and capture payload capacities above that per-rank ceiling before host checks +or Docker access. Zero, the default, disables the ceiling. For example, +`10737418240` admits the 10 GiB DCP4 read-write configuration exactly. The +report includes per-rank and topology totals; DCP shards state but does not +reduce the physical TP rank count. A normal launch also logs the plan. + +This ceiling does not cap total process memory or predict whether serving fits. +It excludes model weights, KV allocation, transient reads, retained shared bases, +CUDA control arrays, compilation and transport buffers, and allocator overhead. +The report lists KV separately. A passing offline plan does not qualify a CUDA +allocation or a serving performance result. When `DFLASH_WARMUP=1`, the readiness entrypoint runs `warmup_dflash.py` before Docker reports rank 0 as healthy. The readiness wrapper, diff --git a/runtime/glm53-flash-jj-r8-gb10/launch-rank.sh b/runtime/glm53-flash-jj-r8-gb10/launch-rank.sh index 01785330..3c35b5a6 100644 --- a/runtime/glm53-flash-jj-r8-gb10/launch-rank.sh +++ b/runtime/glm53-flash-jj-r8-gb10/launch-rank.sh @@ -34,6 +34,11 @@ case "${SPARKRING_OFFLINE_SPEC}" in 1) [[ "${SPARKRING_PRINT_CONTAINER_SPEC}" == 1 ]] || { printf 'offline rendering requires SPARKRING_PRINT_CONTAINER_SPEC=1\n' >&2; exit 78; } ;; *) printf 'SPARKRING_OFFLINE_SPEC must be 0 or 1\n' >&2; exit 78 ;; esac +: "${SPARKRING_PRINT_MEMORY_PLAN:=0}" +case "${SPARKRING_PRINT_MEMORY_PLAN}" in + 0|1) ;; + *) printf 'SPARKRING_PRINT_MEMORY_PLAN must be 0 or 1\n' >&2; exit 78 ;; +esac case "${SPARKRING_PRINT_CONTAINER_SPEC}" in 0|1) ;; *) printf 'SPARKRING_PRINT_CONTAINER_SPEC must be 0 or 1\n' >&2; exit 78 ;; @@ -138,6 +143,7 @@ esac : "${SPARKCACHE_ASYNC_PAGE_CAPTURE:=0}" : "${SPARKCACHE_ASYNC_CAPTURE_SLOT_BYTES:=auto}" : "${SPARKCACHE_ASYNC_CAPTURE_SLOT_COUNT:=2}" +: "${SPARKCACHE_BUFFER_BUDGET_BYTES:=0}" : "${SPARKCACHE_SOURCE_OVERLAY:=}" : "${VLLM_KV_METRICS_OVERLAY:=}" : "${MULTIMODAL_INPUTS:=1}" @@ -190,6 +196,7 @@ do done require_uint SPARKCACHE_LOW_WATERMARK_BYTES require_uint SPARKCACHE_TTL_SECONDS +require_uint SPARKCACHE_BUFFER_BUDGET_BYTES require_uint NCCL_IB_GID_INDEX require_uint MAX_IMAGES_PER_PROMPT require_uint MAX_VIDEOS_PER_PROMPT @@ -303,13 +310,21 @@ case "${ENABLE_PROMPT_TOKENS_DETAILS}" in *) die 'ENABLE_PROMPT_TOKENS_DETAILS must be 0 or 1' ;; esac case "${SPARKCACHE_ASYNC_PAGE_CAPTURE}" in - 0|1) ;; - *) die 'SPARKCACHE_ASYNC_PAGE_CAPTURE must be 0 or 1' ;; + auto|0|1) ;; + *) die 'SPARKCACHE_ASYNC_PAGE_CAPTURE must be auto, 0, or 1' ;; esac case "${SPARKCACHE_ACCESS_MODE}" in read-write|restore-only|store-only|disabled) ;; *) die 'SPARKCACHE_ACCESS_MODE must be read-write, restore-only, store-only, or disabled' ;; esac +if [[ "${SPARKCACHE_ASYNC_PAGE_CAPTURE}" == auto ]]; then + SPARKCACHE_ASYNC_PAGE_CAPTURE=0 + if [[ "${SPARKCACHE_ENABLED}" == 1 ]]; then + case "${SPARKCACHE_ACCESS_MODE}" in + read-write|store-only) SPARKCACHE_ASYNC_PAGE_CAPTURE=1 ;; + esac + fi +fi if [[ "${SPARKCACHE_ASYNC_PAGE_CAPTURE}" == 1 ]]; then [[ "${SPARKCACHE_ENABLED}" == 1 ]] || \ die 'asynchronous page capture requires SPARKCACHE_ENABLED=1' @@ -318,6 +333,69 @@ if [[ "${SPARKCACHE_ASYNC_PAGE_CAPTURE}" == 1 ]]; then *) die 'asynchronous page capture requires a publication-capable access mode' ;; esac fi + +# Resolve the payload buffers before inspecting checkpoints or contacting Docker. +# Python integers avoid overflow when comparing operator-supplied byte budgets. +command -v python3 >/dev/null 2>&1 || die 'python3 is required to resolve the memory plan' +export SPARKCACHE_ENABLED SPARKCACHE_ACCESS_MODE SPARKCACHE_LOAD_THREADS +export SPARKCACHE_CUDA_ARENA_BYTES SPARKCACHE_ASYNC_PAGE_CAPTURE +export SPARKCACHE_ASYNC_CAPTURE_SLOT_BYTES SPARKCACHE_ASYNC_CAPTURE_SLOT_COUNT +export SPARKCACHE_BUFFER_BUDGET_BYTES KV_CACHE_MEMORY_BYTES +export TENSOR_PARALLEL_SIZE PIPELINE_PARALLEL_SIZE DECODE_CONTEXT_PARALLEL_SIZE +memory_plan="$(python3 - <<'PY' +import json +import os +import sys + +def integer(name): + return int(os.environ[name]) + +enabled = os.environ["SPARKCACHE_ENABLED"] == "1" +mode = os.environ["SPARKCACHE_ACCESS_MODE"] +# The pinned manager-page connector caps active placement lanes at eight. +requested_lanes = integer("SPARKCACHE_LOAD_THREADS") +lanes = min(8, requested_lanes) if enabled and mode in ("read-write", "restore-only") else 0 +slots = integer("SPARKCACHE_ASYNC_CAPTURE_SLOT_COUNT") if os.environ["SPARKCACHE_ASYNC_PAGE_CAPTURE"] == "1" else 0 +restore = lanes * 2 * integer("SPARKCACHE_CUDA_ARENA_BYTES") +capture = slots * integer("SPARKCACHE_ASYNC_CAPTURE_SLOT_BYTES") +buffers = restore + capture +budget = integer("SPARKCACHE_BUFFER_BUDGET_BYTES") +ranks = integer("TENSOR_PARALLEL_SIZE") * integer("PIPELINE_PARALLEL_SIZE") +kv = integer("KV_CACHE_MEMORY_BYTES") +print(json.dumps({ + "status": "implemented", + "basis": "configured payload capacities; not measured resident memory", + "access_mode": mode if enabled else "disabled", + "async_page_capture": slots > 0, + "dcp_degree": integer("DECODE_CONTEXT_PARALLEL_SIZE"), + "rank_count": ranks, + "requested_restore_lanes_per_rank": requested_lanes, + "restore_lanes_per_rank": lanes, + "arenas_per_restore_lane": 2, + "capture_slots_per_rank": slots, + "restore_payload_bytes_per_rank": restore, + "capture_payload_bytes_per_rank": capture, + "sparkcache_payload_bytes_per_rank": buffers, + "sparkcache_payload_bytes_all_ranks": buffers * ranks, + "buffer_budget_bytes_per_rank": budget or None, + "within_buffer_budget": budget == 0 or buffers <= budget, + "kv_cache_bytes_per_rank": kv, + "kv_and_payload_bytes_per_rank": kv + buffers, + "kv_and_payload_bytes_all_ranks": (kv + buffers) * ranks, + "excluded": ["model weights", "CUDA control arrays", "Python objects and read buffers", + "shared base retention", "transport", "compilation workspaces", "allocator overhead"], +}, sort_keys=True)) +if budget and buffers > budget: + print(f"SparkCache payload buffers require {buffers} bytes per rank, exceeding " + f"SPARKCACHE_BUFFER_BUDGET_BYTES={budget}", file=sys.stderr) + sys.exit(78) +PY +)" || { printf '%s\n' "${memory_plan}"; exit 78; } +if [[ "${SPARKRING_PRINT_MEMORY_PLAN}" == 1 ]]; then + printf '%s\n' "${memory_plan}" + exit 0 +fi +printf 'sparkcache: memory_plan %s\n' "${memory_plan}" >&2 if [[ -n "${CHAT_TEMPLATE_HOST_PATH}" ]]; then [[ "${CHAT_TEMPLATE_HOST_PATH}" == /* ]] || \ die 'CHAT_TEMPLATE_HOST_PATH must be an absolute host path when set' @@ -770,6 +848,8 @@ kv_transfer_args=() if [[ "${SPARKCACHE_ENABLED}" == 1 ]]; then export SPARKCACHE_CACHE_NAMESPACE SPARKCACHE_CLEAR_ONCE SPARKCACHE_MAX_BYTES export SPARKCACHE_ACCESS_MODE + export SPARKCACHE_PLACEMENT_LIBRARY_SHA256="${SPARKCACHE_PLACEMENT_LIBRARY_SHA256:-d57509052b73853bcc8e3c3f47bb81748d87b9cbd8d908fc20d4c79a09aa400c}" + [[ "${SPARKCACHE_PLACEMENT_LIBRARY_SHA256}" =~ ^[0-9a-f]{64}$ ]] || die 'SPARKCACHE_PLACEMENT_LIBRARY_SHA256 must be a SHA-256 digest' export SPARKCACHE_SHARED_PREFIX_LEASE_TTL_SECONDS export SPARKCACHE_PUBLICATION_SCHEMA export SPARKCACHE_LOW_WATERMARK_BYTES SPARKCACHE_TTL_SECONDS @@ -806,7 +886,7 @@ extra = { "spark_cache_min_span_tokens": integer("SPARKCACHE_MIN_SPAN_TOKENS"), "spark_cache_max_span_tokens": integer("SPARKCACHE_MAX_SPAN_TOKENS"), "spark_cache_cuda_placement_library": "/opt/sparkcache-src/sparkcache/native/build-cuda/libspark_cache_placement.so", - "spark_cache_cuda_placement_library_sha256": "d57509052b73853bcc8e3c3f47bb81748d87b9cbd8d908fc20d4c79a09aa400c", + "spark_cache_cuda_placement_library_sha256": os.environ["SPARKCACHE_PLACEMENT_LIBRARY_SHA256"], "spark_cache_cuda_placement_arena_bytes": integer("SPARKCACHE_CUDA_ARENA_BYTES"), "spark_cache_cuda_restore_io_workers": integer("SPARKCACHE_CUDA_RESTORE_IO_WORKERS"), "spark_cache_load_threads": integer("SPARKCACHE_LOAD_THREADS"), diff --git a/runtime/glm53-flash-jj-r8-gb10/runtime.env.example b/runtime/glm53-flash-jj-r8-gb10/runtime.env.example index 0931cda2..4b59df75 100644 --- a/runtime/glm53-flash-jj-r8-gb10/runtime.env.example +++ b/runtime/glm53-flash-jj-r8-gb10/runtime.env.example @@ -176,9 +176,15 @@ SPARKCACHE_LOAD_THREADS=8 SPARKCACHE_MAX_PENDING_RESTORES=8 SPARKCACHE_CUDA_RESTORE_IO_WORKERS=8 SPARKCACHE_CUDA_ARENA_BYTES=268435456 +# Each load lane owns two mapped arenas: 8 * 2 * 256 MiB = 4 GiB per rank. +# Optional per-rank ceiling for restore and capture payload buffers, in bytes. +# Zero disables the ceiling. This excludes models, KV, control arrays and reads. +SPARKCACHE_BUFFER_BUDGET_BYTES=0 # Two bounded capture slots let the background publisher consume one capture -# while a later request uses the other. Set to 0 for synchronous publication. -SPARKCACHE_ASYNC_PAGE_CAPTURE=1 +# while a later request uses the other. Auto enables capture only for enabled +# read-write/store-only modes; restore-only and disabled modes allocate no slots. +# Set to 0 for synchronous publication or 1 to require asynchronous capture. +SPARKCACHE_ASYNC_PAGE_CAPTURE='auto' SPARKCACHE_ASYNC_CAPTURE_SLOT_BYTES='auto' SPARKCACHE_ASYNC_CAPTURE_SLOT_COUNT=2 diff --git a/runtime/glm53-flash-jj-r8-gb10/test_image_contract.py b/runtime/glm53-flash-jj-r8-gb10/test_image_contract.py index 0a4a8570..25bcbe36 100644 --- a/runtime/glm53-flash-jj-r8-gb10/test_image_contract.py +++ b/runtime/glm53-flash-jj-r8-gb10/test_image_contract.py @@ -470,7 +470,7 @@ def test_launcher_keeps_gather_workspace_below_native_context_limit() -> None: "KV_CACHE_MEMORY_BYTES='auto'", "B12X_MLA_CKV_GATHER_MAX_TOKENS=524288", "SPARKCACHE_MAX_SPAN_TOKENS=1048576", - "SPARKCACHE_ASYNC_PAGE_CAPTURE=1", + "SPARKCACHE_ASYNC_PAGE_CAPTURE='auto'", "SPARKCACHE_ASYNC_CAPTURE_SLOT_BYTES='auto'", "SPARKCACHE_ASYNC_CAPTURE_SLOT_COUNT=2", "SPARKCACHE_SHARED_PREFIX_LEASE_TTL_SECONDS=300", diff --git a/runtime/glm53-flash-jj-r8-gb10/test_launcher_contract.py b/runtime/glm53-flash-jj-r8-gb10/test_launcher_contract.py index 333e5798..279ae9f4 100644 --- a/runtime/glm53-flash-jj-r8-gb10/test_launcher_contract.py +++ b/runtime/glm53-flash-jj-r8-gb10/test_launcher_contract.py @@ -58,6 +58,7 @@ def test_environment_exposes_reproducible_operator_defaults() -> None: assert values["SPARKCACHE_ENABLED"] == "1" assert values["ENABLE_PROMPT_TOKENS_DETAILS"] == "1" assert values["SPARKCACHE_ACCESS_MODE"] == "read-write" + assert values["SPARKCACHE_ASYNC_PAGE_CAPTURE"] == "auto" assert values["SPARKCACHE_SHARED_PREFIX_LEASE_TTL_SECONDS"] == "300" assert values["SPARKCACHE_CACHE_NAMESPACE"] == ( "glm53-flash-vllm-e02b1746-b12x-9ae41c5c-" @@ -74,8 +75,110 @@ def test_environment_exposes_reproducible_operator_defaults() -> None: ) +def _memory_plan(tmp_path: Path, *settings: str) -> subprocess.CompletedProcess[str]: + config = tmp_path / "memory-plan.env" + # Nonexistent checkpoint and cache paths prove this mode needs no model, + # cache directories, Docker daemon, GPU, or serving host. + config.write_text( + "\n".join(( + f"source '{_bash_path(ENVIRONMENT)}'", + "HOST_IP=rank0.example.net", + "MASTER_ADDR=rank0.example.net", + "TARGET_MODEL_HOST_PATH=/nonexistent-memory-plan/target", + "DFLASH_MODEL_HOST_PATH=/nonexistent-memory-plan/draft", + "CACHE_HOST_ROOT=/nonexistent-memory-plan/cache", + "SPARKRING_PRINT_MEMORY_PLAN=1", + *settings, + )), + encoding="utf-8", + newline="\n", + ) + return subprocess.run( + ["bash", _bash_path(LAUNCHER), "0", _bash_path(config)], + cwd=ROOT, text=True, capture_output=True, check=False, + ) + + +@pytest.mark.parametrize("dcp,slot_gib", [(1, 8), (2, 5), (4, 3)]) +def test_memory_plan_resolves_dcp_and_all_lane_allocations( + tmp_path: Path, dcp: int, slot_gib: int, +) -> None: + result = _memory_plan(tmp_path, f"DECODE_CONTEXT_PARALLEL_SIZE={dcp}") + assert result.returncode == 0, result.stderr + plan = json.loads(result.stdout) + gib = 1024 ** 3 + assert plan["restore_payload_bytes_per_rank"] == 4 * gib + assert plan["capture_payload_bytes_per_rank"] == 2 * slot_gib * gib + assert plan["sparkcache_payload_bytes_all_ranks"] == (4 + 2 * slot_gib) * gib * 4 + assert plan["kv_and_payload_bytes_all_ranks"] == (24 + 4 + 2 * slot_gib) * gib * 4 + assert plan["buffer_budget_bytes_per_rank"] is None + + +@pytest.mark.parametrize("mode,enabled,restore_gib,capture_gib", [ + ("read-write", 1, 4, 6), ("restore-only", 1, 4, 0), + ("store-only", 1, 0, 6), ("disabled", 1, 0, 0), + ("read-write", 0, 0, 0), +]) +def test_auto_capture_respects_access_mode( + tmp_path: Path, mode: str, enabled: int, restore_gib: int, capture_gib: int, +) -> None: + result = _memory_plan(tmp_path, f"SPARKCACHE_ACCESS_MODE={mode}", f"SPARKCACHE_ENABLED={enabled}") + assert result.returncode == 0, result.stderr + plan = json.loads(result.stdout) + assert plan["restore_payload_bytes_per_rank"] == restore_gib * 1024 ** 3 + assert plan["capture_payload_bytes_per_rank"] == capture_gib * 1024 ** 3 + assert plan["async_page_capture"] == bool(capture_gib) + + +@pytest.mark.parametrize("setting,error", [ + ("SPARKCACHE_ACCESS_MODE=restore-only", "publication-capable access mode"), + ("SPARKCACHE_ACCESS_MODE=disabled", "publication-capable access mode"), + ("SPARKCACHE_ENABLED=0", "requires SPARKCACHE_ENABLED=1"), +]) +def test_explicit_capture_contradictions_remain_errors( + tmp_path: Path, setting: str, error: str, +) -> None: + result = _memory_plan(tmp_path, setting, "SPARKCACHE_ASYNC_PAGE_CAPTURE=1") + assert result.returncode == 78 + assert error in result.stderr + + +def test_buffer_budget_rejects_overcommit_before_host_checks(tmp_path: Path) -> None: + required = 10 * 1024 ** 3 + result = _memory_plan(tmp_path, f"SPARKCACHE_BUFFER_BUDGET_BYTES={required - 1}") + assert result.returncode == 78 + assert "exceeding SPARKCACHE_BUFFER_BUDGET_BYTES" in result.stderr + assert json.loads(result.stdout)["within_buffer_budget"] is False + accepted = _memory_plan(tmp_path, f"SPARKCACHE_BUFFER_BUDGET_BYTES={required}") + assert accepted.returncode == 0, accepted.stderr + assert json.loads(accepted.stdout)["within_buffer_budget"] is True + + +def test_explicit_synchronous_capture_and_custom_lane_sizes(tmp_path: Path) -> None: + result = _memory_plan(tmp_path, "SPARKCACHE_ASYNC_PAGE_CAPTURE=0", + "SPARKCACHE_LOAD_THREADS=3", "SPARKCACHE_CUDA_ARENA_BYTES=1048576") + assert result.returncode == 0, result.stderr + plan = json.loads(result.stdout) + assert plan["capture_payload_bytes_per_rank"] == 0 + assert plan["restore_payload_bytes_per_rank"] == 6 * 1048576 + + +def test_memory_plan_uses_connector_lane_limit(tmp_path: Path) -> None: + result = _memory_plan(tmp_path, "SPARKCACHE_LOAD_THREADS=64") + assert result.returncode == 0, result.stderr + plan = json.loads(result.stdout) + assert plan["requested_restore_lanes_per_rank"] == 64 + assert plan["restore_lanes_per_rank"] == 8 + assert plan["restore_payload_bytes_per_rank"] == 4 * 1024 ** 3 + + +@pytest.mark.parametrize("capture_mode,access_mode,capture_enabled", [ + ("0", "read-write", False), + ("auto", "read-write", True), + ("auto", "restore-only", False), +]) def test_launcher_resolves_dcp_profiles_and_prompt_token_details( - tmp_path: Path, + tmp_path: Path, capture_mode: str, access_mode: str, capture_enabled: bool, ) -> None: subprocess.run(["bash", "-n", _bash_path(LAUNCHER)], check=True, cwd=ROOT) fake_bin = tmp_path / "bin" @@ -148,6 +251,8 @@ def test_launcher_resolves_dcp_profiles_and_prompt_token_details( "IMAGE_REF=test-image:r8", f"IMAGE_ID={IMAGE_ID}", f"DECODE_CONTEXT_PARALLEL_SIZE={dcp}", + f"SPARKCACHE_ASYNC_PAGE_CAPTURE={capture_mode}", + f"SPARKCACHE_ACCESS_MODE={access_mode}", ) ), encoding="utf-8", @@ -201,7 +306,8 @@ def test_launcher_resolves_dcp_profiles_and_prompt_token_details( extra = connector["kv_connector_extra_config"] assert extra["spark_cache_publication_schema"] == "tail-cow-v2" assert extra["spark_cache_model_profile"] == "glm53-flash-hybrid" - assert extra["spark_cache_access_mode"] == "read-write" + assert extra["spark_cache_access_mode"] == access_mode + assert extra["spark_cache_async_page_capture"] is capture_enabled assert extra["spark_cache_shared_prefix_lease_ttl_seconds"] == 300 assert "spark_cache_store" not in extra assert "spark_cache_restore" not in extra diff --git a/runtime/glm53-spark-mtp3-mesh/Dockerfile b/runtime/glm53-spark-mtp3-mesh/Dockerfile index 6a4982aa..86053fa1 100644 --- a/runtime/glm53-spark-mtp3-mesh/Dockerfile +++ b/runtime/glm53-spark-mtp3-mesh/Dockerfile @@ -5,6 +5,18 @@ ARG PARENT_IMAGE ARG PARENT_IMAGE_ID ARG BUNDLE_MANIFEST_SHA256 ARG SOURCE_RECEIPT_SHA256 +ARG COMPUTE_SOURCE_LOCK_SHA256 + +COPY compute/ /opt/sparkring-compute/ +ENV CUDA_HOME=/opt/cuda-13.3 CUDA_PATH=/opt/cuda-13.3 +ENV PATH="/opt/cuda-13.3/bin:${PATH}" +ENV LD_LIBRARY_PATH="/opt/cuda-13.3/lib:/opt/cuda-13.3/lib64:${LD_LIBRARY_PATH}" +ENV TRITON_PTXAS_PATH=/opt/cuda-13.3/bin/ptxas +ENV VLLM_GDN_SPEC_DECODE_METADATA_FASTPATH=1 VLLM_B12X_DENSE_ACTIVATION_MODE=auto +ENV VLLM_MTP_NVFP4_LM_HEAD=1 VLLM_LM_HEAD_A16=1 VLLM_MXFP8_LM_HEAD=0 +RUN python3 /opt/sparkring-compute/apply_compute.py \ + --prepared /opt/sparkring-compute \ + --site-packages /usr/local/lib/python3.12/dist-packages COPY bundle/ /opt/spark-sircl/ COPY receipts/ /opt/sparkring/receipts/glm53-spark-mtp3-mesh/ @@ -27,4 +39,8 @@ LABEL org.opencontainers.image.title="SparkRing GLM-5.3 Spark MTP3 mesh runtime" org.sparkring.mesh.bundle-manifest-sha256="${BUNDLE_MANIFEST_SHA256}" \ org.sparkring.mesh.source-receipt-sha256="${SOURCE_RECEIPT_SHA256}" \ org.sparkring.mesh.default-speculation="mtp3" \ + org.sparkring.compute.source-lock-sha256="${COMPUTE_SOURCE_LOCK_SHA256}" \ + org.sparkring.b12x.composition="ef308bac0f3b3eb8fea63e4013afc0c2ea1c6301+pr316-9ac14282" \ + org.sparkring.vllm.compute-source="${COMPUTE_SOURCE_LOCK_SHA256}" \ + org.sparkring.mesh.proposal-head="nvfp4-a16" \ org.sparkring.sircl.manifest-sha256="${BUNDLE_MANIFEST_SHA256}" diff --git a/runtime/glm53-spark-mtp3-mesh/IMAGE_BUILD.md b/runtime/glm53-spark-mtp3-mesh/IMAGE_BUILD.md index 0b7ce6ff..8dbb8229 100644 --- a/runtime/glm53-spark-mtp3-mesh/IMAGE_BUILD.md +++ b/runtime/glm53-spark-mtp3-mesh/IMAGE_BUILD.md @@ -1,14 +1,24 @@ -# ARM64 image with the MTP3 mesh bundle - -Status: **research-only**. This child image packages the transport -files used by the GLM-5.3 Spark native-MTP3 profile. It does not change the -parent image's vLLM, B12X model kernels, SparkCache, or NCCL. It replaces the -readiness warmup helper and sets `SPARKRING_WARMUP_TEMPERATURE=1`, so requests -issued before readiness use temperature one. This is an explicit startup -behavior override, not a claim that every parent entrypoint component is -unchanged. The managed quickstart requires this child and its verified -receipt. The parent-image rendering mode is a separate composition interface; -it does not meet the managed marker and temperature-one warmup contract. +# ARM64 image with native-MTP3 compute and mesh bundle + +Status: **research-only**. This child image packages the complete compute and +transport composition used by the GLM-5.3 Spark native-MTP3 profile. On top of +the pinned parent, it installs CUDA 13.3, the checksum-bound vLLM metadata and +proposal-head and loader/RNG patches, B12X revision `ef308bac` with the +source-checked top-k selector, the transport bundle, +the managed marker, and the readiness helper. The proposal head uses runtime +NVFP4 weights with BF16 activations; the target/verifier head retains its BF16 +checkpoint representation. SparkCache and patched NCCL remain inherited. + +The image sets `SPARKRING_WARMUP_TEMPERATURE=1`, so requests issued before +readiness use temperature one. The managed quickstart requires this child and +its verified receipt. The parent-image rendering mode is a separate composition +interface; it does not meet the compute, managed marker, or temperature-one +warmup contract. + +For the pinned GLM vocabulary and TP4 geometry, the proposal head adds about +85.08 MiB of packed NVFP4 values and scales per rank. It does not replace the +retained BF16 target/verifier head, so this is an additional persistent +allocation rather than a net model-memory reduction. The image contains no model weights. It does not provision NIC rules, select network interfaces, install host services, or start a model during construction. @@ -17,16 +27,24 @@ The site plan and native-MTP3 launch configuration remain separate inputs. ## Published image The [public registry receipt](public-image.json) binds the published manifest -to the tested Linux/ARM64 image. Pull before using its local image ID: +to its Linux/ARM64 config-image identity. Pull before using its local image ID: ```bash set -euo pipefail -mtp_image='ghcr.io/fujitsupolycom/sparkring-glm53-sparkcache@sha256:23f00af873ccc784cfb742b7be2a29c6d3c20ebec9741843c025320bb9c04685' -mtp_image_id='sha256:26273b8e358df139ae913610a5d43084ff0fd08aafe282ef633a3bc74afefe47' +mtp_image='ghcr.io/fujitsupolycom/sparkring-glm53-sparkcache@sha256:67dc0ae453baaae6831ccec1d259b4ef8b236a8b0dc9f747d901b95c66ec1987' +mtp_image_id='sha256:2e41b1e934a85ff7c21b780532db2f0a0e978df081e52f4ae2bf11f8992fb24f' docker pull "$mtp_image" test "$(docker image inspect "$mtp_image" --format '{{.Id}}')" = "$mtp_image_id" ``` +The immutable reference is also published as tag +`glm53-spark-mtp3-c139f3670-mesh69313e19`; use the digest above for deployment. +The [compute-image equivalence record](compute-image-equivalence.json) verifies +that all 4,891 vLLM, 385 B12X, and 150 SparkCache package files and the selected +environment match compute-tested image `sha256:3b4768e5ba31cadcc882dffa06d7b667af44abdf157d5c11b7ac7fe962e80c43`. +The published config-image ID also passed its own +[GPU, serving, restart, and persistent-cache checks](../../performance/records/glm53-flash/spark-mtp3-compute-stream-safety-20260906.md). + Use the repository's [content receipt](image-receipt.json) for the renderer, installer, and native qualification runner. This default deployment requires no compiler or local image rebuild. The source reproduction commands below @@ -37,6 +55,8 @@ are optional; skip to **Use and export** when using the published image. | Object | Identity or location | |---|---| | Parent runtime | `operator_image` in `../glm53-flash-jj-r8-gb10/pins.json` | +| Compute source lock and vLLM patch | [`compute/source-lock.json`](compute/source-lock.json), [`compute/vllm-e02-to-compute.patch`](compute/vllm-e02-to-compute.patch) | +| Prepared CUDA 13.3 and B12X source payload | output of [`compute/prepare_compute_source.py`](compute/prepare_compute_source.py) | | Target, speculation, transport, and marker pins | `pins.json` in this directory | | Embedded transport bundle | `/opt/spark-sircl` | | Compiled RDMA transmit marker | `/opt/sparkring/bin/mlx5-rdma-tx-marker` | @@ -61,9 +81,22 @@ Use a Linux/ARM64 Docker host. Check available disk and RAM before pulling the parent image. Do not remove resident model files or stop serving containers to make room for this build. -From the repository root, first pull the pinned parent and compose its SIRCL -bundle. Do not use the already-composed published child as the parent input. -Use absent output paths and an unused temporary container name: +From the repository root, first prepare the checksum-bound compute source. This +step downloads the pinned public CUDA component archives and Local Inference +Lab B12X source before the network-disabled image build. The output and optional +download cache are local build inputs and must use absent or empty paths: + +```bash +python3 runtime/glm53-spark-mtp3-mesh/compute/prepare_compute_source.py \ + --output /var/tmp/mtp3-compute-source \ + --cache /var/tmp/mtp3-compute-downloads +``` + +The preparer verifies every archive checksum, B12X tree, vLLM patch hash, and +package file and writes `prepared-manifest.json`. Then pull the pinned parent +and compose its SIRCL bundle. Do not use the already-composed published child +as the parent input. Use absent output paths and an unused temporary container +name: ```bash mtp_parent='ghcr.io/fujitsupolycom/sparkring-glm53-sparkcache@sha256:0d4029b3b7023cf32c37ac20279469c9a2ee16a057f25aae3bcfee9ee5fb660f' @@ -83,13 +116,14 @@ Then prepare the content-verified build context: ```bash python3 runtime/glm53-spark-mtp3-mesh/build_image.py prepare \ --bundle /var/tmp/mtp3-mesh-bundle \ + --compute-context /var/tmp/mtp3-compute-source \ --context /var/tmp/mtp3-mesh-image-context ``` The context path must not exist. Preparation copies only manifest-listed -transport files, source-pinned marker code, verification code, pins, and the -RoCEnante license and provenance. It rejects unexpected bundle files and writes -a content manifest for every construction input. +compute and transport files, source-pinned marker code, verification code, +pins, and the RoCEnante license and provenance. It rejects unexpected bundle +or compute files and writes a content manifest for every construction input. Build and verify without loading a model: @@ -109,8 +143,9 @@ image to a registry. The local tag is a convenience; use the receipt's full Image construction has networking disabled. Verification runs with no host device mounts, no Linux capabilities, no network, a read-only root filesystem, two CPUs, and a 2 GiB memory limit. It checks the complete parent layer prefix, -package source and native-library hashes, bundle hashes, Python syntax, readiness -warmup helper hash and temperature environment, lazy +package source and native-library hashes, CUDA 13.3 components, the complete +B12X `ef308bac` package with selector overrides, vLLM input/output hashes and proposal-head environment, +bundle hashes, Python syntax, readiness warmup helper hash and temperature, lazy RoCEnante import, and marker linkage. CUDA must remain uninitialized. ## Use and export diff --git a/runtime/glm53-spark-mtp3-mesh/MANAGED_MESH.md b/runtime/glm53-spark-mtp3-mesh/MANAGED_MESH.md index eca6395f..f795f032 100644 --- a/runtime/glm53-spark-mtp3-mesh/MANAGED_MESH.md +++ b/runtime/glm53-spark-mtp3-mesh/MANAGED_MESH.md @@ -362,6 +362,7 @@ The bounded `--run-seconds` mode is for isolated diagnostics only. |---|---| | Child-process and peer checks | 1-second loop; peer HTTP timeout 2 seconds | | Unavailable peer connection | 4-second grace after the first transport failure; degraded health blocks model startup | +| Docker container status | One background query at a time, 3-second timeout; unknown status blocks model startup | | MAC/IP, Ethernet MTU, sysfs GID/netdev, routes, qdiscs, TC state | 5-second periodic check | | Full RDMA active-MTU probe | Startup and approximately every 60 seconds | | Health progress freshness | Readiness rejected after 10 seconds without supervisor progress | @@ -375,6 +376,14 @@ An authentication failure, explicit negative readiness, or changed process generation does not receive transport-error grace: it triggers failure when observed. Local marker exits also trigger failure without that grace. +Docker status queries run outside the fabric-monitor loop. A slow or failed +query reports `docker_status_degraded: true`; it does not declare fabric +failure or interrupt existing serving. Marker, network, and authenticated +peer checks continue. A completed query reporting that the model stopped +still enforces the model-exit policy. Pending queries never reuse a previous +stopped result as proof. Startup and teardown use separate synchronous +checks: unknown Docker state cannot authorize marker or network removal. + These are polling and timeout settings, not zero-window guarantees. Command execution, scheduling, management-network delays, and container-stop time affect detection and containment. In-flight requests can fail; a successful diff --git a/runtime/glm53-spark-mtp3-mesh/README.md b/runtime/glm53-spark-mtp3-mesh/README.md index 74f7b88e..4748f7ec 100644 --- a/runtime/glm53-spark-mtp3-mesh/README.md +++ b/runtime/glm53-spark-mtp3-mesh/README.md @@ -1,21 +1,22 @@ # GLM-5.3 Flash Spark with native MTP3 and hardware-forwarded mesh -Status: **research-only**. Bundle composition, site rendering, managed -host-fabric installation/supervision, and CPU checks are **implemented**. -The [managed functional record](../../performance/records/glm53-flash/spark-mtp3-managed-mesh-functional-20260905.md) -qualifies bounded installer, policy-scoped fault/recovery, post-recovery -readiness, and one persistent-cache recall case. Broader coverage remains -unqualified. The -[sampling-warmup image functional record](../../performance/records/glm53-flash/spark-mtp3-mesh-temperature-one-functional-20260905.md) -qualifies bounded native checks, four-rank startup/restart, and one persistent -recall restoration for its exact image. Broader cache/workload coverage and -failure containment remain unqualified. +Status: **research-only** profile. Bundle composition, site rendering, managed +host services, and CPU checks are **implemented**. The published image has +[qualified bounded GPU, serving, restart, and persistent-cache checks](../../performance/records/glm53-flash/spark-mtp3-compute-stream-safety-20260906.md). +The record identifies the exact image, measurements, and scope; it does not +establish in-flight collective failure containment or unattended availability. This profile serves the `GLM-5.3-Flash-NVFP4-Spark` checkpoint with its built-in -multi-token predictor at depth three. It combines graph-native SIRCL, -dual-rail fused SIRCL, and a modified RoCEnante all-reduce over a four-node -physical ring. Opposite ranks communicate through hardware forwarding in an -intermediate ConnectX-7. No external draft checkpoint is required. +multi-token predictor at depth three. The predictor uses a separate runtime- +NVFP4 proposal head with BF16 activations; the target/verifier head retains its +BF16 checkpoint representation. It combines graph-native SIRCL, dual-rail +fused SIRCL, and a modified RoCEnante all-reduce over a four-node physical +ring. Opposite ranks communicate through hardware forwarding in an intermediate +ConnectX-7. No external draft checkpoint or DFlash model is required. + +At TP4, the proposal head adds approximately 85.08 MiB of packed NVFP4 values +and scales per rank. The BF16 target/verifier head remains allocated. The +proposal-head figure is an added allocation, not a net memory reduction. Follow the [operator quickstart](../../docs/GLM53_SPARK_MTP3_MESH_QUICKSTART.md). It starts from a public checkout and image/model artifacts, explains the @@ -24,56 +25,87 @@ installation commands. No private experiment checkout or existing cache is required. The [managed operations guide](MANAGED_MESH.md) creates the shared authentication inputs and provides model-start, readiness, stop, and recovery commands. Keep private site files and the health key outside this repository. -The [measurement record](../../performance/records/glm53-flash/spark-mtp3-mesh-20260905.md) -contains the bounded throughput observations and their limitations. +The [proposal-head comparison](../../performance/records/glm53-flash/spark-mtp3-nvfp4-proposal-head-20260905.md) +contains the bounded throughput observations for this head configuration and +their limitations. ## Operator benchmark observations +The [compute matrices](../../performance/records/glm53-flash/spark-mtp3-compute-matrices-20260905.md) +record full C1/C2/C4/C8/C12/C16 results for the proposal head, loader/RNG fixes, +shared MoE scales, and top-k selector. Column averages give each of the +8K/32K/64K contexts equal weight. The records identify their measured images; +they do not substitute for validation of the combined stream-safety image. + The [consolidated validation report](../../performance/records/glm53-flash/spark-mtp3-validation-summary-20260905.md) collects the completed tests, three-pass prefill measurements, and remaining work. Use that report to avoid repeating checks already covered by receipts. Status: **research-only** measurements, not general performance guarantees. -The [throughput record](../../performance/records/glm53-flash/spark-mtp3-mesh-20260905.md) -identifies the measured source configuration and its differences from the -packaged image. C denotes concurrent requests; decode values are aggregate -output tokens per second across those requests. - -| Context | C1 | C2 | C4 | C8 | C12 | C16 | -|---:|---:|---:|---:|---:|---:|---:| -| 8K | 48.2 | 75.8 | 112.2 | 168.8 | 193.4 | 231.3 | -| 32K | 49.9 | 76.8 | 119.0 | 164.6 | 197.3 | 222.7 | -| 64K | 43.0 | 76.4 | 119.0 | 165.8 | 192.3 | 220.9 | - -Concurrency-one prefill scouts measured **2,703–2,787 prompt tokens/s** over -8K–128K contexts, with one observation per context. They are not a repeated -cold-cache benchmark. +The [proposal-head comparison](../../performance/records/glm53-flash/spark-mtp3-nvfp4-proposal-head-20260905.md) +reports three proposal-head repetitions and two controls at 8K. C denotes +concurrent requests; decode values are +mean aggregate output tokens per second across those requests. + +| Context | C1 | C2 | C4 | C8 | +|---:|---:|---:|---:|---:| +| 8K | 51.6 | 76.9 | 120.8 | 168.8 | + +Relative to a shared-BF16-proposal-head control with the same CUDA version, +B12X kernels, metadata reuse, and dense-kernel integration, C1 improved 8.22% +in raw output throughput and 4.90% in +acceptance-normalized sequence steps/s. C2/C4/C8 results were mixed. Repeated +prefill means moved by no more than 0.36% over 8K–128K contexts, so no prefill +gain is claimed. The +[broader matrix](../../performance/records/glm53-flash/spark-mtp3-mesh-20260905.md) +records measurements for its explicitly identified image and serving settings. +The head-specific control preserves the verifier implementation; the complete +CUDA 13.3/B12X `b58f34ea` image changes other target computation relative to +the configuration recorded in that matrix, so cross-image gains cannot be assigned only to the +proposal head or dense kernels. + +The published image's +[compute-equivalence record](compute-image-equivalence.json) matches every +vLLM, B12X, and SparkCache package file and selected environment entry to the +serving image identified in that record. This establishes compute-package +content equivalence, not end-to-end performance equivalence: the mounted +transport differs outside that comparison. The published image's separate +[runtime validation](../../performance/records/glm53-flash/spark-mtp3-compute-stream-safety-20260906.md) +includes an 18-cell matrix, two focused repetitions, native and GPU stream +checks, an idle rank-loss test, restart, and a verified persistent-prefix restore. The separate [Estonia accuracy record](../../performance/records/glm53-flash/spark-mtp3-country-recall-20260905.md) reports **30/30 correct** at C8 on one repeated 133,208-token prompt, no output-limit hits, and 1.96 s mean cache-primed TTFT. Its 23.8 tok/s figure -uses summed request times, not cluster wall time. Both records retain the -operator screenshots and metric definitions. +uses summed request times, not cluster wall time. The Estonia record includes +the operator screenshot and metric definitions. The [long-context needle hunt](../../performance/records/glm53-flash/spark-mtp3-needle-20260905.md) passed **4/4** exact-value, revision, and cross-reference checks, reaching -**507,367 actual prompt tokens** on the published image. +**507,367 actual prompt tokens** on serving image +`sha256:26273b8e358df139ae913610a5d43084ff0fd08aafe282ef633a3bc74afefe47`, +as recorded in that report. It does not identify the image pinned by +`public-image.json`. ## Composition | Input | Contract | |---|---| | Model, MTP depth, graph shapes, mesh bundle, marker identity, cache identity | [`pins.json`](pins.json) | -| Linux/ARM64 image, vLLM, B12X kernels, SparkCache, native SIRCL | [`../glm53-flash-jj-r8-gb10/pins.json`](../glm53-flash-jj-r8-gb10/pins.json) | +| Linux/ARM64 parent image, SparkCache, and native SIRCL | [`../glm53-flash-jj-r8-gb10/pins.json`](../glm53-flash-jj-r8-gb10/pins.json) | +| CUDA 13.3, GLM metadata/loader/RNG integration, B12X scale sharing and selector overrides, NVFP4/BF16 proposal head | [`pins.json`](pins.json), [`IMAGE_BUILD.md`](IMAGE_BUILD.md) | | Topology and rank-local filesystem inputs | [`site.example.json`](site.example.json), [`fabric.example.json`](fabric.example.json) | | Source-bound collective dispatch and health checks | [`glm53_rocenante_overlay`](../../spark_transport/experiments/glm53_rocenante_overlay/README.md) | | Hardware-forwarding plan and native source marker | [`cx7_hairpin_diagonal`](../../spark_transport/experiments/cx7_hairpin_diagonal/README.md) | | Modified RoCEnante communication package | [`third_party/b12x_roce`](../../third_party/b12x_roce/README.md) | -The model runtime and kernels come from the pinned parent image. The managed -profile requires the [published child image](IMAGE_BUILD.md), which adds -the verified transport bundle, managed source marker, and temperature-one -readiness helper. Pull the immutable reference in [public-image.json](public-image.json) +The managed profile requires the [published child image](IMAGE_BUILD.md). It +retains the parent runtime while adding CUDA 13.3, the uniform native-MTP3 +metadata and loader/RNG integration, B12X revision `ef308bac` with selector +overrides, the runtime-NVFP4/BF16 +proposal head, the verified transport bundle, the managed source marker, and +the temperature-one readiness helper. The verifier remains BF16. Pull the +immutable reference in [public-image.json](public-image.json) and use [image-receipt.json](image-receipt.json) for rendering and installation. Local source reproduction is optional; distribute identical verified bytes to all ranks. The canonical transport bundle remains mounted read-only. @@ -140,7 +172,8 @@ serving lifecycle. Use authenticated managed readiness for serving. Native MTP uses the target checkpoint as the draft identity. The profile sets SparkCache's `draft_policy=separate` because that describes the registered state layout; it does not request an external model. A dedicated namespace -prevents restoration of entries tagged for an external DFlash checkpoint. +binds compute lock `139f3670` and transport bundle `69313e19` so the +NVFP4-proposal-head composition cannot restore entries from another computation. Do not relabel those entries to avoid cache misses. Persistent restore under the native-MTP identity requires its own qualification. diff --git a/runtime/glm53-spark-mtp3-mesh/build_image.py b/runtime/glm53-spark-mtp3-mesh/build_image.py index 8e48309f..65991aee 100644 --- a/runtime/glm53-spark-mtp3-mesh/build_image.py +++ b/runtime/glm53-spark-mtp3-mesh/build_image.py @@ -67,11 +67,15 @@ def verify_bundle(bundle: Path, expected: str) -> list[dict]: return records -def prepare(bundle: Path, context: Path) -> dict: +def prepare(bundle: Path, context: Path, compute: Path | None = None) -> dict: """Copy only content-verified inputs into a directory that does not exist.""" if context.exists(): raise ValueError(f"Build context already exists: {context}") + if compute is None or not compute.is_dir() or compute.is_symlink(): + raise ValueError("A prepared compute source directory is required; see compute/README.md") profile = read_json(HERE / "pins.json") + if sha256(compute / "source-lock.json") != profile["compute"]["source_lock_sha256"]: + raise ValueError("Prepared compute source lock differs from the mesh profile pin") base_path = (HERE / profile["image_pins"]).resolve() base = read_json(base_path) records = verify_bundle(bundle, profile["canonical_bundle_manifest_sha256"]) @@ -95,6 +99,11 @@ def prepare(bundle: Path, context: Path) -> dict: f"bundle/{MANIFEST}": bundle / MANIFEST, } files.update({f"bundle/{record['path']}": bundle / record["path"] for record in records}) + for source in compute.rglob('*'): + if source.is_symlink(): + raise ValueError(f"Compute inputs cannot contain symlinks: {source}") + if source.is_file(): + files[f"compute/{source.relative_to(compute).as_posix()}"] = source for relative, source in files.items(): destination = context / relative destination.parent.mkdir(parents=True, exist_ok=True) @@ -109,7 +118,7 @@ def prepare(bundle: Path, context: Path) -> dict: "helper_sha256": sha256(context / "warmup_dflash.py"), "temperature_environment": "SPARKRING_WARMUP_TEMPERATURE", "default_temperature": 1.0}, - "scope": "Embedded transport bundle, compiled host marker, and readiness warmup helper with temperature one; target weights are not image contents.", + "scope": "Compute source and CUDA components, embedded transport bundle, compiled host marker, and temperature-one readiness warmup; target weights are not image contents.", } write_json(context / "receipts/source-receipt.json", receipt) return receipt @@ -162,6 +171,7 @@ def build(context: Path, image: str, receipt_path: Path, engine: str, pull: bool "PARENT_IMAGE_ID": parent["image_id"], "BUNDLE_MANIFEST_SHA256": source["bundle_manifest_sha256"], "SOURCE_RECEIPT_SHA256": source_sha, + "COMPUTE_SOURCE_LOCK_SHA256": source["files"]["compute/source-lock.json"], }.items(): argv.extend(["--build-arg", f"{name}={value}"]) argv.extend(["--file", str(context / "Dockerfile"), "--tag", image, str(context)]) @@ -180,6 +190,8 @@ def main() -> int: prepare_parser = sub.add_parser("prepare", help="OFFLINE: construct a content-verified build directory") prepare_parser.add_argument("--bundle", type=Path, required=True) prepare_parser.add_argument("--context", type=Path, required=True) + prepare_parser.add_argument("--compute-context", type=Path, required=True, + help="Prepared, source-pinned compute directory from compute/prepare_compute_source.py") build_parser = sub.add_parser("build", help="MUTATES HOST: build and CPU-check an image; no GPU or fabric access") build_parser.add_argument("--context", type=Path, required=True) build_parser.add_argument("--image", required=True) @@ -188,7 +200,7 @@ def main() -> int: build_parser.add_argument("--pull-parent", action="store_true") args = parser.parse_args() if args.command == "prepare": - result = prepare(args.bundle.resolve(), args.context.resolve()) + result = prepare(args.bundle.resolve(), args.context.resolve(), args.compute_context.resolve()) print(json.dumps({"context": str(args.context.resolve()), "bundle_manifest_sha256": result["bundle_manifest_sha256"]}, indent=2)) else: build(args.context.resolve(), args.image, args.receipt.resolve(), args.engine, args.pull_parent) diff --git a/runtime/glm53-spark-mtp3-mesh/compute-image-equivalence.json b/runtime/glm53-spark-mtp3-mesh/compute-image-equivalence.json new file mode 100644 index 00000000..0e255476 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/compute-image-equivalence.json @@ -0,0 +1,27 @@ +{ + "schema": "sparkring-compute-image-equivalence/v1", + "checks_passed": true, + "tested_serving_image": "sha256:3b4768e5ba31cadcc882dffa06d7b667af44abdf157d5c11b7ac7fe962e80c43", + "published_image": "sha256:2e41b1e934a85ff7c21b780532db2f0a0e978df081e52f4ae2bf11f8992fb24f", + "package_files": { + "b12x": 385, + "sparkcache": 150, + "vllm": 4891 + }, + "package_files_identical": true, + "selected_environment_identical": true, + "environment": { + "CUDA_HOME": "/opt/cuda-13.3", + "SPARKRING_WARMUP_TEMPERATURE": "1", + "TRITON_PTXAS_PATH": "/opt/cuda-13.3/bin/ptxas", + "VLLM_B12X_DENSE_ACTIVATION_MODE": "auto", + "VLLM_GDN_SPEC_DECODE_METADATA_FASTPATH": "1", + "VLLM_LM_HEAD_A16": "1", + "VLLM_MTP_NVFP4_LM_HEAD": "1", + "VLLM_MXFP8_LM_HEAD": "0" + }, + "tested_snapshot_sha256": "7d2413eea1f2da3163f40f2a7cf937aeacb1bca44bdc5477fac87b94594b92b0", + "published_snapshot_sha256": "7d2413eea1f2da3163f40f2a7cf937aeacb1bca44bdc5477fac87b94594b92b0", + "scope": "All installed compute package files except Python bytecode. Mounted transport differs and is validated separately in the linked exact-image record.", + "functional_record": "../../performance/records/glm53-flash/spark-mtp3-compute-stream-safety-20260906.md" +} diff --git a/runtime/glm53-spark-mtp3-mesh/compute/README.md b/runtime/glm53-spark-mtp3-mesh/compute/README.md new file mode 100644 index 00000000..8eab2991 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/compute/README.md @@ -0,0 +1,54 @@ +# GLM-5.3 compute source composition + +Status: **implemented**. These files reproduce the compute source used by the +GLM-5.3 NVFP4-Spark native-MTP3 mesh image. Hardware qualification belongs to +the profile's validation records, not to this source-preparation tooling. + +`source-lock.json` is the authoritative input. It binds: + +- the vLLM source revision already present in the parent image; +- a reviewable vLLM patch, a byte-exact 24-file replacement archive, and every + base and resulting file hash; +- B12X revision `ef308bac0f3b3eb8fea63e4013afc0c2ea1c6301`, its Git tree, + source archive, three source-checked selector overrides, and all 385 installed package-file hashes; +- seven NVIDIA CUDA 13.3 SBSA redistributable archives; and +- the five environment settings that select metadata reuse, dense-kernel + policy, and the NVFP4 native-MTP proposal head. + +The proposal head receives a separate tensor-parallel copy of the checkpoint's +unquantized `lm_head.weight`, converts that copy to NVFP4 during loading, and +uses BF16 activations. `VLLM_MXFP8_LM_HEAD=0` leaves the target/verifier head +unchanged. Rejection sampling therefore retains the target model's sampling +contract, while proposal-head quantization can change acceptance length. + +Deferred GLM weights and scales own their storage when a loader reuses input +buffers. Draft proposal randomness is independent of rejection-sampling +randomness. These changes derive from vLLM revision `17e341b9` and +[PR 653](https://github.com/local-inference-lab/vllm/pull/653), respectively. +B12X shares native MoE scale storage and includes the top-k-512 selector from +[PR 316](https://github.com/local-inference-lab/b12x/pull/316). The source lock +records full donor identities. The selector archive is applied only after all +base and resulting file hashes pass verification. + +The image builder calls `prepare_compute_source.prepare(destination, cache)` +while network access is available. The prepared directory contains the pinned +B12X source and CUDA archives. Docker copies that directory into the build and +runs `apply_compute.py` with network access disabled. The installer verifies +the parent hashes before extracting the replacement archive; this preserves the +mixed line endings bound by the source lock without requiring Git in the image. +`verify_compute.py` +requires exact installed hashes and rejects missing or partial source maps. + +The B12X source is obtained from +[`local-inference-lab/b12x`](https://github.com/local-inference-lab/b12x) and is +licensed under Apache License 2.0. The downloaded source archive includes its +`LICENSE` file. The vLLM patch derives from +[`local-inference-lab/vllm`](https://github.com/local-inference-lab/vllm) +revision `3512b066e7796128c0c380ccc558182960f2f0ea`, with dense-kernel integration +from revision `a8c796f3af74106b2d8d441e9ec54588936a5388`; vLLM is licensed under +Apache License 2.0. + +B12X source archives use LF endings. Source preparation converts Python and C +files to CRLF to reproduce the installed package hashes in `source-lock.json`. +Markdown and compressed profile data retain the archive bytes. This byte-level +contract makes package-content verification independent of checkout settings. diff --git a/runtime/glm53-spark-mtp3-mesh/compute/apply_compute.py b/runtime/glm53-spark-mtp3-mesh/compute/apply_compute.py new file mode 100644 index 00000000..6dda7627 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/compute/apply_compute.py @@ -0,0 +1,168 @@ +"""Install the prepared CUDA, B12X, and vLLM compute composition offline.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import shutil +import subprocess +import tarfile +import tempfile +from pathlib import Path + + +def _sha256(path: Path) -> str: + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def _map_sha256(files: dict[str, str]) -> str: + payload = json.dumps(files, sort_keys=True, separators=(",", ":")).encode() + return hashlib.sha256(payload).hexdigest() + + +def _load(prepared: Path) -> tuple[dict, dict]: + lock_path = prepared / "source-lock.json" + lock = json.loads(lock_path.read_text(encoding="utf-8")) + manifest = json.loads((prepared / "prepared-manifest.json").read_text(encoding="utf-8")) + if hashlib.sha256(lock_path.read_bytes()).hexdigest() != manifest[ + "source_lock_sha256" + ]: + raise ValueError("prepared source lock hash mismatch") + return lock, manifest + + +def _install_cuda(prepared: Path, lock: dict, destination: Path) -> None: + if destination.exists(): + raise ValueError(f"CUDA destination already exists: {destination}") + destination.mkdir(parents=True) + for relative, expected in lock["cuda"]["components"].items(): + archive = prepared / "cuda-archives" / Path(relative).name + if _sha256(archive) != expected: + raise ValueError(f"CUDA archive hash mismatch: {archive.name}") + with tempfile.TemporaryDirectory(prefix="sparkring-cuda-") as temporary: + unpack = Path(temporary) + with tarfile.open(archive) as source: + source.extractall(unpack, filter="data") + entries = list(unpack.iterdir()) + if len(entries) != 1 or not entries[0].is_dir(): + raise ValueError(f"invalid CUDA archive root: {archive.name}") + shutil.copytree(entries[0], destination, dirs_exist_ok=True, symlinks=True) + (destination / "sparkring-component-manifest.json").write_text( + json.dumps(lock["cuda"]["components"], indent=2, sort_keys=True) + "\n", + encoding="utf-8", newline="\n", + ) + + +def _install_vllm(prepared: Path, site: Path, lock: dict) -> dict[str, str]: + entries = lock["vllm"]["files"] + patch = prepared / lock["vllm"]["patch"] + if _sha256(patch) != lock["vllm"]["patch_sha256"]: + raise ValueError("vLLM patch hash mismatch") + archive = prepared / lock["vllm"]["replacement_archive"] + if _sha256(archive) != lock["vllm"]["replacement_archive_sha256"]: + raise ValueError("vLLM replacement archive hash mismatch") + with tempfile.TemporaryDirectory(prefix="sparkring-vllm-") as temporary: + work = Path(temporary) + for relative, base_hash, _ in entries: + installed = site / relative + actual = _sha256(installed) + if actual != base_hash: + raise ValueError(f"vLLM base hash mismatch for {relative}: {actual}") + staged = work / relative + staged.parent.mkdir(parents=True, exist_ok=True) + shutil.copyfile(installed, staged) + with tarfile.open(archive) as source: + names = set(source.getnames()) + expected_names = {entry[0] for entry in entries} + if names != expected_names: + raise ValueError("vLLM replacement archive has an unexpected file set") + source.extractall(work, filter="data") + result: dict[str, str] = {} + for relative, _, expected in entries: + staged = work / relative + actual = _sha256(staged) + if actual != expected: + raise ValueError(f"vLLM result hash mismatch for {relative}: {actual}") + shutil.copyfile(staged, site / relative) + result[relative] = actual + return result + + +def _package_map(root: Path) -> dict[str, str]: + package = root / "b12x" + files = sorted( + path + for path in package.rglob("*") + if path.is_file() + and "__pycache__" not in path.parts + and path.suffix != ".pyc" + ) + if not files: + raise ValueError(f"no installed B12X package files under {package}") + return {path.relative_to(root).as_posix(): _sha256(path) for path in files} + + +def apply( + prepared: Path, + site_packages: Path, + receipt: Path, + cuda_destination: Path, +) -> Path: + prepared = prepared.resolve() + site_packages = site_packages.resolve() + lock, manifest = _load(prepared) + _install_cuda(prepared, lock, cuda_destination) + vllm_files = _install_vllm(prepared, site_packages, lock) + subprocess.run( + [ + "python3", + "-m", + "pip", + "install", + "--no-deps", + "--no-build-isolation", + "--force-reinstall", + str(prepared / "b12x-source"), + ], + check=True, + ) + b12x_files = _package_map(site_packages) + if b12x_files != manifest["b12x_files"]: + raise ValueError("installed B12X package differs from prepared source") + if _map_sha256(b12x_files) != lock["b12x"]["package_files_sha256"]: + raise ValueError("installed B12X package differs from source-lock.json") + output = { + "schema": "sparkring-glm53-compute-installed/v1", + "source_lock_sha256": manifest["source_lock_sha256"], + "vllm_revision": lock["vllm"]["base_revision"], + "vllm_overrides": vllm_files, + "b12x_revision": lock["b12x"]["revision"], + "b12x_tree": lock["b12x"]["tree"], + "b12x_files": b12x_files, + "b12x_package_files_sha256": lock["b12x"]["package_files_sha256"], + "cuda_components": lock["cuda"]["components"], + "environment": lock["environment"], + "target_head_quantization": False, + } + receipt.parent.mkdir(parents=True, exist_ok=True) + receipt.write_text(json.dumps(output, indent=2, sort_keys=True) + "\n", encoding="utf-8", newline="\n") + return receipt + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--prepared", type=Path, required=True) + parser.add_argument("--site-packages", type=Path, required=True) + parser.add_argument( + "--receipt", + type=Path, + default=Path("/opt/sparkring/receipts/glm53-compute-installed.json"), + ) + parser.add_argument("--cuda-destination", type=Path, default=Path("/opt/cuda-13.3")) + args = parser.parse_args() + print(apply(args.prepared, args.site_packages, args.receipt, args.cuda_destination)) + + +if __name__ == "__main__": + main() diff --git a/runtime/glm53-spark-mtp3-mesh/compute/b12x-selector-files.tar.gz b/runtime/glm53-spark-mtp3-mesh/compute/b12x-selector-files.tar.gz new file mode 100644 index 00000000..1c76ea5d Binary files /dev/null and b/runtime/glm53-spark-mtp3-mesh/compute/b12x-selector-files.tar.gz differ diff --git a/runtime/glm53-spark-mtp3-mesh/compute/prepare_compute_source.py b/runtime/glm53-spark-mtp3-mesh/compute/prepare_compute_source.py new file mode 100644 index 00000000..fad54c08 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/compute/prepare_compute_source.py @@ -0,0 +1,188 @@ +"""Prepare the network-fetched compute payload for an offline image build.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import shutil +import tarfile +import tempfile +import urllib.request +from pathlib import Path + +HERE = Path(__file__).resolve().parent +LOCK = json.loads((HERE / "source-lock.json").read_text(encoding="utf-8")) + + +def _sha256(path: Path) -> str: + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def _map_sha256(files: dict[str, str]) -> str: + payload = json.dumps(files, sort_keys=True, separators=(",", ":")).encode() + return hashlib.sha256(payload).hexdigest() + + +def _download(url: str, expected: str, destination: Path) -> None: + if destination.is_file() and _sha256(destination) == expected: + return + destination.parent.mkdir(parents=True, exist_ok=True) + with tempfile.NamedTemporaryFile(dir=destination.parent, delete=False) as stream: + temporary = Path(stream.name) + try: + urllib.request.urlretrieve(url, temporary) + actual = _sha256(temporary) + if actual != expected: + raise ValueError(f"download hash mismatch for {url}: {actual}") + temporary.replace(destination) + finally: + temporary.unlink(missing_ok=True) + + +def _extract_single_root(archive: Path, destination: Path) -> Path: + destination.mkdir(parents=True) + with tarfile.open(archive) as source: + source.extractall(destination, filter="data") + entries = list(destination.iterdir()) + if len(entries) != 1 or not entries[0].is_dir(): + raise ValueError(f"expected one source root in {archive.name}") + return entries[0] + + +def _package_map(root: Path, package: str) -> dict[str, str]: + package_root = root / package + files = sorted( + path + for path in package_root.rglob("*") + if path.is_file() + and "__pycache__" not in path.parts + and path.suffix != ".pyc" + ) + if not files: + raise ValueError(f"no package files found under {package_root}") + return { + path.relative_to(root).as_posix(): _sha256(path) + for path in files + } + + +def _normalize_b12x_bytes(root: Path) -> None: + suffixes = set(LOCK["b12x"]["normalized_text_suffixes"]) + for path in sorted((root / "b12x").rglob("*")): + if path.is_file() and path.suffix in suffixes: + data = path.read_bytes().replace(b"\r\n", b"\n") + path.write_bytes(data.replace(b"\n", b"\r\n")) + + +def _apply_b12x_overrides(root: Path, archive: Path, contract: dict) -> None: + """Apply checksum-bound selector files only over their expected source bytes.""" + if _sha256(archive) != contract["archive_sha256"]: + raise ValueError("B12X override archive hash mismatch") + entries = contract["files"] + expected = {name for name, _, _ in entries} + if len(expected) != len(entries): + raise ValueError("Duplicate B12X override path") + with tarfile.open(archive) as source: + members = source.getmembers() + if (len(members) != len(expected) or {item.name for item in members} != expected + or any(not item.isfile() for item in members)): + raise ValueError("B12X override archive has an unexpected file set") + replacements = {} + for name, base_hash, result_hash in entries: + path = root / name + if not name.startswith("b12x/") or not path.resolve().is_relative_to((root / "b12x").resolve()): + raise ValueError("B12X override escapes package directory") + if _sha256(path) != base_hash: + raise ValueError(f"B12X override base hash mismatch: {name}") + data = source.extractfile(name).read() + if hashlib.sha256(data).hexdigest() != result_hash: + raise ValueError(f"B12X override result hash mismatch: {name}") + replacements[path] = data + for path, data in replacements.items(): + path.write_bytes(data) + + +def prepare(destination: Path, cache: Path | None = None) -> Path: + """Create a complete, checksum-bound context for a network-disabled build.""" + destination = destination.resolve() + if destination.exists() and any(destination.iterdir()): + raise ValueError(f"compute destination is not empty: {destination}") + destination.mkdir(parents=True, exist_ok=True) + cache = (cache or destination.parent / ".compute-downloads").resolve() + cache.mkdir(parents=True, exist_ok=True) + + for name in ( + "source-lock.json", + "vllm-e02-to-compute.patch", + "vllm-compute-files.tar.gz", + "apply_compute.py", + "verify_compute.py", + ): + shutil.copy2(HERE / name, destination / name) + + lock_hash = _sha256(destination / "source-lock.json") + patch = destination / LOCK["vllm"]["patch"] + if _sha256(patch) != LOCK["vllm"]["patch_sha256"]: + raise ValueError("vLLM patch does not match source-lock.json") + replacement = destination / LOCK["vllm"]["replacement_archive"] + if _sha256(replacement) != LOCK["vllm"]["replacement_archive_sha256"]: + raise ValueError("vLLM replacement archive does not match source-lock.json") + + b12x = LOCK["b12x"] + b12x_archive = cache / "b12x.tar.gz" + _download(b12x["archive_url"], b12x["archive_sha256"], b12x_archive) + unpack = destination / ".b12x-unpack" + source_root = _extract_single_root(b12x_archive, unpack) + b12x_destination = destination / "b12x-source" + shutil.move(str(source_root), b12x_destination) + shutil.rmtree(unpack) + _normalize_b12x_bytes(b12x_destination) + if b12x.get("overrides"): + overrides = b12x["overrides"] + archive = destination / overrides["archive"] + shutil.copy2(HERE / overrides["archive"], archive) + _apply_b12x_overrides(b12x_destination, archive, overrides) + b12x_files = _package_map(b12x_destination, "b12x") + if _map_sha256(b12x_files) != b12x["package_files_sha256"]: + raise ValueError("normalized B12X package does not match source-lock.json") + + cuda_archives: dict[str, str] = {} + cuda_dir = destination / "cuda-archives" + cuda_dir.mkdir() + for relative, expected in LOCK["cuda"]["components"].items(): + archive = cache / Path(relative).name + _download(LOCK["cuda"]["base_url"] + relative, expected, archive) + target = cuda_dir / archive.name + shutil.copy2(archive, target) + cuda_archives[f"cuda-archives/{target.name}"] = expected + + prepared = { + "schema": "sparkring-glm53-compute-prepared/v1", + "source_lock_sha256": lock_hash, + "vllm_patch_sha256": LOCK["vllm"]["patch_sha256"], + "vllm_replacement_archive_sha256": LOCK["vllm"][ + "replacement_archive_sha256" + ], + "b12x_revision": b12x["revision"], + "b12x_tree": b12x["tree"], + "b12x_archive_sha256": b12x["archive_sha256"], + "b12x_package_files_sha256": b12x["package_files_sha256"], + "b12x_files": b12x_files, + "cuda_archives": cuda_archives, + } + manifest = destination / "prepared-manifest.json" + manifest.write_text(json.dumps(prepared, indent=2, sort_keys=True) + "\n", encoding="utf-8", newline="\n") + return manifest + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--cache", type=Path) + args = parser.parse_args() + print(prepare(args.output, args.cache)) + + +if __name__ == "__main__": + main() diff --git a/runtime/glm53-spark-mtp3-mesh/compute/source-lock.json b/runtime/glm53-spark-mtp3-mesh/compute/source-lock.json new file mode 100644 index 00000000..8ff0b729 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/compute/source-lock.json @@ -0,0 +1,194 @@ +{ + "schema": "sparkring-glm53-compute-source/v1", + "status": "research-only", + "vllm": { + "base_revision": "e02b174693e13859de61811b5e8cd13d5308e259", + "donor_revision": "3512b066e7796128c0c380ccc558182960f2f0ea", + "dense_donor_revision": "a8c796f3af74106b2d8d441e9ec54588936a5388", + "patch": "vllm-e02-to-compute.patch", + "patch_sha256": "4c20d5874a28cdadea24b8f1deb8799e950786aa39a6e1e866e0f67a1ad6e297", + "replacement_archive": "vllm-compute-files.tar.gz", + "replacement_archive_sha256": "520db578e6be4466d46cdda7698f4044853111d5810ce6da58f81d5a16b6ad88", + "files": [ + [ + "vllm/envs.py", + "22069819122e630dd5131627c6bb8752820be2b66178d7f445eeb6ce93e03d32", + "1fbac28f1a763f9d27845a8e09d59f71ba1740b4491a74d0097a7e9d63763804" + ], + [ + "vllm/model_executor/kernels/linear/mxfp8/b12x.py", + "7abc42bccf03114e880871fa2ffd67d11466483b2ea636d466e762c17417f3d9", + "1b4448f7dbaadcf9ef93b59a6aa54a016653a1019ffa31605722d828ea3d80af" + ], + [ + "vllm/model_executor/kernels/linear/nvfp4/b12x.py", + "4c59ca2067e78a731598a47d146b9e8ddb1b8270f54f6ab19b72c30561f9d53d", + "7a7036e6d0254a0ee6e3265b18685128fc42cb1b896efffdf27c9de038d69fc4" + ], + [ + "vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py", + "efb07f928463fcf12a3c079f4da2ffa330d697a98a6afca3bea94a0a21825f58", + "4a83c1eb9ea9c47c34e8a87b07f087a4cf1384428e74d800e0bdbd667abbb63e" + ], + [ + "vllm/model_executor/layers/quantization/online/nvfp4.py", + "52ce69a32611d8eb00fa2a27cb9b2ec46b8836f34f0a1f8b4834990c3ea228a9", + "8f85d00e6ef0effa873605e506c5c678228cba1459c2245a0770e2a3d0f74450" + ], + [ + "vllm/model_executor/layers/vocab_parallel_embedding.py", + "e86d37bc50b5171e2bc02a8e6428feed34309807f63d20bead6bc5db4ac7fddf", + "0f958f55912dd1f6bf1f902110404018263afd9c0aa34ebb66236780c73b9243" + ], + [ + "vllm/model_executor/models/deepseek_mtp.py", + "e9270724c39a0152dc0a66b94622ebddd384c592534cbbf38d2f43c0ba1592d0", + "ed5b247d014207e81d7d2ddf88e62e82afba70427798428b1b770a4f687ee869" + ], + [ + "vllm/models/glm5next/model_state.py", + "8c66d55da1bcad63f703b1a5463b6969dfa1d3308d3d14abfea778f6c5e34b7f", + "4f884f713335d55fafd729fa4cc9bcd88389ba005323eb79a3f9b02cf5662f36" + ], + [ + "vllm/models/glm5next/nvidia/model.py", + "e992bd796cc8efe0c55656b3f6d858cc2b4737d100d015cb49556d1d53de2288", + "6ffed6a97bf47af392c7eb89b90e384c80a38a3913794876b1827a1b7f3dec69" + ], + [ + "vllm/models/glm5next/nvidia/mtp.py", + "cbff653af56b3589a1ff3d52b5fea660e718cbdb1b278c4fa5dcd8eac0908ee9", + "71e88e3a25d829d52d90a394f7306e29759d3951fb8f912d1279fa0f30496099" + ], + [ + "vllm/utils/b12x.py", + "f610dc19b4dc10d27361b075ff7e8a97a63244f09f770dfc2ff04aeceed4dce1", + "63e7fef8c75f5cd01678b338b9154fc1204b875794d0356904c461f4e9db3076" + ], + [ + "vllm/v1/attention/backends/gdn_attn.py", + "7c325bbcb612aacd2411b3305ee6d068b013f8b6c4445eaea59059ac45881a3d", + "fd75fb72efeb762ec558d332364e889df44a9ea483026bd5ded8f77f42be9b7f" + ], + [ + "vllm/v1/worker/gpu/attn_utils.py", + "012399dde8910deec550df260c132338d20e20543acc7433ace60fec1229814b", + "80108d185a996e52466f2348d6157398248acfd05b735e4a5792949034b70b9d" + ], + [ + "vllm/v1/worker/gpu/model_states/mamba_hybrid.py", + "f2bc9ef85896df2508557bfbbcf2c682f82414595e8a562efdee6ed3dd0515c2", + "b34cb130e233f4d322390acf7fec01a3a368090c7f79da1fa7770c622a2f0dde" + ], + [ + "vllm/v1/worker/gpu/sample/gumbel.py", + "27467f5b570dc1fc2354ca09fbb7cc6a38a1113c81eb99150b94d0c678a2f13c", + "f22d7b5b6b46666ec6d746e49798ba76cb329407cc749a382639dc8ab83f9741" + ], + [ + "vllm/v1/worker/gpu/sample/sampler.py", + "e40830547705a6622fa7c4733791c261c0908245d376514d7ee891e587357cad", + "3a1382ac5d5e0798e1128854b10d4801f1766eb99e90e16e4efa1e8033ef2fa1" + ], + [ + "vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py", + "c2c1f583ea1572c62c2194af70494d44f5843adbc3f94f204b9ac32016097c19", + "d91dc325f52a6999d3c9a204262f8ca200cebdc568c9e9924b76c78ede8aef21" + ], + [ + "vllm/v1/worker/gpu/spec_decode/dflash/speculator.py", + "0f3d16fb90f30ee2fe45c744dac99384ebf081663664bc8464403ddf644203d5", + "9fcb9ddc718656ede02c856c4b224a1a4511e5d955c96f1e25207e4b87a940f4" + ], + [ + "vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py", + "1f6ff5ca9c8f38ff417aafd43bfa3116b5387bf0f7b58721acb2185781879836", + "a140bafbbe25cf689bdae0362232ab103adee0baca374c6d47f7af1789bc312c" + ], + [ + "vllm/v1/worker/gpu/spec_decode/dspark/speculator.py", + "a72c6d0dbf37ed41b901c0a2178de51aca2c74e6dffdffb89ce223bae2f64598", + "5d11c9570897f57069fa0f6325becd1ff393c7bf2526898c1e40b6502ff67e35" + ], + [ + "vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py", + "0dfc31a7e7f0b6247ce8ea7f580dc930828a885b28c9e4a2d4e69f7ba6b516d7", + "6e25b72630c7f87e7a2d8b64d35f2b7569d5d479ba007d4d2b0deb4c72860353" + ], + [ + "vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py", + "68b60488866ab87857197560bd42ba26c2a89da70d539e9ab56f18667da8998c", + "8f411ed35dc0691d9a71a3ae1b9aff5e3f1d4085e4235e6eefd11e8515dc5a66" + ], + [ + "vllm/v1/worker/gpu/spec_decode/speculator.py", + "2a24c693204032b77bf8abe08df916ca20b119cfe60fd91ad33e87920ae37685", + "3a3574b57748be2e4cdab694c4d5dda98fb2107174b34db4cb1db1ee34e2f457" + ], + [ + "vllm/v1/worker/mamba_utils.py", + "b5bcf2c170daefb858b0668b032dd252a379c78332efc91d87c34d7fed2e9373", + "826f30b8719f0c715e9743a71dfe48f2ab60f7e0de97e2d542ccbe867b6f68d3" + ] + ], + "loader_donor_revision": "17e341b9ede04269f81fcac69a29951a0668a94a", + "rng_donor_revision": "44e6766e3397e8fe8ed9c1fa8a8d2783bb4a2ae8", + "rng_donor_pr": "https://github.com/local-inference-lab/vllm/pull/653" + }, + "b12x": { + "repository": "https://github.com/local-inference-lab/b12x", + "revision": "ef308bac0f3b3eb8fea63e4013afc0c2ea1c6301", + "tree": "dcf039e5e754136275835ea997e6b9abbb6b15ae", + "archive_url": "https://github.com/local-inference-lab/b12x/archive/ef308bac0f3b3eb8fea63e4013afc0c2ea1c6301.tar.gz", + "archive_sha256": "029a6047d80f759e964eb302803df3dc5d2ee324b0a0725d38694b5057ed69a6", + "package_files_sha256": "a212d449381174fabfe681c48cde16b22e3b9f9fa7942ab071b0c7384dc78abd", + "installed_text_line_endings": "crlf", + "normalized_text_suffixes": [ + ".c", + ".py" + ], + "overrides": { + "archive": "b12x-selector-files.tar.gz", + "archive_sha256": "31c51e4fc5b334c9e28c469c4b6f50f4a34458bdd39380c6145dc656056ea42f", + "donor_revision": "9ac142824b4edb750892a0fb63d914230086495d", + "donor_pr": "https://github.com/local-inference-lab/b12x/pull/316", + "files": [ + [ + "b12x/attention/dsa_indexer/fused_indexer.py", + "37e948af8b3ee4755fe01f9b853fc95b04afd3bf31a5cc3347f112e3f993aa87", + "893fbcade135b7e1d146b8fb6530cde0650be515f69bf9a17ced0a9c61a141e2" + ], + [ + "b12x/attention/dsa_indexer/paged.py", + "8ceda277ab3ed6d5154b23765f1a8719240bead889c742881413380f006334c9", + "21a037890fb08896d485de71d1f21eb5cedc6540f4063f0fc0c44151e9701249" + ], + [ + "b12x/attention/dsa_indexer/tiled_topk.py", + "1da3541a53db5dec2c724aab13ba79fa5ad5213627b4fc04f4c9b3a1aa8a98b5", + "52fac27a905929131642a085bbb013989dfe1d6eac9b3a39f576dc6f1841cbde" + ] + ] + } + }, + "cuda": { + "version": "13.3", + "base_url": "https://developer.download.nvidia.com/compute/cuda/redist/", + "components": { + "cuda_nvcc/linux-sbsa/cuda_nvcc-linux-sbsa-13.3.33-archive.tar.xz": "b5dde44aadd52234af3944ae3b2e74e811ad8e71fb600bcc9dfe6d8540353499", + "cuda_cudart/linux-sbsa/cuda_cudart-linux-sbsa-13.3.29-archive.tar.xz": "0cdd73d11885062daf3aa98ad4d7b8bd84f89b398be11f7054edea9ed31f597d", + "cuda_nvrtc/linux-sbsa/cuda_nvrtc-linux-sbsa-13.3.33-archive.tar.xz": "d0502b25799be62a50b743c640e94a1722d20b1ee4ab70d697d71750f04d3b8a", + "cuda_crt/linux-sbsa/cuda_crt-linux-sbsa-13.3.33-archive.tar.xz": "6f6194918c00b980d8fd2111bf0aa004977760855c6e1528e0653bf4c889fbef", + "libnvvm/linux-sbsa/libnvvm-linux-sbsa-13.3.33-archive.tar.xz": "5f8ca5c9a10c3c9804b045960ee6192281efec4c7d83d5f3245ec2de8612118e", + "cccl/linux-sbsa/cccl-linux-sbsa-13.3.3.3.1-archive.tar.xz": "37e9024c5e24a9e9d1618c4fb7b36e74a0a68fac91d589867676952204ecde5b", + "libnvjitlink/linux-sbsa/libnvjitlink-linux-sbsa-13.3.33-archive.tar.xz": "6ed3a14646bd53e25ccf03a52586cdd12b07ad48cf81fe79deac49b5d64c2ce6" + } + }, + "environment": { + "VLLM_GDN_SPEC_DECODE_METADATA_FASTPATH": "1", + "VLLM_B12X_DENSE_ACTIVATION_MODE": "auto", + "VLLM_MTP_NVFP4_LM_HEAD": "1", + "VLLM_LM_HEAD_A16": "1", + "VLLM_MXFP8_LM_HEAD": "0" + } +} diff --git a/runtime/glm53-spark-mtp3-mesh/compute/test_compute.py b/runtime/glm53-spark-mtp3-mesh/compute/test_compute.py new file mode 100644 index 00000000..78b7b7e4 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/compute/test_compute.py @@ -0,0 +1,191 @@ +import hashlib +import ast +import io +import importlib.util +import json +import tarfile +from pathlib import Path + +import pytest + +HERE = Path(__file__).resolve().parent + + +def _module(name: str): + spec = importlib.util.spec_from_file_location(name, HERE / f"{name}.py") + result = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(result) + return result + + +apply_compute = _module("apply_compute") +prepare_compute_source = _module("prepare_compute_source") +verify_compute = _module("verify_compute") + + +def test_compute_json_io_declares_utf8(): + for name in ("apply_compute.py", "prepare_compute_source.py", "verify_compute.py"): + tree = ast.parse((HERE / name).read_text(encoding="utf-8")) + for node in ast.walk(tree): + if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) and node.func.attr in ("read_text", "write_text"): + values = {arg.arg: arg.value for arg in node.keywords} + assert ast.literal_eval(values["encoding"]) == "utf-8", (name, node.lineno) + if node.func.attr == "write_text": + assert ast.literal_eval(values["newline"]) == "\n", (name, node.lineno) + + +@pytest.mark.parametrize("relative", ["../outside.py", "/tmp/outside.py", "vllm/../../outside.py", "b12x/not-vllm.py", "vllm\\outside.py"]) +def test_verify_rejects_uncontained_override_paths(tmp_path, relative): + lock = tmp_path / "source-lock.json" + lock.write_text(json.dumps({"vllm": {"files": [[relative, "base", "result"]]}}), encoding="utf-8") + receipt = tmp_path / "receipt.json" + receipt.write_text(json.dumps({ + "source_lock_sha256": hashlib.sha256(lock.read_bytes()).hexdigest(), + "vllm_overrides": {relative: "result"}, + }), encoding="utf-8") + with pytest.raises(ValueError, match="Unsafe vLLM override path"): + verify_compute.verify(tmp_path, receipt, lock) + + +@pytest.mark.parametrize("fault", [None, "base", "result", "archive", "extra", "duplicate"]) +def test_b12x_selector_overrides_fail_closed(tmp_path: Path, fault: str | None) -> None: + root = tmp_path / "source" + path = root / "b12x" / "selector.py" + path.parent.mkdir(parents=True) + path.write_bytes(b"base") + archive = tmp_path / "selector.tar.gz" + with tarfile.open(archive, "w:gz") as output: + names = ["b12x/selector.py"] + if fault == "extra": + names.append("b12x/unexpected.py") + for name in names: + member = tarfile.TarInfo(name) + member.size = len(b"result") + output.addfile(member, io.BytesIO(b"result")) + contract = { + "archive_sha256": hashlib.sha256(archive.read_bytes()).hexdigest(), + "files": [["b12x/selector.py", hashlib.sha256(b"base").hexdigest(), + hashlib.sha256(b"result").hexdigest()]], + } + if fault == "base": + contract["files"][0][1] = "wrong" + elif fault == "result": + contract["files"][0][2] = "wrong" + elif fault == "archive": + contract["archive_sha256"] = "wrong" + elif fault == "duplicate": + contract["files"].append(contract["files"][0]) + if fault: + with pytest.raises(ValueError): + prepare_compute_source._apply_b12x_overrides(root, archive, contract) + assert path.read_bytes() == b"base" + else: + prepare_compute_source._apply_b12x_overrides(root, archive, contract) + assert path.read_bytes() == b"result" + + +def test_source_lock_binds_patch_routes_and_environment() -> None: + lock = json.loads((HERE / "source-lock.json").read_text()) + patch = HERE / lock["vllm"]["patch"] + assert hashlib.sha256(patch.read_bytes()).hexdigest() == lock["vllm"][ + "patch_sha256" + ] + archive = HERE / lock["vllm"]["replacement_archive"] + assert hashlib.sha256(archive.read_bytes()).hexdigest() == lock["vllm"][ + "replacement_archive_sha256" + ] + files = lock["vllm"]["files"] + assert len(files) == 24 + assert len({entry[0] for entry in files}) == len(files) + assert all(base != result for _, base, result in files) + assert lock["b12x"]["revision"] == "ef308bac0f3b3eb8fea63e4013afc0c2ea1c6301" + assert len(lock["b12x"]["overrides"]["files"]) == 3 + assert lock["vllm"]["loader_donor_revision"] == "17e341b9ede04269f81fcac69a29951a0668a94a" + assert lock["vllm"]["rng_donor_revision"] == "44e6766e3397e8fe8ed9c1fa8a8d2783bb4a2ae8" + assert lock["environment"] == { + "VLLM_B12X_DENSE_ACTIVATION_MODE": "auto", + "VLLM_GDN_SPEC_DECODE_METADATA_FASTPATH": "1", + "VLLM_LM_HEAD_A16": "1", + "VLLM_MTP_NVFP4_LM_HEAD": "1", + "VLLM_MXFP8_LM_HEAD": "0", + } + + +def test_package_map_includes_runtime_data(tmp_path: Path) -> None: + package = tmp_path / "b12x" + package.mkdir() + (package / "module.py").write_text("value = 1\n") + (package / "profile.json.gz").write_bytes(b"profile") + (package / "README.md").write_text("runtime data\n") + (package / "ignored.pyc").write_bytes(b"cache") + files = prepare_compute_source._package_map(tmp_path, "b12x") + assert set(files) == { + "b12x/README.md", + "b12x/module.py", + "b12x/profile.json.gz", + } + + +def test_verify_rejects_an_empty_b12x_map(tmp_path: Path) -> None: + source_lock = tmp_path / "source-lock.json" + source_lock.write_text( + json.dumps( + { + "vllm": {"files": []}, + "b12x": {"package_files_sha256": "unused"}, + "environment": {}, + } + ) + ) + receipt = tmp_path / "receipt.json" + receipt.write_text( + json.dumps( + { + "source_lock_sha256": hashlib.sha256( + source_lock.read_bytes() + ).hexdigest(), + "vllm_overrides": {}, + "b12x_files": {}, + "environment": {}, + "target_head_quantization": False, + } + ) + ) + with pytest.raises(ValueError, match="no B12X package map"): + verify_compute.verify(tmp_path, receipt, source_lock) + + +def test_vllm_install_fails_before_patch_when_base_hash_drifts( + tmp_path: Path, +) -> None: + prepared = tmp_path / "prepared" + prepared.mkdir() + patch = prepared / "change.patch" + patch.write_text("") + archive = prepared / "files.tar.gz" + with tarfile.open(archive, "w:gz"): + pass + site = tmp_path / "site" + file = site / "vllm/example.py" + file.parent.mkdir(parents=True) + file.write_text("unexpected\n") + lock = { + "vllm": { + "patch": patch.name, + "patch_sha256": hashlib.sha256(b"").hexdigest(), + "replacement_archive": archive.name, + "replacement_archive_sha256": hashlib.sha256( + archive.read_bytes() + ).hexdigest(), + "files": [ + [ + "vllm/example.py", + hashlib.sha256(b"expected\n").hexdigest(), + "x", + ] + ], + } + } + with pytest.raises(ValueError, match="vLLM base hash mismatch"): + apply_compute._install_vllm(prepared, site, lock) diff --git a/runtime/glm53-spark-mtp3-mesh/compute/verify_compute.py b/runtime/glm53-spark-mtp3-mesh/compute/verify_compute.py new file mode 100644 index 00000000..5435a01a --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/compute/verify_compute.py @@ -0,0 +1,80 @@ +"""Verify the installed compute source against its build receipt.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path, PurePosixPath + + +def _sha256(path: Path) -> str: + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def _map_sha256(files: dict[str, str]) -> str: + payload = json.dumps(files, sort_keys=True, separators=(",", ":")).encode() + return hashlib.sha256(payload).hexdigest() + + +def verify(site_packages: Path, receipt: Path, source_lock: Path) -> dict: + lock = json.loads(source_lock.read_text(encoding="utf-8")) + installed = json.loads(receipt.read_text(encoding="utf-8")) + lock_hash = _sha256(source_lock) + if installed["source_lock_sha256"] != lock_hash: + raise ValueError("installed receipt uses a different compute source lock") + expected_vllm = {path: result for path, _, result in lock["vllm"]["files"]} + if installed["vllm_overrides"] != expected_vllm: + raise ValueError("installed receipt omits or changes vLLM overrides") + for relative, expected in expected_vllm.items(): + path = PurePosixPath(relative) + target = site_packages / path + if (path.is_absolute() or str(path) != relative or ".." in path.parts + or "\\" in relative or not relative.startswith("vllm/") + or not target.resolve().is_relative_to((site_packages / "vllm").resolve()) + or any((site_packages / Path(*path.parts[:index])).is_symlink() + for index in range(1, len(path.parts) + 1))): + raise ValueError(f"Unsafe vLLM override path: {relative}") + actual = _sha256(target) + if actual != expected: + raise ValueError(f"installed vLLM hash mismatch for {relative}: {actual}") + expected_b12x = installed.get("b12x_files") + if not expected_b12x: + raise ValueError("installed receipt has no B12X package map") + actual_b12x = { + path.relative_to(site_packages).as_posix(): _sha256(path) + for path in sorted((site_packages / "b12x").rglob("*")) + if path.is_file() + and "__pycache__" not in path.parts + and path.suffix != ".pyc" + } + if actual_b12x != expected_b12x: + raise ValueError("installed B12X package map mismatch") + expected_b12x_hash = lock["b12x"]["package_files_sha256"] + if installed.get("b12x_package_files_sha256") != expected_b12x_hash: + raise ValueError("installed receipt has a different B12X package-map hash") + if _map_sha256(actual_b12x) != expected_b12x_hash: + raise ValueError("installed B12X package differs from source-lock.json") + if installed["environment"] != lock["environment"]: + raise ValueError("installed compute environment differs from source lock") + if installed.get("target_head_quantization") is not False: + raise ValueError("target LM head must remain unquantized") + return { + "checks_passed": True, + "source_lock_sha256": lock_hash, + "vllm_override_count": len(expected_vllm), + "b12x_file_count": len(expected_b12x), + } + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--site-packages", type=Path, required=True) + parser.add_argument("--receipt", type=Path, required=True) + parser.add_argument("--source-lock", type=Path, required=True) + args = parser.parse_args() + print(json.dumps(verify(args.site_packages, args.receipt, args.source_lock))) + + +if __name__ == "__main__": + main() diff --git a/runtime/glm53-spark-mtp3-mesh/compute/vllm-compute-files.tar.gz b/runtime/glm53-spark-mtp3-mesh/compute/vllm-compute-files.tar.gz new file mode 100644 index 00000000..12aab7ee Binary files /dev/null and b/runtime/glm53-spark-mtp3-mesh/compute/vllm-compute-files.tar.gz differ diff --git a/runtime/glm53-spark-mtp3-mesh/compute/vllm-e02-to-compute.patch b/runtime/glm53-spark-mtp3-mesh/compute/vllm-e02-to-compute.patch new file mode 100644 index 00000000..5b020124 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/compute/vllm-e02-to-compute.patch @@ -0,0 +1,1570 @@ +--- a/vllm/envs.py ++++ b/vllm/envs.py +@@ -190,6 +190,10 @@ + VLLM_HUMMING_USE_F16_ACCUM: bool = False + VLLM_HUMMING_MOE_GEMM_TYPE: Literal["indexed", "grouped", "auto"] | None = None + VLLM_B12X_MOE_FP4_FORCE_A16: bool = False ++ VLLM_GDN_SPEC_DECODE_METADATA_FASTPATH: bool = False ++ VLLM_MTP_NVFP4_LM_HEAD: bool = False ++ VLLM_LM_HEAD_A16: bool = True ++ VLLM_MXFP8_LM_HEAD: bool = False + VLLM_B12X_MLA_CKV_GATHER: bool = False + VLLM_B12X_MLA_CKV_GATHER_MIN_TOKENS: int = 16 + VLLM_B12X_MLA_CKV_GATHER_MAX_TOKENS: int = 524288 +@@ -1641,6 +1645,36 @@ + # Gather DCP-sharded C4 records before B12X sparse-MLA prefill. This avoids + # query replication plus the per-rank LSE combine and is opt-in while the + # path is being qualified on GLM5Next. ++ "VLLM_GDN_SPEC_DECODE_METADATA_FASTPATH": lambda: bool( ++ int(os.getenv("VLLM_GDN_SPEC_DECODE_METADATA_FASTPATH", "0")) ++ ), ++ # Quantize only a GLM5Next native-MTP proposal head at load time. ++ "VLLM_MTP_NVFP4_LM_HEAD": lambda: bool( ++ int(os.getenv("VLLM_MTP_NVFP4_LM_HEAD", "0")) ++ ), ++ # Preserve BF16 activations for a runtime-quantized proposal head. ++ "VLLM_LM_HEAD_A16": lambda: bool( ++ int(os.getenv("VLLM_LM_HEAD_A16", "1")) ++ ), ++ # Reserved for a separate target-head experiment; this arm keeps it off. ++ "VLLM_MXFP8_LM_HEAD": lambda: bool( ++ int(os.getenv("VLLM_MXFP8_LM_HEAD", "0")) ++ ), ++ "VLLM_B12X_DENSE_ACTIVATION_MODE": env_with_choices( ++ "VLLM_B12X_DENSE_ACTIVATION_MODE", ++ "auto", ++ ["auto", "a16", "quantized"], ++ ), ++ "VLLM_B12X_NVFP4_ACTIVATION_MODE": env_with_choices( ++ "VLLM_B12X_NVFP4_ACTIVATION_MODE", ++ None, ++ ["auto", "a16", "quantized"], ++ ), ++ "VLLM_B12X_MXFP8_ACTIVATION_MODE": env_with_choices( ++ "VLLM_B12X_MXFP8_ACTIVATION_MODE", ++ None, ++ ["auto", "a16", "quantized"], ++ ), + "VLLM_B12X_MLA_CKV_GATHER": lambda: ( + os.getenv("VLLM_B12X_MLA_CKV_GATHER", "0").lower() in ("1", "true", "yes", "on") + ), +--- a/vllm/model_executor/kernels/linear/mxfp8/b12x.py ++++ b/vllm/model_executor/kernels/linear/mxfp8/b12x.py +@@ -12,11 +12,14 @@ + ) + from vllm.model_executor.utils import replace_parameter + from vllm.platforms import current_platform +-from vllm.utils.b12x import B12xWarmupUnit, reuse_packed_weight_storage + from vllm.utils.b12x import ( +- get_b12x_mxfp8_linear as _import_b12x_mxfp8, ++ B12xWarmupUnit, ++ get_b12x_dense_activation_mode, ++ reuse_packed_weight_storage, + ) +-from vllm.utils.torch_utils import current_stream ++from vllm.utils.b12x import ( ++ get_b12x_blockscaled as _import_b12x_blockscaled, ++) + + from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig + +@@ -31,13 +34,23 @@ + input_2d = x.reshape(-1, x.shape[-1]).contiguous() + output_shape = [*x.shape[:-1], int(packed_weight.out_features)] + +- mxfp8 = _import_b12x_mxfp8() ++ mode = layer.b12x_activation_mode ++ options = {} ++ if x.dtype == torch.bfloat16 and layer.b12x_bf16_input_supported: ++ options["mode"] = mode ++ else: ++ if mode == "a16": ++ raise ValueError( ++ "b12x MXFP8 A16 requires BF16 on SM120/SM121 with K%128=N%8=0" ++ ) ++ mxfp8 = _import_b12x_blockscaled() + assert mxfp8 is not None + output = mxfp8.mm( + input_2d, + packed_weight, + bias=bias, + expected_m=max(1, int(input_2d.shape[0])), ++ **options, + ) + return output.view(*output_shape) + +@@ -55,11 +68,16 @@ + return False, "b12x MXFP8 kernels are only available on CUDA" + if not current_platform.is_device_capability_family(120): + return False, "b12x MXFP8 kernels require a Blackwell 12x device" +- mxfp8 = _import_b12x_mxfp8() ++ mxfp8 = _import_b12x_blockscaled() + if mxfp8 is None: + return False, "Install the B12X backend with `pip install vllm[b12x]`" + if not mxfp8.is_supported(): +- return False, "b12x.gemm.mxfp8_linear is not supported" ++ return False, "b12x.gemm.blockscaled is not supported" ++ if not hasattr(mxfp8, "w8a16"): ++ return ( ++ False, ++ "b12x MXFP8 requires a source build with dense precision selection", ++ ) + return True, None + + @classmethod +@@ -89,7 +107,7 @@ + f"b12x MXFP8 weight_scale must be 2D, got {weight_scale.ndim}D" + ) + +- mxfp8 = _import_b12x_mxfp8() ++ mxfp8 = _import_b12x_blockscaled() + assert mxfp8 is not None + scale_k = in_features // MXFP8_BLOCK_SIZE + packed_weight = mxfp8.pack_weight( +@@ -99,6 +117,12 @@ + layer.b12x_mxfp8_packed_weight = reuse_packed_weight_storage( + getattr(layer, "b12x_mxfp8_packed_weight", None), + packed_weight, ++ ) ++ layer.b12x_activation_mode = get_b12x_dense_activation_mode("mxfp8") ++ layer.b12x_bf16_input_supported = ( ++ current_platform.is_device_capability_family(120) ++ and in_features % 128 == 0 ++ and out_features % 8 == 0 + ) + replace_parameter(layer, "weight", weight.new_empty((0,))) + replace_parameter(layer, "weight_scale", weight_scale.new_empty((0,))) +@@ -114,20 +138,13 @@ + device = torch.device(packed_weight.weight.values.device) + + def compile() -> None: +- mxfp8 = _import_b12x_mxfp8() +- assert mxfp8 is not None + for tokens in token_counts: + source = torch.zeros( + (tokens, int(packed_weight.in_features)), + dtype=output_dtype, + device=device, + ) +- mxfp8.mm( +- source, +- packed_weight, +- expected_m=max(1, int(tokens)), +- stream=current_stream().cuda_stream, +- ) ++ _apply_b12x_mxfp8_packed_linear(layer, source, None) + + return B12xWarmupUnit( + name="MXFP8", +@@ -137,6 +154,8 @@ + int(packed_weight.in_features), + int(packed_weight.padded_in_features), + int(packed_weight.out_features), ++ layer.b12x_activation_mode, ++ layer.b12x_bf16_input_supported, + output_dtype, + ), + compile=compile, +--- a/vllm/model_executor/kernels/linear/nvfp4/b12x.py ++++ b/vllm/model_executor/kernels/linear/nvfp4/b12x.py +@@ -8,7 +8,7 @@ + from vllm._custom_ops import scaled_fp4_quant + from vllm.model_executor.utils import replace_parameter + from vllm.platforms import current_platform +-from vllm.utils.b12x import B12xWarmupUnit ++from vllm.utils.b12x import B12xWarmupUnit, get_b12x_dense_activation_mode + from vllm.utils.b12x import ( + get_b12x_blockscaled as _import_b12x_blockscaled, + ) +@@ -18,30 +18,40 @@ + + + def _apply_b12x_nvfp4_linear( ++ layer: torch.nn.Module, + x: torch.Tensor, +- weight: torch.Tensor, +- weight_scale_storage: torch.Tensor, +- input_global_scale_inv: torch.Tensor, +- alpha: torch.Tensor, + bias: torch.Tensor | None, + ) -> torch.Tensor: + blockscaled = _import_b12x_blockscaled() + assert blockscaled is not None + +- output_size = int(weight.shape[0]) ++ output_size = int(layer.weight.shape[0]) + output_shape = [*x.shape[:-1], output_size] + x_2d = x.reshape(-1, x.shape[-1]) ++ mode = layer.b12x_activation_mode ++ if x.dtype == torch.bfloat16 and layer.b12x_bf16_input_supported: ++ output = blockscaled.mm( ++ x_2d.contiguous(), ++ layer.b12x_nvfp4_packed_weight, ++ mode=mode, ++ activation_global_scale=layer.input_global_scale_inv, ++ bias=bias, ++ expected_m=max(1, int(x_2d.shape[0])), ++ ) ++ return output.view(*output_shape) ++ if mode == "a16": ++ raise ValueError("b12x NVFP4 A16 requires BF16 on SM120/SM121 with K%128=N%8=0") + x_packed, x_scale_swizzled = scaled_fp4_quant( + x_2d, +- input_global_scale_inv, ++ layer.input_global_scale_inv, + is_sf_swizzled_layout=True, + ) + output = blockscaled.mm_nvfp4( + x_packed, + x_scale_swizzled, +- weight, +- weight_scale_storage, +- alpha, ++ layer.weight, ++ layer.weight_scale, ++ layer.alpha, + out_dtype=x.dtype, + ) + if bias is not None: +@@ -66,6 +76,11 @@ + return False, "Install the B12X backend with `pip install vllm[b12x]`" + if not blockscaled.is_supported(): + return False, "b12x native NVFP4 GEMM is not supported" ++ if not hasattr(blockscaled, "w4a16"): ++ return ( ++ False, ++ "b12x NVFP4 requires a source build with dense precision selection", ++ ) + return True, None + + @classmethod +@@ -80,6 +95,21 @@ + layer, + "weight_scale", + intrinsics.swizzle_block_scale(layer.weight_scale.data), ++ ) ++ blockscaled = _import_b12x_blockscaled() ++ assert blockscaled is not None ++ layer.b12x_nvfp4_packed_weight = blockscaled.pack_weight( ++ layer.weight.data, ++ layer.weight_scale.data, ++ recipe="nvfp4", ++ global_scale=layer.weight_global_scale, ++ ) ++ layer.b12x_activation_mode = get_b12x_dense_activation_mode("nvfp4") ++ n, packed_k = layer.weight.shape ++ layer.b12x_bf16_input_supported = ( ++ current_platform.is_device_capability_family(120) ++ and (packed_k * 2) % 128 == 0 ++ and n % 8 == 0 + ) + layer.b12x_warmup_provider = self + +@@ -100,11 +130,8 @@ + (tokens, k), dtype=output_dtype, device=weight.device + ) + _apply_b12x_nvfp4_linear( ++ layer, + source, +- weight, +- weight_scale, +- layer.input_global_scale_inv, +- layer.alpha, + None, + ) + +@@ -117,6 +144,8 @@ + k, + weight.dtype, + weight_scale.dtype, ++ layer.b12x_activation_mode, ++ layer.b12x_bf16_input_supported, + output_dtype, + ), + compile=compile, +@@ -129,11 +158,8 @@ + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + return _apply_b12x_nvfp4_linear( ++ layer, + x, +- layer.weight, +- layer.weight_scale, +- layer.input_global_scale_inv, +- layer.alpha, + bias, + ) + +--- a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py ++++ b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py +@@ -991,8 +991,7 @@ + """Execute B12X KDA after the convolution projection. + + Args: +- metadata: Forward-context metadata used to share runtime-owned +- packed metadata tensors across compatible layers. ++ metadata: Describes whether packed query boundaries are uniform. + mixed_qkv: Live packed query, key, and value projection. + raw_g: Live unactivated forget gate. + raw_beta: Live unactivated update gate. +@@ -1032,17 +1031,16 @@ + cache = forward_context.additional_kwargs.setdefault( + "b12x_kda_metadata_tensors", {} + ) ++ # Uniform builders own separate buffers with identical fixed boundaries. + cache_key = ( +- id(metadata), ++ None if metadata.is_uniform_spec_decode else query_start_loc.data_ptr(), ++ num_accepted_tokens.data_ptr() if num_accepted_tokens is not None else None, + num_tokens, + num_requests, +- state_columns, +- plan.caps.max_state_slots, + ) + bound_metadata = cache.get(cache_key) + if bound_metadata is None: + query_start_loc = query_start_loc[: num_requests + 1] +- state_indices = state_indices[:num_requests, :state_columns] + if num_accepted_tokens is None: + accepted_tokens = self._b12x_kda_num_accepted_tokens[:num_requests] + accepted_tokens.fill_(1) +@@ -1055,7 +1053,6 @@ + bound_metadata = ( + query_start_loc, + accepted_tokens, +- state_indices, + self._b12x_kda_num_seqs, + self._b12x_kda_num_tokens, + ) +@@ -1063,7 +1060,6 @@ + ( + query_start_loc, + accepted_tokens, +- state_indices, + num_seqs, + num_tokens_tensor, + ) = bound_metadata +@@ -1081,7 +1077,7 @@ + recurrent_state=self.kv_cache[1], + query_start_loc=query_start_loc, + num_accepted_tokens=accepted_tokens, +- state_indices=state_indices, ++ state_indices=state_indices[:num_requests, :state_columns], + num_seqs=num_seqs, + num_tokens=num_tokens_tensor, + output=output, +--- a/vllm/model_executor/layers/quantization/online/nvfp4.py ++++ b/vllm/model_executor/layers/quantization/online/nvfp4.py +@@ -5,6 +5,10 @@ + from torch.nn import Module + + from vllm._custom_ops import scaled_fp4_quant ++from vllm.model_executor.kernels.linear.nvfp4.b12x import ( ++ B12xNvFp4LinearKernel, ++) ++from vllm.model_executor.kernels.linear.nvfp4.base import NvFp4LinearLayerConfig + from vllm.model_executor.layers.fused_moe import RoutedExperts + from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig + from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import ( +@@ -13,6 +17,7 @@ + make_nvfp4_moe_quant_config, + select_nvfp4_moe_backend, + ) ++from vllm.model_executor.layers.quantization.online.fp8 import _Fp8OnlineLinearBase + from vllm.model_executor.layers.quantization.online.moe_base import ( + OnlineMoEMethodBase, + ) +@@ -27,8 +32,72 @@ + ) + from vllm.model_executor.utils import replace_parameter + from vllm.platforms import current_platform ++from vllm.utils.b12x import get_b12x_blockscaled + + FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max ++ ++ ++class Nvfp4OnlineLinearMethod(_Fp8OnlineLinearBase): ++ """Load a BF16 linear weight as NVFP4 for the proposal head.""" ++ ++ def __init__(self, *, use_a16: bool = False): ++ super().__init__() ++ supported, reason = B12xNvFp4LinearKernel.is_supported() ++ if not supported: ++ raise ValueError(f"Online NVFP4 proposal head requires b12x: {reason}") ++ self.kernel = B12xNvFp4LinearKernel(NvFp4LinearLayerConfig()) ++ self.use_a16 = use_a16 ++ if use_a16 and self.input_dtype != torch.bfloat16: ++ raise ValueError("A16 proposal heads require BF16 activations") ++ ++ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: ++ if getattr(layer, "_already_called_process_weights_after_loading", False): ++ return ++ weight = layer.weight.contiguous() ++ if weight.shape[1] % 16: ++ raise ValueError("Online NVFP4 proposal head requires K divisible by 16") ++ amax = weight.abs().amax().float().clamp_min(1e-8) ++ global_scale = (FLOAT4_E2M1_MAX * FLOAT8_E4M3_MAX) / amax ++ packed, scales = scaled_fp4_quant( ++ weight, global_scale, is_sf_swizzled_layout=False ++ ) ++ replace_parameter(layer, "weight", packed) ++ replace_parameter(layer, "weight_scale", scales) ++ replace_parameter(layer, "weight_global_scale", global_scale.reciprocal()) ++ replace_parameter(layer, "input_global_scale_inv", torch.ones_like(amax)) ++ replace_parameter(layer, "alpha", layer.weight_global_scale.clone()) ++ self.kernel.process_weights_after_loading(layer) ++ layer.b12x_activation_mode = "a16" if self.use_a16 else "quantized" ++ layer._already_called_process_weights_after_loading = True ++ ++ def apply( ++ self, ++ layer: torch.nn.Module, ++ x: torch.Tensor, ++ bias: torch.Tensor | None = None, ++ ) -> torch.Tensor: ++ if self.use_a16: ++ return self.kernel.apply_weights(layer, x, bias) ++ amax = x.abs().amax().float().clamp_min(1e-8) ++ input_scale = amax / (FLOAT4_E2M1_MAX * FLOAT8_E4M3_MAX) ++ x_packed, x_scale = scaled_fp4_quant( ++ x.reshape(-1, x.shape[-1]), ++ input_scale.reciprocal(), ++ is_sf_swizzled_layout=True, ++ ) ++ blockscaled = get_b12x_blockscaled() ++ assert blockscaled is not None ++ output = blockscaled.mm_nvfp4( ++ x_packed, ++ x_scale, ++ layer.weight, ++ layer.weight_scale, ++ input_scale * layer.weight_global_scale, ++ out_dtype=x.dtype, ++ ) ++ if bias is not None: ++ output = output + bias ++ return output.view(*x.shape[:-1], layer.weight.shape[0]) + + + def _quantize_moe_weight_to_nvfp4( +--- a/vllm/model_executor/layers/vocab_parallel_embedding.py ++++ b/vllm/model_executor/layers/vocab_parallel_embedding.py +@@ -3,6 +3,7 @@ + + from collections.abc import Sequence + from dataclasses import dataclass ++from typing import Literal + + import torch + import torch.nn.functional as F +@@ -248,6 +249,7 @@ + prefix: str = "", + *, + disable_tp: bool = False, ++ lm_head_quantization: Literal["nvfp4"] | None = None, + ): + super().__init__() + +@@ -291,6 +293,25 @@ + # method must implement the embedding operation. If we are another + # layer type like ParallelLMHead, this is not important. + is_embedding_layer = not isinstance(self, ParallelLMHead) ++ self.runtime_lm_head_quantization: Literal["nvfp4"] | None = None ++ if not is_embedding_layer and lm_head_quantization is not None: ++ from vllm.model_executor.layers.linear import UnquantizedLinearMethod ++ from vllm.model_executor.layers.quantization.online.nvfp4 import ( ++ Nvfp4OnlineLinearMethod, ++ ) ++ ++ if not isinstance( ++ quant_method, ++ (UnquantizedEmbeddingMethod, UnquantizedLinearMethod), ++ ): ++ raise ValueError( ++ "Runtime NVFP4 proposal-head quantization requires an " ++ "unquantized checkpoint head" ++ ) ++ quant_method = Nvfp4OnlineLinearMethod( ++ use_a16=envs.VLLM_LM_HEAD_A16 ++ ) ++ self.runtime_lm_head_quantization = "nvfp4" + quant_method_implements_embedding = method_has_implemented_embedding( + type(quant_method) + ) +@@ -549,6 +570,7 @@ + prefix: str = "", + *, + disable_tp: bool = False, ++ lm_head_quantization: Literal["nvfp4"] | None = None, + ): + super().__init__( + num_embeddings, +@@ -559,6 +581,7 @@ + quant_config, + prefix, + disable_tp=disable_tp, ++ lm_head_quantization=lm_head_quantization, + ) + self.quant_config = quant_config + if bias: +--- a/vllm/model_executor/models/deepseek_mtp.py ++++ b/vllm/model_executor/models/deepseek_mtp.py +@@ -52,6 +52,8 @@ + config: PretrainedConfig, + prefix: str, + quant_config: QuantizationConfig | None = None, ++ *, ++ lm_head_quantization: typing.Literal["nvfp4"] | None = None, + ) -> None: + super().__init__() + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) +@@ -60,6 +62,7 @@ + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "head"), ++ lm_head_quantization=lm_head_quantization, + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: +--- a/vllm/models/glm5next/model_state.py ++++ b/vllm/models/glm5next/model_state.py +@@ -310,25 +310,13 @@ + num_decode_draft_tokens_cpu = torch.from_numpy(num_decode_draft_tokens_np) + + if self._align_mode: +- mamba_group_ids, _ = self._get_mamba_group_info(kv_cache_config) +- aligned_index_builders = [] +- for group_idx, group_id in enumerate(mamba_group_ids): +- for group in attn_groups[group_id]: +- builder = group.get_metadata_builder(0) +- if hasattr(builder, "mamba_aligned_state_indices"): +- aligned_index_builders.append((group_idx, builder)) +- if aligned_index_builders: +- ctx = self._ensure_align_ctx( +- kv_cache_config, +- mamba_group_ids, +- block_tables, +- ) +- all_group_indices = ctx.compute_aligned_state_indices( +- input_batch.seq_lens, +- num_reqs, +- ) +- for group_idx, builder in aligned_index_builders: +- builder.mamba_aligned_state_indices = all_group_indices[group_idx] ++ self._prepare_aligned_state_indices( ++ input_batch.seq_lens, ++ num_reqs, ++ attn_groups, ++ kv_cache_config, ++ block_tables, ++ ) + + model_metadata = Glm5NextAttnMetadata( + is_prefilling=is_prefilling, +--- a/vllm/models/glm5next/nvidia/model.py ++++ b/vllm/models/glm5next/nvidia/model.py +@@ -1334,7 +1334,7 @@ + return False + + entry = buf.setdefault(layer_prefix, {}).setdefault(key, {}) +- entry["weight" if is_weight else "scale"] = tensor ++ entry["weight" if is_weight else "scale"] = tensor.clone() + if "weight" not in entry or "scale" not in entry: + return True + +@@ -1385,7 +1385,7 @@ + return False + + entry = buf.setdefault(layer_prefix, {}).setdefault("indexer_weights", {}) +- entry["weight" if is_weight else "scale"] = tensor ++ entry["weight" if is_weight else "scale"] = tensor.clone() + if "weight" not in entry or "scale" not in entry: + return True + +--- a/vllm/models/glm5next/nvidia/mtp.py ++++ b/vllm/models/glm5next/nvidia/mtp.py +@@ -6,6 +6,7 @@ + import torch + import torch.nn as nn + ++import vllm.envs as envs + from vllm.config import VllmConfig + from vllm.model_executor.layers.fused_moe import ( + fused_moe_make_expert_params_mapping, +@@ -64,7 +65,10 @@ + device=current_platform.device_type, + ) + self.shared_head = SharedHead( +- config=config, prefix=prefix, quant_config=quant_config ++ config=config, ++ prefix=prefix, ++ quant_config=quant_config, ++ lm_head_quantization="nvfp4" if envs.VLLM_MTP_NVFP4_LM_HEAD else None, + ) + # MTP layers sit past the base model's hidden layers; parse the index + # from the prefix (e.g. "...layers.32") so the decoder builds an MLA +@@ -245,14 +249,20 @@ + super().__init__() + self.config = vllm_config.model_config.hf_config + self.quant_config = vllm_config.quant_config +- self.checkpoint_weight_name_prefixes = self._checkpoint_weight_name_prefixes() ++ if envs.VLLM_MTP_NVFP4_LM_HEAD and self.config.tie_word_embeddings: ++ raise ValueError("NVFP4 proposal head requires untied word embeddings") + self.model = Glm5NextMultiTokenPredictor( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) ++ head = self.model._mtp_layers[0].shared_head.head ++ self.has_own_lm_head = head.runtime_lm_head_quantization == "nvfp4" ++ self.checkpoint_weight_name_prefixes = self._checkpoint_weight_name_prefixes() ++ if self.has_own_lm_head: ++ self.lm_head = head + self.set_moe_parameters() + + def _checkpoint_weight_name_prefixes(self) -> tuple[str, ...]: +- return tuple( ++ prefixes = tuple( + prefix + for layer_idx in range( + self.config.num_hidden_layers, +@@ -265,6 +275,14 @@ + f"layers.{layer_idx}.", + ) + ) ++ if self.has_own_lm_head: ++ prefixes += ( ++ "lm_head.", ++ "model.lm_head.", ++ "model.language_model.lm_head.", ++ "language_model.lm_head.", ++ ) ++ return prefixes + + def set_moe_parameters(self): + self.num_moe_layers = self.config.num_nextn_predict_layers +@@ -372,6 +390,19 @@ + # prefix to match. + if name.startswith("model.language_model."): + name = name.replace("model.language_model.", "model.", 1) ++ if name in ( ++ "lm_head.weight", ++ "model.lm_head.weight", ++ "language_model.lm_head.weight", ++ ): ++ if self.has_own_lm_head: ++ for layer_idx in self.model.layers: ++ head_name = f"model.layers.{layer_idx}.shared_head.head.weight" ++ if head_name not in loaded_params: ++ param = params_dict[head_name] ++ param.weight_loader(param, loaded_weight) ++ loaded_params.add(head_name) ++ continue + spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) + if spec_layer is None: + continue +@@ -460,6 +491,8 @@ + + loaded_layers: set[int] = set() + for param_name in loaded_params: ++ if param_name.endswith(".shared_head.head.weight"): ++ continue + spec_layer = get_spec_layer_idx_from_weight_name(self.config, param_name) + if spec_layer is not None: + loaded_layers.add(spec_layer) +@@ -467,6 +500,13 @@ + self.model.mtp_start_layer_idx, + self.model.mtp_start_layer_idx + self.model.num_mtp_layers, + ): ++ if self.has_own_lm_head: ++ head_name = f"model.layers.{layer_idx}.shared_head.head.weight" ++ if head_name not in loaded_params: ++ raise ValueError( ++ f"NVFP4 MTP head {layer_idx} requires an unquantized " ++ "proposal head or target lm_head.weight in the checkpoint." ++ ) + if layer_idx not in loaded_layers: + raise ValueError( + f"MTP speculative decoding layer {layer_idx} weights " +--- a/vllm/utils/b12x.py ++++ b/vllm/utils/b12x.py +@@ -7,9 +7,11 @@ + from collections.abc import Callable, Hashable, Iterable + from dataclasses import dataclass, fields, is_dataclass + from types import ModuleType +-from typing import Any ++from typing import Any, Literal + + import torch ++ ++import vllm.envs as envs + + + @dataclass(frozen=True) +@@ -17,6 +19,12 @@ + name: str + key: Hashable + compile: Callable[[], None] ++ ++ ++def get_b12x_dense_activation_mode(recipe: Literal["nvfp4", "mxfp8"]) -> str: ++ """Resolve the dense precision override once when loading a layer.""" ++ override = getattr(envs, f"VLLM_B12X_{recipe.upper()}_ACTIVATION_MODE") ++ return override if override is not None else envs.VLLM_B12X_DENSE_ACTIVATION_MODE + + + _HAS_B12X = importlib.util.find_spec("b12x") is not None +--- a/vllm/v1/attention/backends/gdn_attn.py ++++ b/vllm/v1/attention/backends/gdn_attn.py +@@ -2,11 +2,13 @@ + # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + """Backend for GatedDeltaNet attention.""" + ++from copy import copy + from dataclasses import dataclass, replace + from typing import Literal + + import torch + ++import vllm.envs as envs + from vllm.config import VllmConfig + from vllm.utils.torch_utils import async_tensor_h2d + from vllm.v1.attention.backend import ( +@@ -101,6 +103,7 @@ + seq_lens: torch.Tensor | None = None + + prefill_checkpoint: GDNPrefillCheckpointMetadata | None = None ++ is_uniform_spec_decode: bool = False + + + class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]): +@@ -108,7 +111,9 @@ + _cudagraph_support = AttentionCGSupport.UNIFORM_BATCH + supports_update_block_table: bool = True + ++ # Runner-owned stable storage, with NULL_BLOCK_ID in padded request rows. + mamba_aligned_state_indices: torch.Tensor | None = None ++ mamba_spec_accepted_tokens: torch.Tensor | None = None + + reorder_batch_threshold: int = 1 + +@@ -189,6 +194,84 @@ + (self.decode_cudagraph_max_bs,), + dtype=torch.int32, + device=device, ++ ) ++ self._decode_state_indices_source: torch.Tensor | None = None ++ self._decode_state_indices_view: torch.Tensor | None = None ++ self._reuse_spec_decode_inputs = envs.VLLM_GDN_SPEC_DECODE_METADATA_FASTPATH ++ self._uniform_spec_masks = torch.ones( ++ self.decode_cudagraph_max_bs, dtype=torch.bool, device=device ++ ) ++ self._uniform_spec_masks_cpu = torch.ones( ++ self.decode_cudagraph_max_bs, dtype=torch.bool ++ ) ++ self._uniform_spec_tokens = torch.arange( ++ self.decode_cudagraph_max_bs, dtype=torch.int32, device=device ++ ) ++ self._uniform_spec_query_start = torch.arange( ++ self.decode_cudagraph_max_bs + 1, dtype=torch.int32, device=device ++ ) * (self.num_spec + 1) ++ self._spec_state_indices_source: torch.Tensor | None = None ++ self._spec_state_indices_view: torch.Tensor | None = None ++ ++ def _can_reuse_spec_inputs( ++ self, ++ m: CommonAttentionMetadata, ++ num_accepted_tokens: torch.Tensor | None, ++ num_decode_draft_tokens_cpu: torch.Tensor | None, ++ ) -> bool: ++ return ( ++ self._reuse_spec_decode_inputs ++ and self.use_spec_decode ++ and self.use_full_cuda_graph ++ and self.vllm_config.cache_config.mamba_cache_mode == "align" ++ and self.mamba_aligned_state_indices is not None ++ and self.mamba_spec_accepted_tokens is not None ++ and num_accepted_tokens is not None ++ and num_decode_draft_tokens_cpu is not None ++ and 0 < m.num_actual_tokens <= self.decode_cudagraph_max_bs ++ and m.num_actual_tokens == m.num_reqs * (self.num_spec + 1) ++ and bool(torch.all(num_decode_draft_tokens_cpu == self.num_spec)) ++ and bool(torch.all(torch.diff(m.query_start_loc_cpu) == self.num_spec + 1)) ++ and (m.is_prefilling is None or not bool(torch.any(m.is_prefilling))) ++ ) ++ ++ def _get_spec_state_indices_view(self, num_reqs: int) -> torch.Tensor: ++ source = self.mamba_aligned_state_indices ++ assert source is not None ++ if ( ++ self._spec_state_indices_source is not source ++ or self._spec_state_indices_view is None ++ or self._spec_state_indices_view.shape[0] != num_reqs ++ ): ++ self._spec_state_indices_source = source ++ self._spec_state_indices_view = source[:num_reqs, : self.num_spec + 1] ++ return self._spec_state_indices_view ++ ++ def _build_uniform_spec_decode( ++ self, m: CommonAttentionMetadata, num_accepted_tokens: torch.Tensor ++ ) -> GDNAttentionMetadata: ++ num_reqs = m.num_reqs ++ assert self.mamba_spec_accepted_tokens is not None ++ accepted = self.mamba_spec_accepted_tokens[:num_reqs] ++ accepted.copy_(num_accepted_tokens[:num_reqs], non_blocking=True) ++ return GDNAttentionMetadata( ++ num_prefills=0, ++ num_prefill_tokens=0, ++ num_decodes=0, ++ num_decode_tokens=0, ++ num_spec_decodes=num_reqs, ++ num_spec_decode_tokens=m.num_actual_tokens, ++ num_actual_tokens=m.num_actual_tokens, ++ spec_query_start_loc=self._uniform_spec_query_start[: num_reqs + 1], ++ spec_state_indices_tensor=self._get_spec_state_indices_view(num_reqs), ++ spec_sequence_masks=self._uniform_spec_masks[:num_reqs], ++ spec_sequence_masks_cpu=self._uniform_spec_masks_cpu[:num_reqs], ++ spec_token_indx=self._uniform_spec_tokens[: m.num_actual_tokens], ++ non_spec_token_indx=self._uniform_spec_tokens[:0], ++ num_accepted_tokens=accepted, ++ num_reqs=num_reqs, ++ seq_lens=m.seq_lens, ++ is_uniform_spec_decode=True, + ) + + def _get_state_indices( +@@ -209,6 +292,13 @@ + self.vllm_config.cache_config.mamba_cache_mode, + ) + ++ def _can_reuse_decode_inputs(self) -> bool: ++ return ( ++ not self.use_spec_decode ++ and self.vllm_config.cache_config.mamba_cache_mode == "align" ++ and self.mamba_aligned_state_indices is not None ++ ) ++ + def _build_chunk_metadata( + self, + prefill_query_start_loc: torch.Tensor, +@@ -260,6 +350,11 @@ + fast_build: bool = False, + ) -> GDNAttentionMetadata: + m = common_attn_metadata ++ if self._can_reuse_spec_inputs( ++ m, num_accepted_tokens, num_decode_draft_tokens_cpu ++ ): ++ assert num_accepted_tokens is not None ++ return self._build_uniform_spec_decode(m, num_accepted_tokens) + + query_start_loc = m.query_start_loc + query_start_loc_cpu = m.query_start_loc_cpu +@@ -599,6 +694,7 @@ + and num_prefills == 0 + and num_spec_decodes == 0 + and num_decodes <= self.decode_cudagraph_max_bs ++ and not self._can_reuse_decode_inputs() + ): + self.non_spec_state_indices_tensor[:num_decodes].copy_( + non_spec_state_indices_tensor, non_blocking=True +@@ -657,6 +753,41 @@ + assert metadata.num_reqs > 0 + assert metadata.seq_lens is not None + ++ if ( ++ metadata.is_uniform_spec_decode ++ and self._reuse_spec_decode_inputs ++ and self.mamba_aligned_state_indices is not None ++ and self.mamba_spec_accepted_tokens is not None ++ ): ++ updated = copy(metadata) ++ updated.spec_state_indices_tensor = self._get_spec_state_indices_view( ++ metadata.num_reqs ++ ) ++ accepted = self.mamba_spec_accepted_tokens[: metadata.num_reqs] ++ assert metadata.num_accepted_tokens is not None ++ if accepted.data_ptr() != metadata.num_accepted_tokens.data_ptr(): ++ accepted.copy_(metadata.num_accepted_tokens, non_blocking=True) ++ updated.num_accepted_tokens = accepted ++ return updated ++ ++ if ( ++ metadata.num_prefills == 0 ++ and metadata.num_spec_decodes == 0 ++ and self._can_reuse_decode_inputs() ++ ): ++ source = self.mamba_aligned_state_indices ++ assert source is not None ++ if ( ++ self._decode_state_indices_source is not source ++ or self._decode_state_indices_view is None ++ or self._decode_state_indices_view.shape[0] != metadata.num_reqs ++ ): ++ self._decode_state_indices_source = source ++ self._decode_state_indices_view = source[: metadata.num_reqs, 0] ++ updated = copy(metadata) ++ updated.non_spec_state_indices_tensor = self._decode_state_indices_view ++ return updated ++ + state_indices = self._get_state_indices( + blk_table, + metadata.seq_lens, +@@ -753,6 +884,7 @@ + and metadata.num_prefills == 0 + and metadata.num_spec_decodes == 0 + and metadata.num_decodes <= self.decode_cudagraph_max_bs ++ and not self._can_reuse_decode_inputs() + ): + self.non_spec_state_indices_tensor[: metadata.num_decodes].copy_( + non_spec_state_indices[: metadata.num_decodes], non_blocking=True +--- a/vllm/v1/worker/gpu/attn_utils.py ++++ b/vllm/v1/worker/gpu/attn_utils.py +@@ -293,43 +293,11 @@ + attn_metadata: dict[str, Any] = {} + cached_attn_metadata: dict[tuple[KVCacheSpec, type], Any] = {} + num_kv_cache_groups = len(kv_cache_config.kv_cache_groups) ++ group_slot_mappings = slot_mappings[:num_kv_cache_groups].unbind(0) + for i in range(num_kv_cache_groups): + block_table = block_tables[i] +- slot_mapping = slot_mappings[i] +- # Per-group causal for hybrid drafters (mixed SWA/full attention). +- group_causal = ( +- causal if isinstance(causal, (bool, torch.Tensor)) else causal.get(i, True) +- ) +- +- common_attn_metadata_extra_kwargs = ( +- model_specific_attn_metadata.get_extra_common_attn_kwargs(i, num_reqs) +- if model_specific_attn_metadata is not None +- else {} +- ) +- # Model-specific metadata (e.g. Mamba hybrid) may supply its own +- # padding-aware is_prefilling, which takes precedence over the default. +- group_is_prefilling = common_attn_metadata_extra_kwargs.pop( +- "is_prefilling", is_prefilling +- ) +- common_attn_metadata = CommonAttentionMetadata( +- query_start_loc=query_start_loc_gpu, +- query_start_loc_cpu=query_start_loc_cpu, +- seq_lens=seq_lens, +- seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, +- max_seq_len=max_seq_len, +- num_reqs=num_reqs, +- num_actual_tokens=num_tokens, +- max_query_len=max_query_len, +- block_table_tensor=block_table, +- slot_mapping=slot_mapping, +- causal=group_causal, +- dcp_local_seq_lens=dcp_local_seq_lens, +- positions=positions, +- is_prefilling=group_is_prefilling, +- mm_req_doc_ranges=mm_req_doc_ranges, +- rswa_prefix_lens=rswa_prefix_lens, +- **common_attn_metadata_extra_kwargs, +- ) ++ slot_mapping = group_slot_mappings[i] ++ common_attn_metadata = None + + for attn_group in attn_groups[i]: + attn_metadata_builder = attn_group.get_metadata_builder(0) +@@ -337,35 +305,74 @@ + if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs): + kv_cache_spec = kv_cache_spec.kv_cache_specs[attn_group.layer_names[0]] + cache_key = (kv_cache_spec, type(attn_metadata_builder)) +- if for_cudagraph_capture: +- metadata = attn_metadata_builder.build_for_cudagraph_capture( +- common_attn_metadata +- ) +- elif ( +- cache_key in cached_attn_metadata ++ if ( ++ not for_cudagraph_capture ++ and cache_key in cached_attn_metadata + and attn_metadata_builder.supports_update_block_table + ): + metadata = attn_metadata_builder.update_block_table( +- cached_attn_metadata[cache_key], +- common_attn_metadata.block_table_tensor, +- common_attn_metadata.slot_mapping, ++ cached_attn_metadata[cache_key], block_table, slot_mapping + ) + else: +- attn_metadata_extra_kwargs = ( +- model_specific_attn_metadata.get_extra_attn_kwargs( +- attn_metadata_builder, +- num_reqs, +- ) +- if model_specific_attn_metadata is not None +- else {} +- ) +- metadata = attn_metadata_builder.build( +- common_prefix_len=0, +- common_attn_metadata=common_attn_metadata, +- **attn_metadata_extra_kwargs, +- ) +- if attn_metadata_builder.supports_update_block_table: +- cached_attn_metadata[cache_key] = metadata ++ if common_attn_metadata is None: ++ # Per-group causal for hybrid drafters (mixed SWA/full attention). ++ group_causal = ( ++ causal ++ if isinstance(causal, (bool, torch.Tensor)) ++ else causal.get(i, True) ++ ) ++ ++ common_attn_metadata_extra_kwargs = ( ++ model_specific_attn_metadata.get_extra_common_attn_kwargs( ++ i, num_reqs ++ ) ++ if model_specific_attn_metadata is not None ++ else {} ++ ) ++ # Model-specific padding takes precedence over the default. ++ group_is_prefilling = common_attn_metadata_extra_kwargs.pop( ++ "is_prefilling", is_prefilling ++ ) ++ common_attn_metadata = CommonAttentionMetadata( ++ query_start_loc=query_start_loc_gpu, ++ query_start_loc_cpu=query_start_loc_cpu, ++ seq_lens=seq_lens, ++ seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, ++ max_seq_len=max_seq_len, ++ num_reqs=num_reqs, ++ num_actual_tokens=num_tokens, ++ max_query_len=max_query_len, ++ block_table_tensor=block_table, ++ slot_mapping=slot_mapping, ++ causal=group_causal, ++ dcp_local_seq_lens=dcp_local_seq_lens, ++ positions=positions, ++ is_prefilling=group_is_prefilling, ++ mm_req_doc_ranges=mm_req_doc_ranges, ++ rswa_prefix_lens=rswa_prefix_lens, ++ **common_attn_metadata_extra_kwargs, ++ ) ++ ++ if for_cudagraph_capture: ++ metadata = attn_metadata_builder.build_for_cudagraph_capture( ++ common_attn_metadata ++ ) ++ else: ++ attn_metadata_extra_kwargs = ( ++ model_specific_attn_metadata.get_extra_attn_kwargs( ++ attn_metadata_builder, ++ num_reqs, ++ ) ++ if model_specific_attn_metadata is not None ++ else {} ++ ) ++ metadata = attn_metadata_builder.build( ++ common_prefix_len=0, ++ common_attn_metadata=common_attn_metadata, ++ **attn_metadata_extra_kwargs, ++ ) ++ if attn_metadata_builder.supports_update_block_table: ++ cached_attn_metadata[cache_key] = metadata + for layer_name in attn_group.layer_names: + attn_metadata[layer_name] = metadata + return attn_metadata +--- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py ++++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +@@ -7,6 +7,7 @@ + import torch + import torch.nn as nn + ++from vllm import envs + from vllm.config import VllmConfig + from vllm.config.compilation import CUDAGraphMode + from vllm.triton_utils import tl, triton +@@ -84,6 +85,11 @@ + self.num_accepted_tokens_gpu = torch.ones( + self.max_num_reqs, dtype=torch.int32, device=self.device + ) ++ self._gdn_spec_accepted_tokens = ( ++ torch.ones_like(self.num_accepted_tokens_gpu) ++ if envs.VLLM_GDN_SPEC_DECODE_METADATA_FASTPATH ++ else None ++ ) + # Pre-copy "align" prefix-cache state (V2). The migration of each + # request's mamba state across block boundaries runs as a fused GPU + # kernel reusing the postprocess copy machinery, so the per-step src +@@ -106,6 +112,9 @@ + self._mamba_group_ids: list[int] = [] + self._mamba_spec: MambaSpec | None = None + self._mamba_copy_funcs_by_type: MambaStateCopyFuncsByType | None = None ++ self._aligned_metadata_groups: list[list[AttentionGroup]] | None = None ++ self._aligned_metadata_builders: list[tuple[int, Any]] = [] ++ self._aligned_metadata_ctx: MambaSpecDecodeGPUContext | None = None + + def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: + super().add_request(req_index, new_req_data) +@@ -122,6 +131,9 @@ + if self._align_mode: + self._mamba_ctx = None + self._mamba_copy_funcs_by_type = None ++ self._aligned_metadata_ctx = None ++ self._aligned_metadata_groups = None ++ self._aligned_metadata_builders = [] + if self.recoverssm is not None: + self.recoverssm.reset() + +@@ -170,6 +182,40 @@ + [block_tables[gid] for gid in mamba_group_ids], + ) + return ctx ++ ++ def _prepare_aligned_state_indices( ++ self, ++ seq_lens: torch.Tensor, ++ num_reqs: int, ++ attn_groups: list[list[AttentionGroup]], ++ kv_cache_config: KVCacheConfig, ++ block_tables: tuple[torch.Tensor, ...], ++ ) -> None: ++ mamba_group_ids, _ = self._get_mamba_group_info(kv_cache_config) ++ if self._aligned_metadata_groups is not attn_groups: ++ self._aligned_metadata_builders = [] ++ for group_idx, group_id in enumerate(mamba_group_ids): ++ for group in attn_groups[group_id]: ++ builder = group.get_metadata_builder(0) ++ if hasattr(builder, "mamba_aligned_state_indices"): ++ self._aligned_metadata_builders.append((group_idx, builder)) ++ if hasattr(builder, "mamba_spec_accepted_tokens"): ++ builder.mamba_spec_accepted_tokens = ( ++ self._gdn_spec_accepted_tokens ++ ) ++ self._aligned_metadata_groups = attn_groups ++ self._aligned_metadata_ctx = None ++ if not self._aligned_metadata_builders: ++ return ++ ++ ctx = self._ensure_align_ctx(kv_cache_config, mamba_group_ids, block_tables) ++ if self._aligned_metadata_ctx is not ctx: ++ assert ctx.aligned_state_indices is not None ++ group_views = ctx.aligned_state_indices.unbind(0) ++ for group_idx, builder in self._aligned_metadata_builders: ++ builder.mamba_aligned_state_indices = group_views[group_idx] ++ self._aligned_metadata_ctx = ctx ++ ctx.compute_aligned_state_indices(seq_lens, num_reqs) + + def preprocess_state( + self, +@@ -277,22 +323,13 @@ + num_decode_draft_tokens_cpu = torch.from_numpy(num_decode_draft_tokens_np) + + if self._align_mode: +- mamba_group_ids, _ = self._get_mamba_group_info(kv_cache_config) +- aligned_index_builders = [] +- for group_idx, group_id in enumerate(mamba_group_ids): +- for group in attn_groups[group_id]: +- builder = group.get_metadata_builder(0) +- if hasattr(builder, "mamba_aligned_state_indices"): +- aligned_index_builders.append((group_idx, builder)) +- if aligned_index_builders: +- ctx = self._ensure_align_ctx( +- kv_cache_config, mamba_group_ids, block_tables +- ) +- all_group_indices = ctx.compute_aligned_state_indices( +- input_batch.seq_lens, num_reqs +- ) +- for group_idx, builder in aligned_index_builders: +- builder.mamba_aligned_state_indices = all_group_indices[group_idx] ++ self._prepare_aligned_state_indices( ++ input_batch.seq_lens, ++ num_reqs, ++ attn_groups, ++ kv_cache_config, ++ block_tables, ++ ) + + mamba_attn_metadata = MambaHybridAttnMetadata( + is_prefilling=is_prefilling, +--- a/vllm/v1/worker/gpu/sample/gumbel.py ++++ b/vllm/v1/worker/gpu/sample/gumbel.py +@@ -12,6 +12,12 @@ + # available — on the CPU worker path `tl` is a placeholder whose `constexpr` + # attribute is `None`, and `tl.constexpr(...)` would crash at import time. + _TL_RAND_MIN = tl.constexpr(4.6566127342e-10) if HAS_TRITON else 4.6566127342e-10 ++ ++# Offset salt keeping the draft's Gumbel noise disjoint from the target's. ++# Verification is a probability-ratio test, not a Gumbel coupling, so a proposal ++# and the residual it is resampled from must not share a noise vector. ++# Positions are int64 and never approach 2**30, so the streams cannot collide. ++_DRAFT_NOISE_SALT = tl.constexpr(1 << 30) if HAS_TRITON else (1 << 30) + + + @triton.jit +@@ -89,6 +95,7 @@ + seed, + pos, + temp, ++ IS_DRAFTING: tl.constexpr, + USE_FP64: tl.constexpr, + APPLY_TEMPERATURE: tl.constexpr = True, + ): +@@ -108,6 +115,8 @@ + if USE_FP64: + logits = logits.to(tl.float64) + if temp != 0.0: ++ if IS_DRAFTING: ++ pos = pos + _DRAFT_NOISE_SALT + gumbel_seed = tl.randint(seed, pos) + if USE_FP64: + u = tl_rand64(gumbel_seed, keys, includes_zero=False) +@@ -137,6 +146,7 @@ + logits_cache_stride_1, + logits_cache_col_ptr, + vocab_size, ++ IS_DRAFTING: tl.constexpr, + APPLY_TEMPERATURE: tl.constexpr, + USE_FP64: tl.constexpr, + PER_TOKEN_COL: tl.constexpr = False, +@@ -173,6 +183,7 @@ + seed, + pos, + temp, ++ IS_DRAFTING=IS_DRAFTING, + USE_FP64=USE_FP64, + APPLY_TEMPERATURE=APPLY_TEMPERATURE, + ) +@@ -197,6 +208,7 @@ + temp_ptr, + vocab_size, + BLOCK_SIZE: tl.constexpr, ++ IS_DRAFTING: tl.constexpr, + APPLY_TEMPERATURE: tl.constexpr, + USE_FP64: tl.constexpr, + PER_TOKEN_COL: tl.constexpr, +@@ -226,6 +238,7 @@ + logits_cache_stride_1, + logits_cache_col_ptr, + vocab_size, ++ IS_DRAFTING=IS_DRAFTING, + APPLY_TEMPERATURE=APPLY_TEMPERATURE, + USE_FP64=USE_FP64, + PER_TOKEN_COL=PER_TOKEN_COL, +@@ -242,6 +255,7 @@ + seed: torch.Tensor, # [max_num_reqs] + pos: torch.Tensor, # [num_tokens] + apply_temperature: bool, ++ is_drafting: bool, + logits_cache: torch.Tensor | None = None, # [max_num_reqs, num_cols, vocab_size] + logits_cache_col: torch.Tensor | None = None, # scalar or [num_tokens] + use_fp64: bool = False, +@@ -281,6 +295,7 @@ + temperature, + vocab_size, + BLOCK_SIZE=BLOCK_SIZE, ++ IS_DRAFTING=is_drafting, + APPLY_TEMPERATURE=apply_temperature, + USE_FP64=use_fp64, + PER_TOKEN_COL=per_token_col, +--- a/vllm/v1/worker/gpu/sample/sampler.py ++++ b/vllm/v1/worker/gpu/sample/sampler.py +@@ -317,6 +317,7 @@ + self.sampling_states.seeds.gpu, + pos, + apply_temperature=False, ++ is_drafting=False, + use_fp64=self.use_fp64_gumbel, + ) + return sampled, processed_logits +--- a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py ++++ b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py +@@ -46,6 +46,9 @@ + ) + self.current_draft_step = torch.tensor(0, dtype=torch.int64, device=device) + self.last_token_indices = torch.zeros( ++ self.max_num_reqs, dtype=torch.int64, device=device ++ ) ++ self.sample_src_positions = torch.zeros( + self.max_num_reqs, dtype=torch.int64, device=device + ) + +@@ -370,6 +373,7 @@ + input_batch.seq_lens, + num_rejected, + self.input_buffers, ++ self.sample_src_positions, + self.max_model_len, + self.max_num_reqs, + advance_draft_positions=self.advance_draft_positions, +@@ -504,6 +508,8 @@ + ) -> None: + last_token_indices = self.last_token_indices[:num_reqs] + positions = self.input_buffers.positions[last_token_indices] ++ # Hidden state P and token P+1 predict P+2; key sampling by P+1. ++ sample_src_positions = positions + 1 + idx_mapping = self.idx_mapping[:num_reqs] + + last_hidden_states, hidden_states = self._run_model( +@@ -523,7 +529,7 @@ + + self.draft_tokens[:num_reqs, 0] = self.sample_draft( + sample_hidden_states, +- positions, ++ sample_src_positions, + idx_mapping, + self.temperature, + self.seeds, +@@ -543,6 +549,7 @@ + num_reqs, + ) + self.input_buffers.positions[:num_reqs] = positions ++ self.sample_src_positions[:num_reqs] = sample_src_positions + + def _multi_step_decode( + self, +@@ -698,7 +705,6 @@ + self._prepare_eplb_forward(num_reqs) + + idx_mapping = self.idx_mapping[:num_reqs] +- positions = self.input_buffers.positions[:num_reqs] + # Run the draft model forward pass. + last_hidden_states, hidden_states = self._run_model( + num_tokens_padded, +@@ -707,18 +713,12 @@ + num_tokens_across_dp, + cudagraph_runtime_mode, + ) +- last_hidden_states = last_hidden_states[:num_reqs] +- +- sample_positions = positions +- if not self.advance_draft_positions: +- # The forward pass holds positions fixed (Q-only, shared target KV), +- # but Gumbel sampling still needs the absolute draft position. +- sample_positions = positions + self.current_draft_step +- + # Sample the draft tokens. ++ sample_hidden_states = last_hidden_states[:num_reqs] ++ sample_src_positions = self.sample_src_positions[:num_reqs] + draft_tokens = self.sample_draft( +- last_hidden_states, +- sample_positions, ++ sample_hidden_states, ++ sample_src_positions, + idx_mapping, + self.temperature, + self.seeds, +@@ -734,6 +734,7 @@ + self.draft_tokens, + self.hidden_states, + self.input_buffers, ++ self.sample_src_positions, + num_reqs, + self.max_model_len, + self.num_speculative_steps, +@@ -1003,6 +1004,7 @@ + num_rejected_ptr, + input_ids_ptr, + positions_ptr, ++ sample_src_positions_ptr, + mrope_positions_ptr, + mrope_positions_stride, + query_start_loc_ptr, +@@ -1034,6 +1036,10 @@ + draft_token = tl.load(draft_tokens_ptr + req_idx * draft_tokens_stride) + tl.store(input_ids_ptr + req_idx, draft_token) + ++ # Sampling advances even when forward positions clamp at max_model_len. ++ sample_position = tl.load(sample_src_positions_ptr + req_idx) ++ tl.store(sample_src_positions_ptr + req_idx, sample_position + 1) ++ + target_seq_len = tl.load(target_seq_lens_ptr + req_idx) + num_rejected = tl.load(num_rejected_ptr + req_idx) + seq_len = target_seq_len - num_rejected +@@ -1061,6 +1067,7 @@ + target_seq_lens: torch.Tensor, + num_rejected: torch.Tensor, + input_buffers: InputBuffers, ++ sample_src_positions: torch.Tensor, + max_model_len: int, + max_num_reqs: int, + advance_draft_positions: bool = True, +@@ -1080,6 +1087,7 @@ + num_rejected, + input_buffers.input_ids, + input_buffers.positions, ++ sample_src_positions, + mrope_positions, + mrope_positions_stride, + input_buffers.query_start_loc, +@@ -1100,6 +1108,7 @@ + next_input_hidden_states_stride, + input_ids_ptr, + positions_ptr, ++ sample_src_positions_ptr, + mrope_positions_ptr, + mrope_positions_stride, + seq_lens_ptr, +@@ -1128,6 +1137,10 @@ + # This is the final step. Skip updating draft forward inputs. + return + ++ # Sampling advances even when forward positions clamp at max_model_len. ++ sample_position = tl.load(sample_src_positions_ptr + req_idx) ++ tl.store(sample_src_positions_ptr + req_idx, sample_position + 1) ++ + # Write the sampled draft token into the input ids tensor for the next + # forward pass. + tl.store(input_ids_ptr + req_idx, draft_token) +@@ -1177,6 +1190,7 @@ + output_draft_tokens: torch.Tensor, + next_input_hidden_states: torch.Tensor, + input_buffers: InputBuffers, ++ sample_src_positions: torch.Tensor, + num_reqs: int, + max_model_len: int, + num_speculative_steps: int, +@@ -1197,6 +1211,7 @@ + next_input_hidden_states.stride(0), + input_buffers.input_ids, + input_buffers.positions, ++ sample_src_positions, + mrope_positions, + mrope_positions_stride, + input_buffers.seq_lens, +--- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py ++++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +@@ -288,11 +288,11 @@ + ) + num_sample = num_reqs * self.num_speculative_steps + sample_hidden_states = last_hidden_states[self.sample_indices[:num_sample]] +- # sample_pos is the predicted token's position Q; verification keys +- # Gumbel by the predecessor (Q-1). sample_draft adds +1, so pass Q-2. ++ # sample_pos is the predicted token's position P. Sampling keys a draw ++ # by the position before the sampled token, P-1. + draft_tokens = self.sample_draft( + sample_hidden_states, +- self.sample_pos[:num_sample] - 2, ++ self.sample_pos[:num_sample] - 1, + self.sample_idx_mapping[:num_sample], + self.temperature, + self.seeds, +--- a/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py ++++ b/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py +@@ -51,15 +51,17 @@ + other=0, + ) + +- # Candidate ids key the noise, matching the target's own sampling. +- position = tl.load(sample_pos_ptr + flat) - 1 ++ # sample_pos is the predicted token's position P. Sampling keys a draw ++ # by the position before the sampled token, P-1. ++ sample_pos = tl.load(sample_pos_ptr + flat) - 1 + _, index = gumbel_noised_argmax( + scores, + candidates, + mask & valid, + seed, +- position, ++ sample_pos, + temperature if SAMPLE_PROBABILISTIC else 0.0, ++ IS_DRAFTING=True, + USE_FP64=USE_FP64, + ) + +--- a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py ++++ b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py +@@ -135,8 +135,8 @@ + buf.index_copy_(1, self._d2t_scatter_index, logits.to(buf.dtype)) + logits = buf + +- # sample_pos is the predicted token's position Q; the target verifies +- # it with the predecessor's Gumbel key (Q-1). Pass Q-1. ++ # sample_pos is the predicted token's position P. Sampling keys a draw ++ # by the position before the sampled token, P-1. + return gumbel_sample( + logits, + idx_map, +@@ -144,6 +144,7 @@ + self.seeds, + sample_pos - 1, + apply_temperature=True, ++ is_drafting=True, + logits_cache=self.draft_logits, + logits_cache_col=self._step_cols[step], + use_fp64=self.use_fp64_gumbel, +--- a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py ++++ b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py +@@ -377,7 +377,11 @@ + cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE, + ) -> None: + last_token_indices = self.last_token_indices[:num_reqs] +- sample_positions = self.input_buffers.positions[last_token_indices] ++ positions = self.input_buffers.positions[last_token_indices] ++ # The output hidden state at position P (= positions) and the token id ++ # at P+1 are used to draft the token at P+2. Sampling keys a draw by the ++ # position before the sampled token, so the net adjustment is +1. ++ sample_src_positions = positions + 1 + idx_mapping = self.idx_mapping[:num_reqs] + + # Cache the trailing token's ids, hidden states (and embeddings for +@@ -415,7 +419,7 @@ + sample_hidden_states = last_hidden_states[last_token_indices] + draft_tokens = self.sample_draft( + sample_hidden_states, +- sample_positions, ++ sample_src_positions, + idx_mapping, + self.temperature, + self.seeds, +@@ -446,7 +450,8 @@ + idx_mapping, + num_reqs, + ) +- sample_positions += 1 ++ # Advance the draft sampling key. ++ sample_src_positions += 1 + + + @triton.jit +--- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py ++++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py +@@ -837,6 +837,7 @@ + 0, # logits_cache_stride_1 + None, # logits_cache_col_ptr + vocab_size, ++ IS_DRAFTING=False, + APPLY_TEMPERATURE=False, + USE_FP64=USE_FP64, + ) +--- a/vllm/v1/worker/gpu/spec_decode/speculator.py ++++ b/vllm/v1/worker/gpu/spec_decode/speculator.py +@@ -391,7 +391,7 @@ + def sample_draft( + self, + hidden_states: torch.Tensor, +- positions: torch.Tensor, ++ sample_src_positions: torch.Tensor, + idx_mapping: torch.Tensor, + temperature: torch.Tensor, + seeds: torch.Tensor, +@@ -400,15 +400,14 @@ + ) -> torch.Tensor: + if draft_logits is not None: + logits = self.model.compute_logits(hidden_states) +- # NOTE(woosuk): We must add 1 to the positions to match the Gumbel noise +- # used for draft and target sampling. + return gumbel_sample( + logits, + idx_mapping, + temperature, + seeds, +- positions + 1, ++ sample_src_positions, + apply_temperature=True, ++ is_drafting=True, + logits_cache=draft_logits, + logits_cache_col=draft_step, + use_fp64=self.use_fp64_gumbel, +--- a/vllm/v1/worker/mamba_utils.py ++++ b/vllm/v1/worker/mamba_utils.py +@@ -85,8 +85,10 @@ + mask=( + valid_group[:, None, None] + & valid_row[None, :, None] ++ & (seq_lens[None, :, None] > 0) + & valid_state_slot[None, None, :] + ), ++ other=-1, + ) + tl.store( + state_indices_ptr diff --git a/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/.gitattributes b/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/.gitattributes new file mode 100644 index 00000000..66e025bc --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/.gitattributes @@ -0,0 +1,3 @@ +# Transform script hashes include their original line endings. +patch_mtp3_*.py -text whitespace=cr-at-eol +fixtures/*.zip -text diff --git a/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/README.md b/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/README.md new file mode 100644 index 00000000..0feb01c6 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/README.md @@ -0,0 +1,173 @@ +# MTP3 cache-reuse source experiment + +Status: **research-only**. These source transforms coordinate recurrent prefix +reuse, GPU-lease accounting, and the fused indexer's histogram barrier for +GLM-5.3 native-MTP3. `compose.py` produces and verifies their exact Python output +from one specified base image. The package includes source inputs and CPU tests +of the allocation, lookup, retention, and accounting rules. + +The composition tool is implemented. Serving and performance qualification +require separate evidence for the complete image, SparkCache package, native +libraries, and workload. Published runtime pins and production Dockerfiles are +outside this experiment's output. + +The scope is GLM-5.3 native-MTP3 on four GB10 ranks, TP4/DCP4, with 512-token +recurrent/hash pages, 2,048-token attention scheduling alignment, one prefill +lookahead token, and aligned recurrent caching. Other speculation methods, +geometries and images require separate evidence. + +## Exact inputs and outputs + +The supported source input is local image ID +`sha256:2e41b1e934a85ff7c21b780532db2f0a0e978df081e52f4ae2bf11f8992fb24f`. +This is a local immutable image identity, not a registry pull reference. + +This composer and the [indexer-barrier image in PR #226](https://github.com/FujitsuPolycom/sparkring/pull/226) +are separate derivatives of that input. Both add histogram synchronization, but +their indexer output bytes and hashes differ. Do not apply this composer to the +barrier child: its preimage checks reject that source. A combined image needs +one selected indexer output, matching source receipts and cache namespace, and +validation of the complete composition. A source-branch merge alone does not +perform that integration. + +`fixtures/manifest.json` lists the ten input Python files, their paths in +that image, byte counts and SHA-256 values. The 155 KiB fixture archive contains +only those files. Source bytes and their existing SPDX/copyright headers are +preserved. Models, compiled libraries, credentials and site configuration are +excluded. Fixtures are test inputs and do not enter a production image build. + +| Component | Output or required-library SHA-256 | +|---|---| +| vLLM scheduler | `75efa57e7ff5a77c76714b85e2e4d8e1d7f456d9a9eec6c67ebb11ca382942f9` | +| vLLM single-type KV manager | `d2e35b012e0cf45ab3771f545c35ca48f2a5858549c574a352975607369124e2` | +| B12X fused indexer | `b43a4a2802c7dfc4a049bbb5751fc7e7688b05cb06d4ece716ab7a1d91d23d2a` | +| Separately supplied SparkCache placement library | `2657cdd2e54a097c9544e4c79ae62c0646db6db123ff24e4f0c384238c3a1e8d` | +| Required SparkCache capture library | `4398f18b8913e743e7bf1ed8fe29560d4580e61b6a1e2ab8b16684b19b6573b5` | + +The fixture manifest attests each patch script's exact bytes. Each transform +accepts only its declared input SHA-256 or its own output SHA-256. Apply them in +this dependency order: + +1. `patch_mtp3_barrier.py` adds the CTA synchronization before the histogram + arrival signal and increments that kernel's compile-cache revision. +2. `patch_mtp3_lease_accounting.py` includes an attached GPU lease in initial + prefill cache statistics without labeling it an external transfer. +3. `patch_mtp3_local_lease_preference.py` prefers a strictly longer converged + local prefix over a shorter GPU lease and reuses that lookup result. +4. `patch_mtp3_sparse_retention.py` pairs an explicit speculative replay + checkpoint stop with sparse retention of its predecessor state. +5. `patch_mtp3_partial_tail_eligibility.py` decides whether a partial recurrent + tail requires a stop from the recurrent page size, rather than attention's + DCP-scaled scheduling unit. + +Every transform rejects an unknown input and checks its complete output. Its +own output is idempotent; source produced by another transform must still match +one of those two declared hashes. Use `compose.py --verify-candidate` to check +the composed output tree. + +The `original/` directory contains the exact source inputs. The `candidate/` +directory contains transformed output, and `composition.json` records output +hashes and transform receipts. These directory names describe the tool's file +interface, not deployment or qualification status. + +## Offline composition and tests + +Python 3.10 or later is sufficient for composition. Pytest is needed for tests. +No Docker, vLLM installation, Torch, model files or GPU is required: + +```bash +python runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/compose.py --check +python -m pytest runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse -q -rs +``` + +To retain the exact source trees and receipt in a fresh local directory: + +```bash +python runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/compose.py \ + --output-root work/mtp3-cache-reuse-composition +python runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/compose.py \ + --verify-candidate work/mtp3-cache-reuse-composition/candidate +python runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/check_mtp3_checkpoint_allocations.py \ + --source-root work/mtp3-cache-reuse-composition/original/vllm \ + --candidate-root work/mtp3-cache-reuse-composition/candidate/vllm \ + --output work/mtp3-checkpoint-allocations.json +``` + +Existing output directories and result files are rejected. `--source-root` +on the composition command accepts an independently extracted original tree, +with `vllm/` and `b12x/` immediately beneath it; every file is checked against +the same manifest before output creation. + +The allocator checker executes the actual allocation and block-registration +methods against planned running-state and GDN checkpoint writes. It covers +336 fresh and 336 resumed empty-table replay schedules with varied shared +boundaries, chunk budgets and speculative buffer counts. A selected null slot +is skipped by the real registration method. A selected non-null slot whose +planned state does not match its hash boundary fails the check. This tests +metadata consistency, not completion of CUDA writes. Some resumed schedules +can still miss a prompt predecessor; the report distinguishes those safe misses +from stale-state registration. + +The tests also cover lease/API accounting, local-versus-lease selection, +speculative backoff, prompt-length boundaries and recurrent partial-tail +eligibility. One optional test uses the companion SparkCache checkout directly: + +```bash +SPARKCACHE_SOURCE_ROOT=/path/to/sparkcache python -m pytest \ + runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse -q -rs +``` + +That optional test needs the companion checkout's CPU development dependencies. +It is explicitly skipped when `SPARKCACHE_SOURCE_ROOT` is absent. + +## Independent fixture extraction + +Use only the exact base image above on a local Docker host. A stopped container +can expose its files without launching Python, CUDA or serving. The following +Python fragment creates that stopped container, copies only manifest-listed +files, and removes the temporary container. It changes local Docker metadata; +do not run it against a serving container or a remote Docker context. + +```python +import json +from pathlib import Path +import subprocess + +package = Path('runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse') +manifest = json.loads((package / 'fixtures/manifest.json').read_text()) +output = Path('work/mtp3-extracted-original') +output.mkdir(parents=True, exist_ok=False) +image = manifest['base_image_id'] +observed = subprocess.check_output(['docker', 'image', 'inspect', image, + '--format', '{{.Id}}'], text=True).strip() +assert observed == image +container = subprocess.check_output(['docker', 'create', '--entrypoint', + '/bin/true', image], text=True).strip() +try: + for name, record in manifest['files'].items(): + target = output / name + target.parent.mkdir(parents=True, exist_ok=True) + subprocess.run(['docker', 'cp', container + ':' + record['image_path'], + str(target)], check=True) +finally: + subprocess.run(['docker', 'rm', container], check=True) +``` + +Then run `compose.py --source-root work/mtp3-extracted-original --check` to +verify the independently copied inputs and complete transform chain. + +## Serving composition boundary + +To serve with these Python sources, select a SparkCache commit and supply both +native libraries with the hashes listed above. The experimental image also +needs a runtime contract that attests the output scheduler and manager hashes. +The contract is +`sparkcache/runtime_patches/vllm-manager-page-async-contract-55969c16.json`; +the exact scheduler and manager preimage hashes are in the fixture manifest. +Record the resulting contract, full SparkCache source tree, native libraries, +image ID, model identity, topology, namespace and workload settings together. +The source composer does not rewrite that contract or published source pins. + +CPU test success establishes the stated source and metadata checks. Throughput, +deployment readiness, and unmeasured preemption/CUDA schedules require hardware +evidence that identifies the complete image and its test conditions. diff --git a/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/check_mtp3_checkpoint_allocations.py b/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/check_mtp3_checkpoint_allocations.py new file mode 100644 index 00000000..a62855da --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/check_mtp3_checkpoint_allocations.py @@ -0,0 +1,211 @@ +"""Reproduce MTP3 aligned checkpoint allocation and cache registration on CPU. + +Execute allocation and BlockPool.cache_full_blocks from the declared source +fixtures together with the transformed split/retention methods. Synthetic block +state labels track the running endpoint and GDN internal checkpoint scheduled +for GPU writes. This checks metadata/allocator consistency; it does not execute +GPU kernels or prove those writes completed on hardware. +Fresh and resumed cases start with an empty block table (preemption replay). +""" +from __future__ import annotations + +import argparse +import ast +import datetime +import hashlib +import json +import os +from pathlib import Path +from types import SimpleNamespace + +from repro_mtp3_checkpoint_materialization import ( + checkpoint_kernel_metadata, extracted_method, load_scheduler, + partial_hit_eligibility, +) +from repro_mtp3_sparse_retention import MambaSpec, ROOT, execute, load_algorithms + +WORK = Path(__file__).resolve().parent +PATCHED = Path(os.environ.get('MTP3_RETENTION_SOURCE', str(WORK / 'mtp3-retention-patched'))) +EXPECTED_SCHEDULER = '75efa57e7ff5a77c76714b85e2e4d8e1d7f456d9a9eec6c67ebb11ca382942f9' +EXPECTED_MANAGER = 'd2e35b012e0cf45ab3771f545c35ca48f2a5858549c574a352975607369124e2' +EXPECTED_ALLOCATOR = '10846c4994e7860deab8b42c8bcd3315ddc96d14a478d4012c398418cc17a04c' + + +def sha256(path): + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def extracted_classes(): + allocation = ast.parse('class Allocator:\n pass').body[0] + allocation.body = [extracted_method(ROOT / 'v1/core/single_type_kv_cache_manager.py', + 'MambaManager', method) + for method in ('_needs_internal_checkpoint', 'allocate_new_blocks')] + scope = {'MambaSpec': MambaSpec, 'cdiv': lambda a, b: (a + b - 1) // b} + execute([allocation], scope) + algorithms = load_algorithms(ROOT / 'v1/core/kv_cache_utils.py', + manager_source=PATCHED / 'v1/core/single_type_kv_cache_manager.py') + registration = ast.parse('class Registrar:\n pass').body[0] + registration.body = [extracted_method(ROOT / 'v1/core/block_pool.py', 'BlockPool', 'cache_full_blocks')] + register_scope = {'resolve_block_hashes': algorithms.resolve_block_hashes, + 'make_block_hash_with_group_id': lambda value, group: (value, group)} + execute([registration], register_scope) + return scope['Allocator'], register_scope['Registrar'], algorithms + + +def run_case(prompt, total, shared, budget, speculative, classes, Scheduler, checkpoint): + Allocator, Registrar, algorithms = classes + scheduler = Scheduler() + scheduler.cache_config = SimpleNamespace(block_size=512) + scheduler.use_eagle = True + scheduler.mamba_has_prefill_checkpoint_blocks = True + scheduler.hash_block_size = 512 + scheduler.block_size = 2048 + scheduler.need_mamba_block_aligned_split = True + scheduler.kv_cache_manager = SimpleNamespace(coordinator=SimpleNamespace(enable_partial_hash_hits=True)) + scheduler.mamba_partial_cache_hit = partial_hit_eligibility(scheduler, + SimpleNamespace(kv_cache_groups=[SimpleNamespace(kv_cache_spec=MambaSpec(512))]), + PATCHED / 'v1/core/sched/scheduler.py') + scheduler.max_num_scheduled_tokens = 8192 + scheduler.scheduler_config = SimpleNamespace(long_prefill_token_threshold=0) + scheduler.num_prefill_lookahead = 1 + scheduler._recurrent_publication_boundaries = lambda request: ((prompt - 1) // 2048 * 2048,) + request = SimpleNamespace(request_id='replay', num_prompt_tokens=prompt, num_tokens=total, + num_computed_tokens=0, shared_prefix_boundary=shared, + block_hashes=list(range(512, total + 1, 512))) + allocator = Allocator() + allocator.kv_cache_spec = MambaSpec(512) + allocator.kv_cache_spec.num_prefill_checkpoint_blocks = 1 + allocator.mamba_cache_mode = 'align' + allocator.block_size = 512 + allocator.num_speculative_blocks = speculative + allocator.req_to_blocks = {'replay': []} + allocator._num_checkpoint_blocks = {} + allocator._partial_hit_reqs = {} + allocator._allocated_block_reqs = set() + allocator.last_state_block_idx = {} + allocator._null_block = SimpleNamespace(is_null=True, state=None, block_hash=None) + allocator.block_pool = SimpleNamespace(get_new_blocks=lambda n: [ + SimpleNamespace(is_null=False, state=None, block_hash=None, block_hash_num_tokens=None) + for _ in range(n)]) + failures, retained, selected_null, steps = [], set(), 0, [] + registrar = Registrar() + registrar.hash_block_size = 512 + registrar.enable_kv_cache_events = False + + def insert(block_hash, block, *, num_tokens): + if block.state != num_tokens: + failures.append({'expected_tokens': num_tokens, 'state_tokens': block.state, + 'step_start': request.num_computed_tokens}) + else: + retained.add(num_tokens) + block.block_hash, block.block_hash_num_tokens = block_hash, num_tokens + + def remove(block): + block.block_hash, block.block_hash_num_tokens = None, None + return [] + + registrar._insert_block_hash = insert + registrar._remove_cached_block_hashes = remove + registrar._emit_block_removed_events = lambda removed: None + cached = 0 + while request.num_computed_tokens < total: + start = request.num_computed_tokens + count = scheduler._mamba_block_aligned_split(request, min(budget, total - start)) + count = scheduler._reserve_prefill_lookahead(request, start, count) + if count <= 0: + raise AssertionError(f'zero progress: {prompt=}, {total=}, {shared=}, {budget=}, {start=}') + end = start + count + last = allocator.last_state_block_idx.get('replay') + # The source fixture's remove_skipped_blocks retires this state slot. + if last is not None and last < (start + 511) // 512 - 1: + allocator.req_to_blocks['replay'][last] = allocator._null_block + allocator._num_checkpoint_blocks['replay'] = int( + allocator._needs_internal_checkpoint('replay', end, start)) + allocator.allocate_new_blocks('replay', end, end) + blocks = allocator.req_to_blocks['replay'] + running = blocks[(end + 511) // 512 - 1] + assert not running.is_null + running.state = end + internal = checkpoint(start, end, 512) + if internal is not None and not blocks[internal // 512 - 1].is_null: + blocks[internal // 512 - 1].state = internal + mask = algorithms.MambaManager.reachable_block_mask( + start_block=cached, end_block=end // 512, alignment_tokens=2048, + kv_cache_spec=allocator.kv_cache_spec, use_eagle=True, retention_interval=0, + reachable_boundaries=(prompt - 1, shared)) + selected_null += sum(keep and blocks[index].is_null for index, keep in enumerate(mask, cached)) + # Execute cache registration so its null-slot handling participates in + # the check of every hash boundary against the planned state writes. + registrar.cache_full_blocks(request, blocks, cached, end // 512, 512, 1, mask) + steps.append(end) + cached, request.num_computed_tokens = end // 512, end + assert len(steps) < 1024 + required_prompt = ((prompt - 1) // 512 - 1) * 512 + return {'prompt_tokens': prompt, 'num_tokens': total, 'shared_boundary': shared, + 'budget': budget, 'speculative_blocks': speculative, 'steps': steps, + 'retained_boundaries': sorted(retained), 'selected_null_slots': selected_null, + 'stale_or_unwritten_registered_states': failures, + 'prompt_predecessor_retained': required_prompt in retained} + + +def main(): + global ROOT, PATCHED + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--output', type=Path, required=True) + parser.add_argument('--source-root', type=Path, default=ROOT, help='Original vllm directory produced by compose.py') + parser.add_argument('--candidate-root', type=Path, default=PATCHED, help='Candidate vllm directory produced by compose.py') + args = parser.parse_args() + ROOT, PATCHED = args.source_root, args.candidate_root + # The extracted helpers resolve their unchanged support files from this root. + import repro_mtp3_sparse_retention as sparse_helpers + import repro_mtp3_checkpoint_materialization as checkpoint_helpers + sparse_helpers.ROOT = ROOT + checkpoint_helpers.ROOT = ROOT + paths = {'scheduler': PATCHED / 'v1/core/sched/scheduler.py', + 'retention_manager': PATCHED / 'v1/core/single_type_kv_cache_manager.py', + 'allocator_manager': ROOT / 'v1/core/single_type_kv_cache_manager.py', + 'block_pool': ROOT / 'v1/core/block_pool.py', + 'gdn_metadata': ROOT / 'v1/attention/backends/gdn_attn.py', + 'hash_utils': ROOT / 'v1/core/kv_cache_utils.py', + 'coordinator': ROOT / 'v1/core/kv_cache_coordinator.py', + 'algorithm_loader': WORK / 'repro_mtp3_sparse_retention.py', + 'checkpoint_loader': WORK / 'repro_mtp3_checkpoint_materialization.py', + 'checker': Path(__file__).resolve()} + identities = {name: {'path': str(path), 'sha256': sha256(path)} for name, path in paths.items()} + for name, expected in [('scheduler', EXPECTED_SCHEDULER), ('retention_manager', EXPECTED_MANAGER), + ('allocator_manager', EXPECTED_ALLOCATOR)]: + if identities[name]['sha256'] != expected: + raise RuntimeError(f'unexpected {name} SHA-256: {identities[name]["sha256"]}') + classes = extracted_classes() + Scheduler = load_scheduler(paths['scheduler']) + checkpoint = checkpoint_kernel_metadata() + populations = {'fresh': [(value, value) for value in (32768, 32789, 33280, 100968)], + 'resumed': [(32789, 33301), (32789, 40000), (32768, 33301), (32768, 40000)]} + result = {'schema': 'mtp3-checkpoint-allocation-registration-check/v1', + 'time_utc': datetime.datetime.now(datetime.timezone.utc).isoformat(), + 'source_inputs': identities, 'gpu_executed': False, 'populations': {}, + 'limits': ['Checks allocator/retention against planned endpoint/checkpoint writes; no GPU execution', + 'Fresh and resumed requests start from empty block tables, not an injected partial local hit', + 'Only align mode with 512-token recurrent/hash pages, DCP4 retention alignment2048, and lookahead1', + 'Missing retained prompt predecessors are reported safe misses, not treated as corruption']} + for label, pairs in populations.items(): + cases = [run_case(prompt, total, shared, budget, speculative, classes, Scheduler, checkpoint) + for prompt, total in pairs for shared in (8192, 12345, 15872, 16384, 16401, 20000, 30000) + for budget in (8192, 7680, 1024, 512) for speculative in (0, 1, 3)] + summary = {'cases': len(cases), 'selected_null_slots': sum(case['selected_null_slots'] for case in cases), + 'stale_or_unwritten_registered_states': sum(len(case['stale_or_unwritten_registered_states']) for case in cases), + 'cases_without_prompt_predecessor_retention': sum(not case['prompt_predecessor_retained'] for case in cases)} + result['populations'][label] = {'summary': summary, 'cases': cases} + assert summary['cases'] == 336 + result['passed'] = all(not item['summary']['stale_or_unwritten_registered_states'] + for item in result['populations'].values()) + with args.output.open('x', encoding='utf-8') as stream: + json.dump(result, stream, indent=2) + stream.write('\n') + print(json.dumps({'passed': result['passed'], 'output': str(args.output), + 'summary': {name: item['summary'] for name, item in result['populations'].items()}}, indent=2)) + return 0 if result['passed'] else 1 + + +if __name__ == '__main__': + raise SystemExit(main()) diff --git a/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/compose.py b/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/compose.py new file mode 100644 index 00000000..49559a85 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/compose.py @@ -0,0 +1,125 @@ +"""Build hash-verified MTP3 Python source output in an unused local directory. + +The candidate directory is transformed source output. Its hashes verify source +composition; serving qualification requires separate complete-image evidence. +The tool performs no Docker, network, or model operations. +""" +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path +import shutil +import tempfile +import zipfile + +import patch_mtp3_barrier as barrier +import patch_mtp3_lease_accounting as accounting +import patch_mtp3_local_lease_preference as preference +import patch_mtp3_partial_tail_eligibility as partial +import patch_mtp3_sparse_retention as retention + +HERE = Path(__file__).resolve().parent +FIXTURES = HERE / 'fixtures' +NATIVE_PLACEMENT_SHA256 = '2657cdd2e54a097c9544e4c79ae62c0646db6db123ff24e4f0c384238c3a1e8d' +FINAL_SHA256 = { + 'vllm/v1/core/sched/scheduler.py': partial.AFTER_SHA256, + 'vllm/v1/core/single_type_kv_cache_manager.py': retention.MANAGER_AFTER, + 'b12x/attention/dsa_indexer/fused_indexer.py': barrier.AFTER_SHA256, +} + + +def sha256(data): + return hashlib.sha256(data).hexdigest() + + +def manifest(): + return json.loads((FIXTURES / 'manifest.json').read_text(encoding='utf-8')) + + +def source_bytes(source_root=None): + """Validate every input before creating any output files.""" + metadata = manifest() + if source_root is None: + archive = FIXTURES / 'original-python-sources.zip' + if sha256(archive.read_bytes()) != metadata['archive_sha256']: + raise ValueError('Source fixture archive checksum differs') + with zipfile.ZipFile(archive) as bundle: + if sorted(bundle.namelist()) != sorted(metadata['files']): + raise ValueError('Source fixture archive members differ') + content = {name: bundle.read(name) for name in metadata['files']} + else: + content = {name: (Path(source_root) / name).read_bytes() for name in metadata['files']} + for name, data in content.items(): + if sha256(data) != metadata['files'][name]['sha256']: + raise ValueError(f'Original source checksum differs: {name}') + for name, expected in metadata['transform_script_sha256'].items(): + if sha256((HERE / name).read_bytes()) != expected: + raise ValueError(f'Attested transform script differs: {name}') + return content + + +def verify_candidate(root): + metadata = manifest() + verified = {} + for name, original in metadata['files'].items(): + expected = FINAL_SHA256.get(name, original['sha256']) + observed = sha256((Path(root) / name).read_bytes()) + if observed != expected: + raise ValueError(f'Candidate source checksum differs: {name}') + verified[name] = observed + return verified + + +def compose(output, *, source_root=None): + content = source_bytes(source_root) + output = Path(output) + output.mkdir(parents=True, exist_ok=False) + original, candidate = output / 'original', output / 'candidate' + for name, data in content.items(): + path = original / name + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(data) + shutil.copytree(original, candidate) + scheduler = candidate / 'vllm/v1/core/sched/scheduler.py' + steps = [barrier.apply_patch(candidate / 'b12x/attention/dsa_indexer/fused_indexer.py'), + accounting.apply_patch(scheduler), preference.apply_patch(scheduler), + retention.apply_patch(candidate / 'vllm'), partial.apply_patch(scheduler)] + result = {'schema': 'sparkring-mtp3-cache-reuse-composition/v1', 'status': 'research-only', + 'base_image_id': manifest()['base_image_id'], 'speculation': 'native MTP3', + 'topology': 'TP4/DCP4, 512-token recurrent/hash pages', + 'candidate_files': verify_candidate(candidate), 'transform_receipts': steps, + 'required_native_placement_sha256': NATIVE_PLACEMENT_SHA256, + 'limits': ['Python source composition only; no SparkCache package or native binary is installed', + 'No serving, transport, GPU, or throughput qualification is implied', + 'Published runtime pins and production Dockerfiles are unchanged']} + (output / 'composition.json').write_text(json.dumps(result, indent=2) + '\n', encoding='utf-8') + return result + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--source-root', type=Path, help='Original fixture layout from the exact base image; defaults to bundled fixtures') + parser.add_argument('--output-root', type=Path, help='Fresh destination; existing directories are rejected') + parser.add_argument('--check', action='store_true', help='Compose and verify only in a temporary directory') + parser.add_argument('--verify-candidate', type=Path, help='Read-only verification of an already composed candidate directory') + args = parser.parse_args() + try: + if args.verify_candidate: + result = {'verified_files': verify_candidate(args.verify_candidate)} + elif args.check: + with tempfile.TemporaryDirectory(prefix='mtp3-cache-reuse-') as temporary: + result = compose(Path(temporary) / 'composition', source_root=args.source_root) + elif args.output_root: + result = compose(args.output_root, source_root=args.source_root) + else: + parser.error('Choose --check, --output-root, or --verify-candidate') + except (ValueError, OSError) as error: + parser.error(str(error)) + print(json.dumps(result, sort_keys=True)) + return 0 + + +if __name__ == '__main__': + raise SystemExit(main()) diff --git a/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/conftest.py b/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/conftest.py new file mode 100644 index 00000000..7c6f36c2 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/conftest.py @@ -0,0 +1,32 @@ +"""Make source-extracted tests self-contained using an attested temporary composition.""" +import os +from pathlib import Path +import sys +import tempfile + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) +import compose # noqa: E402 + + +def pytest_configure(config): + temporary = tempfile.TemporaryDirectory(prefix='mtp3-cache-reuse-tests-') + config._mtp3_cache_reuse_temporary = temporary + output = Path(temporary.name) / 'composition' + compose.compose(output) + values = {'MTP3_ORIGINAL_SOURCE': str(output / 'original/vllm'), + 'MTP3_RETENTION_SOURCE': str(output / 'candidate/vllm'), + 'MTP3_SCHEDULER_SOURCE': str(output / 'candidate/vllm/v1/core/sched/scheduler.py')} + config._mtp3_cache_reuse_environment = {key: os.environ.get(key) for key in values} + os.environ.update(values) + + +def pytest_unconfigure(config): + for key, value in getattr(config, '_mtp3_cache_reuse_environment', {}).items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + temporary = getattr(config, '_mtp3_cache_reuse_temporary', None) + if temporary: + temporary.cleanup() diff --git a/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/fixtures/manifest.json b/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/fixtures/manifest.json new file mode 100644 index 00000000..2d3de924 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/fixtures/manifest.json @@ -0,0 +1,66 @@ +{ + "schema": "sparkring-mtp3-cache-reuse-fixtures/v1", + "status": "research-only", + "base_image_id": "sha256:2e41b1e934a85ff7c21b780532db2f0a0e978df081e52f4ae2bf11f8992fb24f", + "purpose": "Exact original Python files required by source transforms and extracted CPU tests; not model/runtime assets", + "files": { + "b12x/attention/dsa_indexer/fused_indexer.py": { + "sha256": "893fbcade135b7e1d146b8fb6530cde0650be515f69bf9a17ced0a9c61a141e2", + "bytes": 143038, + "image_path": "/usr/local/lib/python3.12/dist-packages/b12x/attention/dsa_indexer/fused_indexer.py" + }, + "vllm/entrypoints/openai/chat_completion/serving.py": { + "sha256": "9982953285e9df469032a82fffa4095d0e9d86278bede6e2b91d03d02373d182", + "bytes": 56414, + "image_path": "/usr/local/lib/python3.12/dist-packages/vllm/entrypoints/openai/chat_completion/serving.py" + }, + "vllm/v1/attention/backends/gdn_attn.py": { + "sha256": "fd75fb72efeb762ec558d332364e889df44a9ea483026bd5ded8f77f42be9b7f", + "bytes": 40208, + "image_path": "/usr/local/lib/python3.12/dist-packages/vllm/v1/attention/backends/gdn_attn.py" + }, + "vllm/v1/core/block_pool.py": { + "sha256": "ddee56dccb2208411b3a035918e917ce8f56a9858471e9ca12b420d5d79bc69c", + "bytes": 33243, + "image_path": "/usr/local/lib/python3.12/dist-packages/vllm/v1/core/block_pool.py" + }, + "vllm/v1/core/kv_cache_coordinator.py": { + "sha256": "149222b2b1de1ee3714c7c4e0843b9c38c730a605740ce7dc913308df906f20f", + "bytes": 41494, + "image_path": "/usr/local/lib/python3.12/dist-packages/vllm/v1/core/kv_cache_coordinator.py" + }, + "vllm/v1/core/kv_cache_utils.py": { + "sha256": "9b5ec588681252832e322387fc7ecea3b146061305f8abb1b02bb55e219c2a8c", + "bytes": 100179, + "image_path": "/usr/local/lib/python3.12/dist-packages/vllm/v1/core/kv_cache_utils.py" + }, + "vllm/v1/core/sched/scheduler.py": { + "sha256": "122d9722b48f8cec267d2136a18c184b8901a5a0a46187eab7030d6227f1983d", + "bytes": 157770, + "image_path": "/usr/local/lib/python3.12/dist-packages/vllm/v1/core/sched/scheduler.py" + }, + "vllm/v1/core/single_type_kv_cache_manager.py": { + "sha256": "10846c4994e7860deab8b42c8bcd3315ddc96d14a478d4012c398418cc17a04c", + "bytes": 89918, + "image_path": "/usr/local/lib/python3.12/dist-packages/vllm/v1/core/single_type_kv_cache_manager.py" + }, + "vllm/v1/engine/output_processor.py": { + "sha256": "f46abafd45792b4816e75c0702113b510f4b3e15ea33016d155ee00abaa3fb59", + "bytes": 33588, + "image_path": "/usr/local/lib/python3.12/dist-packages/vllm/v1/engine/output_processor.py" + }, + "vllm/v1/metrics/stats.py": { + "sha256": "e8c3cbcfb880e7a2a496a1e499d4e4efccf70ba1eef56f4a1a6ee44cc0bf49bc", + "bytes": 22775, + "image_path": "/usr/local/lib/python3.12/dist-packages/vllm/v1/metrics/stats.py" + } + }, + "archive_sha256": "b54fa2cdde26099905be15a3462144877c2c511299642c744daecbcdcdcf7219", + "transform_script_sha256": { + "patch_mtp3_barrier.py": "751cb641cfe42ffa7aa907774a8059edcf69c446d0299612d12b13ed7a1230ec", + "patch_mtp3_lease_accounting.py": "ece6753aa79571eb0f7a9901ac8be69f3bfbf26d10bbdc315134847977d019ae", + "patch_mtp3_local_lease_preference.py": "3489feb646f2b768220422988a6d120a3f82e30972df8c1095d1c4b03fab0bb8", + "patch_mtp3_sparse_retention.py": "a6d27c5804fc4a79b01835beaae01d33a8ef495e98c30b655121d620c62d2afa", + "patch_mtp3_partial_tail_eligibility.py": "2e0609a62a3028841431e92531ea3b5af958cc863274b8518547ce56da7f2c89" + } +} diff --git a/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/fixtures/original-python-sources.zip b/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/fixtures/original-python-sources.zip new file mode 100644 index 00000000..c52e2a83 Binary files /dev/null and b/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/fixtures/original-python-sources.zip differ diff --git a/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/patch_mtp3_barrier.py b/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/patch_mtp3_barrier.py new file mode 100644 index 00000000..dbb88a4a --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/patch_mtp3_barrier.py @@ -0,0 +1,52 @@ +#!/usr/bin/env python3 +"""Patch the attested MTP3 B12X histogram barrier, preserving its source bytes.""" + +from __future__ import annotations + +import argparse +import ast +import hashlib +import json +from pathlib import Path + + +BEFORE_SHA256 = "893fbcade135b7e1d146b8fb6530cde0650be515f69bf9a17ced0a9c61a141e2" +AFTER_SHA256 = "b43a4a2802c7dfc4a049bbb5751fc7e7688b05cb06d4ece716ab7a1d91d23d2a" +_BEFORE = ''' """Grid barrier over the group's CTAs on the arrival counter; returns the next phase.""" + arrival_ptr = _fused_state_ptr(state, group_id, Int32(_FUSED_STATE_ARRIVAL)) +'''.replace("\n", "\r\n").encode() +_AFTER = ''' """Grid barrier over the group's CTAs on the arrival counter; returns the next phase.""" + # Every publishing warp must finish before the leader releases this CTA's + # arrival; otherwise peers can scan partial histograms and diverge in rounds. + cute.arch.sync_threads() + arrival_ptr = _fused_state_ptr(state, group_id, Int32(_FUSED_STATE_ARRIVAL)) +'''.replace("\n", "\r\n").encode() +_OLD_CACHE = b'"attention.indexer.fused_indexer", 1, cache_key, labels=labels' +_NEW_CACHE = b'"attention.indexer.fused_indexer", 2, cache_key, labels=labels' + + +def apply_patch(path: Path, *, check_only: bool = False) -> dict: + source = path.read_bytes() + digest = hashlib.sha256(source).hexdigest() + if digest == AFTER_SHA256: + return {"status": "already_patched", "source_sha256": digest, "changed": False} + if digest != BEFORE_SHA256: + raise RuntimeError(f"unsupported MTP3 B12X source preimage: {digest}") + if source.count(_BEFORE) != 1 or source.count(_OLD_CACHE) != 1: + raise RuntimeError("attested barrier or compile revision anchor differs") + patched = source.replace(_BEFORE, _AFTER, 1).replace(_OLD_CACHE, _NEW_CACHE, 1) + if hashlib.sha256(patched).hexdigest() != AFTER_SHA256: + raise RuntimeError("MTP3 B12X transform differs from its expected postimage") + ast.parse(patched.decode("utf-8"), filename=str(path)) + if not check_only: + path.write_bytes(patched) + return {"status": "checked" if check_only else "patched", "before_sha256": digest, + "after_sha256": AFTER_SHA256, "changed": not check_only, "compile_revision": 2} + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("path", type=Path) + parser.add_argument("--check", action="store_true", help="validate the transform without writing") + args = parser.parse_args() + print(json.dumps(apply_patch(args.path, check_only=args.check), sort_keys=True)) diff --git a/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/patch_mtp3_lease_accounting.py b/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/patch_mtp3_lease_accounting.py new file mode 100644 index 00000000..e1db6902 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/patch_mtp3_lease_accounting.py @@ -0,0 +1,60 @@ +#!/usr/bin/env python3 +"""Account for resident GPU lease reuse in the pinned MTP3 prefill statistics.""" + +from __future__ import annotations + +import argparse +import ast +import hashlib +import json +from pathlib import Path + + +BEFORE_SHA256 = "122d9722b48f8cec267d2136a18c184b8901a5a0a46187eab7030d6227f1983d" +AFTER_SHA256 = "3ec6f357fdc770212528ca3f5dfd48234186f972ec5657a514e8b841df61e48b" +_BEFORE = b''' request.num_computed_tokens = min( + attached_tokens, request.num_tokens - 1 + ) + attached = getattr( +''' +_AFTER = b''' request.num_computed_tokens = min( + attached_tokens, request.num_tokens - 1 + ) + # Lease attachment skips hash lookup but still reuses + # resident GPU state. Include it in prefill/API totals + # without reporting a fresh external KV transfer. + if request.prefill_stats and request.num_preemptions <= 0: + request.prefill_stats.set( + num_prompt_tokens=request.num_prompt_tokens, + num_local_cached_tokens=request.num_computed_tokens, + num_external_cached_tokens=0, + ) + attached = getattr( +''' + + +def apply_patch(path: Path, *, check_only: bool = False) -> dict: + source = path.read_bytes() + digest = hashlib.sha256(source).hexdigest() + if digest == AFTER_SHA256: + return {"status": "already_patched", "source_sha256": digest, "changed": False} + if digest != BEFORE_SHA256: + raise RuntimeError(f"unsupported MTP3 scheduler preimage: {digest}") + if source.count(_BEFORE) != 1: + raise RuntimeError("MTP3 lease attachment anchor differs") + patched = source.replace(_BEFORE, _AFTER, 1) + if hashlib.sha256(patched).hexdigest() != AFTER_SHA256: + raise RuntimeError("MTP3 lease accounting postimage differs") + ast.parse(patched.decode("utf-8"), filename=str(path)) + if not check_only: + path.write_bytes(patched) + return {"status": "checked" if check_only else "patched", "changed": not check_only, + "before_sha256": digest, "after_sha256": AFTER_SHA256} + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("path", type=Path) + parser.add_argument("--check", action="store_true") + args = parser.parse_args() + print(json.dumps(apply_patch(args.path, check_only=args.check), sort_keys=True)) diff --git a/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/patch_mtp3_local_lease_preference.py b/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/patch_mtp3_local_lease_preference.py new file mode 100644 index 00000000..9ab4ab94 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/patch_mtp3_local_lease_preference.py @@ -0,0 +1,77 @@ +#!/usr/bin/env python3 +"""Prefer a strictly longer converged GPU prefix over a shorter shared lease.""" + +from __future__ import annotations + +import argparse +import ast +import hashlib +import json +from pathlib import Path + + +BEFORE_SHA256 = "3ec6f357fdc770212528ca3f5dfd48234186f972ec5657a514e8b841df61e48b" +AFTER_SHA256 = "0df01bf90bbe6ab1e6bc127ca7b15f6286ace944b430c994539d3d81049f2b4a" +TRANSFORMS = ( + (b" did_prefix_cache_lookup = False\n", + b" did_prefix_cache_lookup = False\n local_lease_alternative = None\n"), + (b''' candidate = get_lease(request) if get_lease is not None else None + if candidate is not None: + lease_key, lease_tokens = candidate +''', b''' candidate = get_lease(request) if get_lease is not None else None + if candidate is not None: + lease_key, lease_tokens = candidate + if 0 < lease_tokens <= request.num_tokens: + # This lookup reconciles every KV group and applies + # speculative backoff before any blocks are adopted. + alternative = self.kv_cache_manager.get_computed_blocks(request) + if alternative[1] > min(lease_tokens, request.num_tokens - 1): + local_lease_alternative = (*alternative, False) + candidate = None + # Release this request's follower binding only; + # the verified lease and its other users remain. + rejected = getattr( + self.connector, "shared_prefix_lease_rejected", None + ) + if rejected is not None: + rejected(request_id, lease_key) + if candidate is not None: + lease_key, lease_tokens = candidate +'''), + (b''' ) = self._get_local_prefix_cache_hit(request) +''', b''' ) = ( + local_lease_alternative + if local_lease_alternative is not None + else self._get_local_prefix_cache_hit(request) + ) +'''), +) + + +def apply_patch(path: Path, *, check_only: bool = False) -> dict: + source = path.read_bytes() + digest = hashlib.sha256(source).hexdigest() + if digest == AFTER_SHA256: + return {"status": "already_patched", "source_sha256": digest, "changed": False} + if digest != BEFORE_SHA256: + raise RuntimeError(f"unsupported accounted MTP3 scheduler preimage: {digest}") + patched = source + for before, after in TRANSFORMS: + if patched.count(before) != 1: + raise RuntimeError("MTP3 local/lease selection anchor differs") + patched = patched.replace(before, after, 1) + if hashlib.sha256(patched).hexdigest() != AFTER_SHA256: + raise RuntimeError("MTP3 local/lease selection postimage differs") + ast.parse(patched.decode("utf-8"), filename=str(path)) + if not check_only: + path.write_bytes(patched) + return {"status": "checked" if check_only else "patched", "changed": not check_only, + "before_sha256": digest, "after_sha256": AFTER_SHA256} + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("path", type=Path) + parser.add_argument("--check", action="store_true") + args = parser.parse_args() + print(json.dumps(apply_patch(args.path, check_only=args.check), sort_keys=True)) diff --git a/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/patch_mtp3_partial_tail_eligibility.py b/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/patch_mtp3_partial_tail_eligibility.py new file mode 100644 index 00000000..e5ca8802 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/patch_mtp3_partial_tail_eligibility.py @@ -0,0 +1,63 @@ +#!/usr/bin/env python3 +"""Avoid recurrent partial-tail stops created only by DCP attention geometry.""" + +from __future__ import annotations + +import argparse +import ast +import hashlib +import json +from pathlib import Path + + +BEFORE_SHA256 = "5c0bd785d2d17dce39cdf867a55b5327487070d315851f3c67d9c2cf7a1d4c49" +AFTER_SHA256 = "75efa57e7ff5a77c76714b85e2e4d8e1d7f456d9a9eec6c67ebb11ca382942f9" +BEFORE = b''' # A finer prefix_match_unit is configured: a mamba partial tail entry + # can only be registered by a step ending exactly at the prompt's last + # hash boundary, so the split adds that stop. + self.mamba_partial_cache_hit = ( + self.need_mamba_block_aligned_split + and self.hash_block_size < self.block_size + and self.kv_cache_manager.coordinator.enable_partial_hash_hits + ) +''' +AFTER = b''' # An interior recurrent-page hash needs an explicit tail stop. + # DCP can enlarge attention's scheduling unit without making the + # recurrent page larger than a hash, so inspect Mamba specs directly. + self.mamba_partial_cache_hit = ( + self.need_mamba_block_aligned_split + and any( + isinstance(group.kv_cache_spec, MambaSpec) + and self.hash_block_size < group.kv_cache_spec.block_size + for group in kv_cache_config.kv_cache_groups + ) + and self.kv_cache_manager.coordinator.enable_partial_hash_hits + ) +''' + + +def apply_patch(path: Path, *, check_only=False): + source = path.read_bytes() + digest = hashlib.sha256(source).hexdigest() + if digest == AFTER_SHA256: + return {"status": "already_patched", "source_sha256": digest, "changed": False} + if digest != BEFORE_SHA256: + raise RuntimeError(f"unsupported checkpoint-patched MTP3 scheduler: {digest}") + if source.count(BEFORE) != 1: + raise RuntimeError("MTP3 partial-tail eligibility anchor differs") + patched = source.replace(BEFORE, AFTER, 1) + if hashlib.sha256(patched).hexdigest() != AFTER_SHA256: + raise RuntimeError("MTP3 partial-tail eligibility postimage differs") + ast.parse(patched.decode("utf-8")) + if not check_only: + path.write_bytes(patched) + return {"status": "checked" if check_only else "patched", "changed": not check_only, + "before_sha256": digest, "after_sha256": AFTER_SHA256} + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("path", type=Path) + parser.add_argument("--check", action="store_true") + args = parser.parse_args() + print(json.dumps(apply_patch(args.path, check_only=args.check), sort_keys=True)) diff --git a/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/patch_mtp3_sparse_retention.py b/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/patch_mtp3_sparse_retention.py new file mode 100644 index 00000000..63a2f205 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/patch_mtp3_sparse_retention.py @@ -0,0 +1,106 @@ +#!/usr/bin/env python3 +"""Pair speculative replay checkpoint materialization with sparse retention. + +Attested for the installed scheduler after lease accounting/local preference. +Serving qualification targets native MTP3 with 512-token hash/Mamba pages and +DCP4; no GPU or model execution is performed by this source transformer. +""" + +from __future__ import annotations + +import argparse +import ast +import hashlib +import json +from pathlib import Path + + +SCHEDULER_BEFORE = "0df01bf90bbe6ab1e6bc127ca7b15f6286ace944b430c994539d3d81049f2b4a" +SCHEDULER_AFTER = "5c0bd785d2d17dce39cdf867a55b5327487070d315851f3c67d9c2cf7a1d4c49" +MANAGER_BEFORE = "10846c4994e7860deab8b42c8bcd3315ddc96d14a478d4012c398418cc17a04c" +MANAGER_AFTER = "d2e35b012e0cf45ab3771f545c35ca48f2a5858549c574a352975607369124e2" +SCHEDULER_TRANSFORMS = ( + (b''' if self.use_eagle: + last_cache_position = max(last_cache_position - block_size, 0) +''', b''' if self.use_eagle: + # Lookup excludes the final prompt token, then drops its proof + # block. The corresponding recurrent state needs its own stop. + last_cache_position = max( + (request.num_tokens - 1) // block_size * block_size - block_size, 0 + ) +'''), + (b''' if use_internal_checkpoint: + last_cache_position = 0 +''', b''' if use_internal_checkpoint and not self.use_eagle: + # The internal checkpoint covers the final aligned state only; + # speculative replay may require the preceding state as well. + last_cache_position = 0 +'''), +) +MANAGER_TRANSFORMS = ( + (b''' if start_block <= boundary_block < end_block: + mask[boundary_block - start_block] = True + + return mask + + def remove_skipped_blocks( +''', b''' if start_block <= boundary_block < end_block: + mask[boundary_block - start_block] = True + + if use_eagle: + # Retain the predecessor when materialized, alongside the + # scheduler-aligned fallback. Unmaterialized shared-junction + # slots remain null and cache_full_blocks skips them. + predecessor_block = boundary_tokens // block_size - 2 + if start_block <= predecessor_block < end_block: + mask[predecessor_block - start_block] = True + + return mask + + def remove_skipped_blocks( +'''), +) + + +def transform(source, expected_before, expected_after, transforms): + digest = hashlib.sha256(source).hexdigest() + if digest == expected_after: + return source + if digest != expected_before: + raise RuntimeError(f"unsupported MTP3 checkpoint preimage: {digest}") + for before, after in transforms: + if source.count(before) != 1: + raise RuntimeError("MTP3 checkpoint source anchor differs") + source = source.replace(before, after, 1) + if hashlib.sha256(source).hexdigest() != expected_after: + raise RuntimeError("MTP3 checkpoint postimage differs") + ast.parse(source.decode("utf-8")) + return source + + +def apply_patch(root: Path, *, check_only=False): + files = ( + (root / "v1/core/sched/scheduler.py", SCHEDULER_BEFORE, SCHEDULER_AFTER, SCHEDULER_TRANSFORMS), + (root / "v1/core/single_type_kv_cache_manager.py", MANAGER_BEFORE, MANAGER_AFTER, MANAGER_TRANSFORMS), + ) + prepared = [] + for path, before, after, transforms in files: + original = path.read_bytes() + patched = transform(original, before, after, transforms) + prepared.append((path, original, patched)) + # Validate both preimages before changing either source file. + if not check_only: + for path, original, patched in prepared: + if original != patched: + path.write_bytes(patched) + return {str(path.relative_to(root)): dict(before_sha256=hashlib.sha256(original).hexdigest(), + after_sha256=hashlib.sha256(patched).hexdigest(), changed=original != patched and not check_only) + for path, original, patched in prepared} + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("vllm_root", type=Path) + parser.add_argument("--check", action="store_true") + args = parser.parse_args() + print(json.dumps(apply_patch(args.vllm_root, check_only=args.check), sort_keys=True)) diff --git a/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/repro_mtp3_checkpoint_materialization.py b/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/repro_mtp3_checkpoint_materialization.py new file mode 100644 index 00000000..35993c4c --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/repro_mtp3_checkpoint_materialization.py @@ -0,0 +1,140 @@ +"""Trace MTP3 checkpoint scheduling and retention from exact source fixtures. + +The fixture manifest identifies the image and source inputs. The checker uses +512-token recurrent/hash pages and 2,048-token DCP4 scheduling alignment. +Controlled scheduling budgets expose produced and retained checkpoint +boundaries without GPU execution or a reconstruction of a live request trace. +""" + +import ast +import json +import math +from types import SimpleNamespace + +from repro_mtp3_sparse_retention import ROOT, FullAttentionSpec, MambaSpec, execute, load_algorithms + + +def extracted_method(path, cls_name, method_name): + tree = ast.parse(path.read_text()) + cls = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == cls_name) + return next(node for node in cls.body if isinstance(node, ast.FunctionDef) and node.name == method_name) + + +def load_scheduler(path=None): + path = path or ROOT / "v1/core/sched/scheduler.py" + shell = ast.parse("class Scheduler:\n pass\n").body[0] + shell.body = [extracted_method(path, "Scheduler", name) for name in + ("_mamba_block_aligned_split", "_reserve_prefill_lookahead")] + scope = {} + execute([shell], scope) + return scope["Scheduler"] + + +def partial_hit_eligibility(scheduler, kv_cache_config, path=None): + """Execute the source fixture's recurrent partial-tail eligibility rule.""" + tree = ast.parse((path or ROOT / "v1/core/sched/scheduler.py").read_text()) + nodes = [node for node in ast.walk(tree) if isinstance(node, ast.Assign) + and any(ast.unparse(target) == "self.mamba_partial_cache_hit" for target in node.targets)] + assert len(nodes) == 1 + execute(nodes, dict(self=scheduler, kv_cache_config=kv_cache_config, MambaSpec=MambaSpec)) + return scheduler.mamba_partial_cache_hit + + +def checkpoint_kernel_metadata(): + tree = ast.parse((ROOT / "v1/attention/backends/gdn_attn.py").read_text()) + loops = [node for node in ast.walk(tree) if isinstance(node, ast.For) + and ast.unparse(node.target) == "row" and ast.unparse(node.iter) == "request_rows"] + assert len(loops) == 1 + function = ast.parse("def checkpoint(start, end, block_size):\n pass\n").body[0] + function.body = ast.parse("all_query_lens=[end-start]\nseq_lens=[end]\nrequest_rows=[0]\ncheckpoint_offsets=[]\ncheckpoint_columns=[]").body + function.body.append(loops[0]) + function.body.extend(ast.parse("return start+checkpoint_offsets[0] if checkpoint_offsets[0] else None").body) + scope = {} + execute([function], scope) + return scope["checkpoint"] + + +def main(): + algorithms = load_algorithms(ROOT / "v1/core/kv_cache_utils.py") + utils = ast.parse((ROOT / "v1/core/kv_cache_utils.py").read_text()) + resolve = next(node for node in utils.body if isinstance(node, ast.FunctionDef) + and node.name == "resolve_kv_cache_block_sizes") + scope = dict(math=math, AttentionSpec=FullAttentionSpec, MambaSpec=MambaSpec) + execute([resolve], scope) + recurrent = MambaSpec(512) + recurrent.mamba_cache_mode = "align" + cache = SimpleNamespace(block_size=512, prefix_match_unit=None, enable_prefix_caching=True) + config = SimpleNamespace(cache_config=cache, + parallel_config=SimpleNamespace(decode_context_parallel_size=4), + kv_transfer_config=object()) + groups = SimpleNamespace(kv_cache_groups=[SimpleNamespace(kv_cache_spec=FullAttentionSpec(512)), + SimpleNamespace(kv_cache_spec=recurrent)]) + scheduler_size, hash_size = scope[resolve.name](groups, config) + assert (scheduler_size, hash_size) == (2048, 512) + checkpoint = checkpoint_kernel_metadata() + Scheduler = load_scheduler() + cases = [] + for label, budgets in ( + ("full_8192_token_budgets", [8192] * 8), + ("reduced_budget_materializes_predecessor", [8192, 8192, 8192, 7680, 8192, 8192]), + ): + scheduler = Scheduler() + scheduler.cache_config = cache + scheduler.use_eagle = True + scheduler.mamba_has_prefill_checkpoint_blocks = True + scheduler.mamba_partial_cache_hit = True + scheduler.hash_block_size = hash_size + scheduler.max_num_scheduled_tokens = 8192 + scheduler.scheduler_config = SimpleNamespace(long_prefill_token_threshold=0) + # One native MTP module, repeated three times, uses one prefill lookahead. + scheduler.num_prefill_lookahead = 1 + scheduler._recurrent_publication_boundaries = lambda request: (32768,) + request = SimpleNamespace(num_prompt_tokens=32789, num_tokens=32789, + num_computed_tokens=0, shared_prefix_boundary=0) + steps = [] + materialized = set() + retained = set() + cached_blocks = 0 + for budget in budgets: + start = request.num_computed_tokens + if start == request.num_prompt_tokens: + break + count = scheduler._mamba_block_aligned_split(request, min(budget, request.num_tokens-start)) + count = scheduler._reserve_prefill_lookahead(request, start, count) + assert count > 0 + end = start + count + internal = checkpoint(start, end, recurrent.block_size) + produced = {end} if end % recurrent.block_size == 0 else set() + if internal is not None: + produced.add(internal) + materialized.update(produced) + num_full = end // recurrent.block_size + mask = algorithms.MambaManager.reachable_block_mask( + start_block=cached_blocks, end_block=num_full, + alignment_tokens=scheduler_size, kv_cache_spec=recurrent, + use_eagle=True, retention_interval=0, + reachable_boundaries=(request.num_prompt_tokens-1,), + ) + registered = {(cached_blocks + index + 1) * recurrent.block_size + for index, keep in enumerate(mask) if keep} + retained.update(registered.intersection(materialized)) + steps.append(dict(start=start, end=end, internal_checkpoint=internal, + produced=sorted(produced), newly_retained=sorted(registered))) + cached_blocks = num_full + request.num_computed_tokens = end + assert request.num_computed_tokens == request.num_prompt_tokens + assert retained == {32768} + cases.append(dict(case=label, steps=steps, materialized=sorted(materialized), + retained=sorted(retained), required_predecessor=32256, + predecessor_materialized=32256 in materialized, + predecessor_retained=32256 in retained)) + assert not cases[0]["predecessor_materialized"] + assert cases[1]["predecessor_materialized"] + print(json.dumps(dict(schema="mtp3-installed-checkpoint-trace/v1", gpu_executed=False, + unique_effective_group_sizes=[2048,512], scheduler_block_size=scheduler_size, + hash_block_size=hash_size, lookup_alignment=hash_size, retention_alignment=scheduler_size, + retention_interval=0, num_prefill_checkpoint_blocks=1, cases=cases), sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/repro_mtp3_sparse_retention.py b/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/repro_mtp3_sparse_retention.py new file mode 100644 index 00000000..58a70962 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/repro_mtp3_sparse_retention.py @@ -0,0 +1,168 @@ +"""Check prefix lookup and retention using exact vLLM source fixtures on CPU. + +The block pool contains explicitly described synthetic materialized checkpoints. +The result identifies lookup/retention mismatches under the supplied geometry. +It does not establish which intermediate states a live model materializes. +""" + +from __future__ import annotations + +import argparse +import ast +import copy +import hashlib +import itertools +import json +import os +from collections import namedtuple +from collections.abc import Sequence +from pathlib import Path +from types import SimpleNamespace +from typing import overload + + +ROOT = Path(os.environ.get("MTP3_ORIGINAL_SOURCE", str(Path(__file__).resolve().parent / "mtp3-vllm-source"))) + + +class FullAttentionSpec: + def __init__(self, block_size, dcp_replicated=False): + self.block_size = block_size + self.dcp_replicated = dcp_replicated + + +class MambaSpec: + def __init__(self, block_size): + self.block_size = block_size + + +class ChunkedLocalAttentionSpec: + pass + + +def execute(nodes, scope): + tree = ast.Module(body=[ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0), + *copy.deepcopy(nodes)], type_ignores=[]) + exec(compile(ast.fix_missing_locations(tree), "", "exec"), scope) + + +def load_algorithms(hash_utils, *, manager_source=None): + scope = dict(FullAttentionSpec=FullAttentionSpec, MambaSpec=MambaSpec, + ChunkedLocalAttentionSpec=ChunkedLocalAttentionSpec, Sequence=Sequence, + itertools=itertools, overload=overload, cdiv=lambda a, b: (a + b - 1) // b) + utils_tree = ast.parse(hash_utils.read_text()) + execute([node for node in utils_tree.body if isinstance(node, (ast.FunctionDef, ast.ClassDef)) + and node.name in {"BlockHashListWithBlockSize", "resolve_block_hashes"}], scope) + manager_tree = ast.parse((manager_source or ROOT / "v1/core/single_type_kv_cache_manager.py").read_text()) + for name, methods in (("FullAttentionManager", {"find_longest_cache_hit"}), + ("MambaManager", {"find_longest_cache_hit", "reachable_block_mask"})): + original = next(node for node in manager_tree.body if isinstance(node, ast.ClassDef) and node.name == name) + shell = ast.parse(f"class {name}:\n supports_fine_grained_hash_lookup = True\n").body[0] + shell.body += [node for node in original.body if isinstance(node, ast.FunctionDef) and node.name in methods] + execute([shell], scope) + coordinator = ast.parse((ROOT / "v1/core/kv_cache_coordinator.py").read_text()) + original = next(node for node in coordinator.body if isinstance(node, ast.ClassDef) + and node.name == "HybridKVCacheCoordinator") + shell = ast.parse("class HybridKVCacheCoordinator:\n pass\n").body[0] + shell.body = [node for node in original.body if isinstance(node, ast.FunctionDef) + and node.name in {"find_longest_cache_hit", "_cache_hit_alignment_tokens"}] + execute([shell], scope) + return SimpleNamespace(**scope) + + +class Pool: + def __init__(self, hash_block_size, retained_mamba, shared_boundary): + self.hash_block_size = hash_block_size + self.null_block = SimpleNamespace(is_null=True, block_hash=None) + self.retained_mamba = set(retained_mamba) + self.shared_boundary = shared_boundary + self.queries = [] + + def get_cached_block(self, block_hash, groups): + self.queries.append((block_hash, tuple(groups))) + if block_hash > self.shared_boundary or any( + group == 1 and block_hash not in self.retained_mamba for group in groups + ): + return None + return [SimpleNamespace(is_null=False, block_hash=(block_hash, group)) for group in groups] + + +def run_case(algorithms, args, *, drop=True, extra_checkpoint=None): + Full = algorithms.FullAttentionManager + Mamba = algorithms.MambaManager + target = FullAttentionSpec(args.page) + state = MambaSpec(args.page) + draft = FullAttentionSpec(args.page, dcp_replicated=True) + mask = Mamba.reachable_block_mask( + start_block=0, end_block=args.shared // args.page, + alignment_tokens=args.scheduler_alignment, kv_cache_spec=state, + use_eagle=False, retention_interval=0, reachable_boundaries=(args.prompt - 1,), + ) + retained = [(index + 1) * args.page for index, keep in enumerate(mask) if keep] + if extra_checkpoint is not None: + retained.append(extra_checkpoint) + pool = Pool(args.hash_unit, retained, args.shared) + coordinator = algorithms.HybridKVCacheCoordinator() + coordinator.kv_cache_config = SimpleNamespace(kv_cache_groups=(target, state, draft)) + coordinator.single_type_managers = ( + SimpleNamespace(block_size=args.page * args.dcp), + SimpleNamespace(block_size=args.page), SimpleNamespace(block_size=args.page), + ) + spec_group = namedtuple("SpecGroup", "spec group_ids manager_cls use_eagle") + coordinator.attention_groups = [spec_group(target, [0], Full, False), spec_group(state, [1], Mamba, False), + spec_group(draft, [2], Full, drop)] + coordinator.block_pool = pool + coordinator.hash_block_size = args.hash_unit + coordinator.scheduler_block_size = args.scheduler_alignment + coordinator.enable_partial_hash_hits = args.fine_hits + coordinator.dcp_world_size = args.dcp + hashes = list(range(args.hash_unit, args.continuation, args.hash_unit)) + separate = [] + for spec, groups, manager, use_eagle in coordinator.attention_groups: + _, hit = manager.find_longest_cache_hit( + block_hashes=hashes, max_length=args.continuation - 1, + kv_cache_group_ids=groups, block_pool=pool, kv_cache_spec=spec, + drop_eagle_block=use_eagle, alignment_tokens=coordinator._cache_hit_alignment_tokens, + dcp_world_size=args.dcp if isinstance(spec, FullAttentionSpec) else 1, + ) + separate.append(hit) + _, reconciled, uncached = coordinator.find_longest_cache_hit(hashes, args.continuation - 1) + return dict(drop_speculative_draft=drop, retained_mamba_tokens=retained, + extra_materialized_checkpoint=extra_checkpoint, + per_group_hit_tokens=separate, reconciled_hit_tokens=reconciled, + uncached_shared_prefix_tokens=uncached) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--hash-utils-source", type=Path, default=ROOT / "v1/core/kv_cache_utils.py") + parser.add_argument("--hash-unit", type=int, default=256) + parser.add_argument("--scheduler-alignment", type=int, default=2048) + parser.add_argument("--page", type=int, default=512) + parser.add_argument("--dcp", type=int, default=4) + parser.add_argument("--prompt", type=int, default=32789) + parser.add_argument("--continuation", type=int, default=34856) + parser.add_argument("--shared", type=int, default=32768) + parser.add_argument("--fine-hits", action=argparse.BooleanOptionalAction, default=True) + args = parser.parse_args() + algorithms = load_algorithms(args.hash_utils_source) + baseline = run_case(algorithms, args) + controls = [run_case(algorithms, args, drop=False)] + draft_hit = baseline["per_group_hit_tokens"][-1] + for unit in sorted({args.page, args.scheduler_alignment}): + checkpoint = draft_hit // unit * unit + if checkpoint > 0: + controls.append(run_case(algorithms, args, extra_checkpoint=checkpoint)) + paths = [ROOT / "v1/core/single_type_kv_cache_manager.py", ROOT / "v1/core/kv_cache_coordinator.py", + args.hash_utils_source] + print(json.dumps(dict(schema="mtp3-sparse-retention-metadata-reproducer/v1", + geometry={key: value for key, value in vars(args).items() if key != "hash_utils_source"}, + synthetic_checkpoint_metadata=True, gpu_executed=False, + sources={str(path): hashlib.sha256(path.read_bytes()).hexdigest() for path in paths}, + baseline=baseline, controls=controls), sort_keys=True)) + if baseline["reconciled_hit_tokens"] == 0 and baseline["per_group_hit_tokens"][0] > 0: + assert controls[0]["reconciled_hit_tokens"] > 0 + assert all(case["reconciled_hit_tokens"] > 0 for case in controls[1:]) + + +if __name__ == "__main__": + main() diff --git a/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/test_composition.py b/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/test_composition.py new file mode 100644 index 00000000..1b4a4d20 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/test_composition.py @@ -0,0 +1,70 @@ +"""Composition tests use bundled source fixtures; no Docker, Torch or GPU required.""" +import json +from pathlib import Path +import subprocess +import sys + +import pytest + +import compose + + +def test_composition_entrypoint_is_self_contained(tmp_path): + output = tmp_path / 'composition' + result = subprocess.run([sys.executable, str(Path(compose.__file__)), '--output-root', str(output)], + text=True, capture_output=True, check=False) + assert result.returncode == 0, result.stderr + receipt = json.loads(result.stdout) + assert receipt['candidate_files']['vllm/v1/core/sched/scheduler.py'].startswith('75efa57e') + assert receipt['candidate_files']['vllm/v1/core/single_type_kv_cache_manager.py'].startswith('d2e35b01') + assert compose.verify_candidate(output / 'candidate') == receipt['candidate_files'] + with pytest.raises(FileExistsError): + compose.compose(output) + + +def test_original_drift_rejected_before_output_created(tmp_path): + output = tmp_path / 'prepared' + compose.compose(output) + source = output / 'original' + (source / 'vllm/v1/core/sched/scheduler.py').write_bytes(b'not the attested source') + rejected = tmp_path / 'rejected' + with pytest.raises(ValueError, match='Original source checksum'): + compose.compose(rejected, source_root=source) + assert not rejected.exists() + + +def test_candidate_drift_rejected(tmp_path): + output = tmp_path / 'prepared' + compose.compose(output) + path = output / 'candidate/vllm/v1/core/sched/scheduler.py' + path.write_bytes(path.read_bytes() + b'\n') + with pytest.raises(ValueError, match='Candidate source checksum'): + compose.verify_candidate(output / 'candidate') + + +def test_each_final_patch_remains_idempotent(tmp_path): + output = tmp_path / 'prepared' + compose.compose(output) + candidate = output / 'candidate' + assert not compose.barrier.apply_patch(candidate / 'b12x/attention/dsa_indexer/fused_indexer.py')['changed'] + assert not compose.partial.apply_patch(candidate / 'vllm/v1/core/sched/scheduler.py')['changed'] + # Earlier chain stages intentionally reject later scheduler postimages. + with pytest.raises(RuntimeError, match='preimage'): + compose.accounting.apply_patch(candidate / 'vllm/v1/core/sched/scheduler.py') + + +def test_actual_allocator_checker_accepts_final_source_composition(tmp_path): + output = tmp_path / 'prepared' + compose.compose(output) + receipt = tmp_path / 'allocator.json' + result = subprocess.run([sys.executable, str(Path(compose.__file__).with_name('check_mtp3_checkpoint_allocations.py')), + '--source-root', str(output / 'original/vllm'), + '--candidate-root', str(output / 'candidate/vllm'), '--output', str(receipt)], + text=True, capture_output=True, check=False) + assert result.returncode == 0, result.stderr + data = json.loads(receipt.read_text()) + assert data['passed'] + assert data['source_inputs']['scheduler']['sha256'] == compose.partial.AFTER_SHA256 + for name in ('fresh', 'resumed'): + assert data['populations'][name]['summary']['cases'] == 336 + assert data['populations'][name]['summary']['stale_or_unwritten_registered_states'] == 0 diff --git a/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/test_mtp3_lease_accounting.py b/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/test_mtp3_lease_accounting.py new file mode 100644 index 00000000..3929aae4 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/test_mtp3_lease_accounting.py @@ -0,0 +1,118 @@ +"""Exercise extracted installed vLLM lease/statistics/API code without a GPU.""" + +import ast +import copy +import os +from dataclasses import dataclass +from pathlib import Path +from types import SimpleNamespace + +import pytest + + +ROOT = Path(os.environ["MTP3_ORIGINAL_SOURCE"]) + + +def compile_nodes(nodes, namespace=None): + scope = {} if namespace is None else dict(namespace) + tree = ast.fix_missing_locations(ast.Module(body=copy.deepcopy(nodes), type_ignores=[])) + exec(compile(tree, "", "exec"), scope) + return scope + + +def extract_class(path, name): + tree = ast.parse(path.read_text()) + return next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == name) + + +def wrapper(name, arguments, statements): + tree = ast.parse(f"def {name}({arguments}):\n pass\n") + tree.body[0].body = statements + return tree.body[0] + + +@pytest.fixture +def runtime(): + stats = compile_nodes([extract_class(ROOT / "v1/metrics/stats.py", "PrefillStats")], + {"dataclass": dataclass})["PrefillStats"] + path = Path(os.environ.get("MTP3_SCHEDULER_SOURCE", str(ROOT / "v1/core/sched/scheduler.py"))) + tree = ast.parse(path.read_text()) + candidates = [node for node in ast.walk(tree) if isinstance(node, ast.If) + and ast.unparse(node.test) == "request.num_computed_tokens == 0 and self.connector is not None"] + assert len(candidates) == 1 + setup = ast.parse("request_id = request.request_id").body + apply_lease = compile_nodes([wrapper("apply_lease", "self, request", setup + candidates)])["apply_lease"] + processor = ast.parse((ROOT / "v1/engine/output_processor.py").read_text()) + output = [node for node in ast.walk(processor) if isinstance(node, ast.If) + and ast.unparse(node.test) == "req_state.is_prefilling"] + assert len(output) == 1 + receive = compile_nodes([wrapper("receive", "req_state, engine_core_output", output)])["receive"] + serving = ast.parse((ROOT / "entrypoints/openai/chat_completion/serving.py").read_text()) + helper = next(node for node in serving.body if isinstance(node, ast.FunctionDef) + and node.name == "_make_prompt_tokens_details") + details = compile_nodes([helper], {"PromptTokenUsageInfo": SimpleNamespace})[helper.name] + return SimpleNamespace(Stats=stats, apply=apply_lease, receive=receive, details=details) + + +def attach(runtime, prompt, lease, *, accepted=True, preempted=0, stats=True): + events = [] + request = SimpleNamespace(request_id="request", num_computed_tokens=0, + num_tokens=prompt, num_prompt_tokens=prompt, + num_preemptions=preempted, + prefill_stats=runtime.Stats() if stats else None) + connector = SimpleNamespace(get_shared_prefix_lease_candidate=lambda request: ("lease", lease), + shared_prefix_lease_attached=lambda *args: events.append("attached"), + shared_prefix_lease_rejected=lambda *args: events.append("rejected")) + scheduler = SimpleNamespace(connector=connector, kv_cache_manager=SimpleNamespace( + attach_shared_prefix_lease=lambda *args: lease if accepted else 0, + get_computed_blocks=lambda request: (None, 0, 0))) + runtime.apply(scheduler, request) + return scheduler, request, events + + +@pytest.mark.parametrize("prompt,lease,expected", [(40000, 32768, 32768), (32768, 32768, 32767)]) +def test_attached_gpu_prefix_reaches_api_total_without_external_transfer(runtime, prompt, lease, expected): + scheduler, request, events = attach(runtime, prompt, lease) + stats = request.prefill_stats + assert stats.num_cached_tokens == expected + assert stats.num_local_cached_tokens == expected + assert stats.num_external_cached_tokens == 0 + assert stats.num_computed_tokens == prompt - expected + assert request.num_computed_tokens == expected + # Reentering scheduling cannot double-count the attached prefix. + runtime.apply(scheduler, request) + assert events == ["attached"] + stats.finalize(prompt) + state = SimpleNamespace(is_prefilling=True, num_cached_tokens=0, num_cache_creation_tokens=0) + runtime.receive(state, SimpleNamespace(prefill_stats=stats)) + details = runtime.details(True, state.num_cached_tokens, state.num_cache_creation_tokens, None) + assert details.cached_tokens == expected + assert details.created_cache_tokens == prompt - expected + + +def test_rejected_lease_leaves_normal_lookup_accounting_untouched(runtime): + _, request, events = attach(runtime, 40000, 32768, accepted=False) + assert events == ["rejected"] + assert request.num_computed_tokens == 0 + assert vars(request.prefill_stats) == vars(runtime.Stats()) + + +def test_preemption_does_not_rewrite_first_prefill_stats(runtime): + _, request, events = attach(runtime, 40000, 32768, preempted=1) + assert events == ["attached"] + assert vars(request.prefill_stats) == vars(runtime.Stats()) + + +def test_absent_stats_does_not_change_lease_attachment(runtime): + _, request, events = attach(runtime, 40000, 32768, stats=False) + assert request.prefill_stats is None + assert request.num_computed_tokens == 32768 + assert events == ["attached"] + + +def test_existing_local_plus_external_total_remains_distinct(runtime): + stats = runtime.Stats() + stats.set(num_prompt_tokens=40000, num_local_cached_tokens=8000, num_external_cached_tokens=12000) + details = runtime.details(True, stats.num_cached_tokens, 0, None) + assert (stats.num_local_cached_tokens, stats.num_external_cached_tokens) == (8000, 12000) + assert details.cached_tokens == 20000 diff --git a/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/test_mtp3_local_lease_preference.py b/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/test_mtp3_local_lease_preference.py new file mode 100644 index 00000000..9302f4b2 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/test_mtp3_local_lease_preference.py @@ -0,0 +1,169 @@ +"""GPU-free execution of installed scheduler selection with allocator stubs.""" + +import ast +import os +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from test_mtp3_lease_accounting import ROOT, compile_nodes, wrapper + + +@pytest.fixture +def selection(): + path = Path(os.environ.get("MTP3_SCHEDULER_SOURCE", str(Path(__file__).parent / "mtp3_scheduler_accounted.py"))) + tree = ast.parse(path.read_text()) + lease = next(node for node in ast.walk(tree) if isinstance(node, ast.If) + and ast.unparse(node.test) == "request.num_computed_tokens == 0 and self.connector is not None") + ordinary = min((node for node in ast.walk(tree) if isinstance(node, ast.If) + and ast.unparse(node.test) == "request.num_computed_tokens == 0" + and node.lineno > lease.lineno), key=lambda node: node.lineno) + # Execute the real selection/lease branch and the real local-lookup assignment. + ordinary.body = ordinary.body[:2] + ordinary.orelse = [] + setup = ast.parse("request_id = request.request_id\nlocal_lease_alternative = None\ndid_prefix_cache_lookup = False\nnum_new_local_computed_tokens = 0\nhit_diverged = False").body + result = ast.parse("return did_prefix_cache_lookup, num_new_local_computed_tokens, hit_diverged").body + return compile_nodes([wrapper("select", "self, request", setup + [lease, ordinary] + result)])["select"] + + +def case(selection, *, local, lease=1024, attach=True, stale=False, preempted=False, candidate=True): + calls = [] + members = {"request", "other-follower"} + current = [local] + request = SimpleNamespace(request_id="request", num_computed_tokens=0, num_tokens=2048, + num_prompt_tokens=2048, prefill_stats=None, + num_preemptions=int(preempted), shared_prefix_boundary=0) + + def lookup(request): + calls.append("converged_lookup") + return ("local-blocks", current[0], current[0]) + + def ordinary(request): + calls.append("ordinary_lookup") + return (*lookup(request), False) + + def attach_lease(*args): + calls.append("attach") + if stale: + current[0] = 0 + return lease if attach else 0 + + def reject(*args): + calls.append("reject") + members.discard("request") + + connector = SimpleNamespace(get_shared_prefix_lease_candidate=lambda request: ("lease", lease) if candidate else None, + shared_prefix_lease_attached=lambda *args: calls.append("attached"), + shared_prefix_lease_rejected=reject) + scheduler = SimpleNamespace(connector=connector, + kv_cache_manager=SimpleNamespace(get_computed_blocks=lookup, attach_shared_prefix_lease=attach_lease), + _get_local_prefix_cache_hit=ordinary) + result = selection(scheduler, request) + return result, request, calls, members + + +def test_strictly_longer_converged_local_hit_skips_attachment_and_reuses_lookup(selection): + result, request, calls, members = case(selection, local=1536) + assert result == (True, 1536, False) + assert calls == ["converged_lookup", "reject"] + assert request.num_computed_tokens == 0 # Normal scheduler allocation owns adoption. + assert members == {"other-follower"} + + +@pytest.mark.parametrize("local", [0, 512, 1024]) +def test_equal_or_shorter_local_hit_keeps_lease(selection, local): + result, request, calls, members = case(selection, local=local) + assert not result[0] + assert request.num_computed_tokens == 1024 + assert calls == ["converged_lookup", "attach", "attached"] + assert members == {"request", "other-follower"} + + +def test_rejected_lease_returns_to_normal_lookup(selection): + result, request, calls, members = case(selection, local=512, attach=False) + assert result == (True, 512, False) + assert calls == ["converged_lookup", "attach", "reject", "ordinary_lookup", "converged_lookup"] + assert request.num_computed_tokens == 0 + assert members == {"other-follower"} + + +def test_failed_attachment_does_not_reuse_potentially_stale_probe_blocks(selection): + result, _, calls, _ = case(selection, local=512, attach=False, stale=True) + assert result == (True, 0, False) + assert calls[-2:] == ["ordinary_lookup", "converged_lookup"] + + +def test_preempted_request_still_uses_normal_allocation_for_a_longer_local_hit(selection): + result, request, calls, _ = case(selection, local=1536, preempted=True) + assert result == (True, 1536, False) + assert request.num_computed_tokens == 0 + assert "attach" not in calls + + +def test_no_lease_performs_only_the_ordinary_lookup(selection): + result, _, calls, _ = case(selection, local=1536, candidate=False) + assert result == (True, 1536, False) + assert calls == ["ordinary_lookup", "converged_lookup"] + + +@pytest.mark.skipif(not os.environ.get("SPARKCACHE_SOURCE_ROOT"), reason="optional companion SparkCache source integration") +def test_actual_connector_decline_preserves_verified_lease_and_other_followers(selection, tmp_path): + import sys + sys.path.insert(0, os.environ["SPARKCACHE_SOURCE_ROOT"]) + from sparkcache import test_spark_context_cache_connector as fixtures + + helper = fixtures.AsyncRestoreTests() + connector = helper._cohort_connector(tmp_path) + tokens = list(range(1600)) + digest = helper._offer(connector, tokens) + leader = SimpleNamespace(request_id="leader", prompt_token_ids=tokens) + assert connector.get_num_new_matched_tokens(leader, 0) == (1024, True) + connector.update_state_after_alloc(leader, helper._blocks_stub(), 1024) + connector.build_connector_meta(fixtures._empty_scheduler_output()) + connector.update_connector_output(SimpleNamespace(invalid_block_ids=set(), finished_recving={"leader"})) + assert connector.shared_prefix_lease_published("leader", digest) + other = SimpleNamespace(request_id="other", prompt_token_ids=tokens) + assert connector.get_shared_prefix_lease_candidate(other) == (digest, 1024) + request = SimpleNamespace(request_id="prefer-local", prompt_token_ids=tokens, + num_computed_tokens=0, num_tokens=len(tokens), num_prompt_tokens=len(tokens), + num_preemptions=0, prefill_stats=None, shared_prefix_boundary=0) + scheduler = SimpleNamespace(connector=connector, kv_cache_manager=SimpleNamespace( + get_computed_blocks=lambda request: ("local-blocks", 1280, 1280), + attach_shared_prefix_lease=lambda *args: pytest.fail("short lease was attached")), + _get_local_prefix_cache_hit=lambda request: pytest.fail("lookup repeated")) + try: + assert selection(scheduler, request) == (True, 1280, False) + assert "prefer-local" not in connector._restore_flight_followers + assert "other" in connector._restore_flight_followers + assert connector._restore_flights[digest].lease_published + assert connector.get_shared_prefix_lease_candidate(other) == (digest, 1024) + finally: + connector.shutdown() + + +def test_speculative_backoff_is_applied_before_comparing_with_lease(selection): + # Execute the installed full-attention finder with equal hash/page units. + # In this geometry resolve_block_hashes is its identity branch. + import itertools + from collections.abc import Sequence + from types import SimpleNamespace + + tree = ast.parse((ROOT / "v1/core/single_type_kv_cache_manager.py").read_text()) + original = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "FullAttentionManager") + method = next(node for node in original.body if isinstance(node, ast.FunctionDef) and node.name == "find_longest_cache_hit") + shell = ast.parse("class FullAttentionManager:\n supports_fine_grained_hash_lookup = True\n").body[0] + shell.body.append(method) + future = ast.parse("from __future__ import annotations").body[0] + scope = compile_nodes([future, shell], dict(FullAttentionSpec=SimpleNamespace, + ChunkedLocalAttentionSpec=type("ChunkedLocalAttentionSpec", (), {}), + resolve_block_hashes=lambda values, *args, **kwargs: values, + itertools=itertools, Sequence=Sequence, cdiv=lambda a, b: (a + b - 1) // b)) + pool = SimpleNamespace(hash_block_size=256, get_cached_block=lambda block_hash, groups: [block_hash]) + _, backed_off = scope["FullAttentionManager"].find_longest_cache_hit( + [1, 2, 3, 4, 5], max_length=1280, kv_cache_group_ids=[0], block_pool=pool, + kv_cache_spec=SimpleNamespace(block_size=256), drop_eagle_block=True, alignment_tokens=256) + assert backed_off == 1024 # Raw 1280-token evidence is not a reusable 1280-token hit. + _, request, calls, _ = case(selection, local=backed_off, lease=1024) + assert request.num_computed_tokens == 1024 + assert "attached" in calls diff --git a/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/test_mtp3_partial_tail_eligibility.py b/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/test_mtp3_partial_tail_eligibility.py new file mode 100644 index 00000000..835ddb13 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/test_mtp3_partial_tail_eligibility.py @@ -0,0 +1,42 @@ +"""Execute the actual recurrent partial-tail predicate and scheduling path.""" + +from types import SimpleNamespace + +import pytest + +from repro_mtp3_checkpoint_materialization import partial_hit_eligibility +from repro_mtp3_sparse_retention import MambaSpec +from test_mtp3_sparse_retention import SOURCE, trace + + +@pytest.mark.parametrize("hash_size,page_sizes,alignment,fine,expected", [ + (512, [512], 2048, True, False), + (256, [512], 2048, True, True), + (512, [1024], 2048, True, True), + (512, [], 2048, True, False), + (256, [512], 2048, False, False), +]) +def test_predicate_uses_actual_mamba_pages(hash_size, page_sizes, alignment, fine, expected): + scheduler = SimpleNamespace(need_mamba_block_aligned_split=True, + hash_block_size=hash_size, block_size=alignment, + kv_cache_manager=SimpleNamespace(coordinator=SimpleNamespace(enable_partial_hash_hits=fine))) + config = SimpleNamespace(kv_cache_groups=[SimpleNamespace(kv_cache_spec=MambaSpec(size)) for size in page_sizes]) + assert partial_hit_eligibility(scheduler, config, SOURCE / "v1/core/sched/scheduler.py") is expected + + +@pytest.mark.parametrize("prompt,expected", [ + (5306, [4096, 4608, 5306]), + (32789, [8192, 16384, 24576, 32256, 32768, 32789]), +]) +def test_unneeded_stop_removed_but_publication_and_predecessor_remain(prompt, expected): + _, _, retained, steps = trace(prompt) + assert steps == expected + assert (prompt - 1) // 2048 * 2048 in retained + assert ((prompt - 1) // 512 - 1) * 512 in retained + + +def test_successful_100k_case_drops_only_the_unneeded_fine_tail_stop(): + _, _, retained, steps = trace(100968) + assert 100352 in steps and 100864 not in steps + assert steps[-1] == 100968 + assert 100352 in retained diff --git a/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/test_mtp3_sparse_retention.py b/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/test_mtp3_sparse_retention.py new file mode 100644 index 00000000..7a162478 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/experiments/cache-reuse/test_mtp3_sparse_retention.py @@ -0,0 +1,128 @@ +"""Installed-algorithm tests for MTP3 checkpoint production and reuse.""" + +import os +from collections import namedtuple +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from repro_mtp3_checkpoint_materialization import checkpoint_kernel_metadata, load_scheduler, partial_hit_eligibility +from repro_mtp3_sparse_retention import ROOT, FullAttentionSpec, MambaSpec, Pool, load_algorithms + + +SOURCE = Path(os.environ.get("MTP3_RETENTION_SOURCE", str(Path(__file__).parent / "mtp3-retention-patched"))) + + +def trace(prompt, *, budget=8192, eagle=True, scheduler_source=None, manager_source=None): + algorithms = load_algorithms(ROOT / "v1/core/kv_cache_utils.py", + manager_source=manager_source or SOURCE / "v1/core/single_type_kv_cache_manager.py") + Scheduler = load_scheduler(scheduler_source or SOURCE / "v1/core/sched/scheduler.py") + scheduler = Scheduler() + scheduler.cache_config = SimpleNamespace(block_size=512) + scheduler.use_eagle = eagle + scheduler.mamba_has_prefill_checkpoint_blocks = True + scheduler.hash_block_size = 512 + scheduler.block_size = 2048 + scheduler.need_mamba_block_aligned_split = True + scheduler.kv_cache_manager = SimpleNamespace(coordinator=SimpleNamespace(enable_partial_hash_hits=True)) + partial_hit_eligibility(scheduler, + SimpleNamespace(kv_cache_groups=[SimpleNamespace(kv_cache_spec=MambaSpec(512))]), + scheduler_source or SOURCE / "v1/core/sched/scheduler.py") + scheduler.max_num_scheduled_tokens = 8192 + scheduler.scheduler_config = SimpleNamespace(long_prefill_token_threshold=0) + scheduler.num_prefill_lookahead = 1 + publication = (prompt - 1) // 2048 * 2048 + scheduler._recurrent_publication_boundaries = lambda request: (publication,) if publication else () + request = SimpleNamespace(num_prompt_tokens=prompt, num_tokens=prompt, + num_computed_tokens=0, shared_prefix_boundary=0) + spec = MambaSpec(512) + produced = set() + retained = set() + steps = [] + cached = 0 + checkpoint = checkpoint_kernel_metadata() + while request.num_computed_tokens < prompt: + start = request.num_computed_tokens + n = scheduler._mamba_block_aligned_split(request, min(budget, prompt - start)) + n = scheduler._reserve_prefill_lookahead(request, start, n) + assert n > 0 + end = start + n + point = checkpoint(start, end, 512) + if end % 512 == 0: + produced.add(end) + if point is not None: + produced.add(point) + num_full = end // 512 + mask = algorithms.MambaManager.reachable_block_mask(start_block=cached, + end_block=num_full, alignment_tokens=2048, kv_cache_spec=spec, + use_eagle=eagle, retention_interval=0, reachable_boundaries=(prompt - 1,)) + retained.update((cached + index + 1) * 512 for index, keep in enumerate(mask) + if keep and (cached + index + 1) * 512 in produced) + cached = num_full + request.num_computed_tokens = end + steps.append(end) + assert len(steps) < 100 + return algorithms, produced, retained, steps + + +def replay(algorithms, prompt, retained, *, eagle=True): + target, recurrent = FullAttentionSpec(512), MambaSpec(512) + pool = Pool(512, retained, prompt // 512 * 512) + coordinator = algorithms.HybridKVCacheCoordinator() + coordinator.kv_cache_config = SimpleNamespace(kv_cache_groups=(target, recurrent)) + coordinator.single_type_managers = (SimpleNamespace(block_size=2048), SimpleNamespace(block_size=512)) + group = namedtuple("SpecGroup", "spec group_ids manager_cls use_eagle") + coordinator.attention_groups = [group(target, [0], algorithms.FullAttentionManager, eagle), + group(recurrent, [1], algorithms.MambaManager, eagle)] + coordinator.block_pool = pool + coordinator.hash_block_size = 512 + coordinator.scheduler_block_size = 2048 + coordinator.enable_partial_hash_hits = True + coordinator.dcp_world_size = 4 + hashes = list(range(512, prompt + 1, 512)) + return coordinator.find_longest_cache_hit(hashes, prompt - 1)[1] + + +@pytest.mark.parametrize("prompt", [32768, 32769, 32789, 33280, 33281, 100968]) +@pytest.mark.parametrize("budget", [8192, 7680]) +def test_speculative_predecessor_is_materialized_retained_and_reusable(prompt, budget): + algorithms, produced, retained, _ = trace(prompt, budget=budget) + required = ((prompt - 1) // 512 - 1) * 512 + assert required in produced + assert required in retained + assert replay(algorithms, prompt, retained) == required + # Preserve the scheduler-aligned publication checkpoint as well. + assert (prompt - 1) // 2048 * 2048 in retained + + +def test_known_successful_100k_prompt_keeps_its_existing_reuse(): + algorithms, _, retained, _ = trace(100968) + assert replay(algorithms, 100968, retained) == 100352 + + +def test_non_speculative_internal_checkpoint_behavior_is_unchanged(): + _, produced, retained, steps = trace(32789, eagle=False) + assert steps == [8192, 16384, 24576, 32768, 32789] + assert 32256 not in produced + assert retained == {32768} + + +def test_dense_retention_remains_dense(): + algorithms = load_algorithms(ROOT / "v1/core/kv_cache_utils.py", + manager_source=SOURCE / "v1/core/single_type_kv_cache_manager.py") + assert algorithms.MambaManager.reachable_block_mask(start_block=0, end_block=64, + alignment_tokens=2048, kv_cache_spec=MambaSpec(512), use_eagle=True, + retention_interval=None, reachable_boundaries=(32788,)) is None + + +def test_retention_change_alone_cannot_create_the_missing_checkpoint(): + _, produced, retained, _ = trace(32789, scheduler_source=ROOT / "v1/core/sched/scheduler.py") + assert 32256 not in produced + assert 32256 not in retained + + +def test_materialization_change_alone_does_not_preserve_the_checkpoint(): + _, produced, retained, _ = trace(32789, manager_source=ROOT / "v1/core/single_type_kv_cache_manager.py") + assert 32256 in produced + assert 32256 not in retained diff --git a/runtime/glm53-spark-mtp3-mesh/image-receipt.json b/runtime/glm53-spark-mtp3-mesh/image-receipt.json index a54d72c9..be2b9b5d 100644 --- a/runtime/glm53-spark-mtp3-mesh/image-receipt.json +++ b/runtime/glm53-spark-mtp3-mesh/image-receipt.json @@ -1,16 +1,36 @@ { - "added_layers": 6, - "bundle_manifest_sha256": "4204fabc93303226b9a120b094ef3c82ed4aadd1d7f97cfbe291204c027ed45f", + "added_layers": 8, + "bundle_manifest_sha256": "69313e19e881ec93e9ed3bd150d2f24fc6b444488ac729a69f45d038e2243500", "checks_passed": true, - "image": "sparkring-glm53-spark-mtp3-mesh:managed-8684a696", - "image_id": "sha256:26273b8e358df139ae913610a5d43084ff0fd08aafe282ef633a3bc74afefe47", - "image_reference": "sha256:26273b8e358df139ae913610a5d43084ff0fd08aafe282ef633a3bc74afefe47", - "image_size_bytes": 21076522030, + "image": "ghcr.io/fujitsupolycom/sparkring-glm53-sparkcache@sha256:67dc0ae453baaae6831ccec1d259b4ef8b236a8b0dc9f747d901b95c66ec1987", + "image_id": "sha256:2e41b1e934a85ff7c21b780532db2f0a0e978df081e52f4ae2bf11f8992fb24f", + "image_reference": "sha256:2e41b1e934a85ff7c21b780532db2f0a0e978df081e52f4ae2bf11f8992fb24f", + "image_size_bytes": 22108005586, "inside_image": { - "b12x_commit": "9ae41c5cb9935d740456479954b0089f80bd2ef2", + "b12x_commit": "ef308bac0f3b3eb8fea63e4013afc0c2ea1c6301", "bundle_files": 28, - "bundle_manifest_sha256": "4204fabc93303226b9a120b094ef3c82ed4aadd1d7f97cfbe291204c027ed45f", + "bundle_manifest_sha256": "69313e19e881ec93e9ed3bd150d2f24fc6b444488ac729a69f45d038e2243500", "checks_passed": true, + "compute": { + "b12x_files": 385, + "b12x_revision": "ef308bac0f3b3eb8fea63e4013afc0c2ea1c6301", + "b12x_tree": "dcf039e5e754136275835ea997e6b9abbb6b15ae", + "cuda_version": "13.3", + "environment": { + "CUDA_HOME": "/opt/cuda-13.3", + "TRITON_PTXAS_PATH": "/opt/cuda-13.3/bin/ptxas", + "VLLM_B12X_DENSE_ACTIVATION_MODE": "auto", + "VLLM_GDN_SPEC_DECODE_METADATA_FASTPATH": "1", + "VLLM_LM_HEAD_A16": "1", + "VLLM_MTP_NVFP4_LM_HEAD": "1", + "VLLM_MXFP8_LM_HEAD": "0" + }, + "proposal_head_nvfp4": true, + "source_lock_sha256": "139f36701e0e47f45bf99fba2cc2fa59b417f2ee801dad3a064455d5b464a459", + "target_head_quantization": false, + "vllm_overrides": 24, + "vllm_parent_files": 2905 + }, "cuda_initialized": false, "device_access": false, "limitation": "Content and CPU checks do not qualify CUDA graphs, RDMA forwarding, native MTP, cache restoration, or model performance.", @@ -18,7 +38,7 @@ "marker_source_sha256": "8684a6961b8e86aa474fa2310ff71e4cdf219a63a72ceb5593b2f95e54812792", "model_loaded": false, "parent_package_files": { - "b12x": 374, + "b12x": 385, "sparkcache": 150, "vllm": 2905 }, @@ -26,22 +46,22 @@ "readiness_warmup": { "environment": "SPARKRING_WARMUP_TEMPERATURE", "helper_sha256": "f41c38eef41d15d63dcfc49cd6643357ca1a3ae18200ddbe4f8692d0b767ee79", - "temperature": 1 + "temperature": 1.0 }, "rocenante_lazy_import": "/opt/spark-sircl/b12x_overlay/b12x/comm/roce/__init__.py", "sircl_native_sha256": "61aa0ec56a1b438439bed8611dab0353d2c72c10af02bbd917fb77c87b33e5fc", - "source_receipt_sha256": "9578a15a76820c9253bc55df0a8d4c34d884cd540aeba5726796a0124062b361", + "source_receipt_sha256": "45b935dd4075c7c4e208c5d3568771860bfad5dcbb07a954ac664d03e19f6eeb", "sparkcache_commit": "66057174301a4759ca3a45207ea41016689449cb", "status": "research-only", "vllm_commit": "e02b174693e13859de61811b5e8cd13d5308e259", "vllm_native_extensions": 15 }, - "limitation": "This receipt covers device-free image-content verification. Four-rank serving and managed-fabric checks are recorded separately in performance/records/glm53-flash/spark-mtp3-managed-mesh-functional-20260905.md.", + "limitation": "This receipt records build/content checks. The linked functional record identifies exact-image GPU, serving, restart and cache checks.", "parent_image_id": "sha256:5e32aaa1bbe3559e81db7706ed4286248f18d27cfdb186f6b851bf786eb43075", "parent_layers_retained": 81, "platform": "linux/arm64", "schema": "sparkring-mtp3-mesh-image-receipt/v1", - "source_receipt_sha256": "9578a15a76820c9253bc55df0a8d4c34d884cd540aeba5726796a0124062b361", + "source_receipt_sha256": "45b935dd4075c7c4e208c5d3568771860bfad5dcbb07a954ac664d03e19f6eeb", "status": "research-only", "verification_command": [ "docker", @@ -66,9 +86,10 @@ "PYTHONDONTWRITEBYTECODE=1", "--entrypoint", "python3", - "sha256:26273b8e358df139ae913610a5d43084ff0fd08aafe282ef633a3bc74afefe47", + "sha256:2e41b1e934a85ff7c21b780532db2f0a0e978df081e52f4ae2bf11f8992fb24f", "-I", "/opt/sparkring/bin/verify-mtp3-mesh-image.py", "--inside-image" - ] + ], + "functional_record": "../../performance/records/glm53-flash/spark-mtp3-compute-stream-safety-20260906.md" } diff --git a/runtime/glm53-spark-mtp3-mesh/indexer-barrier/Dockerfile b/runtime/glm53-spark-mtp3-mesh/indexer-barrier/Dockerfile new file mode 100644 index 00000000..fc58fea0 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/indexer-barrier/Dockerfile @@ -0,0 +1,13 @@ +FROM ghcr.io/fujitsupolycom/sparkring-glm53-sparkcache@sha256:67dc0ae453baaae6831ccec1d259b4ef8b236a8b0dc9f747d901b95c66ec1987 +ARG SOURCE_RECEIPT_SHA256 +ARG COMPUTE_SOURCE_LOCK_SHA256 +ARG SPARKRING_REVISION +COPY patch_indexer_barrier.py install.py /opt/sparkring/mesh-indexer-barrier/ +RUN python3 /opt/sparkring/mesh-indexer-barrier/install.py --expected-receipt ${SOURCE_RECEIPT_SHA256} +LABEL org.opencontainers.image.revision="${SPARKRING_REVISION}" \ + org.opencontainers.image.base.name="ghcr.io/fujitsupolycom/sparkring-glm53-sparkcache@sha256:67dc0ae453baaae6831ccec1d259b4ef8b236a8b0dc9f747d901b95c66ec1987" \ + org.sparkring.mesh.parent-image="sha256:2e41b1e934a85ff7c21b780532db2f0a0e978df081e52f4ae2bf11f8992fb24f" \ + org.sparkring.mesh.source-receipt-sha256="${SOURCE_RECEIPT_SHA256}" \ + org.sparkring.compute.source-lock-sha256="${COMPUTE_SOURCE_LOCK_SHA256}" \ + org.sparkring.vllm.compute-source="${COMPUTE_SOURCE_LOCK_SHA256}" \ + org.sparkring.mesh.indexer-barrier="35a7564ec8bf18f1b7dc5562b21616437b65e7b0689c56341b6436c8d202637b" diff --git a/runtime/glm53-spark-mtp3-mesh/indexer-barrier/README.md b/runtime/glm53-spark-mtp3-mesh/indexer-barrier/README.md new file mode 100644 index 00000000..7a48d9fd --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/indexer-barrier/README.md @@ -0,0 +1,106 @@ +# Mesh indexer publication-barrier image + +Status: research-only serving artifact with qualified bounded indexer GPU tests. +The child image preserves the published native-MTP3 mesh parent's compute, +CUDA 13.3, transport bundle, marker, and native libraries. Its only executable +runtime change adds block-wide synchronization before histogram publication +and increases the fused-indexer compile revision from 1 to 2. + +## Artifact identities + +Tag: `ghcr.io/fujitsupolycom/sparkring-glm53-sparkcache:20260906-mtp3-mesh-indexer-barrier` + +Immutable reference: +`ghcr.io/fujitsupolycom/sparkring-glm53-sparkcache@sha256:7aa57ed1901b83c8581c53cb233db930bb1a133d74be01168d967b9dd56b3e49` + +Config image ID: +`sha256:995a42b4b08f525813c5783f78f34571b6bd1c197fccee53617ff8060574a16d` + +Parent reference: +`ghcr.io/fujitsupolycom/sparkring-glm53-sparkcache@sha256:67dc0ae453baaae6831ccec1d259b4ef8b236a8b0dc9f747d901b95c66ec1987` + +The image retains all 89 parent layers and adds two layers totaling 637,767 +uncompressed bytes. Its B12X selector retains the mesh parent's PR316 changes; +it is not replaced with the SIRCL operator image's selector. + +The compute source lock is +`57b6438ca014a493380fb874a84977d35b0983e9ee224a62ade3415fd9c7ce46`. +The image source receipt is +`030e9f3749672d186b66e41336980af1abb5c207f734feafa677ba82456c5d43`. +The unchanged transport bundle is +`69313e19e881ec93e9ed3bd150d2f24fc6b444488ac729a69f45d038e2243500`. + +## Conditions + +Validation used one NVIDIA GB10 SM121 on 2026-09-06, the image above, and +PyTorch 2.13.0+cu130. Tests ran in isolated containers without model weights, +serving traffic, or replacement of the running telemetry stack. + +## Measurement + +The delayed-publisher probe executes the mesh parent's barrier helper in a +three-block kernel, then compares it with an entry-synchronized variant. +Ten eager and ten graph-replay runs per variant check complete publication. +The probe executes only one barrier and does not deliberately hang the GPU. + +The built child image runs fifteen fused-indexer correctness cases and 2,000 +graph replays, checking exact selected-index sets, top-k scores within 0.01 +absolute tolerance, and cleared merge state. Workloads use 32 heads, top-k 512, +3/4/8/16 query rows, lengths alternating between 4,097 and 65,536 tokens +(200,000 for the four-row case), and concurrent 64 MiB device copies. +The harnesses and their dependency requirements are documented in SparkRing's +`performance/harnesses/indexer_barrier/` directory on main. + +## Result + +- Parent helper: incomplete publication in 20/20 probes. +- Entry synchronization: complete publication in 20/20 probes. +- Built mesh image: fifteen GPU correctness tests and 2,000 graph replays passed. +- Local mesh suite: 348 tests passed; three documented platform skips. +- Unmodified mesh content verifier passed before and after installation. + It verifies 385 B12X files, the 24 vLLM overrides, CUDA components, and the + transport/marker/native-library contracts. +- External verification passed the parent-layer, image-label, source-receipt, + and embedded-content checks. + +## Conclusion + +The barrier correction prevents the measured publication race in the mesh +parent's selector. The built image preserves the tested numerical behavior +and completes the bounded graph stress cases. These results support a model +qualification run, not an assertion that every serving stall is eliminated. + +## Limitations and deployment boundary + +No four-rank model startup, collective execution, actual SparkCache restore, +or serving soak was performed with this child image. The running telemetry +image was not patched or replaced. Its additional SparkCache changes are not +included in this artifact. + +The managed profile renderer still pins the parent image and compute lock. +Do not bypass its receipt checks or pass this child's receipt as though it +were the parent's. Managed deployment requires a profile update binding the +child image, source lock, and coordinated startup inputs, followed by an +authorized stop/start window. The artifact is available for that validation; +this change does not silently promote it into the default managed profile. + +The embedded profile records cache namespace +`glm53-spark-df116c4f-mtp3-nvfp4-a16-c57b6438c-mesh69313e19-tail-cow-v2`. +Use a distinct JIT namespace for qualification. Do not rename existing cache +entries into the source-bound namespace. + +## Rebuild and verify + +Place `Dockerfile`, `install.py`, and `patch_indexer_barrier.py` in an empty +context. Run `install.py --output /probe/build-input.json` in an isolated +container of the pinned parent with the context mounted at `/probe`. Build +with that result's `source_receipt_sha256` and `source_lock_sha256` supplied +as `SOURCE_RECEIPT_SHA256` and `COMPUTE_SOURCE_LOCK_SHA256`; set +`SPARKRING_REVISION` to the source commit. The published build uses +`8e54e22477907935c92c52a1c3d7a68e364bd776`. + +The installer preserves the parent source receipt, records the checked source +transform, and updates compute package maps and the embedded profile. The +existing mesh verifier runs before and after installation without modified +verification logic. External verification uses the exact image ID, parent +config ID, source-receipt hash, and bundle hash recorded above. diff --git a/runtime/glm53-spark-mtp3-mesh/indexer-barrier/install.py b/runtime/glm53-spark-mtp3-mesh/indexer-barrier/install.py new file mode 100644 index 00000000..4032e7e9 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/indexer-barrier/install.py @@ -0,0 +1,110 @@ +"""Install a checked mesh indexer patch while retaining compute/transport provenance.""" +import argparse +import hashlib +import json +from pathlib import Path +import subprocess +import sys + +from patch_indexer_barrier import AFTER, BEFORE, RELATIVE, apply + +PARENT = 'ghcr.io/fujitsupolycom/sparkring-glm53-sparkcache@sha256:67dc0ae453baaae6831ccec1d259b4ef8b236a8b0dc9f747d901b95c66ec1987' +PARENT_ID = 'sha256:2e41b1e934a85ff7c21b780532db2f0a0e978df081e52f4ae2bf11f8992fb24f' +PARENT_LOCK = '139f36701e0e47f45bf99fba2cc2fa59b417f2ee801dad3a064455d5b464a459' +RECEIPTS = Path('/opt/sparkring/receipts/glm53-spark-mtp3-mesh') +COMPUTE = Path('/opt/sparkring-compute') +INSTALLED = Path('/opt/sparkring/receipts/glm53-compute-installed.json') +SITE = Path('/usr/local/lib/python3.12/dist-packages') +VERIFY = '/opt/sparkring/bin/verify-mtp3-mesh-image.py' + + +def sha(path): + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def load(path): + return json.loads(path.read_text(encoding='utf-8')) + + +def write(path, value): + path.write_text(json.dumps(value, indent=2, sort_keys=True) + '\n', encoding='utf-8') + + +def install(): + subprocess.run([sys.executable, '-I', VERIFY, '--inside-image'], check=True) + lock_path = COMPUTE / 'source-lock.json' + if sha(lock_path) != PARENT_LOCK: + raise ValueError('Mesh compute source lock differs from the supported parent') + source_path = RECEIPTS / 'source-receipt.json' + source = load(source_path) + parent_source_sha = sha(source_path) + installed = load(INSTALLED) + if installed['b12x_files'][RELATIVE] != BEFORE: + raise ValueError('Installed mesh indexer receipt differs from its preimage') + for target in (SITE / RELATIVE, COMPUTE / 'b12x-source' / RELATIVE): + apply(target) + for cached in (SITE / RELATIVE).parent.glob('__pycache__/fused_indexer.*.pyc'): + cached.unlink() + installed['b12x_files'][RELATIVE] = AFTER + package_hash = hashlib.sha256(json.dumps( + installed['b12x_files'], sort_keys=True, separators=(',', ':')).encode()).hexdigest() + lock = load(lock_path) + lock['b12x']['package_files_sha256'] = package_hash + transform = {'path': RELATIVE, 'preimage_sha256': BEFORE, 'result_sha256': AFTER, + 'script_sha256': sha(Path(__file__).with_name('patch_indexer_barrier.py'))} + lock['b12x']['source_transforms'] = {'histogram_publication_barrier': transform} + write(lock_path, lock) + lock_hash = sha(lock_path) + installed['source_lock_sha256'] = lock_hash + installed['b12x_package_files_sha256'] = package_hash + write(INSTALLED, installed) + prepared_path = COMPUTE / 'prepared-manifest.json' + prepared = load(prepared_path) + prepared['source_lock_sha256'] = lock_hash + prepared['b12x_files'] = installed['b12x_files'] + prepared['b12x_package_files_sha256'] = package_hash + write(prepared_path, prepared) + profile_path = RECEIPTS / 'profile-pins.json' + profile = load(profile_path) + profile['compute']['source_lock_sha256'] = lock_hash + profile['cache_identity']['namespace'] = profile['cache_identity']['namespace'].replace( + PARENT_LOCK[:8], lock_hash[:8]) + profile['cache_identity']['compatibility'] = ( + f"Cache entries require compute source lock {lock_hash} and transport bundle " + f"{profile['canonical_bundle_manifest_sha256']}. Do not relabel entries from another composition.") + write(profile_path, profile) + source_path.with_name('source-receipt.parent.json').write_bytes(source_path.read_bytes()) + source['files']['compute/b12x-source/' + RELATIVE] = AFTER + for relative, path in ( + ('compute/source-lock.json', lock_path), + ('compute/prepared-manifest.json', prepared_path), + ('receipts/profile-pins.json', profile_path), + ): + source['files'][relative] = sha(path) + source['derived_from_mesh_image'] = { + 'reference': PARENT, 'image_id': PARENT_ID, 'source_receipt_sha256': parent_source_sha} + source['source_transforms'] = {'histogram_publication_barrier': transform} + write(source_path, source) + result = {'schema': 'sparkring-mesh-indexer-barrier/v1', 'status': 'research-only', + 'parent_image': PARENT, 'parent_image_id': PARENT_ID, + 'parent_source_lock_sha256': PARENT_LOCK, + 'source_lock_sha256': lock_hash, 'source_receipt_sha256': sha(source_path), + 'bundle_manifest_sha256': profile['canonical_bundle_manifest_sha256'], + 'cache_namespace': profile['cache_identity']['namespace'], + 'indexer': transform, 'installer_sha256': sha(Path(__file__))} + write(RECEIPTS / 'indexer-barrier.json', result) + subprocess.run([sys.executable, '-I', VERIFY, '--inside-image'], check=True) + return result + + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--output', type=Path) + parser.add_argument('--expected-receipt') + args = parser.parse_args() + result = install() + if args.expected_receipt and result['source_receipt_sha256'] != args.expected_receipt: + raise ValueError('Mesh patch receipt differs from the prepared build input') + if args.output: + write(args.output, result) + print(json.dumps(result), flush=True) diff --git a/runtime/glm53-spark-mtp3-mesh/indexer-barrier/patch_indexer_barrier.py b/runtime/glm53-spark-mtp3-mesh/indexer-barrier/patch_indexer_barrier.py new file mode 100644 index 00000000..04d795a8 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/indexer-barrier/patch_indexer_barrier.py @@ -0,0 +1,31 @@ +"""Apply histogram publication ordering to the pinned mesh indexer source.""" +import hashlib +from pathlib import Path + +RELATIVE = 'b12x/attention/dsa_indexer/fused_indexer.py' +BEFORE = '893fbcade135b7e1d146b8fb6530cde0650be515f69bf9a17ced0a9c61a141e2' +AFTER = '35a7564ec8bf18f1b7dc5562b21616437b65e7b0689c56341b6436c8d202637b' + + +def patch_bytes(source: bytes) -> bytes: + digest = hashlib.sha256(source).hexdigest() + if digest == AFTER: + return source + if digest != BEFORE: + raise ValueError(f'Unsupported mesh indexer source SHA-256: {digest}') + arrival = b' arrival_ptr = _fused_state_ptr(state, group_id, Int32(_FUSED_STATE_ARRIVAL))' + barrier = (b' # Publish only after every warp has finished its histogram writes.\r\n' + b' cute.arch.sync_threads()\r\n' + arrival) + cache = b'"attention.indexer.fused_indexer", 1, cache_key, labels=labels' + if source.count(arrival) != 1 or source.count(cache) != 1: + raise ValueError('Mesh indexer barrier or compile-revision anchor differs') + result = source.replace(arrival, barrier, 1).replace( + cache, b'"attention.indexer.fused_indexer", 2, cache_key, labels=labels', 1) + if hashlib.sha256(result).hexdigest() != AFTER: + raise ValueError('Mesh indexer transform produced unexpected source bytes') + return result + + +def apply(path: Path) -> None: + result = patch_bytes(path.read_bytes()) + path.write_bytes(result) diff --git a/runtime/glm53-spark-mtp3-mesh/indexer-barrier/test_patch.py b/runtime/glm53-spark-mtp3-mesh/indexer-barrier/test_patch.py new file mode 100644 index 00000000..11e7cac9 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/indexer-barrier/test_patch.py @@ -0,0 +1,30 @@ +import hashlib +import importlib.util +from pathlib import Path +import tarfile + +import pytest + +HERE = Path(__file__).resolve().parent +spec = importlib.util.spec_from_file_location('mesh_indexer_patch', HERE / 'patch_indexer_barrier.py') +patch = importlib.util.module_from_spec(spec) +spec.loader.exec_module(patch) + + +def test_pinned_selector_archive_produces_the_checked_barrier_source(): + with tarfile.open(HERE.parent / 'compute/b12x-selector-files.tar.gz') as archive: + source = archive.extractfile(patch.RELATIVE).read() + assert hashlib.sha256(source).hexdigest() == patch.BEFORE + result = patch.patch_bytes(source) + assert hashlib.sha256(result).hexdigest() == patch.AFTER + assert result.count(b'\r\n') == source.count(b'\r\n') + 2 + compile(result.decode('utf-8'), patch.RELATIVE, 'exec') + assert patch.patch_bytes(result) == result + + +def test_unknown_source_is_rejected_without_mutation(tmp_path): + target = tmp_path / 'indexer.py' + target.write_bytes(b'unsupported source\n') + with pytest.raises(ValueError, match='Unsupported mesh indexer'): + patch.apply(target) + assert target.read_bytes() == b'unsupported source\n' diff --git a/runtime/glm53-spark-mtp3-mesh/indexer-barrier/validation.json b/runtime/glm53-spark-mtp3-mesh/indexer-barrier/validation.json new file mode 100644 index 00000000..58983ef0 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/indexer-barrier/validation.json @@ -0,0 +1,408 @@ +{ + "bundle_manifest_sha256": "69313e19e881ec93e9ed3bd150d2f24fc6b444488ac729a69f45d038e2243500", + "cache_namespace": "glm53-spark-df116c4f-mtp3-nvfp4-a16-c57b6438c-mesh69313e19-tail-cow-v2", + "indexer": { + "path": "b12x/attention/dsa_indexer/fused_indexer.py", + "preimage_sha256": "893fbcade135b7e1d146b8fb6530cde0650be515f69bf9a17ced0a9c61a141e2", + "result_sha256": "35a7564ec8bf18f1b7dc5562b21616437b65e7b0689c56341b6436c8d202637b", + "script_sha256": "70ac64c8e58c448eb2b4184395c55821d371d5654298a4fb7c5f55a76ab884ff" + }, + "installer_sha256": "af89efb28e34b2035f608ca1a2edf20a34d3f397f52a0e5e65d24f9b2b37667f", + "parent_image": "ghcr.io/fujitsupolycom/sparkring-glm53-sparkcache@sha256:67dc0ae453baaae6831ccec1d259b4ef8b236a8b0dc9f747d901b95c66ec1987", + "parent_image_id": "sha256:2e41b1e934a85ff7c21b780532db2f0a0e978df081e52f4ae2bf11f8992fb24f", + "parent_source_lock_sha256": "139f36701e0e47f45bf99fba2cc2fa59b417f2ee801dad3a064455d5b464a459", + "schema": "sparkring-mesh-indexer-barrier/v1", + "source_lock_sha256": "57b6438ca014a493380fb874a84977d35b0983e9ee224a62ade3415fd9c7ce46", + "source_receipt_sha256": "030e9f3749672d186b66e41336980af1abb5c207f734feafa677ba82456c5d43", + "status": "research-only", + "artifact": { + "reference": "ghcr.io/fujitsupolycom/sparkring-glm53-sparkcache@sha256:7aa57ed1901b83c8581c53cb233db930bb1a133d74be01168d967b9dd56b3e49", + "image_id": "sha256:995a42b4b08f525813c5783f78f34571b6bd1c197fccee53617ff8060574a16d" + }, + "verification": { + "added_layers": 2, + "bundle_manifest_sha256": "69313e19e881ec93e9ed3bd150d2f24fc6b444488ac729a69f45d038e2243500", + "checks_passed": true, + "image": "sha256:995a42b4b08f525813c5783f78f34571b6bd1c197fccee53617ff8060574a16d", + "image_id": "sha256:995a42b4b08f525813c5783f78f34571b6bd1c197fccee53617ff8060574a16d", + "image_reference": "sha256:995a42b4b08f525813c5783f78f34571b6bd1c197fccee53617ff8060574a16d", + "image_size_bytes": 22108643353, + "inside_image": { + "b12x_commit": "ef308bac0f3b3eb8fea63e4013afc0c2ea1c6301", + "bundle_files": 28, + "bundle_manifest_sha256": "69313e19e881ec93e9ed3bd150d2f24fc6b444488ac729a69f45d038e2243500", + "checks_passed": true, + "compute": { + "b12x_files": 385, + "b12x_revision": "ef308bac0f3b3eb8fea63e4013afc0c2ea1c6301", + "b12x_tree": "dcf039e5e754136275835ea997e6b9abbb6b15ae", + "cuda_version": "13.3", + "environment": { + "CUDA_HOME": "/opt/cuda-13.3", + "TRITON_PTXAS_PATH": "/opt/cuda-13.3/bin/ptxas", + "VLLM_B12X_DENSE_ACTIVATION_MODE": "auto", + "VLLM_GDN_SPEC_DECODE_METADATA_FASTPATH": "1", + "VLLM_LM_HEAD_A16": "1", + "VLLM_MTP_NVFP4_LM_HEAD": "1", + "VLLM_MXFP8_LM_HEAD": "0" + }, + "proposal_head_nvfp4": true, + "source_lock_sha256": "57b6438ca014a493380fb874a84977d35b0983e9ee224a62ade3415fd9c7ce46", + "target_head_quantization": false, + "vllm_overrides": 24, + "vllm_parent_files": 2905 + }, + "cuda_initialized": false, + "device_access": false, + "limitation": "Content and CPU checks do not qualify CUDA graphs, RDMA forwarding, native MTP, cache restoration, or model performance.", + "marker_binary_sha256": "2828c07e4255c4962c77425be2c88969e7eb7dd4b1bf9e36485bc705bb5d6d64", + "marker_source_sha256": "8684a6961b8e86aa474fa2310ff71e4cdf219a63a72ceb5593b2f95e54812792", + "model_loaded": false, + "parent_package_files": { + "b12x": 385, + "sparkcache": 150, + "vllm": 2905 + }, + "python_syntax_files": 25, + "readiness_warmup": { + "environment": "SPARKRING_WARMUP_TEMPERATURE", + "helper_sha256": "f41c38eef41d15d63dcfc49cd6643357ca1a3ae18200ddbe4f8692d0b767ee79", + "temperature": 1.0 + }, + "rocenante_lazy_import": "/opt/spark-sircl/b12x_overlay/b12x/comm/roce/__init__.py", + "sircl_native_sha256": "61aa0ec56a1b438439bed8611dab0353d2c72c10af02bbd917fb77c87b33e5fc", + "source_receipt_sha256": "030e9f3749672d186b66e41336980af1abb5c207f734feafa677ba82456c5d43", + "sparkcache_commit": "66057174301a4759ca3a45207ea41016689449cb", + "status": "research-only", + "vllm_commit": "e02b174693e13859de61811b5e8cd13d5308e259", + "vllm_native_extensions": 15 + }, + "limitation": "Content-verifier scope excludes GPU and model execution; separate GPU evidence follows.", + "parent_image_id": "sha256:2e41b1e934a85ff7c21b780532db2f0a0e978df081e52f4ae2bf11f8992fb24f", + "parent_layers_retained": 89, + "platform": "linux/arm64", + "schema": "sparkring-mtp3-mesh-image-receipt/v1", + "source_receipt_sha256": "030e9f3749672d186b66e41336980af1abb5c207f734feafa677ba82456c5d43", + "status": "research-only", + "verification_command": [ + "docker", + "run", + "--rm", + "--network", + "none", + "--read-only", + "--cap-drop", + "ALL", + "--security-opt", + "no-new-privileges", + "--cpus", + "2", + "--memory", + "2g", + "--pids-limit", + "128", + "--tmpfs", + "/tmp:rw,nosuid,nodev,size=128m", + "--env", + "PYTHONDONTWRITEBYTECODE=1", + "--entrypoint", + "python3", + "sha256:995a42b4b08f525813c5783f78f34571b6bd1c197fccee53617ff8060574a16d", + "-I", + "/opt/sparkring/bin/verify-mtp3-mesh-image.py", + "--inside-image" + ] + }, + "gpu_probe": [ + { + "gpu": "NVIDIA GB10", + "sm": [ + 12, + 1 + ], + "source": "/usr/local/lib/python3.12/dist-packages/b12x/attention/dsa_indexer/fused_indexer.py", + "source_sha256": "893fbcade135b7e1d146b8fb6530cde0650be515f69bf9a17ced0a9c61a141e2", + "torch": "2.13.0+cu130" + }, + { + "entry_sync": false, + "eager": [ + [ + 512, + 513, + 512 + ], + [ + 512, + 513, + 512 + ], + [ + 512, + 513, + 512 + ], + [ + 512, + 513, + 512 + ], + [ + 512, + 513, + 512 + ], + [ + 512, + 513, + 512 + ], + [ + 512, + 513, + 512 + ], + [ + 512, + 513, + 512 + ], + [ + 512, + 513, + 512 + ], + [ + 512, + 513, + 512 + ] + ], + "graph": [ + [ + 512, + 513, + 512 + ], + [ + 512, + 513, + 512 + ], + [ + 512, + 513, + 512 + ], + [ + 512, + 513, + 512 + ], + [ + 512, + 513, + 512 + ], + [ + 512, + 513, + 512 + ], + [ + 512, + 513, + 512 + ], + [ + 512, + 513, + 512 + ], + [ + 512, + 513, + 512 + ], + [ + 512, + 513, + 512 + ] + ], + "incomplete_reads": 20 + }, + { + "entry_sync": true, + "eager": [ + [ + 513, + 513, + 513 + ], + [ + 513, + 513, + 513 + ], + [ + 513, + 513, + 513 + ], + [ + 513, + 513, + 513 + ], + [ + 513, + 513, + 513 + ], + [ + 513, + 513, + 513 + ], + [ + 513, + 513, + 513 + ], + [ + 513, + 513, + 513 + ], + [ + 513, + 513, + 513 + ], + [ + 513, + 513, + 513 + ] + ], + "graph": [ + [ + 513, + 513, + 513 + ], + [ + 513, + 513, + 513 + ], + [ + 513, + 513, + 513 + ], + [ + 513, + 513, + 513 + ], + [ + 513, + 513, + 513 + ], + [ + 513, + 513, + 513 + ], + [ + 513, + 513, + 513 + ], + [ + 513, + 513, + 513 + ], + [ + 513, + 513, + 513 + ], + [ + 513, + 513, + 513 + ] + ], + "incomplete_reads": 0 + } + ], + "stress": [ + { + "source_sha256": "35a7564ec8bf18f1b7dc5562b21616437b65e7b0689c56341b6436c8d202637b", + "gpu": "NVIDIA GB10", + "torch": "2.13.0+cu130" + }, + { + "rows": 3, + "heads": 32, + "topk": 512, + "lengths": [ + 4097, + 65536 + ], + "graph_replays": 500, + "concurrent_copy_bytes": 67108864, + "wall_seconds": 0.42007936000300106, + "result": "passed" + }, + { + "rows": 4, + "heads": 32, + "topk": 512, + "lengths": [ + 4097, + 200000 + ], + "graph_replays": 500, + "concurrent_copy_bytes": 67108864, + "wall_seconds": 0.5173187279942795, + "result": "passed" + }, + { + "rows": 8, + "heads": 32, + "topk": 512, + "lengths": [ + 4097, + 65536 + ], + "graph_replays": 500, + "concurrent_copy_bytes": 67108864, + "wall_seconds": 0.5455467720021261, + "result": "passed" + }, + { + "rows": 16, + "heads": 32, + "topk": 512, + "lengths": [ + 4097, + 65536 + ], + "graph_replays": 500, + "concurrent_copy_bytes": 67108864, + "wall_seconds": 0.7676780149995466, + "result": "passed" + } + ], + "gpu_tests_passed": 15, + "local_tests": { + "passed": 348, + "skipped": 3 + } +} diff --git a/runtime/glm53-spark-mtp3-mesh/managed_install.py b/runtime/glm53-spark-mtp3-mesh/managed_install.py index 4a751cba..3330622b 100644 --- a/runtime/glm53-spark-mtp3-mesh/managed_install.py +++ b/runtime/glm53-spark-mtp3-mesh/managed_install.py @@ -30,6 +30,7 @@ 'runtime/glm53-spark-mtp3-mesh/profile.py', 'runtime/glm53-spark-mtp3-mesh/inspect_fabric.py', 'runtime/glm53-spark-mtp3-mesh/pins.json', + 'runtime/glm53-spark-mtp3-mesh/compute/source-lock.json', 'runtime/glm53-flash-jj-r8-gb10/pins.json', 'runtime/glm53-flash-jj-r8-gb10/warmup_dflash.py', 'runtime/glm53-flash-jj-r8-gb10/launch-rank.sh', @@ -254,7 +255,7 @@ def prepare_plan(launch, image_receipt, rank, epoch, health_port, key_file): if '--managed' not in help_result.stdout + help_result.stderr: raise ValueError('Extracted helper does not expose managed lifetime') bundle = Path(site['bundle_root']) - if profile.sha(bundle / 'sparkring-overlay-manifest.json') != profile.PINS['canonical_bundle_manifest_sha256']: + if profile.sha(bundle / 'sparkring-overlay-manifest.json') != receipt['bundle_manifest_sha256']: raise ValueError('Host transport bundle manifest differs') manifest = json.loads((bundle / 'sparkring-overlay-manifest.json').read_text()) for item in manifest['files']: diff --git a/runtime/glm53-spark-mtp3-mesh/managed_service.py b/runtime/glm53-spark-mtp3-mesh/managed_service.py index 3aab600b..9f1076cb 100644 --- a/runtime/glm53-spark-mtp3-mesh/managed_service.py +++ b/runtime/glm53-spark-mtp3-mesh/managed_service.py @@ -126,7 +126,8 @@ def validate_group(rows): generations = {str(row['rank']): row['generation'] for row in rows} view = digest(generations) if any(row.get('phase') != 'armed' or row.get('view_digest') != view - or row.get('peer_health_degraded', False) for row in rows): + or row.get('peer_health_degraded', False) + or row.get('docker_status_degraded', False) for row in rows): raise RuntimeError('Mesh ranks have not armed the same process generation set') return view @@ -167,10 +168,41 @@ def docker_running(name): result = subprocess.run(['docker', 'inspect', '--format', '{{.State.Running}}', name], capture_output=True, text=True, timeout=3) if result.returncode: - if 'No such' in result.stderr: + if re.search(r'No such (?:object|container): ' + re.escape(name) + r'(?:\s|$)', result.stderr): return False raise RuntimeError('Cannot establish the dependent container state') - return result.stdout.strip() == 'true' + state = result.stdout.strip() + if state not in ('true', 'false'): + raise RuntimeError('Docker returned an unknown dependent container state') + return state == 'true' + + +class DockerStatePoll: + """One bounded background query; unknown state never proves model exit.""" + + def __init__(self, name): + self.name = name + self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=1) + self.pending = None + self.error = None + + def poll(self): + running = None + if self.pending is not None and self.pending.done(): + try: + running = self.pending.result() + self.error = None + except (subprocess.TimeoutExpired, OSError, RuntimeError) as error: + self.error = str(error) + self.pending = None + if self.pending is None: + self.pending = self.executor.submit(docker_running, self.name) + # Do not reuse a previous False while a query is pending: a container + # can start between samples. Cleanup uses a separate synchronous proof. + return running + + def close(self): + self.executor.shutdown(wait=False, cancel_futures=True) def stop_model(name): @@ -405,6 +437,7 @@ def run(self): fcntl.flock(service_lock, fcntl.LOCK_EX | fcntl.LOCK_NB) for sig in (signal.SIGTERM, signal.SIGINT): signal.signal(sig, lambda *_: self.stop.set()) + docker_watch = None try: if docker_running(self.model): raise RuntimeError('Stop the dependent model before starting mesh ownership') @@ -420,6 +453,7 @@ def run(self): peer_watch = PeerWatch() last_network = time.monotonic() last_full_network = last_network + docker_watch = DockerStatePoll(self.model) while not self.stop.is_set(): if any(child.poll() is not None for child in self.children): raise RuntimeError('A managed source marker exited') @@ -429,6 +463,9 @@ def run(self): if full: last_full_network = time.monotonic() last_network = time.monotonic() + model_running = docker_watch.poll() + self.publish(docker_status_degraded=model_running is None, + docker_status_error=docker_watch.error) try: rows = group_check(self.site, self.config, self.key, self.identity) view = peer_watch.observe(rows) @@ -436,21 +473,22 @@ def run(self): except OSError: if peer_watch.generations is not None: peer_watch.transport_error(time.monotonic()) - elif docker_running(self.model): + elif model_running is True: raise self.publish(peer_health_degraded=True) except Exception: - if peer_watch.generations is not None or docker_running(self.model): + if peer_watch.generations is not None or model_running is True: raise - if self.state['phase'] != 'armed' and docker_running(self.model): + if self.state['phase'] != 'armed' and model_running is True: raise RuntimeError('Dependent model started before four-rank readiness') intent_path = self.state_dir / 'model-intent.json' if intent_path.exists(): intent = json.loads(intent_path.read_text()) if intent.get('generation') == self.generation and intent.get('active') is True: - if docker_running(self.model): + if model_running is True: self.model_seen = True - elif self.model_seen or time.monotonic() > intent['deadline_monotonic']: + elif model_running is False and ( + self.model_seen or time.monotonic() > intent['deadline_monotonic']): raise RuntimeError('Dependent model exited or failed to start') else: self.model_seen = False @@ -462,6 +500,8 @@ def run(self): self.publish(best_effort=True, local_ready=False, phase='failed', error=str(error)) print(json.dumps({'event': 'mesh_failure', 'rank': self.rank, 'error': str(error)}), flush=True) finally: + if docker_watch is not None: + docker_watch.close() self.publish(best_effort=True, local_ready=False, phase='failed' if self.failed else 'stopping') while self.owns_guard: try: @@ -534,7 +574,8 @@ def model_intent(config_path, active): raise RuntimeError('Mesh service has no readiness state') return status = json.loads(status_path.read_text()) - if active and (status.get('phase') != 'armed' or status.get('local_ready') is not True): + if active and (status.get('phase') != 'armed' or status.get('local_ready') is not True + or status.get('docker_status_degraded', False)): raise RuntimeError('Mesh is not armed for model startup') record = {'generation': status['generation'], 'active': active, 'deadline_monotonic': time.monotonic() + 15} diff --git a/runtime/glm53-spark-mtp3-mesh/performance/.gitattributes b/runtime/glm53-spark-mtp3-mesh/performance/.gitattributes new file mode 100644 index 00000000..afce3926 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/.gitattributes @@ -0,0 +1,6 @@ +*.gz -text +*.json -text whitespace=cr-at-eol +checkpoints/payload-by-sha/** -text whitespace=-trailing-space +checkpoints/ownership-contract.json -text +transport/bundle-source/** -text whitespace=cr-at-eol +transport/B12X-LICENSE -text whitespace=cr-at-eol diff --git a/runtime/glm53-spark-mtp3-mesh/performance/Dockerfile b/runtime/glm53-spark-mtp3-mesh/performance/Dockerfile new file mode 100644 index 00000000..d248198b --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/Dockerfile @@ -0,0 +1,10 @@ +ARG PARENT_IMAGE=ghcr.io/fujitsupolycom/sparkring-glm53-sparkcache@sha256:67dc0ae453baaae6831ccec1d259b4ef8b236a8b0dc9f747d901b95c66ec1987 +FROM ${PARENT_IMAGE} +RUN --mount=type=bind,target=/performance python3 -S -B /performance/install.py +ENV SPARK_GDN_PREFILL_CHECKPOINTS=2 +LABEL org.sparkring.runtime.status="research-only" \ + org.sparkring.runtime.profile="glm53-mtp3-cache-checkpoints" \ + org.sparkcache.commit="48bbd2be4a7b972e56632a2d7b934bac5460f272" \ + org.sparkcache.cuda-placement-sha256="2657cdd2e54a097c9544e4c79ae62c0646db6db123ff24e4f0c384238c3a1e8d" \ + org.sparkring.sircl.manifest-sha256="c0fd5567442b08b908cc193f36d0864e262573c7e5d232509479a823cface742" +ENTRYPOINT ["python3", "-S", "-B", "/opt/sparkring/bin/start-performance.py"] diff --git a/runtime/glm53-spark-mtp3-mesh/performance/README.md b/runtime/glm53-spark-mtp3-mesh/performance/README.md new file mode 100644 index 00000000..22270cec --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/README.md @@ -0,0 +1,64 @@ +# Native MTP3 cache and checkpoint performance composition + +Status: **implemented** source composition; performance is **research-only**. +The build combines GLM-5.3 native MTP3 compute, verified persistent caching, +explicit recurrent checkpoints, and stream-ordered hardware mesh transport. +The recipe preserves the parent model weights and does not change host fabric. + +## Build inputs + +Use a clean SparkCache checkout at `48bbd2be4a7b972e56632a2d7b934bac5460f272`. +It contains the merged restore/publication improvements, periodic-capture +option, and backlog gauges. Periodic full capture defaults off; enabling it +trades more writes for shorter history reconstruction. + +The compute parent is the published native-MTP3 image with config ID +`2e41b1e934a85ff7c21b780532db2f0a0e978df081e52f4ae2bf11f8992fb24f`. +`Dockerfile` identifies its immutable registry reference. + +Supply the verified native placement and transport libraries named by SHA-256 +in `prepare.py`. The transport [source package](transport/README.md) contains +build instructions, licensing, and the full source inventory. SparkCache's +`sparkcache/native/README.md` describes placement builds. Compiler/toolchain +differences can change binary hashes; a different binary requires separate +qualification and must not bypass the expected-artifact checks. + +For exact artifact replay, the published image in `public-image.json` contains +both libraries. Use `docker create` and `docker cp` to extract +`/opt/sparkcache-src/sparkcache/native/build-cuda/libspark_cache_placement.so` +and `/opt/spark-sircl/libspark_transport_capi.so` without running a model. +Verify the digests in `prepare.py` before using them as build inputs. + +```bash +python runtime/glm53-spark-mtp3-mesh/performance/prepare.py \ + --sparkcache /path/to/sparkcache \ + --placement-library /path/to/libspark_cache_placement.so \ + --transport-library /path/to/libspark_transport_capi.so \ + --output /path/to/absent-build-context +docker build -t sparkring-mtp3-cache-checkpoints /path/to/absent-build-context +docker run --rm --network none --entrypoint python3 \ + sparkring-mtp3-cache-checkpoints -S -B /opt/sparkring/bin/verify-performance.py +``` + +The preparer performs no remote action. The installer verifies source inputs, +applies strict runtime preimages, installs the complete checkpoint ownership +contract, and generates a file inventory checked before serving. A mismatched +source or ownership dependency fails the build rather than weakening checks. + +## Behavior + +- Bounded restore memory, authenticated history reconstruction, and protected + publication lifetimes come from SparkCache's pinned source. +- Recurrent checkpoints are materialized and retained through the full + scheduler/runner/model-state path. Cache reuse retains verification backoff. +- The native transport preserves CTA publication ordering, constructs all CPU + proxy state before thread startup, and links payload/doorbell submissions. +- Graph-only collective geometry uses the content-addressed bundle described + in `transport/README.md`; eager geometry is unchanged. +- GLM-5.3 chat requests cannot disable thinking with an unsupported flag. + The API rejects that request before generation; warmup uses supported low + reasoning effort. This does not implement reasoning-free generation. + +The image verifier establishes file identity, not a throughput result. +Source equivalence with a serving image and bounded test evidence must be +recorded separately from any claim that this image completed a serving soak. diff --git a/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/README.md b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/README.md new file mode 100644 index 00000000..611222c7 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/README.md @@ -0,0 +1,56 @@ +# Explicit recurrent checkpoints during prefill + +Status: **implemented**. This source payload lets one eligible fresh prompt +prefill export up to two interior recurrent states. The scheduler specifies +positions explicitly; the model runner, GLM model-state adapter, GDN metadata, +and B12X recurrence kernel carry and materialize the same plan. Convolution +checkpoints preserve logical history rather than physical speculative slots. + +The enabled path requires `SPARK_GDN_PREFILL_CHECKPOINTS=2`. It accepts complete, +block-aligned fresh prompts of at most 8192 tokens and at most two required +interior checkpoints. Continuations and ineligible shapes retain the scheduler's +ordinary path. It does not generalize checkpoint coalescing to arbitrary lengths. + +## Build interface + +`patch-manifest.json` maps image paths to required preimage hashes and exact +replacement payloads. The filesystem identified by image digest +`sha256:3882eccf0b42e26dad399a3ea89a45988e04413802fccfb146d1486f8bb2fc13` +defines the measured parent composition. A source-built parent must satisfy +every runtime preimage even if its image digest differs. + +Run these commands from this directory. Installation is a Docker build operation, +not a command to run against a serving container or host installation. + +```sh +python3 install.py verify-context --manifest-sha256 0970d29ec33e9f8525a2cc55989ab0deb937bff88b5f16d4d5280035359e4c55 +python3 install.py verify-preimages --manifest-sha256 0970d29ec33e9f8525a2cc55989ab0deb937bff88b5f16d4d5280035359e4c55 +python3 install.py apply --manifest-sha256 0970d29ec33e9f8525a2cc55989ab0deb937bff88b5f16d4d5280035359e4c55 +``` + +All preimages, payload hashes, Python syntax, and required symbols are checked +before source replacement. Installation preserves file modes and ownership, +refreshes existing bytecode, and normalizes affected timestamps. Its JSON output +includes `postimages`, the map the composition builder must incorporate into +SparkCache's runtime ownership contract. The builder must regenerate the package +source receipt and execute the ownership verifier before declaring the image +usable. This installer does not replace SparkCache package files or its receipts. + +`attestation_reference_hashes` records the measured composition's attestations +for provenance; those hashes are not replacement payloads for another SparkCache +revision. `native_hashes` records its cache-library dependencies for the enclosing +image verifier. The original payload manifest is identified by +`source_manifest_sha256`; all 18 retained source payloads preserve its exact bytes. + +## Evidence and limits + +A four-Spark TP4/DCP4 native-MTP3 deployment with an 8192-token prefill budget +passed 17 semantic requests covering fresh/repeated pairs, extended triples, +and mixed requests. Repeated requests reported 6144 cached tokens with exact +expected answers. All four ranks reported the 8192-token prompt's checkpoint +positions at 6144 and 7168 across 34 GDN layers. This qualifies those cases; +it does not qualify arbitrary models, context lengths, or rebuilt image compositions. + +Offline checks run with `python -m pytest test_checkpoint_package.py`. They check source +identity rejection and execute the packaged checkpoint planner. GPU recurrence, +convolution state, and live semantic validation remain hardware requirements. diff --git a/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/install.py b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/install.py new file mode 100644 index 00000000..0bdcc5e4 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/install.py @@ -0,0 +1,144 @@ +"""Build-time only: verify all preimages, apply pinned files, verify attestations.""" +from __future__ import annotations +import ast +import hashlib +import json +import os +import py_compile +import re +from pathlib import Path +import stat + +def sha(data): + return hashlib.sha256(data).hexdigest() + + +def verify_files(records, *, root=Path("/")): + for absolute, expected in records.items(): + path = root / absolute.lstrip("/") + if any(p.is_symlink() for p in (path, *path.parents)): + raise RuntimeError(f"Symlink target refused: {absolute}") + if expected is None: + if path.exists(): + raise RuntimeError(f"Expected absent file: {absolute}") + elif not path.is_file() or sha(path.read_bytes()) != expected: + raise RuntimeError(f"Source identity mismatch: {absolute}") + + +def verify_symbols(data, symbols): + tree = ast.parse(data) + classes = {node.name: node for node in tree.body if isinstance(node, ast.ClassDef)} + top = {node.name for node in tree.body if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef))} + for symbol in symbols: + if "." not in symbol: + if symbol not in top: + raise RuntimeError(f"Required top-level symbol missing: {symbol}") + continue + cls, member = symbol.split(".", 1) + members = set() + for node in classes.get(cls, ast.ClassDef(name="", bases=[], keywords=[], body=[], decorator_list=[])).body: + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): + members.add(node.name) + elif isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name): + members.add(node.target.id) + elif isinstance(node, ast.Assign): + members.update(t.id for t in node.targets if isinstance(t, ast.Name)) + if member not in members: + raise RuntimeError(f"Required class member missing: {symbol}") + + +def apply_manifest(expected_manifest_sha256): + context = Path(__file__).parent + encoded = (context / "patch-manifest.json").read_bytes() + if sha(encoded) != expected_manifest_sha256: + raise RuntimeError("Frozen build manifest hash mismatch") + manifest = json.loads(encoded) + verify_files(manifest["preimages"]) + payload = {} + for target, record in manifest["replacements"].items(): + source = context / record["payload"] + data = source.read_bytes() + if sha(data) != record["sha256"]: + raise RuntimeError(f"Frozen payload mismatch: {target}") + if target.endswith(".py"): + compile(data, target, "exec") + verify_symbols(data, record.get("required_symbols", [])) + payload[target] = data + # Validate every input and replacement before modifying the image. + compiled_caches = set() + for target, data in payload.items(): + path = Path(target) + previous = path.stat() if path.exists() else None + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(data) + os.chmod(path, stat.S_IMODE(previous.st_mode) if previous else 0o644) + if previous: + os.chown(path, previous.st_uid, previous.st_gid) + if path.suffix == ".py": + cache = path.parent / "__pycache__" + if cache.is_dir(): + for pyc in cache.glob(path.stem + ".*.pyc"): + if pyc.is_file() and not pyc.is_symlink(): + match = re.search(r"\.opt-(\d+)\.pyc$", pyc.name) + py_compile.compile(str(path), cfile=str(pyc), doraise=True, + optimize=int(match.group(1)) if match else 0, + invalidation_mode=py_compile.PycInvalidationMode.CHECKED_HASH) + compiled_caches.add(pyc) + verify_files({name: row["sha256"] for name, row in manifest["replacements"].items()}) + verify_files({name: expected for name, expected in manifest["preimages"].items() + if name not in manifest["replacements"]}) + postimages = {name: row["sha256"] for name, row in manifest["replacements"].items()} + # Make the filesystem diff independent of host build time. This touches + # only timestamps, including directories affected by file/bytecode updates. + touched = {Path(name) for name in payload} | compiled_caches + parents = set() + for path in touched: + os.utime(path, ns=(0, 0)) + parents.update(path.parents) + cache = path.parent / "__pycache__" + if cache.is_dir(): + parents.add(cache) + for path in sorted(parents, key=lambda p: len(p.parts), reverse=True): + os.utime(path, ns=(0, 0)) + print(json.dumps({"patch_manifest_sha256": sha(encoded), "files_replaced": len(payload), + "postimages": postimages, "attestation_regeneration_required": True})) + return postimages + + +def apply(site_packages: Path) -> dict[str, str]: + """Apply pinned files in the declared image layout and return ownership hashes.""" + if site_packages != Path("/usr/local/lib/python3.12/dist-packages"): + raise ValueError("Checkpoint payload requires /usr/local/lib/python3.12/dist-packages") + return apply_manifest("0970d29ec33e9f8525a2cc55989ab0deb937bff88b5f16d4d5280035359e4c55") + + + +def main(): + import argparse + parser = argparse.ArgumentParser(description="Validate or install the explicit recurrent checkpoint payload.") + parser.add_argument("command", choices=("verify-context", "verify-preimages", "apply")) + parser.add_argument("--manifest-sha256", required=True) + args = parser.parse_args() + context = Path(__file__).resolve().parent + encoded = (context / "patch-manifest.json").read_bytes() + if sha(encoded) != args.manifest_sha256: + raise RuntimeError("Manifest identity mismatch") + manifest = json.loads(encoded) + for target, record in manifest["replacements"].items(): + source = (context / record["payload"]).resolve() + if not source.is_relative_to(context) or sha(source.read_bytes()) != record["sha256"]: + raise RuntimeError(f"Payload identity mismatch: {target}") + if target.endswith(".py"): + compile(source.read_bytes(), target, "exec") + verify_symbols(source.read_bytes(), record.get("required_symbols", [])) + if args.command == "verify-preimages": + verify_files(manifest["preimages"]) + if args.command == "apply": + apply_manifest(args.manifest_sha256) + else: + print(json.dumps({"manifest_sha256": args.manifest_sha256, + "payloads_verified": len(manifest["replacements"])})) + + +if __name__ == "__main__": + main() diff --git a/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/ownership-contract.json b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/ownership-contract.json new file mode 100644 index 00000000..a7481f7f --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/ownership-contract.json @@ -0,0 +1,203 @@ +{ + "base_contract_sha256": "aaa1e1bfd01a013f7d02df7cc75e5eac01a1643fce3e5651a4574da7ce5190c4", + "checkpoint_runner_support": "explicit-plan-v1-and-v2", + "files": [ + { + "path": "vllm/distributed/kv_transfer/kv_connector/v1/base.py", + "required_symbols": [ + "SupportsHMA.request_finished_all_groups", + "KVConnectorBase_V1.handle_preemptions", + "KVConnectorBase_V1.get_finished", + "KVConnectorBase_V1.get_recurrent_publication_boundaries" + ], + "sha256": "7460e0638c0c81808dbdbbad7114db2e840e531e9ce83360c55b9171fcc872b1" + }, + { + "path": "vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py", + "required_symbols": [ + "MultiConnector.supports_recurrent_boundary_blocks", + "MultiConnector.get_recurrent_publication_boundaries" + ], + "sha256": "8741b86b0f3e91af06d01a52240cdc988e3073bf049f72b506f56d6ceebb1ac0" + }, + { + "path": "vllm/distributed/kv_transfer/kv_connector/utils.py", + "required_symbols": [ + "KVOutputAggregator.from_connector", + "KVOutputAggregator.aggregate" + ], + "sha256": "115512dca36b0711223f55f5ff304abffccc6d078923ba4c6bbb3cdb7fdd39a2" + }, + { + "path": "vllm/v1/core/sched/scheduler.py", + "required_symbols": [ + "Scheduler._connector_finished", + "Scheduler._free_blocks", + "Scheduler._free_request_blocks", + "Scheduler._recurrent_checkpoint_plan", + "Scheduler._recurrent_publication_boundaries", + "Scheduler._update_from_kv_xfer_finished", + "Scheduler.schedule" + ], + "sha256": "ce9460834e08f97dbbfeb3f1238b78ee6a3363dd59a2aef8ceeb715857385895" + }, + { + "path": "vllm/v1/core/kv_cache_manager.py", + "required_symbols": [ + "KVCacheManager._pin_recurrent_boundary", + "KVCacheManager.allocate_slots", + "KVCacheManager.get_block_ids_for_computed_tokens", + "KVCacheManager.pop_blocks_for_free", + "KVCacheManager.take_recurrent_boundary_blocks" + ], + "sha256": "f4ccf9da197eb68fd21b23e8c03615cd0d4c2f871da0dc32a9c6ec82e44d3931" + }, + { + "path": "vllm/v1/core/sched/output.py", + "required_symbols": [ + "SchedulerOutput.preempted_req_ids", + "SchedulerOutput.recurrent_boundary_blocks", + "SchedulerOutput.recurrent_prefill_checkpoint_plans" + ], + "sha256": "0e1d3e4e1284fc2ad5b0cde9be9884eec61186e7731750f22ec25d828bb4e0f1" + }, + { + "path": "vllm/v1/worker/gpu_model_runner.py", + "required_symbols": [ + "GPUModelRunner._update_states", + "GPUModelRunner.execute_model" + ], + "sha256": "6888b3c5e16b4ef56ddc6d7bb02aa9736894a758574abeebe7e4ff6e9a4ef295" + }, + { + "path": "vllm/v1/worker/kv_connector_model_runner_mixin.py", + "required_symbols": [ + "KVConnectorModelRunnerMixin._get_kv_connector_output", + "KVConnectorModelRunnerMixin.finalize_kv_connector" + ], + "sha256": "3a3fa444f671956600c864f3eb82fbdd76b606535f1fdb651222dd1a8f9016d1" + }, + { + "path": "vllm/v1/core/block_pool.py", + "required_symbols": [ + "BlockPool.touch", + "BlockPool.free_blocks" + ], + "sha256": "ddee56dccb2208411b3a035918e917ce8f56a9858471e9ca12b420d5d79bc69c" + }, + { + "path": "vllm/v1/core/kv_cache_coordinator.py", + "required_symbols": [ + "KVCacheCoordinator.free", + "KVCacheCoordinator.pop_blocks_for_free", + "KVCacheCoordinator.get_blocks" + ], + "sha256": "149222b2b1de1ee3714c7c4e0843b9c38c730a605740ce7dc913308df906f20f" + }, + { + "path": "vllm/v1/core/single_type_kv_cache_manager.py", + "required_symbols": [ + "MambaManager._queue_aligned_recurrent_boundary", + "MambaManager.allocate_new_blocks", + "SingleTypeKVCacheManager._apply_cow", + "SingleTypeKVCacheManager.pop_blocks_for_free", + "SingleTypeKVCacheManager.take_pending_aligned_recurrent_boundaries", + "SingleTypeKVCacheManager.take_pending_partial_tail_offloads" + ], + "sha256": "10a6a60f31c67b59825aa827084f3d25416128ba5dc7c14a774675d64bbec40b" + }, + { + "path": "vllm/v1/attention/backends/gdn_attn.py", + "required_symbols": [ + "GDNAttentionMetadataBuilder.build", + "GDNPrefillCheckpointMetadata.required_mask" + ], + "sha256": "54707b0b7669642a55af60bdba0d9be5b5715cd3e0a68ca44021f11e4b8e9c9f" + }, + { + "path": "vllm/v1/attention/backend.py", + "required_symbols": [ + "CommonAttentionMetadata.recurrent_prefill_checkpoint_plans_cpu" + ], + "sha256": "c79f0f62c92963959131fa47ac63002b3109a379abbccabd26834d224ccb044d" + }, + { + "path": "vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py", + "required_symbols": [ + "KimiGatedDeltaNetAttention._run_b12x_kda_prefill", + "KimiGatedDeltaNetAttention._store_kda_conv_checkpoint" + ], + "sha256": "347a15ec42368b0c2b35b31cef7c9f58da97814a53ee76d0ee5526aac6b15c1c" + }, + { + "path": "vllm/v1/core/recurrent_prefill_checkpoint.py", + "required_symbols": [], + "sha256": "51ef011d11bb374d61840f1ff819759b6e3dd325da41856b7b0f01b013b54713" + }, + { + "path": "b12x/sequence/kda_prefill/__init__.py", + "required_symbols": [], + "sha256": "d415d6262885da11437e4a343da5ae04447273e529195b87547507f941fac0b8" + }, + { + "path": "b12x/sequence/kda_prefill/_cute_kernels.py", + "required_symbols": [ + "_PrologueKernel.kernel", + "_RecurrenceKernel.kernel" + ], + "sha256": "df12ee84e1677ce1a200d1aab488c9ce5ae774d21f70b7f15b39c636623ab753" + }, + { + "path": "b12x/sequence/kda_prefill/_impl.py", + "required_symbols": [ + "Caps.max_checkpoints" + ], + "sha256": "0d0fc64b5fddd269d823a57f14c09a9e0a1b38a4f802c39df07353f34873a270" + }, + { + "path": "b12x/sequence/kda_prefill/_policy.py", + "required_symbols": [ + "KdaPrefillQuery.max_checkpoints" + ], + "sha256": "ec64f07f841675501ff279bae3074d374ca8015a982ab12dc837ea2d03720a8e" + }, + { + "path": "b12x/sequence/kda_prefill/metadata.py", + "required_symbols": [], + "sha256": "7a9ef9ef3c59c88e4816bea6c9a74ebadd91a85632b42e221a5dd4009c5bb65c" + }, + { + "path": "b12x/sequence/kda_prefill/reference.py", + "required_symbols": [], + "sha256": "daa8e63b76bbbe2acf0b71954958523d6771e84fb5a932c838d88f81110e3816" + }, + { + "path": "vllm/models/glm5next/model_state.py", + "required_symbols": [ + "Glm5NextModelState.prepare_attn", + "Glm5NextAttnMetadata.get_extra_attn_kwargs" + ], + "sha256": "dc627424cd5814db0945b4c787e33fb6d035c2a2dd3d0a4ade4dadce61697810" + }, + { + "path": "vllm/v1/worker/gpu/model_runner.py", + "required_symbols": [ + "GPUModelRunner.execute_model" + ], + "sha256": "5ba22653a26444a989ce0942b5fbdb12cda4fa2722f818c8e199cd50843e6565" + }, + { + "path": "vllm/v1/worker/gpu/model_states/mamba_hybrid.py", + "required_symbols": [ + "MambaHybridAttnMetadata.get_extra_common_attn_kwargs", + "MambaHybridModelState.prepare_attn" + ], + "sha256": "90ae5dd89e808c9b13fd24f203feebd56e728b2928c13547e131d1ac7ca35965" + } + ], + "qualification": "Hash-attested prototype; unchanged schema and verifier; serving qualification separate", + "schema": "sparkring-vllm-kv-block-lease-contract/v1", + "source_state": "bounded-prefill-two-checkpoint-v1", + "status": "research-only", + "vllm_commit": "55969c16d4da57da76ee5729f3102d4b2003833c" +} diff --git a/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/patch-manifest.json b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/patch-manifest.json new file mode 100644 index 00000000..58505cc0 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/patch-manifest.json @@ -0,0 +1,182 @@ +{ + "attestation_reference_hashes": { + "/opt/sparkring/receipts/sparkcache-mtp3-candidate.json": "5d71b421449424d6ec4f3c3f2dee7987a451f338844c5ff844a56cb368e8a4ca", + "/usr/local/lib/python3.12/dist-packages/sparkcache/runtime_patches/vllm-manager-page-async-contract-55969c16.json": "9fbd8e2596287fd4e12aaf042e189bccfee2ea92126de42f58da8e44fb4d9c55" + }, + "base_image": "sha256:3882eccf0b42e26dad399a3ea89a45988e04413802fccfb146d1486f8bb2fc13", + "native_hashes": { + "/opt/sparkcache-src/sparkcache/native/build-cuda/libspark_cache_placement.so": "2657cdd2e54a097c9544e4c79ae62c0646db6db123ff24e4f0c384238c3a1e8d", + "/opt/sparkcache-src/sparkcache/native/build-cuda/libspark_cache_snapshot.so": "4398f18b8913e743e7bf1ed8fe29560d4580e61b6a1e2ab8b16684b19b6573b5" + }, + "parent_requirement": "The runtime filesystem must satisfy every preimage. The base image digest identifies the measured composition; a rebuilt parent with identical runtime preimages is accepted. The composition builder must regenerate SparkCache ownership attestations and verify them after applying the runtime files.", + "preimages": { + "/usr/local/lib/python3.12/dist-packages/b12x/sequence/kda_prefill/__init__.py": "2517a96e9674631eeada5a29ad694f36d771ae495fd0f1480b5090185fbbc8f0", + "/usr/local/lib/python3.12/dist-packages/b12x/sequence/kda_prefill/_cute_kernels.py": "d8c6887bcd2ffb3e95b78b4db82fdb1aeb105a97545fe9848fada9640dbf184f", + "/usr/local/lib/python3.12/dist-packages/b12x/sequence/kda_prefill/_impl.py": "1d10706f43abdc0521ee1ee3cbbf934d80c9c398585a57087ddb9ce38caf31b2", + "/usr/local/lib/python3.12/dist-packages/b12x/sequence/kda_prefill/_policy.py": "3113d0694170ff1175eff930dcb65a8fbc10e21a06701aa72ec7e57c49c861c5", + "/usr/local/lib/python3.12/dist-packages/b12x/sequence/kda_prefill/api.py": "56665da32751056b22c57d9f3995fb828a051377185ce80b812ba1f5e1fbb830", + "/usr/local/lib/python3.12/dist-packages/b12x/sequence/kda_prefill/metadata.py": null, + "/usr/local/lib/python3.12/dist-packages/b12x/sequence/kda_prefill/reference.py": "3f9db6d995d92cbf1ede446344d16fc42221d6f13585dcaf8955f71a51b4e91e", + "/usr/local/lib/python3.12/dist-packages/vllm/distributed/kv_transfer/kv_connector/utils.py": "115512dca36b0711223f55f5ff304abffccc6d078923ba4c6bbb3cdb7fdd39a2", + "/usr/local/lib/python3.12/dist-packages/vllm/distributed/kv_transfer/kv_connector/v1/base.py": "7460e0638c0c81808dbdbbad7114db2e840e531e9ce83360c55b9171fcc872b1", + "/usr/local/lib/python3.12/dist-packages/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py": "8741b86b0f3e91af06d01a52240cdc988e3073bf049f72b506f56d6ceebb1ac0", + "/usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py": "4a83c1eb9ea9c47c34e8a87b07f087a4cf1384428e74d800e0bdbd667abbb63e", + "/usr/local/lib/python3.12/dist-packages/vllm/models/glm5next/model_state.py": "4f884f713335d55fafd729fa4cc9bcd88389ba005323eb79a3f9b02cf5662f36", + "/usr/local/lib/python3.12/dist-packages/vllm/v1/attention/backend.py": "7be97ede36b98e571f9608207cbed8def0ae425c3e5246d50d7ada40d3c3d016", + "/usr/local/lib/python3.12/dist-packages/vllm/v1/attention/backends/gdn_attn.py": "fd75fb72efeb762ec558d332364e889df44a9ea483026bd5ded8f77f42be9b7f", + "/usr/local/lib/python3.12/dist-packages/vllm/v1/core/block_pool.py": "ddee56dccb2208411b3a035918e917ce8f56a9858471e9ca12b420d5d79bc69c", + "/usr/local/lib/python3.12/dist-packages/vllm/v1/core/kv_cache_coordinator.py": "149222b2b1de1ee3714c7c4e0843b9c38c730a605740ce7dc913308df906f20f", + "/usr/local/lib/python3.12/dist-packages/vllm/v1/core/kv_cache_manager.py": "417664f0709b6d432d8eed42c15efdb438be06316ae16224caa34b0870dff5ae", + "/usr/local/lib/python3.12/dist-packages/vllm/v1/core/recurrent_prefill_checkpoint.py": null, + "/usr/local/lib/python3.12/dist-packages/vllm/v1/core/sched/output.py": "9911b3f9d21815a185285852b5a6176e5484e1ab0ff5c30f7caaa68ea0fab543", + "/usr/local/lib/python3.12/dist-packages/vllm/v1/core/sched/scheduler.py": "75efa57e7ff5a77c76714b85e2e4d8e1d7f456d9a9eec6c67ebb11ca382942f9", + "/usr/local/lib/python3.12/dist-packages/vllm/v1/core/single_type_kv_cache_manager.py": "d2e35b012e0cf45ab3771f545c35ca48f2a5858549c574a352975607369124e2", + "/usr/local/lib/python3.12/dist-packages/vllm/v1/worker/gpu/model_runner.py": "f232398c93cf3193166136e823343840a41ebde8669253f56ce63a0023930cb4", + "/usr/local/lib/python3.12/dist-packages/vllm/v1/worker/gpu/model_states/mamba_hybrid.py": "b34cb130e233f4d322390acf7fec01a3a368090c7f79da1fa7770c622a2f0dde", + "/usr/local/lib/python3.12/dist-packages/vllm/v1/worker/gpu_model_runner.py": "5ac63bf43acc6b254282426a8a9b989cc07fcdbd07223101bdc99f91f4e4ce5f", + "/usr/local/lib/python3.12/dist-packages/vllm/v1/worker/kv_connector_model_runner_mixin.py": "3a3fa444f671956600c864f3eb82fbdd76b606535f1fdb651222dd1a8f9016d1" + }, + "replacements": { + "/usr/local/lib/python3.12/dist-packages/b12x/sequence/kda_prefill/__init__.py": { + "payload": "payload-by-sha/d415d6262885da11437e4a343da5ae04447273e529195b87547507f941fac0b8/__init__.py", + "required_symbols": [], + "sha256": "d415d6262885da11437e4a343da5ae04447273e529195b87547507f941fac0b8" + }, + "/usr/local/lib/python3.12/dist-packages/b12x/sequence/kda_prefill/_cute_kernels.py": { + "payload": "payload-by-sha/df12ee84e1677ce1a200d1aab488c9ce5ae774d21f70b7f15b39c636623ab753/_cute_kernels.py", + "required_symbols": [ + "_PrologueKernel.kernel", + "_RecurrenceKernel.kernel" + ], + "sha256": "df12ee84e1677ce1a200d1aab488c9ce5ae774d21f70b7f15b39c636623ab753" + }, + "/usr/local/lib/python3.12/dist-packages/b12x/sequence/kda_prefill/_impl.py": { + "payload": "payload-by-sha/0d0fc64b5fddd269d823a57f14c09a9e0a1b38a4f802c39df07353f34873a270/_impl.py", + "required_symbols": [ + "Caps.max_checkpoints", + "bind", + "run" + ], + "sha256": "0d0fc64b5fddd269d823a57f14c09a9e0a1b38a4f802c39df07353f34873a270" + }, + "/usr/local/lib/python3.12/dist-packages/b12x/sequence/kda_prefill/_policy.py": { + "payload": "payload-by-sha/ec64f07f841675501ff279bae3074d374ca8015a982ab12dc837ea2d03720a8e/_policy.py", + "required_symbols": [ + "KdaPrefillQuery.max_checkpoints" + ], + "sha256": "ec64f07f841675501ff279bae3074d374ca8015a982ab12dc837ea2d03720a8e" + }, + "/usr/local/lib/python3.12/dist-packages/b12x/sequence/kda_prefill/metadata.py": { + "payload": "payload-by-sha/7a9ef9ef3c59c88e4816bea6c9a74ebadd91a85632b42e221a5dd4009c5bb65c/metadata.py", + "required_symbols": [ + "validate_metadata" + ], + "sha256": "7a9ef9ef3c59c88e4816bea6c9a74ebadd91a85632b42e221a5dd4009c5bb65c" + }, + "/usr/local/lib/python3.12/dist-packages/b12x/sequence/kda_prefill/reference.py": { + "payload": "payload-by-sha/daa8e63b76bbbe2acf0b71954958523d6771e84fb5a932c838d88f81110e3816/reference.py", + "required_symbols": [], + "sha256": "daa8e63b76bbbe2acf0b71954958523d6771e84fb5a932c838d88f81110e3816" + }, + "/usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py": { + "payload": "payload-by-sha/347a15ec42368b0c2b35b31cef7c9f58da97814a53ee76d0ee5526aac6b15c1c/kimi_gdn_linear_attn.py", + "required_symbols": [ + "KimiGatedDeltaNetAttention._run_b12x_kda_prefill", + "KimiGatedDeltaNetAttention._store_kda_conv_checkpoint" + ], + "sha256": "347a15ec42368b0c2b35b31cef7c9f58da97814a53ee76d0ee5526aac6b15c1c" + }, + "/usr/local/lib/python3.12/dist-packages/vllm/models/glm5next/model_state.py": { + "payload": "payload-by-sha/dc627424cd5814db0945b4c787e33fb6d035c2a2dd3d0a4ade4dadce61697810/model_state.py", + "required_symbols": [ + "Glm5NextModelState.prepare_attn", + "Glm5NextAttnMetadata.get_extra_attn_kwargs" + ], + "sha256": "dc627424cd5814db0945b4c787e33fb6d035c2a2dd3d0a4ade4dadce61697810" + }, + "/usr/local/lib/python3.12/dist-packages/vllm/v1/attention/backend.py": { + "payload": "payload-by-sha/c79f0f62c92963959131fa47ac63002b3109a379abbccabd26834d224ccb044d/backend.py", + "required_symbols": [ + "CommonAttentionMetadata.recurrent_prefill_checkpoint_plans_cpu" + ], + "sha256": "c79f0f62c92963959131fa47ac63002b3109a379abbccabd26834d224ccb044d" + }, + "/usr/local/lib/python3.12/dist-packages/vllm/v1/attention/backends/gdn_attn.py": { + "payload": "payload-by-sha/54707b0b7669642a55af60bdba0d9be5b5715cd3e0a68ca44021f11e4b8e9c9f/gdn_attn.py", + "required_symbols": [ + "GDNPrefillCheckpointMetadata.required_mask", + "GDNAttentionMetadataBuilder.build" + ], + "sha256": "54707b0b7669642a55af60bdba0d9be5b5715cd3e0a68ca44021f11e4b8e9c9f" + }, + "/usr/local/lib/python3.12/dist-packages/vllm/v1/core/kv_cache_manager.py": { + "payload": "payload-by-sha/f4ccf9da197eb68fd21b23e8c03615cd0d4c2f871da0dc32a9c6ec82e44d3931/kv_cache_manager.py", + "required_symbols": [ + "KVCacheManager.allocate_slots" + ], + "sha256": "f4ccf9da197eb68fd21b23e8c03615cd0d4c2f871da0dc32a9c6ec82e44d3931" + }, + "/usr/local/lib/python3.12/dist-packages/vllm/v1/core/recurrent_prefill_checkpoint.py": { + "payload": "payload-by-sha/51ef011d11bb374d61840f1ff819759b6e3dd325da41856b7b0f01b013b54713/recurrent_prefill_checkpoint.py", + "required_symbols": [ + "validate_plan", + "fresh_prompt_plan", + "checkpoint_metadata" + ], + "sha256": "51ef011d11bb374d61840f1ff819759b6e3dd325da41856b7b0f01b013b54713" + }, + "/usr/local/lib/python3.12/dist-packages/vllm/v1/core/sched/output.py": { + "payload": "payload-by-sha/0e1d3e4e1284fc2ad5b0cde9be9884eec61186e7731750f22ec25d828bb4e0f1/output.py", + "required_symbols": [ + "SchedulerOutput.recurrent_prefill_checkpoint_plans" + ], + "sha256": "0e1d3e4e1284fc2ad5b0cde9be9884eec61186e7731750f22ec25d828bb4e0f1" + }, + "/usr/local/lib/python3.12/dist-packages/vllm/v1/core/sched/scheduler.py": { + "payload": "payload-by-sha/ce9460834e08f97dbbfeb3f1238b78ee6a3363dd59a2aef8ceeb715857385895/scheduler.py", + "required_symbols": [ + "Scheduler.schedule", + "Scheduler._recurrent_checkpoint_plan" + ], + "sha256": "ce9460834e08f97dbbfeb3f1238b78ee6a3363dd59a2aef8ceeb715857385895" + }, + "/usr/local/lib/python3.12/dist-packages/vllm/v1/core/single_type_kv_cache_manager.py": { + "payload": "payload-by-sha/10a6a60f31c67b59825aa827084f3d25416128ba5dc7c14a774675d64bbec40b/single_type_kv_cache_manager.py", + "required_symbols": [ + "MambaManager.allocate_new_blocks" + ], + "sha256": "10a6a60f31c67b59825aa827084f3d25416128ba5dc7c14a774675d64bbec40b" + }, + "/usr/local/lib/python3.12/dist-packages/vllm/v1/worker/gpu/model_runner.py": { + "payload": "payload-by-sha/5ba22653a26444a989ce0942b5fbdb12cda4fa2722f818c8e199cd50843e6565/model_runner.py", + "required_symbols": [ + "GPUModelRunner.execute_model" + ], + "sha256": "5ba22653a26444a989ce0942b5fbdb12cda4fa2722f818c8e199cd50843e6565" + }, + "/usr/local/lib/python3.12/dist-packages/vllm/v1/worker/gpu/model_states/mamba_hybrid.py": { + "payload": "payload-by-sha/90ae5dd89e808c9b13fd24f203feebd56e728b2928c13547e131d1ac7ca35965/mamba_hybrid.py", + "required_symbols": [ + "_checkpoint_plan_rows", + "MambaHybridAttnMetadata.get_extra_common_attn_kwargs", + "MambaHybridModelState.prepare_attn" + ], + "sha256": "90ae5dd89e808c9b13fd24f203feebd56e728b2928c13547e131d1ac7ca35965" + }, + "/usr/local/lib/python3.12/dist-packages/vllm/v1/worker/gpu_model_runner.py": { + "payload": "payload-by-sha/6888b3c5e16b4ef56ddc6d7bb02aa9736894a758574abeebe7e4ff6e9a4ef295/gpu_model_runner.py", + "required_symbols": [ + "GPUModelRunner._update_states", + "GPUModelRunner.execute_model" + ], + "sha256": "6888b3c5e16b4ef56ddc6d7bb02aa9736894a758574abeebe7e4ff6e9a4ef295" + } + }, + "runtime_enablement": { + "SPARK_GDN_PREFILL_CHECKPOINTS": "2" + }, + "runtime_environment_changed": false, + "schema": "sparkring-prefill-two-checkpoint-image-patch/v1", + "source_manifest_sha256": "d183057ad4f4579ab2c264cc21c5ce4d3140b2f8052c0be51f7070e09d567b06", + "status": "research-only" +} diff --git a/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/0d0fc64b5fddd269d823a57f14c09a9e0a1b38a4f802c39df07353f34873a270/_impl.py b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/0d0fc64b5fddd269d823a57f14c09a9e0a1b38a4f802c39df07353f34873a270/_impl.py new file mode 100644 index 00000000..1ad0f96d --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/0d0fc64b5fddd269d823a57f14c09a9e0a1b38a4f802c39df07353f34873a270/_impl.py @@ -0,0 +1,487 @@ +"""Planned contract for chunked KDA prefill: caps, plan, bind, run.""" + +from __future__ import annotations + +import math +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Literal + +import torch + +from b12x._lib.scratch import ScratchBufferSpec, scratch_buffer_spec, scratch_tensor +from b12x._lib.scratch_layout import ( + SCRATCH_ALIGN_BYTES, + align_up, + dtype_nbytes, + materialize_scratch_view, +) +from b12x.policy import PolicyContext, get_auto_policy + +from .._shared.kda_math import KDA_HEAD_DIM +from .._shared.tensors import ( + canonical_device, + overlaps, + positive, + require_paged_recurrent_state, + require_row_contiguous, + require_tensor, +) +from ._policy import ( + CHUNK_TOKENS, + KDA_PREFILL_POLICY, + V_SPLIT_CHOICES, + KdaPrefillQuery, + WorkspaceRecord, + tiles_capacity, +) + +MetadataValidation = Literal["transactional", "trusted"] + + +@dataclass(frozen=True, kw_only=True) +class Caps: + """Static geometry and planned capacity of a KDA prefill plan.""" + + device: torch.device | str + max_tokens: int + max_seqs: int + max_state_slots: int + heads: int + head_dim: int = KDA_HEAD_DIM + model_dtype: torch.dtype = torch.bfloat16 + state_dtype: torch.dtype = torch.float32 + qk_l2norm: bool = True + checkpoint_export: bool = False + max_checkpoints: int = 1 + null_state_index: int | None = None + metadata_validation: MetadataValidation = "transactional" + chunk_tokens: int = 16 + + def __post_init__(self) -> None: + device = canonical_device(self.device) + if device.type != "cuda": + raise ValueError(f"KDA prefill requires a CUDA device, got {device}") + object.__setattr__(self, "device", device) + for name in ("max_tokens", "max_seqs", "max_state_slots", "heads"): + object.__setattr__(self, name, positive(name, getattr(self, name))) + if self.max_seqs > 4096: + raise ValueError("max_seqs must be at most 4096") + if int(self.head_dim) != KDA_HEAD_DIM: + raise ValueError(f"head_dim must be {KDA_HEAD_DIM}, got {self.head_dim}") + object.__setattr__(self, "head_dim", KDA_HEAD_DIM) + if self.model_dtype != torch.bfloat16: + raise ValueError("model_dtype must be torch.bfloat16") + if self.state_dtype != torch.float32: + raise ValueError("state_dtype must be torch.float32") + if self.chunk_tokens != CHUNK_TOKENS: + raise ValueError(f"chunk_tokens must be {CHUNK_TOKENS}") + if self.metadata_validation not in ("transactional", "trusted"): + raise ValueError("metadata_validation must be 'transactional' or 'trusted'") + object.__setattr__(self, "qk_l2norm", bool(self.qk_l2norm)) + object.__setattr__(self, "checkpoint_export", bool(self.checkpoint_export)) + if type(self.max_checkpoints) is not int or self.max_checkpoints not in (1, 2): + raise ValueError("max_checkpoints must be 1 or 2") + if self.max_checkpoints == 2 and (not self.checkpoint_export or self.metadata_validation != "transactional"): + raise ValueError("two checkpoints require checkpoint_export and transactional validation") + if self.null_state_index is not None: + null = int(self.null_state_index) + if null < 0 or null >= self.max_state_slots: + raise ValueError("null_state_index must be a valid slot index") + object.__setattr__(self, "null_state_index", null) + + @property + def tiles_capacity(self) -> int: + """Upper bound on packed chunk tiles: one partial tile per sequence.""" + return tiles_capacity(self.max_tokens, self.max_seqs) + + +@dataclass(frozen=True) +class Plan: + """Fixed launch policy and caller-allocated scratch layout for one Caps. + + Chunk tiles are ordered in bands (local tile index major, sequence rank + minor, longest sequences first) and processed in windows of + ``window_tiles`` consecutive positions, so every live sequence advances + in every window. The prepare kernel writes each window into one of two + workspace ring slots and the recurrence kernel of a window consumes that + slot while the next window is being prepared. A sequence that continues + across a window boundary keeps its running state in its final state slot, + which must therefore not be null. ``max_windows`` is the launch count + that covers the full capacity. + """ + + caps: Caps + v_split: int + k_split: int + stages: int + window_tiles: int + max_windows: int + duplicate_table_size: int + offsets: Mapping[str, int] + _scratch_specs: tuple[ScratchBufferSpec, ...] + policy_resolution: object | None = None + + def scratch_specs(self) -> tuple[ScratchBufferSpec, ...]: + return self._scratch_specs + + def shapes_and_dtypes(self) -> tuple[tuple[tuple[int, ...], torch.dtype], ...]: + return tuple((spec.shape, spec.dtype) for spec in self._scratch_specs) + + @property + def recurrence_rows(self) -> int: + """Grid rows of one recurrence launch: sequences a window can intersect.""" + return min(self.caps.max_seqs, self.window_tiles) + + def launched_windows(self, max_live_tokens: int | None, max_live_seqs: int | None) -> int: + """Windows to launch for a run bounded by the given live counts.""" + if max_live_tokens is None and max_live_seqs is None: + return self.max_windows + tokens = self.caps.max_tokens if max_live_tokens is None else int(max_live_tokens) + seqs = self.caps.max_seqs if max_live_seqs is None else int(max_live_seqs) + if tokens < 0 or tokens > self.caps.max_tokens: + raise ValueError(f"max_live_tokens={tokens} exceeds capacity {self.caps.max_tokens}") + if seqs < 0 or seqs > self.caps.max_seqs: + raise ValueError(f"max_live_seqs={seqs} exceeds capacity {self.caps.max_seqs}") + tiles = tiles_capacity(tokens, seqs) + return max(1, min(self.max_windows, -(-tiles // self.window_tiles))) + + def output_shape(self, tokens: int | None = None) -> tuple[int, int, int]: + live_tokens = self.caps.max_tokens if tokens is None else int(tokens) + if live_tokens < 0 or live_tokens > self.caps.max_tokens: + raise ValueError(f"tokens={live_tokens} exceeds capacity {self.caps.max_tokens}") + return (live_tokens, self.caps.heads, KDA_HEAD_DIM) + + def bind(self, **kwargs) -> "Binding": + return bind(self, **kwargs) + + +@dataclass(frozen=True) +class Binding: + """Caller-owned tensors and scratch views for one prefill invocation.""" + + plan: Plan + scratch: torch.Tensor + error_code: torch.Tensor + duplicate_slots: torch.Tensor + band_base: torch.Tensor + sorted_seq: torch.Tensor + rank_of: torch.Tensor + pos_seq: torch.Tensor + pos_local: torch.Tensor + window_table: torch.Tensor + ready_flags: torch.Tensor + ws: torch.Tensor + q: torch.Tensor + k: torch.Tensor + v: torch.Tensor + raw_g: torch.Tensor + raw_beta: torch.Tensor + A_log: torch.Tensor + dt_bias: torch.Tensor + recurrent_state: torch.Tensor + cu_seqlens: torch.Tensor + initial_state_indices: torch.Tensor + final_state_indices: torch.Tensor + checkpoint_state_indices: torch.Tensor + checkpoint_offsets: torch.Tensor + num_seqs: torch.Tensor + num_tokens: torch.Tensor + output: torch.Tensor + token_capacity: int + seq_capacity: int + + +def _next_power_of_two(value: int) -> int: + return 1 << max(0, int(value) - 1).bit_length() + + +def _query(caps: Caps) -> KdaPrefillQuery: + return KdaPrefillQuery( + heads=caps.heads, + head_dim=caps.head_dim, + model_dtype=str(caps.model_dtype).removeprefix("torch."), + state_dtype=str(caps.state_dtype).removeprefix("torch."), + qk_l2norm=caps.qk_l2norm, + checkpoint_export=caps.checkpoint_export, + max_checkpoints=caps.max_checkpoints, + max_tokens=caps.max_tokens, + max_seqs=caps.max_seqs, + ) + + +def _materialize_plan( + caps: Caps, + *, + v_split: int, + k_split: int, + stages: int, + window_tiles: int, + policy_resolution: object | None, +) -> Plan: + if v_split not in V_SPLIT_CHOICES: + raise ValueError(f"v_split must be one of {V_SPLIT_CHOICES}, got {v_split}") + tiles = caps.tiles_capacity + heads = caps.heads + window_tiles = max(1, min(int(window_tiles), tiles)) + max_windows = -(-tiles // window_tiles) + ring_records = 2 * window_tiles * heads + duplicate_table_size = _next_power_of_two(2 * (1 + caps.max_checkpoints) * caps.max_seqs) + regions = ( + ("error_code", 1, torch.int32), + ("duplicate_slots", duplicate_table_size, torch.int32), + ("band_base", tiles + 2, torch.int32), + ("sorted_seq", caps.max_seqs, torch.int32), + ("rank_of", caps.max_seqs, torch.int32), + ("pos_seq", tiles, torch.int32), + ("pos_local", tiles, torch.int32), + ("window_table", 2 * max_windows, torch.int32), + ("ready_flags", ring_records, torch.int32), + ("ws", ring_records * WorkspaceRecord.BYTES, torch.uint8), + ) + offsets: dict[str, int] = {} + cursor = 0 + for name, elements, dtype in regions: + cursor = align_up(cursor, SCRATCH_ALIGN_BYTES) + offsets[name] = cursor + cursor += elements * dtype_nbytes(dtype) + spec = scratch_buffer_spec("kda_prefill", nbytes=cursor, device=caps.device) + return Plan( + caps=caps, + v_split=int(v_split), + k_split=int(k_split), + stages=int(stages), + window_tiles=window_tiles, + max_windows=max_windows, + duplicate_table_size=duplicate_table_size, + offsets=offsets, + _scratch_specs=(spec,), + policy_resolution=policy_resolution, + ) + + +def plan(caps: Caps, *, policy: PolicyContext | None = None) -> Plan: + """Resolve the policy once and lay out the scratch for ``caps``.""" + if not isinstance(caps, Caps): + raise TypeError("caps must be kda_prefill.Caps") + policy = policy or get_auto_policy(caps.device) + if not isinstance(policy, PolicyContext): + raise TypeError("policy must be a PolicyContext") + policy.require_device(caps.device) + resolution = policy.resolve(KDA_PREFILL_POLICY, _query(caps)) + return _materialize_plan( + caps, + v_split=int(resolution.config.v_split), + k_split=int(resolution.config.k_split), + stages=int(resolution.config.stages), + window_tiles=int(resolution.config.window_tiles), + policy_resolution=resolution, + ) + + +def _record_view( + storage: torch.Tensor, plan: Plan, name: str, shape: tuple[int, ...], dtype: torch.dtype +) -> torch.Tensor: + view, _ = materialize_scratch_view( + storage, offset_bytes=plan.offsets[name], shape=shape, dtype=dtype + ) + return view + + +def bind( + plan: Plan, + *, + scratch: torch.Tensor | Mapping[str, torch.Tensor] | Sequence[torch.Tensor], + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + recurrent_state: torch.Tensor, + cu_seqlens: torch.Tensor, + initial_state_indices: torch.Tensor, + final_state_indices: torch.Tensor, + checkpoint_state_indices: torch.Tensor, + checkpoint_offsets: torch.Tensor, + num_seqs: torch.Tensor, + num_tokens: torch.Tensor, + output: torch.Tensor, +) -> Binding: + """Bind live tensors to a plan without allocating or launching work. + + Live capacities come from the bound tensors: ``q.shape[0]`` tokens and + ``cu_seqlens.numel() - 1`` sequences, each at most the planned capacity. + """ + if not isinstance(plan, Plan): + raise TypeError("plan must be kda_prefill.Plan") + caps = plan.caps + device = caps.device + heads = caps.heads + tiles = caps.tiles_capacity + token_capacity = positive("q token capacity", q.shape[0]) if q.dim() == 3 else 0 + if token_capacity > caps.max_tokens: + raise ValueError(f"token capacity {token_capacity} exceeds planned {caps.max_tokens}") + seq_capacity = int(cu_seqlens.numel()) - 1 + if seq_capacity < 1 or seq_capacity > caps.max_seqs: + raise ValueError( + f"cu_seqlens must hold 2..{caps.max_seqs + 1} entries, got {cu_seqlens.numel()}" + ) + row_shape = (token_capacity, heads, KDA_HEAD_DIM) + for name, tensor in (("q", q), ("k", k), ("v", v), ("raw_g", raw_g)): + require_row_contiguous(name, tensor, shape=row_shape, device=device, dtypes=(torch.bfloat16,)) + require_tensor( + "raw_beta", raw_beta, shape=(token_capacity, heads), device=device, + dtypes=(torch.bfloat16,), contiguous=False, + ) + require_tensor("A_log", A_log, shape=(heads,), device=device, dtypes=(torch.bfloat16, torch.float32)) + require_tensor( + "dt_bias", dt_bias, shape=(heads, KDA_HEAD_DIM), device=device, + dtypes=(torch.bfloat16, torch.float32), + ) + require_paged_recurrent_state( + recurrent_state, + shape=(caps.max_state_slots, heads, KDA_HEAD_DIM, KDA_HEAD_DIM), + device=device, + dtype=caps.state_dtype, + ) + require_tensor("cu_seqlens", cu_seqlens, shape=(seq_capacity + 1,), device=device, dtypes=(torch.int32,)) + index_dtypes = (torch.int32, torch.int64) + for name, tensor in ( + ("initial_state_indices", initial_state_indices), + ): + require_tensor(name, tensor, shape=(seq_capacity,), device=device, dtypes=index_dtypes) + require_tensor( + "final_state_indices", + final_state_indices, + shape=(seq_capacity,), + device=device, + dtypes=index_dtypes, + contiguous=False, + ) + if final_state_indices.stride(0) <= 0: + raise ValueError("final_state_indices must have a positive stride") + if not (initial_state_indices.dtype == final_state_indices.dtype == checkpoint_state_indices.dtype): + raise TypeError("state index tensors must share one dtype") + checkpoint_shape = (seq_capacity,) if caps.max_checkpoints == 1 else (seq_capacity, 2) + require_tensor("checkpoint_state_indices", checkpoint_state_indices, shape=checkpoint_shape, + device=device, dtypes=index_dtypes) + require_tensor( + "checkpoint_offsets", checkpoint_offsets, shape=checkpoint_shape, device=device, dtypes=(torch.int32,) + ) + for name, tensor in (("num_seqs", num_seqs), ("num_tokens", num_tokens)): + require_tensor(name, tensor, shape=(1,), device=device, dtypes=(torch.int32,)) + require_row_contiguous("output", output, shape=row_shape, device=device, dtypes=(torch.bfloat16,)) + + storage = scratch_tensor(scratch, plan.scratch_specs(), owner="KDA prefill") + ring = 2 * plan.window_tiles + views = { + "error_code": _record_view(storage, plan, "error_code", (1,), torch.int32), + "duplicate_slots": _record_view(storage, plan, "duplicate_slots", (plan.duplicate_table_size,), torch.int32), + "band_base": _record_view(storage, plan, "band_base", (tiles + 2,), torch.int32), + "sorted_seq": _record_view(storage, plan, "sorted_seq", (caps.max_seqs,), torch.int32), + "rank_of": _record_view(storage, plan, "rank_of", (caps.max_seqs,), torch.int32), + "pos_seq": _record_view(storage, plan, "pos_seq", (tiles,), torch.int32), + "pos_local": _record_view(storage, plan, "pos_local", (tiles,), torch.int32), + "window_table": _record_view(storage, plan, "window_table", (plan.max_windows, 2), torch.int32), + "ready_flags": _record_view(storage, plan, "ready_flags", (ring, heads), torch.int32), + "ws": _record_view(storage, plan, "ws", (ring, heads, WorkspaceRecord.BYTES), torch.uint8), + } + mutable = {"scratch": storage, "recurrent_state": recurrent_state, "output": output} + read_only = { + "q": q, "k": k, "v": v, "raw_g": raw_g, "raw_beta": raw_beta, "A_log": A_log, + "dt_bias": dt_bias, "cu_seqlens": cu_seqlens, + "initial_state_indices": initial_state_indices, + "final_state_indices": final_state_indices, + "checkpoint_state_indices": checkpoint_state_indices, + "checkpoint_offsets": checkpoint_offsets, "num_seqs": num_seqs, "num_tokens": num_tokens, + } + names = list(mutable) + for left in range(len(names)): + for right in range(left + 1, len(names)): + if overlaps(mutable[names[left]], mutable[names[right]]): + raise ValueError(f"{names[left]} and {names[right]} must not overlap") + for name, tensor in mutable.items(): + for other, candidate in read_only.items(): + if overlaps(tensor, candidate): + raise ValueError(f"{name} must not overlap read-only tensor {other}") + return Binding( + plan=plan, + scratch=storage, + **views, + q=q, k=k, v=v, raw_g=raw_g, raw_beta=raw_beta, A_log=A_log, dt_bias=dt_bias, + recurrent_state=recurrent_state, cu_seqlens=cu_seqlens, + initial_state_indices=initial_state_indices, + final_state_indices=final_state_indices, + checkpoint_state_indices=checkpoint_state_indices, + checkpoint_offsets=checkpoint_offsets, num_seqs=num_seqs, num_tokens=num_tokens, + output=output, token_capacity=token_capacity, seq_capacity=seq_capacity, + ) + + +def _check_run_scalars(lower_bound: float, scale: float | None, eps: float) -> tuple[float, float, float]: + lower_bound_value = float(lower_bound) + if not math.isfinite(lower_bound_value) or not -5.0 <= lower_bound_value < 0.0: + raise ValueError(f"lower_bound must be in [-5, 0), got {lower_bound_value}") + scale_value = KDA_HEAD_DIM**-0.5 if scale is None else float(scale) + if not math.isfinite(scale_value) or scale_value <= 0.0: + raise ValueError(f"scale must be finite and positive, got {scale_value}") + eps_value = float(eps) + if not math.isfinite(eps_value) or eps_value <= 0.0: + raise ValueError(f"eps must be finite and positive, got {eps_value}") + return lower_bound_value, scale_value, eps_value + + +def run( + binding: Binding, + *, + lower_bound: float, + scale: float | None = None, + eps: float = 1e-6, + max_live_tokens: int | None = None, + max_live_seqs: int | None = None, +) -> torch.Tensor: + """Run the prologue, prepare, and recurrence kernels; capture safe. + + ``max_live_tokens`` and ``max_live_seqs`` are optional host-side upper + bounds on the device counts; they only limit how many pipeline windows are + launched. Under transactional validation a run whose live tiles exceed the + launched windows fails closed like any other malformed metadata; under + trusted validation the bounds are part of the caller's contract. + + A sequence whose tiles span more than one pipeline window keeps its + running state in its final state slot between windows, so such a + sequence must have a non-null final slot (transactional validation flags + a null one as an invalid slot). + """ + if not isinstance(binding, Binding): + raise TypeError("binding must be kda_prefill.Binding") + lower_bound_value, scale_value, eps_value = _check_run_scalars(lower_bound, scale, eps) + windows = binding.plan.launched_windows(max_live_tokens, max_live_seqs) + from ._cute_kernels import run_prefill + + run_prefill( + binding, lower_bound=lower_bound_value, scale=scale_value, eps=eps_value, windows=windows + ) + return binding.output + + +def prewarm(binding: Binding) -> None: + """Compile every kernel specialization of ``binding`` without launching.""" + if not isinstance(binding, Binding): + raise TypeError("binding must be kda_prefill.Binding") + from ._cute_kernels import prewarm_binding + + prewarm_binding(binding) + + +__all__ = [ + "Binding", + "Caps", + "Plan", + "bind", + "plan", + "prewarm", + "run", +] diff --git a/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/0e1d3e4e1284fc2ad5b0cde9be9884eec61186e7731750f22ec25d828bb4e0f1/output.py b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/0e1d3e4e1284fc2ad5b0cde9be9884eec61186e7731750f22ec25d828bb4e0f1/output.py new file mode 100644 index 00000000..66b3ae8d --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/0e1d3e4e1284fc2ad5b0cde9be9884eec61186e7731750f22ec25d828bb4e0f1/output.py @@ -0,0 +1,324 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from dataclasses import dataclass +from functools import cached_property +from typing import TYPE_CHECKING + +from vllm.config.ec_manager_config import EncoderCacheManagerMetadata +from vllm.multimodal.utils import strip_covered_mm_data + +if TYPE_CHECKING: + import numpy as np + import numpy.typing as npt + import torch + + from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorMetadata + from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata + from vllm.lora.request import LoRARequest + from vllm.multimodal.inputs import MultiModalFeatureSpec + from vllm.pooling_params import PoolingParams + from vllm.sampling_params import SamplingParams + from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy + from vllm.v1.request import Request +else: + ECConnectorMetadata = object + KVConnectorMetadata = object + KVCacheBlockCopy = object + LoRARequest = object + MultiModalFeatureSpec = object + PoolingParams = object + SamplingParams = object + Request = object + + +@dataclass +class NewRequestData: + req_id: str + prompt_token_ids: list[int] | None + mm_features: list[MultiModalFeatureSpec] + sampling_params: SamplingParams | None + pooling_params: PoolingParams | None + block_ids: tuple[list[int], ...] + num_computed_tokens: int + lora_request: LoRARequest | None + prompt_embeds: "torch.Tensor | None" = None + prompt_is_token_ids: list[bool] | None = None + + # Only used for v2 model runner. + prefill_token_ids: list[int] | None = None + + @classmethod + def from_request( + cls, + request: Request, + block_ids: tuple[list[int], ...], + prefill_token_ids: list[int] | None = None, + uses_mrope: bool = False, + uses_xdrope: bool = False, + ) -> "NewRequestData": + return cls( + req_id=request.request_id, + prompt_token_ids=request.prompt_token_ids, + mm_features=strip_covered_mm_data( + request.mm_features, + request.num_computed_tokens, + uses_mrope=uses_mrope, + uses_xdrope=uses_xdrope, + ), + sampling_params=request.sampling_params, + pooling_params=request.pooling_params, + block_ids=block_ids, + num_computed_tokens=request.num_computed_tokens, + lora_request=request.lora_request, + prompt_embeds=request.prompt_embeds, + prompt_is_token_ids=request.prompt_is_token_ids, + prefill_token_ids=prefill_token_ids, + ) + + @property + def prompt_len(self) -> int: + if self.prompt_token_ids is not None: + return len(self.prompt_token_ids) + if self.prompt_embeds is not None: + return self.prompt_embeds.shape[0] + return 0 + + def __repr__(self) -> str: + prompt_embeds_shape = ( + self.prompt_embeds.shape if self.prompt_embeds is not None else None + ) + return ( + f"NewRequestData(" + f"req_id={self.req_id}," + f"prompt_token_ids={self.prompt_token_ids}," + f"prefill_token_ids={self.prefill_token_ids}," + f"mm_features={self.mm_features}," + f"sampling_params={self.sampling_params}," + f"block_ids={self.block_ids}," + f"num_computed_tokens={self.num_computed_tokens}," + f"lora_request={self.lora_request}," + f"prompt_embeds_shape={prompt_embeds_shape}" + ")" + ) + + # Version of __repr__ with the prompt data obfuscated + def anon_repr(self) -> str: + prompt_token_ids_len = ( + len(self.prompt_token_ids) if self.prompt_token_ids is not None else None + ) + prompt_embeds_shape = ( + self.prompt_embeds.shape if self.prompt_embeds is not None else None + ) + prefill_token_ids_len = ( + len(self.prefill_token_ids) if self.prefill_token_ids is not None else None + ) + return ( + f"NewRequestData(" + f"req_id={self.req_id}," + f"prompt_token_ids_len={prompt_token_ids_len}," + f"prefill_token_ids_len={prefill_token_ids_len}," + f"mm_features={self.mm_features}," + f"sampling_params={self.sampling_params}," + f"block_ids={self.block_ids}," + f"num_computed_tokens={self.num_computed_tokens}," + f"lora_request={self.lora_request}," + f"prompt_embeds_shape={prompt_embeds_shape}" + ")" + ) + + +@dataclass +class CachedRequestData: + req_ids: list[str] + # For request ids not in resumed_req_ids, new_block_ids will be appended to + # the request's block IDs. For those in the set, new_block_ids will be used as the + # request's block IDs instead of appending to the existing block IDs. + resumed_req_ids: set[str] + # NOTE(woosuk): new_token_ids is only used for pipeline parallelism. + # When PP is not used, new_token_ids will be empty. + new_token_ids: list[list[int]] + # MRV1-only: For requests not scheduled in the last step, propagate the token ids + # to the connector. Won't contain requests scheduled in the prior step. + all_token_ids: dict[str, list[int]] + new_block_ids: list[tuple[list[int], ...] | None] + num_computed_tokens: list[int] + num_output_tokens: list[int] + + # Version of dataclass repr with token IDs obfuscated. + def anon_repr(self) -> str: + new_token_ids_lens = [len(toks) for toks in self.new_token_ids] + all_token_ids_lens = { + req_id: len(toks) for req_id, toks in self.all_token_ids.items() + } + return ( + f"CachedRequestData(" + f"req_ids={self.req_ids}," + f"resumed_req_ids={self.resumed_req_ids}," + f"new_token_ids_lens={new_token_ids_lens}," + f"all_token_ids_lens={all_token_ids_lens}," + f"new_block_ids={self.new_block_ids}," + f"num_computed_tokens={self.num_computed_tokens}," + f"num_output_tokens={self.num_output_tokens}" + f")" + ) + + def __repr__(self) -> str: + return self.anon_repr() + + @property + def num_reqs(self) -> int: + return len(self.req_ids) + + @cached_property + def _req_id_to_num_output_tokens(self) -> dict[str, int]: + """Cache mapping of req_id to num_output_tokens for O(1) lookup. + + This cached property is safe because CachedRequestData instances + are created fresh each scheduling iteration and not mutated during + computation of iteration details. + """ + return dict(zip(self.req_ids, self.num_output_tokens)) + + def is_context_phase(self, req_id: str) -> bool: + num_output_tokens = self._req_id_to_num_output_tokens.get(req_id) + return num_output_tokens is not None and num_output_tokens == 0 + + @classmethod + def make_empty(cls) -> "CachedRequestData": + return cls( + req_ids=[], + resumed_req_ids=set(), + new_token_ids=[], + all_token_ids={}, + new_block_ids=[], + num_computed_tokens=[], + num_output_tokens=[], + ) + + +@dataclass +class ScheduledEncoderInputStats: + """Stats for encoder inputs scheduled in one iteration.""" + + num_inputs: int = 0 + output_tokens: int = 0 + + +@dataclass +class SchedulerOutput: + # list of the requests that are scheduled for the first time. + # We cache the request's data in each worker process, so that we don't + # need to re-send it every scheduling step. + scheduled_new_reqs: list[NewRequestData] + # list of the requests that have been scheduled before. + # Since the request's data is already cached in the worker processes, + # we only send the diff to minimize the communication cost. + scheduled_cached_reqs: CachedRequestData + + # req_id -> num_scheduled_tokens + # Number of tokens scheduled for each request. + num_scheduled_tokens: dict[str, int] + # Total number of tokens scheduled for all requests. + # Equal to sum(num_scheduled_tokens.values()) + total_num_scheduled_tokens: int + # req_id -> spec_token_ids + # If a request does not have any spec decode tokens, it will not be + # included in the dictionary. + scheduled_spec_decode_tokens: dict[str, list[int]] + # req_id -> encoder input indices that need processing. + # E.g., if a request has [0, 1], it could mean the vision encoder needs + # to process that the request's 0-th and 1-th images in the current step. + scheduled_encoder_inputs: dict[str, list[int]] + # Number of common prefix blocks for all requests in each KV cache group. + # This can be used for cascade attention. + num_common_prefix_blocks: list[int] + + # Request IDs that are finished in between the previous and the current + # steps. This is used to notify the workers about the finished requests + # so that they can free the cached states for those requests. + finished_req_ids: set[str] + # list of mm_hash strings associated with the encoder outputs to be + # freed from the encoder cache. + free_encoder_mm_hashes: list[str] + + scheduled_encoder_input_stats: ScheduledEncoderInputStats | None = None + + # Request IDs that are preempted in this step. + # Only used for v2 model runner. + preempted_req_ids: set[str] | None = None + + # Whether any of the scheduled requests use structured output. + # Set only in async scheduling case. + has_structured_output_requests: bool = False + + # Whether the scheduled requests have all the output tokens they + # need to perform grammar bitmask computation. + pending_structured_output_tokens: bool = False + + # Used for adjusting acceptance rate calculation. + num_invalid_spec_tokens: dict[str, int] | None = None + + # KV Cache Connector metadata. + kv_connector_metadata: KVConnectorMetadata | None = None + + # EC Cache Connector metadata + ec_connector_metadata: ECConnectorMetadata | None = None + # EC Cache Manager metadata + ec_manager_metadata: EncoderCacheManagerMetadata | None = None + # Block IDs freshly allocated from the pool during this scheduling step. + # The worker zeros the corresponding GPU memory before the blocks are used, + # preventing stale NaN/data from corrupting attention or SSM computation. + new_block_ids_to_zero: list[int] | None = None + + # CoW copies to apply after zeroing new blocks and before forward. + kv_cache_block_copies: list[KVCacheBlockCopy] | None = None + + # Producer partial-tail offload hand-off for external KV connectors: + # {request_id: [(group_id, block_id, boundary_tokens), ...]} pointing at + # the durable boundary block of a producer's last-prompt-boundary partial + # tail (mamba "align" CoW target). None unless partial hash hits are active. + partial_tail_offloads: dict[str, list[tuple[int, int, int]]] | None = None + + # Hash-proven recurrent-state blocks for external KV connectors: + # {request_id: [(group_id, block_id, boundary_tokens), ...]}. The union + # includes full aligned Mamba pages and the partial-tail CoW targets above. + # Consumers opting into aligned boundaries must finish their worker-side + # snapshot before request cleanup releases the scheduler pin. + recurrent_boundary_blocks: dict[str, list[tuple[int, int, int]]] | None = None + + # Explicit absolute checkpoint positions, keyed by request identity. The + # worker validates (start, end) against its actual packed query before use. + recurrent_prefill_checkpoint_plans: dict[str, tuple[int, int, tuple[int, ...]]] | None = None + + # Dynamic speculative decoding: optimal K chosen by scheduler. + # Number of spec tokens to schedule for the next step. + num_spec_tokens_to_schedule: int | None = None + + @classmethod + def make_empty(cls) -> "SchedulerOutput": + return cls( + scheduled_new_reqs=[], + scheduled_cached_reqs=CachedRequestData.make_empty(), + num_scheduled_tokens={}, + total_num_scheduled_tokens=0, + scheduled_spec_decode_tokens={}, + scheduled_encoder_inputs={}, + num_common_prefix_blocks=[], + finished_req_ids=set(), + free_encoder_mm_hashes=[], + ) + + def resolve_num_spec_tokens_to_schedule(self, default: int) -> int: + """Resolve the speculative depth for real and synthetic outputs.""" + if self.num_spec_tokens_to_schedule is None: + return default + return self.num_spec_tokens_to_schedule + + +@dataclass +class GrammarOutput: + # ids of structured output requests. + structured_output_request_ids: list[str] + # Bitmask ordered as structured_output_request_ids. + grammar_bitmask: "npt.NDArray[np.int32]" diff --git a/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/10a6a60f31c67b59825aa827084f3d25416128ba5dc7c14a774675d64bbec40b/single_type_kv_cache_manager.py b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/10a6a60f31c67b59825aa827084f3d25416128ba5dc7c14a774675d64bbec40b/single_type_kv_cache_manager.py new file mode 100644 index 00000000..1f606d7e --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/10a6a60f31c67b59825aa827084f3d25416128ba5dc7c14a774675d64bbec40b/single_type_kv_cache_manager.py @@ -0,0 +1,2087 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import itertools +from abc import ABC, abstractmethod +from collections import defaultdict +from collections.abc import Sequence +from typing import ClassVar + +from vllm.utils.math_utils import cdiv +from vllm.v1.core.block_pool import BlockPool +from vllm.v1.core.kv_cache_utils import ( + BlockHashList, + BlockHashListWithBlockSize, + BlockHashWithGroupId, + KVCacheBlock, + get_group_id, + resolve_block_hashes, +) +from vllm.v1.kv_cache_interface import ( + AttentionSpec, + ChunkedLocalAttentionSpec, + CrossAttentionSpec, + FullAttentionSpec, + HiddenStateCacheSpec, + KVCacheSpec, + MambaSpec, + MLAAttentionSpec, + RSWASpec, + SinkFullAttentionSpec, + SlidingWindowMLASpec, + SlidingWindowSpec, +) +from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry +from vllm.v1.request import Request + + +class SingleTypeKVCacheManager(ABC): + """ + An abstract base class for a manager that handle the kv cache management + logic of one specific type of attention layer. + """ + + supports_fine_grained_hash_lookup: ClassVar[bool] = False + + def __init__( + self, + kv_cache_spec: KVCacheSpec, + block_pool: BlockPool, + enable_caching: bool, + kv_cache_group_id: int, + scheduler_block_size: int, + dcp_world_size: int = 1, + pcp_world_size: int = 1, + needs_kv_cache_zeroing: bool = False, + max_admission_blocks_per_request: int | None = None, + ) -> None: + """ + Initializes the SingleTypeKVCacheManager. + Args: + kv_cache_spec: The kv_cache_spec for this manager. + block_pool: The block pool. + kv_cache_group_id: The id of the kv cache group of this manager. + scheduler_block_size: The scheduling granularity (LCM of all group + block sizes); a multiple of this manager's ``block_size``. + needs_kv_cache_zeroing: Whether worker-side KV cache zeroing needs + newly allocated block IDs from this manager. + max_admission_blocks_per_request: Recycling-aware per-request + block cap used by `get_num_blocks_to_allocate`. Only set for + spec types that recycle blocks across chunks (SWA, + chunked-local); `None` (the default) means no cap, which is + correct for full-attention-style specs that hold every + block until the request finishes. + """ + self.scheduler_block_size = scheduler_block_size + # The block size for this manager; used for actual block allocation. + self.block_size = kv_cache_spec.block_size + self.dcp_world_size = dcp_world_size + self.pcp_world_size = pcp_world_size + if dcp_world_size > 1 and not getattr(kv_cache_spec, "dcp_replicated", False): + self.block_size *= dcp_world_size + self.kv_cache_spec = kv_cache_spec + self.block_pool = block_pool + self.enable_caching = enable_caching + self._max_admission_blocks_per_request = max_admission_blocks_per_request + # Record newly allocated block ids only when worker-side zeroing will + # consume them and this manager holds a spec type that gets zeroed. + self._record_new_block_ids = needs_kv_cache_zeroing and isinstance( + kv_cache_spec, AttentionSpec + ) + self.new_block_ids: list[int] = [] + + # Mapping from request ID to blocks to track the blocks allocated + # for each request, so that we can free the blocks when the request + # is finished. + self.req_to_blocks: defaultdict[str, list[KVCacheBlock]] = defaultdict(list) + + # {req_id: The number of cached blocks for this given request} + # This is used to track the number of cached blocks for each request. + # This is only used to track the RUNNING requests, we do not track the + # data for preempted ones. + self.num_cached_block: dict[str, int] = {} + + self.kv_cache_group_id = kv_cache_group_id + self._null_block = block_pool.null_block + + # Whether this group's prefix-cache hits drop the EAGLE/MTP lookahead + # block. Only consulted by managers whose hit logic is sparse within an + # aligned segment (SWA). Initialized lazily by the coordinator after + # determining the attention groups. + self.use_eagle = False + # Request-local external publication target, set only around one cache + # operation. Native prefix-cache hashes retain block_pool.hash_block_size. + self.recurrent_publication_boundary: int | None = None + self._planned_recurrent_checkpoints: dict[str, tuple[int, int, tuple[int, ...]]] = {} + self._planned_recurrent_publications: dict[str, tuple[int, ...]] = {} + + # Partial-hit copy-on-write bookkeeping. Populated only by fine-grained + # managers (full attention, mamba "align"); harmlessly empty elsewhere. + self._partial_hit_reqs: dict[str, tuple[int, KVCacheBlock]] = {} + self._pending_cow_copies: list[tuple[KVCacheBlock, KVCacheBlock]] = [] + # Partial-tail offload hand-off for external KV connectors: when a + # producer registers its last-prompt-boundary partial tail and the + # durable boundary block is not on the append-only request block table + # (mamba "align" CoW target), record the request, group, block, and + # exact token boundary so a connector can offload it under the right + # hash. Populated only by mamba "align". + self._pending_partial_tail_offloads: list[ + tuple[str, int, KVCacheBlock, int, BlockHashWithGroupId] + ] = [] + self._pending_aligned_recurrent_boundaries: list[ + tuple[str, int, KVCacheBlock, int] + ] = [] + + @classmethod + def _get_num_evictable_blocks(cls, blocks: Sequence[KVCacheBlock]): + return sum(blk.ref_cnt == 0 and not blk.is_null for blk in blocks) + + def _has_partial_local_hit( + self, + new_computed_blocks: Sequence[KVCacheBlock], + num_local_computed_tokens: int, + ) -> bool: + # The local prefix-cache hit ends inside one of this manager's + # blocks: the shared tail block needs CoW. + return ( + len(new_computed_blocks) > 0 + and num_local_computed_tokens % self.block_size != 0 + ) + + def get_num_blocks_to_allocate( + self, + request_id: str, + num_tokens: int, + new_computed_blocks: Sequence[KVCacheBlock], + total_computed_tokens: int, + num_local_computed_tokens: int, + num_tokens_main_model: int, + apply_admission_cap: bool = False, + ) -> int: + """ + Get the number of blocks needed to be allocated for the request. + + Args: + request_id: The request ID. + num_tokens: The total number of tokens that need a slot (including + tokens that are already allocated). + new_computed_blocks: The new computed blocks just hitting the + prefix caching. + total_computed_tokens: Include both local and external computed + tokens. + num_local_computed_tokens: The number of local prefix-cache computed + tokens. + num_tokens_main_model: The number of tokens for the main model (aka target + model in spec decode). w/o spec decode, it is num_tokens; + with spec decode, it is num_tokens - num_lookahead_tokens. + apply_admission_cap: If True, clamp by `num_required_blocks` by + `_max_admission_blocks_per_request`for recycling-aware specs + (SWA, chunked-local). + + Returns: + The number of blocks to allocate. + """ + + num_required_blocks = cdiv(num_tokens, self.block_size) + if apply_admission_cap and self._max_admission_blocks_per_request is not None: + # Recycling-aware specs (SWA, chunked-local) cap the per-request + # reservation here so admission matches the startup pool sizer + # (`SlidingWindowSpec.max_admission_blocks_per_request` / its + # chunked-local counterpart). `remove_skipped_blocks` runs from + # `allocate_slots` before each chunk's `get_num_blocks_to_allocate`, + # so per-request peak real-held blocks <= this cap, which keeps + # `sum(reservations) <= pool` <=> `sum(peak_real_held) <= pool`. + # Drift between the two would re-introduce the deadlock from + # issue #39734 or, worse, mid-prefill OOM. + num_required_blocks = min( + num_required_blocks, self._max_admission_blocks_per_request + ) + num_req_blocks = len(self.req_to_blocks.get(request_id, ())) + + if request_id in self.num_cached_block: + # Fast-path: a running request won't have any new prefix-cache hits. + assert len(new_computed_blocks) == 0 + # NOTE: With speculative decoding, request's blocks may be allocated + # for draft tokens which are later rejected. In this case, + # num_required_blocks may be smaller than num_req_blocks. + return max(num_required_blocks - num_req_blocks, 0) + + num_skipped_tokens = self.get_num_skipped_tokens(total_computed_tokens) + num_local_computed_blocks = len(new_computed_blocks) + num_req_blocks + # Number of whole blocks that are skipped by the attention window. + # If nothing is skipped, this is 0. + num_skipped_blocks = num_skipped_tokens // self.block_size + # We need blocks for the non-skipped suffix. If there are still + # local-computed blocks inside the window, they contribute to the + # required capacity; otherwise, skipped blocks dominate. + num_new_blocks = max( + num_required_blocks - max(num_skipped_blocks, num_local_computed_blocks), + 0, + ) + + # Among the `new_computed_blocks`, the first `num_skipped_blocks` worth + # of blocks are skipped; `num_req_blocks` of those may already be in + # `req_to_blocks`, so only skip the remainder from `new_computed_blocks`. + num_skipped_new_computed_blocks = max(0, num_skipped_blocks - num_req_blocks) + + # If a computed block is an eviction candidate (in the free queue and + # ref_cnt == 0), it will be removed from the free queue when touched by + # the allocated request, so we must count it in the free-capacity check. + num_evictable_blocks = self._get_num_evictable_blocks( + new_computed_blocks[num_skipped_new_computed_blocks:] + ) + if self._has_partial_local_hit(new_computed_blocks, num_local_computed_tokens): + # Reserve the extra block that allocate_new_blocks pulls for the + # partial-hit CoW redirect. + num_new_blocks += 1 + return num_new_blocks + num_evictable_blocks + + def add_local_computed_blocks( + self, + request_id: str, + new_computed_blocks: Sequence[KVCacheBlock], + num_local_computed_tokens: int, + num_external_computed_tokens: int, + ) -> None: + """ + Add the locally cached (prefix-hit) blocks to the request: + 1. Touch the computed blocks (paired with adding them to `req_blocks`) + so their ref_cnt exactly tracks the referencing requests. + 1.5. (Optional) For sliding window, skipped blocks are padded with nulls. + 2. Add the remaining computed blocks. + + Args: + request_id: The request ID. + new_computed_blocks: The new computed blocks just hitting the + prefix cache. + num_local_computed_tokens: The number of local computed tokens. + num_external_computed_tokens: The number of external computed tokens. + """ + # The coordinator only calls this for first-time allocations (running + # requests are short-circuited there), so the request has no blocks yet. + req_blocks = self.req_to_blocks[request_id] + assert len(req_blocks) == 0 + num_total_computed_tokens = ( + num_local_computed_tokens + num_external_computed_tokens + ) + num_skipped_tokens = self.get_num_skipped_tokens(num_total_computed_tokens) + num_skipped_blocks = num_skipped_tokens // self.block_size + if num_skipped_blocks > 0: + # It is possible that all new computed blocks are skipped when + # num_skipped_blocks > len(new_computed_blocks). + new_computed_blocks = new_computed_blocks[num_skipped_blocks:] + + # Touch the computed blocks to make sure they won't be evicted. + if self.enable_caching: + self.block_pool.touch(new_computed_blocks) + else: + assert not any(new_computed_blocks), ( + "Computed blocks should be empty when prefix caching is disabled" + ) + + # Skip blocks are padded with null blocks. + req_blocks.extend([self._null_block] * num_skipped_blocks) + # Add the remaining computed blocks. + req_blocks.extend(new_computed_blocks) + # All cached hits (including skipped nulls) are already cached; mark + # them so cache_blocks() will not try to re-cache blocks that already + # have a block_hash set. + self.num_cached_block[request_id] = len(req_blocks) + if self._has_partial_local_hit(new_computed_blocks, num_local_computed_tokens): + # Record the partial tail for the CoW redirect in + # allocate_new_blocks; cap the cached count at the full blocks so + # cache_blocks() re-caches the private copy once full. + block_idx = num_local_computed_tokens // self.block_size + self._partial_hit_reqs[request_id] = (block_idx, new_computed_blocks[-1]) + self.num_cached_block[request_id] = block_idx + + def allocate_external_computed_blocks( + self, + request_id: str, + num_local_computed_tokens: int, + num_external_computed_tokens: int, + ) -> None: + """ + Allocate new blocks for external (KV-connector) computed tokens. + + Must run only after every group's local blocks have been touched via + `add_local_computed_blocks`, so this group's `get_new_blocks` cannot + evict another group's cache-hit blocks (issue #33775). + + Args: + request_id: The request ID. + num_local_computed_tokens: The number of local computed tokens. + num_external_computed_tokens: The number of external computed tokens. + """ + num_total_computed_tokens = ( + num_local_computed_tokens + num_external_computed_tokens + ) + num_skipped_tokens = self.get_num_skipped_tokens(num_total_computed_tokens) + if num_skipped_tokens > 0: + # Some external computed tokens may be skipped too. + num_external_computed_tokens = min( + num_total_computed_tokens - num_skipped_tokens, + num_external_computed_tokens, + ) + if num_external_computed_tokens <= 0: + return + + req_blocks = self.req_to_blocks[request_id] + allocated_blocks = self.block_pool.get_new_blocks( + cdiv(num_total_computed_tokens, self.block_size) - len(req_blocks) + ) + req_blocks.extend(allocated_blocks) + if self._record_new_block_ids: + self.new_block_ids.extend(b.block_id for b in allocated_blocks) + + def allocate_new_blocks( + self, request_id: str, num_tokens: int, num_tokens_main_model: int + ) -> list[KVCacheBlock]: + """ + Allocate new blocks for the request to give it at least `num_tokens` + token slots. + + Args: + request_id: The request ID. + num_tokens: The total number of tokens that need a slot (including + tokens that are already allocated). + num_tokens_main_model: The number of tokens for the main model (aka target + model in spec decode). w/o spec decode, it is num_tokens; + with spec decode, it is num_tokens - num_lookahead_tokens. + Returns: + The new allocated blocks. + """ + cow_blocks: list[KVCacheBlock] = [] + if request_id in self._partial_hit_reqs: + # Partial hit: redirect the shared tail to a private CoW block. + # Replacing in place keeps the length-based allocation below + # correct; the extra block was reserved by + # get_num_blocks_to_allocate. + block_idx, source_block = self._partial_hit_reqs.pop(request_id) + cow_block = self.block_pool.get_new_blocks(1)[0] + self._apply_cow(request_id, block_idx, source_block, cow_block) + self.new_block_ids.append(cow_block.block_id) + cow_blocks.append(cow_block) + + req_blocks = self.req_to_blocks[request_id] + num_required_blocks = cdiv(num_tokens, self.block_size) + num_new_blocks = num_required_blocks - len(req_blocks) + if num_new_blocks <= 0: + return cow_blocks + else: + new_blocks = self.block_pool.get_new_blocks(num_new_blocks) + req_blocks.extend(new_blocks) + if self._record_new_block_ids: + self.new_block_ids.extend(b.block_id for b in new_blocks) + return cow_blocks + new_blocks + + @property + def records_new_block_ids(self) -> bool: + """Whether this manager's new blocks are zeroed by the worker.""" + return self._record_new_block_ids + + def take_new_block_ids(self) -> list[int]: + """Drain and return block IDs allocated since the last call.""" + ids = self.new_block_ids + self.new_block_ids = [] + return ids + + def take_pending_cow_copies( + self, + ) -> list[tuple[KVCacheBlock, KVCacheBlock]]: + """Drain pending CoW source and destination block pairs.""" + pending_copies = self._pending_cow_copies + self._pending_cow_copies = [] + return pending_copies + + def take_pending_partial_tail_offloads( + self, + ) -> list[tuple[str, int, KVCacheBlock, int, BlockHashWithGroupId]]: + """Drain producer partial-tail hand-offs. + + Entries include request, group, block, boundary, and its exact hash. + + Only mamba "align" populates this. The block lives off the request + block table, so the caller must pin it until the connector has read + it — nothing else keeps it alive once the CoW retention is released. + """ + pending = self._pending_partial_tail_offloads + self._pending_partial_tail_offloads = [] + return pending + + def take_pending_aligned_recurrent_boundaries( + self, + ) -> list[tuple[str, int, KVCacheBlock, int]]: + """Drain full-page Mamba boundary hand-offs.""" + pending = self._pending_aligned_recurrent_boundaries + self._pending_aligned_recurrent_boundaries = [] + return pending + + def _apply_cow( + self, + request_id: str, + block_idx: int, + source_block: KVCacheBlock, + cow_block: KVCacheBlock, + ) -> None: + """Redirect a partial prefix-cache hit to a private CoW block. + + Both copy endpoints stay retained until the copy has run on the worker, + so a same-step free cannot recycle them: ``source_block`` keeps its + hit-ref, ``cow_block`` takes an extra ref beyond the one handed to the + request. + """ + req_blocks = self.req_to_blocks[request_id] + assert block_idx < len(req_blocks) + assert req_blocks[block_idx] is source_block + assert not source_block.is_null and source_block.ref_cnt > 0 + req_blocks[block_idx] = cow_block + self._pending_cow_copies.append((source_block, cow_block)) + cow_block.ref_cnt += 1 + + def cache_blocks( + self, + request: Request, + num_tokens: int, + retention_interval: int | None = None, + ) -> None: + """ + Cache the blocks for the request. + + Args: + request: The request. + num_tokens: The total number of tokens that need to be cached + (including tokens that are already cached). + retention_interval: Sparse local-checkpoint granularity. ``None`` + keeps dense checkpointing; ``0`` keeps only the latest replay + boundary; a positive multiple of ``scheduler_block_size`` keeps + a tail once per that-sized segment. Only SWA acts on it. + """ + num_cached_blocks = self.num_cached_block.get(request.request_id, 0) + num_full_blocks = num_tokens // self.block_size + + if num_cached_blocks >= num_full_blocks: + return + + # Token boundaries whose reachable tail must be retained under sparse + # retention: the replay boundary (``num_prompt - 1``, capped by + # ``get_computed_blocks``) and any detected shared-prefix junction. + reachable_boundaries = [request.num_prompt_tokens - 1] + if request.shared_prefix_boundary: + reachable_boundaries.append(request.shared_prefix_boundary) + + block_mask = self.reachable_block_mask( + start_block=num_cached_blocks, + end_block=num_full_blocks, + alignment_tokens=self.scheduler_block_size, + kv_cache_spec=self.kv_cache_spec, + use_eagle=self.use_eagle, + retention_interval=retention_interval, + reachable_boundaries=reachable_boundaries, + ) + planned = self._planned_recurrent_checkpoints.get(request.request_id) + if planned is not None and block_mask is not None: + for position in planned[2]: + index = position // self.block_size - 1 + if not num_cached_blocks <= index < num_full_blocks: + raise ValueError("planned checkpoint outside cache registration span") + block_mask[index - num_cached_blocks] = True + self.block_pool.cache_full_blocks( + request=request, + blocks=self.req_to_blocks[request.request_id], + num_cached_blocks=num_cached_blocks, + num_full_blocks=num_full_blocks, + block_size=self.block_size, + kv_cache_group_id=self.kv_cache_group_id, + block_mask=block_mask, + ) + + self.num_cached_block[request.request_id] = num_full_blocks + + @classmethod + def reachable_block_mask( + cls, + start_block: int, + end_block: int, + alignment_tokens: int | None, + kv_cache_spec: KVCacheSpec, + use_eagle: bool, + retention_interval: int | None = None, + reachable_boundaries: Sequence[int] = (), + ) -> list[bool] | None: + """Per-block mask for ``cache_full_blocks``. ``None`` means cache + every (non-null) block — the default for full attention. + + Subclasses with sparse hit semantics (SWA / Mamba) override this to skip + blocks that can never serve a hit at any alignment-aligned prefix length. + ``reachable_boundaries`` are token positions whose reachable tail must be + retained; the base (dense) policy ignores them. + """ + return None + + def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]: + """ + Pop the request's bookkeeping and return its blocks without yet + returning them to the block pool. The caller is responsible for + eventually passing the returned blocks to `block_pool.free_blocks`, + freeing them in reverse order (so that tail blocks are evicted first). + + Args: + request_id: The request ID. + + Returns: + The request's blocks in allocation order. + """ + # Default to [] in case a request is freed (aborted) before alloc. + req_blocks = self.req_to_blocks.pop(request_id, []) + self.num_cached_block.pop(request_id, None) + self._partial_hit_reqs.pop(request_id, None) + return req_blocks + + def free(self, request_id: str) -> None: + """ + Free the blocks for the request. + + Args: + request_id: The request ID. + """ + # Free blocks in reverse order so that the tail blocks are freed first. + self.block_pool.free_blocks(reversed(self.pop_blocks_for_free(request_id))) + + @abstractmethod + def get_num_common_prefix_blocks(self, running_request_id: str) -> int: + """ + Get the number of common prefix blocks for all requests with allocated + KV cache. + + Args: + running_request_id: The request ID. + + Returns: + The number of common prefix blocks for all requests with allocated + KV cache. + """ + + raise NotImplementedError + + @classmethod + @abstractmethod + def find_longest_cache_hit( + cls, + block_hashes: BlockHashList, + max_length: int, + kv_cache_group_ids: list[int], + block_pool: BlockPool, + kv_cache_spec: KVCacheSpec, + drop_eagle_block: bool, + alignment_tokens: int, + dcp_world_size: int = 1, + pcp_world_size: int = 1, + ) -> tuple[tuple[list[KVCacheBlock], ...], int]: + """ + Get the longest cache hit prefix of the blocks that is not longer than + `max_length`. The prefix should be a common prefix hit for all the + kv cache groups in `kv_cache_group_ids`. If no cache hit is found, + return an empty list. + If eagle is enabled, drop the last matched block to force recompute the + last block to get the required hidden states for eagle drafting head. + For multi-module MTP, this recompute also rewrites the dropped block's + draft-layer KVs, which depend on up to num_speculative_tokens - 1 + tokens past the matched prefix (i.e. on the cache writer's + continuation, which the block hash does not cover); the coordinator + asserts the block size covers that window. + Need to be customized for each attention type. + + Args: + block_hashes: The block hashes of the request. + max_length: The maximum length of the cache hit prefix. + kv_cache_group_ids: The ids of the kv cache groups. + block_pool: The block pool. + kv_cache_spec: The kv cache spec. + drop_eagle_block: Whether to drop the last matched block for EAGLE/MTP. + Always False for non-EAGLE/MTP groups, but can be False for EAGLE/MTP + groups too if the last block is already dropped (e.g., in a + convergence loop in `find_longest_cache_hit`). + alignment_tokens: The returned cache hit length (in tokens) should + be a multiple of this value (in tokens). By default, it should + be set to the block_size. + dcp_world_size: The world size of decode context parallelism. + pcp_world_size: The world size of prefill context parallelism. + + Returns: + A tuple containing cached blocks and the exact cache-hit length in + tokens. The cached block tuple has skipped blocks replaced by null + blocks for each kv cache group in `kv_cache_group_ids`. + For example, sliding window manager should return a list like + ([NULL, NULL, KVCacheBlock(7), KVCacheBlock(8)]) for block size 4 + and sliding window 8 and len(kv_cache_group_ids) = 1. + """ + + raise NotImplementedError + + def _remove_blocks_in_range( + self, + request_id: str, + first_block: int, + last_block: int, + ) -> None: + """Free blocks in ``[first_block, last_block)`` and replace with null_block. + + Iterates backward so newly-evictable tail blocks are reached even after + earlier blocks in the range were nulled in a prior call. + """ + if request_id not in self.req_to_blocks: + return + if first_block >= last_block: + return + blocks = self.req_to_blocks[request_id] + last_block = min(last_block, len(blocks)) + + freed: list[KVCacheBlock] = [] + for i in range(last_block - 1, first_block - 1, -1): + if blocks[i] == self._null_block: + break + freed.append(blocks[i]) + blocks[i] = self._null_block + if freed: + self.block_pool.free_blocks(freed) + + def remove_skipped_blocks( + self, + request_id: str, + processed_computed_tokens: int, + num_prompt_tokens: int | None = None, + ) -> None: + """ + Remove and free the blocks that are no longer needed for attention computation. + The removed blocks should be replaced by null_block. + + This function depends on `get_num_skipped_tokens`, which need to be implemented + differently for each attention type. + + Args: + request_id: The request ID. + processed_computed_tokens: Computed-token prefix length covering + fully processed and committed tokens only (safe to free). + num_prompt_tokens: Optional prompt length for attention types (e.g. + R-SWA) that evict a middle gap rather than a head prefix. Ignored + by the default implementation. + """ + del num_prompt_tokens + # Remove the blocks that will be skipped during attention computation. + num_skipped_tokens = self.get_num_skipped_tokens(processed_computed_tokens) + if num_skipped_tokens <= 0: + # This indicates that ALL tokens are inside attention window. + # Thus we do not need to free any blocks outside attention window. + # A typical case is full attention that we never free any token + # before the request is finished. + return + blocks = self.req_to_blocks[request_id] + num_skipped_blocks = num_skipped_tokens // self.block_size + # `num_skipped_tokens` may include tokens that haven't been allocated yet + # (e.g., when the attention window moves into the external computed tokens + # range), so we must cap to the number of blocks that currently exist for + # this request. + num_skipped_blocks = min(num_skipped_blocks, len(blocks)) + self._remove_blocks_in_range(request_id, 0, num_skipped_blocks) + + def get_num_skipped_tokens(self, num_computed_tokens: int) -> int: + """ + Get the number of tokens that will be skipped for attention computation. + + Args: + num_computed_tokens: The number of tokens that have been computed. + + Returns: + The number of tokens that will be skipped for attention computation. + """ + # The default behavior is to not skip any tokens. + return 0 + + def new_step_starts(self) -> None: + return None + + +class FullAttentionManager(SingleTypeKVCacheManager): + supports_fine_grained_hash_lookup: ClassVar[bool] = True + + @classmethod + def find_longest_cache_hit( + cls, + block_hashes: BlockHashList, + max_length: int, + kv_cache_group_ids: list[int], + block_pool: BlockPool, + kv_cache_spec: KVCacheSpec, + drop_eagle_block: bool, + alignment_tokens: int, + dcp_world_size: int = 1, + pcp_world_size: int = 1, + ) -> tuple[tuple[list[KVCacheBlock], ...], int]: + assert isinstance( + kv_cache_spec, FullAttentionSpec | ChunkedLocalAttentionSpec + ), ( + "FullAttentionManager can only be used for full attention " + "and chunked local attention groups" + ) + block_size = kv_cache_spec.block_size + if dcp_world_size > 1 and not getattr(kv_cache_spec, "dcp_replicated", False): + # DCP shards each block's KV across ranks; hashes must be viewed at + # the sharded block size. + block_size *= dcp_world_size + block_hashes = resolve_block_hashes( + block_hashes, + block_pool.hash_block_size, + block_size, + supports_fine_grained_hash_lookup=cls.supports_fine_grained_hash_lookup, + alignment_tokens=alignment_tokens, + ) + + # Fine-grained mode (alignment_tokens == hash_block_size < + # block_size): resolve_block_hashes kept the raw hash-granularity + # list so interior boundaries can be probed. + fine_grained = ( + alignment_tokens < block_size and block_size % alignment_tokens == 0 + ) + if fine_grained: + # list or lazy BlobBlockHashes view + assert isinstance(block_hashes, Sequence) + full_block_hashes: BlockHashList = BlockHashListWithBlockSize( + block_hashes, alignment_tokens, block_size + ) + else: + full_block_hashes = block_hashes + + computed_blocks: tuple[list[KVCacheBlock], ...] = tuple( + [] for _ in range(len(kv_cache_group_ids)) + ) + # Phase 1: longest run of cached full blocks from the start. A missing + # block implies every later block misses too (chained hashes). + for block_hash in itertools.islice(full_block_hashes, max_length // block_size): + cached_block = block_pool.get_cached_block(block_hash, kv_cache_group_ids) + if not cached_block: + break + for computed, cached in zip(computed_blocks, cached_block): + computed.append(cached) + hit_length = len(computed_blocks[0]) * block_size + + # Phase 2 (fine-grained only): extend into the first non-full block by + # probing its interior hash boundaries high-to-low (longest hit first). + if fine_grained: + # list or lazy BlobBlockHashes view + assert isinstance(block_hashes, Sequence) + scale_factor = block_size // alignment_tokens + first_partial_idx = len(computed_blocks[0]) * scale_factor + max_partial_idx = min( + first_partial_idx + scale_factor - 1, + max_length // alignment_tokens, + len(block_hashes), + ) + for fine_idx in range(max_partial_idx - 1, first_partial_idx - 1, -1): + cached_tail = block_pool.get_cached_block( + block_hashes[fine_idx], kv_cache_group_ids + ) + if not cached_tail: + continue + for computed, cached in zip(computed_blocks, cached_tail): + computed.append(cached) + hit_length = (fine_idx + 1) * alignment_tokens + break + + # Eagle needs the tokens right before the generation point recomputed: + # drop one hash unit when fine-grained (the tail block's KV is + # append-only, so it still covers the reduced length), else one cache + # block. + if drop_eagle_block and hit_length > 0: + hit_length -= min(alignment_tokens, block_size) + # Round down to the alignment; a no-op when fine-grained (hits land on + # hash boundaries by construction) and when alignment_tokens == + # block_size. Then trim blocks past the new tail. + hit_length -= hit_length % alignment_tokens + num_blocks = cdiv(hit_length, block_size) + for computed in computed_blocks: + del computed[num_blocks:] + return computed_blocks, hit_length + + def cache_blocks( + self, + request: Request, + num_tokens: int, + retention_interval: int | None = None, + ) -> None: + super().cache_blocks(request, num_tokens, retention_interval=retention_interval) + hash_block_size = self.block_pool.hash_block_size + if self.block_size == hash_block_size: + return + self._cache_partial_tail_block(request, num_tokens) + + def _cache_partial_tail_block( + self, + request: Request, + num_tokens: int, + ) -> None: + """Cache the prompt tail when it ends inside a cache block. + + Only the final prompt hash boundary is registered as a partial + prefix-cache entry; intermediate hash boundaries inside the same cache + block are intentionally skipped. + """ + hash_block_size = self.block_pool.hash_block_size + boundary_tokens = request.num_prompt_tokens // hash_block_size * hash_block_size + if boundary_tokens == 0 or boundary_tokens > num_tokens: + return + if boundary_tokens % self.block_size == 0: + return + + blocks = self.req_to_blocks[request.request_id] + block_idx = boundary_tokens // self.block_size + if block_idx >= len(blocks): + return + self.block_pool.cache_partial_block( + request=request, + block=blocks[block_idx], + num_tokens=boundary_tokens, + kv_cache_group_id=self.kv_cache_group_id, + block_size=self.block_size, + ) + + def get_num_common_prefix_blocks(self, running_request_id: str) -> int: + blocks = self.req_to_blocks[running_request_id] + num_common_blocks = 0 + for block in blocks: + if block.ref_cnt == len(self.req_to_blocks): + num_common_blocks += 1 + else: + break + return num_common_blocks + + +class RSWAManager(FullAttentionManager): + """KV cache manager for Reference Sliding Window Attention (R-SWA). + + When ``num_prompt_tokens`` is supplied to ``remove_skipped_blocks``, frees + gap blocks between the prefill tail and the current decode window. This + bounds per-request KV memory at O(prefix_len + rswa_window) instead of + growing linearly with decode length. + """ + + def __init__(self, kv_cache_spec: RSWASpec, **kwargs) -> None: + super().__init__(kv_cache_spec, **kwargs) + self.rswa_window: int = kv_cache_spec.rswa_window + + def remove_skipped_blocks( + self, + request_id: str, + processed_computed_tokens: int, + num_prompt_tokens: int | None = None, + ) -> None: + """Free gap blocks that are no longer needed for attention. + + Gap = blocks entirely within + [ceil(prefix_len / block_size) * block_size, + max(prefix_len, processed_computed_tokens - rswa_window)) + + Freed blocks are replaced with null_block in req_to_blocks so the + block_table passed to FA4 is valid (null_block KV is all-zero; + rswa_mask_mod marks gap positions as non-visible so FA4 skips them). + """ + if num_prompt_tokens is None: + super().remove_skipped_blocks( + request_id, processed_computed_tokens, num_prompt_tokens + ) + return + + bs = self.block_size + # First block fully after the prefill boundary. + first_gap_block = cdiv(num_prompt_tokens, bs) + # Decode window start position; blocks before this are evictable. + window_start = max( + num_prompt_tokens, processed_computed_tokens - self.rswa_window + ) + last_gap_block = window_start // bs # exclusive upper bound + self._remove_blocks_in_range(request_id, first_gap_block, last_gap_block) + + +class SlidingWindowManager(SingleTypeKVCacheManager): + def __init__(self, kv_cache_spec: SlidingWindowSpec, **kwargs) -> None: + super().__init__(kv_cache_spec, **kwargs) + self.sliding_window = kv_cache_spec.sliding_window + # Extra trailing tokens to retain below the window (never attended) so a + # multi-module MTP store-side lag can still reconstruct the window from + # cached blocks. + self.extra_retained_tokens = kv_cache_spec.extra_retained_tokens + + @classmethod + def _contiguous_blocks_for_hit( + cls, window_size: int, block_size: int, use_eagle: bool + ) -> int: + blocks = cdiv(window_size - 1, block_size) + if use_eagle: + # Need to drop the last matched block if eagle is enabled. For + # sliding window layer, we achieve this by increasing the number of + # contiguous blocks needed for prefix cache hit by one and dropping + # the last matched block. + blocks += 1 + return blocks + + @classmethod + def find_longest_cache_hit( + cls, + block_hashes: BlockHashList, + max_length: int, + kv_cache_group_ids: list[int], + block_pool: BlockPool, + kv_cache_spec: KVCacheSpec, + drop_eagle_block: bool, + alignment_tokens: int, + dcp_world_size: int = 1, + pcp_world_size: int = 1, + ) -> tuple[tuple[list[KVCacheBlock], ...], int]: + assert isinstance(kv_cache_spec, SlidingWindowSpec), ( + "SlidingWindowManager can only be used for sliding window groups" + ) + assert dcp_world_size == 1 or kv_cache_spec.dcp_replicated, ( + "DCP only supports sliding-window KV when it is replicated." + ) + assert pcp_world_size == 1 or kv_cache_spec.dcp_replicated, ( + "PCP only supports sliding-window KV when it is replicated." + ) + # Fine-grained partial hits are not supported for sliding window now + assert alignment_tokens % kv_cache_spec.block_size == 0, ( + "SlidingWindowManager does not support fine-grained (partial) cache hits" + ) + block_hashes = resolve_block_hashes( + block_hashes, + block_pool.hash_block_size, + kv_cache_spec.block_size, + supports_fine_grained_hash_lookup=cls.supports_fine_grained_hash_lookup, + alignment_tokens=alignment_tokens, + ) + + # The number of contiguous blocks needed for a prefix cache hit. + sliding_window_contiguous_blocks = cls._contiguous_blocks_for_hit( + kv_cache_spec.sliding_window, kv_cache_spec.block_size, drop_eagle_block + ) + + # TODO: reduce i by sliding_window_contiguous_blocks when cache miss, to + # optimize the time complexity from O(max_num_blocks) to + # O(max_num_blocks / sliding_window_contiguous_blocks + + # sliding_window_contiguous_blocks), + # which is good for low cache hit rate scenarios. + max_num_blocks = max_length // kv_cache_spec.block_size + computed_blocks: tuple[list[KVCacheBlock], ...] = tuple( + [block_pool.null_block] * max_num_blocks + for _ in range(len(kv_cache_group_ids)) + ) + block_size = kv_cache_spec.block_size + num_contiguous_blocks = 0 + match_found = False + # Search from right to left and early stop when a match is found. + for i in range(max_num_blocks - 1, -1, -1): + if cached_block := block_pool.get_cached_block( + block_hashes[i], kv_cache_group_ids + ): + # Skip prefix matching check if the block is not aligned with + # `alignment_tokens`. + if num_contiguous_blocks == 0 and block_size != alignment_tokens: + post_pop_blocks = i if drop_eagle_block else i + 1 + if (post_pop_blocks * block_size) % alignment_tokens != 0: + continue + # Add the cached block to the computed blocks. + for computed, cached in zip(computed_blocks, cached_block): + computed[i] = cached + num_contiguous_blocks += 1 + if num_contiguous_blocks >= sliding_window_contiguous_blocks: + # Trim the trailing blocks. + # E.g., [NULL, NULL, 8, 3, NULL, 9] -> [NULL, NULL, 8, 3] + # when sliding_window_contiguous_blocks=2. + for computed in computed_blocks: + del computed[i + num_contiguous_blocks :] + match_found = True + break + else: + num_contiguous_blocks = 0 + if not match_found: + # The first `num_contiguous_blocks` is a cache hit even if + # `num_contiguous_blocks < sliding_window_contiguous_blocks`. + for computed in computed_blocks: + del computed[num_contiguous_blocks:] + while ( + block_size != alignment_tokens # Faster for common case. + and len(computed_blocks[0]) * block_size % alignment_tokens != 0 + ): + for computed in computed_blocks: + computed.pop() + if drop_eagle_block and computed_blocks[0]: + for computed in computed_blocks: + computed.pop() + # Re-align after eagle pop: the pop may break the alignment + # when block_size != alignment_tokens (hybrid models with + # different page sizes, e.g. Gemma4). + while ( + block_size != alignment_tokens + and len(computed_blocks[0]) * block_size % alignment_tokens != 0 + ): + for computed in computed_blocks: + computed.pop() + hit_length = len(computed_blocks[0]) * block_size + return computed_blocks, hit_length + + @classmethod + def reachable_block_mask( + cls, + start_block: int, + end_block: int, + alignment_tokens: int | None, + kv_cache_spec: KVCacheSpec, + use_eagle: bool, + retention_interval: int | None = None, + reachable_boundaries: Sequence[int] = (), + ) -> list[bool] | None: + assert isinstance(kv_cache_spec, SlidingWindowSpec) + if alignment_tokens is None: + # Fast path: when the coordinator imposes no alignment constraint. + return None + assert alignment_tokens % kv_cache_spec.block_size == 0 + + block_size = kv_cache_spec.block_size + # Contiguous blocks a hit needs at a boundary (incl. the EAGLE peek). + need = cls._contiguous_blocks_for_hit( + window_size=kv_cache_spec.sliding_window, + block_size=block_size, + use_eagle=use_eagle, + ) + # The matched run's right edge sits on the aligned boundary block when + # EAGLE peeks one block past it (shift=1), otherwise on the last block + # before the boundary (shift=0). + shift = 1 if use_eagle else 0 + + mask = [False] * (end_block - start_block) + + # (1) Segment-boundary tails. ``retention_interval``: + # None -> dense (a tail at every ``alignment_tokens`` boundary); + # 0 -> no dense tails (only the replay boundary below); + # >0 -> a tail once per ``retention_interval``-sized segment. + segment_tokens = ( + alignment_tokens + if retention_interval is None + else (None if retention_interval == 0 else retention_interval) + ) + if segment_tokens is not None: + per_segment = segment_tokens // block_size + if need >= per_segment: + # Every block is reachable; cache them all. + return None + for i in range(start_block, end_block): + if i >= shift and (i - shift) % per_segment >= per_segment - need: + mask[i - start_block] = True + + # (2) Reachable-boundary tails: the replay boundary (``num_prompt - 1``, + # capped by ``get_computed_blocks``) and any shared-prefix junction. Both + # land before segments would cover them under sparse retention, so keep + # the ``need``-block tail ending on each boundary explicitly. + if retention_interval is not None: + for boundary_tokens in reachable_boundaries: + aligned = boundary_tokens // alignment_tokens * alignment_tokens + end = aligned // block_size + shift + for j in range(max(start_block, end - need), min(end_block, end)): + mask[j - start_block] = True + + return mask + + def get_num_skipped_tokens(self, num_computed_tokens: int) -> int: + """ + Get the number of tokens that will be skipped for attention computation. + + For sliding window, this corresponds to the tokens that are prior to + the current sliding window. + + Example: + sliding_window=4, num_computed_tokens=7 + + Tokens: [ 0 1 2 3 4 5 6 7 ] + | ---- computed -----| + ^ next token to be computed + |-----------| sliding window for next token + |--skipped---| + + The current window contains tokens 4~7. Tokens 0~3 will be skipped for + attention computation since they are outside the sliding window. + Thus, get_num_skipped_tokens(7) == 4. + + The trailing edge of the window is extended by ``extra_retained_tokens`` + so that those extra trailing tokens' blocks are retained (but not + attended). This is needed for multi-module spec decoding which can + re-prefill the last num_spec_prefill_tokens - 1 tokens from the end + of the sequence, and thus needs to delay freeing/caching of blocks. + + Args: + num_computed_tokens: The number of tokens that have been computed. + + Returns: + The number of tokens that will be skipped for attention computation. + """ + return max( + 0, + num_computed_tokens - self.sliding_window + 1 - self.extra_retained_tokens, + ) + + def get_num_common_prefix_blocks(self, running_request_id: str) -> int: + """ + NOTE(Chen): The prefix blocks are null blocks for sliding window layers. + So it's not correct to count ref_cnt like FullAttentionManager. Return + 0 here for correctness. Need to support cascade attention + sliding + window in the future. + """ + return 0 + + +class ChunkedLocalAttentionManager(SingleTypeKVCacheManager): + def __init__(self, kv_cache_spec: ChunkedLocalAttentionSpec, **kwargs) -> None: + super().__init__(kv_cache_spec, **kwargs) + self.attention_chunk_size = kv_cache_spec.attention_chunk_size + + @classmethod + def find_longest_cache_hit( + cls, + block_hashes: BlockHashList, + max_length: int, + kv_cache_group_ids: list[int], + block_pool: BlockPool, + kv_cache_spec: KVCacheSpec, + drop_eagle_block: bool, + alignment_tokens: int, + dcp_world_size: int = 1, + pcp_world_size: int = 1, + ) -> tuple[tuple[list[KVCacheBlock], ...], int]: + """ + For chunked local attention, we need to find the longest cache hit + prefix of the blocks that is not longer than `max_length`. The prefix + should be a common prefix hit for all the kv cache groups in + `kv_cache_group_ids`. If no cache hit is found, return an empty list. + note we mark as computed if the whole block is outside of the local + window, and set the block as null. Examples: + + 1. Attention chunk size of 8, block size of 4, max length of 15 + for next token at 15th (zero-indexed), 8th - 14th tokens are in + the window(needs lookup), 0th - 7th are not in the window, + so they are already marked as computed. We check the complete + block3 (8th - 11th tokens), Assume block 3 is hit, we will return + [null, null, block 3], otherwise, we return [null, null] + + 2. Attention chunk size of 8, block size of 4, max length of 16 + for next token at 16th (zero-indexed), 0th - 15th tokens are not + in the window, so they are already marked as computed. + we return 4 blocks[null, null, null, null] + + Args: + block_hashes: The block hashes of the request. + max_length: The maximum length of the cache hit prefix. + kv_cache_group_ids: The ids of the kv cache groups. + block_pool: The block pool. + kv_cache_spec: The kv cache spec. + drop_eagle_block: Whether to drop the last matched block for EAGLE/MTP. + dcp_world_size: The world size of decode context parallelism. + pcp_world_size: The world size of prefill context parallelism. + alignment_tokens: The returned cache hit length (in tokens) should + be a multiple of this value (in tokens). + + Returns: + A list of cached blocks + """ + assert isinstance(kv_cache_spec, ChunkedLocalAttentionSpec), ( + "ChunkedLocalAttentionManager can only be used for " + "chunked local attention groups" + ) + assert drop_eagle_block is False, ( + "Hybrid KV cache is not supported for " + "eagle + chunked local attention." + ) + assert dcp_world_size == 1, "DCP not support chunked local attn now." + assert pcp_world_size == 1, "PCP not support chunked local attn now." + assert kv_cache_spec.block_size == alignment_tokens, ( + "KV cache groups with different block sizes are not compatible with " + "chunked local attention now" + ) + block_hashes = resolve_block_hashes( + block_hashes, + block_pool.hash_block_size, + kv_cache_spec.block_size, + supports_fine_grained_hash_lookup=cls.supports_fine_grained_hash_lookup, + alignment_tokens=alignment_tokens, + ) + max_num_blocks = max_length // kv_cache_spec.block_size + if max_length > 0: + local_attention_start_idx = ( + max_length + // kv_cache_spec.attention_chunk_size + * kv_cache_spec.attention_chunk_size + ) + else: + local_attention_start_idx = 0 + # we marked blocks out of window as computed + # with null blocks, and blocks inside window based on cache lookup + # result [null] [null] ... [null] [hit block 1 (1st block contain + # last window)] [hit block 2] ... [hit block x] + local_attention_start_block_idx = ( + local_attention_start_idx // kv_cache_spec.block_size + ) + computed_blocks: tuple[list[KVCacheBlock], ...] = tuple( + [block_pool.null_block] * local_attention_start_block_idx + for _ in range(len(kv_cache_group_ids)) + ) + for i in range(local_attention_start_block_idx, max_num_blocks): + block_hash = block_hashes[i] + if cached_block := block_pool.get_cached_block( + block_hash, kv_cache_group_ids + ): + for computed, cached in zip(computed_blocks, cached_block): + computed.append(cached) + else: + break + hit_length = len(computed_blocks[0]) * kv_cache_spec.block_size + return computed_blocks, hit_length + + def get_num_skipped_tokens(self, num_computed_tokens: int) -> int: + """ + Get the number of tokens that will be skipped for attention computation. + + For chunked local attention, this corresponds to the tokens that are on + the left side of the current chunk. + + Example 1: + chunk size = 8, num_computed_tokens = 13 + Tokens: [ 0 1 2 3 4 5 6 7 | 8 9 10 11 12 13 14 15 ] ... + | ----- computed ---------------| + ^^ next token to be computed + |----------------| <-- attention window for + next token + |--- skipped -----| + Output: get_num_skipped_tokens(13) == 8 + + Example 2: + chunk size = 8, num_computed_tokens = 8 + Tokens: [ 0 1 2 3 4 5 6 7 | 8 9 10 11 12 13 14 15 ] ... + | --- computed ---| + ^ next token to be computed + |--| <-- attention window for next token + | --- skipped ----| + Output: get_num_skipped_tokens(8) == 8 + + Example 3: + chunk size = 8, num_computed_tokens = 7 + Tokens: [ 0 1 2 3 4 5 6 7 | 8 9 10 11 12 13 14 15 ] ... + |---computed---| + ^ next token to be computed + |-----------------| <-- attention window for next token + no token should be skipped. + Output: get_num_skipped_tokens(7) == 0 + + Args: + num_computed_tokens: The number of tokens that have been computed. + + Returns: + The number of tokens that will be skipped for attention computation. + """ + num_skipped_tokens = ( + num_computed_tokens // self.attention_chunk_size + ) * self.attention_chunk_size + return num_skipped_tokens + + def get_num_common_prefix_blocks(self, running_request_id: str) -> int: + """ + cascade attention is not supported by chunked local attention. + """ + return 0 + + +class MambaManager(SingleTypeKVCacheManager): + supports_fine_grained_hash_lookup: ClassVar[bool] = True + + def __init__( + self, kv_cache_spec: MambaSpec, block_pool: BlockPool, **kwargs + ) -> None: + super().__init__(kv_cache_spec, block_pool, **kwargs) + # Mamba layers use TP instead of DCP, so each rank holds the full + # recurrent state. Undo the DCP/PCP block_size scaling that the base + # class applies for attention groups whose KV cache is partitioned. + self.block_size = kv_cache_spec.block_size + self.mamba_cache_mode = kv_cache_spec.mamba_cache_mode + self.num_speculative_blocks: int = kv_cache_spec.num_speculative_blocks + self.cached_blocks_this_step: set[BlockHashWithGroupId] = set() + if self.mamba_cache_mode == "align": + # Mapping from request ID to the index of the block + # allocated in the previous step + self.last_state_block_idx: dict[str, int] = {} + # The set of the requests that have been allocated blocks + self._allocated_block_reqs: set[str] = set() + # Number of internal checkpoint blocks required by each request's + # current allocation. + self._num_checkpoint_blocks: dict[str, int] = {} + # Requests that registered their own last-prompt-boundary partial + # tail (producers). On the next step's CoW the boundary state moves + # into a private cow_block; we record that block for connector + # offload (see _pending_partial_tail_offloads). + self._producer_partial_tail_reqs: dict[ + str, tuple[int, BlockHashWithGroupId] + ] = {} + + @classmethod + def find_longest_cache_hit( + cls, + block_hashes: BlockHashList, + max_length: int, + kv_cache_group_ids: list[int], + block_pool: BlockPool, + kv_cache_spec: KVCacheSpec, + drop_eagle_block: bool, + alignment_tokens: int, + dcp_world_size: int = 1, + pcp_world_size: int = 1, + ) -> tuple[tuple[list[KVCacheBlock], ...], int]: + assert isinstance(kv_cache_spec, MambaSpec), ( + "MambaManager can only be used for mamba groups" + ) + assert dcp_world_size == 1, "DCP not support mamba now." + assert pcp_world_size == 1, "PCP not support mamba now." + block_hashes = resolve_block_hashes( + block_hashes, + block_pool.hash_block_size, + kv_cache_spec.block_size, + supports_fine_grained_hash_lookup=cls.supports_fine_grained_hash_lookup, + alignment_tokens=alignment_tokens, + ) + computed_blocks: tuple[list[KVCacheBlock], ...] = tuple( + [] for _ in range(len(kv_cache_group_ids)) + ) + hit_length = 0 + + block_size = kv_cache_spec.block_size + if alignment_tokens < block_size and block_size % alignment_tokens == 0: + # list or lazy BlobBlockHashes view + assert isinstance(block_hashes, Sequence) + hash_block_size = alignment_tokens + scale_factor = block_size // hash_block_size + max_num_partial_units = min( + max_length // hash_block_size, len(block_hashes) + ) + for fine_idx in range(max_num_partial_units - 1, -1, -1): + num_tokens = (fine_idx + 1) * hash_block_size + block_hash = block_hashes[fine_idx] + if cached_block := block_pool.get_cached_block( + block_hash, kv_cache_group_ids + ): + block_idx = fine_idx // scale_factor + for computed, cached in zip(computed_blocks, cached_block): + computed.extend([block_pool.null_block] * block_idx) + computed.append(cached) + hit_length = num_tokens + break + return computed_blocks, hit_length + + max_num_blocks = max_length // block_size + # Search from right to left and early stop when a match is found. + for i in range(max_num_blocks - 1, -1, -1): + if cached_block := block_pool.get_cached_block( + block_hashes[i], kv_cache_group_ids + ): + # When enable Mamba prefix caching, `block_size` will be aligned + # across full attention layers and Mamba layers to ensure the + # prefix hit length aligned at block + if ( + block_size != alignment_tokens # Faster for common case. + and (i + 1) * block_size % alignment_tokens != 0 + ): + continue + for computed, cached in zip(computed_blocks, cached_block): + # the hit length logic later assumes: + # hit_length = len(hit_blocks_other_attn[0]) + # * self.other_block_size + # so we insert dummy blocks at the beginning: + computed.extend([block_pool.null_block] * i) + computed.append(cached) + hit_length = (i + 1) * block_size + break # we just need the last match - early stopping + + return computed_blocks, hit_length + + @classmethod + def reachable_block_mask( + cls, + start_block: int, + end_block: int, + alignment_tokens: int | None, + kv_cache_spec: KVCacheSpec, + use_eagle: bool, + retention_interval: int | None = None, + reachable_boundaries: Sequence[int] = (), + ) -> list[bool] | None: + """Sparse Mamba state-snapshot retention. + + ``retention_interval``: + + ``None`` -> dense (cache every block; default, unchanged behavior) + ``0`` -> keep only the ``reachable_boundaries`` states + ``> 0`` -> keep one state per ``retention_interval``-sized segment + + ``reachable_boundaries`` are proven reuse points (the replay boundary and + any cross-request shared-prefix junction, Marconi-style APC); their + boundary state is always kept so sparse retention does not defeat reuse. + """ + if retention_interval is None or alignment_tokens is None: + # Dense caching (default) or no alignment constraint imposed. + return None + assert isinstance(kv_cache_spec, MambaSpec) + block_size = kv_cache_spec.block_size + mask = [False] * (end_block - start_block) + + # (1) Segment-boundary states. A Mamba hit needs exactly the single + # state block ending on the boundary (no window, and draft models have + # no mamba layers, so no eagle shift). Block ``i`` ends at token + # ``(i + 1) * block_size``. + segment_tokens = None if retention_interval == 0 else retention_interval + if segment_tokens is not None: + per_segment = segment_tokens // block_size + if per_segment <= 1: + # Interval at/below the block size: every block is a boundary. + return None + first_boundary = ( + start_block + per_segment + ) // per_segment * per_segment - 1 + for i in range(first_boundary - start_block, len(mask), per_segment): + mask[i] = True + + # (2) Reachable-boundary states: the replay boundary (``num_prompt - 1``, + # capped by ``get_computed_blocks``) and any shared-prefix junction, both + # of which segments would otherwise skip under sparse retention. A Mamba + # hit needs exactly the single state block ending on the boundary. + for boundary_tokens in reachable_boundaries: + aligned = boundary_tokens // alignment_tokens * alignment_tokens + boundary_block = aligned // block_size - 1 + if start_block <= boundary_block < end_block: + mask[boundary_block - start_block] = True + + if use_eagle: + # Retain the predecessor when materialized, alongside the + # scheduler-aligned fallback. Unmaterialized shared-junction + # slots remain null and cache_full_blocks skips them. + predecessor_block = boundary_tokens // block_size - 2 + if start_block <= predecessor_block < end_block: + mask[predecessor_block - start_block] = True + + return mask + + def remove_skipped_blocks( + self, + request_id: str, + processed_computed_tokens: int, + num_prompt_tokens: int | None = None, + ) -> None: + assert isinstance(self.kv_cache_spec, MambaSpec) + + super().remove_skipped_blocks( + request_id, processed_computed_tokens, num_prompt_tokens + ) + if self.mamba_cache_mode == "align": + # `last_state_block_idx` refers to the block index allocated two steps ago. + # The block allocated in the previous step is used to copy Mamba states + # into the block allocated in the current step; the earlier block is + # no longer needed and should be freed here. + last_state_block_idx = self.last_state_block_idx.get(request_id) + # Blocks allocated during prefill may be non-contiguous. Use + # `last_state_block_idx` to free the appropriate block and replace it + # with a null block. + if ( + last_state_block_idx is not None + and last_state_block_idx + < cdiv(processed_computed_tokens, self.block_size) - 1 + ): + blocks = self.req_to_blocks[request_id] + if blocks[last_state_block_idx] != self._null_block: + self.block_pool.free_blocks([blocks[last_state_block_idx]]) + blocks[last_state_block_idx] = self._null_block + + def get_num_common_prefix_blocks(self, running_request_id: str) -> int: + """ + cascade attention is not supported by mamba + """ + return 0 + + def _needs_internal_checkpoint( + self, + request_id: str, + num_tokens: int, + num_computed_tokens: int, + ) -> bool: + assert isinstance(self.kv_cache_spec, MambaSpec) + checkpoint_idx = cdiv(num_tokens, self.block_size) - 2 + blocks = self.req_to_blocks[request_id] + return ( + self.kv_cache_spec.num_prefill_checkpoint_blocks > 0 + and num_tokens % self.block_size != 0 + and num_computed_tokens % self.block_size == 0 + and checkpoint_idx >= 0 + and (checkpoint_idx >= len(blocks) or blocks[checkpoint_idx].is_null) + ) + + def get_num_blocks_to_allocate( + self, + request_id: str, + num_tokens: int, + new_computed_blocks: Sequence[KVCacheBlock], + total_computed_tokens: int, + num_local_computed_tokens: int, + num_tokens_main_model: int, + apply_admission_cap: bool = False, + ) -> int: + assert isinstance(self.kv_cache_spec, MambaSpec) + planned = self._planned_recurrent_checkpoints.get(request_id) + if planned is not None: + if self.req_to_blocks.get(request_id) or new_computed_blocks or total_computed_tokens: + raise ValueError("checkpoint admission requires a fresh recurrent table") + return len(planned[2]) + 1 + self.num_speculative_blocks + if ( + len(new_computed_blocks) > 0 + and new_computed_blocks[-1].block_hash in self.cached_blocks_this_step + ): + # Mamba can't rely on blocks generated by other requests in the current step + # To put it in the next step, we return num_gpu_blocks + 1 so + # that kv_cache_manager will think there is no enough blocks to allocate now + # and don't schedule it in the current step. + return self.block_pool.num_gpu_blocks + 1 + if self.mamba_cache_mode != "align": + # Allocate extra `num_speculative_blocks` blocks for + # speculative decoding (MTP/EAGLE) with linear attention. + if self.num_speculative_blocks > 0: + num_tokens += ( + self.kv_cache_spec.block_size * self.num_speculative_blocks + ) + return super().get_num_blocks_to_allocate( + request_id, + num_tokens, + new_computed_blocks, + total_computed_tokens, + num_local_computed_tokens, + num_tokens_main_model, + apply_admission_cap=apply_admission_cap, + ) + else: + # We don't allocate blocks for lookahead tokens in align mode, because if + # x * block_size tokens are scheduled, num_tokens is + # x * block_size + num_lookahead_tokens and breaks the alignment. + # We can ignore lookahead tokens because current draft models don't have + # mamba layers. + num_tokens = num_tokens_main_model + + # NOTE(tdouble): this is an over-estimate of how many blocks we need because + # num_tokens can include draft tokens that will later be rejected. + num_required_blocks = ( + cdiv(num_tokens, self.block_size) + self.num_speculative_blocks + ) + num_new_blocks = ( + num_required_blocks + - len(new_computed_blocks) + - len(self.req_to_blocks[request_id]) + ) + has_partial_hit = ( + self._has_partial_local_hit( + new_computed_blocks, num_local_computed_tokens + ) + or request_id in self._partial_hit_reqs + ) + if has_partial_hit: + num_new_blocks = max(num_new_blocks, 0) + 1 + checkpoint_block = int( + self._needs_internal_checkpoint( + request_id, num_tokens, total_computed_tokens + ) + ) + if not apply_admission_cap: + self._num_checkpoint_blocks[request_id] = checkpoint_block + if num_new_blocks > 0: + num_new_blocks = 1 + int(has_partial_hit) + checkpoint_block + if request_id not in self._allocated_block_reqs: + num_new_blocks += self.num_speculative_blocks + + num_evictable_computed_blocks = self._get_num_evictable_blocks( + new_computed_blocks + ) + return num_new_blocks + num_evictable_computed_blocks + + def allocate_new_blocks( + self, request_id: str, num_tokens: int, num_tokens_main_model: int + ) -> list[KVCacheBlock]: + assert isinstance(self.kv_cache_spec, MambaSpec) + planned = self._planned_recurrent_checkpoints.get(request_id) + if planned is not None: + start, end, targets = planned + if start != 0 or num_tokens_main_model != end or self.req_to_blocks.get(request_id): + raise ValueError("checkpoint allocation span changed after admission") + final_column = cdiv(end, self.block_size) - 1 + columns = [p // self.block_size - 1 for p in targets] + columns += [final_column] + columns += list(range(final_column + 1, final_column + 1 + self.num_speculative_blocks)) + if len(set(columns)) != len(columns): + raise ValueError("checkpoint output aliases another writable state") + physical = self.block_pool.get_new_blocks(len(columns)) + table = [self._null_block] * (final_column + 1 + self.num_speculative_blocks) + for column, block in zip(columns, physical): + table[column] = block + self.req_to_blocks[request_id].extend(table) + self._allocated_block_reqs.add(request_id) + return table + if self.mamba_cache_mode != "align": + # Allocate extra `num_speculative_blocks` blocks for + # speculative decoding (MTP/EAGLE) with linear attention. + if self.num_speculative_blocks > 0: + num_tokens += self.block_size * self.num_speculative_blocks + return super().allocate_new_blocks( + request_id, num_tokens, num_tokens_main_model + ) + else: + # We don't allocate blocks for lookahead tokens in align mode, because if + # x * block_size tokens are scheduled, num_tokens is + # x * block_size + num_lookahead_tokens and breaks the alignment. + # We can ignore lookahead tokens because current draft models don't have + # mamba layers. + num_tokens = num_tokens_main_model + req_blocks: list[KVCacheBlock] = self.req_to_blocks[request_id] + # NOTE(tdouble): this is an over-estimate of how many blocks we need because + # num_tokens can include draft tokens that will later be rejected. + num_required_blocks = ( + cdiv(num_tokens, self.block_size) + self.num_speculative_blocks + ) + checkpoint_block = self._num_checkpoint_blocks.get(request_id, 0) + partial_hit = self._partial_hit_reqs.get(request_id) + has_partial_hit = partial_hit is not None + # `num_required_blocks` might be less than `len(req_blocks)` if blocks are + # over-allocated at last round. + if num_required_blocks <= len(req_blocks) and not has_partial_hit: + self._allocated_block_reqs.add(request_id) + return [] + else: + prev_block_len = len(req_blocks) + blocks_allocated = request_id in self._allocated_block_reqs + # Record the last state block + if blocks_allocated: + # We always save the running state at the last + # (1 + num_speculative_blocks) block + self.last_state_block_idx[request_id] = ( + prev_block_len - 1 - self.num_speculative_blocks + ) + elif prev_block_len > 0: + # When a new request hits the prefix cache, the last block + # saves the hit state. + self.last_state_block_idx[request_id] = prev_block_len - 1 + + num_skipped_blocks = ( + num_required_blocks - self.num_speculative_blocks - 1 + ) + # null blocks + if prev_block_len < num_skipped_blocks: + # minus the internal checkpoint block + # so we don't set null for that block + null_end = num_skipped_blocks - checkpoint_block + req_blocks.extend( + [self._null_block for _ in range(prev_block_len, null_end)] + ) + + if blocks_allocated: + # reuse previous speculative blocks in this step + for block_idx in range( + prev_block_len - self.num_speculative_blocks, prev_block_len + ): + if block_idx < num_skipped_blocks: + req_blocks.append(req_blocks[block_idx]) + req_blocks[block_idx] = self._null_block + else: + break + num_new_blocks = num_required_blocks - len(req_blocks) + if has_partial_hit: + num_new_blocks = max(num_new_blocks, 0) + 1 + max_new_blocks = 1 + int(has_partial_hit) + checkpoint_block + if not blocks_allocated: + max_new_blocks += self.num_speculative_blocks + assert num_new_blocks <= max_new_blocks + new_blocks = self.block_pool.get_new_blocks(num_new_blocks) + returned_blocks = req_blocks[prev_block_len:] + if partial_hit is not None: + block_idx, source_block = partial_hit + cow_block = new_blocks[0] + new_blocks = new_blocks[1:] + if blocks_allocated: + # The worker block table of a running request is + # append-only, so the request must stay on + # source_block. Move the cache entry to cow_block + # instead; the queued copy fills it before forward + # overwrites source_block. + assert req_blocks[block_idx] is source_block + self.block_pool.move_block_hashes(source_block, cow_block) + self._pending_cow_copies.append((source_block, cow_block)) + source_block.ref_cnt += 1 + boundary_record = self._producer_partial_tail_reqs.pop( + request_id, None + ) + if boundary_record is not None: + boundary_tokens, boundary_hash = boundary_record + # This CoW preserved a producer's own boundary + # state in cow_block; hand it to the connector for + # partial-tail offload once the copy has run. + self._pending_partial_tail_offloads.append( + ( + request_id, + self.kv_cache_group_id, + cow_block, + boundary_tokens, + boundary_hash, + ) + ) + if cow_block.block_hash is not None: + # The moved entry is only filled by this step's + # copy, so defer same-step hits on it. + self.cached_blocks_this_step.add(cow_block.block_hash) + else: + self._apply_cow(request_id, block_idx, source_block, cow_block) + returned_blocks = [cow_block] + returned_blocks + req_blocks.extend(new_blocks) + self._allocated_block_reqs.add(request_id) + self._partial_hit_reqs.pop(request_id, None) + returned_blocks.extend(new_blocks) + return returned_blocks + + def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]: + if self.mamba_cache_mode == "align": + self._allocated_block_reqs.discard(request_id) + self.last_state_block_idx.pop(request_id, None) + self._num_checkpoint_blocks.pop(request_id, None) + self._producer_partial_tail_reqs.pop(request_id, None) + # A hand-off whose request died in this same scheduling pass must + # not reach the connector: its unpin hook (free) has already run. + self._pending_partial_tail_offloads = [ + entry + for entry in self._pending_partial_tail_offloads + if entry[0] != request_id + ] + self._pending_aligned_recurrent_boundaries = [ + entry + for entry in self._pending_aligned_recurrent_boundaries + if entry[0] != request_id + ] + return super().pop_blocks_for_free(request_id) + + def get_num_skipped_tokens(self, num_computed_tokens: int) -> int: + """ + Get the number of tokens whose mamba state are not needed anymore. Mamba only + need to keep the state of the last computed token, so we return + num_computed_tokens - 1. + """ + return num_computed_tokens - 1 + + def cache_blocks( + self, + request: Request, + num_tokens: int, + retention_interval: int | None = None, + ) -> None: + num_cached_blocks_before = self.num_cached_block.get(request.request_id, 0) + super().cache_blocks(request, num_tokens, retention_interval=retention_interval) + num_cached_blocks_after = self.num_cached_block.get(request.request_id, 0) + if self.mamba_cache_mode == "align": + partial_hash = self._cache_partial_tail_block(request, num_tokens) + if partial_hash is not None: + self.cached_blocks_this_step.add(partial_hash) + self._queue_aligned_recurrent_boundary(request, num_tokens) + for boundary in self._planned_recurrent_publications.get(request.request_id, ()): + self._queue_aligned_recurrent_boundary(request, boundary, expected_boundary=boundary) + if num_cached_blocks_after > num_cached_blocks_before: + for block in self.req_to_blocks[request.request_id][ + num_cached_blocks_before:num_cached_blocks_after + ]: + # Skip null blocks (align-mode skipped states) and blocks that + # were not cached this step — with sparse retention + # (reachable_block_mask) the intermediate state snapshots carry + # no hash and must not be recorded as cached-this-step. + if block.is_null or block.block_hash is None: + continue + self.cached_blocks_this_step.add(block.block_hash) + + def new_step_starts(self) -> None: + self.cached_blocks_this_step.clear() + + def _cache_partial_tail_block( + self, + request: Request, + num_tokens: int, + ) -> BlockHashWithGroupId | None: + hash_block_size = self.block_pool.hash_block_size + if self.block_size == hash_block_size: + return None + if num_tokens % self.block_size == 0: + return None + if num_tokens % hash_block_size != 0: + return None + publication_boundary = self.recurrent_publication_boundary + if publication_boundary is None: + publication_boundary = ( + request.num_prompt_tokens // hash_block_size + ) * hash_block_size + if num_tokens != publication_boundary: + return None + + block_idx = num_tokens // self.block_size + blocks = self.req_to_blocks[request.request_id] + if block_idx >= len(blocks): + return None + source_block = blocks[block_idx] + if source_block.is_null: + return None + + partial_hash = self.block_pool.cache_partial_block( + request=request, + block=source_block, + num_tokens=num_tokens, + kv_cache_group_id=self.kv_cache_group_id, + block_size=self.block_size, + ) + if partial_hash is not None: + self._partial_hit_reqs[request.request_id] = (block_idx, source_block) + self.num_cached_block[request.request_id] = block_idx + # Producer of this partial tail: the boundary state currently lives + # in ``source_block`` but the next step's forward overwrites it. The + # upcoming CoW copies it into a durable cow_block; record the req so + # allocate_new_blocks hands that block to the connector for offload. + self._producer_partial_tail_reqs[request.request_id] = ( + num_tokens, + partial_hash, + ) + return partial_hash + + def _queue_aligned_recurrent_boundary( + self, + request: Request, + num_tokens: int, + expected_boundary: int | None = None, + ) -> None: + """Queue a full Mamba page only when its hash proves the replay boundary.""" + if num_tokens <= 0 or num_tokens % self.block_size != 0: + return + replay_boundary = expected_boundary if expected_boundary is not None else ( + self.recurrent_publication_boundary + if self.recurrent_publication_boundary is not None + else ((request.num_prompt_tokens - 1) // self.block_pool.hash_block_size) + * self.block_pool.hash_block_size + ) + if num_tokens != replay_boundary: + return + + block_idx = num_tokens // self.block_size - 1 + blocks = self.req_to_blocks[request.request_id] + if block_idx >= len(blocks): + return + block = blocks[block_idx] + if ( + block.is_null + or block.block_hash is None + or block.block_hash_num_tokens != replay_boundary + or get_group_id(block.block_hash) != self.kv_cache_group_id + ): + return + self._pending_aligned_recurrent_boundaries.append( + ( + request.request_id, + self.kv_cache_group_id, + block, + replay_boundary, + ) + ) + + +class CrossAttentionManager(SingleTypeKVCacheManager): + """Manager for cross-attention KV cache in encoder-decoder models.""" + + def add_local_computed_blocks( + self, + request_id: str, + new_computed_blocks: Sequence[KVCacheBlock], + num_local_computed_tokens: int, + num_external_computed_tokens: int, + ) -> None: + # We do not cache blocks for cross-attention to be shared between + # requests, so `new_computed_blocks` should always be empty. + assert len(new_computed_blocks) == 0 + + def allocate_external_computed_blocks( + self, + request_id: str, + num_local_computed_tokens: int, + num_external_computed_tokens: int, + ) -> None: + # Cross-attention does not use prefix caching / external KV loads. + return + + def cache_blocks( + self, + request: Request, + num_tokens: int, + retention_interval: int | None = None, + ) -> None: + # We do not cache blocks for cross-attention to be shared between + # requests, so this method is not relevant. + raise ValueError("Should not be called as prefix caching is disabled.") + + def get_num_common_prefix_blocks(self, running_request_id: str) -> int: + # Cross-attention blocks contain request-specific encoder states + # and are not shared between different requests + return 0 + + @classmethod + def find_longest_cache_hit( + cls, + block_hashes: BlockHashList, + max_length: int, + kv_cache_group_ids: list[int], + block_pool: BlockPool, + kv_cache_spec: KVCacheSpec, + drop_eagle_block: bool, + alignment_tokens: int, + dcp_world_size: int = 1, + pcp_world_size: int = 1, + ) -> tuple[tuple[list[KVCacheBlock], ...], int]: + assert isinstance(kv_cache_spec, CrossAttentionSpec), ( + "CrossAttentionManager can only be used for cross-attention groups" + ) + # Cross-attention does not benefit from prefix caching since: + # 1. Encoder states are unique per request (different audio/image + # inputs) + # 2. Encoder states are computed once per request, not incrementally + # 3. No reusable prefix exists between different multimodal inputs + # Return empty blocks to indicate no cache hits + raise NotImplementedError("CrossAttentionManager does not support caching") + + +class SinkFullAttentionManager(FullAttentionManager): + def __init__( + self, + kv_cache_spec: SinkFullAttentionSpec, + block_pool: BlockPool, + enable_caching: bool, + kv_cache_group_id: int, + scheduler_block_size: int, + dcp_world_size: int = 1, + pcp_world_size: int = 1, + ): + super().__init__( + kv_cache_spec=kv_cache_spec, + block_pool=block_pool, + enable_caching=enable_caching, + kv_cache_group_id=kv_cache_group_id, + scheduler_block_size=scheduler_block_size, + dcp_world_size=dcp_world_size, + pcp_world_size=pcp_world_size, + ) + sink_len = kv_cache_spec.sink_len + assert sink_len is not None and sink_len > 0 and sink_len % self.block_size == 0 + num_sink_block = sink_len // self.block_size + self.sink_blocks = self.block_pool.free_block_queue.popleft_n(num_sink_block) + + +def get_manager_for_kv_cache_spec( + kv_cache_spec: KVCacheSpec, + max_in_flight_tokens: int, + max_model_len: int, + **kwargs, +) -> SingleTypeKVCacheManager: + """ + Get the appropriate manager for a given KVCacheSpec. + + Uses the KVCacheSpecRegistry to look up the manager class, supporting + both built-in and custom specs registered via @register_kv_cache_spec + and KVCacheSpecRegistry.register. + + Args: + kv_cache_spec: The KVCacheSpec instance + max_in_flight_tokens: The max tokens scheduled but not yet settled + (one batch per concurrent step); see `VllmConfig.max_in_flight_tokens` + max_model_len: The maximum context length the model could serve + Returns: + An instance of the appropriate SingleTypeKVCacheManager subclass + """ + manager_class = KVCacheSpecRegistry.get_manager_class(kv_cache_spec) + assert manager_class is not None, ( + f"No manager registered for KVCacheSpec {type(kv_cache_spec)}" + ) + # SlidingWindow / ChunkedLocalAttention managers recycle blocks; + # the runtime admission cap must match the recycling-aware bound the + # startup pool sizer uses (single source of truth: the spec method). + # R-SWA also recycles gap blocks but peak physical KV still fits the + # full-attention bound (prefix + window <= max_model_len), so it inherits + # FullAttentionSpec sizing without a separate admission cap. + if isinstance( + kv_cache_spec, + (SlidingWindowSpec, ChunkedLocalAttentionSpec), + ): + kwargs["max_admission_blocks_per_request"] = ( + kv_cache_spec.max_admission_blocks_per_request( + max_in_flight_tokens=max_in_flight_tokens, + max_model_len=max_model_len, + ) + ) + manager = manager_class(kv_cache_spec, **kwargs) + return manager + + +def register_all_kvcache_specs(vllm_config): + """Built-in spec registration""" + KVCacheSpecRegistry.register( + FullAttentionSpec, + FullAttentionManager, + uniform_type_base_spec=FullAttentionSpec, + ) + + KVCacheSpecRegistry.register( + SlidingWindowSpec, + SlidingWindowManager, + uniform_type_base_spec=SlidingWindowSpec, + ) + KVCacheSpecRegistry.register( + SlidingWindowMLASpec, + SlidingWindowManager, + uniform_type_base_spec=SlidingWindowMLASpec, + ) + + KVCacheSpecRegistry.register( + MambaSpec, MambaManager, uniform_type_base_spec=MambaSpec + ) + KVCacheSpecRegistry.register( + ChunkedLocalAttentionSpec, + ChunkedLocalAttentionManager, + uniform_type_base_spec=ChunkedLocalAttentionSpec, + ) + KVCacheSpecRegistry.register( + CrossAttentionSpec, + CrossAttentionManager, + uniform_type_base_spec=CrossAttentionSpec, + ) + + # FullAttentionSpec subclasses — grouped with FullAttentionSpec + KVCacheSpecRegistry.register( + MLAAttentionSpec, FullAttentionManager, uniform_type_base_spec=FullAttentionSpec + ) + KVCacheSpecRegistry.register( + RSWASpec, RSWAManager, uniform_type_base_spec=FullAttentionSpec + ) + # NOTE(Mengqing): HiddenStateCacheSpec won't take part in + # grouping, thus the uniform_type_base_spec is just a + # placeholder. + KVCacheSpecRegistry.register( + HiddenStateCacheSpec, + FullAttentionManager, + uniform_type_base_spec=FullAttentionSpec, + ) + KVCacheSpecRegistry.register( + SinkFullAttentionSpec, + SinkFullAttentionManager, + uniform_type_base_spec=FullAttentionSpec, + ) + + from vllm.platforms import current_platform + + current_platform.register_custom_kv_cache_specs(vllm_config) diff --git a/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/347a15ec42368b0c2b35b31cef7c9f58da97814a53ee76d0ee5526aac6b15c1c/kimi_gdn_linear_attn.py b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/347a15ec42368b0c2b35b31cef7c9f58da97814a53ee76d0ee5526aac6b15c1c/kimi_gdn_linear_attn.py new file mode 100644 index 00000000..50e583c5 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/347a15ec42368b0c2b35b31cef7c9f58da97814a53ee76d0ee5526aac6b15c1c/kimi_gdn_linear_attn.py @@ -0,0 +1,1681 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Callable +import os +from dataclasses import replace +from typing import TYPE_CHECKING, Any + + +def _prefill_checkpoint_capacity() -> int: + value = os.environ.get("SPARK_GDN_PREFILL_CHECKPOINTS", "1") + if value not in ("1", "2"): + raise ValueError("SPARK_GDN_PREFILL_CHECKPOINTS must be 1 or 2") + return int(value) + +import torch +from einops import rearrange +from torch import nn +from torch.nn.parameter import Parameter + +from vllm.compilation.breakable_cudagraph import eager_break_during_capture +from vllm.config import VllmConfig +from vllm.distributed import divide, get_tensor_model_parallel_rank +from vllm.forward_context import get_forward_context +from vllm.model_executor.custom_op import PluggableLayer +from vllm.model_executor.layers.mamba.gdn.base import GatedDeltaNetAttention +from vllm.model_executor.model_loader.weight_utils import ( + default_weight_loader, + sharded_weight_loader, +) +from vllm.model_executor.parameter import BasevLLMParameter +from vllm.model_executor.utils import set_weight_attrs +from vllm.platforms import current_platform +from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated +from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig +from vllm.triton_utils import tl, triton +from vllm.utils.b12x import ( + B12xWarmupUnit, + get_b12x_gdn_decode, + get_b12x_kda_prefill, + get_b12x_scratch_buffers, +) +from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID +from vllm.v1.kv_cache_interface import MambaSpec +from vllm.v1.worker.workspace import current_workspace_manager + +from ...linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from ..mamba_utils import ( + MambaStateDtypeCalculator, + MambaStateShapeCalculator, + is_conv_state_dim_first, +) +from ..ops.causal_conv1d import causal_conv1d_fn, causal_conv1d_update +from ..ops.gather_initial_states import gather_initial_states + +# Empirical lower bound for the KDA gate to avoid numerical underflow. +_KDA_GATE_LOGBOUND_MIN = -5.0 + + +def is_flashkda_supported( + head_dim: int, + dtype: torch.dtype, + lower_bound: float | None, +) -> bool: + """Return whether FlashKDA supports the layer's prefill contract.""" + if not current_platform.is_cuda(): + return False + capability = current_platform.get_device_capability() + return ( + capability is not None + and capability.major in (9, 10, 12) + and head_dim == 128 + and dtype == torch.bfloat16 + and lower_bound is not None + ) + + +def is_b12x_kda_prefill_supported( + head_dim: int, + dtype: torch.dtype, + lower_bound: float | None, + state_dtype: torch.dtype, +) -> bool: + """Return whether the b12x KDA prefill op supports the layer's contract.""" + api = get_b12x_kda_prefill() + if api is None or not current_platform.is_cuda(): + return False + return ( + head_dim == 128 + and dtype == torch.bfloat16 + and state_dtype == torch.float32 + and lower_bound is not None + and api.is_supported(torch.device(current_platform.current_device())) + ) + + +def _flashkda_prefill( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + lower_bound: float, + initial_state: torch.Tensor, + cu_seqlens: torch.Tensor, + out: torch.Tensor, + final_state: torch.Tensor, + workspace: torch.Tensor, + checkpoint_state: torch.Tensor | None = None, + checkpoint_offsets: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run packed bounded-gate KDA prefill into caller-owned buffers.""" + import vllm._flashkda_C # noqa: F401 + + torch.ops._flashkda_C.fwd( + q.contiguous(), + k.contiguous(), + v.contiguous(), + g.contiguous(), + beta, + q.shape[-1] ** -0.5, + out, + workspace, + A_log.contiguous(), + dt_bias.view(-1, q.shape[-1]).contiguous(), + lower_bound, + initial_state.contiguous(), + final_state, + cu_seqlens.contiguous(), + checkpoint_state, + checkpoint_offsets.contiguous() if checkpoint_offsets is not None else None, + ) + return out, final_state + + +@triton.jit +def _store_cache_checkpoints_kernel( + x_ptr, + conv_state_ptr, + recurrent_checkpoint_ptr, + recurrent_state_ptr, + query_start_loc_ptr, + checkpoint_offsets_ptr, + checkpoint_state_indices_ptr, + x_stride_0: tl.constexpr, + x_stride_1: tl.constexpr, + state_stride_0: tl.constexpr, + state_stride_1: tl.constexpr, + state_stride_2: tl.constexpr, + checkpoint_stride_0: tl.constexpr, + recurrent_state_stride_0: tl.constexpr, + checkpoint_offset_stride: tl.constexpr, + STATE_LEN: tl.constexpr, + WIDTH: tl.constexpr, + RECURRENT_ROW_SIZE: tl.constexpr, + NULL_STATE_IDX: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + STORE_RECURRENT: tl.constexpr, + CHECKPOINTS: tl.constexpr = 1, + error_code_ptr=None, + CHECK_ERROR: tl.constexpr = False, +): + """Store FlashKDA recurrent and convolution state at an internal boundary.""" + checkpoint_idx = tl.program_id(0) + seq_idx = checkpoint_idx // CHECKPOINTS + cols = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + seq_idx_i64 = seq_idx.to(tl.int64) + cols_i64 = cols.to(tl.int64) + state_idx = tl.load(checkpoint_state_indices_ptr + checkpoint_idx.to(tl.int64)) + state_idx_i64 = state_idx.to(tl.int64) + checkpoint_offset = tl.load( + checkpoint_offsets_ptr + + seq_idx_i64 * checkpoint_offset_stride + + (checkpoint_idx % CHECKPOINTS).to(tl.int64) + ) + valid_checkpoint = (state_idx != NULL_STATE_IDX) & (checkpoint_offset > 0) + if CHECK_ERROR: + valid_checkpoint = valid_checkpoint & (tl.load(error_code_ptr) == 0) + valid_conv = ( + (cols < WIDTH * STATE_LEN) & valid_checkpoint & (checkpoint_offset >= STATE_LEN) + ) + width_idx = cols // STATE_LEN + history_idx = cols % STATE_LEN + checkpoint_end = tl.load(query_start_loc_ptr + seq_idx_i64) + checkpoint_offset + token_idx = checkpoint_end.to(tl.int64) - STATE_LEN + history_idx.to(tl.int64) + values = tl.load( + x_ptr + token_idx * x_stride_0 + width_idx.to(tl.int64) * x_stride_1, + mask=valid_conv, + ) + tl.store( + conv_state_ptr + + state_idx_i64 * state_stride_0 + + width_idx.to(tl.int64) * state_stride_1 + + history_idx.to(tl.int64) * state_stride_2, + values, + mask=valid_conv, + ) + + if STORE_RECURRENT: + valid_recurrent = (cols < RECURRENT_ROW_SIZE) & valid_checkpoint + recurrent = tl.load( + recurrent_checkpoint_ptr + seq_idx_i64 * checkpoint_stride_0 + cols_i64, + mask=valid_recurrent, + ) + tl.store( + recurrent_state_ptr + state_idx_i64 * recurrent_state_stride_0 + cols_i64, + recurrent, + mask=valid_recurrent, + ) + + +def resolve_kda_prefill_backend( + backend: str, + head_dim: int, + dtype: torch.dtype, + lower_bound: float | None, + state_dtype: torch.dtype = torch.float32, +) -> str: + """Resolve the packed KDA prefill implementation for one server. + + ``auto`` never selects ``b12x``; that backend must be requested by name + until its serving qualification lands. + """ + if backend not in ("auto", "triton", "flashkda", "b12x"): + raise ValueError(f"Unsupported KDA prefill backend: {backend}") + if backend == "b12x": + if not is_b12x_kda_prefill_supported(head_dim, dtype, lower_bound, state_dtype): + raise RuntimeError( + "The b12x KDA prefill backend requires the b12x package on a " + "supported CUDA device, bfloat16 activations, head_dim=128, " + "float32 recurrent state, and a bounded KDA gate." + ) + return "b12x" + supported = is_flashkda_supported(head_dim, dtype, lower_bound) + if backend == "flashkda" and not supported: + raise RuntimeError( + "FlashKDA requires CUDA SM90/SM10x/SM12x, bfloat16, " + "head_dim=128, and a bounded KDA gate." + ) + if supported and backend != "triton": + return "flashkda" + return "triton" + + +def a_log_weight_loader( + shard_axis: int, +) -> Callable[[torch.Tensor, torch.Tensor], None]: + """Load KDA A_log stored as either old 4D or current 1D weights.""" + + def loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: + tp_rank = get_tensor_model_parallel_rank() + shard_size = param.data.shape[shard_axis] + start_idx = tp_rank * shard_size + + if loaded_weight.dim() == 4: + assert loaded_weight.shape[:2] == (1, 1), ( + f"Expected old A_log shape (1, 1, H, 1), got {loaded_weight.shape}" + ) + assert loaded_weight.shape[-1] == 1, ( + f"Expected old A_log last dim to be 1, got {loaded_weight.shape}" + ) + loaded_weight = loaded_weight.view(loaded_weight.shape[2]) + + loaded_weight = loaded_weight.narrow(shard_axis, start_idx, shard_size) + return default_weight_loader(param, loaded_weight) + + return loader + + +def _make_fused_conv1d_weight_loader( + dims: list[int], + tp_size: int, + tp_rank: int, +) -> Callable[..., None]: + sharded_dims = [dim // tp_size for dim in dims] + + def weight_loader( + param: torch.Tensor, + loaded_weight: torch.Tensor, + loaded_shard_id: int, + ) -> None: + if loaded_weight.dim() == 2: + loaded_weight = loaded_weight.unsqueeze(1) + shard_size = sharded_dims[loaded_shard_id] + source_start = tp_rank * shard_size + target_start = sum(sharded_dims[:loaded_shard_id]) + loaded_shard = loaded_weight[source_start : source_start + shard_size] + param.data[target_start : target_start + shard_size].copy_(loaded_shard) + + return weight_loader + + +class _KimiGDNMergedColumnParallelLinear(MergedColumnParallelLinear): + """Merged projection with one output replicated across TP ranks. + + The replicated shard is represented as ``size * tp_size`` so the merged + parameter reserves ``size`` local rows on every rank. Loading that shard + from rank zero then gives every rank the complete checkpoint weight. + """ + + def __init__( + self, + input_size: int, + output_sizes: list[int], + replicated_shard_id: int, + tp_size: int, + **kwargs, + ) -> None: + self.replicated_shard_id = replicated_shard_id + output_sizes = output_sizes.copy() + output_sizes[replicated_shard_id] *= tp_size + super().__init__(input_size, output_sizes, **kwargs) + + def weight_loader( + self, + param: Parameter, + loaded_weight: torch.Tensor, + loaded_shard_id: tuple[int, ...] | int | None = None, + ) -> None: + tp_rank = self.tp_rank + param_tp_rank = getattr(param, "tp_rank", None) + if loaded_shard_id == self.replicated_shard_id: + self.tp_rank = 0 + if param_tp_rank is not None: + param.tp_rank = 0 + try: + super().weight_loader(param, loaded_weight, loaded_shard_id) + finally: + self.tp_rank = tp_rank + if param_tp_rank is not None: + param.tp_rank = param_tp_rank + + def weight_loader_v2( + self, + param: BasevLLMParameter, + loaded_weight: torch.Tensor, + loaded_shard_id: tuple[int, ...] | int | None = None, + ) -> None: + tp_rank = self.tp_rank + param_tp_rank = getattr(param, "tp_rank", None) + if loaded_shard_id == self.replicated_shard_id: + self.tp_rank = 0 + if param_tp_rank is not None: + param.tp_rank = 0 + try: + super().weight_loader_v2(param, loaded_weight, loaded_shard_id) + finally: + self.tp_rank = tp_rank + if param_tp_rank is not None: + param.tp_rank = param_tp_rank + + +class _B12xKdaPrefillWarmup: + """Warm-up provider that compiles a layer's b12x KDA prefill kernels.""" + + def get_b12x_warmup_unit( + self, + layer: torch.nn.Module, + token_counts: tuple[int, ...], + output_dtype: torch.dtype, + ) -> B12xWarmupUnit: + del token_counts, output_dtype + + def compile() -> None: + plan = layer._b12x_prefill_plan + scratch = layer._b12x_prefill_scratch + api = layer._b12x_prefill_api + if plan is None or scratch is None or api is None: + # The pool is bound after the memory-profiling pass; the + # post-allocation warmup compiles this layer. + return + caps = plan.caps + device = caps.device + heads, head_dim = caps.heads, caps.head_dim + tokens = caps.chunk_tokens + rows = torch.zeros( + (tokens, heads, head_dim), dtype=caps.model_dtype, device=device + ) + indices = torch.zeros(1, dtype=torch.int32, device=device) + checkpoint_shape = ((1, 2) if getattr(caps, "max_checkpoints", 1) == 2 else (1,)) + binding = api.bind( + plan, + scratch=scratch, + q=rows, + k=torch.zeros_like(rows), + v=torch.zeros_like(rows), + raw_g=torch.zeros_like(rows), + raw_beta=torch.zeros( + (tokens, heads), dtype=caps.model_dtype, device=device + ), + A_log=layer.A_log, + dt_bias=layer.dt_bias.view(-1, head_dim), + recurrent_state=layer.kv_cache[1], + cu_seqlens=torch.tensor([0, tokens], dtype=torch.int32, device=device), + initial_state_indices=indices, + final_state_indices=indices, + checkpoint_state_indices=torch.full(checkpoint_shape, NULL_BLOCK_ID, dtype=torch.int32, device=device), + checkpoint_offsets=torch.zeros(checkpoint_shape, dtype=torch.int32, device=device), + num_seqs=layer._b12x_prefill_num_seqs, + num_tokens=layer._b12x_prefill_num_tokens, + output=torch.zeros_like(rows), + ) + api.prewarm(binding) + + caps = getattr(layer._b12x_prefill_plan, "caps", None) + return B12xWarmupUnit( + name="KDA prefill", + key=( + type(layer), + None if caps is None else caps.device, + layer.local_num_heads, + layer.head_dim, + layer._b12x_prefill_max_tokens, + layer._b12x_prefill_max_seqs, + None if caps is None else caps.max_state_slots, + ), + compile=compile, + ) + + +@PluggableLayer.register("kimi_gated_delta_net_attention") +class KimiGatedDeltaNetAttention(GatedDeltaNetAttention): + enable_b12x_kda_decode = False + b12x_kda_null_state_index: int | None = None + + def get_state_dtype( + self, + ) -> tuple[torch.dtype, torch.dtype]: + if self.model_config is None or self.cache_config is None: + raise ValueError("model_config and cache_config must be set") + return MambaStateDtypeCalculator.kda_state_dtype( + self.model_config.dtype, self.cache_config.mamba_cache_dtype + ) + + def get_state_shape( + self, + ) -> tuple[tuple[int, ...], tuple[int, ...]]: + return MambaStateShapeCalculator.kda_state_shape( + self.tp_size, + self.num_heads, + self.head_dim, + conv_kernel_size=self.conv_size, + num_spec=self.num_spec, + ) + + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> MambaSpec: + spec = super().get_kv_cache_spec(vllm_config) + assert isinstance(spec, MambaSpec) + capacity = _prefill_checkpoint_capacity() + if capacity == 2 and self.kda_prefill_backend != "b12x": + raise ValueError("Two internal checkpoints require the B12X KDA backend") + return replace( + spec, + num_prefill_checkpoint_blocks=capacity * int( + self.kda_prefill_backend in ("flashkda", "b12x") + ), + ) + + def __init__( + self, + config: KimiLinearConfig, + vllm_config: VllmConfig, + prefix: str = "", + ) -> None: + super().__init__(config, vllm_config, prefix) + + kda_config = config.linear_attn_config # type: ignore[attr-defined] + assert kda_config is not None, "linear_attn_config must be set" + self.head_dim = kda_config["head_dim"] + self.num_heads = kda_config["num_heads"] + assert self.num_heads % self.tp_size == 0 + self.local_num_heads = divide(self.num_heads, self.tp_size) + + self.projection_size = self.head_dim * self.num_heads + self.local_projection_size = divide(self.projection_size, self.tp_size) + self.conv_size = kda_config["short_conv_kernel_size"] + self.use_full_rank_gate = kda_config.get("use_full_rank_gate", False) + + if self.use_full_rank_gate: + # Keep f_a before the narrow beta shard, then pad each TP-local row + # to select the aligned BF16 GEMM path. The padding also avoids an + # Inductor correctness issue seen with the row-strided G view. + qkvg_output_sizes = [self.projection_size] * 4 + in_proj_output_sizes = qkvg_output_sizes + [ + self.head_dim, + self.num_heads, + ] + local_output_size = ( + 4 * self.local_projection_size + self.head_dim + self.local_num_heads + ) + self.in_proj_padding = -local_output_size % 16 + if self.in_proj_padding: + in_proj_output_sizes.append(self.in_proj_padding * self.tp_size) + else: + in_proj_output_sizes = [self.projection_size] * 3 + [ + self.num_heads, + self.head_dim, + ] + self.in_proj_padding = 0 + self.in_proj_qkvgfab = _KimiGDNMergedColumnParallelLinear( + self.hidden_size, + in_proj_output_sizes, + replicated_shard_id=4, + tp_size=self.tp_size, + bias=False, + quant_config=self.quant_config, + prefix=f"{prefix}.in_proj_qkvgfab", + ) + if self.in_proj_padding: + self.in_proj_qkvgfab.weight.data[-self.in_proj_padding :].zero_() + + self.f_b_proj = ColumnParallelLinear( + self.head_dim, + self.projection_size, + bias=False, + quant_config=self.quant_config, + prefix=f"{prefix}.f_b_proj", + ) + self.dt_bias = nn.Parameter( + torch.empty(self.local_projection_size, dtype=torch.float32) + ) + + set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)}) + + # One packed parameter and cache let decode run a single conv update. + # Prefill slices them back into Q/K/V to obtain dense outputs cheaply. + self.conv1d = ColumnParallelLinear( + input_size=self.conv_size, + output_size=3 * self.projection_size, + bias=False, + params_dtype=torch.float32, + prefix=f"{prefix}.conv1d", + ) + self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1) + delattr(self.conv1d.weight, "weight_loader") + set_weight_attrs( + self.conv1d.weight, + { + "weight_loader": _make_fused_conv1d_weight_loader( + [self.projection_size] * 3, + self.tp_size, + self.tp_rank, + ) + }, + ) + + self.A_log = nn.Parameter( + torch.empty(self.local_num_heads, dtype=torch.float32) + ) + set_weight_attrs(self.A_log, {"weight_loader": a_log_weight_loader(0)}) + + self.gate_lower_bound: float | None = kda_config.get("gate_lower_bound", None) + if self.gate_lower_bound is not None: + assert _KDA_GATE_LOGBOUND_MIN <= self.gate_lower_bound < 0, ( + "KDA gate lower bound must be in " + f"[{_KDA_GATE_LOGBOUND_MIN}, 0). " + f"Got {self.gate_lower_bound}." + ) + self.use_safe_gate = self.gate_lower_bound is not None + additional_config = vllm_config.additional_config + backend = ( + additional_config.get("kda_prefill_backend", "auto") + if isinstance(additional_config, dict) + else "auto" + ) + self.kda_prefill_backend = resolve_kda_prefill_backend( + backend, + self.head_dim, + vllm_config.model_config.dtype, + self.gate_lower_bound, + self.get_state_dtype()[1], + ) + self._flashkda_buffer_specs: ( + tuple[tuple[tuple[int, ...], torch.dtype], ...] | None + ) = None + if self.kda_prefill_backend == "flashkda": + max_tokens = vllm_config.scheduler_config.max_num_batched_tokens + max_sequences = vllm_config.scheduler_config.max_num_seqs + heads, head_dim = self.local_num_heads, self.head_dim + import vllm._flashkda_C # noqa: F401 + + workspace_size = torch.ops._flashkda_C.get_workspace_size( + max_tokens, + heads, + max_sequences, + ) + self._flashkda_buffer_specs = ( + ((1, max_tokens, heads, head_dim), self.model_config.dtype), + ( + (max_sequences, heads, head_dim, head_dim), + self.get_state_dtype()[1], + ), + ( + (max_sequences, heads, head_dim, head_dim), + self.get_state_dtype()[1], + ), + ((workspace_size,), torch.uint8), + ) + if not self.use_full_rank_gate: + self.g_a_proj = ReplicatedLinear( + self.hidden_size, + self.head_dim, + bias=False, + quant_config=self.quant_config, + prefix=f"{prefix}.g_a_proj", + ) + self.g_b_proj = ColumnParallelLinear( + self.head_dim, + self.projection_size, + bias=False, + quant_config=self.quant_config, + prefix=f"{prefix}.g_b_proj", + ) + self.o_norm = FusedRMSNormGated(self.head_dim, activation="sigmoid") + self._b12x_kda_api: Any | None = None + self._b12x_kda_plan = None + self._initialize_b12x_kda_decode(vllm_config) + self._b12x_prefill_api: Any | None = None + self._b12x_prefill_plan = None + self._b12x_prefill_scratch = None + self._initialize_b12x_kda_prefill(vllm_config) + self.o_proj = RowParallelLinear( + self.projection_size, + self.hidden_size, + bias=False, + quant_config=self.quant_config, + prefix=f"{prefix}.o_proj", + ) + + 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 _initialize_b12x_kda_decode(self, vllm_config: VllmConfig) -> None: + if ( + not self.enable_b12x_kda_decode + or self.gate_lower_bound is None + or self.head_dim != 128 + or self.model_config.dtype != torch.bfloat16 + or self.get_state_dtype()[1] not in (torch.bfloat16, torch.float32) + or not current_platform.is_cuda() + ): + return + + api = get_b12x_gdn_decode() + device = torch.device(current_platform.current_device()) + if ( + api is None + or not hasattr(api, "bind_kda") + or not hasattr(api, "run_kda") + or not api.is_supported(device) + ): + return + + max_seqs = int(vllm_config.scheduler_config.max_num_seqs) + state_index_columns = max(1, self.num_spec + 1) + if state_index_columns > 8: + return + max_tokens = max_seqs * state_index_columns + + self._b12x_kda_api = api + self._b12x_kda_max_tokens = max_tokens + self._b12x_kda_max_seqs = max_seqs + self._b12x_kda_state_index_columns = state_index_columns + + self.register_buffer( + "_b12x_kda_num_accepted_tokens", + torch.ones(max_seqs, dtype=torch.int32, device=device), + persistent=False, + ) + self.register_buffer( + "_b12x_kda_num_seqs", + torch.zeros(1, dtype=torch.int32, device=device), + persistent=False, + ) + self.register_buffer( + "_b12x_kda_num_tokens", + torch.zeros(1, dtype=torch.int32, device=device), + persistent=False, + ) + self.register_buffer("_b12x_kda_scratch", None, persistent=False) + + def _make_b12x_kda_plan(self, max_state_slots: int): + api = self._b12x_kda_api + if api is None: + raise RuntimeError("b12x KDA decode was not initialized") + return api.plan( + api.Caps( + device=current_platform.current_device(), + max_tokens=self._b12x_kda_max_tokens, + max_seqs=self._b12x_kda_max_seqs, + max_state_slots=max_state_slots, + key_heads=self.local_num_heads, + value_heads=self.local_num_heads, + key_head_dim=self.head_dim, + value_head_dim=self.head_dim, + state_index_columns=self._b12x_kda_state_index_columns, + model_dtype=self.model_config.dtype, + state_dtype=self.get_state_dtype()[1], + gate_activation="sigmoid", + qk_l2norm=True, + null_state_index=self.b12x_kda_null_state_index, + kda_metadata_validation="trusted", + ) + ) + + def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: + super().bind_kv_cache(kv_cache) + if self._b12x_prefill_api is not None: + prefill_plan = self._make_b12x_kda_prefill_plan( + max_state_slots=int(self.kv_cache[1].shape[0]) + ) + self._b12x_prefill_plan = prefill_plan + prefill_scratch, _ = self._get_b12x_prefill_workspace() + self._b12x_prefill_scratch = prefill_scratch + api = self._b12x_kda_api + if api is None: + return + recurrent_state = self.kv_cache[1] + plan = self._make_b12x_kda_plan(max_state_slots=recurrent_state.shape[0]) + (scratch,) = get_b12x_scratch_buffers(plan) + self._b12x_kda_scratch = scratch + self._b12x_kda_plan = plan + + def unbind_kv_cache(self) -> None: + self._b12x_kda_plan = None + self._b12x_kda_scratch = None + self._b12x_prefill_plan = None + self._b12x_prefill_scratch = None + super().unbind_kv_cache() + + def _initialize_b12x_kda_prefill(self, vllm_config: VllmConfig) -> None: + """Hold the b12x prefill op and its per-request metadata buffers.""" + if self.kda_prefill_backend != "b12x": + return + api = get_b12x_kda_prefill() + if api is None: + raise RuntimeError( + "The b12x KDA prefill backend requires the b12x package." + ) + device = torch.device(current_platform.current_device()) + scheduler_config = vllm_config.scheduler_config + self._b12x_prefill_api = api + self._b12x_prefill_max_tokens = int(scheduler_config.max_num_batched_tokens) + self._b12x_prefill_max_seqs = int(scheduler_config.max_num_seqs) + max_seqs = self._b12x_prefill_max_seqs + checkpoint_shape = ((max_seqs, 2) if _prefill_checkpoint_capacity() == 2 else (max_seqs,)) + self.register_buffer( + "_b12x_prefill_num_seqs", + torch.zeros(1, dtype=torch.int32, device=device), + persistent=False, + ) + self.register_buffer( + "_b12x_prefill_num_tokens", + torch.zeros(1, dtype=torch.int32, device=device), + persistent=False, + ) + self.register_buffer( + "_b12x_prefill_initial_indices", + torch.zeros(max_seqs, dtype=torch.int32, device=device), + persistent=False, + ) + self.register_buffer( + "_b12x_prefill_null_indices", + torch.full(checkpoint_shape, NULL_BLOCK_ID, dtype=torch.int32, device=device), + persistent=False, + ) + self.register_buffer( + "_b12x_prefill_zero_offsets", + torch.zeros(checkpoint_shape, dtype=torch.int32, device=device), + persistent=False, + ) + self.b12x_warmup_provider = _B12xKdaPrefillWarmup() + + def _make_b12x_kda_prefill_plan(self, max_state_slots: int): + api = self._b12x_prefill_api + if api is None: + raise RuntimeError("b12x KDA prefill was not initialized") + capacity = _prefill_checkpoint_capacity() + if capacity == 2 and "max_checkpoints" not in getattr(api.Caps, "__dataclass_fields__", {}): + raise RuntimeError("Two checkpoints require the matching B12X capability implementation") + return api.plan( + api.Caps( + device=current_platform.current_device(), + max_tokens=self._b12x_prefill_max_tokens, + max_seqs=self._b12x_prefill_max_seqs, + max_state_slots=max_state_slots, + heads=self.local_num_heads, + head_dim=self.head_dim, + model_dtype=self.model_config.dtype, + state_dtype=self.get_state_dtype()[1], + qk_l2norm=True, + checkpoint_export=True, + null_state_index=NULL_BLOCK_ID, + metadata_validation="transactional" if capacity == 2 else "trusted", + **({"max_checkpoints": 2} if capacity == 2 else {}), + ) + ) + + def _get_b12x_prefill_workspace(self) -> tuple[torch.Tensor, torch.Tensor]: + plan = self._b12x_prefill_plan + if plan is None: + raise RuntimeError("b12x KDA prefill KV cache is not bound") + scratch_specs = tuple(plan.scratch_specs()) + if len(scratch_specs) != 1: + raise RuntimeError("b12x KDA prefill requires exactly one scratch buffer") + scratch_spec = scratch_specs[0] + scratch, output = current_workspace_manager().get_simultaneous( + (scratch_spec.shape, scratch_spec.dtype), + ( + ( + self._b12x_prefill_max_tokens, + self.local_num_heads, + self.head_dim, + ), + self.model_config.dtype, + ), + ) + return scratch, output + + def _run_b12x_kda_prefill( + self, + *, + scratch: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + cu_seqlens: torch.Tensor, + state_indices: torch.Tensor, + has_initial_state: torch.Tensor, + checkpoint: Any | None, + recurrent_state: torch.Tensor, + output: torch.Tensor, + ) -> torch.Tensor | None: + """Run packed KDA prefill straight against the recurrent-state pool. + + The op reads each request's initial state and writes its final state, + and any checkpoint, by slot index, so this path neither gathers a + dense initial state nor scatters a dense final one. Requests without + an initial state name the null slot and start from zero. + + Args: + scratch: Workspace that is simultaneously live with ``output``. + q: Live packed query rows, ``[tokens, heads, head_dim]``. + k: Live packed key rows. + v: Live packed value rows. + raw_g: Live unactivated forget gate. + raw_beta: Live unactivated update gate, ``[tokens, heads]``. + cu_seqlens: Packed request boundaries, ``[requests + 1]``. + state_indices: Destination state slot of each request. + has_initial_state: Whether each request continues a cached state. + checkpoint: Mid-sequence checkpoint metadata, or ``None``. + recurrent_state: The caller-owned recurrent-state pool. + output: Destination rows, ``[tokens, heads, head_dim]``. + + Raises: + RuntimeError: If the KDA prefill plan or scratch is unavailable. + ValueError: If the live batch exceeds the planned capacity. + """ + api = self._b12x_prefill_api + plan = self._b12x_prefill_plan + if api is None or plan is None: + raise RuntimeError( + "b12x KDA prefill KV cache was not bound before inference" + ) + num_tokens = int(q.shape[0]) + num_requests = int(state_indices.shape[0]) + if ( + num_tokens > self._b12x_prefill_max_tokens + or num_requests > self._b12x_prefill_max_seqs + ): + raise ValueError( + "b12x KDA prefill capacity exceeded: " + f"tokens={num_tokens}/{self._b12x_prefill_max_tokens}, " + f"requests={num_requests}/{self._b12x_prefill_max_seqs}" + ) + + initial_indices = self._b12x_prefill_initial_indices[:num_requests] + initial_indices.copy_(state_indices) + initial_indices.masked_fill_(~has_initial_state[:num_requests], NULL_BLOCK_ID) + if checkpoint is None: + checkpoint_indices = self._b12x_prefill_null_indices[:num_requests] + checkpoint_offsets = self._b12x_prefill_zero_offsets[:num_requests] + else: + checkpoint_indices = checkpoint.state_indices[:num_requests] + checkpoint_offsets = checkpoint.checkpoint_offsets[:num_requests] + self._b12x_prefill_num_seqs.fill_(num_requests) + self._b12x_prefill_num_tokens.fill_(num_tokens) + + binding = api.bind( + plan, + scratch=scratch, + q=q, + k=k, + v=v, + raw_g=raw_g, + raw_beta=raw_beta, + A_log=self.A_log, + dt_bias=self.dt_bias.view(-1, self.head_dim), + recurrent_state=recurrent_state, + cu_seqlens=cu_seqlens[: num_requests + 1], + initial_state_indices=initial_indices, + final_state_indices=state_indices, + checkpoint_state_indices=checkpoint_indices, + checkpoint_offsets=checkpoint_offsets, + num_seqs=self._b12x_prefill_num_seqs, + num_tokens=self._b12x_prefill_num_tokens, + output=output, + ) + api.run( + binding, + lower_bound=self.gate_lower_bound, + max_live_tokens=num_tokens, + max_live_seqs=num_requests, + ) + if getattr(plan.caps, "max_checkpoints", 1) == 2: + # Validation executes on this stream before checkpoint consumers. + # Abort the CUDA context on invalid internal metadata rather than + # allowing convolution-state writes or cache publication to follow. + torch._assert_async(binding.error_code == 0, "invalid recurrent checkpoint metadata") + return binding.error_code + return None + + def _store_kda_conv_checkpoint( + self, + *, + mixed_qkv: torch.Tensor, + conv_state: torch.Tensor, + recurrent_state: torch.Tensor, + query_start_loc: torch.Tensor, + checkpoint: Any, + error_code: torch.Tensor | None = None, + ) -> None: + """Store the convolution history at each request's checkpoint offset.""" + # Speculative storage includes future-token cells. A reusable prefill + # checkpoint stores only the causal kernel's history in the first cells. + state_len = self.conv1d.weight.shape[-1] - 1 + if not 1 <= state_len <= conv_state.shape[-1]: + raise ValueError("Convolution checkpoint history exceeds state storage") + width = mixed_qkv.shape[-1] + store_block_size = 256 + offsets = checkpoint.checkpoint_offsets + indices = checkpoint.state_indices + checkpoint_count = offsets.shape[1] if offsets.ndim == 2 else 1 + if checkpoint_count not in (1, 2) or not offsets.is_contiguous() or not indices.is_contiguous(): + raise ValueError("Checkpoint convolution metadata must be contiguous with capacity one or two") + if tuple(offsets.shape) != tuple(indices.shape): + raise ValueError("Checkpoint offset and destination shapes differ") + _store_cache_checkpoints_kernel[ + ( + checkpoint.checkpoint_offsets.numel(), + triton.cdiv(width * state_len, store_block_size), + ) + ]( + mixed_qkv, + conv_state, + recurrent_state, + recurrent_state, + query_start_loc, + checkpoint.checkpoint_offsets, + checkpoint.state_indices, + mixed_qkv.stride(0), + mixed_qkv.stride(1), + conv_state.stride(0), + conv_state.stride(1), + conv_state.stride(2), + recurrent_state.stride(0), + recurrent_state.stride(0), + checkpoint.checkpoint_offsets.stride(0), + state_len, + width, + 0, + NULL_BLOCK_ID, + store_block_size, + False, + checkpoint_count, + error_code, + error_code is not None, + ) + + def rearrange_mixed_qkv( + self, mixed_qkv: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + seq_len = mixed_qkv.shape[0] + qkv = mixed_qkv.view(seq_len, 3, self.local_num_heads, self.head_dim) + # Materialize all three row-strided inputs with one token-major to + # QKV-major permutation. Each unbound tensor is then contiguous. + qkv = qkv.permute(1, 0, 2, 3).contiguous().unsqueeze(1) + return qkv.unbind(0) + + def _can_use_b12x_kda_decode(self, m: GDNAttentionMetadata) -> bool: + if ( + self._b12x_kda_plan is None + or self._b12x_kda_scratch is None + or m.num_prefills != 0 + or (m.num_decodes == 0 and m.num_spec_decodes == 0) + ): + return False + if m.spec_sequence_masks is None: + return ( + m.num_spec_decodes == 0 + and m.non_spec_state_indices_tensor is not None + and m.non_spec_query_start_loc is not None + ) + return ( + m.num_decodes == 0 + and m.num_spec_decodes > 0 + and m.spec_state_indices_tensor is not None + and m.spec_query_start_loc is not None + and m.num_accepted_tokens is not None + ) + + def _run_b12x_kda_decode_post_conv( + self, + *, + metadata: GDNAttentionMetadata, + mixed_qkv: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + z: torch.Tensor, + output: torch.Tensor, + state_indices: torch.Tensor, + query_start_loc: torch.Tensor, + num_accepted_tokens: torch.Tensor | None, + num_requests: int, + ) -> None: + """Execute B12X KDA after the convolution projection. + + Args: + metadata: Describes whether packed query boundaries are uniform. + mixed_qkv: Live packed query, key, and value projection. + raw_g: Live unactivated forget gate. + raw_beta: Live unactivated update gate. + z: Live output gate. + output: Caller-owned destination tensor. + state_indices: Packed recurrent-state indices. + query_start_loc: Packed request boundaries. + num_accepted_tokens: Accepted speculative-token counts, or ``None`` + for one-token decode requests. + num_requests: Number of packed requests. + + Raises: + RuntimeError: If the KDA plan or cache is unavailable. + ValueError: If the live batch exceeds the planned capacity. + """ + api = self._b12x_kda_api + plan = self._b12x_kda_plan + scratch = self._b12x_kda_scratch + if api is None or plan is None or scratch is None: + raise RuntimeError("b12x KDA KV cache was not bound before inference") + num_tokens = int(mixed_qkv.shape[0]) + state_columns = int(state_indices.shape[1]) + if ( + num_tokens > self._b12x_kda_max_tokens + or num_requests > self._b12x_kda_max_seqs + or state_columns > self._b12x_kda_state_index_columns + ): + raise ValueError( + "b12x KDA capacity exceeded: " + f"tokens={num_tokens}/{self._b12x_kda_max_tokens}, " + f"requests={num_requests}/{self._b12x_kda_max_seqs}, " + f"state_columns={state_columns}/" + f"{self._b12x_kda_state_index_columns}" + ) + + forward_context = get_forward_context() + cache = forward_context.additional_kwargs.setdefault( + "b12x_kda_metadata_tensors", {} + ) + # Uniform builders own separate buffers with identical fixed boundaries. + cache_key = ( + None if metadata.is_uniform_spec_decode else query_start_loc.data_ptr(), + num_accepted_tokens.data_ptr() if num_accepted_tokens is not None else None, + num_tokens, + num_requests, + ) + bound_metadata = cache.get(cache_key) + if bound_metadata is None: + query_start_loc = query_start_loc[: num_requests + 1] + if num_accepted_tokens is None: + accepted_tokens = self._b12x_kda_num_accepted_tokens[:num_requests] + accepted_tokens.fill_(1) + else: + accepted_tokens = num_accepted_tokens[:num_requests] + self._b12x_kda_num_seqs.fill_(num_requests) + self._b12x_kda_num_tokens.copy_( + query_start_loc[num_requests : num_requests + 1] + ) + bound_metadata = ( + query_start_loc, + accepted_tokens, + self._b12x_kda_num_seqs, + self._b12x_kda_num_tokens, + ) + cache[cache_key] = bound_metadata + ( + query_start_loc, + accepted_tokens, + num_seqs, + num_tokens_tensor, + ) = bound_metadata + + binding = api.bind_kda( + plan, + scratch=scratch, + mixed_qkv=mixed_qkv, + raw_g=raw_g, + raw_beta=raw_beta, + z=z, + A_log=self.A_log, + dt_bias=self.dt_bias.view(self.local_num_heads, self.head_dim), + norm_weight=self.o_norm.weight, + recurrent_state=self.kv_cache[1], + query_start_loc=query_start_loc, + num_accepted_tokens=accepted_tokens, + state_indices=state_indices[:num_requests, :state_columns], + num_seqs=num_seqs, + num_tokens=num_tokens_tensor, + output=output, + ) + api.run_kda( + binding, + lower_bound=self.gate_lower_bound, + eps=self.o_norm.eps, + scale=self.head_dim**-0.5, + ) + + def forward( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + output: torch.Tensor, + ) -> None: + num_tokens = hidden_states.size(0) + projected_qkvgfab = self.in_proj_qkvgfab(hidden_states)[0] + if self.use_full_rank_gate: + split_sizes = [ + 3 * self.local_projection_size, + self.local_projection_size, + self.head_dim, + self.local_num_heads, + ] + if self.in_proj_padding: + split_sizes.append(self.in_proj_padding) + projected = projected_qkvgfab.split(split_sizes, dim=-1) + mixed_qkv, g_proj_states, f_a, beta = projected[:4] + else: + mixed_qkv, beta, f_a = projected_qkvgfab.split( + [ + 3 * self.local_projection_size, + self.local_num_heads, + self.head_dim, + ], + dim=-1, + ) + g_proj_states = self.g_b_proj(self.g_a_proj(hidden_states)[0])[0] + + g1 = self.f_b_proj(f_a)[0] + beta = beta.unsqueeze(0) + g1 = rearrange(g1, "n (h d) -> 1 n h d", d=self.head_dim) + + g2 = rearrange(g_proj_states, "... (h d) -> ... h d", d=self.head_dim) + + core_attn_out = torch.empty( + (1, num_tokens, self.local_num_heads, self.head_dim), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + + self._forward( + mixed_qkv=mixed_qkv, + g1=g1, + g2=g2, + beta=beta, + core_attn_out=core_attn_out, + ) + core_attn_out = rearrange(core_attn_out, "1 n h d -> n (h d)") + output[:] = self.o_proj(core_attn_out)[0] + + @eager_break_during_capture + def _forward( + self, + mixed_qkv: torch.Tensor, + g1: torch.Tensor, + g2: torch.Tensor, + beta: torch.Tensor, + core_attn_out: torch.Tensor, + ) -> None: + forward_context = get_forward_context() + attn_metadata_raw = forward_context.attn_metadata + + if attn_metadata_raw is None: + return + + # Vendor-specific KDA kernels: AMD/ROCm and NVIDIA keep their own copies + # under kimi_k3/{amd,nvidia}/ops so each can diverge independently. + # These copies may have different signatures for Kimi-K3, but they agree + # on the arguments used here. + if TYPE_CHECKING: + from vllm.models.kimi_k3.nvidia.ops.third_party.kda import ( + chunk_kda_with_fused_gate, + fused_recurrent_kda, + fused_recurrent_kda_packed_decode, + ) + elif current_platform.is_rocm(): + from vllm.models.kimi_k3.amd.ops.third_party.kda import ( # type: ignore[assignment] + chunk_kda_with_fused_gate, + fused_recurrent_kda, + fused_recurrent_kda_packed_decode, + ) + else: + from vllm.models.kimi_k3.nvidia.ops.third_party.kda import ( + chunk_kda_with_fused_gate, + fused_recurrent_kda, + fused_recurrent_kda_packed_decode, + ) + + assert isinstance(attn_metadata_raw, dict) + attn_metadata_narrowed = attn_metadata_raw.get(self.prefix) + if attn_metadata_narrowed is None: + # Profile/warmup dummy runs skip mamba-family metadata. + return + assert isinstance(attn_metadata_narrowed, GDNAttentionMetadata) + m = attn_metadata_narrowed + prefill_checkpoint = m.prefill_checkpoint + has_initial_state = m.has_initial_state + non_spec_query_start_loc = m.non_spec_query_start_loc + non_spec_state_indices_tensor = m.non_spec_state_indices_tensor + spec_sequence_masks = m.spec_sequence_masks + spec_token_indx = m.spec_token_indx + non_spec_token_indx = m.non_spec_token_indx + spec_state_indices_tensor = m.spec_state_indices_tensor + spec_query_start_loc = m.spec_query_start_loc + num_accepted_tokens = m.num_accepted_tokens + num_actual_tokens = m.num_actual_tokens + mixed_qkv = mixed_qkv[:num_actual_tokens] + g1 = g1[:, :num_actual_tokens] + beta = beta[:, :num_actual_tokens] + g2_actual = g2[:num_actual_tokens] + use_b12x_kda = self._can_use_b12x_kda_decode(m) + + constant_caches = self.kv_cache + + conv_state, recurrent_state = constant_caches + # conv_state must be (..., dim, width-1) for the conv kernels. + # DS layout stores it that way directly; SD layout needs a transpose. + if not is_conv_state_dim_first(): + conv_state = conv_state.transpose(-1, -2) + + conv_weights = self.conv1d.weight.view( + self.conv1d.weight.size(0), self.conv1d.weight.size(2) + ) + q_conv_weight, k_conv_weight, v_conv_weight = conv_weights.split( + self.local_projection_size, dim=0 + ) + q_conv_state, k_conv_state, v_conv_state = conv_state.split( + self.local_projection_size, dim=-2 + ) + + # Split tokens into the multi-query spec-decode part and the remaining + # (prefill / plain decode) part. + if spec_sequence_masks is not None: + if m.num_prefills == 0 and m.num_decodes == 0: + mixed_qkv_spec = mixed_qkv + g1_spec, beta_spec = g1, beta + mixed_qkv_ns = g1_ns = beta_ns = None + g2_spec, g2_ns = g2_actual, None + else: + mixed_qkv_spec = mixed_qkv.index_select(0, spec_token_indx) + g1_spec = g1.index_select(1, spec_token_indx) + beta_spec = beta.index_select(1, spec_token_indx) + mixed_qkv_ns = mixed_qkv.index_select(0, non_spec_token_indx) + g1_ns = g1.index_select(1, non_spec_token_indx) + beta_ns = beta.index_select(1, non_spec_token_indx) + g2_spec = g2_ns = None + else: + mixed_qkv_spec = g1_spec = beta_spec = None + mixed_qkv_ns, g1_ns, beta_ns = mixed_qkv, g1, beta + g2_spec, g2_ns = None, g2_actual + + # ---------- spec-decode multi-query path ---------- + core_attn_out_spec = None + if spec_sequence_masks is not None: + assert spec_state_indices_tensor is not None + assert spec_query_start_loc is not None + spec_conv_indices = spec_state_indices_tensor[:, 0][: m.num_spec_decodes] + spec_max_query_len = spec_state_indices_tensor.size(-1) + + # Sibling beta and, for full-rank gates, output-gate views remain + # live, so write the convolution output separately. + spec_conv_out = torch.empty( + mixed_qkv_spec.shape, + dtype=mixed_qkv_spec.dtype, + device=mixed_qkv_spec.device, + ) + mixed_qkv_spec = causal_conv1d_update( + mixed_qkv_spec, + conv_state, + conv_weights, + self.conv1d.bias, + activation="silu", + conv_state_indices=spec_conv_indices, + num_accepted_tokens=num_accepted_tokens, + query_start_loc=spec_query_start_loc, + max_query_len=spec_max_query_len, + validate_data=False, + out=spec_conv_out, + ) + spec_cu_seqlens = spec_query_start_loc[: m.num_spec_decodes + 1] + if use_b12x_kda: + assert g2_spec is not None + core_attn_out_spec = core_attn_out[:, : mixed_qkv_spec.shape[0]] + self._run_b12x_kda_decode_post_conv( + metadata=m, + mixed_qkv=mixed_qkv_spec, + raw_g=g1_spec[0], + raw_beta=beta_spec[0], + z=g2_spec, + output=core_attn_out_spec[0], + state_indices=spec_state_indices_tensor, + query_start_loc=spec_cu_seqlens, + num_accepted_tokens=num_accepted_tokens, + num_requests=m.num_spec_decodes, + ) + else: + q_spec, k_spec, v_spec = ( + rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim) + for x in mixed_qkv_spec.split(self.local_projection_size, dim=-1) + ) + # Spec-only batches write directly into core_attn_out. + spec_out = ( + core_attn_out[:, : q_spec.shape[1]] + if m.num_prefills == 0 and m.num_decodes == 0 + else None + ) + core_attn_out_spec, _ = fused_recurrent_kda( + q=q_spec, + k=k_spec, + v=v_spec, + raw_g=g1_spec, + raw_beta=beta_spec, + A_log=self.A_log, + dt_bias=self.dt_bias, + lower_bound=self.gate_lower_bound, + initial_state=recurrent_state, + cu_seqlens=spec_cu_seqlens, + ssm_state_indices=spec_state_indices_tensor, + num_accepted_tokens=num_accepted_tokens, + out=spec_out, + ) + + # ---------- non-spec path (prefill or plain decode) ---------- + core_attn_out_non_spec = None + if mixed_qkv_ns is not None: + assert g1_ns is not None and beta_ns is not None + if m.num_prefills > 0: + q_ns, k_ns, v_ns = mixed_qkv_ns.split( + self.local_projection_size, dim=-1 + ) + prefill_mixed_qkv = mixed_qkv_ns + + # Packed prefill conv would require copying V solely to make + # it dense for KDA. Separate calls accept the strided inputs + # and produce dense Q/K/V without that extra traffic. + # TODO: Use packed conv once every KDA prefill backend accepts + # row-strided Q/K/V directly. + def _prefill_conv( + x: torch.Tensor, + state: torch.Tensor, + weight: torch.Tensor, + ) -> torch.Tensor: + return causal_conv1d_fn( + x.transpose(0, 1), + weight, + None, + activation="silu", + conv_states=state, + has_initial_state=has_initial_state, + cache_indices=non_spec_state_indices_tensor, + query_start_loc=non_spec_query_start_loc, + metadata=m, + ).transpose(0, 1) + + q_ns = _prefill_conv(q_ns, q_conv_state, q_conv_weight) + k_ns = _prefill_conv(k_ns, k_conv_state, k_conv_weight) + v_ns = _prefill_conv(v_ns, v_conv_state, v_conv_weight) + q_ns, k_ns, v_ns = ( + rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim) + for x in (q_ns, k_ns, v_ns) + ) + + assert non_spec_state_indices_tensor is not None + assert has_initial_state is not None + + # Mixed non-spec batches are decode-first. Peel the length-one + # decodes off because the chunk kernel only consumes the + # prefill-tail metadata produced by the GDN builder. + core_attn_out_decode = None + split_non_spec = spec_sequence_masks is None and m.num_decodes > 0 + if split_non_spec: + assert non_spec_query_start_loc is not None + nd_tok = m.num_decode_tokens + prefill_mixed_qkv = prefill_mixed_qkv[nd_tok:] + core_attn_out_decode, _ = fused_recurrent_kda( + q=q_ns[:, :nd_tok], + k=k_ns[:, :nd_tok], + v=v_ns[:, :nd_tok], + raw_g=g1_ns[:, :nd_tok], + raw_beta=beta_ns[:, :nd_tok], + A_log=self.A_log, + dt_bias=self.dt_bias, + lower_bound=self.gate_lower_bound, + initial_state=recurrent_state, + cu_seqlens=non_spec_query_start_loc[: m.num_decodes + 1], + ssm_state_indices=non_spec_state_indices_tensor[ + : m.num_decodes + ], + ) + q_ns = q_ns[:, nd_tok:] + k_ns = k_ns[:, nd_tok:] + v_ns = v_ns[:, nd_tok:] + g1_ns = g1_ns[:, nd_tok:] + beta_ns = beta_ns[:, nd_tok:] + prefill_query_start_loc = m.prefill_query_start_loc + prefill_state_indices = m.prefill_state_indices + prefill_has_initial_state = m.prefill_has_initial_state + assert prefill_query_start_loc is not None + assert prefill_state_indices is not None + assert prefill_has_initial_state is not None + else: + prefill_query_start_loc = non_spec_query_start_loc + prefill_state_indices = non_spec_state_indices_tensor + prefill_has_initial_state = has_initial_state + + use_b12x_prefill = self.kda_prefill_backend == "b12x" + initial_state = ( + None + if use_b12x_prefill + else gather_initial_states( + recurrent_state, + prefill_state_indices, + prefill_has_initial_state, + ) + ) + if use_b12x_prefill: + assert self.gate_lower_bound is not None + assert prefill_query_start_loc is not None + num_prefill_tokens = int(q_ns.shape[1]) + b12x_scratch, b12x_out = self._get_b12x_prefill_workspace() + b12x_out = b12x_out[:num_prefill_tokens] + checkpoint_error = self._run_b12x_kda_prefill( + scratch=b12x_scratch, + q=q_ns[0], + k=k_ns[0], + v=v_ns[0], + raw_g=g1_ns[0], + raw_beta=beta_ns[0], + cu_seqlens=prefill_query_start_loc, + state_indices=prefill_state_indices, + has_initial_state=prefill_has_initial_state, + checkpoint=prefill_checkpoint, + recurrent_state=recurrent_state, + output=b12x_out, + ) + core_attn_out_non_spec = b12x_out.unsqueeze(0) + if prefill_checkpoint is not None: + # The op already wrote the recurrent checkpoint into + # its slot; only the convolution history is left. + self._store_kda_conv_checkpoint( + mixed_qkv=prefill_mixed_qkv, + conv_state=conv_state, + recurrent_state=recurrent_state, + query_start_loc=prefill_query_start_loc, + checkpoint=prefill_checkpoint, + error_code=checkpoint_error, + ) + elif self.kda_prefill_backend == "flashkda": + assert initial_state is not None + assert self.gate_lower_bound is not None + assert self._flashkda_buffer_specs is not None + assert prefill_query_start_loc is not None + workspace_out, final_state, checkpoint_state, workspace = ( + current_workspace_manager().get_simultaneous( + *self._flashkda_buffer_specs + ) + ) + flashkda_out = workspace_out[:, : q_ns.shape[1]] + if prefill_checkpoint is not None: + assert prefill_query_start_loc is not None + num_sequences = initial_state.shape[0] + assert prefill_checkpoint.checkpoint_offsets.shape == ( + num_sequences, + ) + final_state = final_state[:num_sequences] + checkpoint_state = checkpoint_state[:num_sequences] + _flashkda_prefill( + q=q_ns, + k=k_ns, + v=v_ns, + g=g1_ns, + beta=beta_ns, + A_log=self.A_log, + dt_bias=self.dt_bias, + lower_bound=self.gate_lower_bound, + initial_state=initial_state, + cu_seqlens=prefill_query_start_loc, + out=flashkda_out, + final_state=final_state, + workspace=workspace, + checkpoint_state=checkpoint_state, + checkpoint_offsets=(prefill_checkpoint.checkpoint_offsets), + ) + core_attn_out_non_spec = flashkda_out + last_recurrent_state = final_state + + state_len = conv_state.shape[-1] + width = prefill_mixed_qkv.shape[-1] + recurrent_row_size = checkpoint_state[0].numel() + store_block_size = 256 + _store_cache_checkpoints_kernel[ + ( + prefill_checkpoint.checkpoint_offsets.numel(), + triton.cdiv( + max(width * state_len, recurrent_row_size), + store_block_size, + ), + ) + ]( + prefill_mixed_qkv, + conv_state, + checkpoint_state, + recurrent_state, + prefill_query_start_loc, + prefill_checkpoint.checkpoint_offsets, + prefill_checkpoint.state_indices, + prefill_mixed_qkv.stride(0), + prefill_mixed_qkv.stride(1), + conv_state.stride(0), + conv_state.stride(1), + conv_state.stride(2), + checkpoint_state.stride(0), + recurrent_state.stride(0), + prefill_checkpoint.checkpoint_offsets.stride(0), + state_len, + width, + recurrent_row_size, + NULL_BLOCK_ID, + store_block_size, + True, + ) + else: + ( + core_attn_out_non_spec, + last_recurrent_state, + ) = _flashkda_prefill( + q=q_ns, + k=k_ns, + v=v_ns, + g=g1_ns, + beta=beta_ns, + A_log=self.A_log, + dt_bias=self.dt_bias, + lower_bound=self.gate_lower_bound, + initial_state=initial_state, + cu_seqlens=prefill_query_start_loc, + out=flashkda_out, + final_state=final_state[: initial_state.shape[0]], + workspace=workspace, + ) + else: + ( + core_attn_out_non_spec, + last_recurrent_state, + ) = chunk_kda_with_fused_gate( + q=q_ns, + k=k_ns, + v=v_ns, + raw_g=g1_ns, + raw_beta=beta_ns, + A_log=self.A_log, + g_bias=self.dt_bias, + lower_bound=self.gate_lower_bound, + initial_state=initial_state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + cu_seqlens=prefill_query_start_loc, + chunk_indices=m.chunk_indices, + chunk_offsets=m.chunk_offsets, + ) + # Init cache. The b12x op writes the pool in place. + if not use_b12x_prefill: + recurrent_state[prefill_state_indices] = last_recurrent_state + + if split_non_spec: + core_attn_out_non_spec = torch.cat( + [core_attn_out_decode, core_attn_out_non_spec], dim=1 + ) + + else: + # pure-decode non-spec batch + assert non_spec_state_indices_tensor is not None + decode_conv_indices = non_spec_state_indices_tensor[ + : mixed_qkv_ns.size(0) + ] + # Sibling beta and, for full-rank gates, output-gate views + # remain live, so write the conv output separately. + packed_conv_out = torch.empty( + mixed_qkv_ns.shape, + dtype=mixed_qkv_ns.dtype, + device=mixed_qkv_ns.device, + ) + mixed_qkv_ns = causal_conv1d_update( + mixed_qkv_ns, + conv_state, + conv_weights, + self.conv1d.bias, + activation="silu", + conv_state_indices=decode_conv_indices, + validate_data=True, + out=packed_conv_out, + ) + if use_b12x_kda: + assert non_spec_query_start_loc is not None + assert g2_ns is not None + core_attn_out_non_spec = core_attn_out[:, : mixed_qkv_ns.shape[0]] + self._run_b12x_kda_decode_post_conv( + metadata=m, + mixed_qkv=mixed_qkv_ns, + raw_g=g1_ns[0], + raw_beta=beta_ns[0], + z=g2_ns, + output=core_attn_out_non_spec[0], + state_indices=non_spec_state_indices_tensor[ + : m.num_decodes, None + ], + query_start_loc=non_spec_query_start_loc, + num_accepted_tokens=None, + num_requests=m.num_decodes, + ) + else: + core_attn_out_non_spec, _ = fused_recurrent_kda_packed_decode( + mixed_qkv=mixed_qkv_ns, + raw_g=g1_ns, + raw_beta=beta_ns, + A_log=self.A_log, + dt_bias=self.dt_bias, + lower_bound=self.gate_lower_bound, + initial_state=recurrent_state, + state_indices=decode_conv_indices, + ) + + # ---------- merge spec and non-spec outputs ---------- + if core_attn_out_spec is not None and core_attn_out_non_spec is not None: + # Mixed batches require indexed placement in the original order. + merged = torch.empty( + (1, num_actual_tokens, *core_attn_out_spec.shape[2:]), + dtype=core_attn_out_spec.dtype, + device=core_attn_out_spec.device, + ) + merged.index_copy_(1, spec_token_indx, core_attn_out_spec) + merged.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec) + core_attn_out[0, :num_actual_tokens] = merged[0, :num_actual_tokens] + elif core_attn_out_non_spec is not None: + core_attn_out[0, :num_actual_tokens] = core_attn_out_non_spec[ + 0, :num_actual_tokens + ] + else: + assert core_attn_out_spec is not None + if not use_b12x_kda: + core_attn_out.copy_(self.o_norm(core_attn_out, g2)) diff --git a/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/51ef011d11bb374d61840f1ff819759b6e3dd325da41856b7b0f01b013b54713/recurrent_prefill_checkpoint.py b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/51ef011d11bb374d61840f1ff819759b6e3dd325da41856b7b0f01b013b54713/recurrent_prefill_checkpoint.py new file mode 100644 index 00000000..4e6f8b04 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/51ef011d11bb374d61840f1ff819759b6e3dd325da41856b7b0f01b013b54713/recurrent_prefill_checkpoint.py @@ -0,0 +1,55 @@ +"""Explicit recurrent checkpoint positions for bounded fresh-prefill coalescing.""" +from __future__ import annotations + +CheckpointPlan = tuple[int, int, tuple[int, ...]] + + +def validate_plan(plan: CheckpointPlan, start: int, end: int, block_size: int) -> tuple[int, ...]: + planned_start, planned_end, targets = plan + if (planned_start, planned_end) != (start, end): + raise ValueError('recurrent checkpoint plan does not match the actual query span') + if not isinstance(targets, tuple) or not 1 <= len(targets) <= 2 or targets != tuple(sorted(set(targets))): + raise ValueError('recurrent checkpoint targets must be one or two sorted unique positions') + if any(type(p) is not int or not start < p < end or p % block_size or (p - start) % 16 for p in targets): + raise ValueError('recurrent checkpoint target is not a representable interior state') + return targets + + +def fresh_prompt_plan(*, start: int, end: int, prompt: int, num_tokens: int, + block_size: int, publications: tuple[int, ...], + shared_prefix_boundary: int = 0) -> CheckpointPlan | None: + # Resume, partial-tail and intermediate-prefill behavior stays in the + # existing scheduler. This first path only appends a completely new table. + if start != 0 or end != prompt or num_tokens != prompt or end > 8192 or end % block_size: + return None + required = {p for p in publications if start < p < end} + predecessor = max((num_tokens - 1) // block_size * block_size - block_size, 0) + if start < predecessor < end: + required.add(predecessor) + if start < shared_prefix_boundary < end: + required.add(shared_prefix_boundary // block_size * block_size) + required.discard(0) + targets = tuple(sorted(required)) + if not targets or len(targets) > 2: + return None + plan = (start, end, targets) + try: + validate_plan(plan, start, end, block_size) + except ValueError: + return None + return plan + + +def checkpoint_metadata(plan: CheckpointPlan | None, start: int, end: int, + block_size: int, capacity: int) -> tuple[list[int], list[int]]: + if capacity not in (1, 2): + raise ValueError('unsupported recurrent checkpoint capacity') + if plan is not None: + targets = validate_plan(plan, start, end, block_size) + if capacity < len(targets): + raise ValueError('metadata checkpoint capacity smaller than scheduled plan') + else: + boundary = end // block_size * block_size + targets = ((boundary,) if end % block_size and start < boundary < end and (boundary - start) % 16 == 0 else ()) + return ([p - start for p in targets] + [0] * (capacity - len(targets)), + [p // block_size - 1 for p in targets] + [-1] * (capacity - len(targets))) diff --git a/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/54707b0b7669642a55af60bdba0d9be5b5715cd3e0a68ca44021f11e4b8e9c9f/gdn_attn.py b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/54707b0b7669642a55af60bdba0d9be5b5715cd3e0a68ca44021f11e4b8e9c9f/gdn_attn.py new file mode 100644 index 00000000..f7492fa5 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/54707b0b7669642a55af60bdba0d9be5b5715cd3e0a68ca44021f11e4b8e9c9f/gdn_attn.py @@ -0,0 +1,968 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Backend for GatedDeltaNet attention.""" + +from vllm.v1.core.recurrent_prefill_checkpoint import checkpoint_metadata + +from copy import copy +from dataclasses import dataclass, replace +from typing import Literal + +import torch + +import vllm.envs as envs +from vllm.config import VllmConfig +from vllm.utils.torch_utils import async_tensor_h2d +from vllm.v1.attention.backend import ( + AttentionBackend, + AttentionCGSupport, + AttentionMetadataBuilder, + CommonAttentionMetadata, +) +from vllm.v1.attention.backends.utils import ( + NULL_BLOCK_ID, + compute_causal_conv1d_metadata, + mamba_get_block_table_tensor, + split_decodes_and_prefills, +) +from vllm.v1.kv_cache_interface import MambaSpec + + +class GDNAttentionBackend(AttentionBackend): + @staticmethod + def get_name() -> str: + return "GDN_ATTN" + + @staticmethod + def get_builder_cls() -> type["GDNAttentionMetadataBuilder"]: + return GDNAttentionMetadataBuilder + + @classmethod + def is_ssm(cls) -> bool: + return True + + +@dataclass +class GDNPrefillCheckpointMetadata: + """One recurrent-state checkpoint inside each packed prefill sequence. + + ``checkpoint_offsets`` are relative to the corresponding packed query. + ``request_rows`` and ``block_table_columns`` identify the cache slots that + receive the checkpoint, allowing metadata reuse to refresh physical block + IDs from a replacement block table. + """ + + checkpoint_offsets: torch.Tensor + state_indices: torch.Tensor + request_rows: torch.Tensor + block_table_columns: torch.Tensor + # Required planned targets must never use the legacy NULL/disabled export. + required_mask: torch.Tensor | None = None + + +@dataclass +class GDNAttentionMetadata: + num_prefills: int + num_prefill_tokens: int + num_decodes: int + num_decode_tokens: int + num_spec_decodes: int + num_spec_decode_tokens: int + num_actual_tokens: int + + has_initial_state: torch.Tensor | None = None + + spec_query_start_loc: torch.Tensor | None = None # shape: [num_spec_decodes + 1,] + non_spec_query_start_loc: torch.Tensor | None = ( + None # shape: [batch - num_spec_decodes + 1,] + ) + + spec_state_indices_tensor: torch.Tensor | None = None # shape: [batch, num_spec] + non_spec_state_indices_tensor: torch.Tensor | None = ( + None # shape: [batch - num_spec_decodes,] + ) + spec_sequence_masks: torch.Tensor | None = None # shape: [batch,] + spec_sequence_masks_cpu: torch.Tensor | None = None # shape: [batch,] + spec_token_indx: torch.Tensor | None = None + non_spec_token_indx: torch.Tensor | None = None + + num_accepted_tokens: torch.Tensor | None = None # shape: [batch,] + + # Pre-computed FLA chunk metadata (avoids GPU->CPU sync in prepare_chunk_indices) + chunk_indices: torch.Tensor | None = None + chunk_offsets: torch.Tensor | None = None + # Chunk-kernel inputs for prefill + prefill_query_start_loc: torch.Tensor | None = None + prefill_state_indices: torch.Tensor | None = None + prefill_has_initial_state: torch.Tensor | None = None + + # The following attributes are for triton implementation of causal_conv1d + nums_dict: dict | None = None + batch_ptr: torch.Tensor | None = None + token_chunk_offset_ptr: torch.Tensor | None = None + + # Required when reusing a metadata build across equivalent Mamba cache + # groups whose state block tables differ. + num_reqs: int = 0 + seq_lens: torch.Tensor | None = None + + prefill_checkpoint: GDNPrefillCheckpointMetadata | None = None + is_uniform_spec_decode: bool = False + + +class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]): + kv_cache_spec: MambaSpec + _cudagraph_support = AttentionCGSupport.UNIFORM_BATCH + supports_update_block_table: bool = True + + # Runner-owned stable storage, with NULL_BLOCK_ID in padded request rows. + mamba_aligned_state_indices: torch.Tensor | None = None + mamba_spec_accepted_tokens: torch.Tensor | None = None + + reorder_batch_threshold: int = 1 + + def __init__( + self, + kv_cache_spec: MambaSpec, + layer_names: list[str], + vllm_config: VllmConfig, + device: torch.device, + ): + self.vllm_config = vllm_config + self.compilation_config = vllm_config.compilation_config + self.speculative_config = vllm_config.speculative_config + self.kv_cache_spec = kv_cache_spec + from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( + _resolve_gdn_prefill_backend, + ) + + self.gdn_prefill_backend: Literal["triton", "flashinfer", "cutedsl"] + _, self.gdn_prefill_backend = _resolve_gdn_prefill_backend(vllm_config) + + if self.speculative_config: + assert self.speculative_config.num_speculative_tokens is not None + self.num_spec: int = self.speculative_config.num_speculative_tokens + else: + self.num_spec = 0 + self.use_spec_decode: bool = self.num_spec > 0 + self._init_reorder_batch_threshold(1, self.use_spec_decode) + self.use_full_cuda_graph: bool = ( + self.compilation_config.cudagraph_mode.has_full_cudagraphs() + ) + + self.decode_cudagraph_max_bs: int = ( + self.vllm_config.scheduler_config.max_num_seqs * (self.num_spec + 1) + ) + if self.compilation_config.max_cudagraph_capture_size is not None: + self.decode_cudagraph_max_bs = min( + self.decode_cudagraph_max_bs, + self.compilation_config.max_cudagraph_capture_size, + ) + + self.spec_state_indices_tensor: torch.Tensor = torch.empty( + (self.decode_cudagraph_max_bs, self.num_spec + 1), + dtype=torch.int32, + device=device, + ) + self.non_spec_state_indices_tensor: torch.Tensor = torch.empty( + (self.decode_cudagraph_max_bs,), + dtype=torch.int32, + device=device, + ) + self.spec_sequence_masks: torch.Tensor = torch.empty( + (self.decode_cudagraph_max_bs,), + dtype=torch.bool, + device=device, + ) + self.spec_token_indx: torch.Tensor = torch.empty( + (self.decode_cudagraph_max_bs * (self.num_spec + 1),), + dtype=torch.int32, + device=device, + ) + self.non_spec_token_indx: torch.Tensor = torch.empty( + (self.decode_cudagraph_max_bs * (self.num_spec + 1),), + dtype=torch.int32, + device=device, + ) + self.spec_query_start_loc: torch.Tensor = torch.empty( + (self.decode_cudagraph_max_bs + 1,), + dtype=torch.int32, + device=device, + ) + self.non_spec_query_start_loc: torch.Tensor = torch.empty( + (self.decode_cudagraph_max_bs + 1,), + dtype=torch.int32, + device=device, + ) + self.num_accepted_tokens: torch.Tensor = torch.empty( + (self.decode_cudagraph_max_bs,), + dtype=torch.int32, + device=device, + ) + self._decode_state_indices_source: torch.Tensor | None = None + self._decode_state_indices_view: torch.Tensor | None = None + self._reuse_spec_decode_inputs = envs.VLLM_GDN_SPEC_DECODE_METADATA_FASTPATH + self._uniform_spec_masks = torch.ones( + self.decode_cudagraph_max_bs, dtype=torch.bool, device=device + ) + self._uniform_spec_masks_cpu = torch.ones( + self.decode_cudagraph_max_bs, dtype=torch.bool + ) + self._uniform_spec_tokens = torch.arange( + self.decode_cudagraph_max_bs, dtype=torch.int32, device=device + ) + self._uniform_spec_query_start = torch.arange( + self.decode_cudagraph_max_bs + 1, dtype=torch.int32, device=device + ) * (self.num_spec + 1) + self._spec_state_indices_source: torch.Tensor | None = None + self._spec_state_indices_view: torch.Tensor | None = None + + def _can_reuse_spec_inputs( + self, + m: CommonAttentionMetadata, + num_accepted_tokens: torch.Tensor | None, + num_decode_draft_tokens_cpu: torch.Tensor | None, + ) -> bool: + return ( + self._reuse_spec_decode_inputs + and self.use_spec_decode + and self.use_full_cuda_graph + and self.vllm_config.cache_config.mamba_cache_mode == "align" + and self.mamba_aligned_state_indices is not None + and self.mamba_spec_accepted_tokens is not None + and num_accepted_tokens is not None + and num_decode_draft_tokens_cpu is not None + and 0 < m.num_actual_tokens <= self.decode_cudagraph_max_bs + and m.num_actual_tokens == m.num_reqs * (self.num_spec + 1) + and bool(torch.all(num_decode_draft_tokens_cpu == self.num_spec)) + and bool(torch.all(torch.diff(m.query_start_loc_cpu) == self.num_spec + 1)) + and (m.is_prefilling is None or not bool(torch.any(m.is_prefilling))) + ) + + def _get_spec_state_indices_view(self, num_reqs: int) -> torch.Tensor: + source = self.mamba_aligned_state_indices + assert source is not None + if ( + self._spec_state_indices_source is not source + or self._spec_state_indices_view is None + or self._spec_state_indices_view.shape[0] != num_reqs + ): + self._spec_state_indices_source = source + self._spec_state_indices_view = source[:num_reqs, : self.num_spec + 1] + return self._spec_state_indices_view + + def _build_uniform_spec_decode( + self, m: CommonAttentionMetadata, num_accepted_tokens: torch.Tensor + ) -> GDNAttentionMetadata: + num_reqs = m.num_reqs + assert self.mamba_spec_accepted_tokens is not None + accepted = self.mamba_spec_accepted_tokens[:num_reqs] + accepted.copy_(num_accepted_tokens[:num_reqs], non_blocking=True) + return GDNAttentionMetadata( + num_prefills=0, + num_prefill_tokens=0, + num_decodes=0, + num_decode_tokens=0, + num_spec_decodes=num_reqs, + num_spec_decode_tokens=m.num_actual_tokens, + num_actual_tokens=m.num_actual_tokens, + spec_query_start_loc=self._uniform_spec_query_start[: num_reqs + 1], + spec_state_indices_tensor=self._get_spec_state_indices_view(num_reqs), + spec_sequence_masks=self._uniform_spec_masks[:num_reqs], + spec_sequence_masks_cpu=self._uniform_spec_masks_cpu[:num_reqs], + spec_token_indx=self._uniform_spec_tokens[: m.num_actual_tokens], + non_spec_token_indx=self._uniform_spec_tokens[:0], + num_accepted_tokens=accepted, + num_reqs=num_reqs, + seq_lens=m.seq_lens, + is_uniform_spec_decode=True, + ) + + def _get_state_indices( + self, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + num_reqs: int, + ) -> torch.Tensor: + if ( + self.vllm_config.cache_config.mamba_cache_mode == "align" + and self.mamba_aligned_state_indices is not None + ): + return self.mamba_aligned_state_indices[:num_reqs] + return mamba_get_block_table_tensor( + block_table, + seq_lens, + self.kv_cache_spec, + self.vllm_config.cache_config.mamba_cache_mode, + ) + + def _can_reuse_decode_inputs(self) -> bool: + return ( + not self.use_spec_decode + and self.vllm_config.cache_config.mamba_cache_mode == "align" + and self.mamba_aligned_state_indices is not None + ) + + def _build_chunk_metadata( + self, + prefill_query_start_loc: torch.Tensor, + prefill_query_start_loc_cpu: torch.Tensor, + device: torch.device, + ) -> tuple[torch.Tensor, torch.Tensor]: + from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE + + if self.gdn_prefill_backend == "cutedsl": + from vllm.model_executor.layers.mamba.ops.gdn_chunk_cutedsl import ( + prepare_metadata_cutedsl, + ) + + assert prefill_query_start_loc is not None + assert prefill_query_start_loc_cpu is not None + total_tokens = int(prefill_query_start_loc_cpu[-1].item()) + return prepare_metadata_cutedsl( + prefill_query_start_loc, + total_tokens, + FLA_CHUNK_SIZE, + ) + + # Only prefill batches use FLA chunk ops. + # Pre-compute on CPU and async-copy to GPU to avoid + # GPU→CPU sync (.tolist()) in prepare_chunk_indices. + from vllm.third_party.flash_linear_attention.ops.index import ( + prepare_chunk_indices, + prepare_chunk_offsets, + ) + + assert prefill_query_start_loc_cpu is not None + return ( + async_tensor_h2d( + prepare_chunk_indices(prefill_query_start_loc_cpu, FLA_CHUNK_SIZE), + device=device, + ), + async_tensor_h2d( + prepare_chunk_offsets(prefill_query_start_loc_cpu, FLA_CHUNK_SIZE), + device=device, + ), + ) + + def build( # type: ignore[override] + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + num_accepted_tokens: torch.Tensor | None = None, + num_decode_draft_tokens_cpu: torch.Tensor | None = None, + fast_build: bool = False, + ) -> GDNAttentionMetadata: + m = common_attn_metadata + if self._can_reuse_spec_inputs( + m, num_accepted_tokens, num_decode_draft_tokens_cpu + ): + assert num_accepted_tokens is not None + return self._build_uniform_spec_decode(m, num_accepted_tokens) + + query_start_loc = m.query_start_loc + query_start_loc_cpu = m.query_start_loc_cpu + nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None + block_table_tensor = self._get_state_indices( + m.block_table_tensor, + m.seq_lens, + m.num_reqs, + ) + + spec_sequence_masks_cpu: torch.Tensor | None = None + if not self.use_spec_decode or num_decode_draft_tokens_cpu is None: + spec_sequence_masks = None + num_spec_decodes = 0 + else: + spec_sequence_masks_cpu = num_decode_draft_tokens_cpu >= 0 + num_spec_decodes = spec_sequence_masks_cpu.sum().item() + if ( + num_spec_decodes == 0 + or num_decode_draft_tokens_cpu[spec_sequence_masks_cpu].sum().item() + == 0 + ): + num_spec_decodes = 0 + spec_sequence_masks = None + spec_sequence_masks_cpu = None + else: + spec_sequence_masks = async_tensor_h2d( + spec_sequence_masks_cpu, device=query_start_loc.device + ) + + if spec_sequence_masks is None: + assert m.is_prefilling is not None + # Mamba cache pages are not allocator-zeroed. Fresh one-token + # requests must take the prefill path so has_initial_state=False + # masks both convolution and recurrent state. + num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( + split_decodes_and_prefills( + m, + decode_threshold=1, + treat_short_extends_as_decodes=False, + ) + ) + num_spec_decode_tokens = 0 + spec_token_indx = None + non_spec_token_indx = None + spec_state_indices_tensor = None + non_spec_state_indices_tensor = block_table_tensor[:, 0] + spec_query_start_loc = None + non_spec_query_start_loc = query_start_loc + non_spec_query_start_loc_cpu = query_start_loc_cpu + num_accepted_tokens = None + else: + query_lens = query_start_loc[1:] - query_start_loc[:-1] + assert spec_sequence_masks_cpu is not None + non_spec_sequence_masks_cpu = ~spec_sequence_masks_cpu + query_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1] + + # Use CPU tensors to avoid CPU-GPU sync + non_spec_query_lens_cpu = query_lens_cpu[non_spec_sequence_masks_cpu] + num_decodes = (non_spec_query_lens_cpu == 1).sum().item() + # Exclude zero-length padded sequences from prefill count. + num_zero_len = (non_spec_query_lens_cpu == 0).sum().item() + num_prefills = non_spec_query_lens_cpu.size(0) - num_decodes - num_zero_len + num_decode_tokens = num_decodes + num_prefill_tokens = ( + non_spec_query_lens_cpu.sum().item() - num_decode_tokens + ) + num_spec_decode_tokens = ( + query_lens_cpu.sum().item() - num_prefill_tokens - num_decode_tokens + ) + + # num_decodes and num_spec_decodes are mutually exclusive. + # Reclassify non-spec decodes as prefills when spec decodes + # exist — the prefill kernel handles 1-token sequences with + # initial state correctly, producing identical results. + if num_decodes > 0 and num_spec_decodes > 0: + num_prefills += num_decodes + num_prefill_tokens += num_decode_tokens + num_decodes = 0 + num_decode_tokens = 0 + + if num_prefills == 0 and num_decodes == 0: + spec_token_size = min( + num_spec_decodes * (self.num_spec + 1), + query_start_loc_cpu[-1].item(), + ) + spec_token_indx = torch.arange( + spec_token_size, + dtype=torch.int32, + device=query_start_loc.device, + ) + non_spec_token_indx = torch.empty( + 0, dtype=torch.int32, device=query_start_loc.device + ) + # Filter by spec_sequence_masks to exclude padded sequences + spec_state_indices_tensor = block_table_tensor[ + spec_sequence_masks_cpu, : self.num_spec + 1 + ] + non_spec_state_indices_tensor = None + # Padded sequences are always at the back, so the first + # num_spec_decodes + 1 entries of query_start_loc already + # contain the correct cumulative token counts. + spec_query_start_loc = query_start_loc[: num_spec_decodes + 1] + non_spec_query_start_loc = None + non_spec_query_start_loc_cpu = None + else: + spec_token_masks = torch.repeat_interleave( + spec_sequence_masks, + query_lens, + output_size=query_start_loc_cpu[-1].item(), + ) + index = torch.argsort(spec_token_masks, stable=True) + num_non_spec_tokens = num_prefill_tokens + num_decode_tokens + non_spec_token_indx = index[:num_non_spec_tokens] + spec_token_indx = index[num_non_spec_tokens:] + + spec_state_indices_tensor = block_table_tensor[ + spec_sequence_masks_cpu, : self.num_spec + 1 + ] + non_spec_state_indices_tensor = block_table_tensor[ + non_spec_sequence_masks_cpu, 0 + ] + + spec_query_start_loc = torch.zeros( + num_spec_decodes + 1, + dtype=torch.int32, + device=query_start_loc.device, + ) + torch.cumsum( + query_lens[spec_sequence_masks_cpu], + dim=0, + out=spec_query_start_loc[1:], + ) + non_spec_query_start_loc = torch.zeros( + query_lens.size(0) - num_spec_decodes + 1, + dtype=torch.int32, + device=query_start_loc.device, + ) + torch.cumsum( + query_lens[non_spec_sequence_masks_cpu], + dim=0, + out=non_spec_query_start_loc[1:], + ) + non_spec_query_start_loc_cpu = torch.zeros( + query_lens_cpu.size(0) - num_spec_decodes + 1, + dtype=torch.int32, + ) + torch.cumsum( + query_lens_cpu[non_spec_sequence_masks_cpu], + dim=0, + out=non_spec_query_start_loc_cpu[1:], + ) + + assert num_accepted_tokens is not None + num_accepted_tokens = num_accepted_tokens[spec_sequence_masks_cpu] + + chunk_indices: torch.Tensor | None = None + chunk_offsets: torch.Tensor | None = None + prefill_query_start_loc: torch.Tensor | None = None + prefill_state_indices: torch.Tensor | None = None + prefill_has_initial_state: torch.Tensor | None = None + if num_prefills > 0: + # In a mixed non-spec batch, decodes are peeled off to the recurrent + # kernel (decode-first front slice), so build chunk metadata from the + # rebased prefill-only cu_seqlens; otherwise use the full non-spec one. + # _forward_core keys off the same condition, so they agree. + if spec_sequence_masks is None and num_decodes > 0: + assert non_spec_query_start_loc is not None + assert non_spec_query_start_loc_cpu is not None + assert non_spec_state_indices_tensor is not None + prefill_query_start_loc = ( + non_spec_query_start_loc[num_decodes:] - num_decode_tokens + ) + prefill_query_start_loc_cpu = ( + non_spec_query_start_loc_cpu[num_decodes:] - num_decode_tokens + ) + prefill_state_indices = non_spec_state_indices_tensor[num_decodes:] + else: + prefill_query_start_loc = non_spec_query_start_loc + prefill_query_start_loc_cpu = non_spec_query_start_loc_cpu + prefill_state_indices = non_spec_state_indices_tensor + + chunk_indices, chunk_offsets = self._build_chunk_metadata( + prefill_query_start_loc, + prefill_query_start_loc_cpu, + query_start_loc.device, + ) + + if num_prefills > 0: + context_lens_tensor = m.compute_num_computed_tokens() + has_initial_state = context_lens_tensor > 0 + if spec_sequence_masks_cpu is not None: + has_initial_state = has_initial_state[~spec_sequence_masks_cpu] + assert non_spec_query_start_loc_cpu is not None + nums_dict, batch_ptr, token_chunk_offset_ptr = ( + compute_causal_conv1d_metadata( + non_spec_query_start_loc_cpu, + device=query_start_loc.device, + ) + ) + if spec_sequence_masks is None and num_decodes > 0: + prefill_has_initial_state = has_initial_state[num_decodes:] + else: + prefill_has_initial_state = has_initial_state + else: + has_initial_state = None + + prefill_checkpoint = None + if ( + num_prefills > 0 + and self.kv_cache_spec.num_prefill_checkpoint_blocks > 0 + and self.vllm_config.cache_config.mamba_cache_mode == "align" + ): + # FlashKDA can materialize one state at a cache-block boundary + # without splitting the target-model forward. Only prefill rows + # participate, in the same order as prefill_query_start_loc. + assert m.seq_lens_cpu_upper_bound is not None + all_query_lens = query_start_loc_cpu.diff().tolist() + if spec_sequence_masks_cpu is None: + request_rows = list(range(num_decodes, num_decodes + num_prefills)) + else: + request_rows = [ + row + for row in (~spec_sequence_masks_cpu).nonzero().flatten().tolist() + if all_query_lens[row] > 0 + ] + assert len(request_rows) == num_prefills + + seq_lens = m.seq_lens_cpu_upper_bound.tolist() + block_size = self.kv_cache_spec.block_size + capacity = min(self.kv_cache_spec.num_prefill_checkpoint_blocks, 2) + plans = getattr(m, "recurrent_prefill_checkpoint_plans_cpu", None) + if plans is not None and len(plans) < len(all_query_lens): + raise ValueError("checkpoint plan rows do not cover the current batch") + checkpoint_offsets = [] + checkpoint_columns = [] + checkpoint_required = [] + for row in request_rows: + query_len = all_query_lens[row] + seq_len = seq_lens[row] + offsets, columns = checkpoint_metadata( + None if plans is None else plans[row], seq_len - query_len, + seq_len, block_size, capacity) + checkpoint_offsets.append(offsets if capacity == 2 else offsets[0]) + checkpoint_columns.append(columns if capacity == 2 else columns[0]) + required = [plans is not None and plans[row] is not None and column >= 0 + for column in columns] + checkpoint_required.append(required if capacity == 2 else required[0]) + + any_checkpoint = (any(any(row) for row in checkpoint_offsets) if capacity == 2 + else any(checkpoint_offsets)) + if any_checkpoint: + checkpoint_offsets_tensor = async_tensor_h2d( + checkpoint_offsets, + dtype=torch.int32, + device=query_start_loc.device, + ) + request_rows_tensor = async_tensor_h2d( + request_rows, + dtype=torch.int64, + device=query_start_loc.device, + ) + checkpoint_columns_tensor = async_tensor_h2d( + checkpoint_columns, + dtype=torch.int64, + device=query_start_loc.device, + ) + checkpoint_state_indices = m.block_table_tensor[ + request_rows_tensor[:, None] if capacity == 2 else request_rows_tensor, checkpoint_columns_tensor + ] + checkpoint_state_indices = torch.where( + checkpoint_columns_tensor >= 0, + checkpoint_state_indices, + NULL_BLOCK_ID, + ) + has_required = (any(any(row) for row in checkpoint_required) if capacity == 2 + else any(checkpoint_required)) + required_mask = None + if has_required: + required_mask = async_tensor_h2d( + checkpoint_required, dtype=torch.bool, device=query_start_loc.device) + torch._assert_async( + torch.all(~required_mask | (checkpoint_state_indices != NULL_BLOCK_ID)), + "planned recurrent checkpoint refers to a NULL state block") + prefill_checkpoint = GDNPrefillCheckpointMetadata( + checkpoint_offsets=checkpoint_offsets_tensor, + state_indices=checkpoint_state_indices, + request_rows=request_rows_tensor, + block_table_columns=checkpoint_columns_tensor, + required_mask=required_mask, + ) + + # Function code counted on either presency non-spec decode or spec decode, + # but not both. + assert not (num_decodes > 0 and num_spec_decodes > 0), ( + f"num_decodes: {num_decodes}, num_spec_decodes: {num_spec_decodes}" + ) + + # Prepare per-request tensors for cudagraph. m.num_actual_tokens is + # token-padded for FULL graph replay, but the GDN state/query/accepted + # metadata below is indexed by request. + batch_size = m.num_reqs + + if ( + self.use_full_cuda_graph + and num_prefills == 0 + and num_decodes == 0 + and num_spec_decodes <= self.decode_cudagraph_max_bs + and num_spec_decode_tokens <= self.decode_cudagraph_max_bs + ): + assert spec_sequence_masks is not None + self.spec_state_indices_tensor[:num_spec_decodes].copy_( + spec_state_indices_tensor, non_blocking=True + ) + spec_state_indices_tensor = self.spec_state_indices_tensor[:batch_size] + spec_state_indices_tensor[num_spec_decodes:].fill_(NULL_BLOCK_ID) + + self.spec_sequence_masks[:num_spec_decodes].copy_( + spec_sequence_masks[:num_spec_decodes], non_blocking=True + ) + spec_sequence_masks = self.spec_sequence_masks[:batch_size] + spec_sequence_masks[num_spec_decodes:].fill_(False) + + assert non_spec_token_indx is not None and spec_token_indx is not None + self.non_spec_token_indx[: non_spec_token_indx.size(0)].copy_( + non_spec_token_indx, non_blocking=True + ) + non_spec_token_indx = self.non_spec_token_indx[ + : non_spec_token_indx.size(0) + ] + + self.spec_token_indx[: spec_token_indx.size(0)].copy_( + spec_token_indx, non_blocking=True + ) + spec_token_indx = self.spec_token_indx[: spec_token_indx.size(0)] + + self.spec_query_start_loc[: num_spec_decodes + 1].copy_( + spec_query_start_loc, non_blocking=True + ) + spec_num_query_tokens = spec_query_start_loc[-1] # type: ignore[index] + spec_query_start_loc = self.spec_query_start_loc[: batch_size + 1] + spec_query_start_loc[num_spec_decodes + 1 :].fill_(spec_num_query_tokens) + + self.num_accepted_tokens[:num_spec_decodes].copy_( + num_accepted_tokens, non_blocking=True + ) + num_accepted_tokens = self.num_accepted_tokens[:batch_size] + num_accepted_tokens[num_spec_decodes:].fill_(1) + + if ( + self.use_full_cuda_graph + and num_prefills == 0 + and num_spec_decodes == 0 + and num_decodes <= self.decode_cudagraph_max_bs + and not self._can_reuse_decode_inputs() + ): + self.non_spec_state_indices_tensor[:num_decodes].copy_( + non_spec_state_indices_tensor, non_blocking=True + ) + non_spec_state_indices_tensor = self.non_spec_state_indices_tensor[ + :batch_size + ] + non_spec_state_indices_tensor[num_decodes:].fill_(NULL_BLOCK_ID) + + self.non_spec_query_start_loc[: num_decodes + 1].copy_( + non_spec_query_start_loc, non_blocking=True + ) + non_spec_num_query_tokens = non_spec_query_start_loc[-1] # type: ignore[index] + non_spec_query_start_loc = self.non_spec_query_start_loc[: batch_size + 1] + non_spec_query_start_loc[num_decodes + 1 :].fill_(non_spec_num_query_tokens) + + attn_metadata = GDNAttentionMetadata( + num_prefills=num_prefills, + num_prefill_tokens=num_prefill_tokens, + num_decodes=num_decodes, + num_decode_tokens=num_decode_tokens, + num_spec_decodes=num_spec_decodes, + num_spec_decode_tokens=num_spec_decode_tokens, + num_actual_tokens=m.num_actual_tokens, + has_initial_state=has_initial_state, + chunk_indices=chunk_indices, + chunk_offsets=chunk_offsets, + prefill_query_start_loc=prefill_query_start_loc, + prefill_state_indices=prefill_state_indices, + prefill_has_initial_state=prefill_has_initial_state, + spec_query_start_loc=spec_query_start_loc, + non_spec_query_start_loc=non_spec_query_start_loc, + spec_state_indices_tensor=spec_state_indices_tensor, + non_spec_state_indices_tensor=non_spec_state_indices_tensor, + spec_sequence_masks=spec_sequence_masks, + spec_sequence_masks_cpu=spec_sequence_masks_cpu, + spec_token_indx=spec_token_indx, + non_spec_token_indx=non_spec_token_indx, + num_accepted_tokens=num_accepted_tokens, + nums_dict=nums_dict, + batch_ptr=batch_ptr, + token_chunk_offset_ptr=token_chunk_offset_ptr, + num_reqs=m.num_reqs, + seq_lens=m.seq_lens, + prefill_checkpoint=prefill_checkpoint, + ) + return attn_metadata + + def update_block_table( + self, + metadata: GDNAttentionMetadata, + blk_table: torch.Tensor, + slot_mapping: torch.Tensor, + ) -> GDNAttentionMetadata: + del slot_mapping + assert metadata.num_reqs > 0 + assert metadata.seq_lens is not None + + if ( + metadata.is_uniform_spec_decode + and self._reuse_spec_decode_inputs + and self.mamba_aligned_state_indices is not None + and self.mamba_spec_accepted_tokens is not None + ): + updated = copy(metadata) + updated.spec_state_indices_tensor = self._get_spec_state_indices_view( + metadata.num_reqs + ) + accepted = self.mamba_spec_accepted_tokens[: metadata.num_reqs] + assert metadata.num_accepted_tokens is not None + if accepted.data_ptr() != metadata.num_accepted_tokens.data_ptr(): + accepted.copy_(metadata.num_accepted_tokens, non_blocking=True) + updated.num_accepted_tokens = accepted + return updated + + if ( + metadata.num_prefills == 0 + and metadata.num_spec_decodes == 0 + and self._can_reuse_decode_inputs() + ): + source = self.mamba_aligned_state_indices + assert source is not None + if ( + self._decode_state_indices_source is not source + or self._decode_state_indices_view is None + or self._decode_state_indices_view.shape[0] != metadata.num_reqs + ): + self._decode_state_indices_source = source + self._decode_state_indices_view = source[: metadata.num_reqs, 0] + updated = copy(metadata) + updated.non_spec_state_indices_tensor = self._decode_state_indices_view + return updated + + state_indices = self._get_state_indices( + blk_table, + metadata.seq_lens, + metadata.num_reqs, + ) + spec_sequence_masks_cpu = metadata.spec_sequence_masks_cpu + if spec_sequence_masks_cpu is None: + spec_state_indices = None + non_spec_state_indices = state_indices[:, 0] + else: + non_spec_sequence_masks_cpu = ~spec_sequence_masks_cpu + spec_state_indices = state_indices[ + spec_sequence_masks_cpu, : self.num_spec + 1 + ] + non_spec_state_indices = state_indices[non_spec_sequence_masks_cpu, 0] + prefill_state_indices = metadata.prefill_state_indices + if metadata.num_prefills > 0: + if spec_sequence_masks_cpu is None and metadata.num_decodes > 0: + prefill_state_indices = non_spec_state_indices[metadata.num_decodes :] + else: + prefill_state_indices = non_spec_state_indices + + prefill_checkpoint = metadata.prefill_checkpoint + if prefill_checkpoint is not None: + checkpoint_state_indices = blk_table[ + prefill_checkpoint.request_rows[:, None] + if prefill_checkpoint.block_table_columns.ndim == 2 + else prefill_checkpoint.request_rows, + prefill_checkpoint.block_table_columns, + ] + checkpoint_state_indices = torch.where( + prefill_checkpoint.block_table_columns >= 0, + checkpoint_state_indices, + NULL_BLOCK_ID, + ) + if prefill_checkpoint.required_mask is not None: + torch._assert_async( + torch.all(~prefill_checkpoint.required_mask | (checkpoint_state_indices != NULL_BLOCK_ID)), + "rebound planned recurrent checkpoint refers to a NULL state block") + prefill_checkpoint = replace( + prefill_checkpoint, + state_indices=checkpoint_state_indices, + ) + + spec_sequence_masks = metadata.spec_sequence_masks + spec_token_indx = metadata.spec_token_indx + non_spec_token_indx = metadata.non_spec_token_indx + spec_query_start_loc = metadata.spec_query_start_loc + num_accepted_tokens = metadata.num_accepted_tokens + non_spec_query_start_loc = metadata.non_spec_query_start_loc + if ( + self.use_full_cuda_graph + and metadata.num_prefills == 0 + and metadata.num_decodes == 0 + and metadata.num_spec_decodes <= self.decode_cudagraph_max_bs + and metadata.num_spec_decode_tokens <= self.decode_cudagraph_max_bs + ): + assert spec_state_indices is not None + assert spec_sequence_masks is not None + assert spec_token_indx is not None + assert non_spec_token_indx is not None + assert spec_query_start_loc is not None + assert num_accepted_tokens is not None + + self.spec_state_indices_tensor[: metadata.num_spec_decodes].copy_( + spec_state_indices, non_blocking=True + ) + spec_state_indices = self.spec_state_indices_tensor[: metadata.num_reqs] + spec_state_indices[metadata.num_spec_decodes :].fill_(NULL_BLOCK_ID) + + self.spec_sequence_masks[: metadata.num_reqs].copy_( + spec_sequence_masks[: metadata.num_reqs], non_blocking=True + ) + spec_sequence_masks = self.spec_sequence_masks[: metadata.num_reqs] + + self.non_spec_token_indx[: non_spec_token_indx.size(0)].copy_( + non_spec_token_indx, non_blocking=True + ) + non_spec_token_indx = self.non_spec_token_indx[ + : non_spec_token_indx.size(0) + ] + + self.spec_token_indx[: spec_token_indx.size(0)].copy_( + spec_token_indx, non_blocking=True + ) + spec_token_indx = self.spec_token_indx[: spec_token_indx.size(0)] + + self.spec_query_start_loc[: metadata.num_reqs + 1].copy_( + spec_query_start_loc[: metadata.num_reqs + 1], non_blocking=True + ) + spec_query_start_loc = self.spec_query_start_loc[: metadata.num_reqs + 1] + + self.num_accepted_tokens[: metadata.num_reqs].copy_( + num_accepted_tokens[: metadata.num_reqs], non_blocking=True + ) + num_accepted_tokens = self.num_accepted_tokens[: metadata.num_reqs] + + if ( + self.use_full_cuda_graph + and metadata.num_prefills == 0 + and metadata.num_spec_decodes == 0 + and metadata.num_decodes <= self.decode_cudagraph_max_bs + and not self._can_reuse_decode_inputs() + ): + self.non_spec_state_indices_tensor[: metadata.num_decodes].copy_( + non_spec_state_indices[: metadata.num_decodes], non_blocking=True + ) + non_spec_state_indices = self.non_spec_state_indices_tensor[ + : metadata.num_reqs + ] + non_spec_state_indices[metadata.num_decodes :].fill_(NULL_BLOCK_ID) + + assert non_spec_query_start_loc is not None + self.non_spec_query_start_loc[: metadata.num_reqs + 1].copy_( + non_spec_query_start_loc[: metadata.num_reqs + 1], + non_blocking=True, + ) + non_spec_query_start_loc = self.non_spec_query_start_loc[ + : metadata.num_reqs + 1 + ] + + return replace( + metadata, + spec_state_indices_tensor=spec_state_indices, + non_spec_state_indices_tensor=non_spec_state_indices, + prefill_state_indices=prefill_state_indices, + spec_sequence_masks=spec_sequence_masks, + spec_token_indx=spec_token_indx, + non_spec_token_indx=non_spec_token_indx, + spec_query_start_loc=spec_query_start_loc, + non_spec_query_start_loc=non_spec_query_start_loc, + num_accepted_tokens=num_accepted_tokens, + prefill_checkpoint=prefill_checkpoint, + ) + + def build_for_cudagraph_capture( + self, common_attn_metadata: CommonAttentionMetadata + ): + """ + This method builds the metadata for full cudagraph capture. + Currently, only decode is supported for full cudagraphs with Mamba. + """ + m = common_attn_metadata + + assert ( + m.num_reqs <= self.decode_cudagraph_max_bs + and m.num_actual_tokens <= self.decode_cudagraph_max_bs + ), ( + f"GDN only supports decode-only full CUDAGraph capture. " + f"Make sure batch size ({m.num_reqs}) <= " + f"cudagraph capture sizes ({self.decode_cudagraph_max_bs}), " + f"and number of tokens ({m.num_actual_tokens}) <= " + f"cudagraph capture sizes ({self.decode_cudagraph_max_bs})." + ) + + num_accepted_tokens = torch.diff(m.query_start_loc) + num_decode_draft_tokens_cpu = (num_accepted_tokens - 1).cpu() + + return self.build(0, m, num_accepted_tokens, num_decode_draft_tokens_cpu) diff --git a/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/5ba22653a26444a989ce0942b5fbdb12cda4fa2722f818c8e199cd50843e6565/model_runner.py b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/5ba22653a26444a989ce0942b5fbdb12cda4fa2722f818c8e199cd50843e6565/model_runner.py new file mode 100644 index 00000000..c0285bae --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/5ba22653a26444a989ce0942b5fbdb12cda4fa2722f818c8e199cd50843e6565/model_runner.py @@ -0,0 +1,2143 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +NOTE: Coding style guide for this file: +This model runner is shared by all models: text and multimodal, generative +and embedding, public and private. As a result, this file must only contain +code that is common to every model. Model-specific behavior belongs in the +appropriate model-specific files. + +In other words: +* Be paranoid about changing this file. It should remain stable. +* Be even more paranoid about adding new lines. It should remain minimal. + +Even for shared features (for example, different parallelism modes), keep the +complexity out of this path. The less common the feature, the more it should be +hidden. Prefer utility functions defined elsewhere and call them from here, +instead of embedding feature-specific logic directly. +""" + +import functools +import gc +import time +from copy import deepcopy +from typing import Any, NamedTuple + +import numpy as np +import torch +import torch.nn as nn + +import vllm.envs as envs +from vllm.compilation.counter import compilation_counter +from vllm.config import VllmConfig +from vllm.config.compilation import CUDAGraphMode +from vllm.distributed.parallel_state import ( + get_dcp_group, + get_pp_group, +) +from vllm.forward_context import BatchDescriptor, set_forward_context +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.all2all_utils import get_ep_all2all_manager +from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( + RoutedExpertsCapturer, + bind_routed_experts_capturer, +) +from vllm.model_executor.layers.mamba.ops.ssu_dispatch import ( + initialize_mamba_ssu_backend, +) +from vllm.model_executor.model_loader import get_model_loader +from vllm.model_executor.offloader import ( + create_offloader, + get_offloader, + set_offloader, +) +from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.encoder_budget import ( + MultiModalBudget, + get_dummy_encoder_profile_inputs, +) +from vllm.sequence import IntermediateTensors +from vllm.tasks import SupportedTask +from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib +from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE +from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput +from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec +from vllm.v1.outputs import ( + DraftTokenIds, + ECConnectorOutput, + ModelRunnerOutput, + RoutedExpertsTensors, +) +from vllm.v1.worker.block_table import get_block_table_width +from vllm.v1.worker.cp_utils import check_attention_cp_compatibility +from vllm.v1.worker.gpu import pcp_manager as pcp +from vllm.v1.worker.gpu.async_utils import ( + AsyncOutput, + AsyncPoolingOutput, + StepTimingCollector, +) +from vllm.v1.worker.gpu.attn_utils import ( + build_slot_mappings_by_layer, + get_kv_cache_spec, + init_attn_backend, + init_kv_cache, +) +from vllm.v1.worker.gpu.block_table import BlockTables +from vllm.v1.worker.gpu.buffer_utils import ( + async_copy_to_gpu, + set_default_max_concurrency, +) +from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens +from vllm.v1.worker.gpu.cudagraph_utils import ( + BatchExecutionDescriptor, + ModelCudaGraphManager, + normalize_model_token_inputs, +) +from vllm.v1.worker.gpu.cudagraph_utils import ( + profile_cudagraph_memory as _profile_cudagraph_memory, +) +from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp +from vllm.v1.worker.gpu.ec_connector import get_ec_connector +from vllm.v1.worker.gpu.eplb_utils import EPLBController, step_eplb_after +from vllm.v1.worker.gpu.input_batch import ( + InputBatch, + InputBuffers, + combine_sampled_and_draft_tokens, + expand_idx_mapping, + post_update, + post_update_num_computed_tokens, + prepare_pos_seq_lens, + prepare_prefill_inputs, + set_dummy_context, +) +from vllm.v1.worker.gpu.kv_connector import ( + NO_OP_KV_CONNECTOR, + KVConnector, + get_kv_connector, +) +from vllm.v1.worker.gpu.lora_utils import ( + LoraState, + create_lora_capture_hook, + get_lora_capture_cases, + get_num_active_loras_for_dispatch, +) +from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache +from vllm.v1.worker.gpu.mm.lora import set_active_mm_loras +from vllm.v1.worker.gpu.model_states import init_model_state +from vllm.v1.worker.gpu.pool.pooling_runner import PoolingRunner +from vllm.v1.worker.gpu.pp_utils import PPHandler +from vllm.v1.worker.gpu.sample.batch_shard import ( + BatchSharder, + all_to_all_logits, + gather_sampler_output, +) +from vllm.v1.worker.gpu.sample.output import SamplerOutput +from vllm.v1.worker.gpu.sample.prompt_logprob import PromptLogprobsWorker +from vllm.v1.worker.gpu.sample.sampler import Sampler +from vllm.v1.worker.gpu.shutdown import free_before_shutdown +from vllm.v1.worker.gpu.spec_decode import init_speculator +from vllm.v1.worker.gpu.spec_decode.adaptive_verification import ( + AdaptiveVerificationManager, + maybe_create_adaptive_verification_manager, +) +from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import ( + set_eagle3_aux_hidden_state_layers, +) +from vllm.v1.worker.gpu.spec_decode.rejection_sampler import ( + RejectionSampler, + get_max_chunk_logits, +) +from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator +from vllm.v1.worker.gpu.spec_decode.utils import ( + DraftTokensHandler, + limit_draft_tokens, +) +from vllm.v1.worker.gpu.states import RequestState +from vllm.v1.worker.gpu.structured_outputs import StructuredOutputsWorker +from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin +from vllm.v1.worker.utils import ( + KVBlockZeroer, + copy_kv_cache_blocks_inplace, + get_uniform_decode_token_count, +) +from vllm.v1.worker.workspace import use_workspace_lane + +logger = init_logger(__name__) + + +class GPUModelRunner(LoRAModelRunnerMixin): + def __init__(self, vllm_config: VllmConfig, device: torch.device): + self.vllm_config = vllm_config + self.model_config = vllm_config.model_config + self.cache_config = vllm_config.cache_config + self.compilation_config = vllm_config.compilation_config + self.lora_config = vllm_config.lora_config + self.load_config = vllm_config.load_config + self.parallel_config = vllm_config.parallel_config + self.scheduler_config = vllm_config.scheduler_config + self.speculative_config = vllm_config.speculative_config + self._draft_workspace_lane = int( + self.speculative_config is not None and self.speculative_config.use_dspark() + ) + self.observability_config = vllm_config.observability_config + self.jit_warmup_registry = JitWarmupRegistry(vllm_config) + + self.device = device + self.dtype = self.model_config.dtype + self.kv_cache_dtype = self.dtype + if self.cache_config.cache_dtype != "auto": + # Quantized KV cache. + self.kv_cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[ + self.cache_config.cache_dtype + ] + + # Lazily initialized in _init_kv_zero_meta() when the KV cache needs + # zeroing (e.g. hybrid models with fp8 KV cache). + self.kv_block_zeroer: KVBlockZeroer | None = None + + self.vocab_size = self.model_config.get_vocab_size() + self.max_model_len = self.model_config.max_model_len + self.max_num_tokens = self.scheduler_config.max_num_batched_tokens + self.max_num_reqs = self.scheduler_config.max_num_seqs + self.is_encoder_decoder = self.model_config.is_encoder_decoder + + self.output_copy_stream = torch.cuda.Stream(self.device) + + # Pipeline parallelism. + self.use_pp = self.parallel_config.pipeline_parallel_size > 1 + self.is_first_pp_rank = get_pp_group().is_first_rank + self.is_last_pp_rank = get_pp_group().is_last_rank + + # Size the UVA buffer pools to the max number of concurrent in-flight + # steps. Must run before any pooled buffer is constructed + set_default_max_concurrency(vllm_config.max_concurrent_batches) + + # PP broadcast/recv helper. Runs the collective on a side stream. + self.pp_handler: PPHandler | None = None + + # Persistent buffer for intermediate tensors (non-first PP ranks). + self.intermediate_tensors: IntermediateTensors | None = None + + # Data parallelism. + self.dp_size = self.parallel_config.data_parallel_size + self.dp_rank = self.parallel_config.data_parallel_rank + + # Detect EP all2all peer faults to prevent emitting corrupted output. + # Only meaningful for MoE + DP with an FT-capable all2all backend. + self.check_ep_fault = False + if self.dp_size > 1 and self.model_config.is_moe: + self.check_ep_fault = get_ep_all2all_manager().support_fault_tolerance + + # Decode context parallelism. + self.dcp_size = self.parallel_config.decode_context_parallel_size + self.use_dcp = self.dcp_size > 1 + self.dcp_rank = get_dcp_group().rank_in_group if self.use_dcp else 0 + self.cp_interleave = self.parallel_config.cp_kv_cache_interleave_size + + # Multimodal + self.mm_registry = MULTIMODAL_REGISTRY + self.supports_mm_inputs = self.mm_registry.supports_multimodal_inputs( + self.model_config + ) + self.uses_inputs_embeds = ( + self.supports_mm_inputs or self.model_config.enable_prompt_embeds + ) + self.encoder_cache = None + if self.supports_mm_inputs and self.is_first_pp_rank: + self.encoder_cache = EncoderCache() + self.ec_connector = get_ec_connector(vllm_config, self.encoder_cache) + + # Speculative decoding. + self.speculator = None + self.use_aux_hidden_state_outputs = False + self.num_speculative_steps = vllm_config.num_speculative_tokens + if self.speculative_config is not None: + if self.is_last_pp_rank: + self.speculator = init_speculator(self.vllm_config, self.device) + + if self.speculative_config.method in ( + "eagle3", + "dflash", + "dspark", + "extract_hidden_states", + ): + # Drafting may require auxiliary hidden states from target model outputs + self.use_aux_hidden_state_outputs = True + if self.use_pp: + raise ValueError( + f"{self.speculative_config.method} with pipeline parallel " + "is not supported." + ) + + # Draft tokens propagation - for spec-dec + struct outputs. + self.draft_tokens_handler = DraftTokensHandler(self.device) + + self.pcp_manager: pcp.PCPManager | None = None + + # Pooling models. + self.is_pooling_model = self.model_config.runner_type == "pooling" + self.pooling_runner: PoolingRunner | None = None + + # Multi-module MTP feeds its modules the next num_speculative_steps prefill + # tokens during chunked prefill. Other speculators only read the immediate + # next one. + num_prefill_lookahead = ( + self.num_speculative_steps + if self.speculative_config is not None + and self.speculative_config.use_multi_module_mtp() + else 1 + ) + + self.step_timing = StepTimingCollector() + + # General request states. + self.req_states = RequestState( + max_num_reqs=self.max_num_reqs, + max_model_len=self.max_model_len, + max_num_batched_tokens=self.max_num_tokens, + num_speculative_steps=self.num_speculative_steps, + vocab_size=self.vocab_size, + device=self.device, + num_prefill_lookahead=num_prefill_lookahead, + ) + self.adaptive_verification: AdaptiveVerificationManager | None = None + self.input_buffers = InputBuffers( + max_num_reqs=self.max_num_reqs, + max_num_tokens=self.max_num_tokens, + device=self.device, + ) + if self.use_pp: + self.pp_handler = PPHandler( + max_num_reqs=self.max_num_reqs, + num_speculative_steps=self.num_speculative_steps, + device=self.device, + ) + + # Samplers and decode_query_len created in load_model() after + # model_state exists (num_new_sampled_tokens_per_step from ModelState). + self.sampler: Sampler | None = None + self.rejection_sampler: RejectionSampler | None = None + self.batch_sharder: BatchSharder | None = None + self.prompt_logprobs_worker: PromptLogprobsWorker | None = None + self.structured_outputs_worker: StructuredOutputsWorker | None = None + self.cudagraph_manager: ModelCudaGraphManager | None = None + + # LoRA-related workers. + self.lora_state = LoraState(max_num_reqs=self.max_num_reqs) + self.lora_capture_cases = [0] + if self.lora_config: + self.lora_capture_cases = get_lora_capture_cases( + self.lora_config, self.compilation_config + ) + + # KV Connector if configured. + self.kv_connector: KVConnector = NO_OP_KV_CONNECTOR + + # For transferring state from execute_model to subsequent sample_tokens call. + self.execute_model_state: ExecuteModelState | None = None + + # Expert parallelism load balancer. + self.eplb = EPLBController(self.parallel_config, self.device) + self.routed_experts_capturer: RoutedExpertsCapturer | None = None + + set_offloader(create_offloader(self.vllm_config.offload_config)) + + def update_max_model_len(self, max_model_len: int) -> None: + self.max_model_len = max_model_len + self.req_states.max_model_len = max_model_len + update_model_len = getattr( + getattr(self, "model", None), "update_max_model_len", None + ) + if update_model_len is not None: + update_model_len(max_model_len) + if isinstance(self.speculator, DraftModelSpeculator): + self.speculator.update_max_model_len(max_model_len) + + def init_routed_experts_capturer(self) -> None: + """Initialize target-model capture on every participating worker.""" + self.routed_experts_capturer = RoutedExpertsCapturer( + max_num_batched_tokens=self.max_num_tokens, + vllm_config=self.vllm_config, + kv_cache_config=self.kv_cache_config, + ) + bind_routed_experts_capturer(self.model, self.routed_experts_capturer) + + def get_supported_tasks(self) -> tuple[SupportedTask, ...]: + tasks: list[SupportedTask] = [] + if self.model_config.runner_type == "generate": + tasks.extend(self.model_state.get_supported_generation_tasks()) + if self.is_pooling_model: + # Do not rely on pooling_runner here, since this information is needed + # on the first PP rank, while pooling_runner is only initialized + # on the last PP rank. + tasks.extend(PoolingRunner.get_supported_tasks(self.model)) + return tuple(tasks) + + def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None: + time_before_load = time.perf_counter() + if load_dummy_weights: + self.load_config.load_format = "dummy" + self.eplb.prepare_load() + eplb_models_added = False + with DeviceMemoryProfiler() as m: + model_loader = get_model_loader(self.vllm_config.load_config) + logger.info_once("Loading model from scratch...") + + self.model = model_loader.load_model( + vllm_config=self.vllm_config, model_config=self.vllm_config.model_config + ) + if self.lora_config: + self.model = self.load_lora_model( + self.model, self.vllm_config, self.device + ) + + if self.use_aux_hidden_state_outputs: + assert self.speculative_config is not None + set_eagle3_aux_hidden_state_layers(self.model, self.speculative_config) + if isinstance(self.speculator, DraftModelSpeculator): + with use_workspace_lane(self._draft_workspace_lane): + self.speculator.load_model(self.model) + eplb_models_added = self.eplb.maybe_register_speculator( + self.speculator, self.speculative_config, load_dummy_weights + ) + time_after_load = time.perf_counter() + + self.model_memory_usage = m.consumed_memory + logger.info( + "Model loading took %s GiB memory and %.6f seconds", + format_gib(m.consumed_memory), + time_after_load - time_before_load, + ) + + # Initialize the components that require the model. + self.model_state = init_model_state( + self.vllm_config, self.model, self.encoder_cache, self.device + ) + + self.decode_query_len = ( + self.num_speculative_steps + + self.model_state.num_new_sampled_tokens_per_step + ) + + if self.parallel_config.enable_batch_sharded_sampling: + if hasattr(self.model, "compute_logits_local"): + self.batch_sharder = BatchSharder( + max_num_reqs=self.max_num_reqs, + max_num_logits_per_req=self.decode_query_len, + device=self.device, + ) + logger.info("Batch-sharded sampling enabled.") + else: + logger.warning_once( + "Disabling batch-sharded sampling: %s does not implement " + "compute_logits_local", + type(self.model).__name__, + ) + + # Initialize samplers. Model states may override via custom_sampler(). + if self.is_last_pp_rank and not self.is_pooling_model: + self.sampler = Sampler( + max_num_reqs=self.max_num_reqs, + vocab_size=self.vocab_size, + device=self.device, + req_states=self.req_states, + logprobs_mode=self.model_config.logprobs_mode, + num_speculative_tokens=self.decode_query_len, + use_fp64_gumbel=self.model_config.use_fp64_gumbel, + enable_trace_replay=self.model_config.enable_trace_replay, + reasoning_config=self.vllm_config.reasoning_config, + return_sampling_mask=self.model_config.return_sampling_mask, + ) + custom = self.model_state.custom_sampler(self.sampler) + + if custom: + self.sampler, self.rejection_sampler = custom + elif self.speculative_config is not None: + self.rejection_sampler = RejectionSampler( + self.sampler, + self.speculative_config, + self.device, + ) + self.prompt_logprobs_worker = PromptLogprobsWorker( + self.max_num_reqs, + logprobs_mode=self.model_config.logprobs_mode, + ) + self.structured_outputs_worker = StructuredOutputsWorker( + max_num_logits=self.max_num_reqs * self.decode_query_len, + vocab_size=self.vocab_size, + device=self.device, + mask_stride=self.decode_query_len, + num_bonus_tokens=self.model_state.num_new_sampled_tokens_per_step, + ) + + if self.is_pooling_model and self.is_last_pp_rank: + self.pooling_runner = PoolingRunner(self.model, self.vllm_config) + eplb_models_added |= self.eplb.maybe_register_model( + self.model, + self.model_config, + load_dummy_weights, + ) + self.eplb.maybe_start_async_loop(eplb_models_added) + + if not self.is_first_pp_rank: + # For non-first PP ranks, create intermediate tensors sized + # for the max capture size so they can be sliced per batch. + # Save as persistent member so runtime can copy received data + # into the same addresses that the CUDA graphs captured. + self.intermediate_tensors = self.model.make_empty_intermediate_tensors( + batch_size=self.max_num_tokens, + dtype=self.model_config.dtype, + device=self.device, + ) + + get_offloader().post_init() + + def get_model(self) -> nn.Module: + return self.model + + def get_draft_model(self) -> nn.Module | None: + speculator = self.speculator + if not isinstance(speculator, DraftModelSpeculator): + return None + return speculator.model + + def reload_weights(self, *args, **kwargs) -> None: + # TODO(Wentao): Use full version instead of import when fully migrated to v2 + from vllm.v1.worker.gpu_model_runner import GPUModelRunner as GPUModelRunnerV1 + + GPUModelRunnerV1.reload_weights(self, *args, **kwargs) # type: ignore[arg-type] + + def update_config(self, *args, **kwargs) -> None: + # TODO(Wentao): Use full version instead of import when fully migrated to v2 + from vllm.v1.worker.gpu_model_runner import GPUModelRunner as GPUModelRunnerV1 + + GPUModelRunnerV1.update_config(self, *args, **kwargs) # type: ignore[arg-type] + + # v2 reads config via self.vllm_config (e.g. in load_model), so keep it + # in sync with the attributes the v1 helper just replaced. + self.vllm_config.model_config = self.model_config + self.vllm_config.load_config = self.load_config + + @functools.cached_property + def main_stream(self) -> torch.cuda.Stream: + # Cache the default CUDA stream to avoid lookup overhead. + return torch.cuda.current_stream(self.device) + + def get_encoder_timing_stats(self) -> dict[str, dict[str, float | int]]: + encoder_runner = getattr(self.model_state, "encoder_runner", None) + if encoder_runner is None: + return {} + return encoder_runner.get_encoder_timing_stats() + + def get_kv_cache_spec(self): + return get_kv_cache_spec(self.vllm_config) + + def initialize_kv_cache( + self, kv_cache_config: KVCacheConfig, is_profiling: bool = False + ) -> None: + kv_cache_config = deepcopy(kv_cache_config) + self.kv_cache_config = kv_cache_config + + block_table_max_model_len = self.max_model_len + if self.is_encoder_decoder: + # Cross-attention block tables need to index encoder tokens, which + # can exceed the decoder's max_model_len. + block_table_max_model_len = max( + block_table_max_model_len, + self.scheduler_config.max_num_encoder_input_tokens, + getattr(self.model_config.hf_config, "max_source_positions", 0), + ) + + block_sizes = [] + max_num_blocks_per_group = [] + group_cp_sizes = [] + for kv_cache_group in kv_cache_config.kv_cache_groups: + spec = kv_cache_group.kv_cache_spec + block_sizes.append(spec.block_size) + group_cp_sizes.append( + 1 if getattr(spec, "dcp_replicated", False) else self.dcp_size + ) + # Let each cache type account for CP. Attention KV is DCP-sharded, + # while Mamba/GDN recurrent state is replicated across DCP ranks. + max_num_blocks = spec.max_num_blocks_per_req( + self.vllm_config, block_table_max_model_len + ) + # Preserve each cache type's alignment requirements after applying + # its topology-aware block-table width. + if isinstance(spec, MambaSpec): + max_num_blocks = get_block_table_width( + max_num_blocks, spec.block_size, token_alignment=None + ) + else: + max_num_blocks = get_block_table_width(max_num_blocks, spec.block_size) + max_num_blocks_per_group.append(max_num_blocks) + + target_attn_layer_names = None + if isinstance(self.speculator, DraftModelSpeculator): + # Adaptive verification validates target attention separately. + target_attn_layer_names = { + layer_name + for group in self.kv_cache_config.kv_cache_groups + for layer_name in group.layer_names + } - self.speculator.draft_attn_layer_names + self.attn_groups, attn_cg_support, self.kernel_block_sizes = init_attn_backend( + self.kv_cache_config, + self.vllm_config, + self.device, + ) + additional_attn_cg_support = self.model_state.get_additional_cg_support() + attn_cg_support = attn_cg_support.narrow(*additional_attn_cg_support) + # The speculator clears the flag at load time when the checkpoint has + # no confidence head, so it holds the effective value. + self.adaptive_verification = maybe_create_adaptive_verification_manager( + enable_adaptive_verification=getattr( + self.speculator, "enable_adaptive_verification", False + ), + attn_groups=self.attn_groups, + attn_cg_support=attn_cg_support, + req_states=self.req_states, + query_start_loc=self.input_buffers.query_start_loc, + num_bonus_tokens=self.model_state.num_new_sampled_tokens_per_step, + max_total_logits=get_max_chunk_logits(self.vocab_size), + vllm_config=self.vllm_config, + target_layer_names=target_attn_layer_names, + additional_attn_cg_support=additional_attn_cg_support, + ) + + self.block_tables = BlockTables( + block_sizes=block_sizes, + max_num_reqs=self.max_num_reqs, + max_num_batched_tokens=self.max_num_tokens, + max_num_blocks_per_group=max_num_blocks_per_group, + device=self.device, + kernel_block_sizes=self.kernel_block_sizes, + cp_size=self.dcp_size, + cp_rank=self.dcp_rank, + cp_interleave=self.cp_interleave, + group_cp_sizes=group_cp_sizes, + ) + self.pcp_manager = pcp.maybe_build_pcp_manager( + self.vllm_config, + self.device, + self.supports_mm_inputs, + self.req_states, + self.block_tables, + cls=self.pcp_manager_cls, + ) + initialize_mamba_ssu_backend( + self.vllm_config.mamba_config, self.kv_cache_config + ) + if self.adaptive_verification is not None: + self.compilation_config.cudagraph_mode = CUDAGraphMode.FULL_AND_PIECEWISE + cudagraph_mode = self.compilation_config.resolve_cudagraph_mode_and_sizes( + attn_cg_support.min_cg_support, + attn_cg_support.min_cg_attn_backend, + self.decode_query_len, + use_v2_model_runner=True, + tensor_parallel_size=self.parallel_config.tensor_parallel_size, + kv_cache_config=self.kv_cache_config, + max_num_reqs=self.max_num_reqs, + is_profiling=is_profiling, + ) + self.cudagraph_manager = ModelCudaGraphManager( + self.vllm_config, + self.device, + cudagraph_mode, + decode_query_len=self.decode_query_len, + lora_capture_cases=self.lora_capture_cases, + varlen_decode=self.adaptive_verification is not None, + ) + check_attention_cp_compatibility(self.vllm_config) + if isinstance(self.speculator, DraftModelSpeculator): + # HACK(woosuk) + self.speculator.set_attn( + self.model_state, + self.kv_cache_config, + self.block_tables, + self.input_buffers, + self.attn_groups, + ) + if self.speculator is not None: + # After set_attn, so the speculator can size its cudagraph mode + # to its own attention support. + self.speculator.init_cudagraph_manager(cudagraph_mode) + + self.kv_caches: list[torch.Tensor] = [] + kv_caches_dict = init_kv_cache( + self.kv_caches, + self.compilation_config.static_forward_context, + self.kv_cache_config, + self.device, + self.kernel_block_sizes, + self.vllm_config, + ) + if is_profiling: + self.kv_connector = NO_OP_KV_CONNECTOR + else: + self.kv_connector = get_kv_connector(self.vllm_config, kv_caches_dict) + + def _init_kv_zero_meta(self) -> None: + """Build KV-block zeroing metadata; invoked from gpu_worker.""" + self.kv_block_zeroer = KVBlockZeroer( + self.device, + attn_groups_iter=(g for groups in self.attn_groups for g in groups), + kernel_block_sizes=self.kernel_block_sizes, + static_forward_context=self.compilation_config.static_forward_context, + ) + + @torch.inference_mode() + @step_eplb_after(is_dummy=True) + def _dummy_run( + self, + num_tokens: int, + *args, + skip_attn: bool = False, + uniform_decode: bool = False, + context_len: int = 0, + skip_eplb: bool = False, + is_profile: bool = False, + **kwargs, + ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + if skip_attn and not is_profile: + raise ValueError( + "skip_attn must only be True for initial memory profiling." + ) + + # Create a dummy scheduler output. + num_reqs = min(num_tokens, self.max_num_reqs) + if uniform_decode: + # HACK(lucas): for now since the worker is shared between MRV1 and MRV2, + # and for spec-decode with MTP we want to make sure the dummy runs use + # 1+num_speculative_tokens we use max here, this will likely be eventually + # changed in the worker: https://github.com/vllm-project/vllm/pull/35243 + num_tokens = max(num_tokens, self.decode_query_len) + num_reqs = num_tokens // self.decode_query_len + assert num_tokens % self.decode_query_len == 0 + # Distribute the remainder evenly so no dummy request exceeds + # ceil(num_tokens / num_reqs) <= max_model_len tokens. + num_tokens_per_request = [ + num_tokens // num_reqs + (i >= num_reqs - num_tokens % num_reqs) + for i in range(num_reqs) + ] + + assert sum(num_tokens_per_request) == num_tokens + num_scheduled_tokens = { + f"_dummy_req_{i}": n for i, n in enumerate(num_tokens_per_request) + } + dummy_scheduler_output = SchedulerOutput.make_empty() + dummy_scheduler_output.total_num_scheduled_tokens = num_tokens + dummy_scheduler_output.num_scheduled_tokens = num_scheduled_tokens + + # Disable any use of KVConnector for dummy runs. + self.kv_connector.set_disabled(True) + + # Get the intermediate tensors for the dummy run. + intermediate_tensors = None + if not self.is_first_pp_rank: + assert self.intermediate_tensors is not None + intermediate_tensors = self.intermediate_tensors[:num_tokens] + + max_loras = self.lora_config.max_loras if self.lora_config is not None else 0 + with self.maybe_dummy_run_with_lora( + self.lora_config, + num_scheduled_tokens=np.array(num_tokens_per_request, dtype=np.int32), + num_sampled_tokens=None, + remove_lora=True, + num_active_loras=max_loras, + ): + # Execute the model. + self.execute_model( + dummy_scheduler_output, + intermediate_tensors=intermediate_tensors, + dummy_run=True, + skip_attn_for_dummy_run=skip_attn, + is_profile=is_profile, + context_len=context_len, + ) + self.kv_connector.set_disabled(False) + + # Non-last PP ranks don't produce output for sampling. + if not self.is_last_pp_rank: + return None, None + + assert self.execute_model_state is not None + input_batch = self.execute_model_state.input_batch + attn_metadata = self.execute_model_state.attn_metadata + slot_mappings_by_layer = self.execute_model_state.slot_mappings_by_layer + hidden_states = self.execute_model_state.hidden_states + aux_hidden_states = self.execute_model_state.aux_hidden_states + self.execute_model_state = None + + self.step_timing.forward_end() + + # dummy run the eagle speculator's propose to ensure DP/EP sync. + if self.speculator is not None: + assert self.sampler is not None + self.step_timing.drafter_start() + mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None + if self.speculator.supports_mm_inputs: + mm_inputs = ( + [], + torch.zeros( + input_batch.num_tokens, + dtype=torch.bool, + device="cpu", + ), + ) + + # Let the target override the hidden state fed to the drafter + # (e.g. DeepSeek V4 MTP needs the pre-hc_head residual). The + # target returns a persistent buffer sized at max_num_batched_tokens; + # slice to the active token count that propose() expects. + spec_hidden_states = hidden_states + if hasattr(self.model, "get_mtp_target_hidden_states"): + pre_hc_hidden_states = self.model.get_mtp_target_hidden_states() + spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]] # type: ignore[union-attr] + with use_workspace_lane(self._draft_workspace_lane): + self.speculator.propose( + input_batch=input_batch, + attn_metadata=attn_metadata, + slot_mappings=slot_mappings_by_layer, + last_hidden_states=spec_hidden_states, + aux_hidden_states=aux_hidden_states, + num_sampled=torch.ones( + input_batch.num_reqs, dtype=torch.int32, device=self.device + ), + num_rejected=torch.zeros( + input_batch.num_reqs, dtype=torch.int32, device=self.device + ), + last_sampled=self.req_states.last_sampled_tokens, + next_prefill_tokens=self.req_states.next_prefill_tokens, + temperature=self.sampler.sampling_states.temperature.gpu, + seeds=self.sampler.sampling_states.seeds.gpu, + dummy_run=True, + skip_attn_for_dummy_run=skip_attn, + mm_inputs=mm_inputs, + is_profile=is_profile, + ) + self.step_timing.drafter_end() + + assert hidden_states is not None # Last PP rank always has hidden_states + sample_hidden_states = hidden_states[input_batch.logits_indices] + return hidden_states, sample_hidden_states + + @torch.inference_mode() + def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> None: + num_reqs = hidden_states.shape[0] + logits = self.model.compute_logits(hidden_states) + dummy_input_batch = InputBatch.make_dummy( + num_reqs, num_reqs, self.input_buffers + ) + + # NOTE(woosuk): During the initial memory profiling, the sampler may skip + # top_k, top_p, and logprobs, using less GPU memory than what is possible + # during actual execution. + assert self.sampler is not None + self.sampler(logits, dummy_input_batch) + + @torch.inference_mode() + def _dummy_pooler_run(self, hidden_states: torch.Tensor) -> None: + assert self.pooling_runner is not None + self.pooling_runner.dummy_pooler_run(hidden_states) + + @torch.inference_mode() + def profile_run(self) -> None: + if self.supports_mm_inputs and self.is_first_pp_rank: + mm_config = self.model_config.multimodal_config + if mm_config is not None and not mm_config.skip_mm_profiling: + mm_budget = MultiModalBudget( + self.vllm_config, + self.mm_registry, + enable_cache=False, + ) + dummy_mm_inputs = get_dummy_encoder_profile_inputs( + self.mm_registry, + mm_budget, + ) + self.model_state.encoder_runner.profile_encoder_cache( + dummy_mm_inputs, mm_budget + ) + + hidden_states, sample_hidden_states = self._dummy_run( + self.max_num_tokens, skip_attn=True, is_profile=True + ) + + # Only run sampler/pooler on last PP rank (non-last ranks return None). + if self.is_last_pp_rank: + assert sample_hidden_states is not None + if self.pooling_runner is None: + self._dummy_sampler_run(sample_hidden_states) + else: + self._dummy_pooler_run(hidden_states) + + torch.accelerator.synchronize() + del hidden_states, sample_hidden_states + self.reset_encoder_cache() + gc.collect() + + def post_kv_cache_wake_up(self) -> None: + self.block_tables.init_block_table_layout_tensors() + + def reset_mm_cache(self) -> None: + if self.encoder_cache is not None: + self.encoder_cache.reset_mm_cache() + + def reset_encoder_cache(self) -> None: + if self.encoder_cache is not None: + self.encoder_cache.reset_encoder_cache() + if self.pooling_runner is not None: + self.pooling_runner.clear() + + @torch.inference_mode() + def profile_cudagraph_memory(self) -> int: + """Estimate the GPU memory required to capture CUDA graphs.""" + return _profile_cudagraph_memory(self) + + @torch.inference_mode() + def capture_model(self) -> int: + assert self.cudagraph_manager is not None + capture_encoder = ( + self.model_state.supports_mm_inputs + and self.model_state.encoder_runner.has_cudagraph() + ) + capture_decoder = self.cudagraph_manager.needs_capture() + if not capture_encoder and not capture_decoder: + logger.warning( + "Skipping encoder and decoder CUDA graph capture. To enable " + "encoder capture, ensure `cudagraph_mm_encoder` is enabled; " + "to enable decoder capture, ensure `cudagraph_mode` is not `NONE`." + ) + return 0 + + compilation_counter.num_gpu_runner_capture_triggers += 1 + + start_time = time.perf_counter() + gc.collect() + torch.accelerator.empty_cache() + start_free_gpu_memory = torch.accelerator.get_memory_info()[0] + + with self.maybe_setup_dummy_loras(self.lora_config): + if capture_encoder: + self.model_state.encoder_runner.capture() + + if capture_decoder: + self.cudagraph_manager.capture( + self.model, + self.model_state, + self.input_buffers, + self.intermediate_tensors, + self.block_tables, + self.attn_groups, + self.kv_cache_config, + has_lora=self.lora_config is not None, + use_aux_hidden_state_outputs=self.use_aux_hidden_state_outputs, + lora_capture_hook=create_lora_capture_hook(self.lora_config, self), + ) + if self.speculator is not None: + with use_workspace_lane(self._draft_workspace_lane): + self.speculator.capture() + if self.adaptive_verification is not None: + with self.step_timing.collect() as timings: + for batch in self.adaptive_verification.batches_to_profile( + self.cudagraph_manager.captured_token_counts() + ): + self._dummy_run(**batch) + self.adaptive_verification.set_initial_cost_curves(timings) + + end_time = time.perf_counter() + end_free_gpu_memory = torch.accelerator.get_memory_info()[0] + elapsed_time = end_time - start_time + cuda_graph_size = start_free_gpu_memory - end_free_gpu_memory + # This usually takes 5~20 seconds. + logger.info( + "Graph capturing finished in %.0f secs, took %.2f GiB", + elapsed_time, + cuda_graph_size / (1 << 30), + ) + return cuda_graph_size + + def _remove_request(self, req_id: str) -> bool: + # Call model_state.remove_request *before* req_states.remove_request + # so the model_state can still look up the slot index. + self.model_state.remove_request(req_id) + req_idx = self.req_states.remove_request(req_id) + if req_idx is None: + return False + if self.pooling_runner is not None: + self.pooling_runner.remove_request(req_idx) + if self.pp_handler is not None: + self.pp_handler.on_req_idx_freed(req_idx) + if self.encoder_cache is not None: + self.encoder_cache.remove_request(req_id) + if self.prompt_logprobs_worker is not None: + self.prompt_logprobs_worker.remove_request(req_id) + self.lora_state.remove_request(req_id) + return True + + def finish_requests(self, scheduler_output: SchedulerOutput) -> None: + finished_req_ids = scheduler_output.finished_req_ids + if self.pooling_runner is not None: + # Preempted docs keep their query-use reservation until rescheduled. + self.pooling_runner.on_requests_finished(finished_req_ids) + preempted_req_ids = scheduler_output.preempted_req_ids + if preempted_req_ids: + finished_req_ids = finished_req_ids.union(preempted_req_ids) + # Sorted so every TP rank frees request slots in the same order. + # Features like batch-sharded sampling derive rank request ownership + # from the slot index. + for req_id in sorted(finished_req_ids): + self._remove_request(req_id) + + def free_states(self, scheduler_output: SchedulerOutput) -> None: + if self.encoder_cache is not None: + for mm_hash in scheduler_output.free_encoder_mm_hashes: + self.encoder_cache.free_encoder_cache(mm_hash) + + def update_pp_decode_requests(self): + # For non-last PP ranks, update decode requests with sampler output from + # the prior step in which they were scheduled (pp_size steps ago). + if self.pp_handler is not None: + outputs = self.pp_handler.get_prev_sampled_outputs() + if outputs is not None: + self.postprocess_sampled(**outputs) + + def add_requests(self, scheduler_output: SchedulerOutput) -> None: + for new_req_data in scheduler_output.scheduled_new_reqs: + assert new_req_data.prefill_token_ids is not None + req_id = new_req_data.req_id + + # Streaming input update: request already exists from a prior + # chunk. Remove old state so it can be cleanly re-added below + # with the updated prompt_token_ids and mm_features. + self._remove_request(req_id) + + prompt_len = new_req_data.prompt_len + sampling_params = new_req_data.sampling_params + self.req_states.add_request( + req_id=req_id, + prompt_len=prompt_len, + all_token_ids=new_req_data.prefill_token_ids, + num_computed_tokens=new_req_data.num_computed_tokens, + max_tokens=sampling_params.max_tokens if sampling_params else 1, # type: ignore[arg-type] + ) + req_index = self.req_states.req_id_to_index[req_id] + if self.adaptive_verification is not None: + self.adaptive_verification.add_request(req_index) + + if self.pooling_runner is not None: + assert new_req_data.pooling_params is not None + assert new_req_data.prompt_token_ids is not None + self.pooling_runner.add_request( + req_id, + req_index, + new_req_data.pooling_params, + new_req_data.prompt_token_ids, + ) + + if self.encoder_cache is not None: + self.encoder_cache.add_request(req_id, new_req_data.mm_features) + + self.model_state.add_request(req_index, new_req_data) + self.block_tables.append_block_ids( + req_index, new_req_data.block_ids, overwrite=True + ) + self.lora_state.add_request(req_id, req_index, new_req_data.lora_request) + + if self.is_last_pp_rank and new_req_data.sampling_params is not None: + assert self.sampler is not None + self.sampler.add_request( + req_index, prompt_len, new_req_data.sampling_params + ) + assert self.prompt_logprobs_worker is not None + self.prompt_logprobs_worker.add_request( + req_id, req_index, new_req_data.sampling_params + ) + + if scheduler_output.scheduled_new_reqs: + self.req_states.apply_staged_writes() + self.model_state.apply_staged_writes() + if self.sampler is not None: + self.sampler.apply_staged_writes() + + def update_requests(self, scheduler_output: SchedulerOutput) -> None: + # Add new blocks and update num_computed_tokens for the existing requests. + reqs = scheduler_output.scheduled_cached_reqs + num_computed_tokens_np = self.req_states.num_computed_tokens_np + for req_id, num_computed_tokens, req_new_block_ids in zip( + reqs.req_ids, reqs.num_computed_tokens, reqs.new_block_ids + ): + req_index = self.req_states.req_id_to_index[req_id] + num_computed_tokens_np[req_index] = num_computed_tokens + if req_new_block_ids is not None: + self.block_tables.append_block_ids( + req_index, req_new_block_ids, overwrite=False + ) + + # Update CPU num_computed_prefill_tokens. + np.minimum( + self.req_states.num_computed_tokens_np, + self.req_states.prefill_len.np, + out=self.req_states.num_computed_prefill_tokens, + ) + + # Zero GPU memory for freshly allocated cache blocks to prevent + # stale NaN/data from corrupting attention or SSM computation. + if scheduler_output.new_block_ids_to_zero: + assert self.kv_block_zeroer is not None + self.kv_block_zeroer.zero_block_ids(scheduler_output.new_block_ids_to_zero) + + # Apply copy-on-write block copies for partial prefix-cache hits, after + # zeroing new blocks and before the forward pass reads them. + if scheduler_output.kv_cache_block_copies: + copy_kv_cache_blocks_inplace( + self.kv_caches, + self.kv_cache_config.num_blocks, + scheduler_output.kv_cache_block_copies, + ) + + def gather_batch_req_state( + self, scheduler_output: SchedulerOutput, dummy_run: bool + ) -> tuple["BatchReqState | None", int | None]: + """Gather CPU request state for the scheduled batch, in batch order. + Returns (batch_state, uniform_decode_token_count) + """ + num_tokens_per_req = scheduler_output.num_scheduled_tokens + num_reqs = len(num_tokens_per_req) + num_toks = scheduler_output.total_num_scheduled_tokens + max_query_len = max(scheduler_output.num_scheduled_tokens.values()) + + if dummy_run: + # Dummy batches are uniform by construction. + return None, get_uniform_decode_token_count( + num_reqs, num_toks, max_query_len, has_prefill=False + ) + + draft_tokens = scheduler_output.scheduled_spec_decode_tokens + # batch_idx -> req_id + req_ids = sort_batch_req_ids( + num_tokens_per_req, draft_tokens, self.decode_query_len + ) + + numtoks_iter = map(num_tokens_per_req.__getitem__, req_ids) + num_scheduled_tokens = np.fromiter(numtoks_iter, dtype=np.int32, count=num_reqs) + + idx_mapping_iter = map(self.req_states.req_id_to_index.__getitem__, req_ids) + idx_mapping_np = np.fromiter(idx_mapping_iter, dtype=np.intp, count=num_reqs) + prefill_len_np = self.req_states.prefill_len.np[idx_mapping_np] + num_computed_prefill_tokens_np = self.req_states.num_computed_prefill_tokens[ + idx_mapping_np + ] + is_prefilling_np = num_computed_prefill_tokens_np < prefill_len_np + + if self.adaptive_verification is not None and draft_tokens: + num_toks = self.adaptive_verification.get_num_tokens( + num_tokens_per_req, draft_tokens + ) + + batch_state = BatchReqState( + req_ids=req_ids, + num_scheduled_tokens=num_scheduled_tokens, + num_tokens=num_toks, + idx_mapping_np=idx_mapping_np, + prefill_len_np=prefill_len_np, + num_computed_prefill_tokens_np=num_computed_prefill_tokens_np, + is_prefilling_np=is_prefilling_np, + has_prefill=bool(is_prefilling_np.any()), + ) + return batch_state, get_uniform_decode_token_count( + num_reqs, num_toks, max_query_len, batch_state.has_prefill + ) + + def prepare_inputs( + self, + scheduler_output: SchedulerOutput, + batch_req_state: "BatchReqState", + batch_desc: BatchExecutionDescriptor, + ) -> InputBatch: + num_tokens = batch_req_state.num_tokens + num_tokens_after_padding = batch_desc.num_tokens + assert num_tokens > 0 + if envs.VLLM_MOE_SKIP_PADDING: + # Mark trailing cudagraph-padding rows so kernels can skip work for + # them when supported. + is_padding = self.input_buffers.is_padding + is_padding[:num_tokens].fill_(False) + is_padding[num_tokens:num_tokens_after_padding].fill_(True) + + req_ids = batch_req_state.req_ids + num_scheduled_tokens_np = batch_req_state.num_scheduled_tokens + idx_mapping_np = batch_req_state.idx_mapping_np + idx_mapping = async_copy_to_gpu(idx_mapping_np, device=self.device) + num_reqs = len(req_ids) + + # Get the number of draft tokens for each request. + draft_tokens = scheduler_output.scheduled_spec_decode_tokens + num_draft_tokens_per_req = None + if not draft_tokens: + # No draft token scheduled (common case). + total_num_draft_tokens = 0 + total_num_logits = num_reqs + cu_num_logits_np = np.arange(num_reqs + 1, dtype=np.int32) + cu_num_logits = torch.arange( + num_reqs + 1, device=self.device, dtype=torch.int32 + ) + expanded_idx_mapping = idx_mapping + expanded_local_pos = torch.zeros( + num_reqs, dtype=torch.int32, device=self.device + ) + else: + num_draft_tokens_per_req = np.fromiter( + (len(draft_tokens.get(req_id, ())) for req_id in req_ids), + dtype=np.int32, + count=num_reqs, + ) + num_bonus_tokens = self.model_state.num_new_sampled_tokens_per_step + total_num_draft_tokens = int(num_draft_tokens_per_req.sum()) + total_num_logits = num_reqs * num_bonus_tokens + total_num_draft_tokens + num_logits = num_draft_tokens_per_req + num_bonus_tokens + cu_num_logits_np = np.empty(num_reqs + 1, dtype=np.int32) + cu_num_logits_np[0] = 0 + np.cumsum(num_logits, out=cu_num_logits_np[1:]) + cu_num_logits = async_copy_to_gpu(cu_num_logits_np, device=self.device) + + adaptive_verification = ( + self.adaptive_verification if num_draft_tokens_per_req is not None else None + ) + num_scheduled_tokens_upper_bound = num_scheduled_tokens_np + if adaptive_verification is not None: + # num_scheduled_tokens represents the draft budget evenly distributed across + # all verification requests, `reallocate_drafts` will unevenly assign the + # draft budget to requests on the GPU side only. + num_scheduled_tokens_np, cu_num_logits_np = ( + adaptive_verification.compact_batch( + num_draft_tokens_per_req, + num_scheduled_tokens_np, + cu_num_logits_np, + ) + ) + + # Get query_start_loc. + # num_reqs_padded is None for PIECEWISE graphs (no request padding needed) + num_reqs_padded = batch_desc.num_reqs or num_reqs + query_start_loc_np = np.empty(self.max_num_reqs + 1, dtype=np.int32) + query_start_loc_np[0] = 0 + np.cumsum(num_scheduled_tokens_np, out=query_start_loc_np[1 : num_reqs + 1]) + # Pad for full CUDA graph mode. + # Some attention backends like FA3 require query_start_loc to be non-decreasing. + query_start_loc_np[num_reqs + 1 :] = num_tokens + query_start_loc = self.input_buffers.query_start_loc + async_copy_to_gpu(query_start_loc_np, out=query_start_loc) + if adaptive_verification is not None: + cu_num_logits, query_start_loc, total_num_draft_tokens = ( + adaptive_verification.reallocate_drafts(req_ids, idx_mapping) + ) + total_num_logits = num_reqs * num_bonus_tokens + total_num_draft_tokens + if draft_tokens: + expanded_idx_mapping, expanded_local_pos = expand_idx_mapping( + idx_mapping, total_num_logits, cu_num_logits, self.decode_query_len + ) + query_start_loc_np = query_start_loc_np[: num_reqs_padded + 1] + query_start_loc = query_start_loc[: num_reqs_padded + 1] + + # Get prefill tokens if any. + if batch_req_state.has_prefill: + prepare_prefill_inputs( + self.input_buffers.input_ids, + self.req_states.next_prefill_tokens, + idx_mapping, + query_start_loc, + self.req_states.all_token_ids.gpu, + self.req_states.prefill_len.gpu, + self.req_states.num_computed_tokens.gpu, + ) + + # Prepare positions and seq_lens. + prepare_pos_seq_lens( + idx_mapping, + query_start_loc, + self.req_states.num_computed_tokens.gpu, + self.input_buffers.positions, + self.input_buffers.seq_lens, + ) + seq_lens = self.input_buffers.seq_lens[:num_reqs_padded] + + dcp_local_seq_lens = None + if self.use_dcp: + # Prepare dcp local seq_lens. + prepare_dcp_local_seq_lens( + self.input_buffers.dcp_local_seq_lens, + self.input_buffers.seq_lens, + num_reqs, + self.dcp_size, + self.dcp_rank, + self.cp_interleave, + ) + dcp_local_seq_lens = self.input_buffers.dcp_local_seq_lens[:num_reqs_padded] + + # Some input token ids are directly read from the last sampled tokens + # and draft tokens. Also, get the logits indices to sample tokens from. + logits_indices = combine_sampled_and_draft_tokens( + self.input_buffers.input_ids, + idx_mapping, + self.req_states.last_sampled_tokens, + query_start_loc, + seq_lens, + self.req_states.prefill_len.gpu, + self.req_states.draft_tokens, + cu_num_logits, + total_num_logits, + self.model_state.num_new_sampled_tokens_per_step, + ) + + # CPU upper bound on seq_lens; padded entries left at zero. + num_computed_tokens_np = self.req_states.num_computed_tokens_np[idx_mapping_np] + seq_lens_cpu_upper_bound_np = np.zeros(num_reqs_padded, dtype=np.int32) + np.add( + num_computed_tokens_np, + num_scheduled_tokens_upper_bound, + out=seq_lens_cpu_upper_bound_np[:num_reqs], + ) + seq_lens_cpu_upper_bound = torch.from_numpy(seq_lens_cpu_upper_bound_np) + + max_seq_len_np = None + if self.use_pp: + # max_seq_len is only consumed by the PP `compute_need_sampled_mask` + max_seq_len_np = self.req_states.max_seq_len[idx_mapping_np] + + prompt_lens = None + if self.model_config.rswa_window is not None: + # prompt_lens is only used in R-SWA case. + prompt_lens = self.req_states.prompt_len.gpu[idx_mapping] + + input_batch = InputBatch( + req_ids=req_ids, + num_reqs=num_reqs, + num_reqs_after_padding=num_reqs_padded, + idx_mapping=idx_mapping, + idx_mapping_np=idx_mapping_np, + expanded_idx_mapping=expanded_idx_mapping, + expanded_local_pos=expanded_local_pos, + num_scheduled_tokens=num_scheduled_tokens_upper_bound, + num_tokens=num_tokens, + num_tokens_after_padding=num_tokens_after_padding, + num_draft_tokens=total_num_draft_tokens, + num_draft_tokens_per_req=num_draft_tokens_per_req, + query_start_loc=query_start_loc, + query_start_loc_np=query_start_loc_np, + seq_lens=seq_lens, + seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, + dcp_local_seq_lens=dcp_local_seq_lens, + num_computed_tokens_np=num_computed_tokens_np, + prefill_len_np=batch_req_state.prefill_len_np, + num_computed_prefill_tokens_np=batch_req_state.num_computed_prefill_tokens_np, + is_prefilling_np=batch_req_state.is_prefilling_np, + has_prefill=batch_req_state.has_prefill, + max_seq_len_np=max_seq_len_np, + input_ids=self.input_buffers.input_ids[:num_tokens_after_padding], + positions=self.input_buffers.positions[:num_tokens_after_padding], + is_padding=self.input_buffers.is_padding[:num_tokens_after_padding], + logits_indices=logits_indices, + cu_num_logits=cu_num_logits, + cu_num_logits_np=cu_num_logits_np, + has_structured_output_reqs=scheduler_output.has_structured_output_requests, + prompt_lens=prompt_lens, + max_query_len=( + int(num_scheduled_tokens_upper_bound.max()) + if adaptive_verification is not None + else None + ), + ) + return pcp.maybe_partition_pcp_batch(self.pcp_manager, input_batch) + + def prepare_attn( + self, input_batch: InputBatch + ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]: + if self.pcp_manager is not None: + return self.pcp_manager.prepare_attn(input_batch) + + # Block tables: num_kv_cache_groups x [num_reqs_padded, max_num_blocks]. + block_tables = self.block_tables.gather_block_tables( + input_batch.idx_mapping, + num_reqs_padded=input_batch.num_reqs_after_padding, + ) + # Slot mappings: [num_kv_cache_groups, num_tokens_padded]. + # Kernel pads beyond num_tokens with PAD_SLOT_ID. + slot_mappings = self.block_tables.compute_slot_mappings( + input_batch.idx_mapping, + input_batch.query_start_loc, + input_batch.positions, + num_tokens_padded=input_batch.num_tokens_after_padding, + ) + return block_tables, slot_mappings + + def prepare_dummy_attn( + self, input_batch: InputBatch + ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]: + block_tables = self.block_tables.get_dummy_block_tables(input_batch.num_reqs) + slot_mappings = pcp.maybe_get_pcp_dummy_slot_mappings( + self.pcp_manager, self.block_tables, input_batch.num_tokens + ) + return block_tables, slot_mappings + + def sample( + self, + hidden_states: torch.Tensor, + input_batch: InputBatch, + grammar_output: GrammarOutput | None, + ) -> tuple[SamplerOutput, torch.Tensor, torch.Tensor]: + shard_metadata = None + global_input_batch = input_batch + if self.batch_sharder is not None: + # Shard the inputs along the batch dimension to sample in parallel + # across TP ranks. + input_batch, sorted_logits_indices, grammar_output, shard_metadata = ( + self.batch_sharder.shard_sampler_inputs(input_batch, grammar_output) + ) + # The hidden states must be gathered in rank-owner-sorted order + # before computing the partial-vocab logits, so that the all-to-all + # produces full-vocab logits for just the locally-owned requests. + sample_hidden_states = hidden_states[sorted_logits_indices] + local_logits = self.model.compute_logits_local(sample_hidden_states) + logits = all_to_all_logits(local_logits, shard_metadata) + logits = logits[:, : self.vocab_size] + else: + sample_hidden_states = hidden_states[input_batch.logits_indices] + logits = self.model.compute_logits(sample_hidden_states) + + if grammar_output is not None: + # Apply grammar bitmask to the logits in-place. + assert self.structured_outputs_worker is not None + self.structured_outputs_worker.apply_grammar_bitmask( + logits, + input_batch, + grammar_output.structured_output_request_ids, + grammar_output.grammar_bitmask, + ) + + sampler_output: SamplerOutput | None + if input_batch.num_reqs == 0: + # This rank owns no requests this step. It contributes an + # all-padding block to the gather below. + sampler_output = None + elif input_batch.num_draft_tokens == 0 or self.rejection_sampler is None: + assert self.sampler is not None + sampler_output = self.sampler(logits, input_batch) + else: + # Rejection sampling for spec decoding. + assert self.rejection_sampler is not None + assert self.speculator is not None + sampler_output = self.rejection_sampler( + logits, + input_batch, + # Draft logits are needed for probabilistic rejection sampling. + self.speculator.draft_logits, + ) + + if shard_metadata is not None: + # Gather the sharded sampler outputs from the TP ranks into a single + # sampler output. + assert self.sampler is not None + sampler_output = gather_sampler_output( + sampler_output, + shard_metadata, + device=self.device, + global_batch=global_input_batch, + local_batch=input_batch, + gather_num_nans=self.sampler.compute_nans, + logprobs_dims=self.sampler.get_logprobs_dims( + global_input_batch.idx_mapping_np, + # Rejection sampler does not return logprob token ids. + include_token_ids=( + global_input_batch.num_draft_tokens == 0 + or self.rejection_sampler is None + ), + ), + ) + + assert sampler_output is not None + return sampler_output, sampler_output.num_sampled, sampler_output.num_rejected + + def postprocess_sampled( + self, + idx_mapping: torch.Tensor, # May include -1 for masked entries + sampled_tokens: torch.Tensor, + num_sampled: torch.Tensor, + num_rejected: torch.Tensor, + query_start_loc: torch.Tensor | None = None, + ) -> None: + # Update the number of computed tokens. + if self.is_last_pp_rank: + assert self.sampler is not None + output_bin_counts = self.sampler.penalties_state.output_bin_counts + else: + output_bin_counts = None + post_update( + idx_mapping, + self.req_states.num_computed_tokens.gpu, + self.req_states.last_sampled_tokens, + output_bin_counts, + sampled_tokens, + num_sampled, + num_rejected, + query_start_loc, + self.req_states.all_token_ids.gpu, + self.req_states.total_len.gpu, + ) + + self.model_state.postprocess_state( + idx_mapping, num_sampled, self.req_states.num_computed_tokens.gpu + ) + + def _merge_ec_connector_no_forward( + self, scheduler_output: SchedulerOutput, output: ModelRunnerOutput + ) -> ModelRunnerOutput: + """Let the EC connector send/recv on a step with no work to run.""" + return ModelRunnerOutput.with_ec_conn_output( + output, + self.ec_connector.no_forward(scheduler_output).ec_connector_output, + ) + + @torch.inference_mode() + def execute_model( + self, + scheduler_output: SchedulerOutput, + intermediate_tensors: IntermediateTensors | None = None, + dummy_run: bool = False, + skip_attn_for_dummy_run: bool = False, + is_profile: bool = False, + context_len: int = 0, + ) -> ModelRunnerOutput | IntermediateTensors | None: + if not dummy_run: + # Update the request states. + self.update_pp_decode_requests() + self.finish_requests(scheduler_output) + self.free_states(scheduler_output) + self.add_requests(scheduler_output) + self.update_requests(scheduler_output) + self.block_tables.apply_staged_writes() + if scheduler_output.total_num_scheduled_tokens == 0: + # No need to run the model. + empty_output = self.kv_connector.no_forward(scheduler_output) + return self._merge_ec_connector_no_forward( + scheduler_output, empty_output + ) + + # Get batch descriptor and sync across DP ranks. + num_reqs = len(scheduler_output.num_scheduled_tokens) + num_toks = scheduler_output.total_num_scheduled_tokens + max_query_len = max(scheduler_output.num_scheduled_tokens.values()) + batch_req_state, uniform_tok_count = self.gather_batch_req_state( + scheduler_output, dummy_run + ) + if batch_req_state is not None: + num_toks = batch_req_state.num_tokens + + num_active_loras = 0 + if self.lora_config: + req_ids = list(scheduler_output.num_scheduled_tokens.keys()) + num_active_loras = get_num_active_loras_for_dispatch( + self.lora_config, self.lora_state, req_ids, dummy_run + ) + + skip_compiled = False + if self.is_encoder_decoder and scheduler_output.scheduled_encoder_inputs: + # Encoder-decoder models such as Whisper should run eager/non-compiled + # when encoder inputs are scheduled, because this step updates + # cross-attention cache with dynamic encoder outputs. + skip_compiled = True + + batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp( + self.cudagraph_manager, + num_reqs, + num_toks, + uniform_tok_count, + self.dp_size, + self.dp_rank, + max_query_len=max_query_len, + need_eager=is_profile or skip_compiled, + num_active_loras=num_active_loras, + ) + + if batch_desc.num_tokens == 0: + # All DP ranks have zero tokens to run. + empty_output = self.kv_connector.no_forward(scheduler_output) + return self._merge_ec_connector_no_forward(scheduler_output, empty_output) + + if not dummy_run: + # Common case. + # Prepare all the inputs and copy to the input buffers. + assert batch_req_state is not None + input_batch = self.prepare_inputs( + scheduler_output, batch_req_state, batch_desc + ) + block_tables, slot_mappings = self.prepare_attn(input_batch) + # Mamba "align" pre-copy: migrate recurrent state across block + # boundaries before the forward. Runs only on real batches, and + # before model_state.prepare_attn gathers num_accepted_tokens so the + # boundary reset is visible to the attention metadata. + self.model_state.preprocess_state( + input_batch, + block_tables, + self.kv_cache_config, + self.req_states.num_computed_tokens.gpu, + ) + + if self.lora_config: + # Activate LoRA adapters. + lora_inputs = self.lora_state.make_lora_inputs( + input_batch.req_ids, + input_batch.idx_mapping_np, + input_batch.num_scheduled_tokens, + ) + self._set_active_loras(*lora_inputs) + else: + # No actual tokens to run. A dummy run for DP or memory profiling. + dummy_num_reqs = batch_desc.num_reqs or num_reqs + input_batch = InputBatch.make_dummy( + dummy_num_reqs, + batch_desc.num_tokens, + self.input_buffers, + max_query_len=batch_desc.max_query_len, + ) + if not skip_attn_for_dummy_run: + block_tables, slot_mappings = self.prepare_dummy_attn(input_batch) + if context_len: + set_dummy_context( + input_batch, + self.block_tables, + context_len, + self.kv_cache_config.num_blocks, + self.max_model_len, + ) + else: + assert batch_desc.cg_mode != CUDAGraphMode.FULL, ( + "Attention metadata must be prepared for dummy runs when using " + "FULL cudagraph mode." + ) + block_tables = None + slot_mappings = None + + attn_metadata = None + slot_mappings_by_layer = None + if not (dummy_run and skip_attn_for_dummy_run): + assert slot_mappings is not None + slot_mappings_by_layer = build_slot_mappings_by_layer( + slot_mappings, self.kv_cache_config + ) + assert block_tables is not None + attn_groups = self.attn_groups + if dummy_run and is_profile: + # Mamba layers take a cheap warmup path with no metadata; + # attention metadata is still built so those kernels tune. + attn_groups = [ + [g for g in groups if not isinstance(g.kv_cache_spec, MambaSpec)] + for groups in attn_groups + ] + checkpoint_attn_kwargs = {} + if not dummy_run: + checkpoint_plans = getattr( + scheduler_output, "recurrent_prefill_checkpoint_plans", None + ) + if checkpoint_plans: + checkpoint_attn_kwargs["recurrent_prefill_checkpoint_plans"] = checkpoint_plans + attn_metadata = self.model_state.prepare_attn( + input_batch, + batch_desc.cg_mode, + block_tables, + slot_mappings, + attn_groups, + self.kv_cache_config, + # FULL replay reads capture-time metadata buffers. Re-stage them + # from the zeroed dummy block tables instead of retaining state + # indices from the previous real batch. + for_capture=dummy_run and batch_desc.cg_mode == CUDAGraphMode.FULL, + **checkpoint_attn_kwargs, + ) + + input_ids = input_batch.input_ids + inputs_embeds = None + ec_connector_output = None + if self.uses_inputs_embeds and self.is_first_pp_rank: + # Prepare inputs_embeds (MM encoder outputs and/or prompt_embeds + # overlay). Only first PP rank prepares them. + if dummy_run: + # Obtain embeddings of correct shape for compiled model. + inputs_embeds = self.model_state.dummy_inputs_embeds( + input_batch.num_tokens_after_padding + ) + else: + scheduled_encoder_inputs = scheduler_output.scheduled_encoder_inputs + if self.supports_mm_inputs and self.lora_config is not None: + set_active_mm_loras( + model=self.model, + lora_manager=self.lora_manager, + encoder_cache=self.encoder_cache, + req_id_to_index=self.req_states.req_id_to_index, + lora_state=self.lora_state, + scheduled_encoder_inputs=scheduled_encoder_inputs, + ) + with self.ec_connector.maybe_get_output( + scheduler_output + ) as ec_connector_output: + inputs_embeds = self.model_state.prepare_inputs_embeds( + scheduled_encoder_inputs, input_batch, self.req_states + ) + model_inputs = { + "input_ids": input_ids, + "positions": input_batch.positions, + "inputs_embeds": inputs_embeds, + "intermediate_tensors": None, + # NOTE: Values returned by `prepare_inputs` will override the default + # values above. + **self.model_state.prepare_inputs(input_batch, self.req_states), + } + normalize_model_token_inputs(self.model, model_inputs) + if not self.is_first_pp_rank: + # Update for non-first PP ranks. + model_inputs["input_ids"] = None + model_inputs["inputs_embeds"] = None + + # Prepare the intermediate tensors. + assert intermediate_tensors is not None + assert self.intermediate_tensors is not None + n = input_batch.num_tokens_after_padding + new_tensors = { + k: v[:n] + if dummy_run + else v[:n].copy_(intermediate_tensors.tensors[k][:n]) + for k, v in self.intermediate_tensors.tensors.items() + } + model_inputs["intermediate_tensors"] = IntermediateTensors(new_tensors) + del intermediate_tensors + + # Update the EPLB meta. + self.eplb.prepare_forward(self.model_config, input_batch.num_tokens) + + self.step_timing.record_batch( + input_batch, batch_desc.cg_mode == CUDAGraphMode.FULL + ) + self.step_timing.forward_start() + + # Run model. + if batch_desc.cg_mode == CUDAGraphMode.FULL: + # Use explicit cudagraph replay for FULL mode. + # NOTE(woosuk): Here, we don't need to pass the input tensors, + # because they are already copied to the CUDA graph input buffers. + assert self.cudagraph_manager is not None + self.kv_connector.pre_forward(scheduler_output) + model_output = self.cudagraph_manager.run_fullgraph(batch_desc) + else: + # For piecewise and eager mode, just call model(). + batch_descriptor = BatchDescriptor( + num_tokens=input_batch.num_tokens_after_padding, + has_lora=self.lora_config is not None, + num_active_loras=batch_desc.num_active_loras, + ) + + with set_forward_context( + attn_metadata, + self.vllm_config, + num_tokens=input_batch.num_tokens_after_padding, + cudagraph_runtime_mode=batch_desc.cg_mode, + num_tokens_across_dp=num_tokens_across_dp, + batch_descriptor=batch_descriptor, + slot_mapping=slot_mappings_by_layer, + skip_compiled=skip_compiled, + is_padding=input_batch.is_padding, + ): + self.kv_connector.pre_forward(scheduler_output) + if batch_desc.cg_mode == CUDAGraphMode.PIECEWISE: + # Run the PIECEWISE graph (compiled PW cudagraph or breakable + # cudagraph, chosen inside run_pw_graph). cg_mode is only + # PIECEWISE after the cudagraph manager exists. + assert self.cudagraph_manager is not None + model_output = self.cudagraph_manager.run_pw_graph( + self.model, model_inputs + ) + else: + # Eager (NONE): call the raw model directly. + model_output = self.model(**model_inputs) + + if self.is_last_pp_rank: + if self.use_aux_hidden_state_outputs: + assert isinstance(model_output, tuple) + hidden_states, aux_hidden_states = model_output + else: + assert isinstance(model_output, torch.Tensor) + hidden_states = model_output + aux_hidden_states = None + output_intermediate_tensors = None + else: + assert isinstance(model_output, IntermediateTensors) + hidden_states = None + aux_hidden_states = None + output_intermediate_tensors = model_output + + routed_experts = None + if not dummy_run and (capturer := self.routed_experts_capturer) is not None: + assert slot_mappings is not None + routed_experts = capturer.get_routed_experts(slot_mappings, num_toks) + + finished_req_ids = scheduler_output.finished_req_ids + self.execute_model_state = ExecuteModelState( + input_batch=input_batch, + attn_metadata=attn_metadata, + slot_mappings_by_layer=slot_mappings_by_layer, + hidden_states=hidden_states, + aux_hidden_states=aux_hidden_states, + finished_req_ids=finished_req_ids, + ec_connector_output=ec_connector_output, + routed_experts=routed_experts, + num_spec_tokens_to_schedule=( + scheduler_output.resolve_num_spec_tokens_to_schedule( + self.num_speculative_steps + ) + ), + ) + + if not self.is_last_pp_rank: + # Non-last PP rank: return IntermediateTensors for sending. + return output_intermediate_tensors + return None + + @torch.inference_mode() + @step_eplb_after() + def sample_tokens( + self, grammar_output: GrammarOutput | None + ) -> AsyncOutput | ModelRunnerOutput | None: + if self.execute_model_state is None: + # The prior execute_model call must have failed. + return None + + input_batch = self.execute_model_state.input_batch + attn_metadata = self.execute_model_state.attn_metadata + slot_mappings_by_layer = self.execute_model_state.slot_mappings_by_layer + hidden_states = self.execute_model_state.hidden_states + aux_hidden_states = self.execute_model_state.aux_hidden_states + finished_req_ids = self.execute_model_state.finished_req_ids + ec_connector_output = self.execute_model_state.ec_connector_output + routed_experts = self.execute_model_state.routed_experts + num_spec_tokens_to_schedule = ( + self.execute_model_state.num_spec_tokens_to_schedule + ) + self.execute_model_state = None + + if not self.is_last_pp_rank: + # Non-last PP rank: hidden_states is None because this rank produced + # IntermediateTensors instead of final hidden states. Receive the + # sampled tokens broadcast from the last rank and update local state. + assert self.pp_handler is not None + all_decode_next = self.pp_handler.receive(input_batch) + # Optimistically update num_computed_tokens for entire batch here. + # Will be adjusted for rejections if necessary in update_requests. + self.postprocess_num_computed_tokens(input_batch) + if not all_decode_next: + # Might contain non-final prefill chunks, which will be scheduled + # in the immediate next step (rather than in pp_size steps). + self.model_state.postprocess_state(input_batch.idx_mapping, 0) + + # Post-step KV connector related operations. + kv_connector_output = self.kv_connector.post_forward(finished_req_ids) + # The first PP rank holds the encoder cache, so pass its EC output on. + output = ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) + return ModelRunnerOutput.with_ec_conn_output(output, ec_connector_output) + + # Last rank: sample tokens + hidden_states, input_batch = pcp.maybe_restore_pcp_for_sampling( + self.pcp_manager, hidden_states, input_batch + ) + + sampler_output, num_sampled, num_rejected = self.sample( + hidden_states, input_batch, grammar_output + ) + + if self.pp_handler is not None: + # Broadcast to non-last PP ranks (handles spec decode multi-token). + self.pp_handler.broadcast( + sampler_output.sampled_token_ids, + num_sampled, + num_rejected, + input_batch, + ) + + assert self.prompt_logprobs_worker is not None + prompt_logprobs_dict = self.prompt_logprobs_worker.compute_prompt_logprobs( + self.model.compute_logits, + hidden_states, + input_batch, + self.req_states.all_token_ids.gpu, + self.req_states.num_computed_tokens.gpu, + self.req_states.prompt_len.np, + ) + + # Prepare the model runner output. + model_runner_output = ModelRunnerOutput( + req_ids=input_batch.req_ids, + # NOTE(woosuk): req_id_to_index is unused in this model runner. + # Only for compatibility with the existing model runner and scheduler. + req_id_to_index={req_id: i for i, req_id in enumerate(input_batch.req_ids)}, + sampled_token_ids=None, # type: ignore + prompt_logprobs_dict=prompt_logprobs_dict, # type: ignore[arg-type] + ) + # Start async output copy here so that it can overlap with speculator proposal. + async_output = AsyncOutput( + model_runner_output=model_runner_output, + sampler_output=sampler_output, + num_sampled_tokens=num_sampled, + main_stream=self.main_stream, + copy_stream=self.output_copy_stream, + check_ep_fault=self.check_ep_fault, + routed_experts=routed_experts, + ) + + mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None + if self.speculator is not None and self.speculator.supports_mm_inputs: + # Get cached multimodal embeddings for draft forward. + # NOTE: This is done here because postprocess updates + # num_computed_prefill_tokens. + # The EAGLE/MTP drafter reads one position ahead of the target. + # TODO(TheEpicDolphin): Gather MM embeddings for all speculative + # steps during multi-module MTP. + mm_inputs = self.model_state.gather_mm_embeddings( + input_batch, draft_lookahead=1 + ) + + # Postprocess results and update request states. + # NOTE: This is intentionally done after creating the AsyncOutput, + # ensuring that `copy_event` is recorded before calling postprocess. + # This sequencing may slightly reduce latency as async D2H copy does not + # need to wait for the postprocess to finish. + self.postprocess_sampled( + input_batch.idx_mapping, + sampler_output.sampled_token_ids, + num_sampled, + num_rejected, + input_batch.query_start_loc, + ) + + draft_tokens_for_next_step: torch.Tensor | None = None + if self.speculator is not None and num_spec_tokens_to_schedule > 0: + assert self.sampler is not None + # Let the target override the hidden state fed to the drafter + # (e.g. DeepSeek V4 MTP needs the pre-hc_head residual). The + # target returns a persistent buffer sized at max_num_batched_tokens; + # slice to the active token count that propose() expects. + spec_hidden_states = hidden_states + if hasattr(self.model, "get_mtp_target_hidden_states"): + pre_hc_hidden_states = self.model.get_mtp_target_hidden_states() + spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]] # type: ignore[union-attr] + with use_workspace_lane(self._draft_workspace_lane): + draft_tokens = self.speculator.propose( + input_batch, + attn_metadata, + slot_mappings_by_layer, + spec_hidden_states, + aux_hidden_states, + num_sampled, + num_rejected, + self.req_states.last_sampled_tokens, + self.req_states.next_prefill_tokens, + self.sampler.sampling_states.temperature.gpu, + self.sampler.sampling_states.seeds.gpu, + num_speculative_tokens=num_spec_tokens_to_schedule, + mm_inputs=mm_inputs, + ) + draft_tokens = limit_draft_tokens( + draft_tokens, + num_spec_tokens_to_schedule, + self.num_speculative_steps, + ) + num_draft_tokens = draft_tokens.shape[1] + if num_draft_tokens > 0: + self.req_states.draft_tokens[ + input_batch.idx_mapping, :num_draft_tokens + ] = draft_tokens + draft_tokens_for_next_step = self.req_states.draft_tokens[ + input_batch.idx_mapping, :num_draft_tokens + ] + else: + draft_tokens_for_next_step = draft_tokens + if self.adaptive_verification is not None: + self.adaptive_verification.record_confidences( + self.speculator.draft_token_confidence_probs, input_batch + ) + elif self.speculator is not None: + draft_tokens_for_next_step = self.req_states.draft_tokens[ + input_batch.idx_mapping, :0 + ] + + if self.num_speculative_steps > 0: + # Spec-decode and diffusion LLMs both use draft tokens but the latter does + # not have a speculator (i.e. self.speculator is None) + self.draft_tokens_handler.set_draft_tokens( + input_batch, + ( + draft_tokens_for_next_step + if draft_tokens_for_next_step is not None + else self.req_states.draft_tokens[input_batch.idx_mapping] + ), + ) + + # Post-step KV connector related operations. + kv_connector_output = self.kv_connector.post_forward(finished_req_ids) + model_runner_output.kv_connector_output = kv_connector_output + model_runner_output.ec_connector_output = ec_connector_output + + return async_output + + def take_draft_token_ids(self) -> DraftTokenIds | None: + return self.draft_tokens_handler.get_draft_tokens() + + @torch.inference_mode() + @step_eplb_after() + def pool(self) -> AsyncPoolingOutput | ModelRunnerOutput | None: + if self.execute_model_state is None: + # The prior execute_model call must have failed. + return None + + input_batch = self.execute_model_state.input_batch + hidden_states = self.execute_model_state.hidden_states + finished_req_ids = self.execute_model_state.finished_req_ids + ec_connector_output = self.execute_model_state.ec_connector_output + self.execute_model_state = None + + # Post-step KV connector related operations. + kv_connector_output = self.kv_connector.post_forward(finished_req_ids) + + if not self.is_last_pp_rank: + self.postprocess_num_computed_tokens(input_batch) + output = ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) + return ModelRunnerOutput.with_ec_conn_output(output, ec_connector_output) + + assert self.pooling_runner is not None + pooler_output, finished_mask = self.pooling_runner.pool( + hidden_states, input_batch, self.req_states + ) + + # Build the model runner output. + model_runner_output = ModelRunnerOutput( + req_ids=input_batch.req_ids, + req_id_to_index={req_id: i for i, req_id in enumerate(input_batch.req_ids)}, + kv_connector_output=kv_connector_output, + ec_connector_output=ec_connector_output, + ) + async_output = AsyncPoolingOutput( + model_runner_output=model_runner_output, + pooler_output=pooler_output, + finished_mask=finished_mask, + main_stream=self.main_stream, + copy_stream=self.output_copy_stream, + ) + + self.postprocess_num_computed_tokens(input_batch) + return async_output + + def postprocess_num_computed_tokens(self, input_batch: InputBatch) -> None: + # Update the number of computed tokens. + post_update_num_computed_tokens( + input_batch.idx_mapping, + self.req_states.num_computed_tokens.gpu, + input_batch.query_start_loc, + ) + + def shutdown(self) -> None: + """Release GPU tensors (model weights, KV caches, workspace) so that + memory is reclaimable when running in the same process.""" + torch.accelerator.synchronize() + self.cudagraph_manager = None + if hasattr(self, "kv_caches"): + self.kv_caches.clear() + if hasattr(self, "attn_groups"): + self.attn_groups.clear() + if hasattr(self, "kv_cache_config"): + del self.kv_cache_config + if hasattr(self, "model_state") and self.model_state.supports_mm_inputs: + self.model_state.encoder_runner.clear() + free_before_shutdown(self.vllm_config) + if hasattr(self, "model_state"): + del self.model_state + if getattr(self, "speculator", None) is not None: + self.speculator = None + if hasattr(self, "model"): + del self.model + + gc.collect() + torch.accelerator.empty_cache() + logger.debug("Cleaned up model weights, KV caches, and workspace") + + ########### EPLB methods start ########### + @property + def eplb_state(self): + return self.eplb.state + + @eplb_state.setter + def eplb_state(self, state) -> None: + self.eplb.state = state + + @property + def eep_eplb_suppressed(self) -> bool: + return self.eplb.suppressed + + @eep_eplb_suppressed.setter + def eep_eplb_suppressed(self, suppressed: bool) -> None: + self.eplb.suppressed = suppressed + + def setup_eplb_from_mapping( + self, + expanded_physical_to_logical: torch.Tensor, + ) -> None: + self.eplb.setup_from_mapping( + self.model, + self.model_config, + expanded_physical_to_logical, + ) + + ########### EPLB methods end ########### + + # Out-of-tree hardware runners can select a PCP manager class. + @property + def pcp_manager_cls(self) -> type[pcp.PCPManager]: + return pcp.PCPManager + + +class ExecuteModelState(NamedTuple): + input_batch: InputBatch + attn_metadata: dict[str, Any] | None + slot_mappings_by_layer: dict[str, torch.Tensor] | None + hidden_states: torch.Tensor | None + aux_hidden_states: list[torch.Tensor] | None + finished_req_ids: set[str] + ec_connector_output: ECConnectorOutput | None + routed_experts: RoutedExpertsTensors | None + num_spec_tokens_to_schedule: int + + +class BatchReqState(NamedTuple): + """CPU request state for a scheduled batch, in batch (sorted) order.""" + + req_ids: list[str] + num_scheduled_tokens: np.ndarray # [num_reqs] + # May be less than scheduler_output.total_num_scheduled_tokens: + # adaptive verification trims the draft budget before running. + num_tokens: int + idx_mapping_np: np.ndarray # [num_reqs] + prefill_len_np: np.ndarray # [num_reqs] + num_computed_prefill_tokens_np: np.ndarray # [num_reqs] + is_prefilling_np: np.ndarray # [num_reqs] + has_prefill: bool + + +def sort_batch_req_ids( + num_tokens_per_req: dict[str, int], + draft_tokens: dict[str, list[int]], + decode_query_len: int, +) -> list[str]: + # Order verification/decode -> short_extend -> prefill; + # split_decodes_and_prefills relies on decode-like requests leading. + key = lambda r: ( + not draft_tokens.get(r), + (num := num_tokens_per_req[r]) != decode_query_len, + num, + ) + return sorted(num_tokens_per_req, key=key) diff --git a/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/6888b3c5e16b4ef56ddc6d7bb02aa9736894a758574abeebe7e4ff6e9a4ef295/gpu_model_runner.py b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/6888b3c5e16b4ef56ddc6d7bb02aa9736894a758574abeebe7e4ff6e9a4ef295/gpu_model_runner.py new file mode 100644 index 00000000..ac526325 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/6888b3c5e16b4ef56ddc6d7bb02aa9736894a758574abeebe7e4ff6e9a4ef295/gpu_model_runner.py @@ -0,0 +1,7790 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import functools +import gc +import itertools +import threading +import time +from collections import defaultdict +from collections.abc import Callable, Iterable, Iterator, Sequence +from contextlib import AbstractContextManager, contextmanager, nullcontext +from copy import copy, deepcopy +from dataclasses import replace +from functools import reduce +from typing import TYPE_CHECKING, Any, NamedTuple, TypeAlias, cast + +import numpy as np +import torch +import torch.distributed +import torch.nn as nn +from tqdm import tqdm + +import vllm.envs as envs +from vllm.compilation.breakable_cudagraph import ( + BreakableCUDAGraphWrapper, + is_breakable_cudagraph_enabled, +) +from vllm.compilation.counter import compilation_counter +from vllm.compilation.cuda_graph import CUDAGraphStat, CUDAGraphWrapper +from vllm.compilation.monitor import set_cudagraph_capturing_enabled +from vllm.config import ( + CompilationMode, + CUDAGraphMode, + VllmConfig, + get_layers_from_vllm_config, + set_current_vllm_config, + update_config, +) +from vllm.config.cache import CacheConfig +from vllm.config.ec_manager_config import EncoderCacheManagerMetadata +from vllm.config.model import PROCESSED_LOGPROBS_MODES +from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer +from vllm.distributed.eplb.eplb_state import EplbState +from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group +from vllm.distributed.kv_transfer.kv_connector.utils import ( + copy_kv_blocks, +) +from vllm.distributed.parallel_state import ( + GraphCaptureContext, + get_dcp_group, + get_pp_group, + get_tp_group, + graph_capture, + is_global_first_rank, +) +from vllm.forward_context import ( + BatchDescriptor, + set_forward_context, +) +from vllm.logger import init_logger +from vllm.lora.layers import BaseLayerWithLoRA, LoRAMapping, LoRAMappingType +from vllm.model_executor.layers.attention import Attention, MLAAttention +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase +from vllm.model_executor.layers.fused_moe.all2all_utils import get_ep_all2all_manager +from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( + RoutedExpertsCapturer, + bind_routed_experts_capturer, +) +from vllm.model_executor.layers.mamba.ops.ssu_dispatch import ( + initialize_mamba_ssu_backend, +) +from vllm.model_executor.layers.rotary_embedding import ( + MRotaryEmbedding, + XDRotaryEmbedding, +) +from vllm.model_executor.model_loader import get_model_loader +from vllm.model_executor.model_loader.reload import ( + finalize_layerwise_reload, + initialize_layerwise_reload, +) +from vllm.model_executor.models.interfaces import ( + MixtureOfExperts, + MultiModalEmbeddings, + SupportsMRoPE, + SupportsMultiModal, + SupportsXDRoPE, + get_mixture_of_experts_model, + supports_eagle3, + supports_mrope, + supports_multimodal_pruning, + supports_realtime, + supports_transcription, + supports_xdrope, +) +from vllm.model_executor.models.interfaces_base import ( + VllmModelForPooling, + is_pooling_model, + is_text_generation_model, +) +from vllm.model_executor.offloader import ( + create_offloader, + get_offloader, + set_offloader, +) +from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.encoder_budget import MultiModalBudget +from vllm.multimodal.inputs import ( + BatchedTensorInputs, + MultiModalKwargsItem, + PlaceholderRange, +) +from vllm.multimodal.utils import ( + copy_mm_embedding_modality, + get_mm_features_in_window, + group_and_batch_mm_kwargs, + set_mm_embedding_modality, +) +from vllm.platforms import current_platform +from vllm.pooling_params import PoolingParams +from vllm.sampling_params import SamplingType +from vllm.sequence import IntermediateTensors +from vllm.tasks import GenerationTask, PoolingTask, SupportedTask +from vllm.tracing import instrument +from vllm.utils import length_from_prompt_token_ids_or_embeds +from vllm.utils.gpu_sync_debug import gpu_sync_allowed +from vllm.utils.math_utils import cdiv, round_up +from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib +from vllm.utils.nvtx_pytorch_hooks import PytHooks +from vllm.utils.platform_utils import num_compute_units +from vllm.utils.torch_utils import ( + PIN_MEMORY, + async_tensor_h2d, + current_stream, + is_quantized_kv_cache, + kv_cache_dtype_str_to_dtype, +) +from vllm.v1.attention.backend import ( + AttentionBackend, + AttentionCGSupport, + AttentionMetadata, + AttentionMetadataBuilder, + AttentionType, + CommonAttentionMetadata, +) +from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder +from vllm.v1.attention.backends.linear_attn import ( + BailingLinearAttentionMetadataBuilder, +) +from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadataBuilder +from vllm.v1.attention.backends.short_conv_attn import ShortConvAttentionMetadataBuilder +from vllm.v1.attention.backends.utils import ( + NULL_BLOCK_ID, + create_fast_prefill_custom_backend, + get_dcp_local_seq_lens, + reorder_batch_to_split_decodes_and_prefills, +) +from vllm.v1.core.sched.output import NewRequestData +from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher +from vllm.v1.kv_cache_interface import ( + AttentionSpec, + ChunkedLocalAttentionSpec, + CrossAttentionSpec, + EncoderOnlyAttentionSpec, + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, + KVCacheSpec, + KVCacheSpecKind, + SlidingWindowSpec, + UniformTypeKVCacheSpecs, + get_kv_cache_spec_kind, +) +from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry +from vllm.v1.outputs import ( + EMPTY_MODEL_RUNNER_OUTPUT, + AsyncModelRunnerOutput, + DraftTokenIds, + ECConnectorOutput, + KVConnectorOutput, + LogprobsLists, + LogprobsTensors, + ModelRunnerOutput, + PoolerOutput, + RoutedExpertsLists, + RoutedExpertsTensors, + SamplerOutput, + make_empty_encoder_model_runner_output, +) +from vllm.v1.pool.late_interaction_runner import LateInteractionRunner +from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates +from vllm.v1.sample.logits_processor import LogitsProcessors, build_logitsprocs +from vllm.v1.sample.logits_processor.interface import LogitsProcessor +from vllm.v1.sample.metadata import SamplingMetadata +from vllm.v1.sample.rejection_sampler import RejectionSampler +from vllm.v1.sample.sampler import Sampler +from vllm.v1.spec_decode.custom_class_proposer import create_custom_proposer +from vllm.v1.spec_decode.dflash import DFlashProposer +from vllm.v1.spec_decode.draft_model import DraftModelProposer +from vllm.v1.spec_decode.eagle import EagleProposer +from vllm.v1.spec_decode.extract_hidden_states import ExtractHiddenStatesProposer +from vllm.v1.spec_decode.gemma4 import Gemma4Proposer +from vllm.v1.spec_decode.medusa import MedusaProposer +from vllm.v1.spec_decode.metadata import SpecDecodeMetadata +from vllm.v1.spec_decode.ngram_proposer_gpu import ( + NgramProposerGPU, + copy_num_valid_draft_tokens, + update_ngram_gpu_tensors_incremental, + update_scheduler_for_invalid_drafts, +) +from vllm.v1.spec_decode.step3p5 import Step3p5MTPProposer +from vllm.v1.spec_decode.suffix_decoding import SuffixDecodingProposer +from vllm.v1.spec_decode.utils import update_num_computed_tokens_for_batch_change +from vllm.v1.structured_output.utils import apply_grammar_bitmask +from vllm.v1.utils import CpuGpuBuffer, record_function_or_nullcontext +from vllm.v1.worker import mamba_utils +from vllm.v1.worker.block_table import SlotMappingMode +from vllm.v1.worker.cp_utils import ( + check_attention_cp_compatibility, + get_dcp_dummy_context_len, + prepare_dcp_dummy_context_metadata, +) +from vllm.v1.worker.dp_utils import coordinate_batch_across_dp +from vllm.v1.worker.ec_connector_model_runner_mixin import ECConnectorModelRunnerMixin +from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch +from vllm.v1.worker.gpu_ubatch_wrapper import UBatchWrapper +from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin +from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin +from vllm.v1.worker.ubatch_utils import ( + UBatchSlices, + check_ubatch_thresholds, + maybe_create_ubatch_slices, + split_attn_metadata, +) +from vllm.v1.worker.utils import ( + EncoderTimingStats, + is_residual_scattered_for_sp, + raise_if_nan_logits, +) +from vllm.v1.worker.workspace import lock_workspace + +from .utils import ( + AttentionGroup, + KVBlockZeroer, + add_kv_sharing_layers_to_kv_cache_groups, + allocate_kv_cache, + bind_kv_cache, + copy_kv_cache_blocks_inplace, + prepare_kernel_block_sizes, + sanity_check_mm_encoder_outputs, + unbind_kv_cache, +) + +if TYPE_CHECKING: + from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput + from vllm.v1.spec_decode.ngram_proposer import NgramProposer + from vllm.v1.worker.encoder_cudagraph import EncoderCudaGraphManager + +logger = init_logger(__name__) + + +def _get_parameter_for_reload(model: nn.Module, name: str) -> nn.Parameter: + """Resolve checkpoint names without changing the model's module tree.""" + module_name, _, parameter_name = name.rpartition(".") + module = model.get_submodule(module_name) + if isinstance(module, BaseLayerWithLoRA): + module = module.base_layer + return module.get_parameter(parameter_name) + + +AttnMetadataDict: TypeAlias = dict[str, AttentionMetadata] +# list when ubatching is enabled +PerLayerAttnMetadata: TypeAlias = list[AttnMetadataDict] | AttnMetadataDict + + +def count_nans_per_row(logits: torch.Tensor) -> torch.Tensor: + """Per-row NaN counts, left on device.""" + return logits.isnan().sum(dim=-1, dtype=torch.int32) + + +def nans_to_dict(counts: list[int], req_id_to_index: dict[str, int]) -> dict[str, int]: + """Map per-row NaN counts onto request ids. + + Rows and requests need not correspond one-to-one, so requests without a + row are reported as 0. + """ + return { + req_id: int(counts[i]) if i < len(counts) else 0 + for req_id, i in req_id_to_index.items() + } + + +# Wrapper for ModelRunnerOutput to support overlapped execution. +class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput): + def __init__( + self, + model_runner_output: ModelRunnerOutput, + sampled_token_ids: torch.Tensor, + logprobs_tensors: LogprobsTensors | None, + invalid_req_indices: list[int], + async_output_copy_stream: torch.cuda.Stream, + vocab_size: int, + routed_experts: RoutedExpertsTensors | None = None, + check_ep_fault: bool = False, + num_nans: torch.Tensor | None = None, + ): + self._model_runner_output = model_runner_output + self._invalid_req_indices = invalid_req_indices + + # Event on the copy stream so we can synchronize the non-blocking copy. + # Blocking (sleep) event to avoid busy-polling the CUDA driver lock. + self.async_copy_ready_event = torch.cuda.Event(blocking=True) + + # Keep a reference to the device tensor to avoid it being + # deallocated until we finish copying it to the host. + self._sampled_token_ids = sampled_token_ids + self.vocab_size = vocab_size + self._logprobs_tensors = logprobs_tensors + self._routed_experts = routed_experts + self._num_nans = num_nans + self._has_fault: torch.Tensor | None = None + + # Initiate the copy on a separate stream, but do not synchronize it. + default_stream = torch.cuda.current_stream() + with torch.cuda.stream(async_output_copy_stream): + async_output_copy_stream.wait_stream(default_stream) + self.sampled_token_ids_cpu = self._sampled_token_ids.to( + "cpu", non_blocking=True + ) + self._logprobs_tensors_cpu = ( + self._logprobs_tensors.to_cpu_nonblocking() + if self._logprobs_tensors + else None + ) + self._routed_experts_cpu = ( + self._routed_experts.to_cpu_nonblocking() + if self._routed_experts is not None + else None + ) + self._num_nans_cpu = ( + self._num_nans.to("cpu", non_blocking=True) + if self._num_nans is not None + else None + ) + if check_ep_fault: + has_fault = get_ep_all2all_manager().query_fault() + self._has_fault = has_fault.to("cpu", non_blocking=True) + self.async_copy_ready_event.record() + + def get_output(self) -> ModelRunnerOutput: + """Copy the device tensors to the host and return a ModelRunnerOutput. + + This function blocks until the copy is finished. + """ + max_gen_len = self.sampled_token_ids_cpu.shape[-1] + self.async_copy_ready_event.synchronize() + + # Release the device tensors once the copy has completed. + del self._logprobs_tensors + del self._sampled_token_ids + if max_gen_len == 1: + valid_sampled_token_ids = self.sampled_token_ids_cpu.tolist() + for i in self._invalid_req_indices: + valid_sampled_token_ids[i].clear() + logprobs_lists = None + if self._logprobs_tensors_cpu is not None: + logprobs_lists = self._logprobs_tensors_cpu.tolists() + else: + valid_sampled_token_ids, logprobs_lists = RejectionSampler.parse_output( + self.sampled_token_ids_cpu, + self.vocab_size, + self._invalid_req_indices, + logprobs_tensors=self._logprobs_tensors_cpu, + ) + + output = self._model_runner_output + output.sampled_token_ids = valid_sampled_token_ids + output.logprobs = logprobs_lists + + if self._routed_experts_cpu is not None: + output.routed_experts = self._routed_experts_cpu.tolists() + del self._routed_experts + + if self._num_nans_cpu is not None: + output.num_nans_in_logits = nans_to_dict( + self._num_nans_cpu.tolist(), output.req_id_to_index + ) + if envs.VLLM_RAISE_ON_LOGIT_NANS: + raise_if_nan_logits(output.num_nans_in_logits) + del self._num_nans + + if self._has_fault is not None and self._has_fault.item(): + mask = get_ep_all2all_manager().query_active_mask() + raise RuntimeError( + "Fault detected in EP all2all communication: " + "one or more ranks timed out during dispatch/combine. " + f"Mask: {mask.cpu().tolist()}" + ) + + return output + + +def _copy_pooler_output_to_cpu( + raw_pooler_output: PoolerOutput, finished_mask: list[bool] +) -> list[torch.Tensor | None]: + num_reqs = len(finished_mask) + + if isinstance(raw_pooler_output, torch.Tensor): + if raw_pooler_output.shape[0] != num_reqs: + raise ValueError( + "Pooler output batch size does not match finished mask size: " + f"{raw_pooler_output.shape[0]} != {num_reqs}." + ) + + num_finished = sum(finished_mask) + if num_finished == 0: + return [None] * num_reqs + if num_finished == num_reqs: + return list(raw_pooler_output.to("cpu", non_blocking=True)) + + # partial finished + finished_indices = [i for i, include in enumerate(finished_mask) if include] + index_tensor = torch.tensor( + finished_indices, device=raw_pooler_output.device, dtype=torch.long + ) + finished_outputs = raw_pooler_output.index_select(0, index_tensor).to( + "cpu", non_blocking=True + ) + partial_pooler_output: list[torch.Tensor | None] = [None] * num_reqs + for i, out in zip(finished_indices, finished_outputs): + partial_pooler_output[i] = out + return partial_pooler_output + + assert isinstance(raw_pooler_output, list) + if len(raw_pooler_output) != num_reqs: + raise ValueError( + "Pooler output batch size does not match finished mask size: " + f"{len(raw_pooler_output)} != {num_reqs}." + ) + + pooler_output: list[torch.Tensor | None] = [None] * num_reqs + for i, (out, include) in enumerate(zip(raw_pooler_output, finished_mask)): + if include and out is not None: + pooler_output[i] = out.to("cpu", non_blocking=True) + return pooler_output + + +class AsyncGPUPoolingModelRunnerOutput(AsyncModelRunnerOutput): + def __init__( + self, + model_runner_output: ModelRunnerOutput, + raw_pooler_output: PoolerOutput, + finished_mask: list[bool], + async_output_copy_stream: torch.cuda.Stream, + ): + self._model_runner_output = model_runner_output + + # Event on the copy stream so we can synchronize the non-blocking copy. + # Blocking (sleep) event to avoid busy-polling the CUDA driver lock. + self.async_copy_ready_event = torch.cuda.Event(blocking=True) + + # Keep a reference to the device tensors to avoid them being + # deallocated until we finish copying it to the host. + self._raw_pooler_output = raw_pooler_output + + # Initiate the copy on a separate stream, but do not synchronize it. + default_stream = torch.cuda.current_stream() + with torch.cuda.stream(async_output_copy_stream): + async_output_copy_stream.wait_stream(default_stream) + self._model_runner_output.pooler_output = _copy_pooler_output_to_cpu( + raw_pooler_output=self._raw_pooler_output, + finished_mask=finished_mask, + ) + self.async_copy_ready_event.record() + + def get_output(self) -> ModelRunnerOutput: + """Copy the device tensors to the host and return a ModelRunnerOutput. + This function blocks until the copy is finished. + """ + self.async_copy_ready_event.synchronize() + + # Release the device tensors once the copy has completed. + del self._raw_pooler_output + return self._model_runner_output + + +class ExecuteModelState(NamedTuple): + """Ephemeral cached state transferred between execute_model() and + sample_tokens(), after execute_model() returns None.""" + + scheduler_output: "SchedulerOutput" + logits: torch.Tensor + spec_decode_metadata: SpecDecodeMetadata | None + spec_decode_common_attn_metadata: CommonAttentionMetadata | None + hidden_states: torch.Tensor + sample_hidden_states: torch.Tensor + aux_hidden_states: list[torch.Tensor] | None + ec_connector_output: ECConnectorOutput | None + cudagraph_stats: CUDAGraphStat | None + slot_mappings: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None + + +class GPUModelRunner( + LoRAModelRunnerMixin, KVConnectorModelRunnerMixin, ECConnectorModelRunnerMixin +): + def __init__( + self, + vllm_config: VllmConfig, + device: torch.device, + ): + self.vllm_config = vllm_config + self.model_config = vllm_config.model_config + self.cache_config = vllm_config.cache_config + self.offload_config = vllm_config.offload_config + self.compilation_config = vllm_config.compilation_config + self.lora_config = vllm_config.lora_config + self.load_config = vllm_config.load_config + self.parallel_config = vllm_config.parallel_config + self.scheduler_config = vllm_config.scheduler_config + self.speculative_config = vllm_config.speculative_config + self.observability_config = vllm_config.observability_config + self.jit_warmup_registry = JitWarmupRegistry(vllm_config) + + model_config = self.model_config + cache_config = self.cache_config + scheduler_config = self.scheduler_config + parallel_config = self.parallel_config + self.device = device + self.dtype = self.model_config.dtype + + self.check_ep_fault = False + if parallel_config.data_parallel_size > 1 and self.model_config.is_moe: + self.check_ep_fault = get_ep_all2all_manager().support_fault_tolerance + + self.kv_cache_dtype = kv_cache_dtype_str_to_dtype( + cache_config.cache_dtype, self.model_config + ) + + self.is_pooling_model = model_config.runner_type == "pooling" + self.enable_prompt_embeds = model_config.enable_prompt_embeds + self.is_multimodal_raw_input_only_model = ( + model_config.is_multimodal_raw_input_only_model + ) + # These will be overridden in load_model() + self.is_multimodal_pruning_enabled = False + self.requires_sequential_video_encoding = False + # Set to True after init_routed_experts_capturer() completes. + # Prevents routed experts code from running during profiling/dummy run. + self.routed_experts_initialized = False + self.max_model_len = model_config.max_model_len + + # Always set to false after the first forward pass + self.dcp_world_size = self.parallel_config.decode_context_parallel_size + self.dcp_rank = 0 if self.dcp_world_size <= 1 else get_dcp_group().rank_in_group + self.max_num_tokens = scheduler_config.max_num_batched_tokens + self.max_num_reqs = scheduler_config.max_num_seqs + + # Broadcast PP output for external_launcher (torchrun) + # to make sure we are synced across pp ranks + # TODO: Support overlapping micro-batches + # https://github.com/vllm-project/vllm/issues/18019 + self.broadcast_pp_output = ( + self.parallel_config.distributed_executor_backend == "external_launcher" + and len(get_pp_group().ranks) > 1 + ) + + # Model-related. + self.num_query_heads = model_config.get_num_attention_heads(parallel_config) + self.inputs_embeds_size = model_config.get_inputs_embeds_size() + # Only relevant for models using ALiBi (e.g, MPT) + self.use_alibi = model_config.uses_alibi + + self.cascade_attn_enabled = not self.model_config.disable_cascade_attn + self.is_mm_prefix_lm = self.model_config.is_mm_prefix_lm + + # Multi-modal data support + self.mm_registry = MULTIMODAL_REGISTRY + self.uses_mrope = model_config.uses_mrope + self.uses_xdrope_dim = model_config.uses_xdrope_dim + self.supports_mm_inputs = self.mm_registry.supports_multimodal_inputs( + model_config + ) + + if self.model_config.is_encoder_decoder: + # Maximum length of the encoder input, only for encoder-decoder + # models. + self.max_encoder_len = scheduler_config.max_num_encoder_input_tokens + else: + self.max_encoder_len = 0 + + # Async scheduling + self.use_async_scheduling = self.scheduler_config.async_scheduling + + # Async PP broadcast of sampled token ids, waited on in _prepare_input_ids. + self._pp_recv_work: torch.distributed.Work | None = None + + # Sampler + self.sampler = Sampler( + logprobs_mode=self.model_config.logprobs_mode, + use_fp64_gumbel=self.model_config.use_fp64_gumbel, + ) + + self.eplb_state: EplbState | None = None + self._moe_model: MixtureOfExperts | None = None + # NOTE(yongji): flag to temporarily disable EPLB during scaling up/down + self.eep_eplb_suppressed = False + """ + State of the expert parallelism load balancer. + + Will be lazily initialized when the model is loaded. + """ + + # Lazy initializations + # self.model: nn.Module # Set after load_model + # Initialize in initialize_kv_cache + self.kv_caches: list[torch.Tensor] = [] + # indexes: [kv_cache_group_id][attn_group] + self.attn_groups: list[list[AttentionGroup]] = [] + # self.kv_cache_config: KVCacheConfig + + # mm_hash -> encoder_output + self.encoder_cache: dict[str, torch.Tensor] = {} + self.late_interaction_runner = LateInteractionRunner() + + # Encoder CUDA graph manager (initialized after model load if enabled) + self.encoder_cudagraph_manager: EncoderCudaGraphManager | None = None + + self.use_aux_hidden_state_outputs = False + # Set up speculative decoding. + # NOTE(Jiayi): currently we put the entire draft model on + # the last PP rank. This is not ideal if there are many + # layers in the draft model. + if self.speculative_config and get_pp_group().is_last_rank: + self.drafter: ( + NgramProposer # noqa: F823 + | NgramProposerGPU + | SuffixDecodingProposer + | EagleProposer + | DFlashProposer + | DraftModelProposer + | MedusaProposer + | ExtractHiddenStatesProposer + | Gemma4Proposer + | Step3p5MTPProposer + ) + if self.speculative_config.method == "custom_class": + self.drafter = create_custom_proposer( # type: ignore[assignment] + self.vllm_config + ) + elif self.speculative_config.method == "ngram": + from vllm.v1.spec_decode.ngram_proposer import NgramProposer + + self.drafter = NgramProposer(self.vllm_config) + elif self.speculative_config.uses_draft_model(): + self.drafter = DraftModelProposer( + vllm_config=self.vllm_config, + device=self.device, + runner=self, + ) + elif self.speculative_config.use_ngram_gpu(): + self.drafter = NgramProposerGPU(self.vllm_config, self.device, self) + self.num_tokens_no_spec_gpu = torch.zeros( + self.max_num_reqs, dtype=torch.int32, device=device + ) + self.token_ids_gpu_tensor = torch.zeros( + self.max_num_reqs, + self.max_model_len, + dtype=torch.int32, + device=device, + ) + self._ngram_pinned_idx_buf = torch.zeros( + self.max_num_reqs, dtype=torch.long, pin_memory=True + ) + self._ngram_pinned_val_buf = torch.zeros( + self.max_num_reqs, dtype=torch.int32, pin_memory=True + ) + elif self.speculative_config.use_gemma4_mtp(): + self.drafter = Gemma4Proposer(self.vllm_config, self.device, self) + elif self.speculative_config.use_step3p5_mtp(): + self.drafter = Step3p5MTPProposer(self.vllm_config, self.device, self) + elif self.speculative_config.use_dflash(): + self.drafter = DFlashProposer(self.vllm_config, self.device, self) + self.use_aux_hidden_state_outputs = True + elif self.speculative_config.method == "suffix": + self.drafter = SuffixDecodingProposer(self.vllm_config) + elif self.speculative_config.use_eagle(): + self.drafter = EagleProposer(self.vllm_config, self.device, self) + if self.speculative_config.method == "eagle3": + self.use_aux_hidden_state_outputs = ( + self.drafter.eagle3_use_aux_hidden_state + ) + elif self.speculative_config.method == "medusa": + self.drafter = MedusaProposer( + vllm_config=self.vllm_config, device=self.device + ) + elif self.speculative_config.method == "extract_hidden_states": + self.drafter = ExtractHiddenStatesProposer( + vllm_config=self.vllm_config, device=self.device + ) + self.use_aux_hidden_state_outputs = True + else: + raise ValueError( + "Unknown speculative decoding method: " + f"{self.speculative_config.method}" + ) + self.rejection_sampler = RejectionSampler( + self.sampler, self.speculative_config, self.device + ) + + self.num_spec_tokens = 0 + self.prev_num_spec_tokens = 0 + self.valid_sampled_token_count_gpu: torch.Tensor | None = None + if self.speculative_config: + self.num_spec_tokens = self.speculative_config.num_speculative_tokens + self.prev_num_spec_tokens = self.num_spec_tokens + draft_config = self.speculative_config.draft_model_config + if draft_config is not None and draft_config.max_model_len is not None: + self.effective_drafter_max_model_len = draft_config.max_model_len + else: + self.effective_drafter_max_model_len = self.max_model_len + self.use_async_spec_decode = ( + self.use_async_scheduling and self.num_spec_tokens > 0 + ) + + # Request states. + self.requests: dict[str, CachedRequestState] = {} + # NOTE(rob): num_prompt_logprobs only includes reqs + # that are currently in the prefill phase. + self.num_prompt_logprobs: dict[str, int] = {} + + # Input Batch + # NOTE(Chen): Ideally, we should initialize the input batch inside + # `initialize_kv_cache` based on the kv cache config. However, as in + # https://github.com/vllm-project/vllm/pull/18298, due to some unknown + # reasons, we have to initialize the input batch before `load_model`, + # quantization + weight offloading will fail otherwise. As a temporary + # solution, we initialize the input batch here, and re-initialize it + # in `initialize_kv_cache` if the block_sizes here is different from + # the block_sizes in the kv cache config. + logits_processors = model_config.logits_processors + custom_logitsprocs: Sequence[str | type[LogitsProcessor]] = ( + tuple(logits_processors) if logits_processors is not None else () + ) + placeholder_block_size = ( + self.cache_config.block_size or CacheConfig.DEFAULT_BLOCK_SIZE + ) + placeholder_max_num_blocks = cdiv( + max(self.max_model_len, self.max_encoder_len), placeholder_block_size + ) + self._init_block_sizes = [placeholder_block_size] + self._init_kernel_block_sizes = [placeholder_block_size] + self._init_max_num_blocks = [placeholder_max_num_blocks] + self._init_slot_mapping_modes = [SlotMappingMode.TOKEN_TO_KV_SLOT] + # Capture warmup providers registered by the initial placeholder InputBatch + with self.jit_warmup_registry.activate(): + self.input_batch = InputBatch( + max_num_reqs=self.max_num_reqs, + # We need to use the encoder length for encoder-decoder + # because of KV cache for cross-attention. + max_model_len=max(self.max_model_len, self.max_encoder_len), + max_num_batched_tokens=self.max_num_tokens, + device=self.device, + vocab_size=self.model_config.get_vocab_size(), + block_sizes=[placeholder_block_size], + kernel_block_sizes=[placeholder_block_size], + max_num_blocks_per_req=[placeholder_max_num_blocks], + num_spec_tokens=self.num_spec_tokens, + logitsprocs=build_logitsprocs( + self.vllm_config, + self.device, + PIN_MEMORY, + self.is_pooling_model, + custom_logitsprocs, + ), + # We currently don't know whether a particular custom logits processor + # uses output token ids so we set this conservatively. Thinking-budget + # tracking is requested dynamically when a budgeted request is in the + # batch. + logitsprocs_need_output_token_ids=bool(custom_logitsprocs), + is_pooling_model=self.is_pooling_model, + cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size, + reasoning_config=self.vllm_config.reasoning_config, + use_replayssm=self.cache_config.use_replayssm, + ) + + # Separate cuda stream for overlapping transfer of sampled token ids from + # GPU to CPU when async scheduling is enabled. + self.async_output_copy_stream: torch.cuda.Stream | None = None + # cuda event to synchronize use of reused CPU tensors between steps + # when async scheduling is enabled. + self.prepare_inputs_event: torch.Event | None = None + if self.use_async_scheduling: + self.async_output_copy_stream = torch.cuda.Stream() + # Blocking (sleep) event to avoid busy-polling the CUDA driver lock; + # under TP contention that spin can balloon and make the rank a straggler. + self.prepare_inputs_event = torch.cuda.Event(blocking=True) + + # self.cudagraph_batch_sizes sorts in ascending order. + if ( + self.compilation_config.cudagraph_capture_sizes + and self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE + ): + self.cudagraph_batch_sizes = sorted( + self.compilation_config.cudagraph_capture_sizes + ) + else: + self.cudagraph_batch_sizes = [] + + # Cache the device properties. + self._init_device_properties() + + # Encoder timing registry for observability + self.encoder_timing_registry: dict[str, EncoderTimingStats] = {} + self._encoder_timing_lock = threading.Lock() + + # Persistent buffers for CUDA graphs. + self.input_ids = self._make_buffer(self.max_num_tokens, dtype=torch.int32) + self.positions = torch.zeros( + self.max_num_tokens, dtype=torch.int64, device=self.device + ) + self.query_start_loc = self._make_buffer( + self.max_num_reqs + 1, dtype=torch.int32 + ) + self.seq_lens = torch.zeros( + self.max_num_reqs, dtype=torch.int32, device=self.device + ) + self.optimistic_seq_lens_cpu = torch.zeros( + self.max_num_reqs, dtype=torch.int32, pin_memory=PIN_MEMORY + ) + self.num_computed_tokens = torch.zeros( + self.max_num_reqs, dtype=torch.int32, device=self.device + ) + self.prev_num_draft_tokens = self._make_buffer( + self.max_num_reqs, dtype=torch.int32 + ) + self.req_indices = self._make_buffer(self.max_num_tokens, dtype=torch.int64) + # Maps current batch position -> previous batch position (-1 for new reqs) + self.prev_positions = self._make_buffer(self.max_num_reqs, dtype=torch.int64) + self.num_scheduled_tokens = self._make_buffer( + self.max_num_reqs, dtype=torch.int32 + ) + + self.encoder_seq_lens = self._make_buffer(self.max_num_reqs, dtype=torch.int32) + if self.dcp_world_size > 1: + self.dcp_local_seq_lens = self._make_buffer( + self.max_num_reqs, dtype=torch.int32 + ) + # Because inputs_embeds may be bfloat16 and we don't need a numpy + # version of this tensor, avoid a RuntimeError by not creating a + # numpy buffer. + self.inputs_embeds = self._make_buffer( + self.max_num_tokens, self.inputs_embeds_size, dtype=self.dtype, numpy=False + ) + self.is_token_ids = self._make_buffer(self.max_num_tokens, dtype=torch.bool) + self.discard_request_mask = self._make_buffer( + self.max_num_reqs, dtype=torch.bool + ) + self.num_decode_draft_tokens = self._make_buffer( + self.max_num_reqs, dtype=torch.int32 + ) + self.num_accepted_tokens = self._make_buffer( + self.max_num_reqs, dtype=torch.int32 + ) + + # Only relevant for models using M-RoPE (e.g, Qwen2-VL) + if self.uses_mrope: + # NOTE: `mrope_positions` is implemented with one additional dummy + # position on purpose to make it non-contiguous so that it can work + # with torch compile. + # See detailed explanation in https://github.com/vllm-project/vllm/pull/12128#discussion_r1926431923 + + # NOTE: When M-RoPE is enabled, position ids are 3D regardless of + # the modality of inputs. For text-only inputs, each dimension has + # identical position IDs, making M-RoPE functionally equivalent to + # 1D-RoPE. + # See page 5 of https://arxiv.org/abs/2409.12191 + self.mrope_positions = self._make_buffer( + (3, self.max_num_tokens + 1), dtype=torch.int64 + ) + + # Only relevant for models using XD-RoPE (e.g, HunYuan-VL) + if self.uses_xdrope_dim > 0: + # Similar to mrope but use assigned dimension number for RoPE, 4 as default. + self.xdrope_positions = self._make_buffer( + (self.uses_xdrope_dim, self.max_num_tokens + 1), dtype=torch.int64 + ) + + # None in the first PP rank. The rest are set after load_model. + self.intermediate_tensors: IntermediateTensors | None = None + + # OPTIMIZATION: Cache the arange tensors rather than creating them + # every step. Keep in int64 to avoid overflow with long context. + # - arange_np: immutable [0, 1, 2, ...] used as source for batched computation + # - query_pos: CpuGpuBuffer for the computed batched arange result + arange_size = max(self.max_num_reqs + 1, self.max_num_tokens) + self.arange_np = np.arange(arange_size, dtype=np.int64) + self.query_pos = self._make_buffer(arange_size, dtype=torch.int64) + self._arange_scratch = np.empty(arange_size, dtype=np.int64) + + # Layer pairings for cross-layer KV sharing. + # If an Attention layer `layer_name` is in the keys of this dict, it + # means this layer will perform attention using the keys and values + # from the KV cache of `shared_kv_cache_layers[layer_name]`. + self.shared_kv_cache_layers: dict[str, str] = {} + self.kv_sharing_fast_prefill_eligible_layers: set[str] = set() + + self.kv_sharing_fast_prefill_logits_indices = None + if self.cache_config.kv_sharing_fast_prefill: + self.kv_sharing_fast_prefill_logits_indices = torch.zeros( + self.max_num_tokens, dtype=torch.int32, device=self.device + ) + + self.uniform_decode_query_len = 1 + self.num_spec_tokens + + # Cudagraph dispatcher for runtime cudagraph dispatching. + self.cudagraph_dispatcher = CudagraphDispatcher(self.vllm_config) + + self.mm_budget = ( + MultiModalBudget(self.vllm_config, self.mm_registry) + if self.supports_mm_inputs + else None + ) + + self.reorder_batch_threshold: int | None = None + + # Attention layers that are only in the KVCacheConfig of the runner + # (e.g., KV sharing, encoder-only attention), but not in the + # KVCacheConfig of the scheduler. + self.runner_only_attn_layers: set[str] = set() + + # Cached outputs. + self._draft_token_ids: list[list[int]] | torch.Tensor | None = None + self._draft_probs: torch.Tensor | None = None + self._draft_prob_req_ids: list[str] | None = None + # N-gram GPU path: async D2H buffer/event for per-request valid draft counts. + self._num_valid_draft_tokens: torch.Tensor | None = None + self._num_valid_draft_tokens_cpu: torch.Tensor | None = None + self._num_valid_draft_tokens_event: torch.cuda.Event | None = None + self._num_valid_draft_tokens_copy_stream: torch.cuda.Stream | None = None + if ( + self.speculative_config is not None + and self.speculative_config.use_ngram_gpu() + ): + self._num_valid_draft_tokens_cpu = torch.empty( + self.max_num_reqs, dtype=torch.int32, pin_memory=PIN_MEMORY + ) + self._num_valid_draft_tokens_event = torch.cuda.Event() + self._num_valid_draft_tokens_copy_stream = torch.cuda.Stream() + + self._draft_token_req_ids: list[str] | None = None + self.transfer_event = torch.Event() + self.sampled_token_ids_pinned_cpu = torch.empty( + (self.max_num_reqs, 1), + dtype=torch.int64, + device="cpu", + pin_memory=PIN_MEMORY, + ) + + # Pre-allocated tensor for copying valid sampled token counts to CPU, + # with dedicated stream for overlapping and event for coordination. + self.valid_sampled_token_count_event: torch.Event | None = None + self.valid_sampled_token_count_copy_stream: torch.cuda.Stream | None = None + # We also copy the drafted tokens to the CPU asynchronously, + # in case we need them for structured outputs. + self.draft_token_ids_event: torch.Event | None = None + self.draft_token_ids_copy_stream: torch.cuda.Stream | None = None + self.valid_sampled_token_count_cpu: torch.Tensor | None = None + self.draft_token_ids_cpu: torch.Tensor | None = None + self.num_accepted_tokens_event: torch.Event | None = None + if self.num_spec_tokens: + self.draft_token_ids_event = torch.Event() + self.num_accepted_tokens_event = torch.Event() + self.draft_token_ids_copy_stream = torch.cuda.Stream() + self.draft_token_ids_cpu = torch.empty( + (self.max_num_reqs, self.num_spec_tokens), + dtype=torch.int64, + device="cpu", + pin_memory=PIN_MEMORY, + ) + if self.use_async_scheduling: + self.valid_sampled_token_count_event = torch.Event() + self.valid_sampled_token_count_copy_stream = torch.cuda.Stream() + self.valid_sampled_token_count_cpu = torch.empty( + self.max_num_reqs, + dtype=torch.int32, + device="cpu", + pin_memory=PIN_MEMORY, + ) + + # Model weight offloader + # Make sure this is called before any get_offloader call + set_offloader(create_offloader(self.offload_config)) + + # Ephemeral state transferred between execute_model() and sample_tokens(). + self.execute_model_state: ExecuteModelState | None = None + self.kv_connector_output: KVConnectorOutput | None = None + self.mamba_state_idx: dict[str, int] = {} + self._mamba_bufs: mamba_utils.MambaBuffers | None = None + self.mamba_prev_last_scheduled_idx: CpuGpuBuffer | None = None + if self.cache_config.mamba_cache_mode == "all" and self.num_spec_tokens > 0: + self.mamba_prev_last_scheduled_idx = self._make_buffer( + self.max_num_reqs, dtype=torch.int32 + ) + self.layerwise_nvtx_hooks_registered = False + + def update_max_model_len(self, max_model_len: int) -> None: + self.max_model_len = max_model_len + if self.speculative_config: + draft_config = self.speculative_config.draft_model_config + if draft_config is None or draft_config.max_model_len is None: + self.effective_drafter_max_model_len = self.max_model_len + update_model_len = getattr( + getattr(self, "model", None), "update_max_model_len", None + ) + if update_model_len is not None: + update_model_len(max_model_len) + + def reset_mm_cache(self) -> None: + """ + Clear the multi-modal cache that was used during profiling, + but no longer needed during inference. + """ + if self.mm_budget: + self.mm_budget.reset_cache() + self.late_interaction_runner.clear() + + def reset_encoder_cache(self) -> None: + """Clear the GPU-side encoder cache storing vision embeddings. + + This should be called when model weights are updated to ensure + stale embeddings computed with old weights are not reused. + """ + self.encoder_cache.clear() + self.late_interaction_runner.clear() + + def post_kv_cache_wake_up(self) -> None: + self.init_fp8_kv_scales() + + @torch.inference_mode() + def init_fp8_kv_scales(self) -> None: + """ + Re-initialize the KV cache and FP8 scales after waking from sleep. + 1. Zero out the KV cache tensors to remove garbage data from re-allocation. + 2. Reset Attention layer scaling factors (_k_scale, _v_scale) to 1.0. + If these are left at 0.0 (default after wake_up), all KV cache values + become effectively zero, causing gibberish output. + """ + if not is_quantized_kv_cache(self.cache_config.cache_dtype): + return + + kv_caches = getattr(self, "kv_caches", []) + for cache_entry in kv_caches: + if cache_entry is None: + continue + # Hybrid models (Mamba, DeltaNet) store per-layer state as a + # list of tensors rather than a single tensor. + if isinstance(cache_entry, list): + for t in cache_entry: + t.zero_() + else: + cache_entry.zero_() + + k_attr_names = ("_k_scale", "k_scale") + v_attr_names = ("_v_scale", "v_scale") + + attn_layers = self.compilation_config.static_forward_context + for name, module in attn_layers.items(): + if isinstance(module, (Attention, MLAAttention)): + # TODO: Generally, scale is 1.0 if user uses on-the-fly fp8 + # kvcache quant. However, to get better accuracy, compression + # frameworks like llm-compressors allow users to tune the + # scale. We may need to restore the specific calibrated scales + # here in the future. + k_scale_val, v_scale_val = 1.0, 1.0 + + # Processing K Scale + for attr in k_attr_names: + if hasattr(module, attr): + param = getattr(module, attr) + if isinstance(param, torch.Tensor): + param.fill_(k_scale_val) + + # Processing V Scale + for attr in v_attr_names: + if hasattr(module, attr): + param = getattr(module, attr) + if isinstance(param, torch.Tensor): + param.fill_(v_scale_val) + + def _get_positions(self, num_tokens: Any): + if isinstance(num_tokens, int): + if self.uses_mrope: + return self.mrope_positions.gpu[:, :num_tokens] + if self.uses_xdrope_dim > 0: + return self.xdrope_positions.gpu[:, :num_tokens] + return self.positions[:num_tokens] + else: + if self.uses_mrope: + return self.mrope_positions.gpu[:, num_tokens] + if self.uses_xdrope_dim > 0: + return self.xdrope_positions.gpu[:, num_tokens] + return self.positions[num_tokens] + + def _make_buffer( + self, *size: int | torch.SymInt, dtype: torch.dtype, numpy: bool = True + ) -> CpuGpuBuffer: + return CpuGpuBuffer( + *size, + dtype=dtype, + device=self.device, + with_numpy=numpy, + ) + + def _get_mamba_bufs(self) -> mamba_utils.MambaBuffers: + # Only reachable on the ``mamba_cache_mode == "align"`` path. + # The postprocess sub-object is additionally gated on spec + # decode + hybrid model. + assert self.cache_config.mamba_cache_mode == "align" + if self._mamba_bufs is None: + copy_funcs = mamba_utils.resolve_mamba_state_copy_funcs( + self.model, self.kv_cache_config + ) + self._mamba_bufs = mamba_utils.MambaBuffers.create( + max_num_reqs=self.max_num_reqs, + kv_cache_config=self.kv_cache_config, + copy_funcs=copy_funcs, + make_buffer=self._make_buffer, + device=self.device, + with_postprocess_align=( + self.speculative_config is not None and self.model_config.is_hybrid + ), + ) + return self._mamba_bufs + + def _init_model_kwargs(self): + model_kwargs = dict[str, Any]() + + if not self.is_pooling_model: + return model_kwargs + + num_reqs = self.input_batch.num_reqs + pooling_params = self.input_batch.get_pooling_params() + + token_type_id_requests = dict[int, Any]() + for i, param in enumerate(pooling_params): + if ( + param.extra_kwargs is not None + and (token_types := param.extra_kwargs.get("compressed_token_type_ids")) + is not None + ): + token_type_id_requests[i] = token_types + + if len(token_type_id_requests) == 0: + return model_kwargs + + # Build ids on CPU using the CPU-resident upper bound for seq_lens; + # `torch.arange(seq_lens[i])` with a GPU scalar would force a sync. + seq_lens_cpu = self.optimistic_seq_lens_cpu[:num_reqs].tolist() + token_type_ids = [] + + for i in range(num_reqs): + seq_len_i = seq_lens_cpu[i] + pos = token_type_id_requests.get(i, seq_len_i) + ids = (torch.arange(seq_len_i) >= pos).int() + token_type_ids.append(ids) + + token_type_ids_cpu = torch.empty( + sum(seq_lens_cpu), dtype=torch.int32, pin_memory=PIN_MEMORY + ) + torch.cat(token_type_ids, out=token_type_ids_cpu) + model_kwargs["token_type_ids"] = token_type_ids_cpu.to( + device=self.device, non_blocking=True + ) + return model_kwargs + + def _may_reorder_batch(self, scheduler_output: "SchedulerOutput") -> None: + """ + Update the order of requests in the batch based on the attention + backend's needs. For example, some attention backends (namely MLA) may + want to separate requests based on if the attention computation will be + compute-bound or memory-bound. + + Args: + scheduler_output: The scheduler output. + """ + # Attention free models have zero kv_cache_groups, however models + # like Mamba are also attention free but use the kv_cache for + # keeping its internal state. This is why we check the number + # of kv_cache groups instead of solely checking + # for self.model_config.is_attention_free. + if len(self.kv_cache_config.kv_cache_groups) == 0: + return + + if self.reorder_batch_threshold is not None: + reorder_batch_to_split_decodes_and_prefills( + self.input_batch, + scheduler_output, + decode_threshold=self.reorder_batch_threshold, + ) + + def _init_kv_zero_meta(self) -> None: + """One-time precomputation for _zero_block_ids. + + Called from gpu_worker.py outside the CuMem pool context. + """ + self._kv_block_zeroer = KVBlockZeroer( + self.device, + attn_groups_iter=self._kv_cache_spec_attn_group_iterator(), + kernel_block_sizes=self._kernel_block_sizes, + runner_only_attn_layers=self.runner_only_attn_layers, + static_forward_context=self.compilation_config.static_forward_context, + ) + + def _zero_block_ids(self, block_ids: list[int]) -> None: + """Zero the KV cache memory for the given block IDs.""" + if hasattr(self, "_kv_block_zeroer"): + self._kv_block_zeroer.zero_block_ids(block_ids) + + # Note: used for model runner override. + def _init_device_properties(self) -> None: + """Initialize attributes from torch.cuda.get_device_properties""" + + self.num_sms = num_compute_units(self.device.index) + + # Note: used for model runner override. + def _sync_device(self) -> None: + torch.accelerator.synchronize() + + def _get_or_create_async_output_copy_stream(self) -> torch.cuda.Stream: + stream = self.async_output_copy_stream + if stream is None: + stream = torch.cuda.Stream() + self.async_output_copy_stream = stream + return stream + + def _on_request_state_removed( + self, + req_id: str, + req_state: CachedRequestState | None, + ) -> None: + """Hook for platform runners to clean request-scoped side caches.""" + del req_id, req_state + + def _process_encoder_cache_scheduler_output( + self, + scheduler_output: "SchedulerOutput", + ) -> None: + """Apply scheduler-side encoder cache lifecycle updates.""" + for mm_hash in scheduler_output.free_encoder_mm_hashes: + self.encoder_cache.pop(mm_hash, None) + + def _update_states(self, scheduler_output: "SchedulerOutput") -> Callable | None: + """Update the cached states and the persistent batch with the scheduler + output. + + The updated states are used by the `_prepare_inputs` function to create + the input GPU tensors for the model. + + The SamplingMetadata is updated and copied to the GPU if there is a + new/resumed/paused/finished request in the batch. + """ + # Remove finished requests from the cached states. + for req_id in scheduler_output.finished_req_ids: + req_state = self.requests.pop(req_id, None) + self._on_request_state_removed(req_id, req_state) + self.num_prompt_logprobs.pop(req_id, None) + self.late_interaction_runner.on_requests_finished( + scheduler_output.finished_req_ids + ) + # Remove the finished requests from the persistent batch. + # NOTE(woosuk): There could be an edge case where finished_req_ids and + # scheduled_req_ids overlap. This happens when a request is aborted and + # then resubmitted with the same ID. In this case, we treat them as two + # distinct requests - clearing the cached states for the first request + # and handling the second as a new request. + for req_id in scheduler_output.finished_req_ids: + self.input_batch.remove_request(req_id) + + # Zero GPU memory for freshly allocated cache blocks to prevent + # stale NaN/data from corrupting attention or SSM computation. + if scheduler_output.new_block_ids_to_zero: + self._zero_block_ids(scheduler_output.new_block_ids_to_zero) + if scheduler_output.kv_cache_block_copies: + copy_kv_cache_blocks_inplace( + self.kv_caches, + self.kv_cache_config.num_blocks, + scheduler_output.kv_cache_block_copies, + ) + + # Free the cached encoder outputs. + self._process_encoder_cache_scheduler_output(scheduler_output) + + # Remove the unscheduled requests from the persistent batch. + # NOTE(woosuk): The unscheduled requests are either preempted requests + # or running requests that are not scheduled in this step. We remove + # them from the persistent batch but keep their cached states since + # they will be scheduled again sometime in the future. + scheduled_req_ids = scheduler_output.num_scheduled_tokens.keys() + cached_req_ids = self.input_batch.req_id_to_index.keys() + resumed_req_ids = scheduler_output.scheduled_cached_reqs.resumed_req_ids + # NOTE(zhuohan): cached_req_ids and resumed_req_ids are usually disjoint, + # so `(scheduled_req_ids - resumed_req_ids) == scheduled_req_ids` holds + # apart from the forced-preemption case in reset_prefix_cache. And in + # that case we include the resumed_req_ids in the unscheduled set so + # that they get cleared from the persistent batch before being re-scheduled + # in the normal resumed request path. + unscheduled_req_ids = cached_req_ids - (scheduled_req_ids - resumed_req_ids) + # NOTE(woosuk): The persistent batch optimization assumes that + # consecutive batches contain mostly the same requests. If batches + # have low request overlap (e.g., alternating between two distinct + # sets of requests), this optimization becomes very inefficient. + for req_id in unscheduled_req_ids: + self.input_batch.remove_request(req_id) + + is_ngram_gpu = ( + self.speculative_config is not None + and self.speculative_config.use_ngram_gpu() + ) + if is_ngram_gpu: + ngram_gpu_new_reqs: list[CachedRequestState] = [] + + reqs_to_add: list[CachedRequestState] = [] + deferred_spec_decode_corrections = [] + + # Add new requests to the cached states. + for new_req_data in scheduler_output.scheduled_new_reqs: + req_id = new_req_data.req_id + if req_id in self.requests: + # For streaming case only. + req_state = self._update_streaming_request(req_id, new_req_data) + reqs_to_add.append(req_state) + continue + + sampling_params = new_req_data.sampling_params + pooling_params = new_req_data.pooling_params + + if ( + sampling_params + and sampling_params.sampling_type == SamplingType.RANDOM_SEED + ): + generator = torch.Generator(device=self.device) + generator.manual_seed(sampling_params.seed) + else: + generator = None + + if self.is_pooling_model: + assert pooling_params is not None + task = pooling_params.task + assert task is not None, "You did not set `task` in the API" + + model = cast(VllmModelForPooling, self.get_model()) + to_update = model.pooler.get_pooling_updates(task) + to_update.apply(pooling_params) + + req_state = CachedRequestState( + req_id=req_id, + prompt_token_ids=new_req_data.prompt_token_ids, + prompt_embeds=new_req_data.prompt_embeds, + prompt_is_token_ids=new_req_data.prompt_is_token_ids, + mm_features=new_req_data.mm_features, + sampling_params=sampling_params, + pooling_params=pooling_params, + generator=generator, + block_ids=new_req_data.block_ids, + num_computed_tokens=new_req_data.num_computed_tokens, + output_token_ids=[], + lora_request=new_req_data.lora_request, + ) + self.requests[req_id] = req_state + self.late_interaction_runner.register_request(req_id, pooling_params) + + if sampling_params and sampling_params.prompt_logprobs is not None: + self.num_prompt_logprobs[req_id] = ( + self.input_batch.vocab_size + if sampling_params.prompt_logprobs == -1 + else sampling_params.prompt_logprobs + ) + + # Only relevant for models using M-RoPE (e.g, Qwen2-VL) + if self.uses_mrope: + self._init_mrope_positions(req_state) + + # Only relevant for models using XD-RoPE (e.g, HunYuan-VL) + if self.uses_xdrope_dim > 0: + self._init_xdrope_positions(req_state) + + reqs_to_add.append(req_state) + # Track new requests for ngram_gpu full tensor copy + if is_ngram_gpu: + ngram_gpu_new_reqs.append(req_state) + + # Update the states of the running/resumed requests. + is_last_rank = get_pp_group().is_last_rank + req_data = scheduler_output.scheduled_cached_reqs + scheduled_spec_tokens = scheduler_output.scheduled_spec_decode_tokens + + # Save scheduler-allocated spec lengths before trimming so + # prev_num_draft_len keeps the optimistic count for rejection correction. + original_num_spec_per_req: dict[str, int] = {} + if ( + self.speculative_config is not None + and self.speculative_config.use_ngram_gpu() + ): + for req_id, toks in scheduled_spec_tokens.items(): + original_num_spec_per_req[req_id] = len(toks) + update_scheduler_for_invalid_drafts( + self._num_valid_draft_tokens_event, + self._num_valid_draft_tokens_cpu, + scheduler_output, + self.input_batch.req_id_to_index, + ) + if self.use_async_spec_decode: + self.prev_num_draft_tokens.np.fill(0) + + for i, req_id in enumerate(req_data.req_ids): + req_state = self.requests[req_id] + num_computed_tokens = req_data.num_computed_tokens[i] + new_block_ids = req_data.new_block_ids[i] + resumed_from_preemption = req_id in req_data.resumed_req_ids + num_output_tokens = req_data.num_output_tokens[i] + req_index = self.input_batch.req_id_to_index.get(req_id) + + if req_state.prev_num_draft_len and self.use_async_scheduling: + # prev_num_draft_len is used in async scheduling mode with + # spec decode. it indicates if need to update num_computed_tokens + # of the request. for example: + # first step: num_computed_tokens = 0, spec_tokens = [], + # prev_num_draft_len = 0. + # second step: num_computed_tokens = 100(prompt length), + # spec_tokens = [a,b], prev_num_draft_len = 0. + # third step: num_computed_tokens = 100 + 2, spec_tokens = [c,d], + # prev_num_draft_len = 2. + # num_computed_tokens in first step and second step doesn't contain + # the spec tokens length, but in third step it contains the + # spec tokens length. we only need to update num_computed_tokens + # when prev_num_draft_len > 0. + if req_index is None: + req_state.prev_num_draft_len = 0 + else: + # Optimistically assume all accepted; queue up a correction + # to be called after the model forward to preserve async + # scheduling. Corrected on GPU in _prepare_inputs. + optimistic_num_accepted = req_state.prev_num_draft_len + req_state.output_token_ids.extend([-1] * optimistic_num_accepted) + + deferred_spec_decode_corrections.append( + (req_id, optimistic_num_accepted, req_state) + ) + + prev_req_index = ( + self.input_batch.prev_req_id_to_index.get(req_id) + if self.input_batch.prev_req_id_to_index + else None + ) + if prev_req_index is not None: + self.prev_num_draft_tokens.np[prev_req_index] = ( + optimistic_num_accepted + ) + + if is_ngram_gpu and optimistic_num_accepted > 0: + self.input_batch.num_tokens_no_spec[req_index] += ( + optimistic_num_accepted + ) + + # Update the cached states. + req_state.num_computed_tokens = num_computed_tokens + + if not is_last_rank: + if not req_data.new_token_ids: + # Async scheduled PP: Sampled tokens propagated via GPU broadcast. + new_token_ids: list[int] = [] + else: + # Non-async scheduling with PP: The scheduler sends + # sampled token ids back because there's no direct communication + # between the first-stage worker and the last-stage worker. + new_token_ids = req_data.new_token_ids[i] + # Add the sampled token(s) from the previous step (if any). + # This doesn't include "unverified" tokens like spec tokens. + num_new_tokens = ( + num_computed_tokens + len(new_token_ids) - req_state.num_tokens + ) + if num_new_tokens == 1: + # Avoid slicing list in most common case. + req_state.output_token_ids.append(new_token_ids[-1]) + elif num_new_tokens > 0: + req_state.output_token_ids.extend( + new_token_ids[-num_new_tokens:] + ) + elif num_output_tokens < len(req_state.output_token_ids): + # Some output tokens were discarded due to a sync-KV-load + # failure, or output_token_ids was inflated by the optimistic + # extend above (async spec decode). Align the cached state. + del req_state.output_token_ids[num_output_tokens:] + if req_index is not None: + end_idx = ( + self.input_batch.num_prompt_tokens[req_index] + + num_output_tokens + ) + self.input_batch.num_tokens_no_spec[req_index] = end_idx + + # Update the block IDs. + if not resumed_from_preemption: + if new_block_ids is not None: + # Append the new blocks to the existing block IDs. + for block_ids, new_ids in zip(req_state.block_ids, new_block_ids): + block_ids.extend(new_ids) + else: + assert req_index is None + assert new_block_ids is not None + # The request is resumed from preemption. + # Replace the existing block IDs with the new ones. + req_state.block_ids = new_block_ids + + if req_index is None: + # The request is not in the persistent batch. + # The request was either preempted and resumed later, or was not + # scheduled in the previous step and needs to be added again. + + if self.use_async_scheduling and num_output_tokens > 0: + # We must recover the output token ids for resumed requests in the + # async scheduling case, so that correct input_ids are obtained. + resumed_token_ids = req_data.all_token_ids[req_id] + req_state.output_token_ids = resumed_token_ids[-num_output_tokens:] + + reqs_to_add.append(req_state) + # Track resumed requests for ngram_gpu full tensor copy + if is_ngram_gpu: + ngram_gpu_new_reqs.append(req_state) + continue + + # Update the persistent batch. + self.input_batch.num_computed_tokens_cpu[req_index] = num_computed_tokens + if new_block_ids is not None: + self.input_batch.block_table.append_row(new_block_ids, req_index) + + # For the last rank, we don't need to update the token_ids_cpu + # because the sampled tokens are already cached. + if not is_last_rank: + start_token_index = self.input_batch.num_tokens_no_spec[req_index] + # For chunked prefill, num_computed_tokens may less + # than num_tokens_no_spec. + # Async scheduled PP: no new_token_ids, advance num_tokens_no_spec + # according to num_computed_tokens. + end_token_index = max( + start_token_index, + num_computed_tokens + len(new_token_ids), + ) + if end_token_index > start_token_index: + if new_token_ids: + # Add new_token_ids to token_ids_cpu. + num_new_tokens = end_token_index - start_token_index + tokens_to_append = new_token_ids[-num_new_tokens:] + self.input_batch.token_ids_cpu[ + req_index, start_token_index:end_token_index + ] = tokens_to_append + self.input_batch.is_token_ids[ + req_index, start_token_index:end_token_index + ] = True + self.input_batch.num_tokens_no_spec[req_index] = end_token_index + + # Add spec_token_ids to token_ids_cpu. + self.input_batch.update_req_spec_token_ids(req_state, scheduled_spec_tokens) + # Restore scheduler-side draft count after ngram trimming. + if original_num_spec_per_req: + orig = original_num_spec_per_req.get(req_id, 0) + if orig != req_state.prev_num_draft_len: + req_state.prev_num_draft_len = orig + + # Add the new or resumed requests to the persistent batch. + # The smaller empty indices are filled first. + for request in reqs_to_add: + self.input_batch.add_request(request) + self.input_batch.update_req_spec_token_ids(request, scheduled_spec_tokens) + + # Condense the batched states if there are gaps left by removed requests + self.input_batch.condense() + # Allow attention backend to reorder the batch, potentially + self._may_reorder_batch(scheduler_output) + # Refresh batch metadata with any pending updates. + self.input_batch.refresh_metadata() + + # Incrementally update ngram_gpu tensors after batch is stable + if is_ngram_gpu: + update_ngram_gpu_tensors_incremental( + self.input_batch, + self.token_ids_gpu_tensor, + self.num_tokens_no_spec_gpu, + ngram_gpu_new_reqs, + self.device, + _pinned_idx_buf=self._ngram_pinned_idx_buf, + _pinned_val_buf=self._ngram_pinned_val_buf, + ) + + if deferred_spec_decode_corrections: + + def correct_spec_decode_token_counts(): + valid_sampled_token_count = self._get_valid_sampled_token_count() + if not valid_sampled_token_count: + return + prev_req_id_to_index = self.input_batch.prev_req_id_to_index + if not prev_req_id_to_index: + return + for ( + req_id, + optimistic_num_accepted, + req_state, + ) in deferred_spec_decode_corrections: + prev_req_index = prev_req_id_to_index.get(req_id) + if prev_req_index is None: + continue + num_accepted = valid_sampled_token_count[prev_req_index] - 1 + correction = optimistic_num_accepted - num_accepted + req_state.num_computed_tokens -= correction + cur_req_index = self.input_batch.req_id_to_index.get(req_id) + if cur_req_index is None: + continue + self.input_batch.num_computed_tokens_cpu[cur_req_index] -= ( + correction + ) + if is_ngram_gpu and correction > 0: + self.input_batch.num_tokens_no_spec[cur_req_index] -= correction + self.num_tokens_no_spec_gpu[cur_req_index] -= correction + + return correct_spec_decode_token_counts + else: + return None + + def _update_states_after_model_execute( + self, output_token_ids: torch.Tensor, scheduler_output: "SchedulerOutput" + ) -> None: + """Update the cached states after model execution. + + This is used for MTP/EAGLE for hybrid models, as in linear attention, + only the last token's state is kept. In MTP/EAGLE, for draft tokens + the state are kept util we decide how many tokens are accepted for + each sequence, and a shifting is done during the next iteration + based on the number of accepted tokens. + """ + if not self.speculative_config or not self.model_config.is_hybrid: + return + + # Count the number of accepted tokens for each sequence. + # Valid tokens are contiguous from position 0, so counting non-(-1) + # tokens gives us the first -1 position (i.e., number of accepted). + num_reqs = output_token_ids.size(0) + self.num_accepted_tokens.gpu[:num_reqs] = (output_token_ids != -1).sum(dim=1) + + if self.cache_config.mamba_cache_mode == "align": + # Fused GPU postprocess: state copies + per-request accepted-token + # update without CPU-GPU sync. The metadata + # (num_scheduled_tokens, num_draft_tokens, num_computed_tokens) is + # pre-staged to GPU buffers in _prepare_inputs. + mamba_utils.postprocess_mamba_align_gpu( + bufs=self._get_mamba_bufs(), + num_reqs=num_reqs, + num_accepted_tokens_gpu=self.num_accepted_tokens.gpu, + num_accepted_tokens_cpu_tensor=( + self.input_batch.num_accepted_tokens_cpu_tensor + ), + input_batch=self.input_batch, + kv_cache_config=self.kv_cache_config, + forward_context=self.compilation_config.static_forward_context, + mamba_state_copy_funcs=( + self._get_mamba_bufs().preprocess.copy_funcs_by_type + ), + ) + + assert self.num_accepted_tokens_event is not None + self.num_accepted_tokens_event.record() + else: + self.input_batch.num_accepted_tokens_cpu_tensor[:num_reqs].copy_( + self.num_accepted_tokens.gpu[:num_reqs], non_blocking=True + ) + assert self.num_accepted_tokens_event is not None + self.num_accepted_tokens_event.record() + + if self.cache_config.mamba_cache_mode == "all": + mamba_utils.postprocess_mamba_all( + scheduler_output, + self.kv_cache_config, + self.input_batch, + self.requests, + self.mamba_state_idx, + self.num_spec_tokens, + num_reqs, + ) + + def _update_streaming_request( + self, req_id: str, new_req_data: NewRequestData + ) -> CachedRequestState: + """Updates streaming session request from `scheduled_new_reqs`. + + Removes the request from InputBatch (if present), updates the cached + state, and prepares it for re-addition to the batch. + + NOTE: prompt_token_ids includes intermediate output tokens - tokens + previously generated but now are input context (part of the prompt). + """ + self.input_batch.remove_request(req_id) + req_state = self.requests[req_id] + + req_state.prompt_token_ids = new_req_data.prompt_token_ids + req_state.mm_features = new_req_data.mm_features + req_state.prompt_embeds = new_req_data.prompt_embeds + req_state.sampling_params = new_req_data.sampling_params + req_state.pooling_params = new_req_data.pooling_params + self.late_interaction_runner.register_request(req_id, req_state.pooling_params) + req_state.block_ids = new_req_data.block_ids + req_state.num_computed_tokens = new_req_data.num_computed_tokens + req_state.num_prompt_tokens = length_from_prompt_token_ids_or_embeds( + req_state.prompt_token_ids, req_state.prompt_embeds + ) + + # Clear `output_token_ids` as previous output tokens are now part of + # `prompt_token_ids`. + req_state.output_token_ids.clear() + + if self.uses_mrope: + self._init_mrope_positions(req_state) + + return req_state + + def _init_mrope_positions(self, req_state: CachedRequestState): + model = self.get_model() + assert supports_mrope(model), "M-RoPE support is not implemented." + mrope_model = cast(SupportsMRoPE, model) + + # `prompt_embeds` is a passthrough modality (no grid_thw), models' + # M-RoPE code assumes per-feature grid info, so filter it out. The + # prompt_embeds positions are treated as text positions for M-RoPE. + mrope_features = [ + f for f in req_state.mm_features if f.modality != "prompt_embeds" + ] + + if req_state.prompt_token_ids is not None: + input_tokens = req_state.prompt_token_ids + elif req_state.prompt_embeds is not None: + # For embeddings-only inputs, get_mrope_input_positions only + # needs the sequence length when mm_features is empty (which is + # the case here since prompt_embeds are filtered out above). + seq_len = req_state.prompt_embeds.shape[0] + input_tokens = list(range(seq_len)) + else: + raise ValueError( + "M-RoPE requires either prompt_token_ids or prompt_embeds." + ) + + req_state.mrope_positions, req_state.mrope_position_delta = ( + mrope_model.get_mrope_input_positions( + input_tokens, + mrope_features, + ) + ) + + def _init_xdrope_positions(self, req_state: CachedRequestState): + model = self.get_model() + xdrope_model = cast(SupportsXDRoPE, model) + assert req_state.prompt_token_ids is not None, ( + "XD-RoPE requires prompt_token_ids to be available." + ) + assert supports_xdrope(model), "XD-RoPE support is not implemented." + + req_state.xdrope_positions = xdrope_model.get_xdrope_input_positions( + req_state.prompt_token_ids, + req_state.mm_features, + ) + + def _extract_mm_kwargs( + self, + scheduler_output: "SchedulerOutput", + ) -> BatchedTensorInputs: + if not scheduler_output or not self.is_multimodal_raw_input_only_model: + return {} + + mm_kwargs = list[tuple[str, MultiModalKwargsItem]]() + for req in scheduler_output.scheduled_new_reqs: + for feature in req.mm_features: + if feature.data is not None: + mm_kwargs.append((feature.modality, feature.data)) + + # Input all modalities at once + mm_kwargs_combined: BatchedTensorInputs = {} + for _, _, mm_kwargs_batch in group_and_batch_mm_kwargs( + mm_kwargs, + device=self.device, + pin_memory=PIN_MEMORY, + ): + mm_kwargs_combined.update(mm_kwargs_batch) + + return mm_kwargs_combined + + def _dummy_mm_kwargs(self, num_seqs: int) -> BatchedTensorInputs: + if not self.is_multimodal_raw_input_only_model: + return {} + + mm_budget = self.mm_budget + assert mm_budget is not None + + if not mm_budget.mm_max_toks_per_item: + return {} # No tower modalities (embed-only mode) + + dummy_modality = mm_budget.get_modality_with_max_tokens() + return self._get_mm_dummy_batch(dummy_modality, num_seqs) + + def _get_cumsum_and_arange( + self, + num_tokens: np.ndarray, + arange_out: np.ndarray, + cumsum_dtype: np.dtype | None = None, + ) -> np.ndarray: + """Get the cumulative sum and batched arange of the given array. + E.g., [2, 5, 3] -> [2, 7, 10], arange written to + arange_out[:10] as [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]. + Equivalent to but faster than: + np.concatenate([np.arange(n) for n in num_tokens]) + """ + # Step 1. [2, 5, 3] -> [2, 7, 10] + cu_num_tokens = np.cumsum(num_tokens, dtype=cumsum_dtype) + total_num_tokens = cu_num_tokens[-1] + # Step 2. [2, 7, 10] -> [0, 0, 2, 2, 2, 2, 2, 7, 7, 7] + cumsums_offsets = np.repeat(cu_num_tokens - num_tokens, num_tokens) + # Step 3. [0, 1, 0, 1, 2, 3, 4, 0, 1, 2] + np.subtract( + self.arange_np[:total_num_tokens], + cumsums_offsets, + out=arange_out[:total_num_tokens], + ) + + return cu_num_tokens + + def _compute_prev_positions(self, num_reqs: int) -> None: + """Build prev_positions mapping: current pos -> previous pos (-1 if new). + + Populates self.prev_positions.np[:num_reqs] with the mapping. + """ + prev_req_id_to_index = self.input_batch.prev_req_id_to_index + prev_positions = self.prev_positions.np[:num_reqs] + + if not prev_req_id_to_index: + prev_positions.fill(-1) + return + + for i, req_id in enumerate(self.input_batch.req_ids[:num_reqs]): + prev_positions[i] = prev_req_id_to_index.get(req_id, -1) + + def _prepare_input_ids( + self, + scheduler_output: "SchedulerOutput", + num_reqs: int, + total_num_scheduled_tokens: int, + cu_num_tokens: np.ndarray, + ) -> None: + """Prepare the input IDs for the current batch. + + Carefully handles the `prev_sampled_token_ids` which can be cached + from the previous engine iteration, in which case those tokens on the + GPU need to be copied into the corresponding slots into input_ids. + + Uses self.prev_positions[:num_reqs] which maps current pos -> prev pos + (-1 for new requests). + """ + + # Sync the async PP broadcast before reading sampled tokens. + if self._pp_recv_work is not None: + self._pp_recv_work.wait() + self._pp_recv_work = None + + if self.input_batch.prev_sampled_token_ids is None: + # Normal scheduling case + self.input_ids.copy_to_gpu(total_num_scheduled_tokens) + if self.enable_prompt_embeds: + self.inputs_embeds.copy_to_gpu(total_num_scheduled_tokens) + self.is_token_ids.copy_to_gpu(total_num_scheduled_tokens) + return + + # Async scheduling case, where some decode requests from the previous + # iteration won't have entries in input_ids_cpu and need to be copied + # on the GPU from prev_sampled_token_ids. + prev_positions = self.prev_positions.np[:num_reqs] + scheduled_spec_tokens = scheduler_output.scheduled_spec_decode_tokens + sample_flattened_indices: list[int] = [] + spec_flattened_indices: list[int] = [] + prev_draft_token_indices: list[int] = [] + prev_indices: list[int] = [] + common_indices_match = True + max_flattened_index = -1 + total_num_spec_tokens = 0 + + for cur_index in range(num_reqs): + prev_index = prev_positions[cur_index] + if prev_index < 0: + continue + prev_indices.append(prev_index) + req_id = self.input_batch.req_ids[cur_index] + # We need to compute the flattened input_ids index of the + # last token in each common request. + draft_len = len(scheduled_spec_tokens.get(req_id, ())) + total_num_spec_tokens += draft_len + flattened_index = cu_num_tokens[cur_index].item() - 1 + # example: cu_num_tokens = [2, 5, 8], draft_tokens = [1, 2, 2] + # sample_flattened_indices = [0, 2, 5] + # spec_flattened_indices = [1, 3, 4, 6, 7] + sample_flattened_indices.append(flattened_index - draft_len) + spec_flattened_indices.extend( + range(flattened_index - draft_len + 1, flattened_index + 1) + ) + start = prev_index * self.prev_num_spec_tokens + # prev_draft_token_indices is used to find which draft_tokens_id + # should be copied to input_ids + # example: prev draft_tokens_id [[1,2], [3,4], [5, 6]] + # flatten draft_tokens_id [1,2,3,4,5,6] + # draft_len of each request [1, 2, 1] + # then prev_draft_token_indices is [0, 2, 3, 4] + prev_draft_token_indices.extend(range(start, start + draft_len)) + common_indices_match &= prev_index == flattened_index + max_flattened_index = max(max_flattened_index, flattened_index) + + num_common_tokens = len(sample_flattened_indices) + total_without_spec = total_num_scheduled_tokens - total_num_spec_tokens + if self.enable_prompt_embeds: + # The multimodal embed path reads is_token_ids.gpu; its .cpu copy is + # refreshed every step but the async fast paths below only scatter + # input_ids.gpu, so refresh is_token_ids.gpu here too. + self.is_token_ids.copy_to_gpu(total_num_scheduled_tokens) + if num_common_tokens < total_without_spec: + # If not all requests are decodes from the last iteration, + # we need to copy the input_ids_cpu to the GPU first. + self.input_ids.copy_to_gpu(total_num_scheduled_tokens) + if self.enable_prompt_embeds: + self.inputs_embeds.copy_to_gpu(total_num_scheduled_tokens) + if num_common_tokens == 0: + # No requests in common with the previous iteration + # So input_ids.cpu will have all the input ids. + return + if common_indices_match and max_flattened_index == (num_common_tokens - 1): + # Common-case optimization: the batch is unchanged + # and no reordering happened. + # The indices are both the same permutation of 0..N-1 so + # we can copy directly using a single slice. + self.input_ids.gpu[:num_common_tokens].copy_( + self.input_batch.prev_sampled_token_ids[:num_common_tokens, 0], + non_blocking=True, + ) + return + # Upload the index tensors asynchronously so the scatter can be non-blocking. + sampled_tokens_index_tensor = torch.tensor( + sample_flattened_indices, dtype=torch.int64, pin_memory=PIN_MEMORY + ).to(self.device, non_blocking=True) + prev_common_req_indices_tensor = torch.tensor( + prev_indices, dtype=torch.int64, pin_memory=PIN_MEMORY + ).to(self.device, non_blocking=True) + self.input_ids.gpu.scatter_( + dim=0, + index=sampled_tokens_index_tensor, + src=self.input_batch.prev_sampled_token_ids[ + prev_common_req_indices_tensor, 0 + ], + ) + + # Scatter the draft tokens after the sampled tokens are scattered. + if self._draft_token_ids is None or not spec_flattened_indices: + return + + assert isinstance(self._draft_token_ids, torch.Tensor) + draft_tokens_index_tensor = torch.tensor( + spec_flattened_indices, dtype=torch.int64, pin_memory=PIN_MEMORY + ).to(self.device, non_blocking=True) + prev_draft_token_indices_tensor = torch.tensor( + prev_draft_token_indices, dtype=torch.int64, pin_memory=PIN_MEMORY + ).to(self.device, non_blocking=True) + + # because input_ids dtype is torch.int32, + # so convert draft_token_ids to torch.int32 here. + draft_token_ids = self._draft_token_ids.to(dtype=torch.int32) + + self.input_ids.gpu.scatter_( + dim=0, + index=draft_tokens_index_tensor, + src=draft_token_ids.flatten()[prev_draft_token_indices_tensor], + ) + + def _get_encoder_seq_lens( + self, + num_scheduled_tokens: dict[str, int], + kv_cache_spec: KVCacheSpec, + num_reqs: int, + for_cudagraph_capture: bool = False, + ) -> tuple[torch.Tensor | None, np.ndarray | None]: + if not isinstance(kv_cache_spec, CrossAttentionSpec): + return None, None + + # Zero out buffer for padding requests that are not actually scheduled (CGs) + self.encoder_seq_lens.np[:num_reqs] = 0 + + # Build encoder_seq_lens array mapping request indices to + # encoder lengths for inputs scheduled in this batch + for req_id in num_scheduled_tokens: + req_index = self.input_batch.req_id_to_index[req_id] + req_state = self.requests[req_id] + if req_state.mm_features is None: + self.encoder_seq_lens.np[req_index] = 0 + continue + + # Get the total number of encoder input tokens for running encoder requests + # whether encoding is finished or not so that cross-attention knows how + # many encoder tokens to attend to. + encoder_input_tokens = sum( + feature.mm_position.length for feature in req_state.mm_features + ) + self.encoder_seq_lens.np[req_index] = encoder_input_tokens + if for_cudagraph_capture: + # During CUDA graph capture, we need to use realistic encoder lengths + # so that max_seqlen_k is captured with the correct value. + max_encoder_len = getattr( + self.model_config.hf_config, + "max_source_positions", + self.max_encoder_len, + ) + self.encoder_seq_lens.np[:num_reqs] = max_encoder_len + + self.encoder_seq_lens.copy_to_gpu(num_reqs) + encoder_seq_lens = self.encoder_seq_lens.gpu[:num_reqs] + encoder_seq_lens_cpu = self.encoder_seq_lens.np[:num_reqs] + + return encoder_seq_lens, encoder_seq_lens_cpu + + def _prepare_inputs( + self, + scheduler_output: "SchedulerOutput", + num_scheduled_tokens: np.ndarray, + ) -> tuple[ + torch.Tensor, + SpecDecodeMetadata | None, + int, + ]: + """ + Returns: + tuple[logits_indices, spec_decode_metadata, max_num_sampled_tokens] + """ + total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens + assert total_num_scheduled_tokens > 0 + num_reqs = self.input_batch.num_reqs + assert num_reqs > 0 + + # OPTIMIZATION: Start copying the block table first. + # This way, we can overlap the copy with the following CPU operations. + self.input_batch.block_table.commit_block_table(num_reqs) + + # Get request indices. + # E.g., [2, 5, 3] -> [0, 0, 1, 1, 1, 1, 1, 2, 2, 2] + req_indices = np.repeat(self.arange_np[:num_reqs], num_scheduled_tokens) + + # cu_num_tokens: [2, 5, 3] -> [2, 7, 10] + # self.query_pos.np[:10]: [0, 1, 0, 1, 2, 3, 4, 0, 1, 2] + cu_num_tokens = self._get_cumsum_and_arange( + num_scheduled_tokens, self.query_pos.np + ) + + # Get positions. + positions_np = ( + self.input_batch.num_computed_tokens_cpu[req_indices] + + self.query_pos.np[: cu_num_tokens[-1]] + ) + + # Calculate M-RoPE positions. + # Only relevant for models using M-RoPE (e.g, Qwen2-VL) + if self.uses_mrope: + self._calc_mrope_positions(scheduler_output) + + # Calculate XD-RoPE positions. + # Only relevant for models using XD-RoPE (e.g, HunYuan-VL) + if self.uses_xdrope_dim > 0: + self._calc_xdrope_positions(scheduler_output) + + # Get token indices. + # E.g., [0, 1, 0, 1, 2, 3, 4, 0, 1, 2] + # -> [0, 1, M, M + 1, M + 2, M + 3, M + 4, 2 * M, 2 * M + 1, 2 * M + 2] + # where M is the max_model_len. + token_indices = ( + positions_np + req_indices * self.input_batch.token_ids_cpu.shape[1] + ) + token_indices_tensor = torch.from_numpy(token_indices) + + # NOTE(woosuk): We use torch.index_select instead of np.take here + # because torch.index_select is much faster than np.take for large + # tensors. + torch.index_select( + self.input_batch.token_ids_cpu_tensor.flatten(), + 0, + token_indices_tensor, + out=self.input_ids.cpu[:total_num_scheduled_tokens], + ) + if self.enable_prompt_embeds: + is_token_ids = self.input_batch.is_token_ids_tensor.flatten() + torch.index_select( + is_token_ids, + 0, + token_indices_tensor, + out=self.is_token_ids.cpu[:total_num_scheduled_tokens], + ) + + # Because we did not pre-allocate a massive prompt_embeds CPU tensor on + # the InputBatch, we need to fill in the prompt embeds into the expected + # spots in the GpuModelRunner's pre-allocated prompt_embeds tensor. + if self.input_batch.req_prompt_embeds: + output_idx = 0 + for req_idx in range(num_reqs): + num_sched = num_scheduled_tokens[req_idx] + + # Skip if this request doesn't have embeddings + if req_idx not in self.input_batch.req_prompt_embeds: + output_idx += num_sched + continue + + # Skip if no tokens scheduled + if num_sched <= 0: + output_idx += num_sched + continue + + req_embeds = self.input_batch.req_prompt_embeds[req_idx] + start_pos = self.input_batch.num_computed_tokens_cpu[req_idx] + + # Skip if trying to read beyond available embeddings + if start_pos >= req_embeds.shape[0]: + output_idx += num_sched + continue + + # Copy available embeddings + end_pos = start_pos + num_sched + actual_end = min(end_pos, req_embeds.shape[0]) + actual_num_sched = actual_end - start_pos + + if actual_num_sched > 0: + self.inputs_embeds.cpu[ + output_idx : output_idx + actual_num_sched + ].copy_(req_embeds[start_pos:actual_end]) + + output_idx += num_sched + + # Prepare the attention metadata. + self.query_start_loc.np[0] = 0 + self.query_start_loc.np[1 : num_reqs + 1] = cu_num_tokens + # Note: pad query_start_loc to be non-decreasing, as kernels + # like FlashAttention requires that + self.query_start_loc.np[num_reqs + 1 :].fill(cu_num_tokens[-1]) + self.query_start_loc.copy_to_gpu() + query_start_loc = self.query_start_loc.gpu[: num_reqs + 1] + + # Compute optimistic seq_lens (assumes all draft tokens from previous + # iteration accepted). Store in optimistic_seq_lens_cpu for use by + # _build_attention_metadata (max_seq_len) and discard_request_mask. + # seq_lens (GPU) will be computed later using the same optimistic values. + torch.add( + self.input_batch.num_computed_tokens_cpu_tensor[:num_reqs], + torch.from_numpy(num_scheduled_tokens), + out=self.optimistic_seq_lens_cpu[:num_reqs], + ) + self.optimistic_seq_lens_cpu[num_reqs:].fill_(0) + + # Build prev_positions mapping: current pos -> prev pos (-1 if new). + # Used for gathering from previous iteration's GPU tensors. + prev_req_id_to_index = self.input_batch.prev_req_id_to_index + self._compute_prev_positions(num_reqs) + + num_tokens = [self.requests[r].num_tokens for r in self.input_batch.req_ids] + num_tokens_np = np.array(num_tokens, dtype=np.int32) + + # Record which requests should not be sampled, + # so that we could clear the sampled tokens before returning + self.discard_request_mask.np[:num_reqs] = ( + self.optimistic_seq_lens_cpu[:num_reqs].numpy() < num_tokens_np + ) + self.discard_request_mask.copy_to_gpu(num_reqs) + + # Sync num_accepted_tokens from CPU (set by + # _update_states_after_model_execute for hybrid models). + # Skipped under async scheduling (non-align): the CPU copy races with + # the in-flight D2H copy and with input-batch row moves. + needs_cpu_accepted_counts = self.num_accepted_tokens_event is not None and not ( + self.use_async_scheduling and self.cache_config.mamba_cache_mode != "align" + ) + if needs_cpu_accepted_counts: + assert self.num_accepted_tokens_event is not None + self.num_accepted_tokens_event.synchronize() + # Async mode: condense() reordered indices, use prev_positions mapping + if self.use_async_scheduling and prev_req_id_to_index: + prev_idx = self.prev_positions.np[:num_reqs] + new_mask = prev_idx < 0 + self.num_accepted_tokens.np[:num_reqs] = ( + self.input_batch.num_accepted_tokens_cpu[ + np.where(new_mask, 0, prev_idx) + ] + ) + self.num_accepted_tokens.np[:num_reqs][new_mask] = 1 + self.input_batch.num_accepted_tokens_cpu[:num_reqs] = ( + self.num_accepted_tokens.np[:num_reqs] + ) + else: + # Non-async mode: use values directly + self.num_accepted_tokens.np[:num_reqs] = ( + self.input_batch.num_accepted_tokens_cpu[:num_reqs] + ) + self.num_accepted_tokens.np[num_reqs:].fill(1) + self.num_accepted_tokens.copy_to_gpu() + else: + # Default to 1; update_num_computed_tokens_for_batch_change below + # corrects rows that had drafts from valid_sampled_token_count. + self.num_accepted_tokens.np.fill(1) + self.num_accepted_tokens.gpu.fill_(1) + + if self.mamba_prev_last_scheduled_idx is not None: + mamba_utils.preprocess_mamba_all_specdec( + scheduler_output, + self.input_batch, + self.mamba_state_idx, + num_reqs, + self.mamba_prev_last_scheduled_idx, + ) + + # Update num_computed_tokens on GPU. In async spec decode, + # CPU values are optimistic (all drafts accepted). The kernel + # corrects on GPU using the previous step's + # valid_sampled_token_count_gpu. Otherwise, just copy from CPU. + if ( + self.use_async_spec_decode + and self.valid_sampled_token_count_gpu is not None + and prev_req_id_to_index + ): + self.prev_positions.copy_to_gpu(num_reqs) + self.prev_num_draft_tokens.copy_to_gpu() + cpu_values = self.input_batch.num_computed_tokens_cpu_tensor[:num_reqs].to( + device=self.device, non_blocking=True + ) + update_num_computed_tokens_for_batch_change( + self.num_computed_tokens, + self.num_accepted_tokens.gpu[:num_reqs], + self.prev_positions.gpu[:num_reqs], + self.valid_sampled_token_count_gpu, + self.prev_num_draft_tokens.gpu, + cpu_values, + ) + else: + self.num_computed_tokens[:num_reqs].copy_( + self.input_batch.num_computed_tokens_cpu_tensor[:num_reqs], + non_blocking=True, + ) + + self.req_indices.np[:total_num_scheduled_tokens] = req_indices + self.req_indices.copy_to_gpu(total_num_scheduled_tokens) + req_indices_gpu = self.req_indices.gpu[:total_num_scheduled_tokens] + + self.query_pos.copy_to_gpu(total_num_scheduled_tokens) + self.num_scheduled_tokens.np[:num_reqs] = num_scheduled_tokens + self.num_scheduled_tokens.copy_to_gpu(num_reqs) + num_scheduled_tokens_gpu = self.num_scheduled_tokens.gpu[:num_reqs] + self.positions[:total_num_scheduled_tokens] = ( + self.num_computed_tokens[req_indices_gpu].to(torch.int64) + + self.query_pos.gpu[:total_num_scheduled_tokens] + ) + self.seq_lens[:num_reqs] = ( + self.num_computed_tokens[:num_reqs] + num_scheduled_tokens_gpu + ) + self.seq_lens[num_reqs:].fill_(0) + + self.input_batch.block_table.compute_slot_mapping( + num_reqs, + self.query_start_loc.gpu[: num_reqs + 1], + self.positions[:total_num_scheduled_tokens], + ) + + # Copy the tensors to the GPU. + self._prepare_input_ids( + scheduler_output, + num_reqs, + total_num_scheduled_tokens, + cu_num_tokens, + ) + + if self.uses_mrope: + # Only relevant for models using M-RoPE (e.g, Qwen2-VL) + # Copy one row at a time. mrope_positions is allocated as + # [3, max_num_tokens + 1] with a dummy trailing column to keep it + # non-contiguous for torch.compile, so cpu[:, :N] is a strided view. + # copy_() cannot express a strided source as a single + # cudaMemcpyAsync, so it first gathers into a contiguous *pageable* + # temporary, and a pageable H2D ignores non_blocking=True and + # synchronizes the stream before the transfer starts. Each row is + # contiguous within the pinned allocation, so per-row copies stay on + # the pinned path and are genuinely asynchronous. + for row in range(self.mrope_positions.gpu.shape[0]): + self.mrope_positions.gpu[row, :total_num_scheduled_tokens].copy_( + self.mrope_positions.cpu[row, :total_num_scheduled_tokens], + non_blocking=True, + ) + elif self.uses_xdrope_dim > 0: + # Only relevant for models using XD-RoPE (e.g, HunYuan-VL) + self.xdrope_positions.gpu[:, :total_num_scheduled_tokens].copy_( + self.xdrope_positions.cpu[:, :total_num_scheduled_tokens], + non_blocking=True, + ) + if self.use_async_spec_decode and (self.uses_mrope or self.uses_xdrope_dim > 0): + drift = self.num_computed_tokens[req_indices_gpu].to( + torch.int64 + ) - self.input_batch.num_computed_tokens_cpu_tensor[req_indices].to( + device=self.device, dtype=torch.int64, non_blocking=True + ) + target = self.mrope_positions if self.uses_mrope else self.xdrope_positions + target.gpu[:, :total_num_scheduled_tokens] += drift + + use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0 + if not use_spec_decode: + # NOTE(woosuk): Due to chunked prefills, the batch may contain + # partial requests. While we should not sample any token + # from these partial requests, we do so for simplicity. + # We will ignore the sampled tokens from the partial requests. + # TODO: Support prompt logprobs. + logits_indices = query_start_loc[1:] - 1 + spec_decode_metadata = None + num_sampled_tokens = np.ones(num_reqs, dtype=np.int32) + else: + # Get the number of draft tokens for each request. + # Iterate over the dictionary rather than all requests since not all + # requests have draft tokens. + num_draft_tokens = np.zeros(num_reqs, dtype=np.int32) + # For chunked prefills, use -1 as mask rather than 0, as guided + # decoding may rollback speculative tokens. + num_decode_draft_tokens = np.full(num_reqs, -1, dtype=np.int32) + for ( + req_id, + draft_token_ids, + ) in scheduler_output.scheduled_spec_decode_tokens.items(): + req_idx = self.input_batch.req_id_to_index[req_id] + draft_len = len(draft_token_ids) + num_draft_tokens[req_idx] = draft_len + if num_scheduled_tokens[req_idx] == draft_len + 1: + num_decode_draft_tokens[req_idx] = draft_len + spec_decode_metadata = self._calc_spec_decode_metadata( + num_draft_tokens, cu_num_tokens + ) + logits_indices = spec_decode_metadata.logits_indices + num_sampled_tokens = num_draft_tokens + 1 + # For DECODE only cuda graph of some attention backends (e.g., GDN). + self.num_decode_draft_tokens.np[:num_reqs] = num_decode_draft_tokens + self.num_decode_draft_tokens.np[num_reqs:].fill(-1) + self.num_decode_draft_tokens.copy_to_gpu() + + # Hot-Swap lora model + if self.lora_config: + assert ( + np.sum(num_sampled_tokens) + <= self.vllm_config.scheduler_config.max_num_batched_tokens + ) + self.set_active_loras( + self.input_batch, num_scheduled_tokens, num_sampled_tokens + ) + + return ( + logits_indices, + spec_decode_metadata, + int(num_sampled_tokens.max()), + ) + + def _build_attention_metadata( + self, + num_tokens: int, + num_reqs: int, + max_query_len: int, + num_tokens_padded: int | None = None, + num_reqs_padded: int | None = None, + ubatch_slices: UBatchSlices | None = None, + logits_indices: torch.Tensor | None = None, + max_num_sampled_tokens: int | None = None, + use_spec_decode: bool = False, + for_cudagraph_capture: bool = False, + num_scheduled_tokens: dict[str, int] | None = None, + cascade_attn_prefix_lens: list[list[int]] | None = None, + slot_mappings: dict[int, torch.Tensor] | None = None, + recurrent_prefill_checkpoint_plans: ( + dict[str, tuple[int, int, tuple[int, ...]]] | None + ) = None, + ) -> tuple[PerLayerAttnMetadata, CommonAttentionMetadata | None]: + """ + Returns: + tuple[attn_metadata, spec_decode_common_attn_metadata] + """ + # Attention metadata is not needed for attention free models + if len(self.kv_cache_config.kv_cache_groups) == 0: + return {}, None + + num_tokens_padded = num_tokens_padded or num_tokens + num_reqs_padded = num_reqs_padded or num_reqs + assert num_reqs_padded is not None and num_tokens_padded is not None + + attn_metadata: PerLayerAttnMetadata = {} + if ubatch_slices is not None: + attn_metadata = [dict() for _ in range(len(ubatch_slices))] + + if for_cudagraph_capture: + # For some attention backends (e.g. FA) with sliding window models we need + # to make sure the backend see a max_seq_len that is larger to the sliding + # window size when capturing to make sure the correct kernel is selected. + max_seq_len = self.max_model_len + else: + max_seq_len = self.optimistic_seq_lens_cpu.numpy()[:num_reqs].max().item() + + kv_cache_groups = self.kv_cache_config.kv_cache_groups + + def _get_block_table(kv_cache_gid: int): + assert num_reqs_padded is not None and num_tokens_padded is not None + kv_cache_spec = kv_cache_groups[kv_cache_gid].kv_cache_spec + if isinstance(kv_cache_spec, EncoderOnlyAttentionSpec): + blk_table_tensor = torch.zeros( + (num_reqs_padded, 1), + dtype=torch.int32, + device=self.device, + ) + else: + blk_table = self.input_batch.block_table[kv_cache_gid] + blk_table_tensor = blk_table.get_device_tensor(num_reqs_padded) + + # Fill unused block table entries with NULL_BLOCK_ID (null block) + # for CUDAGraph padding. Block 0 is reserved for padding. + blk_table_tensor[num_reqs:num_reqs_padded].fill_(NULL_BLOCK_ID) + return blk_table_tensor + + assert slot_mappings is not None + block_table_gid_0 = _get_block_table(0) + slot_mapping_gid_0 = slot_mappings[0] + + if self.routed_experts_initialized: + # Copy this step's attention slot_mapping into our private + # device buffer. The shared ``slot_mappings[attn_gid]`` is + # owned by the attention block table and will be overwritten + # by the next ``_prepare_inputs``; we need a stable snapshot + # because the async D2H may still be in flight on the copy + # stream when the next step runs. + slot_mapping_attn = slot_mappings[self.routed_experts_capturer.attn_gid] + self.routed_experts_slot_mapping_device[:num_tokens].copy_( + slot_mapping_attn[:num_tokens] + ) + + num_computed_tokens_cpu = self.input_batch.num_computed_tokens_cpu_tensor[ + :num_reqs_padded + ] + num_prompt_tokens_cpu = self.input_batch.num_prompt_tokens_cpu_tensor[ + :num_reqs_padded + ] + seq_lens_cpu = self.optimistic_seq_lens_cpu[:num_reqs_padded] + seq_lens_cpu_upper_bound = seq_lens_cpu + + # is_prefilling: True if request is still in prefill phase. + # Used by mamba backends to distinguish actual decodes from + # short extends. + is_prefilling = num_computed_tokens_cpu < num_prompt_tokens_cpu + # Zero out padded rows so stale data from condense() doesn't + # misclassify padding as prefill in CUDA graph mode. + is_prefilling[num_reqs:] = False + + if self.use_async_spec_decode: + # GPU tensors are authoritative in async mode. + seq_lens_cpu = None + num_computed_tokens_cpu = None + + # Compute mm_prefix bidirectional ranges before building + # attention metadata so builders handle them during build(). + # By default, ranges exceeding sliding_window are skipped to prevent + # early tokens from attending across the entire image span. Models that + # clamp mm_prefix to the sliding window *in-kernel* (e.g. Gemma4, which + # needs HF's (causal OR blockwise) AND sliding_window on sliding layers) + # opt out of the skip so the bidirectional range survives for images + # larger than the window; the kernel then bounds it per-query. + req_doc_ranges: dict[int, list[tuple[int, int]]] | None = None + if self.is_mm_prefix_lm: + req_doc_ranges = {} + hf_text_config = self.model_config.hf_text_config + _bidi_sw = getattr(hf_text_config, "sliding_window", None) + _clamps_in_kernel = getattr( + self.model, "mm_prefix_clamp_sliding_window", False + ) + for req_id in self.input_batch.req_ids: + image_doc_ranges = [] + req_state = self.requests[req_id] + for mm_feature in req_state.mm_features: + if mm_feature.modality == "audio": + continue + pos_info = mm_feature.mm_position + img_doc_range = pos_info.extract_embeds_range() + for r in img_doc_range: + if ( + not _clamps_in_kernel + and _bidi_sw is not None + and (r[1] - r[0] + 1) > _bidi_sw + ): + continue + image_doc_ranges.append(r) + req_idx = self.input_batch.req_id_to_index[req_id] + req_doc_ranges[req_idx] = image_doc_ranges + + # Reference Sliding Window Attention (R-SWA): pass per-request prompt + # lengths so the attention backend can keep the prefix globally visible. + # The backend owns the persistent CUDA-graph-safe GPU buffer. + rswa_prefix_lens = None + if self.model_config.rswa_window is not None: + rswa_prefix_lens = num_prompt_tokens_cpu + + replayssm_decode_base_cpu = None + if self.cache_config.use_replayssm: + replayssm_decode_base_cpu = ( + self.input_batch.replayssm_decode_base_cpu_tensor[:num_reqs_padded] + ) + + checkpoint_plans_cpu = None + if recurrent_prefill_checkpoint_plans: + if for_cudagraph_capture or ubatch_slices is not None: + raise ValueError("Planned recurrent checkpoints require unsliced eager metadata") + checkpoint_plans_cpu = [None] * num_reqs_padded + for req_id, plan in recurrent_prefill_checkpoint_plans.items(): + row = self.input_batch.req_id_to_index.get(req_id) + if row is None or not 0 <= row < num_reqs: + raise ValueError("Checkpoint plan request is absent from the active batch") + start, end, targets = plan + if ( + start < 0 or end <= start or len(targets) > 2 + or tuple(sorted(set(targets))) != targets + or any(not start < target < end for target in targets) + or num_scheduled_tokens is None + or num_scheduled_tokens.get(req_id) != end - start + ): + raise ValueError("Checkpoint plan does not match the scheduled query") + checkpoint_plans_cpu[row] = plan + + cm_base = CommonAttentionMetadata( + query_start_loc=self.query_start_loc.gpu[: num_reqs_padded + 1], + query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs_padded + 1], + seq_lens=self.seq_lens[:num_reqs_padded], + _seq_lens_cpu=seq_lens_cpu, + _num_computed_tokens_cpu=num_computed_tokens_cpu, + seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, + replayssm_decode_base_cpu=replayssm_decode_base_cpu, + recurrent_prefill_checkpoint_plans_cpu=checkpoint_plans_cpu, + num_reqs=num_reqs_padded, + num_actual_tokens=num_tokens_padded, + max_query_len=max_query_len, + max_seq_len=max_seq_len, + block_table_tensor=block_table_gid_0, + slot_mapping=slot_mapping_gid_0, + causal=True, + is_prefilling=is_prefilling, + positions=self.positions[:num_tokens_padded], + mm_req_doc_ranges=req_doc_ranges, + rswa_prefix_lens=rswa_prefix_lens, + ) + + if self.dcp_world_size > 1: + self.dcp_local_seq_lens.cpu[:num_reqs] = get_dcp_local_seq_lens( + self.optimistic_seq_lens_cpu[:num_reqs], + self.dcp_world_size, + self.dcp_rank, + self.parallel_config.cp_kv_cache_interleave_size, + ) + self.dcp_local_seq_lens.cpu[num_reqs:].fill_(0) + self.dcp_local_seq_lens.copy_to_gpu(num_reqs_padded) + + cm_base.dcp_local_seq_lens = self.dcp_local_seq_lens.gpu[:num_reqs_padded] + cm_base.dcp_local_seq_lens_cpu = self.dcp_local_seq_lens.cpu[ + :num_reqs_padded + ] + + if logits_indices is not None and self.cache_config.kv_sharing_fast_prefill: + cm_base.num_logits_indices = logits_indices.size(0) + cm_base.max_logits_per_req = max_num_sampled_tokens + cm_base.logits_indices_padded = self._prepare_kv_sharing_fast_prefill( + logits_indices + ) + + # Cache attention metadata builds across hybrid KV-cache groups + # The only thing that changes between different hybrid KV-cache groups when the + # same metadata builder and KVCacheSpec is the same is the block table, so we + # can cache the attention metadata builds and just update the block table using + # `builder.update_block_table` if the builder supports it. + cached_attn_metadata: dict[ + tuple[KVCacheSpec, type[AttentionMetadataBuilder]], AttentionMetadata + ] = {} + + def _build_attn_group_metadata( + kv_cache_gid: int, + attn_gid: int, + common_attn_metadata: CommonAttentionMetadata, + ubid: int | None = None, + ) -> None: + attn_group = self.attn_groups[kv_cache_gid][attn_gid] + builder = attn_group.get_metadata_builder(ubid or 0) + kv_cache_spec = kv_cache_groups[kv_cache_gid].kv_cache_spec + if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs): + kv_cache_spec = kv_cache_spec.kv_cache_specs[attn_group.layer_names[0]] + cache_key = (kv_cache_spec, type(builder)) + + cascade_attn_prefix_len = ( + cascade_attn_prefix_lens[kv_cache_gid][attn_gid] + if cascade_attn_prefix_lens + else 0 + ) + + extra_attn_metadata_args = {} + if use_spec_decode and isinstance( + builder, + ( + Mamba2AttentionMetadataBuilder, + GDNAttentionMetadataBuilder, + BailingLinearAttentionMetadataBuilder, + ShortConvAttentionMetadataBuilder, + ), + ): + assert ubid is None, ( + "UBatching not supported with GDN or linear attn yet" + ) + extra_attn_metadata_args = dict( + num_accepted_tokens=self.num_accepted_tokens.gpu[:num_reqs_padded], + num_decode_draft_tokens_cpu=self.num_decode_draft_tokens.cpu[ + :num_reqs_padded + ], + ) + if ( + isinstance(builder, Mamba2AttentionMetadataBuilder) + and self.mamba_prev_last_scheduled_idx is not None + ): + extra_attn_metadata_args["prev_last_scheduled_idx"] = ( + self.mamba_prev_last_scheduled_idx.gpu[:num_reqs_padded] + ) + + if for_cudagraph_capture: + attn_metadata_i = builder.build_for_cudagraph_capture( + common_attn_metadata + ) + elif ( + cache_key in cached_attn_metadata + and builder.supports_update_block_table + ): + attn_metadata_i = builder.update_block_table( + cached_attn_metadata[cache_key], + common_attn_metadata.block_table_tensor, + common_attn_metadata.slot_mapping, + ) + else: + attn_metadata_i = builder.build( + common_prefix_len=cascade_attn_prefix_len, + common_attn_metadata=common_attn_metadata, + **extra_attn_metadata_args, + ) + if builder.supports_update_block_table: + cached_attn_metadata[cache_key] = attn_metadata_i + + if ubid is None: + assert isinstance(attn_metadata, dict) + attn_metadata_dict = attn_metadata + else: + assert isinstance(attn_metadata, list) + attn_metadata_dict = attn_metadata[ubid] + + for layer_name in attn_group.layer_names: + attn_metadata_dict[layer_name] = attn_metadata_i + + # Prepare the attention metadata for each KV cache group and make layers + # in the same group share the same metadata. + spec_decode_common_attn_metadata = None + for kv_cache_gid, kv_cache_group in enumerate(kv_cache_groups): + cm = copy(cm_base) # shallow copy + + # Basically only the encoder seq_lens, block_table and slot_mapping change + # for each kv_cache_group. + cm.encoder_seq_lens, cm.encoder_seq_lens_cpu = self._get_encoder_seq_lens( + num_scheduled_tokens or {}, + kv_cache_group.kv_cache_spec, + num_reqs_padded, + for_cudagraph_capture=for_cudagraph_capture, + ) + if kv_cache_gid > 0: + cm.block_table_tensor = _get_block_table(kv_cache_gid) + cm.slot_mapping = slot_mappings[kv_cache_gid] + + if self.speculative_config and spec_decode_common_attn_metadata is None: + if isinstance( + self.drafter, + ( + EagleProposer, + DFlashProposer, + Gemma4Proposer, + ExtractHiddenStatesProposer, + ), + ): + if self.drafter.kv_cache_gid == kv_cache_gid: + spec_decode_common_attn_metadata = cm + else: + spec_decode_common_attn_metadata = cm + # Capture per-group block tables for multi-group proposers. + if self.speculative_config and isinstance(self.drafter, Step3p5MTPProposer): + self.drafter.set_per_group_attn_metadata( + kv_cache_gid, cm.block_table_tensor, cm.slot_mapping + ) + elif self.speculative_config and isinstance(self.drafter, Gemma4Proposer): + self.drafter.set_per_group_block_table( + kv_cache_gid, cm.block_table_tensor + ) + + for attn_gid in range(len(self.attn_groups[kv_cache_gid])): + if ubatch_slices is not None: + for ubid, _cm in enumerate(split_attn_metadata(ubatch_slices, cm)): + _build_attn_group_metadata(kv_cache_gid, attn_gid, _cm, ubid) + + else: + _build_attn_group_metadata(kv_cache_gid, attn_gid, cm) + + if spec_decode_common_attn_metadata is not None and ( + num_reqs != num_reqs_padded or num_tokens != num_tokens_padded + ): + # Currently the drafter still only uses piecewise cudagraphs (and modifies + # the attention metadata in directly), and therefore does not want to use + # padded attention metadata. + spec_decode_common_attn_metadata = ( + spec_decode_common_attn_metadata.unpadded(num_tokens, num_reqs) + ) + + return attn_metadata, spec_decode_common_attn_metadata + + def _compute_cascade_attn_prefix_lens( + self, + num_scheduled_tokens: np.ndarray, + num_computed_tokens: np.ndarray, + num_common_prefix_blocks: list[int], + ) -> list[list[int]] | None: + """ + Returns: + Optional[cascade_attn_prefix_lens] + cascade_attn_prefix_lens is 2D: + ``[kv_cache_group_id][attn_group_idx]``, + None if we should not use cascade attention + """ + + use_cascade_attn = False + num_kv_cache_groups = len(self.kv_cache_config.kv_cache_groups) + cascade_attn_prefix_lens: list[list[int]] = [ + [] for _ in range(num_kv_cache_groups) + ] + + for kv_cache_gid in range(num_kv_cache_groups): + for attn_group in self.attn_groups[kv_cache_gid]: + if isinstance(attn_group.kv_cache_spec, EncoderOnlyAttentionSpec): + cascade_attn_prefix_len = 0 + else: + # 0 if cascade attention should not be used + cascade_attn_prefix_len = self._compute_cascade_attn_prefix_len( + num_scheduled_tokens, + num_computed_tokens, + num_common_prefix_blocks[kv_cache_gid], + attn_group.kv_cache_spec, + attn_group.get_metadata_builder(), + ) + cascade_attn_prefix_lens[kv_cache_gid].append(cascade_attn_prefix_len) + use_cascade_attn |= cascade_attn_prefix_len > 0 + + return cascade_attn_prefix_lens if use_cascade_attn else None + + def _compute_cascade_attn_prefix_len( + self, + num_scheduled_tokens: np.ndarray, + num_computed_tokens: np.ndarray, + num_common_prefix_blocks: int, + kv_cache_spec: KVCacheSpec, + attn_metadata_builder: AttentionMetadataBuilder, + ) -> int: + """Compute the length of the common prefix for cascade attention. + + NOTE(woosuk): The common prefix length returned by this function + represents the length used specifically for cascade attention, not the + actual number of tokens shared between requests. When cascade attention + is disabled (use_cascade=False), this function returns 0 even if + requests share common tokens. Additionally, the common prefix length is + truncated to a multiple of the block size and may be further truncated + due to implementation details explained below. + + Args: + num_scheduled_tokens: Number of tokens scheduled per request. + num_common_prefix_blocks: Number of shared KV cache blocks. + + Returns: + int: Length of common prefix in tokens. + """ + + common_prefix_len = num_common_prefix_blocks * kv_cache_spec.block_size + if common_prefix_len == 0: + # Common case. + return 0 + + # NOTE(woosuk): Cascade attention uses two attention kernels: one + # for the common prefix and the other for the rest. For the first + # kernel, we concatenate all the query tokens (possibly from + # different requests) and treat them as if they are from the same + # request. Then, we use bi-directional attention to process the + # common prefix in the KV cache. Importantly, this means that the + # first kernel does not do any masking. + + # Consider the following example: + # Request 1's input query: [D, E, X] + # Request 1's kv cache: [A, B, C, D, E, X] + # Request 1's num_computed_tokens: 3 (i.e., [A, B, C]) + # Request 2's input query: [E, Y] + # Request 2's kv cache: [A, B, C, D, E, Y] + # Request 2's num_computed_tokens: 4 (i.e., [A, B, C, D]) + + # If we use [A, B, C, D, E] as the common prefix, then the + # first kernel will compute the bi-directional attention between + # input query [D, E, X, E, Y] and common prefix [A, B, C, D, E]. + # However, this is wrong because D in Request 1 should not attend to + # E in the common prefix (i.e., we need masking). + # To avoid this, [A, B, C, D] should be the common prefix. + # That is, the common prefix should be capped by the minimum + # num_computed_tokens among the requests, and plus one to include + # the first token of the query. + + # In practice, we use [A, B, C] as the common prefix, instead of + # [A, B, C, D] (i.e., the common prefix is capped by the minimum + # num_computed_tokens, without plus one). + # This is because of an implementation detail: We want to always + # use two kernels for cascade attention. Let's imagine: + # Request 3's input query: [D] + # Request 3's kv cache: [A, B, C, D] + # Request 3's num_computed_tokens: 3 (i.e., [A, B, C]) + # If we use [A, B, C, D] as the common prefix for Request 1-3, + # then Request 3 will be processed only by the first kernel, + # and the second kernel will get an empty input. While this is not + # a fundamental problem, our current implementation does not support + # this case. + common_prefix_len = min(common_prefix_len, num_computed_tokens.min()) + # common_prefix_len should be a multiple of the block size. + common_prefix_len = ( + common_prefix_len // kv_cache_spec.block_size * kv_cache_spec.block_size + ) + use_sliding_window = isinstance(kv_cache_spec, SlidingWindowSpec) or ( + isinstance(kv_cache_spec, FullAttentionSpec) + and kv_cache_spec.sliding_window is not None + ) + use_local_attention = isinstance(kv_cache_spec, ChunkedLocalAttentionSpec) or ( + isinstance(kv_cache_spec, FullAttentionSpec) + and kv_cache_spec.attention_chunk_size is not None + ) + assert isinstance(kv_cache_spec, AttentionSpec) + use_cascade = attn_metadata_builder.use_cascade_attention( + common_prefix_len=common_prefix_len, + query_lens=num_scheduled_tokens, + num_query_heads=self.num_query_heads, + num_kv_heads=kv_cache_spec.num_kv_heads, + use_alibi=self.use_alibi, + use_sliding_window=use_sliding_window, + use_local_attention=use_local_attention, + num_sms=self.num_sms, + dcp_world_size=self.dcp_world_size, + ) + return common_prefix_len if use_cascade else 0 + + def _calc_mrope_positions(self, scheduler_output: "SchedulerOutput"): + mrope_pos_ptr = 0 + for index, req_id in enumerate(self.input_batch.req_ids): + req = self.requests[req_id] + assert req.mrope_positions is not None + + num_computed_tokens = self.input_batch.num_computed_tokens_cpu[index] + num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id] + num_prompt_tokens = length_from_prompt_token_ids_or_embeds( + req.prompt_token_ids, req.prompt_embeds + ) + + if num_computed_tokens + num_scheduled_tokens > num_prompt_tokens: + prompt_part_len = max(0, num_prompt_tokens - num_computed_tokens) + completion_part_len = max(0, num_scheduled_tokens - prompt_part_len) + else: + prompt_part_len = num_scheduled_tokens + completion_part_len = 0 + + assert num_scheduled_tokens == prompt_part_len + completion_part_len + + if prompt_part_len > 0: + # prompt's mrope_positions are pre-computed + dst_start = mrope_pos_ptr + dst_end = mrope_pos_ptr + prompt_part_len + src_start = num_computed_tokens + src_end = num_computed_tokens + prompt_part_len + + self.mrope_positions.cpu[:, dst_start:dst_end] = req.mrope_positions[ + :, src_start:src_end + ] + mrope_pos_ptr += prompt_part_len + + if completion_part_len > 0: + # compute completion's mrope_positions on-the-fly + dst_start = mrope_pos_ptr + dst_end = mrope_pos_ptr + completion_part_len + + assert req.mrope_position_delta is not None + MRotaryEmbedding.get_next_input_positions_tensor( + out=self.mrope_positions.np, + out_offset=dst_start, + mrope_position_delta=req.mrope_position_delta, + context_len=num_computed_tokens + prompt_part_len, + num_new_tokens=completion_part_len, + ) + + mrope_pos_ptr += completion_part_len + + def _calc_xdrope_positions(self, scheduler_output: "SchedulerOutput"): + xdrope_pos_ptr = 0 + for index, req_id in enumerate(self.input_batch.req_ids): + req = self.requests[req_id] + assert req.xdrope_positions is not None + + num_computed_tokens = self.input_batch.num_computed_tokens_cpu[index] + num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id] + num_prompt_tokens = length_from_prompt_token_ids_or_embeds( + req.prompt_token_ids, req.prompt_embeds + ) + + if num_computed_tokens + num_scheduled_tokens > num_prompt_tokens: + prompt_part_len = max(0, num_prompt_tokens - num_computed_tokens) + completion_part_len = max(0, num_scheduled_tokens - prompt_part_len) + else: + prompt_part_len = num_scheduled_tokens + completion_part_len = 0 + + assert num_scheduled_tokens == prompt_part_len + completion_part_len + + if prompt_part_len > 0: + # prompt's xdrope_positions are pre-computed + dst_start = xdrope_pos_ptr + dst_end = xdrope_pos_ptr + prompt_part_len + src_start = num_computed_tokens + src_end = num_computed_tokens + prompt_part_len + + self.xdrope_positions.cpu[:, dst_start:dst_end] = req.xdrope_positions[ + :, src_start:src_end + ] + xdrope_pos_ptr += prompt_part_len + + if completion_part_len > 0: + # compute completion's xdrope_positions on-the-fly + dst_start = xdrope_pos_ptr + dst_end = xdrope_pos_ptr + completion_part_len + + XDRotaryEmbedding.get_next_input_positions_tensor( + out=self.xdrope_positions.np, + out_offset=dst_start, + context_len=num_computed_tokens + prompt_part_len, + num_new_tokens=completion_part_len, + ) + + xdrope_pos_ptr += completion_part_len + + def _calc_spec_decode_metadata( + self, + num_draft_tokens: np.ndarray, + cu_num_scheduled_tokens: np.ndarray, + ) -> SpecDecodeMetadata: + # Inputs: + # cu_num_scheduled_tokens: [ 4, 104, 107, 207, 209] + # num_draft_tokens: [ 3, 0, 2, 0, 1] + # Outputs: + # cu_num_draft_tokens: [ 3, 3, 5, 5, 6] + # logits_indices: [ 0, 1, 2, 3, 103, 104, 105, 106, + # 206, 207, 208] + # target_logits_indices: [ 0, 1, 2, 5, 6, 9] + # bonus_logits_indices: [ 3, 4, 7, 8, 10] + + # Compute the logits indices. + # [4, 1, 3, 1, 2] + num_sampled_tokens = num_draft_tokens + 1 + + # Step 1. + # cu_num_sampled_tokens: [4, 5, 8, 9, 11] + # _arange_scratch[:11]: [0, 1, 2, 3, 0, 0, 1, 2, 0, 0, 1] + cu_num_sampled_tokens = self._get_cumsum_and_arange( + num_sampled_tokens, self._arange_scratch, cumsum_dtype=np.int32 + ) + # Step 2. [0, 0, 0, 0, 103, 104, 104, 104, 206, 207, 207] + logits_indices = np.repeat( + cu_num_scheduled_tokens - num_sampled_tokens, num_sampled_tokens + ) + # Step 3. [0, 1, 2, 3, 103, 104, 105, 106, 206, 207, 208] + logits_indices += self._arange_scratch[: cu_num_sampled_tokens[-1]] + + # Compute the bonus logits indices. + bonus_logits_indices = cu_num_sampled_tokens - 1 + + # Compute the draft logits indices. + # cu_num_draft_tokens: [3, 3, 5, 5, 6] + # _arange_scratch[:6]: [0, 1, 2, 0, 1, 0] + cu_num_draft_tokens = self._get_cumsum_and_arange( + num_draft_tokens, self._arange_scratch, cumsum_dtype=np.int32 + ) + # [0, 0, 0, 5, 5, 9] + target_logits_indices = np.repeat( + cu_num_sampled_tokens - num_sampled_tokens, num_draft_tokens + ) + # [0, 1, 2, 5, 6, 9] + target_logits_indices += self._arange_scratch[: cu_num_draft_tokens[-1]] + + cu_num_draft_tokens = async_tensor_h2d(cu_num_draft_tokens, device=self.device) + cu_num_sampled_tokens = async_tensor_h2d( + cu_num_sampled_tokens, device=self.device + ) + logits_indices = async_tensor_h2d(logits_indices, device=self.device) + target_logits_indices = async_tensor_h2d( + target_logits_indices, device=self.device + ) + bonus_logits_indices = async_tensor_h2d( + bonus_logits_indices, device=self.device + ) + + # Compute the draft token ids. + # draft_token_indices: [ 1, 2, 3, 105, 106, 208] + draft_token_ids = self.input_ids.gpu[logits_indices] + draft_token_ids = draft_token_ids[target_logits_indices + 1] + + return SpecDecodeMetadata( + draft_token_ids=draft_token_ids, + num_draft_tokens=num_draft_tokens.tolist(), + cu_num_draft_tokens=cu_num_draft_tokens, + cu_num_sampled_tokens=cu_num_sampled_tokens, + target_logits_indices=target_logits_indices, + bonus_logits_indices=bonus_logits_indices, + logits_indices=logits_indices, + ) + + def _prepare_kv_sharing_fast_prefill( + self, + logits_indices: torch.Tensor, + ) -> torch.Tensor: + assert self.kv_sharing_fast_prefill_logits_indices is not None + num_logits = logits_indices.shape[0] + assert num_logits > 0 + self.kv_sharing_fast_prefill_logits_indices[:num_logits].copy_(logits_indices) + # There might have leftover indices in logits_indices[num_logits:] + # from previous iterations, whose values may be greater than the + # batch size in the current iteration. To ensure indices are always + # valid, fill the padded indices with the last index. Broadcast the + # scalar GPU-side to avoid a D2H sync on `.item()`. + self.kv_sharing_fast_prefill_logits_indices[num_logits:] = logits_indices[-1] + # Dispatch for the decoder portion of the model. + _, batch_desc = self.cudagraph_dispatcher.dispatch( + num_logits, invalid_modes={CUDAGraphMode.FULL} + ) + num_logits_padded = batch_desc.num_tokens + logits_indices_padded = self.kv_sharing_fast_prefill_logits_indices[ + :num_logits_padded + ] + return logits_indices_padded + + def _batch_mm_inputs_from_scheduler( + self, + scheduler_output: "SchedulerOutput", + ) -> tuple[ + list[str], + list[tuple[str, MultiModalKwargsItem]], + list[tuple[str, PlaceholderRange]], + ]: + """Batch multimodal inputs from scheduled encoder inputs. + + Args: + scheduler_output: The scheduler output containing scheduled encoder + inputs. + + Returns: + A tuple of (mm_hashes, mm_kwargs, mm_lora_refs) where: + - mm_hashes: List of multimodal hashes for each item + - mm_kwargs: List of multimodal kwargs for each item + - mm_lora_refs: List of (req_id, placeholder_range) for each item + """ + scheduled_encoder_inputs = scheduler_output.scheduled_encoder_inputs + if not scheduled_encoder_inputs: + return [], [], [] + + mm_hashes = list[str]() + mm_kwargs = list[tuple[str, MultiModalKwargsItem]]() + # Multimodal LoRA reference info to map each multimodal item + # back to its request & position + mm_lora_refs = list[tuple[str, PlaceholderRange]]() + for req_id, encoder_input_ids in scheduled_encoder_inputs.items(): + req_state = self.requests[req_id] + + for mm_input_id in encoder_input_ids: + mm_feature = req_state.mm_features[mm_input_id] + if mm_feature.data is None: + continue + + mm_hashes.append(mm_feature.identifier) + mm_kwargs.append((mm_feature.modality, mm_feature.data)) + mm_lora_refs.append((req_id, mm_feature.mm_position)) + + return mm_hashes, mm_kwargs, mm_lora_refs + + def _cache_encoder_output( + self, + mm_hash: str, + output: torch.Tensor, + ec_manager_metadata: "EncoderCacheManagerMetadata | None", + free_encoder_mm_hashes: list[str], + ) -> None: + """Store an encoder output for later multimodal embedding gather.""" + del ec_manager_metadata, free_encoder_mm_hashes + self.encoder_cache[mm_hash] = output + self.maybe_save_ec_to_connector(self.encoder_cache, mm_hash) + + def _execute_mm_encoder( + self, scheduler_output: "SchedulerOutput" + ) -> list[torch.Tensor]: + mm_hashes, mm_kwargs, mm_lora_refs = self._batch_mm_inputs_from_scheduler( + scheduler_output + ) + + if not mm_kwargs: + return [] + + # `prompt_embeds` is a passthrough modality, the tensor is already in + # the model embedding space, so no encoder runs. Inject each + # `prompt_embeds` tensor directly into the encoder cache here so that + # `_gather_mm_embeddings` can splice it via the standard `is_mm_embed` + # path. + pe_indices = [ + i + for i, (modality, _) in enumerate(mm_kwargs) + if modality == "prompt_embeds" + ] + if pe_indices: + for i in pe_indices: + pe_tensor = mm_kwargs[i][1]["embedding"].data + assert isinstance(pe_tensor, torch.Tensor) + + self._cache_encoder_output( + mm_hashes[i], + async_tensor_h2d(pe_tensor, device=self.device), + scheduler_output.ec_manager_metadata, + scheduler_output.free_encoder_mm_hashes, + ) + # Filter out `prompt_embeds` items from mm_kwargs/mm_hashes/mm_lora_refs + # since they don't require further encoder processing. + mm_hashes = [h for i, h in enumerate(mm_hashes) if i not in pe_indices] + mm_kwargs = [k for i, k in enumerate(mm_kwargs) if i not in pe_indices] + mm_lora_refs = [ + r for i, r in enumerate(mm_lora_refs) if i not in pe_indices + ] + if not mm_kwargs: + return [] # nothing left to encode after filtering out `prompt_embeds` + + should_time = bool( + self.observability_config + and self.observability_config.enable_mm_processor_stats + and scheduler_output.scheduled_encoder_inputs + ) + + # Batch mm inputs as much as we can: if a request in the batch has + # multiple modalities or a different modality than the previous one, + # we process it separately to preserve item order. + # FIXME(ywang96): This is a hacky way to deal with multiple modalities + # in the same batch while still being able to benefit from batching + # multimodal inputs. The proper solution should be reordering the + # encoder outputs. + model = cast(SupportsMultiModal, self.model) + + if self.lora_config and self.lora_manager.supports_tower_connector_lora(): + # Build LoRA mappings independently for encoder inputs + # (encoder batch structure is different from main batch) + prompt_lora_mapping = [] + token_lora_mapping = [] + lora_requests = set() + encoder_token_counts = [] + connector_token_counts = [] + + for (req_id, pos_info), (modality, mm_item) in zip( + mm_lora_refs, + mm_kwargs, + ): + req_idx = self.input_batch.req_id_to_index[req_id] + lora_id = int(self.input_batch.request_lora_mapping[req_idx]) + + tower_tokens, connector_tokens = self.model.get_mm_lora_token_counts( # type: ignore[attr-defined] + modality=modality, + mm_kwargs=mm_item, + num_mm_embeds=pos_info.get_num_embeds(), + ) + prompt_lora_mapping.append(lora_id) + token_lora_mapping.extend([lora_id] * tower_tokens) + encoder_token_counts.append(tower_tokens) + connector_token_counts.append(connector_tokens) + + if lora_id > 0: + lora_request = self.input_batch.lora_id_to_lora_request.get(lora_id) + if lora_request is not None: + lora_requests.add(lora_request) + + # Set tower adapter mapping + tower_mapping = LoRAMapping( + tuple(token_lora_mapping), + tuple(prompt_lora_mapping), + is_prefill=True, + type=LoRAMappingType.TOWER, + ) + self.lora_manager.set_active_adapters(lora_requests, tower_mapping) + + # Only set connector mapping if the model actually has a connector. + # Some multimodal models inherit a stub `get_num_mm_connector_tokens` + # from `SupportsMultiModal`, which returns None and should not be + # treated as a signal that connector LoRA is supported. + mm_mapping = ( + self.model.get_mm_mapping() # type: ignore[attr-defined] + if hasattr(self.model, "get_mm_mapping") + else None + ) + if ( + mm_mapping is not None + and mm_mapping.connector + and all(count is not None for count in connector_token_counts) + ): + connector_token_mapping = np.repeat( + np.array(prompt_lora_mapping, dtype=np.int32), + np.array(connector_token_counts, dtype=np.int32), + ) + connector_mapping = LoRAMapping( + index_mapping=tuple(connector_token_mapping.tolist()), + prompt_mapping=tuple(prompt_lora_mapping), + is_prefill=True, + type=LoRAMappingType.CONNECTOR, + ) + + self.lora_manager.set_active_adapters( + lora_requests, + connector_mapping, + ) + + encoder_outputs: list[torch.Tensor] = [] + # Track the current index in mm_kwargs/mm_lora_refs to map groups to request IDs + current_item_idx = 0 + for modality, num_items, mm_kwargs_batch in group_and_batch_mm_kwargs( + mm_kwargs, device=self.device, pin_memory=PIN_MEMORY + ): + batch_outputs: MultiModalEmbeddings + + # EVS and dynamic res video related change. + # (ekhvedchenia): Temporary hack to limit peak memory usage when + # processing multimodal data. This solves the issue with scheduler + # putting too many video samples into a single batch. Scheduler + # uses pruned vision tokens count to compare it versus compute + # budget which is incorrect (Either input media size or non-pruned + # output vision tokens count should be considered) + # dynamic res video for nemotron temporarily uses this hack via + # requires_sequential_video_encoding + # because it doesn't yet support video batching. + # TODO(ywang96): Fix memory profiling to take EVS into account and + # remove this hack. + if ( + ( + self.is_multimodal_pruning_enabled + or self.requires_sequential_video_encoding + ) + and modality == "video" + and num_items > 1 + ): + batch_outputs_lst = list[torch.Tensor]() + for video_idx in range(num_items): + video_mm_kwargs_item = mm_kwargs[current_item_idx + video_idx] + with self.timed_encoder_operation( + should_time, mm_lora_refs, current_item_idx + video_idx, 1 + ): + _, _, micro_batch_mm_inputs = next( + group_and_batch_mm_kwargs( + [video_mm_kwargs_item], + device=self.device, + pin_memory=PIN_MEMORY, + ) + ) + + micro_batch_outputs = model.embed_multimodal( + **micro_batch_mm_inputs + ) + + batch_outputs_lst.extend(micro_batch_outputs) + + batch_outputs = batch_outputs_lst + else: + # Run the encoder. + # `batch_outputs` is either of the following: + # 1. A tensor of shape (num_items, feature_size, hidden_size) + # in case feature_size is fixed across all multimodal items. + # 2. A list or tuple (length: num_items) of tensors, + # each of shape (feature_size, hidden_size) in case the feature + # size is dynamic depending on the input multimodal items. + + with self.timed_encoder_operation( + should_time, mm_lora_refs, current_item_idx, num_items + ): + cudagraph_output = None + if ( + self.encoder_cudagraph_manager is not None + and self.encoder_cudagraph_manager.supports_modality(modality) + ): + cudagraph_output = self.encoder_cudagraph_manager.execute( + mm_kwargs_batch, + ) + + if cudagraph_output is not None: + batch_outputs = cudagraph_output + else: + batch_outputs = model.embed_multimodal(**mm_kwargs_batch) + + sanity_check_mm_encoder_outputs(batch_outputs, expected_num_items=num_items) + encoder_outputs.extend(batch_outputs) + + current_item_idx += num_items + + # Cache the encoder outputs by mm_hash + for mm_hash, output in zip(mm_hashes, encoder_outputs): + self._cache_encoder_output( + mm_hash, + output, + scheduler_output.ec_manager_metadata, + scheduler_output.free_encoder_mm_hashes, + ) + logger.debug("Finish execute for mm hash %s", mm_hash) + + return encoder_outputs + + def _get_encoder_output_from_cache(self, mm_hash: str) -> torch.Tensor | None: + """Return a cached encoder output for multimodal + embedding gather.""" + return self.encoder_cache.get(mm_hash, None) + + def _gather_mm_embeddings( + self, + scheduler_output: "SchedulerOutput", + shift_computed_tokens: int = 0, + ) -> tuple[list[torch.Tensor], torch.Tensor]: + total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens + + mm_embeds = list[torch.Tensor]() + is_mm_embed = torch.zeros( + total_num_scheduled_tokens, + dtype=torch.bool, + device="cpu", + pin_memory=PIN_MEMORY, + ) + + req_start_idx = 0 + should_sync_mrope_positions = False + should_sync_xdrope_positions = False + + for req_id in self.input_batch.req_ids: + mm_embeds_req: list[torch.Tensor] = [] + + num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id] + req_state = self.requests[req_id] + num_computed_tokens = req_state.num_computed_tokens + shift_computed_tokens + + mm_features = req_state.mm_features + lo, hi = get_mm_features_in_window( + mm_features, + start=num_computed_tokens, + end=num_computed_tokens + num_scheduled_tokens, + ) + for i in range(lo, hi): + mm_feature = mm_features[i] + pos_info = mm_feature.mm_position + start_pos = pos_info.offset + num_encoder_tokens = pos_info.length + + start_idx = max(num_computed_tokens - start_pos, 0) + end_idx = min( + num_computed_tokens - start_pos + num_scheduled_tokens, + num_encoder_tokens, + ) + assert start_idx < end_idx + curr_embeds_start, curr_embeds_end = ( + pos_info.get_embeds_indices_in_range(start_idx, end_idx) + ) + # If there are no embeddings in the current range, we skip + # gathering the embeddings. + if curr_embeds_start == curr_embeds_end: + continue + + mm_hash = mm_feature.identifier + encoder_output = self._get_encoder_output_from_cache(mm_hash) + if encoder_output is None: + # A feature starting at/after the processed boundary is only + # reached via the drafter's +1 look-ahead and might not be + # encoded yet; fall back to the token embedding for drafting. + if ( + start_pos + >= req_state.num_computed_tokens + num_scheduled_tokens + ): + continue + raise RuntimeError(f"Encoder cache miss for {mm_hash}.") + + if (is_embed := pos_info.is_embed) is not None: + is_embed = is_embed[start_idx:end_idx] + mm_embeds_item = encoder_output[curr_embeds_start:curr_embeds_end] + else: + mm_embeds_item = encoder_output[start_idx:end_idx] + + req_start_pos = req_start_idx + start_pos - num_computed_tokens + # OR mask for overlapping mm_features (use_audio_in_video) + if is_embed is None: + is_mm_embed[req_start_pos + start_idx : req_start_pos + end_idx] = ( + True + ) + else: + is_mm_embed[ + req_start_pos + start_idx : req_start_pos + end_idx + ] |= is_embed + set_mm_embedding_modality(mm_embeds_item, mm_feature.modality) + mm_embeds_req.append(mm_embeds_item) + + if self.is_multimodal_pruning_enabled and self.uses_mrope: + assert req_state.mrope_positions is not None + should_sync_mrope_positions = True + old_mm_embeds_req = mm_embeds_req + mm_embeds_req, new_mrope_positions, new_delta = ( + self.model.recompute_mrope_positions( + input_ids=req_state.prompt_token_ids, + multimodal_embeddings=mm_embeds_req, + mrope_positions=req_state.mrope_positions, + num_computed_tokens=req_state.num_computed_tokens, + ) + ) + mm_embeds_req = [ + copy_mm_embedding_modality(src, dst) + for src, dst in zip(old_mm_embeds_req, mm_embeds_req) + ] + req_state.mrope_positions.copy_(new_mrope_positions) + req_state.mrope_position_delta = new_delta + + mm_embeds.extend(mm_embeds_req) + req_start_idx += num_scheduled_tokens + + if should_sync_mrope_positions: + self._calc_mrope_positions(scheduler_output) + self.mrope_positions.copy_to_gpu(total_num_scheduled_tokens) + + if should_sync_xdrope_positions: + self._calc_xdrope_positions(scheduler_output) + self.xdrope_positions.copy_to_gpu(total_num_scheduled_tokens) + + return mm_embeds, is_mm_embed + + def get_model(self) -> nn.Module: + if not hasattr(self, "model"): + raise ValueError("Cannot get model before model has been initialized") + if isinstance( + self.model, (CUDAGraphWrapper, UBatchWrapper, BreakableCUDAGraphWrapper) + ): + # get raw model out of the cudagraph wrapper. + return self.model.unwrap() + return self.model + + def get_draft_model(self) -> nn.Module | None: + drafter = getattr(self, "drafter", None) + if drafter is None: + return None + model = getattr(drafter, "model", None) + if isinstance( + model, (CUDAGraphWrapper, UBatchWrapper, BreakableCUDAGraphWrapper) + ): + return cast(nn.Module, model.unwrap()) + return cast(nn.Module | None, model) + + def get_supported_generation_tasks(self) -> list[GenerationTask]: + model = self.get_model() + supported_tasks = list[GenerationTask]() + + if is_text_generation_model(model): + supported_tasks.append("generate") + + if supports_transcription(model): + if model.supports_transcription_only: + return ["transcription"] + + supported_tasks.append("transcription") + + if supports_realtime(model): + supported_tasks.append("realtime") + + return supported_tasks + + def get_supported_pooling_tasks(self) -> list[PoolingTask]: + model = self.get_model() + if not is_pooling_model(model): + return [] + + return list(model.pooler.get_supported_tasks()) + + def get_supported_tasks(self) -> tuple[SupportedTask, ...]: + tasks = list[SupportedTask]() + + if self.model_config.runner_type == "generate": + tasks.extend(self.get_supported_generation_tasks()) + if self.model_config.runner_type == "pooling": + tasks.extend(self.get_supported_pooling_tasks()) + + return tuple(tasks) + + def sync_and_gather_intermediate_tensors( + self, + num_tokens: int, + intermediate_tensors: IntermediateTensors | None, + sync_self: bool, + ) -> IntermediateTensors: + assert self.intermediate_tensors is not None + + tp = self.vllm_config.parallel_config.tensor_parallel_size + is_rs = is_residual_scattered_for_sp(self.vllm_config, num_tokens) + + # When sequence parallelism is enabled, the "residual" tensor is + # sharded across TP ranks. All-gather it here because downstream + # QKV + Attention needs the full residual before the SP split point. + if sync_self: + assert intermediate_tensors is not None + for k, v in intermediate_tensors.items(): + is_scattered = k == "residual" and is_rs + if is_scattered: + local_len = num_tokens // tp + v = get_tp_group().all_gather(v[:local_len], dim=0) + + self.intermediate_tensors[k][:num_tokens].copy_( + v[:num_tokens], non_blocking=True + ) + + return IntermediateTensors( + {k: v[:num_tokens] for k, v in self.intermediate_tensors.items()} + ) + + def eplb_step(self, is_dummy: bool = False, is_profile: bool = False) -> None: + """ + Step for the EPLB (Expert Parallelism Load Balancing) state. + """ + if not self.parallel_config.enable_eplb or self.eep_eplb_suppressed: + return + + assert self.eplb_state is not None + assert self._moe_model is not None + self.eplb_state.step( + is_dummy, + is_profile, + log_stats=self.parallel_config.eplb_config.log_balancedness, + ) + + def setup_eplb_from_mapping( + self, + expanded_physical_to_logical: torch.Tensor, + ) -> None: + assert self.eplb_state is not None + self.eplb_state.update_mapping( + self.model_config, + expanded_physical_to_logical, + ) + + def _pool( + self, + hidden_states: torch.Tensor, + num_scheduled_tokens: int, + num_scheduled_tokens_np: np.ndarray, + kv_connector_output: KVConnectorOutput | None, + ) -> ModelRunnerOutput | AsyncModelRunnerOutput: + num_reqs = self.input_batch.num_reqs + assert num_reqs == len(self.input_batch.pooling_params), ( + "Either all or none of the requests in a batch must be pooling request" + ) + + hidden_states = hidden_states[:num_scheduled_tokens] + seq_lens_cpu = self.optimistic_seq_lens_cpu[:num_reqs] + + pooling_metadata = self.input_batch.get_pooling_metadata() + pooling_metadata.build_pooling_cursor( + num_scheduled_tokens_np, + seq_lens_cpu, + device=hidden_states.device, + query_start_loc_gpu=self.query_start_loc.gpu[: num_reqs + 1], + ) + + model = cast(VllmModelForPooling, self.model) + raw_pooler_output: PoolerOutput = model.pooler( + hidden_states=hidden_states, pooling_metadata=pooling_metadata + ) + + finished_mask = pooling_metadata.get_pooling_cursor().get_finished_mask() + raw_pooler_output = self.late_interaction_runner.postprocess_pooler_output( + raw_pooler_output=raw_pooler_output, + pooling_params=pooling_metadata.pooling_params, + req_ids=self.input_batch.req_ids, + finished_mask=finished_mask, + ) + + model_runner_output = ModelRunnerOutput( + req_ids=self.input_batch.req_ids.copy(), + req_id_to_index=self.input_batch.req_id_to_index.copy(), + kv_connector_output=kv_connector_output, + ) + + if raw_pooler_output is None or not any(finished_mask): + self._sync_device() + model_runner_output.pooler_output = [None] * num_reqs + return model_runner_output + + if not current_platform.is_cuda_alike(): + # cpu/xpu runners cannot use the CUDA stream/event-based wrapper. + model_runner_output.pooler_output = _copy_pooler_output_to_cpu( + raw_pooler_output=raw_pooler_output, + finished_mask=finished_mask, + ) + self._sync_device() + return model_runner_output + + return AsyncGPUPoolingModelRunnerOutput( + model_runner_output=model_runner_output, + raw_pooler_output=raw_pooler_output, + finished_mask=finished_mask, + async_output_copy_stream=self._get_or_create_async_output_copy_stream(), + ) + + def _pad_for_sequence_parallelism(self, num_scheduled_tokens: int) -> int: + # Pad tokens to multiple of tensor_parallel_size when + # enabled collective fusion for SP + tp_size = self.vllm_config.parallel_config.tensor_parallel_size + if self.compilation_config.pass_config.enable_sp and tp_size > 1: + return round_up(num_scheduled_tokens, tp_size) + return num_scheduled_tokens + + def _prepare_mm_inputs( + self, num_tokens: int + ) -> tuple[torch.Tensor | None, torch.Tensor]: + if self.model.requires_raw_input_tokens: + input_ids = self.input_ids.gpu[:num_tokens] + else: + input_ids = None + + inputs_embeds = self.inputs_embeds.gpu[:num_tokens] + return input_ids, inputs_embeds + + def _preprocess( + self, + scheduler_output: "SchedulerOutput", + num_input_tokens: int, # Padded + intermediate_tensors: IntermediateTensors | None = None, + ) -> tuple[ + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor, + IntermediateTensors | None, + dict[str, Any], + ECConnectorOutput | None, + ]: + num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens + is_first_rank = get_pp_group().is_first_rank + is_encoder_decoder = self.model_config.is_encoder_decoder + + # Clamp speculative scheduler placeholders (-1) before embedding lookup. + if self.speculative_config is not None: + self.input_ids.gpu[:num_input_tokens].clamp_(min=0) + + # _prepare_inputs may reorder the batch, so we must gather multi + # modal outputs after that to ensure the correct order + ec_connector_output = None + + if self.supports_mm_inputs and is_first_rank and not is_encoder_decoder: + # Run the multimodal encoder if any. + with self.maybe_get_ec_connector_output( + scheduler_output, + encoder_cache=self.encoder_cache, + ) as ec_connector_output: + self._execute_mm_encoder(scheduler_output) + mm_embeds, is_mm_embed = self._gather_mm_embeddings(scheduler_output) + + # NOTE(woosuk): To unify token ids and soft tokens (vision + # embeddings), we always use embeddings (rather than token ids) + # as input to the multimodal model, even when the input is text. + if self.enable_prompt_embeds and self.input_batch.req_prompt_embeds: + # Some positions carry precomputed prompt_embeds: they are + # already in self.inputs_embeds and marked is_token_ids=False. + # Embed only the token-id positions (zeroing the placeholder ids + # at prompt_embeds positions so the embedding gather cannot read + # out-of-range ids), and write them back without clobbering the + # prompt_embeds positions. + is_token_ids = self.is_token_ids.gpu[:num_scheduled_tokens] + safe_input_ids = torch.where( + is_token_ids, + self.input_ids.gpu[:num_scheduled_tokens], + 0, + ) + inputs_embeds_scheduled = self.model.embed_input_ids( + safe_input_ids, + multimodal_embeddings=mm_embeds, + is_multimodal=is_mm_embed, + ) + target = self.inputs_embeds.gpu[:num_scheduled_tokens] + self.inputs_embeds.gpu[:num_scheduled_tokens] = torch.where( + is_token_ids.unsqueeze(-1), + inputs_embeds_scheduled, + target, + ) + else: + inputs_embeds_scheduled = self.model.embed_input_ids( + self.input_ids.gpu[:num_scheduled_tokens], + multimodal_embeddings=mm_embeds, + is_multimodal=is_mm_embed, + ) + + # TODO(woosuk): Avoid the copy. Optimize. + self.inputs_embeds.gpu[:num_scheduled_tokens].copy_( + inputs_embeds_scheduled + ) + + input_ids, inputs_embeds = self._prepare_mm_inputs(num_input_tokens) + model_kwargs = { + **self._init_model_kwargs(), + **self._extract_mm_kwargs(scheduler_output), + } + elif self.enable_prompt_embeds and is_first_rank: + # Get the input embeddings for the tokens that are not input embeds, + # then put them into the appropriate positions. + # TODO(qthequartermasterman): Since even when prompt embeds are + # enabled, (a) not all requests will use prompt embeds, and (b) + # after the initial prompt is processed, the rest of the generated + # tokens will be token ids, it is not desirable to have the + # embedding layer outside of the CUDA graph all the time. The v0 + # engine avoids this by "double compiling" the CUDA graph, once + # with input_ids and again with inputs_embeds, for all num_tokens. + # If a batch only has token ids, then including the embedding layer + # in the CUDA graph will be more performant (like in the else case + # below). + is_token_ids = self.is_token_ids.np[:num_scheduled_tokens] + token_ids_idx_np = np.nonzero(is_token_ids)[0] + # Some tokens ids may need to become embeds + if token_ids_idx_np.size > 0: + token_ids_idx = async_tensor_h2d(token_ids_idx_np, device=self.device) + token_ids = self.input_ids.gpu[token_ids_idx] + tokens_to_embeds = self.model.embed_input_ids(input_ids=token_ids) + self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds + + inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens] + model_kwargs = self._init_model_kwargs() + input_ids = None + else: + # For text-only models, we use token ids as input. + # While it is possible to use embeddings as input just like the + # multimodal models, it is not desirable for performance since + # then the embedding layer is not included in the CUDA graph. + input_ids = self.input_ids.gpu[:num_input_tokens] + inputs_embeds = None + model_kwargs = self._init_model_kwargs() + + if self.uses_mrope: + positions = self.mrope_positions.gpu[:, :num_input_tokens] + elif self.uses_xdrope_dim > 0: + positions = self.xdrope_positions.gpu[:, :num_input_tokens] + else: + positions = self.positions[:num_input_tokens] + if num_input_tokens > num_scheduled_tokens: + self.positions[num_scheduled_tokens:num_input_tokens].zero_() + + if is_first_rank: + intermediate_tensors = None + else: + assert intermediate_tensors is not None + intermediate_tensors = self.sync_and_gather_intermediate_tensors( + num_input_tokens, intermediate_tensors, True + ) + + if is_encoder_decoder and scheduler_output.scheduled_encoder_inputs: + # Run the encoder, just like we do with other multimodal inputs. + # For an encoder-decoder model, our processing here is a bit + # simpler, because the outputs are just passed to the decoder. + # We are not doing any prompt replacement. We also will only + # ever have a single encoder input. + encoder_outputs = self._execute_mm_encoder(scheduler_output) + model_kwargs.update({"encoder_outputs": encoder_outputs}) + + return ( + input_ids, + inputs_embeds, + positions, + intermediate_tensors, + model_kwargs, + ec_connector_output, + ) + + def _sample( + self, + logits: torch.Tensor | None, + spec_decode_metadata: SpecDecodeMetadata | None, + ) -> SamplerOutput: + # Sample the next token and get logprobs if needed. + sampling_metadata = self.input_batch.sampling_metadata + # Update output token ids with tokens sampled in last step + # if async scheduling and required by current sampling params. + self.input_batch.update_async_output_token_ids() + if spec_decode_metadata is None: + return self.sampler( + logits=logits, + sampling_metadata=sampling_metadata, + ) + + # Update spec_token_ids with real draft tokens from pre step only when + # output_token_ids is needed (penalties or bad_words are in use). + if self.use_async_scheduling and self._draft_token_req_ids is not None: + draft_token_ids_cpu, _ = self._get_draft_token_ids_cpu() + self.input_batch.update_async_spec_token_ids(draft_token_ids_cpu) + + draft_probs = self._get_spec_decode_draft_probs(spec_decode_metadata) + sampler_output = self.rejection_sampler( + spec_decode_metadata, + draft_probs, + logits, + sampling_metadata, + ) + return sampler_output + + def _bookkeeping_sync( + self, + scheduler_output: "SchedulerOutput", + sampler_output: SamplerOutput, + logits: torch.Tensor | None, + hidden_states: torch.Tensor, + num_scheduled_tokens: int, + ) -> tuple[ + dict[str, int], + torch.Tensor | None, + LogprobsLists | None, + list[list[int]], + dict[str, LogprobsTensors | None], + list[str], + dict[str, int], + list[int], + ]: + num_nans: torch.Tensor | None = None + num_nans_in_logits: dict[str, int] = {} + if envs.VLLM_COMPUTE_NANS_IN_LOGITS: + if self.use_async_scheduling: + # Keep the counts on device; they ride the async output copy + # stream rather than blocking here. + num_nans = None if logits is None else count_nans_per_row(logits) + else: + num_nans_in_logits = self._get_nans_in_logits(logits) + + num_reqs = self.input_batch.num_reqs + discard_sampled_tokens_req_indices = np.nonzero( + self.discard_request_mask.np[:num_reqs] + )[0] + for i in discard_sampled_tokens_req_indices: + gen = self.input_batch.generators.get(int(i)) + if gen is not None: + gen.set_offset(gen.get_offset() - 4) + + # Copy some objects so they don't get modified after returning. + # This is important when using async scheduling. + req_ids_output_copy = self.input_batch.req_ids.copy() + req_id_to_index_output_copy = self.input_batch.req_id_to_index.copy() + + num_sampled_tokens = sampler_output.sampled_token_ids.shape[0] + sampled_token_ids = sampler_output.sampled_token_ids + logprobs_tensors = sampler_output.logprobs_tensors + invalid_req_indices = [] + logprobs_lists = None + if not self.use_async_scheduling: + # Sync scheduling: issue routed experts D2H into the pinned + # CPU buffer BEFORE ``_to_list`` below. ``_to_list`` does + # ``event.synchronize()`` on the async copy stream which + # waits for every D2H queued on the default stream since + # the last sync, so this enqueue is naturally covered + # without requiring its own synchronize. + if self.routed_experts_initialized: + buf = self.routed_experts_capturer.get_device_buffer() + total = scheduler_output.total_num_scheduled_tokens + self.routed_experts_cpu[:total].copy_(buf[:total], non_blocking=True) + self.routed_experts_slot_mapping_cpu[:total].copy_( + self.routed_experts_slot_mapping_device[:total], + non_blocking=True, + ) + with gpu_sync_allowed(): + # Get the valid generated tokens. + max_gen_len = sampled_token_ids.shape[-1] + if max_gen_len == 1: + # No spec decode tokens. + valid_sampled_token_ids = self._to_list(sampled_token_ids) + # Mask out the sampled tokens that should not be sampled. + for i in discard_sampled_tokens_req_indices: + valid_sampled_token_ids[int(i)].clear() + + if logprobs_tensors is not None: + logprobs_lists = logprobs_tensors.tolists() + else: + # Includes spec decode tokens. + valid_sampled_token_ids, logprobs_lists = ( + RejectionSampler.parse_output( + sampled_token_ids, + self.input_batch.vocab_size, + discard_sampled_tokens_req_indices, + logprobs_tensors=logprobs_tensors, + ) + ) + else: + valid_sampled_token_ids = [] + invalid_req_indices = discard_sampled_tokens_req_indices.tolist() + invalid_req_indices_set = set(invalid_req_indices) + + # Cache the sampled tokens on the GPU and avoid CPU sync. + # These will be copied into input_ids in the next step + # when preparing inputs. + # With spec decoding, this is done in propose_draft_token_ids(). + if self.input_batch.prev_sampled_token_ids is None: + assert sampled_token_ids.shape[-1] == 1 + self.input_batch.prev_sampled_token_ids = sampled_token_ids + self.input_batch.prev_req_id_to_index = { + req_id: i + for i, req_id in enumerate(self.input_batch.req_ids) + if i not in invalid_req_indices_set + } + + # Cache the sampled tokens in the model runner, so that the scheduler + # doesn't need to send them back. + # NOTE(woosuk): As an exception, when using PP, the scheduler sends + # the sampled tokens back, because there's no direct communication + # between the first-stage worker and the last-stage worker. + req_ids = self.input_batch.req_ids + for req_idx in range(num_sampled_tokens): + if self.use_async_scheduling: + sampled_ids = [-1] if req_idx not in invalid_req_indices_set else None + else: + sampled_ids = valid_sampled_token_ids[req_idx] + + num_sampled_ids: int = len(sampled_ids) if sampled_ids else 0 + + if not sampled_ids: + continue + + start_idx = self.input_batch.num_tokens_no_spec[req_idx] + end_idx = start_idx + num_sampled_ids + assert end_idx <= self.max_model_len, ( + "Sampled token IDs exceed the max model length. " + f"Total number of tokens: {end_idx} > max_model_len: " + f"{self.max_model_len}" + ) + + self.input_batch.token_ids_cpu[req_idx, start_idx:end_idx] = sampled_ids + self.input_batch.is_token_ids[req_idx, start_idx:end_idx] = True + self.input_batch.num_tokens_no_spec[req_idx] = end_idx + + req_id = req_ids[req_idx] + req_state = self.requests[req_id] + req_state.output_token_ids.extend(sampled_ids) + + # Compute prompt logprobs if needed. + prompt_logprobs_dict = self._get_prompt_logprobs_dict( + hidden_states[:num_scheduled_tokens], + scheduler_output.num_scheduled_tokens, + ) + + return ( + num_nans_in_logits, + num_nans, + logprobs_lists, + valid_sampled_token_ids, + prompt_logprobs_dict, + req_ids_output_copy, + req_id_to_index_output_copy, + invalid_req_indices, + ) + + @contextmanager + def synchronize_input_prep(self): + if self.prepare_inputs_event is None: + yield + return + + # Ensure prior step has finished with reused CPU tensors. + # This is required in the async scheduling case because + # the CPU->GPU transfer happens async. + self.prepare_inputs_event.synchronize() + try: + yield + finally: + self.prepare_inputs_event.record() + + def _model_forward( + self, + input_ids: torch.Tensor | None = None, + positions: torch.Tensor | None = None, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **model_kwargs: dict[str, Any], + ) -> Any: + """Helper method to call the model forward pass. + + This method can be overridden by subclasses for model execution. + Motivation: We can inspect only this method versus + the whole execute_model, which has additional logic. + + Args: + input_ids: Input token IDs + positions: Token positions + intermediate_tensors: Tensors from previous pipeline stages + inputs_embeds: Input embeddings (alternative to input_ids) + **model_kwargs: Additional model arguments + + Returns: + Model output tensor + """ + return self.model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + **model_kwargs, + ) + + @staticmethod + def _is_uniform_decode( + max_num_scheduled_tokens: int, + uniform_decode_query_len: int, + num_tokens: int, + num_reqs: int, + force_uniform_decode: bool | None = None, + ) -> bool: + """ + Checks if it's a decode batch with same amount scheduled tokens + across all requests. + """ + return ( + ( + (max_num_scheduled_tokens == uniform_decode_query_len) + and (num_tokens == max_num_scheduled_tokens * num_reqs) + ) + if force_uniform_decode is None + else force_uniform_decode + ) + + def _allow_microbatching( + self, num_reqs: int, num_scheduled_tokens_np: np.ndarray + ) -> bool: + """Refuse to microbatch a step that splits a prefix from its writer. + + A request can be admitted on a prefix cache hit against blocks another + request in the same batch is only computing now. Run whole, the step + issues every KV cache write before any attention read and the hit + holds; split, a reader in the first half would attend over blocks the + writer in the second half has not filled in yet. Vetoing on one rank + settles it for all, since ranks agree on microbatching collectively. + """ + if not self.parallel_config.use_ubatching or num_reqs < 2: + return True + computed = self.input_batch.num_computed_tokens_cpu[:num_reqs] + query_lens = num_scheduled_tokens_np[:num_reqs] + # A decode reads nothing it was not already given in an earlier step. + readers = np.flatnonzero( + (query_lens > (self.reorder_batch_threshold or 1)) & (computed > 0) + ) + if readers.size == 0: + return True + + for block_table in self.input_batch.block_table.block_tables: + block_size = block_table.block_size + table = block_table.get_numpy_array() + start = computed // block_size + stop = (computed + query_lens + block_size - 1) // block_size + span = int((stop - start).max()) + columns = start[:, None] + np.arange(span)[None, :] + written = np.unique( + np.take_along_axis( + table[:num_reqs], + np.minimum(columns, table.shape[1] - 1), + axis=1, + )[columns < stop[:, None]] + ) + for reader in readers: + # Whole blocks only: a reader ends inside a partly filled one, + # which is private to it and which it fills before reading. + whole = computed[reader] // block_size + if np.isin(table[reader, :whole], written).any(): + return False + return True + + def _determine_batch_execution_and_padding( + self, + num_tokens: int, + num_reqs: int, + num_scheduled_tokens_np: np.ndarray, + max_num_scheduled_tokens: int, + use_cascade_attn: bool, + allow_microbatching: bool = True, + force_eager: bool = False, + # For cudagraph capture TODO(lucas): Refactor how we capture cudagraphs (will + # be improved in model runner v2) + force_uniform_decode: bool | None = None, + force_has_lora: bool | None = None, + force_num_active_loras: int | None = None, + num_encoder_reqs: int = 0, + ) -> tuple[ + CUDAGraphMode, + BatchDescriptor, + bool, + torch.Tensor | None, + CUDAGraphStat | None, + ]: + uniform_decode = self._is_uniform_decode( + max_num_scheduled_tokens=max_num_scheduled_tokens, + uniform_decode_query_len=self.uniform_decode_query_len, + num_tokens=num_tokens, + num_reqs=num_reqs, + force_uniform_decode=force_uniform_decode, + ) + # Encoder-decoder models only support CG for decoder_step > 0 (no enc_output + # is present). Also, chunked-prefill is disabled, so batch are uniform. + has_encoder_output = ( + self.model_config.is_encoder_decoder and num_encoder_reqs > 0 + ) + + # Compute LoRA state for cudagraph dispatch + num_active_loras = ( + force_num_active_loras + if force_num_active_loras is not None + else len(self.input_batch.lora_id_to_lora_request) + ) + has_lora = num_active_loras > 0 if force_has_lora is None else force_has_lora + + num_tokens_padded = self._pad_for_sequence_parallelism(num_tokens) + + def dispatch_cudagraph(num_tokens, disable_full=False, valid_modes=None): + return self.cudagraph_dispatcher.dispatch( + num_tokens=num_tokens, + has_lora=has_lora, + uniform_decode=uniform_decode, + num_active_loras=num_active_loras, + valid_modes={CUDAGraphMode.NONE} if force_eager else valid_modes, + invalid_modes={CUDAGraphMode.FULL} if disable_full else None, + ) + + cudagraph_mode, batch_descriptor = dispatch_cudagraph( + num_tokens_padded, disable_full=use_cascade_attn or has_encoder_output + ) + num_tokens_padded = batch_descriptor.num_tokens + if self.compilation_config.pass_config.enable_sp: + assert ( + batch_descriptor.num_tokens + % self.vllm_config.parallel_config.tensor_parallel_size + == 0 + ), ( + "Sequence parallelism requires num_tokens to be " + "a multiple of tensor parallel size" + ) + + # Extra coordination when running data-parallel since we need to coordinate + # across ranks + should_ubatch, num_tokens_across_dp = False, None + if self.vllm_config.parallel_config.data_parallel_size > 1: + should_ubatch, num_tokens_across_dp, synced_cudagraph_mode = ( + coordinate_batch_across_dp( + num_tokens_unpadded=num_tokens, + parallel_config=self.parallel_config, + allow_microbatching=allow_microbatching, + num_tokens_padded=num_tokens_padded, + uniform_decode=uniform_decode, + cudagraph_mode=cudagraph_mode.value, + ) + ) + + # Extract DP-synced values + if num_tokens_across_dp is not None: + dp_rank = self.parallel_config.data_parallel_rank + num_tokens_padded = int(num_tokens_across_dp[dp_rank].item()) + # Re-dispatch with DP padding so we have the correct batch_descriptor + cudagraph_mode, batch_descriptor = dispatch_cudagraph( + num_tokens_padded, + valid_modes={CUDAGraphMode(synced_cudagraph_mode)}, + ) + # Assert to make sure the agreed upon token count is correct otherwise + # num_tokens_across_dp will no-longer be valid + assert batch_descriptor.num_tokens == num_tokens_padded + + cudagraph_stats = None + if self.vllm_config.observability_config.cudagraph_metrics: + cudagraph_stats = CUDAGraphStat( + num_unpadded_tokens=num_tokens, + num_padded_tokens=batch_descriptor.num_tokens, + num_paddings=batch_descriptor.num_tokens - num_tokens, + runtime_mode=str(cudagraph_mode), + ) + + return ( + cudagraph_mode, + batch_descriptor, + should_ubatch, + num_tokens_across_dp, + cudagraph_stats, + ) + + def _register_layerwise_nvtx_hooks(self) -> None: + """ + Register layerwise NVTX hooks if --enable-layerwise-nvtx-tracing is enabled + to trace detailed information of each layer or module in the model. + """ + + if ( + self.vllm_config.observability_config.enable_layerwise_nvtx_tracing + and not self.layerwise_nvtx_hooks_registered + ): + if self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE: + logger.debug_once( + "layerwise NVTX tracing is not supported when CUDA graph is " + "turned off; you may observe part or all of the model " + "missing NVTX markers" + ) + + # In STOCK_TORCH_COMPILE mode, after registering hooks here, + # the __call__ function of nn.module will be recompiled with + # fullgraph=True. Since nvtx.range_push/pop are not traceable + # by torch dynamo, we can't register hook functions here + # because hook functions will also be traced by torch dynamo. + if ( + self.vllm_config.compilation_config.mode + == CompilationMode.STOCK_TORCH_COMPILE + ): + logger.debug_once( + "layerwise NVTX tracing is not supported when " + "CompilationMode is STOCK_TORCH_COMPILE, skipping " + "function hooks registration" + ) + else: + pyt_hooks = PytHooks() + pyt_hooks.register_hooks(self.model, self.model.__class__.__name__) + self.layerwise_nvtx_hooks_registered = True + + def _get_slot_mappings( + self, + num_tokens_padded: int, + num_reqs_padded: int, + num_tokens_unpadded: int, + ubatch_slices: "UBatchSlices | None" = None, + ) -> tuple[ + dict[int, torch.Tensor] | None, + dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None, + ]: + """ + Build slot mappings in both formats needed by the system. + + Args: + num_tokens_padded: Total number of tokens (padded) + num_reqs_padded: Total number of requests (padded) + num_tokens_unpadded: Actual number of tokens (unpadded) + ubatch_slices: Optional ubatch slicing info for DBO + + Returns: + A tuple of: + - slot_mappings_by_gid: dict[int, torch.Tensor] for attention metadata + - slot_mappings_by_layer: dict[str, torch.Tensor] or list for ForwardContext + """ + if not ( + hasattr(self, "kv_cache_config") + and self.kv_cache_config is not None + and len(self.kv_cache_config.kv_cache_groups) > 0 + ): + return None, None + + def _get_slot_mapping(kv_cache_gid: int): + assert num_reqs_padded is not None and num_tokens_padded is not None + kv_cache_spec = self.kv_cache_config.kv_cache_groups[ + kv_cache_gid + ].kv_cache_spec + if isinstance(kv_cache_spec, EncoderOnlyAttentionSpec): + slot_mapping = torch.zeros( + (num_tokens_padded,), + dtype=torch.int64, + device=self.device, + ) + else: + blk_table = self.input_batch.block_table[kv_cache_gid] + slot_mapping = blk_table.slot_mapping.gpu[:num_tokens_padded] + + # Fill unused with -1. Needed for reshape_and_cache in full cuda + # graph mode. `blk_table_tensor` -1 to match mamba PAD_SLOT_ID + slot_mapping[num_tokens_unpadded:num_tokens_padded].fill_(-1) + + return slot_mapping + + slot_mappings_by_gid = { + gid: _get_slot_mapping(gid) + for gid, _ in enumerate(self.kv_cache_config.kv_cache_groups) + } + + slot_mappings_by_layer: dict[str, torch.Tensor] = {} + for gid, kv_cache_group in enumerate(self.kv_cache_config.kv_cache_groups): + slot_mapping = slot_mappings_by_gid[gid] + for layer_name in kv_cache_group.layer_names: + slot_mappings_by_layer[layer_name] = slot_mapping + + if ubatch_slices is not None: + result: list[dict[str, torch.Tensor]] = [] + for ubatch in ubatch_slices: + sliced_mappings: dict[str, torch.Tensor] = {} + for layer_name, slot_mapping in slot_mappings_by_layer.items(): + sliced_mappings[layer_name] = slot_mapping[ubatch.token_slice] + result.append(sliced_mappings) + return slot_mappings_by_gid, result + + return slot_mappings_by_gid, slot_mappings_by_layer + + def _is_all_reqs_chunked_prefill(self) -> bool: + """Check if all scheduled requests are marked to discard sampled tokens. + + This is true when `discard_request_mask` is set for every scheduled + request (e.g., for chunked prefill requests that are not the last + prefill chunk).""" + num_reqs = self.input_batch.num_reqs + return bool(self.discard_request_mask.np[:num_reqs].all()) + + @torch.inference_mode() + def execute_model( + self, + scheduler_output: "SchedulerOutput", + intermediate_tensors: IntermediateTensors | None = None, + ) -> ModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors | None: + if self.execute_model_state is not None: + raise RuntimeError( + "State error: sample_tokens() must be called " + "after execute_model() returns None." + ) + + # If ngram_gpu is used, we need to copy the scheduler_output to avoid + # the modification has influence on the scheduler_output in engine core process. + # The replace is much faster than deepcopy. + if ( + self.speculative_config is not None + and self.speculative_config.use_ngram_gpu() + ): + num_scheduled_tokens_copy = scheduler_output.num_scheduled_tokens.copy() + spec_decode_tokens_copy = ( + scheduler_output.scheduled_spec_decode_tokens.copy() + ) + scheduler_output = replace( + scheduler_output, + num_scheduled_tokens=num_scheduled_tokens_copy, + scheduled_spec_decode_tokens=spec_decode_tokens_copy, + ) + + if has_kv_transfer_group(): + kv_connector_metadata = scheduler_output.kv_connector_metadata + assert kv_connector_metadata is not None + get_kv_transfer_group().handle_preemptions(kv_connector_metadata) + + num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens + with ( + record_function_or_nullcontext("gpu_model_runner: preprocess"), + self.synchronize_input_prep(), + ): + # Update persistent batch states. + deferred_state_corrections_fn = self._update_states(scheduler_output) + + if has_ec_transfer() and not get_ec_transfer().is_consumer: + with self.maybe_get_ec_connector_output( + scheduler_output, + encoder_cache=self.encoder_cache, + ) as ec_connector_output: + self._execute_mm_encoder(scheduler_output) + return make_empty_encoder_model_runner_output(scheduler_output) + + if not num_scheduled_tokens: + if ( + self.parallel_config.distributed_executor_backend + == "external_launcher" + and self.parallel_config.data_parallel_size > 1 + ): + # this is a corner case when both external launcher + # and DP are enabled, num_scheduled_tokens could be + # 0, and has_unfinished_requests in the outer loop + # returns True. before returning early here we call + # dummy run to ensure coordinate_batch_across_dp + # is called into to avoid out of sync issues. + self._dummy_run(1) + if not has_kv_transfer_group(): + # Return empty ModelRunnerOutput if no work to do. + return EMPTY_MODEL_RUNNER_OUTPUT + return self.kv_connector_no_forward(scheduler_output, self.vllm_config) + + if self.cache_config.kv_sharing_fast_prefill: + assert not self.num_prompt_logprobs, ( + "--kv-sharing-fast-prefill produces incorrect " + "logprobs for prompt tokens, tokens, please disable " + "it when the requests need prompt logprobs" + ) + + num_reqs = self.input_batch.num_reqs + req_ids = self.input_batch.req_ids + tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids] + num_scheduled_tokens_np = np.array(tokens, dtype=np.int32) + max_num_scheduled_tokens = int(num_scheduled_tokens_np.max()) + num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens + + logits_indices, spec_decode_metadata, max_num_sampled_tokens = ( + self._prepare_inputs(scheduler_output, num_scheduled_tokens_np) + ) + + cascade_attn_prefix_lens = None + # Disable cascade attention when using microbatching (DBO) + if self.cascade_attn_enabled and not self.parallel_config.use_ubatching: + # Pre-compute cascade attention prefix lengths + cascade_attn_prefix_lens = self._compute_cascade_attn_prefix_lens( + num_scheduled_tokens_np, + self.input_batch.num_computed_tokens_cpu[:num_reqs], + scheduler_output.num_common_prefix_blocks, + ) + + ( + cudagraph_mode, + batch_desc, + should_ubatch, + num_tokens_across_dp, + cudagraph_stats, + ) = self._determine_batch_execution_and_padding( + num_tokens=num_tokens_unpadded, + num_reqs=num_reqs, + num_scheduled_tokens_np=num_scheduled_tokens_np, + max_num_scheduled_tokens=max_num_scheduled_tokens, + use_cascade_attn=cascade_attn_prefix_lens is not None, + num_encoder_reqs=len(scheduler_output.scheduled_encoder_inputs), + allow_microbatching=self._allow_microbatching( + num_reqs, num_scheduled_tokens_np + ), + ) + + logger.debug( + "Running batch with cudagraph_mode: %s, batch_descriptor: %s, " + "should_ubatch: %s, num_tokens_across_dp: %s", + cudagraph_mode, + batch_desc, + should_ubatch, + num_tokens_across_dp, + ) + + num_tokens_padded = batch_desc.num_tokens + num_reqs_padded = ( + batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs + ) + ubatch_slices, ubatch_slices_padded = maybe_create_ubatch_slices( + should_ubatch, + num_scheduled_tokens_np, + num_tokens_padded, + num_reqs_padded, + self.parallel_config.num_ubatches, + ) + + logger.debug( + "ubatch_slices: %s, ubatch_slices_padded: %s", + ubatch_slices, + ubatch_slices_padded, + ) + + # True if any attention backend handles KV cache update separately + # from forward() (i.e., forward_includes_kv_cache_update=False). When true, + # slot_mappings must use padded dimensions to match the key/value tensors. + has_separate_kv_update = not all( + all( + g.backend.forward_includes_kv_cache_update + for g in self.attn_groups[id] + ) + for id, spec in enumerate(self.kv_cache_config.kv_cache_groups) + if not isinstance(spec.kv_cache_spec, EncoderOnlyAttentionSpec) + ) + pad_attn = cudagraph_mode == CUDAGraphMode.FULL + + if self.cache_config.mamba_cache_mode == "align": + # preprocess_mamba reads req_state.num_computed_tokens (CPU) + # to decide copy operations, so we must apply deferred + # corrections before it runs. + if deferred_state_corrections_fn: + deferred_state_corrections_fn() + deferred_state_corrections_fn = None + mamba_bufs = self._get_mamba_bufs() + mamba_utils.preprocess_mamba( + scheduler_output, + self.kv_cache_config, + self.cache_config, + self.mamba_state_idx, + self.input_batch, + self.requests, + self.compilation_config.static_forward_context, + mamba_bufs.preprocess.copy_funcs_by_type, + mamba_bufs.preprocess, + align_ctx=mamba_bufs.postprocess_align, + ) + # preprocess_mamba resets num_accepted_tokens_cpu to 1 + # for requests whose state was copied to a new block. + # Re-sync to GPU so the mamba kernel reads from the + # correct initial state slot (init_token_idx = 0). + self.num_accepted_tokens.np[:num_reqs] = ( + self.input_batch.num_accepted_tokens_cpu[:num_reqs] + ) + self.num_accepted_tokens.copy_to_gpu(num_reqs) + + # Stage per-request inputs for the fused postprocess kernel + # only when that kernel will actually run. The kernel is + # gated on spec-decode + hybrid (see MambaBuffers.create); + # without it, ``mamba_bufs.postprocess_align`` is None and + # the staging buffers don't exist. + if mamba_bufs.postprocess_align is not None: + mamba_utils.stage_postprocess_inputs_to_gpu( + mamba_bufs.postprocess_align, + scheduler_output, + self.input_batch.req_ids, + num_reqs, + self.requests, + self.mamba_state_idx, + ) + + use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0 + ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices + + slot_mappings_by_group, slot_mappings = self._get_slot_mappings( + num_tokens_padded=num_tokens_padded + if pad_attn or has_separate_kv_update + else num_tokens_unpadded, + num_reqs_padded=( + num_reqs_padded if pad_attn or has_separate_kv_update else num_reqs + ), + num_tokens_unpadded=num_tokens_unpadded, + ubatch_slices=ubatch_slices_padded, + ) + + attn_metadata, spec_decode_common_attn_metadata = ( + self._build_attention_metadata( + num_tokens=num_tokens_unpadded, + num_tokens_padded=num_tokens_padded if pad_attn else None, + num_reqs=num_reqs, + num_reqs_padded=num_reqs_padded if pad_attn else None, + max_query_len=max_num_scheduled_tokens, + ubatch_slices=ubatch_slices_attn, + logits_indices=logits_indices, + max_num_sampled_tokens=max_num_sampled_tokens, + use_spec_decode=use_spec_decode, + num_scheduled_tokens=scheduler_output.num_scheduled_tokens, + recurrent_prefill_checkpoint_plans=getattr( + scheduler_output, "recurrent_prefill_checkpoint_plans", None + ), + cascade_attn_prefix_lens=cascade_attn_prefix_lens, + slot_mappings=slot_mappings_by_group, + ) + ) + + ( + input_ids, + inputs_embeds, + positions, + intermediate_tensors, + model_kwargs, + ec_connector_output, + ) = self._preprocess( + scheduler_output, num_tokens_padded, intermediate_tensors + ) + + # Encoder-decoder models can only compile the pure decode steps where no + # encoder inputs are present. Use eager for the first pass. + num_encoder_reqs = len(scheduler_output.scheduled_encoder_inputs) + has_encoder_input = ( + self.model_config.is_encoder_decoder and num_encoder_reqs > 0 + ) + + # Run the model. + # Use persistent buffers for CUDA graphs. + # When spec decode is enabled, defer connector finalization + # (wait_for_save + clear metadata) until after draft model runs. + defer_kv_connector_finalize = self.speculative_config is not None + # Update the EPLB meta. + if self.eplb_state is not None: + self.eplb_state.prepare_forward( + self.model_config, + num_tokens_unpadded, + ubatch_slices_padded, + ) + with ( + set_forward_context( + attn_metadata, + self.vllm_config, + num_tokens=num_tokens_padded, + num_tokens_across_dp=num_tokens_across_dp, + cudagraph_runtime_mode=cudagraph_mode, + batch_descriptor=batch_desc, + ubatch_slices=ubatch_slices_padded, + slot_mapping=slot_mappings, + skip_compiled=has_encoder_input, + ), + record_function_or_nullcontext("gpu_model_runner: forward"), + self.maybe_get_kv_connector_output( + scheduler_output, + defer_finalize=defer_kv_connector_finalize, + ) as kv_connector_output, + ): + model_output = self._model_forward( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + **model_kwargs, + ) + + with record_function_or_nullcontext("gpu_model_runner: postprocess"): + if self.use_aux_hidden_state_outputs: + # True when EAGLE 3 is used. + hidden_states, aux_hidden_states = model_output + else: + # Common case. + hidden_states = model_output + aux_hidden_states = None + + if not self.broadcast_pp_output: + # Common case. + if not get_pp_group().is_last_rank: + # Return the intermediate tensors. + assert isinstance(hidden_states, IntermediateTensors) + self.kv_connector_output = kv_connector_output + return hidden_states + + if self.is_pooling_model: + # Return the pooling output. + return self._pool( + hidden_states, + num_scheduled_tokens, + num_scheduled_tokens_np, + kv_connector_output, + ) + + sample_hidden_states = hidden_states[logits_indices] + logits = self.model.compute_logits(sample_hidden_states) + else: + # Rare case. + assert not self.is_pooling_model + + sample_hidden_states = hidden_states[logits_indices] + if not get_pp_group().is_last_rank: + all_gather_tensors = { + "residual": not is_residual_scattered_for_sp( + self.vllm_config, num_tokens_padded + ) + } + get_pp_group().send_tensor_dict( + hidden_states.tensors, + all_gather_group=get_tp_group(), + all_gather_tensors=all_gather_tensors, + ) + logits = None + else: + logits = self.model.compute_logits(sample_hidden_states) + + model_output_broadcast_data: dict[str, Any] = {} + if logits is not None: + model_output_broadcast_data["logits"] = logits.contiguous() + + broadcasted = get_pp_group().broadcast_tensor_dict( + model_output_broadcast_data, src=len(get_pp_group().ranks) - 1 + ) + assert broadcasted is not None + logits = broadcasted["logits"] + + self.execute_model_state = ExecuteModelState( + scheduler_output, + logits, + spec_decode_metadata, + spec_decode_common_attn_metadata, + hidden_states, + sample_hidden_states, + aux_hidden_states, + ec_connector_output, + cudagraph_stats, + slot_mappings, + ) + self.kv_connector_output = kv_connector_output + + # Now the batch has been launched we can wait for corrections from the + # previous model forward without breaking async scheduling. + if deferred_state_corrections_fn: + deferred_state_corrections_fn() + + return None + + def _input_fits_in_drafter( + self, common_attn_metadata: CommonAttentionMetadata | None + ) -> bool: + if common_attn_metadata is None: + return False + assert self.speculative_config is not None + # DFlash queries one extra token (the bonus token) beyond num_spec_tokens + num_drafter_query_tokens = self.num_spec_tokens + ( + 1 if self.speculative_config.use_dflash() else 0 + ) + return ( + common_attn_metadata.max_seq_len + num_drafter_query_tokens + <= self.effective_drafter_max_model_len + ) + + @torch.inference_mode + def sample_tokens( + self, grammar_output: "GrammarOutput | None" + ) -> ModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors: + if self.execute_model_state is None: + kv_connector_output = self.kv_connector_output + self.kv_connector_output = None + # receive sampled token ids from the last PP rank. + if self.use_async_scheduling and not get_pp_group().is_last_rank: + self._pp_receive_prev_sampled_token_ids_to_input_batch() + # In case of PP with kv transfer, we need to pass through the + # kv_connector_output + return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) + + # Unpack ephemeral state. + ( + scheduler_output, + logits, + spec_decode_metadata, + spec_decode_common_attn_metadata, + hidden_states, + sample_hidden_states, + aux_hidden_states, + ec_connector_output, + cudagraph_stats, + slot_mappings, + ) = self.execute_model_state + # Clear ephemeral state. + self.execute_model_state = None + + # Apply structured output bitmasks if present. + if grammar_output is not None: + apply_grammar_bitmask( + scheduler_output, grammar_output, self.input_batch, logits + ) + + with record_function_or_nullcontext("gpu_model_runner: sample"): + sampler_output = self._sample(logits, spec_decode_metadata) + + self._update_states_after_model_execute( + sampler_output.sampled_token_ids, scheduler_output + ) + if self.use_async_scheduling: + pp = get_pp_group() + # For torchrun external_launcher PP mode with broadcast_pp_output=True, + # PP outputs have been broadcasted to all ranks at logits computation. + # Therefore, here is no need to send sampled token ids again in this case. + if not self.broadcast_pp_output and pp.world_size > 1 and pp.is_last_rank: + self._pp_broadcast_prev_sampled_token_ids( + sampler_output.sampled_token_ids + ) + + self._draft_token_ids = None + self._draft_probs = None + self._draft_prob_req_ids = None + self._draft_token_req_ids = None + self.valid_sampled_token_count_gpu = None + self.input_batch.prev_sampled_token_ids = None + + def propose_draft_token_ids(sampled_token_ids): + assert spec_decode_common_attn_metadata is not None + with record_function_or_nullcontext("gpu_model_runner: draft"): + self._draft_token_ids = self.propose_draft_token_ids( + scheduler_output, + sampled_token_ids, + self.input_batch.sampling_metadata, + hidden_states, + sample_hidden_states, + aux_hidden_states, + spec_decode_metadata, + spec_decode_common_attn_metadata, + slot_mappings, + ) + self._copy_draft_token_ids_to_cpu(scheduler_output) + + spec_config = self.speculative_config + draft_after_bookkeeping = False + if spec_config is not None: + # Decide whether to run the drafter or zero out draft tokens. + input_fits_in_drafter = self._input_fits_in_drafter( + spec_decode_common_attn_metadata + ) + # Whether the drafter runs a GPU model forward (and thus carries + # TP/EP/DP collectives), independent of padded-batch timing. + drafter_runs_model_forward = ( + spec_config.use_eagle() + or spec_config.uses_draft_model() + or spec_config.uses_extract_hidden_states() + ) + use_gpu_toks = ( + drafter_runs_model_forward + and not spec_config.disable_padded_drafter_batch + ) + if use_gpu_toks: + # EAGLE/DraftModel speculative decoding can use the GPU sampled tokens + # as inputs, and does not need to wait for bookkeeping to finish. + assert isinstance( + self.drafter, + EagleProposer + | DFlashProposer + | DraftModelProposer + | ExtractHiddenStatesProposer + | Gemma4Proposer, + ) + sampled_token_ids = sampler_output.sampled_token_ids + if input_fits_in_drafter: + propose_draft_token_ids(sampled_token_ids) + else: + if self.valid_sampled_token_count_event is not None: + assert spec_decode_common_attn_metadata is not None + next_token_ids, valid_sampled_tokens_count = ( + self.drafter.prepare_next_token_ids_padded( + sampled_token_ids, + self.requests, + self.input_batch, + self.discard_request_mask.gpu, + ) + ) + self._copy_valid_sampled_token_count( + next_token_ids, valid_sampled_tokens_count + ) + if self.parallel_config.data_parallel_size > 1: + # Prevent hang when DP ranks disagree on input_fits_in_drafter + self.drafter.dummy_run(num_tokens=1) + elif ( + spec_config.use_ngram_gpu() + and not spec_config.disable_padded_drafter_batch + ): + assert isinstance(self.drafter, NgramProposerGPU) + sampled_token_ids = sampler_output.sampled_token_ids + if input_fits_in_drafter: + propose_draft_token_ids(sampled_token_ids) + elif self.valid_sampled_token_count_event is not None: + assert spec_decode_common_attn_metadata is not None + next_token_ids, valid_sampled_tokens_count, _ = ( + self.drafter.update_token_ids_ngram( + sampled_token_ids, + self.input_batch, + self.token_ids_gpu_tensor, + self.num_tokens_no_spec_gpu, + self.discard_request_mask.gpu, + ) + ) + self._copy_valid_sampled_token_count( + next_token_ids, valid_sampled_tokens_count + ) + else: + # These drafters consume CPU sampled tokens, so they run + # after bookkeeping. + draft_after_bookkeeping = True + + if not input_fits_in_drafter: + # Zero out draft tokens so the scheduler doesn't schedule + # stale drafts from the previous step. + # For Nemotron-H: it is necessary to zero out the draft tokens, + # otherwise the stale tokens will corrupt Mamba recurrent + # state and logprobs for sequences near max_model_len. + self._draft_token_ids = torch.zeros( + 1, device=self.device, dtype=torch.int32 + ).expand(len(self.input_batch.req_ids), self.num_spec_tokens) + self._draft_probs = None + self._draft_prob_req_ids = None + self._copy_draft_token_ids_to_cpu(scheduler_output, zeros_only=True) + + with record_function_or_nullcontext("gpu_model_runner: bookkeep"): + ( + num_nans_in_logits, + num_nans_device, + logprobs_lists, + valid_sampled_token_ids, + prompt_logprobs_dict, + req_ids_output_copy, + req_id_to_index_output_copy, + invalid_req_indices, + ) = self._bookkeeping_sync( + scheduler_output, + sampler_output, + logits, + hidden_states, + scheduler_output.total_num_scheduled_tokens, + ) + + if draft_after_bookkeeping: + # ngram and other speculative decoding methods use the sampled + # tokens on the CPU, so they are run after bookkeeping. + if input_fits_in_drafter: + propose_draft_token_ids(valid_sampled_token_ids) + elif ( + drafter_runs_model_forward + and self.parallel_config.data_parallel_size > 1 + ): + # Prevent hang when DP ranks disagree on input_fits_in_drafter + assert isinstance( + self.drafter, + EagleProposer + | DFlashProposer + | DraftModelProposer + | ExtractHiddenStatesProposer + | Gemma4Proposer, + ) + self.drafter.dummy_run(num_tokens=1) + + # Finalize KV connector (wait_for_save + clear metadata) after + # draft model runs. Deferred from target model forward to allow + # draft model to also save its KV cache. + if spec_config is not None: + self.finalize_kv_connector() + + with record_function_or_nullcontext("gpu_model_runner: eplb"): + self.eplb_step() + + # self.kv_connector_output may be modified during drafting + kv_connector_output = self.kv_connector_output + self.kv_connector_output = None + + with record_function_or_nullcontext("gpu_model_runner: ModelRunnerOutput"): + output = ModelRunnerOutput( + req_ids=req_ids_output_copy, + req_id_to_index=req_id_to_index_output_copy, + sampled_token_ids=valid_sampled_token_ids, + logprobs=logprobs_lists, + prompt_logprobs_dict=prompt_logprobs_dict, + kv_connector_output=kv_connector_output, + ec_connector_output=ec_connector_output + if self.supports_mm_inputs + else None, + num_nans_in_logits=num_nans_in_logits, + cudagraph_stats=cudagraph_stats, + routed_experts=None, + ) + + if not self.use_async_scheduling: + if self.routed_experts_initialized: + # Sync path: D2H was issued in ``_bookkeeping_sync`` and + # synchronized by ``_to_list``'s event.synchronize(), so + # the pinned buffers are ready to be wrapped as numpy. + total = scheduler_output.total_num_scheduled_tokens + output.routed_experts = RoutedExpertsLists( + routing_data=self.routed_experts_cpu[:total].numpy(), + slot_mapping=self.routed_experts_slot_mapping_cpu[:total].numpy(), + ) + return output + + with record_function_or_nullcontext( + "gpu_model_runner: AsyncGPUModelRunnerOutput" + ): + # Async path: produce a device-side snapshot that the async + # copy stream can D2H later. Both tensors must be private + # clones because: + # - ``routing_data`` source is the shared capturer buffer, + # which the next forward overwrites on the default stream. + # - ``slot_mapping`` source is our own + # ``routed_experts_slot_mapping_device``, which the + # next ``_prepare_inputs`` overwrites on the default + # stream while the D2H is still pending on the copy + # stream. + # Without clones, the copy stream would read torn data. + routed_experts_snapshot = self.get_routed_experts( + scheduler_output.total_num_scheduled_tokens + ) + + async_output = AsyncGPUModelRunnerOutput( + model_runner_output=output, + sampled_token_ids=sampler_output.sampled_token_ids, + logprobs_tensors=sampler_output.logprobs_tensors, + invalid_req_indices=invalid_req_indices, + async_output_copy_stream=self._get_or_create_async_output_copy_stream(), + vocab_size=self.input_batch.vocab_size, + routed_experts=routed_experts_snapshot, + check_ep_fault=self.check_ep_fault, + num_nans=num_nans_device, + ) + with record_function_or_nullcontext( + "gpu_model_runner: set_async_sampled_token_ids" + ): + # Save ref of sampled_token_ids CPU tensor if the batch contains + # any requests with sampling params that require output ids. + self.input_batch.set_async_sampled_token_ids( + async_output.sampled_token_ids_cpu, + async_output.async_copy_ready_event, + ) + + return async_output + + def _pp_broadcast_prev_sampled_token_ids( + self, sampled_token_ids: torch.Tensor + ) -> None: + """Broadcast sampled token ids (GPU) from last PP stage""" + pp = get_pp_group() + assert pp.is_last_rank + # `prev_sampled_token_ids` is expected to have shape [num_reqs, 1]. + assert sampled_token_ids.dim() == 2 and sampled_token_ids.shape[-1] == 1, ( + "PP+async expects sampled_token_ids to have shape [num_reqs, 1]" + ) + # Skip for chunked prefill: sampled tokens are dummy + # and will be discarded, no need to broadcast. + if not self._is_all_reqs_chunked_prefill(): + torch.distributed.broadcast( + sampled_token_ids, src=pp.rank, group=pp.device_group + ) + + def _pp_receive_prev_sampled_token_ids_to_input_batch(self) -> None: + """Receive sampled token ids broadcast from last PP stage""" + pp = get_pp_group() + assert not pp.is_last_rank + num_reqs = self.input_batch.num_reqs + # `prev_sampled_token_ids` is expected to have shape [num_reqs, 1]. + recv = torch.empty((num_reqs, 1), dtype=torch.int32, device=self.device) + # skip for chunked prefill. + if not self._is_all_reqs_chunked_prefill(): + self._pp_recv_work = torch.distributed.broadcast( + recv, src=pp.last_rank, group=pp.device_group, async_op=True + ) + self.input_batch.prev_sampled_token_ids = recv + + # construct `prev_req_id_to_index` here so `_prepare_input_ids` + # can map req_id -> previous batch row + discard_req_indices = np.nonzero(self.discard_request_mask.np[:num_reqs])[0] + discard_req_indices_set = set(discard_req_indices) + prev_req_id_to_index: dict[str, int] = {} + for i, req_id in enumerate(self.input_batch.req_ids): + if i in discard_req_indices_set: + continue + prev_req_id_to_index[req_id] = i + # PP+async scheduling: advance per-request local cached output length by + # appending a placeholder (-1) token id. + if (req_state := self.requests.get(req_id)) is not None: + req_state.output_token_ids.append(-1) + pos = self.input_batch.num_tokens_no_spec[i] + self.input_batch.is_token_ids[i, pos] = True + self.input_batch.num_tokens_no_spec[i] = pos + 1 + self.input_batch.prev_req_id_to_index = prev_req_id_to_index + + def take_draft_token_ids(self) -> DraftTokenIds | None: + if not self.num_spec_tokens or not self._draft_token_req_ids: + return None + draft_token_ids, req_ids = self._get_draft_token_ids_cpu() + return DraftTokenIds(req_ids, draft_token_ids) + + def _copy_draft_token_ids_to_cpu( + self, scheduler_output: "SchedulerOutput", zeros_only: bool = False + ) -> None: + if torch.is_tensor(self._draft_token_ids): + assert isinstance(self._draft_token_ids, torch.Tensor) + self.prev_num_spec_tokens = self._draft_token_ids.shape[1] + # Check if we need to copy draft tokens to CPU. In async scheduling, + # we only copy when needed for structured output, penalties or bad_words. + if self.use_async_scheduling and not ( + scheduler_output.has_structured_output_requests + or self.input_batch.sampling_metadata.output_token_ids + ): + return + # We must also set the corresponding request ids. + self._draft_token_req_ids = self.input_batch.req_ids.copy() + + draft_token_ids: torch.Tensor = self._draft_token_ids + if not torch.is_tensor(draft_token_ids): + return + assert self.draft_token_ids_event is not None + assert self.draft_token_ids_copy_stream is not None + assert self.draft_token_ids_cpu is not None + default_stream = torch.cuda.current_stream() + num_reqs = draft_token_ids.shape[0] + num_spec_tokens = draft_token_ids.shape[1] + with torch.cuda.stream(self.draft_token_ids_copy_stream): + if not zeros_only: + # Trigger async copy of draft token ids to cpu. + self.draft_token_ids_copy_stream.wait_stream(default_stream) + self.draft_token_ids_cpu[:num_reqs, :num_spec_tokens].copy_( + draft_token_ids, non_blocking=True + ) + else: + # No copy needed, just zero-out cpu tensor. + self.draft_token_ids_cpu[:num_reqs, :num_spec_tokens] = 0 + self.draft_token_ids_event.record() + + def _get_draft_token_ids_cpu(self) -> tuple[list[list[int]], list[str]]: + if isinstance(self._draft_token_ids, list): + return self._draft_token_ids, self.input_batch.req_ids + req_ids = self._draft_token_req_ids + if req_ids is None: + return [], [] + assert self.draft_token_ids_event is not None + assert self.draft_token_ids_cpu is not None + self.draft_token_ids_event.synchronize() + assert isinstance(self._draft_token_ids, torch.Tensor) + num_spec_tokens = self._draft_token_ids.shape[1] + return self.draft_token_ids_cpu[ + : len(req_ids), :num_spec_tokens + ].tolist(), req_ids + + def _copy_valid_sampled_token_count( + self, next_token_ids: torch.Tensor, valid_sampled_tokens_count: torch.Tensor + ) -> None: + if self.valid_sampled_token_count_event is None: + return + + default_stream = torch.cuda.current_stream() + # Initialize a new stream to overlap the copy operation with + # prepare_input of draft model. + with torch.cuda.stream(self.valid_sampled_token_count_copy_stream): + self.valid_sampled_token_count_copy_stream.wait_stream(default_stream) # type: ignore + counts = valid_sampled_tokens_count + counts_cpu = self.valid_sampled_token_count_cpu + assert counts_cpu is not None + counts_cpu[: counts.shape[0]].copy_(counts, non_blocking=True) + self.valid_sampled_token_count_event.record() + + if self.use_async_spec_decode: + # Stash for GPU-side correction in _prepare_inputs. + self.valid_sampled_token_count_gpu = valid_sampled_tokens_count + self.input_batch.prev_sampled_token_ids = next_token_ids.unsqueeze(1) + + def _get_valid_sampled_token_count(self) -> list[int]: + # Wait until valid_sampled_tokens_count is copied to cpu, + prev_sampled_token_ids = self.input_batch.prev_sampled_token_ids + sampled_count_event = self.valid_sampled_token_count_event + if sampled_count_event is None or prev_sampled_token_ids is None: + return [] + + counts_cpu = self.valid_sampled_token_count_cpu + assert counts_cpu is not None + sampled_count_event.synchronize() + return counts_cpu[: prev_sampled_token_ids.shape[0]].tolist() + + def _get_spec_decode_draft_probs( + self, spec_decode_metadata: SpecDecodeMetadata + ) -> torch.Tensor | None: + if self._draft_probs is None or self._draft_prob_req_ids is None: + return None + + row_by_req_id = { + req_id: idx for idx, req_id in enumerate(self._draft_prob_req_ids) + } + draft_probs_rows: list[torch.Tensor] = [] + for req_id, num_draft in zip( + self.input_batch.req_ids, spec_decode_metadata.num_draft_tokens + ): + if num_draft == 0: + continue + row_idx = row_by_req_id.get(req_id) + if row_idx is None: + logger.warning( + "Missing cached draft probabilities for request %s; " + "falling back to legacy speculative rejection behavior.", + req_id, + ) + return None + draft_probs_rows.append(self._draft_probs[row_idx, :num_draft]) + + if not draft_probs_rows: + return None + return torch.cat(draft_probs_rows, dim=0).contiguous() + + def propose_draft_token_ids( + self, + scheduler_output: "SchedulerOutput", + sampled_token_ids: torch.Tensor | list[list[int]], + sampling_metadata: SamplingMetadata, + hidden_states: torch.Tensor, + sample_hidden_states: torch.Tensor, + aux_hidden_states: list[torch.Tensor] | None, + spec_decode_metadata: SpecDecodeMetadata | None, + common_attn_metadata: CommonAttentionMetadata, + slot_mappings: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None, + ) -> list[list[int]] | torch.Tensor: + num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens + spec_config = self.speculative_config + assert spec_config is not None + num_spec_tokens_to_schedule = ( + scheduler_output.resolve_num_spec_tokens_to_schedule(self.num_spec_tokens) + ) + self._draft_probs = None + self._draft_prob_req_ids = None + if spec_config.method == "ngram": + from vllm.v1.spec_decode.ngram_proposer import NgramProposer + + assert isinstance(sampled_token_ids, list) + assert isinstance(self.drafter, NgramProposer) + draft_token_ids = self.drafter.propose( + num_spec_tokens_to_schedule, + sampled_token_ids, + self.input_batch.num_tokens_no_spec, + self.input_batch.token_ids_cpu, + slot_mappings=slot_mappings, + ) + elif spec_config.method == "custom_class": + assert isinstance(sampled_token_ids, list) + draft_token_ids = cast(Any, self.drafter).propose( + sampled_token_ids, + self.input_batch.num_tokens_no_spec, + self.input_batch.token_ids_cpu, + slot_mappings=slot_mappings, + ) + elif spec_config.use_ngram_gpu(): + assert isinstance(self.drafter, NgramProposerGPU) + ( + next_token_ids, + valid_sampled_tokens_count, + valid_sampled_token_ids_gpu, + ) = self.drafter.update_token_ids_ngram( + sampled_token_ids, + self.input_batch, + self.token_ids_gpu_tensor, + self.num_tokens_no_spec_gpu, + self.discard_request_mask.gpu, + ) + self._copy_valid_sampled_token_count( + next_token_ids, valid_sampled_tokens_count + ) + + batch_size = next_token_ids.shape[0] + + draft_token_ids, num_valid_draft_tokens = self.drafter.propose( + num_spec_tokens_to_schedule, + self.num_tokens_no_spec_gpu[:batch_size], + self.token_ids_gpu_tensor[:batch_size], + valid_sampled_token_ids_gpu, + valid_sampled_tokens_count, + ) + + # Cache valid draft counts for scheduler-side trimming. + self._num_valid_draft_tokens = num_valid_draft_tokens + + # Async D2H copy on a dedicated stream. + copy_num_valid_draft_tokens( + self._num_valid_draft_tokens_cpu, + self._num_valid_draft_tokens_copy_stream, + self._num_valid_draft_tokens_event, + self._num_valid_draft_tokens, + self.input_batch.num_reqs, + ) + elif spec_config.method == "suffix": + assert isinstance(sampled_token_ids, list) + assert isinstance(self.drafter, SuffixDecodingProposer) + draft_token_ids = self.drafter.propose( + num_spec_tokens_to_schedule, + self.input_batch, + sampled_token_ids, + slot_mappings=slot_mappings, + ) + elif spec_config.method == "medusa": + assert isinstance(sampled_token_ids, list) + assert isinstance(self.drafter, MedusaProposer) + + if sample_hidden_states.shape[0] == len(sampled_token_ids): + # The input to the target model does not include draft tokens. + hidden_states = sample_hidden_states + else: + indices = [] + offset = 0 + assert spec_decode_metadata is not None, ( + "No spec decode metadata for medusa" + ) + for num_draft, tokens in zip( + spec_decode_metadata.num_draft_tokens, sampled_token_ids + ): + indices.append(offset + len(tokens) - 1) + offset += num_draft + 1 + indices = async_tensor_h2d(indices, device=self.device) + hidden_states = sample_hidden_states[indices] + + draft_token_ids = self.drafter.propose( + num_speculative_tokens=num_spec_tokens_to_schedule, + target_hidden_states=hidden_states, + sampling_metadata=sampling_metadata, + slot_mappings=slot_mappings, + ) + elif spec_config.uses_extract_hidden_states(): + assert isinstance(self.drafter, ExtractHiddenStatesProposer) + assert isinstance(sampled_token_ids, torch.Tensor), ( + "sampled_token_ids should be a torch.Tensor for " + "extract_hidden_states method." + ) + if not self.use_aux_hidden_state_outputs or aux_hidden_states is None: + raise ValueError( + "aux_hidden_states are required when using `extract_hidden_states`" + ) + target_hidden_states = [h[:num_scheduled_tokens] for h in aux_hidden_states] + + draft_token_ids = self.drafter.propose( + num_speculative_tokens=num_spec_tokens_to_schedule, + sampled_token_ids=sampled_token_ids, + target_hidden_states=target_hidden_states, + common_attn_metadata=common_attn_metadata, + slot_mappings=slot_mappings, + ) + next_token_ids, valid_sampled_tokens_count = ( + self.drafter.prepare_next_token_ids_padded( + sampled_token_ids, + self.requests, + self.input_batch, + self.discard_request_mask.gpu, + ) + ) + self._copy_valid_sampled_token_count( + next_token_ids, valid_sampled_tokens_count + ) + + elif ( + spec_config.use_eagle() + or spec_config.use_dflash() + or spec_config.uses_draft_model() + ): + assert isinstance( + self.drafter, + EagleProposer | DFlashProposer | DraftModelProposer | Gemma4Proposer, + ) + + if spec_config.disable_padded_drafter_batch: + # When padded-batch is disabled, the sampled_token_ids should be + # the cpu-side list[list[int]] of valid sampled tokens for each + # request, with invalid requests having empty lists. + assert isinstance(sampled_token_ids, list), ( + "sampled_token_ids should be a python list when" + "padded-batch is disabled." + ) + next_token_ids = self.drafter.prepare_next_token_ids_cpu( + sampled_token_ids, + self.requests, + self.input_batch, + scheduler_output.num_scheduled_tokens, + ) + else: + # When using padded-batch, the sampled_token_ids should be + # the gpu tensor of sampled tokens for each request, of shape + # (num_reqs, num_spec_tokens + 1) with rejected tokens having + # value -1. + assert isinstance(sampled_token_ids, torch.Tensor), ( + "sampled_token_ids should be a torch.Tensor when" + "padded-batch is enabled." + ) + next_token_ids, valid_sampled_tokens_count = ( + self.drafter.prepare_next_token_ids_padded( + sampled_token_ids, + self.requests, + self.input_batch, + self.discard_request_mask.gpu, + ) + ) + self._copy_valid_sampled_token_count( + next_token_ids, valid_sampled_tokens_count + ) + + # Let the target override the hidden state fed to the drafter + # (e.g. DeepSeek V4 MTP needs the pre-hc_head residual). Safe to + # rebind here: hidden_states was already consumed for sampling + # above and is not used again in this branch. + alt = getattr( + self.get_model(), "get_mtp_target_hidden_states", lambda: None + )() + if alt is not None: + hidden_states = alt + + num_rejected_tokens_gpu = None + if spec_decode_metadata is None: + token_indices_to_sample = None + # input_ids can be None for multimodal models. + target_token_ids = self.input_ids.gpu[:num_scheduled_tokens] + target_positions = self._get_positions(num_scheduled_tokens) + if self.use_aux_hidden_state_outputs: + assert aux_hidden_states is not None + target_hidden_states = torch.cat( + [h[:num_scheduled_tokens] for h in aux_hidden_states], dim=-1 + ) + else: + target_hidden_states = hidden_states[:num_scheduled_tokens] + else: + if spec_config.disable_padded_drafter_batch: + token_indices_to_sample = None + common_attn_metadata, token_indices = self.drafter.prepare_inputs( + common_attn_metadata, + sampled_token_ids, + spec_decode_metadata.num_draft_tokens, + ) + target_token_ids = self.input_ids.gpu[token_indices] + target_positions = self._get_positions(token_indices) + if self.use_aux_hidden_state_outputs: + assert aux_hidden_states is not None + target_hidden_states = torch.cat( + [h[token_indices] for h in aux_hidden_states], dim=-1 + ) + else: + target_hidden_states = hidden_states[token_indices] + else: + ( + common_attn_metadata, + token_indices_to_sample, + num_rejected_tokens_gpu, + ) = self.drafter.prepare_inputs_padded( + common_attn_metadata, + spec_decode_metadata, + valid_sampled_tokens_count, + ) + total_num_tokens = common_attn_metadata.num_actual_tokens + # When padding the batch, token_indices is just a range + target_token_ids = self.input_ids.gpu[:total_num_tokens] + target_positions = self._get_positions(total_num_tokens) + if self.use_aux_hidden_state_outputs: + assert aux_hidden_states is not None + target_hidden_states = torch.cat( + [h[:total_num_tokens] for h in aux_hidden_states], dim=-1 + ) + else: + target_hidden_states = hidden_states[:total_num_tokens] + + if self.supports_mm_inputs and self.drafter.supports_mm_inputs: + mm_embed_inputs = self._gather_mm_embeddings( + scheduler_output, + shift_computed_tokens=1, + ) + else: + mm_embed_inputs = None + + draft_token_ids = self.drafter.propose( + num_speculative_tokens=num_spec_tokens_to_schedule, + target_token_ids=target_token_ids, + target_positions=target_positions, + target_hidden_states=target_hidden_states, + next_token_ids=next_token_ids, + token_indices_to_sample=token_indices_to_sample, + sampling_metadata=sampling_metadata, + common_attn_metadata=common_attn_metadata, + mm_embed_inputs=mm_embed_inputs, + num_rejected_tokens_gpu=num_rejected_tokens_gpu, + slot_mappings=slot_mappings, + ) + if hasattr(self.drafter, "take_last_draft_probs"): + draft_probs = self.drafter.take_last_draft_probs() + if draft_probs is not None: + self._draft_probs = draft_probs + self._draft_prob_req_ids = self.input_batch.req_ids.copy() + + return draft_token_ids + + def update_config(self, overrides: dict[str, Any]) -> None: + allowed_config_names = {"load_config", "model_config"} + for config_name, config_overrides in overrides.items(): + if config_name not in allowed_config_names: + allowed = ", ".join(sorted(allowed_config_names)) + raise ValueError( + f"Config override '{config_name}' is not supported. " + f"Supported configs: {allowed}" + ) + config = getattr(self, config_name) + new_config = update_config(config, config_overrides) + setattr(self, config_name, new_config) + + @instrument(span_name="Loading (GPU)") + def load_model(self, load_dummy_weights: bool = False) -> None: + """ + Args: + load_dummy_weights: load dummy weights instead of real weights. + """ + logger.info_once( + "Starting to load model %s...", + self.model_config.model, + scope="global", + ) + + if self.parallel_config.enable_eplb: + self.eplb_state = EplbState(self.parallel_config, self.device) + eplb_models = 0 + + try: + with DeviceMemoryProfiler() as m: + time_before_load = time.perf_counter() + if load_dummy_weights: + self.load_config.load_format = "dummy" + model_loader = get_model_loader(self.load_config) + self.model = model_loader.load_model( + vllm_config=self.vllm_config, model_config=self.model_config + ) + if self.lora_config: + self.model = self.load_lora_model( + self.model, self.vllm_config, self.device + ) + if hasattr(self, "drafter"): + logger.info_once("Loading drafter model...") + if hasattr(self.drafter, "load_model"): + self.drafter.load_model(self.model) + if ( + self.parallel_config.enable_eplb + and hasattr(self.drafter, "model") + and ( + drafter_moe_model := get_mixture_of_experts_model( + self.drafter.model + ) + ) + is not None + ): + assert not self.parallel_config.enable_elastic_ep, ( + "Elastic EP is not supported with drafter model." + ) + spec_config = self.vllm_config.speculative_config + assert spec_config is not None + assert spec_config.draft_model_config is not None + logger.info_once( + "EPLB is enabled for MoE part of drafter model %s.", + spec_config.draft_model_config.model, + ) + if self.eplb_state is None: + self.eplb_state = EplbState( + self.parallel_config, self.device + ) + self.eplb_state.add_model( + drafter_moe_model, + spec_config.draft_model_config, + ) + assert hasattr(self.drafter, "set_eplb_state") + self.drafter.set_eplb_state(self.eplb_state) + eplb_models += 1 + + self._setup_eagle3_aux_hidden_state_outputs() + + # Resolve the MoE model, unwrapping VLM wrappers if needed. + # VLM models (e.g. KimiK25ForConditionalGeneration) wrap the + # actual MoE language model but don't implement + # MixtureOfExperts themselves. + self._moe_model = get_mixture_of_experts_model(self.model) + + if self._moe_model is not None and self.parallel_config.enable_eplb: + logger.info_once( + "EPLB is enabled for MoE part of model %s.", + self.model_config.model, + ) + assert self.eplb_state is not None + self.eplb_state.add_model( + self._moe_model, + self.model_config, + ) + eplb_models += 1 + + time_after_load = time.perf_counter() + self.model_memory_usage = m.consumed_memory + except torch.cuda.OutOfMemoryError as e: + msg = ( + "Failed to load model - not enough GPU memory. " + "Try lowering --gpu-memory-utilization to free memory for weights, " + "increasing --tensor-parallel-size, or using --quantization. " + "See https://docs.vllm.ai/en/latest/configuration/conserving_memory/ " + "for more tips." + ) + combined_msg = f"{msg} (original error: {e})" + logger.error(combined_msg) + raise e + logger.info_once( + "Model loading took %s GiB memory and %.6f seconds", + format_gib(self.model_memory_usage), + time_after_load - time_before_load, + ) + + mm_config = self.model_config.multimodal_config + self.is_multimodal_pruning_enabled = ( + supports_multimodal_pruning(self.get_model()) + and mm_config is not None + and mm_config.is_multimodal_pruning_enabled() + ) + self.requires_sequential_video_encoding = hasattr( + self.get_model(), "requires_sequential_video_encoding" + ) # Temporary hack for dynamic res video w/o support for bs>1 yet + + if ( + self.parallel_config.enable_eplb + and self._moe_model is not None + and not load_dummy_weights + and self.eplb_state is not None + and self.eplb_state.is_async + ): + self.eplb_state.start_async_loop() + + if ( + self.vllm_config.compilation_config.mode + == CompilationMode.STOCK_TORCH_COMPILE + ): + from vllm.env_override import _apply_constrain_to_fx_strides_patch + + _apply_constrain_to_fx_strides_patch() + backend = self.vllm_config.compilation_config.init_backend(self.vllm_config) + compilation_counter.stock_torch_compile_count += 1 + self.model.compile(fullgraph=True, backend=backend) + return + # for other compilation modes, cudagraph behavior is controlled by + # CudagraphWrapper and CudagraphDispatcher of vllm. + + # wrap the model with full cudagraph wrapper if needed. + cudagraph_mode = self.compilation_config.cudagraph_mode + assert cudagraph_mode is not None + if ( + is_breakable_cudagraph_enabled() + and cudagraph_mode != CUDAGraphMode.NONE + and not self.parallel_config.use_ubatching + ): + self.model = BreakableCUDAGraphWrapper(self.model, self.vllm_config) + drafter = getattr(self, "drafter", None) + if drafter is not None and hasattr(drafter, "model"): + drafter.model = BreakableCUDAGraphWrapper( + drafter.model, self.vllm_config + ) + elif ( + cudagraph_mode.has_full_cudagraphs() + and not self.parallel_config.use_ubatching + ): + self.model = CUDAGraphWrapper( + self.model, self.vllm_config, runtime_mode=CUDAGraphMode.FULL + ) + elif self.parallel_config.use_ubatching: + if cudagraph_mode.has_full_cudagraphs(): + self.model = UBatchWrapper( + self.model, self.vllm_config, CUDAGraphMode.FULL, self.device + ) + else: + self.model = UBatchWrapper( + self.model, self.vllm_config, CUDAGraphMode.NONE, self.device + ) + + get_offloader().post_init() + + def _setup_eagle3_aux_hidden_state_outputs(self) -> None: + if not self.use_aux_hidden_state_outputs: + return + + if not supports_eagle3(self.get_model()): + raise RuntimeError( + "Model does not support EAGLE3 interface but " + "aux_hidden_state_outputs was requested" + ) + # Try to get auxiliary layers from speculative config, + # otherwise use model's default layers + aux_layers = self._get_eagle3_aux_layers_from_config() + if aux_layers: + logger.info( + "Using auxiliary layers from speculative config: %s", aux_layers + ) + else: + aux_layers = self.model.get_eagle3_default_aux_hidden_state_layers() + + self.model.set_aux_hidden_state_layers(aux_layers) + + def _get_eagle3_aux_layers_from_config(self) -> tuple[int, ...] | None: + """Extract Eagle3 auxiliary layer indices from speculative config. + + These indices specify which hidden states from the base model should + be used as auxiliary inputs for the Eagle3 drafter model during + speculative decoding. + + Returns: + Tuple of layer indices if found in draft model config, + None otherwise. + """ + if not (self.speculative_config and self.speculative_config.draft_model_config): + return None + + hf_config = self.speculative_config.draft_model_config.hf_config + + layer_ids = getattr(hf_config, "eagle_aux_hidden_state_layer_ids", None) + if not layer_ids: + dflash_config = getattr(hf_config, "dflash_config", None) + eagle_config = getattr(hf_config, "eagle_config", None) + + if dflash_config and isinstance(dflash_config, dict): + # Add 1 to convert DFlash's aux layer id semantics + layer_ids = [ + i + 1 for i in (dflash_config.get("target_layer_ids") or []) + ] + + if eagle_config and isinstance(eagle_config, dict): + layer_ids = eagle_config.get("eagle_aux_hidden_state_layer_ids") + + if layer_ids and isinstance(layer_ids, (list, tuple)): + return tuple(layer_ids) + + return None + + def reload_weights( + self, + weights_iterator: Iterable[tuple[str, torch.Tensor]] | None = None, + weights_path: str | None = None, + is_checkpoint_format: bool = True, + ) -> None: + """ + Reload weights from a weights iterator or from disk + + Args: + weights_iterator: weights to load into model + weights_path: path to load weights from if weights_iterator is not + provided. Use path of original model if neither is provided. + is_checkpoint_format: set to False if weights have already been + processed into kernel format (repacking, renaming, etc.) + """ + # TODO(@kylesayrs): generalize to all runners and loaders + # argument validation + if weights_iterator is None and not is_checkpoint_format: + logger.warning( + "Reloading from disk means that weights will be in checkpoint format. " + "Please use `is_checkpoint_format=True` " + "to avoid weight reloading errors" + ) + + model = self.get_model() + weights_to_load = { + name.replace(".base_layer.", ".") if self.lora_config else name + for name, _ in model.named_parameters() + } + counter_before_reloading = time.perf_counter() + + # load weights from disk if none are provided + if weights_iterator is None: + model_loader = get_model_loader(self.load_config) + if not hasattr(model_loader, "get_all_weights"): + raise NotImplementedError( + f"Model reloading with `{self.load_config.load_format}` format" + ) + + if weights_path is not None: + # The revision belongs to the model we are reloading away from, + # so it must not be carried over to the new path. + self.model_config.model = weights_path + self.model_config.revision = None + weights_iterator = model_loader.get_all_weights(self.model_config, model) + weights_iterator = cast( + Iterable[tuple[str, torch.Tensor]], weights_iterator + ) + + # begin loading weights + logger.info_once("Reloading weights inplace...") + if is_checkpoint_format: + # load weights from checkpoint/ original model format + initialize_layerwise_reload(model) + loaded_weights = model.load_weights(weights_iterator) + finalize_layerwise_reload(model, self.model_config) + + else: + # load weights from kernel format + logger.warning_once( + "Reloading with `is_checkpoint_format=True` requires that " + "weights be in kernel format and already sharded", + ) + loaded_weights = set() + for name, loaded_weight in weights_iterator: + param = _get_parameter_for_reload(model, name) # TODO: buffers? + param.copy_(loaded_weight) + loaded_weights.add(name) + + self.reset_lora_state() + + # logging and validation + counter_after_reloading = time.perf_counter() + diff_seconds = counter_after_reloading - counter_before_reloading + logger.info_once( + "Reloading and processing weights took %.2f seconds", + diff_seconds, + ) + if self.model_config.quantization is None and loaded_weights is not None: + weights_not_loaded = weights_to_load - loaded_weights + if weights_not_loaded: + logger.warning( + "Following weights were not loaded from checkpoint: %s", + weights_not_loaded, + ) + + self.reset_encoder_cache() + self.reset_mm_cache() + + def _get_prompt_logprobs_dict( + self, + hidden_states: torch.Tensor, + num_scheduled_tokens: dict[str, int], + ) -> dict[str, LogprobsTensors | None]: + num_prompt_logprobs_dict = self.num_prompt_logprobs + if not num_prompt_logprobs_dict: + return {} + + prompt_logprobs_dict: dict[str, LogprobsTensors | None] = {} + + # Since prompt logprobs are a rare feature, prioritize simple, + # maintainable loop over optimal performance. + completed_prefill_reqs = [] + for req_id, num_prompt_logprobs in num_prompt_logprobs_dict.items(): + num_tokens = num_scheduled_tokens.get(req_id) + if num_tokens is None: + # This can happen if the request was preempted in prefill stage. + continue + + # Get metadata for this request. + request = self.requests[req_id] + if request.prompt_token_ids is None: + # Prompt logprobs is incompatible with prompt embeddings + continue + + num_prompt_tokens = len(request.prompt_token_ids) + prompt_token_ids = async_tensor_h2d( + request.prompt_token_ids, device=self.device + ) + + # Set up target LogprobsTensors object. + logprobs_tensors = request.in_progress_prompt_logprobs_cpu + if logprobs_tensors is None: + # Create empty logprobs CPU tensors for the entire prompt. + # If chunked, we'll copy in slice by slice. + logprobs_tensors = LogprobsTensors.empty_cpu( + num_prompt_tokens - 1, num_prompt_logprobs + 1 + ) + request.in_progress_prompt_logprobs_cpu = logprobs_tensors + + # Determine number of logits to retrieve. + start_idx = request.num_computed_tokens + start_tok = start_idx + 1 + num_remaining_tokens = num_prompt_tokens - start_tok + if num_tokens <= num_remaining_tokens: + # This is a chunk, more tokens remain. + # In the == case, there are no more prompt logprobs to produce + # but we want to defer returning them to the next step where we + # have new generated tokens to return. + num_logits = num_tokens + else: + # This is the last chunk of prompt tokens to return. + num_logits = num_remaining_tokens + completed_prefill_reqs.append(req_id) + prompt_logprobs_dict[req_id] = logprobs_tensors + + if num_logits <= 0: + # This can happen for the final chunk if we prefilled exactly + # (num_prompt_tokens - 1) tokens for this request in the prior + # step. There are no more prompt logprobs to produce. + continue + + # Get the logits corresponding to this req's prompt tokens. + # If this is a partial request (i.e. chunked prefill), + # then there is prompt logprob generated for each index. + req_idx = self.input_batch.req_id_to_index[req_id] + offset = self.query_start_loc.np[req_idx].item() + prompt_hidden_states = hidden_states[offset : offset + num_logits] + logits = self.model.compute_logits(prompt_hidden_states) + + # Get the "target" tokens for each index. For prompt at index i, + # the token at prompt index i+1 is the "sampled" token we want + # to gather the logprob for. + tgt_token_ids = prompt_token_ids[start_tok : start_tok + num_logits] + + # Compute prompt scores respecting logprobs_mode. + # NOTE: prompt tokens skip sampling processors, so + # processed_* and raw_* yield the same scores here. + if self.model_config.logprobs_mode in ("raw_logits", "processed_logits"): + scores = logits.to(torch.float32) + else: + scores = self.sampler.compute_logprobs(logits) + token_ids, logprobs, ranks, *_ = self.sampler.gather_logprobs( + scores, num_prompt_logprobs, tgt_token_ids + ) + + # Transfer GPU->CPU async. + chunk_slice = slice(start_idx, start_idx + num_logits) + logprobs_tensors.logprob_token_ids[chunk_slice].copy_( + token_ids, non_blocking=True + ) + logprobs_tensors.logprobs[chunk_slice].copy_(logprobs, non_blocking=True) + logprobs_tensors.selected_token_ranks[chunk_slice].copy_( + ranks, non_blocking=True + ) + + # Remove requests that have completed prefill from the batch + # num_prompt_logprobs_dict. + for req_id in completed_prefill_reqs: + del num_prompt_logprobs_dict[req_id] + self.requests[req_id].in_progress_prompt_logprobs_cpu = None + + # Must synchronize the non-blocking GPU->CPU transfers. + if prompt_logprobs_dict: + self._sync_device() + + return prompt_logprobs_dict + + def _get_nans_in_logits(self, logits: torch.Tensor | None) -> dict[str, int]: + """Count NaNs per request, reading the result back to the host. + + Only used under sync scheduling, The async path keeps the counts + on device instead; see`AsyncGPUModelRunnerOutput`. + """ + try: + # Reporting per-request NaN counts requires them on the host; this + # path is opt-in diagnostics, so the D2H is intended. + with gpu_sync_allowed(): + counts = [] if logits is None else count_nans_per_row(logits).tolist() + num_nans_in_logits = nans_to_dict(counts, self.input_batch.req_id_to_index) + if envs.VLLM_RAISE_ON_LOGIT_NANS: + raise_if_nan_logits(num_nans_in_logits) + return num_nans_in_logits + except IndexError: + return {} + + @contextmanager + def maybe_randomize_inputs( + self, + input_ids: torch.Tensor | None, + inputs_embeds: torch.Tensor | None, + randomize_inputs: bool = False, + ): + """ + Randomize input_ids if VLLM_RANDOMIZE_DP_DUMMY_INPUTS is set. + This is to help balance expert-selection + - during profile_run + - during DP rank dummy run + """ + + dp_size = self.vllm_config.parallel_config.data_parallel_size + randomize_inputs = randomize_inputs or ( + envs.VLLM_RANDOMIZE_DP_DUMMY_INPUTS and dp_size > 1 + ) + if not randomize_inputs: + yield + elif input_ids is not None: + + @functools.cache + def rand_input_ids() -> torch.Tensor: + return torch.randint_like( + self.input_ids.gpu, + low=0, + high=self.model_config.get_vocab_size(), + ) + + logger.debug_once("Randomizing dummy input_ids for DP Rank") + input_ids.copy_(rand_input_ids()[: input_ids.size(0)], non_blocking=True) + yield + input_ids.fill_(0) + else: + + @functools.cache + def rand_inputs_embeds() -> torch.Tensor: + return torch.randn_like( + self.inputs_embeds.gpu, + ) + + assert inputs_embeds is not None + logger.debug_once("Randomizing dummy inputs_embeds for DP Rank") + inputs_embeds.copy_( + rand_inputs_embeds()[: inputs_embeds.size(0)], non_blocking=True + ) + yield + inputs_embeds.fill_(0) + + def _get_mm_dummy_batch( + self, + modality: str, + max_items_per_batch: int, + ) -> BatchedTensorInputs: + """Dummy data for profiling and precompiling multimodal models.""" + assert self.mm_budget is not None + + # Don't use `max_items_per_batch` here to avoid redundant computation + dummy_mm_inputs = self.mm_registry.get_dummy_mm_inputs( + self.model_config, + mm_counts={modality: 1}, + cache=self.mm_budget.cache, + ) + dummy_mm_item = dummy_mm_inputs["mm_kwargs"][modality][0] + + # We use the cache so that the item is saved to the cache, + # but not read from the cache + assert dummy_mm_item is not None, "Item should not already be cached" + + return next( + mm_kwargs_batch + for _, _, mm_kwargs_batch in group_and_batch_mm_kwargs( + [(modality, dummy_mm_item)] * max_items_per_batch, + device=self.device, + pin_memory=PIN_MEMORY, + ) + ) + + @torch.inference_mode() + def _dummy_run( + self, + num_tokens: int, + cudagraph_runtime_mode: CUDAGraphMode | None = None, + force_attention: bool = False, + uniform_decode: bool = False, + allow_microbatching: bool = True, + skip_eplb: bool = False, + is_profile: bool = False, + create_mixed_batch: bool = False, + remove_lora: bool = True, + is_graph_capturing: bool = False, + num_active_loras: int = 0, + profile_seq_lens: int | None = None, + randomize_inputs: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Run a dummy forward pass to warm up/profile run or capture the + CUDA graph for the model. + + Args: + num_tokens: Number of tokens to run the dummy forward pass. + cudagraph_runtime_mode: used to control the behavior. + - if not set will determine the cudagraph mode based on using + the self.cudagraph_dispatcher. + - CUDAGraphMode.NONE: No cudagraph, for warm up and profile run + - CUDAGraphMode.PIECEWISE: Piecewise cudagraph. + - CUDAGraphMode.FULL: Full cudagraph, attention metadata is + needed. + force_attention: If True, always create attention metadata. Used to + warm up attention backend when mode is NONE. + uniform_decode: If True, the batch is a uniform decode batch. + skip_eplb: If True, skip EPLB state update. + is_profile: If True, this is a profile run. + create_mixed_batch: If True, create a mixed batch with both decode + (1 token) and prefill (multiple tokens) requests. + remove_lora: If False, dummy LoRAs are not destroyed after the run + num_active_loras: Number of distinct active LoRAs to capture for. + LoRA is activated when num_active_loras > 0. + profile_seq_lens: If provided, use this value for seq_lens instead + of max_query_len. Used to profile attention workspace that + scales with context length. + """ + mm_config = self.vllm_config.model_config.multimodal_config + if mm_config and mm_config.mm_encoder_only: + # The current dummy run only covers LM execution, so we can skip it. + # mm encoder dummy run may need to add in the future. + return torch.tensor([]), torch.tensor([]) + + assert ( + cudagraph_runtime_mode is None + or cudagraph_runtime_mode.is_valid_runtime_mode() + ) + + # If cudagraph_mode.decode_mode() == FULL and + # cudagraph_mode.separate_routine(). This means that we are using + # different graphs and/or modes for mixed prefill-decode batches vs. + # uniform decode batches. A uniform decode batch means that all + # requests have identical query length, except a potential virtual + # request (shorter) in the batch account for padding. + # Uniform decode batch could either be common pure decode, where + # max_query_len == 1, or speculative decode, where + # max_query_len == 1 + num_spec_decode_tokens. + + # When setting max_query_len = 1, we switch to and capture the optimized + # routine of FA2 for pure decode, i.e., Flashdecode + an optimization + # for GQA/MQA. + max_query_len = self.uniform_decode_query_len if uniform_decode else num_tokens + + # Set num_scheduled_tokens based on num_tokens and max_num_seqs + # for dummy run with LoRA so that the num_reqs collectively + # has num_tokens in total. + assert num_tokens <= self.max_num_tokens + max_num_reqs = self.scheduler_config.max_num_seqs + if create_mixed_batch: + assert not uniform_decode + # Create mixed batch: + # first half decode tokens, second half one prefill + num_decode_tokens = min(max_num_reqs - 1, num_tokens // 2) + num_prefill_tokens = num_tokens - num_decode_tokens + num_reqs = num_decode_tokens + 1 + + # Create decode requests (1 token each) followed by prefill request + num_scheduled_tokens_list = [1] * num_decode_tokens + [num_prefill_tokens] + # Note: Overriding max_query_len to be the prefill tokens + max_query_len = num_prefill_tokens + elif uniform_decode: + assert not create_mixed_batch + num_reqs = min(max_num_reqs, cdiv(num_tokens, max_query_len)) + num_scheduled_tokens_list = [max_query_len] * num_reqs + if num_tokens % max_query_len != 0: + num_scheduled_tokens_list[-1] = num_tokens % max_query_len + else: + num_reqs = min(num_tokens, max_num_reqs) + min_tokens_per_req = num_tokens // num_reqs + num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs + num_scheduled_tokens_list[-1] += num_tokens % num_reqs + + assert sum(num_scheduled_tokens_list) == num_tokens + assert len(num_scheduled_tokens_list) == num_reqs + num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32) + num_tokens_unpadded = int(num_scheduled_tokens.sum()) + + num_sampled_tokens = np.ones(num_reqs, dtype=np.int32) + + _cudagraph_mode, batch_desc, should_ubatch, num_tokens_across_dp, _ = ( + self._determine_batch_execution_and_padding( + num_tokens=num_tokens_unpadded, + num_reqs=num_reqs, + num_scheduled_tokens_np=num_scheduled_tokens, + max_num_scheduled_tokens=max_query_len, + use_cascade_attn=False, + allow_microbatching=allow_microbatching, + force_eager=is_profile + or (cudagraph_runtime_mode == CUDAGraphMode.NONE), + # `force_uniform_decode` is used for cudagraph capture; because for + # capturing mixed prefill-decode batches, we sometimes use + # num_tokens == num_reqs which looks like a uniform decode batch to the + # dispatcher; but we actually want to capture a piecewise cudagraph + force_uniform_decode=uniform_decode, + # `force_has_lora` is used for cudagraph capture; because LoRA is + # activated later in the context manager, but we need to know the + # LoRA state when determining the batch descriptor for capture + force_has_lora=num_active_loras > 0, + # `force_num_active_loras` is used for cudagraph capture; because we + # need to capture graphs for specific num_active_loras counts + force_num_active_loras=num_active_loras, + ) + ) + + if cudagraph_runtime_mode is None: + cudagraph_runtime_mode = _cudagraph_mode + else: + assert cudagraph_runtime_mode == _cudagraph_mode, ( + f"Cudagraph runtime mode mismatch in dummy_run. " + f"Expected {_cudagraph_mode}, but got {cudagraph_runtime_mode}." + ) + + num_tokens_padded = batch_desc.num_tokens + num_reqs_padded = ( + batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs + ) + dcp_dummy_context_len = get_dcp_dummy_context_len( + self.dcp_world_size, + self.parallel_config.cp_kv_cache_interleave_size, + hasattr(self, "kv_cache_config"), + create_mixed_batch, + is_graph_capturing, + uniform_decode, + ) + ubatch_slices, ubatch_slices_padded = maybe_create_ubatch_slices( + should_ubatch, + num_scheduled_tokens, + num_tokens_padded, + num_reqs_padded, + self.vllm_config.parallel_config.num_ubatches, + ) + logger.debug( + "ubatch_slices: %s, ubatch_slices_padded: %s", + ubatch_slices, + ubatch_slices_padded, + ) + + attn_metadata: PerLayerAttnMetadata | None = None + + slot_mappings_by_group, slot_mappings = self._get_slot_mappings( + num_tokens_padded=num_tokens_padded, + num_reqs_padded=num_reqs_padded, + num_tokens_unpadded=num_tokens_unpadded, + ubatch_slices=ubatch_slices_padded, + ) + + # Dummy runs have no real slot assignments — fill with -1 so + # concat_and_cache kernels skip the KV write. + if slot_mappings_by_group is not None: + for sm in slot_mappings_by_group.values(): + sm.fill_(-1) + + # _dummy_run shares pinned CPU buffers (seq_lens, query_start_loc, + # etc.) with execute_model. It must participate in the same event + # protocol so that back-to-back dummy/real steps don't overwrite + # pinned memory while a prior non_blocking H2D DMA is still reading. + with self.synchronize_input_prep(): + # If force_attention is True, we always capture attention. + # Otherwise, it only happens for cudagraph_runtime_mode=FULL. + if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL: + if profile_seq_lens is not None: + seq_lens = profile_seq_lens # type: ignore[assignment] + elif create_mixed_batch: + # In the mixed batch mode (used for FI warmup), we use + # shorter sequence lengths to run faster. + # TODO(luka) better system for describing dummy batches + if dcp_dummy_context_len > 0: + seq_lens = torch.tensor( # type: ignore[assignment] + [1 + dcp_dummy_context_len] * num_decode_tokens + + [num_prefill_tokens + dcp_dummy_context_len], + dtype=torch.int, + ) + else: + seq_lens = torch.tensor( # type: ignore[assignment] + [1] * num_decode_tokens + [num_prefill_tokens + 1], + dtype=torch.int, + ) + elif dcp_dummy_context_len > 0: + seq_lens = max_query_len + dcp_dummy_context_len # type: ignore[assignment] + else: + seq_lens = max_query_len # type: ignore[assignment] + self.optimistic_seq_lens_cpu[:num_reqs] = seq_lens + self.optimistic_seq_lens_cpu[num_reqs:].fill_(0) + self.seq_lens.copy_(self.optimistic_seq_lens_cpu, non_blocking=True) + + cum_num_tokens = self._get_cumsum_and_arange( + num_scheduled_tokens, self.query_pos.np + ) + self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens + self.query_start_loc.np[num_reqs + 1 : num_reqs_padded + 1].fill( + cum_num_tokens[-1] + ) + self.query_start_loc.copy_to_gpu() + + prepare_dcp_dummy_context_metadata( + input_batch=self.input_batch, + kv_cache_config=getattr(self, "kv_cache_config", None), + query_pos=self.query_pos, + positions=self.positions, + query_start_loc=self.query_start_loc, + num_reqs=num_reqs, + num_tokens_unpadded=num_tokens_unpadded, + dcp_dummy_context_len=dcp_dummy_context_len, + ) + + # Sync block table CPU->GPU so cleared rows from + # remove_request() are visible to the attention metadata + # builder. Without this, stale block IDs from finished + # requests can corrupt Mamba state. + self.input_batch.block_table.commit_block_table(num_reqs_padded) + + pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL + attn_metadata, _ = self._build_attention_metadata( + num_tokens=num_tokens_unpadded, + num_tokens_padded=num_tokens_padded if pad_attn else None, + num_reqs=num_reqs_padded, + max_query_len=max_query_len, + ubatch_slices=(ubatch_slices_padded if pad_attn else ubatch_slices), + # FULL replay reads capture-time metadata buffers. Re-stage them + # from the zeroed dummy block tables instead of retaining state + # indices from the previous real batch. + for_cudagraph_capture=( + is_graph_capturing + or cudagraph_runtime_mode == CUDAGraphMode.FULL + ), + slot_mappings=slot_mappings_by_group, + use_spec_decode=self.speculative_config is not None, + ) + + with self.maybe_dummy_run_with_lora( + self.lora_config, + num_scheduled_tokens, + num_sampled_tokens, + remove_lora, + num_active_loras, + ): + # Make sure padding doesn't exceed max_num_tokens + assert num_tokens_padded <= self.max_num_tokens + model_kwargs = self._init_model_kwargs() + if self.supports_mm_inputs and not self.model_config.is_encoder_decoder: + input_ids, inputs_embeds = self._prepare_mm_inputs(num_tokens_padded) + + model_kwargs = { + **model_kwargs, + **self._dummy_mm_kwargs(num_reqs), + } + elif self.enable_prompt_embeds: + input_ids = None + inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded] + model_kwargs = self._init_model_kwargs() + else: + input_ids = self.input_ids.gpu[:num_tokens_padded] + inputs_embeds = None + + if self.uses_mrope: + positions = self.mrope_positions.gpu[:, :num_tokens_padded] + elif self.uses_xdrope_dim > 0: + positions = self.xdrope_positions.gpu[:, :num_tokens_padded] + else: + positions = self.positions[:num_tokens_padded] + + if get_pp_group().is_first_rank: + intermediate_tensors = None + else: + if self.intermediate_tensors is None: + self.intermediate_tensors = ( + self.model.make_empty_intermediate_tensors( + batch_size=self.max_num_tokens, + dtype=self.model_config.dtype, + device=self.device, + ) + ) + + intermediate_tensors = self.sync_and_gather_intermediate_tensors( + num_tokens_padded, None, False + ) + + if ubatch_slices_padded is not None: + # Adjust values to reflect a single ubatch. + # TODO(sage,lucas): this is cruft that should be addressed in + # the padding refactor. + num_tokens_padded = ubatch_slices_padded[0].num_tokens + if num_tokens_across_dp is not None: + num_tokens_across_dp[:] = num_tokens_padded + + with ( + self.maybe_randomize_inputs( + input_ids, inputs_embeds, randomize_inputs=randomize_inputs + ), + set_forward_context( + attn_metadata, + self.vllm_config, + num_tokens=num_tokens_padded, + num_tokens_across_dp=num_tokens_across_dp, + cudagraph_runtime_mode=cudagraph_runtime_mode, + batch_descriptor=batch_desc, + ubatch_slices=ubatch_slices_padded, + slot_mapping=slot_mappings, + ), + ): + outputs = self.model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + **model_kwargs, + ) + + if self.use_aux_hidden_state_outputs: + hidden_states, _ = outputs + else: + hidden_states = outputs + + if self.speculative_config and ( + self.speculative_config.use_eagle() + or self.speculative_config.uses_draft_model() + or self.speculative_config.uses_extract_hidden_states() + ): + assert isinstance( + self.drafter, + EagleProposer + | DFlashProposer + | DraftModelProposer + | ExtractHiddenStatesProposer + | Gemma4Proposer, + ) + assert self.speculative_config is not None + # Eagle currently only supports PIECEWISE cudagraphs. + # Therefore only use cudagraphs if the main model uses PIECEWISE + # NOTE(lucas): this is a hack, need to clean up. + use_cudagraphs = ( + ( + is_graph_capturing + and cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE + ) + or ( + not is_graph_capturing + and cudagraph_runtime_mode != CUDAGraphMode.NONE + ) + ) and not self.speculative_config.enforce_eager + + # Note(gnovack) - We need to disable cudagraphs for one of the two + # lora cases when cudagraph_specialize_lora is enabled. This is a + # short term mitigation for issue mentioned in + # https://github.com/vllm-project/vllm/issues/28334 + if ( + self.compilation_config.cudagraph_specialize_lora + and num_active_loras > 0 + ): + use_cudagraphs = False + + self.drafter.dummy_run( + num_tokens, + use_cudagraphs=use_cudagraphs, + is_graph_capturing=is_graph_capturing, + slot_mappings=slot_mappings, + ) + + # We register layerwise NVTX hooks here after the first dynamo tracing is + # done to avoid nvtx operations in hook functions being traced by + # torch dynamo and causing graph breaks. + # Note that for DYNAMO_ONCE and VLLM_COMPILE mode, + # compiled model's dynamo tracing is only done once and the compiled model's + # __call__ function is replaced by calling the compiled function. + # So it's safe to register hooks here. Hooks will be registered to + # both compiled and uncompiled models but they will never + # be called on the compiled model execution path. + self._register_layerwise_nvtx_hooks() + + # This is necessary to avoid blocking DP. + # For dummy runs, we typically skip EPLB since we don't have any real + # requests to process. + # However, in DP settings, there may be cases when some DP ranks do + # not have any requests to process, so they're executing dummy batches. + # In such cases, we still have to trigger EPLB to make sure + # ranks execute the rearrangement in synchronization. + if not skip_eplb: + self.eplb_step(is_dummy=True, is_profile=is_profile) + + logit_indices = np.cumsum(num_scheduled_tokens) - 1 + logit_indices_device = torch.from_numpy(logit_indices).to( + self.device, non_blocking=True + ) + return hidden_states, hidden_states[logit_indices_device] + + @torch.inference_mode() + def _dummy_sampler_run( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + # The dummy hidden states may contain special values, + # like `inf` or `nan`. + # To avoid breaking the sampler, we use a random tensor here instead. + + mm_config = self.vllm_config.model_config.multimodal_config + if mm_config and mm_config.mm_encoder_only: + # MM Encoder only model no need to run sampler. + return torch.tensor([]) + + hidden_states = torch.rand_like(hidden_states) + + logits = self.model.compute_logits(hidden_states) + num_reqs = logits.size(0) + + dummy_tensors = lambda v: torch.full((num_reqs,), v, device=self.device) + + dummy_metadata = SamplingMetadata( + temperature=dummy_tensors(0.5), + all_greedy=False, + all_random=False, + top_p=dummy_tensors(0.9), + top_k=dummy_tensors(logits.size(1) - 1), + generators={}, + max_num_logprobs=None, + logprob_token_ids=None, + no_penalties=True, + prompt_token_ids=None, + frequency_penalties=dummy_tensors(0.1), + presence_penalties=dummy_tensors(0.1), + repetition_penalties=dummy_tensors(0.1), + output_token_ids=[[] for _ in range(num_reqs)], + spec_token_ids=[[] for _ in range(num_reqs)], + allowed_token_ids_mask=None, + bad_words_token_ids={}, + logitsprocs=LogitsProcessors(), + ) + try: + sampler_output = self.sampler( + logits=logits, sampling_metadata=dummy_metadata + ) + # Also warm forward_native (taken when generators dict is non-empty), + # but skip the extra call in 'processed_logits' / 'processed_logprobs' + # modes — there TopKTopPSampler binds forward = forward_native at + # init time, so the warmup call is redundant and only inflates peak + # memory during profile_run. + # No .clone() of logits: warmup output is discarded, so any in-place + # mutation by forward_native does not affect correctness. + if self.sampler.logprobs_mode not in PROCESSED_LOGPROBS_MODES: + self.sampler( + logits=logits, + sampling_metadata=replace( + dummy_metadata, + generators={ + 0: torch.Generator(device=self.device).manual_seed(0) + }, + ), + ) + except RuntimeError as e: + if "out of memory" in str(e): + raise RuntimeError( + "CUDA out of memory occurred when warming up sampler with " + f"{num_reqs} dummy requests. Please try lowering " + "`max_num_seqs` or `gpu_memory_utilization` when " + "initializing the engine." + ) from e + else: + raise e + if self.speculative_config: + draft_token_ids = [[0] for _ in range(num_reqs)] + dummy_spec_decode_metadata = SpecDecodeMetadata.make_dummy( + draft_token_ids, self.device + ) + + num_tokens = sum(len(ids) for ids in draft_token_ids) + draft_probs = None + if ( + self.speculative_config.rejection_sample_method == "standard" + and self.speculative_config.draft_sample_method == "probabilistic" + ): + draft_probs = torch.rand( + num_tokens, + logits.shape[-1], + device=self.device, + dtype=torch.float32, + ) + draft_probs = torch.softmax(draft_probs, dim=-1) + logits = torch.randn( + num_tokens + num_reqs, + logits.shape[-1], + device=self.device, + dtype=logits.dtype, + ) + self.rejection_sampler( + dummy_spec_decode_metadata, + draft_probs, + logits, + dummy_metadata, + ) + # All-greedy is a separate Triton specialization from the + # mixed-sampling path above. Compile and exercise it before serving. + all_greedy_metadata = replace( + dummy_metadata, + temperature=None, + all_greedy=True, + all_random=False, + ) + self.rejection_sampler( + dummy_spec_decode_metadata, + draft_probs, + logits, + all_greedy_metadata, + ) + torch.accelerator.synchronize() + return sampler_output + + def _dummy_pooler_run_task( + self, + hidden_states: torch.Tensor, + task: PoolingTask, + ) -> PoolerOutput: + num_tokens = hidden_states.shape[0] + max_num_reqs = self.scheduler_config.max_num_seqs + num_reqs = min(num_tokens, max_num_reqs) + min_tokens_per_req = num_tokens // num_reqs + num_scheduled_tokens_np = np.full(num_reqs, min_tokens_per_req) + num_scheduled_tokens_np[-1] += num_tokens % num_reqs + assert np.sum(num_scheduled_tokens_np) == num_tokens + assert len(num_scheduled_tokens_np) == num_reqs + + req_num_tokens = num_tokens // num_reqs + + dummy_prompt_lens = torch.from_numpy(num_scheduled_tokens_np) + dummy_token_ids = torch.zeros( + (num_reqs, req_num_tokens), dtype=torch.int32, device=self.device + ) + + model = cast(VllmModelForPooling, self.get_model()) + dummy_pooling_params = PoolingParams(task=task) + dummy_pooling_params.verify(self.model_config) + to_update = model.pooler.get_pooling_updates(task) + to_update.apply(dummy_pooling_params) + + dummy_metadata = PoolingMetadata( + prompt_lens=dummy_prompt_lens, + prompt_token_ids=dummy_token_ids, + prompt_token_ids_cpu=dummy_token_ids.cpu(), + pooling_params=[dummy_pooling_params] * num_reqs, + pooling_states=[PoolingStates() for i in range(num_reqs)], + ) + + dummy_metadata.build_pooling_cursor( + num_scheduled_tokens_np, + seq_lens_cpu=dummy_prompt_lens, + device=hidden_states.device, + ) + + try: + return model.pooler( + hidden_states=hidden_states, pooling_metadata=dummy_metadata + ) + except RuntimeError as e: + if "out of memory" in str(e): + raise RuntimeError( + "CUDA out of memory occurred when warming up pooler " + f"({task=}) with {num_reqs} dummy requests. Please try " + "lowering `max_num_seqs` or `gpu_memory_utilization` when " + "initializing the engine." + ) from e + else: + raise e + + @torch.inference_mode() + def _dummy_pooler_run( + self, + hidden_states: torch.Tensor, + ) -> PoolerOutput: + mm_config = self.vllm_config.model_config.multimodal_config + if mm_config and mm_config.mm_encoder_only: + # MM Encoder only model not need to run pooler. + return torch.tensor([]) + + # Find the task that has the largest output for subsequent steps + supported_pooling_tasks = self.get_supported_pooling_tasks() + + if not supported_pooling_tasks: + raise RuntimeError( + f"Model {self.model_config.model} does not support " + "any pooling tasks. See " + "https://docs.vllm.ai/en/latest/models/pooling_models.html " + "to learn more." + ) + + output_size = dict[PoolingTask, float]() + for task in supported_pooling_tasks: + # Run a full batch with each task to ensure none of them OOMs + output = self._dummy_pooler_run_task(hidden_states, task) + output_size[task] = sum(o.nbytes for o in output if o is not None) + del output # Allow GC + + max_task = max(output_size.items(), key=lambda x: x[1])[0] + return self._dummy_pooler_run_task(hidden_states, max_task) + + def profile_run(self) -> None: + # Profile with multimodal encoder & encoder cache. + if self.supports_mm_inputs: + mm_config = self.model_config.multimodal_config + if mm_config is not None and mm_config.skip_mm_profiling: + logger.info( + "Skipping memory profiling for multimodal encoder and " + "encoder cache." + ) + else: + mm_budget = self.mm_budget + assert mm_budget is not None + + if (encoder_budget := mm_budget.get_encoder_budget()) > 0: + if not mm_budget.mm_max_toks_per_item: + # All modality limits are 0 — embedding-only mode. + # Budget is non-zero for embedding storage, but + # there's no encoder to profile. + logger.info( + "Skipping encoder profiling for embedding-only " + "mode (all modality limits=0 with " + "enable_mm_embeds=True).", + ) + else: + # NOTE: Currently model is profiled with a single + # non-text modality with the max possible input + # tokens even when it supports multiple. + dummy_modality = mm_budget.get_modality_with_max_tokens() + max_mm_items_per_batch = mm_budget.mm_max_items_per_batch[ + dummy_modality + ] + + logger.info_once( + "Encoder cache will be initialized with a " + "budget of %s tokens, and profiled with " + "%s %s items of the maximum feature size.", + encoder_budget, + max_mm_items_per_batch, + dummy_modality, + ) + + # Create dummy batch of multimodal inputs. + batched_dummy_mm_inputs = self._get_mm_dummy_batch( + dummy_modality, + max_mm_items_per_batch, + ) + + # Run multimodal encoder. + dummy_encoder_outputs = self.model.embed_multimodal( + **batched_dummy_mm_inputs + ) + + sanity_check_mm_encoder_outputs( + dummy_encoder_outputs, + expected_num_items=max_mm_items_per_batch, + ) + for i, output in enumerate(dummy_encoder_outputs): + self.encoder_cache[f"tmp_{i}"] = output + + # Add `is_profile` here to pre-allocate communication buffers + hidden_states, last_hidden_states = self._dummy_run( + self.max_num_tokens, is_profile=True + ) + if get_pp_group().is_last_rank: + if self.is_pooling_model: + output = self._dummy_pooler_run(hidden_states) + else: + output = self._dummy_sampler_run(last_hidden_states) + else: + output = None + self._sync_device() + del hidden_states, output + self.encoder_cache.clear() + gc.collect() + + def _init_minimal_kv_cache_for_profiling(self) -> None: + from vllm.v1.core.kv_cache_utils import ( + get_kv_cache_config_from_groups, + get_kv_cache_groups, + ) + + kv_cache_spec = self.get_kv_cache_spec() + KVCacheSpecRegistry.check_kv_cache_spec_registry(kv_cache_spec) + kv_cache_groups = get_kv_cache_groups(self.vllm_config, kv_cache_spec) + # the minimum number of blocks required is 1 block *per sequence* + min_blocks = ( + min(self.max_num_reqs, self.compilation_config.max_cudagraph_capture_size) + or 1 + ) + + # Temporarily change num_gpu_blocks_override to allocate a minimal KV cache + saved_override = self.cache_config.num_gpu_blocks_override + self.cache_config.num_gpu_blocks_override = min_blocks + try: + minimal_config = get_kv_cache_config_from_groups( + self.vllm_config, kv_cache_groups, available_memory=0 + ) + finally: + self.cache_config.num_gpu_blocks_override = saved_override + + self.initialize_kv_cache(minimal_config, is_profiling=True) + self.cache_config.num_gpu_blocks = minimal_config.num_blocks + + logger.debug("Initialized minimal KV cache for CUDA graph profiling") + + @staticmethod + @contextmanager + def _freeze_gc(): + gc_was_enabled = gc.isenabled() + gc.collect() + should_freeze = not envs.VLLM_ENABLE_CUDAGRAPH_GC + if should_freeze: + gc.freeze() + # A Triton kernel finalized during stream capture unloads its + # module and invalidates the captured graph. + gc.disable() + try: + yield + finally: + if should_freeze: + try: + gc.unfreeze() + gc.collect() + finally: + if gc_was_enabled: + gc.enable() + else: + gc.disable() + + def shutdown(self) -> None: + """Release GPU tensors (model weights, KV caches, workspace) so that + memory is reclaimable when running in the same process.""" + from vllm.model_executor.layers.rotary_embedding import _ROPE_DICT + from vllm.v1.worker.workspace import reset_workspace_manager + + # Calls torch.accelerator.synchronize() + self._cleanup_profiling_kv_cache() + if current_platform.is_rocm(): + # Drop captured graphs before distributed teardown. On ROCm, delayed + # graph destruction can surface HSA faults in the next engine startup. + CUDAGraphWrapper.clear_all_graphs() + BreakableCUDAGraphWrapper.clear_all_graphs() + self.encoder_cudagraph_manager = None + self.compilation_config.static_forward_context.clear() + self.model = None # type: ignore[assignment] + _ROPE_DICT.clear() + + reset_workspace_manager() + if current_platform.is_rocm() or current_platform.is_xpu(): + gc.collect() + torch.accelerator.empty_cache() + torch.accelerator.synchronize() + + def _cleanup_profiling_kv_cache(self) -> None: + torch.accelerator.synchronize() + if hasattr(self, "kv_caches") and self.kv_caches: + for i in range(len(self.kv_caches)): + self.kv_caches[i] = None # type: ignore + self.kv_caches.clear() + if hasattr(self, "attn_groups"): + self.attn_groups.clear() + if hasattr(self, "drafter") and hasattr(self.drafter, "draft_attn_groups"): + self.drafter.draft_attn_groups.clear() + if hasattr(self, "kv_cache_config"): + delattr(self, "kv_cache_config") + self.cache_config.num_gpu_blocks = None + + unbind_kv_cache(self.compilation_config.static_forward_context) + self._mamba_bufs = None + + gc.collect() + torch.accelerator.empty_cache() + + logger.debug("Cleaned up profiling KV cache and CUDA graphs") + + @torch.inference_mode() + def _create_encoder_cudagraph_manager(self) -> "EncoderCudaGraphManager | None": + if not ( + self.compilation_config.cudagraph_mm_encoder and self.supports_mm_inputs + ): + return None + + # Use get_model() to unwrap CUDAGraphWrapper/UBatchWrapper, because + # @runtime_checkable Protocol isinstance() checks do not work through + # __getattr__ forwarding. + from vllm.model_executor.models.interfaces import ( + SupportsEncoderCudaGraph, + supports_encoder_cudagraph, + ) + from vllm.v1.worker.encoder_cudagraph import ( + EncoderCudaGraphManager, + ) + + raw_model = self.get_model() + if not supports_encoder_cudagraph(raw_model): + return None + + return EncoderCudaGraphManager( + vllm_config=self.vllm_config, + device=self.device, + dtype=self.dtype, + model=cast(SupportsEncoderCudaGraph, raw_model), + ) + + @torch.inference_mode() + def _maybe_init_encoder_cudagraph_manager(self) -> None: + if self.encoder_cudagraph_manager is None: + self.encoder_cudagraph_manager = self._create_encoder_cudagraph_manager() + if self.encoder_cudagraph_manager is not None: + logger.info("Initialized EncoderCudaGraphManager for vision encoder") + + @torch.inference_mode() + def profile_cudagraph_memory(self) -> int: + profiling_state_initialized = False + try: + with set_current_vllm_config(self.vllm_config): + self._init_minimal_kv_cache_for_profiling() + profiling_state_initialized = True + finally: + if not profiling_state_initialized: + self._cleanup_profiling_kv_cache() + + saved_num_cudagraph_captured = compilation_counter.num_cudagraph_captured + + capture_descs = self.cudagraph_dispatcher.get_capture_descs() + # Use a temporary manager for memory profiling. The persistent manager + # is initialized later so it does not keep profiling-only graph state. + encoder_cudagraph_manager = self._create_encoder_cudagraph_manager() + + decoder_graphs = sum(len(descs) for _, descs in capture_descs) + encoder_graphs = ( + encoder_cudagraph_manager.get_num_graphs_to_capture() + if encoder_cudagraph_manager is not None + else 0 + ) + total_graphs = decoder_graphs + encoder_graphs + if total_graphs == 0: + logger.debug("No CUDA graphs will be captured, skipping profiling") + self._cleanup_profiling_kv_cache() + return 0 + + graph_groups = [ + *( + f"{mode.name}={len(descs)} (largest={descs[0].num_tokens})" + for mode, descs in capture_descs + if descs + ), + ] + if encoder_graphs > 0: + graph_groups.append( + f"ENCODER={encoder_graphs} " + f"(largest={encoder_cudagraph_manager.token_budgets[-1]})" + ) + + logger.info("Profiling CUDA graph memory: %s", ", ".join(graph_groups)) + + # Use a temporary pool for profiling to avoid fragmentation in the main pool. + profiling_pool = current_platform.graph_pool_handle() + encoder_profiling_pool = current_platform.graph_pool_handle() + original_pools: dict[int, Any] = {} + all_wrappers = list(CUDAGraphWrapper._all_instances) + list( + BreakableCUDAGraphWrapper._all_instances + ) + for instance in all_wrappers: + original_pools[id(instance)] = instance.graph_pool + instance.graph_pool = profiling_pool + + shared_memory_estimate = {} + per_graph_estimate = {} + encoder_memory_estimate = 0 + + # On ROCm, capture these throwaway profiling graphs on vLLM's dedicated + # compute stream instead of the fresh side stream graph_capture() + # allocates by default. torch's allocator pools free blocks per stream, + # so a side-stream forward strands a persistent aiter scratch buffer in + # a separate pool, shifting the physical placement of the real KV cache + # allocated afterward and slowing bandwidth-bound decode ~20%. The + # graphs are discarded, so a side stream is unnecessary here. + # Use current_stream(), not torch.cuda.current_stream(): before vLLM + # initializes its dedicated stream, torch returns the per-thread default + # stream (cuda_stream=0), which cannot be used for cudagraph capture. + # cap_ctx=None keeps the side-stream path on CUDA. + cap_ctx = ( + GraphCaptureContext(current_stream()) + if current_platform.is_rocm() + else None + ) + + # Cleanup-only guard: CUDA graph capture errors should still propagate + # because encoder graph capture is opt-in. + try: + set_cudagraph_capturing_enabled(True) + with ( + self._freeze_gc(), + graph_capture(device=self.device, graph_capture_context=cap_ctx), + ): + torch.accelerator.synchronize() + torch.accelerator.empty_cache() + + for mode, descs in capture_descs: + profile_descs = descs[:2] + mem_samples: list[int] = [] + + for i, desc in enumerate(profile_descs): + mem_before = torch.accelerator.get_memory_info()[0] + self._warmup_and_capture( + desc, + cudagraph_runtime_mode=mode, + profile_seq_lens=( + min( + self.max_model_len, + self.max_num_tokens // desc.num_tokens, + ) + if mode == CUDAGraphMode.FULL and i == 0 + else None + ), + ) + torch.accelerator.synchronize() + free_after = torch.accelerator.get_memory_info()[0] + mem_samples.append(mem_before - free_after) + + first_capture = mem_samples[0] + # Use at least 1 MiB per graph for driver overhead + per_graph = max( + mem_samples[1] if len(mem_samples) > 1 else 0, 1 << 20 + ) + + shared_memory_estimate[mode] = first_capture + per_graph_estimate[mode] = per_graph * (len(descs) - 1) + + logger.debug( + "Estimated %s CUDA graph memory: " + "%.2f MiB first-capture + (%d-1) × %.2f MiB per-graph", + mode.name, + first_capture / (1 << 20), + len(descs), + per_graph / (1 << 20), + ) + + if encoder_cudagraph_manager is not None: + mem_before = torch.accelerator.get_memory_info()[0] + encoder_cudagraph_manager.capture(graph_pool=encoder_profiling_pool) + torch.accelerator.synchronize() + free_after = torch.accelerator.get_memory_info()[0] + encoder_memory_estimate = max(mem_before - free_after, 0) + + logger.debug( + "Estimated encoder CUDA graph memory: %.2f MiB for %d graphs", + encoder_memory_estimate / (1 << 20), + encoder_graphs, + ) + finally: + set_cudagraph_capturing_enabled(False) + CUDAGraphWrapper.clear_all_graphs() + BreakableCUDAGraphWrapper.clear_all_graphs() + if encoder_cudagraph_manager is not None: + encoder_cudagraph_manager.clear() + all_wrappers = list(CUDAGraphWrapper._all_instances) + list( + BreakableCUDAGraphWrapper._all_instances + ) + for instance in all_wrappers: + if id(instance) in original_pools: + instance.graph_pool = original_pools[id(instance)] + for key_set in self.cudagraph_dispatcher.cudagraph_keys.values(): + key_set.clear() + self.cudagraph_dispatcher.keys_initialized = False + self.maybe_remove_all_loras(self.lora_config) + self._cleanup_profiling_kv_cache() + compilation_counter.num_cudagraph_captured = saved_num_cudagraph_captured + + # FULL and PIECEWISE graphs share the global pool at runtime and are + # never replayed concurrently, so the pool overlays their memory. + # Take the max to avoid double-counting the overlap. + decoder_estimate = max(shared_memory_estimate.values(), default=0) + sum( + per_graph_estimate.values() + ) + # Encoder graphs use a manager-local pool at runtime, separate from the + # decoder pool, so add their estimate instead of overlaying it. + total_estimate = decoder_estimate + encoder_memory_estimate + logger.info( + "Estimated CUDA graph memory: %.2f GiB total", + total_estimate / (1 << 30), + ) + + return int(total_estimate) + + @instrument(span_name="Capture model") + def capture_model(self) -> int: + if self.compilation_config.cudagraph_mode == CUDAGraphMode.NONE: + logger.warning( + "Skipping CUDA graph capture. To turn on CUDA graph capture, " + "ensure `cudagraph_mode` was not manually set to `NONE`" + ) + return 0 + + # Initialize encoder CUDA graph manager if enabled. + self._maybe_init_encoder_cudagraph_manager() + + compilation_counter.num_gpu_runner_capture_triggers += 1 + + start_time = time.perf_counter() + + # Trigger CUDA graph capture for specific shapes. + # Capture the large shapes first so that the smaller shapes + # can reuse the memory pool allocated for the large shapes. + set_cudagraph_capturing_enabled(True) + + # Setup torch profiler for graph capture traces (conditional) + from vllm.distributed.parallel_state import get_world_group + + local_rank = get_world_group().local_rank + enable_profiler = ( + local_rank == 0 + ) and self.vllm_config.profiler_config.capture_torch_profiler + if enable_profiler: + trace_dir = ( + self.vllm_config.profiler_config.torch_profiler_dir + "/capture_traces" + ) + profiler = torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ], + record_shapes=True, + profile_memory=True, + with_stack=True, + on_trace_ready=torch.profiler.tensorboard_trace_handler( + trace_dir, + worker_name=f"graph_capture_rank_{local_rank}", + use_gzip=True, + ), + ) + logger.info_once( + "Rank %d: Torch profiler enabled for CUDA graph capture, " + "traces will be saved to: %s", + local_rank, + trace_dir, + ) + else: + profiler = nullcontext() + logger.info_once( + "Rank %d: Torch profiler disabled for CUDA graph capture", local_rank + ) + + with self._freeze_gc(), graph_capture(device=self.device): + torch.accelerator.synchronize() + torch.accelerator.empty_cache() + start_free_gpu_memory = torch.accelerator.get_memory_info()[0] + + for ( + runtime_mode, + batch_descs, + ) in self.cudagraph_dispatcher.get_capture_descs(): + self._capture_cudagraphs( + batch_descriptors=batch_descs, + cudagraph_runtime_mode=runtime_mode, + profiler=profiler, + ) + torch.accelerator.synchronize() + + # Capture encoder CUDA graphs if enabled + if self.encoder_cudagraph_manager is not None: + encoder_graph_pool = current_platform.graph_pool_handle() + self.encoder_cudagraph_manager.capture(graph_pool=encoder_graph_pool) + + torch.accelerator.synchronize() + end_free_gpu_memory = torch.accelerator.get_memory_info()[0] + + # Disable cudagraph capturing globally, so any unexpected cudagraph + # capturing will be detected and raise an error after here. + # Note: We don't put it into graph_capture context manager because + # we may do lazy capturing in future that still allows capturing + # after here. + set_cudagraph_capturing_enabled(False) + + torch.accelerator.synchronize() + torch.accelerator.empty_cache() + + # Lock workspace to prevent resizing during execution. + # Max workspace sizes should have been captured during warmup/profiling. + lock_workspace() + + end_time = time.perf_counter() + elapsed_time = end_time - start_time + cuda_graph_size = start_free_gpu_memory - end_free_gpu_memory + # This usually takes 5~20 seconds. + logger.info_once( + "Graph capturing finished in %.0f secs, took %.2f GiB", + elapsed_time, + cuda_graph_size / (1 << 30), + ) + return cuda_graph_size + + def _warmup_and_capture( + self, + desc: BatchDescriptor, + cudagraph_runtime_mode: CUDAGraphMode, + profile_seq_lens: int | None = None, + allow_microbatching: bool = False, + num_warmups: int | None = None, + profiler: AbstractContextManager[Any] | None = None, + ): + if profiler is None: + profiler = nullcontext() + if num_warmups is None: + num_warmups = self.compilation_config.cudagraph_num_of_warmups + force_attention = cudagraph_runtime_mode == CUDAGraphMode.FULL + for _ in range(num_warmups): + self._dummy_run( + desc.num_tokens, + cudagraph_runtime_mode=CUDAGraphMode.NONE, + force_attention=force_attention, + uniform_decode=desc.uniform, + allow_microbatching=allow_microbatching, + skip_eplb=True, + remove_lora=False, + num_active_loras=desc.num_active_loras, + profile_seq_lens=profile_seq_lens, + ) + if num_warmups > 0: + # Warmups may use auxiliary streams. Ensure all of their work has + # completed before beginning CUDA graph capture. + torch.accelerator.synchronize() + with ( + profiler, + torch.profiler.record_function( + f"capture_{desc.num_tokens}_{cudagraph_runtime_mode.name}" + ), + ): + self._dummy_run( + desc.num_tokens, + cudagraph_runtime_mode=cudagraph_runtime_mode, + uniform_decode=desc.uniform, + allow_microbatching=allow_microbatching, + skip_eplb=True, + remove_lora=False, + num_active_loras=desc.num_active_loras, + is_graph_capturing=True, + profile_seq_lens=profile_seq_lens, + ) + + def _capture_cudagraphs( + self, + batch_descriptors: list[BatchDescriptor], + cudagraph_runtime_mode: CUDAGraphMode, + profiler: AbstractContextManager[Any] | None = None, + ): + assert ( + cudagraph_runtime_mode != CUDAGraphMode.NONE + and cudagraph_runtime_mode.is_valid_runtime_mode() + ), f"Invalid cudagraph runtime mode: {cudagraph_runtime_mode}" + + if not batch_descriptors: + return + + uniform_decode = batch_descriptors[0].uniform + + # Only rank 0 should print progress bar during capture + if is_global_first_rank(): + batch_descriptors = tqdm( + batch_descriptors, + disable=not self.load_config.use_tqdm_on_load, + desc="Capturing CUDA graphs ({}, {})".format( + "decode" if uniform_decode else "mixed prefill-decode", + cudagraph_runtime_mode.name, + ), + ) + + # We skip EPLB here since we don't want to record dummy metrics + for batch_desc in batch_descriptors: + # We currently only capture ubatched graphs when its a FULL + # cudagraph, a uniform decode batch, and the number of tokens + # is above the threshold. Otherwise we just capture a non-ubatched + # version of the graph + allow_microbatching = ( + self.parallel_config.use_ubatching + and cudagraph_runtime_mode == CUDAGraphMode.FULL + and uniform_decode + and check_ubatch_thresholds( + config=self.vllm_config.parallel_config, + num_tokens=batch_desc.num_tokens, + uniform_decode=uniform_decode, + ) + ) + self._warmup_and_capture( + batch_desc, + cudagraph_runtime_mode=cudagraph_runtime_mode, + allow_microbatching=allow_microbatching, + profiler=profiler, + ) + torch.accelerator.synchronize() + self.maybe_remove_all_loras(self.lora_config) + + def initialize_attn_backend( + self, + kv_cache_config: KVCacheConfig, + is_profiling: bool = False, + ) -> None: + """ + Initialize the attention backends and attention metadata builders. + """ + assert len(self.attn_groups) == 0, "Attention backends are already initialized" + + class AttentionGroupKey(NamedTuple): + """Deduplication key for attention groups within a KV cache group. + + Splits on per-rank ``num_heads_q`` in addition to backend + spec + so layers with different Q-head counts (e.g. a spec-decode draft + with fewer attention heads than its target) get separate metadata + builders. The builders' scratch (e.g. ``softmax_segm_*`` in + ``triton_attn``, ``num_qo_heads`` in FlashInfer) is sized by + ``num_heads_q`` and assumes uniformity within the group; see + ``get_num_attention_heads_from_layers`` in + ``vllm/v1/attention/backends/utils.py``. + """ + + attn_backend: type[AttentionBackend] + kv_cache_spec: KVCacheSpec + num_heads_q: int + + def get_attn_backends_for_group( + kv_cache_group_spec: KVCacheGroupSpec, + ) -> tuple[dict[AttentionGroupKey, list[str]], set[type[AttentionBackend]]]: + layer_type = cast(type[Any], AttentionLayerBase) + layers = get_layers_from_vllm_config( + self.vllm_config, layer_type, kv_cache_group_spec.layer_names + ) + attn_backends = {} + attn_backend_layers = defaultdict(list) + # Dedupe based on full class name; this is a bit safer than + # using the class itself as the key because when we create dynamic + # attention backend subclasses (e.g. ChunkedLocalAttention) unless + # they are cached correctly, there will be different objects per + # layer. + for layer_name in kv_cache_group_spec.layer_names: + attn_backend = layers[layer_name].get_attn_backend() + + if layer_name in self.kv_sharing_fast_prefill_eligible_layers: + attn_backend = create_fast_prefill_custom_backend( + "FastPrefill", + attn_backend, # type: ignore[arg-type] + ) + + full_cls_name = attn_backend.full_cls_name() + layer_kv_cache_spec = kv_cache_group_spec.kv_cache_spec + if isinstance(layer_kv_cache_spec, UniformTypeKVCacheSpecs): + layer_kv_cache_spec = layer_kv_cache_spec.kv_cache_specs[layer_name] + # Non-Attention layer types (e.g. Mamba1, ShortConv) do not + # expose ``num_heads``; fall back to 0 so they cluster as + # before. Such layers never coexist with Attention in a + # single KV cache group (different KVCacheSpec), so the + # fallback can never spuriously merge them with attention + # layers. + num_heads_q = getattr(layers[layer_name], "num_heads", 0) + key = (full_cls_name, layer_kv_cache_spec, num_heads_q) + attn_backends[key] = AttentionGroupKey( + attn_backend, layer_kv_cache_spec, num_heads_q + ) + attn_backend_layers[key].append(layer_name) + return ( + {attn_backends[k]: v for k, v in attn_backend_layers.items()}, + set(group_key.attn_backend for group_key in attn_backends.values()), + ) + + def create_attn_groups( + attn_backends_map: dict[AttentionGroupKey, list[str]], + kv_cache_group_id: int, + ) -> list[AttentionGroup]: + attn_groups: list[AttentionGroup] = [] + for key, layer_names in attn_backends_map.items(): + attn_group = AttentionGroup( + key.attn_backend, + layer_names, + key.kv_cache_spec, + kv_cache_group_id, + ) + + attn_groups.append(attn_group) + return attn_groups + + attention_backend_maps = [] + attention_backend_list = [] + for kv_cache_group_spec in kv_cache_config.kv_cache_groups: + attn_backends = get_attn_backends_for_group(kv_cache_group_spec) + attention_backend_maps.append(attn_backends[0]) + attention_backend_list.append(attn_backends[1]) + + # Resolve cudagraph_mode before actually initialize metadata_builders + self._check_and_update_cudagraph_mode( + attention_backend_list, + kv_cache_config.kv_cache_groups, + is_profiling=is_profiling, + ) + + # Check if attention backend supports PCP&DCP and related features. + check_attention_cp_compatibility(self.vllm_config) + + for i, attn_backend_map in enumerate(attention_backend_maps): + self.attn_groups.append(create_attn_groups(attn_backend_map, i)) + + def initialize_metadata_builders( + self, kv_cache_config: KVCacheConfig, kernel_block_sizes: list[int] + ) -> None: + """ + Create the metadata builders for all KV cache groups and attn groups. + """ + for kv_cache_group_id in range(len(kv_cache_config.kv_cache_groups)): + for attn_group in self.attn_groups[kv_cache_group_id]: + attn_group.create_metadata_builders( + self.vllm_config, + self.device, + kernel_block_sizes[kv_cache_group_id] + if kv_cache_group_id < len(kernel_block_sizes) + else None, + num_metadata_builders=1 + if not self.parallel_config.use_ubatching + else self.parallel_config.num_ubatches, + ) + # Calculate reorder batch threshold (if needed) + # Note (tdoublep): do this *after* constructing builders, + # because some of them change the threshold at init time. + self.calculate_reorder_batch_threshold() + + # Initialize drafter attention backend + if self.speculative_config and ( + self.speculative_config.use_eagle() + or self.speculative_config.uses_draft_model() + ): + assert isinstance( + self.drafter, + EagleProposer | DFlashProposer | DraftModelProposer | Gemma4Proposer, + ) + self.drafter.initialize_attn_backend(kv_cache_config, kernel_block_sizes) + + def _check_and_update_cudagraph_mode( + self, + attention_backends: list[set[type[AttentionBackend]]], + kv_cache_groups: list[KVCacheGroupSpec], + is_profiling: bool = False, + ) -> None: + """ + Resolve the cudagraph_mode when there are multiple attention + groups with potential conflicting CUDA graph support. + Then initialize the cudagraph_dispatcher based on the resolved + cudagraph_mode. + """ + min_cg_support = AttentionCGSupport.ALWAYS + min_cg_attn_backend = None + + for attn_backend_set, kv_cache_group in zip( + attention_backends, kv_cache_groups + ): + for attn_backend in attn_backend_set: + builder_cls = attn_backend.get_builder_cls() + + cg_support = builder_cls.get_cudagraph_support( + self.vllm_config, kv_cache_group.kv_cache_spec + ) + if cg_support.value < min_cg_support.value: + min_cg_support = cg_support + min_cg_attn_backend = attn_backend.__name__ + cudagraph_mode = self.compilation_config.resolve_cudagraph_mode_and_sizes( + min_cg_support, + min_cg_attn_backend, + self.uniform_decode_query_len, + use_v2_model_runner=False, + tensor_parallel_size=self.parallel_config.tensor_parallel_size, + kv_cache_config=self.kv_cache_config, + max_num_reqs=self.max_num_reqs, + is_profiling=is_profiling, + ) + # Trigger cudagraph dispatching keys initialization after + # resolved cudagraph mode. + self.cudagraph_dispatcher.initialize_cudagraph_keys( + cudagraph_mode, self.uniform_decode_query_len + ) + + # Initialize drafter's cudagraph dispatcher if using spec decode. + if self.speculative_config and ( + self.speculative_config.use_eagle() + or self.speculative_config.uses_draft_model() + or self.speculative_config.uses_extract_hidden_states() + ): + assert isinstance( + self.drafter, + EagleProposer + | DFlashProposer + | DraftModelProposer + | ExtractHiddenStatesProposer + | Gemma4Proposer, + ) + self.drafter.initialize_cudagraph_keys(cudagraph_mode) + + def calculate_reorder_batch_threshold(self) -> None: + """ + Choose the minimum reorder batch threshold from all attention groups. + Backends should be able to support lower threshold then what they request + just may have a performance penalty due to that backend treating decodes + as prefills. + """ + min_none_high = lambda a, b: a if b is None else b if a is None else min(a, b) + + reorder_batch_thresholds: list[int | None] = [ + group.get_metadata_builder().reorder_batch_threshold + for group in self._attn_group_iterator() + ] + # If there are no attention groups (attention-free model) or no backend + # reports a threshold, leave reordering disabled. + if len(reorder_batch_thresholds) == 0: + self.reorder_batch_threshold = None + return + self.reorder_batch_threshold = reduce(min_none_high, reorder_batch_thresholds) # type: ignore[assignment] + + def may_reinitialize_input_batch( + self, kv_cache_config: KVCacheConfig, kernel_block_sizes: list[int] + ) -> None: + """ + Re-initialize the input batch if the block sizes are different from + what it was originally created with. This happens when the final + block size (determined after model loading) differs from the + placeholder used during __init__, or when there are multiple + KV cache groups. + + Args: + kv_cache_config: The KV cache configuration. + kernel_block_sizes: The kernel block sizes for each KV cache group. + """ + block_sizes = [] + max_num_blocks = [] + slot_mapping_modes = [] + max_model_len = max(self.max_model_len, self.max_encoder_len) + for kv_cache_group in kv_cache_config.kv_cache_groups: + kv_cache_spec = kv_cache_group.kv_cache_spec + kv_cache_spec_kind = get_kv_cache_spec_kind(kv_cache_spec) + if kv_cache_spec_kind == KVCacheSpecKind.ENCODER_ONLY_ATTENTION: + continue + block_size = kv_cache_spec.block_size + block_sizes.append(block_size) + if kv_cache_spec_kind == KVCacheSpecKind.MAMBA: + slot_mapping_modes.append(SlotMappingMode.NONE) + else: + slot_mapping_modes.append(SlotMappingMode.TOKEN_TO_KV_SLOT) + max_num_blocks_per_req = kv_cache_spec.max_num_blocks_per_req( + self.vllm_config, max_model_len + ) + max_num_blocks.append(max_num_blocks_per_req) + + if ( + block_sizes != self._init_block_sizes + or kernel_block_sizes != self._init_kernel_block_sizes + or max_num_blocks != self._init_max_num_blocks + or slot_mapping_modes != self._init_slot_mapping_modes + ): + self._init_block_sizes = block_sizes + self._init_kernel_block_sizes = kernel_block_sizes + self._init_max_num_blocks = max_num_blocks + self._init_slot_mapping_modes = slot_mapping_modes + # Capture warmup providers registered after final KV-cache geometry is known + with self.jit_warmup_registry.activate(): + self.input_batch = InputBatch( + max_num_reqs=self.max_num_reqs, + max_model_len=max_model_len, + max_num_batched_tokens=self.max_num_tokens, + device=self.device, + vocab_size=self.model_config.get_vocab_size(), + block_sizes=block_sizes, + kernel_block_sizes=kernel_block_sizes, + max_num_blocks_per_req=max_num_blocks, + num_spec_tokens=self.num_spec_tokens, + logitsprocs=self.input_batch.logitsprocs, + logitsprocs_need_output_token_ids=self.input_batch.logitsprocs_need_output_token_ids, + is_pooling_model=self.is_pooling_model, + cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size, + reasoning_config=self.vllm_config.reasoning_config, + use_replayssm=self.cache_config.use_replayssm, + slot_mapping_modes=slot_mapping_modes, + ) + + assert self._init_block_sizes == block_sizes, ( + f"InputBatch block_sizes {self._init_block_sizes} != " + f"kv_cache block_sizes {block_sizes}" + ) + assert self._init_kernel_block_sizes == kernel_block_sizes, ( + f"InputBatch kernel_block_sizes {self._init_kernel_block_sizes} " + f"!= kv_cache kernel_block_sizes {kernel_block_sizes}" + ) + + def _attn_group_iterator(self) -> Iterator[AttentionGroup]: + return itertools.chain.from_iterable(self.attn_groups) + + def _kv_cache_spec_attn_group_iterator(self) -> Iterator[AttentionGroup]: + if not self.kv_cache_config.kv_cache_groups: + return + for attn_groups in self.attn_groups: + yield from attn_groups + + def initialize_kv_cache_tensors( + self, kv_cache_config: KVCacheConfig, kernel_block_sizes: list[int] + ) -> dict[str, torch.Tensor]: + """ + Initialize the memory buffer for KV cache. + + Args: + kv_cache_config: The KV cache config + kernel_block_sizes: The kernel block sizes for each KV cache group. + + Returns: + Dict[str, torch.Tensor]: A map between layer names to their + corresponding memory buffer for KV cache. + """ + + kv_caches = allocate_kv_cache( + kv_cache_config, + self.device, + self.cache_config.get_resolved_kv_cache_layout(), + kernel_block_sizes, + ) + + # Set up cross-layer KV cache sharing + for layer_name, target_layer_name in self.shared_kv_cache_layers.items(): + logger.debug("%s reuses KV cache of %s", layer_name, target_layer_name) + kv_caches[layer_name] = kv_caches[target_layer_name] + + num_attn_module = ( + 2 if self.model_config.hf_config.model_type == "longcat_flash" else 1 + ) + bind_kv_cache( + kv_caches, + self.compilation_config.static_forward_context, + self.kv_caches, + num_attn_module, + ) + return kv_caches + + def maybe_add_kv_sharing_layers_to_kv_cache_groups( + self, kv_cache_config: KVCacheConfig + ) -> None: + """ + Add layers that re-use KV cache to KV cache group of its target layer. + Mapping of KV cache tensors happens in `initialize_kv_cache_tensors()` + """ + if not self.shared_kv_cache_layers: + # No cross-layer KV sharing, return + return + + add_kv_sharing_layers_to_kv_cache_groups( + self.shared_kv_cache_layers, + kv_cache_config.kv_cache_groups, + self.runner_only_attn_layers, + ) + + if self.cache_config.kv_sharing_fast_prefill: + # In You Only Cache Once (https://arxiv.org/abs/2405.05254) or other + # similar KV sharing setups, only the layers that generate KV caches + # are involved in the prefill phase, enabling prefill to early exit. + attn_layers = get_layers_from_vllm_config(self.vllm_config, Attention) + for layer_name in reversed(attn_layers): + if layer_name in self.shared_kv_cache_layers: + self.kv_sharing_fast_prefill_eligible_layers.add(layer_name) + else: + break + + def initialize_kv_cache( + self, + kv_cache_config: KVCacheConfig, + is_profiling: bool = False, + ) -> None: + """ + Initialize KV cache based on `kv_cache_config`. + Args: + kv_cache_config: Configuration for the KV cache, including the KV + cache size of each layer + """ + kv_cache_config = deepcopy(kv_cache_config) + self.kv_cache_config = kv_cache_config + self._mamba_bufs = None + self.may_add_encoder_only_layers_to_kv_cache_config() + self.maybe_add_kv_sharing_layers_to_kv_cache_groups(kv_cache_config) + self.initialize_attn_backend(kv_cache_config, is_profiling=is_profiling) + initialize_mamba_ssu_backend( + self.vllm_config.mamba_config, self.kv_cache_config + ) + # The kernel block size for all KV cache groups. For example, if + # kv_cache_manager uses block_size 256 for a given group, but the attention + # backends for that group only supports block_size 64, we will return + # kernel_block_size 64 and split the 256-token-block to 4 blocks with 64 + # tokens each. + kernel_block_sizes = prepare_kernel_block_sizes( + kv_cache_config, self.attn_groups + ) + self._kernel_block_sizes = kernel_block_sizes + + # create metadata builders + self.initialize_metadata_builders(kv_cache_config, kernel_block_sizes) + + # Reinitialize need to after initialize_attn_backend + self.may_reinitialize_input_batch(kv_cache_config, kernel_block_sizes) + kv_caches = self.initialize_kv_cache_tensors( + kv_cache_config, kernel_block_sizes + ) + + if ( + self.speculative_config + and self.speculative_config.uses_extract_hidden_states() + ): + assert isinstance(self.drafter, ExtractHiddenStatesProposer) + # validate all draft model layers belong to the same kv cache + # group + self.drafter.validate_same_kv_cache_group(kv_cache_config) + + if has_kv_transfer_group() and not is_profiling: + kv_transfer_group = get_kv_transfer_group() + kv_transfer_group.register_kv_caches(kv_caches) + kv_transfer_group.set_host_xfer_buffer_ops(copy_kv_blocks) + + def get_routed_experts( + self, + num_tokens: int, + ) -> RoutedExpertsTensors | None: + if not self.routed_experts_initialized: + return None + + device_buffer = self.routed_experts_capturer.get_device_buffer() + return RoutedExpertsTensors( + routing_data=device_buffer[:num_tokens].clone(), + slot_mapping=self.routed_experts_slot_mapping_device[:num_tokens].clone(), + ) + + def init_routed_experts_capturer(self): + logger.info( + "Initializing routed experts capturer, enable_return_routed_experts: %s", + self.model_config.enable_return_routed_experts, + ) + self.routed_experts_capturer = RoutedExpertsCapturer( + max_num_batched_tokens=self.scheduler_config.max_num_batched_tokens, + vllm_config=self.vllm_config, + kv_cache_config=self.kv_cache_config, + ) + bind_routed_experts_capturer(self.model, self.routed_experts_capturer) + + # Pinned CPU buffer for non-blocking D2H of ``routing_data`` on + # the sync scheduling path. Shape / dtype mirror the device + # capturer exactly so ``copy_`` is a straight memcpy. + self.routed_experts_cpu = torch.empty( + self.routed_experts_capturer.device_buffer.shape, + dtype=self.routed_experts_capturer.device_buffer.dtype, + device="cpu", + pin_memory=PIN_MEMORY, + ) + # ``slot_mapping`` dtype is fixed to int64 by + # ``block_table.slot_mapping``; we mirror that here. + max_tokens = self.scheduler_config.max_num_batched_tokens + self.routed_experts_slot_mapping_cpu = torch.empty( + (max_tokens,), + dtype=torch.int64, + device="cpu", + pin_memory=PIN_MEMORY, + ) + # Private device buffer so the shared ``block_table.slot_mapping`` + # can be overwritten by the next ``_prepare_inputs`` while the + # D2H is still pending on the copy stream. Written in + # ``_prepare_inputs``, read in ``_bookkeeping_sync`` (sync path) + # or cloned into a snapshot (async path). + self.routed_experts_slot_mapping_device = torch.empty( + (max_tokens,), + dtype=torch.int64, + device=self.device, + ) + self.routed_experts_initialized = True + + def may_add_encoder_only_layers_to_kv_cache_config(self) -> None: + """ + Add encoder-only layers to the KV cache config. + """ + block_size = self.vllm_config.cache_config.block_size + encoder_only_attn_specs: dict[AttentionSpec, list[str]] = defaultdict(list) + attn_layers = get_layers_from_vllm_config(self.vllm_config, Attention) + for layer_name, attn_module in attn_layers.items(): + if attn_module.attn_type == AttentionType.ENCODER_ONLY: + attn_spec: AttentionSpec = EncoderOnlyAttentionSpec( + block_size=block_size, + num_kv_heads=attn_module.num_kv_heads, + head_size=attn_module.head_size, + dtype=self.kv_cache_dtype, + ) + encoder_only_attn_specs[attn_spec].append(layer_name) + self.runner_only_attn_layers.add(layer_name) + if len(encoder_only_attn_specs) > 0: + assert len(encoder_only_attn_specs) == 1, ( + "Only support one encoder-only attention spec now" + ) + spec, layer_names = encoder_only_attn_specs.popitem() + self.kv_cache_config.kv_cache_groups.append( + KVCacheGroupSpec(layer_names=layer_names, kv_cache_spec=spec) + ) + + def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]: + """ + Generates the KVCacheSpec by parsing the kv cache format from each + Attention module in the static forward context. + Returns: + KVCacheSpec: A dictionary mapping layer names to their KV cache + format. Layers that do not need KV cache are not included. + """ + if has_ec_transfer() and not get_ec_transfer().is_consumer: + return {} + kv_cache_spec: dict[str, KVCacheSpec] = {} + layer_type = cast(type[Any], AttentionLayerBase) + attn_layers = get_layers_from_vllm_config(self.vllm_config, layer_type) + for layer_name, attn_module in attn_layers.items(): + if isinstance(attn_module, Attention) and ( + kv_tgt_layer := attn_module.kv_sharing_target_layer_name + ): + # The layer doesn't need its own KV cache and will use that of + # the target layer. We skip creating a KVCacheSpec for it, so + # that KV cache management logic will act as this layer does + # not exist, and doesn't allocate KV cache for the layer. This + # enables the memory saving of cross-layer kv sharing, allowing + # a given amount of memory to accommodate longer context lengths + # or enable more requests to be processed simultaneously. + self.shared_kv_cache_layers[layer_name] = kv_tgt_layer + continue + # Skip modules that don't need KV cache (eg encoder-only attention) + if spec := attn_module.get_kv_cache_spec(self.vllm_config): + if isinstance(spec, AttentionSpec): + spec = attn_module.get_attn_backend().customize_spec(spec) + kv_cache_spec[layer_name] = spec + + return kv_cache_spec + + def _to_list(self, sampled_token_ids: torch.Tensor) -> list[list[int]]: + # This is a short term mitigation for issue mentioned in + # https://github.com/vllm-project/vllm/issues/22754. + # `tolist` would trigger a cuda wise stream sync, which + # would block other copy ops from other cuda streams. + # A cuda event sync would avoid such a situation. Since + # this is in the critical path of every single model + # forward loop, this has caused perf issue for a disagg + # setup. + pinned = self.sampled_token_ids_pinned_cpu[: sampled_token_ids.shape[0]] + pinned.copy_(sampled_token_ids, non_blocking=True) + self.transfer_event.record() + self.transfer_event.synchronize() + return pinned.tolist() + + def get_encoder_timing_stats(self) -> dict[str, dict[str, float | int]]: + """ + Get encoder timing stats for all requests and clear the registry. + + Returns: + Dictionary mapping request_id to stats dict. + """ + with self._encoder_timing_lock: + stats = { + req_id: stats_obj.to_dict() + for req_id, stats_obj in self.encoder_timing_registry.items() + } + self.encoder_timing_registry.clear() + return stats + + @contextmanager + def timed_encoder_operation( + self, + should_time: bool, + group_lora_refs: list[tuple[str, Any]], + current_item_idx: int, + num_items: int, + ): + """ + Context manager to time encoder forward operations. + + Args: + should_time: Whether timing is enabled + group_lora_refs: Full list of (request_id, pos_info) tuples + current_item_idx: Starting index for this group + num_items: Number of items in this group + """ + if not should_time: + yield + return + + group_refs = group_lora_refs[current_item_idx : current_item_idx + num_items] + group_request_ids = {req_id for req_id, _ in group_refs} + + torch.accelerator.synchronize() + start_time = time.perf_counter() + + try: + yield + finally: + torch.accelerator.synchronize() + elapsed = time.perf_counter() - start_time + + per_request_time = elapsed / max(len(group_request_ids), 1) + + with self._encoder_timing_lock: + for req_id in group_request_ids: + if req_id not in self.encoder_timing_registry: + self.encoder_timing_registry[req_id] = EncoderTimingStats() + + stats = self.encoder_timing_registry[req_id] + stats.encoder_forward_secs += per_request_time + stats.num_encoder_calls += 1 diff --git a/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/7a9ef9ef3c59c88e4816bea6c9a74ebadd91a85632b42e221a5dd4009c5bb65c/metadata.py b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/7a9ef9ef3c59c88e4816bea6c9a74ebadd91a85632b42e221a5dd4009c5bb65c/metadata.py new file mode 100644 index 00000000..b446a858 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/7a9ef9ef3c59c88e4816bea6c9a74ebadd91a85632b42e221a5dd4009c5bb65c/metadata.py @@ -0,0 +1,64 @@ +"""GPU-free checkpoint metadata oracle for transactional KDA prefill.""" +from __future__ import annotations + + +def validate_metadata(*, cu_seqlens, initial_state_indices, final_state_indices, + checkpoint_state_indices, checkpoint_offsets, num_seqs, + num_tokens, token_capacity, seq_capacity, state_slots, + chunk=16, null_state_index=None, max_checkpoints=1): + """Validate packed ownership and return spans without reading state tensors. + + Nonpositive offsets disable an entry. A null checkpoint never owns storage; + its positive offset still must be aligned and within the request, matching + the original device validator. Two enabled, non-null exports need distinct + offsets and globally unique destinations. Initial may alias only own final. + """ + if type(max_checkpoints) is not int or max_checkpoints not in (1, 2): + raise ValueError("max_checkpoints must be 1 or 2") + if not 0 <= num_seqs <= seq_capacity or not 0 <= num_tokens <= token_capacity: + raise ValueError("live counts exceed capacities") + if int(cu_seqlens[0]) != 0 or int(cu_seqlens[num_seqs]) != num_tokens: + raise ValueError("packed boundaries do not match live tokens") + spans = [] + writes = set() + + def null(slot): + return null_state_index is not None and slot == null_state_index + + def check_slot(slot): + if not null(slot) and not 0 <= slot < state_slots: + raise IndexError("state slot out of range") + + def write(slot): + if null(slot): + return + check_slot(slot) + if slot in writes: + raise ValueError("duplicate checkpoint/final write slot") + writes.add(slot) + + for seq in range(num_seqs): + start, end = int(cu_seqlens[seq]), int(cu_seqlens[seq + 1]) + if not 0 <= start <= end <= num_tokens: + raise ValueError("invalid packed sequence interval") + spans.append((start, end)) + check_slot(int(initial_state_indices[seq])) + write(int(final_state_indices[seq])) + seen_offsets = set() + for cp in range(max_checkpoints): + slot = int(checkpoint_state_indices[seq] if max_checkpoints == 1 + else checkpoint_state_indices[seq][cp]) + offset = int(checkpoint_offsets[seq] if max_checkpoints == 1 + else checkpoint_offsets[seq][cp]) + if offset > end - start or (offset > 0 and offset % chunk): + raise ValueError("checkpoint offset out of bounds or unaligned") + if offset > 0 and not null(slot): + if offset in seen_offsets: + raise ValueError("two active checkpoints have the same offset") + seen_offsets.add(offset) + write(slot) + for seq in range(num_seqs): + initial, final = int(initial_state_indices[seq]), int(final_state_indices[seq]) + if not null(initial) and initial in writes - ({final} if not null(final) else set()): + raise ValueError("initial slot conflicts with a checkpoint or another request write") + return spans diff --git a/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/90ae5dd89e808c9b13fd24f203feebd56e728b2928c13547e131d1ac7ca35965/mamba_hybrid.py b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/90ae5dd89e808c9b13fd24f203feebd56e728b2928c13547e131d1ac7ca35965/mamba_hybrid.py new file mode 100644 index 00000000..812301e8 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/90ae5dd89e808c9b13fd24f203feebd56e728b2928c13547e131d1ac7ca35965/mamba_hybrid.py @@ -0,0 +1,507 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from dataclasses import dataclass +from typing import Any + +import numpy as np +import torch +import torch.nn as nn + +from vllm import envs +from vllm.logger import init_logger +from vllm.config import VllmConfig +from vllm.config.compilation import CUDAGraphMode +from vllm.triton_utils import tl, triton +from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder +from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadataBuilder +from vllm.v1.attention.backends.short_conv_attn import ShortConvAttentionMetadataBuilder +from vllm.v1.core.sched.output import NewRequestData +from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec +from vllm.v1.utils import CpuGpuBuffer +from vllm.v1.worker.gpu.attn_utils import build_attn_metadata +from vllm.v1.worker.gpu.input_batch import InputBatch +from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache +from vllm.v1.worker.gpu.model_states.default import DefaultModelState +from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata +from vllm.v1.worker.gpu.model_states.recoverssm import RecoverSSMState +from vllm.v1.worker.mamba_utils import ( + MambaSpecDecodeGPUContext, + MambaStateCopyFuncsByType, + get_mamba_groups, + preprocess_mamba_align_fused_kernel, + resolve_mamba_state_copy_funcs, +) +from vllm.v1.worker.utils import AttentionGroup + + +logger = init_logger(__name__) + + +def _checkpoint_plan_rows(input_batch, plans, num_reqs, *, for_capture=False): + """Map immutable scheduler plans to this actual V2 batch's request order.""" + if not plans: + return None + if for_capture: + raise ValueError("Planned recurrent checkpoints cannot enter graph capture") + if len(input_batch.req_ids) != input_batch.num_reqs or num_reqs < input_batch.num_reqs: + raise ValueError("Checkpoint batch request geometry differs") + rows = [None] * num_reqs + positions = {request_id: row for row, request_id in enumerate(input_batch.req_ids)} + if len(positions) != input_batch.num_reqs: + raise ValueError("Checkpoint batch has duplicate request IDs") + for request_id, plan in plans.items(): + row = positions.get(request_id) + if row is None: + raise ValueError("Checkpoint plan request is absent from the active V2 batch") + start, end, targets = plan + if (type(start) is not int or type(end) is not int or start < 0 or end <= start + or not isinstance(targets, tuple) or not 1 <= len(targets) <= 2 + or any(type(target) is not int for target in targets) + or targets != tuple(sorted(set(targets))) + or any(not start < target < end for target in targets) + or int(input_batch.num_computed_tokens_np[row]) != start + or int(input_batch.num_scheduled_tokens[row]) != end - start + or int(input_batch.query_start_loc_np[row + 1] - input_batch.query_start_loc_np[row]) != end - start): + raise ValueError("Checkpoint plan does not match the actual V2 query span") + rows[row] = (start, end, targets) + return rows + + +@dataclass +class MambaHybridAttnMetadata(ModelSpecificAttnMetadata): + is_prefilling: torch.Tensor + num_accepted_tokens: torch.Tensor | None = None + num_decode_draft_tokens_cpu: torch.Tensor | None = None + recurrent_prefill_checkpoint_plans_cpu: list | None = None + + def get_extra_common_attn_kwargs( + self, + kv_cache_group_id: int, + num_reqs: int, + ) -> dict[str, Any]: + values = {"is_prefilling": self.is_prefilling[:num_reqs]} + if self.recurrent_prefill_checkpoint_plans_cpu is not None: + values["recurrent_prefill_checkpoint_plans_cpu"] = ( + self.recurrent_prefill_checkpoint_plans_cpu[:num_reqs] + ) + return values + + def get_extra_attn_kwargs( + self, + attn_metadata_builder: Any, + num_reqs: int, + ) -> dict[str, Any]: + if not isinstance( + attn_metadata_builder, + ( + Mamba2AttentionMetadataBuilder, + GDNAttentionMetadataBuilder, + ShortConvAttentionMetadataBuilder, + ), + ): + return {} + return { + "num_accepted_tokens": None + if self.num_accepted_tokens is None + else self.num_accepted_tokens[:num_reqs], + "num_decode_draft_tokens_cpu": None + if self.num_decode_draft_tokens_cpu is None + else self.num_decode_draft_tokens_cpu[:num_reqs], + } + + +class MambaHybridModelState(DefaultModelState): + """Model state for hybrid attention + Mamba / linear-attention models.""" + + def __init__( + self, + vllm_config: VllmConfig, + model: nn.Module, + encoder_cache: EncoderCache | None, + device: torch.device, + ) -> None: + super().__init__(vllm_config, model, encoder_cache, device) + self.cache_config = vllm_config.cache_config + self.num_accepted_tokens_gpu = torch.ones( + self.max_num_reqs, dtype=torch.int32, device=self.device + ) + self._gdn_spec_accepted_tokens = ( + torch.ones_like(self.num_accepted_tokens_gpu) + if envs.VLLM_GDN_SPEC_DECODE_METADATA_FASTPATH + else None + ) + # Pre-copy "align" prefix-cache state (V2). The migration of each + # request's mamba state across block boundaries runs as a fused GPU + # kernel reusing the postprocess copy machinery, so the per-step src + # columns and the running state_idx are kept GPU-resident. + self._align_mode = self.cache_config.mamba_cache_mode == "align" + self.recoverssm = ( + RecoverSSMState() if self.cache_config.use_kda_recoverssm else None + ) + if self._align_mode: + self._mamba_state_idx_gpu = torch.zeros( + self.max_num_reqs, dtype=torch.int32, device=self.device + ) + self._mamba_src_col_gpu = torch.full( + (self.max_num_reqs,), -1, dtype=torch.int32, device=self.device + ) + self._mamba_src_off_gpu = torch.zeros( + self.max_num_reqs, dtype=torch.int32, device=self.device + ) + self._mamba_ctx: MambaSpecDecodeGPUContext | None = None + self._mamba_group_ids: list[int] = [] + self._mamba_spec: MambaSpec | None = None + self._mamba_copy_funcs_by_type: MambaStateCopyFuncsByType | None = None + self._aligned_metadata_groups: list[list[AttentionGroup]] | None = None + self._aligned_metadata_builders: list[tuple[int, Any]] = [] + self._aligned_metadata_ctx: MambaSpecDecodeGPUContext | None = None + + def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: + super().add_request(req_index, new_req_data) + # Must reset the speculative acceptance count in this idx which could be stale. + self.num_accepted_tokens_gpu[req_index].fill_(1) + if self._align_mode: + # Seed the running state block from the resumed/prefilled position. + self._mamba_state_idx_gpu[req_index].fill_( + (new_req_data.num_computed_tokens - 1) // self.cache_config.block_size + ) + + def reset_kv_cache_state(self) -> None: + """Release pointer metadata derived from Mamba cache allocations.""" + if self._align_mode: + self._mamba_ctx = None + self._mamba_copy_funcs_by_type = None + self._aligned_metadata_ctx = None + self._aligned_metadata_groups = None + self._aligned_metadata_builders = [] + if self.recoverssm is not None: + self.recoverssm.reset() + + def _get_mamba_group_info( + self, kv_cache_config: KVCacheConfig + ) -> tuple[list[int], MambaSpec]: + if self._mamba_spec is None: + self._mamba_group_ids, self._mamba_spec = get_mamba_groups(kv_cache_config) + return self._mamba_group_ids, self._mamba_spec + + def _ensure_align_ctx( + self, + kv_cache_config: KVCacheConfig, + mamba_group_ids: list[int], + block_tables: tuple[torch.Tensor, ...], + ) -> MambaSpecDecodeGPUContext: + if self._mamba_ctx is None: + self._mamba_copy_funcs_by_type = resolve_mamba_state_copy_funcs( + self.model, kv_cache_config + ) + # Both SD and DS conv layouts support a >0 spec-decode shift: the + # fused pre-copy kernel (``_copy_mamba_state_block``) applies the + # ``token_bias = num_accepted - 1`` window shift per conv layout + # (SD: contiguous slice; DS: per-dim-row strided slice), matching + # the V1 ``get_conv_copy_spec`` semantics. + self._mamba_ctx = MambaSpecDecodeGPUContext.create( + max_num_reqs=self.max_num_reqs, + kv_cache_config=kv_cache_config, + num_state_types=None, + device=self.device, + make_buffer=lambda n, dtype: CpuGpuBuffer( + n, dtype=dtype, device=self.device + ), + copy_funcs_by_type=self._mamba_copy_funcs_by_type, + ) + ctx = self._mamba_ctx + if not ctx.is_initialized: + forward_context = self.vllm_config.compilation_config.static_forward_context + # block_tables are batch-order slices of the persistent + # input_block_tables (stable data_ptr), so the metadata is captured + # once here and reused across steps. + ctx.initialize_from_forward_context( + kv_cache_config, + forward_context, + self._mamba_copy_funcs_by_type, + [block_tables[gid] for gid in mamba_group_ids], + ) + return ctx + + def _prepare_aligned_state_indices( + self, + seq_lens: torch.Tensor, + num_reqs: int, + attn_groups: list[list[AttentionGroup]], + kv_cache_config: KVCacheConfig, + block_tables: tuple[torch.Tensor, ...], + ) -> None: + mamba_group_ids, _ = self._get_mamba_group_info(kv_cache_config) + if self._aligned_metadata_groups is not attn_groups: + self._aligned_metadata_builders = [] + for group_idx, group_id in enumerate(mamba_group_ids): + for group in attn_groups[group_id]: + builder = group.get_metadata_builder(0) + if hasattr(builder, "mamba_aligned_state_indices"): + self._aligned_metadata_builders.append((group_idx, builder)) + if hasattr(builder, "mamba_spec_accepted_tokens"): + builder.mamba_spec_accepted_tokens = ( + self._gdn_spec_accepted_tokens + ) + self._aligned_metadata_groups = attn_groups + self._aligned_metadata_ctx = None + if not self._aligned_metadata_builders: + return + + ctx = self._ensure_align_ctx(kv_cache_config, mamba_group_ids, block_tables) + if self._aligned_metadata_ctx is not ctx: + assert ctx.aligned_state_indices is not None + group_views = ctx.aligned_state_indices.unbind(0) + for group_idx, builder in self._aligned_metadata_builders: + builder.mamba_aligned_state_indices = group_views[group_idx] + self._aligned_metadata_ctx = ctx + ctx.compute_aligned_state_indices(seq_lens, num_reqs) + + def preprocess_state( + self, + input_batch: InputBatch, + block_tables: tuple[torch.Tensor, ...], + kv_cache_config: KVCacheConfig, + num_computed_tokens: torch.Tensor, + ) -> None: + """Migrate each request's mamba state across block boundaries before the + forward (V1 align semantics, done on GPU). Runs on real batches only + (dummy DP/profiling runs skip preprocess_state), and before + ``prepare_attn`` gathers ``num_accepted_tokens``, so the boundary reset + is visible to the forward kernels. + """ + if not self._align_mode: + return + num_reqs = input_batch.num_reqs + if num_reqs == 0: + return + mamba_group_ids, mamba_spec = self._get_mamba_group_info(kv_cache_config) + ctx = self._ensure_align_ctx(kv_cache_config, mamba_group_ids, block_tables) + + # The state-advance + pre-copy kernels run every step; they fast-exit per + # request when src_col < 0 or src_col == dst_col, so no copy happens on + # steps that don't cross a block boundary. (Skipping the launch entirely + # would need a V1-style async-D2H of the actual num_computed, since + # num_computed_tokens_np is an optimistic mirror under async scheduling; + # the launch cost is ~0.3% of TPOT, so the GPU fast-exit suffices.) + block = 256 + grid = (triton.cdiv(num_reqs, block),) + preprocess_mamba_align_fused_kernel[grid]( + input_batch.idx_mapping, + self._mamba_state_idx_gpu, + num_computed_tokens, + input_batch.query_start_loc, + self.num_accepted_tokens_gpu, + self._mamba_src_col_gpu, + self._mamba_src_off_gpu, + num_reqs, + BLOCK_SIZE=block, + MAMBA_BLOCK_SIZE=mamba_spec.block_size, + ) + ctx.run_fused_precopy( + num_reqs, + self._mamba_state_idx_gpu, + self._mamba_src_col_gpu, + self._mamba_src_off_gpu, + input_batch.idx_mapping, + ) + + def prepare_attn( + self, + input_batch: InputBatch, + cudagraph_mode: CUDAGraphMode, + block_tables: tuple[torch.Tensor, ...], + slot_mappings: torch.Tensor, + attn_groups: list[list[AttentionGroup]], + kv_cache_config: KVCacheConfig, + for_capture: bool = False, + recurrent_prefill_checkpoint_plans: dict | None = None, + ) -> dict[str, Any]: + if cudagraph_mode == CUDAGraphMode.FULL: + num_reqs = input_batch.num_reqs_after_padding + num_tokens = input_batch.num_tokens_after_padding + else: + num_reqs = input_batch.num_reqs + num_tokens = input_batch.num_tokens + query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np) + max_query_len = input_batch.num_scheduled_tokens.max().item() + seq_lens_cpu_upper_bound = input_batch.seq_lens_cpu_upper_bound + if for_capture: + # Capture with worst-case max_seq_len so the graph is valid at any replay. + max_seq_len = self.max_model_len + else: + max_seq_len = seq_lens_cpu_upper_bound[:num_reqs].max().item() + + is_prefilling = torch.zeros(num_reqs, dtype=torch.bool, device="cpu") + is_prefilling[: input_batch.num_reqs] = torch.from_numpy( + input_batch.is_prefilling_np + ) + # During CUDAGraph capture, num_decode_draft_tokens_cpu and num_accepted_tokens + # are created by attn_metadata_builder.build_for_cudagraph_capture, so we only + # compute them during actual (non-capture) forward execution. + num_accepted_tokens = None + num_decode_draft_tokens_cpu = None + if not for_capture and self.vllm_config.num_speculative_tokens > 0: + num_accepted_tokens = self.num_accepted_tokens_gpu.new_ones(num_reqs) + num_accepted_tokens[: input_batch.num_reqs] = self.num_accepted_tokens_gpu[ + input_batch.idx_mapping + ] + + # GDN uses >= 0 to select spec-decode rows, so non-decode rows + # need the -1 sentinel rather than a raw zero draft count. + num_decode_draft_tokens_np = np.full(num_reqs, -1, dtype=np.int32) + num_draft_tokens_per_req = input_batch.num_draft_tokens_per_req + if num_draft_tokens_per_req is not None: + # A row is a spec-decode row only when its whole prompt is already + # computed, i.e. exactly one non-draft (decode) token is scheduled. + is_decode = ( + input_batch.num_scheduled_tokens == num_draft_tokens_per_req + 1 + ) + spec_decode_mask = (num_draft_tokens_per_req > 0) & is_decode + num_decode_draft_tokens_np[: input_batch.num_reqs] = np.where( + spec_decode_mask, num_draft_tokens_per_req, -1 + ) + num_decode_draft_tokens_cpu = torch.from_numpy(num_decode_draft_tokens_np) + + if self._align_mode: + self._prepare_aligned_state_indices( + input_batch.seq_lens, + num_reqs, + attn_groups, + kv_cache_config, + block_tables, + ) + + checkpoint_plans_cpu = _checkpoint_plan_rows( + input_batch, recurrent_prefill_checkpoint_plans, num_reqs, + for_capture=for_capture, + ) + mamba_attn_metadata = MambaHybridAttnMetadata( + recurrent_prefill_checkpoint_plans_cpu=checkpoint_plans_cpu, + is_prefilling=is_prefilling, + num_accepted_tokens=num_accepted_tokens, + num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, + ) + attn_metadata = build_attn_metadata( + attn_groups=attn_groups, + num_reqs=num_reqs, + num_tokens=num_tokens, + query_start_loc_gpu=input_batch.query_start_loc, + query_start_loc_cpu=query_start_loc_cpu, + max_query_len=max_query_len, + seq_lens=input_batch.seq_lens, + max_seq_len=max_seq_len, + block_tables=block_tables, + slot_mappings=slot_mappings, + kv_cache_config=kv_cache_config, + seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, + dcp_local_seq_lens=input_batch.dcp_local_seq_lens, + model_specific_attn_metadata=mamba_attn_metadata, + for_cudagraph_capture=for_capture, + rswa_prefix_lens=input_batch.prompt_lens, + ) + if checkpoint_plans_cpu is not None: + checked_layers = 0 + for groups in attn_groups: + for group in groups: + if not isinstance(group.get_metadata_builder(0), GDNAttentionMetadataBuilder): + continue + for layer_name in group.layer_names: + checkpoint = getattr(attn_metadata[layer_name], "prefill_checkpoint", None) + if checkpoint is None or checkpoint.required_mask is None: + raise RuntimeError("V2 scheduled a checkpoint plan without GDN export metadata") + checked_layers += 1 + if not checked_layers: + raise RuntimeError("V2 scheduled recurrent checkpoints without an active GDN layer") + if not getattr(self, "_explicit_checkpoint_plan_logged", False): + logger.info("V2 explicit recurrent checkpoint plan active: rows=%s layers=%d", + checkpoint_plans_cpu, checked_layers) + self._explicit_checkpoint_plan_logged = True + if self.recoverssm is not None: + self.recoverssm.record_step( + attn_metadata, + attn_groups, + for_capture=for_capture, + ) + return attn_metadata + + def postprocess_state( + self, + idx_mapping: torch.Tensor, + num_sampled: torch.Tensor | int, + num_computed_tokens: torch.Tensor | None = None, + ) -> None: + # Chunked prefill does not sample a token, so num_sampled can be 0. + # Mamba treats num_accepted_tokens=1 as the neutral non-spec value. + num_reqs = idx_mapping.shape[0] + if num_reqs: + if not isinstance(num_sampled, int): + # idx_mapping may contain -1 sentinels (filtered rows) under PP; the + # kernel skips them rather than scattering with a host-side gather. + _scatter_num_accepted_kernel[(num_reqs,)]( + idx_mapping, + num_sampled, + self.num_accepted_tokens_gpu, + ) + else: + # Fill with single value. + _fill_num_accepted_kernel[(num_reqs,)]( + idx_mapping, + self.num_accepted_tokens_gpu, + max(num_sampled, 1), + ) + + if self.recoverssm is not None: + self.recoverssm.commit_step( + num_sampled, + idx_mapping, + state_indices=(self._mamba_state_idx_gpu if self._align_mode else None), + num_accepted_tokens=self.num_accepted_tokens_gpu, + ) + + if not num_reqs: + return + + # Align: save the running state to the block-aligned position when + # spec-decode acceptance leaves the sequence non-block-aligned (mirrors + # the V1 align postprocess). num_computed_tokens already holds the + # post-step advanced count. + if ( + self._align_mode + and num_computed_tokens is not None + and self._mamba_ctx is not None + ): + self._mamba_ctx.run_fused_postprocess_align( + num_reqs, + self.num_accepted_tokens_gpu, + self._mamba_state_idx_gpu, + num_computed_tokens, + idx_mapping, + ) + + +@triton.jit +def _scatter_num_accepted_kernel( + idx_mapping_ptr, # [num_reqs] batch_idx -> req_state_idx (-1 to skip) + num_sampled_ptr, # [num_reqs] + num_accepted_ptr, # [max_num_reqs] +): + row = tl.program_id(0) + req_state_idx = tl.load(idx_mapping_ptr + row) + if req_state_idx < 0: + return + num_sampled = tl.load(num_sampled_ptr + row) + tl.store(num_accepted_ptr + req_state_idx, tl.maximum(num_sampled, 1)) + + +@triton.jit +def _fill_num_accepted_kernel( + idx_mapping_ptr, # [num_reqs] batch_idx -> req_state_idx (-1 to skip) + num_accepted_ptr, # [max_num_reqs] + num_sampled, +): + row = tl.program_id(0) + req_state_idx = tl.load(idx_mapping_ptr + row) + if req_state_idx < 0: + return + tl.store(num_accepted_ptr + req_state_idx, num_sampled) diff --git a/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/c79f0f62c92963959131fa47ac63002b3109a379abbccabd26834d224ccb044d/backend.py b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/c79f0f62c92963959131fa47ac63002b3109a379abbccabd26834d224ccb044d/backend.py new file mode 100644 index 00000000..9d0856de --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/c79f0f62c92963959131fa47ac63002b3109a379abbccabd26834d224ccb044d/backend.py @@ -0,0 +1,1115 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from abc import ABC, abstractmethod +from dataclasses import dataclass, replace +from enum import Enum +from typing import TYPE_CHECKING, Any, ClassVar, Generic, Protocol, TypeVar + +import numpy as np +import torch +from typing_extensions import deprecated + +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8Dynamic64Sym, + kFp8Dynamic128Sym, + kFp8StaticTensorSym, + kNvfp4Dynamic, +) + +if TYPE_CHECKING: + from vllm.config import VllmConfig + from vllm.config.cache import CacheDType + from vllm.model_executor.layers.linear import ColumnParallelLinear + from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey + from vllm.platforms.interface import DeviceCapability + from vllm.v1.kv_cache_interface import ( + AttentionSpec, + KVCacheLayout, + KVCacheSpec, + KVQuantMode, + ) + +from vllm.v1.kv_cache_interface import KVCacheLayout, get_kv_quant_mode + + +class AttentionType(str, Enum): + """ + Attention type. + Use string to be compatible with `torch.compile`. + """ + + DECODER = "decoder" + """Decoder attention between previous layer Q/K/V.""" + ENCODER = "encoder" + """Encoder attention between previous layer Q/K/V for encoder-decoder.""" + ENCODER_ONLY = "encoder_only" + """Encoder attention between previous layer Q/K/V.""" + ENCODER_DECODER = "encoder_decoder" + """Attention between dec. Q and enc. K/V for encoder-decoder.""" + + +class MultipleOf: + base: int + + def __init__(self, base: int): + self.base = base + + +class AttentionBackend(ABC): + """Abstract class for attention backends.""" + + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16] + supported_kv_cache_dtypes: ClassVar[list["CacheDType"]] = [ + "auto", + "float16", + "bfloat16", + ] + + # Does attention's forward() include kv cache update? + forward_includes_kv_cache_update: bool = True + + # Whether metadata builders and kernels can execute a DCP-replicated cache + # group as a local DCP1 operation inside a larger DCP world. + supports_dcp_replicated: ClassVar[bool] = False + + @staticmethod + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: + return [MultipleOf(1)] + + @staticmethod + @abstractmethod + def get_name() -> str: + raise NotImplementedError + + @staticmethod + @abstractmethod + def get_impl_cls() -> type["AttentionImplBase"]: + raise NotImplementedError + + @staticmethod + @abstractmethod + def get_builder_cls(): # -> Type["AttentionMetadataBuilder"]: + raise NotImplementedError + + @classmethod + def full_cls_name(cls) -> tuple[str, str]: + return (cls.__module__, cls.__qualname__) + + @classmethod + def get_supported_head_sizes(cls) -> list[int]: + return [] + + @classmethod + def supports_head_size(cls, head_size: int) -> bool: + supported_head_sizes = cls.get_supported_head_sizes() + return (not supported_head_sizes) or head_size in supported_head_sizes + + @classmethod + def supports_dtype(cls, dtype: torch.dtype) -> bool: + return dtype in cls.supported_dtypes + + @classmethod + def supports_kv_cache_dtype(cls, kv_cache_dtype: "CacheDType | None") -> bool: + if kv_cache_dtype is None: + return True + return (not cls.supported_kv_cache_dtypes) or ( + kv_cache_dtype in cls.supported_kv_cache_dtypes + ) + + @classmethod + def supports_block_size(cls, block_size: int | None) -> bool: + if block_size is None: + return True + + supported_kernel_block_sizes = cls.get_supported_kernel_block_sizes() + if not supported_kernel_block_sizes: + return True + + for supported_size in supported_kernel_block_sizes: + if isinstance(supported_size, MultipleOf): + supported_size = supported_size.base + # With hybrid_blocks feature, the framework-level block size + # only needs to be a multiple of the kernel's requirement, + # even if the kernel requires a fixed block_size. + if block_size % supported_size == 0: + return True + return False + + @classmethod + def customize_spec(cls, spec: "AttentionSpec") -> "AttentionSpec": + """Adjust the layer's KV cache spec for this backend's kernels. Used when the + kernels want KV packed in a specific way. + + NOTE: temporary compatibility API. Today the Attention layer builds the spec + from the model config and the backend only gets to adjust it post-hoc; the end + state is for the backend to build and return the spec directly, at which point + this hook goes away. + + (see: https://github.com/vllm-project/vllm/issues/42449) + """ + return spec + + @classmethod + def get_preferred_block_size(cls, default_block_size: int) -> int: + supported_sizes = cls.get_supported_kernel_block_sizes() + if not supported_sizes: + return default_block_size + + if cls.supports_block_size(default_block_size): + return default_block_size + + return min(s.base if isinstance(s, MultipleOf) else s for s in supported_sizes) + + @classmethod + def is_mla(cls) -> bool: + return False + + @classmethod + def supports_sink(cls) -> bool: + return False + + @classmethod + def supports_alibi_sqrt(cls) -> bool: + return False + + @classmethod + def supports_mm_prefix(cls) -> bool: + return False + + @classmethod + def is_sparse(cls) -> bool: + return False + + @classmethod + def supports_per_head_quant_scales(cls) -> bool: + return False + + @classmethod + def supports_sliding_window(cls) -> bool: + return False + + @classmethod + def supports_non_causal(cls) -> bool: + """Check if backend supports non-causal (bidirectional) attention + for decoder models. + + Unlike ENCODER_ONLY attention type which implies a different + execution model, this refers to non-causal attention within the + standard paged-KV-cache decoder path. + """ + return False + + @classmethod + def supports_batch_invariance(cls) -> bool: + return False + + @classmethod + def supports_kv_connector(cls) -> bool: + return True + + @classmethod + def supports_device_cpu_query_lens_mismatch(cls) -> bool: + """Whether this backend can run a batch whose device query_start_loc disagrees + with the CPU one; backends that plan off the CPU query lengths must opt out. + + Currently only verification requests are affected: adaptive verification trims + their drafts on device. On the CPU the draft budget is evenly distributed across + requests, so the total draft budget, the decode/prefill split point and the CPU + prefill query lengths all stay correct. + + SSM backends opt out: their recurrent-state planning is built from the CPU + per-request boundaries, which the trimmed batch no longer matches. + """ + return not cls.is_ssm() + + @classmethod + def supports_pcp(cls) -> bool: + try: + return cls.get_impl_cls().supports_pcp + except NotImplementedError: + return False + + @classmethod + def supports_non_causal_dcp(cls) -> bool: + builder_cls = cls.get_builder_cls() + return bool(getattr(builder_cls, "supports_non_causal_multi_token_dcp", False)) + + @classmethod + def supports_attn_type(cls, attn_type: str) -> bool: + """Check if backend supports a given attention type. + + By default, only supports decoder attention. + Backends should override this to support other attention types. + """ + return attn_type == AttentionType.DECODER + + @classmethod + def supports_compute_capability(cls, capability: "DeviceCapability") -> bool: + return True + + @classmethod + def supports_combination( + cls, + head_size: int, + dtype: torch.dtype, + kv_cache_dtype: "CacheDType | None", + block_size: int | None, + use_mla: bool, + has_sink: bool, + use_sparse: bool, + use_mm_prefix: bool, + device_capability: "DeviceCapability", + ) -> str | None: + return None + + @classmethod + def validate_configuration( + cls, + head_size: int, + dtype: torch.dtype, + kv_cache_dtype: "CacheDType | None", + block_size: int | None, + use_mla: bool, + has_sink: bool, + use_sparse: bool, + use_mm_prefix: bool, + use_per_head_quant_scales: bool, + device_capability: "DeviceCapability", + attn_type: str, + has_sliding_window: bool = False, + use_non_causal: bool = False, + use_batch_invariant: bool = False, + use_kv_connector: bool = False, + use_pcp: bool = False, + use_adaptive_verification: bool = False, + use_dcp: bool = False, + ) -> list[str]: + invalid_reasons = [] + if not cls.supports_head_size(head_size): + invalid_reasons.append("head_size not supported") + if not cls.supports_dtype(dtype): + invalid_reasons.append("dtype not supported") + if not cls.supports_kv_cache_dtype(kv_cache_dtype): + invalid_reasons.append("kv_cache_dtype not supported") + if not cls.supports_block_size(block_size): + invalid_reasons.append("block_size not supported") + if use_mm_prefix and not cls.supports_mm_prefix(): + invalid_reasons.append( + "partial multimodal token full attention not supported" + ) + if use_mla != cls.is_mla(): + if use_mla: + invalid_reasons.append("MLA not supported") + else: + invalid_reasons.append("non-MLA not supported") + if has_sink and not cls.supports_sink(): + invalid_reasons.append("attention sinks not supported") + if use_sparse != cls.is_sparse(): + if use_sparse: + invalid_reasons.append("sparse not supported") + else: + invalid_reasons.append("non-sparse not supported") + if use_per_head_quant_scales and not cls.supports_per_head_quant_scales(): + invalid_reasons.append("per-head quant scales not supported") + if not cls.supports_compute_capability(device_capability): + invalid_reasons.append("compute capability not supported") + if not cls.supports_attn_type(attn_type): + invalid_reasons.append(f"attention type {attn_type} not supported") + if has_sliding_window and not cls.supports_sliding_window(): + invalid_reasons.append("sliding window not supported") + if use_non_causal and not cls.supports_non_causal(): + invalid_reasons.append("non-causal attention not supported") + if use_mla and use_non_causal and use_dcp and not cls.supports_non_causal_dcp(): + invalid_reasons.append("non-causal MLA attention with DCP not supported") + if use_batch_invariant and not cls.supports_batch_invariance(): + invalid_reasons.append("batch invariance not supported") + if use_kv_connector and not cls.supports_kv_connector(): + invalid_reasons.append("KV connector not supported") + if use_pcp and not cls.supports_pcp(): + invalid_reasons.append("PCP not supported") + if ( + use_adaptive_verification + and not cls.supports_device_cpu_query_lens_mismatch() + ): + invalid_reasons.append( + "device-cpu query lens mismatch not supported, " + "this is needed for adaptive verification" + ) + combination_reason = cls.supports_combination( + head_size, + dtype, + kv_cache_dtype, + block_size, + use_mla, + has_sink, + use_sparse, + use_mm_prefix, + device_capability, + ) + if combination_reason is not None: + invalid_reasons.append(combination_reason) + return invalid_reasons + + @classmethod + def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...] | None: + """Layouts this backend's kernels can consume, most preferred first, or + None when the kernels consume any layout and express no preference.""" + return None + + @classmethod + def is_ssm(cls) -> bool: + return False + + +class AttentionMetadata: + pass + + +T = TypeVar("T", bound=AttentionMetadata) + + +@dataclass +class CommonAttentionMetadata: + """ + Per-batch attention metadata, shared across layers and backends. + AttentionMetadataBuilder instances use it to construct per-layer metadata. + + For many of the tensors we keep both GPU and CPU versions. + """ + + query_start_loc: torch.Tensor + query_start_loc_cpu: torch.Tensor + """(batch_size + 1,), the start location of each request in query Tensor""" + + seq_lens: torch.Tensor + """(batch_size,), the number of computed tokens for each request""" + + num_reqs: int + """Number of requests""" + # TODO(lucas): rename to num_tokens since it may be padded and this is misleading + num_actual_tokens: int + """Total number of tokens in batch""" + max_query_len: int + """Longest query in batch""" + max_seq_len: int + """Longest context length (may be an upper bound)""" + + block_table_tensor: torch.Tensor + slot_mapping: torch.Tensor + + causal: bool | torch.Tensor = True + + # Needed by FastPrefillAttentionBuilder + logits_indices_padded: torch.Tensor | None = None + num_logits_indices: int | None = None + max_logits_per_req: int | None = None + + # Needed by CrossAttentionBuilder + encoder_seq_lens: torch.Tensor | None = None + encoder_seq_lens_cpu: np.ndarray | None = None + + dcp_local_seq_lens: torch.Tensor | None = None + dcp_local_seq_lens_cpu: torch.Tensor | None = None + """Sequence lengths of the local rank in decode context parallelism world""" + + positions: torch.Tensor | None = None + """(num_actual_tokens,) token positions. Optional; set when the caller + has positions available so that builders can pre-compute position-dependent + sparse metadata for DeepSeek V4 C128A layers.""" + + is_prefilling: torch.Tensor | None = None + """(batch_size,) bool tensor: True if request is still in prefill phase + (num_computed_tokens < num_prompt_tokens). Used by some backends to + distinguish actual decodes from short extends.""" + + seq_lens_cpu_upper_bound: torch.Tensor | None = None + """(batch_size,) CPU upper bound on seq_lens. Precise for prefill rows + and for all rows outside async spec decode; optimistic for async-spec + decode rows (assumes every draft was accepted). Not safe for kernels + that need exact per-row context lengths on decode rows.""" + + mm_req_doc_ranges: dict[int, list[tuple[int, int]]] | None = None + """PrefixLM bidirectional ranges for multimodal tokens. Maps + request index to list of (start, end) token position ranges + where bidirectional attention should apply. None for text-only + batches or non-PrefixLM models. A request's ranges must not overlap.""" + + rswa_prefix_lens: torch.Tensor | None = None + """(batch_size,) per-request prefix length (prompt/image token count) for + Reference Sliding Window Attention (R-SWA). Tokens with logical index below + this stay globally visible; later (generated) tokens additionally see a + fixed sliding window. None disables R-SWA. The attention backend copies this + into its own persistent buffer and reads ``rswa_window`` from model config.""" + + replayssm_decode_base_cpu: torch.Tensor | None = None + """(batch_size,) CPU ring origin for Mamba2 ReplaySSM decode: num_computed + at the current decode run's last full-state write. write_pos counts from + here, so a preemption-resumed request re-anchors past the prompt boundary.""" + + recurrent_prefill_checkpoint_plans_cpu: ( + list[tuple[int, int, tuple[int, ...]] | None] | None + ) = None + """Current-batch-row plans: absolute start, end and checkpoint positions. + Populated only for explicitly planned eager recurrent prefills; padding + rows are None. Targets must refer to allocated, request-owned state slots. + """ + + # WARNING: Deprecated fields. Will be removed in a future release (v0.15.0) + _seq_lens_cpu: torch.Tensor | None = None + _num_computed_tokens_cpu: torch.Tensor | None = None + + _num_computed_tokens_cache: torch.Tensor | None = None + _token_to_req_indices_cache: torch.Tensor | None = None + + def batch_size(self) -> int: + return self.seq_lens.shape[0] + + def naive_query_lens(self) -> torch.Tensor: + """Naive because it assumes that query ends where the next query starts.""" + return self.query_start_loc[1:] - self.query_start_loc[:-1] + + def replace(self, **kwargs) -> "CommonAttentionMetadata": + return replace(self, **kwargs) + + @property + @deprecated( + """ + Prefer using device seq_lens directly to avoid implicit H<>D sync. + If a CPU copy is needed, use `seq_lens.cpu()` instead. + Will be removed in a future release, please migrate as soon as possible. + """ + ) + def seq_lens_cpu(self) -> torch.Tensor: + if self._seq_lens_cpu is None: + self._seq_lens_cpu = self.seq_lens.to("cpu") + return self._seq_lens_cpu + + @property + @deprecated( + """ + Prefer using device seq_lens directly to avoid implicit H<>D sync which breaks full + async scheduling. If a CPU copy is needed, it can be derived from + query_start_loc_cpu and seq_lens. + Will be removed in a future release, please migrate as soon as possible. + """ + ) + def num_computed_tokens_cpu(self) -> torch.Tensor: + if self._num_computed_tokens_cpu is None: + query_seq_lens = ( + self.query_start_loc_cpu[1:] - self.query_start_loc_cpu[:-1] + ) + self._num_computed_tokens_cpu = self.seq_lens_cpu - query_seq_lens + return self._num_computed_tokens_cpu + + def compute_num_computed_tokens(self) -> torch.Tensor: + """Compute num_computed_tokens on device (seq_lens - query_lens).""" + if self._num_computed_tokens_cache is None: + query_lens = self.query_start_loc[1:] - self.query_start_loc[:-1] + self._num_computed_tokens_cache = self.seq_lens - query_lens + return self._num_computed_tokens_cache + + def token_to_req_indices(self, buffer: torch.Tensor) -> torch.Tensor: + """Build or reuse the per-token request index mapping.""" + num_tokens = self.num_actual_tokens + if self._token_to_req_indices_cache is not None: + assert self._token_to_req_indices_cache.device == buffer.device + assert self._token_to_req_indices_cache.dtype == torch.int32 + assert self._token_to_req_indices_cache.shape[0] >= num_tokens + return self._token_to_req_indices_cache[:num_tokens] + + # Built from the device query_start_loc: adaptive verification decides the + # per-request draft split on device, so the CPU copy carries the right total + # but not the right per-request boundaries. Padding requests have a query + # length of zero and drop out of the repeat. + num_mapped_tokens = int(self.query_start_loc_cpu[-1]) + query_lens = self.query_start_loc[1:] - self.query_start_loc[:-1] + assert buffer.shape[0] >= max(num_mapped_tokens, num_tokens) + token_to_req_indices = torch.repeat_interleave( + torch.arange(query_lens.shape[0], dtype=torch.int32, device=buffer.device), + query_lens, + output_size=num_mapped_tokens, + ) + buffer[:num_mapped_tokens].copy_(token_to_req_indices) + if num_mapped_tokens < num_tokens: + buffer[num_mapped_tokens:num_tokens].zero_() + self._token_to_req_indices_cache = buffer[: max(num_mapped_tokens, num_tokens)] + return self._token_to_req_indices_cache[:num_tokens] + + # TODO(lucas): remove once we have FULL-CG spec-decode support + def unpadded( + self, num_actual_tokens: int, num_actual_reqs: int + ) -> "CommonAttentionMetadata": + maybe_slice_reqs = lambda x: x[:num_actual_reqs] if x is not None else None + return CommonAttentionMetadata( + query_start_loc=self.query_start_loc[: num_actual_reqs + 1], + query_start_loc_cpu=self.query_start_loc_cpu[: num_actual_reqs + 1], + seq_lens=self.seq_lens[:num_actual_reqs], + _seq_lens_cpu=self._seq_lens_cpu[:num_actual_reqs] + if self._seq_lens_cpu is not None + else None, + _num_computed_tokens_cpu=self._num_computed_tokens_cpu[:num_actual_reqs] + if self._num_computed_tokens_cpu is not None + else None, + num_reqs=num_actual_reqs, + num_actual_tokens=num_actual_tokens, + max_query_len=self.max_query_len, + max_seq_len=self.max_seq_len, + block_table_tensor=self.block_table_tensor[:num_actual_reqs], + slot_mapping=self.slot_mapping[:num_actual_tokens], + causal=self.causal[:num_actual_reqs] + if isinstance(self.causal, torch.Tensor) + else self.causal, + logits_indices_padded=self.logits_indices_padded, + num_logits_indices=self.num_logits_indices, + max_logits_per_req=self.max_logits_per_req, + encoder_seq_lens=maybe_slice_reqs(self.encoder_seq_lens), + encoder_seq_lens_cpu=maybe_slice_reqs(self.encoder_seq_lens_cpu), + dcp_local_seq_lens=maybe_slice_reqs(self.dcp_local_seq_lens), + dcp_local_seq_lens_cpu=maybe_slice_reqs(self.dcp_local_seq_lens_cpu), + is_prefilling=maybe_slice_reqs(self.is_prefilling), + rswa_prefix_lens=maybe_slice_reqs(self.rswa_prefix_lens), + replayssm_decode_base_cpu=maybe_slice_reqs(self.replayssm_decode_base_cpu), + recurrent_prefill_checkpoint_plans_cpu=maybe_slice_reqs( + self.recurrent_prefill_checkpoint_plans_cpu + ), + ) + + +M = TypeVar("M") + + +class AttentionCGSupport(Enum): + """Constants for the cudagraph support of the attention backend + Here we do not consider the cascade attention, as currently + it is never cudagraph supported.""" + + ALWAYS = 3 + """Cudagraph always supported; supports mixed-prefill-decode""" + UNIFORM_BATCH = 2 + """Cudagraph supported for batches the only contain query lengths that are + the same, this can be used for spec-decode + i.e. "decodes" are 1 + num_speculative_tokens""" + UNIFORM_SINGLE_TOKEN_DECODE = 1 + """Cudagraph supported for batches the only contain query_len==1 decodes""" + NEVER = 0 + """NO cudagraph support""" + + +class AttentionMetadataBuilder(ABC, Generic[M]): + # Does this backend/builder support CUDA Graphs for attention (default: no). + # Do not access directly. Call get_cudagraph_support() instead. + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.NEVER + # Does this backend/builder reorder the batch? + # If not, set this to None. Otherwise set it to the query + # length that will be pulled into the front of the batch. + reorder_batch_threshold: int | None = None + # Does this backend/builder support updating the block table in existing + # metadata + supports_update_block_table: bool = False + # Whether the builder constructor requires the block-table width. + requires_block_table_width: ClassVar[bool] = False + # Whether all step-dependent draft decode metadata can be updated in place, + # allowing one metadata build to be reused across autoregressive draft steps. + supports_draft_decode_metadata_update: bool = False + + @abstractmethod + def __init__( + self, + kv_cache_spec: "KVCacheSpec", + layer_names: list[str], + vllm_config: "VllmConfig", + device: torch.device, + ): + self.kv_cache_spec = kv_cache_spec + self.layer_names = layer_names + self.vllm_config = vllm_config + self.device = device + + @classmethod + def get_cudagraph_support( + cls: type["AttentionMetadataBuilder"], + vllm_config: "VllmConfig", + kv_cache_spec: "KVCacheSpec", + ) -> AttentionCGSupport: + """Get the cudagraph support level of this builder class.""" + return cls._cudagraph_support + + def _init_reorder_batch_threshold( + self, + reorder_batch_threshold: int | None = 1, + supports_spec_as_decode: bool = False, + supports_dcp_with_varlen: bool = False, + ) -> None: + self.reorder_batch_threshold = reorder_batch_threshold + if self.reorder_batch_threshold is not None and supports_spec_as_decode: + # If the backend supports spec-as-decode kernels, then we can set + # the reorder_batch_threshold based on the number of speculative + # tokens from the config. + speculative_config = self.vllm_config.speculative_config + if ( + speculative_config is not None + and speculative_config.num_speculative_tokens is not None + ): + max_num_queries_for_spec = ( + 1 + + (2 if speculative_config.parallel_drafting else 1) + * speculative_config.num_speculative_tokens + ) + self.reorder_batch_threshold = max( + self.reorder_batch_threshold, + max_num_queries_for_spec, + ) + + if ( + self.vllm_config.parallel_config.decode_context_parallel_size > 1 + and not supports_dcp_with_varlen + ): + self.reorder_batch_threshold = 1 + + @abstractmethod + def build( + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + fast_build: bool = False, + ) -> M: + """ + Central method that builds attention metadata. + Some builders (MLA) require reorder_batch to be called prior to build. + + Args: + common_prefix_len: The length of the common prefix of the batch. + common_attn_metadata: The common attention metadata. + fast_build: The meta-data will prioritize speed of building over + then speed at execution. Can be used for spec-decode where the + result of a build call may only be used for few layers/iters. + """ + raise NotImplementedError + + def update_block_table( + self, + metadata: M, + blk_table: torch.Tensor, + slot_mapping: torch.Tensor, + ) -> M: + """ + Update the block table for the attention metadata. + Faster when theres multiple kv-cache groups that create virtually the + same metadata but just with different block tables. + + Only needs to be implemented if supports_update_block_table is True. + """ + raise NotImplementedError + + def build_for_cudagraph_capture( + self, common_attn_metadata: CommonAttentionMetadata + ) -> M: + """ + Build attention metadata for CUDA graph capture. Uses build by default. + Subclasses that override this method should call self.build or + super().build_for_cudagraph_capture. + """ + return self.build( + common_prefix_len=0, common_attn_metadata=common_attn_metadata + ) + + def build_for_drafting( + self, + common_attn_metadata: CommonAttentionMetadata, + draft_index: int, + ) -> M: + """ + Build attention metadata for draft model. Uses build by default. + + Args: + common_attn_metadata: The common attention metadata. + draft_index: The index of the current draft operation. + When speculating a chain of tokens, this index refers to the + draft attempt for the i-th token. + For tree-based attention, this index instead refers to the + draft attempt for the i-th level in the tree of tokens. + """ + return self.build( + common_prefix_len=0, + common_attn_metadata=common_attn_metadata, + fast_build=True, + ) + + def update_draft_decode_metadata(self, metadata: M) -> None: + """Update step-dependent draft decode metadata in place. + + The fused draft loop may call this method during full CUDA graph + capture. CUDA graph replay does not run this Python method, so + implementations must emit capture-safe operations and keep replayed + tensor state in persistent storage. + """ + raise NotImplementedError + + def use_cascade_attention( + self, + common_prefix_len: int, + query_lens: np.ndarray, + num_query_heads: int, + num_kv_heads: int, + use_alibi: bool, + use_sliding_window: bool, + use_local_attention: bool, + num_sms: int, + dcp_world_size: int, + ) -> bool: + return False + + +class AttentionLayer(Protocol): + _q_scale: torch.Tensor + _k_scale: torch.Tensor + _k_scale_cpu: torch.Tensor + _v_scale: torch.Tensor + _v_scale_cpu: torch.Tensor + _q_scale_float: float + _k_scale_float: float + _v_scale_float: float + _prob_scale: torch.Tensor + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + kv_cache: torch.Tensor, + attn_metadata: AttentionMetadata, + ) -> torch.Tensor: ... + + +class AttentionImplBase(ABC, Generic[T]): + """Base class for attention implementations. + + Contains common attributes and initialization logic shared by both + standard AttentionImpl and MLAAttentionImpl. Does not define a forward + method - subclasses define their own forward interfaces. + """ + + # Whether this impl uses a sparse (top-k) attention path. Used by MLA to + # route between the dense-MHA prefill and sparse-MQA paths. + is_sparse: ClassVar[bool] = False + + # Whether this impl provides a dense-MHA prefill path (forward_mha). Sparse + # impls without one run the top-k MQA path for all requests. + supports_dense_mha_prefill: ClassVar[bool] = True + + # Required attributes that all impls should have + num_heads: int + head_size: int + scale: float + + # Whether the attention impl can return the softmax lse for decode. + # Some features like decode context parallelism require the softmax lse. + can_return_lse_for_decode: bool = False + + # Base of the logarithm used by this backend when returning softmax lse. + # True => natural log (lse = ln(sum(exp(qk)))) + # -- e.g. Triton MLA, FlashAttention, FlashMLA, Cutlass MLA + # False => base 2 (lse = log2(sum(exp(qk)))) + # -- e.g. FlashInfer trtllm-gen MLA + # The DCP combine kernel (cp_lse_ag_out_rs / dcp_a2a_lse_reduce in + # vllm/v1/attention/ops/dcp.py) branches on this via its IS_BASE_E + # constexpr; getting it wrong silently corrupts the cross-shard + # softmax denominator. + lse_base_on_e: bool = True + + # Whether the attention impl supports Prefill Context Parallelism. + supports_pcp: bool = False + # Whether the attention impl supports Decode Context Parallelism. + supports_dcp: bool = True + # Whether the attention impl(or ops) supports MTP + # when cp_kv_cache_interleave_size > 1 + supports_mtp_with_cp_non_trivial_interleave_size: bool = False + + # some attention backends might not always want to return lse + # even if they can return lse (for efficiency reasons) + need_to_return_lse_for_decode: bool = False + + # Whether this attention implementation supports pre-quantized query input. + # When True, the attention layer will quantize queries before passing them + # to this backend, allowing torch.compile to fuse the quantization with + # previous operations. This is typically supported when using FP8 KV cache + # with compatible attention kernels (e.g., TRT-LLM). + # Subclasses should set this in __init__. + # TODO add support to more backends: + # https://github.com/vllm-project/vllm/issues/25584 + supports_quant_query_input: bool = False + + dcp_world_size: int + dcp_rank: int + + pcp_world_size: int + pcp_rank: int + + total_cp_world_size: int + total_cp_rank: int + + def __new__(cls, *args, **kwargs): + # use __new__ so that all subclasses will call this + self = super().__new__(cls) + try: + from vllm.distributed.parallel_state import get_dcp_group + + self.dcp_world_size = get_dcp_group().world_size + self.dcp_rank = get_dcp_group().rank_in_group + except AssertionError: + # DCP might not be initialized in testing + self.dcp_world_size = 1 + self.dcp_rank = 0 + try: + from vllm.distributed.parallel_state import get_pcp_group + + self.pcp_world_size = get_pcp_group().world_size + self.pcp_rank = get_pcp_group().rank_in_group + except AssertionError: + self.pcp_world_size = 1 + self.pcp_rank = 0 + self.total_cp_world_size = self.dcp_world_size + self.total_cp_rank = self.dcp_rank + + self.need_to_return_lse_for_decode = ( + self.dcp_world_size > 1 and self.can_return_lse_for_decode + ) + return self + + def process_weights_after_loading(self, act_dtype: torch.dtype): + pass + + +class AttentionImpl(AttentionImplBase[T], Generic[T]): + """Standard attention implementation with forward method.""" + + kv_cache_dtype: str + + @property + def kv_quant_mode(self) -> "KVQuantMode": + """Return the KV cache quantization mode for this layer.""" + return get_kv_quant_mode(self.kv_cache_dtype) + + @abstractmethod + def __init__( + self, + num_heads: int, + head_size: int, + scale: float, + num_kv_heads: int | None = None, + alibi_slopes: list[float] | None = None, + sliding_window: int | None = None, + kv_cache_dtype: str = "auto", + logits_soft_cap: float | None = None, + attn_type: str = AttentionType.DECODER, + kv_sharing_target_layer_name: str | None = None, + ) -> None: + raise NotImplementedError + + @abstractmethod + def forward( + self, + layer: AttentionLayer, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + kv_cache: torch.Tensor, + attn_metadata: T, + output: torch.Tensor, + output_scale: torch.Tensor | None = None, + output_block_scale: torch.Tensor | None = None, + ) -> torch.Tensor: + raise NotImplementedError + + def fused_output_quant_supported(self, quant_key: "QuantKey") -> bool: + """ + Does this attention implementation support fused output quantization. + This is used by the AttnFusionPass to only fuse output quantization + onto implementations that support it. + + Args: + quant_key: QuantKey object that describes the quantization op + + Returns: + is fusion supported for this type of quantization + """ + return False + + def fused_qk_norm_rope_kvcache_supported(self): + """ + Does this attention implementation support fused QKNorm+RoPE+KVCache fusion. + This is used by the QkNormRopeKvCachePattern to only fuse the QKNorm ops + with the RoPE ops and the KV cache update for implementations that support it. + """ + return False + + def fused_rope_kvcache_supported(self): + """ + Does this attention implementation support RoPE+KVCache fusion. + This is used by the RopeKVCacheFusionPass to only fuse the RoPE ops + with the KV cache update for implementations that support it. + """ + return False + + def do_qk_norm_rope_kvcache_update( + self, + layer: AttentionLayer, + qkv: torch.Tensor, + q_out: torch.Tensor, + k_out: torch.Tensor, + positions: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + rms_norm_eps: float, + cos_sin_cache: torch.Tensor, + is_neox: bool, + kv_cache: torch.Tensor, + layer_slot_mapping: torch.Tensor, + ): + """ + If `fused_qk_norm_rope_kvcache_supported` returns True, this method + will be called by the fused custom op. Applies QK-norm + RoPE and + writes K/V to the KV cache. Results are written to the pre-allocated + q_out and k_out tensors; V is split from QKV at the graph level. + """ + raise NotImplementedError + + def do_rope_and_kv_cache_update( + self, + layer: AttentionLayer, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + is_neox: bool, + kv_cache: torch.Tensor, + layer_slot_mapping: torch.Tensor, + ): + """ + If `fused_rope_kvcache_supported` returns True, this method will be called + by torch.ops.vllm.fused_rope_and_unified_kv_cache_update + to perform the inplace RoPE and KV cache update. + """ + raise NotImplementedError + + +class MLAAttentionImpl(AttentionImplBase[T], Generic[T]): + """MLA attention implementation with forward_mqa and forward_mha methods.""" + + supports_pcp: bool = True + + def uses_full_ckv_dcp(self, attn_metadata: T, num_tokens: int) -> bool: + """Whether this call attends a transient globally gathered DCP cache.""" + return False + + @abstractmethod + def __init__( + self, + num_heads: int, + head_size: int, + scale: float, + num_kv_heads: int, + alibi_slopes: list[float] | None, + sliding_window: int | None, + kv_cache_dtype: str, + logits_soft_cap: float | None, + attn_type: str, + kv_sharing_target_layer_name: str | None, + # MLA Specific Arguments + q_lora_rank: int | None, + kv_lora_rank: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + qk_head_dim: int, + v_head_dim: int, + kv_b_proj: "ColumnParallelLinear", + indexer: object | None = None, + q_pad_num_heads: int | None = None, + ) -> None: + raise NotImplementedError + + def forward_mha( + self, + q: torch.Tensor, + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + kv_c_and_k_pe_cache: torch.Tensor, + attn_metadata: T, + k_scale: torch.Tensor, + output: torch.Tensor, + output_scale: torch.Tensor | None = None, + ) -> None: + """MHA-style prefill forward pass.""" + raise NotImplementedError + + @abstractmethod + def forward_mqa( + self, + q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + kv_c_and_k_pe_cache: torch.Tensor, + attn_metadata: T, + layer: AttentionLayer, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """MQA-style decode forward pass.""" + raise NotImplementedError + + def fused_output_quant_supported(self, quant_key: "QuantKey"): + """ + Does this attention implementation support fused output quantization. + Since MLA quantization is done manually in forward_impl (common code), + all MLA backends support it by default. + """ + return quant_key in ( + kFp8StaticTensorSym, + kNvfp4Dynamic, + kFp8Dynamic128Sym, + kFp8Dynamic64Sym, + ) + + def do_kv_cache_update( + self, + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + kv_cache: torch.Tensor, + slot_mapping: torch.Tensor, + kv_cache_dtype: str, + k_scale: torch.Tensor, + ) -> None: + if kv_cache.numel() == 0: + return + from vllm import _custom_ops as ops + + ops.concat_and_cache_mla( + kv_c_normed, + k_pe.squeeze(1), + kv_cache, + slot_mapping.flatten(), + kv_cache_dtype=kv_cache_dtype, + scale=k_scale, + ) + + +def subclass_attention_backend( + name_prefix: str, + attention_backend_cls: type[AttentionBackend], + builder_cls: type[AttentionMetadataBuilder[M]], +) -> type[AttentionBackend]: + """ + Return a new subclass where `get_builder_cls` returns `builder_cls`. + """ + name: str = name_prefix + attention_backend_cls.__name__ # type: ignore + + return type( + name, (attention_backend_cls,), {"get_builder_cls": lambda: builder_cls} + ) + + +def subclass_attention_backend_with_overrides( + name_prefix: str, + attention_backend_cls: type[AttentionBackend], + overrides: dict[str, Any], +) -> type[AttentionBackend]: + name: str = name_prefix + attention_backend_cls.__name__ # type: ignore + return type(name, (attention_backend_cls,), overrides) diff --git a/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/ce9460834e08f97dbbfeb3f1238b78ee6a3363dd59a2aef8ceeb715857385895/scheduler.py b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/ce9460834e08f97dbbfeb3f1238b78ee6a3363dd59a2aef8ceeb715857385895/scheduler.py new file mode 100644 index 00000000..a8e3b379 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/ce9460834e08f97dbbfeb3f1238b78ee6a3363dd59a2aef8ceeb715857385895/scheduler.py @@ -0,0 +1,3442 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import itertools +import time +import os +from vllm.v1.core.recurrent_prefill_checkpoint import fresh_prompt_plan + +from collections import defaultdict, deque +from collections.abc import Iterable +from dataclasses import replace +from typing import Any + +from vllm.compilation.cuda_graph import CUDAGraphStat +from vllm.config import KVEventsConfig, VllmConfig +from vllm.distributed.ec_transfer.ec_connector.base import ( + ECConnectorBase, + ECConnectorMetadata, + ECConnectorRole, +) +from vllm.distributed.ec_transfer.ec_connector.factory import ECConnectorFactory +from vllm.distributed.kv_events import EventPublisherFactory, KVEventBatch +from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory +from vllm.distributed.kv_transfer.kv_connector.v1 import ( + KVConnectorBase_V1, + KVConnectorRole, + SupportsHMA, +) +from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata +from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( + RoutedExpertsManager, +) +from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry +from vllm.multimodal.encoder_budget import MultiModalBudget +from vllm.multimodal.utils import get_mm_features_in_window +from vllm.v1.core.encoder_cache_manager import ( + EncoderCacheManager, + EncoderDecoderCacheManager, +) +from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager +from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector +from vllm.v1.core.kv_cache_utils import KVCacheBlock +from vllm.v1.core.sched.interface import PauseState, SchedulerInterface +from vllm.v1.core.sched.output import ( + CachedRequestData, + GrammarOutput, + NewRequestData, + ScheduledEncoderInputStats, + SchedulerOutput, +) +from vllm.v1.core.sched.request_queue import ( + RequestQueue, + SchedulingPolicy, + create_request_queue, +) +from vllm.v1.core.sched.utils import check_stop, remove_all +from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs +from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec +from vllm.v1.metrics.perf import ModelMetrics, PerfStats +from vllm.v1.metrics.stats import ( + PrefixCacheStats, + RequestSpecDecodeMetrics, + SchedulerStats, +) +from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput +from vllm.v1.request import Request, RequestStatus, StreamingUpdate +from vllm.v1.spec_decode.dynamic.acceptance_length import ( + BatchSizeAcceptanceLengthController, +) +from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup +from vllm.v1.spec_decode.metrics import SpecDecodingStats +from vllm.v1.structured_output import StructuredOutputGrammar, StructuredOutputManager +from vllm.v1.utils import record_function_or_nullcontext + +logger = init_logger(__name__) + + +class Scheduler(SchedulerInterface): + def __init__( + self, + vllm_config: VllmConfig, + kv_cache_config: KVCacheConfig, + structured_output_manager: StructuredOutputManager, + block_size: int, + hash_block_size: int | None = None, + mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY, + include_finished_set: bool = False, + log_stats: bool = False, + ) -> None: + self.vllm_config = vllm_config + self.scheduler_config = vllm_config.scheduler_config + self.cache_config = vllm_config.cache_config + self.lora_config = vllm_config.lora_config + self.model_uses_mrope = vllm_config.model_config.uses_mrope + self.model_uses_xdrope = vllm_config.model_config.uses_xdrope + self.kv_cache_config = kv_cache_config + self.kv_events_config = vllm_config.kv_events_config + self.parallel_config = vllm_config.parallel_config + self.log_stats = log_stats + self.observability_config = vllm_config.observability_config + self.spec_decode_metrics_level = ( + self.observability_config.per_request_spec_decode_metrics + ) + self.kv_metrics_collector: KVCacheMetricsCollector | None = None + if self.observability_config.kv_cache_metrics: + self.kv_metrics_collector = KVCacheMetricsCollector( + self.observability_config.kv_cache_metrics_sample, + ) + self.structured_output_manager = structured_output_manager + self.is_encoder_decoder = vllm_config.model_config.is_encoder_decoder + self.is_mm_encoder_only = vllm_config.is_mm_encoder_only + + # include_finished_set controls whether a separate set of finished + # request ids should be included in the EngineCoreOutputs returned + # by update_from_outputs(). This is currently used in the multi-engine + # case to track request lifetimes efficiently. + self.finished_req_ids_dict: dict[int, set[str]] | None = ( + defaultdict(set) if include_finished_set else None + ) + # Track requests scheduled in prior step (MRV1-only). + self.prev_step_scheduled_req_ids: set[str] = set() + + # Scheduling constraints. + self.max_num_running_reqs = self.scheduler_config.max_num_seqs + self.max_num_scheduled_tokens = ( + self.scheduler_config.max_num_scheduled_tokens + if self.scheduler_config.max_num_scheduled_tokens is not None + else self.scheduler_config.max_num_batched_tokens + ) + self.max_model_len = vllm_config.model_config.max_model_len + self.enable_kv_cache_events = ( + self.kv_events_config is not None + and self.kv_events_config.enable_kv_cache_events + ) + # Diffusion models may not sample any tokens for a denoising step. + self.num_sampled_tokens_per_step = ( + 1 if not vllm_config.model_config.is_diffusion else 0 + ) + + # Create KVConnector for the Scheduler. Note that each Worker + # will have a corresponding KVConnector with Role=WORKER. + # KV Connector pushes/pull of remote KVs for P/D and offloading. + self.connector = None + self.connector_prefix_cache_stats: PrefixCacheStats | None = None + self.recompute_kv_load_failures = True + self.defer_block_free = False + # Whether a preempted request's in-flight output must be dropped; see + # KVConnectorBase_V1.requires_kv_delivery. + self.requires_kv_delivery = False + kv_transfer_config = self.vllm_config.kv_transfer_config + if kv_transfer_config is not None: + assert not self.is_encoder_decoder, ( + "Encoder-decoder models are not currently supported with KV connectors" + ) + self.connector = KVConnectorFactory.create_connector( + config=self.vllm_config, + role=KVConnectorRole.SCHEDULER, + kv_cache_config=self.kv_cache_config, + ) + if self.log_stats: + self.connector_prefix_cache_stats = PrefixCacheStats() + kv_load_failure_policy = kv_transfer_config.kv_load_failure_policy + self.recompute_kv_load_failures = kv_load_failure_policy == "recompute" + + # With overlapping batches (async scheduling or PP), a step may + # still be writing a freed request's KV blocks. A consumer KV + # Connector can reallocate and fill those blocks via a load that + # isn't ordered against that write, so defer freeing them. + multiple_inflight_batches = self.vllm_config.max_concurrent_batches > 1 + if multiple_inflight_batches and ( + kv_transfer_config.is_kv_consumer + or bool( + getattr(self.connector, "supports_recurrent_boundary_blocks", False) + ) + ): + self.defer_block_free = True + + self.requires_kv_delivery = self.connector.requires_kv_delivery + + self.kv_event_publisher = EventPublisherFactory.create( + self.kv_events_config, + self.parallel_config.data_parallel_index, + ) + self.ec_connector = None + if self.vllm_config.ec_transfer_config is not None: + self.ec_connector = ECConnectorFactory.create_connector( + config=self.vllm_config, role=ECConnectorRole.SCHEDULER + ) + + num_gpu_blocks = self.cache_config.num_gpu_blocks + assert num_gpu_blocks is not None and num_gpu_blocks > 0 + + self.block_size = block_size + self.dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size + self.pcp_world_size = vllm_config.parallel_config.prefill_context_parallel_size + + # req_id -> Request + self.requests: dict[str, Request] = {} + # Scheduling policy + try: + self.policy = SchedulingPolicy(self.scheduler_config.policy) + except ValueError as e: + raise ValueError( + f"Unknown scheduling policy: {self.scheduler_config.policy}" + ) from e + # Priority queues for requests. + self.waiting = create_request_queue(self.policy) + # requests skipped in waiting flow due async deps or constraints. + self.skipped_waiting = create_request_queue(self.policy) + self.running: list[Request] = [] + + # The request IDs that are finished in between the previous and the + # current steps. This is used to notify the workers about the finished + # requests so that they can free the cached states for those requests. + # This is flushed at the end of each scheduling step. + self.finished_req_ids: set[str] = set() + + # IDs of requests preempted since the last call to schedule(). + self.reset_preempted_req_ids: set[str] = set() + + # Counter for requests waiting for streaming input. Used to calculate + # number of unfinished requests + self.num_waiting_for_streaming_input: int = 0 + + # KV Connector: requests in process of async KV loading or recving + self.finished_recving_kv_req_ids: set[str] = set() + self.failed_recving_kv_req_ids: set[str] = set() + # lease_key -> (leader request id, non-empty scheduler step whose + # worker completion proves every dedicated partial-page copy ran) + self._pending_shared_prefix_leases: dict[str, tuple[str, int]] = {} + + # Grammar compilation failures to finish as per-request errors in + # update_from_output. + self.grammar_compile_error_reqs: set[str] = set() + + # Encoder-related. + # Calculate encoder cache size if applicable + supports_mm_inputs = mm_registry.supports_multimodal_inputs( + vllm_config.model_config + ) + mm_budget = ( + MultiModalBudget(vllm_config, mm_registry) if supports_mm_inputs else None + ) + + # NOTE: Text-only encoder-decoder models are implemented as + # multi-modal models for convenience + # Example: https://github.com/vllm-project/bart-plugin + if self.is_encoder_decoder: + assert mm_budget and len(mm_budget.mm_max_toks_per_item) <= 1, ( + "Encoder-decoder models are expected to implement the " + "multimodal interface with at most one modality." + ) + + self.max_num_encoder_input_tokens = ( + mm_budget.encoder_compute_budget if mm_budget else 0 + ) + encoder_cache_size = mm_budget.encoder_cache_size if mm_budget else 0 + manager_cls_obj = vllm_config.ec_manager_config.get_encoder_cache_manager_obj() + if manager_cls_obj is None: + manager_cls_obj = ( + EncoderDecoderCacheManager + if self.is_encoder_decoder + else EncoderCacheManager + ) + self.encoder_cache_manager = manager_cls_obj.create_manager( + cache_size=encoder_cache_size, vllm_config=vllm_config + ) + speculative_config = vllm_config.speculative_config + self.use_eagle = False + self.num_spec_tokens = vllm_config.num_speculative_tokens + self.num_lookahead_tokens = vllm_config.num_lookahead_tokens + # Positions past the computed tokens that the drafter reads mid-prefill. + # Eagle-family drafters read 1 ahead, but multi-module MTP reads + # num_spec_tokens ahead at chunked-prefill boundaries. Determines the + # encoder scheduling shift, the deferred encoder free, the KV cache + # manager's re-prefillable window (this minus 1), and how many tokens to + # reserve between a chunk boundary and the prefill end. + self.num_prefill_lookahead = 0 + self.dynamic_sd_lookup: list[int] | None = None + self.acceptance_length_controller: ( + BatchSizeAcceptanceLengthController | None + ) = None + if speculative_config is not None: + if speculative_config.num_speculative_tokens_per_batch_size: + self.dynamic_sd_lookup = build_dynamic_sd_schedule_lookup( + speculative_config.num_speculative_tokens_per_batch_size, + vllm_max_batch_size=self.scheduler_config.max_num_seqs, + vllm_num_speculative_tokens=self.num_spec_tokens, + ) + if ( + observation_window := ( + speculative_config.adaptive_speculative_tokens_window + ) + ) is not None: + self.acceptance_length_controller = BatchSizeAcceptanceLengthController( + max_num_spec_tokens=self.num_spec_tokens, + observation_window=observation_window, + initial_num_spec_tokens=( + speculative_config.adaptive_speculative_tokens_initial + ), + num_spec_tokens_by_batch_size=self.dynamic_sd_lookup, + ) + self.use_eagle = speculative_config.use_eagle() + if self.use_eagle: + self.num_prefill_lookahead = ( + self.num_spec_tokens + if speculative_config.use_multi_module_mtp() + else 1 + ) + + # Create the KV cache manager. + if hash_block_size is None: + hash_block_size = block_size + self.hash_block_size = hash_block_size + self.kv_cache_manager = KVCacheManager( + kv_cache_config=kv_cache_config, + max_model_len=self.max_model_len, + max_in_flight_tokens=vllm_config.max_in_flight_tokens, + enable_caching=self.cache_config.enable_prefix_caching, + use_eagle=self.use_eagle, + num_prefill_lookahead=self.num_prefill_lookahead, + log_stats=self.log_stats, + enable_kv_cache_events=self.enable_kv_cache_events, + dcp_world_size=self.dcp_world_size, + pcp_world_size=1, + scheduler_block_size=self.block_size, + hash_block_size=hash_block_size, + metrics_collector=self.kv_metrics_collector, + watermark=self.scheduler_config.watermark, + ) + # Bind GPU block pool to the KV connector. This must happen after + # kv_cache_manager is constructed so block_pool is available. + if self.connector is not None: + self.connector.bind_gpu_block_pool(self.kv_cache_manager.block_pool) + + self.use_pp = self.parallel_config.pipeline_parallel_size > 1 + self.use_v2_model_runner = vllm_config.use_v2_model_runner + # Scheduler iteration counter. Drives the V2+PP+async decode-throttle + # cadence (`next_decode_eligible_step`). + self.current_step = 0 + # DP prefill balancing: Flag to track whether the last cadence-aligned + # prefill batch fully drained the waiting queue. Prefill throttling + # is disabled in this case. + self.prefill_capacity_bound = False + self.scheduler_reserve_full_isl = ( + self.scheduler_config.scheduler_reserve_full_isl + ) + + self.has_mamba_layers = kv_cache_config.has_mamba_layers + self.needs_kv_cache_zeroing = kv_cache_config.needs_kv_cache_zeroing + # Blocks that async KV loads will overwrite this step, skipped from + # zeroing since the zeroing could race the out-of-band write. + self._skip_zero_block_ids: set[int] = set() + self.need_mamba_block_aligned_split = ( + self.has_mamba_layers and self.cache_config.mamba_cache_mode == "align" + ) + glm5_next_mtp_has_independent_draft_state = ( + speculative_config is not None + and speculative_config.method == "mtp" + and speculative_config.draft_model_config is not None + and "Glm5NextMTPModel" + in (speculative_config.draft_model_config.hf_config.architectures or ()) + ) + self.mamba_has_prefill_checkpoint_blocks = ( + self.has_mamba_layers + # DFlash and GLM-5.3 MTP keep draft attention state independently + # from the target GDN recurrent state. Publishing a target GDN + # checkpoint therefore does not mutate draft KV. Prefix-cache + # lookup still drops and re-prefills the lookahead-dependent MTP + # draft tail through the EAGLE group policy. + and ( + not self.use_eagle + or (speculative_config is not None and speculative_config.use_dflash()) + or glm5_next_mtp_has_independent_draft_state + ) + and all( + not isinstance(group.kv_cache_spec, MambaSpec) + or group.kv_cache_spec.num_prefill_checkpoint_blocks > 0 + for group in kv_cache_config.kv_cache_groups + ) + ) + self._two_checkpoint_prefill_enabled = ( + os.environ.get("SPARK_GDN_PREFILL_CHECKPOINTS") == "2" + and glm5_next_mtp_has_independent_draft_state + and self.mamba_has_prefill_checkpoint_blocks + and all(not isinstance(group.kv_cache_spec, MambaSpec) + or (group.kv_cache_spec.num_prefill_checkpoint_blocks >= 2 + and group.kv_cache_spec.block_size == self.cache_config.block_size) + for group in kv_cache_config.kv_cache_groups) + ) + # An interior recurrent-page hash needs an explicit tail stop. + # DCP can enlarge attention's scheduling unit without making the + # recurrent page larger than a hash, so inspect Mamba specs directly. + self.mamba_partial_cache_hit = ( + self.need_mamba_block_aligned_split + and any( + isinstance(group.kv_cache_spec, MambaSpec) + and self.hash_block_size < group.kv_cache_spec.block_size + for group in kv_cache_config.kv_cache_groups + ) + and self.kv_cache_manager.coordinator.enable_partial_hash_hits + ) + + # Counts of non-empty steps scheduled / processed. update_from_output + # is called once per scheduled step in FIFO order, so these stay in sync. + self.sched_step_seq = 0 + self.processed_step_seq = 0 + # FIFO of (fence_seq, blocks): blocks become safe to free once + # processed_step_seq >= fence_seq. + self.deferred_frees: deque[tuple[int, list[KVCacheBlock]]] = deque() + + self.perf_metrics: ModelMetrics | None = None + if self.log_stats and vllm_config.observability_config.enable_mfu_metrics: + self.perf_metrics = ModelMetrics(vllm_config) + + self.enable_return_routed_experts = ( + vllm_config.model_config.enable_return_routed_experts + ) + self.return_sampling_mask = vllm_config.model_config.return_sampling_mask + + if self.enable_return_routed_experts: + assert self.dcp_world_size == 1 and self.pcp_world_size == 1, ( + "enable_return_routed_experts does not support context parallelism " + "(dcp_world_size > 1 or pcp_world_size > 1)" + ) + + self.routed_experts_mgr = RoutedExpertsManager( + vllm_config=vllm_config, + kv_cache_config=kv_cache_config, + ) + # Block-ID snapshot taken at schedule time (before forward), + # so update_from_output can read slot data even if a later + # schedule() frees the blocks (async scheduling race). + self._re_block_ids: dict[str, list[int]] = {} + + self._pause_state: PauseState = PauseState.UNPAUSED + + # In-flight requests still prefilling (prefill chunks + in-progress + # async KV loads). Their remaining-block reservation gates async loads. + self._inflight_prefills: set[Request] = set() + + def _recurrent_publication_boundaries(self, request: Request) -> tuple[int, ...]: + connector = self.connector + if connector is None: + return () + boundaries = connector.get_recurrent_publication_boundaries(request) + if not isinstance(boundaries, (list, tuple)): + raise ValueError("recurrent publication boundaries must be a sequence") + normalized = tuple(sorted(set(boundaries))) + if any( + type(boundary) is not int + or boundary <= 0 + or boundary >= request.num_prompt_tokens + or boundary % self.hash_block_size != 0 + for boundary in normalized + ): + raise ValueError( + "recurrent publication boundaries must be positive internal-hash" + " multiples below the prompt length" + ) + return normalized + + def _recurrent_publication_boundary_at( + self, request: Request, end: int + ) -> int | None: + return end if end in self._recurrent_publication_boundaries(request) else None + + def _recurrent_checkpoint_plan(self, request, start, end): + if (not self._two_checkpoint_prefill_enabled or start != 0 + or request.num_computed_tokens != 0 or request.has_encoder_inputs): + return None + for manager in self.kv_cache_manager.coordinator.single_type_managers: + if isinstance(manager.kv_cache_spec, MambaSpec) and ( + manager.req_to_blocks.get(request.request_id) + or request.request_id in manager._allocated_block_reqs + or request.request_id in manager._partial_hit_reqs + ): + return None + return fresh_prompt_plan( + start=start, end=end, prompt=request.num_prompt_tokens, + num_tokens=request.num_tokens, block_size=self.cache_config.block_size, + publications=self._recurrent_publication_boundaries(request), + shared_prefix_boundary=request.shared_prefix_boundary, + ) + + def _mamba_block_aligned_split( + self, + request: Request, + num_new_tokens: int, + num_new_local_computed_tokens: int = 0, + num_external_computed_tokens: int = 0, + ) -> int: + """Clip a prefill chunk so it ends where Mamba state must be cached. + + In "align" cache mode reusable SSM states are materialized at block + boundaries, plus mandatory early stops (the prompt's partial-tail hash + boundary, a detected shared-prefix junction). If a block is larger + than the configured prefill chunk limit, intermediate chunks keep + private running state until they reach the next cacheable position. + """ + start = ( + request.num_computed_tokens + + num_new_local_computed_tokens + + num_external_computed_tokens + ) + # Split only during prefill: `request.num_tokens - 1` extends this to + # resumed requests replaying their output tokens. + prefill_end = max(request.num_prompt_tokens, request.num_tokens - 1) + if start >= prefill_end: + return num_new_tokens + + checkpoint_plan = self._recurrent_checkpoint_plan(request, start, start + num_new_tokens) + if checkpoint_plan is not None: + return checkpoint_plan[1] - checkpoint_plan[0] + block_size = self.cache_config.block_size + # The last block-aligned position whose state can be cached. With + # Eagle, FullAttn prunes the last matching block, so back off one + # block to avoid a Mamba cache miss. + last_cache_position = request.num_tokens - request.num_tokens % block_size + if self.use_eagle: + # Lookup excludes the final prompt token, then drops its proof + # block. The corresponding recurrent state needs its own stop. + last_cache_position = max( + (request.num_tokens - 1) // block_size * block_size - block_size, 0 + ) + + end = start + num_new_tokens + use_internal_checkpoint = ( + self.mamba_has_prefill_checkpoint_blocks and start % block_size == 0 + ) + if use_internal_checkpoint and not self.use_eagle: + # The internal checkpoint covers the final aligned state only; + # speculative replay may require the preceding state as well. + last_cache_position = 0 + # Invariant: slot p holds the state after exactly (p + 1) * block_size + # tokens. State is written at chunk ends, so chunk ends must be block + # aligned. Exempt: the prompt's last chunk, whose slot decode advances + # to the boundary. A block too wide for one chunk advances sub-block + # and re-aligns at the next boundary. + checkpoint_covers_prompt_tail = ( + use_internal_checkpoint and end >= request.num_prompt_tokens + ) + if ( + end < prefill_end + and not checkpoint_covers_prompt_tail + and ( + not use_internal_checkpoint + # DFlash reserves draft input slots, so its target-token budget + # is not necessarily block aligned. Keep intermediate target + # chunks aligned; only the prompt tail uses the checkpoint. + or self.use_eagle + ) + ): + max_prefill_tokens = self.max_num_scheduled_tokens + long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold + if long_prefill_threshold > 0: + max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) + aligned_end = end // block_size * block_size + if aligned_end > start or block_size <= max_prefill_tokens: + end = aligned_end + + next_block_boundary = (start // block_size + 1) * block_size + tail_boundary = ( + request.num_prompt_tokens // self.hash_block_size * self.hash_block_size + if self.mamba_partial_cache_hit + else 0 + ) + boundary_proposer = getattr(self, "_recurrent_publication_boundaries", None) + publication_boundaries = ( + boundary_proposer(request) if boundary_proposer is not None else () + ) + publication_boundary = min( + (boundary for boundary in publication_boundaries if start < boundary < end), + default=0, + ) + stops = ( + # Same invariant: a chunk starting mid-block stops at the boundary + # rather than running past it. + next_block_boundary if start % block_size != 0 else 0, + # Never run past the last cacheable block boundary mid-chunk. + last_cache_position, + # Fine-grained hits: the prompt's partial-tail entry can only be + # registered by a chunk ending exactly at its last hash boundary. + tail_boundary + if last_cache_position < tail_boundary < request.num_prompt_tokens + else 0, + # Connector publication uses its own wire-identity boundary while + # native prefix-cache hashes retain hash_block_size. + publication_boundary, + # Marconi shared-prefix junction, block-floored (a sub-block + # junction's state is not separately cacheable): cache its state + # so sibling requests sharing the prefix can reuse it. + start + (request.shared_prefix_boundary - start) // block_size * block_size + if start < request.shared_prefix_boundary < end + else 0, + ) + # Stop at the earliest mandatory position strictly inside the chunk. + end = min((s for s in stops if start < s < end), default=end) + return max(end - start, 0) + + def _get_local_prefix_cache_hit( + self, request: Request + ) -> tuple[KVCacheBlocks, int, int, bool]: + connector = self.connector + if connector is not None and connector.supports_divergent_local_hybrid_hits: + return self.kv_cache_manager.get_computed_blocks_for_connector(request) + + blocks, num_local, shared_prefix_boundary = ( + self.kv_cache_manager.get_computed_blocks(request) + ) + return blocks, num_local, shared_prefix_boundary, False + + def _reserve_prefill_lookahead( + self, + request: Request, + num_computed_tokens: int, + num_new_tokens: int, + ) -> int: + """Never end a prefill chunk within num_prefill_lookahead of the + prefill end. + + At a chunked-prefill boundary, the multi-module MTP drafter consumes + the next num_prefill_lookahead known prefill tokens as draft inputs. A + boundary closer to the end than that would make it fall back to + sampled drafts, permanently polluting the trailing modules' KV caches. + Either finish the prefill or leave at least num_prefill_lookahead for + the next chunk. No-op for eagle-family drafters (lookahead 1). + """ + remaining = request.num_tokens - num_computed_tokens - num_new_tokens + if 0 < remaining < self.num_prefill_lookahead: + num_new_tokens -= self.num_prefill_lookahead - remaining + return max(num_new_tokens, 0) + + def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: + self.current_step += 1 + # NOTE(woosuk) on the scheduling algorithm: + # There's no "decoding phase" nor "prefill phase" in the scheduler. + # Each request just has the num_computed_tokens and + # num_tokens_with_spec. num_tokens_with_spec = + # len(prompt_token_ids) + len(output_token_ids) + len(spec_token_ids). + # At each step, the scheduler tries to assign tokens to the requests + # so that each request's num_computed_tokens can catch up its + # num_tokens_with_spec. This is general enough to cover + # chunked prefills, prefix caching, speculative decoding, + # and the "jump decoding" optimization in the future. + + scheduled_new_reqs: list[Request] = [] + scheduled_resumed_reqs: list[Request] = [] + scheduled_running_reqs: list[Request] = [] + preempted_reqs: list[Request] = [] + + req_to_new_blocks: dict[str, KVCacheBlocks] = {} + recurrent_prefill_checkpoint_plans = {} + num_scheduled_tokens: dict[str, int] = {} + token_budget = self.max_num_scheduled_tokens + spec = self.vllm_config.speculative_config + draft_slots = spec.max_num_new_slots_for_drafting if spec is not None else 0 + input_budget = self.scheduler_config.max_num_batched_tokens + if self._pause_state == PauseState.PAUSED_ALL: + # Do not schedule any requests when paused. + token_budget = 0 + + # Encoder-related. + scheduled_encoder_inputs: dict[str, list[int]] = {} + encoder_compute_budget = self.max_num_encoder_input_tokens + # Spec decode-related. + scheduled_spec_decode_tokens: dict[str, list[int]] = {} + # Whether the running batch contains any prefill requests. + prefill_scheduled = False + + # For logging. + scheduled_timestamp = time.monotonic() + + self.kv_cache_manager.new_step_starts() + # Shared-prefix leases are optional, bounded cache work. Expiring + # them is O(2) and releases only the scheduler-owned pin; requests + # already attached to those blocks keep their ordinary references. + self.kv_cache_manager.expire_shared_prefix_leases() + + # On a throttled cadence step, defer prefill compute only when this + # decision can schedule decode work. Pipeline-parallel async requests + # can be temporarily ineligible; admitting prefill in that case avoids + # an empty model-executor step. + has_eligible_decode = any( + not request.is_prefill_chunk + and self.current_step >= request.next_decode_eligible_step + for request in self.running + ) + defer_prefills = ( + throttle_prefills and not self.prefill_capacity_bound + ) and has_eligible_decode + + # First, schedule the RUNNING requests. + req_index = 0 + while req_index < len(self.running) and token_budget > 0: + request = self.running[req_index] + if input_budget <= draft_slots: + break + + if ( + request.num_output_placeholders > 0 + # This is (num_computed_tokens + 1) - (num_output_placeholders - 1). + # Since output placeholders are also included in the computed tokens + # count, we subtract (num_output_placeholders - 1) to remove any draft + # tokens, so that we can be sure no further steps are needed even if + # they are all rejected. + and request.num_computed_tokens + 2 - request.num_output_placeholders + >= request.num_prompt_tokens + request.max_tokens + ): + # Async scheduling: Avoid scheduling an extra step when we are sure that + # the previous step has reached request.max_tokens. We don't schedule + # partial draft tokens since this prevents uniform decode optimizations. + req_index += 1 + continue + + if self.current_step < request.next_decode_eligible_step: + # V2+PP+async: enforce `pp_size` steps between same-req decodes + # to match worker-side sampled-tokens broadcast slot ring cadence. + req_index += 1 + continue + + if defer_prefills and request.is_prefill_chunk: + # Defer this in-progress chunk to a cadence-aligned step; + # decodes still run to fill this step. + req_index += 1 + continue + + num_new_tokens = ( + request.num_tokens_with_spec + + request.num_output_placeholders + - request.num_computed_tokens + ) + if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: + num_new_tokens = self.scheduler_config.long_prefill_token_threshold + num_new_tokens = min( + num_new_tokens, token_budget, input_budget - draft_slots + ) + + # Make sure the input position does not exceed the max model len. + # This is necessary when using spec decoding. + num_new_tokens = min( + num_new_tokens, + self.max_model_len + - request.num_computed_tokens + - self.num_sampled_tokens_per_step, + ) + + # Apply Mamba alignment before encoder caps. + if self.need_mamba_block_aligned_split: + num_new_tokens = self._mamba_block_aligned_split( + request, num_new_tokens + ) + + # Schedule encoder inputs. + encoder_inputs_to_schedule = None + external_load_encoder_input: list[int] = [] + new_encoder_compute_budget = encoder_compute_budget + if request.has_encoder_inputs: + ( + encoder_inputs_to_schedule, + num_new_tokens, + new_encoder_compute_budget, + external_load_encoder_input, + ) = self._try_schedule_encoder_inputs( + request, + request.num_computed_tokens, + num_new_tokens, + encoder_compute_budget, + shift_computed_tokens=self.num_prefill_lookahead, + ) + + # Multi-module MTP: avoid ending a prefill chunk within + # num_prefill_lookahead of the prefill end. + num_new_tokens = self._reserve_prefill_lookahead( + request, request.num_computed_tokens, num_new_tokens + ) + + if num_new_tokens == 0: + # The request cannot be scheduled because one of the following + # reasons: + # 1. No new tokens to schedule. This may happen when + # (1) PP>1 and we have already scheduled all prompt tokens + # but they are not finished yet. + # (2) Async scheduling and the request has reached to either + # its max_total_tokens or max_model_len. + # 2. The encoder budget is exhausted. + # 3. The encoder cache is exhausted. + # 4. Insufficient budget for a block-aligned chunk in hybrid + # models with mamba cache mode \"align\". + # 5. Insufficient budget to keep a multi-module MTP prefill + # chunk out of the prefill-lookahead window. + # NOTE(woosuk): Here, by doing `continue` instead of `break`, + # we do not strictly follow the FCFS scheduling policy and + # allow the lower-priority requests to be scheduled. + req_index += 1 + continue + + # Schedule newly needed KV blocks for the request. + with record_function_or_nullcontext("schedule: allocate_slots"): + recurrent_publication_boundary = ( + self._recurrent_publication_boundary_at( + request, + request.num_computed_tokens + num_new_tokens, + ) + ) + checkpoint_plan = self._recurrent_checkpoint_plan( + request, request.num_computed_tokens, + request.num_computed_tokens + num_new_tokens) + while True: + new_blocks = self.kv_cache_manager.allocate_slots( + request, + num_new_tokens, + num_lookahead_tokens=self.num_lookahead_tokens, + recurrent_publication_boundary=(recurrent_publication_boundary), + recurrent_prefill_checkpoint_plan=checkpoint_plan, + recurrent_checkpoint_publications=( + self._recurrent_publication_boundaries(request) if checkpoint_plan else ()), + ) + + if new_blocks is not None: + if checkpoint_plan is not None: + recurrent_prefill_checkpoint_plans[request.request_id] = checkpoint_plan + # The request can be scheduled. + break + + # The request cannot be scheduled. + # Preempt the lowest-priority request. + if self.policy == SchedulingPolicy.PRIORITY: + preempted_req = max( + self.running, + key=lambda r: (r.priority, r.arrival_time), + ) + # Record the index of the preemption victim to + # maintain accurate loop state. + victim_index = self.running.index(preempted_req) + del self.running[victim_index] + # Decrement the loop cursor if the removed request + # preceded the current iteration, preventing the + # silent omission of the subsequent request. + if victim_index < req_index: + req_index -= 1 + + if preempted_req in scheduled_running_reqs: + preempted_req_id = preempted_req.request_id + scheduled_running_reqs.remove(preempted_req) + restored = num_scheduled_tokens.pop(preempted_req_id) + token_budget += restored + input_budget += restored + draft_slots + req_to_new_blocks.pop(preempted_req_id) + scheduled_spec_decode_tokens.pop(preempted_req_id, None) + preempted_encoder_inputs = scheduled_encoder_inputs.pop( + preempted_req_id, None + ) + if preempted_encoder_inputs: + # Restore encoder compute budget if the preempted + # request had encoder inputs scheduled in this step. + num_embeds_to_restore = sum( + preempted_req.get_num_encoder_embeds(i) + for i in preempted_encoder_inputs + ) + encoder_compute_budget += num_embeds_to_restore + else: + preempted_req = self.running.pop() + + self._preempt_request( + preempted_req, + scheduled_timestamp, + drop_stale_output=self.requires_kv_delivery, + ) + preempted_reqs.append(preempted_req) + if preempted_req == request: + # No more request to preempt. Cannot schedule this request. + break + + if new_blocks is None: + # Cannot schedule this request. + break + + # Schedule the request. + scheduled_running_reqs.append(request) + prefill_scheduled |= request.is_prefill_chunk + request_id = request.request_id + req_to_new_blocks[request_id] = new_blocks + num_scheduled_tokens[request_id] = num_new_tokens + token_budget -= num_new_tokens + input_budget -= num_new_tokens + draft_slots + req_index += 1 + + # Speculative decode related. + if request.spec_token_ids: + num_scheduled_spec_tokens = ( + num_new_tokens + + request.num_computed_tokens + - request.num_tokens + - request.num_output_placeholders + ) + if num_scheduled_spec_tokens > 0: + spec_token_ids = request.spec_token_ids + if len(spec_token_ids) > num_scheduled_spec_tokens: + spec_token_ids = spec_token_ids[:num_scheduled_spec_tokens] + scheduled_spec_decode_tokens[request.request_id] = spec_token_ids + + # New spec tokens will be set in `update_draft_token_ids` before the + # next step when applicable. + request.spec_token_ids = [] + + # Encoder-related. + if encoder_inputs_to_schedule: + scheduled_encoder_inputs[request_id] = encoder_inputs_to_schedule + # Allocate the encoder cache. + for i in encoder_inputs_to_schedule: + self.encoder_cache_manager.allocate(request, i) + if self.ec_connector is not None: + self.ec_connector.update_state_after_alloc(request, i) + encoder_compute_budget = new_encoder_compute_budget + if external_load_encoder_input: + for i in external_load_encoder_input: + self.encoder_cache_manager.allocate(request, i) + if self.ec_connector is not None: + self.ec_connector.update_state_after_alloc(request, i) + + # Record the LoRAs in scheduled_running_reqs + scheduled_loras: set[int] = set() + if self.lora_config: + scheduled_loras = set( + req.lora_request.lora_int_id + for req in scheduled_running_reqs + if req.lora_request and req.lora_request.lora_int_id > 0 + ) + assert len(scheduled_loras) <= self.lora_config.max_loras + + # Next, schedule the WAITING requests. + if not preempted_reqs and self._pause_state == PauseState.UNPAUSED: + step_skipped_waiting = create_request_queue(self.policy) + + while (self.waiting or self.skipped_waiting) and token_budget > 0: + if input_budget <= draft_slots: + break + # Paused streaming sessions (WAITING_FOR_STREAMING_REQ) are not + # in `running` but still hold a model-runner request slot. + num_running = len(self.running) + self.num_waiting_for_streaming_input + if num_running >= self.max_num_running_reqs: + break + + request_queue = self._select_waiting_queue_for_scheduling() + assert request_queue is not None + + request = request_queue.peek_request() + request_id = request.request_id + + # try to promote blocked statuses while traversing skipped queue. + if self._is_blocked_waiting_status( + request.status + ) and not self._try_promote_blocked_waiting_request(request): + if request.status == RequestStatus.WAITING_FOR_REMOTE_KVS: + logger.debug( + "%s is still in WAITING_FOR_REMOTE_KVS state.", + request_id, + ) + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + + if ( + request.num_stale_output_tokens > 0 + and not request.drop_stale_output + ): + # Deliverable stale output still in flight: resuming now + # could resample a position that output later delivers. + # It drains within the pipeline depth. + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + + # Check that adding the request still respects the max_loras + # constraint. + if ( + self.lora_config + and request.lora_request + and ( + len(scheduled_loras) == self.lora_config.max_loras + and request.lora_request.lora_int_id not in scheduled_loras + ) + ): + # Scheduling would exceed max_loras, skip. + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + + num_external_computed_tokens = 0 + load_kv_async = False + connector_prefix_cache_queries, connector_prefix_cache_hits = 0, 0 + did_prefix_cache_lookup = False + local_lease_alternative = None + + # A connector may offer an exact block-table lease after a + # prior async restore completed on every worker. Attach it + # before hash lookup: hybrid Mamba/EAGLE tables are not always + # rediscoverable as an ordinary local prefix hit even though + # the verified blocks are still resident. + if request.num_computed_tokens == 0 and self.connector is not None: + get_lease = getattr( + self.connector, "get_shared_prefix_lease_candidate", None + ) + candidate = get_lease(request) if get_lease is not None else None + if candidate is not None: + lease_key, lease_tokens = candidate + if 0 < lease_tokens <= request.num_tokens: + # This lookup reconciles every KV group and applies + # speculative backoff before any blocks are adopted. + alternative = self.kv_cache_manager.get_computed_blocks(request) + if alternative[1] > min(lease_tokens, request.num_tokens - 1): + local_lease_alternative = (*alternative, False) + candidate = None + # Release this request's follower binding only; + # the verified lease and its other users remain. + rejected = getattr( + self.connector, "shared_prefix_lease_rejected", None + ) + if rejected is not None: + rejected(request_id, lease_key) + if candidate is not None: + lease_key, lease_tokens = candidate + attached_tokens = 0 + if 0 < lease_tokens <= request.num_tokens: + attached_tokens = ( + self.kv_cache_manager.attach_shared_prefix_lease( + lease_key, request_id + ) + ) + if attached_tokens: + # Preserve vLLM's full-hit rule: the final prompt + # token is recomputed to produce sampling logits. + request.num_computed_tokens = min( + attached_tokens, request.num_tokens - 1 + ) + # Lease attachment skips hash lookup but still reuses + # resident GPU state. Include it in prefill/API totals + # without reporting a fresh external KV transfer. + if request.prefill_stats and request.num_preemptions <= 0: + request.prefill_stats.set( + num_prompt_tokens=request.num_prompt_tokens, + num_local_cached_tokens=request.num_computed_tokens, + num_external_cached_tokens=0, + ) + attached = getattr( + self.connector, + "shared_prefix_lease_attached", + None, + ) + if attached is not None: + attached(request_id, lease_key) + else: + rejected = getattr( + self.connector, + "shared_prefix_lease_rejected", + None, + ) + if rejected is not None: + rejected(request_id, lease_key) + + # Get already-cached tokens. + if request.num_computed_tokens == 0: + did_prefix_cache_lookup = True + ( + new_computed_blocks, + num_new_local_computed_tokens, + request.shared_prefix_boundary, + hit_diverged, + ) = ( + local_lease_alternative + if local_lease_alternative is not None + else self._get_local_prefix_cache_hit(request) + ) + + # Get externally-cached tokens if using a KVConnector. + if self.connector is not None: + # Present a block-aligned local hit to the connector so + # a strictly longer remote hit can supersede a local + # sub-block tail without racing its copy-on-write. + partial_tail = num_new_local_computed_tokens % self.block_size + block_aligned_local = ( + num_new_local_computed_tokens - partial_tail + ) + ext_tokens, load_kv_async = ( + self.connector.get_num_new_matched_tokens( + request, block_aligned_local + ) + ) + + if ext_tokens is None: + # The request cannot be scheduled because + # the KVConnector couldn't determine + # the number of matched tokens. + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + + if partial_tail and ext_tokens > partial_tail: + # Remote strictly exceeds the full local hit: drop the + # sub-block tail so no CoW is needed, and let the load + # cover it. Trim the partial block out of the local + # computed blocks so it is not adopted from the cache. + new_computed_blocks = ( + self.kv_cache_manager.truncate_computed_blocks( + new_computed_blocks, block_aligned_local + ) + ) + num_new_local_computed_tokens = block_aligned_local + num_external_computed_tokens = ext_tokens + elif partial_tail: + # Remote does not exceed the full local hit: keep the + # local sub-block tail and load nothing external. + num_external_computed_tokens = 0 + # Nothing to load remotely -> not an async-load step; + # clearing avoids the `load_kv_async` assert below. + load_kv_async = False + else: + num_external_computed_tokens = ext_tokens + + if hit_diverged and num_external_computed_tokens == 0: + # No external tokens back the deeper local hit, so its + # resume boundary would have no valid Mamba state. + # Reconcile to the boundary every group agrees on. + ( + new_computed_blocks, + num_new_local_computed_tokens, + request.shared_prefix_boundary, + ) = self.kv_cache_manager.get_computed_blocks(request) + + connector_prefix_cache_queries = ( + request.num_tokens - num_new_local_computed_tokens + ) + connector_prefix_cache_hits = num_external_computed_tokens + + # Total computed tokens (local + external). + num_computed_tokens = ( + num_new_local_computed_tokens + num_external_computed_tokens + ) + assert num_computed_tokens <= request.num_tokens + + # Skip request with pending mm encoding prefetches + if ( + self.ec_connector is not None + and request.mm_features + and not self.ec_connector.ensure_cache_available( + request, num_computed_tokens + ) + ): + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + + # Track first scheduled prefill, not post-preemption repeat prefills + if request.prefill_stats and request.num_preemptions <= 0: + assert num_computed_tokens <= request.num_prompt_tokens + request.prefill_stats.set( + num_prompt_tokens=request.num_prompt_tokens, + num_local_cached_tokens=num_new_local_computed_tokens, + num_external_cached_tokens=num_external_computed_tokens, + ) + else: + # KVTransfer: WAITING reqs have num_computed_tokens > 0 + # after async KV recvs are completed. + new_computed_blocks = self.kv_cache_manager.empty_kv_cache_blocks + num_new_local_computed_tokens = 0 + num_computed_tokens = request.num_computed_tokens + + encoder_inputs_to_schedule = None + external_load_encoder_input = [] + new_encoder_compute_budget = encoder_compute_budget + pad_spec_decode = False + + if load_kv_async: + # KVTransfer: loading remote KV, do not allocate for new work. + assert num_external_computed_tokens > 0 + num_new_tokens = 0 + elif defer_prefills and num_computed_tokens < request.num_tokens - 1: + # DP prefill balancing: defer this step's local prefill + # compute to a cadence-aligned step. + break + else: + request_token_budget = min(token_budget, input_budget - draft_slots) + # Number of tokens to be scheduled. + # We use `request.num_tokens` instead of + # `request.num_prompt_tokens` to consider the resumed + # requests, which have output tokens. + num_new_tokens = request.num_tokens - num_computed_tokens + + # Pad new decode requests to uniform spec decoding size to + # preserve full cudagraph for this step. + # Not for diffusion where draft tokens can't be padded. + if ( + ( + self.num_spec_tokens > 0 + and self.dynamic_sd_lookup is None + and self.acceptance_length_controller is None + ) + and self.num_sampled_tokens_per_step > 0 + and num_new_tokens == 1 + and (scheduled_running_reqs and not prefill_scheduled) + ): + num_new_tokens = 1 + self.num_spec_tokens + if ( + num_new_tokens > request_token_budget + or num_computed_tokens + num_new_tokens > self.max_model_len + ): + # Prefer to not schedule than schedule un-padded here. + break + pad_spec_decode = True + + threshold = self.scheduler_config.long_prefill_token_threshold + if 0 < threshold < num_new_tokens: + num_new_tokens = threshold + + # chunked prefill has to be enabled explicitly to allow + # pooling requests to be chunked + if ( + not self.scheduler_config.enable_chunked_prefill + and num_new_tokens > request_token_budget + ): + # If chunked_prefill is disabled, + # we can stop the scheduling here. + break + + num_new_tokens = min(num_new_tokens, request_token_budget) + assert num_new_tokens > 0 + + # Apply Mamba alignment before encoder caps. + if self.need_mamba_block_aligned_split: + num_new_tokens = self._mamba_block_aligned_split( + request, + num_new_tokens, + num_new_local_computed_tokens, + num_external_computed_tokens, + ) + if num_new_tokens == 0: + break + + # Schedule encoder inputs. + if request.has_encoder_inputs: + ( + encoder_inputs_to_schedule, + num_new_tokens, + new_encoder_compute_budget, + external_load_encoder_input, + ) = self._try_schedule_encoder_inputs( + request, + num_computed_tokens, + num_new_tokens, + encoder_compute_budget, + shift_computed_tokens=self.num_prefill_lookahead, + ) + + # Multi-module MTP: avoid ending a prefill chunk within + # num_prefill_lookahead of the prefill end. + num_new_tokens = self._reserve_prefill_lookahead( + request, num_computed_tokens, num_new_tokens + ) + + if num_new_tokens == 0: + # The request cannot be scheduled. + break + + # During async KV load, no forward pass is run yet. + # Allocate speculative lookahead slots later to avoid + # mismatching local and remote block counts. + limit_lookahead_tokens = load_kv_async and self.num_lookahead_tokens > 0 + effective_lookahead_tokens = ( + 0 if limit_lookahead_tokens else self.num_lookahead_tokens + ) + + # Determine if we need to allocate cross-attention blocks. + num_encoder_tokens = 0 + if ( + self.is_encoder_decoder + and request.has_encoder_inputs + and encoder_inputs_to_schedule + ): + num_encoder_tokens = sum( + request.get_num_encoder_embeds(i) + for i in encoder_inputs_to_schedule + ) + + reserved_blocks = 0 + if load_kv_async: + # An async load holds its blocks for the whole transfer with + # no forward progress and isn't preemptible here. Admit it + # only if it fits in (free - other in-flight reservations), to + # avoid deadlock and predictable preemptions. + reserved_blocks = self._inflight_prefill_reserved_blocks() + + checkpoint_plan = (None if load_kv_async else self._recurrent_checkpoint_plan( + request, num_computed_tokens, num_computed_tokens + num_new_tokens)) + new_blocks = self.kv_cache_manager.allocate_slots( + request, + num_new_tokens, + num_new_computed_tokens=num_new_local_computed_tokens, + new_computed_blocks=new_computed_blocks, + num_lookahead_tokens=effective_lookahead_tokens, + num_external_computed_tokens=num_external_computed_tokens, + delay_cache_blocks=load_kv_async, + num_encoder_tokens=num_encoder_tokens, + full_sequence_must_fit=self.scheduler_reserve_full_isl, + reserved_blocks=reserved_blocks, + has_scheduled_reqs=bool(self.running), + recurrent_prefill_checkpoint_plan=checkpoint_plan, + recurrent_checkpoint_publications=( + self._recurrent_publication_boundaries(request) if checkpoint_plan else ()) , + recurrent_publication_boundary=( + self._recurrent_publication_boundary_at( + request, + num_computed_tokens + num_new_tokens, + ) + ), + ) + + if new_blocks is None: + # The request cannot be scheduled. + + # NOTE: we need to untouch the request from the encode cache + # manager + if request.has_encoder_inputs: + self.encoder_cache_manager.free(request) + break + + if checkpoint_plan is not None: + recurrent_prefill_checkpoint_plans[request.request_id] = checkpoint_plan + # KVTransfer: the connector uses this info to determine + # if a load is needed. Note that + # This information is used to determine if a load is + # needed for this request. + if self.connector is not None: + self.connector.update_state_after_alloc( + request, + self.kv_cache_manager.get_blocks(request_id), + num_external_computed_tokens, + ) + if ( + self.connector_prefix_cache_stats is not None + and connector_prefix_cache_queries != 0 + ): + self.connector_prefix_cache_stats.record( + num_tokens=connector_prefix_cache_queries, + num_hits=connector_prefix_cache_hits, + preempted=request.num_preemptions > 0, + ) + + # Record at admission so unscheduled lookups are not counted. + if did_prefix_cache_lookup: + self.kv_cache_manager.record_prefix_cache_stats( + request, num_new_local_computed_tokens + ) + + request = request_queue.pop_request() + if load_kv_async: + # If loading async, allocate memory and put request + # into the WAITING_FOR_REMOTE_KV state. + request.status = RequestStatus.WAITING_FOR_REMOTE_KVS + step_skipped_waiting.prepend_request(request) + # Set num_computed_tokens even though KVs are not yet loaded. + # request.num_computed_tokens will not be used anywhere until + # the request finished the KV transfer. + # + # If a transfer error is reported by the connector, + # request.num_computed_tokens will be re-set accordingly in + # _update_requests_with_invalid_blocks. + # + # When the transfer is finished, either successfully or not, + # request.num_computed_tokens will correctly reflect the number + # of computed tokens. + # _update_waiting_for_remote_kv will then cache + # only the successfully loaded tokens. + request.num_computed_tokens = num_computed_tokens + self._inflight_prefills.add(request) + if self.needs_kv_cache_zeroing: + # Skip zeroing of the blocks the async load will + # overwrite; the zeroing could race the write. + self._skip_zero_block_ids.update( + self.kv_cache_manager.get_zeroing_block_ids_in_range( + request.request_id, + num_new_local_computed_tokens, + num_computed_tokens, + ) + ) + continue + + self.running.append(request) + if self.log_stats: + request.record_event( + EngineCoreEventType.SCHEDULED, scheduled_timestamp + ) + if request.status == RequestStatus.WAITING: + scheduled_new_reqs.append(request) + elif request.status == RequestStatus.PREEMPTED: + scheduled_resumed_reqs.append(request) + else: + raise RuntimeError(f"Invalid request status: {request.status}") + + if self.lora_config and request.lora_request: + scheduled_loras.add(request.lora_request.lora_int_id) + req_to_new_blocks[request_id] = self.kv_cache_manager.get_blocks( + request_id + ) + num_scheduled_tokens[request_id] = num_new_tokens + token_budget -= num_new_tokens + input_budget -= num_new_tokens + draft_slots + request.status = RequestStatus.RUNNING + request.num_computed_tokens = num_computed_tokens + if pad_spec_decode: + scheduled_spec_decode_tokens[request_id] = [ + -1 + ] * self.num_spec_tokens + # Only track requests that will still be prefilling after this chunk. + if num_computed_tokens + num_new_tokens < request.num_tokens: + self._inflight_prefills.add(request) + # Encoder-related. + if encoder_inputs_to_schedule: + scheduled_encoder_inputs[request_id] = encoder_inputs_to_schedule + # Allocate the encoder cache. + for i in encoder_inputs_to_schedule: + self.encoder_cache_manager.allocate(request, i) + if self.ec_connector is not None: + self.ec_connector.update_state_after_alloc(request, i) + encoder_compute_budget = new_encoder_compute_budget + # Allocate for external load encoder cache + if external_load_encoder_input: + for i in external_load_encoder_input: + self.encoder_cache_manager.allocate(request, i) + if self.ec_connector is not None: + self.ec_connector.update_state_after_alloc(request, i) + + # re-queue requests skipped in this pass ahead of older skipped items. + if step_skipped_waiting: + self.skipped_waiting.prepend_requests(step_skipped_waiting) + + # DP prefill balancing: on a step that admitted prefills (release), + # record whether it was capacity-bound. + if not defer_prefills: + self.prefill_capacity_bound = bool(self.waiting) + + # Check if the scheduling constraints are satisfied. + total_num_scheduled_tokens = sum(num_scheduled_tokens.values()) + assert total_num_scheduled_tokens <= self.max_num_scheduled_tokens + + assert token_budget >= 0 + assert input_budget >= 0 + assert len(self.running) <= self.max_num_running_reqs + # Since some requests in the RUNNING queue may not be scheduled in + # this step, the total number of scheduled requests can be smaller than + # len(self.running). + assert len(scheduled_new_reqs) + len(scheduled_resumed_reqs) + len( + scheduled_running_reqs + ) <= len(self.running) + + # Get the longest common prefix among all requests in the running queue. + # This can be potentially used for cascade attention. + num_common_prefix_blocks = [0] * len(self.kv_cache_config.kv_cache_groups) + with record_function_or_nullcontext("schedule: get_num_common_prefix_blocks"): + if self.running: + any_request_id = self.running[0].request_id + num_common_prefix_blocks = ( + self.kv_cache_manager.get_num_common_prefix_blocks(any_request_id) + ) + + # Construct the scheduler output. + if self.use_v2_model_runner: + scheduled_new_reqs.extend(scheduled_resumed_reqs) + scheduled_resumed_reqs.clear() + new_reqs_data = [ + NewRequestData.from_request( + req, + req_to_new_blocks[req.request_id].get_block_ids(), + req._all_token_ids, + uses_mrope=self.model_uses_mrope, + uses_xdrope=self.model_uses_xdrope, + ) + for req in scheduled_new_reqs + ] + else: + new_reqs_data = [ + NewRequestData.from_request( + req, + req_to_new_blocks[req.request_id].get_block_ids(), + uses_mrope=self.model_uses_mrope, + uses_xdrope=self.model_uses_xdrope, + ) + for req in scheduled_new_reqs + ] + + with record_function_or_nullcontext("schedule: make_cached_request_data"): + cached_reqs_data = self._make_cached_request_data( + scheduled_running_reqs, + scheduled_resumed_reqs, + num_scheduled_tokens, + scheduled_spec_decode_tokens, + req_to_new_blocks, + ) + + # Record the request ids that were scheduled in this step (MRV1-only). + if not self.use_v2_model_runner: + self.prev_step_scheduled_req_ids.clear() + self.prev_step_scheduled_req_ids.update(num_scheduled_tokens.keys()) + + # Producer partial-tail hand-off for external KV connectors. Drained + # before the CoW retentions are released below, so the pin lands while + # the cow block still holds a retention ref. Without a producer-side + # connector nothing consumes the hand-off, so skip the drain (and its + # pin); the manager drops stale entries when the request's blocks are + # popped for free. + pending_partial_tail_offloads = None + pending_recurrent_boundary_blocks = None + if ( + self.connector is not None + and self.vllm_config.kv_transfer_config is not None + and self.vllm_config.kv_transfer_config.is_kv_producer + ): + pending_partial_tail_offloads = ( + self.kv_cache_manager.take_partial_tail_offloads() or None + ) + if bool( + getattr(self.connector, "supports_recurrent_boundary_blocks", False) + ): + # Opting in requires the connector's worker snapshot to finish + # before request cleanup releases the boundary pins. + pending_recurrent_boundary_blocks = ( + self.kv_cache_manager.take_recurrent_boundary_blocks( + pending_partial_tail_offloads + ) + or None + ) + else: + self.kv_cache_manager.discard_aligned_recurrent_boundaries() + + kv_cache_block_copies, cow_retained_blocks = ( + self.kv_cache_manager.take_kv_cache_block_copies() + ) + if kv_cache_block_copies: + # The copies run with this step's execution; the first non-empty + # step at or after it gets seq `sched_step_seq + 1` (0-token steps + # do not advance the seq), and its completion implies the copies + # have run. + self._free_cow_retained_blocks(cow_retained_blocks, self.sched_step_seq + 1) + pending_kv_cache_block_copies = kv_cache_block_copies or None + + # Dynamic speculative decoding: compute optimal K + num_spec_tokens_to_schedule = self.num_spec_tokens + batch_size = len(num_scheduled_tokens) + if self.acceptance_length_controller is not None and batch_size > 0: + num_spec_tokens_to_schedule = ( + self.acceptance_length_controller.num_spec_tokens_for_batch_size( + batch_size + ) + ) + elif self.dynamic_sd_lookup is not None and batch_size > 0: + num_spec_tokens_to_schedule = self.dynamic_sd_lookup[batch_size] + + scheduled_encoder_input_stats = None + if ( + self.log_stats + and self.observability_config.enable_logging_iteration_details + ): + scheduled_encoder_input_stats = self._make_scheduled_encoder_input_stats( + scheduled_encoder_inputs + ) + + scheduler_output = SchedulerOutput( + scheduled_new_reqs=new_reqs_data, + scheduled_cached_reqs=cached_reqs_data, + num_scheduled_tokens=num_scheduled_tokens, + total_num_scheduled_tokens=total_num_scheduled_tokens, + scheduled_spec_decode_tokens=scheduled_spec_decode_tokens, + scheduled_encoder_inputs=scheduled_encoder_inputs, + scheduled_encoder_input_stats=scheduled_encoder_input_stats, + num_common_prefix_blocks=num_common_prefix_blocks, + preempted_req_ids=self.reset_preempted_req_ids, + # finished_req_ids is an existing state in the scheduler, + # instead of being newly scheduled in this step. + # It contains the request IDs that are finished in between + # the previous and the current steps. + finished_req_ids=self.finished_req_ids, + free_encoder_mm_hashes=self.encoder_cache_manager.get_freed_mm_hashes(), + new_block_ids_to_zero=self._get_new_block_ids_to_zero(), + kv_cache_block_copies=pending_kv_cache_block_copies, + partial_tail_offloads=pending_partial_tail_offloads, + recurrent_boundary_blocks=pending_recurrent_boundary_blocks, + recurrent_prefill_checkpoint_plans=( + {req_id: plan for req_id, plan in recurrent_prefill_checkpoint_plans.items() + if req_id in num_scheduled_tokens} or None), + num_spec_tokens_to_schedule=num_spec_tokens_to_schedule, + ec_manager_metadata=self.encoder_cache_manager.get_manager_metadata(), + ) + + # NOTE(Kuntai): this function is designed for multiple purposes: + # 1. Plan the KV cache store + # 2. Wrap up all the KV cache load / save ops into an opaque object + # 3. Clear the internal states of the connector + if self.connector is not None: + meta = self._build_kv_connector_meta(self.connector, scheduler_output) + scheduler_output.kv_connector_metadata = meta + + # Build the connector meta for ECConnector + if self.ec_connector is not None: + ec_meta: ECConnectorMetadata = self.ec_connector.build_connector_meta( + scheduler_output + ) + scheduler_output.ec_connector_metadata = ec_meta + + # Advance the fence only for non-empty steps (those that actually + # write KV and have their output processed later in update_from_output). + if total_num_scheduled_tokens > 0: + self.sched_step_seq += 1 + + with record_function_or_nullcontext("schedule: update_after_schedule"): + self._update_after_schedule(scheduler_output) + return scheduler_output + + def _build_kv_connector_meta( + self, connector: KVConnectorBase_V1, scheduler_output: SchedulerOutput + ) -> KVConnectorMetadata: + return connector.build_connector_meta(scheduler_output) + + def _get_new_block_ids_to_zero(self) -> list[int] | None: + # Drain new attention block ids every step so the manager-side list + # does not grow unbounded; only kv-cache zeroing consumes them. + new_block_ids_to_zero = self.kv_cache_manager.take_new_block_ids() + if not self.needs_kv_cache_zeroing: + return None + + if self._skip_zero_block_ids: + skip = self._skip_zero_block_ids + new_block_ids_to_zero = [b for b in new_block_ids_to_zero if b not in skip] + skip.clear() + + return new_block_ids_to_zero or None + + def _preempt_request( + self, request: Request, timestamp: float, drop_stale_output: bool = False + ) -> None: + """Preempt a request and put it back to the waiting queue. + + NOTE: The request should be popped from the running queue outside of this + method. + + drop_stale_output: drop (rather than deliver) any in-flight output; used + by reset_prefix_cache, whose same-step resume would otherwise deliver + tokens out of order, and for connectors with a pending KV hand-off, + which the preemption's block free would leave without valid KV. + """ + assert request.status == RequestStatus.RUNNING, ( + "Only running requests can be preempted" + ) + self._free_request_blocks(request) + self.encoder_cache_manager.free(request) + self._inflight_prefills.discard(request) + request.status = RequestStatus.PREEMPTED + request.num_computed_tokens = 0 + if request.spec_token_ids: + request.spec_token_ids = [] + # Async scheduling: mark all in-flight output as stale. Its tokens are + # still delivered on return (dropping them would perturb spec-decode + # acceptance) but must not mutate the reset counters; each step drains + # its share in update_from_output. num_in_flight_tokens already + # includes any undrained stale share, so assign rather than accumulate. + # An undrained drop-mode share stays dropped: its positions have + # already been resampled. + request.drop_stale_output = drop_stale_output or ( + request.drop_stale_output and request.num_stale_output_tokens > 0 + ) + request.num_stale_output_tokens = request.num_in_flight_tokens + request.num_output_placeholders = 0 + request.num_preemptions += 1 + if self.log_stats: + request.record_event(EngineCoreEventType.PREEMPTED, timestamp) + + # Put the request back to the waiting queue. + self.waiting.prepend_request(request) + self.reset_preempted_req_ids.add(request.request_id) + + def _update_after_schedule(self, scheduler_output: SchedulerOutput) -> None: + # Advance the number of computed tokens for the request AFTER + # the request is scheduled. + # 1. The scheduler_output of the current step has to include the + # original number of scheduled tokens to determine input IDs. + # 2. Advance the number of computed tokens here allowing us to + # schedule the prefill request again immediately in the next + # scheduling step. + # 3. If some tokens (e.g. spec tokens) are rejected later, the number of + # computed tokens will be adjusted in update_from_output. + num_scheduled_tokens = scheduler_output.num_scheduled_tokens + for req_id, num_scheduled_token in num_scheduled_tokens.items(): + request = self.requests[req_id] + request.num_computed_tokens += num_scheduled_token + request.num_in_flight_tokens += num_scheduled_token + if self.defer_block_free: + # Record the in-flight step, to fence deferred block freeing. + request.last_sched_seq = self.sched_step_seq + request.is_prefill_chunk = request.num_computed_tokens < ( + request.num_tokens + request.num_output_placeholders + ) + scheduler_output.has_structured_output_requests |= ( + request.use_structured_output and not request.is_prefill_chunk + ) + # Drop from the in-flight-prefill set once it's no longer prefilling. + if not request.is_prefill_chunk: + self._inflight_prefills.discard(request) + + # Snapshot block IDs for routed experts before forward starts. + # A concurrent schedule() may preempt requests and free blocks + # before update_from_output runs; the snapshot survives that. + # Use update() to preserve entries from the previous step that + # have not yet been consumed by update_from_output (async + # scheduling may call _update_after_schedule again before the + # prior update_from_output runs). + if self.enable_return_routed_experts: + gid = self.routed_experts_mgr.attn_gid + self._re_block_ids.update( + { + rid: self.kv_cache_manager.get_blocks(rid).get_block_ids()[gid] + for rid in num_scheduled_tokens + } + ) + + # Clear the finished and preempted request IDs. + # NOTE: We shouldn't just clear() here because it will also affect + # the scheduler output. + self.finished_req_ids = set() + self.reset_preempted_req_ids = set() + + def _update_request_as_session( + self, session: Request, update: StreamingUpdate + ) -> None: + """ + Updates the waiting session with the next streaming update. + + Discards the last sampled output token from the prior input chunk. + """ + + # Current streaming input behaviour: Keep only computed output tokens + # (discard final sampled output token). + num_computed_tokens = session.num_computed_tokens + kept_output_tokens = session._all_token_ids[ + session.num_prompt_tokens : num_computed_tokens + ] + del session._all_token_ids[num_computed_tokens:] + session._output_token_ids.clear() + assert session.prompt_token_ids is not None + # Extend prompt with kept output tokens. + session.prompt_token_ids.extend(kept_output_tokens) + + if update.mm_features: + base = session.num_tokens + for mm_feature in update.mm_features: + mm_feature.mm_position = replace( + mm_feature.mm_position, offset=mm_feature.mm_position.offset + base + ) + session.mm_features.extend(update.mm_features) + + session._all_token_ids.extend(update.prompt_token_ids or ()) + session.prompt_token_ids.extend(update.prompt_token_ids or ()) + # Update block hashes for the new tokens. + session.update_block_hashes() + session.num_prompt_tokens = len(session.prompt_token_ids) + session.arrival_time = update.arrival_time + session.sampling_params = update.sampling_params + if session.status == RequestStatus.WAITING_FOR_STREAMING_REQ: + self.num_waiting_for_streaming_input -= 1 + session.status = RequestStatus.WAITING + + if self.log_stats: + session.record_event(EngineCoreEventType.QUEUED) + + def _make_cached_request_data( + self, + running_reqs: list[Request], + resumed_reqs: list[Request], + num_scheduled_tokens: dict[str, int], + spec_decode_tokens: dict[str, list[int]], + req_to_new_blocks: dict[str, KVCacheBlocks], + ) -> CachedRequestData: + req_ids: list[str] = [] + new_token_ids: list[list[int]] = [] + new_block_ids: list[tuple[list[int], ...] | None] = [] + all_token_ids: dict[str, list[int]] = {} + num_computed_tokens: list[int] = [] + num_output_tokens: list[int] = [] + resumed_req_ids = set() + + num_running_reqs = len(running_reqs) + for idx, req in enumerate(itertools.chain(running_reqs, resumed_reqs)): + req_id = req.request_id + req_ids.append(req_id) + # NOTE: In PP+async scheduling, we consume token ids via a direct GPU + # broadcast path (`input_batch.prev_sampled_token_ids`), so we can + # omit this payload. + if self.use_pp and not self.scheduler_config.async_scheduling: + # When using PP, the scheduler sends the sampled tokens back, + # because there's no direct communication between the first- + # stage worker and the last-stage worker. Otherwise, we don't + # need to send the sampled tokens back because the model runner + # will cache them. + num_tokens = num_scheduled_tokens[req_id] - len( + spec_decode_tokens.get(req_id, ()) + ) + token_ids = req.all_token_ids[ + req.num_computed_tokens : req.num_computed_tokens + num_tokens + ] + new_token_ids.append(token_ids) + if idx >= num_running_reqs: + resumed_req_ids.add(req_id) + if not self.use_v2_model_runner: # noqa: SIM102 + if req_id not in self.prev_step_scheduled_req_ids: + all_token_ids[req_id] = req.all_token_ids.copy() + new_block_ids.append( + req_to_new_blocks[req_id].get_block_ids(allow_none=True) + ) + num_computed_tokens.append(req.num_computed_tokens) + num_output_tokens.append( + req.num_output_tokens + req.num_output_placeholders + ) + + return CachedRequestData( + req_ids=req_ids, + resumed_req_ids=resumed_req_ids, + new_token_ids=new_token_ids, + all_token_ids=all_token_ids, + new_block_ids=new_block_ids, + num_computed_tokens=num_computed_tokens, + num_output_tokens=num_output_tokens, + ) + + def _try_schedule_encoder_inputs( + self, + request: Request, + num_computed_tokens: int, + num_new_tokens: int, + encoder_compute_budget: int, + shift_computed_tokens: int = 0, + ) -> tuple[list[int], int, int, list[int]]: + """ + Determine which encoder inputs need to be scheduled in the current step, + and update `num_new_tokens` and encoder token budget accordingly. + + An encoder input will be scheduled if: + - Its output tokens overlap with the range of tokens being computed + in this step, i.e., + [num_computed_tokens, num_computed_tokens + num_new_tokens). + - It is not already computed and stored in the encoder cache. + - It is not exist on remote encoder cache (via ECConnector) + - There is sufficient encoder token budget to process it. + - The encoder cache has space to store it. + + If an encoder input cannot be scheduled due to cache or budget + limitations, the method adjusts `num_new_tokens` to schedule only the + decoder tokens up to just before the unschedulable encoder input. + + Note that num_computed_tokens includes both locally cached + blocks and externally cached blocks (via KVConnector). + """ + if num_new_tokens == 0 or not request.has_encoder_inputs: + return [], num_new_tokens, encoder_compute_budget, [] + encoder_inputs_to_schedule: list[int] = [] + mm_features = request.mm_features + assert mm_features is not None + assert len(mm_features) > 0 + external_load_encoder_input = [] + + # NOTE: since scheduler operates on the request level (possibly with + # multiple encoder inputs per request), we need to create temporary + # trackers for accounting at the encoder input level. + mm_hashes_to_schedule = set() + num_embeds_to_schedule = 0 + + encoder_window_end = ( + num_computed_tokens + num_new_tokens + shift_computed_tokens + ) + lo, hi = get_mm_features_in_window( + mm_features, + start=num_computed_tokens, + end=encoder_window_end, + ) + # For encoder-decoder, all inputs sit at start_pos=0, so lo=0 always. + if self.is_encoder_decoder: + lo = 0 + + for i in range(lo, hi): + mm_feature = mm_features[i] + start_pos = mm_feature.mm_position.offset + num_encoder_tokens = mm_feature.mm_position.length + num_encoder_embeds = mm_feature.mm_position.get_num_embeds() + item_identifier = mm_feature.identifier + + if self.is_encoder_decoder and num_computed_tokens > 0: + assert start_pos == 0, ( + "Encoder input should be processed at the beginning of " + "the sequence when encoder-decoder models are used." + ) + # Encoder input has already been computed + # The calculation here is a bit different. We don't turn encoder + # output into tokens that get processed by the decoder and + # reflected in num_computed_tokens. Instead, start_pos reflects + # the position where we need to ensure we calculate encoder + # inputs. This should always be 0 to ensure we calculate encoder + # inputs before running the decoder. Once we've calculated some + # decoder tokens (num_computed_tokens > 0), then we know we + # already calculated encoder inputs and can skip here. + continue + + if not self.is_encoder_decoder: + # We are not using the encoder cache for encoder-decoder models, + # yet. + if item_identifier in mm_hashes_to_schedule: + # The same encoder input has already been scheduled in the + # current step. + continue + + if self.encoder_cache_manager.check_and_update_cache(request, i): + # The encoder input is already computed and cached from a + # previous step. + continue + + # If no encoder input chunking is allowed, we do not want to + # partially schedule a multimodal item. If the scheduled range would + # only cover part of the mm input, roll back to before the mm item. + if ( + self.scheduler_config.disable_chunked_mm_input + and num_computed_tokens < start_pos + and (num_computed_tokens + num_new_tokens) + < (start_pos + num_encoder_tokens) + ): + # Account for EAGLE shift when rolling back to avoid + # encoder cache miss. This ensures the scheduled range + # stops before start_pos even with the shift. + num_new_tokens = max( + 0, start_pos - (num_computed_tokens + shift_computed_tokens) + ) + break + if not self.encoder_cache_manager.can_allocate( + request, i, encoder_compute_budget, num_embeds_to_schedule + ): + # The encoder cache is full or the encoder budget is exhausted. + # NOTE(woosuk): We assume that the encoder input tokens should + # be processed altogether, as the encoder usually uses + # bidirectional attention. + if num_computed_tokens + shift_computed_tokens < start_pos: + # We only schedule the decoder tokens just before the + # encoder input. + num_new_tokens = start_pos - ( + num_computed_tokens + shift_computed_tokens + ) + else: + # Because of prefix caching, num_computed_tokens is greater + # than start_pos even though its encoder input is not + # available. In this case, we can't schedule any token for + # the request in this step. + num_new_tokens = 0 + break + + # Calculate the number of embeddings to schedule in the current range + # of scheduled encoder placeholder tokens. + start_idx_rel = max(0, num_computed_tokens - start_pos) + end_idx_rel = min(num_encoder_tokens, encoder_window_end - start_pos) + curr_embeds_start, curr_embeds_end = ( + mm_feature.mm_position.get_embeds_indices_in_range( + start_idx_rel, end_idx_rel + ) + ) + # There's no embeddings in the current range of encoder placeholder tokens + # so we can skip the encoder input. + if curr_embeds_end - curr_embeds_start == 0: + continue + + if self.ec_connector is not None and self.ec_connector.has_cache_item( + item_identifier + ): + mm_hashes_to_schedule.add(item_identifier) + external_load_encoder_input.append(i) + num_embeds_to_schedule += num_encoder_embeds + continue + + num_embeds_to_schedule += num_encoder_embeds + encoder_compute_budget -= num_encoder_embeds + mm_hashes_to_schedule.add(item_identifier) + encoder_inputs_to_schedule.append(i) + + return ( + encoder_inputs_to_schedule, + num_new_tokens, + encoder_compute_budget, + external_load_encoder_input, + ) + + def _make_scheduled_encoder_input_stats( + self, scheduled_encoder_inputs: dict[str, list[int]] + ) -> ScheduledEncoderInputStats | None: + stats = ScheduledEncoderInputStats() + + for req_id, input_ids in scheduled_encoder_inputs.items(): + request = self.requests.get(req_id) + if request is None: + continue + + for input_id in input_ids: + mm_feature = request.mm_features[input_id] + stats.num_inputs += 1 + stats.output_tokens += mm_feature.mm_position.get_num_embeds() + + return stats if stats.num_inputs else None + + def get_grammar_bitmask( + self, scheduler_output: SchedulerOutput + ) -> GrammarOutput | None: + # Collect list of scheduled request ids that use structured output. + # The corresponding rows of the bitmask will be in this order. + if not scheduler_output.has_structured_output_requests: + return None + + structured_output_request_ids = [ + req_id + for req_id in scheduler_output.num_scheduled_tokens + if (req := self.requests.get(req_id)) + and (req.use_structured_output and not req.is_prefill_chunk) + ] + if not structured_output_request_ids: + return None + + bitmask = self.structured_output_manager.grammar_bitmask( + self.requests, + structured_output_request_ids, + scheduler_output.scheduled_spec_decode_tokens, + ) + return GrammarOutput(structured_output_request_ids, bitmask) + + def update_from_output( + self, + scheduler_output: SchedulerOutput, + model_runner_output: ModelRunnerOutput, + ) -> dict[int, EngineCoreOutputs]: + sampled_token_ids = model_runner_output.sampled_token_ids + logprobs = model_runner_output.logprobs + prompt_logprobs_dict = model_runner_output.prompt_logprobs_dict + num_scheduled_tokens = scheduler_output.num_scheduled_tokens + pooler_outputs = model_runner_output.pooler_output + num_nans_in_logits = model_runner_output.num_nans_in_logits + kv_connector_output = model_runner_output.kv_connector_output + ec_connector_output = model_runner_output.ec_connector_output + cudagraph_stats = model_runner_output.cudagraph_stats + + # Every GPU write enqueued by this and earlier steps has completed, so it is + # safe to return deferred-free blocks to the pool. + if scheduler_output.total_num_scheduled_tokens > 0: + self.processed_step_seq += 1 + if self.defer_block_free: + self._drain_deferred_frees() + self._finalize_shared_prefix_leases() + + perf_stats: PerfStats | None = None + if self.perf_metrics and self.perf_metrics.is_enabled(): + perf_stats = self.perf_metrics.get_step_perf_stats_per_gpu(scheduler_output) + + outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list) + spec_decoding_stats: SpecDecodingStats | None = None + + failed_kv_load_req_ids = None + if kv_connector_output and kv_connector_output.invalid_block_ids: + # These blocks contain externally computed tokens that failed to + # load. Identify affected requests and adjust their computed token + # count to trigger recomputation of the invalid blocks. + failed_kv_load_req_ids = self._handle_invalid_blocks( + kv_connector_output.invalid_block_ids, + num_scheduled_tokens, + ) + + # Persist per-step routed experts into the scheduler-side slot + # buffer (CPU->CPU fancy-index assign; ~few MB per step). + # MUST precede the per-request routing reads below: stopped + # requests may terminate on tokens generated in this very step, + # whose routing was just D2H'd into model_runner_output. + routing_data = None + routing_offsets: dict[str, int] = {} + if model_runner_output.routed_experts is not None: + re = model_runner_output.routed_experts + self.routed_experts_mgr.store_batch(re.routing_data, re.slot_mapping) + routing_data = re.routing_data.astype( + self.routed_experts_mgr.routed_experts_by_slot.dtype, + copy=False, + ) + # Build offset map using model runner's request order + # (input_batch ordering), NOT scheduler dict order. + offset = 0 + for rid in model_runner_output.req_ids: + routing_offsets[rid] = offset + offset += num_scheduled_tokens[rid] + + # NOTE(woosuk): As len(num_scheduled_tokens) can be up to 1K or more, + # the below loop can be a performance bottleneck. We should do our best + # to avoid expensive operations inside the loop. + stopped_running_reqs: set[Request] = set() + stopped_preempted_reqs: set[Request] = set() + adaptive_num_drafts = 0 + adaptive_num_draft_tokens = 0 + adaptive_num_accepted_tokens = 0 + for req_id, num_tokens_scheduled in num_scheduled_tokens.items(): + assert num_tokens_scheduled > 0 + request = self.requests.get(req_id) + output_is_stale = False + if request is not None: + request.num_in_flight_tokens -= num_tokens_scheduled + # Drain any stale share (see _preempt_request) in lockstep. + if request.num_stale_output_tokens > 0: + output_is_stale = True + request.num_stale_output_tokens -= num_tokens_scheduled + assert request.num_stale_output_tokens >= 0 + if failed_kv_load_req_ids and req_id in failed_kv_load_req_ids: + # skip failed or rescheduled requests from KV load failure + continue + if request is None or request.is_finished(): + # The request is already finished. This can happen if the + # request is aborted while the model is executing it (e.g., + # in pipeline parallelism or in async scheduling). + # NOTE(Kuntai): When delay_free_blocks=True (for async KV + # cache transfer in KV connector), the aborted request will not + # be set to None (in order to finish async KV transfer). + # In this case, we use is_finished() to check. + continue + + # Drop-mode stale output (same-step resume) is discarded entirely. + if output_is_stale and request.drop_stale_output: + continue + + req_index = model_runner_output.req_id_to_index[req_id] + generated_token_ids = ( + sampled_token_ids[req_index] if sampled_token_ids else [] + ) + + scheduled_spec_token_ids = ( + scheduler_output.scheduled_spec_decode_tokens.get(req_id) + ) + if scheduled_spec_token_ids and ( + generated_token_ids or self.num_sampled_tokens_per_step == 0 + ): + num_draft_tokens = len(scheduled_spec_token_ids) + num_sampled = self.num_sampled_tokens_per_step + num_accepted = max(len(generated_token_ids) - num_sampled, 0) + num_rejected = num_draft_tokens - num_accepted + if self.acceptance_length_controller is not None: + adaptive_num_drafts += 1 + adaptive_num_draft_tokens += num_draft_tokens + adaptive_num_accepted_tokens += num_accepted + # Rejections roll back num_computed_tokens (and, under async + # scheduling, num_output_placeholders, which covers the spec + # tokens). A stale rejection count predates the preemption + # rollback and must not apply. + if not output_is_stale: + if request.num_computed_tokens > 0: + request.num_computed_tokens -= num_rejected + if request.num_output_placeholders > 0: + request.num_output_placeholders -= num_rejected + spec_decoding_stats = self.make_spec_decoding_stats( + spec_decoding_stats, + num_draft_tokens=num_draft_tokens, + num_accepted_tokens=num_accepted, + num_invalid_spec_tokens=scheduler_output.num_invalid_spec_tokens, + request_id=req_id, + ) + if request.spec_decode_metrics is not None: + # Exclude grammar-invalidated drafts from the proposed + # count, mirroring make_spec_decoding_stats; the accepted + # bucket (j) is unaffected. + adj_draft_tokens = num_draft_tokens + if scheduler_output.num_invalid_spec_tokens: + adj_draft_tokens -= ( + scheduler_output.num_invalid_spec_tokens.get(req_id, 0) + ) + request.spec_decode_metrics.observe( + num_draft_tokens=adj_draft_tokens, + num_accepted=num_accepted, + detailed=self.spec_decode_metrics_level == "detailed", + ) + + # Free encoder inputs only after the step has actually executed. + if request.has_encoder_inputs: + self._free_encoder_inputs(request) + + stopped = False + new_logprobs = None + new_sampling_mask = None + new_token_ids = generated_token_ids + pooler_output = pooler_outputs[req_index] if pooler_outputs else None + kv_transfer_params = None + ec_transfer_params = None + prefill_stats = None + status_before_stop = request.status + num_output_tokens_before = len(request._output_token_ids) + + # Check for stop and update request status. + if new_token_ids: + new_token_ids, stopped = self._update_request_with_output( + request, new_token_ids, is_stale=output_is_stale + ) + elif request.pooling_params and pooler_output is not None: + # Pooling stops as soon as there is output. + request.status = RequestStatus.FINISHED_STOPPED + stopped = True + elif ( + self.is_mm_encoder_only + and request.num_computed_tokens >= request.num_prompt_tokens + ): + # An encoder instance runs the encoder and publishes the + # embeddings instead of sampling, so it stops as soon as the + # whole prompt is consumed. Encoder inputs are never scheduled + # past a multi-modal item the encoder cache could not admit, so + # a consumed prompt also means every item in it was encoded. + request.status = RequestStatus.FINISHED_STOPPED + stopped = True + + if new_token_ids and self.structured_output_manager.should_advance( + request, new_token_ids=new_token_ids + ): + struct_output_request = request.structured_output_request + assert struct_output_request is not None + grammar = struct_output_request.grammar + assert isinstance(grammar, StructuredOutputGrammar) + # new_token_ids can be a mixed block of reasoning content, then + # the reasoning end marker, then the start of the grammar content. + # Trim the reasoning content so the grammar only sees grammar content. + advance_token_ids = ( + self.structured_output_manager.trim_reasoning_for_advance( + request, new_token_ids + ) + ) + if advance_token_ids and not grammar.accept_tokens( + req_id, advance_token_ids + ): + logger.error( + "Unexpected: grammar rejected tokens %s for request %s. " + "Terminating request.", + advance_token_ids, + req_id, + ) + request.status = RequestStatus.FINISHED_ERROR + request.resumable = False + stopped = True + + routed_experts = None + if ( + self.enable_return_routed_experts + and routing_data is not None + and new_token_ids + ): + req_offset = routing_offsets[req_id] + end = req_offset + num_tokens_scheduled + block_ids = self._re_block_ids.pop(req_id, []) + if num_output_tokens_before == 0: + # Prefill completed: read full prompt routing from + # slot buffer using the block-ID snapshot taken at + # schedule time (immune to async preemption). + if ( + request.sampling_params is not None + and request.sampling_params.routed_experts_prompt_start + is not None + ): + prompt_start = ( + request.sampling_params.routed_experts_prompt_start + ) + assert prompt_start < request.num_prompt_tokens + else: + prompt_start = 0 + routed_experts = self.routed_experts_mgr.get( + block_ids, + request.num_prompt_tokens, + token_start=prompt_start, + ) + else: + if scheduled_spec_token_ids: + # Spec decode: accepted tokens at the START of + # the scheduled range, rejected at the end. + routed_experts = routing_data[ + req_offset : req_offset + len(new_token_ids) + ] + else: + # Normal decode / re-prefill: token(s) at the END. + routed_experts = routing_data[end - len(new_token_ids) : end] + + should_emit_output = bool( + new_token_ids or pooler_output is not None or stopped + ) + if should_emit_output: + prefill_stats = request.take_prefill_stats() + if prefill_stats is not None: + prefill_stats.finalize( + self.kv_cache_manager.estimate_cached_tokens(request) + ) + + finish_reason = None + if stopped: + # Capture finish_reason BEFORE _handle_stopped_request, which may + # reset the status to WAITING for streaming requests that continue. + finish_reason = request.get_finished_reason() + finished = self._handle_stopped_request(request) + if finished: + kv_transfer_params, ec_transfer_params = self._free_request(request) + + if status_before_stop == RequestStatus.RUNNING: + stopped_running_reqs.add(request) + else: + stopped_preempted_reqs.add(request) + + # Extract sample logprobs if needed. + if ( + request.sampling_params is not None + and request.sampling_params.num_logprobs is not None + and logprobs + ): + new_logprobs = logprobs.slice_request(req_index, len(new_token_ids)) + + if self.return_sampling_mask: + sampling_masks = model_runner_output.sampling_masks + if new_token_ids and sampling_masks is not None: + new_sampling_mask = sampling_masks.slice_request( + req_index, len(new_token_ids) + ) + + if num_nans_in_logits is not None and req_id in num_nans_in_logits: + request.num_nans_in_logits = num_nans_in_logits[req_id] + + # Get prompt logprobs for this request. + prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id) + if should_emit_output: + # Add EngineCoreOutput for this Request. + outputs[request.client_index].append( + EngineCoreOutput( + request_id=req_id, + new_token_ids=new_token_ids, + finish_reason=finish_reason, + new_logprobs=new_logprobs, + new_sampling_mask=new_sampling_mask, + new_prompt_logprobs_tensors=prompt_logprobs_tensors, + pooling_output=pooler_output, + stop_reason=request.stop_reason, + events=request.take_events(), + prefill_stats=prefill_stats, + spec_decode_metrics=( + request.spec_decode_metrics + if finish_reason is not None + else None + ), + kv_transfer_params=kv_transfer_params, + ec_transfer_params=ec_transfer_params, + trace_headers=request.trace_headers, + routed_experts=routed_experts, + num_nans_in_logits=request.num_nans_in_logits, + ) + ) + else: + # Invariant: EngineCore returns no partial prefill outputs. + assert not prompt_logprobs_tensors + + batch_size = len(num_scheduled_tokens) + if self.acceptance_length_controller is not None and batch_size > 0: + update = self.acceptance_length_controller.observe_batch( + batch_size=batch_size, + num_drafts=adaptive_num_drafts, + num_draft_tokens=adaptive_num_draft_tokens, + num_accepted_tokens=adaptive_num_accepted_tokens, + ) + if ( + update is not None + and update.previous_num_spec_tokens != update.num_spec_tokens + ): + logger.debug( + "Adaptive speculative depth for batch size %d changed " + "from %d to %d (mean accepted drafts: %.2f, mean " + "attempted drafts: %.2f, window: %d steps).", + batch_size, + update.previous_num_spec_tokens, + update.num_spec_tokens, + update.mean_num_accepted_tokens, + update.mean_num_draft_tokens, + self.acceptance_length_controller.observation_window, + ) + + if spec_decoding_stats is not None: + spec_decoding_stats.current_num_spec_tokens = ( + self.acceptance_length_controller.num_spec_tokens_for_batch_size( + batch_size + ) + if self.acceptance_length_controller is not None + else scheduler_output.resolve_num_spec_tokens_to_schedule( + self.num_spec_tokens + ) + ) + + # Remove the stopped requests from the running and waiting queues. + if stopped_running_reqs: + self.running = remove_all(self.running, stopped_running_reqs) + if stopped_preempted_reqs: + # This is a rare case and unlikely to impact performance. + self.waiting.remove_requests(stopped_preempted_reqs) + self.skipped_waiting.remove_requests(stopped_preempted_reqs) + + error_req_ids = set(self.grammar_compile_error_reqs) + self.grammar_compile_error_reqs.clear() + if failed_kv_load_req_ids and not self.recompute_kv_load_failures: + error_req_ids.update(failed_kv_load_req_ids) + + if error_req_ids: + error_reqs = self.finish_requests( + error_req_ids, RequestStatus.FINISHED_ERROR + ) + for request in error_reqs: + outputs[request.client_index].append( + EngineCoreOutput( + request_id=request.request_id, + new_token_ids=[], + finish_reason=request.get_finished_reason(), + events=request.take_events(), + trace_headers=request.trace_headers, + ) + ) + + # KV Connector: update state for finished KV Transfers. + if kv_connector_output: + self._update_from_kv_xfer_finished(kv_connector_output) + + # EC Connector: update state from worker-side EC connector output. + if self.ec_connector is not None and ec_connector_output: + self.ec_connector.update_connector_output(ec_connector_output) + + # Worker-side KV connector stats from the model runner output. + kv_connector_stats: KVConnectorStats | None = ( + kv_connector_output.kv_connector_stats if kv_connector_output else None + ) + if self.connector: + # Scheduler-side KV connector stats collected after connector update. + scheduler_kv_connector_stats = self.connector.get_kv_connector_stats() + if ( + scheduler_kv_connector_stats is not None + and not scheduler_kv_connector_stats.is_empty() + ): + kv_connector_stats = ( + kv_connector_stats.aggregate(scheduler_kv_connector_stats) + if kv_connector_stats is not None + else scheduler_kv_connector_stats + ) + + # collect KV cache events from KV cache manager + events = self.kv_cache_manager.take_events() + + # collect KV cache events from connector + if self.connector is not None: + connector_events = self.connector.take_events() + if connector_events: + if events is None: + events = list(connector_events) + else: + events.extend(connector_events) + + # publish collected KV cache events + if events: + batch = KVEventBatch(ts=time.time(), events=events) + self.kv_event_publisher.publish(batch) + + # Create EngineCoreOutputs for all clients that have requests with + # outputs in this step. + engine_core_outputs = { + client_index: EngineCoreOutputs(outputs=outs) + for client_index, outs in outputs.items() + } + + finished_req_ids = self.finished_req_ids_dict + if finished_req_ids: + # Include ids of requests that finished since last outputs + # were sent. + for client_index, finished_set in finished_req_ids.items(): + # Set finished request set in EngineCoreOutputs for this client. + if (eco := engine_core_outputs.get(client_index)) is not None: + eco.finished_requests = finished_set + else: + engine_core_outputs[client_index] = EngineCoreOutputs( + finished_requests=finished_set + ) + finished_req_ids.clear() + + if ( + stats := self.make_stats( + spec_decoding_stats, + kv_connector_stats, + cudagraph_stats, + perf_stats, + ) + ) is not None: + # Return stats to only one of the front-ends. + if (eco := next(iter(engine_core_outputs.values()), None)) is None: + # We must return the stats even if there are no request + # outputs this step. + engine_core_outputs[0] = eco = EngineCoreOutputs() + eco.scheduler_stats = stats + + return engine_core_outputs + + @staticmethod + def _is_blocked_waiting_status(status: RequestStatus) -> bool: + return status in ( + RequestStatus.WAITING_FOR_STRUCTURED_OUTPUT_GRAMMAR, + RequestStatus.WAITING_FOR_REMOTE_KVS, + RequestStatus.WAITING_FOR_STREAMING_REQ, + ) + + def _enqueue_waiting_request(self, request: Request) -> None: + if self._is_blocked_waiting_status(request.status): + self.skipped_waiting.add_request(request) + else: + self.waiting.add_request(request) + + def _select_waiting_queue_for_scheduling(self) -> RequestQueue | None: + if self.policy == SchedulingPolicy.FCFS: + return self.skipped_waiting or self.waiting or None + + # PRIORITY mode: compare queue heads when both queues are non-empty. + if self.waiting and self.skipped_waiting: + waiting_req = self.waiting.peek_request() + skipped_req = self.skipped_waiting.peek_request() + return self.waiting if waiting_req < skipped_req else self.skipped_waiting + + return self.waiting or self.skipped_waiting or None + + def _handle_stopped_request(self, request: Request) -> bool: + """Return True if finished (can be False for resumable requests).""" + if not request.resumable: + return True + + if request.streaming_queue: + update = request.streaming_queue.popleft() + if update is None: + # Streaming request finished. + return True + self._update_request_as_session(request, update) + else: + request.status = RequestStatus.WAITING_FOR_STREAMING_REQ + self.num_waiting_for_streaming_input += 1 + + self._enqueue_waiting_request(request) + return False + + def _update_request_with_output( + self, request: Request, new_token_ids: list[int], is_stale: bool = False + ) -> tuple[list[int], bool]: + # is_stale is only used by the AsyncScheduler override. + # Append generated tokens and check for stop. Note that if + # a request is still being prefilled, we expect the model runner + # to return empty token ids for the request. + stopped = False + for num_new, output_token_id in enumerate(new_token_ids, 1): + request.append_output_token_ids(output_token_id) + + # Check for stop and update request state. + # This must be called before we make the EngineCoreOutput. + stopped = check_stop(request, self.max_model_len) + if stopped: + del new_token_ids[num_new:] # Trim new tokens if needed. + break + return new_token_ids, stopped + + def _free_encoder_inputs(self, request: Request) -> None: + cached_encoder_input_ids = self.encoder_cache_manager.get_cached_input_ids( + request + ) + # OPTIMIZATION: Avoid list(set) if the set is empty. + if not cached_encoder_input_ids: + return + + # Defer the free by the drafter's look-ahead so an entry stays + # referenced until the drafter's read-ahead has also passed it, + # mirroring the shift the encoder scheduling path applies. + spec_lookahead = self.num_prefill_lookahead + + # Here, we use list(set) to avoid modifying the set while iterating + # over it. + for input_id in list(cached_encoder_input_ids): + mm_feature = request.mm_features[input_id] + start_pos = mm_feature.mm_position.offset + num_tokens = mm_feature.mm_position.length + if self.is_encoder_decoder and request.num_computed_tokens > 0: + # With Whisper, as soon as we've generated a single token, + # we know we're done with the encoder input. Cross Attention + # KVs have been calculated and cached already. + self.encoder_cache_manager.free_encoder_input(request, input_id) + elif ( + start_pos + num_tokens + spec_lookahead + <= request.num_computed_tokens - request.num_output_placeholders + ): + # Processed, stored in the decoder KV cache, and far enough past + # the placeholder range (plus the drafter's look-ahead) that no + # rejection or drafter gather can reference it. + self.encoder_cache_manager.free_encoder_input(request, input_id) + + def update_draft_token_ids(self, draft_token_ids: DraftTokenIds) -> None: + for req_id, spec_token_ids in zip( + draft_token_ids.req_ids, + draft_token_ids.draft_token_ids, + ): + request = self.requests.get(req_id) + if request is None or request.is_finished(): + # The request may have been finished. Skip. + continue + + if request.is_prefill_chunk: + # Ignore draft tokens for prefill chunks. + if request.spec_token_ids: + request.spec_token_ids = [] + continue + + # Add newly generated spec token ids to the request. + if self.structured_output_manager.should_advance(request): + metadata = request.structured_output_request + spec_token_ids = metadata.grammar.validate_tokens(spec_token_ids) # type: ignore[union-attr] + request.spec_token_ids = spec_token_ids + + def update_draft_token_ids_in_output( + self, draft_token_ids: DraftTokenIds, scheduler_output: SchedulerOutput + ) -> None: + num_invalid_spec_tokens: dict[str, int] = {} + + sched_spec_tokens = scheduler_output.scheduled_spec_decode_tokens + for req_id, spec_token_ids in zip( + draft_token_ids.req_ids, + draft_token_ids.draft_token_ids, + ): + request = self.requests.get(req_id) + if request is None or request.is_finished(): + # The request may have been finished. Skip. + continue + + placeholder_spec_tokens = sched_spec_tokens.get(req_id) + if not placeholder_spec_tokens: + continue + + orig_num_spec_tokens = len(placeholder_spec_tokens) + # Trim drafts to scheduled number of spec tokens + # (needed for chunked prefill case for example). + del spec_token_ids[orig_num_spec_tokens:] + # Filter out spec tokens which do not adhere to the grammar. + if self.structured_output_manager.should_advance(request): + metadata = request.structured_output_request + spec_token_ids = metadata.grammar.validate_tokens(spec_token_ids) # type: ignore[union-attr] + # Pad to original number of spec tokens. + num_invalid_tokens = orig_num_spec_tokens - len(spec_token_ids) + if num_invalid_tokens: + spec_token_ids.extend([-1] * num_invalid_tokens) + num_invalid_spec_tokens[req_id] = num_invalid_tokens + + sched_spec_tokens[req_id] = spec_token_ids + + scheduler_output.num_invalid_spec_tokens = num_invalid_spec_tokens + + def get_request_counts(self) -> tuple[int, int]: + """Returns (num_running_reqs, num_waiting_reqs).""" + return len(self.running), len(self.waiting) + len(self.skipped_waiting) + + def get_kv_cache_usage(self) -> float: + """Returns the fraction of the KV cache currently in use (0.0-1.0).""" + return self.kv_cache_manager.usage + + def add_request(self, request: Request) -> None: + existing = self.requests.get(request.request_id) + if existing is not None: + update = StreamingUpdate.from_request(request) + if existing.status != RequestStatus.WAITING_FOR_STREAMING_REQ: + assert existing.streaming_queue is not None, "duplicate request id" + # Queue next input chunk (or finished sentinel). + existing.streaming_queue.append(update) + elif update is not None: + # Commence next input chunk. + self._update_request_as_session(existing, update) + else: + # Streaming-input session finished. + self.finish_requests(request.request_id, RequestStatus.FINISHED_ABORTED) + else: + if request.resumable: + request.streaming_queue = deque() + self._enqueue_waiting_request(request) + self.requests[request.request_id] = request + if self.spec_decode_metrics_level != "none": + request.spec_decode_metrics = RequestSpecDecodeMetrics.new( + self.num_spec_tokens + ) + if self.connector is not None: + self.connector.on_new_request(request) + if self.log_stats: + request.record_event(EngineCoreEventType.QUEUED) + + def finish_requests( + self, request_ids: str | Iterable[str] | None, finished_status: RequestStatus + ) -> list[Request]: + """Handles the finish signal from outside the scheduler. + + For example, the API server can abort a request when the client + disconnects. + + If request_ids is None, all requests will be finished. + + Returns: + List of requests that were aborted. Will not include any that were + already finished. + """ + assert RequestStatus.is_finished(finished_status) + if isinstance(request_ids, str): + request_ids = (request_ids,) + elif request_ids is not None: + request_ids = set(request_ids) + else: + request_ids = self.requests.keys() + + running_requests_to_remove = set() + waiting_requests_to_remove = [] + valid_requests = [] + + # First pass: collect requests to remove from queues + for req_id in request_ids: + request = self.requests.get(req_id) + if request is None or request.is_finished(): + # Invalid request ID. + continue + + valid_requests.append(request) + if request.status == RequestStatus.RUNNING: + running_requests_to_remove.add(request) + else: + if request.status == RequestStatus.WAITING_FOR_STREAMING_REQ: + self.num_waiting_for_streaming_input -= 1 + waiting_requests_to_remove.append(request) + + # Remove all requests from queues at once for better efficiency + if running_requests_to_remove: + self.running = remove_all(self.running, running_requests_to_remove) + if waiting_requests_to_remove: + self.waiting.remove_requests(waiting_requests_to_remove) + self.skipped_waiting.remove_requests(waiting_requests_to_remove) + + # Second pass: set status and free requests + for request in valid_requests: + delay_free_blocks = False + if request.status == RequestStatus.WAITING_FOR_REMOTE_KVS: + delay_free_blocks = ( + request.request_id not in self.finished_recving_kv_req_ids + ) + self.finished_recving_kv_req_ids.discard(request.request_id) + self.failed_recving_kv_req_ids.discard(request.request_id) + + request.status = finished_status + self._free_request(request, delay_free_blocks=delay_free_blocks) + + return valid_requests + + def _free_request( + self, request: Request, delay_free_blocks: bool = False + ) -> tuple[dict[str, Any] | None, dict[str, Any] | None]: + assert request.is_finished() + + self._inflight_prefills.discard(request) + connector_delay_free_blocks, kv_xfer_params = self._connector_finished(request) + + # EC Connector: mirror the KV hook. The contract requires firing + # before the encoder cache is freed so the connector can inspect + # per-request state (e.g. which mm_hashes it recorded during + # save_caches()) and emit ec_transfer_params for the response body. + ec_xfer_params: dict[str, Any] | None = None + if self.ec_connector is not None: + ec_delay_free, ec_xfer_params = self.ec_connector.request_finished(request) + connector_delay_free_blocks |= ec_delay_free + + self.encoder_cache_manager.free(request) + request_id = request.request_id + self.finished_req_ids.add(request_id) + if self.finished_req_ids_dict is not None: + self.finished_req_ids_dict[request.client_index].add(request_id) + + delay_free_blocks |= connector_delay_free_blocks + if not delay_free_blocks: + self._free_blocks(request) + + return kv_xfer_params, ec_xfer_params + + def _free_blocks(self, request: Request): + assert request.is_finished() + self._free_request_blocks(request) + del self.requests[request.request_id] + + @property + def pause_state(self) -> PauseState: + return self._pause_state + + def set_pause_state(self, pause_state: PauseState) -> None: + logger.info("setting pause state to %s", pause_state.name) + self._pause_state = pause_state + + def _free_request_blocks(self, request: Request): + """Free the request's KV blocks, deferring the return to the block + pool when an in-flight GPU step may still write them. + """ + if not self.defer_block_free or ( + # Last scheduled step already processed: no in-flight write remains + # (always the case for a normal finish), so free now. + request.last_sched_seq <= self.processed_step_seq + ): + self.kv_cache_manager.free(request) + return + blocks = self.kv_cache_manager.pop_blocks_for_free(request) + if blocks: + self.deferred_frees.append((self.sched_step_seq, blocks)) + + def _free_cow_retained_blocks( + self, blocks: list[KVCacheBlock], fence_seq: int + ) -> None: + """Release CoW copy retentions, deferring their return to the block + pool while the step that runs the copy may still be in flight. + """ + if not self.defer_block_free or fence_seq <= self.processed_step_seq: + self.kv_cache_manager.block_pool.free_blocks(blocks) + return + self.deferred_frees.append((fence_seq, blocks[::-1])) + + def _drain_deferred_frees(self): + """Return deferred blocks whose fence step has completed. + + Fences are appended in near-monotonic order (a CoW retention fence + can lead request-free fences by one step), so stop at the first + pending one; any satisfied entry behind it is merely freed later. + """ + while self.deferred_frees: + fence, _ = self.deferred_frees[0] + if fence > self.processed_step_seq: + break + _, blocks = self.deferred_frees.popleft() + # Free in reverse order so that the tail blocks are evicted first. + self.kv_cache_manager.block_pool.free_blocks(reversed(blocks)) + + def get_num_unfinished_requests(self) -> int: + if self._pause_state == PauseState.PAUSED_ALL: + return 0 + if self._pause_state == PauseState.PAUSED_NEW: + return len(self.running) + num_waiting = ( + len(self.waiting) + + len(self.skipped_waiting) + - self.num_waiting_for_streaming_input + ) + return num_waiting + len(self.running) + + def has_finished_requests(self) -> bool: + if self.finished_req_ids: + return True + if self.connector is None: + return False + # Finished requests waiting on delayed connector cleanup remain in + # self.requests after they have been removed from scheduling queues. + num_in_queues = ( + len(self.waiting) + len(self.skipped_waiting) + len(self.running) + ) + return len(self.requests) > num_in_queues + + def has_requests(self) -> bool: + # Override the interface default to also keep the engine alive while a + # connector still has pending push work (e.g. push-mode WRITE transfers + # in flight after all "live" requests have finished). Without this hook + # the engine would quiesce before the connector can drain completions. + # TODO: replace with a more general mechanism for connectors to keep + # the scheduler alive. + return ( + self.has_unfinished_requests() + or self.has_finished_requests() + or (self.connector is not None and self.connector.has_pending_push_work()) + or ( + self.ec_connector is not None + and self.ec_connector.has_pending_push_work() + ) + ) + + def reset_prefix_cache( + self, reset_running_requests: bool = False, reset_connector: bool = False + ) -> bool: + """Reset the KV prefix cache. + + If reset_running_requests is True, all the running requests will be + preempted and moved to the waiting queue. + Otherwise, this method will only reset the KV prefix cache when there + is no running requests taking KV cache. + """ + if reset_running_requests: + # For logging. + timestamp = time.monotonic() + # Invalidate all the current running requests KV's by pushing them to + # the waiting queue. In this case, we can reduce the ref count of all + # the kv blocks to 0 and thus we can make sure the reset is successful. + # Preempt in reverse order so the requests will be added back to the + # running queue in FIFO order. + while self.running: + request = self.running.pop() + self._preempt_request(request, timestamp, drop_stale_output=True) + + # Clear scheduled request ids cache. Since we are forcing preemption + # + resumption in the same step, we must act as if these requests were + # not scheduled in the prior step. They will be flushed from the + # persistent batch in the model runner. + self.prev_step_scheduled_req_ids.clear() + + reset_successful = self.kv_cache_manager.reset_prefix_cache() + if reset_running_requests and not reset_successful: + raise RuntimeError( + "Failed to reset KV cache even when all the running requests are " + "preempted and moved to the waiting queue. This is likely due to " + "the presence of running requests waiting for remote KV transfer, " + "which is not supported yet." + ) + + if reset_connector: + reset_successful = self.reset_connector_cache() and reset_successful + + return reset_successful + + def reset_connector_cache(self) -> bool: + if self.connector is None: + # No connector attached -> nothing to reset, treat as success so + # callers that unconditionally request a connector reset (e.g. as + # part of a cache-clearing cascade after a weight update) don't + # see reset_prefix_cache() flip to False purely because they + # didn't configure a connector. + logger.debug( + "reset_connector requested but no KV connector is configured; " + "treating as no-op success." + ) + return True + + if self.connector.reset_cache() is False: + return False + + if self.log_stats: + assert self.connector_prefix_cache_stats is not None + self.connector_prefix_cache_stats.reset = True + + return True + + def reset_encoder_cache(self) -> None: + """Reset the encoder cache to invalidate all cached encoder outputs. + + This should be called when model weights are updated to ensure + stale vision embeddings are not reused. + """ + self.encoder_cache_manager.reset() + + def make_stats( + self, + spec_decoding_stats: SpecDecodingStats | None = None, + kv_connector_stats: KVConnectorStats | None = None, + cudagraph_stats: CUDAGraphStat | None = None, + perf_stats: PerfStats | None = None, + ) -> SchedulerStats | None: + if not self.log_stats: + return None + prefix_cache_stats = self.kv_cache_manager.make_prefix_cache_stats() + assert prefix_cache_stats is not None + connector_prefix_cache_stats: PrefixCacheStats | None = None + if self.connector_prefix_cache_stats is not None: + connector_prefix_cache_stats = self.connector_prefix_cache_stats + self.connector_prefix_cache_stats = PrefixCacheStats() + eviction_events = ( + self.kv_metrics_collector.drain_events() + if self.kv_metrics_collector is not None + else [] + ) + spec_stats = spec_decoding_stats + connector_stats_payload = ( + kv_connector_stats.data if kv_connector_stats else None + ) + return SchedulerStats( + num_running_reqs=len(self.running), + num_waiting_reqs=len(self.waiting), + num_skipped_waiting_reqs=len(self.skipped_waiting), + kv_cache_usage=self.kv_cache_manager.usage, + prefix_cache_stats=prefix_cache_stats, + connector_prefix_cache_stats=connector_prefix_cache_stats, + kv_cache_eviction_events=eviction_events, + spec_decoding_stats=spec_stats, + kv_connector_stats=connector_stats_payload, + cudagraph_stats=cudagraph_stats, + perf_stats=perf_stats, + ) + + def make_spec_decoding_stats( + self, + spec_decoding_stats: SpecDecodingStats | None, + num_draft_tokens: int, + num_accepted_tokens: int, + num_invalid_spec_tokens: dict[str, int] | None, + request_id: str, + ) -> SpecDecodingStats | None: + if not self.log_stats or not num_draft_tokens: + return None + if spec_decoding_stats is None: + spec_decoding_stats = SpecDecodingStats.new(self.num_spec_tokens) + if num_invalid_spec_tokens: + num_draft_tokens -= num_invalid_spec_tokens.get(request_id, 0) + spec_decoding_stats.observe_draft( + num_draft_tokens=num_draft_tokens, num_accepted_tokens=num_accepted_tokens + ) + return spec_decoding_stats + + def shutdown(self) -> None: + logger.debug_once("[shutdown] Scheduler: start") + if self.kv_event_publisher: + self.kv_event_publisher.shutdown() + if self.connector is not None: + self.connector.shutdown() + + if self.ec_connector is not None: + self.ec_connector.shutdown() + + logger.debug_once("[shutdown] Scheduler: complete") + + ######################################################################## + # KV Connector Related Methods + ######################################################################## + + def get_kv_connector(self) -> KVConnectorBase_V1 | None: + return self.connector + + def get_ec_connector(self) -> ECConnectorBase | None: + return self.ec_connector + + def get_kv_event_publisher_config(self) -> KVEventsConfig | None: + return self.kv_event_publisher.get_publisher_config() + + def _connector_finished( + self, request: Request + ) -> tuple[bool, dict[str, Any] | None]: + """ + Invoke the KV connector request_finished() method if applicable. + + Returns optional kv transfer parameters to be included with the + request outputs. + """ + if self.connector is None: + return False, None + + # Free any out-of-window prefix blocks before we hand the block table to + # the connector, on the processed-token basis (see `allocate_slots`). + self.kv_cache_manager.remove_skipped_blocks( + request_id=request.request_id, + processed_computed_tokens=max( + 0, request.num_computed_tokens - request.num_in_flight_tokens + ), + num_prompt_tokens=request.num_prompt_tokens, + ) + + block_ids = self.kv_cache_manager.get_block_ids_for_computed_tokens( + request_id=request.request_id, + num_computed_tokens=request.num_computed_tokens, + ) + + if not isinstance(self.connector, SupportsHMA): + # NOTE(Kuntai): We should deprecate this code path after we enforce + # all connectors to support HMA. + # Hybrid memory allocator should be already turned off for this + # code path, but let's double-check here. + assert len(self.kv_cache_config.kv_cache_groups) == 1 + return self.connector.request_finished(request, block_ids[0]) + + return self.connector.request_finished_all_groups(request, block_ids) + + def _request_remaining_blocks(self, request: Request) -> int: + """Blocks `request` still needs to allocate to hold its full sequence.""" + full_num_tokens = min(request.num_tokens, self.max_model_len) + return self.kv_cache_manager.coordinator.get_num_blocks_to_allocate( + request_id=request.request_id, + num_tokens=full_num_tokens, + new_computed_blocks=self.kv_cache_manager.empty_kv_cache_blocks.blocks, + num_encoder_tokens=0, + total_computed_tokens=request.num_computed_tokens, + num_local_computed_tokens=request.num_computed_tokens, + num_tokens_main_model=full_num_tokens, + apply_admission_cap=True, + ) + + def _inflight_prefill_reserved_blocks(self) -> int: + """Num blocks in-flight prefills still need to finish (their reservation).""" + + return sum( + self._request_remaining_blocks(req) for req in self._inflight_prefills + ) + + def _update_waiting_for_remote_kv(self, request: Request) -> None: + """ + KV Connector: update request state after async recv is finished. + + When the kv transfer is ready, we cache the blocks + and the request state will be moved back to WAITING from + WAITING_FOR_REMOTE_KV. + """ + assert self.connector is not None + + if request.request_id in self.failed_recving_kv_req_ids: + # Request had KV load failures; num_computed_tokens was already + # updated in _update_requests_with_invalid_blocks + if request.num_computed_tokens: + # Cache any valid computed tokens. + self.kv_cache_manager.cache_blocks(request, request.num_computed_tokens) + if self.needs_kv_cache_zeroing: + # The failed load left the blocks beyond the valid + # prefix unwritten and their zeroing was skipped; zero + # them before they are recomputed locally. + self.kv_cache_manager.record_blocks_for_zeroing( + request.request_id, request.num_computed_tokens + ) + else: + # No valid computed tokens, release allocated blocks. + # There may be a local cache hit on retry. + # (Freed blocks are re-recorded for zeroing when + # reallocated, so the skipped blocks need no handling.) + self.kv_cache_manager.free(request) + + self.failed_recving_kv_req_ids.remove(request.request_id) + else: + # Now that the blocks are ready, actually cache them. + # This will cache the blocks iff caching is enabled. + self.kv_cache_manager.cache_blocks(request, request.num_computed_tokens) + + # Publication is deliberately after both the all-worker receive + # completion and cache_blocks(). A connector can now name this + # exact verified prefix for bounded by-reference sharing. + publish = getattr( + self.connector, "get_shared_prefix_lease_to_publish", None + ) + lease = publish(request) if publish is not None else None + if lease is not None: + lease_key, lease_tokens, ttl_seconds = lease + published = self.kv_cache_manager.publish_shared_prefix_lease( + lease_key, + request.request_id, + lease_tokens, + ttl_seconds, + max_entries=2, + reserved_blocks=self._request_remaining_blocks(request), + ) + if published: + self._pending_shared_prefix_leases[lease_key] = ( + request.request_id, + self.sched_step_seq + 1, + ) + else: + callback = getattr( + self.connector, "shared_prefix_lease_rejected", None + ) + if callback is not None: + callback(request.request_id, lease_key) + + # on a full prompt hit, we need to re-compute the last token + # in order to be able to sample the next token + if request.num_computed_tokens == request.num_tokens: + request.num_computed_tokens = request.num_tokens - 1 + + self.finished_recving_kv_req_ids.remove(request.request_id) + + def _finalize_shared_prefix_leases(self) -> None: + """Make copied hot pages attachable only after worker completion.""" + connector = self.connector + for lease_key, (request_id, fence_step) in tuple( + self._pending_shared_prefix_leases.items() + ): + if fence_step > self.processed_step_seq: + continue + self._pending_shared_prefix_leases.pop(lease_key, None) + ready = self.kv_cache_manager.mark_shared_prefix_lease_ready(lease_key) + accepted = False + if ready and connector is not None: + callback = getattr(connector, "shared_prefix_lease_published", None) + if callback is not None: + accepted = callback(request_id, lease_key) is not False + if not ready or not accepted: + self.kv_cache_manager.discard_shared_prefix_lease(lease_key) + if connector is not None: + callback = getattr(connector, "shared_prefix_lease_rejected", None) + if callback is not None: + callback(request_id, lease_key) + + def _try_promote_blocked_waiting_request(self, request: Request) -> bool: + """ + Try to promote a blocked waiting request back to schedulable states. + """ + if request.status == RequestStatus.WAITING_FOR_REMOTE_KVS: + # finished_recving_kv_req_ids is populated during + # update_from_output(), based on worker-side connector signals + # in KVConnectorOutput.finished_recving + if request.request_id not in self.finished_recving_kv_req_ids: + return False + self._update_waiting_for_remote_kv(request) + if request.num_preemptions: + request.status = RequestStatus.PREEMPTED + else: + request.status = RequestStatus.WAITING + return True + + if request.status == RequestStatus.WAITING_FOR_STRUCTURED_OUTPUT_GRAMMAR: + structured_output_req = request.structured_output_request + if not structured_output_req or structured_output_req.grammar is None: + return False + if isinstance(structured_output_req.grammar, Exception): + self.grammar_compile_error_reqs.add(request.request_id) + return False + request.status = RequestStatus.WAITING + return True + + if request.status == RequestStatus.WAITING_FOR_STREAMING_REQ: + assert not request.streaming_queue + return False + + raise AssertionError( + "Unexpected blocked waiting status in promotion: " + f"{request.status.name} for request {request.request_id}" + ) + + def _update_from_kv_xfer_finished(self, kv_connector_output: KVConnectorOutput): + """ + KV Connector: update the scheduler state based on the output. + + The Worker side connectors add finished_recving and + finished_sending reqs to the output. + * if finished_sending: free the blocks + # if finished_recving: add to state so we can + schedule the request during the next step. + """ + + if self.connector is not None: + self.connector.update_connector_output(kv_connector_output) + + # KV Connector:: update recv and send status from last step. + for req_id in kv_connector_output.finished_recving or (): + logger.debug("Finished recving KV transfer for request %s", req_id) + assert req_id in self.requests + req = self.requests[req_id] + if req.status == RequestStatus.WAITING_FOR_REMOTE_KVS: + self.finished_recving_kv_req_ids.add(req_id) + else: + assert RequestStatus.is_finished(req.status) + self._free_blocks(self.requests[req_id]) + for req_id in kv_connector_output.finished_sending or (): + logger.debug("Finished sending KV transfer for request %s", req_id) + assert req_id in self.requests + self._free_blocks(self.requests[req_id]) + + def _update_requests_with_invalid_blocks( + self, + requests: Iterable[Request], + invalid_block_ids: set[int], + num_scheduled_tokens: dict[str, int], + evict_blocks: bool = True, + ) -> tuple[set[str], int, set[int]]: + """ + Identify and update requests affected by invalid KV cache blocks. + + This method scans the given requests, detects those with invalid blocks + and adjusts their `num_computed_tokens` to the longest valid prefix. + For observability, it also accumulates the total number of tokens that + will need to be recomputed across all affected requests. + + Args: + requests: The set of requests to scan for invalid blocks. + invalid_block_ids: IDs of invalid blocks. + num_scheduled_tokens: req_id -> number of scheduled tokens. + evict_blocks: Whether to collect blocks for eviction (False for + async requests which aren't cached yet). + + Returns: + tuple: + - affected_req_ids (set[str]): IDs of requests impacted by + invalid blocks. + - total_affected_tokens (int): Total number of tokens that must + be recomputed across all affected requests. + - blocks_to_evict (set[int]): Block IDs to evict from cache, + including invalid blocks and downstream dependent blocks. + """ + affected_req_ids: set[str] = set() + total_affected_tokens = 0 + blocks_to_evict: set[int] = set() + # If a block is invalid and shared by multiple requests in the batch, + # these requests must be rescheduled, but only the first will recompute + # it. This set tracks blocks already marked for recomputation. + marked_invalid_block_ids: set[int] = set() + for request in requests: + is_affected = False + marked_invalid_block = False + req_id = request.request_id + req_num_computed_tokens = ( + request.num_computed_tokens - num_scheduled_tokens.get(req_id, 0) + ) + req_block_groups = self.kv_cache_manager.get_block_ids(req_id) + if len(req_block_groups) != 1: + request_block_ids = { + block_id + for group_block_ids in req_block_groups + for block_id in group_block_ids + } + if request_block_ids.isdisjoint(invalid_block_ids): + continue + # SparkCache restores every HMA group as one verified + # transaction. A failure in any group invalidates the whole + # external prefix, so restart this request from token zero. + affected_req_ids.add(req_id) + total_affected_tokens += req_num_computed_tokens + request.num_computed_tokens = 0 + if evict_blocks: + blocks_to_evict.update(request_block_ids) + continue + + (req_block_ids,) = req_block_groups + # We iterate only over blocks that may contain externally computed + # tokens + + req_num_computed_blocks = ( + req_num_computed_tokens + self.block_size - 1 + ) // self.block_size + for idx, block_id in zip(range(req_num_computed_blocks), req_block_ids): + if block_id not in invalid_block_ids: + continue + + is_affected = True + + if block_id in marked_invalid_block_ids: + # This invalid block is shared with a previous request + # and was already marked for recomputation. + # This means this request can still consider this block + # as computed when rescheduled. + # Currently this only applies to sync loading; Async + # loading does not yet support block sharing + continue + + marked_invalid_block_ids.add(block_id) + + if marked_invalid_block: + # This request has already marked an invalid block for + # recomputation and updated its num_computed_tokens. + continue + + marked_invalid_block = True + # Truncate the computed tokens at the first failed block + request.num_computed_tokens = idx * self.block_size + num_affected_tokens = ( + req_num_computed_tokens - request.num_computed_tokens + ) + total_affected_tokens += num_affected_tokens + + # collect invalid block and all downstream dependent blocks + if evict_blocks: + blocks_to_evict.update(req_block_ids[idx:]) + + if is_affected: + if not marked_invalid_block: + # All invalid blocks of this request are shared with + # previous requests and will be recomputed by them. + # Revert to considering only cached tokens as computed. + # Currently this only applies to sync loading; Async + # loading does not yet support block sharing + total_affected_tokens += ( + request.num_computed_tokens - req_num_computed_tokens + ) + request.num_computed_tokens = req_num_computed_tokens + + affected_req_ids.add(request.request_id) + + return affected_req_ids, total_affected_tokens, blocks_to_evict + + def _handle_invalid_blocks( + self, invalid_block_ids: set[int], num_scheduled_tokens: dict[str, int] + ) -> set[str]: + """ + Handle requests affected by invalid KV cache blocks. + + Returns: + Set of affected request IDs to skip in update_from_output main loop. + """ + should_fail = not self.recompute_kv_load_failures + + # handle async KV loads (not cached yet, evict_blocks=False) + async_load_reqs = ( + req + for req in self.skipped_waiting + if req.status == RequestStatus.WAITING_FOR_REMOTE_KVS + ) + async_failed_req_ids, num_failed_tokens, _ = ( + self._update_requests_with_invalid_blocks( + async_load_reqs, + invalid_block_ids, + num_scheduled_tokens, + evict_blocks=False, + ) + ) + + total_failed_requests = len(async_failed_req_ids) + total_failed_tokens = num_failed_tokens + + # handle sync loads (may be cached, collect blocks for eviction) + sync_failed_req_ids, num_failed_tokens, sync_blocks_to_evict = ( + self._update_requests_with_invalid_blocks( + self.running, invalid_block_ids, num_scheduled_tokens, evict_blocks=True + ) + ) + + total_failed_requests += len(sync_failed_req_ids) + total_failed_tokens += num_failed_tokens + + if not total_failed_requests: + return set() + + # evict invalid blocks and downstream dependent blocks from cache + # only when not using recompute policy (where blocks will be recomputed + # and reused by other requests sharing them) + if sync_blocks_to_evict and not self.recompute_kv_load_failures: + self.kv_cache_manager.evict_blocks(sync_blocks_to_evict) + + if should_fail: + all_failed_req_ids = async_failed_req_ids | sync_failed_req_ids + logger.error( + "Failing %d request(s) due to KV load failure " + "(failure_policy=fail, %d tokens affected). Request IDs: %s", + total_failed_requests, + total_failed_tokens, + all_failed_req_ids, + ) + return all_failed_req_ids + + logger.warning( + "Recovered from KV load failure: " + "%d request(s) rescheduled (%d tokens affected).", + total_failed_requests, + total_failed_tokens, + ) + + # Mark async requests with KV load failures for retry once loading completes + self.failed_recving_kv_req_ids |= async_failed_req_ids + # Return sync affected IDs to skip in update_from_output + return sync_failed_req_ids diff --git a/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/d415d6262885da11437e4a343da5ae04447273e529195b87547507f941fac0b8/__init__.py b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/d415d6262885da11437e4a343da5ae04447273e529195b87547507f941fac0b8/__init__.py new file mode 100644 index 00000000..aaf3ab92 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/d415d6262885da11437e4a343da5ae04447273e529195b87547507f941fac0b8/__init__.py @@ -0,0 +1,93 @@ +"""Chunked lower-bounded KDA prefill over a paged recurrent-state pool. + +The op consumes already-projected and convolved packed Q/K/V, the raw gate and +beta projections, and the per-head ``A_log`` and ``dt_bias`` parameters. It +computes the chunked gated delta rule with a lower-bounded gate, writes each +token's output, and advances a caller-owned recurrent-state pool in place. +Projection GEMMs, the causal convolution, and output gating are outside this +package. + +The recurrent-state pool uses the ``gdn_decode`` physical layout +``[slot, head, value_dim, key_dim]`` in fp32, so a prefill and a decode of the +same request share one pool without conversion. State slots are addressed by +index rather than gathered: a request names its initial slot, its final slot, +and optionally checkpoint slots with chunk-aligned token offsets, and the +op reads and writes those slots directly. ``Caps.max_checkpoints`` defaults to +one with legacy vector metadata. Explicit ``max_checkpoints=2`` uses contiguous +``[sequence_capacity, 2]`` checkpoint indices/offsets, requires checkpoint export +and transactional validation, and exports both states during the same recurrence. ``Caps.null_state_index`` may reserve +one index meaning "zero initial state" and "do not write". + +Requests are packed. Request ``r`` covers tokens +``cu_seqlens[r]:cu_seqlens[r + 1]``; ``num_seqs`` and ``num_tokens`` are device +scalars, so one plan serves every batch shape within its capacity. Tokens are +processed in sixteen-token chunks, ordered so that one pipeline window advances +every live request, which keeps the prepare and recurrence kernels overlapped. + +Planned lifecycle: ``plan(Caps(...))`` -> ``bind`` -> ``run``. Runtime launches +use caller-owned scratch, allocate no tensor storage, and are capture safe. +Device-side validation is transactional: bit 0 reports a duplicate or +conflicting write slot, bit 1 malformed packed metadata, bit 2 an invalid state +slot, and bit 3 an unusable checkpoint offset. Any error poisons the live +output rows without mutating recurrent state. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from ..._lib.meta import OpMeta, Provenance, install_lazy_api + +META = OpMeta( + name="kda_prefill", + group="sequence", + api_style="planned", + entry_points=( + "Binding", + "Caps", + "KdaPrefillConfig", + "KdaPrefillQuery", + "Plan", + "bind", + "clear_caches", + "is_supported", + "plan", + "prewarm", + "reference", + "run", + ), + dtypes=("bf16", "fp32", "int32", "int64"), + recipes=("lower_bounded_kda",), + provenance=Provenance( + repo="https://github.com/lukealonso/b12x", + commit="f6a46f4cc", + paths=("b12x/sequence/kda_prefill/_cute_kernels.py",), + ), + test_path="tests/sequence/test_kda_prefill.py", + since="1.4.0", + notes=( + "Chunk size is sixteen tokens, head dim 128, bf16 activations, and " + "fp32 recurrent state. Checkpoint offsets must be chunk aligned. The " + "gate lower bound must lie in [-5, 0). Requests whose chunks span more " + "than one pipeline window keep their running state in their final " + "slot, so those requests require a non-null final slot." + ), +) + +if TYPE_CHECKING: + from .api import ( # noqa: F401 + Binding, + Caps, + KdaPrefillConfig, + KdaPrefillQuery, + Plan, + bind, + clear_caches, + is_supported, + plan, + prewarm, + reference, + run, + ) + +install_lazy_api(globals(), META) diff --git a/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/daa8e63b76bbbe2acf0b71954958523d6771e84fb5a932c838d88f81110e3816/reference.py b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/daa8e63b76bbbe2acf0b71954958523d6771e84fb5a932c838d88f81110e3816/reference.py new file mode 100644 index 00000000..7ab86ba9 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/daa8e63b76bbbe2acf0b71954958523d6771e84fb5a932c838d88f81110e3816/reference.py @@ -0,0 +1,573 @@ +"""PyTorch oracles for chunked lower-bounded KDA prefill. + +Three oracles share one contract: + +``recurrent_kda`` runs the fp32 token recurrence for one packed sequence and is +the ground truth. ``prefill_kda`` applies it to a packed batch over a +recurrent-state pool, honouring the same metadata the kernel validates. +``prefill_kda_chunk_mirror`` implements the kernel's chunked algorithm with its +exact rounding points so kernel stages can be compared tensor by tensor; a +``MirrorPolicy`` selects alternative precisions for offline studies and a +``MirrorTrace`` records every per-tile intermediate. + +State layout: the pool is ``[slot, head, value_dim, key_dim]`` (the transpose of +the mathematical ``[key_dim, value_dim]`` state); oracles keep that orientation. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, field +from typing import Literal + +import torch + +from .._shared.kda_math import KDA_HEAD_DIM, kda_beta, kda_log_decay, l2_normalize + +LOG2E = 1.4426950408889634 + + +def _scalar(value: torch.Tensor | int) -> int: + if isinstance(value, torch.Tensor): + if value.numel() != 1: + raise ValueError("device count tensors must contain one element") + return int(value.item()) + return int(value) + + +def _bf16(x: torch.Tensor) -> torch.Tensor: + return x.to(torch.bfloat16).float() + + +def recurrent_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + *, + lower_bound: float, + initial_state: torch.Tensor, + checkpoint_offset: int = -1, + checkpoint_offsets: tuple[int, ...] | None = None, + scale: float | None = None, + eps: float = 1e-6, + qk_l2norm: bool = True, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | dict[int, torch.Tensor] | None]: + """Run the fp32 token recurrence for one sequence. + + ``q, k, v, raw_g`` are ``[T, heads, 128]``, ``raw_beta`` is ``[T, heads]``, + ``initial_state`` is ``[heads, 128, 128]`` in ``[value_dim, key_dim]`` + order. Returns the bf16 output ``[T, heads, 128]``, the fp32 final state in + the same orientation, and the state after ``checkpoint_offset`` tokens + (``None`` unless ``0 <= checkpoint_offset <= T``). With explicit plural + ``checkpoint_offsets``, the third result is an offset-to-state mapping + captured by that same recurrence; the legacy single-offset return is unchanged. + """ + tokens = int(q.shape[0]) + heads = int(q.shape[1]) + scale_value = KDA_HEAD_DIM**-0.5 if scale is None else float(scale) + qf = l2_normalize(q, eps) if qk_l2norm else q.float() + kf = l2_normalize(k, eps) if qk_l2norm else k.float() + qf = qf * scale_value + vf = v.float() + log_decay = kda_log_decay(raw_g, dt_bias, A_log, lower_bound) + beta = kda_beta(raw_beta) + state = initial_state.float().transpose(-1, -2).clone() # [heads, K, V] + output = torch.empty( + (tokens, heads, KDA_HEAD_DIM), dtype=torch.bfloat16, device=q.device + ) + checkpoints = {} + requested = (checkpoint_offset,) if checkpoint_offsets is None else checkpoint_offsets + checkpoint = None + if checkpoint_offset == 0: + checkpoint = state.transpose(-1, -2).contiguous() + for t in range(tokens): + state = state * torch.exp(log_decay[t])[:, :, None] + k_t = kf[t] + delta = vf[t] - torch.einsum("hk,hkv->hv", k_t, state) + state = state + (beta[t][:, None] * k_t)[:, :, None] * delta[:, None, :] + output[t] = torch.einsum("hk,hkv->hv", qf[t], state).to(torch.bfloat16) + if t + 1 in requested: + saved = state.transpose(-1, -2).contiguous().clone() + checkpoints[t + 1] = saved + if t + 1 == checkpoint_offset: + checkpoint = saved + return output, state.transpose(-1, -2).contiguous(), (checkpoint if checkpoint_offsets is None else checkpoints) + + +def _validate_packed( + *, + cu_seqlens: torch.Tensor, + initial_state_indices: torch.Tensor, + final_state_indices: torch.Tensor, + checkpoint_state_indices: torch.Tensor, + checkpoint_offsets: torch.Tensor, + num_seqs: int, + num_tokens: int, + token_capacity: int, + seq_capacity: int, + state_slots: int, + chunk: int, + null_state_index: int | None, + max_checkpoints: int = 1, +) -> list[tuple[int, int]]: + """Raise on every condition the device validator flags; return spans.""" + if max_checkpoints != 1: + from .metadata import validate_metadata + return validate_metadata(cu_seqlens=cu_seqlens, initial_state_indices=initial_state_indices, + final_state_indices=final_state_indices, checkpoint_state_indices=checkpoint_state_indices, + checkpoint_offsets=checkpoint_offsets, num_seqs=num_seqs, num_tokens=num_tokens, + token_capacity=token_capacity, seq_capacity=seq_capacity, state_slots=state_slots, + chunk=chunk, null_state_index=null_state_index, max_checkpoints=max_checkpoints) + if num_seqs < 0 or num_seqs > seq_capacity: + raise ValueError(f"num_seqs={num_seqs} exceeds capacity {seq_capacity}") + if num_tokens < 0 or num_tokens > token_capacity: + raise ValueError(f"num_tokens={num_tokens} exceeds capacity {token_capacity}") + if int(cu_seqlens[0]) != 0: + raise ValueError("cu_seqlens[0] must be zero") + if int(cu_seqlens[num_seqs]) != num_tokens: + raise ValueError("cu_seqlens[num_seqs] must equal num_tokens") + spans: list[tuple[int, int]] = [] + write_slots: set[int] = set() + read_slots: set[int] = set() + for request in range(num_seqs): + start = int(cu_seqlens[request]) + end = int(cu_seqlens[request + 1]) + if start < 0 or end < start or end > num_tokens: + raise ValueError(f"invalid query interval [{start}, {end})") + spans.append((start, end)) + + def is_null(slot: int) -> bool: + return null_state_index is not None and slot == null_state_index + + for request, (start, end) in enumerate(spans): + initial = int(initial_state_indices[request]) + final = int(final_state_indices[request]) + checkpoint = int(checkpoint_state_indices[request]) + offset = int(checkpoint_offsets[request]) + for slot, role in ((initial, "initial"), (final, "final"), (checkpoint, "checkpoint")): + if is_null(slot): + continue + if slot < 0 or slot >= state_slots: + raise IndexError(f"{role} state index {slot} is out of range") + if not is_null(initial): + read_slots.add(initial) + if offset > end - start: + raise ValueError("checkpoint offset exceeds the sequence length") + if offset > 0 and offset % chunk != 0: + raise ValueError(f"checkpoint offset {offset} is not a multiple of {chunk}") + for slot in (final, checkpoint if offset > 0 else None): + if slot is None or is_null(slot): + continue + if slot in write_slots: + raise ValueError(f"duplicate write state index {slot}") + write_slots.add(slot) + for request in range(len(spans)): + initial = int(initial_state_indices[request]) + final = int(final_state_indices[request]) + if is_null(initial): + continue + conflicting = write_slots - ({final} if not is_null(final) else set()) + if initial in conflicting: + raise ValueError(f"initial state index {initial} is written by another sequence") + return spans + + +def prefill_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + recurrent_state: torch.Tensor, + cu_seqlens: torch.Tensor, + initial_state_indices: torch.Tensor, + final_state_indices: torch.Tensor, + checkpoint_state_indices: torch.Tensor, + checkpoint_offsets: torch.Tensor, + num_seqs: torch.Tensor | int, + num_tokens: torch.Tensor | int, + *, + lower_bound: float = -5.0, + scale: float | None = None, + eps: float = 1e-6, + qk_l2norm: bool = True, + null_state_index: int | None = None, + chunk: int = 16, + max_checkpoints: int = 1, + output: torch.Tensor | None = None, +) -> torch.Tensor: + """Run the fp32 recurrence for a packed batch over a state pool. + + Sequence ``i`` consumes tokens ``cu_seqlens[i]:cu_seqlens[i+1]``, reads its + initial state from ``initial_state_indices[i]`` (zero when null), writes its + final state to ``final_state_indices[i]`` and, when ``checkpoint_offsets[i]`` + is a positive multiple of ``chunk``, the state after that many tokens to + ``checkpoint_state_indices[i]``. Output rows at or beyond ``num_tokens`` are + left untouched when ``output`` is supplied. + """ + lower_bound_value = float(lower_bound) + if not math.isfinite(lower_bound_value) or not -5.0 <= lower_bound_value < 0.0: + raise ValueError("lower_bound must be in [-5, 0)") + heads = int(q.shape[1]) + token_capacity = int(q.shape[0]) + seq_capacity = int(cu_seqlens.numel()) - 1 + live_seqs = _scalar(num_seqs) + live_tokens = _scalar(num_tokens) + spans = _validate_packed( + cu_seqlens=cu_seqlens, + initial_state_indices=initial_state_indices, + final_state_indices=final_state_indices, + checkpoint_state_indices=checkpoint_state_indices, + checkpoint_offsets=checkpoint_offsets, + num_seqs=live_seqs, + num_tokens=live_tokens, + token_capacity=token_capacity, + seq_capacity=seq_capacity, + state_slots=int(recurrent_state.shape[0]), + chunk=chunk, + null_state_index=null_state_index, + max_checkpoints=max_checkpoints, + ) + if output is None: + output = torch.zeros( + (token_capacity, heads, KDA_HEAD_DIM), dtype=torch.bfloat16, device=q.device + ) + + def is_null(slot: int) -> bool: + return null_state_index is not None and slot == null_state_index + + for request, (start, end) in enumerate(spans): + initial = int(initial_state_indices[request]) + final = int(final_state_indices[request]) + slots = [int(checkpoint_state_indices[request])] if max_checkpoints == 1 else [int(x) for x in checkpoint_state_indices[request]] + offsets = [int(checkpoint_offsets[request])] if max_checkpoints == 1 else [int(x) for x in checkpoint_offsets[request]] + if is_null(initial): + state = torch.zeros( + (heads, KDA_HEAD_DIM, KDA_HEAD_DIM), dtype=torch.float32, device=q.device + ) + else: + state = recurrent_state[initial].float() + out, final_state, checkpoint = recurrent_kda( + q[start:end], + k[start:end], + v[start:end], + raw_g[start:end], + raw_beta[start:end], + A_log, + dt_bias, + lower_bound=lower_bound_value, + initial_state=state, + checkpoint_offsets=tuple(offset for offset in offsets if offset > 0), + scale=scale, + eps=eps, + qk_l2norm=qk_l2norm, + ) + output[start:end] = out + for checkpoint_slot, offset in zip(slots, offsets): + if offset in checkpoint and not is_null(checkpoint_slot): + recurrent_state[checkpoint_slot].copy_(checkpoint[offset].to(recurrent_state.dtype)) + if not is_null(final): + recurrent_state[final].copy_(final_state.to(recurrent_state.dtype)) + return output + + +@dataclass(frozen=True) +class MirrorPolicy: + """Rounding points of the chunk mirror; the default is the kernel's policy.""" + + state_master: Literal["fp32", "bf16"] = "fp32" + shadow: bool = True + inv_operand: Literal["bf16", "fp32"] = "bf16" + u_operand: Literal["bf16", "fp32"] = "bf16" + single_rounding: bool = True + scale_dtype: Literal["fp32", "bf16"] = "fp32" + operands: Literal["bf16", "fp32"] = "bf16" + + +@dataclass +class MirrorTrace: + """Per-(sequence, tile) intermediates recorded by the chunk mirror.""" + + k1: dict[tuple[int, int], dict[str, torch.Tensor]] = field(default_factory=dict) + k2: dict[tuple[int, int], dict[str, torch.Tensor]] = field(default_factory=dict) + checkpoints: dict[int | tuple[int, int], torch.Tensor] = field(default_factory=dict) + + +def _neumann_inverse(lower: torch.Tensor, chunk: int) -> torch.Tensor: + """Return ``(I + L)^{-1}`` for strictly lower-triangular ``L`` in fp32. + + ``-L`` is nilpotent, so the product ``(I - L)(I + L^2)(I + L^4)...`` over + ``log2(chunk)`` factors is the exact inverse. + """ + eye = torch.eye(chunk, dtype=torch.float32, device=lower.device) + inverse = eye - lower + power = lower + steps = int(math.log2(chunk)) + for _ in range(1, steps): + power = power @ power + inverse = inverse + inverse @ power + return inverse + + +def _prepare_tile( + q: torch.Tensor, + k: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + *, + rows: int, + chunk: int, + lower_bound: float, + scale: float, + eps: float, + qk_l2norm: bool, + policy: MirrorPolicy, +) -> dict[str, torch.Tensor]: + """K1 mirror for one tile of all heads; inputs are ``[chunk, heads, ...]``.""" + heads = int(q.shape[1]) + mask = torch.zeros((chunk, 1, 1), dtype=torch.bool, device=q.device) + mask[:rows] = True + g2 = kda_log_decay(raw_g, dt_bias, A_log, lower_bound) * LOG2E + g2 = torch.where(mask, g2, torch.zeros_like(g2)) + cum = torch.empty_like(g2) + running = torch.zeros_like(g2[0]) + for t in range(chunk): + running = running + g2[t] + cum[t] = running + lam = torch.exp2(cum) + lam_inv = torch.exp2(-cum) + lam_c = torch.exp2(cum[chunk - 1]) + lam_r = torch.exp2(cum[chunk - 1][None] - cum) + qn = l2_normalize(q, eps) if qk_l2norm else q.float() + kn = l2_normalize(k, eps) if qk_l2norm else k.float() + qn = torch.where(mask, qn, torch.zeros_like(qn)) + kn = torch.where(mask, kn, torch.zeros_like(kn)) + scale_value = _bf16(torch.tensor(scale)).item() if policy.scale_dtype == "bf16" else scale + if policy.operands == "fp32": + q_tilde = qn * lam * scale_value + k_tilde = kn * lam + k_inv = kn * lam_inv + k_r = kn * lam_r + elif policy.single_rounding: + q_tilde = _bf16(qn * lam * scale_value) + k_tilde = _bf16(kn * lam) + k_inv = _bf16(kn * lam_inv) + k_r = _bf16(kn * lam_r) + else: + qb, kb = _bf16(qn), _bf16(kn) + q_tilde = _bf16(_bf16(qb * _bf16(lam)) * scale_value) + k_tilde = _bf16(kb * _bf16(lam)) + k_inv = _bf16(kb * _bf16(lam_inv)) + k_r = _bf16(k_inv * _bf16(lam_c)[None]) + beta = kda_beta(raw_beta) + beta = torch.where(mask[:, :, 0], beta, torch.zeros_like(beta)) + # [heads, chunk, 128] operands for the per-head GEMMs. + q_h, k_h, kinv_h = (x.transpose(0, 1) for x in (q_tilde, k_tilde, k_inv)) + beta_h = beta.transpose(0, 1) + causal = torch.tril(torch.ones((chunk, chunk), dtype=torch.bool, device=q.device), -1) + lower = beta_h[:, :, None] * (k_h @ kinv_h.transpose(-1, -2)) + lower = torch.where(causal, lower, torch.zeros_like(lower)) + inverse = _neumann_inverse(lower, chunk) + inclusive = torch.tril(torch.ones((chunk, chunk), dtype=torch.bool, device=q.device)) + mqk = q_h @ kinv_h.transpose(-1, -2) + mqk = torch.where(inclusive, mqk, torch.zeros_like(mqk)) + return { + "g_cum": cum.transpose(0, 1), + "lambda_c": lam_c, + "q_tilde": q_h, + "k_tilde": k_h, + "k_inv": kinv_h, + "k_r": k_r.transpose(0, 1), + "beta": beta_h, + "L": lower, + "inv": inverse, + "inv_op": _bf16(inverse) if policy.inv_operand == "bf16" else inverse, + "mqk": _bf16(mqk) if policy.operands == "bf16" else mqk, + "heads": torch.tensor(heads), + } + + +def _recur_tile( + state: torch.Tensor, + v: torch.Tensor, + prep: dict[str, torch.Tensor], + *, + rows: int, + chunk: int, + policy: MirrorPolicy, +) -> tuple[torch.Tensor, torch.Tensor, dict[str, torch.Tensor]]: + """K2 mirror for one tile; ``state`` is ``[heads, V, K]`` fp32 (master).""" + mask = torch.zeros((chunk, 1, 1), dtype=torch.bool, device=v.device) + mask[:rows] = True + v_h = torch.where(mask, v.float(), torch.zeros_like(v.float())).transpose(0, 1) + shadow = _bf16(state) if policy.shadow else state + v_prime = (v_h - prep["k_tilde"] @ shadow.transpose(-1, -2)) * prep["beta"][:, :, None] + v_prime_op = _bf16(v_prime) if policy.operands == "bf16" else v_prime + u = prep["inv_op"] @ v_prime_op + u_op = _bf16(u) if policy.u_operand == "bf16" else u + out = prep["q_tilde"] @ shadow.transpose(-1, -2) + prep["mqk"] @ u_op + delta_t = u_op.transpose(-1, -2) @ prep["k_r"] + new_state = state * prep["lambda_c"][:, None, :] + delta_t + if policy.state_master == "bf16": + new_state = _bf16(new_state) + trace = { + "v_prime": v_prime, + "u": u, + "out_tile": out, + "delta_state": delta_t, + "state_after": new_state, + } + return out.transpose(0, 1).to(torch.bfloat16), new_state, trace + + +def prefill_kda_chunk_mirror( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + recurrent_state: torch.Tensor, + cu_seqlens: torch.Tensor, + initial_state_indices: torch.Tensor, + final_state_indices: torch.Tensor, + checkpoint_state_indices: torch.Tensor, + checkpoint_offsets: torch.Tensor, + num_seqs: torch.Tensor | int, + num_tokens: torch.Tensor | int, + *, + lower_bound: float = -5.0, + scale: float | None = None, + eps: float = 1e-6, + qk_l2norm: bool = True, + null_state_index: int | None = None, + chunk: int = 16, + max_checkpoints: int = 1, + policy: MirrorPolicy | None = None, + trace: bool = False, + output: torch.Tensor | None = None, +) -> torch.Tensor | tuple[torch.Tensor, MirrorTrace]: + """Run the kernel's chunked algorithm with its rounding points. + + Same contract as :func:`prefill_kda`. With ``trace=True`` the per-tile + K1 and K2 intermediates are returned alongside the output. + """ + lower_bound_value = float(lower_bound) + if not math.isfinite(lower_bound_value) or not -5.0 <= lower_bound_value < 0.0: + raise ValueError("lower_bound must be in [-5, 0)") + if chunk & (chunk - 1) or chunk < 2: + raise ValueError("chunk must be a power of two") + policy = MirrorPolicy() if policy is None else policy + heads = int(q.shape[1]) + scale_value = KDA_HEAD_DIM**-0.5 if scale is None else float(scale) + token_capacity = int(q.shape[0]) + seq_capacity = int(cu_seqlens.numel()) - 1 + live_seqs = _scalar(num_seqs) + live_tokens = _scalar(num_tokens) + spans = _validate_packed( + cu_seqlens=cu_seqlens, + initial_state_indices=initial_state_indices, + final_state_indices=final_state_indices, + checkpoint_state_indices=checkpoint_state_indices, + checkpoint_offsets=checkpoint_offsets, + num_seqs=live_seqs, + num_tokens=live_tokens, + token_capacity=token_capacity, + seq_capacity=seq_capacity, + state_slots=int(recurrent_state.shape[0]), + chunk=chunk, + null_state_index=null_state_index, + max_checkpoints=max_checkpoints, + ) + if output is None: + output = torch.zeros( + (token_capacity, heads, KDA_HEAD_DIM), dtype=torch.bfloat16, device=q.device + ) + record = MirrorTrace() + + def is_null(slot: int) -> bool: + return null_state_index is not None and slot == null_state_index + + def padded(x: torch.Tensor, start: int, rows: int) -> torch.Tensor: + tile = torch.zeros((chunk, *x.shape[1:]), dtype=x.dtype, device=x.device) + tile[:rows] = x[start : start + rows] + return tile + + previous_tf32 = torch.backends.cuda.matmul.allow_tf32 + torch.backends.cuda.matmul.allow_tf32 = False + try: + for request, (start, end) in enumerate(spans): + initial = int(initial_state_indices[request]) + final = int(final_state_indices[request]) + slots = [int(checkpoint_state_indices[request])] if max_checkpoints == 1 else [int(x) for x in checkpoint_state_indices[request]] + offsets = [int(checkpoint_offsets[request])] if max_checkpoints == 1 else [int(x) for x in checkpoint_offsets[request]] + if is_null(initial): + state = torch.zeros( + (heads, KDA_HEAD_DIM, KDA_HEAD_DIM), + dtype=torch.float32, + device=q.device, + ) + else: + state = recurrent_state[initial].float() + length = end - start + for local in range((length + chunk - 1) // chunk): + tile_start = start + local * chunk + rows = min(chunk, end - tile_start) + prep = _prepare_tile( + padded(q, tile_start, rows), + padded(k, tile_start, rows), + padded(raw_g, tile_start, rows), + padded(raw_beta, tile_start, rows), + A_log, + dt_bias, + rows=rows, + chunk=chunk, + lower_bound=lower_bound_value, + scale=scale_value, + eps=eps, + qk_l2norm=qk_l2norm, + policy=policy, + ) + out_tile, state, step = _recur_tile( + state, + padded(v, tile_start, rows), + prep, + rows=rows, + chunk=chunk, + policy=policy, + ) + output[tile_start : tile_start + rows] = out_tile[:rows] + if trace: + record.k1[(request, local)] = prep + record.k2[(request, local)] = step + for checkpoint_index, (checkpoint_slot, offset) in enumerate(zip(slots, offsets)): + if offset > 0 and (local + 1) * chunk == offset: + if trace: + key = request if max_checkpoints == 1 else (request, checkpoint_index) + record.checkpoints[key] = state.clone() + if not is_null(checkpoint_slot): + recurrent_state[checkpoint_slot].copy_(state.to(recurrent_state.dtype)) + if not is_null(final): + recurrent_state[final].copy_(state.to(recurrent_state.dtype)) + finally: + torch.backends.cuda.matmul.allow_tf32 = previous_tf32 + return (output, record) if trace else output + + +__all__ = [ + "LOG2E", + "MirrorPolicy", + "MirrorTrace", + "prefill_kda", + "prefill_kda_chunk_mirror", + "recurrent_kda", +] diff --git a/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/dc627424cd5814db0945b4c787e33fb6d035c2a2dd3d0a4ade4dadce61697810/model_state.py b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/dc627424cd5814db0945b4c787e33fb6d035c2a2dd3d0a4ade4dadce61697810/model_state.py new file mode 100644 index 00000000..ee32274c --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/dc627424cd5814db0945b4c787e33fb6d035c2a2dd3d0a4ade4dadce61697810/model_state.py @@ -0,0 +1,434 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Model-runner state for the GLM5Next pooled sparse-attention selector.""" + +from dataclasses import dataclass +from typing import Any + +import numpy as np +import torch +import torch.nn as nn + +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.config.compilation import CUDAGraphMode +from vllm.triton_utils import tl, triton +from vllm.v1.core.sched.output import NewRequestData +from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.utils import CpuGpuBuffer +from vllm.v1.worker.gpu.attn_utils import build_attn_metadata +from vllm.v1.worker.gpu.input_batch import InputBatch +from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache +from vllm.v1.worker.gpu.model_states.mamba_hybrid import ( + _checkpoint_plan_rows, + MambaHybridAttnMetadata, + MambaHybridModelState, +) +from vllm.v1.worker.utils import AttentionGroup + + +logger = init_logger(__name__) + + +@dataclass +class Glm5NextAttnMetadata(MambaHybridAttnMetadata): + """Per-request state consumed by GLM5Next's b12x selector builder.""" + + selector_state_slot_ids: torch.Tensor | None = None + selector_state_is_fresh: torch.Tensor | None = None + selector_num_accepted_tokens: torch.Tensor | None = None + selector_is_prefilling: torch.Tensor | None = None + + def get_extra_attn_kwargs( + self, + attn_metadata_builder: Any, + num_reqs: int, + ) -> dict[str, Any]: + kwargs = super().get_extra_attn_kwargs(attn_metadata_builder, num_reqs) + if not getattr( + attn_metadata_builder, + "requires_glm_next_selector_metadata", + False, + ): + return kwargs + assert self.selector_state_slot_ids is not None + assert self.selector_state_is_fresh is not None + assert self.selector_num_accepted_tokens is not None + assert self.selector_is_prefilling is not None + kwargs.update( + selector_state_slot_ids=self.selector_state_slot_ids[:num_reqs], + selector_state_is_fresh=self.selector_state_is_fresh[:num_reqs], + selector_num_accepted_tokens=(self.selector_num_accepted_tokens[:num_reqs]), + selector_is_prefilling=self.selector_is_prefilling[:num_reqs], + ) + return kwargs + + +class Glm5NextModelState(MambaHybridModelState): + """Add persistent request identity for GLM5Next pooled selector state.""" + + def __init__( + self, + vllm_config: VllmConfig, + model: nn.Module, + encoder_cache: EncoderCache | None, + device: torch.device, + ) -> None: + super().__init__(vllm_config, model, encoder_cache, device) + config = self.model_config.hf_text_config + self.uses_pooled_selector = getattr(config, "index_topk", None) is not None + self.selector_pool_size = int(getattr(config, "index_kpool", 1) or 1) + + # These are fixed-capacity staging buffers. Their addresses remain stable + # across request reordering and CUDA-graph capture/replay. + self.selector_state_slot_ids = torch.full( + (self.max_num_reqs,), + -1, + dtype=torch.int32, + device=self.device, + ) + self.selector_state_is_fresh = torch.ones( + self.max_num_reqs, + dtype=torch.bool, + device=self.device, + ) + self.selector_num_accepted_tokens = torch.ones( + self.max_num_reqs, + dtype=torch.int32, + device=self.device, + ) + self.mamba_num_accepted_tokens = torch.ones( + self.max_num_reqs, + dtype=torch.int32, + device=self.device, + ) + self.selector_committed_num_accepted_tokens_gpu = torch.ones( + self.max_num_reqs, + dtype=torch.int32, + device=self.device, + ) + self.selector_state_is_fresh_gpu = torch.ones( + self.max_num_reqs, + dtype=torch.bool, + device=self.device, + ) + self.selector_is_prefilling = CpuGpuBuffer( + self.max_num_reqs, + dtype=torch.bool, + device=self.device, + ) + self._selector_draft_is_prefilling = torch.zeros( + self.max_num_reqs, + dtype=torch.bool, + device="cpu", + ) + self._selector_draft_is_prefilling_gpu = torch.zeros( + self.max_num_reqs, + dtype=torch.bool, + device=self.device, + ) + + def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: + prefix_length = int(new_req_data.num_computed_tokens) + if self.uses_pooled_selector and prefix_length % self.selector_pool_size != 0: + raise ValueError( + "GLM5Next pooled selector cannot resume a fresh request from " + f"num_computed_tokens={prefix_length}; the prefix length must be " + f"divisible by index_kpool={self.selector_pool_size}." + ) + super().add_request(req_index, new_req_data) + if self.uses_pooled_selector: + # The scheduler may recycle this request-state slot while selector + # raw-ring tags and its interval anchor still belong to the prior owner. + self.selector_state_is_fresh_gpu[req_index].fill_(True) + self.selector_committed_num_accepted_tokens_gpu[req_index].fill_(1) + + def reset_kv_cache_state(self) -> None: + super().reset_kv_cache_state() + if self.uses_pooled_selector: + self.selector_state_is_fresh_gpu.fill_(True) + self.selector_committed_num_accepted_tokens_gpu.fill_(1) + + def _prepare_selector_state( + self, + input_batch: InputBatch, + num_reqs: int, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + slots = self.selector_state_slot_ids[:num_reqs] + fresh = self.selector_state_is_fresh[:num_reqs] + accepted = self.selector_num_accepted_tokens[:num_reqs] + slots.fill_(-1) + fresh.fill_(True) + accepted.fill_(1) + + num_actual_reqs = input_batch.num_reqs + if num_actual_reqs: + idx_mapping = input_batch.idx_mapping[:num_actual_reqs] + slots[:num_actual_reqs].copy_(idx_mapping) + torch.index_select( + self.selector_state_is_fresh_gpu, + 0, + idx_mapping, + out=fresh[:num_actual_reqs], + ) + torch.index_select( + self.selector_committed_num_accepted_tokens_gpu, + 0, + idx_mapping, + out=accepted[:num_actual_reqs], + ) + return slots, fresh, accepted + + def _prepare_mamba_acceptance( + self, + input_batch: InputBatch, + num_reqs: int, + ) -> torch.Tensor: + accepted = self.mamba_num_accepted_tokens[:num_reqs] + accepted.fill_(1) + num_actual_reqs = input_batch.num_reqs + if num_actual_reqs: + torch.index_select( + self.num_accepted_tokens_gpu, + 0, + input_batch.idx_mapping[:num_actual_reqs], + out=accepted[:num_actual_reqs], + ) + return accepted + + def prepare_draft_attn_metadata( + self, + *, + idx_mapping: torch.Tensor, + num_reqs: int, + num_reqs_padded: int, + draft_index: int, + ) -> Glm5NextAttnMetadata | None: + if not self.uses_pooled_selector: + return None + if draft_index < 1: + raise RuntimeError( + "GLM5Next supports pooled-selector draft metadata only for " + "autoregressive MTP lookahead (draft_index >= 1)" + ) + if not 0 <= num_reqs <= num_reqs_padded <= self.max_num_reqs: + raise ValueError( + "draft request counts must satisfy " + "0 <= num_reqs <= num_reqs_padded <= max_num_reqs" + ) + if idx_mapping.numel() < num_reqs: + raise ValueError("idx_mapping does not cover every active draft request") + + slots = self.selector_state_slot_ids[:num_reqs_padded] + fresh = self.selector_state_is_fresh[:num_reqs_padded] + accepted = self.selector_num_accepted_tokens[:num_reqs_padded] + slots.fill_(-1) + fresh.fill_(True) + accepted.fill_(1) + if num_reqs: + active_slots = idx_mapping[:num_reqs] + slots[:num_reqs].copy_(active_slots) + # Draft prefill uses the target metadata immediately before the + # lookahead loop and has initialized these MTP selector slots. + fresh[:num_reqs].fill_(False) + if draft_index == 1: + torch.index_select( + self.selector_committed_num_accepted_tokens_gpu, + 0, + active_slots, + out=accepted[:num_reqs], + ) + + return Glm5NextAttnMetadata( + is_prefilling=self._selector_draft_is_prefilling[:num_reqs_padded], + num_accepted_tokens=accepted, + selector_state_slot_ids=slots, + selector_state_is_fresh=fresh, + selector_num_accepted_tokens=accepted, + selector_is_prefilling=( + self._selector_draft_is_prefilling_gpu[:num_reqs_padded] + ), + ) + + def prepare_attn( + self, + input_batch: InputBatch, + cudagraph_mode: CUDAGraphMode, + block_tables: tuple[torch.Tensor, ...], + slot_mappings: torch.Tensor, + attn_groups: list[list[AttentionGroup]], + kv_cache_config: KVCacheConfig, + for_capture: bool = False, + recurrent_prefill_checkpoint_plans: dict | None = None, + ) -> dict[str, Any]: + # This is the MambaHybridModelState construction with only the metadata + # object specialized. Keeping it package-local avoids a GLM hook in the + # generic model runner. + if cudagraph_mode == CUDAGraphMode.FULL: + num_reqs = input_batch.num_reqs_after_padding + num_tokens = input_batch.num_tokens_after_padding + else: + num_reqs = input_batch.num_reqs + num_tokens = input_batch.num_tokens + query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np) + max_query_len = input_batch.num_scheduled_tokens.max().item() + seq_lens_cpu_upper_bound = input_batch.seq_lens_cpu_upper_bound + if for_capture: + max_seq_len = self.max_model_len + else: + max_seq_len = seq_lens_cpu_upper_bound[:num_reqs].max().item() + + self.selector_is_prefilling.np[:num_reqs] = False + self.selector_is_prefilling.np[: input_batch.num_reqs] = ( + input_batch.is_prefilling_np + ) + is_prefilling = self.selector_is_prefilling.cpu[:num_reqs] + selector_is_prefilling = self.selector_is_prefilling.copy_to_gpu(num_reqs) + ( + selector_state_slot_ids, + selector_state_is_fresh, + selector_num_accepted_tokens, + ) = self._prepare_selector_state(input_batch, num_reqs) + + # During CUDA-graph capture the builders create their own neutral + # speculative metadata. Runtime calls stage the persistent buffers. + num_accepted_tokens = None + num_decode_draft_tokens_cpu = None + if not for_capture and self.vllm_config.num_speculative_tokens > 0: + # Mamba page alignment can reset its accepted-token offset after + # moving state. The selector advances an independent pooled interval. + num_accepted_tokens = self._prepare_mamba_acceptance( + input_batch, + num_reqs, + ) + num_decode_draft_tokens_np = np.full(num_reqs, -1, dtype=np.int32) + num_draft_tokens_per_req = input_batch.num_draft_tokens_per_req + if num_draft_tokens_per_req is not None: + is_decode = ( + input_batch.num_scheduled_tokens == num_draft_tokens_per_req + 1 + ) + spec_decode_mask = (num_draft_tokens_per_req > 0) & is_decode + num_decode_draft_tokens_np[: input_batch.num_reqs] = np.where( + spec_decode_mask, + num_draft_tokens_per_req, + -1, + ) + num_decode_draft_tokens_cpu = torch.from_numpy(num_decode_draft_tokens_np) + + if self._align_mode: + self._prepare_aligned_state_indices( + input_batch.seq_lens, + num_reqs, + attn_groups, + kv_cache_config, + block_tables, + ) + + checkpoint_plans_cpu = _checkpoint_plan_rows( + input_batch, recurrent_prefill_checkpoint_plans, num_reqs, + for_capture=for_capture, + ) + model_metadata = Glm5NextAttnMetadata( + recurrent_prefill_checkpoint_plans_cpu=checkpoint_plans_cpu, + is_prefilling=is_prefilling, + num_accepted_tokens=num_accepted_tokens, + num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, + selector_state_slot_ids=selector_state_slot_ids, + selector_state_is_fresh=selector_state_is_fresh, + selector_num_accepted_tokens=selector_num_accepted_tokens, + selector_is_prefilling=selector_is_prefilling, + ) + attn_metadata = build_attn_metadata( + attn_groups=attn_groups, + num_reqs=num_reqs, + num_tokens=num_tokens, + query_start_loc_gpu=input_batch.query_start_loc, + query_start_loc_cpu=query_start_loc_cpu, + max_query_len=max_query_len, + seq_lens=input_batch.seq_lens, + max_seq_len=max_seq_len, + block_tables=block_tables, + slot_mappings=slot_mappings, + kv_cache_config=kv_cache_config, + seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, + dcp_local_seq_lens=input_batch.dcp_local_seq_lens, + model_specific_attn_metadata=model_metadata, + for_cudagraph_capture=for_capture, + rswa_prefix_lens=input_batch.prompt_lens, + ) + if checkpoint_plans_cpu is not None: + gdn_metadata = [metadata for metadata in attn_metadata.values() + if hasattr(metadata, "prefill_checkpoint")] + if not gdn_metadata or any(metadata.prefill_checkpoint is None + or metadata.prefill_checkpoint.required_mask is None + for metadata in gdn_metadata): + raise RuntimeError("GLM V2 scheduled checkpoint plans without GDN export metadata") + if not getattr(self, "_explicit_checkpoint_plan_logged", False): + logger.info("GLM V2 explicit recurrent checkpoint plan active: rows=%s layers=%d", + checkpoint_plans_cpu, len(gdn_metadata)) + self._explicit_checkpoint_plan_logged = True + if self.recoverssm is not None: + self.recoverssm.record_step( + attn_metadata, + attn_groups, + for_capture=for_capture, + ) + return attn_metadata + + def postprocess_state( + self, + idx_mapping: torch.Tensor, + num_sampled: torch.Tensor | int, + num_computed_tokens: torch.Tensor | None = None, + ) -> None: + if self.uses_pooled_selector and idx_mapping.numel(): + if isinstance(num_sampled, int): + _fill_selector_request_state_kernel[(idx_mapping.numel(),)]( + idx_mapping, + self.selector_committed_num_accepted_tokens_gpu, + self.selector_state_is_fresh_gpu, + max(num_sampled, 1), + ) + else: + _commit_selector_request_state_kernel[(idx_mapping.numel(),)]( + idx_mapping, + num_sampled, + self.selector_committed_num_accepted_tokens_gpu, + self.selector_state_is_fresh_gpu, + ) + super().postprocess_state(idx_mapping, num_sampled, num_computed_tokens) + + +@triton.jit +def _commit_selector_request_state_kernel( + idx_mapping_ptr, + num_sampled_ptr, + selector_num_accepted_ptr, + state_is_fresh_ptr, +): + row = tl.program_id(0) + state_slot = tl.load(idx_mapping_ptr + row) + if state_slot >= 0: + num_sampled = tl.load(num_sampled_ptr + row) + tl.store( + selector_num_accepted_ptr + state_slot, + tl.maximum(num_sampled, 1), + ) + tl.store(state_is_fresh_ptr + state_slot, 0) + + +@triton.jit +def _fill_selector_request_state_kernel( + idx_mapping_ptr, + selector_num_accepted_ptr, + state_is_fresh_ptr, + num_sampled, +): + row = tl.program_id(0) + state_slot = tl.load(idx_mapping_ptr + row) + if state_slot >= 0: + tl.store(selector_num_accepted_ptr + state_slot, num_sampled) + tl.store(state_is_fresh_ptr + state_slot, 0) + + +__all__ = ["Glm5NextAttnMetadata", "Glm5NextModelState"] diff --git a/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/df12ee84e1677ce1a200d1aab488c9ce5ae774d21f70b7f15b39c636623ab753/_cute_kernels.py b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/df12ee84e1677ce1a200d1aab488c9ce5ae774d21f70b7f15b39c636623ab753/_cute_kernels.py new file mode 100644 index 00000000..5982483e --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/df12ee84e1677ce1a200d1aab488c9ce5ae774d21f70b7f15b39c636623ab753/_cute_kernels.py @@ -0,0 +1,2283 @@ +"""CuTe DSL kernels for chunked KDA prefill: prologue, prepare, recurrence. + +The prologue (one CTA) validates the packed metadata, builds the tile tables, +and zeroes the error code. The prepare kernel (one CTA per chunk tile and +head) turns raw projections into the per-tile operands of the chunked delta +rule. The recurrence kernel (one CTA per sequence, head, and value split) +walks a sequence's tiles with the state resident in registers. + +Workspace tile layout (private to these kernels): the ``[16 x 128]`` bf16 tiles +are stored with their 16-byte chunks XOR-swizzled by ``row & 7`` so the +recurrence kernel's ``ldmatrix`` reads are bank-conflict free; +:func:`workspace_tiles` returns de-swizzled views for tests. +""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +from cutlass import BFloat16, Float32, Int32, Int64, Uint32 +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import T, dsl_user_op + +from b12x._lib.compiler import KernelCompileSpec +from b12x._lib.compiler import compile as b12x_compile +from b12x._lib.intrinsics import ( + atomic_cas_global_i32, + bf16_mma_m16n8k16_f32, + cp_async_bulk_g2s_mbar, + ld_shared_v4_f32, + ld_shared_v4_u32, + ldmatrix_m8n8x4_b16, + ldmatrix_m8n8x4_trans_b16, + pack_f32x2_to_bfloat2, + shared_ptr_to_u32, + st_global_v4_u32, + st_shared_v4_f32, + warp_reduce, +) +from b12x._lib.runtime_control import raise_if_kernel_resolution_frozen +from b12x._lib.utils import current_cuda_stream, make_ptr + +from ._impl import Binding +from ._policy import WorkspaceRecord as REC + +_HEAD_DIM = 128 +_CHUNK = 16 +_PROLOGUE_THREADS = 256 +_PREPARE_THREADS = 128 +_LOG2E = 1.4426950408889634 +# Tiles of lookahead for the value-row L2 prefetch issued by the producer. +_V_PREFETCH_TILES = 6 + +_PROLOGUE_CACHE: dict[tuple, Callable[..., None]] = {} +_PREPARE_CACHE: dict[tuple, Callable[..., None]] = {} +_RECURRENCE_CACHE: dict[tuple, Callable[..., None]] = {} +_WARMED: set[tuple] = set() + + +def _add(left: Float32, right: Float32) -> Float32: + return left + right + + +@dsl_user_op +def _exp2_approx_ftz_f32(a: Float32, *, loc=None, ip=None) -> Float32: + """``ex2.approx.ftz.f32``; every argument here is at least -116.""" + return Float32( + llvm.inline_asm( + T.f32(), + [Float32(a).ir_value(loc=loc, ip=ip)], + "ex2.approx.ftz.f32 $0, $1;", + "=f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _cp_async_16_zfill(smem_addr: Int32, gmem_addr: Int64, src_bytes: Int32, *, loc=None, ip=None): + """16-byte ``cp.async.cg`` that zero-fills the bytes past ``src_bytes``.""" + llvm.inline_asm( + None, + [ + Int32(smem_addr).ir_value(loc=loc, ip=ip), + Int64(gmem_addr).ir_value(loc=loc, ip=ip), + Int32(src_bytes).ir_value(loc=loc, ip=ip), + ], + "cp.async.cg.shared.global [$0], [$1], 16, $2;", + "r,l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _ld_shared_v2_f32(smem_addr: Int32, *, loc=None, ip=None) -> tuple[Float32, Float32]: + """Load two consecutive fp32 values from a shared-memory byte address.""" + result = llvm.inline_asm( + llvm.StructType.get_literal([T.f32(), T.f32()]), + [Int32(smem_addr).ir_value(loc=loc, ip=ip)], + "ld.shared.v2.f32 {$0, $1}, [$2];", + "=f,=f,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return ( + Float32(llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)), + Float32(llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)), + ) + + +@dsl_user_op +def _bf16x2_to_f32x2(packed: Uint32, *, loc=None, ip=None) -> tuple[Float32, Float32]: + """Unpack a bf16x2 register into (low element, high element) fp32 values.""" + result = llvm.inline_asm( + llvm.StructType.get_literal([T.f32(), T.f32()]), + [Uint32(packed).ir_value(loc=loc, ip=ip)], + "{ .reg .b32 t; shl.b32 t, $2, 16; mov.b32 $0, t; and.b32 t, $2, 0xffff0000; mov.b32 $1, t; }", + "=f,=f,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return ( + Float32(llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)), + Float32(llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)), + ) + + +@dsl_user_op +def _stmatrix_x4_trans(smem_addr: Int32, r0: Uint32, r1: Uint32, r2: Uint32, r3: Uint32, *, loc=None, ip=None): + """``stmatrix.sync.aligned.m8n8.x4.trans.shared.b16`` from four fragment registers.""" + llvm.inline_asm( + None, + [ + Int32(smem_addr).ir_value(loc=loc, ip=ip), + Uint32(r0).ir_value(loc=loc, ip=ip), + Uint32(r1).ir_value(loc=loc, ip=ip), + Uint32(r2).ir_value(loc=loc, ip=ip), + Uint32(r3).ir_value(loc=loc, ip=ip), + ], + "stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 [$0], {$1, $2, $3, $4};", + "r,r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _ld_acquire_gpu_i32(gmem_addr: Int64, *, loc=None, ip=None) -> Int32: + """GPU-scope acquire load of one int32.""" + result = llvm.inline_asm( + T.i32(), + [Int64(gmem_addr).ir_value(loc=loc, ip=ip)], + "ld.acquire.gpu.global.b32 $0, [$1];", + "=r,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + return Int32(result) + + +@dsl_user_op +def _st_release_gpu_i32(gmem_addr: Int64, value: Int32, *, loc=None, ip=None): + """GPU-scope release store of one int32.""" + llvm.inline_asm( + None, + [Int64(gmem_addr).ir_value(loc=loc, ip=ip), Int32(value).ir_value(loc=loc, ip=ip)], + "st.release.gpu.global.b32 [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@dsl_user_op +def _cp_async_mbarrier_arrive_noinc(mbar_addr: Int32, *, loc=None, ip=None): + """Arrive on ``mbar`` once every prior cp.async of this thread has landed. + + The arrival is one of the barrier's expected arrivals (``.noinc``), so the + barrier must be initialized with a count that includes it. + """ + llvm.inline_asm( + None, + [Int32(mbar_addr).ir_value(loc=loc, ip=ip)], + "cp.async.mbarrier.arrive.noinc.shared::cta.b64 [$0];", + "r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@dsl_user_op +def _prefetch_l2(gmem_addr: Int64, *, loc=None, ip=None): + """Prefetch the L2 line holding ``gmem_addr``.""" + llvm.inline_asm( + None, + [Int64(gmem_addr).ir_value(loc=loc, ip=ip)], + "prefetch.global.L2 [$0];", + "l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@dsl_user_op +def _nanosleep(nanoseconds: Int32, *, loc=None, ip=None): + llvm.inline_asm( + None, + [Int32(nanoseconds).ir_value(loc=loc, ip=ip)], + "nanosleep.u32 $0;", + "r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@dsl_user_op +def _pointer_address(ptr: cute.Pointer, offset, *, loc=None, ip=None) -> Int64: + """Return the global address of ``ptr[offset]`` as an Int64.""" + element = ptr + offset + return Int64(llvm.ptrtoint(T.i64(), element.llvm_ptr, loc=loc, ip=ip)) + + +def _numeric_type(dtype: torch.dtype) -> type[cutlass.Numeric]: + if dtype == torch.bfloat16: + return BFloat16 + if dtype == torch.float32: + return Float32 + if dtype == torch.int32: + return Int32 + if dtype == torch.int64: + return Int64 + if dtype == torch.int8: + return cutlass.Int8 + raise TypeError(f"unsupported KDA prefill dtype {dtype}") + + +def _fake_pointer(dtype: type[cutlass.Numeric]) -> cute.Pointer: + return make_ptr(dtype, 16, cute.AddressSpace.gmem, assumed_align=max(1, dtype.width // 8)) + + +def _pointer(tensor: torch.Tensor, dtype: type[cutlass.Numeric]) -> cute.Pointer: + return make_ptr( + dtype, tensor.data_ptr(), cute.AddressSpace.gmem, assumed_align=max(1, dtype.width // 8) + ) + + +class _PrologueKernel: + """Validate packed metadata and build the banded tile tables in one CTA. + + Tiles are laid out in bands: band ``l`` holds local tile ``l`` of every + sequence that has one, ordered by sequence rank (longest sequence first, + ties by sequence index). Consecutive tile positions therefore interleave + the live sequences, so a window of positions advances every sequence at + once. The kernel produces ``band_base[l]`` (first position of band ``l``, + with ``band_base[l + 1] - band_base[l]`` sequences in the band), + ``sorted_seq[rank]``, ``rank_of[seq]``, and per position ``pos_seq`` and + ``pos_local`` (``pos_seq`` is -1 past the live tiles). + """ + + def __init__( + self, + *, + max_seqs: int, + tiles_capacity: int, + window_tiles: int, + max_windows: int, + table_size: int, + flag_count: int, + max_state_slots: int, + validate: bool, + max_checkpoints: int, + null_state_index: int | None, + index_type: type[cutlass.Numeric], + ) -> None: + self.max_seqs = int(max_seqs) + self.tiles_capacity = int(tiles_capacity) + self.window_tiles = int(window_tiles) + self.max_windows = int(max_windows) + self.table_size = int(table_size) + self.flag_count = int(flag_count) + self.max_state_slots = int(max_state_slots) + self.validate = bool(validate) + self.max_checkpoints = int(max_checkpoints) + self.has_null = null_state_index is not None + self.null_state_index = 0 if null_state_index is None else int(null_state_index) + self.index_type = index_type + self.seq_block = (self.max_seqs + _PROLOGUE_THREADS - 1) // _PROLOGUE_THREADS + # Band tables have tiles_capacity + 2 entries (bands 0..tiles_capacity + # plus the total). + self.band_entries = self.tiles_capacity + 2 + self.band_block = (self.band_entries + _PROLOGUE_THREADS - 1) // _PROLOGUE_THREADS + + @cute.jit + def __call__( + self, + cu_seqlens: cute.Pointer, + initial_indices: cute.Pointer, + final_indices: cute.Pointer, + checkpoint_indices: cute.Pointer, + checkpoint_offsets: cute.Pointer, + num_seqs: cute.Pointer, + num_tokens: cute.Pointer, + error_code: cute.Pointer, + table: cute.Pointer, + band_base: cute.Pointer, + sorted_seq: cute.Pointer, + rank_of: cute.Pointer, + pos_seq: cute.Pointer, + pos_local: cute.Pointer, + window_table: cute.Pointer, + ready: cute.Pointer, + final_stride: Int64, + seq_capacity: Int32, + token_capacity: Int32, + launched_tiles: Int32, + stream: cuda.CUstream, + ): + self.kernel( + cu_seqlens, initial_indices, final_indices, checkpoint_indices, checkpoint_offsets, + num_seqs, num_tokens, error_code, table, band_base, sorted_seq, rank_of, pos_seq, + pos_local, window_table, ready, final_stride, seq_capacity, token_capacity, + launched_tiles, + ).launch(grid=(1, 1, 1), block=(_PROLOGUE_THREADS, 1, 1), stream=stream) + + @cute.jit + def _is_null(self, slot: Int64) -> cutlass.Boolean: + result = slot != slot + if cutlass.const_expr(self.has_null): + result = slot == Int64(self.null_state_index) + return result + + @cute.jit + def _insert(self, table: cute.Pointer, slot: Int64) -> Int32: + """Insert ``slot`` into the open-addressing table; 1 when present.""" + key = slot.to(Int32) & Int32(0x7FFFFFFF) + stored = key + Int32(1) + position = key & Int32(self.table_size - 1) + duplicate = Int32(0) + done = Int32(0) + while done == Int32(0): + previous = atomic_cas_global_i32(_pointer_address(table, position), Int32(0), stored) + if previous == Int32(0): + done = Int32(1) + elif previous == stored: + duplicate = Int32(1) + done = Int32(1) + else: + position = (position + Int32(1)) & Int32(self.table_size - 1) + return duplicate + + @cute.jit + def _contains(self, table: cute.Pointer, slot: Int64) -> Int32: + key = slot.to(Int32) & Int32(0x7FFFFFFF) + stored = key + Int32(1) + position = key & Int32(self.table_size - 1) + found = Int32(0) + done = Int32(0) + while done == Int32(0): + current = table[position].to(Int32) + if current == Int32(0): + done = Int32(1) + elif current == stored: + found = Int32(1) + done = Int32(1) + else: + position = (position + Int32(1)) & Int32(self.table_size - 1) + return found + + @cute.jit + def _exclusive_scan( + self, values: cute.Tensor, block_sums: cute.Tensor, thread: Int32, length: Int32 + ) -> Int32: + """In-place exclusive prefix scan of ``values[0:length]``; returns the total. + + Each thread scans one contiguous block of ``band_block`` entries, the + block totals are scanned across the CTA, and the block offsets are + applied. The caller must synchronize before reading the result. + """ + running = Int32(0) + for item in cutlass.range_constexpr(self.band_block): + index = thread * Int32(self.band_block) + Int32(item) + if index < length: + current = values[index] + values[index] = running + running += current + block_sums[thread] = running + cute.arch.sync_threads() + for step in cutlass.range_constexpr(8): + distance = Int32(1 << step) + addend = Int32(0) + if thread >= distance: + addend = block_sums[thread - distance] + cute.arch.sync_threads() + block_sums[thread] = block_sums[thread] + addend + cute.arch.sync_threads() + offset = block_sums[thread] - running + for item in cutlass.range_constexpr(self.band_block): + index = thread * Int32(self.band_block) + Int32(item) + if index < length: + values[index] = values[index] + offset + total = block_sums[Int32(_PROLOGUE_THREADS - 1)] + cute.arch.sync_threads() + return total + + @cute.kernel + def kernel( + self, + cu_seqlens: cute.Pointer, + initial_indices: cute.Pointer, + final_indices: cute.Pointer, + checkpoint_indices: cute.Pointer, + checkpoint_offsets: cute.Pointer, + num_seqs: cute.Pointer, + num_tokens: cute.Pointer, + error_code: cute.Pointer, + table: cute.Pointer, + band_base: cute.Pointer, + sorted_seq: cute.Pointer, + rank_of: cute.Pointer, + pos_seq: cute.Pointer, + pos_local: cute.Pointer, + window_table: cute.Pointer, + ready: cute.Pointer, + final_stride: Int64, + seq_capacity: Int32, + token_capacity: Int32, + launched_tiles: Int32, + ): + thread, _, _ = cute.arch.thread_idx() + thread = Int32(thread) + allocator = cutlass.utils.SmemAllocator() + counts = allocator.allocate_tensor( + element_type=Int32, + layout=cute.make_layout((self.max_seqs,), stride=(1,)), + byte_alignment=16, + ) + # hist[c]: sequences with c tiles; becomes its exclusive prefix. + hist = allocator.allocate_tensor( + element_type=Int32, + layout=cute.make_layout((self.band_entries,), stride=(1,)), + byte_alignment=16, + ) + # more[l]: sequences with more than l tiles; bands becomes band_base. + more = allocator.allocate_tensor( + element_type=Int32, + layout=cute.make_layout((self.band_entries,), stride=(1,)), + byte_alignment=16, + ) + bands = allocator.allocate_tensor( + element_type=Int32, + layout=cute.make_layout((self.band_entries,), stride=(1,)), + byte_alignment=16, + ) + cursor = allocator.allocate_tensor( + element_type=Int32, + layout=cute.make_layout((self.band_entries,), stride=(1,)), + byte_alignment=16, + ) + block_sums = allocator.allocate_tensor( + element_type=Int32, + layout=cute.make_layout((_PROLOGUE_THREADS + 1,), stride=(1,)), + byte_alignment=16, + ) + flags = allocator.allocate_tensor( + element_type=Int32, + layout=cute.make_layout((8,), stride=(1,)), + byte_alignment=16, + ) + if thread < Int32(8): + flags[thread] = Int32(0) + position = thread + while position < Int32(self.table_size): + table[position] = Int32(0) + position += Int32(_PROLOGUE_THREADS) + # Ready flags of both workspace ring slots start clear every run so a + # recurrence launch can only match flags published by this run. + position = thread + while position < Int32(self.flag_count): + ready[position] = Int32(0) + position += Int32(_PROLOGUE_THREADS) + entry = thread + while entry < Int32(self.band_entries): + hist[entry] = Int32(0) + cursor[entry] = Int32(0) + entry += Int32(_PROLOGUE_THREADS) + cute.arch.sync_threads() + + live_seqs = num_seqs[Int32(0)].to(Int32) + live_tokens = num_tokens[Int32(0)].to(Int32) + bounded_seqs = cutlass.max(Int32(0), cutlass.min(live_seqs, seq_capacity)) + if cutlass.const_expr(self.validate): + if thread == Int32(0): + bad_counts = ( + (live_seqs < Int32(0)) + | (live_seqs > seq_capacity) + | (live_tokens < Int32(0)) + | (live_tokens > token_capacity) + ) + if bad_counts: + flags[1] = Int32(1) + if cu_seqlens[Int32(0)].to(Int32) != Int32(0): + flags[1] = Int32(1) + if cu_seqlens[bounded_seqs].to(Int32) != live_tokens: + flags[1] = Int32(1) + + # Per-sequence pass: tile counts, slot checks, write-slot insertion, + # and the tile-count histogram. + seq = thread + while seq < Int32(self.max_seqs): + count = Int32(0) + if seq < bounded_seqs: + start = cu_seqlens[seq].to(Int32) + end = cu_seqlens[seq + Int32(1)].to(Int32) + length = cutlass.max(Int32(0), end - start) + count = cutlass.min((length + Int32(_CHUNK - 1)) // Int32(_CHUNK), Int32(self.tiles_capacity)) + if cutlass.const_expr(self.validate): + if (start < Int32(0)) | (end < start) | (end > live_tokens): + flags[1] = Int32(1) + initial = Int64(initial_indices[seq]) + final = Int64(final_indices[seq.to(Int64) * final_stride]) + slot_limit = Int64(self.max_state_slots) + if not self._is_null(initial): + if (initial < Int64(0)) | (initial >= slot_limit): + flags[2] = Int32(1) + if not self._is_null(final): + if (final < Int64(0)) | (final >= slot_limit): + flags[2] = Int32(1) + elif self._insert(table, final) != Int32(0): + flags[0] = Int32(1) + for cp in cutlass.range_constexpr(self.max_checkpoints): + cp_index = seq * Int32(self.max_checkpoints) + Int32(cp) + checkpoint = Int64(checkpoint_indices[cp_index]) + offset = checkpoint_offsets[cp_index].to(Int32) + if offset > length: + flags[3] = Int32(1) + if (offset > Int32(0)) & ((offset % Int32(_CHUNK)) != Int32(0)): + flags[3] = Int32(1) + if offset > Int32(0): + if not self._is_null(checkpoint): + if (checkpoint < Int64(0)) | (checkpoint >= slot_limit): + flags[2] = Int32(1) + elif self._insert(table, checkpoint) != Int32(0): + flags[0] = Int32(1) + if cutlass.const_expr(cp > 0): + previous_index = cp_index - Int32(1) + previous_slot = Int64(checkpoint_indices[previous_index]) + previous_offset = checkpoint_offsets[previous_index].to(Int32) + if not self._is_null(previous_slot): + if previous_offset == offset: + flags[3] = Int32(1) + cute.arch.atomic_add(hist.iterator + count, Int32(1)) + counts[seq] = count + seq += Int32(_PROLOGUE_THREADS) + cute.arch.sync_threads() + + # more[l] = sequences with more than l tiles = live - prefix(hist)[l + 1]. + self._exclusive_scan(hist, block_sums, thread, Int32(self.tiles_capacity + 1)) + entry = thread + while entry < Int32(self.band_entries): + value = Int32(0) + if entry < Int32(self.tiles_capacity): + value = bounded_seqs - hist[entry + Int32(1)] + more[entry] = value + bands[entry] = value + entry += Int32(_PROLOGUE_THREADS) + cute.arch.sync_threads() + # bands[l] = first position of band l; bands[tiles_capacity + 1] = total. + total_tiles = self._exclusive_scan(bands, block_sums, thread, Int32(self.band_entries)) + entry = thread + while entry < Int32(self.band_entries): + band_base[entry] = bands[entry] + entry += Int32(_PROLOGUE_THREADS) + + # Ranks: sequences with more tiles first, ties in arrival order. + seq = thread + while seq < bounded_seqs: + count = counts[seq] + rank = more[count] + cute.arch.atomic_add(cursor.iterator + count, Int32(1)) + rank_of[seq] = rank + sorted_seq[rank] = seq + seq += Int32(_PROLOGUE_THREADS) + if thread == Int32(0): + if cutlass.const_expr(self.validate): + if total_tiles > launched_tiles: + flags[1] = Int32(1) + cute.arch.sync_threads() + + # Per window: the band holding its first position and that position's + # rank within the band. + if thread < Int32(self.max_windows): + window_begin = thread * Int32(self.window_tiles) + low = Int32(0) + high = Int32(self.tiles_capacity) + while high - low > Int32(1): + mid = (low + high) >> Int32(1) + if bands[mid] <= window_begin: + low = mid + else: + high = mid + window_table[thread * Int32(2)] = low + window_table[thread * Int32(2) + Int32(1)] = window_begin - bands[low] + # Position tables (binary search of the band) and the unused tail. + bounded_tiles = cutlass.min(total_tiles, Int32(self.tiles_capacity)) + tile = thread + while tile < bounded_tiles: + low = Int32(0) + high = Int32(self.tiles_capacity) + while high - low > Int32(1): + mid = (low + high) >> Int32(1) + if bands[mid] <= tile: + low = mid + else: + high = mid + pos_seq[tile] = sorted_seq[tile - bands[low]] + pos_local[tile] = low + tile += Int32(_PROLOGUE_THREADS) + tile = bounded_tiles + thread + while tile < Int32(self.tiles_capacity): + pos_seq[tile] = Int32(-1) + tile += Int32(_PROLOGUE_THREADS) + # Initial-slot conflicts and the running-state slot of sequences that + # span pipeline windows. + seq = thread + while seq < bounded_seqs: + if cutlass.const_expr(self.validate): + initial = Int64(initial_indices[seq]) + final = Int64(final_indices[seq.to(Int64) * final_stride]) + if not self._is_null(initial): + if (initial >= Int64(0)) & (initial < Int64(self.max_state_slots)): + if initial != final: + if self._contains(table, initial) != Int32(0): + flags[0] = Int32(1) + count = counts[seq] + if count > Int32(0): + rank = rank_of[seq].to(Int32) + first_window = rank // Int32(self.window_tiles) + last_window = (bands[count - Int32(1)] + rank) // Int32(self.window_tiles) + if first_window != last_window: + if self._is_null(final): + flags[2] = Int32(1) + seq += Int32(_PROLOGUE_THREADS) + cute.arch.sync_threads() + # Always publish the code, including the zero of a trusted run, so a + # run never inherits a stale or uninitialized word from the scratch. + if thread == Int32(0): + code = Int32(0) + if cutlass.const_expr(self.validate): + code = flags[0] | (flags[1] << Int32(1)) | (flags[2] << Int32(2)) | (flags[3] << Int32(3)) + error_code[Int32(0)] = code + + +class _PrepareKernel: + """Per (tile, head): gates, norms, decayed operands, WY inverse, Mqk.""" + + def __init__( + self, + *, + heads: int, + tiles_capacity: int, + window_tiles: int, + qk_l2norm: bool, + a_log_type: type[cutlass.Numeric], + dt_bias_type: type[cutlass.Numeric], + ) -> None: + self.heads = int(heads) + self.tiles_capacity = int(tiles_capacity) + self.window_tiles = int(window_tiles) + self.qk_l2norm = bool(qk_l2norm) + self.a_log_type = a_log_type + self.dt_bias_type = dt_bias_type + + @cute.jit + def __call__( + self, + q: cute.Pointer, + k: cute.Pointer, + raw_g: cute.Pointer, + raw_beta: cute.Pointer, + A_log: cute.Pointer, + dt_bias: cute.Pointer, + cu_seqlens: cute.Pointer, + pos_seq: cute.Pointer, + pos_local: cute.Pointer, + error_code: cute.Pointer, + ready: cute.Pointer, + ws_bf16: cute.Pointer, + ws_f32: cute.Pointer, + q_stride: Int64, + k_stride: Int64, + g_stride: Int64, + beta_token_stride: Int64, + beta_head_stride: Int64, + scale: Float32, + gate_scale: Float32, + eps: Float32, + window: Int32, + stream: cuda.CUstream, + ): + self.kernel( + q, k, raw_g, raw_beta, A_log, dt_bias, cu_seqlens, pos_seq, pos_local, + error_code, ready, ws_bf16, ws_f32, q_stride, k_stride, g_stride, + beta_token_stride, beta_head_stride, scale, gate_scale, eps, window, + ).launch( + grid=(self.window_tiles, self.heads, 1), + block=(_PREPARE_THREADS, 1, 1), + stream=stream, + ) + + @cute.kernel + def kernel( + self, + q: cute.Pointer, + k: cute.Pointer, + raw_g: cute.Pointer, + raw_beta: cute.Pointer, + A_log: cute.Pointer, + dt_bias: cute.Pointer, + cu_seqlens: cute.Pointer, + pos_seq: cute.Pointer, + pos_local: cute.Pointer, + error_code: cute.Pointer, + ready: cute.Pointer, + ws_bf16: cute.Pointer, + ws_f32: cute.Pointer, + q_stride: Int64, + k_stride: Int64, + g_stride: Int64, + beta_token_stride: Int64, + beta_head_stride: Int64, + scale: Float32, + gate_scale: Float32, + eps: Float32, + window: Int32, + ): + local_tile, head, _ = cute.arch.block_idx() + thread, _, _ = cute.arch.thread_idx() + local_tile = Int32(local_tile) + tile = window * Int32(self.window_tiles) + local_tile + head = Int32(head) + column = Int32(thread) + warp = column // Int32(32) + lane = Int32(cute.arch.lane_idx()) + error = error_code[Int32(0)].to(Int32) + seq = Int32(-1) + local = Int32(0) + if tile < Int32(self.tiles_capacity): + seq = pos_seq[tile].to(Int32) + local = pos_local[tile].to(Int32) + if (error == Int32(0)) & (seq >= Int32(0)): + allocator = cutlass.utils.SmemAllocator() + tile_elements = _CHUNK * _HEAD_DIM + s_part = allocator.allocate_tensor( + element_type=Float32, + layout=cute.make_layout((2 * _CHUNK,), stride=(1,)), + byte_alignment=16, + ) + s_beta = allocator.allocate_tensor( + element_type=Float32, + layout=cute.make_layout((_CHUNK,), stride=(1,)), + byte_alignment=16, + ) + # Raw q, k, g rows. Once the gate loop has consumed them, their + # regions hold the swizzled q~ and k~ tiles and the transposed + # k_inv tile for the tensor-core products. + s_raw = allocator.allocate_tensor( + element_type=BFloat16, + layout=cute.make_layout((3 * tile_elements,), stride=(1,)), + byte_alignment=128, + ) + s_squares = allocator.allocate_tensor( + element_type=Float32, + layout=cute.make_layout((4 * _CHUNK * _CHUNK,), stride=(1,)), + byte_alignment=16, + ) + square_layout = cute.make_layout((_CHUNK * _CHUNK,), stride=(1,)) + squares = s_squares.iterator + s_p = cute.make_tensor(squares, square_layout) + s_p2 = cute.make_tensor(squares + _CHUNK * _CHUNK, square_layout) + s_inv = cute.make_tensor(squares + 2 * _CHUNK * _CHUNK, square_layout) + s_inv2 = cute.make_tensor(squares + 3 * _CHUNK * _CHUNK, square_layout) + + start = cu_seqlens[seq].to(Int32) + local * Int32(_CHUNK) + end = cu_seqlens[seq + Int32(1)].to(Int32) + rows = cutlass.min(Int32(_CHUNK), end - start) + head_elements = head.to(Int64) * Int64(_HEAD_DIM) + + # Stage the raw rows: 16 rows x 16 chunks of 16 bytes per tensor, + # zero-filled past the sequence tail. + raw_addr = shared_ptr_to_u32(s_raw.iterator) + for item in cutlass.range_constexpr(6): + chunk = column + Int32(item * _PREPARE_THREADS) + tensor_index = chunk // Int32(256) + local_chunk = chunk % Int32(256) + row = local_chunk // Int32(16) + col_chunk = local_chunk % Int32(16) + live_row = cutlass.min(row, cutlass.max(rows - Int32(1), Int32(0))) + token = (start + live_row).to(Int64) + src_bytes = Int32(16) + if row >= rows: + src_bytes = Int32(0) + col_elements = (col_chunk * Int32(8)).to(Int64) + if tensor_index == Int32(0): + _cp_async_16_zfill( + raw_addr + chunk * Int32(16), + _pointer_address(q, token * q_stride + head_elements + col_elements), + src_bytes, + ) + elif tensor_index == Int32(1): + _cp_async_16_zfill( + raw_addr + chunk * Int32(16), + _pointer_address(k, token * k_stride + head_elements + col_elements), + src_bytes, + ) + else: + _cp_async_16_zfill( + raw_addr + chunk * Int32(16), + _pointer_address(raw_g, token * g_stride + head_elements + col_elements), + src_bytes, + ) + cute.arch.cp_async_commit_group() + # Parameter and beta loads overlap the staging copies. + rate = cute.math.exp(Float32(A_log[head]), fastmath=False) + bias = Float32(dt_bias[head * Int32(_HEAD_DIM) + column]) + beta_raw = Float32(0.0) + if column < rows: + beta_offset = ( + (start + column).to(Int64) * beta_token_stride + + head.to(Int64) * beta_head_stride + ) + beta_raw = Float32(raw_beta[beta_offset]) + cute.arch.cp_async_wait_group(0) + cute.arch.sync_threads() + + # Row sums of squares: eight lanes per row, sixteen strided elements each. + sum_row = column >> Int32(3) + sum_part = column & Int32(7) + q_sq = Float32(0.0) + k_sq = Float32(0.0) + for item in cutlass.range_constexpr(_HEAD_DIM // 8): + element = sum_row * Int32(_HEAD_DIM) + Int32(item * 8) + sum_part + q_value = Float32(s_raw[element]) + k_value = Float32(s_raw[Int32(tile_elements) + element]) + q_sq += q_value * q_value + k_sq += k_value * k_value + q_sq = warp_reduce(q_sq, _add, 8) + k_sq = warp_reduce(k_sq, _add, 8) + if sum_part == Int32(0): + s_part[sum_row] = q_sq + s_part[Int32(_CHUNK) + sum_row] = k_sq + + q_values = cute.make_rmem_tensor((_CHUNK,), Float32) + k_values = cute.make_rmem_tensor((_CHUNK,), Float32) + g_cum = cute.make_rmem_tensor((_CHUNK,), Float32) + running = Float32(0.0) + for t in cutlass.range_constexpr(_CHUNK): + q_value = Float32(0.0) + k_value = Float32(0.0) + g2 = Float32(0.0) + if Int32(t) < rows: + q_value = Float32(s_raw[Int32(t * _HEAD_DIM) + column]) + k_value = Float32(s_raw[Int32(tile_elements + t * _HEAD_DIM) + column]) + g_value = Float32(s_raw[Int32(2 * tile_elements + t * _HEAD_DIM) + column]) + z = rate * (g_value + bias) + sigmoid = cute.arch.rcp_approx( + Float32(1.0) + _exp2_approx_ftz_f32(-z * Float32(_LOG2E)) + ) + g2 = gate_scale * sigmoid + q_values[t] = q_value + k_values[t] = k_value + running += g2 + g_cum[t] = running + if column < Int32(_CHUNK): + beta = Float32(0.0) + if column < rows: + beta = cute.arch.rcp_approx( + Float32(1.0) + _exp2_approx_ftz_f32(-beta_raw * Float32(_LOG2E)) + ) + s_beta[column] = beta + cute.arch.sync_threads() + + ring_index = (window & Int32(1)) * Int32(self.window_tiles) + local_tile + record = ring_index.to(Int64) * Int64(self.heads) + head.to(Int64) + rec_bf16 = record * Int64(REC.BYTES // 2) + rec_f32 = record * Int64(REC.BYTES // 4) + q_base = rec_bf16 + Int64(REC.Q_TILDE // 2) + k_base = rec_bf16 + Int64(REC.K_TILDE // 2) + kr_base = rec_bf16 + Int64(REC.K_R // 2) + last = g_cum[_CHUNK - 1] + lambda_c = _exp2_approx_ftz_f32(last) + ws_f32[rec_f32 + Int64(REC.LAMBDA_C // 4) + column.to(Int64)] = lambda_c + if column < Int32(_CHUNK): + ws_f32[rec_f32 + Int64(REC.BETA // 4) + column.to(Int64)] = s_beta[column] + for t in cutlass.range_constexpr(_CHUNK): + rinv_q = Float32(1.0) + rinv_k = Float32(1.0) + if cutlass.const_expr(self.qk_l2norm): + rinv_q = cute.math.rsqrt(s_part[Int32(t)] + eps, fastmath=False) + rinv_k = cute.math.rsqrt(s_part[Int32(_CHUNK + t)] + eps, fastmath=False) + lam = _exp2_approx_ftz_f32(g_cum[t]) + lam_inv = _exp2_approx_ftz_f32(-g_cum[t]) + lam_r = _exp2_approx_ftz_f32(last - g_cum[t]) + q_tilde = BFloat16(q_values[t] * rinv_q * lam * scale) + k_tilde = BFloat16(k_values[t] * rinv_k * lam) + k_inv = BFloat16(k_values[t] * rinv_k * lam_inv) + k_r = BFloat16(k_values[t] * rinv_k * lam_r) + physical = ( + Int32(t * _HEAD_DIM) + + ((((column >> Int32(3)) ^ Int32(t & 7)) << Int32(3)) | (column & Int32(7))) + ) + s_raw[physical] = q_tilde + s_raw[Int32(tile_elements) + physical] = k_tilde + s_raw[Int32(2 * tile_elements) + column * Int32(_CHUNK) + Int32(t)] = k_inv + ws_bf16[q_base + physical.to(Int64)] = q_tilde + ws_bf16[k_base + physical.to(Int64)] = k_tilde + ws_bf16[kr_base + physical.to(Int64)] = k_r + cute.arch.sync_threads() + + # L = beta_i (j < i) on warp 0 and Mqk = + # (j <= i) on warp 1, each one m16n16k128 tensor-core product. + if warp < Int32(2): + a_base = raw_addr + Int32(tile_elements * 2) + if warp == Int32(1): + a_base = raw_addr + b_base = raw_addr + Int32(2 * tile_elements * 2) + gid = lane >> Int32(2) + tid = lane & Int32(3) + matrix = lane >> Int32(3) + matrix_row = lane & Int32(7) + a_row = (matrix & Int32(1)) * Int32(8) + matrix_row + prod = cute.make_rmem_tensor((2, 4), Float32) + for half in cutlass.range_constexpr(2): + for item in cutlass.range_constexpr(4): + prod[half, item] = Float32(0.0) + for kb in cutlass.range_constexpr(8): + a_chunk = (Int32(kb * 2) + (matrix >> Int32(1))) ^ (a_row & Int32(7)) + a0, a1, a2, a3 = ldmatrix_m8n8x4_b16(a_base + a_row * Int32(256) + a_chunk * Int32(16)) + b_row = Int32(kb * 16) + (matrix & Int32(1)) * Int32(8) + matrix_row + b0, b1, b2, b3 = ldmatrix_m8n8x4_trans_b16( + b_base + b_row * Int32(32) + (matrix >> Int32(1)) * Int32(16) + ) + prod[0, 0], prod[0, 1], prod[0, 2], prod[0, 3] = bf16_mma_m16n8k16_f32( + prod[0, 0], prod[0, 1], prod[0, 2], prod[0, 3], a0, a1, a2, a3, b0, b1 + ) + prod[1, 0], prod[1, 1], prod[1, 2], prod[1, 3] = bf16_mma_m16n8k16_f32( + prod[1, 0], prod[1, 1], prod[1, 2], prod[1, 3], a0, a1, a2, a3, b2, b3 + ) + for half in cutlass.range_constexpr(2): + for item in cutlass.range_constexpr(4): + row_i = gid + Int32((item >> 1) * 8) + col_j = Int32(half * 8) + tid * Int32(2) + Int32(item & 1) + index = row_i * Int32(_CHUNK) + col_j + value = prod[half, item] + if warp == Int32(0): + lower = Float32(0.0) + if col_j < row_i: + lower = s_beta[row_i] * value + identity = Float32(0.0) + if col_j == row_i: + identity = Float32(1.0) + s_p[index] = lower + s_inv[index] = identity - lower + else: + mqk = Float32(0.0) + if col_j <= row_i: + mqk = value + ws_bf16[rec_bf16 + Int64(REC.MQK // 2) + index.to(Int64)] = BFloat16(mqk) + cute.arch.sync_threads() + + # Neumann series: INV = (I - L)(I + L^2)(I + L^4)(I + L^8). + for _step in cutlass.range_constexpr(3): + for entry in cutlass.range_constexpr(2): + index = column + Int32(entry * _PREPARE_THREADS) + row = index // Int32(_CHUNK) + col = index % Int32(_CHUNK) + acc = Float32(0.0) + for j in cutlass.range_constexpr(_CHUNK): + acc += s_p[row * Int32(_CHUNK) + Int32(j)] * s_p[Int32(j * _CHUNK) + col] + s_p2[index] = acc + cute.arch.sync_threads() + for entry in cutlass.range_constexpr(2): + index = column + Int32(entry * _PREPARE_THREADS) + row = index // Int32(_CHUNK) + col = index % Int32(_CHUNK) + acc = s_inv[index] + for j in cutlass.range_constexpr(_CHUNK): + acc += s_inv[row * Int32(_CHUNK) + Int32(j)] * s_p2[Int32(j * _CHUNK) + col] + s_inv2[index] = acc + cute.arch.sync_threads() + for entry in cutlass.range_constexpr(2): + index = column + Int32(entry * _PREPARE_THREADS) + s_p[index] = s_p2[index] + s_inv[index] = s_inv2[index] + cute.arch.sync_threads() + for entry in cutlass.range_constexpr(2): + index = column + Int32(entry * _PREPARE_THREADS) + ws_bf16[rec_bf16 + Int64(REC.INV // 2) + index.to(Int64)] = BFloat16(s_inv[index]) + # Publish the record: every thread's stores are ordered before the + # barrier, and the fence makes them visible at GPU scope before + # the release store of the flag. + cute.arch.sync_threads() + if column == Int32(0): + cute.arch.fence_acq_rel_gpu() + _st_release_gpu_i32(_pointer_address(ready, record), window + Int32(1)) + + +class _RecurrenceKernel: + """Per (window, sequence, head, value split): walk the tiles with the state in registers. + + A CTA holds ``v_split`` rows of the transposed state for one head. Its MMA + warps form ``v_split // 16`` row groups of ``k_split`` warps; each warp owns + ``128 // k_split`` key columns of its group's sixteen rows as fp32 m16n8 + accumulator fragments plus a bf16 shadow packed as k16 A fragments. Every + per-tile product reuses the previous accumulator as the next A operand; the + two products that contract over the key axis are reduced across the + group's warps through shared memory when ``k_split > 1``. One producer warp + streams the prepared tiles through a ``stages``-deep shared-memory ring + with bulk async copies and mbarriers, so the MMA warps never wait on loads + or on a CTA-wide barrier inside the tile loop. + + One launch covers one window of ``window_tiles`` consecutive banded tile + positions (see the prologue), so every live sequence advances in every + window. CTA row ``by`` maps to one sequence rank with positions in the + window; a sequence that began in an earlier window resumes from its final + state slot, and one that continues past the window leaves its running + state there. Empty sequences copy initial to final in window 0. The + producer polls the per-tile ready flags of the prepare kernel, so the two + kernels of a window overlap. + """ + + def __init__( + self, + *, + heads: int, + tiles_capacity: int, + window_tiles: int, + rows: int, + v_split: int, + k_split: int, + stages: int, + checkpoint_export: bool, + max_checkpoints: int, + null_state_index: int | None, + index_type: type[cutlass.Numeric], + ) -> None: + self.heads = int(heads) + self.tiles_capacity = int(tiles_capacity) + self.window_tiles = int(window_tiles) + self.rows = int(rows) + self.v_split = int(v_split) + self.k_split = int(k_split) + self.stages = int(stages) + if self.stages < 2: + raise ValueError("the recurrence pipeline needs at least two stages") + if self.k_split not in (1, 2, 4): + raise ValueError("k_split must be 1, 2, or 4") + self.splits = _HEAD_DIM // self.v_split + self.row_groups = self.v_split // 16 + self.mma_warps = self.row_groups * self.k_split + self.mma_threads = 32 * self.mma_warps + self.threads = self.mma_threads + 32 + self.cols = _HEAD_DIM // self.k_split + self.kb_steps = self.cols // 16 + self.nb_blocks = self.cols // 8 + self.checkpoint_export = bool(checkpoint_export) + self.max_checkpoints = int(max_checkpoints) + self.has_null = null_state_index is not None + self.null_state_index = 0 if null_state_index is None else int(null_state_index) + self.index_type = index_type + self.v_chunks_per_row = self.v_split // 8 + + @cute.jit + def __call__( + self, + v: cute.Pointer, + cu_seqlens: cute.Pointer, + band_base: cute.Pointer, + sorted_seq: cute.Pointer, + window_table: cute.Pointer, + initial_indices: cute.Pointer, + final_indices: cute.Pointer, + checkpoint_indices: cute.Pointer, + checkpoint_offsets: cute.Pointer, + num_seqs: cute.Pointer, + error_code: cute.Pointer, + ready: cute.Pointer, + ws: cute.Pointer, + recurrent_state: cute.Pointer, + output: cute.Pointer, + v_stride: Int64, + out_stride: Int64, + slot_stride: Int64, + final_stride: Int64, + token_capacity: Int32, + window: Int32, + stream: cuda.CUstream, + ): + self.kernel( + v, cu_seqlens, band_base, sorted_seq, window_table, initial_indices, final_indices, + checkpoint_indices, checkpoint_offsets, num_seqs, error_code, ready, ws, + recurrent_state, output, v_stride, out_stride, slot_stride, final_stride, + token_capacity, window, + ).launch( + grid=(self.heads * self.splits, self.rows, 1), + block=(self.threads, 1, 1), + stream=stream, + ) + + @cute.jit + def _is_null(self, slot: Int64) -> cutlass.Boolean: + result = slot != slot + if cutlass.const_expr(self.has_null): + result = slot == Int64(self.null_state_index) + return result + + @cute.jit + def _store_state( + self, + target: cute.Pointer, + acc: cute.Tensor, + base: Int64, + row0: Int32, + row1: Int32, + col_base: Int32, + tid: Int32, + ): + for nb in cutlass.range_constexpr(self.nb_blocks): + kcol = col_base + Int32(nb * 8) + tid * Int32(2) + offset0 = base + row0.to(Int64) * Int64(_HEAD_DIM) + kcol.to(Int64) + offset1 = base + row1.to(Int64) * Int64(_HEAD_DIM) + kcol.to(Int64) + target[offset0] = acc[nb, 0] + target[offset0 + Int64(1)] = acc[nb, 1] + target[offset1] = acc[nb, 2] + target[offset1 + Int64(1)] = acc[nb, 3] + + @cute.jit + def _load_state( + self, + source: cute.Pointer, + acc: cute.Tensor, + base: Int64, + row0: Int32, + row1: Int32, + col_base: Int32, + tid: Int32, + ): + for nb in cutlass.range_constexpr(self.nb_blocks): + kcol = col_base + Int32(nb * 8) + tid * Int32(2) + offset0 = base + row0.to(Int64) * Int64(_HEAD_DIM) + kcol.to(Int64) + offset1 = base + row1.to(Int64) * Int64(_HEAD_DIM) + kcol.to(Int64) + acc[nb, 0] = Float32(source[offset0]) + acc[nb, 1] = Float32(source[offset0 + Int64(1)]) + acc[nb, 2] = Float32(source[offset1]) + acc[nb, 3] = Float32(source[offset1 + Int64(1)]) + + @cute.jit + def _refresh_shadow(self, acc: cute.Tensor, shadow: cute.Tensor): + for kb in cutlass.range_constexpr(self.kb_steps): + shadow[kb, 0] = pack_f32x2_to_bfloat2(acc[2 * kb, 0], acc[2 * kb, 1]) + shadow[kb, 1] = pack_f32x2_to_bfloat2(acc[2 * kb, 2], acc[2 * kb, 3]) + shadow[kb, 2] = pack_f32x2_to_bfloat2(acc[2 * kb + 1, 0], acc[2 * kb + 1, 1]) + shadow[kb, 3] = pack_f32x2_to_bfloat2(acc[2 * kb + 1, 2], acc[2 * kb + 1, 3]) + + @cute.jit + def _issue_tile( + self, + ws: cute.Pointer, + v: cute.Pointer, + stage_addr: Int32, + ring_index: Int32, + head: Int32, + token_base: Int32, + rows_live: Int32, + head_elements: Int64, + split_elements: Int64, + v_stride: Int64, + full_bar_u32: Int32, + full_bar_ptr: cute.Pointer, + lane: Int32, + ): + """Issue the copies of ring record ``ring_index`` and its value rows into a stage. + + Every lane copies its share of the ``[16 x v_split]`` value rows with + 16-byte cp.async (zero-filled past the live rows) into the stage's + sixteen-column-group layout and arrives on the stage barrier when + they land; lane 0 then posts the record head's byte count and copies + it with one bulk copy. The barrier expects the 32 lane arrivals plus + lane 0's, so the phase cannot complete before every value row landed. + """ + vd = self.v_split + record = ring_index.to(Int64) * Int64(self.heads) + head.to(Int64) + record_base = record * Int64(REC.BYTES) + chunks_per_row = vd // 8 + v_addr = stage_addr + Int32(REC.V) + for item in cutlass.range_constexpr((_CHUNK * chunks_per_row) // 32): + chunk = lane + Int32(item * 32) + row = chunk // Int32(chunks_per_row) + col8 = chunk % Int32(chunks_per_row) + src_bytes = Int32(16) + if row >= rows_live: + src_bytes = Int32(0) + live_row = cutlass.min(row, cutlass.max(rows_live - Int32(1), Int32(0))) + element = (token_base + live_row).to(Int64) * v_stride + head_elements + split_elements + (col8 * Int32(8)).to(Int64) + dst = v_addr + (col8 >> Int32(1)) * Int32(512) + row * Int32(32) + (col8 & Int32(1)) * Int32(16) + _cp_async_16_zfill(dst, _pointer_address(v, element), src_bytes) + _cp_async_mbarrier_arrive_noinc(full_bar_u32) + if lane == Int32(0): + cute.arch.fence_acq_rel_cta() + cute.arch.mbarrier_arrive_and_expect_tx(full_bar_ptr, Int32(REC.HEAD_BYTES)) + # Orders the async-proxy copy after the acquire of the ready flag. + cute.arch.fence_proxy("async.global") + cp_async_bulk_g2s_mbar( + stage_addr, _pointer_address(ws, record_base), Int32(REC.HEAD_BYTES), full_bar_u32 + ) + + @cute.jit + def _group_reduce( + self, + partial: cute.Tensor, + red_addr: Int32, + group: Int32, + kq: Int32, + lane: Int32, + ): + """Sum a [16 x 16] fp32 fragment across the group's warps.""" + slot = (group * Int32(self.k_split) + kq) * Int32(256) + lane * Int32(8) + mine = red_addr + slot * Int32(4) + st_shared_v4_f32(mine, partial[0, 0], partial[0, 1], partial[0, 2], partial[0, 3]) + st_shared_v4_f32(mine + Int32(16), partial[1, 0], partial[1, 1], partial[1, 2], partial[1, 3]) + cute.arch.barrier(barrier_id=1, number_of_threads=self.mma_threads) + for other in cutlass.range_constexpr(self.k_split): + if Int32(other) != kq: + theirs = red_addr + ((group * Int32(self.k_split) + Int32(other)) * Int32(256) + lane * Int32(8)) * Int32(4) + a0, a1, a2, a3 = ld_shared_v4_f32(theirs) + b0, b1, b2, b3 = ld_shared_v4_f32(theirs + Int32(16)) + partial[0, 0] += a0 + partial[0, 1] += a1 + partial[0, 2] += a2 + partial[0, 3] += a3 + partial[1, 0] += b0 + partial[1, 1] += b1 + partial[1, 2] += b2 + partial[1, 3] += b3 + + @cute.kernel + def kernel( + self, + v: cute.Pointer, + cu_seqlens: cute.Pointer, + band_base: cute.Pointer, + sorted_seq: cute.Pointer, + window_table: cute.Pointer, + initial_indices: cute.Pointer, + final_indices: cute.Pointer, + checkpoint_indices: cute.Pointer, + checkpoint_offsets: cute.Pointer, + num_seqs: cute.Pointer, + error_code: cute.Pointer, + ready: cute.Pointer, + ws: cute.Pointer, + recurrent_state: cute.Pointer, + output: cute.Pointer, + v_stride: Int64, + out_stride: Int64, + slot_stride: Int64, + final_stride: Int64, + token_capacity: Int32, + window: Int32, + ): + bx, by, _ = cute.arch.block_idx() + thread, _, _ = cute.arch.thread_idx() + thread = Int32(thread) + head = Int32(bx) // Int32(self.splits) + split = Int32(bx) % Int32(self.splits) + warp = thread // Int32(32) + lane = Int32(cute.arch.lane_idx()) + gid = lane >> Int32(2) + tid = lane & Int32(3) + matrix = lane >> Int32(3) + matrix_row = lane & Int32(7) + is_producer = warp == Int32(self.mma_warps) + group = cutlass.min(warp, Int32(self.mma_warps - 1)) // Int32(self.k_split) + kq = cutlass.min(warp, Int32(self.mma_warps - 1)) % Int32(self.k_split) + col_base = kq * Int32(self.cols) + chunk_base = kq * Int32(self.nb_blocks) + group_lead = kq == Int32(0) + if cutlass.const_expr(self.k_split == 1): + group_lead = cutlass.Boolean(True) + error = error_code[Int32(0)].to(Int32) + live_seqs = num_seqs[Int32(0)].to(Int32) + vd = self.v_split + if error != Int32(0): + # Transactional failure: poison every output row of this CTA's + # value columns; no state is written. + if (window == Int32(0)) & (Int32(by) == Int32(0)): + nan_pair = Uint32(0x7FC07FC0) + chunk = thread + while chunk < token_capacity * Int32(self.v_chunks_per_row): + row = chunk // Int32(self.v_chunks_per_row) + col_chunk = chunk % Int32(self.v_chunks_per_row) + element = ( + row.to(Int64) * out_stride + + head.to(Int64) * Int64(_HEAD_DIM) + + (split * Int32(vd) + col_chunk * Int32(8)).to(Int64) + ) + st_global_v4_u32( + _pointer_address(output, element), nan_pair, nan_pair, nan_pair, nan_pair + ) + chunk += Int32(self.threads) + else: + window_begin = window * Int32(self.window_tiles) + window_end = window_begin + Int32(self.window_tiles) + total_tiles = band_base[Int32(self.tiles_capacity + 1)].to(Int32) + tiled_seqs = band_base[Int32(1)].to(Int32) + # The band holding the window's first position and the rank of + # that position within it (from the prologue's window table). + # Rows first cover the ranks from there to the end of that band, + # then the ranks before it, which reach the window in the next + # band. + band_a = window_table[window * Int32(2)].to(Int32) + first_rank = window_table[window * Int32(2) + Int32(1)].to(Int32) + seg1 = band_base[band_a + Int32(1)].to(Int32) - window_begin + if window_begin >= total_tiles: + first_rank = Int32(0) + seg1 = Int32(0) + rank = Int32(-1) + l_begin = Int32(0) + if Int32(by) < seg1: + rank = first_rank + Int32(by) + l_begin = band_a + elif Int32(by) - seg1 < first_rank: + rank = Int32(by) - seg1 + l_begin = band_a + Int32(1) + seq = Int32(0) + start = Int32(0) + end = Int32(0) + tiles_s = Int32(0) + l_end = Int32(0) + if rank >= Int32(0): + seq = sorted_seq[rank].to(Int32) + start = cu_seqlens[seq].to(Int32) + end = cu_seqlens[seq + Int32(1)].to(Int32) + tiles_s = cutlass.min( + (cutlass.max(Int32(0), end - start) + Int32(_CHUNK - 1)) // Int32(_CHUNK), + Int32(self.tiles_capacity), + ) + l_begin = cutlass.min(l_begin, tiles_s) + stage_bytes = REC.V + vd * 32 + allocator = cutlass.utils.SmemAllocator() + # Band positions of this CTA's local tiles l_begin .. l_begin + WT. + s_band = allocator.allocate_tensor( + element_type=Int32, + layout=cute.make_layout((self.window_tiles + 1,), stride=(1,)), + byte_alignment=16, + ) + s_stage = allocator.allocate_tensor( + element_type=cutlass.Int8, + layout=cute.make_layout((self.stages * stage_bytes,), stride=(1,)), + byte_alignment=128, + ) + s_out = allocator.allocate_tensor( + element_type=BFloat16, + layout=cute.make_layout((_CHUNK * vd,), stride=(1,)), + byte_alignment=128, + ) + # Cross-warp reduction buffers, only with a key split. + s_red = allocator.allocate_tensor( + element_type=Float32, + layout=cute.make_layout((2 * self.mma_warps * 256 if self.k_split > 1 else 4,), stride=(1,)), + byte_alignment=128, + ) + mbar = allocator.allocate_tensor( + element_type=Int64, + layout=cute.make_layout((2 * self.stages,), stride=(1,)), + byte_alignment=8, + ) + stage_base = shared_ptr_to_u32(s_stage.iterator) + out_addr = shared_ptr_to_u32(s_out.iterator) + red_base = shared_ptr_to_u32(s_red.iterator) + red_a = red_base + red_c = red_base + Int32(self.mma_warps * 256 * 4) + full_bar = mbar.iterator + empty_bar = mbar.iterator + self.stages + if thread == Int32(0): + for stage_index in cutlass.range_constexpr(self.stages): + cute.arch.mbarrier_init(full_bar + stage_index, Int32(33)) + cute.arch.mbarrier_init(empty_bar + stage_index, Int32(self.mma_warps)) + if rank >= Int32(0): + entry = thread + while entry < Int32(self.window_tiles + 1): + band = cutlass.min(l_begin + entry, Int32(self.tiles_capacity + 1)) + s_band[entry] = band_base[band].to(Int32) + entry += Int32(self.threads) + cute.arch.sync_threads() + # Last local tile of this window: the first band whose position + # reaches the window end (at most WT bands lie in a window). + l_end = l_begin + if rank >= Int32(0): + # Smallest staged entry whose position reaches the window end; + # entries past the sequence's tiles are excluded by the clamp. + span = cutlass.min(tiles_s - l_begin, Int32(self.window_tiles)) + low = Int32(0) + high = span + while high > low: + mid = (low + high) >> Int32(1) + if s_band[mid] + rank >= window_end: + high = mid + else: + low = mid + Int32(1) + l_end = l_begin + low + has_tiles = l_end > l_begin + + head_elements = head.to(Int64) * Int64(_HEAD_DIM) + split_elements = (split * Int32(vd)).to(Int64) + v_chunks = self.v_chunks_per_row + ring_base = (window & Int32(1)) * Int32(self.window_tiles) - window_begin + rank + expected_flag = window + Int32(1) + head_base = head.to(Int64) * Int64(_HEAD_DIM * _HEAD_DIM) + + if is_producer: + # Producer: refill ring stages as they free up, once the + # prepare kernel has published the tile. Lane 0 waits and + # polls; every lane then issues its share of the copies. + prod_phase = Int32(1) + count = Int32(0) + if has_tiles: + # The flag of the next tile is read right after the copies + # of the current one are issued, so its latency overlaps the + # wait for a free stage; a stale value falls back to a spin. + seen = Int32(0) + if lane == Int32(0): + first_flag = _pointer_address( + ready, (ring_base + s_band[Int32(0)]).to(Int64) * Int64(self.heads) + head.to(Int64) + ) + seen = _ld_acquire_gpu_i32(first_flag) + for step in cutlass.range(l_end - l_begin, unroll=1): + local = l_begin + step + stage = count % Int32(self.stages) + ring_index = ring_base + s_band[step] + if lane == Int32(0): + cute.arch.mbarrier_wait(empty_bar + stage, phase=prod_phase) + if lane == Int32(0): + flag = _pointer_address( + ready, ring_index.to(Int64) * Int64(self.heads) + head.to(Int64) + ) + while seen != expected_flag: + _nanosleep(Int32(128)) + seen = _ld_acquire_gpu_i32(flag) + cute.arch.sync_warp() + token_base = start + local * Int32(_CHUNK) + rows_live = cutlass.min(Int32(_CHUNK), end - token_base) + self._issue_tile( + ws, v, stage_base + stage * Int32(stage_bytes), ring_index, head, + token_base, rows_live, head_elements, split_elements, v_stride, + shared_ptr_to_u32(full_bar + stage), full_bar + stage, lane, + ) + if lane == Int32(0): + if step + Int32(1) < l_end - l_begin: + next_flag = _pointer_address( + ready, + (ring_base + s_band[step + Int32(1)]).to(Int64) * Int64(self.heads) + + head.to(Int64), + ) + seen = _ld_acquire_gpu_i32(next_flag) + # Warm L2 with the value rows of a tile a few steps + # ahead: one lane per row, one line per v_split * 2 bytes. + ahead = local + Int32(_V_PREFETCH_TILES) + if ahead < l_end: + if lane < Int32(_CHUNK): + row_token = start + ahead * Int32(_CHUNK) + lane + if row_token < end: + line = row_token.to(Int64) * v_stride + head_elements + split_elements + for part in cutlass.range_constexpr(max(1, (self.v_split * 2) // 128)): + _prefetch_l2(_pointer_address(v, line + Int64(part * 64))) + if stage == Int32(self.stages - 1): + prod_phase = prod_phase ^ Int32(1) + count += Int32(1) + else: + # State fragments: rows row0/row1, columns col_base.. of this head. + row_local0 = group * Int32(16) + gid + row0 = split * Int32(vd) + row_local0 + row1 = row0 + Int32(8) + acc = cute.make_rmem_tensor((self.nb_blocks, 4), Float32) + shadow = cute.make_rmem_tensor((self.kb_steps, 4), Uint32) + parts = cute.make_rmem_tensor((4, 4), Float32) + bfrag = cute.make_rmem_tensor((self.kb_steps, 4), Uint32) + lam = cute.make_rmem_tensor((self.nb_blocks, 2), Float32) + vp = cute.make_rmem_tensor((2, 4), Float32) + u = cute.make_rmem_tensor((2, 4), Float32) + out = cute.make_rmem_tensor((2, 4), Float32) + cons_phase = Int32(0) + count = Int32(0) + if has_tiles: + initial = Int64(initial_indices[seq]) + final = Int64(final_indices[seq.to(Int64) * final_stride]) + for nb in cutlass.range_constexpr(self.nb_blocks): + acc[nb, 0] = Float32(0.0) + acc[nb, 1] = Float32(0.0) + acc[nb, 2] = Float32(0.0) + acc[nb, 3] = Float32(0.0) + if l_begin == Int32(0): + if not self._is_null(initial): + self._load_state( + recurrent_state, acc, initial * slot_stride + head_base, + row0, row1, col_base, tid, + ) + else: + # Resume the running state left in the final slot. + self._load_state( + recurrent_state, acc, final * slot_stride + head_base, + row0, row1, col_base, tid, + ) + self._refresh_shadow(acc, shadow) + # The guard is load-bearing: the loop body and its zero + # trip count are the same without it, but the generated + # schedule is far worse (the sixteen-head 4096-token case + # measures 428 us instead of 220 us on RTX PRO 6000). + if has_tiles: + for step in cutlass.range(l_end - l_begin, unroll=1): + local = l_begin + step + stage = count % Int32(self.stages) + token_base = start + local * Int32(_CHUNK) + rows_live = cutlass.min(Int32(_CHUNK), end - token_base) + cute.arch.mbarrier_wait(full_bar + stage, phase=cons_phase) + stage_addr = stage_base + stage * Int32(stage_bytes) + qt_addr = stage_addr + Int32(REC.Q_TILDE) + kt_addr = stage_addr + Int32(REC.K_TILDE) + kr_addr = stage_addr + Int32(REC.K_R) + inv_addr = stage_addr + Int32(REC.INV) + mqk_addr = stage_addr + Int32(REC.MQK) + lam_addr = stage_addr + Int32(REC.LAMBDA_C) + beta_addr = stage_addr + Int32(REC.BETA) + v_addr = stage_addr + Int32(REC.V) + + # Phase A: partial v'^T over this warp's key columns. + # Every operand of the phase is loaded first; the + # sixteen products run as four accumulator chains + # whose dependent steps are four issues apart. + tok_a = (matrix >> Int32(1)) * Int32(8) + matrix_row + for kb in cutlass.range_constexpr(self.kb_steps): + logical_chunk = chunk_base + Int32(kb * 2) + (matrix & Int32(1)) + physical = logical_chunk ^ (tok_a & Int32(7)) + bfrag[kb, 0], bfrag[kb, 1], bfrag[kb, 2], bfrag[kb, 3] = ldmatrix_m8n8x4_b16( + kt_addr + tok_a * Int32(256) + physical * Int32(16) + ) + # v^T in accumulator layout: one transposed ldmatrix of the + # group's sixteen value columns (rows tok, cols v). + v_tok = (matrix & Int32(1)) * Int32(8) + matrix_row + v_col = group * Int32(16) + (matrix >> Int32(1)) * Int32(8) + r0, r1, r2, r3 = ldmatrix_m8n8x4_trans_b16( + v_addr + group * Int32(512) + v_tok * Int32(32) + (matrix >> Int32(1)) * Int32(16) + ) + beta_lo = beta_addr + tid * Int32(8) + beta00, beta01 = _ld_shared_v2_f32(beta_lo) + beta10, beta11 = _ld_shared_v2_f32(beta_lo + Int32(32)) + square_row = (matrix >> Int32(1)) * Int32(8) + matrix_row + square_addr = square_row * Int32(32) + (matrix & Int32(1)) * Int32(16) + inv0, inv1, inv2, inv3 = ldmatrix_m8n8x4_b16(inv_addr + square_addr) + for chain in cutlass.range_constexpr(4): + for item in cutlass.range_constexpr(4): + parts[chain, item] = Float32(0.0) + for kb in cutlass.range_constexpr(self.kb_steps): + for half in cutlass.range_constexpr(2): + chain = 2 * (kb % 2) + half + parts[chain, 0], parts[chain, 1], parts[chain, 2], parts[chain, 3] = ( + bf16_mma_m16n8k16_f32( + parts[chain, 0], parts[chain, 1], parts[chain, 2], parts[chain, 3], + shadow[kb, 0], shadow[kb, 1], shadow[kb, 2], shadow[kb, 3], + bfrag[kb, 2 * half], bfrag[kb, 2 * half + 1], + ) + ) + for half in cutlass.range_constexpr(2): + for item in cutlass.range_constexpr(4): + vp[half, item] = parts[half, item] + parts[2 + half, item] + if cutlass.const_expr(self.k_split > 1): + self._group_reduce(vp, red_a, group, kq, lane) + v00, v01 = _bf16x2_to_f32x2(r0) + v10, v11 = _bf16x2_to_f32x2(r1) + v02, v03 = _bf16x2_to_f32x2(r2) + v12, v13 = _bf16x2_to_f32x2(r3) + vp[0, 0] = (v00 - vp[0, 0]) * beta00 + vp[0, 1] = (v01 - vp[0, 1]) * beta01 + vp[0, 2] = (v02 - vp[0, 2]) * beta00 + vp[0, 3] = (v03 - vp[0, 3]) * beta01 + vp[1, 0] = (v10 - vp[1, 0]) * beta10 + vp[1, 1] = (v11 - vp[1, 1]) * beta11 + vp[1, 2] = (v12 - vp[1, 2]) * beta10 + vp[1, 3] = (v13 - vp[1, 3]) * beta11 + a_vp0 = pack_f32x2_to_bfloat2(vp[0, 0], vp[0, 1]) + a_vp1 = pack_f32x2_to_bfloat2(vp[0, 2], vp[0, 3]) + a_vp2 = pack_f32x2_to_bfloat2(vp[1, 0], vp[1, 1]) + a_vp3 = pack_f32x2_to_bfloat2(vp[1, 2], vp[1, 3]) + + # Phase B: U^T = v'^T INV^T (every warp of the group). + u[0, 0], u[0, 1], u[0, 2], u[0, 3] = bf16_mma_m16n8k16_f32( + Float32(0.0), Float32(0.0), Float32(0.0), Float32(0.0), + a_vp0, a_vp1, a_vp2, a_vp3, inv0, inv1, + ) + u[1, 0], u[1, 1], u[1, 2], u[1, 3] = bf16_mma_m16n8k16_f32( + Float32(0.0), Float32(0.0), Float32(0.0), Float32(0.0), + a_vp0, a_vp1, a_vp2, a_vp3, inv2, inv3, + ) + a_u0 = pack_f32x2_to_bfloat2(u[0, 0], u[0, 1]) + a_u1 = pack_f32x2_to_bfloat2(u[0, 2], u[0, 3]) + a_u2 = pack_f32x2_to_bfloat2(u[1, 0], u[1, 1]) + a_u3 = pack_f32x2_to_bfloat2(u[1, 2], u[1, 3]) + + # Phase C: out^T = U^T Mqk^T + S^T q~^T over this warp's + # columns. The U^T Mqk^T product seeds two of the four + # chains (kq == 0 only); the rest is as in phase A. + for kb in cutlass.range_constexpr(self.kb_steps): + logical_chunk = chunk_base + Int32(kb * 2) + (matrix & Int32(1)) + physical = logical_chunk ^ (tok_a & Int32(7)) + bfrag[kb, 0], bfrag[kb, 1], bfrag[kb, 2], bfrag[kb, 3] = ldmatrix_m8n8x4_b16( + qt_addr + tok_a * Int32(256) + physical * Int32(16) + ) + for chain in cutlass.range_constexpr(4): + for item in cutlass.range_constexpr(4): + parts[chain, item] = Float32(0.0) + if group_lead: + b0, b1, b2, b3 = ldmatrix_m8n8x4_b16(mqk_addr + square_addr) + parts[0, 0], parts[0, 1], parts[0, 2], parts[0, 3] = bf16_mma_m16n8k16_f32( + Float32(0.0), Float32(0.0), Float32(0.0), Float32(0.0), + a_u0, a_u1, a_u2, a_u3, b0, b1, + ) + parts[1, 0], parts[1, 1], parts[1, 2], parts[1, 3] = bf16_mma_m16n8k16_f32( + Float32(0.0), Float32(0.0), Float32(0.0), Float32(0.0), + a_u0, a_u1, a_u2, a_u3, b2, b3, + ) + for kb in cutlass.range_constexpr(self.kb_steps): + for half in cutlass.range_constexpr(2): + chain = 2 * (kb % 2) + half + parts[chain, 0], parts[chain, 1], parts[chain, 2], parts[chain, 3] = ( + bf16_mma_m16n8k16_f32( + parts[chain, 0], parts[chain, 1], parts[chain, 2], parts[chain, 3], + shadow[kb, 0], shadow[kb, 1], shadow[kb, 2], shadow[kb, 3], + bfrag[kb, 2 * half], bfrag[kb, 2 * half + 1], + ) + ) + for half in cutlass.range_constexpr(2): + for item in cutlass.range_constexpr(4): + out[half, item] = parts[half, item] + parts[2 + half, item] + if cutlass.const_expr(self.k_split > 1): + self._group_reduce(out, red_c, group, kq, lane) + if group_lead: + # out^T fragments -> s_out[tok][v] through stmatrix.trans, + # then the live rows go out as 16-byte chunks, one per + # lane; the store's latency overlaps phase D. + _stmatrix_x4_trans( + out_addr + (v_tok * Int32(vd) + v_col) * Int32(2), + pack_f32x2_to_bfloat2(out[0, 0], out[0, 1]), + pack_f32x2_to_bfloat2(out[1, 0], out[1, 1]), + pack_f32x2_to_bfloat2(out[0, 2], out[0, 3]), + pack_f32x2_to_bfloat2(out[1, 2], out[1, 3]), + ) + cute.arch.sync_warp() + store_row = lane >> Int32(1) + store_chunk = group * Int32(2) + (lane & Int32(1)) + if store_row < rows_live: + c0, c1, c2, c3 = ld_shared_v4_u32( + out_addr + (store_row * Int32(v_chunks) + store_chunk) * Int32(16) + ) + element = ( + (token_base + store_row).to(Int64) * out_stride + + head_elements + + split_elements + + (store_chunk * Int32(8)).to(Int64) + ) + st_global_v4_u32(_pointer_address(output, element), c0, c1, c2, c3) + + # Phase D: S^T <- S^T * lambda_c[k] + U^T k_r over this warp's columns. + tok_d = (matrix & Int32(1)) * Int32(8) + matrix_row + for pair in cutlass.range_constexpr(self.nb_blocks // 2): + logical_chunk = chunk_base + Int32(pair * 2) + (matrix >> Int32(1)) + physical = logical_chunk ^ (tok_d & Int32(7)) + bfrag[pair, 0], bfrag[pair, 1], bfrag[pair, 2], bfrag[pair, 3] = ( + ldmatrix_m8n8x4_trans_b16( + kr_addr + tok_d * Int32(256) + physical * Int32(16) + ) + ) + for nb in cutlass.range_constexpr(self.nb_blocks): + kcol = col_base + Int32(nb * 8) + tid * Int32(2) + lam[nb, 0], lam[nb, 1] = _ld_shared_v2_f32(lam_addr + kcol * Int32(4)) + # Pair p: scale its two blocks, issue their MMAs, then + # refresh the shadow of pair p - 2 (whose MMAs are done). + pairs = self.nb_blocks // 2 + for pair in cutlass.range_constexpr(pairs): + for nb in cutlass.range_constexpr(2 * pair, 2 * pair + 2): + acc[nb, 0] = acc[nb, 0] * lam[nb, 0] + acc[nb, 1] = acc[nb, 1] * lam[nb, 1] + acc[nb, 2] = acc[nb, 2] * lam[nb, 0] + acc[nb, 3] = acc[nb, 3] * lam[nb, 1] + acc[2 * pair, 0], acc[2 * pair, 1], acc[2 * pair, 2], acc[2 * pair, 3] = ( + bf16_mma_m16n8k16_f32( + acc[2 * pair, 0], acc[2 * pair, 1], acc[2 * pair, 2], acc[2 * pair, 3], + a_u0, a_u1, a_u2, a_u3, bfrag[pair, 0], bfrag[pair, 1], + ) + ) + ( + acc[2 * pair + 1, 0], + acc[2 * pair + 1, 1], + acc[2 * pair + 1, 2], + acc[2 * pair + 1, 3], + ) = bf16_mma_m16n8k16_f32( + acc[2 * pair + 1, 0], acc[2 * pair + 1, 1], + acc[2 * pair + 1, 2], acc[2 * pair + 1, 3], + a_u0, a_u1, a_u2, a_u3, bfrag[pair, 2], bfrag[pair, 3], + ) + if cutlass.const_expr(pair >= 2): + done = pair - 2 + shadow[done, 0] = pack_f32x2_to_bfloat2(acc[2 * done, 0], acc[2 * done, 1]) + shadow[done, 1] = pack_f32x2_to_bfloat2(acc[2 * done, 2], acc[2 * done, 3]) + shadow[done, 2] = pack_f32x2_to_bfloat2(acc[2 * done + 1, 0], acc[2 * done + 1, 1]) + shadow[done, 3] = pack_f32x2_to_bfloat2(acc[2 * done + 1, 2], acc[2 * done + 1, 3]) + for done in cutlass.range_constexpr(max(0, pairs - 2), pairs): + shadow[done, 0] = pack_f32x2_to_bfloat2(acc[2 * done, 0], acc[2 * done, 1]) + shadow[done, 1] = pack_f32x2_to_bfloat2(acc[2 * done, 2], acc[2 * done, 3]) + shadow[done, 2] = pack_f32x2_to_bfloat2(acc[2 * done + 1, 0], acc[2 * done + 1, 1]) + shadow[done, 3] = pack_f32x2_to_bfloat2(acc[2 * done + 1, 2], acc[2 * done + 1, 3]) + # Every read of this stage is done: release it to the producer. + cute.arch.sync_warp() + if lane == Int32(0): + cute.arch.mbarrier_arrive(empty_bar + stage) + if stage == Int32(self.stages - 1): + cons_phase = cons_phase ^ Int32(1) + count += Int32(1) + + if cutlass.const_expr(self.checkpoint_export): + for cp in cutlass.range_constexpr(self.max_checkpoints): + cp_index = seq * Int32(self.max_checkpoints) + Int32(cp) + checkpoint = Int64(checkpoint_indices[cp_index]) + offset = checkpoint_offsets[cp_index].to(Int32) + if (offset > Int32(0)) & ((local + Int32(1)) * Int32(_CHUNK) == offset): + if not self._is_null(checkpoint): + self._store_state( + recurrent_state, acc, checkpoint * slot_stride + head_base, + row0, row1, col_base, tid, + ) + # Final state, or the running state for the next window. + if not self._is_null(final): + self._store_state( + recurrent_state, acc, final * slot_stride + head_base, + row0, row1, col_base, tid, + ) + # Empty sequences (ranks past the tiled ones) copy initial to + # final in window 0. + if window == Int32(0): + empty_rank = tiled_seqs + Int32(by) + while empty_rank < live_seqs: + empty_seq = sorted_seq[empty_rank].to(Int32) + empty_initial = Int64(initial_indices[empty_seq]) + empty_final = Int64( + final_indices[empty_seq.to(Int64) * final_stride] + ) + if not self._is_null(empty_final): + for nb in cutlass.range_constexpr(self.nb_blocks): + acc[nb, 0] = Float32(0.0) + acc[nb, 1] = Float32(0.0) + acc[nb, 2] = Float32(0.0) + acc[nb, 3] = Float32(0.0) + if not self._is_null(empty_initial): + self._load_state( + recurrent_state, acc, empty_initial * slot_stride + head_base, + row0, row1, col_base, tid, + ) + self._store_state( + recurrent_state, acc, empty_final * slot_stride + head_base, + row0, row1, col_base, tid, + ) + empty_rank += Int32(self.rows) + + +def _recurrence_key(binding: Binding) -> tuple[object, ...]: + caps = binding.plan.caps + plan = binding.plan + return ( + "recurrence", + binding.output.device.index, + caps.heads, + caps.tiles_capacity, + plan.window_tiles, + plan.max_windows, + plan.recurrence_rows, + plan.v_split, + plan.k_split, + plan.stages, + caps.checkpoint_export, + caps.max_checkpoints, + caps.null_state_index, + binding.initial_state_indices.dtype, + ) + + +def _compile_recurrence(binding: Binding) -> tuple[tuple[object, ...], Callable[..., None]]: + key = _recurrence_key(binding) + cached = _RECURRENCE_CACHE.get(key) + if cached is not None: + return key, cached + caps = binding.plan.caps + index_type = _numeric_type(binding.initial_state_indices.dtype) + kernel = _RecurrenceKernel( + heads=caps.heads, + tiles_capacity=caps.tiles_capacity, + window_tiles=binding.plan.window_tiles, + rows=binding.plan.recurrence_rows, + v_split=binding.plan.v_split, + k_split=binding.plan.k_split, + stages=binding.plan.stages, + checkpoint_export=caps.checkpoint_export, + max_checkpoints=caps.max_checkpoints, + null_state_index=caps.null_state_index, + index_type=index_type, + ) + raise_if_kernel_resolution_frozen("cute.compile", target=kernel, cache_key=key) + raw = b12x_compile( + kernel, + _fake_pointer(BFloat16), + _fake_pointer(Int32), + _fake_pointer(Int32), + _fake_pointer(Int32), + _fake_pointer(Int32), + _fake_pointer(index_type), + _fake_pointer(index_type), + _fake_pointer(index_type), + _fake_pointer(Int32), + _fake_pointer(Int32), + _fake_pointer(Int32), + _fake_pointer(Int32), + _fake_pointer(cutlass.Int8), + _fake_pointer(Float32), + _fake_pointer(BFloat16), + Int64(1), + Int64(1), + Int64(1), + Int64(1), + Int32(1), + Int32(0), + current_cuda_stream(), + compile_spec=KernelCompileSpec.from_key("sequence.kda_prefill.recurrence", 9, key), + ) + + def launch(active: Binding, window: int) -> None: + if _recurrence_key(active) != key: + raise ValueError("compiled KDA recurrence kernel does not match the binding") + raw( + _pointer(active.v, BFloat16), + _pointer(active.cu_seqlens, Int32), + _pointer(active.band_base, Int32), + _pointer(active.sorted_seq, Int32), + _pointer(active.window_table, Int32), + _pointer(active.initial_state_indices, index_type), + _pointer(active.final_state_indices, index_type), + _pointer(active.checkpoint_state_indices, index_type), + _pointer(active.checkpoint_offsets, Int32), + _pointer(active.num_seqs, Int32), + _pointer(active.error_code, Int32), + _pointer(active.ready_flags, Int32), + _pointer(active.ws.view(torch.int8), cutlass.Int8), + _pointer(active.recurrent_state, Float32), + _pointer(active.output, BFloat16), + int(active.v.stride(0)), + int(active.output.stride(0)), + int(active.recurrent_state.stride(0)), + int(active.final_state_indices.stride(0)), + int(active.token_capacity), + int(window), + current_cuda_stream(), + ) + + _RECURRENCE_CACHE[key] = launch + return key, launch + + +def run_recurrence(binding: Binding, *, window: int = 0) -> None: + """Walk the live tiles of one window (stage 2); requires stages 0 and 1. + + The launch polls the window's ready flags, so it may be issued on a stream + running concurrently with the same window's prepare launch. + """ + with torch.cuda.device(binding.output.device): + _launch_stage( + lambda b: (_recurrence_key(b), _RECURRENCE_CACHE.get(_recurrence_key(b))), + _compile_recurrence, + binding, + int(window), + ) + + +def _prologue_key(binding: Binding) -> tuple[object, ...]: + caps = binding.plan.caps + return ( + "prologue", + binding.output.device.index, + caps.max_seqs, + caps.tiles_capacity, + binding.plan.window_tiles, + binding.plan.max_windows, + binding.plan.duplicate_table_size, + 2 * binding.plan.window_tiles * caps.heads, + caps.max_state_slots, + caps.metadata_validation, + caps.max_checkpoints, + caps.null_state_index, + binding.initial_state_indices.dtype, + ) + + +def _compile_prologue(binding: Binding) -> tuple[tuple[object, ...], Callable[..., None]]: + key = _prologue_key(binding) + cached = _PROLOGUE_CACHE.get(key) + if cached is not None: + return key, cached + caps = binding.plan.caps + index_type = _numeric_type(binding.initial_state_indices.dtype) + kernel = _PrologueKernel( + max_seqs=caps.max_seqs, + tiles_capacity=caps.tiles_capacity, + window_tiles=binding.plan.window_tiles, + max_windows=binding.plan.max_windows, + table_size=binding.plan.duplicate_table_size, + flag_count=2 * binding.plan.window_tiles * caps.heads, + max_state_slots=caps.max_state_slots, + validate=caps.metadata_validation == "transactional", + max_checkpoints=caps.max_checkpoints, + null_state_index=caps.null_state_index, + index_type=index_type, + ) + raise_if_kernel_resolution_frozen("cute.compile", target=kernel, cache_key=key) + raw = b12x_compile( + kernel, + _fake_pointer(Int32), + _fake_pointer(index_type), + _fake_pointer(index_type), + _fake_pointer(index_type), + _fake_pointer(Int32), + _fake_pointer(Int32), + _fake_pointer(Int32), + _fake_pointer(Int32), + _fake_pointer(Int32), + _fake_pointer(Int32), + _fake_pointer(Int32), + _fake_pointer(Int32), + _fake_pointer(Int32), + _fake_pointer(Int32), + _fake_pointer(Int32), + _fake_pointer(Int32), + Int64(1), + Int32(1), + Int32(1), + Int32(1), + current_cuda_stream(), + compile_spec=KernelCompileSpec.from_key("sequence.kda_prefill.prologue", 6, key), + ) + + def launch(active: Binding, launched_tiles: int) -> None: + if _prologue_key(active) != key: + raise ValueError("compiled KDA prologue does not match the binding") + raw( + _pointer(active.cu_seqlens, Int32), + _pointer(active.initial_state_indices, index_type), + _pointer(active.final_state_indices, index_type), + _pointer(active.checkpoint_state_indices, index_type), + _pointer(active.checkpoint_offsets, Int32), + _pointer(active.num_seqs, Int32), + _pointer(active.num_tokens, Int32), + _pointer(active.error_code, Int32), + _pointer(active.duplicate_slots, Int32), + _pointer(active.band_base, Int32), + _pointer(active.sorted_seq, Int32), + _pointer(active.rank_of, Int32), + _pointer(active.pos_seq, Int32), + _pointer(active.pos_local, Int32), + _pointer(active.window_table, Int32), + _pointer(active.ready_flags, Int32), + int(active.final_state_indices.stride(0)), + int(active.seq_capacity), + int(active.token_capacity), + int(launched_tiles), + current_cuda_stream(), + ) + + _PROLOGUE_CACHE[key] = launch + return key, launch + + +def _prepare_key(binding: Binding) -> tuple[object, ...]: + caps = binding.plan.caps + return ( + "prepare", + binding.output.device.index, + caps.heads, + caps.tiles_capacity, + binding.plan.window_tiles, + caps.qk_l2norm, + binding.A_log.dtype, + binding.dt_bias.dtype, + ) + + +def _compile_prepare(binding: Binding) -> tuple[tuple[object, ...], Callable[..., None]]: + key = _prepare_key(binding) + cached = _PREPARE_CACHE.get(key) + if cached is not None: + return key, cached + caps = binding.plan.caps + a_log_type = _numeric_type(binding.A_log.dtype) + dt_bias_type = _numeric_type(binding.dt_bias.dtype) + kernel = _PrepareKernel( + heads=caps.heads, + tiles_capacity=caps.tiles_capacity, + window_tiles=binding.plan.window_tiles, + qk_l2norm=caps.qk_l2norm, + a_log_type=a_log_type, + dt_bias_type=dt_bias_type, + ) + raise_if_kernel_resolution_frozen("cute.compile", target=kernel, cache_key=key) + raw = b12x_compile( + kernel, + _fake_pointer(BFloat16), + _fake_pointer(BFloat16), + _fake_pointer(BFloat16), + _fake_pointer(BFloat16), + _fake_pointer(a_log_type), + _fake_pointer(dt_bias_type), + _fake_pointer(Int32), + _fake_pointer(Int32), + _fake_pointer(Int32), + _fake_pointer(Int32), + _fake_pointer(Int32), + _fake_pointer(BFloat16), + _fake_pointer(Float32), + Int64(1), + Int64(1), + Int64(1), + Int64(1), + Int64(1), + Float32(1.0), + Float32(1.0), + Float32(1.0), + Int32(0), + current_cuda_stream(), + compile_spec=KernelCompileSpec.from_key("sequence.kda_prefill.prepare", 6, key), + ) + + def launch(active: Binding, scale: float, gate_scale: float, eps: float, window: int) -> None: + if _prepare_key(active) != key: + raise ValueError("compiled KDA prepare kernel does not match the binding") + raw( + _pointer(active.q, BFloat16), + _pointer(active.k, BFloat16), + _pointer(active.raw_g, BFloat16), + _pointer(active.raw_beta, BFloat16), + _pointer(active.A_log, a_log_type), + _pointer(active.dt_bias, dt_bias_type), + _pointer(active.cu_seqlens, Int32), + _pointer(active.pos_seq, Int32), + _pointer(active.pos_local, Int32), + _pointer(active.error_code, Int32), + _pointer(active.ready_flags, Int32), + _pointer(active.ws.view(torch.bfloat16), BFloat16), + _pointer(active.ws.view(torch.float32), Float32), + int(active.q.stride(0)), + int(active.k.stride(0)), + int(active.raw_g.stride(0)), + int(active.raw_beta.stride(0)), + int(active.raw_beta.stride(1)), + float(scale), + float(gate_scale), + float(eps), + int(window), + current_cuda_stream(), + ) + + _PREPARE_CACHE[key] = launch + return key, launch + + +def _launch_stage(cache_lookup, compile_fn, binding: Binding, *args) -> None: + capturing = torch.cuda.is_current_stream_capturing() + key, launch = cache_lookup(binding) + if capturing and (launch is None or key not in _WARMED): + raise RuntimeError( + "KDA prefill kernels must be compiled and warm-run before CUDA graph capture" + ) + if launch is None: + key, launch = compile_fn(binding) + launch(binding, *args) + if not capturing: + _WARMED.add(key) + + +def run_prologue(binding: Binding, *, windows: int | None = None) -> None: + """Validate metadata, clear the ready flags, and build the tile tables (stage 0).""" + plan = binding.plan + launched = plan.max_windows if windows is None else int(windows) + launched_tiles = min(plan.caps.tiles_capacity, launched * plan.window_tiles) + with torch.cuda.device(binding.output.device): + _launch_stage( + lambda b: (_prologue_key(b), _PROLOGUE_CACHE.get(_prologue_key(b))), + _compile_prologue, + binding, + launched_tiles, + ) + + +def run_prepare( + binding: Binding, *, lower_bound: float, scale: float, eps: float, window: int = 0 +) -> None: + """Fill the workspace ring slot of one window (stage 1); requires the prologue first.""" + with torch.cuda.device(binding.output.device): + _launch_stage( + lambda b: (_prepare_key(b), _PREPARE_CACHE.get(_prepare_key(b))), + _compile_prepare, + binding, + float(scale), + float(lower_bound) * _LOG2E, + float(eps), + int(window), + ) + + +@dataclass +class _SideResources: + """Per-device side stream and event pool for the window pipeline.""" + + stream: torch.cuda.Stream + events: list[torch.cuda.Event] + + +_SIDE: dict[int, _SideResources] = {} + + +def _side_resources(device: torch.device, windows: int) -> _SideResources: + """Return the side stream and at least ``2 * windows + 1`` initialized events.""" + needed = 2 * int(windows) + 1 + resources = _SIDE.get(device.index) + capturing = torch.cuda.is_current_stream_capturing() + if resources is None or len(resources.events) < needed: + if capturing: + raise RuntimeError( + "KDA prefill pipeline resources must be created by a warm run before CUDA graph capture" + ) + if resources is None: + resources = _SideResources(stream=torch.cuda.Stream(device=device), events=[]) + _SIDE[device.index] = resources + current = torch.cuda.current_stream(device) + while len(resources.events) < needed: + event = torch.cuda.Event() + event.record(current) + resources.events.append(event) + return resources + + +def run_prefill( + binding: Binding, *, lower_bound: float, scale: float, eps: float, windows: int | None = None +) -> None: + """Launch the window pipeline: prologue, then prepare and recurrence per window. + + Prepare launches run on a per-device side stream, recurrence launches on + the current stream. Recurrence waits for its window's prepare, while the + next prepare can overlap the previous recurrence. Prepare of window ``w`` + waits for the recurrence of window ``w - 2`` before reusing that workspace + ring slot. Under stream capture the fork and join are recorded as graph + dependencies. + """ + device = binding.output.device + plan = binding.plan + launched = plan.max_windows if windows is None else int(windows) + if launched < 1 or launched > plan.max_windows: + raise ValueError(f"windows must be in 1..{plan.max_windows}, got {launched}") + with torch.cuda.device(device): + main = torch.cuda.current_stream(device) + resources = _side_resources(device, launched) + side = resources.stream + fork = resources.events[2 * launched] + prepared = resources.events[:launched] + consumed = resources.events[launched : 2 * launched] + run_prologue(binding, windows=launched) + # Enqueue window by window so every event is recorded before a stream + # waits on it (a wait binds to the event's most recent record). + fork.record(main) + side.wait_event(fork) + for window in range(launched): + with torch.cuda.stream(side): + if window >= 2: + side.wait_event(consumed[window - 2]) + run_prepare(binding, lower_bound=lower_bound, scale=scale, eps=eps, window=window) + prepared[window].record(side) + main.wait_event(prepared[window]) + run_recurrence(binding, window=window) + consumed[window].record(main) + main.wait_event(prepared[launched - 1]) + + +def prewarm_binding(binding: Binding) -> None: + """Compile the three stages for ``binding`` and create the pipeline resources.""" + with torch.cuda.device(binding.output.device): + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("KDA prefill compilation is forbidden during CUDA capture") + _compile_prologue(binding) + _compile_prepare(binding) + _compile_recurrence(binding) + _side_resources(binding.output.device, binding.plan.max_windows) + + +def workspace_tiles(binding: Binding, window: int = 0) -> dict[str, torch.Tensor]: + """Return logical views of the ring slot holding ``window``. + + Operand tiles are de-swizzled copies; ``inv``, ``mqk``, ``lambda_c``, + ``beta``, and ``ready`` are views into the buffer. + """ + tiles = binding.plan.window_tiles + heads = binding.plan.caps.heads + slot = slice((window & 1) * tiles, (window & 1) * tiles + tiles) + ws = binding.ws[slot] + rows = torch.arange(_CHUNK).view(_CHUNK, 1) + cols = torch.arange(_HEAD_DIM).view(1, _HEAD_DIM) + physical = ((((cols >> 3) ^ (rows & 7)) << 3) | (cols & 7)).to(ws.device) + index = physical.view(1, 1, _CHUNK, _HEAD_DIM).expand(tiles, heads, _CHUNK, _HEAD_DIM) + + def bf16(offset: int, elements: int, *shape: int) -> torch.Tensor: + return ws[..., offset : offset + 2 * elements].view(torch.bfloat16).view(tiles, heads, *shape) + + def f32(offset: int, elements: int) -> torch.Tensor: + return ws[..., offset : offset + 4 * elements].view(torch.float32).view(tiles, heads, elements) + + def logical(tile: torch.Tensor) -> torch.Tensor: + return torch.gather(tile, 3, index) + + return { + "q_tilde": logical(bf16(REC.Q_TILDE, _CHUNK * _HEAD_DIM, _CHUNK, _HEAD_DIM)), + "k_tilde": logical(bf16(REC.K_TILDE, _CHUNK * _HEAD_DIM, _CHUNK, _HEAD_DIM)), + "k_r": logical(bf16(REC.K_R, _CHUNK * _HEAD_DIM, _CHUNK, _HEAD_DIM)), + "lambda_c": f32(REC.LAMBDA_C, _HEAD_DIM), + "beta": f32(REC.BETA, _CHUNK), + "inv": bf16(REC.INV, _CHUNK * _CHUNK, _CHUNK, _CHUNK), + "mqk": bf16(REC.MQK, _CHUNK * _CHUNK, _CHUNK, _CHUNK), + "ready": binding.ready_flags[slot], + } + + +def clear_caches() -> None: + _PROLOGUE_CACHE.clear() + _PREPARE_CACHE.clear() + _RECURRENCE_CACHE.clear() + _WARMED.clear() + _SIDE.clear() + + +__all__ = [ + "clear_caches", + "prewarm_binding", + "run_prefill", + "run_prepare", + "run_prologue", + "run_recurrence", + "workspace_tiles", +] diff --git a/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/ec64f07f841675501ff279bae3074d374ca8015a982ab12dc837ea2d03720a8e/_policy.py b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/ec64f07f841675501ff279bae3074d374ca8015a982ab12dc837ea2d03720a8e/_policy.py new file mode 100644 index 00000000..2023e55f --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/ec64f07f841675501ff279bae3074d374ca8015a982ab12dc837ea2d03720a8e/_policy.py @@ -0,0 +1,219 @@ +"""Policy contract for chunked KDA prefill: query, config, heuristic.""" + +from __future__ import annotations + +from dataclasses import dataclass + +from b12x.policy import ComponentPolicy +from b12x.policy.components import KDA_PREFILL +from b12x.policy.types import FrozenMapping + +BACKEND = "cutedsl" +V_SPLIT_CHOICES = (16, 32, 64, 128) +K_SPLIT_CHOICES = (1, 2, 4) +STAGE_CHOICES = (2, 3, 4) +CHUNK_TOKENS = 16 + + +class WorkspaceRecord: + """Byte layout of one prepared (tile, head) record in the workspace ring. + + The recurrence kernel copies ``[0, HEAD_BYTES)`` into a pipeline stage + with one bulk copy; the stage then holds this CTA's value rows from + offset ``V`` as ``v_split // 16`` groups of ``[16 tokens x 16 values]``, + copied straight from the value tensor. The operand tiles are stored in + the swizzled 16-byte-chunk order the consumer's ldmatrix reads. + """ + + Q_TILDE = 0 + K_TILDE = 4096 + K_R = 8192 + INV = 12288 + MQK = 12800 + LAMBDA_C = 13312 + BETA = 13824 + HEAD_BYTES = 13952 + BYTES = 14080 + V = 14080 + + +WORKSPACE_RECORD_BYTES = WorkspaceRecord.BYTES +# Prepared-tile bytes one window may occupy; two windows stay L2 resident. +WINDOW_BYTES_BUDGET = 36 << 20 + + +def tiles_capacity(max_tokens: int, max_seqs: int) -> int: + """Upper bound on packed chunk tiles: one partial tile per sequence.""" + return -(-int(max_tokens) // CHUNK_TOKENS) + int(max_seqs) + + +def default_window_tiles(heads: int, max_tokens: int, max_seqs: int) -> int: + """Tiles per pipeline window so a window's prepared tiles fit the L2 budget.""" + per_row = int(heads) * WORKSPACE_RECORD_BYTES + return max(1, min(tiles_capacity(max_tokens, max_seqs), WINDOW_BYTES_BUDGET // per_row)) + + +@dataclass(frozen=True, kw_only=True) +class KdaPrefillQuery: + """Immutable geometry and planned capacity of one KDA prefill plan.""" + + heads: int + head_dim: int + model_dtype: str + state_dtype: str + qk_l2norm: bool + checkpoint_export: bool + max_tokens: int + max_seqs: int + max_checkpoints: int = 1 + + def profile_fields(self) -> dict[str, object]: + fields = { + "heads": int(self.heads), + "head_dim": int(self.head_dim), + "model_dtype": str(self.model_dtype), + "state_dtype": str(self.state_dtype), + "qk_l2norm": bool(self.qk_l2norm), + "checkpoint_export": bool(self.checkpoint_export), + "max_tokens": int(self.max_tokens), + "max_seqs": int(self.max_seqs), + } + # Checkpoint capacity changes compiled export/validation code, whose + # cache keys include it, but not the registry's launch-geometry schema. + return fields + + +@dataclass(frozen=True) +class KdaPrefillConfig: + """Backend selection plus the recurrence kernel's launch geometry. + + ``v_split`` is the number of value rows one recurrence CTA owns (smaller + splits launch more CTAs per sequence and head at the cost of re-reading + the prepared tiles from L2); ``k_split`` is how many warps share each + sixteen-row group by splitting its key columns (more warps shorten the + per-tile tensor-core chain at the cost of shared-memory reductions); + ``stages`` is the tile prefetch depth; ``window_tiles`` is the number of + consecutive banded tile positions one pipeline window covers. The prepare + kernel of a window runs concurrently with the recurrence of that window + and the next window's prepare, and two windows of prepared records form + the workspace ring, so the window size bounds the ring's footprint. + """ + + backend: str = BACKEND + v_split: int = 64 + k_split: int = 1 + stages: int = 3 + window_tiles: int = 64 + + @classmethod + def from_profile(cls, payload: FrozenMapping) -> "KdaPrefillConfig": + keys = set(payload.keys()) + if "backend" not in keys or not keys <= { + "backend", "v_split", "k_split", "stages", "window_tiles" + }: + raise ValueError( + "KDA prefill profiles require backend and accept only v_split, " + "k_split, stages, and window_tiles" + ) + backend = payload["backend"] + if not isinstance(backend, str): + raise TypeError("backend must be a string") + values = {} + for name, default in (("v_split", 64), ("k_split", 1), ("stages", 3), ("window_tiles", 64)): + value = payload.get(name, default) + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"{name} must be an integer") + values[name] = int(value) + return cls(backend=backend, **values) + + def to_dict(self) -> dict[str, object]: + return { + "backend": self.backend, + "v_split": int(self.v_split), + "k_split": int(self.k_split), + "stages": int(self.stages), + "window_tiles": int(self.window_tiles), + } + + +def _heuristic(query: KdaPrefillQuery, device) -> KdaPrefillConfig: + del device + return KdaPrefillConfig( + backend=BACKEND, + v_split=64, + k_split=1, + stages=3, + window_tiles=default_window_tiles(query.heads, query.max_tokens, query.max_seqs), + ) + + +def _validate(query: KdaPrefillQuery, config: KdaPrefillConfig, device) -> None: + del device + if config.backend != BACKEND: + raise ValueError(f"unsupported {KDA_PREFILL} backend {config.backend!r}") + if config.v_split not in V_SPLIT_CHOICES: + raise ValueError( + f"unsupported {KDA_PREFILL} v_split {config.v_split!r}; expected one " + f"of {V_SPLIT_CHOICES}" + ) + if config.k_split not in K_SPLIT_CHOICES: + raise ValueError( + f"unsupported {KDA_PREFILL} k_split {config.k_split!r}; expected one " + f"of {K_SPLIT_CHOICES}" + ) + if config.stages not in STAGE_CHOICES: + raise ValueError( + f"unsupported {KDA_PREFILL} stages {config.stages!r}; expected one of " + f"{STAGE_CHOICES}" + ) + if 2 * config.v_split * config.k_split + 32 > 1024: + raise ValueError(f"{KDA_PREFILL} v_split x k_split exceeds the thread limit") + if isinstance(config.window_tiles, bool) or int(config.window_tiles) < 1: + raise ValueError(f"{KDA_PREFILL} window_tiles must be a positive integer") + if query.head_dim != 128: + raise ValueError(f"{KDA_PREFILL} requires head_dim 128, got {query.head_dim}") + if query.model_dtype != "bfloat16" or query.state_dtype != "float32": + raise ValueError( + f"{KDA_PREFILL} requires bfloat16 activations and float32 state, got " + f"{query.model_dtype}/{query.state_dtype}" + ) + + +KDA_PREFILL_POLICY = ComponentPolicy( + component_id=KDA_PREFILL, + query_schema_version=1, + config_schema_version=1, + query_fields=frozenset( + { + "heads", + "head_dim", + "model_dtype", + "state_dtype", + "qk_l2norm", + "checkpoint_export", + "max_tokens", + "max_seqs", + } + ), + config_fields=frozenset({"backend", "v_split", "k_split", "stages", "window_tiles"}), + encode_query=KdaPrefillQuery.profile_fields, + decode_profile=KdaPrefillConfig.from_profile, + heuristic=_heuristic, + validate_config=_validate, +) + +__all__ = [ + "BACKEND", + "CHUNK_TOKENS", + "K_SPLIT_CHOICES", + "KDA_PREFILL_POLICY", + "KdaPrefillConfig", + "KdaPrefillQuery", + "STAGE_CHOICES", + "V_SPLIT_CHOICES", + "WINDOW_BYTES_BUDGET", + "WORKSPACE_RECORD_BYTES", + "WorkspaceRecord", + "default_window_tiles", + "tiles_capacity", +] diff --git a/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/f4ccf9da197eb68fd21b23e8c03615cd0d4c2f871da0dc32a9c6ec82e44d3931/kv_cache_manager.py b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/f4ccf9da197eb68fd21b23e8c03615cd0d4c2f871da0dc32a9c6ec82e44d3931/kv_cache_manager.py new file mode 100644 index 00000000..73295359 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/payload-by-sha/f4ccf9da197eb68fd21b23e8c03615cd0d4c2f871da0dc32a9c6ec82e44d3931/kv_cache_manager.py @@ -0,0 +1,1309 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import itertools +import time +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Literal, overload + +from vllm.distributed.kv_events import BlockStored, KVCacheEvent +from vllm.logger import init_logger +from vllm.utils.math_utils import cdiv +from vllm.v1.core.kv_cache_coordinator import ( + HybridKVCacheCoordinator, + get_kv_cache_coordinator, +) +from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector +from vllm.v1.core.kv_cache_utils import ( + BlockHashWithGroupId, + KVCacheBlock, + KVCacheBlockCopy, +) +from vllm.v1.kv_cache_interface import ( + AttentionSpec, + CrossAttentionSpec, + EncoderOnlyAttentionSpec, + KVCacheConfig, + MambaSpec, + get_kv_cache_spec_kind, + get_kv_cache_spec_sliding_window, +) +from vllm.v1.metrics.stats import PrefixCacheStats +from vllm.v1.request import Request, RequestStatus + +logger = init_logger(__name__) + + +@dataclass +class KVCacheBlocks: + """ + The allocation result of KVCacheManager, work as the interface between + Scheduler and KVCacheManager, to hide KVCacheManager's internal data + structure from the Scheduler. + """ + + blocks: tuple[Sequence[KVCacheBlock], ...] + """ + `blocks[i][j]` refers to the i-th kv_cache_group + and the j-th block of tokens.We don't use block of + tokens as the outer dimension because it assumes all + kv_cache_groups have the same number of blocks, which is true for now but + will be broken if we want to give different block_size to different + kv_cache_groups in the future. + + Each single type KVCacheBlocks could be represented as: + - list[KVCacheBlock] for more than one KVCacheBlock + - an empty tuple for requests without KVCacheBlock + (a precomputed KVCacheBlocks is in KVCacheManager to avoid GC overhead) + """ + + def __add__(self, other: "KVCacheBlocks") -> "KVCacheBlocks": + """Adds two KVCacheBlocks instances.""" + return KVCacheBlocks( + tuple( + list(itertools.chain(blk1, blk2)) + for blk1, blk2 in zip(self.blocks, other.blocks) + ) + ) + + @overload + def get_block_ids( + self, + allow_none: Literal[False] = False, + ) -> tuple[list[int], ...]: ... + + @overload + def get_block_ids( + self, + allow_none: Literal[True] = True, + ) -> tuple[list[int], ...] | None: ... + + def get_block_ids( + self, + allow_none: bool = False, + ) -> tuple[list[int], ...] | None: + """ + Converts the KVCacheBlocks instance to block_ids. + + Returns: + tuple[list[int], ...]: A tuple of lists where: + - the outer tuple corresponds to KV cache groups + - each inner list contains the block_ids of the blocks in that + group + """ + if allow_none and all(len(group) == 0 for group in self.blocks): + return None + return tuple([blk.block_id for blk in group] for group in self.blocks) + + def get_unhashed_block_ids(self) -> list[int]: + """Get block_ids of unhashed blocks from KVCacheBlocks instance.""" + assert len(self.blocks) == 1, "Only one group is supported" + return [block.block_id for block in self.blocks[0] if block.block_hash is None] + + def get_unhashed_block_ids_all_groups(self) -> list[list[int]]: + """Get block_ids of unhashed blocks from KVCacheBlocks instance.""" + # Skip padding blocks. + return [ + [ + block.block_id + for block in group + if block.block_hash is None and not block.is_null + ] + for group in self.blocks + ] + + def new_empty(self) -> "KVCacheBlocks": + """ + Creates a new KVCacheBlocks instance with no blocks. + """ + return KVCacheBlocks(tuple(() for _ in range(len(self.blocks)))) + + +@dataclass(frozen=True) +class _SharedPrefixLease: + """Scheduler-owned reference to one verified multi-group prefix.""" + + request_id: str + num_computed_tokens: int + expires_at: float + ready: bool = False + + +class KVCacheManager: + def __init__( + self, + kv_cache_config: KVCacheConfig, + max_model_len: int, + scheduler_block_size: int, + hash_block_size: int, + max_in_flight_tokens: int | None = None, + enable_caching: bool = True, + use_eagle: bool = False, + num_prefill_lookahead: int = 0, + log_stats: bool = False, + enable_kv_cache_events: bool = False, + dcp_world_size: int = 1, + pcp_world_size: int = 1, + metrics_collector: KVCacheMetricsCollector | None = None, + watermark: float = 0.0, + ) -> None: + self.max_model_len = max_model_len + # When unset, fall back to `max_model_len` so the recycling-aware cap + # collapses to the prior (uncapped) admission behavior. The scheduler + # always supplies the real value at runtime. + if max_in_flight_tokens is None: + max_in_flight_tokens = max_model_len + + self.enable_caching = enable_caching + self.enable_kv_cache_events = enable_kv_cache_events + self.use_eagle = use_eagle + self.log_stats = log_stats + self.metrics_collector = metrics_collector + # FIXME: make prefix cache stats conditional on log_stats. We still need + # this comment because when the log stats is enabled there are still + # potential configs we could expose in the future. + self.prefix_cache_stats = PrefixCacheStats() if log_stats else None + + self.coordinator = get_kv_cache_coordinator( + kv_cache_config=kv_cache_config, + max_model_len=self.max_model_len, + max_in_flight_tokens=max_in_flight_tokens, + use_eagle=self.use_eagle, + enable_caching=self.enable_caching, + enable_kv_cache_events=enable_kv_cache_events, + dcp_world_size=dcp_world_size, + pcp_world_size=pcp_world_size, + scheduler_block_size=scheduler_block_size, + hash_block_size=hash_block_size, + metrics_collector=self.metrics_collector, + num_prefill_lookahead=num_prefill_lookahead, + ) + self.num_kv_cache_groups = len(kv_cache_config.kv_cache_groups) + self.block_pool = self.coordinator.block_pool + self.kv_cache_config = kv_cache_config + + # Watermark: minimum number of KV cache blocks to keep free when + # admitting waiting/preempted requests, to avoid frequent preemptions. + assert watermark >= 0.0, "watermark must be non-negative" + self.watermark_blocks = int(watermark * kv_cache_config.num_blocks) + self.kv_cache_event_metadata = tuple( + ( + get_kv_cache_spec_kind(group.kv_cache_spec).value, + get_kv_cache_spec_sliding_window(group.kv_cache_spec), + ) + for group in kv_cache_config.kv_cache_groups + ) + + # Pre-constructed KVCacheBlocks with no blocks, callers should use this + # via create_kv_cache_blocks instead of creating new ones to avoid GC + # overhead. + # + # We use nested tuples to ensure the empty KVCacheBlocks is immutable. + self.empty_kv_cache_blocks = KVCacheBlocks( + tuple(() for _ in range(self.num_kv_cache_groups)) + ) + # Short-lived leases let a KV connector hand a verified external + # prefix to requests that arrive after the original loader. The + # lease is represented as an ordinary hidden request in every + # single-type manager, so BlockPool remains the sole refcount owner. + self._shared_prefix_leases: dict[str, _SharedPrefixLease] = {} + + # Off-table cow blocks handed to a KV connector for partial-tail + # offload; pinned until the request's blocks are freed. + self._partial_tail_pins: dict[str, list[KVCacheBlock]] = {} + + @property + def usage(self) -> float: + """Get the KV cache usage. + + Returns: + The KV cache usage (between 0.0 and 1.0). + """ + return self.block_pool.get_usage() + + def make_prefix_cache_stats(self) -> PrefixCacheStats | None: + """Get (and reset) the prefix cache stats. + + Returns: + The current prefix caching stats, or None if logging is disabled. + """ + if not self.log_stats: + return None + stats = self.prefix_cache_stats + self.prefix_cache_stats = PrefixCacheStats() + return stats + + def prefix_cache_lookup_enabled(self, request: Request) -> bool: + """Whether a local prefix cache lookup may be run for this request.""" + return self.enable_caching and not request.skip_reading_prefix_cache + + def record_prefix_cache_stats(self, request: Request, num_hits: int) -> None: + # Don't count a request that skipped the cache lookup. + if not self.log_stats or not self.prefix_cache_lookup_enabled(request): + return + assert self.prefix_cache_stats is not None + self.prefix_cache_stats.record( + num_tokens=request.num_tokens, + num_hits=num_hits, + preempted=request.num_preemptions > 0, + ) + + def get_computed_blocks(self, request: Request) -> tuple[KVCacheBlocks, int, int]: + """Get the computed (cached) blocks for the request. + Note that the computed blocks must be full. + + Args: + request: The request to get the computed blocks. + + Returns: + A tuple containing: + - A list of blocks that are computed for the request. + - The number of computed tokens. + - ``shared_prefix_boundary``: the block-aligned token position of + a shared prefix that a sparse-retention group (Mamba / sliding + window) has not cached yet (Marconi-style APC), or 0 if none. + Pinned so sparse prefix-cache retention does not drop + the junction and defeat cross-request reuse. + """ + # We skip finding the prefix cache hit when prefix caching is + # disabled or the request is marked as skipping kv cache read + # (which happens when the request requires prompt logprobs + # or calls a pooling model with all pooling). + if not self.prefix_cache_lookup_enabled(request): + return self.empty_kv_cache_blocks, 0, 0 + + # NOTE: When all tokens hit the cache, we must recompute the last token + # to obtain logits. Thus, set max_cache_hit_length to prompt_length - 1. + # This can trigger recomputation of an entire block, rather than just + # the single last token, because allocate_slots() requires + # num_computed_tokens to be block-size aligned. Removing this limitation + # could slightly improve performance in the future. + max_cache_hit_length = request.num_tokens - 1 + computed_blocks, num_new_computed_tokens, num_uncached = ( + self.coordinator.find_longest_cache_hit( + request.block_hashes, max_cache_hit_length + ) + ) + + # When kv_cache_report_mode is "full", emit BlockStored events + # for the reused prefix cache blocks so that external consumers + # (e.g. gateway) can learn about them. + if ( + num_new_computed_tokens > 0 + and self.enable_kv_cache_events + and getattr(request, "kv_cache_report_mode", "incremental") == "full" + ): + for group_idx, group_blocks in enumerate(computed_blocks): + num_blocks = len(group_blocks) + if num_blocks > 0: + group = self.kv_cache_config.kv_cache_groups[group_idx] + block_size = group.kv_cache_spec.block_size + self.block_pool.emit_cached_block_events( + request, + num_blocks, + block_size, + group_idx, + ) + + # The junction to pin is where the lagging sparse-retention group stops + # (``num_new_computed_tokens``) plus the uncached shared prefix -- i.e. + # the longest single-group hit. Sub-block gaps are left to the mask, + # which floors to the alignment boundary (a no-op there). + shared_prefix_boundary = ( + num_new_computed_tokens + num_uncached if num_uncached else 0 + ) + + blocks = self.create_kv_cache_blocks(computed_blocks) + return blocks, num_new_computed_tokens, shared_prefix_boundary + + def get_computed_blocks_for_connector( + self, request: Request + ) -> tuple[KVCacheBlocks, int, int, bool]: + """Local prefix-cache lookup for a request scheduled with a KV connector. + + Hybrid (Mamba + full-attention) models can have per-group prefix hits + diverge under block pressure: the full-attention tail may be evicted + while a deeper Mamba state survives, or vice versa. Report the + full-attention hit as the local prefix - the connector transfers the + remaining suffix and the Mamba state is transferred unconditionally by + nixl's ``_apply_prefix_caching`` - and flag when that hit ran deeper + than a lagging group. Such a hit only has a valid Mamba state at its + boundary if the connector supplies it, so the caller must fall back to + ``get_computed_blocks`` to reconcile when no external tokens are found. + + Non-hybrid models and already-convergent hits use ``get_computed_blocks``. + + Returns: + The ``get_computed_blocks`` triple (blocks, number of local computed + tokens, shared-prefix boundary) plus ``hit_diverged``. + """ + coordinator = self.coordinator + if not ( + self.kv_cache_config.has_mamba_layers + and isinstance(coordinator, HybridKVCacheCoordinator) + and coordinator.full_attention_group_id is not None + ): + return *self.get_computed_blocks(request), False + + if not self.prefix_cache_lookup_enabled(request): + return self.empty_kv_cache_blocks, 0, 0, False + + fa_group_id = coordinator.full_attention_group_id + computed, per_group_hits = coordinator.find_longest_cache_hit_per_group( + request.block_hashes, request.num_tokens - 1 + ) + if any(hit > per_group_hits[fa_group_id] for hit in per_group_hits): + # A lagging group hit deeper than full attention means its + # full-attention blocks were evicted; use the reconciled boundary + # that every group agrees on. + return *self.get_computed_blocks(request), False + + num_local = per_group_hits[fa_group_id] + blocks = self.create_kv_cache_blocks(computed) + # Per-group lookups do not detect an uncached shared prefix (boundary 0). + return blocks, num_local, 0, min(per_group_hits) < num_local + + def allocate_slots( + self, + request: Request, + num_new_tokens: int, + num_new_computed_tokens: int = 0, + new_computed_blocks: KVCacheBlocks | None = None, + num_lookahead_tokens: int = 0, + num_external_computed_tokens: int = 0, + delay_cache_blocks: bool = False, + num_encoder_tokens: int = 0, + full_sequence_must_fit: bool = False, + reserved_blocks: int = 0, + has_scheduled_reqs: bool = True, + recurrent_publication_boundary: int | None = None, + recurrent_prefill_checkpoint_plan: tuple[int, int, tuple[int, ...]] | None = None, + recurrent_checkpoint_publications: tuple[int, ...] = (), + ) -> KVCacheBlocks | None: + from vllm.v1.core.recurrent_prefill_checkpoint import validate_plan + managers = self.coordinator.single_type_managers + plan = recurrent_prefill_checkpoint_plan + if plan is None: + if recurrent_checkpoint_publications: + raise ValueError("interior publication requires an explicit checkpoint plan") + return self._allocate_slots_without_checkpoint_plan(request=request, num_new_tokens=num_new_tokens, num_new_computed_tokens=num_new_computed_tokens, new_computed_blocks=new_computed_blocks, num_lookahead_tokens=num_lookahead_tokens, num_external_computed_tokens=num_external_computed_tokens, delay_cache_blocks=delay_cache_blocks, num_encoder_tokens=num_encoder_tokens, full_sequence_must_fit=full_sequence_must_fit, reserved_blocks=reserved_blocks, has_scheduled_reqs=has_scheduled_reqs, recurrent_publication_boundary=recurrent_publication_boundary) + start, end, targets = plan + if (start != 0 or request.num_computed_tokens != 0 or end != request.num_prompt_tokens + or end != request.num_tokens or num_new_tokens != end + or num_new_computed_tokens or num_external_computed_tokens or delay_cache_blocks + or num_encoder_tokens or (new_computed_blocks is not None and any(new_computed_blocks.blocks))): + raise ValueError("coalesced checkpoint allocation requires a fresh complete prompt") + if any(p not in targets for p in recurrent_checkpoint_publications): + raise ValueError("publication boundary is not an explicitly produced internal state") + selected = [manager for manager in managers if hasattr(manager, "_allocated_block_reqs")] + if not selected: + raise ValueError("checkpoint plan has no recurrent manager") + for manager in selected: + validate_plan(plan, start, end, manager.block_size) + if (manager.kv_cache_spec.num_prefill_checkpoint_blocks < len(targets) + or manager.req_to_blocks.get(request.request_id) + or request.request_id in manager._allocated_block_reqs + or request.request_id in manager._partial_hit_reqs): + raise ValueError("checkpoint plan cannot append a fresh recurrent table") + try: + for manager in selected: + manager._planned_recurrent_checkpoints[request.request_id] = plan + manager._planned_recurrent_publications[request.request_id] = recurrent_checkpoint_publications + return self._allocate_slots_without_checkpoint_plan(request=request, num_new_tokens=num_new_tokens, num_new_computed_tokens=num_new_computed_tokens, new_computed_blocks=new_computed_blocks, num_lookahead_tokens=num_lookahead_tokens, num_external_computed_tokens=num_external_computed_tokens, delay_cache_blocks=delay_cache_blocks, num_encoder_tokens=num_encoder_tokens, full_sequence_must_fit=full_sequence_must_fit, reserved_blocks=reserved_blocks, has_scheduled_reqs=has_scheduled_reqs, recurrent_publication_boundary=recurrent_publication_boundary) + finally: + for manager in selected: + manager._planned_recurrent_checkpoints.pop(request.request_id, None) + manager._planned_recurrent_publications.pop(request.request_id, None) + + def _allocate_slots_without_checkpoint_plan( + self, + request: Request, + num_new_tokens: int, + num_new_computed_tokens: int = 0, + new_computed_blocks: KVCacheBlocks | None = None, + num_lookahead_tokens: int = 0, + num_external_computed_tokens: int = 0, + delay_cache_blocks: bool = False, + num_encoder_tokens: int = 0, + full_sequence_must_fit: bool = False, + reserved_blocks: int = 0, + has_scheduled_reqs: bool = True, + recurrent_publication_boundary: int | None = None, + ) -> KVCacheBlocks | None: + """Add slots for a request with new tokens to append. + + Args: + request: The request to allocate slots. + num_new_tokens: The number of new tokens to be allocated and computed. + num_new_computed_tokens: The number of new computed tokens just + hitting the prefix caching, excluding external tokens. + new_computed_blocks: The cached blocks for the above new computed + tokens, grouped as a tuple by kv cache groups. + num_lookahead_tokens: The number of speculative tokens to allocate. + This is used by spec decode proposers with kv-cache such + as eagle. + num_external_computed_tokens: The number of tokens that their + KV caches are not cached by vLLM but cached by the connector. + delay_cache_blocks: Whether to skip caching the blocks. This is + used by P/D when allocating blocks used in a KV transfer + which will complete in a future step. + num_encoder_tokens: The number of encoder tokens to allocate for + cross-attention in encoder-decoder models(e.g., Whisper). + For decoder-only models, this should be 0. + full_sequence_must_fit: Only allocate blocks if the KV cache has enough + free blocks to hold the full sequence, accounting for prefix cache hits + and sliding window. Used as an admission gate to prevent over-admitting + requests when chunked prefill would otherwise only check the first chunk + reserved_blocks: Number of free blocks that must be left available for + other in-flight sequences to complete. The actual allocation is only + made if it fits within (free blocks - reserved_blocks). Used to gate + async KV-connector loads so their initial allocation cannot consume + blocks an already in-flight (prefilling) sequence is relying on. + has_scheduled_reqs: Whether any requests are already scheduled to run + this step, controls whether watermark is applied. + recurrent_publication_boundary: Exact connector boundary materialized + by this allocation; it does not change native hash geometry. + + Blocks layout: + ``` + ---------------------------------------------------------------------- + | < comp > | < new_comp > | < ext_comp > | < new > | < lookahead > | + ---------------------------------------------------------------------- + | < to be computed > | + ---------------------------------------------------------------------- + | < to be allocated > | + ---------------------------------------------------------------------- + | < to be cached (roughly, | + | details below)> | + ---------------------------------------------------------------------- + | Prefix-cached tokens from either vLLM | + | or connector. Can be safely removed if | + | they are outside sliding window. | + ---------------------------------------------------------------------- + | < cached by vLLM > | not cached by | + | vLLM, but | + | ref_cnt | ref_cnt not | cached by | + | increased| increased yet| connector | + ---------------------------------------------------------------------- + ``` + + Abbrivations: + + ``` + comp = request.num_computed_tokens + new_comp = num_new_computed_tokens + = len(new_computed_blocks) * block_size + ext_comp = num_external_computed_tokens, cached by the connector + new = num_new_tokens, including unverified draft tokens + lookahead = num_lookahead_tokens + ``` + + NOTE: for new tokens which include both verified and unverified draft + tokens, we only cache the verified tokens (by capping the number at + `request.num_tokens`). + + The allocation has three stages: + - Free unnecessary blocks in `comp` and check + if we have sufficient free blocks (return None if not). + - Handle prefix tokens (`comp + new_comp + ext_comp`): + - Free unnecessary blocks (e.g. outside sliding window) + - Allocate new blocks for `ext_comp` tokens inside + sliding window + - Allocate new blocks for tokens to be computed (`new + lookahead`) + + Returns: + A list of new allocated blocks. + """ + # When loading KV data asynchronously, we may have zero new tokens to + # compute while still allocating slots for externally computed tokens. + if num_new_tokens == 0 and num_external_computed_tokens == 0: + raise ValueError( + "num_new_tokens must be greater than 0 when there are no " + "external computed tokens" + ) + + if new_computed_blocks is not None: + new_computed_block_list = new_computed_blocks.blocks + else: + new_computed_block_list = self.empty_kv_cache_blocks.blocks + + # The number of computed tokens is the number of computed tokens plus + # the new prefix caching hits + num_local_computed_tokens = ( + request.num_computed_tokens + num_new_computed_tokens + ) + total_computed_tokens = min( + num_local_computed_tokens + num_external_computed_tokens, + self.max_model_len, + ) + if recurrent_publication_boundary is not None: + expected_boundary = min( + total_computed_tokens + num_new_tokens, + request.num_tokens, + ) + if ( + type(recurrent_publication_boundary) is not int + or recurrent_publication_boundary <= 0 + or recurrent_publication_boundary % self.block_pool.hash_block_size != 0 + or recurrent_publication_boundary != expected_boundary + ): + raise ValueError( + "recurrent publication boundary must equal this allocation's" + " finalized internal-hash boundary" + ) + + watermark_blocks = 0 + # The watermark is applied to waiting/preempted requests only, and only + # when there's at least one request already scheduled. + if has_scheduled_reqs and request.status in ( + RequestStatus.WAITING, + RequestStatus.PREEMPTED, + ): + watermark_blocks = self.watermark_blocks + + if full_sequence_must_fit: + # First check and fail if the full request sequence won't fit. + full_num_tokens = min(request.num_tokens, self.max_model_len) + + num_blocks_to_allocate = self.coordinator.get_num_blocks_to_allocate( + request_id=request.request_id, + num_tokens=full_num_tokens, + new_computed_blocks=new_computed_block_list, + num_encoder_tokens=num_encoder_tokens, + total_computed_tokens=total_computed_tokens, + num_local_computed_tokens=num_local_computed_tokens, + num_tokens_main_model=full_num_tokens, + apply_admission_cap=True, + ) + required_blocks = num_blocks_to_allocate + watermark_blocks + self.evict_shared_prefix_leases_until_free(required_blocks) + if required_blocks > self.block_pool.get_num_free_blocks(): + return None + + num_tokens_main_model = total_computed_tokens + num_new_tokens + num_tokens_need_slot = min( + num_tokens_main_model + num_lookahead_tokens, self.max_model_len + ) + + # Free the blocks that are skipped during the attention computation + # (e.g., tokens outside the sliding window). + # We can do this even if we cannot schedule this request due to + # insufficient free blocks. + # Should call this function before allocating new blocks to reduce + # the number of evicted blocks. + # Free on the processed-token basis: in-flight steps' attention windows + # still read blocks below the optimistic boundary, and rejected spec + # tokens can roll it back. + self.coordinator.remove_skipped_blocks( + request.request_id, + max(0, total_computed_tokens - request.num_in_flight_tokens), + num_prompt_tokens=request.num_prompt_tokens, + ) + + num_blocks_to_allocate = self.coordinator.get_num_blocks_to_allocate( + request_id=request.request_id, + num_tokens=num_tokens_need_slot, + new_computed_blocks=new_computed_block_list, + num_encoder_tokens=num_encoder_tokens, + total_computed_tokens=num_local_computed_tokens + + num_external_computed_tokens, + num_local_computed_tokens=num_local_computed_tokens, + num_tokens_main_model=num_tokens_main_model, + ) + + # Keep `reserved_blocks` free for other in-flight sequences, and an + # additional watermark of headroom for waiting/preempted admissions. + available_blocks = self.block_pool.get_num_free_blocks() - reserved_blocks + required_blocks = num_blocks_to_allocate + watermark_blocks + self.evict_shared_prefix_leases_until_free(required_blocks + reserved_blocks) + available_blocks = self.block_pool.get_num_free_blocks() - reserved_blocks + if required_blocks > available_blocks: + # Cannot allocate new blocks + return None + + if ( + new_computed_block_list is not self.empty_kv_cache_blocks.blocks + or num_external_computed_tokens > 0 + ): + # Append the new computed blocks to the request blocks until now to + # avoid the case where the new blocks cannot be allocated. + self.coordinator.allocate_new_computed_blocks( + request_id=request.request_id, + new_computed_blocks=new_computed_block_list, + num_local_computed_tokens=num_local_computed_tokens, + num_external_computed_tokens=num_external_computed_tokens, + ) + + new_blocks = self.coordinator.allocate_new_blocks( + request.request_id, + num_tokens_need_slot, + num_tokens_main_model, + num_encoder_tokens, + ) + + # P/D: delay caching blocks if we have to recv from + # remote. Update state for locally cached blocks. + if not self.enable_caching or delay_cache_blocks: + return self.create_kv_cache_blocks(new_blocks) + + # NOTE(woosuk): We want to commit (cache) up to num_local_computed_tokens + # + num_external_computed_tokens + num_new_tokens, but must exclude + # "non-committable" tokens (e.g., draft tokens that could be rejected). + # Therefore, we cap the number at `request.num_tokens`, ensuring only + # "finalized" tokens are cached. + num_tokens_to_cache = min( + total_computed_tokens + num_new_tokens, + request.num_tokens, + ) + if recurrent_publication_boundary is not None: + for manager in self.coordinator.single_type_managers: + manager.recurrent_publication_boundary = recurrent_publication_boundary + try: + self.coordinator.cache_blocks(request, num_tokens_to_cache) + finally: + if recurrent_publication_boundary is not None: + for manager in self.coordinator.single_type_managers: + manager.recurrent_publication_boundary = None + + return self.create_kv_cache_blocks(new_blocks) + + def free(self, request: Request) -> None: + """Free the blocks allocated for the request. + We free the blocks in reverse order so that the tail blocks are evicted + first when caching is enabled. + + Args: + request: The request to free the blocks. + """ + pins = self._partial_tail_pins.pop(request.request_id, None) + if pins: + self.block_pool.free_blocks(pins) + self.coordinator.free(request.request_id) + + def remove_skipped_blocks( + self, + request_id: str, + processed_computed_tokens: int, + num_prompt_tokens: int | None = None, + ) -> None: + """Remove the blocks that are no longer needed from `blocks` and replace + the removed blocks with null_block. + + Args: + request_id: The request ID. + processed_computed_tokens: Computed-token prefix length covering + fully processed and committed tokens only (safe to free). + num_prompt_tokens: Optional prompt length for R-SWA gap eviction. + """ + self.coordinator.remove_skipped_blocks( + request_id, processed_computed_tokens, num_prompt_tokens + ) + + def pop_blocks_for_free(self, request: Request) -> list[KVCacheBlock]: + """Pop the request's bookkeeping and return its blocks without + returning them to the block pool. The caller must eventually free + them in reverse order (so that tail blocks are evicted first). + + Args: + request: The request to pop the blocks for. + + Returns: + The request's blocks in allocation order. + """ + blocks = self.coordinator.pop_blocks_for_free(request.request_id) + # Pins ride the same (possibly deferred) free as the request blocks. + # Preemption may release a pin under a still-queued offload — the same + # exposure normal saves of table blocks already have. + pins = self._partial_tail_pins.pop(request.request_id, None) + if pins: + blocks = pins + blocks + return blocks + + def evict_blocks(self, block_ids: set[int]) -> None: + """evict blocks from the prefix cache by their block IDs. + + Args: + block_ids: Set of block IDs to evict from cache. + """ + self.block_pool.evict_blocks(block_ids) + + def reset_prefix_cache(self) -> bool: + """Reset prefix cache. This function may be used in RLHF + flows to invalidate prefix caching after the weights are updated, + or used for resetting prefix caching status for benchmarking. + + Returns: + bool: True if the prefix cache is successfully reset, + False otherwise. + """ + if not self.block_pool.reset_prefix_cache(): + return False + if self.log_stats: + assert self.prefix_cache_stats is not None + self.prefix_cache_stats.reset = True + return True + + def get_num_common_prefix_blocks(self, running_request_id: str) -> list[int]: + """Calculate the number of common prefix blocks for each kv cache group. + + The function selects a running request and iterates through its blocks. + A block is considered a common prefix block if ALL requests with + allocated KV cache share it (i.e., ref_cnt equals the number of entries + in req_to_blocks). + + NOTE(woosuk): The number of requests with allocated KV cache is **greater + than or equal to** the number of requests scheduled in the current step. + This is because having allocated KV cache only indicates that: + 1. The request has not yet finished, and + 2. The request holds its blocks unfreed. + + While all scheduled requests must have allocated KV cache, the inverse + is not necessarily true. There may be requests with allocated KV cache + that are not scheduled in the current step. + + This can result in an edge case where the number of common prefix blocks + is 0, even though all scheduled requests share a common prefix. This + occurs because there may be unscheduled requests that do not share the + common prefix. Currently, this case cannot be easily detected, so the + function returns 0 in such cases. + + Args: + running_request_id: The request ID of any running request, used to + identify the common prefix blocks. + + Returns: + list[int]: The number of common prefix blocks for each kv cache + group. + """ + return self.coordinator.get_num_common_prefix_blocks(running_request_id) + + def take_events(self) -> list[KVCacheEvent]: + """Take the KV cache events from the block pool. + + Returns: + A list of KV cache events. + """ + events = self.block_pool.take_events() + for event in events: + if not isinstance(event, BlockStored): + continue + if event.group_idx is None: + continue + if event.group_idx < 0 or event.group_idx >= len( + self.kv_cache_event_metadata + ): + logger.warning( + "Group index `%s` not in KV cache metadata", event.group_idx + ) + continue + # Annotate here so BlockPool can keep emitting structural cache + # events without owning semantic KV cache spec metadata. + kind, sliding_window = self.kv_cache_event_metadata[event.group_idx] + event.kv_cache_spec_kind = kind + event.kv_cache_spec_sliding_window = sliding_window + return events + + def get_blocks(self, request_id: str) -> KVCacheBlocks: + """Get the blocks of a request.""" + return self.create_kv_cache_blocks(self.coordinator.get_blocks(request_id)) + + def _release_shared_prefix_lease(self, lease_key: str) -> None: + lease = self._shared_prefix_leases.pop(lease_key, None) + if lease is not None: + # The hidden request owns exactly one ordinary BlockPool reference + # to every live block in the leased prefix. Coordinator.free() + # releases those references through the same path as any request. + self.coordinator.free(lease.request_id) + + def expire_shared_prefix_leases(self, now: float | None = None) -> None: + """Release connector prefix leases whose bounded grace period elapsed.""" + now = time.monotonic() if now is None else now + for lease_key, lease in tuple(self._shared_prefix_leases.items()): + if lease.expires_at <= now: + self._release_shared_prefix_lease(lease_key) + + def publish_shared_prefix_lease( + self, + lease_key: str, + source_request_id: str, + num_computed_tokens: int, + ttl_seconds: float, + *, + max_entries: int = 2, + reserved_blocks: int = 0, + now: float | None = None, + ) -> bool: + """Pin a verified request prefix for bounded by-reference attachment. + + This must be called only after the connector reports every worker's + receive complete and ``cache_blocks`` has published the source. A + hidden request uses the existing per-group allocation bookkeeping and + BlockPool refcounts; no parallel ownership system is introduced. + """ + if not lease_key or num_computed_tokens <= 0 or ttl_seconds <= 0: + return False + if max_entries <= 0 or max_entries > 2: + raise ValueError("shared prefix leases are bounded to at most two entries") + + now = time.monotonic() if now is None else now + self.expire_shared_prefix_leases(now) + expires_at = now + ttl_seconds + existing = self._shared_prefix_leases.get(lease_key) + if existing is not None: + if existing.num_computed_tokens != num_computed_tokens: + self._release_shared_prefix_lease(lease_key) + else: + self._shared_prefix_leases[lease_key] = _SharedPrefixLease( + existing.request_id, + existing.num_computed_tokens, + expires_at, + existing.ready, + ) + return True + + source_blocks = self.coordinator.get_blocks(source_request_id) + managers = self.coordinator.single_type_managers + if ( + len(source_blocks) != self.num_kv_cache_groups + or len(managers) != self.num_kv_cache_groups + or not any(source_blocks) + ): + return False + + partial_page_count = sum( + num_computed_tokens % manager.block_size != 0 for manager in managers + ) + required_free_blocks = reserved_blocks + partial_page_count + self.evict_shared_prefix_leases_until_free(required_free_blocks, now=now) + if self.block_pool.get_num_free_blocks() < required_free_blocks: + return False + + while len(self._shared_prefix_leases) >= max_entries: + oldest_key = min( + self._shared_prefix_leases, + key=lambda key: self._shared_prefix_leases[key].expires_at, + ) + self._release_shared_prefix_lease(oldest_key) + + lease_request_id = f"\x00sparkcache-shared-prefix:{lease_key}" + if any(self.coordinator.get_blocks(lease_request_id)): + # A request ID collision or stale bookkeeping must never alias two + # different prefixes. Refuse the optimization and recompute. + return False + try: + normalized_blocks: list[Sequence[KVCacheBlock]] = [] + for manager, group_blocks in zip(managers, source_blocks): + skipped_blocks = ( + manager.get_num_skipped_tokens(num_computed_tokens) + // manager.block_size + ) + required_blocks = cdiv(num_computed_tokens, manager.block_size) + if ( + skipped_blocks < 0 + or skipped_blocks > required_blocks + or required_blocks > len(group_blocks) + ): + raise RuntimeError("shared prefix lease source table is incomplete") + selected = list(group_blocks[skipped_blocks:required_blocks]) + if len(selected) != required_blocks - skipped_blocks: + raise RuntimeError( + "shared prefix lease logical source slice is incomplete" + ) + if num_computed_tokens % manager.block_size != 0: + if not selected: + raise RuntimeError( + "shared prefix lease lacks its physical partial page" + ) + logical_boundary_idx = required_blocks - 1 + partial = manager._partial_hit_reqs.get(source_request_id) + if partial is None: + source_block = selected[-1] + else: + partial_idx, source_block = partial + source_occurrences = sum( + block is source_block for block in group_blocks + ) + if ( + partial_idx != logical_boundary_idx + or source_occurrences != 1 + ): + raise RuntimeError( + "shared prefix lease source metadata is invalid" + ) + selected[-1] = source_block + if ( + source_block.is_null + or sum(block is source_block for block in selected) != 1 + ): + raise RuntimeError( + "shared prefix lease source boundary is null or duplicated" + ) + normalized_blocks.append( + [manager._null_block] * skipped_blocks + selected + ) + self.coordinator.allocate_new_computed_blocks( + request_id=lease_request_id, + new_computed_blocks=tuple(normalized_blocks), + num_local_computed_tokens=num_computed_tokens, + num_external_computed_tokens=0, + ) + # A prefix ending inside any group's physical page cannot remain + # on the leader's mutable page. Move every such lease tail to a + # dedicated block through vLLM's existing CoW copy pipeline. The + # scheduler marks the lease ready only after the copy step fences. + for manager in managers: + if num_computed_tokens % manager.block_size == 0: + continue + lease_blocks = manager.req_to_blocks[lease_request_id] + partial = manager._partial_hit_reqs.pop(lease_request_id, None) + if partial is None: + block_idx = num_computed_tokens // manager.block_size + if block_idx >= len(lease_blocks): + raise RuntimeError( + "shared prefix lease lacks its physical partial page" + ) + source_block = lease_blocks[block_idx] + else: + block_idx, source_block = partial + if ( + not 0 <= block_idx < len(lease_blocks) + or lease_blocks[block_idx] is not source_block + ): + raise RuntimeError( + "shared prefix lease partial-page metadata is invalid" + ) + if source_block.is_null: + raise RuntimeError( + "shared prefix lease partial page is a null placeholder" + ) + hot_block = self.block_pool.get_new_blocks(1)[0] + manager._apply_cow(lease_request_id, block_idx, source_block, hot_block) + manager.new_block_ids.append(hot_block.block_id) + except Exception as error: + self.coordinator.free(lease_request_id) + logger.warning("Lease pin skipped key=%s reason=%s", lease_key, error) + return False + + self._shared_prefix_leases[lease_key] = _SharedPrefixLease( + lease_request_id, + num_computed_tokens, + expires_at, + False, + ) + return True + + def mark_shared_prefix_lease_ready( + self, lease_key: str, *, now: float | None = None + ) -> bool: + """Publish a lease only after its partial-page copies have fenced.""" + self.expire_shared_prefix_leases(now) + lease = self._shared_prefix_leases.get(lease_key) + if lease is None: + return False + self._shared_prefix_leases[lease_key] = _SharedPrefixLease( + lease.request_id, + lease.num_computed_tokens, + lease.expires_at, + True, + ) + return True + + def attach_shared_prefix_lease( + self, + lease_key: str, + target_request_id: str, + *, + now: float | None = None, + ) -> int: + """Attach a request to a verified leased prefix without copying KV.""" + now = time.monotonic() if now is None else now + self.expire_shared_prefix_leases(now) + lease = self._shared_prefix_leases.get(lease_key) + if lease is None or not lease.ready: + return 0 + if any(self.coordinator.get_blocks(target_request_id)): + return 0 + + lease_blocks = self.coordinator.get_blocks(lease.request_id) + try: + self.coordinator.allocate_new_computed_blocks( + request_id=target_request_id, + new_computed_blocks=lease_blocks, + num_local_computed_tokens=lease.num_computed_tokens, + num_external_computed_tokens=0, + ) + except Exception as error: + self.coordinator.free(target_request_id) + logger.warning( + "Lease attach skipped key=%s request=%s reason=%s", + lease_key, + target_request_id, + error, + ) + return 0 + return lease.num_computed_tokens + + def discard_shared_prefix_lease(self, lease_key: str) -> None: + """Release only the lease pin; attached requests retain their refs.""" + self._release_shared_prefix_lease(lease_key) + + def evict_shared_prefix_leases_until_free( + self, min_free_blocks: int, *, now: float | None = None + ) -> None: + """Drop oldest lease pins under allocation pressure, never request refs.""" + self.expire_shared_prefix_leases(now) + while ( + self._shared_prefix_leases + and self.block_pool.get_num_free_blocks() < min_free_blocks + ): + oldest_key = min( + self._shared_prefix_leases, + key=lambda key: self._shared_prefix_leases[key].expires_at, + ) + self._release_shared_prefix_lease(oldest_key) + + def get_block_ids(self, request_id: str) -> tuple[list[int], ...]: + """Get the block ids of a request.""" + return self.get_blocks(request_id).get_block_ids() + + def get_block_ids_for_computed_tokens( + self, + request_id: str, + num_computed_tokens: int, + ) -> tuple[list[int], ...]: + """Get block ids covering the request's computed tokens.""" + block_ids = self.get_block_ids(request_id) + clipped_block_ids: list[list[int]] = [] + for group, ids in zip(self.kv_cache_config.kv_cache_groups, block_ids): + spec = group.kv_cache_spec + if not isinstance(spec, AttentionSpec) or isinstance( + spec, (CrossAttentionSpec, EncoderOnlyAttentionSpec) + ): + clipped_block_ids.append(ids) + continue + + num_valid_blocks = cdiv(num_computed_tokens, spec.block_size) + clipped_block_ids.append(ids[:num_valid_blocks]) + return tuple(clipped_block_ids) + + def estimate_cached_tokens(self, request: Request) -> int: + """Estimate the number of tokens cached by the request.""" + cached_tokens: int | None = None + for group, blocks in zip( + self.kv_cache_config.kv_cache_groups, + self.get_blocks(request.request_id).blocks, + ): + if isinstance( + group.kv_cache_spec, + (CrossAttentionSpec, EncoderOnlyAttentionSpec), + ): + # Cross-attention and encoder-only groups are not prefix cached. + continue + + group_cached_tokens = 0 + for block in blocks: + group_cached_tokens = max( + group_cached_tokens, + block.block_hash_num_tokens or 0, + ) + + cached_tokens = ( + group_cached_tokens + if cached_tokens is None + else min(cached_tokens, group_cached_tokens) + ) + + return cached_tokens or 0 + + def cache_blocks(self, request: Request, num_computed_tokens: int) -> None: + """Cache the blocks for the request, if enabled. + + Args: + request: The request to cache the blocks. + num_computed_tokens: The number of computed tokens, including tokens + that are already cached and tokens to be cached. + """ + if self.enable_caching: + self.coordinator.cache_blocks(request, num_computed_tokens) + + def create_kv_cache_blocks( + self, blocks: tuple[list[KVCacheBlock], ...] + ) -> KVCacheBlocks: + # Only create new KVCacheBlocks for non-empty blocks + return KVCacheBlocks(blocks) if any(blocks) else self.empty_kv_cache_blocks + + def truncate_computed_blocks( + self, blocks: KVCacheBlocks, num_computed_tokens: int + ) -> KVCacheBlocks: + """Return a lookup-result view truncated at an aligned token endpoint. + + An external hit can supply the final Mamba state even when the local + Mamba group ends before this endpoint. Other groups must cover it. + Pure slicing: refcounts are untouched and ``blocks`` is not mutated. + """ + truncated: list[list[KVCacheBlock]] = [] + for group_blocks, manager, group in zip( + blocks.blocks, + self.coordinator.single_type_managers, + self.kv_cache_config.kv_cache_groups, + strict=True, + ): + assert num_computed_tokens % manager.block_size == 0 + num_blocks = num_computed_tokens // manager.block_size + if isinstance(group.kv_cache_spec, MambaSpec): + num_blocks = min(num_blocks, len(group_blocks)) + else: + assert num_blocks <= len(group_blocks) + truncated.append(list(group_blocks[:num_blocks])) + return self.create_kv_cache_blocks(tuple(truncated)) + + def take_new_block_ids(self) -> list[int]: + """Drain and return new attention block IDs for zeroing.""" + ids: list[int] = [] + for mgr in self.coordinator.single_type_managers: + ids.extend(mgr.take_new_block_ids()) + return ids + + def get_zeroing_block_ids_in_range( + self, request_id: str, start_token: int, end_token: int + ) -> list[int]: + """The request's block ids covering [start_token, end_token), from + the groups whose new blocks are zeroed by the worker.""" + ids: list[int] = [] + for mgr in self.coordinator.single_type_managers: + if mgr.records_new_block_ids: + start_idx = start_token // mgr.block_size + end_idx = cdiv(end_token, mgr.block_size) + blocks = mgr.req_to_blocks[request_id] + ids.extend(blk.block_id for blk in blocks[start_idx:end_idx]) + return ids + + def record_blocks_for_zeroing(self, request_id: str, start_token: int) -> None: + """Re-record the request's blocks from start_token onwards for + zeroing, e.g. blocks a failed async KV load left unwritten. + + start_token must be block-aligned: zeroing a partially-valid block + would wipe its valid prefix. + """ + for mgr in self.coordinator.single_type_managers: + if mgr.records_new_block_ids: + assert start_token % mgr.block_size == 0 + start_idx = start_token // mgr.block_size + blocks = mgr.req_to_blocks[request_id] + mgr.new_block_ids.extend(blk.block_id for blk in blocks[start_idx:]) + + def take_kv_cache_block_copies( + self, + ) -> tuple[list[KVCacheBlockCopy], list[KVCacheBlock]]: + """Drain pending copies and return their retained endpoints.""" + pending_copies: list[tuple[KVCacheBlock, KVCacheBlock]] = [] + for mgr in self.coordinator.single_type_managers: + pending_copies.extend(mgr.take_pending_cow_copies()) + copies = [ + KVCacheBlockCopy( + src_block_id=source_block.block_id, + dst_block_id=cow_block.block_id, + ) + for source_block, cow_block in pending_copies + ] + retained_blocks = [block for pair in pending_copies for block in pair] + return copies, retained_blocks + + def take_partial_tail_offloads(self) -> dict[str, list[tuple[int, int, int]]]: + """Drain producer partial-tail offload hand-offs per request. + + Returns ``{request_id: [(group_id, block_id, boundary_tokens), ...]}`` + for the durable boundary blocks of producers' last-prompt-boundary + partial tails. Only mamba "align" groups contribute; empty otherwise. + A KV connector reads the referenced blocks and offloads them so a later + request can hit the sub-block prefix. + + Each handed-off block lives off the request block table, so it is + pinned here and unpinned when the request's blocks are freed — for a + producer with saved tokens, after the connector reports sends done. + """ + offloads: dict[str, list[tuple[int, int, int]]] = {} + for mgr in self.coordinator.single_type_managers: + for ( + req_id, + group_id, + block, + boundary_tokens, + boundary_hash, + ) in mgr.take_pending_partial_tail_offloads(): + block = self._pin_recurrent_boundary( + req_id, + block, + boundary_tokens, + boundary_hash, + ) + offloads.setdefault(req_id, []).append( + (group_id, block.block_id, boundary_tokens) + ) + return offloads + + def take_recurrent_boundary_blocks( + self, + partial_tail_offloads: dict[str, list[tuple[int, int, int]]] | None = None, + ) -> dict[str, list[tuple[int, int, int]]]: + """Drain hash-proven recurrent boundaries for connector metadata. + + The returned union preserves the partial-tail entries already exposed + to connectors and adds full Mamba pages at the scheduler replay + boundary. Pins use the existing request-cleanup lifetime. + """ + boundaries = { + req_id: list(entries) + for req_id, entries in (partial_tail_offloads or {}).items() + } + for mgr in self.coordinator.single_type_managers: + for ( + req_id, + group_id, + block, + boundary_tokens, + ) in mgr.take_pending_aligned_recurrent_boundaries(): + assert block.block_hash is not None + block = self._pin_recurrent_boundary( + req_id, + block, + boundary_tokens, + block.block_hash, + ) + entry = (group_id, block.block_id, boundary_tokens) + request_entries = boundaries.setdefault(req_id, []) + if entry not in request_entries: + request_entries.append(entry) + return boundaries + + def _pin_recurrent_boundary( + self, + request_id: str, + block: KVCacheBlock, + boundary_tokens: int, + boundary_hash: BlockHashWithGroupId, + ) -> KVCacheBlock: + """Pin one exact boundary without duplicating a prior hand-off.""" + assert not block.is_null + assert self.block_pool.cached_block_hash_to_block.contain( + boundary_hash, block.block_id + ) + pins = self._partial_tail_pins.setdefault(request_id, []) + for pinned in pins: + if pinned is block and self.block_pool.cached_block_hash_to_block.contain( + boundary_hash, pinned.block_id + ): + return pinned + self.block_pool.touch((block,)) + pins.append(block) + return block + + def discard_aligned_recurrent_boundaries(self) -> None: + """Drop aligned hand-offs when the producer connector did not opt in.""" + for mgr in self.coordinator.single_type_managers: + mgr.take_pending_aligned_recurrent_boundaries() + + def new_step_starts(self) -> None: + """Notify the coordinator that a new step is starting.""" + self.coordinator.new_step_starts() diff --git a/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/test_checkpoint_package.py b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/test_checkpoint_package.py new file mode 100644 index 00000000..7ba8cb87 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/checkpoints/test_checkpoint_package.py @@ -0,0 +1,59 @@ +"""Source identity and planner boundary checks for the checkpoint payload.""" +import importlib.util +import json +from pathlib import Path +import subprocess +import sys + +import pytest + +ROOT = Path(__file__).resolve().parent +spec = importlib.util.spec_from_file_location("checkpoint_installer", ROOT / "install.py") +installer = importlib.util.module_from_spec(spec) +spec.loader.exec_module(installer) +MANIFEST_SHA = "0970d29ec33e9f8525a2cc55989ab0deb937bff88b5f16d4d5280035359e4c55" + + +def test_packaged_payload_identity_and_syntax(): + result = subprocess.run([sys.executable, str(ROOT / "install.py"), "verify-context", + "--manifest-sha256", MANIFEST_SHA], capture_output=True, text=True) + assert result.returncode == 0, result.stderr + assert json.loads(result.stdout)["payloads_verified"] == 18 + + +def test_preimage_mismatch_and_unexpected_file_are_rejected(tmp_path): + path = tmp_path / "module.py" + path.write_bytes(b"source") + installer.verify_files({"/module.py": installer.sha(b"source")}, root=tmp_path) + with pytest.raises(RuntimeError, match="identity mismatch"): + installer.verify_files({"/module.py": installer.sha(b"different")}, root=tmp_path) + with pytest.raises(RuntimeError, match="Expected absent"): + installer.verify_files({"/module.py": None}, root=tmp_path) + + +def planner(): + manifest = json.loads((ROOT / "patch-manifest.json").read_bytes()) + row = next(v for k, v in manifest["replacements"].items() + if k.endswith("/recurrent_prefill_checkpoint.py")) + namespace = {} + exec(compile((ROOT / row["payload"]).read_bytes(), "checkpoint_planner", "exec"), namespace) + return namespace + + +def test_fresh_prompt_exports_publication_and_predecessor(): + functions = planner() + plan = functions["fresh_prompt_plan"](start=0, end=8192, prompt=8192, + num_tokens=8192, block_size=512, publications=(6144,), shared_prefix_boundary=0) + assert plan == (0, 8192, (6144, 7168)) + assert functions["checkpoint_metadata"](plan, 0, 8192, 512, 2) == ([6144, 7168], [11, 13]) + with pytest.raises(ValueError, match="actual query span"): + functions["checkpoint_metadata"](plan, 512, 8192, 512, 2) + with pytest.raises(ValueError, match="capacity"): + functions["checkpoint_metadata"](plan, 0, 8192, 512, 1) + + +@pytest.mark.parametrize("start,end,prompt,tokens", [(512, 8192, 8192, 7680), + (0, 16384, 16384, 16384), (0, 8191, 8191, 8191), (0, 4096, 8192, 4096)]) +def test_ineligible_prefill_uses_ordinary_scheduler(start, end, prompt, tokens): + assert planner()["fresh_prompt_plan"](start=start, end=end, prompt=prompt, + num_tokens=tokens, block_size=512, publications=()) is None diff --git a/runtime/glm53-spark-mtp3-mesh/performance/install.py b/runtime/glm53-spark-mtp3-mesh/performance/install.py new file mode 100644 index 00000000..7b173a94 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/install.py @@ -0,0 +1,105 @@ +"""Install verified cache, recurrent checkpoints, and reasoning contracts.""" + +# ruff: noqa: E402 +import hashlib +import json +from pathlib import Path +import shutil +import subprocess +import sys + +SOURCE = Path(__file__).resolve().parent +SITE = Path("/usr/local/lib/python3.12/dist-packages") +context = json.loads((SOURCE / "context.json").read_text()) +for name, expected in context["files"].items(): + if hashlib.sha256((SOURCE / name).read_bytes()).hexdigest() != expected: + raise ValueError(f"Build input checksum differs: {name}") +sys.path.insert(0, str(SOURCE / "reuse")) +import compose + +compose.source_bytes(SITE) +import patch_mtp3_barrier as barrier +import patch_mtp3_lease_accounting as accounting +import patch_mtp3_local_lease_preference as preference +import patch_mtp3_sparse_retention as retention +import patch_mtp3_partial_tail_eligibility as eligibility + +scheduler = SITE / "vllm/v1/core/sched/scheduler.py" +barrier.apply_patch(SITE / "b12x/attention/dsa_indexer/fused_indexer.py") +accounting.apply_patch(scheduler) +preference.apply_patch(scheduler) +retention.apply_patch(SITE / "vllm") +eligibility.apply_patch(scheduler) +subprocess.run( + [ + sys.executable, + "-S", + "-B", + str(SOURCE / "checkpoints/install.py"), + "apply", + "--manifest-sha256", + "0970d29ec33e9f8525a2cc55989ab0deb937bff88b5f16d4d5280035359e4c55", + ], + check=True, +) +sys.path.insert(0, str(SOURCE / "reasoning")) +from patch_contract import apply + +apply(SITE, Path("/opt/sparkring/bin/warmup_dflash.py")) +shutil.copytree(SOURCE / "sparkcache", SITE / "sparkcache", dirs_exist_ok=True) +contract = ( + SITE / "sparkcache/runtime_patches/vllm-manager-page-async-contract-55969c16.json" +) +contract_source = SOURCE / "checkpoints/ownership-contract.json" +if ( + hashlib.sha256(contract_source.read_bytes()).hexdigest() + != "9fbd8e2596287fd4e12aaf042e189bccfee2ea92126de42f58da8e44fb4d9c55" +): + raise ValueError("Checkpoint ownership-contract template differs") +data = json.loads(contract_source.read_text()) +# Keep the complete checkpoint ownership surface and symbol requirements. +# Every expected runtime digest must match before the contract is installed. +for row in data["files"]: + path = SITE / row["path"] + if row["sha256"] != hashlib.sha256(path.read_bytes()).hexdigest(): + raise ValueError(f"Ownership dependency differs: {row['path']}") +contract.write_text(json.dumps(data, indent=2) + "\n", encoding="utf-8") +shutil.copytree(SOURCE / "bundle", Path("/opt/spark-sircl"), dirs_exist_ok=True) +shutil.copyfile( + SOURCE / "libspark_cache_placement.so", + Path( + "/opt/sparkcache-src/sparkcache/native/build-cuda/libspark_cache_placement.so" + ), +) +files = {} +for root in ( + SITE / "sparkcache", + SITE / "vllm", + SITE / "b12x", + Path("/opt/spark-sircl"), +): + for path in root.rglob("*"): + if path.is_file() and "__pycache__" not in path.parts and path.suffix != ".pyc": + files[str(path)] = hashlib.sha256(path.read_bytes()).hexdigest() +for path in ( + Path("/opt/sparkring/bin/warmup_dflash.py"), + Path( + "/opt/sparkcache-src/sparkcache/native/build-cuda/libspark_cache_placement.so" + ), +): + files[str(path)] = hashlib.sha256(path.read_bytes()).hexdigest() +receipt = Path("/opt/sparkring/receipts/mtp3-performance.json") +receipt.write_text( + json.dumps( + { + "schema": "sparkring-mtp3-performance-image/v1", + "status": "research-only", + "sparkcache_commit": context["sparkcache_commit"], + "files": files, + }, + sort_keys=True, + ) + + "\n" +) +shutil.copyfile(SOURCE / "verify.py", Path("/opt/sparkring/bin/verify-performance.py")) +shutil.copyfile(SOURCE / "start.py", Path("/opt/sparkring/bin/start-performance.py")) diff --git a/runtime/glm53-spark-mtp3-mesh/performance/prepare.py b/runtime/glm53-spark-mtp3-mesh/performance/prepare.py new file mode 100644 index 00000000..e2d56e53 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/prepare.py @@ -0,0 +1,114 @@ +"""Prepare a source-pinned GLM MTP3 performance image without remote actions.""" + +import argparse +import hashlib +import io +import json +from pathlib import Path +import shutil +import subprocess +import tarfile + +HERE = Path(__file__).resolve().parent +CACHE_COMMIT = "48bbd2be4a7b972e56632a2d7b934bac5460f272" +BASE_IMAGE = "ghcr.io/fujitsupolycom/sparkring-glm53-sparkcache@sha256:67dc0ae453baaae6831ccec1d259b4ef8b236a8b0dc9f747d901b95c66ec1987" +BASE_ID = "sha256:2e41b1e934a85ff7c21b780532db2f0a0e978df081e52f4ae2bf11f8992fb24f" +PLACEMENT = "2657cdd2e54a097c9544e4c79ae62c0646db6db123ff24e4f0c384238c3a1e8d" +TRANSPORT = "056243fad27d224b82e437925ffa2aed42037e6bd29f239f56076a832f6ca5cb" + + +def prepare(cache, placement, transport, output): + if output.exists(): + raise ValueError("Image context destination must not exist") + for path, expected in ((placement, PLACEMENT), (transport, TRANSPORT)): + if hashlib.sha256(path.read_bytes()).hexdigest() != expected: + raise ValueError( + f"Native artifact differs from qualified input: {path.name}" + ) + revision = subprocess.check_output( + ["git", "-C", str(cache), "rev-parse", "HEAD"], text=True + ).strip() + if revision != CACHE_COMMIT or subprocess.check_output( + ["git", "-C", str(cache), "status", "--porcelain"] + ): + raise ValueError("SparkCache checkout must be clean at the pinned revision") + archive = subprocess.check_output( + ["git", "-C", str(cache), "archive", "HEAD", "sparkcache"] + ) + output.mkdir(parents=True) + with tarfile.open(fileobj=io.BytesIO(archive)) as bundle: + for member in bundle.getmembers(): + name = Path(member.name) + if ( + name.is_absolute() + or ".." in name.parts + or not name.parts + or name.parts[0] != "sparkcache" + ): + raise ValueError("Source archive path escaped package") + if member.isfile(): + path = output / name + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(bundle.extractfile(member).read()) + for name in ("checkpoints", "reasoning"): + shutil.copytree( + HERE / name, + output / name, + ignore=shutil.ignore_patterns("__pycache__", ".pytest_cache"), + ) + shutil.copytree( + HERE.parent / "experiments/cache-reuse", + output / "reuse", + ignore=shutil.ignore_patterns("__pycache__", ".pytest_cache"), + ) + subprocess.run( + [ + "python", + str(HERE / "transport/package.py"), + "bundle", + "--native-library", + str(transport), + "--native-sha256", + TRANSPORT, + "--output", + str(output / "bundle"), + ], + check=True, + ) + shutil.copyfile(placement, output / "libspark_cache_placement.so") + for name in ("install.py", "verify.py", "start.py", "Dockerfile"): + shutil.copyfile(HERE / name, output / name) + files = { + p.relative_to(output).as_posix(): hashlib.sha256(p.read_bytes()).hexdigest() + for p in output.rglob("*") + if p.is_file() + } + record = { + "schema": "sparkring-mtp3-performance-context/v1", + "status": "implemented", + "sparkcache_commit": CACHE_COMMIT, + "base_reference": BASE_IMAGE, + "base_image_id": BASE_ID, + "files": files, + } + (output / "context.json").write_text( + json.dumps(record, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + return record + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--sparkcache", type=Path, required=True) + parser.add_argument("--placement-library", type=Path, required=True) + parser.add_argument("--transport-library", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + result = prepare( + args.sparkcache, args.placement_library, args.transport_library, args.output + ) + print( + json.dumps( + {"source": result["sparkcache_commit"], "files": len(result["files"])} + ) + ) diff --git a/runtime/glm53-spark-mtp3-mesh/performance/public-image.json b/runtime/glm53-spark-mtp3-mesh/performance/public-image.json new file mode 100644 index 00000000..90524347 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/public-image.json @@ -0,0 +1,35 @@ +{ + "schema": "sparkring-mtp3-performance-public-image/v1", + "status": "research-only", + "image_reference": "ghcr.io/fujitsupolycom/sparkring-glm53-sparkcache@sha256:11a556a54041fd823d152a7f051ac4f7c617dc539030df26e93008392fee0746", + "image_id": "sha256:6921a6c163ea40b603e19a0332330efe3dbccbf4dce9f6cbbf6b756c9231835a", + "platform": "linux/arm64", + "anonymous_manifest_verified": true, + "sparkcache_commit": "48bbd2be4a7b972e56632a2d7b934bac5460f272", + "bundle_manifest_sha256": "c0fd5567442b08b908cc193f36d0864e262573c7e5d232509479a823cface742", + "native_placement_sha256": "2657cdd2e54a097c9544e4c79ae62c0646db6db123ff24e4f0c384238c3a1e8d", + "source_equivalence": { + "checked_files": 5308, + "different_files": 0, + "comparison_image_id": "sha256:84c4546d37e8504fc98a359435b262a5d8e7bfe9c892af7429e876db5bc98422", + "scope": "vLLM, B12X, transport bundle, and warmup files; excludes SparkCache main source changes" + }, + "verification": { + "files": 5472, + "model_loaded": false, + "cuda_initialized": false, + "marker_source_sha256": "8684a6961b8e86aa474fa2310ff71e4cdf219a63a72ceb5593b2f95e54812792", + "marker_binary_sha256": "2828c07e4255c4962c77425be2c88969e7eb7dd4b1bf9e36485bc705bb5d6d64", + "readiness_warmup": { + "helper_sha256": "959cd1601f9a8c16124957aa03a94e2f3df2efab431a0779b46df1f8a868609c", + "temperature": 1.0, + "reasoning_effort": "low" + }, + "independent_host_file_check": true + }, + "limits": [ + "Image file verification and source equivalence do not constitute a serving soak of this rebuilt image" + ], + "cache_namespace": "glm53-spark-mtp3-cache-checkpoints-c0fd5567-tail-cow-v2", + "registry_pull_verified": true +} diff --git a/runtime/glm53-spark-mtp3-mesh/performance/reasoning/chat-serving-source.py.gz b/runtime/glm53-spark-mtp3-mesh/performance/reasoning/chat-serving-source.py.gz new file mode 100644 index 00000000..0938bec5 Binary files /dev/null and b/runtime/glm53-spark-mtp3-mesh/performance/reasoning/chat-serving-source.py.gz differ diff --git a/runtime/glm53-spark-mtp3-mesh/performance/reasoning/patch_contract.py b/runtime/glm53-spark-mtp3-mesh/performance/reasoning/patch_contract.py new file mode 100644 index 00000000..f5a738c6 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/reasoning/patch_contract.py @@ -0,0 +1,58 @@ +"""Reject unsupported reasoning suppression before GLM-5.3 prompt generation.""" + +import ast +import hashlib +from pathlib import Path + +ANCHOR = ( + " chat_template_kwargs = self._effective_chat_template_kwargs(request)\n" +) +GUARD = """ # GLM-5.3's template always opens a reasoning segment. Disabling + # its parser would expose reasoning as answer text without changing + # model generation. Reject the unsupported mode before rendering. + if getattr(self.model_config.hf_config, "model_type", None) == "glm5_next": + disabled = any( + chat_template_kwargs.get(key) is False + for key in ("enable_thinking", "thinking") + ) + if disabled or request.reasoning_effort == "none": + return self.create_error_response( + "GLM-5.3-Flash does not support disabling thinking with its " + "chat template. Use reasoning_effort='low', 'high', or 'max'." + ) +""" + + +def patched(source): + assert source.count(ANCHOR) == 1 + result = source.replace(ANCHOR, ANCHOR + GUARD, 1) + ast.parse(result) + return result + + +def apply(site_packages, warmup_path): + path = Path(site_packages) / "vllm/entrypoints/openai/chat_completion/serving.py" + source = path.read_text(encoding="utf-8") + if ( + hashlib.sha256(source.encode()).hexdigest() + != "9982953285e9df469032a82fffa4095d0e9d86278bede6e2b91d03d02373d182" + ): + raise ValueError("GLM chat serving source differs from the attested preimage") + result = patched(source).encode() + warmup = Path(warmup_path) + data = warmup.read_text(encoding="utf-8") + needle = '"chat_template_kwargs": {"enable_thinking": False},' + if data.count(needle) != 1: + raise ValueError("GLM warmup request differs from the supported source") + data = data.replace( + needle, + '"chat_template_kwargs": {"enable_thinking": True}, "reasoning_effort": "low",', + 1, + ) + compile(data, str(warmup), "exec") + path.write_bytes(result) + warmup.write_bytes(data.encode()) + return { + str(path): hashlib.sha256(result).hexdigest(), + str(warmup): hashlib.sha256(data.encode()).hexdigest(), + } diff --git a/runtime/glm53-spark-mtp3-mesh/performance/reasoning/test_contract.py b/runtime/glm53-spark-mtp3-mesh/performance/reasoning/test_contract.py new file mode 100644 index 00000000..89f508b0 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/reasoning/test_contract.py @@ -0,0 +1,93 @@ +import ast +import asyncio +from pathlib import Path +from types import SimpleNamespace + +import pytest +from patch_contract import patched + +import gzip + +SOURCE = gzip.decompress( + (Path(__file__).parent / "chat-serving-source.py.gz").read_bytes() +).decode() + + +def request_path(source): + tree = ast.parse(source) + function = next( + node + for node in ast.walk(tree) + if isinstance(node, ast.AsyncFunctionDef) + and node.name == "_create_chat_completion" + ) + # Execute the actual serving prefix through parser construction. A renderer + # sentinel identifies whether an unsupported request reaches model work. + stop = next( + i + for i, node in enumerate(function.body) + if isinstance(node, ast.Assign) and isinstance(node.value, ast.Await) + ) + body = function.body[:stop] + ast.parse("return 'rendering'").body + wrapper = ast.parse("async def invoke(self, request):\n pass").body[0] + wrapper.body = body + scope = {"Parser": object} + exec( + compile( + ast.fix_missing_locations(ast.Module(body=[wrapper], type_ignores=[])), + "", + "exec", + ), + scope, + ) + return scope["invoke"] + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize( + "kwargs,effort", + [({"enable_thinking": False}, None), ({"thinking": False}, None), ({}, "none")], +) +def test_glm53_rejects_unsupported_mode_before_parser_or_rendering( + kwargs, effort, stream +): + server = SimpleNamespace( + renderer=SimpleNamespace(tokenizer=object()), + model_config=SimpleNamespace(hf_config=SimpleNamespace(model_type="glm5_next")), + _effective_chat_template_kwargs=lambda request: kwargs, + parser_cls=None, + create_error_response=lambda message: {"error": message}, + ) + request = SimpleNamespace(reasoning_effort=effort, stream=stream) + assert asyncio.run(request_path(SOURCE)(server, request)) == "rendering" + result = asyncio.run(request_path(patched(SOURCE))(server, request)) + assert "does not support disabling thinking" in result["error"] + + +@pytest.mark.parametrize( + "model,kwargs,effort", + [ + ("glm5_next", {}, None), + ("glm5_next", {}, "low"), + ("glm5_next", {}, "high"), + ("glm5_next", {}, "max"), + ("glm5_next", {"enable_thinking": True}, None), + ("glm4_moe", {"enable_thinking": False}, "none"), + ], +) +def test_supported_modes_and_other_models_preserve_rendering(model, kwargs, effort): + server = SimpleNamespace( + renderer=SimpleNamespace(tokenizer=object()), + model_config=SimpleNamespace(hf_config=SimpleNamespace(model_type=model)), + _effective_chat_template_kwargs=lambda request: kwargs, + parser_cls=None, + create_error_response=lambda message: {"error": message}, + ) + assert ( + asyncio.run( + request_path(patched(SOURCE))( + server, SimpleNamespace(reasoning_effort=effort) + ) + ) + == "rendering" + ) diff --git a/runtime/glm53-spark-mtp3-mesh/performance/ruff.toml b/runtime/glm53-spark-mtp3-mesh/performance/ruff.toml new file mode 100644 index 00000000..b1b87383 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/ruff.toml @@ -0,0 +1,2 @@ +# Content-addressed upstream inputs are checked byte-for-byte by package tests. +extend-exclude = ["checkpoints/payload-by-sha", "transport/bundle-source"] diff --git a/runtime/glm53-spark-mtp3-mesh/performance/start.py b/runtime/glm53-spark-mtp3-mesh/performance/start.py new file mode 100644 index 00000000..bab8d72f --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/start.py @@ -0,0 +1,13 @@ +"""Verify immutable runtime files before executing model warmup and serving.""" + +import os +import subprocess +import sys + +subprocess.run( + [sys.executable, "-S", "-B", "/opt/sparkring/bin/verify-performance.py"], check=True +) +os.execv( + "/opt/sparkring/bin/serve-with-warmup.py", + ["/opt/sparkring/bin/serve-with-warmup.py", *sys.argv[1:]], +) diff --git a/runtime/glm53-spark-mtp3-mesh/performance/test_performance_profile.py b/runtime/glm53-spark-mtp3-mesh/performance/test_performance_profile.py new file mode 100644 index 00000000..bf0aec28 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/test_performance_profile.py @@ -0,0 +1,80 @@ +"""Verify profile selection without host operations or GPU access.""" + +import importlib.util +import json +from pathlib import Path +import shutil +import sys + +import pytest + +HERE = Path(__file__).resolve().parent + + +def load(name, path): + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +profile = load("performance_profile_contract", HERE.parent / "profile.py") +example = load("performance_profile_examples", HERE.parent / "make_example.py") + + +def test_performance_receipt_rejects_modified_identity(tmp_path): + receipt = json.loads((HERE / "public-image.json").read_text()) + receipt["image_id"] = "sha256:" + "0" * 64 + path = tmp_path / "receipt.json" + path.write_text(json.dumps(receipt)) + with pytest.raises(ValueError, match="repository pin"): + profile.load_image_receipt(path) + + +def test_renderer_uses_matched_image_native_and_namespace(tmp_path, monkeypatch): + receipt = profile.load_image_receipt(HERE / "public-image.json") + bundle = tmp_path / "bundle" + shutil.copytree(HERE / "transport/bundle-source", bundle) + native = bundle / "libspark_transport_capi.so" + native.write_bytes(b"fixture") + original = profile.sha + + def digest(path): + if Path(path) == native: + return "056243fad27d224b82e437925ffa2aed42037e6bd29f239f56076a832f6ca5cb" + return original(path) + + monkeypatch.setattr(profile, "sha", digest) + (tmp_path / "fabric.example.json").write_text( + json.dumps(example.topology_example()) + ) + site = tmp_path / "site.json" + site.write_text(json.dumps(example.site_example())) + output = tmp_path / "launch" + result = profile.render(site, bundle, output, HERE / "public-image.json") + assert result["image"]["image_id"] == receipt["image_id"] + assert result["bundle_manifest_sha256"] == receipt["bundle_manifest_sha256"] + for rank in range(4): + env = profile.defaults(output / f"rank{rank}.env") + assert env["IMAGE_REF"] == receipt["image_reference"] + assert ( + env["SPARKCACHE_PLACEMENT_LIBRARY_SHA256"] + == receipt["native_placement_sha256"] + ) + assert env["SPARKCACHE_CACHE_NAMESPACE"] == receipt["cache_namespace"] + assert env["SPARKRING_WARMUP_TEMPERATURE"] == "1" + + +def test_performance_recipe_and_guide_use_same_contract(): + root = HERE.parents[2] + recipe = json.loads( + (root / "recipes/glm53-mtp3-cache-checkpoints-tp4.json").read_text() + ) + receipt = json.loads((HERE / "public-image.json").read_text()) + assert recipe["runtime"]["image"] == receipt["image_reference"] + assert recipe["runtime"]["image_id"] == receipt["image_id"] + assert recipe["sparkcache"]["source_commit"] == receipt["sparkcache_commit"] + assert recipe["sparkcache"]["periodic_full_capture_interval_tokens"] == 0 + guide = (root / recipe["runtime"]["quickstart"]).read_text() + assert receipt["image_reference"] in guide and receipt["image_id"] in guide diff --git a/runtime/glm53-spark-mtp3-mesh/performance/transport/B12X-LICENSE b/runtime/glm53-spark-mtp3-mesh/performance/transport/B12X-LICENSE new file mode 100644 index 00000000..261eeb9e --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/transport/B12X-LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/runtime/glm53-spark-mtp3-mesh/performance/transport/README.md b/runtime/glm53-spark-mtp3-mesh/performance/transport/README.md new file mode 100644 index 00000000..3bf2803e --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/transport/README.md @@ -0,0 +1,85 @@ +# MTP3 mesh transport sources + +Status: **implemented**. This package prepares source trees and immutable overlay +directories offline. It does not install services, contact hosts, start models, +or qualify a rebuilt binary. + +The fused-prefill transport synchronizes each CTA before its leader enters the +cross-CTA barrier. Its CPU proxy starts only after all queue, mutex, and atomic +members have finished construction. Payload and doorbell work requests use one +linked provider submission on each QP. The doorbell alone is signaled; its +completion retires both requests. Partial posting errors remain fatal and must +not be retried. + +The CUDA graph all-reduce adapter uses two blocks for payloads through 16 KiB, +four through 64 KiB, and eight for larger payloads. Eager calls use eight blocks. +Weighted arrival counters preserve eight arrivals per operation. All ranks must +use the same source identity and geometry contract. + +## Source and artifact identities + +`source-manifest.json` lists every packaged native and overlay source file with +its SHA256. `native-source.tar.gz` contains the CMake project, compilation +sources, headers, test sources, and Apache license. It excludes Python bytecode +and unrelated research directories. The native source bytes derive from the +archive identified by `source_archive_provenance_sha256`; per-file identities +are preserved, while archive layout and compression are normalized. + +The native library used for four-rank checks has SHA256 +`056243fad27d224b82e437925ffa2aed42037e6bd29f239f56076a832f6ca5cb`. +The overlay manifest used with that library has SHA256 +`c0fd5567442b08b908cc193f36d0864e262573c7e5d232509479a823cface742`. +`bundle-source/` preserves its text/source inventory; the native binary is not +versioned. The B12X RoCE sources derive from commit +`eac260a8257cc6b14e7d4ad674f51e9a09b8790f` with graph-geometry modifications +recorded by the bundle manifest. `B12X-LICENSE` supplies their Apache license. + +The bundle configuration specifies the four-rank Spark mesh routing contract. +It contains no host credentials or deployment paths. It is not a topology +discovery tool: operators must qualify the physical mesh against that contract. + +## Build and prepare + +Run these commands from this directory on an ARM64 Linux build host with CMake +3.24 or later, a CUDA toolkit supporting architecture 121, a C++17 compiler, +and libibverbs development headers/libraries: + +```sh +python3 package.py verify +python3 package.py extract --output /tmp/mtp3-transport-source +cmake -S /tmp/mtp3-transport-source/spark_transport -B /tmp/mtp3-transport-build \ + -DCMAKE_BUILD_TYPE=Release -DCMAKE_CUDA_ARCHITECTURES=121 \ + -DBUILD_TESTING=ON -DSPARK_TP4_ENABLE_FUSED_STREAM_SWITCH_SMOKE=ON +cmake --build /tmp/mtp3-transport-build -j4 +ctest --test-dir /tmp/mtp3-transport-build --output-on-failure +sha256sum /tmp/mtp3-transport-build/libspark_transport_capi.so +python3 package.py bundle \ + --native-library /tmp/mtp3-transport-build/libspark_transport_capi.so \ + --native-sha256 SHA256_FROM_PRECEDING_COMMAND \ + --output /tmp/mtp3-transport-bundle +``` + +Output directories must not exist. All source hashes are checked before writing +output. The library digest must match the explicit argument. With the reference +library, the builder reproduces the reference overlay manifest byte for byte. +With a different library, it records that library's digest and emits a distinct +manifest with `research-only` status. Compilers and linkers can change binary +bytes; source availability does not establish bit-for-bit binary reproducibility. +Record the build environment and run native and serving checks before publishing +a rebuilt artifact as qualified. Generating a manifest alone proves no runtime +or ABI property of the supplied library. + +## Evidence and limits + +The reference library passed 28 native CTests and four four-rank sessions totaling +256 operations. Sessions covered changing exact and noninteger inputs, query +sizes 128/512/2048/8192, alternating streams, both operation slots, numerical +results, input and tail guards, and health completion. Two sessions used tracing +and therefore do not support timing claims. Provider-call tracing counted 144 +calls per operation with linked posting versus 240 separate calls; work-request +and completion counts were unchanged. Those observations establish bounded +correctness and submission-count evidence, not a model-throughput speedup. + +Offline package checks run with `python3 -m pytest test_transport_package.py -q`. They +validate inventories, corruption refusal, path restrictions, destination +preservation, and rebuilt-library manifest handling without CUDA or networking. diff --git a/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/b12x_overlay/b12x/comm/roce/__init__.py b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/b12x_overlay/b12x/comm/roce/__init__.py new file mode 100644 index 00000000..4eff8774 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/b12x_overlay/b12x/comm/roce/__init__.py @@ -0,0 +1,65 @@ +"""``b12x.comm.roce`` (RoCEnante): one-shot all-reduce and all-gather over RoCE for multi-node TP. + +Target: clusters of DGX Spark nodes joined by their ConnectX-7 200 GbE ports, +one GPU per node. The GB10's unified memory lets the NIC RDMA-write straight +into pinned host memory that the GPU kernel then reads in place, so no +GPUDirect RDMA (dmabuf/peermem) support is required. + +``AllReduce`` mirrors the ``comm.pcie.AllReduce`` surface (``from_exchange_group``, +``should_allreduce``, ``all_reduce``, ``for_stream``, ``capture``, ``close``) plus +``should_all_gather``/``all_gather`` for dim-0 and last-dim concatenation, so +integrations can dispatch to it behind the same adapter. See +``roce_oneshot.py`` for the protocol and constraints. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from ..._lib.meta import OpMeta, Provenance, install_lazy_api + +META = OpMeta( + name="roce", + group="comm", + api_style="stateful", + entry_points=( + "API_VERSION", + "AllReduce", + "DEFAULT_MAX_GATHER_BYTES", + "DEFAULT_MAX_SIZE", + "SUPPORTED_DTYPES", + "SUPPORTED_WORLD_SIZES", + "default_gid_index", + "discover_hcas", + "is_supported", + ), + archs=("sm121a",), + dtypes=("bf16", "fp16", "fp32", "int32", "int64"), + requires=("multi_node", "rdma"), + provenance=Provenance( + repo="https://github.com/local-inference-lab/b12x", + commit="b9e450f1", + paths=("b12x/comm/roce/",), + ), + test_path="tests/comm/test_roce_oneshot_gpu.py", + since="1.3.0", + notes=( + "RoCEnante. Python/CuTe DSL kernels; the RDMA proxy is a small C file " + "built with the host C compiler at first use (no CUDA extension build)." + ), +) + +if TYPE_CHECKING: # static analysis only; runtime resolution is lazy + from .api import ( # noqa: F401 + API_VERSION, + DEFAULT_MAX_GATHER_BYTES, + DEFAULT_MAX_SIZE, + SUPPORTED_DTYPES, + SUPPORTED_WORLD_SIZES, + AllReduce, + default_gid_index, + discover_hcas, + is_supported, + ) + +install_lazy_api(globals(), META) diff --git a/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/b12x_overlay/b12x/comm/roce/_allgather_cute.py b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/b12x_overlay/b12x/comm/roce/_allgather_cute.py new file mode 100644 index 00000000..c173aa16 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/b12x_overlay/b12x/comm/roce/_allgather_cute.py @@ -0,0 +1,394 @@ +"""CuTe DSL kernel for the RoCE one-shot all-gather. + +Same transport and protocol as the one-shot all-reduce (stage, doorbell, +wait on per-peer flags, advance the epoch) with the reduction replaced by a +strided copy that writes the concatenated output directly: + +* ``dim == 0`` concat: ``rows == 1``, output is shard 0, shard 1, ... in order; +* last-dim concat: each shard is ``rows`` rows of ``row_packs`` 16-byte packs + and shard ``s`` lands at column block ``s`` of every output row, so no + separate reshape/copy is needed after the collective. + +The local shard is copied from the input; peer shards are read in place from +the NIC-written slots with system-scope loads. Every launch of one runtime +uses the same grid (shared counters with the all-reduce kernel). +""" + +from __future__ import annotations + +import functools +from collections.abc import Callable + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +from cutlass import Int32, Int64, Uint32 + +from b12x._lib.compiler import KernelCompileSpec +from b12x._lib.compiler import compile as b12x_compile +from b12x._lib.runtime_control import raise_if_kernel_resolution_frozen +from b12x._lib.utils import current_cuda_stream, make_ptr + +from ._cute_intrinsics import ( + atomic_add_relaxed_gpu_u32, + fence_sc_gpu, + fence_sc_sys, + ld_global_v4_u32, + ld_relaxed_gpu_u32, + ld_relaxed_sys_u32, + ld_relaxed_sys_v4_u32, + spin_until_eq_acquire_sys, + st_global_v4_u32, + st_release_gpu_u32, + st_relaxed_sys_u32, +) + +PACK_BYTES = 16 +PATH_COUNT = 2 +_PREPARED_LAUNCHERS: set[tuple[object, ...]] = set() + + +class _RoceAllGatherLaunch: + def __init__( + self, + world_size: int, + rank: int, + threads: int, + slots: int, + flag_stride: int, + opposite_paths: int, + ) -> None: + """Bind one kernel specialization: world size, rank, and layout constants.""" + if int(threads) < int(world_size): + raise ValueError( + f"RoCE kernels need threads >= world_size (one thread waits on one " + f"peer flag), got threads={threads} world_size={world_size}" + ) + self._world_size = int(world_size) + self._rank = int(rank) + self._threads = int(threads) + self._slots = int(slots) + self._flag_stride = int(flag_stride) + self._opposite_paths = int(opposite_paths) + if self._opposite_paths not in (PATH_COUNT, 4): + raise ValueError("opposite_paths must be 2 or 4") + if self._opposite_paths == 4 and self._world_size != 4: + raise ValueError("four opposite paths require a four-rank group") + + @cute.jit + def __call__( + self, + input_ptr: cute.Pointer, + output_ptr: cute.Pointer, + shard_packs: Int32, + nbytes: Int32, + row_packs: Int32, + recv_base: Int64, + flag_base: Int64, + send_base: Int64, + ctrl_base: Int64, + slot_bytes: Int64, + epoch_ptr: Int64, + spin_limit: Uint32, + grid_x: Int32, + stream: cuda.CUstream, + ) -> None: + """Host entry: launch the all-gather kernel with runtime scalars.""" + self.kernel( + input_ptr, + output_ptr, + shard_packs, + nbytes, + row_packs, + recv_base, + flag_base, + send_base, + ctrl_base, + slot_bytes, + epoch_ptr, + spin_limit, + ).launch( + grid=(grid_x, 1, 1), + block=[self._threads, 1, 1], + cluster=(1, 1, 1), + stream=stream, + ) + + @cute.kernel + def kernel( + self, + input_ptr: cute.Pointer, + output_ptr: cute.Pointer, + shard_packs: Int32, + nbytes: Int32, + row_packs: Int32, + recv_base: Int64, + flag_base: Int64, + send_base: Int64, + ctrl_base: Int64, + slot_bytes: Int64, + epoch_ptr: Int64, + spin_limit: Uint32, + ) -> None: + """Device kernel: stage, doorbell, wait for peer flags, strided copy, advance the epoch.""" + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + gdim, _, _ = cute.arch.grid_dim() + input_base = Int64(input_ptr.toint()) + output_base = Int64(output_ptr.toint()) + stage_counter_ptr = epoch_ptr + Int64(4) + tail_counter_ptr = epoch_ptr + Int64(8) + + epoch = ld_relaxed_gpu_u32(epoch_ptr) + seq = epoch + Uint32(1) + slot = Int64(seq & Uint32(1)) + send_slot = send_base + slot * slot_bytes + + index = Int32(bidx) * Int32(self._threads) + Int32(tidx) + stride = Int32(gdim) * Int32(self._threads) + + # A recorded timeout poisons the runtime: later launches do nothing so + # the host sees the failure without waiting another spin limit per op. + # The device poison word (fourth counter) is written by the same waiting + # threads that write the host error word and only ever goes from 0 to + # the failed sequence, so a cheap GPU-scope load is enough here. + poison_ptr = epoch_ptr + Int64(12) + poisoned = ld_relaxed_gpu_u32(poison_ptr) + if poisoned == Uint32(0): + # 1. stage the local shard into the pinned send slot + stage_index = index + while stage_index < shard_packs: + words = ld_global_v4_u32(input_base + Int64(stage_index) * Int64(PACK_BYTES)) + st_global_v4_u32( + send_slot + Int64(stage_index) * Int64(PACK_BYTES), + words[0], + words[1], + words[2], + words[3], + ) + stage_index += stride + fence_sc_sys() + cute.arch.sync_threads() + + # 2. the last block to finish staging rings the proxy doorbell + if Int32(tidx) == Int32(0): + prior = atomic_add_relaxed_gpu_u32(stage_counter_ptr, Uint32(1)) + if (prior + Uint32(1)) % Uint32(gdim) == Uint32(0): + st_relaxed_sys_u32(ctrl_base + Int64(4), Uint32(nbytes)) + st_relaxed_sys_u32(ctrl_base + Int64(16) + slot * Int64(4), Uint32(nbytes)) + fence_sc_sys() + st_relaxed_sys_u32(ctrl_base, seq) + + # 3. wait for every path flag. Four-path mode assigns one flag to + # each of eight threads; its extra opposite flags occupy the two + # otherwise-unused columns in the receiver-local source row. + if cutlass.const_expr(self._opposite_paths == 4): + waiter = 0 + for peer in cutlass.range_constexpr(self._world_size): + if cutlass.const_expr(peer != self._rank): + peer_paths = ( + 4 if (peer - self._rank) % self._world_size == 2 else 2 + ) + for path in cutlass.range_constexpr(peer_paths): + if Int32(tidx) == Int32(waiter): + flag_row = peer if path < PATH_COUNT else self._rank + flag_column = path if path < PATH_COUNT else path - PATH_COUNT + flag_addr = flag_base + ( + (Int64(flag_row) * Int64(self._slots) + slot) + * Int64(PATH_COUNT) + + Int64(flag_column) + ) * Int64(self._flag_stride) + timed_out = spin_until_eq_acquire_sys( + flag_addr, seq, spin_limit + ) + if timed_out != Uint32(0): + st_relaxed_sys_u32( + ctrl_base + Int64(12), Uint32(peer) + ) + st_relaxed_sys_u32(ctrl_base + Int64(8), seq) + st_release_gpu_u32(poison_ptr, seq) + waiter += 1 + else: + if Int32(tidx) < Int32(self._world_size): + if Int32(tidx) != Int32(self._rank): + for path in cutlass.range_constexpr(PATH_COUNT): + flag_addr = flag_base + ( + (Int64(tidx) * Int64(self._slots) + slot) + * Int64(PATH_COUNT) + + Int64(path) + ) * Int64(self._flag_stride) + timed_out = spin_until_eq_acquire_sys( + flag_addr, seq, spin_limit + ) + if timed_out != Uint32(0): + st_relaxed_sys_u32( + ctrl_base + Int64(12), Uint32(tidx) + ) + st_relaxed_sys_u32(ctrl_base + Int64(8), seq) + st_release_gpu_u32(poison_ptr, seq) + cute.arch.sync_threads() + # A wait that timed out in this block leaves the peer slot unreliable: + # skip the data phase so nothing derived from it is stored. + failed = ld_relaxed_gpu_u32(poison_ptr) + if failed == Uint32(0): + # 4. concatenate: shard s occupies column block s of every output row + out_row_packs = Int32(self._world_size) * row_packs + for source in cutlass.range_constexpr(self._world_size): + copy_index = index + while copy_index < shard_packs: + row = copy_index // row_packs + col = copy_index - row * row_packs + dest = output_base + ( + Int64(row) * Int64(out_row_packs) + + Int64(source) * Int64(row_packs) + + Int64(col) + ) * Int64(PACK_BYTES) + if cutlass.const_expr(source == self._rank): + words = ld_global_v4_u32( + input_base + Int64(copy_index) * Int64(PACK_BYTES) + ) + else: + peer_slot = recv_base + ( + Int64(source) * Int64(self._slots) + slot + ) * slot_bytes + words = ld_relaxed_sys_v4_u32( + peer_slot + Int64(copy_index) * Int64(PACK_BYTES) + ) + st_global_v4_u32(dest, words[0], words[1], words[2], words[3]) + copy_index += stride + + # 5. the last block to finish publishes the next epoch + fence_sc_gpu() + cute.arch.sync_threads() + if Int32(tidx) == Int32(0): + prior = atomic_add_relaxed_gpu_u32(tail_counter_ptr, Uint32(1)) + if (prior + Uint32(1)) % Uint32(gdim) == Uint32(0): + fence_sc_gpu() + # Every block's timeout store precedes its tail arrival, so the + # error word is final here. A failed sequence keeps the epoch, + # which makes every later launch a no-op until the host raises. + if ld_relaxed_sys_u32(ctrl_base + Int64(8)) == Uint32(0): + st_release_gpu_u32(epoch_ptr, seq) + + +def _dummy(dtype, alignment: int): + """A CUDA tensor of ``dtype`` used to trace launcher argument types.""" + return make_ptr(dtype, 16, cute.AddressSpace.gmem, assumed_align=alignment) + + +def _process_key( + world_size: int, + rank: int, + threads: int, + slots: int, + flag_stride: int, + opposite_paths: int, + device_index: int, +) -> tuple[object, ...]: + """Cache key of one compiled launcher specialization.""" + return ( + int(world_size), + int(rank), + int(threads), + int(slots), + int(flag_stride), + int(opposite_paths), + int(device_index), + ) + + +def is_launcher_prepared(*key) -> bool: + """True when the launcher for ``key`` is already compiled.""" + return _process_key(*key) in _PREPARED_LAUNCHERS + + +@functools.cache +def get_launcher( + world_size: int, + rank: int, + threads: int, + slots: int, + flag_stride: int, + opposite_paths: int, + device_index: int, +) -> Callable[..., None]: + """Compile the launcher for ``key`` once and return it.""" + process_key = _process_key( + world_size, + rank, + threads, + slots, + flag_stride, + opposite_paths, + device_index, + ) + del device_index + launch = _RoceAllGatherLaunch( + world_size, rank, threads, slots, flag_stride, opposite_paths + ) + cache_key = ( + int(world_size), + int(rank), + int(threads), + int(slots), + int(flag_stride), + int(opposite_paths), + ) + raise_if_kernel_resolution_frozen("cute.compile", target=launch, cache_key=cache_key) + raw = b12x_compile( + launch, + _dummy(cutlass.Uint32, 16), + _dummy(cutlass.Uint32, 16), + 1, + 16, + 1, + 16, + 16, + 16, + 16, + 4096, + 16, + 1, + 1, + current_cuda_stream(), + compile_spec=KernelCompileSpec.from_key("comm.roce.allgather", 2, cache_key), + ) + + def run( + input_address: int, + output_address: int, + shard_packs: int, + nbytes: int, + row_packs: int, + recv_base: int, + flag_base: int, + send_base: int, + ctrl_base: int, + slot_bytes: int, + epoch_address: int, + spin_limit: int, + grid_x: int, + ) -> None: + """Launch the compiled kernel with runtime scalar arguments.""" + raw( + make_ptr(cutlass.Uint32, input_address, cute.AddressSpace.gmem, assumed_align=16), + make_ptr(cutlass.Uint32, output_address, cute.AddressSpace.gmem, assumed_align=16), + int(shard_packs), + int(nbytes), + int(row_packs), + int(recv_base), + int(flag_base), + int(send_base), + int(ctrl_base), + int(slot_bytes), + int(epoch_address), + int(spin_limit), + int(grid_x), + current_cuda_stream(), + ) + + _PREPARED_LAUNCHERS.add(process_key) + return run + + +__all__ = ["PACK_BYTES", "get_launcher", "is_launcher_prepared"] diff --git a/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/b12x_overlay/b12x/comm/roce/_cute_intrinsics.py b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/b12x_overlay/b12x/comm/roce/_cute_intrinsics.py new file mode 100644 index 00000000..b7b639b4 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/b12x_overlay/b12x/comm/roce/_cute_intrinsics.py @@ -0,0 +1,332 @@ +"""PTX intrinsics for the RoCE one-shot all-reduce kernel. + +Payload slots and flags live in pinned host memory that the NIC writes by +RDMA and the GPU reads in place, so every access to them is spelled out with +system scope: ``ld.relaxed.sys`` for payload packs, ``ld.acquire.sys`` for the +arrival flag, ``st.relaxed.sys`` plus ``fence.sc.sys`` for the doorbell the +proxy thread polls. Keeping them as small user ops makes the protocol +explicit and independent of compiler defaults. +""" + +from __future__ import annotations + +from typing import Tuple + +from cutlass import Float32, Int64, Uint32 +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import T, dsl_user_op + + +def _asm(result_type, operands, text, constraints, *, side_effects=True, loc=None, ip=None): + """Emit one inline PTX statement through the CuTe DSL and return its result.""" + return llvm.inline_asm( + result_type, + operands, + text, + constraints, + has_side_effects=side_effects, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def ld_relaxed_gpu_u32(addr: Int64, *, loc=None, ip=None) -> Uint32: + """GPU-scope relaxed 32-bit load.""" + return Uint32( + _asm( + T.i32(), + [Int64(addr).ir_value(loc=loc, ip=ip)], + "ld.relaxed.gpu.global.u32 $0, [$1];", + "=r,l", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def ld_relaxed_sys_u32(addr: Int64, *, loc=None, ip=None) -> Uint32: + """System-scope relaxed 32-bit load (sees host and NIC writes).""" + return Uint32( + _asm( + T.i32(), + [Int64(addr).ir_value(loc=loc, ip=ip)], + "ld.relaxed.sys.global.u32 $0, [$1];", + "=r,l", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def atomic_add_relaxed_gpu_u32(addr: Int64, value: Uint32, *, loc=None, ip=None) -> Uint32: + """GPU-scope relaxed atomic add; returns the prior value.""" + return Uint32( + _asm( + T.i32(), + [Int64(addr).ir_value(loc=loc, ip=ip), Uint32(value).ir_value(loc=loc, ip=ip)], + "atom.relaxed.gpu.global.add.u32 $0, [$1], $2;", + "=r,l,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def st_release_gpu_u32(addr: Int64, value: Uint32, *, loc=None, ip=None) -> None: + """GPU-scope release 32-bit store.""" + _asm( + None, + [Int64(addr).ir_value(loc=loc, ip=ip), Uint32(value).ir_value(loc=loc, ip=ip)], + "st.release.gpu.global.u32 [$0], $1;", + "l,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def st_relaxed_sys_u32(addr: Int64, value: Uint32, *, loc=None, ip=None) -> None: + """System-scope relaxed 32-bit store, visible to the host and the NIC.""" + _asm( + None, + [Int64(addr).ir_value(loc=loc, ip=ip), Uint32(value).ir_value(loc=loc, ip=ip)], + "st.relaxed.sys.global.u32 [$0], $1;", + "l,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def fence_sc_sys(*, loc=None, ip=None) -> None: + """Sequentially consistent system-scope fence.""" + _asm(None, [], "fence.sc.sys;", "", loc=loc, ip=ip) + + +@dsl_user_op +def fence_sc_gpu(*, loc=None, ip=None) -> None: + """Sequentially consistent GPU-scope fence.""" + _asm(None, [], "fence.sc.gpu;", "", loc=loc, ip=ip) + + +@dsl_user_op +def spin_until_eq_acquire_sys( + addr: Int64, expected: Uint32, limit: Uint32, *, loc=None, ip=None +) -> Uint32: + """Spin until the word at ``addr`` equals ``expected`` (system scope). + + Returns 0 on success and 1 after ``limit`` polls without a match, so a + dead peer or proxy surfaces as an error instead of a hung kernel. + """ + return Uint32( + _asm( + T.i32(), + [ + Int64(addr).ir_value(loc=loc, ip=ip), + Uint32(expected).ir_value(loc=loc, ip=ip), + Uint32(limit).ir_value(loc=loc, ip=ip), + ], + """ + { + .reg .pred pending, expired; + .reg .b32 seen, polls; + mov.u32 polls, 0; + mov.u32 $0, 0; + roce_wait: + ld.acquire.sys.global.u32 seen, [$1]; + setp.ne.u32 pending, seen, $2; + @!pending bra roce_done; + add.u32 polls, polls, 1; + setp.ge.u32 expired, polls, $3; + @!expired bra roce_wait; + mov.u32 $0, 1; + roce_done: + } + """, + "=r,l,r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def ld_relaxed_sys_v4_u32( + addr: Int64, *, loc=None, ip=None +) -> Tuple[Uint32, Uint32, Uint32, Uint32]: + """Load one 16-byte pack from NIC-written pinned memory without caching it.""" + result = _asm( + llvm.StructType.get_literal([T.i32(), T.i32(), T.i32(), T.i32()]), + [Int64(addr).ir_value(loc=loc, ip=ip)], + "ld.relaxed.sys.global.v4.u32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,l", + loc=loc, + ip=ip, + ) + return tuple( + Uint32(llvm.extractvalue(T.i32(), result, [i], loc=loc, ip=ip)) for i in range(4) + ) + + +@dsl_user_op +def ld_global_v4_u32(addr: Int64, *, loc=None, ip=None) -> Tuple[Uint32, Uint32, Uint32, Uint32]: + """Plain global 16-byte load as four 32-bit words.""" + result = _asm( + llvm.StructType.get_literal([T.i32(), T.i32(), T.i32(), T.i32()]), + [Int64(addr).ir_value(loc=loc, ip=ip)], + "ld.global.v4.u32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,l", + loc=loc, + ip=ip, + ) + return tuple( + Uint32(llvm.extractvalue(T.i32(), result, [i], loc=loc, ip=ip)) for i in range(4) + ) + + +@dsl_user_op +def st_global_v4_u32( + addr: Int64, v0: Uint32, v1: Uint32, v2: Uint32, v3: Uint32, *, loc=None, ip=None +) -> None: + """Plain global 16-byte store of four 32-bit words.""" + _asm( + None, + [ + Int64(addr).ir_value(loc=loc, ip=ip), + Uint32(v0).ir_value(loc=loc, ip=ip), + Uint32(v1).ir_value(loc=loc, ip=ip), + Uint32(v2).ir_value(loc=loc, ip=ip), + Uint32(v3).ir_value(loc=loc, ip=ip), + ], + "st.global.v4.u32 [$0], {$1, $2, $3, $4};", + "l,r,r,r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def u32_as_f32(value: Uint32, *, loc=None, ip=None) -> Float32: + """Reinterpret the bits of a 32-bit word as float32.""" + return Float32( + _asm( + T.f32(), + [Uint32(value).ir_value(loc=loc, ip=ip)], + "mov.b32 $0, $1;", + "=f,r", + side_effects=False, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def f32_as_u32(value: Float32, *, loc=None, ip=None) -> Uint32: + """Reinterpret the bits of a float32 as a 32-bit word.""" + return Uint32( + _asm( + T.i32(), + [Float32(value).ir_value(loc=loc, ip=ip)], + "mov.b32 $0, $1;", + "=r,f", + side_effects=False, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def unpack_bf16x2(value: Uint32, *, loc=None, ip=None) -> Tuple[Float32, Float32]: + """Split a packed bf16 pair into two float32 values, low half first.""" + result = _asm( + llvm.StructType.get_literal([T.f32(), T.f32()]), + [Uint32(value).ir_value(loc=loc, ip=ip)], + """ + { + .reg .b16 lo, hi; + mov.b32 {lo, hi}, $2; + cvt.f32.bf16 $0, lo; + cvt.f32.bf16 $1, hi; + } + """, + "=f,=f,r", + side_effects=False, + loc=loc, + ip=ip, + ) + return ( + Float32(llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)), + Float32(llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)), + ) + + +@dsl_user_op +def unpack_f16x2(value: Uint32, *, loc=None, ip=None) -> Tuple[Float32, Float32]: + """Split a packed fp16 pair into two float32 values, low half first.""" + result = _asm( + llvm.StructType.get_literal([T.f32(), T.f32()]), + [Uint32(value).ir_value(loc=loc, ip=ip)], + """ + { + .reg .b16 lo, hi; + mov.b32 {lo, hi}, $2; + cvt.f32.f16 $0, lo; + cvt.f32.f16 $1, hi; + } + """, + "=f,=f,r", + side_effects=False, + loc=loc, + ip=ip, + ) + return ( + Float32(llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)), + Float32(llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)), + ) + + +@dsl_user_op +def pack_f32x2_to_bf16x2(lo: Float32, hi: Float32, *, loc=None, ip=None) -> Uint32: + """Match two scalar ``__float2bfloat16`` conversions without saturation.""" + return Uint32( + _asm( + T.i32(), + [Float32(lo).ir_value(loc=loc, ip=ip), Float32(hi).ir_value(loc=loc, ip=ip)], + """ + { + .reg .b16 blo, bhi; + cvt.rn.bf16.f32 blo, $1; + cvt.rn.bf16.f32 bhi, $2; + mov.b32 $0, {blo, bhi}; + } + """, + "=r,f,f", + side_effects=False, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def pack_f32x2_to_f16x2(lo: Float32, hi: Float32, *, loc=None, ip=None) -> Uint32: + """Round two float32 values into a packed fp16 pair with ``lo`` in the low half.""" + return Uint32( + _asm( + T.i32(), + [Float32(lo).ir_value(loc=loc, ip=ip), Float32(hi).ir_value(loc=loc, ip=ip)], + "cvt.rn.f16x2.f32 $0, $2, $1;", + "=r,f,f", + side_effects=False, + loc=loc, + ip=ip, + ) + ) diff --git a/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/b12x_overlay/b12x/comm/roce/_oneshot_cute.py b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/b12x_overlay/b12x/comm/roce/_oneshot_cute.py new file mode 100644 index 00000000..47c0538b --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/b12x_overlay/b12x/comm/roce/_oneshot_cute.py @@ -0,0 +1,470 @@ +"""CuTe DSL kernel for the RoCE one-shot all-reduce. + +One launch performs a complete all-reduce for one message: + +1. stage: copy the device input into ``send[seq & 1]`` in pinned host memory; +2. doorbell: the last block to finish staging publishes ``nbytes`` (in the + per-slot word ``nbytes[seq & 1]``) and then ``seq`` in the control record + that the RDMA proxy thread polls. The doorbell is a level, not a queue: + a proxy that was descheduled across two doorbells finds ``seq`` two ahead + and posts both slots, which is why the byte count lives per slot; +3. wait: spin on both ``flag[peer][seq & 1][path]`` values for every peer (the + peer's proxy writes each flag after its stripe on the same reliable QP); a + wait that exceeds ``spin_limit`` polls records ``seq`` in the control + record's error word and the host raises instead of hanging; +4. reduce: sum the local input and every peer slot in fixed rank order, so all + ranks produce bit-identical output, and store the result; +5. epoch: the last block to finish reduction advances the device-resident + epoch, which makes the sequence number a runtime value rather than a launch + argument and keeps CUDA-graph replay correct. + +Staging arrivals and tail arrivals use two separate counters. A block that +stages nothing can pass the peer wait (peers do not depend on our doorbell) +and reach the tail before a slower block has staged, so one shared counter +would ring the doorbell early and publish stale bytes. + +Each all-reduce uses 2, 4, or 8 CTAs and increments both counters by 8/grid +per arriving CTA. Every successful operation contributes eight to each counter, +so arbitrary grid transitions and the unchanged eight-CTA all-gather preserve +the modulo-eight boundary. Message size is a runtime scalar. +""" + +from __future__ import annotations + +import functools +from collections.abc import Callable + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +from cutlass import Int32, Int64, Uint32 + +from b12x._lib.compiler import KernelCompileSpec +from b12x._lib.compiler import compile as b12x_compile +from b12x._lib.runtime_control import raise_if_kernel_resolution_frozen +from b12x._lib.utils import current_cuda_stream, make_ptr + +from ._cute_intrinsics import ( + atomic_add_relaxed_gpu_u32, + f32_as_u32, + fence_sc_gpu, + fence_sc_sys, + ld_global_v4_u32, + ld_relaxed_gpu_u32, + ld_relaxed_sys_u32, + ld_relaxed_sys_v4_u32, + pack_f32x2_to_bf16x2, + pack_f32x2_to_f16x2, + spin_until_eq_acquire_sys, + st_global_v4_u32, + st_release_gpu_u32, + st_relaxed_sys_u32, + u32_as_f32, + unpack_bf16x2, + unpack_f16x2, +) + +PACK_BYTES = 16 +PATH_COUNT = 2 +_DTYPE_PACK_ELEMS = {"float32": 4, "float16": 8, "bfloat16": 8} +_PREPARED_LAUNCHERS: set[tuple[object, ...]] = set() + + +class _RoceOneshotLaunch: + def __init__( + self, + dtype_name: str, + world_size: int, + rank: int, + threads: int, + slots: int, + flag_stride: int, + opposite_paths: int, + ) -> None: + """Bind one kernel specialization: dtype, world size, rank, and layout constants.""" + if dtype_name not in _DTYPE_PACK_ELEMS: + raise ValueError(f"unsupported RoCE one-shot dtype {dtype_name!r}") + if int(threads) < int(world_size): + raise ValueError( + f"RoCE kernels need threads >= world_size (one thread waits on one " + f"peer flag), got threads={threads} world_size={world_size}" + ) + self._dtype_name = dtype_name + self._pack_elems = _DTYPE_PACK_ELEMS[dtype_name] + self._world_size = int(world_size) + self._rank = int(rank) + self._threads = int(threads) + self._slots = int(slots) + self._flag_stride = int(flag_stride) + self._opposite_paths = int(opposite_paths) + if self._opposite_paths not in (PATH_COUNT, 4): + raise ValueError("opposite_paths must be 2 or 4") + if self._opposite_paths == 4 and self._world_size != 4: + raise ValueError("four opposite paths require a four-rank group") + + @cute.jit + def _accumulate_words( + self, + accumulator: cute.Tensor, + words, + initialize: cutlass.Constexpr[bool], + ) -> None: + """Add one 16-byte pack of ``dtype`` values to the float32 accumulator.""" + if cutlass.const_expr(self._dtype_name == "float32"): + for word in cutlass.range_constexpr(4): + value = u32_as_f32(words[word]) + if cutlass.const_expr(initialize): + accumulator[word] = value + else: + accumulator[word] = accumulator[word] + value + else: + for word in cutlass.range_constexpr(4): + if cutlass.const_expr(self._dtype_name == "float16"): + lo, hi = unpack_f16x2(words[word]) + else: + lo, hi = unpack_bf16x2(words[word]) + lane = word * 2 + if cutlass.const_expr(initialize): + accumulator[lane] = lo + accumulator[lane + 1] = hi + else: + accumulator[lane] = accumulator[lane] + lo + accumulator[lane + 1] = accumulator[lane + 1] + hi + + @cute.jit + def _store_accumulator(self, address: Int64, accumulator: cute.Tensor) -> None: + """Convert the accumulator back to ``dtype`` and store one 16-byte pack.""" + packed = cute.make_rmem_tensor((4,), cutlass.Uint32) + if cutlass.const_expr(self._dtype_name == "float32"): + for word in cutlass.range_constexpr(4): + packed[word] = f32_as_u32(accumulator[word]) + else: + for word in cutlass.range_constexpr(4): + lane = word * 2 + if cutlass.const_expr(self._dtype_name == "float16"): + packed[word] = pack_f32x2_to_f16x2( + accumulator[lane], accumulator[lane + 1] + ) + else: + packed[word] = pack_f32x2_to_bf16x2( + accumulator[lane], accumulator[lane + 1] + ) + st_global_v4_u32(address, packed[0], packed[1], packed[2], packed[3]) + + @cute.jit + def __call__( + self, + input_ptr: cute.Pointer, + output_ptr: cute.Pointer, + size_packs: Int32, + nbytes: Int32, + recv_base: Int64, + flag_base: Int64, + send_base: Int64, + ctrl_base: Int64, + slot_bytes: Int64, + epoch_ptr: Int64, + spin_limit: Uint32, + grid_x: Int32, + stream: cuda.CUstream, + ) -> None: + """Host entry: launch the all-reduce kernel with runtime scalars.""" + self.kernel( + input_ptr, + output_ptr, + size_packs, + nbytes, + recv_base, + flag_base, + send_base, + ctrl_base, + slot_bytes, + epoch_ptr, + spin_limit, + ).launch( + grid=(grid_x, 1, 1), + block=[self._threads, 1, 1], + cluster=(1, 1, 1), + stream=stream, + ) + + @cute.kernel + def kernel( + self, + input_ptr: cute.Pointer, + output_ptr: cute.Pointer, + size_packs: Int32, + nbytes: Int32, + recv_base: Int64, + flag_base: Int64, + send_base: Int64, + ctrl_base: Int64, + slot_bytes: Int64, + epoch_ptr: Int64, + spin_limit: Uint32, + ) -> None: + """Device kernel: stage, doorbell, wait for peer flags, reduce, advance the epoch.""" + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + gdim, _, _ = cute.arch.grid_dim() + input_base = Int64(input_ptr.toint()) + output_base = Int64(output_ptr.toint()) + stage_counter_ptr = epoch_ptr + Int64(4) + tail_counter_ptr = epoch_ptr + Int64(8) + + # Every block reads the epoch before any block can advance it: the + # advance happens only after all blocks arrived at the tail counter. + epoch = ld_relaxed_gpu_u32(epoch_ptr) + seq = epoch + Uint32(1) + slot = Int64(seq & Uint32(1)) + send_slot = send_base + slot * slot_bytes + + index = Int32(bidx) * Int32(self._threads) + Int32(tidx) + stride = Int32(gdim) * Int32(self._threads) + arrival_step = Uint32(8) // Uint32(gdim) + + # A recorded timeout poisons the runtime: later launches do nothing so + # the host sees the failure without waiting another spin limit per op. + # The device poison word (fourth counter) is written by the same waiting + # threads that write the host error word and only ever goes from 0 to + # the failed sequence, so a cheap GPU-scope load is enough here. + poison_ptr = epoch_ptr + Int64(12) + poisoned = ld_relaxed_gpu_u32(poison_ptr) + if poisoned == Uint32(0): + # 1. stage the input into the pinned send slot + stage_index = index + while stage_index < size_packs: + words = ld_global_v4_u32(input_base + Int64(stage_index) * Int64(PACK_BYTES)) + st_global_v4_u32( + send_slot + Int64(stage_index) * Int64(PACK_BYTES), + words[0], + words[1], + words[2], + words[3], + ) + stage_index += stride + fence_sc_sys() + cute.arch.sync_threads() + + # 2. the last block to finish staging rings the proxy doorbell + if Int32(tidx) == Int32(0): + prior = atomic_add_relaxed_gpu_u32(stage_counter_ptr, arrival_step) + if (prior + arrival_step) % Uint32(8) == Uint32(0): + st_relaxed_sys_u32(ctrl_base + Int64(4), Uint32(nbytes)) + st_relaxed_sys_u32(ctrl_base + Int64(16) + slot * Int64(4), Uint32(nbytes)) + fence_sc_sys() + st_relaxed_sys_u32(ctrl_base, seq) + + # 3. wait for every path flag. Four-path mode assigns one flag to + # each of eight threads; its extra opposite flags occupy the two + # otherwise-unused columns in the receiver-local source row. + if cutlass.const_expr(self._opposite_paths == 4): + waiter = 0 + for peer in cutlass.range_constexpr(self._world_size): + if cutlass.const_expr(peer != self._rank): + peer_paths = ( + 4 if (peer - self._rank) % self._world_size == 2 else 2 + ) + for path in cutlass.range_constexpr(peer_paths): + if Int32(tidx) == Int32(waiter): + flag_row = peer if path < PATH_COUNT else self._rank + flag_column = path if path < PATH_COUNT else path - PATH_COUNT + flag_addr = flag_base + ( + (Int64(flag_row) * Int64(self._slots) + slot) + * Int64(PATH_COUNT) + + Int64(flag_column) + ) * Int64(self._flag_stride) + timed_out = spin_until_eq_acquire_sys( + flag_addr, seq, spin_limit + ) + if timed_out != Uint32(0): + st_relaxed_sys_u32( + ctrl_base + Int64(12), Uint32(peer) + ) + st_relaxed_sys_u32(ctrl_base + Int64(8), seq) + st_release_gpu_u32(poison_ptr, seq) + waiter += 1 + else: + if Int32(tidx) < Int32(self._world_size): + if Int32(tidx) != Int32(self._rank): + for path in cutlass.range_constexpr(PATH_COUNT): + flag_addr = flag_base + ( + (Int64(tidx) * Int64(self._slots) + slot) + * Int64(PATH_COUNT) + + Int64(path) + ) * Int64(self._flag_stride) + timed_out = spin_until_eq_acquire_sys( + flag_addr, seq, spin_limit + ) + if timed_out != Uint32(0): + st_relaxed_sys_u32( + ctrl_base + Int64(12), Uint32(tidx) + ) + st_relaxed_sys_u32(ctrl_base + Int64(8), seq) + st_release_gpu_u32(poison_ptr, seq) + cute.arch.sync_threads() + # A wait that timed out in this block leaves the peer slot unreliable: + # skip the data phase so nothing derived from it is stored. + failed = ld_relaxed_gpu_u32(poison_ptr) + if failed == Uint32(0): + # 4. reduce in fixed rank order so every rank stores identical bits + reduce_index = index + while reduce_index < size_packs: + accumulator = cute.make_rmem_tensor((self._pack_elems,), cutlass.Float32) + offset = Int64(reduce_index) * Int64(PACK_BYTES) + for source in cutlass.range_constexpr(self._world_size): + if cutlass.const_expr(source == self._rank): + words = ld_global_v4_u32(input_base + offset) + else: + peer_slot = recv_base + ( + Int64(source) * Int64(self._slots) + slot + ) * slot_bytes + words = ld_relaxed_sys_v4_u32(peer_slot + offset) + self._accumulate_words(accumulator, words, source == 0) + self._store_accumulator(output_base + offset, accumulator) + reduce_index += stride + + # 5. the last block to finish reduction publishes the next epoch + fence_sc_gpu() + cute.arch.sync_threads() + if Int32(tidx) == Int32(0): + prior = atomic_add_relaxed_gpu_u32(tail_counter_ptr, arrival_step) + if (prior + arrival_step) % Uint32(8) == Uint32(0): + fence_sc_gpu() + # Every block's timeout store precedes its tail arrival, so the + # error word is final here. A failed sequence keeps the epoch, + # which makes every later launch a no-op until the host raises. + if ld_relaxed_sys_u32(ctrl_base + Int64(8)) == Uint32(0): + st_release_gpu_u32(epoch_ptr, seq) + + +def _dummy(dtype, alignment: int): + """A CUDA tensor of ``dtype`` used to trace launcher argument types.""" + return make_ptr(dtype, 16, cute.AddressSpace.gmem, assumed_align=alignment) + + +def _process_key( + dtype_name: str, + world_size: int, + rank: int, + threads: int, + slots: int, + flag_stride: int, + opposite_paths: int, + device_index: int, +) -> tuple[object, ...]: + """Cache key of one compiled launcher specialization.""" + return ( + str(dtype_name), + int(world_size), + int(rank), + int(threads), + int(slots), + int(flag_stride), + int(opposite_paths), + int(device_index), + ) + + +def is_launcher_prepared(*key) -> bool: + """Return whether this exact process-local launcher is already loaded.""" + + return _process_key(*key) in _PREPARED_LAUNCHERS + + +@functools.cache +def get_launcher( + dtype_name: str, + world_size: int, + rank: int, + threads: int, + slots: int, + flag_stride: int, + opposite_paths: int, + device_index: int, +) -> Callable[..., None]: + """Compile the launcher for ``key`` once and return it.""" + process_key = _process_key( + dtype_name, + world_size, + rank, + threads, + slots, + flag_stride, + opposite_paths, + device_index, + ) + del device_index # retained in the functools and preparation keys only + launch = _RoceOneshotLaunch( + dtype_name, + world_size, + rank, + threads, + slots, + flag_stride, + opposite_paths, + ) + cache_key = ( + str(dtype_name), + int(world_size), + int(rank), + int(threads), + int(slots), + int(flag_stride), + int(opposite_paths), + ) + raise_if_kernel_resolution_frozen("cute.compile", target=launch, cache_key=cache_key) + raw = b12x_compile( + launch, + _dummy(cutlass.Uint32, 16), + _dummy(cutlass.Uint32, 16), + 1, + 16, + 16, + 16, + 16, + 16, + 4096, + 16, + 1, + 1, + current_cuda_stream(), + compile_spec=KernelCompileSpec.from_key("comm.roce.oneshot", 2, cache_key), + ) + + def run( + input_address: int, + output_address: int, + size_packs: int, + nbytes: int, + recv_base: int, + flag_base: int, + send_base: int, + ctrl_base: int, + slot_bytes: int, + epoch_address: int, + spin_limit: int, + grid_x: int, + ) -> None: + """Launch the compiled kernel with runtime scalar arguments.""" + raw( + make_ptr(cutlass.Uint32, input_address, cute.AddressSpace.gmem, assumed_align=16), + make_ptr(cutlass.Uint32, output_address, cute.AddressSpace.gmem, assumed_align=16), + int(size_packs), + int(nbytes), + int(recv_base), + int(flag_base), + int(send_base), + int(ctrl_base), + int(slot_bytes), + int(epoch_address), + int(spin_limit), + int(grid_x), + current_cuda_stream(), + ) + + _PREPARED_LAUNCHERS.add(process_key) + return run + + +__all__ = ["PACK_BYTES", "get_launcher", "is_launcher_prepared"] diff --git a/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/b12x_overlay/b12x/comm/roce/_path_config.py b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/b12x_overlay/b12x/comm/roce/_path_config.py new file mode 100644 index 00000000..bb5a75c0 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/b12x_overlay/b12x/comm/roce/_path_config.py @@ -0,0 +1,167 @@ +"""Pure configuration contract for RoCEnante peer paths.""" + +from __future__ import annotations + +import os +from collections.abc import Sequence + + +PATH_COUNT = 2 +MAX_PATH_COUNT = 4 +PEER_HCA_MAP_ENV = "B12X_ROCE_PEER_HCA_MAP" +OPPOSITE_PATHS_ENV = "B12X_ROCE_OPPOSITE_PATHS" +CANONICAL_TP4_HCAS = ( + "rocep1s0f0", + "rocep1s0f1", + "roceP2p1s0f0", + "roceP2p1s0f1", +) +CANONICAL_TP4_FOUR_PATH_MAP = ( + ((-1, -1), (0, 2), (0, 3, 2, 1), (1, 3)), + ((1, 3), (-1, -1), (0, 2), (0, 3, 2, 1)), + ((1, 2, 3, 0), (1, 3), (-1, -1), (0, 2)), + ((0, 2), (1, 2, 3, 0), (1, 3), (-1, -1)), +) + + +def opposite_path_count(world_size: int, hca_count: int) -> int: + """Resolve the research-only opposite-rank path count from the environment.""" + + raw = os.getenv(OPPOSITE_PATHS_ENV, "").strip() + if not raw: + return PATH_COUNT + try: + count = int(raw) + except ValueError as exc: + raise ValueError(f"{OPPOSITE_PATHS_ENV} must be 2 or 4") from exc + if count not in (PATH_COUNT, MAX_PATH_COUNT): + raise ValueError(f"{OPPOSITE_PATHS_ENV} must be 2 or 4") + if count == MAX_PATH_COUNT and (world_size != 4 or hca_count != 4): + raise RuntimeError( + f"{OPPOSITE_PATHS_ENV}=4 requires four ranks and four active RDMA devices" + ) + return count + + +def peer_path_count( + world_size: int, rank: int, peer: int, opposite_paths: int +) -> int: + """Return two paths for a neighbor and the configured count for the TP4 opposite.""" + + if peer == rank: + return 0 + if world_size == 4 and (peer - rank) % world_size == 2: + return opposite_paths + return PATH_COUNT + + +def peer_hca_map( + world_size: int, + rank: int, + hca_count: int, + opposite_paths: int = PATH_COUNT, +) -> tuple[tuple[int, ...], ...]: + """Resolve local HCA indices per peer from ``B12X_ROCE_PEER_HCA_MAP``. + + Entries use ``peer=hca0/hca1`` for neighbors. With four opposite paths, + the opposite-rank entry has four distinct HCA indices. A two-HCA runtime + defaults every non-local peer to ``0/1``; runtimes with more HCAs require + an explicit mapping because local and reciprocal indices can differ. + """ + + if hca_count < PATH_COUNT: + raise RuntimeError( + f"RoCE two-path transport needs at least {PATH_COUNT} active RDMA devices, " + f"got {hca_count}" + ) + raw = os.getenv(PEER_HCA_MAP_ENV, "").strip() + if not raw: + if hca_count == PATH_COUNT: + return tuple( + (-1, -1) if peer == rank else (0, 1) + for peer in range(world_size) + ) + raise RuntimeError( + f"{PEER_HCA_MAP_ENV} is required when more than two RDMA devices are selected; " + "use peer=path0/path1 entries such as 1=0/2,2=0/3,3=1/3" + ) + parsed: dict[int, tuple[int, ...]] = {} + for entry in raw.split(","): + item = entry.strip() + try: + peer_text, paths_text = item.split("=", 1) + path_text = paths_text.split("/") + peer = int(peer_text) + expected_paths = peer_path_count( + world_size, rank, peer, opposite_paths + ) + if len(path_text) != expected_paths: + raise ValueError + paths = tuple(int(value) for value in path_text) + except ValueError as exc: + raise ValueError( + f"invalid {PEER_HCA_MAP_ENV} entry {item!r}; expected the configured " + "number of slash-separated HCA indices" + ) from exc + if peer < 0 or peer >= world_size or peer == rank: + raise ValueError( + f"{PEER_HCA_MAP_ENV} peer {peer} is not a non-local rank in " + f"[0,{world_size})" + ) + if peer in parsed: + raise ValueError(f"{PEER_HCA_MAP_ENV} repeats peer {peer}") + if len(set(paths)) != len(paths) or any( + hca < 0 or hca >= hca_count for hca in paths + ): + raise ValueError( + f"{PEER_HCA_MAP_ENV} peer {peer} needs distinct HCA indices in " + f"[0,{hca_count}) for every configured path" + ) + parsed[peer] = paths + expected = set(range(world_size)) - {rank} + if set(parsed) != expected: + missing = sorted(expected - set(parsed)) + unexpected = sorted(set(parsed) - expected) + raise ValueError( + f"{PEER_HCA_MAP_ENV} must map every non-local rank exactly once; " + f"missing={missing} unexpected={unexpected}" + ) + return tuple( + (-1, -1) if peer == rank else parsed[peer] for peer in range(world_size) + ) + + +def validate_four_path_tp4_mapping( + hca_names: Sequence[str], + world_size: int, + rank: int, + mapping: tuple[tuple[int, ...], ...], + opposite_paths: int, +) -> None: + """Require the measured four-rank HCA order and reciprocal four-path map.""" + + if opposite_paths != MAX_PATH_COUNT: + return + if ( + world_size != 4 + or tuple(hca_names) != CANONICAL_TP4_HCAS + or mapping != CANONICAL_TP4_FOUR_PATH_MAP[rank] + ): + raise RuntimeError( + f"{OPPOSITE_PATHS_ENV}=4 requires the canonical four-rank HCA order " + "and peer-path mapping" + ) + + +__all__ = [ + "CANONICAL_TP4_FOUR_PATH_MAP", + "CANONICAL_TP4_HCAS", + "MAX_PATH_COUNT", + "OPPOSITE_PATHS_ENV", + "PATH_COUNT", + "PEER_HCA_MAP_ENV", + "opposite_path_count", + "peer_hca_map", + "peer_path_count", + "validate_four_path_tp4_mapping", +] diff --git a/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/b12x_overlay/b12x/comm/roce/_proxy.py b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/b12x_overlay/b12x/comm/roce/_proxy.py new file mode 100644 index 00000000..965b4f4e --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/b12x_overlay/b12x/comm/roce/_proxy.py @@ -0,0 +1,397 @@ +"""Build and bind the RDMA proxy shared library (``_roce_proxy.c``). + +The proxy is plain C over libibverbs. It is compiled once per source hash +with the host C compiler into the b12x cache directory, so the package stays +pure Python for packaging purposes while the RDMA posting loop runs without +the interpreter on its critical path. +""" + +from __future__ import annotations + +import contextlib +import ctypes +import hashlib +import os +import shutil +import subprocess +import threading +from pathlib import Path + +_SOURCE = Path(__file__).with_name("_roce_proxy.c") +_LOCK = threading.Lock() +_LIB: ctypes.CDLL | None = None + +BLOB_STRUCT_ERR = "roce proxy blob size mismatch" + + +def _peer_path_count(world_size: int, rank: int, peer: int, opposite_paths: int) -> int: + """Return the active path count for one local peer relation.""" + + if peer == rank: + return 0 + if world_size == 4 and (peer - rank) % world_size == 2: + return opposite_paths + return 2 + + +def _flatten_peer_hca_map( + peer_hca_map: tuple[tuple[int, ...], ...], + *, + world_size: int, + rank: int, + opposite_paths: int, +) -> tuple[int, ...]: + """Encode a variable path map as four HCA bytes per peer for ABI 4.""" + + if int(opposite_paths) not in (2, 4): + raise ValueError("opposite_paths must be 2 or 4") + if len(peer_hca_map) != int(world_size): + raise ValueError("peer_hca_map must contain one entry per rank") + flattened = [] + for peer, paths in enumerate(peer_hca_map): + expected = _peer_path_count(world_size, rank, peer, opposite_paths) + if peer == rank: + if any(int(hca) >= 0 for hca in paths): + raise ValueError("the local peer_hca_map entry must contain sentinels") + elif len(paths) != expected: + raise ValueError( + f"peer_hca_map peer {peer} must contain {expected} HCA indices" + ) + values = [255 if int(hca) < 0 else int(hca) for hca in paths] + if len(values) > 4: + raise ValueError("peer_hca_map entries cannot exceed four paths") + flattened.extend(values + [255] * (4 - len(values))) + return tuple(flattened) + + +def _cache_dir() -> Path: + """Directory for the compiled proxy library: ``B12X_ROCE_CACHE_DIR``, else ``/b12x/roce``.""" + override = os.getenv("B12X_ROCE_CACHE_DIR") + if override: + return Path(override) + root = os.getenv("XDG_CACHE_HOME") or os.path.join(os.path.expanduser("~"), ".cache") + return Path(root) / "b12x" / "roce" + + +def _compiler() -> str: + """Host C compiler used to build the proxy (``CC``, gcc, cc, or clang).""" + for candidate in (os.getenv("CC"), "gcc", "cc", "clang"): + if candidate and shutil.which(candidate): + return candidate + raise RuntimeError( + "b12x.comm.roce needs a C compiler (gcc/cc/clang) and libibverbs headers " + "to build its RDMA proxy" + ) + + +def _build() -> Path: + """Compile the proxy source into the cache directory keyed by source hash; returns the .so path.""" + source = _SOURCE.read_bytes() + digest = hashlib.sha256(source).hexdigest()[:16] + cache = _cache_dir() + cache.mkdir(parents=True, exist_ok=True) + target = cache / f"roce_proxy-{digest}.so" + if target.exists(): + return target + tmp = cache / f".roce_proxy-{digest}-{os.getpid()}.so" + cmd = [ + _compiler(), + "-O2", + "-std=gnu11", + "-shared", + "-fPIC", + "-o", + str(tmp), + str(_SOURCE), + "-libverbs", + "-lpthread", + ] + proc = subprocess.run(cmd, capture_output=True, text=True) + if proc.returncode != 0: + raise RuntimeError( + "failed to build the b12x RoCE proxy: " + " ".join(cmd) + "\n" + proc.stderr + ) + os.replace(tmp, target) + return target + + +def load() -> ctypes.CDLL: + """Return the bound proxy library, building it on first use.""" + + global _LIB + with _LOCK: + if _LIB is not None: + return _LIB + lib = ctypes.CDLL(str(_build()), use_errno=True) + u64 = ctypes.c_uint64 + p = ctypes.c_void_p + lib.roce_abi_version.restype = ctypes.c_int + lib.roce_abi_version.argtypes = [] + lib.roce_layout.restype = ctypes.c_int + lib.roce_layout.argtypes = [ctypes.c_int, u64, ctypes.POINTER(u64)] + lib.roce_blob_bytes.restype = u64 + lib.roce_blob_bytes.argtypes = [] + lib.roce_create.restype = p + lib.roce_create.argtypes = [ + ctypes.c_int, + ctypes.c_int, + ctypes.POINTER(ctypes.c_char_p), + ctypes.c_int, + ctypes.c_int, + p, + u64, + u64, + ctypes.c_int, + ctypes.POINTER(ctypes.c_uint8), + u64, + ctypes.c_char_p, + u64, + ] + lib.roce_local_blob.restype = ctypes.c_int + lib.roce_local_blob.argtypes = [p, p, u64] + lib.roce_connect.restype = ctypes.c_int + lib.roce_connect.argtypes = [p, p, u64] + lib.roce_start.restype = ctypes.c_int + lib.roce_start.argtypes = [p] + lib.roce_stop.restype = None + lib.roce_stop.argtypes = [p] + lib.roce_failed.restype = ctypes.c_int + lib.roce_failed.argtypes = [p] + lib.roce_error.restype = ctypes.c_char_p + lib.roce_error.argtypes = [p] + lib.roce_stat.restype = u64 + lib.roce_stat.argtypes = [p, ctypes.c_int] + lib.roce_two_wave_threshold_bytes.restype = u64 + lib.roce_two_wave_threshold_bytes.argtypes = [p] + lib.roce_wave_mode.restype = u64 + lib.roce_wave_mode.argtypes = [p] + lib.roce_peer_hca.restype = ctypes.c_int + lib.roce_peer_hca.argtypes = [p, ctypes.c_int, ctypes.c_int] + lib.roce_path_stat.restype = u64 + lib.roce_path_stat.argtypes = [p, ctypes.c_int, ctypes.c_int, ctypes.c_int] + lib.roce_destroy.restype = None + lib.roce_destroy.argtypes = [p] + if lib.roce_abi_version() != 4: + raise RuntimeError("unexpected b12x RoCE proxy ABI version") + _LIB = lib + return lib + + +class Layout: + """Byte offsets of the pinned region shared by the kernel and the proxy.""" + + __slots__ = ( + "recv_off", + "flag_off", + "send_off", + "ctrl_off", + "total_bytes", + "flag_stride", + "slots", + "paths", + ) + + def __init__(self, world_size: int, slot_bytes: int) -> None: + """Create the proxy context: open the HCAs, register the pinned region, create the queue pairs.""" + out = (ctypes.c_uint64 * 8)() + if load().roce_layout(int(world_size), int(slot_bytes), out) != 0: + raise ValueError( + f"unsupported RoCE geometry: world_size={world_size} slot_bytes={slot_bytes} " + "(2..16 ranks, slot_bytes a positive multiple of 4096)" + ) + ( + self.recv_off, + self.flag_off, + self.send_off, + self.ctrl_off, + self.total_bytes, + self.flag_stride, + self.slots, + self.paths, + ) = (int(v) for v in out) + + +class Proxy: + """Owns one RDMA proxy context for one rank.""" + + def __init__( + self, + *, + world_size: int, + rank: int, + hca_names: tuple[str, ...], + gid_index: int, + region_ptr: int, + region_bytes: int, + slot_bytes: int, + peer_hca_map: tuple[tuple[int, ...], ...], + opposite_paths: int = 2, + ) -> None: + """Create the proxy context: open the HCAs, register the pinned region, create the queue pairs.""" + self._lib = load() + names = (ctypes.c_char_p * len(hca_names))(*[n.encode() for n in hca_names]) + flattened = _flatten_peer_hca_map( + peer_hca_map, + world_size=int(world_size), + rank=int(rank), + opposite_paths=int(opposite_paths), + ) + hca_map = (ctypes.c_uint8 * len(flattened))(*flattened) + err = ctypes.create_string_buffer(512) + self._ctx = self._lib.roce_create( + int(world_size), + int(rank), + names, + len(hca_names), + int(gid_index), + ctypes.c_void_p(int(region_ptr)), + int(region_bytes), + int(slot_bytes), + int(opposite_paths), + hca_map, + len(flattened), + err, + len(err), + ) + if not self._ctx: + raise RuntimeError(f"RoCE proxy setup failed: {err.value.decode(errors='replace')}") + self.world_size = int(world_size) + self.rank = int(rank) + self.hca_names = tuple(hca_names) + self.opposite_paths = int(opposite_paths) + + def local_blob(self) -> bytes: + """Serialized connection record (region address, keys, queue-pair numbers) to send to every peer.""" + n = int(self._lib.roce_blob_bytes()) + buf = ctypes.create_string_buffer(n) + if self._lib.roce_local_blob(self._ctx, buf, n) != 0: + raise RuntimeError(BLOB_STRUCT_ERR) + return buf.raw + + def connect(self, blobs: list[bytes]) -> None: + """Connect the queue pairs from every rank's ``local_blob``, in rank order.""" + n = int(self._lib.roce_blob_bytes()) + if len(blobs) != self.world_size or any(len(b) != n for b in blobs): + raise RuntimeError(BLOB_STRUCT_ERR) + joined = b"".join(blobs) + buf = ctypes.create_string_buffer(joined, len(joined)) + if self._lib.roce_connect(self._ctx, buf, len(joined)) != 0: + raise RuntimeError(f"RoCE queue-pair connect failed: {self.error()}") + + def start(self) -> None: + """Start the proxy thread; a restart resumes from the last posted sequence.""" + if self._lib.roce_start(self._ctx) != 0: + raise RuntimeError(f"RoCE proxy thread failed to start: {self.error()}") + + def stop(self) -> None: + """Stop the proxy thread; ``start`` resumes from the last posted op.""" + self._lib.roce_stop(self._ctx) + + def failed(self) -> bool: + """True once the proxy thread has stopped on an error.""" + return bool(self._lib.roce_failed(self._ctx)) + + def error(self) -> str: + """The proxy's last error message, empty when none.""" + raw = self._lib.roce_error(self._ctx) + return raw.decode(errors="replace") if raw else "" + + def stats(self) -> dict[str, int | str]: + """Operation, completion, sequence, and two-wave scheduling counters.""" + return { + "ops_posted": int(self._lib.roce_stat(self._ctx, 0)), + "writes_completed": int(self._lib.roce_stat(self._ctx, 1)), + "last_seq": int(self._lib.roce_stat(self._ctx, 2)), + "two_wave_activations": int(self._lib.roce_stat(self._ctx, 3)), + "two_wave_threshold_bytes": int( + self._lib.roce_two_wave_threshold_bytes(self._ctx) + ), + "wave_mode": self.wave_mode(), + } + + def two_wave_threshold_bytes(self) -> int: + """Payload-size threshold for direct-then-forwarded posting; zero disables it.""" + return int(self._lib.roce_two_wave_threshold_bytes(self._ctx)) + + def wave_mode(self) -> str: + """Configured large-payload posting schedule.""" + value = int(self._lib.roce_wave_mode(self._ctx)) + return { + 0: "two", + 1: "mixed2", + 2: "opposite_first", + 3: "strict3", + 4: "balanced32", + }.get(value, f"unknown-{value}") + + def peer_hca(self, peer: int) -> tuple[int, ...]: + """Local HCA indices for every active path to ``peer``.""" + count = ( + self.opposite_paths + if self.world_size == 4 + and (int(peer) - self.rank) % self.world_size == 2 + else 2 + ) + return tuple( + int(self._lib.roce_peer_hca(self._ctx, int(peer), path)) + for path in range(count) + ) + + def path_counters(self) -> list[dict[str, int | str | None]]: + """Absolute origin-QP counters since this proxy context was created.""" + names = ( + "payload_writes", + "payload_bytes", + "physical_hop_payload_bytes", + "flag_writes", + "send_completions", + "completion_errors", + "qp_number", + "remote_qp_number", + "local_hca_index", + "remote_hca_index", + "physical_hops", + ) + result: list[dict[str, int | str | None]] = [] + for peer in range(self.world_size): + if peer == self.rank: + continue + path_count = ( + self.opposite_paths + if self.world_size == 4 and (peer - self.rank) % self.world_size == 2 + else 2 + ) + for path in range(path_count): + values = { + name: int(self._lib.roce_path_stat(self._ctx, peer, path, index)) + for index, name in enumerate(names) + } + local_hca = int(values["local_hca_index"]) + result.append( + { + "peer_rank": peer, + "path_index": path, + "device": self.hca_names[local_hca], + **values, + # Standard verbs reports retry exhaustion as a failed + # completion but does not expose the number of RC retry + # packets for an individual QP. + "retries": None, + "retry_events": None, + } + ) + return result + + def close(self) -> None: + """Stop the thread and release the RDMA resources.""" + ctx, self._ctx = self._ctx, None + if ctx: + self._lib.roce_destroy(ctx) + + def __del__(self) -> None: # pragma: no cover - defensive teardown + """Release the RDMA resources if ``close`` was never called.""" + with contextlib.suppress(Exception): + self.close() + + +__all__ = ["Layout", "Proxy", "load"] diff --git a/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/b12x_overlay/b12x/comm/roce/_roce_proxy.c b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/b12x_overlay/b12x/comm/roce/_roce_proxy.c new file mode 100644 index 00000000..d8043504 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/b12x_overlay/b12x/comm/roce/_roce_proxy.c @@ -0,0 +1,1159 @@ +// RDMA proxy for the b12x RoCE one-shot all-reduce. +// +// One rank owns one pinned host region laid out as: +// +// recv[src][slot] (world * SLOTS * slot_bytes) filled by peers' RDMA writes +// flag[row][slot][lane] (world * SLOTS * 2 * FLAG_STRIDE) sequence number +// written after one stripe. Four-path TP4 mode keeps +// lanes 0/1 in the source row and stores opposite-path +// lanes 2/3 in the otherwise-unused receiver row. +// send[slot] (SLOTS * slot_bytes) staged by the local GPU kernel +// ctrl (FLAG_STRIDE) {u32 seq, u32 nbytes, u32 error, +// u32 missing_peer} doorbell; the +// last two are set by the kernel +// when a wait times out +// +// The GPU kernel stages its input into send[seq & 1], publishes nbytes and seq +// in ctrl, then spins on every active path flag for every peer. The proxy uses +// two half-payload paths per neighbor. Research-only TP4 mode uses four +// quarter-payload paths to the opposite rank. Each stripe is followed by its +// 4-byte sequence flag on the same reliable QP, so a path flag cannot become +// visible before its stripe. Nothing on the receive path involves the host. +// +// This file is compiled by b12x.comm.roce._proxy at first use with the host +// gcc and libibverbs; it must stay plain C with no CUDA dependency. + +#define _GNU_SOURCE +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#define ROCE_MAX_PEERS 16 +#define ROCE_MAX_HCAS 4 +#define ROCE_SLOTS 2 +#define ROCE_LAYOUT_PATHS 2 +#define ROCE_MAX_PATHS 4 +#define ROCE_FLAG_STRIDE 128 +#define ROCE_PORT 1 +#define ROCE_SEND_DEPTH 256 +#define ROCE_ABI_VERSION 4 +#define ROCE_IDLE_SPINS 200000 +#define ROCE_DEFAULT_TWO_WAVE_THRESHOLD_BYTES 131072u +#define ROCE_WAVE_MODE_TWO 0u +#define ROCE_WAVE_MODE_MIXED_TWO 1u +#define ROCE_WAVE_MODE_OPPOSITE_FIRST 2u +#define ROCE_WAVE_MODE_STRICT_THREE 3u +#define ROCE_WAVE_MODE_BALANCED32 4u + +typedef struct { + uint32_t abi_version; + uint32_t world; + uint32_t rank; + uint32_t n_hca; + uint32_t layout_paths; + uint32_t opposite_paths; + uint64_t region_addr; + uint32_t rkey[ROCE_MAX_HCAS]; + uint16_t lid[ROCE_MAX_HCAS]; + uint8_t gid[ROCE_MAX_HCAS][16]; + uint32_t mtu[ROCE_MAX_HCAS]; + uint32_t qp_num[ROCE_MAX_HCAS][ROCE_MAX_PEERS]; + uint8_t peer_hca[ROCE_MAX_PEERS][ROCE_MAX_PATHS]; +} roce_blob_t; + +typedef struct { + struct ibv_context *ctx; + struct ibv_pd *pd; + struct ibv_mr *mr; + struct ibv_cq *cq; + struct ibv_qp *qp[ROCE_MAX_PEERS]; + uint32_t outstanding[ROCE_MAX_PEERS]; + union ibv_gid gid; + uint16_t lid; + enum ibv_mtu mtu; +} roce_hca_t; + +typedef struct { + int world; + int rank; + int n_hca; + int gid_index; + roce_hca_t hca[ROCE_MAX_HCAS]; + uint8_t *region; + size_t region_bytes; + size_t slot_bytes; + size_t recv_off; + size_t flag_off; + size_t send_off; + size_t ctrl_off; + int started; + uint64_t peer_addr[ROCE_MAX_PEERS]; + uint32_t peer_rkey[ROCE_MAX_PEERS][ROCE_MAX_PATHS]; + uint32_t remote_qp[ROCE_MAX_PEERS][ROCE_MAX_PATHS]; + int peer_hca[ROCE_MAX_PEERS][ROCE_MAX_PATHS]; + int remote_hca[ROCE_MAX_PEERS][ROCE_MAX_PATHS]; + uint32_t peer_path_count[ROCE_MAX_PEERS]; + uint32_t opposite_paths; + uint32_t physical_hops[ROCE_MAX_PEERS]; + atomic_uint_fast64_t payload_writes[ROCE_MAX_PEERS][ROCE_MAX_PATHS]; + atomic_uint_fast64_t payload_bytes[ROCE_MAX_PEERS][ROCE_MAX_PATHS]; + atomic_uint_fast64_t flag_writes[ROCE_MAX_PEERS][ROCE_MAX_PATHS]; + atomic_uint_fast64_t send_completions[ROCE_MAX_PEERS][ROCE_MAX_PATHS]; + atomic_uint_fast64_t completion_errors[ROCE_MAX_PEERS][ROCE_MAX_PATHS]; + int direct_peer_by_hca[ROCE_MAX_HCAS]; + int direct_path_by_hca[ROCE_MAX_HCAS]; + int opposite_peer_by_hca[ROCE_MAX_HCAS]; + int opposite_path_by_hca[ROCE_MAX_HCAS]; + pthread_t thread; + atomic_int running; + atomic_int failed; + uint32_t last_seq; + uint64_t ops_posted; + uint64_t writes_completed; + uint32_t two_wave_threshold_bytes; + uint32_t wave_mode; + atomic_uint_fast64_t two_wave_activations; + char err[512]; +} roce_ctx_t; + +void roce_destroy(roce_ctx_t *c); + +static void set_err(roce_ctx_t *c, const char *what, int e) { + snprintf(c->err, sizeof(c->err), "%s: %s", what, e ? strerror(e) : "failed"); +} + +int roce_abi_version(void) { return ROCE_ABI_VERSION; } + +int roce_layout(int world, uint64_t slot_bytes, uint64_t *out) { + // out = {recv_off, flag_off, send_off, ctrl_off, total_bytes, + // flag_stride, slots, paths} + if (world < 2 || world > ROCE_MAX_PEERS || slot_bytes == 0 || (slot_bytes % 4096) != 0) { + return -1; + } + // Reject a layout whose arithmetic would wrap; the caller sizes slots from + // configuration, so a wrapped region must fail here rather than at the NIC. + uint64_t recv_bytes, send_bytes, flag_bytes, flag_off, send_off, ctrl_off, total; + if (slot_bytes > ((uint64_t)1 << 40) || + __builtin_mul_overflow((uint64_t)world * ROCE_SLOTS, slot_bytes, &recv_bytes) || + __builtin_mul_overflow((uint64_t)ROCE_SLOTS, slot_bytes, &send_bytes) || + __builtin_mul_overflow((uint64_t)world * ROCE_SLOTS * ROCE_LAYOUT_PATHS, + (uint64_t)ROCE_FLAG_STRIDE, &flag_bytes) || + __builtin_add_overflow(recv_bytes, flag_bytes, &send_off) || + __builtin_add_overflow(send_off, send_bytes, &ctrl_off) || + __builtin_add_overflow(ctrl_off, (uint64_t)ROCE_FLAG_STRIDE, &total)) { + return -1; + } + uint64_t recv_off = 0; + flag_off = recv_off + recv_bytes; + send_off = flag_off + flag_bytes; + out[0] = recv_off; + out[1] = flag_off; + out[2] = send_off; + out[3] = ctrl_off; + out[4] = total; + out[5] = ROCE_FLAG_STRIDE; + out[6] = ROCE_SLOTS; + out[7] = ROCE_LAYOUT_PATHS; + return 0; +} + +uint64_t roce_blob_bytes(void) { return sizeof(roce_blob_t); } + +static int open_hca(roce_ctx_t *c, int h, const char *name) { + int num = 0; + struct ibv_device **list = ibv_get_device_list(&num); + if (list == NULL) { + set_err(c, "ibv_get_device_list", errno); + return -1; + } + struct ibv_device *dev = NULL; + for (int i = 0; i < num; i++) { + if (strcmp(ibv_get_device_name(list[i]), name) == 0) { + dev = list[i]; + break; + } + } + if (dev == NULL) { + ibv_free_device_list(list); + snprintf(c->err, sizeof(c->err), "RDMA device %s not found", name); + return -1; + } + roce_hca_t *hca = &c->hca[h]; + hca->ctx = ibv_open_device(dev); + ibv_free_device_list(list); + if (hca->ctx == NULL) { + set_err(c, "ibv_open_device", errno); + return -1; + } + struct ibv_port_attr port; + if (ibv_query_port(hca->ctx, ROCE_PORT, &port) != 0) { + set_err(c, "ibv_query_port", errno); + return -1; + } + if (port.state != IBV_PORT_ACTIVE) { + snprintf(c->err, sizeof(c->err), "RDMA device %s port %d is not active", name, ROCE_PORT); + return -1; + } + hca->lid = port.lid; + hca->mtu = port.active_mtu; + if (ibv_query_gid(hca->ctx, ROCE_PORT, c->gid_index, &hca->gid) != 0) { + set_err(c, "ibv_query_gid", errno); + return -1; + } + hca->pd = ibv_alloc_pd(hca->ctx); + if (hca->pd == NULL) { + set_err(c, "ibv_alloc_pd", errno); + return -1; + } + hca->mr = ibv_reg_mr(hca->pd, c->region, c->region_bytes, + IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_WRITE); + if (hca->mr == NULL) { + set_err(c, "ibv_reg_mr(pinned region)", errno); + return -1; + } + hca->cq = ibv_create_cq(hca->ctx, ROCE_SEND_DEPTH * ROCE_MAX_PEERS, NULL, NULL, 0); + if (hca->cq == NULL) { + set_err(c, "ibv_create_cq", errno); + return -1; + } + for (int p = 0; p < c->world; p++) { + if (p == c->rank) { + continue; + } + struct ibv_qp_init_attr attr; + memset(&attr, 0, sizeof(attr)); + attr.send_cq = hca->cq; + attr.recv_cq = hca->cq; + attr.qp_type = IBV_QPT_RC; + attr.cap.max_send_wr = ROCE_SEND_DEPTH; + attr.cap.max_recv_wr = 1; + attr.cap.max_send_sge = 1; + attr.cap.max_recv_sge = 1; + attr.cap.max_inline_data = 16; + hca->qp[p] = ibv_create_qp(hca->pd, &attr); + if (hca->qp[p] == NULL) { + set_err(c, "ibv_create_qp", errno); + return -1; + } + struct ibv_qp_attr init; + memset(&init, 0, sizeof(init)); + init.qp_state = IBV_QPS_INIT; + init.pkey_index = 0; + init.port_num = ROCE_PORT; + init.qp_access_flags = IBV_ACCESS_REMOTE_WRITE; + int rc = ibv_modify_qp(hca->qp[p], &init, + IBV_QP_STATE | IBV_QP_PKEY_INDEX | IBV_QP_PORT | IBV_QP_ACCESS_FLAGS); + if (rc != 0) { + set_err(c, "ibv_modify_qp(INIT)", rc); + return -1; + } + } + return 0; +} + +roce_ctx_t *roce_create(int world, int rank, const char *const *hca_names, int n_hca, + int gid_index, void *region, uint64_t region_bytes, + uint64_t slot_bytes, int opposite_paths, + const uint8_t *peer_hca_map, + uint64_t peer_hca_count, char *err, uint64_t err_len) { + uint64_t layout[8]; + if (roce_layout(world, slot_bytes, layout) != 0 || layout[4] > region_bytes || + rank < 0 || rank >= world || n_hca < ROCE_LAYOUT_PATHS || + n_hca > ROCE_MAX_HCAS || + (opposite_paths != ROCE_LAYOUT_PATHS && opposite_paths != ROCE_MAX_PATHS) || + (opposite_paths == ROCE_MAX_PATHS && (world != 4 || n_hca != 4)) || + peer_hca_map == NULL || + peer_hca_count != (uint64_t)world * ROCE_MAX_PATHS) { + snprintf(err, err_len, "invalid roce runtime geometry"); + return NULL; + } + roce_ctx_t *c = calloc(1, sizeof(*c)); + if (c == NULL) { + snprintf(err, err_len, "out of memory"); + return NULL; + } + c->world = world; + c->rank = rank; + c->n_hca = n_hca; + c->gid_index = gid_index; + c->region = region; + c->region_bytes = region_bytes; + c->slot_bytes = slot_bytes; + c->recv_off = layout[0]; + c->flag_off = layout[1]; + c->send_off = layout[2]; + c->ctrl_off = layout[3]; + c->opposite_paths = (uint32_t)opposite_paths; + c->two_wave_threshold_bytes = ROCE_DEFAULT_TWO_WAVE_THRESHOLD_BYTES; + c->wave_mode = opposite_paths == ROCE_MAX_PATHS + ? ROCE_WAVE_MODE_BALANCED32 + : ROCE_WAVE_MODE_TWO; + const char *wave_mode_text = getenv("B12X_ROCE_WAVE_MODE"); + if (wave_mode_text != NULL && wave_mode_text[0] != '\0') { + if (opposite_paths == ROCE_MAX_PATHS && + strcmp(wave_mode_text, "balanced32") == 0) { + c->wave_mode = ROCE_WAVE_MODE_BALANCED32; + } else if (opposite_paths == ROCE_MAX_PATHS) { + snprintf(err, err_len, + "B12X_ROCE_OPPOSITE_PATHS=4 requires " + "B12X_ROCE_WAVE_MODE=balanced32 when the wave mode is set"); + roce_destroy(c); + return NULL; + } else if (strcmp(wave_mode_text, "two") == 0) { + c->wave_mode = ROCE_WAVE_MODE_TWO; + } else if (strcmp(wave_mode_text, "mixed2") == 0) { + c->wave_mode = ROCE_WAVE_MODE_MIXED_TWO; + } else if (strcmp(wave_mode_text, "opposite_first") == 0) { + c->wave_mode = ROCE_WAVE_MODE_OPPOSITE_FIRST; + } else if (strcmp(wave_mode_text, "strict3") == 0) { + c->wave_mode = ROCE_WAVE_MODE_STRICT_THREE; + } else { + snprintf(err, err_len, + "B12X_ROCE_WAVE_MODE must be 'two', 'mixed2', " + "'opposite_first', or 'strict3'"); + roce_destroy(c); + return NULL; + } + } + const char *threshold_text = getenv("B12X_ROCE_TWO_WAVE_THRESHOLD_BYTES"); + if (threshold_text != NULL && threshold_text[0] != '\0') { + char *end = NULL; + errno = 0; + unsigned long long value = strtoull(threshold_text, &end, 10); + if (errno != 0 || end == threshold_text || *end != '\0' || + value > UINT32_MAX || (value != 0 && value % 16u != 0)) { + snprintf(err, err_len, + "B12X_ROCE_TWO_WAVE_THRESHOLD_BYTES must be zero or a " + "16-byte-aligned integer no larger than %u", + UINT32_MAX); + roce_destroy(c); + return NULL; + } + c->two_wave_threshold_bytes = (uint32_t)value; + } + for (int h = 0; h < ROCE_MAX_HCAS; h++) { + c->direct_peer_by_hca[h] = -1; + c->direct_path_by_hca[h] = -1; + c->opposite_peer_by_hca[h] = -1; + c->opposite_path_by_hca[h] = -1; + } + for (int p = 0; p < world; p++) { + for (int path = 0; path < ROCE_MAX_PATHS; path++) { + c->peer_hca[p][path] = -1; + c->remote_hca[p][path] = -1; + } + if (p == rank) { + continue; + } + int distance = (p - rank + world) % world; + int count = world == 4 && distance == 2 + ? opposite_paths + : ROCE_LAYOUT_PATHS; + c->peer_path_count[p] = (uint32_t)count; + uint32_t seen_hcas = 0; + for (int path = 0; path < count; path++) { + int h = (int)peer_hca_map[p * ROCE_MAX_PATHS + path]; + if (h < 0 || h >= n_hca || (seen_hcas & (1u << h)) != 0) { + snprintf(err, err_len, + "rank %d peer %d needs %d distinct HCA indices in [0,%d)", + rank, p, count, n_hca); + roce_destroy(c); + return NULL; + } + seen_hcas |= 1u << h; + c->peer_hca[p][path] = h; + } + for (int path = count; path < ROCE_MAX_PATHS; path++) { + if (peer_hca_map[p * ROCE_MAX_PATHS + path] != UINT8_MAX) { + snprintf(err, err_len, + "rank %d peer %d publishes an inactive path %d", rank, p, path); + roce_destroy(c); + return NULL; + } + } + c->physical_hops[p] = world == 4 && distance == 2 ? 2u : 1u; + } + if (opposite_paths == ROCE_MAX_PATHS || c->wave_mode != ROCE_WAVE_MODE_TWO) { + static const char *const canonical_hcas[4] = { + "rocep1s0f0", "rocep1s0f1", "roceP2p1s0f0", "roceP2p1s0f1"}; + static const int canonical_map[4][4][ROCE_MAX_PATHS] = { + {{-1, -1, -1, -1}, {0, 2, -1, -1}, {0, 3, 2, 1}, {1, 3, -1, -1}}, + {{1, 3, -1, -1}, {-1, -1, -1, -1}, {0, 2, -1, -1}, {0, 3, 2, 1}}, + {{1, 2, 3, 0}, {1, 3, -1, -1}, {-1, -1, -1, -1}, {0, 2, -1, -1}}, + {{0, 2, -1, -1}, {1, 2, 3, 0}, {1, 3, -1, -1}, {-1, -1, -1, -1}}, + }; + int canonical = world == 4 && n_hca == 4; + for (int h = 0; canonical && h < 4; h++) { + canonical = strcmp(hca_names[h], canonical_hcas[h]) == 0; + } + for (int p = 0; canonical && p < world; p++) { + int count = c->peer_path_count[p]; + for (int path = 0; canonical && path < ROCE_MAX_PATHS; path++) { + int expected = path < count ? canonical_map[rank][p][path] : -1; + canonical = c->peer_hca[p][path] == expected; + } + } + if (!canonical) { + snprintf(err, err_len, + "the selected RoCEnante path schedule requires the canonical " + "four-rank HCA order and peer-path mapping"); + roce_destroy(c); + return NULL; + } + } + if (opposite_paths == ROCE_MAX_PATHS) { + int opposite = (rank + 2) % world; + for (int p = 0; p < world; p++) { + if (p == rank) { + continue; + } + int is_opposite = p == opposite; + for (int path = 0; path < (int)c->peer_path_count[p]; path++) { + int h = c->peer_hca[p][path]; + int *peer_slot = is_opposite ? &c->opposite_peer_by_hca[h] + : &c->direct_peer_by_hca[h]; + int *path_slot = is_opposite ? &c->opposite_path_by_hca[h] + : &c->direct_path_by_hca[h]; + if (*peer_slot != -1) { + snprintf(err, err_len, + "rank %d HCA %d has multiple %s paths", rank, h, + is_opposite ? "opposite" : "direct"); + roce_destroy(c); + return NULL; + } + *peer_slot = p; + *path_slot = path; + } + } + for (int h = 0; h < n_hca; h++) { + if (c->direct_peer_by_hca[h] < 0 || c->opposite_peer_by_hca[h] < 0) { + snprintf(err, err_len, + "rank %d HCA %d lacks one direct or opposite path", rank, h); + roce_destroy(c); + return NULL; + } + } + } + for (int h = 0; h < n_hca; h++) { + if (open_hca(c, h, hca_names[h]) != 0) { + snprintf(err, err_len, "%s", c->err); + roce_destroy(c); + return NULL; + } + } + return c; +} + +int roce_local_blob(roce_ctx_t *c, void *out, uint64_t out_len) { + if (out_len < sizeof(roce_blob_t)) { + return -1; + } + roce_blob_t blob; + memset(&blob, 0, sizeof(blob)); + blob.abi_version = ROCE_ABI_VERSION; + blob.world = (uint32_t)c->world; + blob.rank = (uint32_t)c->rank; + blob.n_hca = (uint32_t)c->n_hca; + blob.layout_paths = ROCE_LAYOUT_PATHS; + blob.opposite_paths = c->opposite_paths; + blob.region_addr = (uint64_t)(uintptr_t)c->region; + for (int h = 0; h < c->n_hca; h++) { + blob.rkey[h] = c->hca[h].mr->rkey; + blob.lid[h] = c->hca[h].lid; + blob.mtu[h] = (uint32_t)c->hca[h].mtu; + memcpy(blob.gid[h], c->hca[h].gid.raw, 16); + for (int p = 0; p < c->world; p++) { + blob.qp_num[h][p] = (p == c->rank) ? 0 : c->hca[h].qp[p]->qp_num; + } + } + memset(blob.peer_hca, UINT8_MAX, sizeof(blob.peer_hca)); + for (int p = 0; p < c->world; p++) { + if (p == c->rank) { + continue; + } + for (int path = 0; path < (int)c->peer_path_count[p]; path++) { + blob.peer_hca[p][path] = (uint8_t)c->peer_hca[p][path]; + } + } + memcpy(out, &blob, sizeof(blob)); + return 0; +} + +static int connect_qp(roce_ctx_t *c, int local_h, int remote_h, int p, + const roce_blob_t *peer) { + roce_hca_t *hca = &c->hca[local_h]; + struct ibv_qp_attr rtr; + memset(&rtr, 0, sizeof(rtr)); + rtr.qp_state = IBV_QPS_RTR; + rtr.path_mtu = (enum ibv_mtu)(peer->mtu[remote_h] < (uint32_t)hca->mtu + ? peer->mtu[remote_h] + : (uint32_t)hca->mtu); + rtr.dest_qp_num = peer->qp_num[remote_h][c->rank]; + rtr.rq_psn = 0; + rtr.max_dest_rd_atomic = 1; + rtr.min_rnr_timer = 12; + rtr.ah_attr.is_global = 1; + rtr.ah_attr.dlid = peer->lid[remote_h]; + rtr.ah_attr.sl = 0; + rtr.ah_attr.src_path_bits = 0; + rtr.ah_attr.port_num = ROCE_PORT; + memcpy(rtr.ah_attr.grh.dgid.raw, peer->gid[remote_h], 16); + rtr.ah_attr.grh.sgid_index = (uint8_t)c->gid_index; + rtr.ah_attr.grh.hop_limit = 64; + rtr.ah_attr.grh.traffic_class = 0; + // A four-rank cycle has no physical link between opposite ranks. + // Hardware-forwarding rules for opposite-rank paths match UDP source port + // 65535, which mlx5 derives from flow label 16383, and rewrite only the + // Ethernet header at the intermediate ConnectX device. Neighbor QPs use + // flow label zero and do not match those rules. + rtr.ah_attr.grh.flow_label = + (c->world == 4 && ((p - c->rank + c->world) % c->world) == 2) + ? 16383 + : 0; + int rc = ibv_modify_qp(hca->qp[p], &rtr, + IBV_QP_STATE | IBV_QP_AV | IBV_QP_PATH_MTU | IBV_QP_DEST_QPN | + IBV_QP_RQ_PSN | IBV_QP_MAX_DEST_RD_ATOMIC | IBV_QP_MIN_RNR_TIMER); + if (rc != 0) { + set_err(c, "ibv_modify_qp(RTR)", rc); + return -1; + } + struct ibv_qp_attr rts; + memset(&rts, 0, sizeof(rts)); + rts.qp_state = IBV_QPS_RTS; + rts.timeout = 14; + rts.retry_cnt = 7; + rts.rnr_retry = 7; + rts.sq_psn = 0; + rts.max_rd_atomic = 1; + rc = ibv_modify_qp(hca->qp[p], &rts, + IBV_QP_STATE | IBV_QP_TIMEOUT | IBV_QP_RETRY_CNT | IBV_QP_RNR_RETRY | + IBV_QP_SQ_PSN | IBV_QP_MAX_QP_RD_ATOMIC); + if (rc != 0) { + set_err(c, "ibv_modify_qp(RTS)", rc); + return -1; + } + return 0; +} + +int roce_connect(roce_ctx_t *c, const void *blobs, uint64_t blobs_len) { + if (blobs_len < sizeof(roce_blob_t) * (uint64_t)c->world) { + snprintf(c->err, sizeof(c->err), "peer blob buffer too small"); + return -1; + } + const roce_blob_t *all = (const roce_blob_t *)blobs; + for (int p = 0; p < c->world; p++) { + if (all[p].abi_version != ROCE_ABI_VERSION || + all[p].world != (uint32_t)c->world || all[p].rank != (uint32_t)p || + all[p].layout_paths != ROCE_LAYOUT_PATHS || + all[p].opposite_paths != c->opposite_paths || + all[p].n_hca < ROCE_LAYOUT_PATHS || all[p].n_hca > ROCE_MAX_HCAS) { + snprintf(c->err, sizeof(c->err), + "rank %d published incompatible RoCE connection metadata", p); + return -1; + } + if (p == c->rank) { + continue; + } + c->peer_addr[p] = all[p].region_addr; + uint32_t remote_seen_hcas = 0; + for (int path = 0; path < (int)c->peer_path_count[p]; path++) { + int local_h = c->peer_hca[p][path]; + int remote_h = (int)all[p].peer_hca[c->rank][path]; + if (remote_h < 0 || remote_h >= (int)all[p].n_hca || + (remote_seen_hcas & (1u << remote_h)) != 0 || + all[p].region_addr == 0 || all[p].rkey[remote_h] == 0 || + all[p].qp_num[remote_h][c->rank] == 0) { + snprintf(c->err, sizeof(c->err), + "rank %d path %d published incomplete RoCE connection metadata", + p, path); + return -1; + } + remote_seen_hcas |= 1u << remote_h; + c->remote_hca[p][path] = remote_h; + c->peer_rkey[p][path] = all[p].rkey[remote_h]; + c->remote_qp[p][path] = all[p].qp_num[remote_h][c->rank]; + if (connect_qp(c, local_h, remote_h, p, &all[p]) != 0) { + return -1; + } + } + } + return 0; +} + +static int drain_cq(roce_ctx_t *c, int h) { + struct ibv_wc wc[32]; + int n = ibv_poll_cq(c->hca[h].cq, 32, wc); + if (n < 0) { + set_err(c, "ibv_poll_cq", errno); + return -1; + } + for (int i = 0; i < n; i++) { + uint64_t wr_id = wc[i].wr_id; + int peer = (int)(wr_id / ROCE_MAX_PATHS); + int path = (int)(wr_id % ROCE_MAX_PATHS); + if (peer < 0 || peer >= c->world || peer == c->rank || + path < 0 || path >= (int)c->peer_path_count[peer] || + c->peer_hca[peer][path] != h) { + snprintf(c->err, sizeof(c->err), + "RDMA completion has invalid path identifier %llu", + (unsigned long long)wr_id); + return -1; + } + if (wc[i].status != IBV_WC_SUCCESS) { + atomic_fetch_add(&c->completion_errors[peer][path], 1); + snprintf(c->err, sizeof(c->err), + "RDMA write to rank %d path %d failed: %s (vendor_err 0x%x, seq %u)", + peer, path, ibv_wc_status_str(wc[i].status), wc[i].vendor_err, + c->last_seq); + return -1; + } + c->hca[h].outstanding[peer] -= 1; + atomic_fetch_add(&c->send_completions[peer][path], 1); + c->writes_completed += 1; + } + return 0; +} + +static int post_path(roce_ctx_t *c, uint32_t seq, uint32_t slot, uint8_t *send, + int peer, int path, uint32_t stripe_offset, + uint32_t stripe_bytes) { + int h = c->peer_hca[peer][path]; + roce_hca_t *hca = &c->hca[h]; + // Each QP gets one signaled completion per operation. Keep its queue + // below one quarter of the configured depth. + while (hca->outstanding[peer] >= ROCE_SEND_DEPTH / 4) { + if (drain_cq(c, h) != 0) { + return -1; + } + if (!atomic_load_explicit(&c->running, memory_order_relaxed)) { + snprintf(c->err, sizeof(c->err), + "RoCE proxy stopped with %u writes outstanding to rank %d path %d", + hca->outstanding[peer], peer, path); + return -1; + } + } + uint32_t seq_copy = seq; + uint64_t remote = c->peer_addr[peer]; + struct ibv_sge data_sge = { + .addr = (uint64_t)(uintptr_t)(send + stripe_offset), + .length = stripe_bytes, + .lkey = hca->mr->lkey, + }; + struct ibv_sge flag_sge = { + .addr = (uint64_t)(uintptr_t)&seq_copy, + .length = 4, + .lkey = 0, + }; + struct ibv_send_wr flag_wr; + memset(&flag_wr, 0, sizeof(flag_wr)); + flag_wr.wr_id = (uint64_t)peer * ROCE_MAX_PATHS + (uint64_t)path; + flag_wr.sg_list = &flag_sge; + flag_wr.num_sge = 1; + flag_wr.opcode = IBV_WR_RDMA_WRITE; + flag_wr.send_flags = IBV_SEND_SIGNALED | IBV_SEND_INLINE; + uint32_t flag_row = path < ROCE_LAYOUT_PATHS ? (uint32_t)c->rank + : (uint32_t)peer; + uint32_t flag_column = path < ROCE_LAYOUT_PATHS + ? (uint32_t)path + : (uint32_t)(path - ROCE_LAYOUT_PATHS); + flag_wr.wr.rdma.remote_addr = + remote + c->flag_off + + (((uint64_t)flag_row * ROCE_SLOTS + slot) * ROCE_LAYOUT_PATHS + + flag_column) * ROCE_FLAG_STRIDE; + flag_wr.wr.rdma.rkey = c->peer_rkey[peer][path]; + struct ibv_send_wr data_wr; + memset(&data_wr, 0, sizeof(data_wr)); + data_wr.wr_id = flag_wr.wr_id; + data_wr.next = &flag_wr; + data_wr.sg_list = &data_sge; + data_wr.num_sge = 1; + data_wr.opcode = IBV_WR_RDMA_WRITE; + data_wr.send_flags = 0; + data_wr.wr.rdma.remote_addr = + remote + c->recv_off + + ((uint64_t)c->rank * ROCE_SLOTS + slot) * c->slot_bytes + stripe_offset; + data_wr.wr.rdma.rkey = c->peer_rkey[peer][path]; + struct ibv_send_wr *first_wr = stripe_bytes == 0 ? &flag_wr : &data_wr; + struct ibv_send_wr *bad = NULL; + int rc = ibv_post_send(hca->qp[peer], first_wr, &bad); + if (rc != 0) { + atomic_fetch_add(&c->completion_errors[peer][path], 1); + set_err(c, "ibv_post_send", rc); + return -1; + } + if (stripe_bytes != 0) { + atomic_fetch_add(&c->payload_writes[peer][path], 1); + atomic_fetch_add(&c->payload_bytes[peer][path], stripe_bytes); + } + atomic_fetch_add(&c->flag_writes[peer][path], 1); + hca->outstanding[peer] += 1; + return 0; +} + +static int post_peer(roce_ctx_t *c, uint32_t seq, uint32_t slot, uint8_t *send, + int peer, const uint32_t stripe_offset[ROCE_MAX_PATHS], + const uint32_t stripe_bytes[ROCE_MAX_PATHS]) { + for (int path = 0; path < (int)c->peer_path_count[peer]; path++) { + if (post_path(c, seq, slot, send, peer, path, stripe_offset[path], + stripe_bytes[path]) != 0) { + return -1; + } + } + return 0; +} + +static int drain_path(roce_ctx_t *c, int peer, int path) { + int h = c->peer_hca[peer][path]; + while (c->hca[h].outstanding[peer] != 0) { + if (drain_cq(c, h) != 0) { + return -1; + } + if (!atomic_load_explicit(&c->running, memory_order_relaxed)) { + snprintf(c->err, sizeof(c->err), + "RoCE proxy stopped while draining rank %d path %d", peer, path); + return -1; + } + } + return 0; +} + +static int drain_peer_paths(roce_ctx_t *c, int peer) { + for (;;) { + int pending = 0; + for (int path = 0; path < (int)c->peer_path_count[peer]; path++) { + int h = c->peer_hca[peer][path]; + if (c->hca[h].outstanding[peer] != 0) { + pending = 1; + if (drain_cq(c, h) != 0) { + return -1; + } + } + } + if (!pending) { + return 0; + } + if (!atomic_load_explicit(&c->running, memory_order_relaxed)) { + snprintf(c->err, sizeof(c->err), + "RoCE proxy stopped while draining direct paths to rank %d", peer); + return -1; + } + } +} + +static int mixed_two_wave(roce_ctx_t *c, uint32_t seq, uint32_t slot, + uint8_t *send, + const uint32_t stripe_offset[ROCE_MAX_PATHS], + const uint32_t stripe_bytes[ROCE_MAX_PATHS]) { + // Bits enumerate (peer, path) in peer-rank order while omitting the local + // rank. Each mask contains three origin QPs and balances every directed + // physical edge across the two waves for the canonical four-rank mapping. + static const uint8_t masks[4][2] = { + {0x0bu, 0x34u}, + {0x2cu, 0x13u}, + {0x31u, 0x0eu}, + {0x07u, 0x38u}, + }; + for (int wave = 0; wave < 2; wave++) { + int bit = 0; + for (int peer = 0; peer < c->world; peer++) { + if (peer == c->rank) { + continue; + } + for (int path = 0; path < ROCE_LAYOUT_PATHS; path++, bit++) { + if ((masks[c->rank][wave] & (1u << bit)) != 0 && + post_path(c, seq, slot, send, peer, path, + stripe_offset[path], stripe_bytes[path]) != 0) { + return -1; + } + } + } + bit = 0; + for (int peer = 0; peer < c->world; peer++) { + if (peer == c->rank) { + continue; + } + for (int path = 0; path < ROCE_LAYOUT_PATHS; path++, bit++) { + if ((masks[c->rank][wave] & (1u << bit)) != 0 && + drain_path(c, peer, path) != 0) { + return -1; + } + } + } + } + return 0; +} + +static int strict_three_wave(roce_ctx_t *c, uint32_t seq, uint32_t slot, + uint8_t *send, + const uint32_t stripe_offset[ROCE_MAX_PATHS], + const uint32_t stripe_bytes[ROCE_MAX_PATHS]) { + // Wave zero contains all four direct-link origins. Waves one and two each + // contain one reciprocal opposite-rank path, with the path assignment + // reversed between the two diagonal rank pairs. + static const uint8_t masks[4][3] = { + {0x33u, 0x04u, 0x08u}, + {0x0fu, 0x20u, 0x10u}, + {0x3cu, 0x01u, 0x02u}, + {0x33u, 0x08u, 0x04u}, + }; + for (int wave = 0; wave < 3; wave++) { + int bit = 0; + for (int peer = 0; peer < c->world; peer++) { + if (peer == c->rank) { + continue; + } + for (int path = 0; path < ROCE_LAYOUT_PATHS; path++, bit++) { + if ((masks[c->rank][wave] & (1u << bit)) != 0 && + post_path(c, seq, slot, send, peer, path, + stripe_offset[path], stripe_bytes[path]) != 0) { + return -1; + } + } + } + bit = 0; + for (int peer = 0; peer < c->world; peer++) { + if (peer == c->rank) { + continue; + } + for (int path = 0; path < ROCE_LAYOUT_PATHS; path++, bit++) { + if ((masks[c->rank][wave] & (1u << bit)) != 0 && + drain_path(c, peer, path) != 0) { + return -1; + } + } + } + } + return 0; +} + +static void split_stripes(uint32_t nbytes, int count, + uint32_t stripe_offset[ROCE_MAX_PATHS], + uint32_t stripe_bytes[ROCE_MAX_PATHS]) { + uint32_t packs = nbytes / 16u; + uint32_t base = packs / (uint32_t)count; + uint32_t remainder = packs % (uint32_t)count; + uint32_t offset = 0; + for (int path = 0; path < ROCE_MAX_PATHS; path++) { + uint32_t path_packs = path < count + ? base + ((uint32_t)path < remainder ? 1u : 0u) + : 0u; + stripe_offset[path] = offset; + stripe_bytes[path] = path_packs * 16u; + offset += stripe_bytes[path]; + } +} + +static int balanced32_post(roce_ctx_t *c, uint32_t seq, uint32_t slot, + uint8_t *send, + const uint32_t half_offset[ROCE_MAX_PATHS], + const uint32_t half_bytes[ROCE_MAX_PATHS], + const uint32_t quarter_offset[ROCE_MAX_PATHS], + const uint32_t quarter_bytes[ROCE_MAX_PATHS]) { + // Queue one path on every HCA before queueing a second. Rank and generation + // rotate the first HCA; generation parity alternates direct and opposite + // priority without a completion boundary or feedback loop. + int start_hca = (c->rank + (int)(seq & 3u)) % ROCE_MAX_HCAS; + int direct_first = (seq & 1u) == 0; + for (int round = 0; round < 2; round++) { + int post_direct = round == 0 ? direct_first : !direct_first; + for (int ordinal = 0; ordinal < ROCE_MAX_HCAS; ordinal++) { + int h = (start_hca + ordinal) % ROCE_MAX_HCAS; + int peer = post_direct ? c->direct_peer_by_hca[h] + : c->opposite_peer_by_hca[h]; + int path = post_direct ? c->direct_path_by_hca[h] + : c->opposite_path_by_hca[h]; + const uint32_t *offset = post_direct ? half_offset : quarter_offset; + const uint32_t *bytes = post_direct ? half_bytes : quarter_bytes; + if (peer < 0 || path < 0 || c->peer_hca[peer][path] != h) { + snprintf(c->err, sizeof(c->err), + "balanced path schedule is incomplete for HCA %d", h); + return -1; + } + if (post_path(c, seq, slot, send, peer, path, offset[path], + bytes[path]) != 0) { + return -1; + } + } + } + return 0; +} + +static int post_op(roce_ctx_t *c, uint32_t seq, uint32_t nbytes) { + uint32_t slot = seq & 1u; + uint8_t *send = c->region + c->send_off + (size_t)slot * c->slot_bytes; + if (nbytes == 0 || nbytes > c->slot_bytes || (nbytes % 16u) != 0) { + snprintf(c->err, sizeof(c->err), + "RoCE payload bytes must be a positive 16-byte multiple within the slot"); + return -1; + } + uint32_t stripe_offset[ROCE_MAX_PATHS] = {0}; + uint32_t stripe_bytes[ROCE_MAX_PATHS] = {0}; + uint32_t quarter_offset[ROCE_MAX_PATHS] = {0}; + uint32_t quarter_bytes[ROCE_MAX_PATHS] = {0}; + split_stripes(nbytes, ROCE_LAYOUT_PATHS, stripe_offset, stripe_bytes); + split_stripes(nbytes, ROCE_MAX_PATHS, quarter_offset, quarter_bytes); + if (c->opposite_paths == ROCE_MAX_PATHS) { + if (balanced32_post(c, seq, slot, send, stripe_offset, stripe_bytes, + quarter_offset, quarter_bytes) != 0) { + return -1; + } + goto posted; + } + int two_wave = c->world == 4 && c->n_hca == 4 && + c->two_wave_threshold_bytes != 0 && + nbytes >= c->two_wave_threshold_bytes; + if (two_wave) { + atomic_fetch_add(&c->two_wave_activations, 1); + if (c->wave_mode == ROCE_WAVE_MODE_MIXED_TWO) { + if (mixed_two_wave(c, seq, slot, send, stripe_offset, + stripe_bytes) != 0) { + return -1; + } + } else if (c->wave_mode == ROCE_WAVE_MODE_STRICT_THREE) { + if (strict_three_wave(c, seq, slot, send, stripe_offset, + stripe_bytes) != 0) { + return -1; + } + } else if (c->wave_mode == ROCE_WAVE_MODE_OPPOSITE_FIRST) { + int opposite = (c->rank + 2) % c->world; + if (post_peer(c, seq, slot, send, opposite, stripe_offset, + stripe_bytes) != 0 || + drain_peer_paths(c, opposite) != 0) { + return -1; + } + for (int peer = 0; peer < c->world; peer++) { + int distance = (peer - c->rank + c->world) % c->world; + if ((distance == 1 || distance == 3) && + post_peer(c, seq, slot, send, peer, stripe_offset, + stripe_bytes) != 0) { + return -1; + } + } + } else { + // Direct-link QPs complete before the two hardware-forwarded QPs + // are submitted, preventing both traffic classes from competing + // in the same ConnectX reliability window for larger payloads. + for (int peer = 0; peer < c->world; peer++) { + int distance = (peer - c->rank + c->world) % c->world; + if (distance == 1 || distance == 3) { + if (post_peer(c, seq, slot, send, peer, stripe_offset, + stripe_bytes) != 0) { + return -1; + } + } + } + for (int peer = 0; peer < c->world; peer++) { + int distance = (peer - c->rank + c->world) % c->world; + if ((distance == 1 || distance == 3) && + drain_peer_paths(c, peer) != 0) { + return -1; + } + } + int opposite = (c->rank + 2) % c->world; + if (post_peer(c, seq, slot, send, opposite, stripe_offset, + stripe_bytes) != 0) { + return -1; + } + } + } else { + for (int peer = 0; peer < c->world; peer++) { + if (peer != c->rank && + post_peer(c, seq, slot, send, peer, stripe_offset, + stripe_bytes) != 0) { + return -1; + } + } + } +posted: + c->ops_posted += 1; + for (int h = 0; h < c->n_hca; h++) { + if (drain_cq(c, h) != 0) { + return -1; + } + } + return 0; +} + +static void *proxy_main(void *arg) { + roce_ctx_t *c = (roce_ctx_t *)arg; + volatile uint32_t *ctrl = (volatile uint32_t *)(c->region + c->ctrl_off); + // Spin while ops are flowing. After ROCE_IDLE_SPINS polls without a + // doorbell, request a short nanosleep between polls (the OS decides the + // actual delay) so an idle runtime does not hold a core next to the + // serving process. The missed-doorbell catch-up below keeps the protocol + // correct however long the thread is away. + uint64_t idle = 0; + const struct timespec nap = {0, 20000}; + while (atomic_load_explicit(&c->running, memory_order_relaxed)) { + uint32_t seq = __atomic_load_n(&ctrl[0], __ATOMIC_ACQUIRE); + if (seq == c->last_seq) { + idle++; + if (idle % 64 == 0) { + for (int h = 0; h < c->n_hca; h++) { + if (drain_cq(c, h) != 0) { + atomic_store(&c->failed, 1); + return NULL; + } + } + } + if (idle >= ROCE_IDLE_SPINS) { + nanosleep(&nap, NULL); + } + continue; + } + idle = 0; + // The doorbell holds only the newest sequence. Our kernel for op N + // completes on the peers' payloads alone, so op N+1 can ring before + // this thread has seen op N (it slept, or the scheduler moved it). + // Peers cannot get further than one op ahead of us, so at most + // ROCE_SLOTS doorbells are pending and every send slot is intact: + // post each missed sequence in order using its per-slot byte count. + uint32_t pending = seq - c->last_seq; + if (pending > ROCE_SLOTS) { + snprintf(c->err, sizeof(c->err), + "doorbell skipped %u ops (last %u, now %u)", pending, c->last_seq, seq); + atomic_store(&c->failed, 1); + return NULL; + } + for (uint32_t s = c->last_seq + 1; pending > 0; s++, pending--) { + uint32_t nbytes = ctrl[4 + (s & 1u)]; + if (post_op(c, s, nbytes) != 0) { + atomic_store(&c->failed, 1); + return NULL; + } + c->last_seq = s; + } + } + return NULL; +} + +int roce_start(roce_ctx_t *c) { + if (atomic_load(&c->running)) { + return 0; + } + if (!c->started) { + // A restart continues from the last posted sequence so ops that rang + // the doorbell while the thread was stopped are still posted. + volatile uint32_t *ctrl = (volatile uint32_t *)(c->region + c->ctrl_off); + c->last_seq = ctrl[0]; + c->started = 1; + } + atomic_store(&c->failed, 0); + atomic_store(&c->running, 1); + int rc = pthread_create(&c->thread, NULL, proxy_main, c); + if (rc != 0) { + atomic_store(&c->running, 0); + set_err(c, "pthread_create", rc); + return -1; + } + return 0; +} + +void roce_stop(roce_ctx_t *c) { + if (atomic_exchange(&c->running, 0)) { + pthread_join(c->thread, NULL); + } +} + +int roce_failed(roce_ctx_t *c) { return atomic_load(&c->failed); } + +const char *roce_error(roce_ctx_t *c) { return c->err; } + +uint64_t roce_stat(roce_ctx_t *c, int which) { + switch (which) { + case 0: + return c->ops_posted; + case 1: + return c->writes_completed; + case 2: + return c->last_seq; + case 3: + return atomic_load(&c->two_wave_activations); + default: + return 0; + } +} + +uint64_t roce_two_wave_threshold_bytes(roce_ctx_t *c) { + return c->two_wave_threshold_bytes; +} + +uint64_t roce_wave_mode(roce_ctx_t *c) { + return c->wave_mode; +} + +int roce_peer_hca(roce_ctx_t *c, int peer, int path) { + if (peer < 0 || peer >= c->world || path < 0 || + path >= (int)c->peer_path_count[peer]) { + return -1; + } + return c->peer_hca[peer][path]; +} + +uint64_t roce_path_stat(roce_ctx_t *c, int peer, int path, int which) { + if (peer < 0 || peer >= c->world || peer == c->rank || + path < 0 || path >= (int)c->peer_path_count[peer]) { + return UINT64_MAX; + } + switch (which) { + case 0: + return atomic_load(&c->payload_writes[peer][path]); + case 1: + return atomic_load(&c->payload_bytes[peer][path]); + case 2: + return atomic_load(&c->payload_bytes[peer][path]) * c->physical_hops[peer]; + case 3: + return atomic_load(&c->flag_writes[peer][path]); + case 4: + return atomic_load(&c->send_completions[peer][path]); + case 5: + return atomic_load(&c->completion_errors[peer][path]); + case 6: + return c->hca[c->peer_hca[peer][path]].qp[peer]->qp_num; + case 7: + return c->remote_qp[peer][path]; + case 8: + return (uint64_t)c->peer_hca[peer][path]; + case 9: + return (uint64_t)c->remote_hca[peer][path]; + case 10: + return c->physical_hops[peer]; + default: + return UINT64_MAX; + } +} + +void roce_destroy(roce_ctx_t *c) { + if (c == NULL) { + return; + } + roce_stop(c); + for (int h = 0; h < ROCE_MAX_HCAS; h++) { + roce_hca_t *hca = &c->hca[h]; + for (int p = 0; p < ROCE_MAX_PEERS; p++) { + if (hca->qp[p] != NULL) { + ibv_destroy_qp(hca->qp[p]); + } + } + if (hca->cq != NULL) { + ibv_destroy_cq(hca->cq); + } + if (hca->mr != NULL) { + ibv_dereg_mr(hca->mr); + } + if (hca->pd != NULL) { + ibv_dealloc_pd(hca->pd); + } + if (hca->ctx != NULL) { + ibv_close_device(hca->ctx); + } + } + free(c); +} diff --git a/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/b12x_overlay/b12x/comm/roce/api.py b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/b12x_overlay/b12x/comm/roce/api.py new file mode 100644 index 00000000..b6be3ef1 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/b12x_overlay/b12x/comm/roce/api.py @@ -0,0 +1,27 @@ +"""Public surface for comm.roce (docs in the op ``__init__``).""" + +from __future__ import annotations + +from .roce_oneshot import ( + API_VERSION, + DEFAULT_MAX_GATHER_BYTES, + DEFAULT_MAX_SIZE, + SUPPORTED_DTYPES, + SUPPORTED_WORLD_SIZES, + RoceOneshotAllReduce as AllReduce, + default_gid_index, + discover_hcas, + is_supported, +) + +__all__ = [ + "API_VERSION", + "AllReduce", + "DEFAULT_MAX_GATHER_BYTES", + "DEFAULT_MAX_SIZE", + "SUPPORTED_DTYPES", + "SUPPORTED_WORLD_SIZES", + "default_gid_index", + "discover_hcas", + "is_supported", +] diff --git a/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/b12x_overlay/b12x/comm/roce/roce_oneshot.py b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/b12x_overlay/b12x/comm/roce/roce_oneshot.py new file mode 100644 index 00000000..a97ad426 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/b12x_overlay/b12x/comm/roce/roce_oneshot.py @@ -0,0 +1,942 @@ +"""RoCEnante: one-shot RoCE all-reduce and all-gather runtime for multi-node tensor parallelism. + +Designed for DGX Spark clusters, whose integrated GPU can read pinned host +memory in place and whose ConnectX-7 can RDMA-write into the same memory. +Each rank owns one pinned region (see ``_roce_proxy.c``); every all-reduce is +one kernel launch that stages the input, rings the proxy thread, waits for +every peer's RDMA-written payload, and reduces in fixed rank order. + +Runtime constraints: + +* one all-reduce in flight per runtime (single channel, one stream); +* all-reduce messages up to ``max_size`` bytes and all-gather shards up to + ``max_gather_bytes``, both multiples of 16 bytes; +* two paths per peer by default; the research-only TP4 cycle mode selected by + ``B12X_ROCE_OPPOSITE_PATHS=4`` keeps two neighbor paths and uses four paths + to the opposite rank; +* every rank of the exchange group must construct the runtime collectively. +""" + +from __future__ import annotations + +import contextlib +import logging +import os +import threading +from contextlib import contextmanager +from pathlib import Path +from typing import Any, Optional, Sequence + +import torch +import torch.distributed as dist +from torch.distributed import ProcessGroup + +from . import _allgather_cute +from ._oneshot_cute import PACK_BYTES, get_launcher, is_launcher_prepared +from ._path_config import ( + OPPOSITE_PATHS_ENV, + PATH_COUNT, + opposite_path_count as _opposite_path_count, + peer_hca_map as _peer_hca_map, + validate_four_path_tp4_mapping as _validate_four_path_tp4_mapping, +) +from ._proxy import Layout, Proxy, load as _load_proxy_library + +logger = logging.getLogger(__name__) + +# Bumped when the adapter-visible surface (constructor keywords, methods, their +# semantics) changes incompatibly; integrations pin the value they were built for. +API_VERSION = 1 +SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32) +SUPPORTED_WORLD_SIZES = tuple(range(2, 17)) +DEFAULT_MAX_SIZE = 2 * 1024 * 1024 +DEFAULT_MAX_GATHER_BYTES = 16 * 1024 * 1024 +DEFAULT_THREADS = 512 +DEFAULT_BLOCKS = 8 +ADAPTIVE_GEOMETRY = ("weighted8-capture-only-v1", ((16384, 2), (65536, 4)), 8) + + +def _allreduce_blocks(nbytes: int) -> int: + """Select a divisor-of-eight grid; each operation contributes eight arrivals.""" + return 2 if nbytes <= 16384 else 4 if nbytes <= 65536 else 8 + +DEFAULT_GID_INDEX = 3 +# Polls of a peer flag before the kernel gives up (each poll is a system-scope +# load of host memory, roughly a microsecond): about 20 s. +DEFAULT_SPIN_LIMIT = 20_000_000 +_SLOT_ALIGNMENT = 4096 +_DTYPE_NAMES = { + torch.float16: "float16", + torch.bfloat16: "bfloat16", + torch.float32: "float32", +} + + +def _env_list(*names: str) -> tuple[str, ...]: + """First non-empty comma-separated list among the environment variables ``names``.""" + for name in names: + raw = os.getenv(name) + if raw: + items = [] + for item in raw.split(","): + item = item.strip().lstrip("=^") + if item: + items.append(item.split(":")[0]) + if items: + return tuple(items) + return () + + +def _env_int(*names: str, default: int) -> int: + """First integer-valued environment variable among ``names``, else ``default``.""" + for name in names: + raw = os.getenv(name) + if raw: + return int(raw) + return default + + +def default_gid_index() -> int: + """``B12X_ROCE_GID_INDEX``, else NCCL's ``NCCL_IB_GID_INDEX``, else 3.""" + + return _env_int("B12X_ROCE_GID_INDEX", "NCCL_IB_GID_INDEX", default=DEFAULT_GID_INDEX) + + +def discover_hcas(gid_index: Optional[int] = None) -> tuple[str, ...]: + """Return the RDMA devices to use, at most four. + + ``B12X_ROCE_HCA`` (or NCCL's ``NCCL_IB_HCA``) selects explicitly; otherwise + every active device with a populated GID at ``gid_index`` is used. + """ + + explicit = _env_list("B12X_ROCE_HCA", "NCCL_IB_HCA") + if explicit: + return explicit[:4] + gid_index = default_gid_index() if gid_index is None else int(gid_index) + found = [] + root = Path("/sys/class/infiniband") + for dev in sorted(root.glob("*")): + state = dev / "ports" / "1" / "state" + gid = dev / "ports" / "1" / "gids" / str(gid_index) + try: + if "ACTIVE" not in state.read_text(): + continue + if gid.read_text().strip().replace(":", "").strip("0") == "": + continue + except OSError: + continue + found.append(dev.name) + return tuple(found[:4]) + + +def is_supported(device: torch.device | int | str | None = None) -> bool: + """True on an integrated GPU with at least one active RDMA device. + + The kernel reads pinned host memory in place, which needs an integrated + (unified-memory) GPU such as the DGX Spark GB10. + """ + + if not torch.cuda.is_available(): + return False + index = torch.cuda.current_device() if device is None else torch.device(device).index + props = torch.cuda.get_device_properties(index if index is not None else 0) + if not getattr(props, "is_integrated", False): + return False + return len(discover_hcas()) > 0 + + +def _normalize_device(device: torch.device | int | str) -> torch.device: + """Coerce an int, string, or device into a CUDA ``torch.device`` with an index.""" + if isinstance(device, int): + device = torch.device("cuda", device) + elif not isinstance(device, torch.device): + device = torch.device(device) + if device.type != "cuda": + raise ValueError("RoCE all-reduce requires a CUDA device") + if device.index is None: + device = torch.device("cuda", torch.cuda.current_device()) + return device + + +def _capture_id(stream: torch.cuda.Stream) -> int: + """CUDA's id of the capture ``stream`` is in, 0 when it is not capturing. + + Two captures are distinct even on the same stream, and one capture can be + joined by several streams; the id identifies the capture itself. + """ + + from cuda.bindings import runtime as cudart + + info = cudart.cudaStreamGetCaptureInfo(stream.cuda_stream) + if info[0] != cudart.cudaError_t.cudaSuccess: + return 0 + if info[1] != cudart.cudaStreamCaptureStatus.cudaStreamCaptureStatusActive: + return 0 + return int(info[2]) + + +def _align_up(value: int, alignment: int) -> int: + """Round ``value`` up to a multiple of ``alignment``.""" + return (int(value) + alignment - 1) // alignment * alignment + + +def _exchange(local: object, group: ProcessGroup) -> list[object]: + """All-gather one picklable object from every rank of ``group``.""" + gathered: list[object] = [None] * dist.get_world_size(group=group) + dist.all_gather_object(gathered, local, group=group) + return gathered + + +class RoceOneshotAllReduce: + """One-shot RDMA all-reduce over the DGX Spark 200 GbE fabric.""" + + algorithm = "rocenante" + + def __init__( + self, + *, + exchange_group: ProcessGroup, + device: torch.device | int | str, + max_size: int = DEFAULT_MAX_SIZE, + max_gather_bytes: int = DEFAULT_MAX_GATHER_BYTES, + hca_names: Optional[Sequence[str]] = None, + gid_index: Optional[int] = None, + threads: int = DEFAULT_THREADS, + blocks: int = DEFAULT_BLOCKS, + ) -> None: + """Allocate the pinned region, build and connect the proxy, and start it.""" + self.device = _normalize_device(device) + self.rank = dist.get_rank(group=exchange_group) + self.world_size = dist.get_world_size(group=exchange_group) + self._group = exchange_group + self._closed = False + self._lock = threading.Lock() + self._proxy: Optional[Proxy] = None + self._gather_buffers: Optional[tuple[torch.Tensor, torch.Tensor]] = None + self._align_buffers: Optional[tuple[torch.Tensor, torch.Tensor]] = None + self._stream_event = torch.cuda.Event() + self._last_stream: Optional[torch.cuda.Stream] = None + self._capture_stream: Optional[torch.cuda.Stream] = None + self._capture_id = 0 + if self.world_size not in SUPPORTED_WORLD_SIZES: + raise ValueError( + f"unsupported RoCE all-reduce world size {self.world_size}; " + f"supported sizes are {SUPPORTED_WORLD_SIZES}" + ) + if int(max_size) < PACK_BYTES: + raise ValueError("max_size must hold at least one 16-byte pack") + if int(threads) % 32 != 0 or int(threads) < 32 or int(threads) > 1024: + raise ValueError("threads must be a multiple of 32 between 32 and 1024") + if int(blocks) < 1 or int(blocks) & (int(blocks) - 1) != 0: + # The kernels find the last arriving block with a modulo of a + # free-running 32-bit counter, which survives the counter wrap only + # for a power-of-two block count. + raise ValueError("blocks must be a positive power of two") + if self.world_size != 4 or int(threads) != 512 or int(blocks) != 8: + raise ValueError("adaptive weighted8 requires world_size=4, threads=512, base blocks=8") + self.max_size = int(max_size) + self.max_gather_bytes = int(max_gather_bytes) + self._threads = int(threads) + self._blocks = int(blocks) + self.gid_index = default_gid_index() if gid_index is None else int(gid_index) + self.spin_limit = _env_int("B12X_ROCE_SPIN_LIMIT", default=DEFAULT_SPIN_LIMIT) + names = tuple(hca_names) if hca_names else discover_hcas(self.gid_index) + if not names: + raise RuntimeError("no active RDMA device found for the RoCE all-reduce") + self.hca_names = names[:4] + path_error: Optional[str] = None + self._opposite_paths = PATH_COUNT + self._peer_hca_map: tuple[tuple[int, ...], ...] = tuple() + try: + self._opposite_paths = _opposite_path_count( + self.world_size, len(self.hca_names) + ) + self._peer_hca_map = _peer_hca_map( + self.world_size, + self.rank, + len(self.hca_names), + self._opposite_paths, + ) + _validate_four_path_tp4_mapping( + self.hca_names, + self.world_size, + self.rank, + self._peer_hca_map, + self._opposite_paths, + ) + except Exception as exc: # noqa: BLE001 - voted across every rank below + path_error = str(exc) + path_statuses = _exchange( + (path_error, self._opposite_paths), exchange_group + ) + path_failures = [ + f"rank {index}: {status[0]}" + for index, status in enumerate(path_statuses) + if status[0] is not None + ] + path_counts = {int(status[1]) for status in path_statuses} + if len(path_counts) != 1: + path_failures.append( + f"ranks selected different {OPPOSITE_PATHS_ENV} values: " + f"{sorted(path_counts)}" + ) + if path_failures: + raise RuntimeError( + "RoCE all-reduce path preflight failed: " + "; ".join(path_failures) + ) + + slot_bytes = _align_up(max(self.max_size, self.max_gather_bytes), _SLOT_ALIGNMENT) + self._layout = Layout(self.world_size, slot_bytes) + self._slot_bytes = slot_bytes + + with torch.cuda.device(self.device): + # Pinned, zero-initialised: flags and the control record start at 0 + # and the first sequence number is 1. + self._region = torch.zeros( + self._layout.total_bytes, dtype=torch.uint8, pin_memory=True + ) + self._counters = torch.zeros(4, dtype=torch.int32, device=self.device) + host_ptr = self._region.data_ptr() + device_ptr = self._device_pointer(host_ptr) + if device_ptr != host_ptr: + raise RuntimeError( + "RoCE all-reduce needs host pointers that are directly device " + "accessible (integrated GPU with unified addressing)" + ) + self._recv_base = host_ptr + self._layout.recv_off + self._flag_base = host_ptr + self._layout.flag_off + self._send_base = host_ptr + self._layout.send_off + self._ctrl_base = host_ptr + self._layout.ctrl_off + # ctrl record (kernel-written): seq, nbytes, error seq, missing peer, + # nbytes per slot (the proxy uses these when it has to catch up) + self._ctrl_words = self._region[ + self._layout.ctrl_off : self._layout.ctrl_off + 24 + ].view(torch.int32) + self._error_word = self._ctrl_words[2:3] + # numpy view of the control words: reading it costs nanoseconds, so the + # health check before and after every launch stays off the profile. + self._ctrl_np = self._ctrl_words.numpy() + self._epoch_address = self._counters.data_ptr() + + error: Optional[str] = None + try: + self._proxy = Proxy( + world_size=self.world_size, + rank=self.rank, + hca_names=self.hca_names, + gid_index=self.gid_index, + region_ptr=host_ptr, + region_bytes=self._layout.total_bytes, + slot_bytes=slot_bytes, + peer_hca_map=self._peer_hca_map, + opposite_paths=self._opposite_paths, + ) + blob = self._proxy.local_blob() + except Exception as exc: # noqa: BLE001 - reported collectively below + error = str(exc) + blob = b"" + # Every rank publishes the configuration the protocol depends on; the + # ranks must agree exactly, and all of them see the same verdict. + config = { + "api_version": API_VERSION, + "proxy_abi": _load_proxy_library().roce_abi_version() if error is None else None, + "world_size": self.world_size, + "hca_count": len(self.hca_names), + "slot_bytes": slot_bytes, + "slots": self._layout.slots, + "paths": self._layout.paths, + "opposite_paths": self._opposite_paths, + "origin_paths_per_rank": 2 * (self.world_size - 2) + self._opposite_paths + if self.world_size == 4 + else PATH_COUNT * (self.world_size - 1), + "two_wave_threshold_bytes": ( + self._proxy.two_wave_threshold_bytes() + if self._proxy is not None + else None + ), + "wave_mode": self._proxy.wave_mode() if self._proxy is not None else None, + "flag_stride": self._layout.flag_stride, + "max_size": self.max_size, + "max_gather_bytes": self.max_gather_bytes, + "spin_limit": self.spin_limit, + "threads": self._threads, + # Reuse an existing compared key so an unmodified rank-zero peer + # also rejects this adaptive package before connecting the QPs. + "blocks": (self._blocks, ADAPTIVE_GEOMETRY), + } + statuses = _exchange((error, blob, config), exchange_group) + failures = [f"rank {i}: {s[0]}" for i, s in enumerate(statuses) if s[0] is not None] + if not failures: + reference = statuses[0][2] + for i, s in enumerate(statuses): + differing = {k: (reference[k], s[2].get(k)) for k in reference if s[2].get(k) != reference[k]} + if differing: + failures.append(f"rank {i} configuration differs from rank 0: {differing}") + if failures: + self.close() + raise RuntimeError("RoCE all-reduce setup failed: " + "; ".join(failures)) + try: + self._proxy.connect([s[1] for s in statuses]) + self._proxy.start() + except Exception as exc: # noqa: BLE001 + error = str(exc) + verdicts = _exchange(error, exchange_group) + failures = [f"rank {i}: {v}" for i, v in enumerate(verdicts) if v is not None] + if failures: + self.close() + raise RuntimeError("RoCE all-reduce connect failed: " + "; ".join(failures)) + if self.rank == 0: + logger.info( + "RoCEnante ready: world=%d hcas=%s gid_index=%d max_size=%d", + self.world_size, + ",".join(self.hca_names), + self.gid_index, + self.max_size, + ) + + @staticmethod + def _device_pointer(host_ptr: int) -> int: + """Device address of pinned host memory (``cudaHostGetDevicePointer``).""" + from cuda.bindings import runtime as cudart + + err, ptr = cudart.cudaHostGetDevicePointer(host_ptr, 0) + if err != cudart.cudaError_t.cudaSuccess: + raise RuntimeError(f"cudaHostGetDevicePointer failed: {err}") + return int(ptr) + + @classmethod + def from_exchange_group( + cls, + *, + exchange_group: ProcessGroup, + device: torch.device | int | str, + max_size: int = DEFAULT_MAX_SIZE, + eager_buffer_bytes: Optional[int] = None, + max_gather_bytes: int = DEFAULT_MAX_GATHER_BYTES, + **_ignored: Any, + ) -> "RoceOneshotAllReduce": + """Mirror ``comm.pcie.AllReduce.from_exchange_group``; PCIe-only knobs are ignored.""" + + capacity = max(int(max_size), int(eager_buffer_bytes or 0)) + return cls( + exchange_group=exchange_group, + device=device, + max_size=capacity, + max_gather_bytes=max_gather_bytes, + ) + + @classmethod + def from_process_group( + cls, + *, + process_group: ProcessGroup, + device: torch.device | int | str, + max_size: int = DEFAULT_MAX_SIZE, + max_input_bytes: Optional[int] = None, + **_ignored: Any, + ) -> "RoceOneshotAllReduce": + """Build a runtime from a process group used for both rank identity and the setup exchange.""" + capacity = max(int(max_size), int(max_input_bytes or 0)) + return cls(exchange_group=process_group, device=device, max_size=capacity) + + # -- policy ----------------------------------------------------------------- + + @property + def supports_all_peer_auxiliary(self) -> bool: + """False: peer auxiliary inputs (fused residual paths) are not implemented.""" + return False + + def should_allreduce(self, inp: torch.Tensor) -> bool: + """Eligible: contiguous CUDA tensor of a supported dtype, 16-byte total size within ``max_size``. + + The decision depends only on dtype, shape, contiguity and size, which + tensor-parallel ranks share, never on the pointer value; a closed runtime + raises rather than declining, so no rank can silently fall back alone. + """ + + if self._closed or self._proxy is None: + raise RuntimeError("RoCE runtime is closed") + if inp.dtype not in SUPPORTED_DTYPES or not inp.is_cuda: + return False + if inp.device != self.device or not inp.is_contiguous(): + return False + nbytes = inp.numel() * inp.element_size() + return 0 < nbytes <= self.max_size and nbytes % PACK_BYTES == 0 + + # -- channels / streams (single channel runtime) --------------------------- + + def prepare_channels(self, channel_ids: Sequence[str]) -> None: + """No-op: the runtime has a single channel.""" + return None + + def for_stream(self, stream: object = None, *, channel_id: Optional[str] = None): + """Return this runtime; the collectives take ``stream`` per call.""" + return self + + # -- compilation ------------------------------------------------------------ + + def _launcher_key(self, dtype: torch.dtype) -> tuple[object, ...]: + """Cache key of the all-reduce launcher for ``dtype``.""" + return ( + _DTYPE_NAMES[dtype], + self.world_size, + self.rank, + self._threads, + self._layout.slots, + self._layout.flag_stride, + self._opposite_paths, + self.device.index, + ) + + def _gather_launcher_key(self) -> tuple[object, ...]: + """Cache key of the all-gather launcher.""" + return ( + self.world_size, + self.rank, + self._threads, + self._layout.slots, + self._layout.flag_stride, + self._opposite_paths, + self.device.index, + ) + + def prepare( + self, dtypes: Sequence[torch.dtype] = (torch.bfloat16,), *, padded_gather: bool = False + ) -> None: + """Compile the launchers for ``dtypes`` and allocate scratch ahead of capture. + + Always allocates the two ``max_size`` alignment buffers. The padded + all-gather scratch (``max_gather_bytes`` plus ``world_size`` times that) + is allocated only with ``padded_gather=True`` or on first eager use, so + a workload whose shards have 16-byte rows never pays for it. + Compilation and allocation are refused inside a capture. + """ + + with torch.cuda.device(self.device): + for dtype in dtypes: + get_launcher(*self._launcher_key(dtype)) + _allgather_cute.get_launcher(*self._gather_launcher_key()) + self._aligned_scratch(0, self._region[:PACK_BYTES]) + if padded_gather: + self._gather_scratch(PACK_BYTES) + + def prepare_graph_all_reduce(self, inp: torch.Tensor, *, stream: object = None) -> None: + """Compile the launcher for ``inp.dtype`` ahead of CUDA graph capture.""" + self.prepare((inp.dtype,)) + + # -- execution ---------------------------------------------------------------- + + def all_reduce( + self, + inp: torch.Tensor, + *, + out: Optional[torch.Tensor] = None, + stream: object = None, + channel_id: Optional[str] = None, + peer_input_ptrs: Optional[Sequence[int]] = None, + ) -> torch.Tensor: + """Sum ``inp`` across ranks into ``out`` (allocated like ``inp`` when omitted). + + Eligibility (``should_allreduce``) depends only on dtype, shape and size, + so every tensor-parallel rank takes the same decision; a pointer that is + not 16-byte aligned is staged through runtime scratch instead of being + rejected. A poisoned runtime (a previous wait timed out) raises here + before anything is launched: a failed collective is fatal, never a + fallback. Admission through kernel enqueue holds the lifecycle lock so + ``close`` cannot release the proxy meanwhile. Launches on different + streams are ordered with an event; under CUDA graph capture every + collective must be captured on one stream. Without ``out`` the output + comes from the caller's allocator; under capture that is the graph's + private pool, which replays at a fixed address. + """ + + with self._lock: + self.check_health() + if not self.should_allreduce(inp): + raise ValueError("input is not eligible for the RoCE one-shot all-reduce") + if out is not None and ( + out.shape != inp.shape + or out.dtype != inp.dtype + or out.device != inp.device + or not out.is_contiguous() + ): + raise ValueError("out must be a contiguous tensor on the input's device matching the input") + key = self._launcher_key(inp.dtype) + nbytes = inp.numel() * inp.element_size() + context = torch.cuda.stream(stream) if stream is not None else _nullcontext() + with torch.cuda.device(self.device), context: + # Capture state belongs to the target stream, so query it here. + capturing = torch.cuda.is_current_stream_capturing() + if capturing and not is_launcher_prepared(*key): + raise RuntimeError( + "RoCE all-reduce launcher must be prepared before CUDA graph capture" + ) + launcher = get_launcher(*key) + if out is None: + out = torch.empty_like(inp) + # Shared staging buffers belong to the preceding operation until + # its completion event, including its output copy, has fired. + self._order_stream(capturing) + src = inp + if inp.data_ptr() % PACK_BYTES != 0: + src = self._aligned_scratch(0, inp) + src.copy_(inp) + dst = out if out.data_ptr() % PACK_BYTES == 0 else self._aligned_scratch(1, out) + launcher( + src.data_ptr(), + dst.data_ptr(), + nbytes // PACK_BYTES, + nbytes, + self._recv_base, + self._flag_base, + self._send_base, + self._ctrl_base, + self._slot_bytes, + self._epoch_address, + self.spin_limit, + _allreduce_blocks(nbytes) if capturing else self._blocks, + ) + if dst is not out: + out.copy_(dst) + self._mark_stream(capturing) + if not capturing: + self.check_health() + return out + + def _aligned_scratch(self, which: int, like: torch.Tensor) -> torch.Tensor: + """16-byte-aligned device scratch shaped like ``like`` (input 0, output 1). + + Allocated by ``prepare`` (two ``max_size`` buffers); a first use under + CUDA graph capture is refused. + """ + + if self._align_buffers is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "RoCE all-reduce alignment scratch must be allocated before CUDA graph " + "capture; call prepare() first" + ) + self._align_buffers = ( + torch.empty(self.max_size, dtype=torch.uint8, device=self.device), + torch.empty(self.max_size, dtype=torch.uint8, device=self.device), + ) + nbytes = like.numel() * like.element_size() + return self._align_buffers[which][:nbytes].view(like.dtype).view(like.shape) + + def _order_stream(self, capturing: bool) -> None: + """Make the current stream wait for the previous collective's stream. + + The runtime has one epoch, one doorbell and two transport slots, so + collectives must execute in launch order even across streams. Under + capture no event is used: every collective must be on the capture + stream, which keeps the graph's own ordering; captures are told apart + by CUDA's capture id, so consecutive captures may use different streams. + """ + + current = torch.cuda.current_stream(self.device) + if capturing: + capture_id = _capture_id(current) + if capture_id != self._capture_id: + self._capture_id = capture_id + self._capture_stream = current + elif current != self._capture_stream: + raise RuntimeError("RoCE collectives must all be captured on one stream") + return + if self._last_stream is not None and current != self._last_stream: + current.wait_event(self._stream_event) + + def _mark_stream(self, capturing: bool) -> None: + """Record the ordering event on the current stream after a launch.""" + + if capturing: + return + current = torch.cuda.current_stream(self.device) + self._stream_event.record(current) + self._last_stream = current + + def check_health(self) -> None: + """Raise if the proxy thread died or a kernel wait timed out. + + Cheap (two host memory reads); call it after graph replays, which + cannot check inline. + """ + + if self._proxy is not None and self._proxy.failed(): + raise RuntimeError(f"RoCE proxy failed: {self._proxy.error()}") + failed_seq = int(self._ctrl_np[2]) + if failed_seq != 0: + peer = int(self._ctrl_np[3]) + raise RuntimeError( + f"RoCE collective on rank {self.rank} timed out waiting for rank {peer} " + f"at sequence {failed_seq}; the runtime is poisoned (its epoch stopped at " + f"{failed_seq - 1}, later launches do nothing) and rank data is no " + "longer trustworthy" + ) + + @property + def poisoned(self) -> bool: + """True once a wait timed out or the proxy failed; the runtime cannot be reused.""" + + return (self._proxy is not None and self._proxy.failed()) or int(self._ctrl_np[2]) != 0 + + # -- all-gather --------------------------------------------------------------- + + @staticmethod + def _normalize_dim(inp: torch.Tensor, dim: int) -> int: + """Map a negative ``dim`` to its positive index.""" + if dim < 0: + dim += inp.dim() + return dim + + def should_all_gather(self, inp: torch.Tensor, dim: int = -1) -> bool: + """Eligible: contiguous CUDA tensor of any dtype, concat along dim 0 or the last dim. + + 16-byte-aligned rows take the direct-layout kernel; anything else goes + through a padded contiguous gather plus a torch reshape, so shape never + forces a fallback to another backend. + """ + + if self._closed or self._proxy is None: + raise RuntimeError("RoCE runtime is closed") + if not inp.is_cuda or inp.dim() == 0: + return False + if inp.device != self.device or not inp.is_contiguous(): + return False + if inp.is_complex() or inp.is_sparse or inp.dtype == torch.bool: + return False + dim = self._normalize_dim(inp, dim) + if dim not in (0, inp.dim() - 1): + return False + nbytes = inp.numel() * inp.element_size() + return 0 < nbytes <= self.max_gather_bytes + + def _direct_gather_layout(self, inp: torch.Tensor, dim: int) -> bool: + """True when ``inp`` can be gathered straight into the concatenated layout (16-byte rows and pointer).""" + nbytes = inp.numel() * inp.element_size() + if nbytes % PACK_BYTES != 0 or inp.data_ptr() % PACK_BYTES != 0: + return False + return dim == 0 or (inp.shape[-1] * inp.element_size()) % PACK_BYTES == 0 + + def all_gather( + self, + inp: torch.Tensor, + *, + dim: int = -1, + out: Optional[torch.Tensor] = None, + stream: object = None, + ) -> torch.Tensor: + """Concatenate every rank's ``inp`` along ``dim`` (0 or the last dim). + + With 16-byte-aligned rows the kernel writes the concatenated layout + directly (no reshape or copy afterwards). Otherwise the shards are + gathered contiguously with 16-byte padding and finished with a torch + reshape, which still keeps the collective on RDMA. + """ + + with self._lock: + self.check_health() + if not self.should_all_gather(inp, dim): + raise ValueError("input is not eligible for the RoCE all-gather") + dim = self._normalize_dim(inp, dim) + shape = list(inp.shape) + shape[dim] *= self.world_size + if out is not None and ( + list(out.shape) != shape + or out.dtype != inp.dtype + or out.device != inp.device + or not out.is_contiguous() + ): + raise ValueError("out must be a contiguous tensor on the input's device of the gathered shape") + context = torch.cuda.stream(stream) if stream is not None else _nullcontext() + with torch.cuda.device(self.device), context: + capturing = torch.cuda.is_current_stream_capturing() + # Both paths stage the shard contiguously and only the reader is + # strided, so ranks may take different paths for one collective. + direct = self._direct_gather_layout(inp, dim) and ( + out is None or out.data_ptr() % PACK_BYTES == 0 + ) + if direct: + if out is None: + out = torch.empty(shape, dtype=inp.dtype, device=inp.device) + nbytes = inp.numel() * inp.element_size() + row_packs = ( + nbytes // PACK_BYTES + if dim == 0 + else (inp.shape[-1] * inp.element_size()) // PACK_BYTES + ) + self._order_stream(capturing) + self._launch_gather(inp.data_ptr(), out.data_ptr(), nbytes, row_packs) + self._mark_stream(capturing) + return out + # Padded path: pad each shard to a whole number of packs, gather + # contiguously into fixed scratch, then let torch produce the + # requested layout. + nbytes = inp.numel() * inp.element_size() + padded = _align_up(nbytes, PACK_BYTES) + staged, gathered = self._gather_scratch(padded) + self._order_stream(capturing) + staged[:nbytes].copy_(inp.reshape(-1).view(torch.uint8)) + self._launch_gather( + staged.data_ptr(), gathered.data_ptr(), padded, padded // PACK_BYTES + ) + stacked = ( + gathered.view(self.world_size, padded)[:, :nbytes] + .reshape(-1) + .view(inp.dtype) + .reshape(self.world_size, *inp.shape) + ) + result = stacked.movedim(0, dim).reshape(shape) + if out is None: + result = result.contiguous() + else: + out.copy_(result) + result = out + self._mark_stream(capturing) + return result + + def _gather_scratch(self, padded: int) -> tuple[torch.Tensor, torch.Tensor]: + """Fixed device scratch for the padded all-gather path, allocated once. + + ``prepare(padded_gather=True)`` or the first eager padded gather + allocates it; a first use under CUDA graph capture is refused so a + captured graph never owns the scratch. + """ + + if self._gather_buffers is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "RoCE all-gather scratch must be allocated before CUDA graph capture; " + "call prepare(padded_gather=True) first" + ) + capacity = _align_up(self.max_gather_bytes, PACK_BYTES) + self._gather_buffers = ( + torch.empty(capacity, dtype=torch.uint8, device=self.device), + torch.empty(self.world_size * capacity, dtype=torch.uint8, device=self.device), + ) + staged, gathered = self._gather_buffers + return staged[:padded], gathered[: self.world_size * padded] + + def _launch_gather(self, input_address: int, output_address: int, nbytes: int, row_packs: int) -> None: + """Launch the all-gather kernel for ``nbytes`` per rank with ``row_packs`` 16-byte packs per row.""" + key = self._gather_launcher_key() + capturing = torch.cuda.is_current_stream_capturing() + if capturing and not _allgather_cute.is_launcher_prepared(*key): + raise RuntimeError( + "RoCE all-gather launcher must be prepared before CUDA graph capture" + ) + launcher = _allgather_cute.get_launcher(*key) + launcher( + input_address, + output_address, + nbytes // PACK_BYTES, + nbytes, + row_packs, + self._recv_base, + self._flag_base, + self._send_base, + self._ctrl_base, + self._slot_bytes, + self._epoch_address, + self.spin_limit, + self._blocks, + ) + if not capturing: + self.check_health() + + @contextmanager + def capture(self, stream: object = None, *, channel_id: Optional[str] = None): + """Context manager for CUDA graph capture; the launchers must already be prepared. + + The cross-stream ordering event is dropped for the capture: an event + recorded outside a capture must not be waited on inside one, and + ``torch.cuda.graph`` has already synchronized the device before the + capture began. Inside, every collective must use one stream. + """ + + self._last_stream = None + # CUDA's capture ID owns stream admission, not this Python context. + # Nested or per-call contexts must retain the same-capture guard. + try: + yield self + finally: + self._last_stream = None + + # -- diagnostics / lifecycle -------------------------------------------------- + + def stats(self) -> dict[str, Any]: + """Runtime, control-record, and proxy counters for diagnostics.""" + info: dict[str, Any] = { + "world_size": self.world_size, + "rank": self.rank, + "hcas": list(self.hca_names), + "max_size": self.max_size, + "max_gather_bytes": self.max_gather_bytes, + "slot_bytes": self._slot_bytes, + "epoch": int(self._counters[0].item()), + "error_seq": int(self._error_word.item()), + "error_peer": int(self._ctrl_words[3].item()), + "ctrl_seq": int(self._ctrl_words[0].item()), + "spin_limit": self.spin_limit, + "opposite_paths": self._opposite_paths, + } + if self._proxy is not None: + info.update(self._proxy.stats()) + info["peer_hca"] = { + peer: self._proxy.peer_hca(peer) + for peer in range(self.world_size) + if peer != self.rank + } + return info + + def benchmark_counters(self) -> list[dict[str, int | str | None]]: + """Return absolute per-origin-QP counters since runtime construction.""" + + if self._proxy is None: + return [] + return self._proxy.path_counters() + + def close(self) -> None: + """Stop the proxy and release the transport. + + Waits for the device first so no in-flight kernel still reads the + pinned region or rings the doorbell after the proxy is gone. Peers + that write to this rank afterwards see a remote-access completion + error and raise on their side, which is the intended shutdown signal. + """ + + with self._lock: + if self._closed: + return + self._closed = True + with contextlib.suppress(Exception): + torch.cuda.synchronize(self.device) + if self._proxy is not None: + self._proxy.close() + self._proxy = None + + def __del__(self) -> None: # pragma: no cover - defensive teardown + """Release the transport if ``close`` was never called.""" + with contextlib.suppress(Exception): + self.close() + + +class _nullcontext: + def __enter__(self): + """No-op context entry.""" + return None + + def __exit__(self, *exc): + """No-op context exit.""" + return False + + +__all__ = [ + "DEFAULT_MAX_GATHER_BYTES", + "DEFAULT_MAX_SIZE", + "SUPPORTED_DTYPES", + "SUPPORTED_WORLD_SIZES", + "RoceOneshotAllReduce", + "default_gid_index", + "discover_hcas", + "is_supported", +] diff --git a/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/rocenante-overlay-config.json b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/rocenante-overlay-config.json new file mode 100644 index 00000000..f6838f8d --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/rocenante-overlay-config.json @@ -0,0 +1,220 @@ +{ + "artifacts": { + "b12x_git": { + "commit": "eac260a8257cc6b14e7d4ad674f51e9a09b8790f", + "roce_source_dirty": true, + "roce_status": [ + " M b12x/comm/roce/_oneshot_cute.py", + " M b12x/comm/roce/roce_oneshot.py" + ] + }, + "b12x_roce_tree_sha256": "8a64063506c9c96d3fd3b004c3ab51dda88408f7246f5c89d0a0b6713065fca6", + "base_sircl_manifest": { + "files": [ + { + "path": "sitecustomize.py", + "sha256": "4efc3fdfe4874820064dbcb713c02271df4ca41f4f2bbcf4683a1548ca02f2b1", + "source": "spark_transport/integrations/vllm/sitecustomize.py" + }, + { + "path": "spark_collective_audit.py", + "sha256": "8fa7aa9747845a84b40e8d3a0d72c39f8a562d6490764b777e5dcbe1cbe47721", + "source": "spark_transport/integrations/vllm/spark_collective_audit.py" + }, + { + "path": "spark_cudagraph_bucket_contract.py", + "sha256": "e7411ff9cde9d5ed4a7ecd6e78a3158512c188c254baad10554b72229f14d272", + "source": "spark_transport/integrations/vllm/spark_cudagraph_bucket_contract.py" + }, + { + "path": "spark_cudagraph_replay_timing.py", + "sha256": "dc3ee0fd5da52e79b097e9edc6191a880ca5a3f4e6d83add2a4acd30966b2175", + "source": "spark_transport/integrations/vllm/spark_cudagraph_replay_timing.py" + }, + { + "path": "spark_dcp_collective_audit.py", + "sha256": "077a234e4edff8b8dd44784953aef713884b4dd7a3f7c46589b14c6bb8b40745", + "source": "spark_transport/integrations/vllm/spark_dcp_collective_audit.py" + }, + { + "path": "spark_graph_status_reporter.py", + "sha256": "4bd50d1be42e2ac750b382825bec8a49564ce288d376852ee7c51e33a303f277", + "source": "spark_transport/integrations/vllm/spark_graph_status_reporter.py" + }, + { + "path": "spark_persistent_output_ring.py", + "sha256": "03220c78ed87651f93c3165f1eedb843ffcad6a116126e2a9b89beab299bd2f3", + "source": "spark_transport/integrations/vllm/spark_persistent_output_ring.py" + }, + { + "path": "spark_tp4_backend.py", + "sha256": "5999814740ddfa1c35e925ed0f3bf43edd8335b30c9b4fe9b78006c618d18cb9", + "source": "spark_transport/integrations/vllm/spark_tp4_backend.py" + }, + { + "path": "spark_tp4_capability.py", + "sha256": "d90f6fc07ff5255b2df25b3ffc1c52c83f9eb866ca68666bc0f273db1c5e83b7", + "source": "spark_transport/integrations/vllm/spark_tp4_capability.py" + }, + { + "path": "spark_tp4_health_gate.py", + "sha256": "7651b88e9bff40193891a0810ee8a920e3548a17de3a0dc477a30db779738930", + "source": "spark_transport/integrations/vllm/spark_tp4_health_gate.py" + }, + { + "path": "spark_tp4_port_namespace.py", + "sha256": "c789739811def9722753bb128e69e7bd5342bab827b2cfd99cd8a4e847d18d4f", + "source": "spark_transport/integrations/vllm/spark_tp4_port_namespace.py" + }, + { + "path": "spark_tp4_query_contract.py", + "sha256": "978df868fd83a6a780111092e99de2a3f07a5df375cf52d4b2358ccf4dcd0c31", + "source": "spark_transport/integrations/vllm/spark_tp4_query_contract.py" + }, + { + "path": "spark_tp4_query_row_provider.py", + "sha256": "09c9e7af5d14b5c6930e15c999ada2231d639eb59609d705aadecdd7497ddec7", + "source": "spark_transport/integrations/vllm/spark_tp4_query_row_provider.py" + }, + { + "path": "spark_tp4_vocab_allgather_backend.py", + "sha256": "fd44dea558c115b17f17bb129033345af4543702fbbdaf1fa4f186b8f242e697", + "source": "spark_transport/integrations/vllm/spark_tp4_vocab_allgather_backend.py" + } + ], + "schema": "sparkring-public-overlay/v1" + }, + "base_sircl_manifest_name": "sparkring-overlay-manifest.json", + "base_sircl_manifest_sha256": "85a231e6d2a290f7d6cccbc2cc6b1ccad7a6adbefc7ce4dde05b158f249aadd4" + }, + "canonical_hca_order": [ + "rocep1s0f0", + "rocep1s0f1", + "roceP2p1s0f0", + "roceP2p1s0f1" + ], + "dispatch": { + "candidate": { + "contiguous": true, + "cuda": true, + "dtype": "bfloat16", + "group_name_prefix": "tp", + "maximum_query_rows": 32, + "minimum_query_rows": 1, + "tensor_dimensions": 2, + "width": 4096, + "world_size": 4 + }, + "fallback": "Call the already-installed SIRCL wrapper, which retains its own graph, fused eager, and NCCL routing rules.", + "known_sircl_graph_query_rows": [ + 64, + 128 + ] + }, + "excluded_performance_arms": { + "four_opposite_paths_mesh32": { + "disposition": "rejected", + "status": "research-only" + }, + "hardware_qos": { + "disposition": "rejected", + "status": "research-only" + } + }, + "hca_index_semantics": { + "0": "physical port f0 through PCIe-domain-0", + "1": "physical port f1 through PCIe-domain-0", + "2": "physical port f0 through PCIe-domain-2", + "3": "physical port f1 through PCIe-domain-2" + }, + "health": { + "check_boundary": "after vLLM's synchronized synchronous output or asynchronous get_output", + "post_publication_fallback": false, + "worker_exit_code": 70 + }, + "metadata": { + "create_device_process_group": false, + "create_nccl_communicator": false, + "process_group": "the TP GroupCoordinator cpu_group", + "rank_wide_vote": true, + "required_backend": "gloo" + }, + "peer_hca_maps": { + "0": "1=0/2,2=0/3,3=1/3", + "1": "0=1/3,2=0/2,3=0/3", + "2": "0=1/2,1=1/3,3=0/2", + "3": "0=0/2,1=1/2,2=1/3" + }, + "purpose": "Route eligible GLM-5.3-Flash TP4 all-reduces through the six-QP B12X virtual-diagonal runtime while preserving the mounted SIRCL and NCCL dispatch chain for every rejected signature.", + "runtime": { + "all_gather_enabled": false, + "captured_sircl_query_rows": [ + 16, + 20, + 24, + 28, + 32 + ], + "decode_context_parallel_size": 4, + "direct_then_diagonal_threshold_bytes": 196608, + "dtype": "bfloat16", + "execution_mode": "both", + "forbidden_proxy_cpus": [ + 10, + 11 + ], + "gid_index": 3, + "maximum_candidate_payload_bytes": 262144, + "maximum_candidate_query_rows": 32, + "opposite_rank_paths": 2, + "origin_queue_pairs_per_rank": 6, + "origin_queue_pairs_total": 24, + "proxy_cpu": 13, + "target_shape": "[Q,4096]", + "tensor_parallel_size": 4, + "transport_slots": 2, + "wave_mode": "two", + "world_size": 4 + }, + "schema": "sparkring.glm53-rocenante-overlay-contract/v1", + "sidecars": { + "automatic_cleanup_on_apply_failure": true, + "automatic_cleanup_on_exit": true, + "cleanup_order": [ + "source_markers", + "intermediate_rules", + "intermediate_qdiscs", + "endpoint_neighbors", + "endpoint_routes", + "mtu_restore", + "mtu_verify" + ], + "intermediate_rules": 8, + "qualified_topology_gate_required": true, + "source_marker_ether_type": "0x88b5", + "source_marker_match_udp_port": 65535, + "source_marker_runtime_seconds": 7200, + "source_marker_source": "spark_transport/experiments/cx7_hairpin_diagonal/native/mlx5_rdma_tx_rewrite_probe.c" + }, + "source_references": { + "b12x_rocenante_pull_request": { + "base_commit": "ffb7442d04a9f50b950df1fb17280acad881b7d5", + "url": "https://github.com/local-inference-lab/b12x/pull/295" + }, + "published_glm_runtime": { + "b12x_commit": "9ae41c5c", + "runtime_path": "runtime/glm53-flash-jj-r8-gb10", + "vllm_commit": "e02b174693e13859de61811b5e8cd13d5308e259", + "vllm_python_blobs": { + "vllm/distributed/device_communicators/cuda_communicator.py": "260dda5d899a319c767970c25a53c6ba2e2c0715", + "vllm/distributed/parallel_state.py": "9353dbfe36097fc12348cb055127e0caf798e2ac", + "vllm/v1/worker/gpu_worker.py": "1f6336478c3456138e37fd007917a25afab22b49" + } + }, + "vllm_rocenante_pull_request": { + "commit": "a7935eb1d8fa51400cd13452b1d988197deef3b8", + "url": "https://github.com/local-inference-lab/vllm/pull/597" + } + }, + "status": "research-only" +} diff --git a/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/rocenante_health_gate.py b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/rocenante_health_gate.py new file mode 100644 index 00000000..f2cc95e7 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/rocenante_health_gate.py @@ -0,0 +1,98 @@ +"""Check B12X virtual-diagonal health after vLLM output synchronization. + +Status: research-only. These wrappers have a distinct marker from the SIRCL +health gate, so both checks remain installed and execute in wrapper order. +""" + +from __future__ import annotations + +import importlib +import os + +from rocenante_vllm_overlay import require_health + + +_installed = False + + +def _checked(callable_): + try: + return callable_() + except BaseException: + os._exit(70) + raise RuntimeError("virtual-diagonal health exit unexpectedly returned") + + +def _wrap_get_output(output_type: type) -> bool: + original = output_type.get_output + if getattr(original, "_rocenante_health_gate", False): + return False + + def get_output(self): + output = original(self) + _checked(require_health) + return output + + get_output._rocenante_health_gate = True + output_type.get_output = get_output + return True + + +def _wrap_worker_output(worker_type: type, method_name: str) -> None: + original = getattr(worker_type, method_name) + if getattr(original, "_rocenante_health_gate", False): + return + + def worker_output(self, *args, **kwargs): + output = original(self, *args, **kwargs) + if hasattr(output, "get_output"): + return output + _checked(require_health) + return output + + worker_output._rocenante_health_gate = True + setattr(worker_type, method_name, worker_output) + + +def install() -> None: + """Install B12X checks outside any SIRCL output wrappers already present.""" + + global _installed + if _installed: + return + wrapped = 0 + for module_name, class_names in ( + ("vllm.v1.worker.gpu.async_utils", ("AsyncOutput", "AsyncPoolingOutput")), + ( + "vllm.v1.worker.gpu_model_runner", + ("AsyncGPUModelRunnerOutput", "AsyncGPUPoolingModelRunnerOutput"), + ), + ): + try: + module = importlib.import_module(module_name) + except ImportError: + continue + for class_name in class_names: + output_type = getattr(module, class_name, None) + if output_type is not None: + wrapped += int(_wrap_get_output(output_type)) + try: + worker_type = importlib.import_module("vllm.v1.worker.gpu_worker").Worker + except (AttributeError, ImportError) as error: + raise RuntimeError( + "virtual-diagonal health gate found no pinned vLLM GPU Worker" + ) from error + for method_name in ("execute_model", "sample_tokens"): + if not hasattr(worker_type, method_name): + raise RuntimeError( + f"virtual-diagonal health gate found no Worker.{method_name} boundary" + ) + _wrap_worker_output(worker_type, method_name) + if wrapped == 0: + raise RuntimeError( + "virtual-diagonal health gate found no asynchronous output type" + ) + _installed = True + + +__all__ = ["install"] diff --git a/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/rocenante_vllm_overlay.py b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/rocenante_vllm_overlay.py new file mode 100644 index 00000000..068209d2 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/rocenante_vllm_overlay.py @@ -0,0 +1,599 @@ +"""Route bounded GLM-5.3 TP4 shapes through B12X virtual diagonals. + +Status: research-only. The adapter wraps an already-installed vLLM all-reduce +chain. Eligible Q1-Q32 BF16 ``[Q,4096]`` tensors use one six-origin-QP B12X +runtime per rank. Every rejected signature calls the saved SIRCL/NCCL chain. +""" + +from __future__ import annotations + +import contextlib +import hashlib +import importlib +import json +import os +import sys +import threading +import weakref +from collections.abc import Mapping +from pathlib import Path +from typing import Any + + +CONFIG_SCHEMA = "sparkring.glm53-rocenante-overlay-contract/v1" +DEFAULT_CONFIG = Path("/opt/spark-sircl/rocenante-overlay-config.json") +_installed = False +_adapters: weakref.WeakSet[VirtualDiagonalAdapter] = weakref.WeakSet() +_registry_lock = threading.Lock() +_b12x_path_installed = False + + +class OverlayError(RuntimeError): + """The private overlay cannot prove a rank-invariant safe configuration.""" + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as source: + for block in iter(lambda: source.read(1 << 20), b""): + digest.update(block) + return digest.hexdigest() + + +def _canonical_sha256(value: object) -> str: + encoded = json.dumps( + value, sort_keys=True, separators=(",", ":"), ensure_ascii=True + ).encode() + return hashlib.sha256(encoded).hexdigest() + + +def _mapping(value: object, name: str) -> Mapping[str, object]: + if not isinstance(value, Mapping): + raise OverlayError(f"{name} must be an object") + return value + + +def _integer(value: object, name: str, minimum: int, maximum: int) -> int: + if isinstance(value, bool) or not isinstance(value, int): + raise OverlayError(f"{name} must be an integer") + if not minimum <= value <= maximum: + raise OverlayError(f"{name} must be in [{minimum}, {maximum}]") + return value + + +def load_contract(path: Path = DEFAULT_CONFIG) -> dict[str, object]: + """Load and strictly validate the mounted overlay contract.""" + + try: + value = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + raise OverlayError(f"overlay contract cannot be read: {error}") from error + root = dict(_mapping(value, "overlay contract")) + if root.get("schema") != CONFIG_SCHEMA or root.get("status") != "research-only": + raise OverlayError(f"overlay contract must use {CONFIG_SCHEMA}") + runtime = _mapping(root.get("runtime"), "runtime") + if _integer(runtime.get("world_size"), "runtime.world_size", 1, 64) != 4: + raise OverlayError("runtime.world_size must be four") + if ( + _integer( + runtime.get("opposite_rank_paths"), "runtime.opposite_rank_paths", 2, 4 + ) + != 2 + ): + raise OverlayError("runtime.opposite_rank_paths must be two") + if ( + _integer( + runtime.get("origin_queue_pairs_per_rank"), + "runtime.origin_queue_pairs_per_rank", + 1, + 64, + ) + != 6 + ): + raise OverlayError("runtime.origin_queue_pairs_per_rank must be six") + if ( + _integer( + runtime.get("direct_then_diagonal_threshold_bytes"), + "runtime.direct_then_diagonal_threshold_bytes", + 16, + 1 << 30, + ) + != 196608 + ): + raise OverlayError( + "runtime.direct_then_diagonal_threshold_bytes must be 196608" + ) + if runtime.get("wave_mode") != "two": + raise OverlayError("runtime.wave_mode must be 'two'") + if runtime.get("execution_mode", "both") not in { + "both", "eager_only", "graph_only", "disabled" + }: + raise OverlayError( + "runtime.execution_mode must be both, eager_only, graph_only, or disabled" + ) + captured_sircl_rows = runtime.get("captured_sircl_query_rows", []) + if ( + not isinstance(captured_sircl_rows, list) + or any(type(q) is not int or not 1 <= q <= 32 for q in captured_sircl_rows) + or captured_sircl_rows != sorted(set(captured_sircl_rows)) + ): + raise OverlayError( + "runtime.captured_sircl_query_rows must be sorted unique integers in [1,32]" + ) + proxy_cpu = _integer(runtime.get("proxy_cpu"), "runtime.proxy_cpu", 0, 4095) + forbidden = runtime.get("forbidden_proxy_cpus") + if forbidden != [10, 11] or proxy_cpu in forbidden: + raise OverlayError("runtime.proxy_cpu must differ from SIRCL CPUs 10 and 11") + hcas = root.get("canonical_hca_order") + expected_hcas = [ + "rocep1s0f0", + "rocep1s0f1", + "roceP2p1s0f0", + "roceP2p1s0f1", + ] + if hcas != expected_hcas: + raise OverlayError("canonical_hca_order differs from the measured HCA order") + maps = _mapping(root.get("peer_hca_maps"), "peer_hca_maps") + expected_maps = { + "0": "1=0/2,2=0/3,3=1/3", + "1": "0=1/3,2=0/2,3=0/3", + "2": "0=1/2,1=1/3,3=0/2", + "3": "0=0/2,1=1/2,2=1/3", + } + if dict(maps) != expected_maps: + raise OverlayError("peer_hca_maps differs from the reciprocal six-QP map") + dispatch = _mapping(root.get("dispatch"), "dispatch") + candidate = _mapping(dispatch.get("candidate"), "dispatch.candidate") + expected_candidate = { + "group_name_prefix": "tp", + "world_size": 4, + "cuda": True, + "contiguous": True, + "dtype": "bfloat16", + "tensor_dimensions": 2, + "width": 4096, + "minimum_query_rows": 1, + "maximum_query_rows": 32, + } + if dict(candidate) != expected_candidate: + raise OverlayError("dispatch.candidate differs from the Q1-Q32 contract") + metadata = _mapping(root.get("metadata"), "metadata") + if ( + metadata.get("required_backend") != "gloo" + or metadata.get("create_device_process_group") is not False + or metadata.get("create_nccl_communicator") is not False + ): + raise OverlayError("metadata must use Gloo without a device communicator") + return root + + +def _gid_available(device: str, gid_index: int) -> bool: + root = Path("/sys/class/infiniband") / device / "ports" + if not root.is_dir(): + return False + for port in root.iterdir(): + path = port / "gids" / str(gid_index) + if not path.is_file(): + continue + value = path.read_text(encoding="ascii").strip().replace(":", "") + if value and any(character != "0" for character in value): + return True + return False + + +def _set_exact_environment(name: str, value: str) -> None: + present = os.environ.get(name) + if present is not None and present != value: + raise OverlayError(f"{name} conflicts with the mounted overlay contract") + os.environ[name] = value + + +def _install_b12x_roce_path() -> None: + """Prepend only the private ``b12x.comm.roce`` package to installed B12X.""" + + global _b12x_path_installed + if _b12x_path_installed: + return + if any( + name == "b12x.comm.roce" or name.startswith("b12x.comm.roce.") + for name in sys.modules + ): + raise OverlayError("b12x.comm.roce was imported before the private source path") + import b12x.comm + + root = Path("/opt/spark-sircl/b12x_overlay/b12x/comm") + if not (root / "roce" / "__init__.py").is_file(): + raise OverlayError(f"private b12x.comm.roce package is missing: {root}") + paths = list(b12x.comm.__path__) + if str(root) not in paths: + b12x.comm.__path__.insert(0, str(root)) + importlib.invalidate_caches() + _b12x_path_installed = True + + +@contextlib.contextmanager +def _inherited_thread_affinity(cpu: int): + """Create the B12X proxy with a one-CPU inherited pthread mask.""" + + if not hasattr(os, "sched_getaffinity") or not hasattr(os, "sched_setaffinity"): + raise OverlayError("proxy CPU isolation requires Linux scheduler affinity") + original = os.sched_getaffinity(0) + if cpu not in original: + raise OverlayError(f"proxy CPU {cpu} is outside the worker affinity mask") + os.sched_setaffinity(0, {cpu}) + try: + yield + finally: + os.sched_setaffinity(0, original) + + +class VirtualDiagonalAdapter: + """One TP rank's six-QP B12X runtime and exact Q1-Q32 policy.""" + + def __init__(self, communicator: Any, contract: Mapping[str, object]) -> None: + import torch + import torch.distributed as dist + + self.communicator = communicator + self.group = communicator.cpu_group + self.device = communicator.device + self.rank = int(communicator.rank) + self.world_size = int(communicator.world_size) + self._runtime = None + self._closed = False + self._candidate_calls = 0 + self._captured_nodes = 0 + self._fallback_calls = 0 + runtime = _mapping(contract["runtime"], "runtime") + self.execution_mode = str(runtime.get("execution_mode", "both")) + self.captured_sircl_rows = frozenset(runtime.get("captured_sircl_query_rows", [])) + dispatch = _mapping(contract["dispatch"], "dispatch") + candidate = _mapping(dispatch["candidate"], "dispatch.candidate") + self.minimum_query_rows = int(candidate["minimum_query_rows"]) + self.maximum_query_rows = int(candidate["maximum_query_rows"]) + self.width = int(candidate["width"]) + self.proxy_cpu = int(runtime["proxy_cpu"]) + self.gid_index = int(runtime["gid_index"]) + self.hcas = tuple(contract["canonical_hca_order"]) + self.peer_map = str(contract["peer_hca_maps"][str(self.rank)]) + + errors: list[str] = [] + if self.world_size != 4: + errors.append(f"TP world size is {self.world_size}, expected 4") + if not str(communicator.unique_name).startswith( + str(candidate["group_name_prefix"]) + ): + errors.append(f"communicator {communicator.unique_name!r} is not TP") + try: + backend = str(dist.get_backend(self.group)).lower() + except Exception as error: # noqa: BLE001 - included in unanimous vote + backend = f"error:{error}" + if backend != "gloo": + errors.append( + f"metadata process group backend is {backend!r}, expected gloo" + ) + if self.device.type != "cuda": + errors.append(f"communicator device is {self.device}, expected CUDA") + if self.proxy_cpu in (10, 11): + errors.append("proxy CPU overlaps a SIRCL graph CPU") + if not hasattr(os, "sched_getaffinity") or not hasattr(os, "sched_setaffinity"): + errors.append("Linux scheduler affinity is unavailable") + else: + try: + if self.proxy_cpu not in os.sched_getaffinity(0): + errors.append( + f"proxy CPU {self.proxy_cpu} is outside the worker affinity mask" + ) + except OSError as error: + errors.append(f"worker affinity cannot be read: {error}") + for device in self.hcas: + try: + available = _gid_available(device, self.gid_index) + except OSError as error: + available = False + errors.append( + f"RDMA device/GID probe failed: {device}:{self.gid_index}: {error}" + ) + if not available: + errors.append( + f"RDMA device/GID is unavailable: {device}:{self.gid_index}" + ) + try: + _install_b12x_roce_path() + from b12x.comm import roce + + if getattr(roce, "API_VERSION", None) != 1: + errors.append( + f"b12x.comm.roce API version is {getattr(roce, 'API_VERSION', None)!r}, expected 1" + ) + except Exception as error: + errors.append(f"b12x.comm.roce cannot import: {error}") + + expected_environment = { + "B12X_ROCE_HCA": ",".join(self.hcas), + "B12X_ROCE_GID_INDEX": str(self.gid_index), + "B12X_ROCE_OPPOSITE_PATHS": "2", + "B12X_ROCE_PEER_HCA_MAP": self.peer_map, + "B12X_ROCE_WAVE_MODE": "two", + "B12X_ROCE_TWO_WAVE_THRESHOLD_BYTES": "196608", + } + for name, expected in expected_environment.items(): + present = os.environ.get(name) + if present is not None and present != expected: + errors.append(f"{name} conflicts with the mounted overlay contract") + + config_digest = _canonical_sha256(contract) + manifest_path = Path("/opt/spark-sircl/sparkring-overlay-manifest.json") + try: + manifest_digest = _sha256(manifest_path) if manifest_path.is_file() else "" + except OSError as error: + manifest_digest = "" + errors.append(f"private bundle manifest cannot be read: {error}") + local_vote = { + "rank": self.rank, + "errors": errors, + "config_sha256": config_digest, + "manifest_sha256": manifest_digest, + "hcas": list(self.hcas), + "peer_map": self.peer_map, + "proxy_cpu": self.proxy_cpu, + } + votes: list[object] = [None] * self.world_size + dist.all_gather_object(votes, local_vote, group=self.group) + self._validate_votes(votes, contract) + + if self.execution_mode == "disabled": + return + + for name, expected in expected_environment.items(): + _set_exact_environment(name, expected) + + from b12x.comm import roce + + max_bytes = self.maximum_query_rows * self.width * 2 + with _inherited_thread_affinity(self.proxy_cpu): + self._runtime = roce.AllReduce.from_exchange_group( + exchange_group=self.group, + device=self.device, + max_size=max_bytes, + max_gather_bytes=16, + ) + self._runtime.prepare((torch.bfloat16,)) + stats = self._runtime.stats() + if ( + stats.get("hcas") != list(self.hcas) + or stats.get("wave_mode") != "two" + or stats.get("two_wave_threshold_bytes") != 196608 + ): + self.close() + raise OverlayError("B12X runtime attestation differs from the contract") + counters = self._runtime.benchmark_counters() + if len(counters) != 6: + self.close() + raise OverlayError("B12X runtime must expose six active origin-QP counters") + + def _validate_votes( + self, votes: list[object], contract: Mapping[str, object] + ) -> None: + if len(votes) != 4 or any(not isinstance(vote, Mapping) for vote in votes): + raise OverlayError("Gloo capability vote must contain four rank records") + by_rank = {vote.get("rank"): vote for vote in votes} + if set(by_rank) != {0, 1, 2, 3}: + raise OverlayError("Gloo capability vote must contain ranks 0, 1, 2, and 3") + digests = {vote.get("config_sha256") for vote in votes} + manifests = {vote.get("manifest_sha256") for vote in votes} + if len(digests) != 1 or len(manifests) != 1 or "" in manifests: + raise OverlayError("overlay source or configuration differs across ranks") + failures = [ + f"rank {rank}: {error}" + for rank, vote in sorted(by_rank.items()) + for error in vote.get("errors", []) + ] + if failures: + raise OverlayError("; ".join(failures)) + expected_maps = _mapping(contract["peer_hca_maps"], "peer_hca_maps") + for rank, vote in sorted(by_rank.items()): + if vote.get("peer_map") != expected_maps[str(rank)]: + raise OverlayError(f"rank {rank} peer map differs from the contract") + + def eligible(self, tensor: Any) -> bool: + """Return a rank-invariant decision for one TP all-reduce signature.""" + + import torch + + shape = tuple(int(value) for value in tensor.shape) + valid_signature = ( + not self._closed + and tensor.is_cuda + and tensor.device == self.device + and tensor.dtype == torch.bfloat16 + and tensor.is_contiguous() + and len(shape) == 2 + and shape[1] == self.width + and self.minimum_query_rows <= shape[0] <= self.maximum_query_rows + ) + if not valid_signature: + return False + if self.execution_mode == "disabled": + return False + if shape[0] in self.captured_sircl_rows: + if bool(torch.cuda.is_current_stream_capturing()): + return False + if self.execution_mode == "both": + return True + capturing = bool(torch.cuda.is_current_stream_capturing()) + return capturing if self.execution_mode == "graph_only" else not capturing + + def all_reduce(self, tensor: Any) -> Any: + """Run one accepted B12X collective or terminate on transport failure.""" + + import torch + + if self._runtime is None: + raise OverlayError("B12X runtime is unavailable") + capturing = bool(torch.cuda.is_current_stream_capturing()) + try: + if capturing: + stream = torch.cuda.current_stream(self.device) + result = self._runtime.all_reduce(tensor, stream=stream) + self._captured_nodes += 1 + else: + result = self._runtime.all_reduce(tensor) + self._candidate_calls += 1 + return result + except BaseException: + os._exit(70) + raise AssertionError("worker exit unexpectedly returned") + + def record_fallback(self) -> None: + self._fallback_calls += 1 + + def check_health(self) -> None: + if not self._closed and self._runtime is not None: + self._runtime.check_health() + + def diagnostic_snapshot(self) -> dict[str, object]: + return { + "rank": self.rank, + "execution_mode": self.execution_mode, + "candidate_calls": self._candidate_calls, + "captured_nodes": self._captured_nodes, + "fallback_calls": self._fallback_calls, + "runtime": None if self._runtime is None else self._runtime.stats(), + "origin_qps": [] + if self._runtime is None + else self._runtime.benchmark_counters(), + } + + def close(self) -> None: + if self._closed: + return + self._closed = True + runtime, self._runtime = self._runtime, None + if runtime is not None: + runtime.close() + + +def require_health() -> None: + """Raise when any process-local virtual-diagonal runtime is unhealthy.""" + + with _registry_lock: + adapters = list(_adapters) + for adapter in adapters: + adapter.check_health() + + +def diagnostic_snapshot() -> list[dict[str, object]]: + with _registry_lock: + adapters = list(_adapters) + return [adapter.diagnostic_snapshot() for adapter in adapters] + + +def host_route_snapshot() -> list[dict[str, object]]: + """Read Python routing counters without CUDA or native transport calls. + + Captured-node counts describe graph construction, not replay executions. + """ + with _registry_lock: + return [ + { + "rank": adapter.rank, + "execution_mode": adapter.execution_mode, + "captured_sircl_query_rows": sorted(adapter.captured_sircl_rows), + "eager_calls": adapter._candidate_calls - adapter._captured_nodes, + "captured_nodes": adapter._captured_nodes, + "fallback_calls": adapter._fallback_calls, + } + for adapter in sorted(_adapters, key=lambda item: item.rank) + ] + + +def _install_status_reporting() -> None: + """Append CPU-only routing counters to the existing low-rate status file.""" + import spark_graph_status_reporter as reporter + + original = reporter.collect_graph_status + if getattr(original, "_rocenante_route_status", False): + return + + def collect(): + value = original() + value["rocenante_routing"] = host_route_snapshot() + return value + + collect._rocenante_route_status = True + reporter.collect_graph_status = collect + + +def install(config_path: Path = DEFAULT_CONFIG) -> None: + """Wrap vLLM after SIRCL so rejected signatures retain its dispatch chain.""" + + global _installed + if _installed: + return + contract = load_contract(config_path) + from vllm.distributed.device_communicators.cuda_communicator import ( + CudaCommunicator, + ) + + original_init = CudaCommunicator.__init__ + original_all_reduce = CudaCommunicator.all_reduce + original_destroy = CudaCommunicator.destroy + if getattr(original_all_reduce, "_rocenante_virtual_diagonal", False): + _installed = True + return + + def wrapped_init(self, *args, **kwargs): + original_init(self, *args, **kwargs) + prefix = str(contract["dispatch"]["candidate"]["group_name_prefix"]) + if str(self.unique_name).startswith(prefix): + if int(self.world_size) != 4: + raise OverlayError( + f"virtual-diagonal TP group requires four ranks, got {self.world_size}" + ) + with _registry_lock: + if list(_adapters): + raise OverlayError( + "one process cannot construct more than one virtual-diagonal TP runtime" + ) + adapter = VirtualDiagonalAdapter(self, contract) + self._rocenante_virtual_diagonal_adapter = adapter + with _registry_lock: + _adapters.add(adapter) + + def wrapped_all_reduce(self, tensor): + adapter = getattr(self, "_rocenante_virtual_diagonal_adapter", None) + if adapter is not None and adapter.eligible(tensor): + return adapter.all_reduce(tensor) + if adapter is not None: + adapter.record_fallback() + return original_all_reduce(self, tensor) + + def wrapped_destroy(self): + adapter = getattr(self, "_rocenante_virtual_diagonal_adapter", None) + if adapter is not None: + adapter.close() + with _registry_lock: + _adapters.discard(adapter) + self._rocenante_virtual_diagonal_adapter = None + return original_destroy(self) + + wrapped_all_reduce._rocenante_virtual_diagonal = True + wrapped_all_reduce._rocenante_saved_all_reduce = original_all_reduce + CudaCommunicator.__init__ = wrapped_init + CudaCommunicator.all_reduce = wrapped_all_reduce + CudaCommunicator.destroy = wrapped_destroy + _install_status_reporting() + _installed = True + + +__all__ = [ + "CONFIG_SCHEMA", + "DEFAULT_CONFIG", + "OverlayError", + "VirtualDiagonalAdapter", + "diagnostic_snapshot", + "install", + "load_contract", + "require_health", +] diff --git a/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/sircl_sitecustomize.py b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/sircl_sitecustomize.py new file mode 100644 index 00000000..43c3b16e --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/sircl_sitecustomize.py @@ -0,0 +1,72 @@ +"""Install the vLLM adapters used by supported SparkRing profiles.""" + +import os +import sys +import traceback +from collections.abc import Callable +from typing import Any + + +def _install_required(label: str, operation: Callable[[], Any]) -> Any: + """Install one enabled hook or terminate before vLLM can serve traffic. + + CPython's ``site`` module reports and suppresses ordinary exceptions raised + while importing ``sitecustomize``. Explicitly exiting the process here is + therefore part of the feature contract, not merely nicer error handling. + """ + + try: + return operation() + except BaseException: + try: + print( + f"FATAL: required Spark startup hook failed: {label}", + file=sys.stderr, + flush=True, + ) + traceback.print_exc(file=sys.stderr) + sys.stderr.flush() + finally: + # Even a closed/broken stderr must not let CPython's site module + # suppress this required-hook failure and continue serving. + os._exit(78) + raise RuntimeError("os._exit unexpectedly returned") + + +if os.getenv("VLLM_SPARK_TP4_MODE"): + from spark_tp4_backend import install as install_tp4 + + _install_required("TP4 all-reduce backend", install_tp4) + +if os.getenv("SPARK_TP4_HEALTH_GATE") == "1": + from spark_tp4_health_gate import install as install_tp4_health_gate + + _install_required("TP4 post-output health gate", install_tp4_health_gate) + + +if os.getenv("VLLM_SPARK_TP4_VOCAB_MODE"): + from spark_tp4_vocab_allgather_backend import ( + install as install_tp4_vocab_allgather, + ) + + _install_required( + "TP4 vocabulary all-gather backend", + install_tp4_vocab_allgather, + ) + +if os.getenv("SPARK_CUDAGRAPH_REPLAY_TIMING") == "1": + from spark_cudagraph_replay_timing import ( + install as install_cudagraph_replay_timing, + ) + + _install_required( + "CUDA graph replay timing", + install_cudagraph_replay_timing, + ) + +if os.getenv("SPARK_TP4_DCP_COLLECTIVE_AUDIT") == "1": + from spark_dcp_collective_audit import ( + install as install_dcp_collective_audit, + ) + + _install_required("DCP collective audit", install_dcp_collective_audit) diff --git a/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/sitecustomize.py b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/sitecustomize.py new file mode 100644 index 00000000..5fc7b15e --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/sitecustomize.py @@ -0,0 +1,59 @@ +"""Compose the mounted SIRCL hooks with the virtual-diagonal adapter. + +Status: research-only. The private bundle builder installs this file as the +process ``sitecustomize`` and preserves the SIRCL bundle's entry point as +``sircl_sitecustomize.py``. +""" + +from __future__ import annotations + +import os +import runpy +import sys +import traceback +from pathlib import Path + + +BASE_SITECUSTOMIZE = Path("/opt/spark-sircl/sircl_sitecustomize.py") + + +def _required(label: str, operation) -> None: + try: + operation() + except BaseException: + try: + print( + f"FATAL: required private startup hook failed: {label}", + file=sys.stderr, + flush=True, + ) + traceback.print_exc(file=sys.stderr) + sys.stderr.flush() + finally: + os._exit(78) + + +def _install() -> None: + if os.getenv("SPARK_TP4_HEALTH_GATE") != "1": + raise RuntimeError( + "virtual-diagonal full-model testing requires SPARK_TP4_HEALTH_GATE=1" + ) + if not BASE_SITECUSTOMIZE.is_file(): + raise RuntimeError( + f"preserved SIRCL sitecustomize is missing: {BASE_SITECUSTOMIZE}" + ) + + # The base entry point installs every receipt-bound SIRCL, vocabulary, + # timing, audit, and health hook. The B12X wrapper uses a distinct marker + # and is installed outside those wrappers. + runpy.run_path(str(BASE_SITECUSTOMIZE), run_name="_sparkring_sircl_sitecustomize") + + from rocenante_vllm_overlay import install as install_rocenante + + install_rocenante() + from rocenante_health_gate import install as install_rocenante_health + + install_rocenante_health() + + +_required("SIRCL and RoCEnante composition", _install) diff --git a/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/spark_collective_audit.py b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/spark_collective_audit.py new file mode 100644 index 00000000..d23ae248 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/spark_collective_audit.py @@ -0,0 +1,318 @@ +"""Bounded counters proving that live graph requests avoid stock collectives.""" + +from __future__ import annotations + +import os +import threading +from collections import defaultdict +from dataclasses import dataclass + +_PHASES = ("capture", "eager") +_SIGNATURE_LIMIT_ENV = "SPARK_COLLECTIVE_AUDIT_SIGNATURE_LIMIT_PER_PHASE" +_DEFAULT_SIGNATURES_PER_PHASE = 512 +_HARD_MAX_SIGNATURES_PER_PHASE = 2048 +_MAX_SHAPE_RANK = 8 +_MAX_DTYPE_CHARS = 64 +_MAX_UNIQUE_NAME_CHARS = 128 +_MAX_DIMENSION = (1 << 63) - 1 + + +def _signature_limit_from_env() -> int: + raw = os.getenv(_SIGNATURE_LIMIT_ENV) + if raw is None: + return _DEFAULT_SIGNATURES_PER_PHASE + try: + value = int(raw, 10) + except ValueError as exc: + raise RuntimeError( + f"{_SIGNATURE_LIMIT_ENV} must be an integer between 1 and " + f"{_HARD_MAX_SIGNATURES_PER_PHASE}; got {raw!r}" + ) from exc + if not 1 <= value <= _HARD_MAX_SIGNATURES_PER_PHASE: + raise RuntimeError( + f"{_SIGNATURE_LIMIT_ENV} must be an integer between 1 and " + f"{_HARD_MAX_SIGNATURES_PER_PHASE}; got {raw!r}" + ) + return value + + +_MAX_SIGNATURES_PER_PHASE = _signature_limit_from_env() +_lock = threading.Lock() +_counts: dict[str, dict[str, int]] = { + phase: defaultdict(int) for phase in _PHASES +} +_signature_counts: dict[ + str, + dict[tuple[str, str, "StockCollectiveSignature"], int], +] = {phase: {} for phase in _PHASES} +_signature_dropped_calls: dict[str, int] = {phase: 0 for phase in _PHASES} +_signature_dropped_calls_by_family_reason: dict[str, dict[str, int]] = { + phase: defaultdict(int) for phase in _PHASES +} + + +@dataclass(frozen=True) +class StockCollectiveSignature: + """Pointer-free, cardinality-safe evidence for one stock fallback.""" + + shape: tuple[int, ...] + dtype: str + is_cuda: bool + contiguous: bool + world_size: int | None + unique_name: str + + def __post_init__(self) -> None: + if not isinstance(self.shape, tuple): + raise ValueError( + "stock collective signature shape must be a tuple" + ) + if len(self.shape) > _MAX_SHAPE_RANK: + raise ValueError( + "stock collective signature shape rank exceeds " + f"{_MAX_SHAPE_RANK}" + ) + if any( + not isinstance(dimension, int) + or isinstance(dimension, bool) + or dimension < 0 + or dimension > _MAX_DIMENSION + for dimension in self.shape + ): + raise ValueError( + "stock collective signature shape dimensions must be " + "nonnegative signed 64-bit integers" + ) + if ( + not isinstance(self.dtype, str) + or not self.dtype + or len(self.dtype) > _MAX_DTYPE_CHARS + ): + raise ValueError( + "stock collective signature dtype must contain between " + f"1 and {_MAX_DTYPE_CHARS} characters" + ) + if not isinstance(self.is_cuda, bool) or not isinstance( + self.contiguous, bool + ): + raise ValueError( + "stock collective signature CUDA and contiguous fields " + "must be booleans" + ) + if ( + not isinstance(self.unique_name, str) + or len(self.unique_name) > _MAX_UNIQUE_NAME_CHARS + ): + raise ValueError( + "stock collective signature unique_name exceeds " + f"{_MAX_UNIQUE_NAME_CHARS} characters" + ) + if self.world_size is not None and ( + not isinstance(self.world_size, int) + or isinstance(self.world_size, bool) + or self.world_size < 0 + or self.world_size > _MAX_DIMENSION + ): + raise ValueError( + "stock collective signature world_size must be a " + "nonnegative signed 64-bit integer or None" + ) + + def to_dict(self) -> dict[str, object]: + return { + "shape": list(self.shape), + "dtype": self.dtype, + "is_cuda": self.is_cuda, + "contiguous": self.contiguous, + "world_size": self.world_size, + "unique_name": self.unique_name, + } + + +def classify_stock_family( + seam: str, + signature: StockCollectiveSignature, + *, + dim: int | None = None, +) -> str: + """Map a pointer-free call signature to its semantic collective family.""" + shape = signature.shape + dtype = signature.dtype + group = signature.unique_name + + if seam == "group_all_gather": + if group.startswith("dcp:"): + exact_dcp4_tensor = ( + signature.world_size == 4 + and signature.is_cuda + and signature.contiguous + and bool(shape) + and shape[0] > 0 + ) + if ( + exact_dcp4_tensor + and len(shape) == 3 + and shape[1:] == (16, 576) + and dtype == "torch.bfloat16" + and dim in {1, -2} + ): + return "dcp_query_all_gather" + if ( + exact_dcp4_tensor + and len(shape) == 2 + and shape[1] == 64 + and dtype == "torch.float32" + and dim in {0, -2} + ): + return "dcp_lse_all_gather" + return "dcp_all_gather" + if ( + group.startswith("tp:") + and len(shape) == 2 + and shape[1] == 38720 + and dtype == "torch.bfloat16" + and dim in {1, -1} + ): + return "vocabulary_all_gather" + return "group_all_gather" + + if seam == "pynccl_all_gather": + if ( + len(shape) == 3 + and shape[1:] == (2, 2048) + and dtype == "torch.int32" + ): + return "dcp_owner_topk_all_gather" + if ( + len(shape) == 1 + and shape[0] in {23552, 753664} + and dtype == "torch.uint8" + ): + return "dcp_ckv_all_gather" + if ( + shape == (1, 38720) + and dtype == "torch.bfloat16" + ): + return "vocabulary_all_gather" + return "pynccl_all_gather" + + if seam == "group_reduce_scatter": + if ( + group.startswith("dcp:") + and len(shape) == 3 + and shape[1] == 64 + and shape[2] in {256, 512} + and dtype == "torch.bfloat16" + and dim in {1, -2} + ): + return "dcp_output_reduce_scatter" + return "group_reduce_scatter" + + raise ValueError(f"unknown stock collective seam: {seam}") + + +def enabled() -> bool: + """Return whether the live graph status surface requested auditing.""" + return bool(os.getenv("SPARK_TP4_GRAPH_STATUS_PATH")) + + +def record_stock( + family: str, + *, + capturing: bool, + reason: str, + signature: StockCollectiveSignature | None = None, +) -> None: + """Record one call through an original/stock collective implementation.""" + if not enabled(): + return + if not family or not reason: + raise ValueError("stock collective family and reason must be nonempty") + phase = "capture" if capturing else "eager" + key = f"{family}:{reason}" + with _lock: + _counts[phase][key] += 1 + if signature is not None: + signature_key = (family, reason, signature) + phase_signatures = _signature_counts[phase] + if signature_key in phase_signatures: + phase_signatures[signature_key] += 1 + elif len(phase_signatures) < _MAX_SIGNATURES_PER_PHASE: + phase_signatures[signature_key] = 1 + else: + _signature_dropped_calls[phase] += 1 + _signature_dropped_calls_by_family_reason[phase][key] += 1 + + +def _signature_sort_key( + item: tuple[tuple[str, str, StockCollectiveSignature], int], +) -> tuple[object, ...]: + (family, reason, signature), _count = item + return ( + family, + reason, + signature.shape, + signature.dtype, + signature.is_cuda, + signature.contiguous, + -1 if signature.world_size is None else signature.world_size, + signature.unique_name, + ) + + +def stock_collective_snapshot() -> dict[str, object]: + """Return a stable, bounded copy for the low-rate status reporter.""" + with _lock: + phases = { + phase: dict(sorted(_counts[phase].items())) + for phase in _PHASES + } + signatures = { + phase: [ + { + "family": family, + "reason": reason, + "count": count, + **signature.to_dict(), + } + for (family, reason, signature), count in sorted( + _signature_counts[phase].items(), + key=_signature_sort_key, + ) + ] + for phase in _PHASES + } + dropped_calls = dict(_signature_dropped_calls) + dropped_calls_by_family_reason = { + phase: dict( + sorted(_signature_dropped_calls_by_family_reason[phase].items()) + ) + for phase in _PHASES + } + snapshot: dict[str, object] = { + "capture": phases["capture"], + "eager": phases["eager"], + "capture_total": sum(phases["capture"].values()), + "eager_total": sum(phases["eager"].values()), + } + if any(signatures.values()) or any(dropped_calls.values()): + snapshot.update( + { + "signatures": signatures, + "signature_limit_per_phase": _MAX_SIGNATURES_PER_PHASE, + "signature_dropped_calls": dropped_calls, + "signature_dropped_calls_by_family_reason": ( + dropped_calls_by_family_reason + ), + } + ) + return snapshot + + +def _reset_for_tests() -> None: + with _lock: + for phase in _PHASES: + _counts[phase].clear() + _signature_counts[phase].clear() + _signature_dropped_calls[phase] = 0 + _signature_dropped_calls_by_family_reason[phase].clear() diff --git a/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/spark_cudagraph_bucket_contract.py b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/spark_cudagraph_bucket_contract.py new file mode 100644 index 00000000..406bcfe1 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/spark_cudagraph_bucket_contract.py @@ -0,0 +1,226 @@ +"""Separate CUDA-graph bucket contracts for decode and prefill. + +The native TP4 query transport remains bounded at Q40. PIECEWISE prefill +graphs may pad larger observed query shapes into a small, separately attested +bucket set through Q512; those buckets must never become FULL speculative +decode widths. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass + + +DECODE_CAPTURE_ENV = "VLLM_SPARK_DECODE_CAPTURE_SIZES" +FULL_DECODE_CAPTURE_ENV = "VLLM_SPARK_FULL_DECODE_CAPTURE_SIZES" +PREFILL_PIECEWISE_CAPTURE_ENV = ( + "VLLM_SPARK_PREFILL_PIECEWISE_CAPTURE_SIZES" +) +COMBINED_CAPTURE_ENV = "VLLM_SPARK_GRAPH_CAPTURE_SIZES" + +MAX_FULL_DECODE_QUERY_ROWS = 40 +MAX_PREFILL_PIECEWISE_QUERY_ROWS = 512 +MAX_PREFILL_BUCKET_COUNT = 16 +MAX_OBSERVED_PREFILL_PADDING_ROWS = 32 + +# Existing C8 adaptive-MTP4 graph plan. Only its Q5 multiples are eligible +# for FULL speculative decode; the other values remain decode padding buckets. +DECODE_CAPTURE_BUCKETS = ( + 1, + 2, + 3, + 4, + 5, + 6, + 8, + 10, + 12, + 15, + 16, + 20, + 24, + 25, + 30, + 32, + 35, + 40, +) +FULL_SPECULATIVE_DECODE_BUCKETS = tuple( + range(5, MAX_FULL_DECODE_QUERY_ROWS + 1, 5) +) + +OBSERVED_PREFILL_QUERY_ROWS = (48, 69, 72, 143, 210, 279, 348, 417, 486) +DEFAULT_PREFILL_PIECEWISE_BUCKETS = ( + 48, + 72, + 144, + 224, + 288, + 352, + 432, + 512, +) + + +@dataclass(frozen=True) +class CUDAGraphBucketContract: + decode: tuple[int, ...] + full_speculative_decode: tuple[int, ...] + prefill_piecewise: tuple[int, ...] + combined: tuple[int, ...] + + +def _strict_ordered_buckets( + buckets: Sequence[int], *, field: str +) -> tuple[int, ...]: + values = tuple(buckets) + if ( + not values + or any( + isinstance(value, bool) or not isinstance(value, int) + for value in values + ) + or tuple(sorted(set(values))) != values + ): + raise ValueError( + f"{field} must be nonempty, strictly increasing unique integers" + ) + return values + + +def validate_full_speculative_decode_buckets( + buckets: Sequence[int], +) -> tuple[int, ...]: + values = _strict_ordered_buckets( + buckets, field="FULL speculative decode buckets" + ) + if values != FULL_SPECULATIVE_DECODE_BUCKETS: + raise ValueError( + "FULL speculative decode buckets must be Q5 multiples through " + f"Q{MAX_FULL_DECODE_QUERY_ROWS}" + ) + return values + + +def prefill_padding_plan( + buckets: Sequence[int], +) -> dict[int, int]: + values = tuple(buckets) + plan: dict[int, int] = {} + for query_rows in OBSERVED_PREFILL_QUERY_ROWS: + padded = next( + (bucket for bucket in values if bucket >= query_rows), + None, + ) + if padded is None: + raise ValueError( + f"PIECEWISE prefill buckets do not cover Q{query_rows}" + ) + plan[query_rows] = padded + return plan + + +def maximum_observed_padding(buckets: Sequence[int]) -> int: + return max( + padded - query_rows + for query_rows, padded in prefill_padding_plan(buckets).items() + ) + + +def validate_prefill_piecewise_buckets( + buckets: Sequence[int], +) -> tuple[int, ...]: + values = _strict_ordered_buckets( + buckets, field="PIECEWISE prefill buckets" + ) + if len(values) > MAX_PREFILL_BUCKET_COUNT: + raise ValueError( + "PIECEWISE prefill bucket count exceeds " + f"{MAX_PREFILL_BUCKET_COUNT}" + ) + if values[0] <= MAX_FULL_DECODE_QUERY_ROWS: + raise ValueError( + "PIECEWISE prefill buckets must remain above the Q40 decode " + "contract" + ) + if values[-1] != MAX_PREFILL_PIECEWISE_QUERY_ROWS: + raise ValueError( + "PIECEWISE prefill buckets must terminate exactly at Q512" + ) + if any(value > MAX_PREFILL_PIECEWISE_QUERY_ROWS for value in values): + raise ValueError("PIECEWISE prefill buckets may not exceed Q512") + padding = maximum_observed_padding(values) + if padding > MAX_OBSERVED_PREFILL_PADDING_ROWS: + raise ValueError( + "PIECEWISE prefill buckets exceed observed-shape padding bound: " + f"{padding} > {MAX_OBSERVED_PREFILL_PADDING_ROWS}" + ) + return values + + +def combined_capture_buckets( + prefill_piecewise: Sequence[int], +) -> tuple[int, ...]: + prefill = validate_prefill_piecewise_buckets(prefill_piecewise) + return DECODE_CAPTURE_BUCKETS + prefill + + +def format_buckets(buckets: Sequence[int]) -> str: + return ",".join(str(value) for value in buckets) + + +def parse_bucket_csv(raw: object, *, field: str) -> tuple[int, ...]: + if not isinstance(raw, str) or not raw: + raise ValueError(f"{field} must be an explicit canonical CSV list") + try: + values = tuple(int(item) for item in raw.split(",")) + except ValueError as error: + raise ValueError( + f"{field} must be an explicit canonical CSV list" + ) from error + if format_buckets(values) != raw: + raise ValueError(f"{field} must be an explicit canonical CSV list") + return values + + +def contract_from_environment( + environment: Mapping[str, str], +) -> CUDAGraphBucketContract: + decode = parse_bucket_csv( + environment.get(DECODE_CAPTURE_ENV), + field=DECODE_CAPTURE_ENV, + ) + if decode != DECODE_CAPTURE_BUCKETS: + raise ValueError( + f"{DECODE_CAPTURE_ENV} must preserve the current 18 decode " + "buckets" + ) + full = validate_full_speculative_decode_buckets( + parse_bucket_csv( + environment.get(FULL_DECODE_CAPTURE_ENV), + field=FULL_DECODE_CAPTURE_ENV, + ) + ) + prefill = validate_prefill_piecewise_buckets( + parse_bucket_csv( + environment.get(PREFILL_PIECEWISE_CAPTURE_ENV), + field=PREFILL_PIECEWISE_CAPTURE_ENV, + ) + ) + combined = parse_bucket_csv( + environment.get(COMBINED_CAPTURE_ENV), + field=COMBINED_CAPTURE_ENV, + ) + expected_combined = combined_capture_buckets(prefill) + if combined != expected_combined: + raise ValueError( + f"{COMBINED_CAPTURE_ENV} does not equal decode + PIECEWISE " + "prefill buckets" + ) + return CUDAGraphBucketContract( + decode=decode, + full_speculative_decode=full, + prefill_piecewise=prefill, + combined=combined, + ) diff --git a/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/spark_cudagraph_replay_timing.py b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/spark_cudagraph_replay_timing.py new file mode 100644 index 00000000..0f4011b1 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/spark_cudagraph_replay_timing.py @@ -0,0 +1,296 @@ +"""Arm-file-gated CUDA timing for vLLM FULL graph replays. + +The serving thread only records CUDA events. Completion polling and elapsed +time calculation happen from the existing low-rate graph status reporter, so +the measured request is not host-synchronized by this diagnostic. +""" + +from __future__ import annotations + +import hashlib +import inspect +import math +import os +import threading +from collections import defaultdict, deque +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +_SUPPORTED_RUN_FULLGRAPH_SHA256 = { + ( + "0.11.2.dev279+eldritch.final.fcc6141.b12x284a2ea." + "fi25dd814.cu132.20260626" + ): { + "4d58b8ef1a5023af0c11eb7a659620faca15f8a0303b37774ed0d28f4a5919db" + }, + # Public GLM-5.3 Flash operator image. Keep the source hash gate: the + # version string alone does not prove that replay still has the semantics + # this diagnostic wraps. + "0.1.dev1+gd377796e8": { + "af0da51c4ca27ed318af8d50baaaec3074de33820b2f1be4a6b75de1510a9bb3" + }, +} +_DEFAULT_SAMPLE_LIMIT = 512 +_MAXIMUM_SAMPLE_LIMIT = 65536 +_installed = False +_collector: ReplayTimingCollector | None = None +_timing_reporter: Any | None = None + + +@dataclass(frozen=True) +class _PendingSample: + key: str + start: Any + end: Any + + +def _percentile(values: list[float], percentile: float) -> float: + ordered = sorted(values) + index = max(0, math.ceil(percentile * len(ordered)) - 1) + return ordered[index] + + +class ReplayTimingCollector: + def __init__( + self, + *, + event_factory: Callable[[], Any], + arm_path: Path, + sample_limit: int, + ) -> None: + if not 1 <= sample_limit <= _MAXIMUM_SAMPLE_LIMIT: + raise ValueError( + "sample_limit must be in " + f"[1, {_MAXIMUM_SAMPLE_LIMIT}]" + ) + self._event_factory = event_factory + self._arm_path = arm_path + self._sample_limit = sample_limit + self._lock = threading.Lock() + self._pending: deque[_PendingSample] = deque() + self._durations: dict[str, list[float]] = defaultdict(list) + self._reserved = 0 + self._dropped = 0 + self._errors = 0 + + def measure( + self, + key: str, + stream: Any, + operation: Callable[[], Any], + ) -> Any: + if not self._arm_path.is_file(): + return operation() + + with self._lock: + if self._reserved >= self._sample_limit: + self._dropped += 1 + return operation() + self._reserved += 1 + + start = self._event_factory() + end = self._event_factory() + start.record(stream) + try: + return operation() + finally: + end.record(stream) + with self._lock: + self._pending.append( + _PendingSample(key=key, start=start, end=end) + ) + + def snapshot(self) -> dict[str, Any]: + with self._lock: + pending = list(self._pending) + self._pending.clear() + + completed_samples: list[tuple[str, float]] = [] + remaining: list[_PendingSample] = [] + errors = 0 + for index, sample in enumerate(pending): + try: + if not sample.end.query(): + remaining = pending[index:] + break + completed_samples.append( + ( + sample.key, + float(sample.start.elapsed_time(sample.end)), + ) + ) + except Exception: + errors += 1 + + with self._lock: + if remaining: + self._pending = deque(remaining) + self._pending + for key, elapsed_ms in completed_samples: + self._durations[key].append(elapsed_ms) + self._errors += errors + durations = { + key: list(values) + for key, values in self._durations.items() + } + reserved = self._reserved + pending_count = len(self._pending) + dropped = self._dropped + error_count = self._errors + + descriptors: dict[str, dict[str, float | int]] = {} + total_ms = 0.0 + completed = 0 + for key, values in sorted(durations.items()): + count = len(values) + subtotal = sum(values) + completed += count + total_ms += subtotal + descriptors[key] = { + "count": count, + "total_ms": round(subtotal, 6), + "mean_ms": round(subtotal / count, 6), + "p50_ms": round(_percentile(values, 0.50), 6), + "p90_ms": round(_percentile(values, 0.90), 6), + "max_ms": round(max(values), 6), + } + + return { + "enabled": True, + "armed": self._arm_path.is_file(), + "arm_path": str(self._arm_path), + "sample_limit": self._sample_limit, + "reserved": reserved, + "completed": completed, + "pending": pending_count, + "dropped": dropped, + "errors": error_count, + "total_completed_ms": round(total_ms, 6), + "descriptors": descriptors, + } + + +def _sample_limit() -> int: + text = os.getenv( + "SPARK_CUDAGRAPH_REPLAY_TIMING_SAMPLES", + str(_DEFAULT_SAMPLE_LIMIT), + ) + try: + value = int(text) + except ValueError as error: + raise RuntimeError( + "SPARK_CUDAGRAPH_REPLAY_TIMING_SAMPLES must be an integer" + ) from error + if not 1 <= value <= _MAXIMUM_SAMPLE_LIMIT: + raise RuntimeError( + "SPARK_CUDAGRAPH_REPLAY_TIMING_SAMPLES must be in " + f"[1, {_MAXIMUM_SAMPLE_LIMIT}]" + ) + return value + + +def _arm_path() -> Path: + value = os.getenv("SPARK_CUDAGRAPH_REPLAY_TIMING_ARM_PATH") + if not value: + raise RuntimeError( + "SPARK_CUDAGRAPH_REPLAY_TIMING_ARM_PATH is required" + ) + return Path(value) + + +def _descriptor_key(descriptor: Any) -> str: + mode = getattr(getattr(descriptor, "cg_mode", None), "name", "unknown") + fields = ( + ("num_tokens", getattr(descriptor, "num_tokens", None)), + ("num_reqs", getattr(descriptor, "num_reqs", None)), + ( + "uniform_token_count", + getattr(descriptor, "uniform_token_count", None), + ), + ("max_query_len", getattr(descriptor, "max_query_len", None)), + ( + "num_active_loras", + getattr(descriptor, "num_active_loras", 0), + ), + ) + encoded = ",".join( + f"{name}={'none' if value is None else int(value)}" + for name, value in fields + ) + return f"mode={mode},{encoded}" + + +def install() -> None: + global _collector, _installed, _timing_reporter + if _installed: + return + + import torch + import vllm + from vllm.v1.worker.gpu.cudagraph_utils import CudaGraphManager + + expected_hashes = _SUPPORTED_RUN_FULLGRAPH_SHA256.get(vllm.__version__) + if expected_hashes is None: + raise RuntimeError( + "unsupported vLLM version for graph replay timing: " + f"{vllm.__version__}" + ) + original = CudaGraphManager.run_fullgraph + actual_hash = hashlib.sha256( + inspect.getsource(original).encode("utf-8") + ).hexdigest() + if actual_hash not in expected_hashes: + raise RuntimeError( + "unsupported CudaGraphManager.run_fullgraph source: " + f"{actual_hash}" + ) + if getattr(original, "_spark_replay_timing", False): + _installed = True + return + + _collector = ReplayTimingCollector( + event_factory=lambda: torch.cuda.Event(enable_timing=True), + arm_path=_arm_path(), + sample_limit=_sample_limit(), + ) + + status_path = os.getenv("SPARK_CUDAGRAPH_REPLAY_TIMING_STATUS_PATH") + if status_path: + try: + rank = int(os.environ["SPARKRING_NODE_RANK"]) + except (KeyError, ValueError) as error: + raise RuntimeError( + "SPARKRING_NODE_RANK must be an integer when the replay " + "timing status path is enabled" + ) from error + from spark_graph_status_reporter import start_status_reporter + + _timing_reporter = start_status_reporter( + status_path, + snapshot_provider=lambda: { + "cudagraph_replay_timing": graph_replay_timing_snapshot() + }, + interval_seconds=0.25, + rank=rank, + ) + + def timed_run_fullgraph(self: Any, descriptor: Any) -> Any: + assert _collector is not None + stream = torch.cuda.current_stream(self.device) + return _collector.measure( + _descriptor_key(descriptor), + stream, + lambda: original(self, descriptor), + ) + + timed_run_fullgraph._spark_replay_timing = True # type: ignore[attr-defined] + timed_run_fullgraph._spark_original = original # type: ignore[attr-defined] + CudaGraphManager.run_fullgraph = timed_run_fullgraph + _installed = True + + +def graph_replay_timing_snapshot() -> dict[str, Any]: + if _collector is None: + return {"enabled": False} + return _collector.snapshot() diff --git a/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/spark_dcp_collective_audit.py b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/spark_dcp_collective_audit.py new file mode 100644 index 00000000..dfb1ad0d --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/spark_dcp_collective_audit.py @@ -0,0 +1,204 @@ +"""Audit-only DCP combine/reduce-scatter accounting at vLLM seams.""" + +from __future__ import annotations + +import importlib.abc +import importlib.machinery +import os +import sys +from types import ModuleType +from typing import Any + +_installed = False +_COMBINE_TARGET = "vllm.v1.attention.ops.common" + + +def _enabled() -> bool: + value = os.getenv("SPARK_TP4_DCP_COLLECTIVE_AUDIT", "0") + if value not in {"0", "1"}: + raise ValueError( + "SPARK_TP4_DCP_COLLECTIVE_AUDIT must be '0' or '1'" + ) + return value == "1" + + +def _is_stream_capturing(torch_module: Any) -> bool: + checker = getattr(torch_module.cuda, "is_current_stream_capturing", None) + return bool(checker is not None and checker()) + + +def _signature(group: Any, tensor: Any) -> Any: + from spark_collective_audit import StockCollectiveSignature + + world_size = getattr(group, "world_size", None) + return StockCollectiveSignature( + shape=tuple(int(value) for value in tensor.shape), + dtype=str(tensor.dtype), + is_cuda=bool(tensor.is_cuda), + contiguous=bool(tensor.is_contiguous()), + world_size=None if world_size is None else int(world_size), + unique_name=str(getattr(group, "unique_name", "")), + ) + + +def _standalone_combine_audit_enabled() -> bool: + """Use the stock-only hook only when the native DCP adapter is absent.""" + + return not bool(os.getenv("VLLM_SPARK_TP4_DCP_MODE", "")) + + +def _patch_existing_combine_aliases(original: Any, replacement: Any) -> None: + """Repair aliases bound before the defining common module was patched.""" + + for loaded in tuple(sys.modules.values()): + if not isinstance(loaded, ModuleType) or loaded is sys.modules.get( + _COMBINE_TARGET + ): + continue + if vars(loaded).get("cp_lse_ag_out_rs") is original: + loaded.cp_lse_ag_out_rs = replacement + + +def _patch_combine(module: ModuleType) -> None: + """Wrap the stock combine without constructing any transport state.""" + + current = module.cp_lse_ag_out_rs + if getattr(current, "_spark_dcp_collective_audit", False): + _patch_existing_combine_aliases(current._spark_original, current) + return + # A selected native DCP adapter owns this seam and records its own stock + # calls. Wrapping it here would falsely label custom calls as stock. + if getattr(current, "_spark_tp4_dcp_backend", False): + return + original = current + + def audited_combine( + cp_attn_out: Any, + cp_attn_lse: Any, + cp_group: Any, + ctx: Any = None, + return_lse: bool = False, + is_lse_base_on_e: bool = True, + head_major_output: bool = False, + ) -> Any: + import torch + from spark_collective_audit import enabled, record_stock + + signature = _signature(cp_group, cp_attn_out) if enabled() else None + record_stock( + "dcp_combine", + capturing=_is_stream_capturing(torch), + reason="original", + signature=signature, + ) + return original( + cp_attn_out, + cp_attn_lse, + cp_group, + ctx=ctx, + return_lse=return_lse, + is_lse_base_on_e=is_lse_base_on_e, + head_major_output=head_major_output, + ) + + audited_combine._spark_dcp_collective_audit = True # type: ignore[attr-defined] + audited_combine._spark_original = original # type: ignore[attr-defined] + module.cp_lse_ag_out_rs = audited_combine + _patch_existing_combine_aliases(original, audited_combine) + + +class _CombineLoader(importlib.abc.Loader): + def __init__(self, delegate: importlib.abc.Loader) -> None: + self._delegate = delegate + + def create_module(self, spec: Any) -> ModuleType | None: + create = getattr(self._delegate, "create_module", None) + return None if create is None else create(spec) + + def exec_module(self, module: ModuleType) -> None: + self._delegate.exec_module(module) + _patch_combine(module) + + +class _CombineFinder(importlib.abc.MetaPathFinder): + def find_spec( + self, + fullname: str, + path: Any, + target: ModuleType | None = None, + ) -> Any: + if fullname != _COMBINE_TARGET: + return None + try: + sys.meta_path.remove(self) + except ValueError: + pass + spec = importlib.machinery.PathFinder.find_spec(fullname, path, target) + if spec is None or spec.loader is None: + return spec + spec.loader = _CombineLoader(spec.loader) + return spec + + +def _install_combine_import_hook() -> None: + loaded = sys.modules.get(_COMBINE_TARGET) + if loaded is not None: + _patch_combine(loaded) + return + if not any(isinstance(finder, _CombineFinder) for finder in sys.meta_path): + sys.meta_path.insert(0, _CombineFinder()) + + +def install() -> None: + """Install pointer-free accounting wrappers without changing execution.""" + global _installed + if _installed or not _enabled(): + return + + if _standalone_combine_audit_enabled(): + # Install before importing vLLM so a deferred common-module import is + # captured. Native/custom DCP mode uses spark_tp4_dcp_backend instead; + # two independent import hooks on this seam would not be fail-closed. + _install_combine_import_hook() + + from vllm.distributed.parallel_state import GroupCoordinator + + original = GroupCoordinator._reduce_scatter_out_place + if getattr(original, "_spark_dcp_collective_audit", False): + _installed = True + return + + def audited_reduce_scatter( + self: Any, + input_tensor: Any, + dim: int, + ) -> Any: + import torch + from spark_collective_audit import ( + classify_stock_family, + enabled, + record_stock, + ) + + signature = _signature(self, input_tensor) if enabled() else None + family = ( + "group_reduce_scatter" + if signature is None + else classify_stock_family( + "group_reduce_scatter", + signature, + dim=dim, + ) + ) + record_stock( + family, + capturing=_is_stream_capturing(torch), + reason="original", + signature=signature, + ) + return original(self, input_tensor, dim) + + audited_reduce_scatter._spark_dcp_collective_audit = True # type: ignore[attr-defined] + audited_reduce_scatter._spark_original = original # type: ignore[attr-defined] + GroupCoordinator._reduce_scatter_out_place = audited_reduce_scatter + _installed = True diff --git a/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/spark_graph_status_reporter.py b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/spark_graph_status_reporter.py new file mode 100644 index 00000000..fcf05edd --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/spark_graph_status_reporter.py @@ -0,0 +1,145 @@ +"""Low-rate atomic status files for live graph-native transport gates.""" + +from __future__ import annotations + +import json +import os +import threading +import time +from collections.abc import Callable +from pathlib import Path +from typing import Any + +SnapshotProvider = Callable[[], dict[str, Any]] +_process_reporter: GraphStatusReporter | None = None +_process_reporter_lock = threading.Lock() + + +def collect_graph_status() -> dict[str, object]: + from spark_collective_audit import stock_collective_snapshot + from spark_tp4_backend import graph_q1_diagnostic_snapshot + from spark_tp4_vocab_allgather_backend import ( + vocab_graph_diagnostic_snapshot, + ) + + snapshot: dict[str, object] = { + "all_reduce": graph_q1_diagnostic_snapshot(), + "vocabulary": vocab_graph_diagnostic_snapshot(), + "stock_collectives": stock_collective_snapshot(), + } + if os.getenv("SPARK_CUDAGRAPH_REPLAY_TIMING") == "1": + from spark_cudagraph_replay_timing import ( + graph_replay_timing_snapshot, + ) + + snapshot["cudagraph_replay_timing"] = ( + graph_replay_timing_snapshot() + ) + return snapshot + + +class GraphStatusReporter: + def __init__( + self, + path: Path, + snapshot_provider: SnapshotProvider, + interval_seconds: float, + rank: int, + ) -> None: + self._path = path + self._snapshot_provider = snapshot_provider + self._interval_seconds = interval_seconds + self._rank = rank + self._stop = threading.Event() + self._thread = threading.Thread( + target=self._run, + name=f"spark-graph-status-rank{rank}", + daemon=True, + ) + self._thread.start() + + def _publish(self) -> None: + snapshot_start_unix_ns = time.time_ns() + snapshot = self._snapshot_provider() + snapshot_end_unix_ns = time.time_ns() + payload = { + "schema_version": 3, + # Kept as the compatibility timestamp; v3 consumers should use + # the explicit collection interval to bind before/after deltas. + "unix_ns": snapshot_end_unix_ns, + "snapshot_start_unix_ns": snapshot_start_unix_ns, + "snapshot_end_unix_ns": snapshot_end_unix_ns, + "pid": os.getpid(), + "rank": self._rank, + "snapshot": snapshot, + } + encoded = json.dumps( + payload, sort_keys=True, separators=(",", ":") + ) + self._path.parent.mkdir(parents=True, exist_ok=True) + temporary = self._path.with_name( + f".{self._path.name}.{os.getpid()}.tmp" + ) + temporary.write_text(encoded + "\n", encoding="utf-8") + deadline = time.monotonic() + 0.25 + while True: + try: + os.replace(temporary, self._path) + break + except PermissionError: + if time.monotonic() >= deadline: + raise + time.sleep(0.001) + + def _run(self) -> None: + while not self._stop.is_set(): + try: + self._publish() + except Exception: + pass + self._stop.wait(self._interval_seconds) + + def stop(self) -> None: + self._stop.set() + self._thread.join(timeout=max(1.0, self._interval_seconds * 2)) + + +def start_status_reporter( + path: str | Path, + *, + snapshot_provider: SnapshotProvider, + interval_seconds: float = 0.25, + rank: int, +) -> GraphStatusReporter: + return GraphStatusReporter( + Path(path), snapshot_provider, interval_seconds, rank + ) + + +def ensure_status_reporter( + *, + rank: int, + interval_seconds: float = 0.25, +) -> GraphStatusReporter | None: + path = os.getenv("SPARK_TP4_GRAPH_STATUS_PATH") + if not path: + return None + global _process_reporter + with _process_reporter_lock: + if _process_reporter is None: + _process_reporter = start_status_reporter( + path, + snapshot_provider=collect_graph_status, + interval_seconds=interval_seconds, + rank=rank, + ) + return _process_reporter + + +def stop_status_reporter() -> None: + global _process_reporter + with _process_reporter_lock: + reporter = _process_reporter + _process_reporter = None + if reporter is not None: + reporter.stop() diff --git a/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/spark_persistent_output_ring.py b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/spark_persistent_output_ring.py new file mode 100644 index 00000000..f63864bd --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/spark_persistent_output_ring.py @@ -0,0 +1,77 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + + +@dataclass(frozen=True) +class OutputSignature: + shape: tuple[int, ...] + dtype: object + device: object + + +class PersistentOutputRing: + """Lazily allocate and cycle stable CUDA output buffers. + + One ring is owned by one fixed-payload native transport session. The + caller remains responsible for submitting producers and consumers in + CUDA stream order before a slot is reused. + """ + + def __init__(self, slots: int) -> None: + if slots <= 0: + raise ValueError("persistent output ring slots must be positive") + self._slots = slots + self._signature: OutputSignature | None = None + self._storage: Any | None = None + self._outputs: tuple[Any, ...] = () + self._next = 0 + self._acquires = 0 + + @property + def slots(self) -> int: + return self._slots + + @property + def acquires(self) -> int: + return self._acquires + + @property + def wraps(self) -> int: + return self._acquires // self._slots + + @staticmethod + def _signature_for(tensor: Any) -> OutputSignature: + return OutputSignature( + shape=tuple(int(dimension) for dimension in tensor.shape), + dtype=tensor.dtype, + device=tensor.device, + ) + + def acquire(self, tensor: Any, torch_module: Any) -> Any: + signature = self._signature_for(tensor) + if self._signature is None: + storage = torch_module.empty( + (self._slots, *signature.shape), + dtype=signature.dtype, + device=signature.device, + ) + outputs = tuple(storage.unbind(0)) + if len(outputs) != self._slots: + raise RuntimeError( + "persistent output storage returned the wrong slot count" + ) + self._storage = storage + self._outputs = outputs + self._signature = signature + elif signature != self._signature: + raise RuntimeError( + "persistent output session signature changed: " + f"expected={self._signature} observed={signature}" + ) + + output = self._outputs[self._next] + self._next = (self._next + 1) % self._slots + self._acquires += 1 + return output diff --git a/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/spark_tp4_backend.py b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/spark_tp4_backend.py new file mode 100644 index 00000000..14381a54 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/spark_tp4_backend.py @@ -0,0 +1,2315 @@ +"""Signature-gated vLLM adapter for the direct-cable TP4 transport.""" + +from __future__ import annotations + +import ctypes +import logging +import os +import sys +from dataclasses import asdict, dataclass +from typing import Any + +from spark_persistent_output_ring import PersistentOutputRing +from spark_tp4_port_namespace import ( + bidirectional_prefill_control_ports, + bidirectional_prefill_secondary_control_ports, + eager_allreduce_admitted_widths, + eager_allreduce_ports_for_payload, + graph_allreduce_control_ports, + graph_dual_port_q40_control_ports, + validate_active_port_namespace, + validate_control_port_pair, +) +from spark_tp4_query_row_provider import resolve_query_rows +from spark_tp4_query_contract import ( + ABSOLUTE_MAX_QUERY_ROWS, + MAX_QUERY_ROWS, +) +from spark_tp4_capability import ensure_capability_vote + +logger = logging.getLogger(__name__) + +_installed = False +_VALID_MODES = {"shadow", "custom", "disabled"} +_TARGET_WIDTH = 6144 +_BF16_BYTES = 2 +_BYTES_PER_ROW = _TARGET_WIDTH * _BF16_BYTES +_ALLREDUCE_PREFILL_MAX_QUERY_ROWS = 512 +_ALLREDUCE_PREFILL_CAPACITY_BYTES = ( + _ALLREDUCE_PREFILL_MAX_QUERY_ROWS * _BYTES_PER_ROW +) +_DUAL_PORT_Q40_CAPACITY_BYTES = ABSOLUTE_MAX_QUERY_ROWS * _BYTES_PER_ROW +_GRAPH_CAPACITY_BYTES = MAX_QUERY_ROWS * _BYTES_PER_ROW +# Research-only width-4096 graph geometry (DeepSeek decode shapes). One +# maximum-capacity session serves every [Q <= 512, 4096] bucket; the +# command ring carries the active Q. +_RESEARCH_GRAPH_WIDTH = 4096 +_RESEARCH_GRAPH_ROW_BYTES = _RESEARCH_GRAPH_WIDTH * _BF16_BYTES +_RESEARCH_GRAPH_MAX_QUERY_ROWS = 512 +_RESEARCH_GRAPH_CAPACITY_BYTES = ( + _RESEARCH_GRAPH_MAX_QUERY_ROWS * _RESEARCH_GRAPH_ROW_BYTES +) +_BIDIRECTIONAL_PREFILL_WIDTH = 4096 +_BIDIRECTIONAL_PREFILL_ROW_BYTES = _BIDIRECTIONAL_PREFILL_WIDTH * _BF16_BYTES +_BIDIRECTIONAL_PREFILL_QUERY_ROWS = frozenset({1024, 2048, 4096, 8192}) +# Keep these research C ABI names synchronized with tp4_c_api.h; no dispatch +# code embeds a symbol name directly. +_BIDIRECTIONAL_PREFILL_C_ABI_SYMBOLS = { + "create": "spark_tp4_bidirectional_prefill_create", + "all_reduce": "spark_tp4_bidirectional_prefill_all_reduce", + "health": "spark_tp4_bidirectional_prefill_get_health_status", + "destroy": "spark_tp4_bidirectional_prefill_destroy", +} +_FUSED_PREFILL_C_ABI_SYMBOLS = { + "create": "spark_tp4_fused_prefill_create", + "all_reduce": "spark_tp4_fused_prefill_all_reduce_rows", + "health": "spark_tp4_fused_prefill_get_health_status", + "destroy": "spark_tp4_fused_prefill_destroy", +} +_GRAPH_STATUS_CAPTURE_CONFIGURED = 1 << 0 +_GRAPH_STATUS_POLLING_ENABLED = 1 << 1 +_GRAPH_STATUS_HOST_NATIVE_ATOMICS = 1 << 2 +_GRAPH_STATUS_SUBMIT_AFFINITY_VERIFIED = 1 << 3 +_GRAPH_STATUS_PROGRESS_AFFINITY_VERIFIED = 1 << 4 +_GRAPH_STATUS_OVERFLOW_FATAL = 1 << 5 +_GRAPH_STATUS_TWO_SLOT_DEFERRED_ACK = 1 << 7 +_GRAPH_STATUS_SPLIT_64K = 1 << 8 +_GRAPH_STATUS_TIERED_64K = 1 << 9 +_GRAPH_STATUS_DUAL_PORT_STRIPED = 1 << 10 +_GRAPH_STATUS_DIRECT_DOORBELL = 1 << 11 +_HEALTHY = 1 << 0 +_HEALTH_POISONED = 1 << 1 +_HEALTH_PROGRESS_THREAD_RUNNING = 1 << 2 +_HEALTH_STOPPING = 1 << 3 +_SERIAL_ACK_PROTOCOL = "serial_ack" +_TWO_SLOT_DEFERRED_ACK_PROTOCOL = "two_slot_deferred_ack" +_ALLREDUCE_PROTOCOL_WIRE = { + _SERIAL_ACK_PROTOCOL: 0, + _TWO_SLOT_DEFERRED_ACK_PROTOCOL: 1, +} +_FUSED_GRAPH_KERNEL = "fused" +_SPLIT_64K_GRAPH_KERNEL = "split_64k" +_TIERED_64K_GRAPH_KERNEL = "tiered_64k" +_GRAPH_KERNEL_WIRE = { + _FUSED_GRAPH_KERNEL: 0, + _SPLIT_64K_GRAPH_KERNEL: 1, + _TIERED_64K_GRAPH_KERNEL: 2, +} +_SEQUENTIAL_WIRE_SCHEDULE = "sequential" +_DUAL_PORT_STRIPED_WIRE_SCHEDULE = "dual_port_striped" +_WIRE_SCHEDULE_WIRE = { + _SEQUENTIAL_WIRE_SCHEDULE: 0, + _DUAL_PORT_STRIPED_WIRE_SCHEDULE: 1, +} +_MAX_PERSISTENT_OUTPUT_SLOTS = 4096 +_graph_q1_sessions: dict[int, "_NativeSession"] = {} +_graph_dual_port_q40_sessions: dict[int, "_NativeSession"] = {} +_graph_width4096_sessions: dict[int, "_NativeSession"] = {} +_backends: dict[int, "_Backend"] = {} +_graph_event_counts: dict[str, int] = {} +# PLACEHOLDER ring peers (RFC 5737 TEST-NET-1): 192.0.2.N stands in for +# rank N-1's direct-cable address. These are NOT routable and MUST be +# replaced for any live run by setting SPARK_TP4_PEER0 / SPARK_TP4_PEER1 +# (the authoritative per-rank overrides) or by editing this table. +_DEFAULT_PEERS = { + 0: ("192.0.2.2", "192.0.2.4"), + 1: ("192.0.2.1", "192.0.2.3"), + 2: ("192.0.2.4", "192.0.2.2"), + 3: ("192.0.2.3", "192.0.2.1"), +} + + +def _abort_after_native_failure() -> None: + """Terminate a worker whose CUDA stream may contain an unfulfillable wait.""" + os._exit(70) + + +class _NativeConfig(ctypes.Structure): + _fields_ = [ + ("rank", ctypes.c_uint32), + ("peer0", ctypes.c_char_p), + ("peer1", ctypes.c_char_p), + ("device0", ctypes.c_char_p), + ("device1", ctypes.c_char_p), + ("gid0", ctypes.c_uint8), + ("gid1", ctypes.c_uint8), + ("control_port0", ctypes.c_uint16), + ("control_port1", ctypes.c_uint16), + ("payload_bytes", ctypes.c_size_t), + ("graph_submit_cpu_plus_one", ctypes.c_uint32), + ("graph_progress_cpu_plus_one", ctypes.c_uint32), + ] + + +class _NativeConfigV2(ctypes.Structure): + _fields_ = [ + ("struct_size", ctypes.c_uint32), + ("base", _NativeConfig), + ("elements_per_row", ctypes.c_uint32), + ("bytes_per_row", ctypes.c_uint32), + ] + + +class _BidirectionalPrefillConfigV1(ctypes.Structure): + _fields_ = [ + ("struct_size", ctypes.c_uint32), + ("primary", _NativeConfigV2), + ("rail_count", ctypes.c_uint32), + ("query_rows", ctypes.c_uint32), + ("secondary_peer0", ctypes.c_char_p), + ("secondary_peer1", ctypes.c_char_p), + ("secondary_device0", ctypes.c_char_p), + ("secondary_device1", ctypes.c_char_p), + ("secondary_gid0", ctypes.c_uint8), + ("secondary_gid1", ctypes.c_uint8), + ("secondary_control_port0", ctypes.c_uint16), + ("secondary_control_port1", ctypes.c_uint16), + ("timeout_seconds", ctypes.c_uint32), + ] + + +class _NativeGraphStatus(ctypes.Structure): + _fields_ = [ + ("struct_size", ctypes.c_uint32), + ("flags", ctypes.c_uint32), + ("captured_nodes", ctypes.c_uint64), + ("published_sequence", ctypes.c_uint64), + ("consumed_sequence", ctypes.c_uint64), + ("completed_sequence", ctypes.c_uint64), + ("overflow_sequence", ctypes.c_uint64), + ("graph_submit_cpu_plus_one", ctypes.c_uint32), + ("graph_progress_cpu_plus_one", ctypes.c_uint32), + ] + + +class _NativeHealthStatus(ctypes.Structure): + _fields_ = [ + ("struct_size", ctypes.c_uint32), + ("flags", ctypes.c_uint32), + ("submitted_sequence", ctypes.c_uint64), + ("completed_sequence", ctypes.c_uint64), + ("failing_sequence", ctypes.c_uint64), + ("error_code", ctypes.c_int32), + ("failing_stage", ctypes.c_int32), + ("failing_rail", ctypes.c_int32), + ("failing_peer", ctypes.c_int32), + ] + + +@dataclass(frozen=True) +class NativeHealthStatus: + healthy: bool + poisoned: bool + progress_thread_running: bool + stopping: bool + submitted_sequence: int + completed_sequence: int + failing_sequence: int + error_code: int + failing_stage: int + failing_rail: int + failing_peer: int + + +@dataclass(frozen=True) +class GraphReplayStatus: + captured_nodes: int + published_sequence: int + consumed_sequence: int + completed_sequence: int + overflow_sequence: int + capture_configured: bool + polling_enabled: bool + host_native_atomics: bool + submit_affinity_verified: bool + progress_affinity_verified: bool + submit_cpu: int | None + progress_cpu: int | None + two_slot_deferred_ack: bool = False + direct_doorbell: bool = False + graph_kernel_strategy: str = _FUSED_GRAPH_KERNEL + wire_schedule: str = _SEQUENTIAL_WIRE_SCHEDULE + + @property + def replay_advanced(self) -> bool: + """True only after a replay command has reached the native publisher.""" + return self.published_sequence > 0 + + @property + def command_caught_up(self) -> bool: + return ( + self.replay_advanced + and self.published_sequence == self.consumed_sequence + and self.published_sequence == self.completed_sequence + ) + + @property + def replay_caught_up(self) -> bool: + """Compatibility alias for command-ring completion. + + In two-slot mode this proves both outgoing credits were posted, not + that reciprocal credits retired both physical payload slots. + """ + return self.command_caught_up + + @property + def payload_slots_retired(self) -> bool | None: + """Return retirement proof when the selected protocol exposes it.""" + if self.two_slot_deferred_ack: + return None + return self.command_caught_up + + @property + def fatal(self) -> bool: + return self.overflow_sequence != 0 + + def to_dict(self) -> dict[str, object]: + snapshot = asdict(self) + snapshot["command_caught_up"] = self.command_caught_up + snapshot["replay_advanced"] = self.replay_advanced + snapshot["replay_caught_up"] = self.replay_caught_up + snapshot["payload_slots_retired"] = self.payload_slots_retired + snapshot["fatal"] = self.fatal + return snapshot + + +def _mode() -> str: + mode = os.getenv("VLLM_SPARK_TP4_MODE", "").lower() + if mode and mode not in _VALID_MODES: + raise ValueError("VLLM_SPARK_TP4_MODE must be 'shadow', 'custom', 'disabled', or unset") + return mode + + +def _graph_allreduce_protocol() -> str: + value = os.getenv( + "VLLM_SPARK_TP4_GRAPH_ALLREDUCE_PROTOCOL", + _SERIAL_ACK_PROTOCOL, + ) + if value not in _ALLREDUCE_PROTOCOL_WIRE: + raise ValueError( + "VLLM_SPARK_TP4_GRAPH_ALLREDUCE_PROTOCOL must be " + "'serial_ack', 'two_slot_deferred_ack', or unset" + ) + return value + + +def _graph_kernel_strategy() -> str: + value = os.getenv( + "VLLM_SPARK_TP4_GRAPH_KERNEL_STRATEGY", + _FUSED_GRAPH_KERNEL, + ) + if value not in _GRAPH_KERNEL_WIRE: + raise ValueError( + "VLLM_SPARK_TP4_GRAPH_KERNEL_STRATEGY must be " + "'fused', 'split_64k', 'tiered_64k', or unset" + ) + return value + + +def _prefill_q512_enabled() -> bool: + value = os.getenv("VLLM_SPARK_TP4_PREFILL_Q512", "0") + if value not in {"0", "1"}: + raise ValueError( + "VLLM_SPARK_TP4_PREFILL_Q512 must be '0', '1', or unset" + ) + return value == "1" + + +def _bidirectional_prefill_enabled() -> bool: + value = os.getenv("VLLM_SPARK_TP4_BIDIRECTIONAL_PREFILL", "0") + if value not in {"0", "1"}: + raise ValueError( + "VLLM_SPARK_TP4_BIDIRECTIONAL_PREFILL must be '0', '1', or unset" + ) + return value == "1" + + +def _bidirectional_prefill_rail_mode() -> str: + value = os.getenv( + "VLLM_SPARK_TP4_BIDIRECTIONAL_PREFILL_RAIL_MODE", "single" + ).lower() + if value not in {"single", "dual"}: + raise ValueError( + "VLLM_SPARK_TP4_BIDIRECTIONAL_PREFILL_RAIL_MODE must be " + "'single', 'dual', or unset" + ) + return value + + +def _bidirectional_prefill_exposure() -> str: + value = os.getenv( + "VLLM_SPARK_TP4_BIDIRECTIONAL_PREFILL_EXPOSURE", "sync" + ).lower() + if value not in {"sync", "fused"}: + raise ValueError( + "VLLM_SPARK_TP4_BIDIRECTIONAL_PREFILL_EXPOSURE must be " + "'sync', 'fused', or unset" + ) + return value + + +def _validate_bidirectional_prefill_dual_environment() -> None: + if _bidirectional_prefill_rail_mode() != "dual": + return + names = ( + "SPARK_TP4_PEER0", + "SPARK_TP4_PEER1", + "SPARK_TP4_DEVICE0", + "SPARK_TP4_DEVICE1", + "SPARK_TP4_BIDIRECTIONAL_PREFILL_SECONDARY_PEER0", + "SPARK_TP4_BIDIRECTIONAL_PREFILL_SECONDARY_PEER1", + "SPARK_TP4_BIDIRECTIONAL_PREFILL_SECONDARY_DEVICE0", + "SPARK_TP4_BIDIRECTIONAL_PREFILL_SECONDARY_DEVICE1", + ) + values = tuple(os.getenv(name, "").strip() for name in names) + if not all(values): + missing = [name for name, value in zip(names, values) if not value] + raise ValueError( + "dual-rail bidirectional prefill requires: " + + ", ".join(missing) + ) + if len(set(values[:2] + values[4:6])) != 4: + raise ValueError( + "dual-rail primary/secondary peer addresses must be distinct" + ) + if len(set(values[2:4] + values[6:8])) != 4: + raise ValueError( + "dual-rail primary/secondary devices must be distinct" + ) + + +def _bidirectional_prefill_eligible( + communicator: Any, + tensor: Any, + *, + mode: str, + capturing: bool, +) -> bool: + """Rank-invariant admission for the eager width-4096 research lane. + + The decision uses only configuration and tensor properties shared by TP + ranks. Pointer values and rank-local allocator state are deliberately not + consulted, so an ineligible call falls through to NCCL on every rank. + """ + + shape = _tensor_shape(tensor) + return ( + _bidirectional_prefill_enabled() + and mode in {"custom", "shadow"} + and not capturing + and getattr(communicator, "world_size", None) == 4 + and getattr(communicator, "unique_name", "") == "tp:0" + and len(shape) == 2 + and shape[0] in _BIDIRECTIONAL_PREFILL_QUERY_ROWS + and shape[1] == _BIDIRECTIONAL_PREFILL_WIDTH + and str(tensor.dtype) == "torch.bfloat16" + and bool(tensor.is_cuda) + and bool(tensor.is_contiguous()) + ) + + +def _fused_prefill_eligible( + communicator: Any, + tensor: Any, + *, + mode: str, + capturing: bool, +) -> bool: + shape = _tensor_shape(tensor) + return ( + _bidirectional_prefill_exposure() == "fused" + and _bidirectional_prefill_enabled() + and mode in {"custom", "shadow"} + and not capturing + and getattr(communicator, "world_size", None) == 4 + and getattr(communicator, "unique_name", "") == "tp:0" + and len(shape) == 2 + and 128 <= shape[0] <= 8192 + and shape[1] == _BIDIRECTIONAL_PREFILL_WIDTH + and str(tensor.dtype) == "torch.bfloat16" + and bool(tensor.is_cuda) + and bool(tensor.is_contiguous()) + ) + + +def _maximum_allreduce_query_rows() -> int: + return resolve_query_rows()[-1] + + +def _graph_capacity_bytes() -> int: + return _maximum_allreduce_query_rows() * _BYTES_PER_ROW + + +def _admitted_default_width_rows() -> tuple[int, ...]: + """Row set admitted at the default width. + + Admission and reservation share one resolution: + spark_tp4_query_row_provider owns the row policy for both this gate + and the port namespace, so the two can never disagree. Ambient + (process-environment) resolutions are cached, keeping this hot-path + lookup allocation-free. + """ + return resolve_query_rows() + + +def _target_shape_eligible(shape: tuple[int, ...]) -> bool: + return ( + len(shape) == 2 + and shape[1] == _TARGET_WIDTH + and shape[0] in _admitted_default_width_rows() + ) + + +def _shadow_result_passed( + outside_tolerance: int, + nonfinite_mismatches: int, + max_ulp: int, +) -> bool: + """Apply numerical gates without making near-zero BF16 ULP diagnostic fatal.""" + + max_allowed_ulp = int(os.getenv("SPARK_TP4_SHADOW_MAX_ULP", "0")) + if max_allowed_ulp < 0: + raise ValueError("SPARK_TP4_SHADOW_MAX_ULP must be nonnegative") + ulp_failed = max_allowed_ulp > 0 and max_ulp > max_allowed_ulp + return not (outside_tolerance or nonfinite_mismatches or ulp_failed) + + +def _tensor_shape(tensor: Any) -> tuple[int, ...]: + return tuple(int(size) for size in tensor.shape) + + +def _eager_admitted_widths() -> tuple[int, ...]: + """Admitted eager all-reduce widths, read from the port namespace.""" + return eager_allreduce_admitted_widths() + + +def _eager_shape_eligible(shape: tuple[int, ...]) -> bool: + """Eager admission across widths. + + The default width follows the same row set as the graph gate, + including the sparse provider contract when enabled. Non-default + widths use the contiguous row range that the port namespace + enumerates for extension payloads; the sparse contract is a + default-width serving constraint and does not govern them. + """ + if len(shape) != 2 or shape[1] not in _eager_admitted_widths(): + return False + if shape[1] == _TARGET_WIDTH: + return shape[0] in _admitted_default_width_rows() + extension_maximum = ( + _ALLREDUCE_PREFILL_MAX_QUERY_ROWS + if _prefill_q512_enabled() + else MAX_QUERY_ROWS + ) + return 1 <= shape[0] <= extension_maximum + + +def _eligible(communicator: Any, tensor: Any, mode: str | None = None) -> bool: + active_mode = _mode() if mode is None else mode + return ( + active_mode in _VALID_MODES + and getattr(communicator, "world_size", None) == 4 + and getattr(communicator, "unique_name", "") == "tp:0" + and _eager_shape_eligible(_tensor_shape(tensor)) + and str(tensor.dtype) == "torch.bfloat16" + and bool(tensor.is_cuda) + and bool(tensor.is_contiguous()) + ) + + +def _is_stream_capturing(torch_module: Any) -> bool: + checker = getattr(torch_module.cuda, "is_current_stream_capturing", None) + return bool(checker is not None and checker()) + + +def _record_stock_path( + *, + capturing: bool, + reason: str, + communicator: Any | None = None, + tensor: Any | None = None, +) -> None: + from spark_collective_audit import ( + StockCollectiveSignature, + enabled, + record_stock, + ) + + signature = None + if enabled() and communicator is not None and tensor is not None: + world_size = getattr(communicator, "world_size", None) + signature = StockCollectiveSignature( + shape=_tensor_shape(tensor), + dtype=str(tensor.dtype), + is_cuda=bool(tensor.is_cuda), + contiguous=bool(tensor.is_contiguous()), + world_size=( + None if world_size is None else int(world_size) + ), + unique_name=str(getattr(communicator, "unique_name", "")), + ) + if enabled(): + # Invariant: an armed audit always produces a status file, + # regardless of model width or graph eligibility. The recording + # path must start the reporter itself because no graph-session + # preparation is guaranteed to run in the serving profile. + from spark_graph_status_reporter import ensure_status_reporter + + reporter_rank = getattr(communicator, "rank_in_group", None) + if reporter_rank is None: + reporter_rank = 0 + ensure_status_reporter(rank=int(reporter_rank)) + record_stock( + "all_reduce", + capturing=capturing, + reason=reason, + signature=signature, + ) + + +def _research_graph_all_reduce( + communicator: Any, tensor: Any, capturing: bool +) -> Any | None: + """Route research-admitted width-4096 collectives. + + Returns the collective output when this path handled the call, or + None to fall through to the unchanged default dispatch. Eager calls + are never handled here: the eager pass only prepares the session so + it exists before vLLM's first capture. + """ + if ( + getattr(communicator, "world_size", None) != 4 + or getattr(communicator, "unique_name", "") != "tp:0" + ): + return None + if os.getenv("SPARK_TP4_CAPABILITY_VOTE", "0") == "1": + ensure_capability_vote(communicator) + backend = getattr(communicator, "_spark_tp4_native", None) + if backend is None: + backend = _Backend(int(communicator.rank_in_group)) + communicator._spark_tp4_native = backend + if not capturing: + if backend.graph_width4096_session is None: + backend.prepare_graph_width4096() + return None + shape = _tensor_shape(tensor) + if ( + not _research_graph_shape_eligible(shape) + or str(tensor.dtype) != "torch.bfloat16" + or not bool(tensor.is_cuda) + or not bool(tensor.is_contiguous()) + ): + return None + session = backend.graph_width4096_session + if session is None: + # Fail closed: an admitted capture-phase collective must never + # silently take the stock path in the research profile. + logger.error( + "research width-4096 graph capture reached before session " + "preparation; terminating worker" + ) + _abort_after_native_failure() + raise AssertionError("unreachable after worker termination") + try: + output = session.capture(tensor) + _record_graph_event(communicator, "width4096_captured_nodes") + return output + except BaseException: + logger.exception( + "fatal Spark TP4 width-4096 graph-capture error; terminating " + "worker because a partially captured native graph cannot " + "safely fall back" + ) + _abort_after_native_failure() + raise AssertionError("unreachable after worker termination") + + +def _payload_bytes(tensor: Any) -> int: + rows, width = _tensor_shape(tensor) + return rows * width * _BF16_BYTES + + +def _collective_signature(tensor: Any) -> tuple[int, tuple[int, ...], str]: + return ( + _payload_bytes(tensor), + _tensor_shape(tensor), + str(tensor.dtype), + ) + + +def _control_ports(payload_bytes: int) -> tuple[int, int]: + return eager_allreduce_ports_for_payload(payload_bytes) + + +def _graph_q1_enabled() -> bool: + value = os.getenv("VLLM_SPARK_TP4_GRAPH_Q1", "0") + if value not in {"0", "1"}: + raise ValueError("VLLM_SPARK_TP4_GRAPH_Q1 must be '0' or '1'") + return value == "1" + + +def _graph_width4096_research_enabled() -> bool: + """Research-only graph admission for BF16 [Q <= 512, 4096]. + + Explicitly opt-in and unqualified: enabling it routes captured + width-4096 all-reduces through one maximum-capacity native graph + session (sequential tiered_64k, two-slot deferred ACK) instead of + the stock path. Default behavior is byte-identical when unset. + """ + value = os.getenv("VLLM_SPARK_TP4_GRAPH_WIDTH4096_RESEARCH", "0") + if value not in {"0", "1"}: + raise ValueError( + "VLLM_SPARK_TP4_GRAPH_WIDTH4096_RESEARCH must be '0' or '1'" + ) + return value == "1" + + +def _research_graph_shape_eligible(shape: tuple[int, ...]) -> bool: + return ( + len(shape) == 2 + and shape[1] == _RESEARCH_GRAPH_WIDTH + and 1 <= shape[0] <= _RESEARCH_GRAPH_MAX_QUERY_ROWS + ) + + +def _graph_dual_port_q40_enabled() -> bool: + value = os.getenv("VLLM_SPARK_TP4_GRAPH_DUAL_PORT_Q40", "0") + if value not in {"0", "1"}: + raise ValueError( + "VLLM_SPARK_TP4_GRAPH_DUAL_PORT_Q40 must be '0' or '1'" + ) + return value == "1" + + +def _persistent_output_slots() -> int: + text = os.getenv("SPARK_TP4_PERSISTENT_OUTPUT_SLOTS", "0") + try: + slots = int(text) + except ValueError as error: + raise ValueError( + "SPARK_TP4_PERSISTENT_OUTPUT_SLOTS must be an integer" + ) from error + if slots < 0 or slots > _MAX_PERSISTENT_OUTPUT_SLOTS: + raise ValueError( + "SPARK_TP4_PERSISTENT_OUTPUT_SLOTS must be in " + f"[0, {_MAX_PERSISTENT_OUTPUT_SLOTS}]" + ) + return slots + + +def _graph_control_ports() -> tuple[int, int]: + return graph_allreduce_control_ports() + + +def _graph_dual_port_q40_control_ports() -> tuple[int, int]: + return graph_dual_port_q40_control_ports() + + +def _fixed_kv_cache_bytes(argv: list[str] | None = None) -> int: + arguments = sys.argv if argv is None else argv + values: list[str] = [] + for index, argument in enumerate(arguments): + if argument == "--kv-cache-memory-bytes": + if index + 1 >= len(arguments): + raise RuntimeError( + "graph TP4 requires a value for --kv-cache-memory-bytes" + ) + values.append(arguments[index + 1]) + elif argument.startswith("--kv-cache-memory-bytes="): + values.append(argument.partition("=")[2]) + if len(values) != 1: + raise RuntimeError( + "graph TP4 requires exactly one positive " + "--kv-cache-memory-bytes value so vLLM skips its throwaway " + "memory-profiling capture stream" + ) + try: + parsed = int(values[0]) + except ValueError as error: + raise RuntimeError( + "graph TP4 --kv-cache-memory-bytes must be a positive integer" + ) from error + if parsed <= 0: + raise RuntimeError( + "graph TP4 --kv-cache-memory-bytes must be a positive integer" + ) + return parsed + + +def _graph_cpu_affinity() -> tuple[int, int]: + names = ( + "SPARK_TP4_GRAPH_SUBMIT_CPU", + "SPARK_TP4_GRAPH_PROGRESS_CPU", + ) + values: list[int] = [] + for name in names: + text = os.getenv(name) + if text is None: + raise RuntimeError( + f"graph TP4 requires explicit {name}; submission and " + "transport progress must use distinct CPUs" + ) + try: + value = int(text) + except ValueError as error: + raise RuntimeError(f"graph TP4 {name} must be a nonnegative integer") from error + if value < 0: + raise RuntimeError(f"graph TP4 {name} must be a nonnegative integer") + values.append(value) + submit_cpu, progress_cpu = values + if submit_cpu == progress_cpu: + raise RuntimeError( + "graph TP4 submit/progress CPUs must be distinct" + ) + return submit_cpu, progress_cpu + + +def _graph_preflight(argv: list[str] | None = None) -> tuple[int, int]: + _fixed_kv_cache_bytes(argv) + if os.getenv("VLLM_SPARK_SHARED_CAPTURE_STREAM", "0") != "1": + raise RuntimeError( + "graph TP4 requires VLLM_SPARK_SHARED_CAPTURE_STREAM=1" + ) + if not sys.platform.startswith("linux"): + raise RuntimeError( + "graph TP4 CPU affinity requires Linux pthread affinity support" + ) + return _graph_cpu_affinity() + + +def _graph_dual_port_q40_cpu_affinity( + argv: list[str] | None = None, +) -> tuple[int, int]: + submit_cpu, tp_progress_cpu = _graph_preflight(argv) + name = "SPARK_TP4_GRAPH_DUAL_PORT_Q40_PROGRESS_CPU" + text = os.getenv(name) + if text is None: + raise RuntimeError( + f"exact-Q40 dual-port graph TP4 requires explicit {name}" + ) + try: + progress_cpu = int(text) + except ValueError as error: + raise RuntimeError(f"{name} must be a nonnegative integer") from error + if progress_cpu < 0: + raise RuntimeError(f"{name} must be a nonnegative integer") + + occupied = { + submit_cpu: "shared graph submit", + tp_progress_cpu: "sequential TP graph progress", + } + if os.getenv("VLLM_SPARK_TP4_VOCAB_MODE", "") == "custom": + occupied[ + int(os.getenv("SPARK_TP4_GRAPH_VOCAB_PROGRESS_CPU", "12")) + ] = "vocabulary graph progress" + if ( + os.getenv("VLLM_SPARK_TP4_DCP_GRAPH_CUSTOM", "0") == "1" + or os.getenv("VLLM_SPARK_TP4_DCP_GRAPH_SHADOW", "0") == "1" + ): + occupied[ + int(os.getenv("SPARK_TP4_GRAPH_DCP_PROGRESS_CPU", "13")) + ] = "DCP graph progress" + if os.getenv("VLLM_SPARK_TP4_INDEXER_GRAPH_CUSTOM", "0") == "1": + occupied[ + int(os.getenv("SPARK_TP4_GRAPH_INDEXER_PROGRESS_CPU", "14")) + ] = "indexer graph progress" + if progress_cpu in occupied: + raise RuntimeError( + f"{name}={progress_cpu} collides with " + f"{occupied[progress_cpu]} CPU ownership" + ) + return submit_cpu, progress_cpu + + +def _record_graph_event(communicator: Any, event: str) -> int: + attribute = f"_spark_tp4_graph_q1_{event}" + count = int(getattr(communicator, attribute, 0)) + 1 + setattr(communicator, attribute, count) + _graph_event_counts[event] = _graph_event_counts.get(event, 0) + 1 + if count == 1 or count % 128 == 0: + logger.warning( + "Spark TP4 mixed-Q graph %s on rank %d: count=%d", + event, + int(communicator.rank_in_group), + count, + ) + return count + + +def _bind_bidirectional_prefill_native_api( + library: Any, +) -> tuple[Any, Any, Any, Any]: + """Bind the isolated bidirectional-prefill research ABI.""" + + try: + create = getattr( + library, _BIDIRECTIONAL_PREFILL_C_ABI_SYMBOLS["create"] + ) + all_reduce = getattr( + library, _BIDIRECTIONAL_PREFILL_C_ABI_SYMBOLS["all_reduce"] + ) + health = getattr(library, _BIDIRECTIONAL_PREFILL_C_ABI_SYMBOLS["health"]) + destroy = getattr( + library, _BIDIRECTIONAL_PREFILL_C_ABI_SYMBOLS["destroy"] + ) + except AttributeError as error: + raise RuntimeError( + "Spark TP4 native library lacks the bidirectional prefill " + "research ABI; rebuild and deploy a matching native library" + ) from error + create.argtypes = [ + ctypes.POINTER(_BidirectionalPrefillConfigV1), + ctypes.c_char_p, + ctypes.c_size_t, + ] + create.restype = ctypes.c_void_p + all_reduce.argtypes = [ + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_char_p, + ctypes.c_size_t, + ] + all_reduce.restype = ctypes.c_int + health.argtypes = [ + ctypes.c_void_p, + ctypes.POINTER(_NativeHealthStatus), + ctypes.c_size_t, + ctypes.c_char_p, + ctypes.c_size_t, + ] + health.restype = ctypes.c_int + destroy.argtypes = [ctypes.c_void_p] + destroy.restype = None + return create, all_reduce, health, destroy + + +def _bind_fused_prefill_native_api(library: Any) -> tuple[Any, Any, Any, Any]: + """Bind the caller-stream fused Q8192 ABI.""" + + try: + create = getattr(library, _FUSED_PREFILL_C_ABI_SYMBOLS["create"]) + all_reduce = getattr( + library, _FUSED_PREFILL_C_ABI_SYMBOLS["all_reduce"] + ) + health = getattr(library, _FUSED_PREFILL_C_ABI_SYMBOLS["health"]) + destroy = getattr(library, _FUSED_PREFILL_C_ABI_SYMBOLS["destroy"]) + except AttributeError as error: + raise RuntimeError( + "Spark TP4 native library lacks the fused prefill ABI; " + "rebuild and deploy a matching native library" + ) from error + create.argtypes = [ + ctypes.POINTER(_BidirectionalPrefillConfigV1), + ctypes.c_char_p, + ctypes.c_size_t, + ] + create.restype = ctypes.c_void_p + all_reduce.argtypes = [ + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_uint32, + ctypes.c_void_p, + ctypes.c_char_p, + ctypes.c_size_t, + ] + all_reduce.restype = ctypes.c_int + health.argtypes = [ + ctypes.c_void_p, + ctypes.POINTER(_NativeHealthStatus), + ctypes.c_size_t, + ctypes.c_char_p, + ctypes.c_size_t, + ] + health.restype = ctypes.c_int + destroy.argtypes = [ctypes.c_void_p] + destroy.restype = None + return create, all_reduce, health, destroy + + +def _validate_fused_prefill_native_api() -> None: + library_path = os.getenv("SPARK_TP4_LIBRARY", "").strip() + if not library_path: + raise ValueError( + "fused prefill requires SPARK_TP4_LIBRARY before installation" + ) + _bind_fused_prefill_native_api(ctypes.CDLL(library_path)) + + +class _BidirectionalPrefillNativeSession: + """Lazy one-shape binding for the research-only eager prefill C ABI.""" + + _valid_query_rows = _BIDIRECTIONAL_PREFILL_QUERY_ROWS + _bind_native_api = staticmethod(_bind_bidirectional_prefill_native_api) + _session_label = "bidirectional prefill" + _exposure = "sync" + + def __init__(self, rank: int, shape: tuple[int, ...]) -> None: + if ( + rank not in _DEFAULT_PEERS + or len(shape) != 2 + or shape[0] not in self._valid_query_rows + or shape[1] != _BIDIRECTIONAL_PREFILL_WIDTH + ): + raise ValueError("invalid bidirectional prefill session shape") + query_rows = shape[0] + payload_bytes = query_rows * _BIDIRECTIONAL_PREFILL_ROW_BYTES + control_port0, control_port1 = bidirectional_prefill_control_ports( + query_rows + ) + default_peer0, default_peer1 = _DEFAULT_PEERS[rank] + peer0 = os.getenv("SPARK_TP4_PEER0", default_peer0) + peer1 = os.getenv("SPARK_TP4_PEER1", default_peer1) + device0 = os.getenv("SPARK_TP4_DEVICE0", "rocep1s0f0") + device1 = os.getenv("SPARK_TP4_DEVICE1", "rocep1s0f1") + base = _NativeConfig( + rank=rank, + peer0=peer0.encode(), + peer1=peer1.encode(), + device0=device0.encode(), + device1=device1.encode(), + gid0=int(os.getenv("SPARK_TP4_GID0", "3")), + gid1=int(os.getenv("SPARK_TP4_GID1", "3")), + control_port0=control_port0, + control_port1=control_port1, + payload_bytes=payload_bytes, + graph_submit_cpu_plus_one=0, + graph_progress_cpu_plus_one=0, + ) + primary = _NativeConfigV2( + struct_size=ctypes.sizeof(_NativeConfigV2), + base=base, + elements_per_row=_BIDIRECTIONAL_PREFILL_WIDTH, + bytes_per_row=_BIDIRECTIONAL_PREFILL_ROW_BYTES, + ) + rail_mode = _bidirectional_prefill_rail_mode() + rail_count = 2 if rail_mode == "dual" else 1 + secondary_peer0 = secondary_peer1 = None + secondary_device0 = secondary_device1 = None + secondary_gid0 = secondary_gid1 = 0 + secondary_port0 = secondary_port1 = 0 + if rail_count == 2: + names = ( + "SPARK_TP4_BIDIRECTIONAL_PREFILL_SECONDARY_PEER0", + "SPARK_TP4_BIDIRECTIONAL_PREFILL_SECONDARY_PEER1", + "SPARK_TP4_BIDIRECTIONAL_PREFILL_SECONDARY_DEVICE0", + "SPARK_TP4_BIDIRECTIONAL_PREFILL_SECONDARY_DEVICE1", + ) + values = tuple(os.getenv(name, "").strip() for name in names) + if not all(values): + missing = [name for name, value in zip(names, values) if not value] + raise ValueError( + "dual-rail bidirectional prefill requires: " + + ", ".join(missing) + ) + secondary_peer0, secondary_peer1, secondary_device0, secondary_device1 = values + if len({peer0, peer1, secondary_peer0, secondary_peer1}) != 4: + raise ValueError( + "dual-rail primary/secondary peer addresses must be distinct" + ) + if len({device0, device1, secondary_device0, secondary_device1}) != 4: + raise ValueError( + "dual-rail primary/secondary devices must be distinct" + ) + try: + secondary_gid0 = int(os.getenv( + "SPARK_TP4_BIDIRECTIONAL_PREFILL_SECONDARY_GID0", "3" + )) + secondary_gid1 = int(os.getenv( + "SPARK_TP4_BIDIRECTIONAL_PREFILL_SECONDARY_GID1", "3" + )) + except ValueError as error: + raise ValueError("dual-rail secondary GIDs must be integers") from error + if not (0 <= secondary_gid0 <= 255 and 0 <= secondary_gid1 <= 255): + raise ValueError("dual-rail secondary GIDs must be in [0, 255]") + secondary_port0, secondary_port1 = ( + bidirectional_prefill_secondary_control_ports(query_rows) + ) + try: + timeout_seconds = int(os.getenv( + "SPARK_TP4_BIDIRECTIONAL_PREFILL_TIMEOUT_SECONDS", "120" + )) + except ValueError as error: + raise ValueError("bidirectional prefill timeout must be an integer") from error + if not (1 <= timeout_seconds <= 0xFFFFFFFF): + raise ValueError("bidirectional prefill timeout must be in [1, 4294967295]") + config = _BidirectionalPrefillConfigV1( + struct_size=ctypes.sizeof(_BidirectionalPrefillConfigV1), + primary=primary, + rail_count=rail_count, + query_rows=query_rows, + secondary_peer0=(None if secondary_peer0 is None else secondary_peer0.encode()), + secondary_peer1=(None if secondary_peer1 is None else secondary_peer1.encode()), + secondary_device0=(None if secondary_device0 is None else secondary_device0.encode()), + secondary_device1=(None if secondary_device1 is None else secondary_device1.encode()), + secondary_gid0=secondary_gid0, + secondary_gid1=secondary_gid1, + secondary_control_port0=secondary_port0, + secondary_control_port1=secondary_port1, + timeout_seconds=timeout_seconds, + ) + self._library = ctypes.CDLL(os.environ["SPARK_TP4_LIBRARY"]) + create, self._all_reduce, self._health, self._destroy = ( + self._bind_native_api(self._library) + ) + error = ctypes.create_string_buffer(512) + self._handle = create(ctypes.byref(config), error, len(error)) + if not self._handle: + message = error.value.decode(errors="replace") + raise RuntimeError( + f"failed to create Spark TP4 {self._session_label} session: " + f"{message}" + ) + self.rank = rank + self.shape = shape + self.payload_bytes = payload_bytes + logger.warning( + "Spark TP4 %s session ready: rank=%d shape=%s rails=%d " + "exposure=%s primary_ports=%s secondary_ports=%s", + self._session_label, + rank, + shape, + rail_count, + self._exposure, + (control_port0, control_port1), + ( + None + if rail_count == 1 + else (secondary_port0, secondary_port1) + ), + ) + + def all_reduce(self, tensor: Any) -> Any: + import torch + + if _tensor_shape(tensor) != self.shape: + raise ValueError("bidirectional prefill tensor shape changed") + output = torch.empty_like(tensor) + stream = torch.cuda.current_stream(device=tensor.device) + error = ctypes.create_string_buffer(512) + result = self._all_reduce( + self._handle, + ctypes.c_void_p(tensor.data_ptr()), + ctypes.c_void_p(output.data_ptr()), + ctypes.c_void_p(stream.cuda_stream), + error, + len(error), + ) + if result != 0: + message = error.value.decode(errors="replace") + raise RuntimeError( + f"Spark TP4 {self._session_label} all-reduce failed: {message}" + ) + return output + + def health_status(self) -> NativeHealthStatus: + native = _NativeHealthStatus() + error = ctypes.create_string_buffer(512) + result = self._health( + self._handle, + ctypes.byref(native), + ctypes.sizeof(native), + error, + len(error), + ) + if result != 0: + message = error.value.decode(errors="replace") + raise RuntimeError( + f"Spark TP4 {self._session_label} health failed: {message}" + ) + if native.struct_size != ctypes.sizeof(_NativeHealthStatus): + raise RuntimeError( + f"Spark TP4 {self._session_label} health status ABI mismatch: " + f"native={native.struct_size} python=" + f"{ctypes.sizeof(_NativeHealthStatus)}" + ) + flags = int(native.flags) + return NativeHealthStatus( + healthy=bool(flags & _HEALTHY), + poisoned=bool(flags & _HEALTH_POISONED), + progress_thread_running=bool( + flags & _HEALTH_PROGRESS_THREAD_RUNNING + ), + stopping=bool(flags & _HEALTH_STOPPING), + submitted_sequence=int(native.submitted_sequence), + completed_sequence=int(native.completed_sequence), + failing_sequence=int(native.failing_sequence), + error_code=int(native.error_code), + failing_stage=int(native.failing_stage), + failing_rail=int(native.failing_rail), + failing_peer=int(native.failing_peer), + ) + + def close(self) -> None: + if self._handle: + self._destroy(self._handle) + self._handle = None + + def __del__(self) -> None: # pragma: no cover - defensive teardown + try: + self.close() + except Exception: + pass + + +class _FusedPrefillNativeSession(_BidirectionalPrefillNativeSession): + """One-launch caller-stream fused session with Q8192 capacity.""" + + _valid_query_rows = frozenset({8192}) + _bind_native_api = staticmethod(_bind_fused_prefill_native_api) + _session_label = "fused prefill" + _exposure = "fused" + + def __init__(self, rank: int, shape: tuple[int, ...]) -> None: + if _bidirectional_prefill_rail_mode() != "dual": + raise ValueError("fused prefill requires strict dual-rail mode") + super().__init__(rank, shape) + + def all_reduce(self, tensor: Any) -> Any: + import torch + + shape = _tensor_shape(tensor) + if ( + len(shape) != 2 + or not 128 <= shape[0] <= 8192 + or shape[1] != _BIDIRECTIONAL_PREFILL_WIDTH + ): + raise ValueError("fused prefill tensor exceeds session capacity") + output = torch.empty_like(tensor) + stream = torch.cuda.current_stream(device=tensor.device) + error = ctypes.create_string_buffer(512) + result = self._all_reduce( + self._handle, + ctypes.c_void_p(tensor.data_ptr()), + ctypes.c_void_p(output.data_ptr()), + ctypes.c_uint32(shape[0]), + ctypes.c_void_p(stream.cuda_stream), + error, + len(error), + ) + if result != 0: + message = error.value.decode(errors="replace") + raise RuntimeError( + f"Spark TP4 fused prefill all-reduce failed: {message}" + ) + return output + +class _NativeSession: + def __init__( + self, + rank: int, + payload_bytes: int, + *, + control_ports: tuple[int, int] | None = None, + graph_only: bool = False, + graph_cpu_affinity: tuple[int, int] | None = None, + allreduce_protocol: str = _SERIAL_ACK_PROTOCOL, + graph_kernel_strategy: str = _FUSED_GRAPH_KERNEL, + wire_schedule: str = _SEQUENTIAL_WIRE_SCHEDULE, + row_bytes: int = _BYTES_PER_ROW, + ) -> None: + if rank not in _DEFAULT_PEERS: + raise ValueError(f"TP4 rank must be in [0, 3], got {rank}") + if row_bytes != _BYTES_PER_ROW: + if not graph_only: + raise ValueError( + "non-default TP4 row geometry requires a graph-only " + "session" + ) + if row_bytes != _RESEARCH_GRAPH_ROW_BYTES: + raise ValueError( + "TP4 row geometry must be the default width or the " + f"research width ({_RESEARCH_GRAPH_ROW_BYTES} bytes " + "per row)" + ) + if payload_bytes != _RESEARCH_GRAPH_CAPACITY_BYTES: + raise ValueError( + "research row-geometry graph session requires " + f"Q{_RESEARCH_GRAPH_MAX_QUERY_ROWS} capacity" + ) + expected_graph_capacity = _graph_capacity_bytes() + if ( + graph_only + and row_bytes == _BYTES_PER_ROW + and wire_schedule == _SEQUENTIAL_WIRE_SCHEDULE + and payload_bytes != expected_graph_capacity + ): + raise ValueError( + "Spark TP4 graph session requires " + f"Q{expected_graph_capacity // _BYTES_PER_ROW} capacity" + ) + if graph_only != (graph_cpu_affinity is not None): + raise ValueError( + "Spark TP4 graph session requires an explicit graph CPU " + "affinity pair; eager sessions cannot set one" + ) + if allreduce_protocol not in _ALLREDUCE_PROTOCOL_WIRE: + raise ValueError("invalid Spark TP4 all-reduce protocol") + if graph_kernel_strategy not in _GRAPH_KERNEL_WIRE: + raise ValueError("invalid Spark TP4 graph kernel strategy") + if wire_schedule not in _WIRE_SCHEDULE_WIRE: + raise ValueError("invalid Spark TP4 wire schedule") + if graph_kernel_strategy != _FUSED_GRAPH_KERNEL and not graph_only: + raise ValueError( + "research graph kernel strategies require a graph-only session" + ) + if ( + allreduce_protocol == _TWO_SLOT_DEFERRED_ACK_PROTOCOL + and not graph_only + ): + raise ValueError( + "two_slot_deferred_ack requires a graph-only session" + ) + if wire_schedule == _DUAL_PORT_STRIPED_WIRE_SCHEDULE: + if not graph_only: + raise ValueError( + "dual_port_striped requires a graph-only session" + ) + if allreduce_protocol != _TWO_SLOT_DEFERRED_ACK_PROTOCOL: + raise ValueError( + "dual_port_striped requires two_slot_deferred_ack" + ) + if graph_kernel_strategy != _FUSED_GRAPH_KERNEL: + raise ValueError( + "dual_port_striped requires the fused graph kernel" + ) + if payload_bytes not in { + _DUAL_PORT_Q40_CAPACITY_BYTES, + _ALLREDUCE_PREFILL_CAPACITY_BYTES, + }: + raise ValueError( + "dual_port_striped requires Q40 or Q512 capacity" + ) + if control_ports is None: + control_port0, control_port1 = _control_ports(payload_bytes) + else: + control_port0, control_port1 = validate_control_port_pair( + control_ports, + owner=( + "graph all-reduce" if graph_only else "eager all-reduce" + ), + ) + validate_active_port_namespace() + library_path = os.environ["SPARK_TP4_LIBRARY"] + self._library = ctypes.CDLL(library_path) + self._library.spark_tp4_create.argtypes = [ + ctypes.POINTER(_NativeConfig), + ctypes.c_char_p, + ctypes.c_size_t, + ] + self._library.spark_tp4_create.restype = ctypes.c_void_p + create_session = self._library.spark_tp4_create + use_v2_geometry = row_bytes != _BYTES_PER_ROW + if use_v2_geometry: + try: + create_session = self._library.spark_tp4_create_v2 + except AttributeError as error: + raise RuntimeError( + "Spark TP4 native library lacks the v2 row-geometry " + "ABI; rebuild and deploy a matching library" + ) from error + create_session.argtypes = [ + ctypes.POINTER(_NativeConfigV2), + ctypes.c_uint32, + ctypes.c_uint32, + ctypes.c_uint32, + ctypes.c_char_p, + ctypes.c_size_t, + ] + create_session.restype = ctypes.c_void_p + elif wire_schedule != _SEQUENTIAL_WIRE_SCHEDULE: + try: + create_session = ( + self._library + .spark_tp4_create_with_protocol_graph_kernel_and_schedule + ) + except AttributeError as error: + raise RuntimeError( + "Spark TP4 native library lacks the wire-schedule ABI; " + "rebuild and deploy a matching library" + ) from error + create_session.argtypes = [ + ctypes.POINTER(_NativeConfig), + ctypes.c_uint32, + ctypes.c_uint32, + ctypes.c_uint32, + ctypes.c_char_p, + ctypes.c_size_t, + ] + create_session.restype = ctypes.c_void_p + elif graph_kernel_strategy != _FUSED_GRAPH_KERNEL: + try: + create_session = ( + self._library.spark_tp4_create_with_protocol_and_graph_kernel + ) + except AttributeError as error: + raise RuntimeError( + "Spark TP4 native library lacks the graph-kernel ABI; " + "rebuild and deploy a matching library" + ) from error + create_session.argtypes = [ + ctypes.POINTER(_NativeConfig), + ctypes.c_uint32, + ctypes.c_uint32, + ctypes.c_char_p, + ctypes.c_size_t, + ] + create_session.restype = ctypes.c_void_p + elif allreduce_protocol == _TWO_SLOT_DEFERRED_ACK_PROTOCOL: + try: + create_session = ( + self._library.spark_tp4_create_with_protocol + ) + except AttributeError as error: + raise RuntimeError( + "Spark TP4 native library lacks the deferred-ACK ABI; " + "rebuild and deploy a matching library" + ) from error + create_session.argtypes = [ + ctypes.POINTER(_NativeConfig), + ctypes.c_uint32, + ctypes.c_char_p, + ctypes.c_size_t, + ] + create_session.restype = ctypes.c_void_p + self._library.spark_tp4_all_reduce.argtypes = [ + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_char_p, + ctypes.c_size_t, + ] + self._library.spark_tp4_all_reduce.restype = ctypes.c_int + if graph_only: + self._library.spark_tp4_capture_all_reduce.argtypes = [ + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_uint32, + ctypes.c_void_p, + ctypes.c_char_p, + ctypes.c_size_t, + ] + self._library.spark_tp4_capture_all_reduce.restype = ctypes.c_int + self._library.spark_tp4_get_graph_status.argtypes = [ + ctypes.c_void_p, + ctypes.POINTER(_NativeGraphStatus), + ctypes.c_size_t, + ctypes.c_char_p, + ctypes.c_size_t, + ] + self._library.spark_tp4_get_graph_status.restype = ctypes.c_int + self._library.spark_tp4_get_health_status.argtypes = [ + ctypes.c_void_p, + ctypes.POINTER(_NativeHealthStatus), + ctypes.c_size_t, + ctypes.c_char_p, + ctypes.c_size_t, + ] + self._library.spark_tp4_get_health_status.restype = ctypes.c_int + self._library.spark_tp4_destroy.argtypes = [ctypes.c_void_p] + self._library.spark_tp4_destroy.restype = None + self._graph_only = graph_only + self._allreduce_protocol = allreduce_protocol + self._graph_kernel_strategy = graph_kernel_strategy + self._wire_schedule = wire_schedule + self._row_bytes = row_bytes + self._elements_per_row = row_bytes // _BF16_BYTES + self._graph_max_query_rows = ( + payload_bytes // row_bytes if graph_only else 0 + ) + persistent_slots = 0 if graph_only else _persistent_output_slots() + self._persistent_outputs = ( + PersistentOutputRing(persistent_slots) + if persistent_slots + else None + ) + if persistent_slots: + logger.warning( + "Spark TP4 eager persistent output ring enabled on rank %d " + "for %d bytes with %d slots", + rank, + payload_bytes, + persistent_slots, + ) + + default_peer0, default_peer1 = _DEFAULT_PEERS[rank] + submit_cpu, progress_cpu = graph_cpu_affinity or (-1, -1) + config = _NativeConfig( + rank=rank, + peer0=os.getenv("SPARK_TP4_PEER0", default_peer0).encode(), + peer1=os.getenv("SPARK_TP4_PEER1", default_peer1).encode(), + device0=os.getenv("SPARK_TP4_DEVICE0", "rocep1s0f0").encode(), + device1=os.getenv("SPARK_TP4_DEVICE1", "rocep1s0f1").encode(), + gid0=int(os.getenv("SPARK_TP4_GID0", "3")), + gid1=int(os.getenv("SPARK_TP4_GID1", "3")), + control_port0=control_port0, + control_port1=control_port1, + payload_bytes=payload_bytes, + graph_submit_cpu_plus_one=submit_cpu + 1, + graph_progress_cpu_plus_one=progress_cpu + 1, + ) + error = ctypes.create_string_buffer(512) + if use_v2_geometry: + config_v2 = _NativeConfigV2( + struct_size=ctypes.sizeof(_NativeConfigV2), + base=config, + elements_per_row=self._elements_per_row, + bytes_per_row=row_bytes, + ) + self._handle = create_session( + ctypes.byref(config_v2), + _ALLREDUCE_PROTOCOL_WIRE[allreduce_protocol], + _GRAPH_KERNEL_WIRE[graph_kernel_strategy], + _WIRE_SCHEDULE_WIRE[wire_schedule], + error, + len(error), + ) + elif wire_schedule != _SEQUENTIAL_WIRE_SCHEDULE: + self._handle = create_session( + ctypes.byref(config), + _ALLREDUCE_PROTOCOL_WIRE[allreduce_protocol], + _GRAPH_KERNEL_WIRE[graph_kernel_strategy], + _WIRE_SCHEDULE_WIRE[wire_schedule], + error, + len(error), + ) + elif graph_kernel_strategy != _FUSED_GRAPH_KERNEL: + self._handle = create_session( + ctypes.byref(config), + _ALLREDUCE_PROTOCOL_WIRE[allreduce_protocol], + _GRAPH_KERNEL_WIRE[graph_kernel_strategy], + error, + len(error), + ) + elif allreduce_protocol == _SERIAL_ACK_PROTOCOL: + self._handle = create_session( + ctypes.byref(config), error, len(error) + ) + else: + self._handle = create_session( + ctypes.byref(config), + _ALLREDUCE_PROTOCOL_WIRE[allreduce_protocol], + error, + len(error), + ) + if not self._handle: + message = error.value.decode(errors="replace") + raise RuntimeError(f"failed to create Spark TP4 session: {message}") + logger.warning( + "Spark TP4 %s session ready on rank %d for %d bytes " + "(control ports %d/%d, protocol %s, graph kernel %s, " + "wire schedule %s)", + "graph-only" if graph_only else "eager", + rank, + payload_bytes, + control_port0, + control_port1, + allreduce_protocol, + graph_kernel_strategy, + wire_schedule, + ) + + def all_reduce(self, tensor: Any) -> Any: + import torch + + output = ( + torch.empty_like(tensor) + if self._persistent_outputs is None + else self._persistent_outputs.acquire(tensor, torch) + ) + stream = torch.cuda.current_stream(device=tensor.device) + error = ctypes.create_string_buffer(512) + result = self._library.spark_tp4_all_reduce( + self._handle, + ctypes.c_void_p(tensor.data_ptr()), + ctypes.c_void_p(output.data_ptr()), + ctypes.c_void_p(stream.cuda_stream), + error, + len(error), + ) + if result != 0: + message = error.value.decode(errors="replace") + raise RuntimeError(f"Spark TP4 all-reduce failed: {message}") + return output + + def graph_status(self) -> GraphReplayStatus: + if not self._graph_only: + raise RuntimeError( + "Spark TP4 eager session has no graph replay status" + ) + native = _NativeGraphStatus() + error = ctypes.create_string_buffer(512) + result = self._library.spark_tp4_get_graph_status( + self._handle, + ctypes.byref(native), + ctypes.sizeof(native), + error, + len(error), + ) + if result != 0: + message = error.value.decode(errors="replace") + raise RuntimeError(f"Spark TP4 graph status failed: {message}") + if native.struct_size != ctypes.sizeof(_NativeGraphStatus): + raise RuntimeError( + "Spark TP4 graph status ABI mismatch: " + f"native={native.struct_size} python=" + f"{ctypes.sizeof(_NativeGraphStatus)}" + ) + flags = int(native.flags) + native_deferred = bool( + flags & _GRAPH_STATUS_TWO_SLOT_DEFERRED_ACK + ) + requested_deferred = ( + self._allreduce_protocol == _TWO_SLOT_DEFERRED_ACK_PROTOCOL + ) + if native_deferred != requested_deferred: + raise RuntimeError( + "Spark TP4 graph protocol attestation mismatch: " + f"requested={self._allreduce_protocol} " + f"native_deferred={native_deferred}" + ) + native_split = bool(flags & _GRAPH_STATUS_SPLIT_64K) + native_tiered = bool(flags & _GRAPH_STATUS_TIERED_64K) + if native_split and native_tiered: + raise RuntimeError( + "Spark TP4 graph kernel attestation is contradictory" + ) + native_graph_kernel = ( + _TIERED_64K_GRAPH_KERNEL + if native_tiered + else _SPLIT_64K_GRAPH_KERNEL + if native_split + else _FUSED_GRAPH_KERNEL + ) + if native_graph_kernel != self._graph_kernel_strategy: + raise RuntimeError( + "Spark TP4 graph kernel attestation mismatch: " + f"requested={self._graph_kernel_strategy} " + f"native={native_graph_kernel}" + ) + native_wire_schedule = ( + _DUAL_PORT_STRIPED_WIRE_SCHEDULE + if flags & _GRAPH_STATUS_DUAL_PORT_STRIPED + else _SEQUENTIAL_WIRE_SCHEDULE + ) + if native_wire_schedule != self._wire_schedule: + raise RuntimeError( + "Spark TP4 wire schedule attestation mismatch: " + f"requested={self._wire_schedule} " + f"native={native_wire_schedule}" + ) + native_direct_doorbell = bool( + flags & _GRAPH_STATUS_DIRECT_DOORBELL + ) + requested_direct_doorbell = ( + os.getenv("SPARK_TP4_GRAPH_DIRECT_DOORBELL", "0") == "1" + ) + if native_direct_doorbell != requested_direct_doorbell: + raise RuntimeError( + "Spark TP4 direct-doorbell attestation mismatch: " + f"requested={requested_direct_doorbell} " + f"native={native_direct_doorbell}" + ) + return GraphReplayStatus( + captured_nodes=int(native.captured_nodes), + published_sequence=int(native.published_sequence), + consumed_sequence=int(native.consumed_sequence), + completed_sequence=int(native.completed_sequence), + overflow_sequence=int(native.overflow_sequence), + capture_configured=bool( + flags & _GRAPH_STATUS_CAPTURE_CONFIGURED + ), + polling_enabled=bool(flags & _GRAPH_STATUS_POLLING_ENABLED), + host_native_atomics=bool( + flags & _GRAPH_STATUS_HOST_NATIVE_ATOMICS + ), + submit_affinity_verified=bool( + flags & _GRAPH_STATUS_SUBMIT_AFFINITY_VERIFIED + ), + progress_affinity_verified=bool( + flags & _GRAPH_STATUS_PROGRESS_AFFINITY_VERIFIED + ), + submit_cpu=( + int(native.graph_submit_cpu_plus_one) - 1 + if native.graph_submit_cpu_plus_one + else None + ), + progress_cpu=( + int(native.graph_progress_cpu_plus_one) - 1 + if native.graph_progress_cpu_plus_one + else None + ), + two_slot_deferred_ack=native_deferred, + direct_doorbell=native_direct_doorbell, + graph_kernel_strategy=native_graph_kernel, + wire_schedule=native_wire_schedule, + ) + + def health_status(self) -> NativeHealthStatus: + native = _NativeHealthStatus() + error = ctypes.create_string_buffer(512) + result = self._library.spark_tp4_get_health_status( + self._handle, + ctypes.byref(native), + ctypes.sizeof(native), + error, + len(error), + ) + if result != 0: + message = error.value.decode(errors="replace") + raise RuntimeError(f"Spark TP4 health snapshot failed: {message}") + if native.struct_size != ctypes.sizeof(_NativeHealthStatus): + raise RuntimeError( + "Spark TP4 health status ABI mismatch: " + f"native={native.struct_size} python=" + f"{ctypes.sizeof(_NativeHealthStatus)}" + ) + flags = int(native.flags) + return NativeHealthStatus( + healthy=bool(flags & _HEALTHY), + poisoned=bool(flags & _HEALTH_POISONED), + progress_thread_running=bool( + flags & _HEALTH_PROGRESS_THREAD_RUNNING + ), + stopping=bool(flags & _HEALTH_STOPPING), + submitted_sequence=int(native.submitted_sequence), + completed_sequence=int(native.completed_sequence), + failing_sequence=int(native.failing_sequence), + error_code=int(native.error_code), + failing_stage=int(native.failing_stage), + failing_rail=int(native.failing_rail), + failing_peer=int(native.failing_peer), + ) + + def capture(self, tensor: Any) -> Any: + if not self._graph_only: + raise RuntimeError("Spark TP4 eager session cannot define graph nodes") + shape = _tensor_shape(tensor) + dual_port_session = ( + self._wire_schedule == _DUAL_PORT_STRIPED_WIRE_SCHEDULE + ) + exact_dual_port_shape = ( + not dual_port_session + or shape == (self._graph_max_query_rows, _TARGET_WIDTH) + ) + if self._row_bytes == _BYTES_PER_ROW: + shape_admitted = _target_shape_eligible(shape) + else: + shape_admitted = ( + len(shape) == 2 + and shape[1] == self._elements_per_row + and 1 <= shape[0] <= self._graph_max_query_rows + ) + if ( + not shape_admitted + or shape[0] > self._graph_max_query_rows + or not exact_dual_port_shape + or str(tensor.dtype) != "torch.bfloat16" + or not bool(tensor.is_cuda) + or not bool(tensor.is_contiguous()) + ): + query_requirement = ( + f"Q={self._graph_max_query_rows}" + if dual_port_session + else f"Q in [1, {self._graph_max_query_rows}]" + ) + raise ValueError( + "Spark TP4 graph capture requires contiguous CUDA BF16 " + f"[Q, {self._elements_per_row}] with {query_requirement}; " + "dual-port striped sessions require their exact fixed Q" + ) + q = shape[0] + + import torch + + output = torch.empty_like(tensor) + stream = torch.cuda.current_stream(device=tensor.device) + error = ctypes.create_string_buffer(512) + result = self._library.spark_tp4_capture_all_reduce( + self._handle, + ctypes.c_void_p(tensor.data_ptr()), + ctypes.c_void_p(output.data_ptr()), + q, + ctypes.c_void_p(stream.cuda_stream), + error, + len(error), + ) + if result != 0: + message = error.value.decode(errors="replace") + raise RuntimeError(f"Spark TP4 graph capture failed: {message}") + return output + + def capture_q1(self, tensor: Any) -> Any: + if _tensor_shape(tensor) != (1, _TARGET_WIDTH): + raise ValueError("Spark TP4 capture_q1 requires BF16 [1, 6144]") + return self.capture(tensor) + + +class _ShadowStats: + def __init__(self) -> None: + self.count = 0 + self.validated = False + self.exact_mismatches: Any | None = None + self.outside_tolerance: Any | None = None + self.nonfinite_mismatches: Any | None = None + self.maximum_absolute_error: Any | None = None + self.maximum_ulp_error: Any | None = None + self.ulp_over_one: Any | None = None + self.ulp_over_two: Any | None = None + self.ulp_over_four: Any | None = None + + def observe(self, candidate: Any, reference: Any) -> None: + import torch + + candidate_raw = candidate.view(torch.int16) + reference_raw = reference.view(torch.int16) + exact = torch.count_nonzero(candidate_raw != reference_raw) + close = torch.isclose( + candidate, reference, rtol=0.01, atol=0.01, equal_nan=True + ) + outside = torch.count_nonzero(~close) + candidate_finite = torch.isfinite(candidate) + reference_finite = torch.isfinite(reference) + both_finite = candidate_finite & reference_finite + same_nan = torch.isnan(candidate) & torch.isnan(reference) + same_infinity = ( + torch.isinf(candidate) & torch.isinf(reference) & (candidate == reference) + ) + nonfinite = torch.count_nonzero(~(both_finite | same_nan | same_infinity)) + + absolute = (candidate.float() - reference.float()).abs() + finite_absolute = torch.where(both_finite, absolute, torch.zeros_like(absolute)) + delta = finite_absolute.max() + + candidate_bits = candidate_raw.to(torch.int32) & 0xFFFF + reference_bits = reference_raw.to(torch.int32) & 0xFFFF + candidate_ordered = torch.where( + (candidate_bits & 0x8000) != 0, + 0x8000 - (candidate_bits & 0x7FFF), + 0x8000 + candidate_bits, + ) + reference_ordered = torch.where( + (reference_bits & 0x8000) != 0, + 0x8000 - (reference_bits & 0x7FFF), + 0x8000 + reference_bits, + ) + ulp_distance = torch.where( + both_finite, + (candidate_ordered - reference_ordered).abs(), + torch.zeros_like(candidate_ordered), + ) + ulp = ulp_distance.max() + ulp_over_one = torch.count_nonzero(ulp_distance > 1) + ulp_over_two = torch.count_nonzero(ulp_distance > 2) + ulp_over_four = torch.count_nonzero(ulp_distance > 4) + + if self.exact_mismatches is None: + self.exact_mismatches = exact + self.outside_tolerance = outside + self.nonfinite_mismatches = nonfinite + self.maximum_absolute_error = delta + self.maximum_ulp_error = ulp + self.ulp_over_one = ulp_over_one + self.ulp_over_two = ulp_over_two + self.ulp_over_four = ulp_over_four + else: + self.exact_mismatches += exact + self.outside_tolerance += outside + self.nonfinite_mismatches += nonfinite + self.maximum_absolute_error = torch.maximum( + self.maximum_absolute_error, delta + ) + self.maximum_ulp_error = torch.maximum(self.maximum_ulp_error, ulp) + self.ulp_over_one += ulp_over_one + self.ulp_over_two += ulp_over_two + self.ulp_over_four += ulp_over_four + self.count += 1 + + def report(self) -> tuple[int, int, int, float, int, int, int, int]: + return ( + int(self.exact_mismatches.item()), + int(self.outside_tolerance.item()), + int(self.nonfinite_mismatches.item()), + float(self.maximum_absolute_error.item()), + int(self.maximum_ulp_error.item()), + int(self.ulp_over_one.item()), + int(self.ulp_over_two.item()), + int(self.ulp_over_four.item()), + ) + + +class _Backend: + def __init__(self, rank: int) -> None: + self.rank = rank + _backends[rank] = self + self.native_sessions: dict[int, _NativeSession] = {} + self.graph_q1_session: _NativeSession | None = None + self.graph_dual_port_q40_session: _NativeSession | None = None + self.graph_width4096_session: _NativeSession | None = None + self.bidirectional_prefill_sessions: dict[ + tuple[int, ...], _BidirectionalPrefillNativeSession + ] = {} + self.fused_prefill_sessions: dict[ + tuple[int, ...], _FusedPrefillNativeSession + ] = {} + self.shadow_stats: dict[tuple[int, tuple[int, ...], str], _ShadowStats] = {} + + def native_for(self, payload_bytes: int) -> _NativeSession: + session = self.native_sessions.get(payload_bytes) + if session is None: + session = _NativeSession(self.rank, payload_bytes) + self.native_sessions[payload_bytes] = session + return session + + def bidirectional_prefill_for( + self, shape: tuple[int, ...] + ) -> _BidirectionalPrefillNativeSession: + session = self.bidirectional_prefill_sessions.get(shape) + if session is None: + session = _BidirectionalPrefillNativeSession(self.rank, shape) + self.bidirectional_prefill_sessions[shape] = session + return session + + def fused_prefill_for(self) -> _FusedPrefillNativeSession: + capacity_shape = (8192, _BIDIRECTIONAL_PREFILL_WIDTH) + session = self.fused_prefill_sessions.get(capacity_shape) + if session is None: + session = _FusedPrefillNativeSession(self.rank, capacity_shape) + self.fused_prefill_sessions[capacity_shape] = session + return session + + def prepare_graph_q1(self) -> _NativeSession: + session = self.graph_q1_session + if session is None: + graph_cpu_affinity = _graph_preflight() + session = _NativeSession( + self.rank, + _graph_capacity_bytes(), + control_ports=_graph_control_ports(), + graph_only=True, + graph_cpu_affinity=graph_cpu_affinity, + allreduce_protocol=_graph_allreduce_protocol(), + graph_kernel_strategy=_graph_kernel_strategy(), + ) + self.graph_q1_session = session + _graph_q1_sessions[self.rank] = session + from spark_graph_status_reporter import ( + ensure_status_reporter, + ) + + ensure_status_reporter(rank=self.rank) + if ( + _graph_dual_port_q40_enabled() + and self.graph_dual_port_q40_session is None + ): + dual_port_session = _NativeSession( + self.rank, + _DUAL_PORT_Q40_CAPACITY_BYTES, + control_ports=_graph_dual_port_q40_control_ports(), + graph_only=True, + graph_cpu_affinity=_graph_dual_port_q40_cpu_affinity(), + allreduce_protocol=_TWO_SLOT_DEFERRED_ACK_PROTOCOL, + graph_kernel_strategy=_FUSED_GRAPH_KERNEL, + wire_schedule=_DUAL_PORT_STRIPED_WIRE_SCHEDULE, + ) + self.graph_dual_port_q40_session = dual_port_session + _graph_dual_port_q40_sessions[self.rank] = dual_port_session + return session + + def prepare_graph_width4096(self) -> _NativeSession: + """Research-only [Q <= 512, 4096] graph session. + + Sequential tiered_64k on two-slot deferred ACK: the split-class + kernels are the measured fix for the fused kernel's + large-payload pathology, and tiered keeps fused behavior for + small Q (see docs/DUAL_PORT_STRIPING_PROBE_20260818.md). + """ + session = self.graph_width4096_session + if session is None: + graph_cpu_affinity = _graph_preflight() + session = _NativeSession( + self.rank, + _RESEARCH_GRAPH_CAPACITY_BYTES, + control_ports=_graph_control_ports(), + graph_only=True, + graph_cpu_affinity=graph_cpu_affinity, + allreduce_protocol=_TWO_SLOT_DEFERRED_ACK_PROTOCOL, + graph_kernel_strategy=_TIERED_64K_GRAPH_KERNEL, + row_bytes=_RESEARCH_GRAPH_ROW_BYTES, + ) + self.graph_width4096_session = session + _graph_width4096_sessions[self.rank] = session + from spark_graph_status_reporter import ( + ensure_status_reporter, + ) + + ensure_status_reporter(rank=self.rank) + return session + + def graph_session_for_capture( + self, tensor: Any + ) -> _NativeSession | None: + if ( + _tensor_shape(tensor) == (ABSOLUTE_MAX_QUERY_ROWS, _TARGET_WIDTH) + and self.graph_dual_port_q40_session is not None + ): + return self.graph_dual_port_q40_session + return self.graph_q1_session + + def shadow_for(self, signature: tuple[int, tuple[int, ...], str]) -> _ShadowStats: + stats = self.shadow_stats.get(signature) + if stats is None: + stats = _ShadowStats() + self.shadow_stats[signature] = stats + return stats + + +def graph_q1_status_snapshot() -> dict[int, dict[str, object]]: + """Return nonblocking native graph progress for process-local TP ranks.""" + return { + rank: session.graph_status().to_dict() + for rank, session in sorted(_graph_q1_sessions.items()) + } + + +def graph_dual_port_q40_status_snapshot() -> dict[int, dict[str, object]]: + """Return exact-Q40 striped graph progress for process-local TP ranks.""" + return { + rank: session.graph_status().to_dict() + for rank, session in sorted(_graph_dual_port_q40_sessions.items()) + } + + +def graph_width4096_status_snapshot() -> dict[int, dict[str, object]]: + """Return research width-4096 graph progress for process-local ranks.""" + return { + rank: session.graph_status().to_dict() + for rank, session in sorted(_graph_width4096_sessions.items()) + } + + +def graph_q1_diagnostic_snapshot() -> dict[str, object]: + return { + "sessions": graph_q1_status_snapshot(), + "dual_port_q40_sessions": graph_dual_port_q40_status_snapshot(), + "width4096_sessions": graph_width4096_status_snapshot(), + "events": dict(sorted(_graph_event_counts.items())), + } + + +def native_health_snapshot() -> dict[str, NativeHealthStatus]: + """Read process-local native graph-session health without CUDA sync.""" + + sessions = { + **{f"graph-q1-rank-{rank}": session for rank, session in _graph_q1_sessions.items()}, + **{ + f"graph-q40-rank-{rank}": session + for rank, session in _graph_dual_port_q40_sessions.items() + }, + **{ + f"graph-width4096-rank-{rank}": session + for rank, session in _graph_width4096_sessions.items() + }, + **{ + f"eager-rank-{rank}-bytes-{payload_bytes}": session + for rank, backend in _backends.items() + for payload_bytes, session in backend.native_sessions.items() + }, + **{ + f"bidirectional-prefill-rank-{rank}-q{shape[0]}x{shape[1]}": session + for rank, backend in _backends.items() + for shape, session in backend.bidirectional_prefill_sessions.items() + }, + **{ + f"fused-prefill-rank-{rank}-q{shape[0]}x{shape[1]}": session + for rank, backend in _backends.items() + for shape, session in backend.fused_prefill_sessions.items() + }, + } + return {name: session.health_status() for name, session in sessions.items()} + + +def require_native_health() -> None: + failures = [ + (name, status) + for name, status in native_health_snapshot().items() + if not status.healthy + ] + if failures: + name, status = failures[0] + raise RuntimeError( + "SIRCL native session is unhealthy before model output publication: " + f"session={name} submitted={status.submitted_sequence} " + f"completed={status.completed_sequence} " + f"failing={status.failing_sequence} error={status.error_code}" + ) + + +def install() -> None: + global _installed + mode = _mode() + bidirectional_prefill = _bidirectional_prefill_enabled() + bidirectional_rail_mode = _bidirectional_prefill_rail_mode() + bidirectional_exposure = _bidirectional_prefill_exposure() + if bidirectional_exposure == "fused" and not bidirectional_prefill: + raise ValueError( + "fused prefill exposure requires " + "VLLM_SPARK_TP4_BIDIRECTIONAL_PREFILL=1" + ) + if bidirectional_exposure == "fused" and bidirectional_rail_mode != "dual": + raise ValueError("fused prefill exposure requires strict dual-rail mode") + if bidirectional_rail_mode == "dual" and not bidirectional_prefill: + raise ValueError( + "dual-rail bidirectional prefill requires " + "VLLM_SPARK_TP4_BIDIRECTIONAL_PREFILL=1" + ) + if bidirectional_prefill: + _validate_bidirectional_prefill_dual_environment() + if bidirectional_exposure == "fused": + _validate_fused_prefill_native_api() + if bidirectional_prefill and mode not in {"custom", "shadow"}: + raise ValueError( + "VLLM_SPARK_TP4_BIDIRECTIONAL_PREFILL requires " + "VLLM_SPARK_TP4_MODE=custom or shadow" + ) + graph_protocol = _graph_allreduce_protocol() + graph_kernel = _graph_kernel_strategy() + graph_dual_port_q40 = _graph_dual_port_q40_enabled() + if _graph_width4096_research_enabled(): + if mode != "custom": + raise ValueError( + "VLLM_SPARK_TP4_GRAPH_WIDTH4096_RESEARCH requires " + "VLLM_SPARK_TP4_MODE=custom" + ) + if _graph_q1_enabled() or graph_dual_port_q40: + raise ValueError( + "VLLM_SPARK_TP4_GRAPH_WIDTH4096_RESEARCH is mutually " + "exclusive with the width-6144 graph paths: they share " + "the graph control-port pair" + ) + if graph_protocol == _TWO_SLOT_DEFERRED_ACK_PROTOCOL and ( + mode != "custom" or not _graph_q1_enabled() + ): + raise ValueError( + "two_slot_deferred_ack requires custom graph all-reduce" + ) + if graph_kernel != _FUSED_GRAPH_KERNEL and ( + mode != "custom" or not _graph_q1_enabled() + ): + raise ValueError( + "research graph kernel strategies require custom graph all-reduce" + ) + if graph_dual_port_q40 and ( + mode != "custom" + or not _graph_q1_enabled() + or MAX_QUERY_ROWS != ABSOLUTE_MAX_QUERY_ROWS + ): + raise ValueError( + "exact-Q40 dual-port graph TP4 requires custom graph all-reduce " + f"with VLLM_SPARK_MAX_QUERY_ROWS={ABSOLUTE_MAX_QUERY_ROWS}" + ) + if _installed or not mode or mode == "disabled": + return + _eager_admitted_widths() + _prefill_q512_enabled() + # Resolving validates any configured query-row provider (import, + # interface, row bounds) so a broken provider fails installation + # rather than the first collective. + resolve_query_rows() + validate_active_port_namespace() + + from vllm.distributed.device_communicators.cuda_communicator import ( + CudaCommunicator, + ) + + original = CudaCommunicator.all_reduce + if getattr(original, "_spark_tp4_backend", False): + _installed = True + return + + def spark_all_reduce(self: Any, input_: Any) -> Any: + mode = _mode() + import torch + + capturing = _is_stream_capturing(torch) + bidirectional_exposure = _bidirectional_prefill_exposure() + bidirectional_eligible = _bidirectional_prefill_eligible( + self, input_, mode=mode, capturing=capturing + ) + use_fused_prefill = _fused_prefill_eligible( + self, input_, mode=mode, capturing=capturing + ) + use_bidirectional_prefill = ( + bidirectional_exposure == "sync" and bidirectional_eligible + ) + if mode == "custom" and _graph_width4096_research_enabled(): + handled = _research_graph_all_reduce( + self, input_, capturing + ) + if handled is not None: + return handled + if ( + not use_fused_prefill + and not use_bidirectional_prefill + and not _eligible(self, input_, mode) + ): + _record_stock_path( + capturing=capturing, + reason="ineligible_signature", + communicator=self, + tensor=input_, + ) + return original(self, input_) + + if capturing: + if ( + mode == "custom" + and _graph_q1_enabled() + ): + if not _target_shape_eligible(_tensor_shape(input_)): + _record_stock_path( + capturing=True, + reason="graph_width_ineligible", + communicator=self, + tensor=input_, + ) + return original(self, input_) + backend = getattr(self, "_spark_tp4_native", None) + graph_session = ( + None + if backend is None + else backend.graph_session_for_capture(input_) + ) + if graph_session is not None: + try: + output = graph_session.capture(input_) + _record_graph_event(self, "captured_nodes") + return output + except BaseException: + logger.exception( + "fatal Spark TP4 graph-capture error; terminating " + "worker because a partially captured native graph " + "cannot safely fall back" + ) + _abort_after_native_failure() + raise AssertionError("unreachable after worker termination") + _record_graph_event(self, "fallbacks") + _record_stock_path( + capturing=True, + reason="graph_session_unprepared", + ) + else: + _record_stock_path( + capturing=True, + reason="graph_transport_disabled", + ) + return original(self, input_) + + payload_bytes = _payload_bytes(input_) + signature = _collective_signature(input_) + if os.getenv("SPARK_TP4_CAPABILITY_VOTE", "0") == "1": + ensure_capability_vote(self) + backend = getattr(self, "_spark_tp4_native", None) + if backend is None: + backend = _Backend(int(self.rank_in_group)) + self._spark_tp4_native = backend + + shadow = None + promoted = False + if mode == "shadow": + shadow = backend.shadow_for(signature) + promoted = shadow.validated and ( + os.getenv("SPARK_TP4_SHADOW_PROMOTE", "0") == "1" + ) + + try: + if mode == "custom" and _graph_q1_enabled(): + backend.prepare_graph_q1() + if use_fused_prefill: + native_session = backend.fused_prefill_for() + elif use_bidirectional_prefill: + native_session = backend.bidirectional_prefill_for( + _tensor_shape(input_) + ) + else: + native_session = backend.native_for(payload_bytes) + candidate = native_session.all_reduce(input_) + except BaseException: + logger.exception( + "fatal Spark TP4 error; terminating worker because its " + "CUDA stream may be poisoned" + ) + _abort_after_native_failure() + raise AssertionError("unreachable after worker termination") + if mode == "custom" or promoted: + return candidate + + _record_stock_path( + capturing=False, + reason="shadow_reference", + ) + reference = original(self, input_) + assert shadow is not None + limit = int(os.getenv("SPARK_TP4_SHADOW_COLLECTIVES", "10000")) + if shadow.count < limit: + shadow.observe(candidate, reference) + if shadow.count == limit: + ( + exact, + outside, + nonfinite, + maximum, + max_ulp, + ulp_gt1, + ulp_gt2, + ulp_gt4, + ) = shadow.report() + logger.warning( + "Spark TP4 shadow complete: payload_bytes=%d " + "shape=%s dtype=%s collectives=%d " + "exact_mismatches=%d outside_tolerance=%d " + "nonfinite_mismatches=%d max_abs=%g max_ulp=%d " + "ulp_gt1=%d ulp_gt2=%d ulp_gt4=%d", + payload_bytes, + _tensor_shape(input_), + str(input_.dtype), + limit, + exact, + outside, + nonfinite, + maximum, + max_ulp, + ulp_gt1, + ulp_gt2, + ulp_gt4, + ) + passed = _shadow_result_passed(outside, nonfinite, max_ulp) + shadow.validated = passed + if not passed and os.getenv("SPARK_TP4_SHADOW_STRICT") == "1": + raise RuntimeError( + "Spark TP4 shadow result exceeded correctness limits" + ) + if passed and os.getenv("SPARK_TP4_SHADOW_PROMOTE", "0") == "1": + logger.warning( + "Spark TP4 payload_bytes=%d shape=%s dtype=%s " + "will promote to custom on its next call", + payload_bytes, + _tensor_shape(input_), + str(input_.dtype), + ) + return reference + + spark_all_reduce._spark_tp4_backend = True # type: ignore[attr-defined] + spark_all_reduce._spark_original = original # type: ignore[attr-defined] + CudaCommunicator.all_reduce = spark_all_reduce + _installed = True + logger.warning( + "installed Spark TP4 vLLM backend in %s mode: " + "bidirectional_prefill_exposure=%s", + _mode(), + bidirectional_exposure, + ) diff --git a/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/spark_tp4_capability.py b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/spark_tp4_capability.py new file mode 100644 index 00000000..ca86c794 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/spark_tp4_capability.py @@ -0,0 +1,363 @@ +"""Rank-wide capability agreement before SIRCL native construction.""" + +from __future__ import annotations + +import ctypes +import hashlib +import os +from pathlib import Path +from typing import Any, Callable + + +ADAPTER_ABI = "sparkring-sircl-capability/v1" +NATIVE_ABI_VERSION = 1 +FUSED_PREFILL_OPERATION_SLOTS = 2 +REQUIRED_SYMBOLS = ( + "spark_tp4_get_abi_version", + "spark_tp4_create", + "spark_tp4_create_v2", + "spark_tp4_all_reduce", + "spark_tp4_capture_all_reduce", + "spark_tp4_get_graph_status", + "spark_tp4_get_health_status", + "spark_tp4_destroy", + "spark_tp4_bidirectional_prefill_create", + "spark_tp4_bidirectional_prefill_all_reduce", + "spark_tp4_bidirectional_prefill_get_health_status", + "spark_tp4_bidirectional_prefill_destroy", + "spark_tp4_fused_prefill_create", + "spark_tp4_fused_prefill_all_reduce_rows", + "spark_tp4_fused_prefill_get_health_status", + "spark_tp4_fused_prefill_destroy", +) + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as source: + for block in iter(lambda: source.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def _device_gid_available( + device: str, + gid: str, + *, + sysfs_root: Path = Path("/sys/class/infiniband"), +) -> bool: + root = sysfs_root / device / "ports" + if not root.is_dir(): + return False + for port in root.iterdir(): + gid_path = port / "gids" / gid + if not gid_path.is_file(): + continue + value = gid_path.read_text(encoding="ascii").strip().replace(":", "") + if value and any(character != "0" for character in value): + return True + return False + + +def _integer_setting( + name: str, + default: int, + errors: list[str], + *, + minimum: int = 1, + maximum: int = 0xFFFFFFFF, +) -> int | str: + raw = os.environ.get(name, str(default)) + try: + value = int(raw) + except (TypeError, ValueError): + errors.append(f"{name} must be an integer: {raw!r}") + return raw + if not minimum <= value <= maximum: + errors.append(f"{name} must be in [{minimum}, {maximum}]: {value}") + return value + + +def _shared_capability(errors: list[str]) -> dict[str, Any]: + rail_mode = os.environ.get( + "VLLM_SPARK_TP4_BIDIRECTIONAL_PREFILL_RAIL_MODE", "single" + ) + return { + "mode": os.environ.get("VLLM_SPARK_TP4_MODE", ""), + "graph_protocol": os.environ.get( + "VLLM_SPARK_TP4_GRAPH_ALLREDUCE_PROTOCOL", "serial_ack" + ), + "direct_doorbell": os.environ.get("SPARK_TP4_GRAPH_DIRECT_DOORBELL", "0"), + "prefill_enabled": os.environ.get( + "VLLM_SPARK_TP4_BIDIRECTIONAL_PREFILL", "0" + ), + "prefill_exposure": os.environ.get( + "VLLM_SPARK_TP4_BIDIRECTIONAL_PREFILL_EXPOSURE", "sync" + ), + "prefill_rail_mode": rail_mode, + "rail_count": 2 if rail_mode == "dual" else 1, + "operation_slots": FUSED_PREFILL_OPERATION_SLOTS, + "max_inflight": _integer_setting("SPARK_TP4_MAX_INFLIGHT", 64, errors), + "admission": { + "collective_min_query_rows": 1, + "collective_max_query_rows": _integer_setting( + "VLLM_SPARK_MAX_QUERY_ROWS", 6, errors, maximum=40 + ), + "bidirectional_query_rows": (1024, 2048, 4096, 8192), + "fused_min_query_rows": 128, + "fused_max_query_rows": 8192, + }, + "control_ports": { + "eager": ( + _integer_setting( + "SPARK_TP4_CONTROL_PORT0", 11000, errors, maximum=65535 + ), + _integer_setting( + "SPARK_TP4_CONTROL_PORT1", 11001, errors, maximum=65535 + ), + ), + "graph": ( + _integer_setting( + "SPARK_TP4_GRAPH_CONTROL_PORT0", 9970, errors, maximum=65535 + ), + _integer_setting( + "SPARK_TP4_GRAPH_CONTROL_PORT1", 9971, errors, maximum=65535 + ), + ), + "graph_dual_port_q40": ( + _integer_setting( + "SPARK_TP4_GRAPH_DUAL_PORT_Q40_CONTROL_PORT0", + 9972, + errors, + maximum=65535, + ), + _integer_setting( + "SPARK_TP4_GRAPH_DUAL_PORT_Q40_CONTROL_PORT1", + 9973, + errors, + maximum=65535, + ), + ), + "bidirectional": ( + _integer_setting( + "SPARK_TP4_BIDIRECTIONAL_PREFILL_CONTROL_PORT0", + 19000, + errors, + maximum=65535, + ), + _integer_setting( + "SPARK_TP4_BIDIRECTIONAL_PREFILL_CONTROL_PORT1", + 19001, + errors, + maximum=65535, + ), + ), + "bidirectional_secondary": ( + _integer_setting( + "SPARK_TP4_BIDIRECTIONAL_PREFILL_SECONDARY_CONTROL_PORT0", + 19100, + errors, + maximum=65535, + ), + _integer_setting( + "SPARK_TP4_BIDIRECTIONAL_PREFILL_SECONDARY_CONTROL_PORT1", + 19101, + errors, + maximum=65535, + ), + ), + }, + "timeouts": { + "control_connect_seconds": _integer_setting( + "SPARK_TP4_CONTROL_CONNECT_TIMEOUT_SECONDS", 10, errors + ), + "bidirectional_seconds": _integer_setting( + "SPARK_TP4_BIDIRECTIONAL_PREFILL_TIMEOUT_SECONDS", 120, errors + ), + }, + } + + +def _cuda_capability(errors: list[str]) -> dict[str, Any]: + try: + import torch + + available = bool(torch.cuda.is_available()) + device_count = int(torch.cuda.device_count()) if available else 0 + except Exception as error: + errors.append(f"CUDA probe failed: {type(error).__name__}: {error}") + return {"available": False, "device_count": 0} + if not available or device_count < 1: + errors.append("CUDA is unavailable") + return {"available": available, "device_count": device_count} + + +def _read_sha256(path: Path, label: str, errors: list[str]) -> str: + try: + return _sha256(path) + except OSError as error: + errors.append(f"{label} cannot be read: {error}") + return "" + + +def _local_capability(rank: int) -> dict[str, Any]: + errors: list[str] = [] + shared = _shared_capability(errors) + cuda = _cuda_capability(errors) + + library_value = os.environ.get("SPARK_TP4_LIBRARY", "") + library = Path(library_value) if library_value else None + native_sha256 = "" + native_abi_version: int | None = None + if library is None or not library.is_file(): + errors.append("native library is missing") + else: + native_sha256 = _read_sha256(library, "native library", errors) + expected_native = os.environ.get("SPARKRING_SIRCL_NATIVE_SHA256", "") + if expected_native and native_sha256 and native_sha256 != expected_native: + errors.append("native library digest does not match the launcher") + try: + loaded = ctypes.CDLL(str(library)) + missing = [name for name in REQUIRED_SYMBOLS if not hasattr(loaded, name)] + if missing: + errors.append("native ABI is missing: " + ",".join(missing)) + else: + abi_version = loaded.spark_tp4_get_abi_version + abi_version.argtypes = [] + abi_version.restype = ctypes.c_uint32 + native_abi_version = int(abi_version()) + if native_abi_version != NATIVE_ABI_VERSION: + errors.append( + "native ABI version is unsupported: " + f"{native_abi_version} != {NATIVE_ABI_VERSION}" + ) + except (OSError, TypeError, ValueError) as error: + errors.append(f"native library cannot load: {error}") + + device_specs = ( + (os.environ.get("SPARK_TP4_DEVICE0", ""), os.environ.get("SPARK_TP4_GID0", "")), + (os.environ.get("SPARK_TP4_DEVICE1", ""), os.environ.get("SPARK_TP4_GID1", "")), + ) + if os.environ.get("VLLM_SPARK_TP4_BIDIRECTIONAL_PREFILL_RAIL_MODE") == "dual": + device_specs += ( + ( + os.environ.get("SPARK_TP4_BIDIRECTIONAL_PREFILL_SECONDARY_DEVICE0", ""), + os.environ.get("SPARK_TP4_BIDIRECTIONAL_PREFILL_SECONDARY_GID0", ""), + ), + ( + os.environ.get("SPARK_TP4_BIDIRECTIONAL_PREFILL_SECONDARY_DEVICE1", ""), + os.environ.get("SPARK_TP4_BIDIRECTIONAL_PREFILL_SECONDARY_GID1", ""), + ), + ) + rdma = [] + for device, gid in device_specs: + available = False + try: + available = bool(device and gid and _device_gid_available(device, gid)) + except OSError as error: + errors.append( + f"RDMA device/GID probe failed: {device or '-'}:{gid or '-'}: {error}" + ) + if not available: + errors.append(f"RDMA device/GID is unavailable: {device or '-'}:{gid or '-'}") + rdma.append({"device": device, "gid": gid, "available": available}) + + manifest_path = Path( + os.environ.get( + "SPARKRING_SIRCL_MANIFEST_PATH", + "/opt/spark-sircl/sparkring-overlay-manifest.json", + ) + ) + manifest_sha256 = ( + _read_sha256(manifest_path, "overlay manifest", errors) + if manifest_path.is_file() + else "" + ) + expected_manifest = os.environ.get("SPARKRING_SIRCL_MANIFEST_SHA256", "") + if not manifest_sha256: + errors.append("overlay manifest is missing") + elif expected_manifest and manifest_sha256 != expected_manifest: + errors.append("overlay manifest digest does not match the launcher") + + return { + "rank": rank, + "adapter_abi": ADAPTER_ABI, + "native_abi_version": native_abi_version, + "native_sha256": native_sha256, + "manifest_sha256": manifest_sha256, + "shared": shared, + "local": {"cuda": cuda, "rdma": tuple(rdma)}, + "errors": tuple(errors), + } + + +def local_capability(rank: int) -> dict[str, Any]: + """Build a non-throwing local record so every rank reaches the vote.""" + + try: + return _local_capability(rank) + except Exception as error: + errors = [f"capability probe failed: {type(error).__name__}: {error}"] + return { + "rank": rank, + "adapter_abi": ADAPTER_ABI, + "native_abi_version": None, + "native_sha256": "", + "manifest_sha256": "", + "shared": _shared_capability(errors), + "local": { + "cuda": {"available": False, "device_count": 0}, + "rdma": (), + }, + "errors": tuple(errors), + } + + +def validate_capabilities(records: list[dict[str, Any]]) -> None: + if not records: + raise RuntimeError("SIRCL capability vote returned no rank records") + expected = records[0] + failures = [] + for record in records: + rank = record.get("rank", "?") + errors = tuple(record.get("errors", ())) + if errors: + failures.append(f"rank {rank}: " + "; ".join(errors)) + for field in ( + "adapter_abi", + "native_abi_version", + "native_sha256", + "manifest_sha256", + "shared", + ): + if record.get(field) != expected.get(field): + failures.append(f"rank {rank}: {field} disagrees with rank 0") + if failures: + raise RuntimeError("SIRCL capability vote failed: " + " | ".join(failures)) + + +def _exchange(communicator: Any, local: dict[str, Any]) -> list[dict[str, Any]]: + import torch.distributed as dist + + records: list[dict[str, Any] | None] = [None] * int(communicator.world_size) + dist.all_gather_object(records, local, group=communicator.cpu_group) + if any(record is None for record in records): + raise RuntimeError("SIRCL capability vote omitted a physical rank") + return [record for record in records if record is not None] + + +def ensure_capability_vote( + communicator: Any, + *, + exchange: Callable[[Any, dict[str, Any]], list[dict[str, Any]]] = _exchange, +) -> None: + if getattr(communicator, "_sparkring_sircl_capability_voted", False): + return + local = local_capability(int(communicator.rank_in_group)) + records = exchange(communicator, local) + validate_capabilities(records) + communicator._sparkring_sircl_capability_voted = True + print( + f"SIRCL capability vote accepted: physical_ranks={len(records)}", + flush=True, + ) diff --git a/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/spark_tp4_health_gate.py b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/spark_tp4_health_gate.py new file mode 100644 index 00000000..b97f6fb9 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/spark_tp4_health_gate.py @@ -0,0 +1,98 @@ +"""Check SIRCL host health after vLLM's existing output synchronization.""" + +from __future__ import annotations + +import importlib +import os +from typing import Any, Callable + +from spark_tp4_backend import require_native_health + + +_installed = False + + +def wrap_get_output( + output_type: type, + *, + check: Callable[[], None] = require_native_health, + abort: Callable[[int], Any] = os._exit, +) -> None: + original = output_type.get_output + if getattr(original, "_sparkring_health_gate", False): + return + + def get_output(self): + output = original(self) + try: + check() + except BaseException: + abort(70) + raise RuntimeError("SIRCL health abort unexpectedly returned") + return output + + get_output._sparkring_health_gate = True + output_type.get_output = get_output + + +def wrap_worker_output( + worker_type: type, + method_name: str, + *, + check: Callable[[], None] = require_native_health, + abort: Callable[[int], Any] = os._exit, +) -> None: + original = getattr(worker_type, method_name) + if getattr(original, "_sparkring_health_gate", False): + return + + def worker_output(self, *args, **kwargs): + output = original(self, *args, **kwargs) + if output is None or hasattr(output, "get_output"): + return output + try: + check() + except BaseException: + abort(70) + raise RuntimeError("SIRCL health abort unexpectedly returned") + return output + + worker_output._sparkring_health_gate = True + setattr(worker_type, method_name, worker_output) + + +def install() -> None: + global _installed + if _installed: + return + wrapped = 0 + for module_name, class_names in ( + ("vllm.v1.worker.gpu.async_utils", ("AsyncOutput", "AsyncPoolingOutput")), + ( + "vllm.v1.worker.gpu_model_runner", + ("AsyncGPUModelRunnerOutput", "AsyncGPUPoolingModelRunnerOutput"), + ), + ): + try: + module = importlib.import_module(module_name) + except ImportError: + continue + for class_name in class_names: + output_type = getattr(module, class_name, None) + if output_type is not None: + wrap_get_output(output_type) + wrapped += 1 + try: + worker_module = importlib.import_module("vllm.v1.worker.gpu_worker") + worker_type = worker_module.Worker + except (AttributeError, ImportError) as error: + raise RuntimeError("SIRCL health gate found no pinned vLLM GPU Worker") from error + for method_name in ("execute_model", "sample_tokens"): + if not hasattr(worker_type, method_name): + raise RuntimeError( + f"SIRCL health gate found no Worker.{method_name} output boundary" + ) + wrap_worker_output(worker_type, method_name) + if not wrapped: + raise RuntimeError("SIRCL health gate found no asynchronous output type") + _installed = True diff --git a/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/spark_tp4_port_namespace.py b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/spark_tp4_port_namespace.py new file mode 100644 index 00000000..95b5e37f --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/spark_tp4_port_namespace.py @@ -0,0 +1,593 @@ +"""Deterministic control-port namespace for process-local SIRCL sessions. + +Every selected TP4 adapter derives this complete plan from the process +environment before it creates a native session. Reserving every session that +the executable configuration can instantiate prevents a later lazy bind from +colliding with an already-running transport family. +""" + +from __future__ import annotations + +import os +from collections import defaultdict +from dataclasses import dataclass +from typing import Mapping + +from spark_tp4_query_contract import MAX_QUERY_ROWS +from spark_tp4_query_row_provider import resolve_query_rows + +_MIN_CONTROL_PORT = 1 +_MAX_CONTROL_PORT = 65535 + +_EAGER_ALLREDUCE_DEFAULT_PORTS = (11000, 11001) +_EAGER_ALLREDUCE_PORT_STRIDE = 2 +_EAGER_ALLREDUCE_PREFILL_MAX_QUERY_ROWS = 512 +_EAGER_ALLREDUCE_DEFAULT_WIDTH_ELEMENTS = 6144 + +_GRAPH_ALLREDUCE_DEFAULT_PORTS = (9970, 9971) +_GRAPH_DUAL_PORT_Q40_DEFAULT_PORTS = (9972, 9973) +_EAGER_VOCAB_DEFAULT_PORTS = (9990, 9991) +_GRAPH_VOCAB_DEFAULT_PORTS = (10110, 10111) +_BIDIRECTIONAL_PREFILL_DEFAULT_PORTS = (19000, 19001) +_BIDIRECTIONAL_PREFILL_SECONDARY_DEFAULT_PORTS = (19100, 19101) +_BIDIRECTIONAL_PREFILL_PORT_STRIDE = 2 +_BIDIRECTIONAL_PREFILL_QUERY_ROWS = (1024, 2048, 4096, 8192) + + +@dataclass(frozen=True) +class PortReservation: + """One native session's two process-local TCP control ports.""" + + owner: str + ports: tuple[int, int] + + +def _environment(environ: Mapping[str, str] | None) -> Mapping[str, str]: + return os.environ if environ is None else environ + + +def _integer( + environ: Mapping[str, str], name: str, default: int +) -> int: + value = environ.get(name, str(default)) + try: + return int(value) + except ValueError as error: + raise ValueError(f"{name} must be an integer, got {value!r}") from error + + +def _flag( + environ: Mapping[str, str], name: str, default: str = "0" +) -> bool: + value = environ.get(name, default) + if value not in {"0", "1"}: + raise ValueError(f"{name} must be '0', '1', or unset") + return value == "1" + + +def _mode( + environ: Mapping[str, str], + name: str, + allowed: frozenset[str], +) -> str: + value = environ.get(name, "").lower() + if value and value not in allowed: + choices = ", ".join(repr(choice) for choice in sorted(allowed)) + raise ValueError(f"{name} must be one of {choices}, or unset") + return value + + +def validate_control_port_pair( + ports: tuple[int, int], *, owner: str +) -> tuple[int, int]: + """Validate one pair independently of whether its family is selected.""" + + if len(ports) != 2: + raise ValueError(f"Spark TP4 {owner} requires exactly two control ports") + port0, port1 = ports + if not all( + isinstance(port, int) + and not isinstance(port, bool) + and _MIN_CONTROL_PORT <= port <= _MAX_CONTROL_PORT + for port in ports + ): + raise ValueError( + f"Spark TP4 {owner} control ports must be in " + f"[{_MIN_CONTROL_PORT}, {_MAX_CONTROL_PORT}]: {ports}" + ) + if port0 == port1: + raise ValueError( + f"Spark TP4 {owner} control ports must be distinct: {ports}" + ) + return ports + + +def _configured_pair( + environ: Mapping[str, str], + *, + owner: str, + name0: str, + name1: str, + defaults: tuple[int, int], +) -> tuple[int, int]: + return validate_control_port_pair( + ( + _integer(environ, name0, defaults[0]), + _integer(environ, name1, defaults[1]), + ), + owner=owner, + ) + + +def _extension_maximum_query_rows(environ: Mapping[str, str]) -> int: + """Row bound for non-default-width extension payloads. + + Extensions enumerate the contiguous row range regardless of any + configured row-policy provider: a provider constrains default-width + serving, not the extension widths. + """ + return ( + _EAGER_ALLREDUCE_PREFILL_MAX_QUERY_ROWS + if _flag(environ, "VLLM_SPARK_TP4_PREFILL_Q512") + else MAX_QUERY_ROWS + ) + + +def _supported_allreduce_query_rows( + environ: Mapping[str, str], +) -> tuple[int, ...]: + return resolve_query_rows(environ) + + +_EAGER_ALLREDUCE_WIDTH_ENV = "VLLM_SPARK_TP4_EAGER_WIDTHS" +# Width cap: 1_048_576 elements x 512 max rows x 2 bytes = 1 GiB, which +# stays under the native single-RDMA-write bound of UINT32_MAX bytes by +# construction. +_EAGER_ALLREDUCE_MAX_WIDTH_ELEMENTS = 1_048_576 +_EAGER_ALLREDUCE_BF16_BYTES = 2 + + +def eager_allreduce_admitted_widths( + environ: Mapping[str, str] | None = None, +) -> tuple[int, ...]: + """Parse VLLM_SPARK_TP4_EAGER_WIDTHS into a sorted ascending tuple. + + Unset or empty -> (6144,). Otherwise comma-separated integer widths + (elements per row). Fail-closed ValueError names the env var on any + empty token, non-integer, out-of-range, or duplicated width. + """ + environment = _environment(environ) + raw = environment.get(_EAGER_ALLREDUCE_WIDTH_ENV, "") + if not raw: + return (_EAGER_ALLREDUCE_DEFAULT_WIDTH_ELEMENTS,) + tokens = [token.strip() for token in raw.split(",")] + widths: list[int] = [] + seen: set[int] = set() + for token in tokens: + if not token: + raise ValueError( + f"{_EAGER_ALLREDUCE_WIDTH_ENV} contains an empty token" + ) + try: + width = int(token) + except ValueError as error: + raise ValueError( + f"{_EAGER_ALLREDUCE_WIDTH_ENV} token {token!r} is not " + "an integer" + ) from error + if width < 1 or width > _EAGER_ALLREDUCE_MAX_WIDTH_ELEMENTS: + raise ValueError( + f"{_EAGER_ALLREDUCE_WIDTH_ENV} width {width} is out of " + f"range [1, {_EAGER_ALLREDUCE_MAX_WIDTH_ELEMENTS}]" + ) + if width in seen: + raise ValueError( + f"{_EAGER_ALLREDUCE_WIDTH_ENV} duplicates width {width}" + ) + seen.add(width) + widths.append(width) + return tuple(sorted(widths)) + + +def eager_allreduce_payload_sizes( + environ: Mapping[str, str] | None = None, +) -> tuple[int, ...]: + """Sorted unique payload byte counts for all admitted widths and rows. + + Each entry is rows * width * 2 bytes for every admitted width and every + rows in [1, _extension_maximum_query_rows(environ)]. Cross-width + duplicates collapse to one entry: the same byte count is the same + native operation. + """ + environment = _environment(environ) + legacy, extensions = _eager_allreduce_size_regimes(environment) + return tuple(sorted(legacy | set(extensions))) + + +def _eager_allreduce_size_regimes( + environ: Mapping[str, str], +) -> tuple[set[int], tuple[int, ...]]: + """Split admissible payload sizes into their two port regimes. + + Legacy sizes are the default width's row-denominated payloads over + the resolved query-row set (spark_tp4_query_row_provider owns that + resolution) and keep the row-slot port formula, holes included. + Extension sizes come from non-default admitted widths over the + contiguous row range; a row-policy provider is a default-width + serving constraint and does not restrict them. A non-default-width + payload equal to a supported legacy payload belongs to the legacy + regime: an identical byte count is the same native operation. + """ + default_bytes_per_row = ( + _EAGER_ALLREDUCE_DEFAULT_WIDTH_ELEMENTS + * _EAGER_ALLREDUCE_BF16_BYTES + ) + supported = _supported_allreduce_query_rows(environ) + legacy = {rows * default_bytes_per_row for rows in supported} + maximum = _extension_maximum_query_rows(environ) + extensions: set[int] = set() + for width in eager_allreduce_admitted_widths(environ): + if width == _EAGER_ALLREDUCE_DEFAULT_WIDTH_ELEMENTS: + continue + bytes_per_row = width * _EAGER_ALLREDUCE_BF16_BYTES + for rows in range(1, maximum + 1): + size = rows * bytes_per_row + if size not in legacy: + extensions.add(size) + return legacy, tuple(sorted(extensions)) + + +def eager_allreduce_ports_for_payload( + payload_bytes: int, + environ: Mapping[str, str] | None = None, +) -> tuple[int, int]: + """Control-port pair for one payload size under two slot regimes. + + A legacy payload (default width, resolved row set) occupies slot + ``row - 1``; unsupported rows leave permanent holes. An extension + payload (non-default width) occupies slot ``512 + i``, where ``i`` + is its position in the sorted extension-size tuple, past the + largest legacy span. The pair is + ``(base0 + 2*slot, base1 + 2*slot)``, validated as a pair. + """ + environment = _environment(environ) + legacy, extensions = _eager_allreduce_size_regimes(environment) + default_bytes_per_row = ( + _EAGER_ALLREDUCE_DEFAULT_WIDTH_ELEMENTS + * _EAGER_ALLREDUCE_BF16_BYTES + ) + if payload_bytes in legacy: + # Deployed row-slot formula: slot = row - 1, so unsupported rows + # leave permanent holes in the port sequence. The exact-state + # arena accounting depends on these exact pairs. + slot = payload_bytes // default_bytes_per_row - 1 + elif payload_bytes in extensions: + # Width extensions occupy slots past the largest legacy span + # (Q512 ends at slot 511) so no row-policy or Q512 toggle + # moves them. + slot = ( + _EAGER_ALLREDUCE_PREFILL_MAX_QUERY_ROWS + + extensions.index(payload_bytes) + ) + else: + raise ValueError( + "unsupported Spark TP4 eager all-reduce payload size: " + f"{payload_bytes} bytes" + ) + base0 = _integer( + environment, + "SPARK_TP4_CONTROL_PORT0", + _EAGER_ALLREDUCE_DEFAULT_PORTS[0], + ) + base1 = _integer( + environment, + "SPARK_TP4_CONTROL_PORT1", + _EAGER_ALLREDUCE_DEFAULT_PORTS[1], + ) + offset = slot * _EAGER_ALLREDUCE_PORT_STRIDE + return validate_control_port_pair( + (base0 + offset, base1 + offset), + owner=f"eager all-reduce payload {payload_bytes}", + ) + + +def _eager_allreduce_pair( + query_rows: int, environ: Mapping[str, str] +) -> tuple[int, int]: + supported = _supported_allreduce_query_rows(environ) + if ( + not isinstance(query_rows, int) + or isinstance(query_rows, bool) + or query_rows not in supported + ): + raise ValueError( + "Spark TP4 eager all-reduce query rows must be in the exact " + f"supported set {supported}: {query_rows}" + ) + payload_bytes = ( + query_rows + * _EAGER_ALLREDUCE_DEFAULT_WIDTH_ELEMENTS + * _EAGER_ALLREDUCE_BF16_BYTES + ) + return eager_allreduce_ports_for_payload(payload_bytes, environ) + + + + +def _graph_allreduce_pair(environ: Mapping[str, str]) -> tuple[int, int]: + return _configured_pair( + environ, + owner="graph all-reduce", + name0="SPARK_TP4_GRAPH_CONTROL_PORT0", + name1="SPARK_TP4_GRAPH_CONTROL_PORT1", + defaults=_GRAPH_ALLREDUCE_DEFAULT_PORTS, + ) + + +def _graph_dual_port_q40_pair( + environ: Mapping[str, str], +) -> tuple[int, int]: + return _configured_pair( + environ, + owner="graph exact-Q40 dual-port all-reduce", + name0="SPARK_TP4_GRAPH_DUAL_PORT_Q40_CONTROL_PORT0", + name1="SPARK_TP4_GRAPH_DUAL_PORT_Q40_CONTROL_PORT1", + defaults=_GRAPH_DUAL_PORT_Q40_DEFAULT_PORTS, + ) + + +def _bidirectional_prefill_pair( + query_rows: int, environ: Mapping[str, str] +) -> tuple[int, int]: + if query_rows not in _BIDIRECTIONAL_PREFILL_QUERY_ROWS: + raise ValueError( + "Spark TP4 bidirectional prefill Q must be one of " + f"{_BIDIRECTIONAL_PREFILL_QUERY_ROWS}: {query_rows}" + ) + base = _configured_pair( + environ, + owner="bidirectional prefill base", + name0="SPARK_TP4_BIDIRECTIONAL_PREFILL_CONTROL_PORT0", + name1="SPARK_TP4_BIDIRECTIONAL_PREFILL_CONTROL_PORT1", + defaults=_BIDIRECTIONAL_PREFILL_DEFAULT_PORTS, + ) + offset = ( + _BIDIRECTIONAL_PREFILL_QUERY_ROWS.index(query_rows) + * _BIDIRECTIONAL_PREFILL_PORT_STRIDE + ) + return validate_control_port_pair( + (base[0] + offset, base[1] + offset), + owner=f"bidirectional prefill q={query_rows}", + ) + + +def _bidirectional_prefill_secondary_pair( + query_rows: int, environ: Mapping[str, str] +) -> tuple[int, int]: + if query_rows not in _BIDIRECTIONAL_PREFILL_QUERY_ROWS: + raise ValueError( + "Spark TP4 bidirectional prefill Q must be one of " + f"{_BIDIRECTIONAL_PREFILL_QUERY_ROWS}: {query_rows}" + ) + base = _configured_pair( + environ, + owner="bidirectional prefill secondary base", + name0="SPARK_TP4_BIDIRECTIONAL_PREFILL_SECONDARY_CONTROL_PORT0", + name1="SPARK_TP4_BIDIRECTIONAL_PREFILL_SECONDARY_CONTROL_PORT1", + defaults=_BIDIRECTIONAL_PREFILL_SECONDARY_DEFAULT_PORTS, + ) + offset = ( + _BIDIRECTIONAL_PREFILL_QUERY_ROWS.index(query_rows) + * _BIDIRECTIONAL_PREFILL_PORT_STRIDE + ) + return validate_control_port_pair( + (base[0] + offset, base[1] + offset), + owner=f"bidirectional prefill secondary q={query_rows}", + ) + + + + +def _eager_vocab_pair(environ: Mapping[str, str]) -> tuple[int, int]: + return _configured_pair( + environ, + owner="eager vocabulary all-gather", + name0="SPARK_TP4_VOCAB_CONTROL_PORT0", + name1="SPARK_TP4_VOCAB_CONTROL_PORT1", + defaults=_EAGER_VOCAB_DEFAULT_PORTS, + ) + + +def _graph_vocab_pair(environ: Mapping[str, str]) -> tuple[int, int]: + return _configured_pair( + environ, + owner="graph vocabulary all-gather", + name0="SPARK_TP4_GRAPH_VOCAB_CONTROL_PORT0", + name1="SPARK_TP4_GRAPH_VOCAB_CONTROL_PORT1", + defaults=_GRAPH_VOCAB_DEFAULT_PORTS, + ) + + +def active_port_reservations( + environ: Mapping[str, str] | None = None, +) -> tuple[PortReservation, ...]: + """Return every control-port pair the selected adapters may bind.""" + + environment = _environment(environ) + reservations: list[PortReservation] = [] + + allreduce_mode = _mode( + environment, + "VLLM_SPARK_TP4_MODE", + frozenset({"custom", "disabled", "shadow"}), + ) + bidirectional_rail_mode = _mode( + environment, + "VLLM_SPARK_TP4_BIDIRECTIONAL_PREFILL_RAIL_MODE", + frozenset({"single", "dual"}), + ) or "single" + if allreduce_mode in {"custom", "shadow"}: + for query_rows in _supported_allreduce_query_rows(environment): + reservations.append( + PortReservation( + f"eager_allreduce:q={query_rows}", + _eager_allreduce_pair(query_rows, environment), + ) + ) + _, extension_sizes = _eager_allreduce_size_regimes(environment) + for payload_bytes in extension_sizes: + reservations.append( + PortReservation( + f"eager_allreduce:payload={payload_bytes}", + eager_allreduce_ports_for_payload( + payload_bytes, environment + ), + ) + ) + if allreduce_mode == "custom" and _flag( + environment, "VLLM_SPARK_TP4_GRAPH_Q1" + ): + reservations.append( + PortReservation( + "graph_allreduce", _graph_allreduce_pair(environment) + ) + ) + if _flag( + environment, "VLLM_SPARK_TP4_GRAPH_DUAL_PORT_Q40" + ): + reservations.append( + PortReservation( + "graph_dual_port_q40_allreduce", + _graph_dual_port_q40_pair(environment), + ) + ) + if _flag(environment, "VLLM_SPARK_TP4_BIDIRECTIONAL_PREFILL"): + for query_rows in _BIDIRECTIONAL_PREFILL_QUERY_ROWS: + reservations.append( + PortReservation( + f"bidirectional_prefill:q={query_rows}", + _bidirectional_prefill_pair(query_rows, environment), + ) + ) + if bidirectional_rail_mode == "dual": + reservations.append( + PortReservation( + f"bidirectional_prefill_secondary:q={query_rows}", + _bidirectional_prefill_secondary_pair( + query_rows, environment + ), + ) + ) + + vocab_mode = _mode( + environment, + "VLLM_SPARK_TP4_VOCAB_MODE", + frozenset({"custom", "shadow"}), + ) + if vocab_mode: + reservations.append( + PortReservation("eager_vocab", _eager_vocab_pair(environment)) + ) + if vocab_mode == "custom" and _flag( + environment, "VLLM_SPARK_TP4_GRAPH_Q1" + ): + reservations.append( + PortReservation("graph_vocab", _graph_vocab_pair(environment)) + ) + + return tuple(reservations) + + +def validate_active_port_namespace( + environ: Mapping[str, str] | None = None, +) -> tuple[PortReservation, ...]: + """Reject out-of-range, internally duplicated, or shared active ports.""" + + reservations = active_port_reservations(environ) + owners_by_port: dict[int, list[str]] = defaultdict(list) + for reservation in reservations: + validate_control_port_pair( + reservation.ports, owner=reservation.owner + ) + for port in reservation.ports: + owners_by_port[port].append(reservation.owner) + + collisions = { + port: tuple(owners) + for port, owners in sorted(owners_by_port.items()) + if len(owners) > 1 + } + if collisions: + details = "; ".join( + f"port {port}: {', '.join(owners)}" + for port, owners in collisions.items() + ) + raise ValueError( + "Spark TP4 active control-port namespaces collide: " + details + ) + return reservations + + +def eager_allreduce_control_ports( + query_rows: int, environ: Mapping[str, str] | None = None +) -> tuple[int, int]: + environment = _environment(environ) + ports = _eager_allreduce_pair(query_rows, environment) + validate_active_port_namespace(environment) + return ports + + +def graph_allreduce_control_ports( + environ: Mapping[str, str] | None = None, +) -> tuple[int, int]: + environment = _environment(environ) + ports = _graph_allreduce_pair(environment) + validate_active_port_namespace(environment) + return ports + + +def graph_dual_port_q40_control_ports( + environ: Mapping[str, str] | None = None, +) -> tuple[int, int]: + environment = _environment(environ) + ports = _graph_dual_port_q40_pair(environment) + validate_active_port_namespace(environment) + return ports + + +def bidirectional_prefill_control_ports( + query_rows: int, environ: Mapping[str, str] | None = None +) -> tuple[int, int]: + environment = _environment(environ) + ports = _bidirectional_prefill_pair(query_rows, environment) + validate_active_port_namespace(environment) + return ports + + +def bidirectional_prefill_secondary_control_ports( + query_rows: int, environ: Mapping[str, str] | None = None +) -> tuple[int, int]: + environment = _environment(environ) + ports = _bidirectional_prefill_secondary_pair(query_rows, environment) + validate_active_port_namespace(environment) + return ports + + + + +def eager_vocab_control_ports( + environ: Mapping[str, str] | None = None, +) -> tuple[int, int]: + environment = _environment(environ) + ports = _eager_vocab_pair(environment) + validate_active_port_namespace(environment) + return ports + + +def graph_vocab_control_ports( + environ: Mapping[str, str] | None = None, +) -> tuple[int, int]: + environment = _environment(environ) + ports = _graph_vocab_pair(environment) + validate_active_port_namespace(environment) + return ports diff --git a/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/spark_tp4_query_contract.py b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/spark_tp4_query_contract.py new file mode 100644 index 00000000..48531aea --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/spark_tp4_query_contract.py @@ -0,0 +1,81 @@ +"""Shared GLM-5.2 TP4 query-width contract. + +One scheduler step can mix ordinary decode rows and speculative target rows. +With at most eight active sequences and maximum adaptive MTP depth four, the +total target-query width is bounded by 8 * (4 + 1) = 40. Every custom +collective must accept that complete interval; otherwise a legal mixed batch +can silently fall back to NCCL Socket. +""" + +from __future__ import annotations + +import os + +MAX_CONCURRENT_SEQUENCES = 8 +MAX_SPECULATIVE_TOKENS = 4 +ABSOLUTE_MAX_QUERY_ROWS = MAX_CONCURRENT_SEQUENCES * ( + MAX_SPECULATIVE_TOKENS + 1 +) +_DEFAULT_MAX_QUERY_ROWS = 6 +_BASE_CAPTURE_BUCKETS = (1, 2, 3, 4, 5, 6, 8, 10, 12, 16, 20, 24, 32, 40) + + +def _configured_max_query_rows() -> int: + raw = os.getenv( + "VLLM_SPARK_MAX_QUERY_ROWS", + str(_DEFAULT_MAX_QUERY_ROWS), + ) + try: + value = int(raw) + except ValueError as error: + raise ValueError( + "VLLM_SPARK_MAX_QUERY_ROWS must be an integer" + ) from error + if not 1 <= value <= ABSOLUTE_MAX_QUERY_ROWS: + raise ValueError( + "VLLM_SPARK_MAX_QUERY_ROWS must be in " + f"[1, {ABSOLUTE_MAX_QUERY_ROWS}]" + ) + return value + + +MAX_QUERY_ROWS = _configured_max_query_rows() +SUPPORTED_QUERY_ROWS = frozenset(range(1, MAX_QUERY_ROWS + 1)) + + +def capture_buckets( + max_query_rows: int, + *, + uniform_query_rows: int, +) -> list[int]: + """Return mixed-batch buckets plus exact FULL-decode widths.""" + + if not 1 <= max_query_rows <= ABSOLUTE_MAX_QUERY_ROWS: + raise ValueError( + f"max_query_rows must be in [1, {ABSOLUTE_MAX_QUERY_ROWS}]" + ) + if not 1 <= uniform_query_rows <= max_query_rows: + raise ValueError("uniform_query_rows must be in [1, max_query_rows]") + + buckets = {value for value in _BASE_CAPTURE_BUCKETS if value <= max_query_rows} + buckets.add(max_query_rows) + buckets.update(range(uniform_query_rows, max_query_rows + 1, uniform_query_rows)) + return sorted(buckets) + + +def capture_sizes(max_num_seqs: int) -> list[int]: + """Return bounded-padding CUDA-graph coverage for adaptive MTP2/4.""" + + if ( + isinstance(max_num_seqs, bool) + or not isinstance(max_num_seqs, int) + or not 1 <= max_num_seqs <= MAX_CONCURRENT_SEQUENCES + ): + raise ValueError( + "max_num_seqs must be an integer in " + f"[1, {MAX_CONCURRENT_SEQUENCES}]" + ) + maximum_q = max_num_seqs * (MAX_SPECULATIVE_TOKENS + 1) + if max_num_seqs == 1: + return [1, 3, 5] + return capture_buckets(maximum_q, uniform_query_rows=5) diff --git a/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/spark_tp4_query_row_provider.py b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/spark_tp4_query_row_provider.py new file mode 100644 index 00000000..703bf5ec --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/spark_tp4_query_row_provider.py @@ -0,0 +1,164 @@ +"""Generic query-row policy resolver for the TP4 adapters. + +The eager TP4 all-reduce admits payloads by query-row count at the +default width. Which row counts exist is policy, and this module is the +single owner of that policy for both the admission gate +(spark_tp4_backend) and the reservation namespace +(spark_tp4_port_namespace), so the two can never disagree. + +Three mutually exclusive sources, resolved per call: + +- ``VLLM_SPARK_TP4_PREFILL_Q512=1``: the broad prefill geometry, + rows 1..512. +- ``VLLM_SPARK_TP4_QUERY_ROW_PROVIDER=``: an external provider + module owns the row set. The module is imported lazily, only when the + variable is set, and must expose + ``provider_query_rows(environ) -> Iterable[int]``. The returned rows + are validated here; the provider owns any policy of its own + (including reading provider-specific environment variables). +- Neither: the contiguous range ``1..MAX_QUERY_ROWS`` from + spark_tp4_query_contract. + +Failure behavior is closed: a configured provider that cannot be +imported, lacks the interface, or returns an invalid row set raises +ValueError naming ``VLLM_SPARK_TP4_QUERY_ROW_PROVIDER`` and the module. +Configuring both Q512 and a provider raises rather than letting two +geometry sources compete silently. + +The provider module name is part of the launch identity: every rank +must configure the same value (or none). A mismatch produces differing +port reservations and fails native session establishment; it is not +detectable from a single rank. +""" + +from __future__ import annotations + +import importlib +import os +from collections.abc import Iterable, Mapping + +from spark_tp4_query_contract import MAX_QUERY_ROWS + +PROVIDER_ENV = "VLLM_SPARK_TP4_QUERY_ROW_PROVIDER" +_PREFILL_Q512_ENV = "VLLM_SPARK_TP4_PREFILL_Q512" + +# Rows occupy row-denominated reservation slots below the extension +# span, which begins at slot 512; a provider may not exceed it. +MAX_PROVIDER_QUERY_ROW = 512 + +_PREFILL_Q512_ROWS = tuple(range(1, 513)) + +# Cache for ambient-environment resolutions only. os.environ is +# immutable for the life of a launch, so those resolutions are +# per-process constants and the cache keeps the admission hot path +# allocation-free. Explicit environment mappings are never cached: a +# provider may read arbitrary variables from the mapping it is given, +# so two different mappings naming the same provider are distinct +# resolutions. +_ambient_cache: dict[tuple[str | None, bool, int], tuple[int, ...]] = {} + + +def _environment(environ: Mapping[str, str] | None) -> Mapping[str, str]: + return os.environ if environ is None else environ + + +def _prefill_q512(environ: Mapping[str, str]) -> bool: + value = environ.get(_PREFILL_Q512_ENV, "0") + if value not in {"0", "1"}: + raise ValueError( + f"{_PREFILL_Q512_ENV} must be '0', '1', or unset" + ) + return value == "1" + + +def provider_module_name( + environ: Mapping[str, str] | None = None, +) -> str | None: + """The configured provider module name, or None when unset/empty.""" + + return _environment(environ).get(PROVIDER_ENV, "").strip() or None + + +def _validated_rows(rows: Iterable[int], *, module_name: str) -> tuple[int, ...]: + prefix = f"{PROVIDER_ENV}={module_name}" + try: + candidates = list(rows) + except TypeError as error: + raise ValueError( + f"{prefix}: provider_query_rows must return an iterable of " + f"integers, got {type(rows).__name__}" + ) from error + if not candidates: + raise ValueError(f"{prefix}: provider returned no query rows") + for row in candidates: + if isinstance(row, bool) or not isinstance(row, int): + raise ValueError( + f"{prefix}: query rows must be integers, got {row!r}" + ) + if not 1 <= row <= MAX_PROVIDER_QUERY_ROW: + raise ValueError( + f"{prefix}: query rows must be in " + f"[1, {MAX_PROVIDER_QUERY_ROW}], got {row}" + ) + if len(set(candidates)) != len(candidates): + raise ValueError(f"{prefix}: query rows must be unique") + return tuple(sorted(candidates)) + + +def _provider_rows( + module_name: str, environ: Mapping[str, str] +) -> tuple[int, ...]: + try: + module = importlib.import_module(module_name) + except ImportError as error: + raise ValueError( + f"{PROVIDER_ENV} names {module_name!r}, which is not " + f"importable: {error}" + ) from error + interface = getattr(module, "provider_query_rows", None) + if not callable(interface): + raise ValueError( + f"{PROVIDER_ENV} names {module_name!r}, which does not " + "define provider_query_rows(environ)" + ) + return _validated_rows(interface(environ), module_name=module_name) + + +def resolve_query_rows( + environ: Mapping[str, str] | None = None, +) -> tuple[int, ...]: + """The admitted query-row set for this environment, sorted ascending. + + Raises ValueError on conflicting geometry sources or on a configured + provider that is absent, interface-less, or returns invalid rows. + """ + + ambient = environ is None + environment = _environment(environ) + q512 = _prefill_q512(environment) + module_name = provider_module_name(environment) + if q512 and module_name is not None: + raise ValueError( + f"{_PREFILL_Q512_ENV}=1 and {PROVIDER_ENV}={module_name} " + "are mutually exclusive geometry sources" + ) + key = (module_name, q512, MAX_QUERY_ROWS) + if ambient: + cached = _ambient_cache.get(key) + if cached is not None: + return cached + if q512: + rows = _PREFILL_Q512_ROWS + elif module_name is not None: + rows = _provider_rows(module_name, environment) + else: + rows = tuple(range(1, MAX_QUERY_ROWS + 1)) + if ambient: + _ambient_cache[key] = rows + return rows + + +def maximum_query_rows(environ: Mapping[str, str] | None = None) -> int: + """The largest admitted query-row count for this environment.""" + + return resolve_query_rows(environ)[-1] diff --git a/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/spark_tp4_vocab_allgather_backend.py b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/spark_tp4_vocab_allgather_backend.py new file mode 100644 index 00000000..1a399bd5 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/spark_tp4_vocab_allgather_backend.py @@ -0,0 +1,848 @@ +"""Exact-signature vLLM adapter for Spark TP4 vocabulary all-gather.""" + +from __future__ import annotations + +import ctypes +import logging +import os +from typing import Any + +from spark_tp4_port_namespace import ( + eager_vocab_control_ports, + graph_vocab_control_ports, + validate_active_port_namespace, + validate_control_port_pair, +) +from spark_tp4_query_contract import SUPPORTED_QUERY_ROWS + +logger = logging.getLogger(__name__) + +_installed = False +_VALID_MODES = {"shadow", "custom"} +_TP_GROUP = "tp:0" +_WORLD_SIZE = 4 +_VOCAB_PER_RANK = 38720 +_OUTPUT_VOCAB = _VOCAB_PER_RANK * _WORLD_SIZE +_BF16_BYTES = 2 +_SUPPORTED_Q = SUPPORTED_QUERY_ROWS +_VOCAB_GRAPH_DEFAULT_PROGRESS_CPU = 12 +_GRAPH_STATUS_CAPTURE_CONFIGURED = 1 << 0 +_GRAPH_STATUS_POLLING_ENABLED = 1 << 1 +_GRAPH_STATUS_HOST_NATIVE_ATOMICS = 1 << 2 +_GRAPH_STATUS_SUBMIT_AFFINITY_VERIFIED = 1 << 3 +_GRAPH_STATUS_PROGRESS_AFFINITY_VERIFIED = 1 << 4 +_vocab_graph_sessions: dict[int, "_NativeVocabSession"] = {} +_graph_event_counts: dict[str, int] = {} +# PLACEHOLDER ring peers (RFC 5737 TEST-NET-1): 192.0.2.N stands in for +# rank N-1's direct-cable address. These are NOT routable and MUST be +# replaced for any live run by setting SPARK_TP4_PEER0 / SPARK_TP4_PEER1 +# (the authoritative per-rank overrides) or by editing this table. +_DEFAULT_PEERS = { + 0: ("192.0.2.2", "192.0.2.4"), + 1: ("192.0.2.1", "192.0.2.3"), + 2: ("192.0.2.4", "192.0.2.2"), + 3: ("192.0.2.3", "192.0.2.1"), +} + +_Signature = int + + +def _abort_after_native_failure() -> None: + """Terminate a worker whose CUDA stream may contain a native wait.""" + + os._exit(73) + + +def _mode() -> str: + mode = os.getenv("VLLM_SPARK_TP4_VOCAB_MODE", "").lower() + if mode and mode not in _VALID_MODES: + raise ValueError( + "VLLM_SPARK_TP4_VOCAB_MODE must be 'shadow', 'custom', or unset" + ) + return mode + + +def _positive_integer(name: str, default: int) -> int: + value = int(os.getenv(name, str(default))) + if value <= 0: + raise ValueError(f"{name} must be a positive integer") + return value + + +def _shape(tensor: Any) -> tuple[int, ...]: + return tuple(int(value) for value in tensor.shape) + + +def _signature( + group: Any, input_tensor: Any, dim: int, mode: str +) -> _Signature | None: + shape = _shape(input_tensor) + normalized_dim = dim + len(shape) if dim < 0 else dim + if ( + mode not in _VALID_MODES + or getattr(group, "unique_name", None) != _TP_GROUP + or int(getattr(group, "world_size", 0)) != _WORLD_SIZE + or int(getattr(group, "rank_in_group", -1)) not in range(_WORLD_SIZE) + or normalized_dim != 1 + or len(shape) != 2 + or shape[0] not in _SUPPORTED_Q + or shape[1] != _VOCAB_PER_RANK + or str(input_tensor.dtype) != "torch.bfloat16" + or not bool(input_tensor.is_cuda) + or not bool(input_tensor.is_contiguous()) + ): + return None + return shape[0] + + +def _is_stream_capturing(torch_module: Any) -> bool: + checker = getattr(torch_module.cuda, "is_current_stream_capturing", None) + return bool(checker is not None and checker()) + + +def _record_stock_path( + *, + capturing: bool, + reason: str, + group: Any | None = None, + tensor: Any | None = None, + dim: int | None = None, +) -> None: + from spark_collective_audit import ( + StockCollectiveSignature, + classify_stock_family, + enabled, + record_stock, + ) + + signature = None + family = "vocabulary_all_gather" + if enabled() and group is not None and tensor is not None: + world_size = getattr(group, "world_size", None) + signature = StockCollectiveSignature( + shape=_shape(tensor), + dtype=str(tensor.dtype), + is_cuda=bool(tensor.is_cuda), + contiguous=bool(tensor.is_contiguous()), + world_size=( + None if world_size is None else int(world_size) + ), + unique_name=str(getattr(group, "unique_name", "")), + ) + family = classify_stock_family( + "group_all_gather", + signature, + dim=dim, + ) + record_stock( + family, + capturing=capturing, + reason=reason, + signature=signature, + ) + + +def _graph_enabled() -> bool: + value = os.getenv("VLLM_SPARK_TP4_GRAPH_Q1", "0") + if value not in {"0", "1"}: + raise ValueError("VLLM_SPARK_TP4_GRAPH_Q1 must be '0' or '1'") + return value == "1" + + +def _graph_preflight() -> tuple[int, int]: + from spark_tp4_backend import _graph_preflight as tp_graph_preflight + + submit_cpu, tp_progress_cpu = tp_graph_preflight() + progress_cpu = int( + os.getenv( + "SPARK_TP4_GRAPH_VOCAB_PROGRESS_CPU", + str(_VOCAB_GRAPH_DEFAULT_PROGRESS_CPU), + ) + ) + if progress_cpu < 0: + raise RuntimeError( + "Spark TP4 vocabulary graph progress CPU must be nonnegative" + ) + if progress_cpu in {submit_cpu, tp_progress_cpu}: + raise RuntimeError( + "Spark TP4 vocabulary graph progress CPU must differ from " + "the shared submit and TP progress CPUs" + ) + return submit_cpu, progress_cpu + + +def _graph_control_ports() -> tuple[int, int]: + return graph_vocab_control_ports() + + +def _eager_control_ports() -> tuple[int, int]: + return eager_vocab_control_ports() + + +def _validate_control_ports(ports: tuple[int, int]) -> None: + validate_control_port_pair(ports, owner="vocabulary all-gather") + validate_active_port_namespace() + + +def _record_graph_event(group: Any, event: str) -> int: + attribute = f"_spark_tp4_vocab_graph_{event}" + count = int(getattr(group, attribute, 0)) + 1 + setattr(group, attribute, count) + _graph_event_counts[event] = _graph_event_counts.get(event, 0) + 1 + if count == 1 or count % 128 == 0: + logger.warning( + "Spark TP4 vocabulary graph %s on rank %d: count=%d", + event, + int(group.rank_in_group), + count, + ) + return count + + +class _NativeVocabConfig(ctypes.Structure): + _fields_ = [ + ("rank", ctypes.c_uint32), + ("peer0", ctypes.c_char_p), + ("peer1", ctypes.c_char_p), + ("device0", ctypes.c_char_p), + ("device1", ctypes.c_char_p), + ("gid0", ctypes.c_uint8), + ("gid1", ctypes.c_uint8), + ("control_port0", ctypes.c_uint16), + ("control_port1", ctypes.c_uint16), + ] + + +class _NativeVocabGraphConfig(ctypes.Structure): + _fields_ = [ + ("rank", ctypes.c_uint32), + ("peer0", ctypes.c_char_p), + ("peer1", ctypes.c_char_p), + ("device0", ctypes.c_char_p), + ("device1", ctypes.c_char_p), + ("gid0", ctypes.c_uint8), + ("gid1", ctypes.c_uint8), + ("control_port0", ctypes.c_uint16), + ("control_port1", ctypes.c_uint16), + ("graph_submit_cpu_plus_one", ctypes.c_uint32), + ("graph_progress_cpu_plus_one", ctypes.c_uint32), + ] + + +class _NativeVocabGraphStatus(ctypes.Structure): + _fields_ = [ + ("struct_size", ctypes.c_uint32), + ("flags", ctypes.c_uint32), + ("captured_nodes", ctypes.c_uint64), + ("published_sequence", ctypes.c_uint64), + ("consumed_sequence", ctypes.c_uint64), + ("completed_sequence", ctypes.c_uint64), + ("overflow_sequence", ctypes.c_uint64), + ("graph_submit_cpu_plus_one", ctypes.c_uint32), + ("graph_progress_cpu_plus_one", ctypes.c_uint32), + ] + + +class _NativeVocabSession: + def __init__( + self, + rank: int, + *, + graph_only: bool = False, + control_ports: tuple[int, int] | None = None, + graph_cpu_affinity: tuple[int, int] | None = None, + ) -> None: + if rank not in _DEFAULT_PEERS: + raise ValueError(f"TP rank must be in [0, 3], got {rank}") + if graph_only != (graph_cpu_affinity is not None): + raise ValueError( + "Spark TP4 vocabulary graph session requires an explicit " + "CPU pair; eager sessions cannot set one" + ) + self._library = ctypes.CDLL(os.environ["SPARK_TP4_LIBRARY"]) + self._library.spark_tp4_vocab_allgather_create.argtypes = [ + ctypes.POINTER(_NativeVocabConfig), + ctypes.c_char_p, + ctypes.c_size_t, + ] + self._library.spark_tp4_vocab_allgather_create.restype = ( + ctypes.c_void_p + ) + if graph_only: + self._library.spark_tp4_vocab_graph_create.argtypes = [ + ctypes.POINTER(_NativeVocabGraphConfig), + ctypes.c_char_p, + ctypes.c_size_t, + ] + self._library.spark_tp4_vocab_graph_create.restype = ( + ctypes.c_void_p + ) + self._library.spark_tp4_vocab_allgather.argtypes = [ + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_uint32, + ctypes.c_void_p, + ctypes.c_char_p, + ctypes.c_size_t, + ] + self._library.spark_tp4_vocab_allgather.restype = ctypes.c_int + if graph_only: + self._library.spark_tp4_vocab_capture_allgather.argtypes = [ + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_uint32, + ctypes.c_void_p, + ctypes.c_char_p, + ctypes.c_size_t, + ] + self._library.spark_tp4_vocab_capture_allgather.restype = ( + ctypes.c_int + ) + self._library.spark_tp4_vocab_get_graph_status.argtypes = [ + ctypes.c_void_p, + ctypes.POINTER(_NativeVocabGraphStatus), + ctypes.c_size_t, + ctypes.c_char_p, + ctypes.c_size_t, + ] + self._library.spark_tp4_vocab_get_graph_status.restype = ( + ctypes.c_int + ) + self._library.spark_tp4_vocab_allgather_destroy.argtypes = [ + ctypes.c_void_p + ] + self._library.spark_tp4_vocab_allgather_destroy.restype = None + self._graph_only = graph_only + self._capture_stream: int | None = None + + default_peer0, default_peer1 = _DEFAULT_PEERS[rank] + ports = ( + _eager_control_ports() + if control_ports is None + else control_ports + ) + _validate_control_ports(ports) + port0, port1 = ports + submit_cpu, progress_cpu = graph_cpu_affinity or (-1, -1) + common_config = { + "rank": rank, + "peer0": os.getenv( + "SPARK_TP4_PEER0", default_peer0 + ).encode(), + "peer1": os.getenv( + "SPARK_TP4_PEER1", default_peer1 + ).encode(), + "device0": os.getenv( + "SPARK_TP4_DEVICE0", "rocep1s0f0" + ).encode(), + "device1": os.getenv( + "SPARK_TP4_DEVICE1", "rocep1s0f1" + ).encode(), + "gid0": int(os.getenv("SPARK_TP4_GID0", "3")), + "gid1": int(os.getenv("SPARK_TP4_GID1", "3")), + "control_port0": port0, + "control_port1": port1, + } + if graph_only: + config = _NativeVocabGraphConfig( + **common_config, + graph_submit_cpu_plus_one=submit_cpu + 1, + graph_progress_cpu_plus_one=progress_cpu + 1, + ) + create = self._library.spark_tp4_vocab_graph_create + else: + config = _NativeVocabConfig(**common_config) + create = self._library.spark_tp4_vocab_allgather_create + error = ctypes.create_string_buffer(512) + self._handle = create( + ctypes.byref(config), error, len(error) + ) + if not self._handle: + message = error.value.decode(errors="replace") + raise RuntimeError( + "failed to create Spark TP4 vocabulary session: " + f"{message}" + ) + logger.warning( + "Spark TP4 vocabulary %s session ready: rank=%d ports=%d/%d", + "graph-only" if graph_only else "eager", + rank, + port0, + port1, + ) + + def all_gather( + self, + input_tensor: Any, + output_tensor: Any, + query_rows: int, + stream: Any, + ) -> None: + error = ctypes.create_string_buffer(512) + result = self._library.spark_tp4_vocab_allgather( + self._handle, + ctypes.c_void_p(input_tensor.data_ptr()), + ctypes.c_void_p(output_tensor.data_ptr()), + query_rows, + ctypes.c_void_p(stream.cuda_stream), + error, + len(error), + ) + if result != 0: + message = error.value.decode(errors="replace") + raise RuntimeError( + f"Spark TP4 vocabulary all-gather failed: {message}" + ) + + def capture( + self, + input_tensor: Any, + output_tensor: Any, + query_rows: int, + stream: Any, + ) -> None: + if not self._graph_only: + raise RuntimeError( + "Spark TP4 eager vocabulary session cannot capture" + ) + input_shape = _shape(input_tensor) + output_shape = _shape(output_tensor) + if ( + query_rows not in _SUPPORTED_Q + or input_shape != (query_rows, _VOCAB_PER_RANK) + or output_shape != (query_rows, _OUTPUT_VOCAB) + or str(input_tensor.dtype) != "torch.bfloat16" + or str(output_tensor.dtype) != "torch.bfloat16" + or not bool(input_tensor.is_cuda) + or not bool(output_tensor.is_cuda) + or not bool(input_tensor.is_contiguous()) + or not bool(output_tensor.is_contiguous()) + ): + raise ValueError( + "Spark TP4 vocabulary graph capture requires contiguous " + "CUDA BF16 [Q,38720] -> [Q,154880], Q in [1,6]" + ) + stream_pointer = int(stream.cuda_stream) + if ( + self._capture_stream is not None + and stream_pointer != self._capture_stream + ): + raise ValueError( + "Spark TP4 vocabulary graph session requires one stable " + "capture stream" + ) + error = ctypes.create_string_buffer(512) + result = self._library.spark_tp4_vocab_capture_allgather( + self._handle, + ctypes.c_void_p(input_tensor.data_ptr()), + ctypes.c_void_p(output_tensor.data_ptr()), + query_rows, + ctypes.c_void_p(stream.cuda_stream), + error, + len(error), + ) + if result != 0: + message = error.value.decode(errors="replace") + raise RuntimeError( + "Spark TP4 vocabulary graph capture failed: " + f"{message}" + ) + self._capture_stream = stream_pointer + + def graph_status(self) -> dict[str, object]: + if not self._graph_only: + raise RuntimeError( + "Spark TP4 eager vocabulary session has no graph status" + ) + status = _NativeVocabGraphStatus() + error = ctypes.create_string_buffer(512) + result = self._library.spark_tp4_vocab_get_graph_status( + self._handle, + ctypes.byref(status), + ctypes.sizeof(status), + error, + len(error), + ) + if result != 0: + message = error.value.decode(errors="replace") + raise RuntimeError( + f"Spark TP4 vocabulary graph status failed: {message}" + ) + if status.struct_size != ctypes.sizeof(_NativeVocabGraphStatus): + raise RuntimeError( + "Spark TP4 vocabulary graph status ABI mismatch: " + f"native={status.struct_size} python=" + f"{ctypes.sizeof(_NativeVocabGraphStatus)}" + ) + flags = int(status.flags) + published = int(status.published_sequence) + consumed = int(status.consumed_sequence) + completed = int(status.completed_sequence) + overflow = int(status.overflow_sequence) + return { + "captured_nodes": int(status.captured_nodes), + "published_sequence": published, + "consumed_sequence": consumed, + "completed_sequence": completed, + "overflow_sequence": overflow, + "capture_configured": bool( + flags & _GRAPH_STATUS_CAPTURE_CONFIGURED + ), + "polling_enabled": bool( + flags & _GRAPH_STATUS_POLLING_ENABLED + ), + "host_native_atomics": bool( + flags & _GRAPH_STATUS_HOST_NATIVE_ATOMICS + ), + "submit_affinity_verified": bool( + flags & _GRAPH_STATUS_SUBMIT_AFFINITY_VERIFIED + ), + "progress_affinity_verified": bool( + flags & _GRAPH_STATUS_PROGRESS_AFFINITY_VERIFIED + ), + "submit_cpu": ( + int(status.graph_submit_cpu_plus_one) - 1 + if status.graph_submit_cpu_plus_one + else None + ), + "progress_cpu": ( + int(status.graph_progress_cpu_plus_one) - 1 + if status.graph_progress_cpu_plus_one + else None + ), + "replay_advanced": published > 0, + "replay_caught_up": ( + published > 0 + and published == consumed + and published == completed + ), + "fatal": overflow != 0, + } + + +class _ShadowState: + def __init__(self, output_tensor: Any) -> None: + self.candidate = output_tensor + self.mismatches: Any | None = None + self.count = 0 + self.validated = False + + def observe(self, reference: Any) -> None: + import torch + + mismatch = torch.count_nonzero( + self.candidate.view(torch.uint8) != reference.view(torch.uint8) + ) + if self.mismatches is None: + self.mismatches = mismatch + else: + self.mismatches += mismatch + self.count += 1 + + def mismatch_count(self) -> int: + if self.mismatches is None: + return 0 + return int(self.mismatches.item()) + + +class _Backend: + def __init__(self, rank: int) -> None: + self.rank = rank + self._session: _NativeVocabSession | None = None + self._graph_session: _NativeVocabSession | None = None + self.disabled = False + self.graph_disabled = False + self.shadows: dict[_Signature, _ShadowState] = {} + + def session(self) -> _NativeVocabSession | None: + if self.disabled: + return None + if self._session is None: + try: + self._session = _NativeVocabSession(self.rank) + except Exception: + self.disabled = True + logger.exception( + "disabling Spark TP4 vocabulary before native enqueue " + "because session creation failed" + ) + return None + return self._session + + def prepare_graph(self) -> _NativeVocabSession | None: + if self.graph_disabled: + return None + if self._graph_session is None: + try: + graph_cpu_affinity = _graph_preflight() + self._graph_session = _NativeVocabSession( + self.rank, + graph_only=True, + control_ports=_graph_control_ports(), + graph_cpu_affinity=graph_cpu_affinity, + ) + _vocab_graph_sessions[self.rank] = self._graph_session + from spark_graph_status_reporter import ( + ensure_status_reporter, + ) + + ensure_status_reporter(rank=self.rank) + except Exception: + self.graph_disabled = True + logger.exception( + "disabling Spark TP4 vocabulary graph capture before " + "native enqueue because session creation failed" + ) + return None + return self._graph_session + + def shadow( + self, signature: _Signature, output_tensor: Any + ) -> _ShadowState: + state = self.shadows.get(signature) + if state is None: + state = _ShadowState(output_tensor) + self.shadows[signature] = state + return state + + +def _new_output(torch_module: Any, input_tensor: Any, q: int) -> Any: + return torch_module.empty( + (q, _OUTPUT_VOCAB), + dtype=input_tensor.dtype, + device=input_tensor.device, + ) + + +def vocab_graph_status_snapshot() -> dict[int, dict[str, object]]: + """Return process-local vocabulary graph replay status.""" + return { + rank: session.graph_status() + for rank, session in sorted(_vocab_graph_sessions.items()) + } + + +def vocab_graph_diagnostic_snapshot() -> dict[str, object]: + return { + "sessions": vocab_graph_status_snapshot(), + "events": dict(sorted(_graph_event_counts.items())), + } + + +def install() -> None: + global _installed + mode = _mode() + if _installed or not mode: + return + if mode == "custom": + _graph_enabled() + if mode == "shadow": + _positive_integer("SPARK_TP4_VOCAB_SHADOW_COLLECTIVES", 8) + validate_active_port_namespace() + + from vllm.distributed.parallel_state import GroupCoordinator + + original = GroupCoordinator._all_gather_out_place + if getattr(original, "_spark_tp4_vocab_backend", False): + _installed = True + return + + def spark_vocab_all_gather( + self: Any, input_tensor: Any, dim: int + ) -> Any: + mode = _mode() + signature = _signature(self, input_tensor, dim, mode) + if signature is None: + import torch + + _record_stock_path( + capturing=_is_stream_capturing(torch), + reason="ineligible_signature", + group=self, + tensor=input_tensor, + dim=dim, + ) + return original(self, input_tensor, dim) + + import torch + + capturing = _is_stream_capturing(torch) + if capturing: + if mode == "custom" and _graph_enabled(): + backend = getattr( + self, "_spark_tp4_vocab_native", None + ) + graph_session = ( + None if backend is None else backend._graph_session + ) + if graph_session is not None: + candidate = _new_output( + torch, input_tensor, signature + ) + stream = torch.cuda.current_stream( + device=input_tensor.device + ) + try: + graph_session.capture( + input_tensor, + candidate, + signature, + stream, + ) + _record_graph_event(self, "captured_nodes") + return candidate + except BaseException: + logger.exception( + "fatal Spark TP4 vocabulary graph-capture " + "failure; terminating worker because a partial " + "native graph cannot safely fall back" + ) + _abort_after_native_failure() + raise AssertionError( + "unreachable after worker termination" + ) + _record_graph_event(self, "cold_fallbacks") + logger.critical( + "fatal Spark TP4 vocabulary graph session is absent " + "during custom capture; terminating worker to prevent " + "a rank-split collective" + ) + _abort_after_native_failure() + raise AssertionError( + "unreachable after worker termination" + ) + else: + _record_stock_path( + capturing=True, + reason="graph_transport_disabled", + ) + return original(self, input_tensor, dim) + + backend = getattr(self, "_spark_tp4_vocab_native", None) + if backend is None: + backend = _Backend(int(self.rank_in_group)) + self._spark_tp4_vocab_native = backend + if mode == "custom" and _graph_enabled(): + graph_session = backend.prepare_graph() + if graph_session is None: + logger.critical( + "fatal Spark TP4 vocabulary graph session creation " + "failed in custom mode; terminating worker to prevent " + "a rank-split collective" + ) + _abort_after_native_failure() + raise AssertionError( + "unreachable after worker termination" + ) + + session = backend.session() + if session is None: + if mode == "custom": + logger.critical( + "fatal Spark TP4 vocabulary eager session creation " + "failed in custom mode; terminating worker to prevent " + "a rank-split collective" + ) + _abort_after_native_failure() + raise AssertionError( + "unreachable after worker termination" + ) + _record_stock_path( + capturing=False, + reason="native_session_unavailable", + ) + return original(self, input_tensor, dim) + + shadow_limit = 0 + shadow = None + promoted = False + if mode == "shadow": + shadow_limit = _positive_integer( + "SPARK_TP4_VOCAB_SHADOW_COLLECTIVES", 8 + ) + shadow = backend.shadows.get(signature) + if shadow is None: + template = _new_output(torch, input_tensor, signature) + shadow = backend.shadow(signature, template) + promoted = shadow.validated and ( + os.getenv("SPARK_TP4_VOCAB_SHADOW_PROMOTE", "0") == "1" + ) + if not promoted and shadow.count >= shadow_limit: + _record_stock_path( + capturing=False, + reason="shadow_reference_only", + ) + return original(self, input_tensor, dim) + + candidate = ( + _new_output(torch, input_tensor, signature) + if mode == "custom" or promoted + else shadow.candidate + ) + stream = torch.cuda.current_stream(device=input_tensor.device) + try: + session.all_gather( + input_tensor, candidate, signature, stream + ) + except BaseException: + logger.exception( + "fatal Spark TP4 vocabulary failure; terminating worker " + "because native enqueue may have poisoned its CUDA stream" + ) + _abort_after_native_failure() + raise AssertionError("unreachable after worker termination") + + if mode == "custom" or promoted: + return candidate + + try: + _record_stock_path( + capturing=False, + reason="shadow_reference", + ) + reference = original(self, input_tensor, dim) + assert shadow is not None + shadow.observe(reference) + except BaseException: + logger.exception( + "fatal failure after Spark TP4 vocabulary enqueue; " + "terminating worker" + ) + _abort_after_native_failure() + raise AssertionError("unreachable after worker termination") + + if shadow.count == shadow_limit: + mismatches = shadow.mismatch_count() + logger.warning( + "Spark TP4 vocabulary shadow complete: q=%d " + "input_bytes=%d collectives=%d byte_mismatches=%d", + signature, + signature * _VOCAB_PER_RANK * _BF16_BYTES, + shadow_limit, + mismatches, + ) + if mismatches: + raise RuntimeError( + "Spark TP4 vocabulary shadow found byte mismatches" + ) + shadow.validated = True + if ( + os.getenv("SPARK_TP4_VOCAB_SHADOW_PROMOTE", "0") + == "1" + ): + logger.warning( + "Spark TP4 vocabulary Q%d will promote to custom on " + "its next call", + signature, + ) + return reference + + spark_vocab_all_gather._spark_tp4_vocab_backend = True # type: ignore[attr-defined] + spark_vocab_all_gather._spark_original = original # type: ignore[attr-defined] + GroupCoordinator._all_gather_out_place = spark_vocab_all_gather + _installed = True + logger.warning( + "installed Spark TP4 vocabulary backend in %s mode", _mode() + ) diff --git a/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/sparkring-overlay-manifest.json b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/sparkring-overlay-manifest.json new file mode 100644 index 00000000..069c85bd --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/transport/bundle-source/sparkring-overlay-manifest.json @@ -0,0 +1,157 @@ +{ + "b12x_git": { + "commit": "eac260a8257cc6b14e7d4ad674f51e9a09b8790f", + "roce_source_dirty": true, + "roce_status": [ + " M b12x/comm/roce/_oneshot_cute.py", + " M b12x/comm/roce/roce_oneshot.py" + ] + }, + "b12x_roce_tree_sha256": "8a64063506c9c96d3fd3b004c3ab51dda88408f7246f5c89d0a0b6713065fca6", + "base_sircl_manifest_name": "sparkring-overlay-manifest.json", + "base_sircl_manifest_sha256": "85a231e6d2a290f7d6cccbc2cc6b1ccad7a6adbefc7ce4dde05b158f249aadd4", + "files": [ + { + "path": "b12x_overlay/b12x/comm/roce/__init__.py", + "role": "b12x_comm_roce", + "sha256": "d1eabaa9395bbdc568d503c79cc1a57ab74626feed0c67851a00be0ba2718a7e" + }, + { + "path": "b12x_overlay/b12x/comm/roce/_allgather_cute.py", + "role": "b12x_comm_roce", + "sha256": "e26ab2e7cce1cff334452ea237b1c0db8ab11ecba8e0a933f5cad35075a37f0f" + }, + { + "path": "b12x_overlay/b12x/comm/roce/_cute_intrinsics.py", + "role": "b12x_comm_roce", + "sha256": "ff0847d986243bf92444cd4c3db9156a6fc10d8c2e1a16d54761375e6aa7a78e" + }, + { + "path": "b12x_overlay/b12x/comm/roce/_oneshot_cute.py", + "role": "b12x_comm_roce", + "sha256": "1dd563020d15e50c7b391aac8bbee2e2760e66ced1bb891de7e59da5cc477173" + }, + { + "path": "b12x_overlay/b12x/comm/roce/_path_config.py", + "role": "b12x_comm_roce", + "sha256": "7c5550b16a70d74216970cd08a8d969a9928a975d1efcb9fb70e247eb0322e3a" + }, + { + "path": "b12x_overlay/b12x/comm/roce/_proxy.py", + "role": "b12x_comm_roce", + "sha256": "3ace33c00f64ce43317a00b19879bbc6562bb3e466aa94683ed95ce211ec25c0" + }, + { + "path": "b12x_overlay/b12x/comm/roce/_roce_proxy.c", + "role": "b12x_comm_roce", + "sha256": "b208f07d4bb12613a6aef3e8e59ac334b53028655b76f3f5c6e531bd354b369e" + }, + { + "path": "b12x_overlay/b12x/comm/roce/api.py", + "role": "b12x_comm_roce", + "sha256": "02becd3757e71c917a71e98ab3f779d52bb18e6801b8efebd792ab186277ea46" + }, + { + "path": "b12x_overlay/b12x/comm/roce/roce_oneshot.py", + "role": "b12x_comm_roce", + "sha256": "2e431e7c5ef714636247b81602fdc1b85b38ebe1cef7665bc8d486f7841a65e1" + }, + { + "path": "libspark_transport_capi.so", + "role": "base_sircl", + "sha256": "056243fad27d224b82e437925ffa2aed42037e6bd29f239f56076a832f6ca5cb" + }, + { + "path": "rocenante-overlay-config.json", + "role": "generated_contract", + "sha256": "69535f64b484699695cea587b24343af9346149223b50e17b28a8c56d9ad3dee" + }, + { + "path": "rocenante_health_gate.py", + "role": "glm53_rocenante_overlay", + "sha256": "aac0ad50f2065f596e8173ddb4d07984a644fbd786d38604ab5f0b4a042feade" + }, + { + "path": "rocenante_vllm_overlay.py", + "role": "glm53_rocenante_overlay", + "sha256": "0cf6c21cc0baa06dd141fe9660081fbbf8aa56821fc2fe43e393f8b5c8877b82" + }, + { + "path": "sircl_sitecustomize.py", + "role": "base_sircl", + "sha256": "4efc3fdfe4874820064dbcb713c02271df4ca41f4f2bbcf4683a1548ca02f2b1" + }, + { + "path": "sitecustomize.py", + "role": "glm53_rocenante_overlay", + "sha256": "ebd7f5e23de68ee0062c1c2c022a7d4332aa3c998a835e7aff91fdd6be954211" + }, + { + "path": "spark_collective_audit.py", + "role": "base_sircl", + "sha256": "8fa7aa9747845a84b40e8d3a0d72c39f8a562d6490764b777e5dcbe1cbe47721" + }, + { + "path": "spark_cudagraph_bucket_contract.py", + "role": "base_sircl", + "sha256": "e7411ff9cde9d5ed4a7ecd6e78a3158512c188c254baad10554b72229f14d272" + }, + { + "path": "spark_cudagraph_replay_timing.py", + "role": "base_sircl", + "sha256": "dc3ee0fd5da52e79b097e9edc6191a880ca5a3f4e6d83add2a4acd30966b2175" + }, + { + "path": "spark_dcp_collective_audit.py", + "role": "base_sircl", + "sha256": "077a234e4edff8b8dd44784953aef713884b4dd7a3f7c46589b14c6bb8b40745" + }, + { + "path": "spark_graph_status_reporter.py", + "role": "base_sircl", + "sha256": "4bd50d1be42e2ac750b382825bec8a49564ce288d376852ee7c51e33a303f277" + }, + { + "path": "spark_persistent_output_ring.py", + "role": "base_sircl", + "sha256": "03220c78ed87651f93c3165f1eedb843ffcad6a116126e2a9b89beab299bd2f3" + }, + { + "path": "spark_tp4_backend.py", + "role": "base_sircl", + "sha256": "5999814740ddfa1c35e925ed0f3bf43edd8335b30c9b4fe9b78006c618d18cb9" + }, + { + "path": "spark_tp4_capability.py", + "role": "base_sircl", + "sha256": "d90f6fc07ff5255b2df25b3ffc1c52c83f9eb866ca68666bc0f273db1c5e83b7" + }, + { + "path": "spark_tp4_health_gate.py", + "role": "base_sircl", + "sha256": "7651b88e9bff40193891a0810ee8a920e3548a17de3a0dc477a30db779738930" + }, + { + "path": "spark_tp4_port_namespace.py", + "role": "base_sircl", + "sha256": "c789739811def9722753bb128e69e7bd5342bab827b2cfd99cd8a4e847d18d4f" + }, + { + "path": "spark_tp4_query_contract.py", + "role": "base_sircl", + "sha256": "978df868fd83a6a780111092e99de2a3f07a5df375cf52d4b2358ccf4dcd0c31" + }, + { + "path": "spark_tp4_query_row_provider.py", + "role": "base_sircl", + "sha256": "09c9e7af5d14b5c6930e15c999ada2231d639eb59609d705aadecdd7497ddec7" + }, + { + "path": "spark_tp4_vocab_allgather_backend.py", + "role": "base_sircl", + "sha256": "fd44dea558c115b17f17bb129033345af4543702fbbdaf1fa4f186b8f242e697" + } + ], + "schema": "sparkring.glm53-rocenante-private-bundle/v1", + "status": "research-only" +} diff --git a/runtime/glm53-spark-mtp3-mesh/performance/transport/native-source.tar.gz b/runtime/glm53-spark-mtp3-mesh/performance/transport/native-source.tar.gz new file mode 100644 index 00000000..06ac7066 Binary files /dev/null and b/runtime/glm53-spark-mtp3-mesh/performance/transport/native-source.tar.gz differ diff --git a/runtime/glm53-spark-mtp3-mesh/performance/transport/package.py b/runtime/glm53-spark-mtp3-mesh/performance/transport/package.py new file mode 100644 index 00000000..b35af1ae --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/transport/package.py @@ -0,0 +1,119 @@ +#!/usr/bin/env python3 +"""Verify transport sources or prepare an offline source tree and overlay bundle.""" +from __future__ import annotations + +import argparse +import hashlib +import io +import json +from pathlib import Path, PurePosixPath +import re +import tarfile + +ROOT = Path(__file__).resolve().parent +MANIFEST = "sparkring-overlay-manifest.json" +LIBRARY = "libspark_transport_capi.so" + + +def digest(data: bytes) -> str: + return hashlib.sha256(data).hexdigest() + + +def checked_name(name: str) -> str: + path = PurePosixPath(name) + if not name or path.is_absolute() or ".." in path.parts or "\\" in name or str(path) != name: + raise ValueError("Archive or manifest contains an unsafe path") + return name + + +def verified_sources(root: Path = ROOT) -> tuple[dict, dict[str, bytes], dict[str, bytes]]: + receipt = json.loads((root / "source-manifest.json").read_bytes()) + archive = (root / "native-source.tar.gz").read_bytes() + if digest(archive) != receipt["native_source_archive_sha256"]: + raise ValueError("Native source archive digest differs") + sources = {} + with tarfile.open(fileobj=io.BytesIO(archive), mode="r:gz") as stream: + for member in stream: + name = checked_name(member.name) + if not member.isfile() or name in sources: + raise ValueError("Archive contains duplicate or non-regular entries") + sources[name] = stream.extractfile(member).read() + if {n: digest(b) for n, b in sources.items()} != receipt["source_files"]: + raise ValueError("Native source inventory differs") + bundle_root = root / "bundle-source" + bundle = {} + for path in bundle_root.rglob("*"): + if path.is_symlink(): + raise ValueError("Bundle source contains a symbolic link") + if path.is_file(): + bundle[checked_name(path.relative_to(bundle_root).as_posix())] = path.read_bytes() + if {n: digest(b) for n, b in bundle.items()} != receipt["bundle_files"]: + raise ValueError("Bundle source inventory differs") + if digest(bundle[MANIFEST]) != receipt["reference_bundle_manifest_sha256"]: + raise ValueError("Reference bundle manifest digest differs") + document = json.loads(bundle[MANIFEST]) + listed = {row["path"]: row["sha256"] for row in document["files"]} + actual = {n: digest(b) for n, b in bundle.items() if n != MANIFEST} + actual[LIBRARY] = receipt["reference_native_sha256"] + if len(listed) != len(document["files"]) or listed != actual: + raise ValueError("Reference bundle content differs from its manifest") + return receipt, sources, bundle + + +def write_tree(destination: Path, files: dict[str, bytes]) -> None: + if destination.exists(): + raise ValueError("Output directory must not exist") + destination.mkdir(parents=True) + for name, data in sorted(files.items()): + path = destination / checked_name(name) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(data) + + +def prepare_bundle(library: Path, expected_sha256: str, destination: Path, + root: Path = ROOT) -> dict: + receipt, _, bundle = verified_sources(root) + if not re.fullmatch(r"[0-9a-f]{64}", expected_sha256): + raise ValueError("Native library digest must be a full lowercase SHA256") + data = library.read_bytes() + if digest(data) != expected_sha256: + raise ValueError("Native library digest differs") + bundle[LIBRARY] = data + matched = expected_sha256 == receipt["reference_native_sha256"] + if not matched: + document = json.loads(bundle[MANIFEST]) + for row in document["files"]: + if row["path"] == LIBRARY: + row["sha256"] = expected_sha256 + document["status"] = "research-only" + bundle[MANIFEST] = (json.dumps(document, indent=2, sort_keys=True) + "\n").encode() + write_tree(destination, bundle) + return {"status": "implemented" if matched else "research-only", + "reference_artifact_match": matched, + "native_sha256": expected_sha256, + "bundle_manifest_sha256": digest(bundle[MANIFEST]), + "hardware_qualification_performed": False} + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + commands = parser.add_subparsers(dest="command", required=True) + commands.add_parser("verify") + extract = commands.add_parser("extract") + extract.add_argument("--output", type=Path, required=True) + bundle = commands.add_parser("bundle") + bundle.add_argument("--native-library", type=Path, required=True) + bundle.add_argument("--native-sha256", required=True) + bundle.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + receipt, sources, _ = verified_sources() + if args.command == "extract": + write_tree(args.output, sources) + elif args.command == "bundle": + print(json.dumps(prepare_bundle(args.native_library, args.native_sha256, args.output), sort_keys=True)) + return + print(json.dumps({"verified": True, "source_archive_sha256": receipt["native_source_archive_sha256"]}, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/runtime/glm53-spark-mtp3-mesh/performance/transport/source-manifest.json b/runtime/glm53-spark-mtp3-mesh/performance/transport/source-manifest.json new file mode 100644 index 00000000..6da6f3ac --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/transport/source-manifest.json @@ -0,0 +1,187 @@ +{ + "schema": "sparkring-mtp3-transport-source/v1", + "status": "implemented", + "native_source_archive_sha256": "104096807f3510fe7d01f6ba455a435d2893f901a72c3424d3e30fe9a3adc79e", + "source_files": { + "LICENSE": "3ddf9be5c28fe27dad143a5dc76eea25222ad1dd68934a047064e56ed2fa40c5", + "spark_transport/CMakeLists.txt": "f0ace9039108d1444c9928a8c8f99789b3d8d0522612e31343742703751e8048", + "spark_transport/app/tp4_bidirectional_prefill_probe.cu": "113452bf0c3e111f8f3d84bd1336eebc7c5519b124b4b5833279a59bea78f1d9", + "spark_transport/app/tp4_fused_prefill_probe.cu": "bd57e42e249f1c59031b11d10b3beed569d0ec7eeeb1334774361913be01698f", + "spark_transport/app/tp4_fused_ring_occupancy_probe.cu": "88d48ca15858d1b2e92987544a8e297c32a885db7619b32f3f64758154bdd8bd", + "spark_transport/app/tp4_fused_stream_switch_smoke.cu": "03ef4477467e60979ae569498c13786caa9b12aea46a6139312ee63d7d082759", + "spark_transport/app/tp4_graph_q1_probe.cu": "a34323764f00d418c8768d65b58b0166571b5e929c147d75ea978cd7643ceef3", + "spark_transport/app/tp4_probe.cpp": "aab6fd79810c705ae8750993738feabe00bf83f0aadf4a24b038f36cc3a2b0cf", + "spark_transport/app/tp4_tensor_probe.cu": "3f24db64387df43ad456f25a0186a2c97cf74afa22fd7bd3eb4140c2fd7109b2", + "spark_transport/app/tp4_tiled_prefill_probe.cu": "d8449543bad6305da059b984922962e2065c70a30993a9b680cdc787d00ca44f", + "spark_transport/app/tp4_vocab_allgather_probe.cu": "c02997123f25c6cf3f332d76c88ea2be560011ba8bdf61bcae679dc69b647cd4", + "spark_transport/app/tp4_vocab_graph_probe.cu": "391a515069b7643b3ccf1dd36eb4ac7bd49368b90c83c71b823505941d985cd0", + "spark_transport/app/transport_probe.cpp": "8c80d2a14991af428536e3e7a1d798ce4b4e5d9f82d455d1870eeeeeca3e5eab", + "spark_transport/experiments/tiled_prefill/CMakeLists.txt": "248a4c1d2406826448cec5357e9666b6ff1e010e9ba1f8fc214e4e8098bc7923", + "spark_transport/experiments/tiled_prefill/__init__.py": "b7df0bd191d2ba1216d74586913ffc58c75efb90738d2229cac975c9200d2955", + "spark_transport/experiments/tiled_prefill/bidirectional_bulk_abi.hpp": "bf92c5a91fd3759fa3545cde9681e3189db5d81f8abf2da0733b434e37deb314", + "spark_transport/experiments/tiled_prefill/bidirectional_bulk_cuda_smoke_test.cu": "facb960531b5f85ece5511d66d8acebd7585c4ae60d112bd445341a2f2a69daa", + "spark_transport/experiments/tiled_prefill/bidirectional_bulk_kernels.cu": "a7bafbfb67300c085aa6145ee52d5c39c25d7431a35bc56ac8eae0c533d49eb3", + "spark_transport/experiments/tiled_prefill/bidirectional_ring_executor.cpp": "2e2f436fcc7dff4a3f8cf29aadeda2fb28295792db3bbe0a42f3a67b6b1bc2cf", + "spark_transport/experiments/tiled_prefill/bidirectional_ring_executor.hpp": "310d9fcdf479c36ad4474f72bd1f314c2686f41d5dc2417eabec7e41068aead0", + "spark_transport/experiments/tiled_prefill/correctness_oracle.py": "b5ee829b5904502734be999470446a7a772bbba1deb06c5ad1391f64c068c5a8", + "spark_transport/experiments/tiled_prefill/fused_prefill_abi.hpp": "548ee9e7d1619ba3b0a681d184e60a996dac5d7149ea8cfd419a29dcf68d89c2", + "spark_transport/experiments/tiled_prefill/fused_prefill_kernels.cu": "9b885cad147415613b5d44c020d077559b212c1fb70044e632dbbdeaf250cf77", + "spark_transport/experiments/tiled_prefill/fused_prefill_proxy_smoke_test.cu": "543b4cc5a825241cb1fa501d619e8a2f562740f3c65ae7aa36aa930330466b9d", + "spark_transport/experiments/tiled_prefill/fused_prefill_verbs_proxy.cpp": "8cfe27bbf0ac0134feb74e06a2eed2ddc61ff905779c0b6ab243a04a2d30b811", + "spark_transport/experiments/tiled_prefill/fused_prefill_verbs_proxy.hpp": "09bdf9f6300f212a10c6c7eef6a6e0bc38ba102c5be005a39f5e453f900e4343", + "spark_transport/experiments/tiled_prefill/fused_ring_simulator.py": "23e0478b279a277316aecf03770ba3c19935da0b924aeef8bfa21e1c3111038f", + "spark_transport/experiments/tiled_prefill/gpu_harness.py": "ac339de838c8aab0af7a1446e4ec40ff3b1b9cb0326d5459f295e304b22a9629", + "spark_transport/experiments/tiled_prefill/qualification.py": "6200d59b7892e7042876e259172985cd68024c6490b013aa05d68518e1f22346", + "spark_transport/experiments/tiled_prefill/substrate.py": "991ba949131e62b8fe7c103972dc93ca0372403e98fbb44482ba3d8c823294f2", + "spark_transport/experiments/tiled_prefill/test_fused_ring_simulator.py": "7db080e743383a879bf23ff0bc94e395fea3fe820139317df1a3dc44eea26e4c", + "spark_transport/experiments/tiled_prefill/test_gpu_harness.py": "cc50565c584ffac60e1ef4e64a314c5557a230f088797b9917b63e3a3f9b39e8", + "spark_transport/experiments/tiled_prefill/test_qualification.py": "a54dcad85fb5c8193b0e8485e1380214f51ee5875530aa90374a79e7529b7dd8", + "spark_transport/experiments/tiled_prefill/test_substrate.py": "9f45e9e68e16baca198f657b0573bf9f7d04d9a15f82cc4d41496e5353e69725", + "spark_transport/experiments/tiled_prefill/test_tiled_correctness_oracle.py": "ff085c9675e89ca4523b4c0deea463c02527348a4abd550645c367ede4149588", + "spark_transport/experiments/tiled_prefill/tiled_bulk_abi.hpp": "82dea293646b9df378477c7fa33475d7f15f0162f79be78d046323ed2416c030", + "spark_transport/experiments/tiled_prefill/tiled_bulk_kernels.cu": "de4266c99f5dca47c096bec4f90eb7e1a3edc9a02037fa8a6f3e681514eebd29", + "spark_transport/experiments/tiled_prefill/tiled_correctness_kernels.cu": "b92ece744a553fcb7cf3950632c21a264c7df160abec2f590d3944c6da0d3f67", + "spark_transport/experiments/tiled_prefill/tiled_correctness_oracle.hpp": "927e2c58be2b99783b866896663a5984f71d784afcd5006188dcce6624f29421", + "spark_transport/experiments/tiled_prefill/tiled_correctness_oracle_test.cpp": "6113e25248e98ea69f8b7a45f4c5c40e6401f595245021393513ef39c209f999", + "spark_transport/experiments/tiled_prefill/tiled_cuda_smoke_test.cu": "fddab7aef0eea5582d008225321f11febdd2626b57e87d7690377550f312ecf0", + "spark_transport/experiments/tiled_prefill/tiled_dual_rail_edge.cpp": "f28423206025b9dae8574c14331c8c8e6dfff2b8afdf35ee26fc50c974fbd393", + "spark_transport/experiments/tiled_prefill/tiled_dual_rail_edge.hpp": "b1c2986e0191471c93947f516290296fce0b6e1f86139b2c32ba0a07e95a7625", + "spark_transport/experiments/tiled_prefill/tiled_executor.cpp": "6a8de0feed23cfd862dbe61d450c76b13d340193d9ef98d27d3f64f862c23501", + "spark_transport/experiments/tiled_prefill/tiled_executor.hpp": "da9ac848ea0483d9b4424dae2a0c1354402c9051a759af1ec0fc57a0241ba17a", + "spark_transport/experiments/tiled_prefill/tiled_signaled_work_gate.hpp": "8d55253f091fa915d6af6d59e4cfc29ccee870e2cdeda35ba6c9ecce593083c8", + "spark_transport/include/spark_transport/control_channel.hpp": "bc0c3820a9821669d3f9a99623cca2782acbbbdba3b274babbc985c6f457bb70", + "spark_transport/include/spark_transport/eager_staging_timeout.hpp": "d502761e2492c9a92ebd3d8e45cbd4f75dcb48430f4e5a42e19d71aa9dcbb9fb", + "spark_transport/include/spark_transport/gpu_doorbell.hpp": "89efdeceaa925e461201e08af090fc4293abf9261d596cc2363332eed58956ef", + "spark_transport/include/spark_transport/gpu_tp4.hpp": "e1e4bfb3d5aa2c5050d41d7473add1c132ea7ff4aa2f6dd40db261d145335fc5", + "spark_transport/include/spark_transport/gpu_tp4_tensor.hpp": "159fce82e5be6cd24606d0e75369326ce703a0777d722736bb43227ade896d69", + "spark_transport/include/spark_transport/gpu_tp4_vocab_allgather.hpp": "c2057fd97d3cf7a76c81538bc6b474cbf2c56bbfc393b15a6b0a9b7e0c945c0a", + "spark_transport/include/spark_transport/graph_poll_policy.hpp": "0ee9f166120714a999433d2c9d67ff30df3d85151a0e932118b725c40d60a2c7", + "spark_transport/include/spark_transport/memory_buffer.hpp": "c9ab8d42545f9f529c1a29b156e34f46d36a1d0d2d466dbe981ea9509188aff5", + "spark_transport/include/spark_transport/statistics.hpp": "8d5441e2389b4e60fe9e44ca400e51074cdc0f46ace85c4c56e7ee1cb3a8112e", + "spark_transport/include/spark_transport/topology.hpp": "57834f046c43dc261bc3feeffd11764f02e43580170c1c03c0507e2330dc8a83", + "spark_transport/include/spark_transport/tp4_allreduce_protocol.hpp": "555049e41220b52ae373114450301343783fc3ad7e1360ebcc3ccafa3ab8f20c", + "spark_transport/include/spark_transport/tp4_bidirectional_prefill.hpp": "4c1f95342a843dd52b1bb5b7cd5ecbccad7975d625475a01dedfeb248bfa0540", + "spark_transport/include/spark_transport/tp4_bidirectional_prefill_c_api.h": "29e2beddd894e8734799170c90235f197245464a20d3369ae12a3d3bc1fee373", + "spark_transport/include/spark_transport/tp4_bidirectional_prefill_session.hpp": "dbae7cdc3fb98e7a7ac84f8f2006e631c53226448a8e2f7ea46a89d68e5bb15b", + "spark_transport/include/spark_transport/tp4_c_api.h": "4981cb036669fff2efaa1a3ff8731e204cf7d2b37f4c4ee7285a5e0fc85762ea", + "spark_transport/include/spark_transport/tp4_dual_port_striped_allreduce.hpp": "3f3d3d819478d61be542958fccff433ef3fb89e34518887fb621437ded79b65e", + "spark_transport/include/spark_transport/tp4_fused_prefill_c_api.h": "b9b416d4d13e75025024bb18f14f27202a22f7f26ccf9821eba4c680a8e7a5e5", + "spark_transport/include/spark_transport/tp4_fused_prefill_session.hpp": "c2e63e46c7ce501ba828cc41eed6bf863e8bd8c97f4afe7cad7a18e4000affa5", + "spark_transport/include/spark_transport/tp4_fused_ring_contract.hpp": "3f8a911f3baafc5e98ae60003bc39c80e43f4aac44b680898ca929e92f08ddc3", + "spark_transport/include/spark_transport/tp4_graph_command.hpp": "016931e746d09e3ffa4fd85f7c9b310a1b4439220d7ed01b49deb71ba3ae54fe", + "spark_transport/include/spark_transport/tp4_graph_kernel_strategy.hpp": "4b7816875b645cef8c899033dff793c29930babe1f86f6b15707aac974977423", + "spark_transport/include/spark_transport/tp4_schedule.hpp": "bc028046bd997d2b597b8514e03f2aa5779e908dddae9451ed43e9fd6dc3514d", + "spark_transport/include/spark_transport/tp4_session.hpp": "4bd3611bf846fcda7af03c5f36e4ef266afde6494e2ea919cc57bb1d28abb148", + "spark_transport/include/spark_transport/tp4_tiled_session.hpp": "1cef667950abbd450dde446c11627473acfa9628a1bdb47a66995a3beeec555c", + "spark_transport/include/spark_transport/tp4_vocab_allgather_c_api.h": "64ab46e4a7fa34708a7448384a62711b5e4276304c94cd5dbec4f70119084ff8", + "spark_transport/include/spark_transport/tp4_vocab_allgather_session.hpp": "c2d90e49f746ae693a8939090b437865fa0a61381cbec359dc9d2d78d6edf11c", + "spark_transport/include/spark_transport/verbs_endpoint.hpp": "d7fe42545d045c74103a11a0ec763e1ed3c58532c4b5638fbddd08c2695320ab", + "spark_transport/src/control_channel.cpp": "a900f4a3490d86366b7865070e6e7dd4ad911d21aa1ea9474cf648f9ab0d42d0", + "spark_transport/src/cuda_event_gate.hpp": "dc499572d8ffe13324d8beaedcf62e6f2d975363262267b83a97e9e18df8e9e1", + "spark_transport/src/cuda_stream_handoff.hpp": "2e06c478d5f50b00777537edc130b322c09f4c24cc81defe1ca95115c3fc916d", + "spark_transport/src/gpu_doorbell.cu": "f59f31b574b0e6858f3ce53817b94941ae99781c419cf9d0e54dc1dfa2958158", + "spark_transport/src/gpu_tp4.cu": "706dbcbf4f8f736b5006a2e8dbef5a14c3bed53d4ee5510e037e90c774ab770f", + "spark_transport/src/gpu_tp4_tensor.cu": "5f9ac0028260d7d369cf2e3213fabeb362677efecb7b50b66c3f19e7776cf610", + "spark_transport/src/gpu_tp4_vocab_allgather.cu": "7247d48b15fdad25108d3aaa8458172fe2f40837bd57577f542d64f0a5b0e578", + "spark_transport/src/memory_buffer.cu": "b09c2a208c0c4bd0c86e6f3283f56644f52205e395d0309712be0ef064084b1c", + "spark_transport/src/statistics.cpp": "2868979eda911798899e1d9fd515eb1eca8a7b37e35d3c06cc88d0520f6f789c", + "spark_transport/src/topology.cpp": "c29f9b8b59cbdd624d2bb21d52c43bf73922e3fc291db3de62bb001550089e52", + "spark_transport/src/tp4_bidirectional_prefill_c_api.cpp": "0f9dd3135c4cf940b79e9f84a5a7d9ebeaa7053d316d2e5de1b7ad964b99754b", + "spark_transport/src/tp4_bidirectional_prefill_session.cpp": "1533bf3d7f39136b9638ae1268a7901cf2d0400db98836ffbf30dacf053800a9", + "spark_transport/src/tp4_c_api.cpp": "afbbcc7387beec1305d4deb133ba16386cf802e5be8efa3fd4e9415191828c17", + "spark_transport/src/tp4_dual_port_striped_host.hpp": "75b4f2ef6eddfc9992efae6ab7a306641baae43e50153979c9a1659c16430a69", + "spark_transport/src/tp4_fused_prefill_c_api.cpp": "a268b6294a2e6fea880b1b9fe5f19bf95eb3de6b3b38aca5723a7f5ed5e697ba", + "spark_transport/src/tp4_fused_prefill_session.cpp": "588a79fd5ec3197668477855bfc9c24bd793f0bf426ab0fe7928ea4b10163967", + "spark_transport/src/tp4_graph_command.cpp": "fd2b5306ce00d00a4d893a96e909173f822932739baf0953225369c74a62da38", + "spark_transport/src/tp4_schedule.cpp": "fd4f5de1c3ae4d3d3e3626800e174585a43762817beaf49e4f792a05bc66ba73", + "spark_transport/src/tp4_session.cpp": "4a2c0c688de179b68d2dc6dfb1ab52e62efdfb7c19e4a2efd22cb6c5069440ba", + "spark_transport/src/tp4_vocab_allgather_c_api.cpp": "19598f1dd38df80166a618ffea9101d9a2f9d50134ee61835cc95df9369f98c6", + "spark_transport/src/tp4_vocab_allgather_session.cpp": "ad99401ea483b139f3fa2037a8693690c913a70287462590cf8adbe5fc1c2738", + "spark_transport/src/verbs_endpoint.cpp": "aa46336db9a38892583d0e8c4cc4c0408456fe0f37eda0a6a88d60cb31db2e8a", + "spark_transport/tests/cuda_event_gate_test.cu": "0f115d69e21e20238d8478a66bc4f4f3cd76f67d1e6f6ca06f88c6a188ce79b3", + "spark_transport/tests/eager_staging_timeout_test.cpp": "024f3e78d0f1cb1898c7884fd86f66ef19783b4fe5094346d92c2392a2482600", + "spark_transport/tests/gpu_doorbell_layout_test.cpp": "4b3b541cf61e444262ff1c8bee5898655f498197b5022fd4bb8c99f084925965", + "spark_transport/tests/graph_poll_policy_test.cpp": "fbefe13f27831135d88da909803d35e3dbd3b21089c948d8a247f4f86390c86a", + "spark_transport/tests/statistics_test.cpp": "5ff8d739be15275e88b2f334998f42240bd1e665e9279c0baabc491d43a5b21d", + "spark_transport/tests/syntax_stubs/cuda_runtime.h": "15e1617d4f8ef7a2d1473a4e79f3c511fac6e6d38ad99bb7f73108d2720fe338", + "spark_transport/tests/syntax_stubs/infiniband/verbs.h": "2af1c2d5bf7aeca0ef0fe43a90c07eb366f9f1d9e0156cb1f1178cdc89a51458", + "spark_transport/tests/test_topology_schedule_source_contract.py": "a1d4c1d415fa955fa567d9f130821b2bf46c2694f86809d6fd9fb5f0d03c59b3", + "spark_transport/tests/test_tp4_bidirectional_bulk_cuda_source_contract.py": "018f1fc3eee8329b361d007351b6a6dd7bf4138067d4cb9a1e31606915f6f8e3", + "spark_transport/tests/test_tp4_bidirectional_prefill_probe_source_contract.py": "f02d37491eaf158e699ac0a981527687685ce036d713da8b4a989958a0865679", + "spark_transport/tests/test_tp4_bidirectional_prefill_session_source_contract.py": "c64d51ab885f5ec04b0341f103d5bc89d0fd486ed027c3d2eda554bbf7b27355", + "spark_transport/tests/test_tp4_deferred_ack_source_contract.py": "ea565b6092f771890b6277918723af59606f21115eea62fec5c7faa1331df793", + "spark_transport/tests/test_tp4_dual_port_striped_gpu_source_contract.py": "b28f99be72b5c3efb8080806a41dac315c92d9095786f652c2412a825ade9e56", + "spark_transport/tests/test_tp4_dual_port_striped_host_source_contract.py": "9bcb976ecd71a4ec6e86ee0119a544aa7e3d9a5348e5a4565f697e5185e45558", + "spark_transport/tests/test_tp4_fused_prefill_cuda_source_contract.py": "ed0af4f0736b6f622d8067701863a1268bafc458ed656158103440459d5ad7bd", + "spark_transport/tests/test_tp4_fused_prefill_probe_source_contract.py": "f8b725db1af42554b57f961b7affcac42b8e67d941660811e0b18a9089ba4102", + "spark_transport/tests/test_tp4_fused_prefill_proxy_smoke_source_contract.py": "66efb6f94bdbe96d7255903847de1fab31ce12ea06c939411bc2600329bd923d", + "spark_transport/tests/test_tp4_fused_prefill_session_source_contract.py": "8c570fde32a14bb73fe8afca8db7ebbcd372ff211de514132c804dd140665fa2", + "spark_transport/tests/test_tp4_fused_prefill_verbs_proxy_source_contract.py": "f4b889f5709dfb06750b45fd9ef750c58db9b9b8c8123bf14bad4759f33b0981", + "spark_transport/tests/test_tp4_graph_prefill_probe_contract.py": "96eb38e1d68eae6c442e1c7cbb2c3fcb6d9d330d3428511cbf2833a209416808", + "spark_transport/tests/test_tp4_health_source_contract.py": "f70c6ef55ffdbbfa5925652e91797ca2a0575e0db1f1dce3f0db20fe169137af", + "spark_transport/tests/test_tp4_split_graph_source_contract.py": "268c5a25cc284e7dc52be513d09699c1853a158b1992b0cc9ea938d7f81b352c", + "spark_transport/tests/test_tp4_tiled_prefill_native_probe_source_contract.py": "3494978e8fa23074d5076ffbe21847ef9be8921cf11600ccd2fded752f0b1b50", + "spark_transport/tests/topology_test.cpp": "72d65643df8cc3bb845404b45f24356c6d9ae761702832f6b99a31a1a3af43c5", + "spark_transport/tests/tp4_allreduce_protocol_test.cpp": "3fe05ca4cf25899efbbeceee37708ae636a36d2a9a1b09d74182a4bc5eeb8af4", + "spark_transport/tests/tp4_bidirectional_prefill_c_api_layout_test.cpp": "28b4509cccefba0d8fb9cdef20bc19e2448905f2d8caeec2d6b01a6408350507", + "spark_transport/tests/tp4_bidirectional_prefill_test.cpp": "1d31e58fa1b290efe5e4ca24487cb799016f1e312e4cd8440d8ea80a526e1248", + "spark_transport/tests/tp4_bidirectional_ring_executor_test.cpp": "638ceeccf51b0c6a5abc0e79a88ab3ba5e721ee4cde31467a44fcdbbf0778f64", + "spark_transport/tests/tp4_c_api_layout_test.cpp": "306b8273899de96640d9d2e9d592435f2b100e9126a7368ad5adc9233c097498", + "spark_transport/tests/tp4_c_api_test.cpp": "e89f5195f0dd6cbc1ecc325639cf962e4ca9c4d0abd3ee7f96f55079083da739", + "spark_transport/tests/tp4_dual_port_striped_allreduce_test.cpp": "24492536225001852f3d4d415c243601e8a3a3381505655ca8e51841a3eeddd7", + "spark_transport/tests/tp4_dual_port_striped_host_test.cpp": "0d963fe393fc06e72a3dbd4ef21ab357a0bbd180124894aa32e1e30a12bd8126", + "spark_transport/tests/tp4_fused_prefill_c_api_layout_test.cpp": "019e6cd909a638197498b7686f9ff17c99bfd5d14ac235ae707115fe8592e085", + "spark_transport/tests/tp4_fused_prefill_layout_test.cpp": "fdeee67a1159856ab51942287c66015a83bf40c25dbda72f8a18947431644f21", + "spark_transport/tests/tp4_fused_ring_contract_test.cpp": "45c9ba24938f68e76c066536766afeafdea922b2186e991dc1d7ad0b8f405c67", + "spark_transport/tests/tp4_graph_command_test.cpp": "486d312bb11b0f9219b2298ff6d026f4b17ae67523ac3d9d730a6b0b6ca3a0b3", + "spark_transport/tests/tp4_graph_kernel_strategy_test.cpp": "39f3dc393ada915bf330ad9814f907bbfdc9e7c67a40fc254332030c9de2c876", + "spark_transport/tests/tp4_schedule_test.cpp": "8f55dcd67e8e48373a8876619111e31a527bf089d25412776e2293580a6cbcac", + "spark_transport/tests/tp4_tiled_dual_rail_edge_test.cpp": "4e0797096ebc4157dca32512272b0a98c65ec1e2487b9b10f62190edcee8a06d", + "spark_transport/tests/tp4_tiled_executor_test.cpp": "9bf61f37d2e9e099e748f94fd4d587b7f19ff8638004cb415655df61727f0fc7", + "spark_transport/tests/tp4_tiled_session_test.cpp": "4f501abd36b4a7cfa73ea1866754892d7deda6b2ff8f4a0cbd8652dfacd97fec", + "spark_transport/tests/tp4_tiled_signaled_work_gate_test.cpp": "876df67017459680dfbe0196ed81b4faa8ec148e5edaa53920682952d457fa7c", + "spark_transport/tests/tp4_vocab_allgather_c_api_test.cpp": "3392518d6cca3a620f911250630a38ba67f8ff8050811f9ff2c124998ca81f29", + "spark_transport/tests/tp4_vocab_allgather_layout_test.cpp": "abbb0aa7e2e9c14cb321ce77e093a0f406f4347207c483ab3fb280c2e2675364", + "spark_transport/tests/verbs_completion_test.cpp": "92a93fa696c2fa48bb48edc48c122c98eea21703e7396b218c27cd43a28a92eb", + "spark_transport/tests/wire_protocol_test.cpp": "55e60ae68b3ce70f6e1c58d7da1198b0357bf6c0cdb624e3c6632595a750886c" + }, + "bundle_files": { + "b12x_overlay/b12x/comm/roce/__init__.py": "d1eabaa9395bbdc568d503c79cc1a57ab74626feed0c67851a00be0ba2718a7e", + "b12x_overlay/b12x/comm/roce/_allgather_cute.py": "e26ab2e7cce1cff334452ea237b1c0db8ab11ecba8e0a933f5cad35075a37f0f", + "b12x_overlay/b12x/comm/roce/_cute_intrinsics.py": "ff0847d986243bf92444cd4c3db9156a6fc10d8c2e1a16d54761375e6aa7a78e", + "b12x_overlay/b12x/comm/roce/_oneshot_cute.py": "1dd563020d15e50c7b391aac8bbee2e2760e66ced1bb891de7e59da5cc477173", + "b12x_overlay/b12x/comm/roce/_path_config.py": "7c5550b16a70d74216970cd08a8d969a9928a975d1efcb9fb70e247eb0322e3a", + "b12x_overlay/b12x/comm/roce/_proxy.py": "3ace33c00f64ce43317a00b19879bbc6562bb3e466aa94683ed95ce211ec25c0", + "b12x_overlay/b12x/comm/roce/_roce_proxy.c": "b208f07d4bb12613a6aef3e8e59ac334b53028655b76f3f5c6e531bd354b369e", + "b12x_overlay/b12x/comm/roce/api.py": "02becd3757e71c917a71e98ab3f779d52bb18e6801b8efebd792ab186277ea46", + "b12x_overlay/b12x/comm/roce/roce_oneshot.py": "2e431e7c5ef714636247b81602fdc1b85b38ebe1cef7665bc8d486f7841a65e1", + "rocenante-overlay-config.json": "69535f64b484699695cea587b24343af9346149223b50e17b28a8c56d9ad3dee", + "rocenante_health_gate.py": "aac0ad50f2065f596e8173ddb4d07984a644fbd786d38604ab5f0b4a042feade", + "rocenante_vllm_overlay.py": "0cf6c21cc0baa06dd141fe9660081fbbf8aa56821fc2fe43e393f8b5c8877b82", + "sircl_sitecustomize.py": "4efc3fdfe4874820064dbcb713c02271df4ca41f4f2bbcf4683a1548ca02f2b1", + "sitecustomize.py": "ebd7f5e23de68ee0062c1c2c022a7d4332aa3c998a835e7aff91fdd6be954211", + "spark_collective_audit.py": "8fa7aa9747845a84b40e8d3a0d72c39f8a562d6490764b777e5dcbe1cbe47721", + "spark_cudagraph_bucket_contract.py": "e7411ff9cde9d5ed4a7ecd6e78a3158512c188c254baad10554b72229f14d272", + "spark_cudagraph_replay_timing.py": "dc3ee0fd5da52e79b097e9edc6191a880ca5a3f4e6d83add2a4acd30966b2175", + "spark_dcp_collective_audit.py": "077a234e4edff8b8dd44784953aef713884b4dd7a3f7c46589b14c6bb8b40745", + "spark_graph_status_reporter.py": "4bd50d1be42e2ac750b382825bec8a49564ce288d376852ee7c51e33a303f277", + "spark_persistent_output_ring.py": "03220c78ed87651f93c3165f1eedb843ffcad6a116126e2a9b89beab299bd2f3", + "spark_tp4_backend.py": "5999814740ddfa1c35e925ed0f3bf43edd8335b30c9b4fe9b78006c618d18cb9", + "spark_tp4_capability.py": "d90f6fc07ff5255b2df25b3ffc1c52c83f9eb866ca68666bc0f273db1c5e83b7", + "spark_tp4_health_gate.py": "7651b88e9bff40193891a0810ee8a920e3548a17de3a0dc477a30db779738930", + "spark_tp4_port_namespace.py": "c789739811def9722753bb128e69e7bd5342bab827b2cfd99cd8a4e847d18d4f", + "spark_tp4_query_contract.py": "978df868fd83a6a780111092e99de2a3f07a5df375cf52d4b2358ccf4dcd0c31", + "spark_tp4_query_row_provider.py": "09c9e7af5d14b5c6930e15c999ada2231d639eb59609d705aadecdd7497ddec7", + "spark_tp4_vocab_allgather_backend.py": "fd44dea558c115b17f17bb129033345af4543702fbbdaf1fa4f186b8f242e697", + "sparkring-overlay-manifest.json": "c0fd5567442b08b908cc193f36d0864e262573c7e5d232509479a823cface742" + }, + "reference_native_sha256": "056243fad27d224b82e437925ffa2aed42037e6bd29f239f56076a832f6ca5cb", + "reference_bundle_manifest_sha256": "c0fd5567442b08b908cc193f36d0864e262573c7e5d232509479a823cface742", + "source_archive_provenance_sha256": "1bb902137576b4ff09064e0af9f6c648ca4503158c4a04d6fb7b00d6bb752531", + "b12x_source_commit": "eac260a8257cc6b14e7d4ad674f51e9a09b8790f", + "b12x_roce_tree_sha256": "8a64063506c9c96d3fd3b004c3ab51dda88408f7246f5c89d0a0b6713065fca6", + "licenses": { + "native": "Apache-2.0 (LICENSE inside source archive)", + "b12x": "Apache-2.0 (B12X-LICENSE)" + } +} diff --git a/runtime/glm53-spark-mtp3-mesh/performance/transport/test_transport_package.py b/runtime/glm53-spark-mtp3-mesh/performance/transport/test_transport_package.py new file mode 100644 index 00000000..c51323b4 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/transport/test_transport_package.py @@ -0,0 +1,69 @@ +"""Offline integrity and output-boundary checks for transport packaging.""" +import hashlib +import importlib.util +import json +from pathlib import Path +import shutil + +import pytest + +ROOT = Path(__file__).resolve().parent +spec = importlib.util.spec_from_file_location("transport_package", ROOT / "package.py") +package = importlib.util.module_from_spec(spec) +spec.loader.exec_module(package) + + +def test_source_inventory_and_correctness_repairs(): + receipt, source, _ = package.verified_sources() + assert receipt["reference_native_sha256"].startswith("056243fa") + kernel = source["spark_transport/experiments/tiled_prefill/fused_prefill_kernels.cu"].decode() + barrier = kernel.split("__device__ bool flow_barrier", 1)[1] + assert barrier.index("__syncthreads();") < barrier.index("atomicAdd(") + worker = source["spark_transport/src/tp4_fused_prefill_session.cpp"].decode() + assert "thread_ = std::thread([this] { loop(); });" in worker + assert "write_payload_and_doorbell" in source["spark_transport/src/verbs_endpoint.cpp"].decode() + + +@pytest.mark.parametrize("name", ["../escape", "/absolute", "a\\b", "a/../b"]) +def test_unsafe_paths_refused(name): + with pytest.raises(ValueError): + package.checked_name(name) + + +def test_tampering_is_rejected_before_output(tmp_path): + copied = tmp_path / "package" + shutil.copytree(ROOT, copied, ignore=shutil.ignore_patterns("__pycache__", ".pytest_cache")) + with (copied / "native-source.tar.gz").open("ab") as stream: + stream.write(b"tampered") + with pytest.raises(ValueError, match="archive digest"): + package.verified_sources(copied) + + +def test_existing_destination_is_not_modified(tmp_path): + sentinel = tmp_path / "sentinel" + sentinel.write_text("preserve") + with pytest.raises(ValueError, match="must not exist"): + package.write_tree(tmp_path, {"sentinel": b"replace"}) + assert sentinel.read_text() == "preserve" + + +def test_rebuilt_library_has_distinct_manifest_and_no_qualification_claim(tmp_path): + library = tmp_path / "fixture.so" + library.write_bytes(b"offline fixture, not a loadable library") + digest = hashlib.sha256(library.read_bytes()).hexdigest() + result = package.prepare_bundle(library, digest, tmp_path / "bundle") + assert result["status"] == "research-only" + assert not result["reference_artifact_match"] + assert not result["hardware_qualification_performed"] + manifest = json.loads((tmp_path / "bundle" / package.MANIFEST).read_bytes()) + rows = {row["path"]: row["sha256"] for row in manifest["files"]} + assert rows[package.LIBRARY] == digest + assert result["bundle_manifest_sha256"] != package.verified_sources()[0]["reference_bundle_manifest_sha256"] + + +def test_incorrect_library_digest_creates_no_output(tmp_path): + library = tmp_path / "fixture.so" + library.write_bytes(b"fixture") + with pytest.raises(ValueError, match="library digest differs"): + package.prepare_bundle(library, "0" * 64, tmp_path / "bundle") + assert not (tmp_path / "bundle").exists() diff --git a/runtime/glm53-spark-mtp3-mesh/performance/verify.py b/runtime/glm53-spark-mtp3-mesh/performance/verify.py new file mode 100644 index 00000000..99d729f9 --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/performance/verify.py @@ -0,0 +1,18 @@ +"""Verify the complete installed performance image file inventory.""" + +import hashlib +import json +from pathlib import Path + +receipt = json.loads(Path("/opt/sparkring/receipts/mtp3-performance.json").read_text()) +for name, expected in receipt["files"].items(): + if hashlib.sha256(Path(name).read_bytes()).hexdigest() != expected: + raise RuntimeError(f"Runtime source differs from image receipt: {name}") +print( + json.dumps( + { + "verified_files": len(receipt["files"]), + "sparkcache_commit": receipt["sparkcache_commit"], + } + ) +) diff --git a/runtime/glm53-spark-mtp3-mesh/pins.json b/runtime/glm53-spark-mtp3-mesh/pins.json index 59a1b680..3267552b 100644 --- a/runtime/glm53-spark-mtp3-mesh/pins.json +++ b/runtime/glm53-spark-mtp3-mesh/pins.json @@ -2,6 +2,14 @@ "schema": "sparkring-glm53-spark-mtp3-mesh/v1", "status": "research-only", "image_pins": "../glm53-flash-jj-r8-gb10/pins.json", + "compute": { + "source_lock": "compute/source-lock.json", + "source_lock_sha256": "139f36701e0e47f45bf99fba2cc2fa59b417f2ee801dad3a064455d5b464a459", + "vllm_base_revision": "e02b174693e13859de61811b5e8cd13d5308e259", + "b12x_revision": "ef308bac0f3b3eb8fea63e4013afc0c2ea1c6301", + "b12x_tree": "dcf039e5e754136275835ea997e6b9abbb6b15ae", + "cuda_version": "13.3" + }, "target": { "repository": "local-inference-lab/GLM-5.3-Flash-NVFP4-Spark", "revision": "df116c4fb16b1d37ae43d2cfd624de26ffbc832e", @@ -9,12 +17,42 @@ "index_sha256": "db30fc7c5a70ccfb3b1c46637bb4ddb04226b95a5dfc451dffccb96a4f0ff544", "checkpoint_identity": "357f6a86160ebd5caff25d9a10d9f29e8547b16c6c73e78751fa69fde11ac4e4" }, - "speculation": {"method": "mtp", "num_speculative_tokens": 3, "attention_backend": "B12X", "draft_tensor_parallel_size": 4, "draft_sample_method": "probabilistic", "rejection_sample_method": "standard"}, - "capture_sizes": [4, 8, 12, 16, 20, 24, 28, 32, 36, 40, 44, 48, 52, 56, 60, 64], - "captured_sircl_query_rows": [16, 20, 24, 28, 32], - "canonical_bundle_manifest_sha256": "4204fabc93303226b9a120b094ef3c82ed4aadd1d7f97cfbe291204c027ed45f", + "speculation": { + "method": "mtp", + "num_speculative_tokens": 3, + "attention_backend": "B12X", + "draft_tensor_parallel_size": 4, + "draft_sample_method": "probabilistic", + "rejection_sample_method": "standard" + }, + "capture_sizes": [ + 4, + 8, + 12, + 16, + 20, + 24, + 28, + 32, + 36, + 40, + 44, + 48, + 52, + 56, + 60, + 64 + ], + "captured_sircl_query_rows": [ + 16, + 20, + 24, + 28, + 32 + ], + "canonical_bundle_manifest_sha256": "69313e19e881ec93e9ed3bd150d2f24fc6b444488ac729a69f45d038e2243500", "measured_bundle_manifest_sha256": "701bdc42069a97492981b8f34e006ebfa9e68c2160472cba631b56965efae226", - "bundle_difference": "Configuration/manifest line endings differ; parsed configuration and every executable file are identical.", + "bundle_difference": "The serving measurements use the recorded measured bundle. The canonical bundle additionally preserves capture-ID stream guards and orders shared staging/output copies across streams; routing and collective algorithms are unchanged.", "marker": { "source": "../../spark_transport/experiments/cx7_hairpin_diagonal/native/mlx5_rdma_tx_rewrite_probe.c", "source_sha256": "8684a6961b8e86aa474fa2310ff71e4cdf219a63a72ceb5593b2f95e54812792", @@ -28,7 +66,7 @@ "cache_identity": { "draft_policy": "separate", "draft_checkpoint_source": "target.checkpoint_identity", - "namespace": "glm53-spark-df116c4f-native-mtp3-mesh-4204fabc-tail-cow-v2", - "compatibility": "The native predictor is identified by its target checkpoint. This intentionally misses the measured deployment's external-draft-tagged entries. Cold-start and persistent-restore qualification of this namespace remain required." + "namespace": "glm53-spark-df116c4f-mtp3-nvfp4-a16-c139f3670-mesh69313e19-tail-cow-v2", + "compatibility": "The namespace binds the NVFP4/BF16 proposal head, compute source lock139f3670, and mesh bundle69313e19. Do not relabel cache entries from another compute composition. Persistent restore requires matching checkpoint, computation, and cache geometry." } } diff --git a/runtime/glm53-spark-mtp3-mesh/profile.py b/runtime/glm53-spark-mtp3-mesh/profile.py index 258383ad..237d0b38 100644 --- a/runtime/glm53-spark-mtp3-mesh/profile.py +++ b/runtime/glm53-spark-mtp3-mesh/profile.py @@ -104,6 +104,11 @@ def load_site(path: Path): def load_image_receipt(path: Path) -> dict: document = json.loads(path.read_text()) + if document.get("schema") == "sparkring-mtp3-performance-public-image/v1": + pinned = json.loads((HERE / "performance/public-image.json").read_text()) + if document != pinned or not document.get("anonymous_manifest_verified"): + raise ValueError("Performance image receipt differs from the repository pin") + return dict(document, inside_image=document["verification"]) expected = PINS["canonical_bundle_manifest_sha256"] image_id = document.get("image_id", "") inside = document.get("inside_image", {}) @@ -127,6 +132,21 @@ def load_image_receipt(path: Path) -> dict: "helper_sha256": sha(BASE / "warmup_dflash.py"), "temperature": 1.0, }: raise ValueError("Image receipt does not verify the sampling warmup helper") + compute = inside.get("compute") + required = PINS.get("compute", {}) + if required: + if not isinstance(compute, dict): + raise ValueError("Image receipt lacks the required compute attestation") + for field in ("source_lock_sha256", "b12x_revision", "b12x_tree", "cuda_version"): + if compute.get(field) != required[field]: + raise ValueError(f"Image receipt compute identity differs: {field}") + lock = json.loads((HERE / required["source_lock"]).read_text()) + environment = compute.get("environment", {}) + if (not isinstance(environment, dict) + or any(environment.get(k) != v for k, v in lock["environment"].items()) + or compute.get("proposal_head_nvfp4") is not True + or compute.get("target_head_quantization") is not False): + raise ValueError("Image receipt does not attest the required proposal and verifier paths") return document @@ -145,7 +165,9 @@ def manifest_file(root: Path, value: object) -> Path: def render(site_path: Path, bundle: Path, output: Path, image_receipt: Path | None = None) -> dict: if output.exists(): raise ValueError("Output directory exists; use an absent directory") - if sha(bundle / "sparkring-overlay-manifest.json") != PINS["canonical_bundle_manifest_sha256"]: + image_record = load_image_receipt(image_receipt) if image_receipt else None + expected_bundle = image_record["bundle_manifest_sha256"] if image_record else PINS["canonical_bundle_manifest_sha256"] + if sha(bundle / "sparkring-overlay-manifest.json") != expected_bundle: raise ValueError("Bundle manifest does not match the MTP3 mesh profile") manifest = json.loads((bundle / "sparkring-overlay-manifest.json").read_text()) for item in manifest["files"]: @@ -166,12 +188,14 @@ def render(site_path: Path, bundle: Path, output: Path, image_receipt: Path | No "MASTER_ADDR": site["management_addresses"][0], "DFLASH_WARMUP": "1", "SPARKRING_WARMUP_TEMPERATURE": "0", }) - image_record = load_image_receipt(image_receipt) if image_receipt else None if image_record is not None: values["IMAGE_ID"] = image_record["image_id"] values["IMAGE_REF"] = image_record["image_reference"] if image_record["inside_image"].get("readiness_warmup") is not None: values["SPARKRING_WARMUP_TEMPERATURE"] = "1" + if image_record.get("schema") == "sparkring-mtp3-performance-public-image/v1": + values["SPARKCACHE_PLACEMENT_LIBRARY_SHA256"] = image_record["native_placement_sha256"] + values["SPARKCACHE_CACHE_NAMESPACE"] = image_record["cache_namespace"] output.mkdir(parents=True) ranks = [] for rank in range(4): @@ -216,7 +240,7 @@ def render(site_path: Path, bundle: Path, output: Path, image_receipt: Path | No shutil.copyfile(topology.source_path, output / "fabric.json") result = {"schema": "sparkring-mtp3-mesh-render/v1", "status": "research-only", "execution_authorized": False, "site_sha256": sha(site_path), "topology_sha256": topology.sha256, - "bundle_manifest_sha256": PINS["canonical_bundle_manifest_sha256"], + "bundle_manifest_sha256": expected_bundle, "image": IMAGE["operator_image"], "marker_binary": site["marker_binary"], "marker_binary_sha256": site["marker_binary_sha256"], "state_root": site["state_root"], "marker_scope": "All RDMA-TX packets with reserved UDP source port 65535 on each selected function; not an IP/QPN-scoped rule.", diff --git a/runtime/glm53-spark-mtp3-mesh/public-image.json b/runtime/glm53-spark-mtp3-mesh/public-image.json index 7ecf739a..7f19c99a 100644 --- a/runtime/glm53-spark-mtp3-mesh/public-image.json +++ b/runtime/glm53-spark-mtp3-mesh/public-image.json @@ -1,18 +1,21 @@ { + "schema": "sparkring-managed-mesh-public-image/v1", + "status": "research-only", "checks_passed": true, "anonymous_manifest_read": true, "anonymous_config_read": true, "anonymous_pull": true, - "public_reference": "ghcr.io/fujitsupolycom/sparkring-glm53-sparkcache@sha256:23f00af873ccc784cfb742b7be2a29c6d3c20ebec9741843c025320bb9c04685", - "tag": "glm53-spark-mtp3-managed-8684a696", - "manifest_digest": "sha256:23f00af873ccc784cfb742b7be2a29c6d3c20ebec9741843c025320bb9c04685", - "config_image_id": "sha256:26273b8e358df139ae913610a5d43084ff0fd08aafe282ef633a3bc74afefe47", + "anonymous_pull_method": "Empty Docker client credential directory; local image layers already present", + "public_reference": "ghcr.io/fujitsupolycom/sparkring-glm53-sparkcache@sha256:67dc0ae453baaae6831ccec1d259b4ef8b236a8b0dc9f747d901b95c66ec1987", + "tag": "glm53-spark-mtp3-c139f3670-mesh69313e19", + "manifest_digest": "sha256:67dc0ae453baaae6831ccec1d259b4ef8b236a8b0dc9f747d901b95c66ec1987", + "config_image_id": "sha256:2e41b1e934a85ff7c21b780532db2f0a0e978df081e52f4ae2bf11f8992fb24f", "platform": "linux/arm64", - "layer_count": 87, + "layer_count": 89, "all_layer_diff_ids_match_tested_image": true, + "layer_identity_reference": "The tested config-image ID is identical to the published config-image ID.", "package_url": "https://github.com/FujitsuPolycom/sparkring/pkgs/container/sparkring-glm53-sparkcache", - "schema": "sparkring-managed-mesh-public-image/v1", - "status": "research-only", "content_receipt": "image-receipt.json", - "functional_record": "../../performance/records/glm53-flash/spark-mtp3-managed-mesh-functional-20260905.md" + "compute_equivalence": "compute-image-equivalence.json", + "functional_record": "../../performance/records/glm53-flash/spark-mtp3-compute-stream-safety-20260906.md" } diff --git a/runtime/glm53-spark-mtp3-mesh/qualification/run_native.py b/runtime/glm53-spark-mtp3-mesh/qualification/run_native.py index e7f945c0..96745d03 100644 --- a/runtime/glm53-spark-mtp3-mesh/qualification/run_native.py +++ b/runtime/glm53-spark-mtp3-mesh/qualification/run_native.py @@ -85,9 +85,13 @@ def main() -> None: parser.add_argument("--output", type=Path, required=True) parser.add_argument("--rows", nargs="+", type=int, default=[4, 20, 28, 64]) parser.add_argument("--port", type=int, default=29960) + parser.add_argument("--mode", choices=("correctness", "streams"), default="correctness") parser.add_argument("--execute-authorized", action="store_true") args = parser.parse_args() plan = make_plan(args.launch, args.image_receipt, args.rows, args.port) + source_path = HERE / ("stream_roce.py" if args.mode == "streams" else "native_roce.py") + plan["mode"] = args.mode + plan["source_sha256"] = profile.sha(source_path) if not args.execute_authorized: print(json.dumps(plan, indent=2)) return @@ -100,7 +104,7 @@ def main() -> None: (args.output / f"preflight-r{rank['rank']}.json").write_text(json.dumps({"containers": state, "image": image}, indent=2)) if state["returncode"] or state["stdout"].strip() or image["returncode"] or image["stdout"].strip() != plan["image"]: raise RuntimeError("Require no running containers and the exact image on every rank") - source = (HERE / "native_roce.py").read_text() + source = source_path.read_text() for cell in plan["cells"]: with concurrent.futures.ThreadPoolExecutor(max_workers=4) as pool: results = list(pool.map(lambda rank: remote(rank["host"], rank["argv"], source), cell["ranks"])) diff --git a/runtime/glm53-spark-mtp3-mesh/qualification/stream_roce.py b/runtime/glm53-spark-mtp3-mesh/qualification/stream_roce.py new file mode 100644 index 00000000..ef8960ba --- /dev/null +++ b/runtime/glm53-spark-mtp3-mesh/qualification/stream_roce.py @@ -0,0 +1,81 @@ +"""Validate bundled RoCEnante staging across CUDA streams with exact BF16 sums.""" +import argparse +import json +import torch +import torch.distributed as dist + +parser=argparse.ArgumentParser() +parser.add_argument('--bytes',type=int,required=True) +args,_=parser.parse_known_args() +dist.init_process_group('gloo') +rank=dist.get_rank() +assert dist.get_world_size()==4 +torch.cuda.set_device(0) +import b12x.comm # noqa: E402 -- device selection precedes transport import +b12x.comm.__path__.insert(0,'/opt/spark-sircl/b12x_overlay/b12x/comm') +from b12x.comm import roce # noqa: E402 -- resolve only the image-bundled transport +runtime=roce.AllReduce.from_exchange_group(exchange_group=dist.group.WORLD,device=torch.device('cuda',0), + max_size=2<<20,max_gather_bytes=2<<20) +runtime.prepare((torch.bfloat16,)) +numel=args.bytes//2 +streams=[torch.cuda.Stream(),torch.cuda.Stream()] +inputs=[torch.empty(numel+1,device='cuda',dtype=torch.bfloat16)[1:] for _ in range(16)] +outputs=[torch.empty_like(inp.new_empty(numel+1))[1:] for inp in inputs] +assert all(t.data_ptr()%16 for t in inputs+outputs) +# Establish shared scratch before exercising asynchronous alternating callers. +inputs[0].fill_(rank+1) +runtime.all_reduce(inputs[0],out=outputs[0]) +torch.cuda.synchronize() +dist.barrier() +for i,(inp,out) in enumerate(zip(inputs,outputs)): + with torch.cuda.stream(streams[i%2]): + inp.fill_(rank+1+i) + out.fill_(-123) + runtime.all_reduce(inp,out=out) +torch.cuda.synchronize() +runtime.check_health() +assert all(torch.equal(out,torch.full_like(out,10+4*i)) for i,out in enumerate(outputs)) +dist.barrier() +graph=torch.cuda.CUDAGraph() +streams[0].wait_stream(torch.cuda.current_stream()) +with torch.cuda.graph(graph,stream=streams[0]): + runtime.all_reduce(inputs[0],out=outputs[0]) +for i in (20,21): + inputs[0].fill_(rank+1+i) + outputs[0].fill_(-321) + graph.replay() + torch.cuda.synchronize() + runtime.check_health() + assert torch.equal(outputs[0],torch.full_like(outputs[0],10+4*i)) +dist.barrier() +# A separate Python capture context must not admit another stream in one CUDA capture. +probe=torch.zeros(1,device='cuda') +probe.add_(1) +torch.cuda.synchronize() +guard_graph=torch.cuda.CUDAGraph() +rejected=False +with torch.cuda.graph(guard_graph,stream=streams[0]): + probe.add_(1) + with runtime.capture(): + runtime._order_stream(True) + streams[1].wait_stream(streams[0]) + with torch.cuda.stream(streams[1]): + with runtime.capture(): + try: + runtime._order_stream(True) + except RuntimeError as error: + assert 'one stream' in str(error) + rejected=True + streams[0].wait_stream(streams[1]) +assert rejected +torch.cuda.synchronize() +runtime.check_health() +record={'rank':rank,'payload_bytes':args.bytes,'alternating_stream_calls':16, + 'misaligned_input_output':True,'changed_input_graph_replays':2, + 'capture_stream_rejected':rejected,'passed':True} +rows=[None]*4 +dist.all_gather_object(rows,record) +if rank==0: + print('EVIDENCE_JSON '+json.dumps({'checks':rows,'passed':all(row['passed'] for row in rows)}),flush=True) +runtime.close() +dist.destroy_process_group() diff --git a/runtime/glm53-spark-mtp3-mesh/qualification/test_checks.py b/runtime/glm53-spark-mtp3-mesh/qualification/test_checks.py index a70ed499..e043ce57 100644 --- a/runtime/glm53-spark-mtp3-mesh/qualification/test_checks.py +++ b/runtime/glm53-spark-mtp3-mesh/qualification/test_checks.py @@ -33,7 +33,8 @@ def inputs(tmp_path): "image_id": "sha256:" + "a" * 64, "image_reference": "sha256:" + "a" * 64, "bundle_manifest_sha256": bundle_sha, "source_receipt_sha256": "b" * 64, "inside_image": {"checks_passed": True, "bundle_manifest_sha256": bundle_sha, - "source_receipt_sha256": "b" * 64, "cuda_initialized": False, "model_loaded": False}} + "source_receipt_sha256": "b" * 64, "cuda_initialized": False, "model_loaded": False, + "compute": json.loads((HERE.parent / "image-receipt.json").read_text())["inside_image"]["compute"]}} receipt_path = tmp_path / "image.json" receipt_path.write_text(json.dumps(receipt)) plan = {"schema": "sparkring-mtp3-mesh-render/v1", "image": receipt, diff --git a/runtime/glm53-spark-mtp3-mesh/test_image.py b/runtime/glm53-spark-mtp3-mesh/test_image.py index a7926873..75532a53 100644 --- a/runtime/glm53-spark-mtp3-mesh/test_image.py +++ b/runtime/glm53-spark-mtp3-mesh/test_image.py @@ -3,6 +3,7 @@ import importlib.util import json from pathlib import Path +from types import SimpleNamespace import pytest @@ -84,6 +85,239 @@ def test_file_map_rejects_parent_traversal(tmp_path): verifier.verify_file_map(tmp_path, {"../test": "0" * 64}) +def test_layered_file_map_checks_unchanged_parent_and_exact_overrides(tmp_path): + (tmp_path / "vllm").mkdir() + unchanged = tmp_path / "vllm" / "unchanged.py" + replaced = tmp_path / "vllm" / "replaced.py" + unchanged.write_text("parent\n", encoding="utf-8") + replaced.write_text("result\n", encoding="utf-8") + parent = { + "vllm/unchanged.py": verifier.sha256(unchanged), + "vllm/replaced.py": "0" * 64, + } + overrides = { + "vllm/replaced.py": { + "base_sha256": "0" * 64, + "result_sha256": verifier.sha256(replaced), + } + } + assert verifier.verify_layered_file_map(tmp_path, parent, overrides) == { + "parent_files": 2, + "overrides": 1, + } + unchanged.write_text("changed\n", encoding="utf-8") + with pytest.raises(ValueError, match="content pin"): + verifier.verify_layered_file_map(tmp_path, parent, overrides) + + +def test_layered_file_map_rejects_unbound_or_wrong_base_override(tmp_path): + (tmp_path / "vllm").mkdir() + source = tmp_path / "vllm" / "source.py" + source.write_text("result\n", encoding="utf-8") + parent = {"vllm/source.py": "1" * 64} + with pytest.raises(ValueError, match="base identity"): + verifier.verify_layered_file_map( + tmp_path, + parent, + {"vllm/source.py": { + "base_sha256": "2" * 64, + "result_sha256": verifier.sha256(source), + }}, + ) + with pytest.raises(ValueError, match="absent from the parent"): + verifier.verify_layered_file_map( + tmp_path, + parent, + {"vllm/other.py": { + "base_sha256": "1" * 64, + "result_sha256": verifier.sha256(source), + }}, + ) + + +def test_complete_package_file_map_rejects_unmanifested_module(tmp_path): + package = tmp_path / "b12x" + package.mkdir() + declared = package / "declared.py" + declared.write_text("VALUE = 1\n", encoding="utf-8") + records = {"b12x/declared.py": verifier.sha256(declared)} + assert verifier.verify_complete_package_file_map(tmp_path, "b12x", records) == 1 + (package / "unmanifested.py").write_text("VALUE = 0\n", encoding="utf-8") + with pytest.raises(ValueError, match="complete manifest"): + verifier.verify_complete_package_file_map(tmp_path, "b12x", records) + + +def test_compute_environment_requires_cuda_and_quantization_contract(): + expected = { + "CUDA_HOME": "/opt/cuda-13.3", + "TRITON_PTXAS_PATH": "/opt/cuda-13.3/bin/ptxas", + "VLLM_GDN_SPEC_DECODE_METADATA_FASTPATH": "1", + "VLLM_B12X_DENSE_ACTIVATION_MODE": "auto", + "VLLM_MTP_NVFP4_LM_HEAD": "1", + "VLLM_LM_HEAD_A16": "1", + "VLLM_MXFP8_LM_HEAD": "0", + } + assert verifier.verify_required_environment(expected, expected) == expected + for name in expected: + changed = dict(expected) + changed[name] = "wrong" + with pytest.raises(ValueError, match=name): + verifier.verify_required_environment(changed, expected) + + +def test_compute_verifier_composes_parent_vllm_and_complete_b12x(tmp_path): + compute = tmp_path / "compute" + receipts = tmp_path / "receipts" + site = tmp_path / "site" + cuda = tmp_path / "cuda" + for path in (compute, receipts, site / "vllm", site / "b12x", cuda / "bin"): + path.mkdir(parents=True, exist_ok=True) + unchanged = site / "vllm" / "unchanged.py" + override = site / "vllm" / "override.py" + b12x_python = site / "b12x" / "__init__.py" + b12x_notice = site / "b12x" / "README.md" + unchanged.write_text("unchanged\n", encoding="utf-8") + override.write_text("result\n", encoding="utf-8") + b12x_python.write_text("", encoding="utf-8") + b12x_notice.write_text("source\n", encoding="utf-8") + (cuda / "bin" / "ptxas").write_text("tool\n", encoding="utf-8") + components = {"cuda_nvcc/archive.tar.xz": "a" * 64} + (cuda / "sparkring-component-manifest.json").write_text( + json.dumps(components), encoding="utf-8" + ) + base_hash = "b" * 64 + lock = { + "schema": "sparkring-glm53-compute-source/v1", + "vllm": { + "base_revision": "e02", + "files": [["vllm/override.py", base_hash, verifier.sha256(override)]], + }, + "b12x": { + "revision": "b58", + "tree": "tree", + "package_files_sha256": verifier.file_map_sha256({ + "b12x/__init__.py": verifier.sha256(b12x_python), + "b12x/README.md": verifier.sha256(b12x_notice), + }), + }, + "cuda": {"version": "13.3", "components": components}, + "environment": { + "VLLM_GDN_SPEC_DECODE_METADATA_FASTPATH": "1", + "VLLM_B12X_DENSE_ACTIVATION_MODE": "auto", + "VLLM_MTP_NVFP4_LM_HEAD": "1", + "VLLM_LM_HEAD_A16": "1", + "VLLM_MXFP8_LM_HEAD": "0", + }, + } + lock_path = compute / "source-lock.json" + lock_path.write_text(json.dumps(lock), encoding="utf-8") + lock_hash = verifier.sha256(lock_path) + (receipts / "vllm-source-manifest.json").write_text( + json.dumps({ + "commit": "e02", + "files": { + "vllm/unchanged.py": verifier.sha256(unchanged), + "vllm/override.py": base_hash, + }, + }), + encoding="utf-8", + ) + b12x_files = { + "b12x/__init__.py": verifier.sha256(b12x_python), + "b12x/README.md": verifier.sha256(b12x_notice), + } + installed_path = tmp_path / "installed.json" + installed_path.write_text( + json.dumps({ + "schema": "sparkring-glm53-compute-installed/v1", + "source_lock_sha256": lock_hash, + "vllm_revision": "e02", + "vllm_overrides": {"vllm/override.py": verifier.sha256(override)}, + "b12x_revision": "b58", + "b12x_tree": "tree", + "b12x_files": b12x_files, + "cuda_components": components, + "environment": lock["environment"], + "target_head_quantization": False, + }), + encoding="utf-8", + ) + profile = {"compute": { + "source_lock": "compute/source-lock.json", + "source_lock_sha256": lock_hash, + "vllm_base_revision": "e02", + "b12x_revision": "b58", + "b12x_tree": "tree", + "cuda_version": "13.3", + }} + source = {"files": { + "compute/source-lock.json": lock_hash, + "compute/b12x-source/b12x/__init__.py": verifier.sha256(b12x_python), + "compute/b12x-source/b12x/README.md": verifier.sha256(b12x_notice), + "compute/b12x-source/pyproject.toml": "a" * 64, + "compute/b12x-source/tests/test_example.py": "b" * 64, + }} + environment = { + **lock["environment"], + "CUDA_HOME": "/opt/cuda-13.3", + "TRITON_PTXAS_PATH": "/opt/cuda-13.3/bin/ptxas", + } + result = verifier.verify_compute( + profile, + {"vllm": {"commit": "e02"}}, + source, + environment, + compute_root=compute, + receipt_path=installed_path, + site=site, + base_receipts=receipts, + cuda_root_override=cuda, + ptxas_runner=lambda *args, **kwargs: SimpleNamespace( + stdout="ptxas release 13.3", stderr="" + ), + ) + assert result["vllm_parent_files"] == 2 + assert result["vllm_overrides"] == 1 + assert result["b12x_files"] == 2 + assert result["proposal_head_nvfp4"] is True + assert result["target_head_quantization"] is False + + changed_source = json.loads(json.dumps(source)) + changed_source["files"]["compute/b12x-source/b12x/README.md"] = "f" * 64 + with pytest.raises(ValueError, match="source-receipt-bound"): + verifier.verify_compute( + profile, + {"vllm": {"commit": "e02"}}, + changed_source, + environment, + compute_root=compute, + receipt_path=installed_path, + site=site, + base_receipts=receipts, + cuda_root_override=cuda, + ptxas_runner=lambda *args, **kwargs: SimpleNamespace( + stdout="ptxas release 13.3", stderr="" + ), + ) + + unchanged.write_text("changed\n", encoding="utf-8") + with pytest.raises(ValueError, match="content pin"): + verifier.verify_compute( + profile, + {"vllm": {"commit": "e02"}}, + source, + environment, + compute_root=compute, + receipt_path=installed_path, + site=site, + base_receipts=receipts, + cuda_root_override=cuda, + ptxas_runner=lambda *args, **kwargs: SimpleNamespace( + stdout="ptxas release 13.3", stderr="" + ), + ) + + def test_container_verification_has_no_device_or_network_access(): source = (HERE / "verify_mesh_image.py").read_text(encoding="utf-8") assert '"--network", "none"' in source @@ -111,6 +345,27 @@ def test_pins_use_native_mtp_only(): assert pins["target"]["repository"].endswith("-Spark") +def test_profile_pins_exact_compute_source_and_quantization_environment(): + pins = json.loads((HERE / "pins.json").read_text(encoding="utf-8")) + compute = pins["compute"] + lock_path = HERE / compute["source_lock"] + lock = json.loads(lock_path.read_text(encoding="utf-8")) + assert verifier.sha256(lock_path) == compute["source_lock_sha256"] + assert lock["schema"] == "sparkring-glm53-compute-source/v1" + assert lock["vllm"]["base_revision"] == compute["vllm_base_revision"] + assert lock["b12x"]["revision"] == compute["b12x_revision"] + assert lock["b12x"]["tree"] == compute["b12x_tree"] + assert lock["cuda"]["version"] == compute["cuda_version"] == "13.3" + assert lock["environment"] == { + "VLLM_GDN_SPEC_DECODE_METADATA_FASTPATH": "1", + "VLLM_B12X_DENSE_ACTIVATION_MODE": "auto", + "VLLM_MTP_NVFP4_LM_HEAD": "1", + "VLLM_LM_HEAD_A16": "1", + "VLLM_MXFP8_LM_HEAD": "0", + } + assert len(lock["vllm"]["files"]) == 24 + + def test_schema_accepts_research_only_status(): schema = json.loads((HERE / "image-receipt.schema.json").read_text(encoding="utf-8")) assert schema["properties"]["status"]["const"] == "research-only" diff --git a/runtime/glm53-spark-mtp3-mesh/test_managed_service.py b/runtime/glm53-spark-mtp3-mesh/test_managed_service.py index f2b6db54..fcdbbac4 100644 --- a/runtime/glm53-spark-mtp3-mesh/test_managed_service.py +++ b/runtime/glm53-spark-mtp3-mesh/test_managed_service.py @@ -3,6 +3,8 @@ from pathlib import Path import threading import time +from concurrent.futures import Future +from types import SimpleNamespace import pytest @@ -167,6 +169,197 @@ class Result: assert calls == [['docker', 'kill', 'a' * 64]] +def test_docker_inspect_timeout_is_unknown_without_stalling_fabric_monitor(monkeypatch): + pending = Future() + submitted = [] + + class Executor: + def submit(self, function, name): + submitted.append((function, name)) + return pending + + def shutdown(self, **kwargs): + pass + + monkeypatch.setattr(service.concurrent.futures, 'ThreadPoolExecutor', lambda **kw: Executor()) + watcher = service.DockerStatePoll('a' * 64) + assert watcher.poll() is None + for _ in range(20): + assert watcher.poll() is None + assert len(submitted) == 1 + command = ['docker', 'inspect', '--format', '{{.State.Running}}', 'a' * 64] + pending.set_exception(service.subprocess.TimeoutExpired(command, 3)) + assert watcher.poll() is None + assert '3 seconds' in watcher.error + assert len(submitted) == 2 + watcher.close() + + +def test_docker_status_recovers_without_reusing_stale_stopped_evidence(monkeypatch): + futures = [Future(), Future(), Future()] + pending = iter(futures) + executor = SimpleNamespace(submit=lambda *args: next(pending), shutdown=lambda **kw: None) + monkeypatch.setattr(service.concurrent.futures, 'ThreadPoolExecutor', lambda **kw: executor) + watcher = service.DockerStatePoll('a' * 64) + assert watcher.poll() is None + futures[0].set_result(False) + assert watcher.poll() is False + assert watcher.poll() is None + futures[1].set_result(True) + assert watcher.poll() is True + assert watcher.error is None + watcher.close() + + +def test_unknown_docker_status_blocks_model_admission(): + changed = rows() + changed[0]['docker_status_degraded'] = True + with pytest.raises(RuntimeError): + service.validate_group(changed) + + +def test_model_arm_rechecks_local_docker_status_after_group_gate(tmp_path, monkeypatch): + monkeypatch.setattr(service, 'load_config', lambda path: ({'state_dir': str(tmp_path)},)) + (tmp_path / 'status.json').write_bytes(service.canonical({ + 'phase': 'armed', 'local_ready': True, 'generation': 'g', 'docker_status_degraded': True, + })) + with pytest.raises(RuntimeError, match='not armed'): + service.model_intent('unused', True) + assert not (tmp_path / 'model-intent.json').exists() + + +@pytest.mark.parametrize('message, missing', [ + ('Error: No such object: ' + 'a' * 64, True), + ('Error: No such object: ' + 'b' * 64, False), + ('No such file or directory: Docker socket', False), +]) +def test_only_pinned_missing_container_proves_absence(monkeypatch, message, missing): + monkeypatch.setattr(service.subprocess, 'run', lambda *args, **kw: + SimpleNamespace(returncode=1, stdout='', stderr=message)) + if missing: + assert service.docker_running('a' * 64) is False + else: + with pytest.raises(RuntimeError): + service.docker_running('a' * 64) + + +@pytest.mark.parametrize('output', ['', 'unexpected', 'False']) +def test_unknown_docker_reply_cannot_prove_model_stopped(monkeypatch, output): + monkeypatch.setattr(service.subprocess, 'run', lambda *args, **kw: + SimpleNamespace(returncode=0, stdout=output, stderr='')) + with pytest.raises(RuntimeError): + service.stop_model('a' * 64) + + +def test_docker_timeout_cannot_pass_model_stop_barrier(monkeypatch): + calls = [] + + def timeout(argv, **kwargs): + calls.append((argv, kwargs['timeout'])) + raise service.subprocess.TimeoutExpired(argv, kwargs['timeout']) + + monkeypatch.setattr(service.subprocess, 'run', timeout) + with pytest.raises(service.subprocess.TimeoutExpired): + service.stop_model('a' * 64) + assert len(calls) == 1 + assert calls[0][1] == 3 + + +@pytest.mark.parametrize('peer_failure', [False, True]) +def test_monitor_keeps_fabric_checks_live_when_docker_is_unknown(tmp_path, monkeypatch, peer_failure): + result = owner() + result.rank, result.generation, result.model = 0, 'g', 'a' * 64 + result.config = {'site_path': '/unused'} + result.site, result.identity, result.key = {}, 'identity', b'k' * 32 + result.state_dir, result.network, result.server = tmp_path, None, None + result.marker_records, result.logfiles = [], [] + result.failed, result.owns_guard, result.model_seen = False, False, True + events = [] + samples = iter([True, None, True]) + clock = [100.0] + monitor_round = [0] + + class Stop: + def is_set(self): + return monitor_round[0] >= 3 + + def wait(self, seconds): + monitor_round[0] += 1 + clock[0] += 6 + + result.stop = Stop() + result.start_markers = lambda: None + result.start_server = lambda: None + result.publish = lambda **changes: result.state.update(changes) + result.children = [SimpleNamespace( + poll=lambda: None, + terminate=lambda: events.append('marker-stop'), + wait=lambda **kw: None, + ) for _ in range(2)] + (tmp_path / 'model-intent.json').write_bytes(service.canonical({ + 'generation': 'g', 'active': True, 'deadline_monotonic': 0, + })) + monkeypatch.setattr(service.os, 'geteuid', lambda: 0, raising=False) + original_lstat = type(tmp_path).lstat + + def fixture_lstat(path): + info = original_lstat(path) + if path == tmp_path: + return SimpleNamespace(st_mode=info.st_mode, st_uid=0) + return info + + # Model the root-owned service directory without changing host ownership. + monkeypatch.setattr(type(tmp_path), 'lstat', fixture_lstat) + monkeypatch.setattr(service.signal, 'signal', lambda *args: None) + monkeypatch.setitem(service.sys.modules, 'fcntl', SimpleNamespace( + flock=lambda *args: None, LOCK_EX=1, LOCK_NB=2, + )) + monkeypatch.setitem(service.sys.modules, 'managed_network', SimpleNamespace( + NetworkManager=lambda *args: SimpleNamespace( + up=lambda: None, + check=lambda **kw: events.append('network-check'), + down=lambda: events.append('network-down') or {'clean': True}, + ) + )) + monkeypatch.setattr(service.time, 'monotonic', lambda: clock[0]) + monkeypatch.setattr(service.time, 'sleep', lambda seconds: events.append('retain-markers')) + monkeypatch.setattr(service, 'notify', lambda message: events.append(message)) + monkeypatch.setattr(service, 'docker_running', lambda name: events.append('startup-inspect') or False) + monkeypatch.setattr(service, 'DockerStatePoll', lambda name: SimpleNamespace( + poll=lambda: next(samples), error='Docker inspect timed out after 3 seconds', close=lambda: None, + )) + + def group_check(*args): + events.append('peer-check') + if peer_failure and monitor_round[0] == 1: + raise RuntimeError('Authenticated peer is not locally ready') + return rows() + + stop_attempts = [0] + + def stop_model(name): + stop_attempts[0] += 1 + if stop_attempts[0] == 1: + assert 'marker-stop' not in events and 'network-down' not in events + raise service.subprocess.TimeoutExpired(['docker', 'inspect', name], 3) + events.append('model-stop-confirmed') + + monkeypatch.setattr(service, 'group_check', group_check) + monkeypatch.setattr(service, 'stop_model', stop_model) + assert result.run() == 1 # The explicit stop barrier needed a retry. + assert events.count('startup-inspect') == 1 + assert events.count('peer-check') == (2 if peer_failure else 3) + assert 'network-check' in events + assert events.index('retain-markers') < events.index('model-stop-confirmed') + assert events.index('model-stop-confirmed') < events.index('marker-stop') + assert events.index('marker-stop') < events.index('network-down') + if peer_failure: + assert 'not locally ready' in result.state['error'] + else: + assert 'error' not in result.state + assert result.state['docker_status_degraded'] is False + + def test_units_bind_model_and_disable_automatic_recovery(): rendered = units.unit_text('/opt/sparkring/managed-mesh', '/etc/sparkring/managed-mesh', 'a' * 64) mesh, model = rendered['sparkring-mesh.service'], rendered['sparkring-mesh-model.service'] diff --git a/runtime/glm53-spark-mtp3-mesh/test_profile.py b/runtime/glm53-spark-mtp3-mesh/test_profile.py index 237a364a..c94543d6 100644 --- a/runtime/glm53-spark-mtp3-mesh/test_profile.py +++ b/runtime/glm53-spark-mtp3-mesh/test_profile.py @@ -330,10 +330,24 @@ def _image_receipt_document(): "image_id": "sha256:" + "a" * 64, "image_reference": "sha256:" + "a" * 64, "bundle_manifest_sha256": bundle_sha, "source_receipt_sha256": source_sha, "inside_image": {"checks_passed": True, "bundle_manifest_sha256": bundle_sha, - "source_receipt_sha256": source_sha, "cuda_initialized": False, "model_loaded": False}, + "source_receipt_sha256": source_sha, "cuda_initialized": False, "model_loaded": False, + "compute": json.loads((mesh_profile.HERE / "image-receipt.json").read_text())["inside_image"]["compute"]}, } +@pytest.mark.parametrize('field', ['compute', 'source_lock_sha256', 'b12x_revision', 'environment', 'proposal_head_nvfp4', 'target_head_quantization']) +def test_receipt_requires_profile_compute(tmp_path, field): + document = _image_receipt_document() + if field == 'compute': + del document['inside_image']['compute'] + else: + document['inside_image']['compute'][field] = None + path = tmp_path / 'wrong-compute.json' + path.write_text(json.dumps(document)) + with pytest.raises(ValueError): + mesh_profile.load_image_receipt(path) + + def test_verified_image_receipt_changes_only_image_selection(tmp_path, manifest_bundle): site_path = _site(tmp_path) document = _image_receipt_document() diff --git a/runtime/glm53-spark-mtp3-mesh/verify_mesh_image.py b/runtime/glm53-spark-mtp3-mesh/verify_mesh_image.py index 84bd2608..4d22cb27 100644 --- a/runtime/glm53-spark-mtp3-mesh/verify_mesh_image.py +++ b/runtime/glm53-spark-mtp3-mesh/verify_mesh_image.py @@ -18,6 +18,8 @@ BASE_RECEIPTS = Path("/opt/sparkring/receipts/jj-r8-sparkcache-arm64") SITE = Path("/usr/local/lib/python3.12/dist-packages") BUNDLE = Path("/opt/spark-sircl") +COMPUTE = Path("/opt/sparkring-compute") +COMPUTE_RECEIPT = Path("/opt/sparkring/receipts/glm53-compute-installed.json") def sha256(path: Path) -> str: @@ -28,6 +30,11 @@ def sha256(path: Path) -> str: return digest.hexdigest() +def file_map_sha256(records: dict) -> str: + payload = json.dumps(records, sort_keys=True, separators=(",", ":")).encode() + return hashlib.sha256(payload).hexdigest() + + def load(path: Path) -> dict: value = json.loads(path.read_text(encoding="utf-8")) if not isinstance(value, dict): @@ -51,6 +58,181 @@ def verify_file_map(root: Path, records: dict) -> int: return len(records) +def verify_layered_file_map(root: Path, parent_records: dict, + overrides: dict) -> dict: + """Verify every parent file, substituting only manifest-bound overrides.""" + if not isinstance(parent_records, dict) or not parent_records: + raise ValueError("Parent source file manifest is empty") + if not isinstance(overrides, dict) or not overrides: + raise ValueError("Compute override manifest is empty") + expected = dict(parent_records) + for relative, record in overrides.items(): + if relative not in parent_records: + raise ValueError(f"Compute override is absent from the parent manifest: {relative}") + if not isinstance(record, dict): + raise ValueError(f"Compute override record is invalid: {relative}") + if record.get("base_sha256") != parent_records[relative]: + raise ValueError(f"Compute override base identity differs from the parent: {relative}") + result = record.get("result_sha256") + if not isinstance(result, str) or len(result) != 64: + raise ValueError(f"Compute override result identity is invalid: {relative}") + expected[relative] = result + verify_file_map(root, expected) + return {"parent_files": len(parent_records), "overrides": len(overrides)} + + +def verify_complete_package_file_map(root: Path, package: str, + records: dict) -> int: + """Verify all installed package files and reject inherited stale files.""" + prefix = f"{package}/" + if not isinstance(records, dict) or not records: + raise ValueError(f"{package} source manifest is empty") + if any(not relative.startswith(prefix) for relative in records): + raise ValueError(f"{package} source manifest contains an invalid path") + verify_file_map(root, records) + observed = { + path.relative_to(root).as_posix() + for path in (root / package).rglob("*") + if (path.is_file() and not path.is_symlink() + and "__pycache__" not in path.parts and path.suffix != ".pyc") + } + if observed != set(records): + missing = sorted(set(records) - observed) + extra = sorted(observed - set(records)) + raise ValueError( + f"{package} source set differs from its complete manifest; " + f"missing={missing}, extra={extra}" + ) + return len(records) + + +def verify_required_environment(environment: dict, expected: dict) -> dict: + """Require exact construction-time values for compute-selection settings.""" + if not isinstance(expected, dict) or not expected: + raise ValueError("Required compute environment is empty") + for name, value in expected.items(): + if environment.get(name) != value: + raise ValueError( + f"Mesh compute environment differs from its required value: {name}" + ) + return dict(sorted(expected.items())) + + +def verify_compute(profile: dict, base: dict, source: dict, environment: dict, + *, compute_root: Path = COMPUTE, + receipt_path: Path = COMPUTE_RECEIPT, + site: Path = SITE, + base_receipts: Path = BASE_RECEIPTS, + cuda_root_override: Path | None = None, + ptxas_runner=subprocess.run) -> dict: + """Verify the manifest-bound vLLM, B12X, and CUDA compute composition.""" + pin = profile.get("compute") + required_pin_fields = { + "source_lock", "source_lock_sha256", "vllm_base_revision", + "b12x_revision", "b12x_tree", "cuda_version", + } + if not isinstance(pin, dict) or not required_pin_fields.issubset(pin): + raise ValueError("Mesh profile does not bind the required compute source") + if pin["source_lock"] != "compute/source-lock.json": + raise ValueError("Mesh profile compute source-lock locator is unsupported") + lock_path = compute_root / "source-lock.json" + check_file(lock_path, pin["source_lock_sha256"]) + if source["files"].get("compute/source-lock.json") != pin["source_lock_sha256"]: + raise ValueError("Image source receipt does not bind the compute source lock") + lock = load(lock_path) + if lock.get("schema") != "sparkring-glm53-compute-source/v1": + raise ValueError("Compute source lock uses an unsupported schema") + if lock["vllm"]["base_revision"] != pin["vllm_base_revision"]: + raise ValueError("Compute vLLM base revision differs from the profile pin") + if lock["b12x"]["revision"] != pin["b12x_revision"]: + raise ValueError("Compute B12X revision differs from the profile pin") + if lock["b12x"]["tree"] != pin["b12x_tree"]: + raise ValueError("Compute B12X tree differs from the profile pin") + if lock["cuda"]["version"] != pin["cuda_version"]: + raise ValueError("Compute CUDA version differs from the profile pin") + + installed = load(receipt_path) + if installed.get("schema") != "sparkring-glm53-compute-installed/v1": + raise ValueError("Installed compute receipt uses an unsupported schema") + if installed.get("source_lock_sha256") != pin["source_lock_sha256"]: + raise ValueError("Installed compute receipt uses a different source lock") + if installed.get("vllm_revision") != lock["vllm"]["base_revision"]: + raise ValueError("Installed vLLM base revision differs from the source lock") + if (installed.get("b12x_revision") != lock["b12x"]["revision"] + or installed.get("b12x_tree") != lock["b12x"]["tree"]): + raise ValueError("Installed B12X identity differs from the source lock") + + parent_vllm = load(base_receipts / "vllm-source-manifest.json") + if parent_vllm.get("commit") != base["vllm"]["commit"]: + raise ValueError("Parent source receipt has the wrong vllm revision") + if base["vllm"]["commit"] != lock["vllm"]["base_revision"]: + raise ValueError("Compute source does not extend the pinned parent vLLM") + overrides = { + relative: {"base_sha256": parent_hash, "result_sha256": result_hash} + for relative, parent_hash, result_hash in lock["vllm"]["files"] + } + expected_results = { + relative: record["result_sha256"] for relative, record in overrides.items() + } + if installed.get("vllm_overrides") != expected_results: + raise ValueError("Installed vLLM override map differs from the source lock") + vllm = verify_layered_file_map(site, parent_vllm["files"], overrides) + + b12x_prefix = "compute/b12x-source/" + expected_b12x_files = { + relative.removeprefix(b12x_prefix): expected + for relative, expected in source["files"].items() + if relative.startswith(b12x_prefix + "b12x/") + } + b12x_files = installed.get("b12x_files") + if b12x_files != expected_b12x_files: + raise ValueError( + "Installed B12X file map differs from the source-receipt-bound source" + ) + if file_map_sha256(b12x_files) != lock["b12x"]["package_files_sha256"]: + raise ValueError("Installed B12X file-map identity differs from the source lock") + b12x_count = verify_complete_package_file_map(site, "b12x", b12x_files) + if installed.get("cuda_components") != lock["cuda"]["components"]: + raise ValueError("Installed CUDA component map differs from the source lock") + cuda_root = cuda_root_override or Path(f"/opt/cuda-{lock['cuda']['version']}") + cuda_manifest = load(cuda_root / "sparkring-component-manifest.json") + if cuda_manifest != lock["cuda"]["components"]: + raise ValueError("CUDA component manifest differs from the source lock") + ptxas = cuda_root / "bin/ptxas" + if ptxas.is_symlink() or not ptxas.is_file(): + raise ValueError("Pinned CUDA toolkit does not contain ptxas") + version_result = ptxas_runner( + [str(ptxas), "--version"], capture_output=True, text=True, check=True + ) + version_text = version_result.stdout + version_result.stderr + if f"release {lock['cuda']['version']}" not in version_text: + raise ValueError("Installed ptxas version differs from the CUDA source lock") + expected_environment = { + **lock["environment"], + "CUDA_HOME": f"/opt/cuda-{lock['cuda']['version']}", + "TRITON_PTXAS_PATH": f"/opt/cuda-{lock['cuda']['version']}/bin/ptxas", + } + if installed.get("environment") != lock["environment"]: + raise ValueError("Installed compute environment differs from the source lock") + verified_environment = verify_required_environment( + environment, expected_environment + ) + if installed.get("target_head_quantization") is not False: + raise ValueError("Target LM head must remain unquantized") + return { + "source_lock_sha256": pin["source_lock_sha256"], + "vllm_parent_files": vllm["parent_files"], + "vllm_overrides": vllm["overrides"], + "b12x_revision": lock["b12x"]["revision"], + "b12x_tree": lock["b12x"]["tree"], + "b12x_files": b12x_count, + "cuda_version": lock["cuda"]["version"], + "environment": verified_environment, + "proposal_head_nvfp4": True, + "target_head_quantization": False, + } + + def verify_warmup(path: Path, expected: str, environment: dict) -> dict: """Check the readiness-only helper override and its explicit temperature.""" check_file(path, expected) @@ -83,12 +265,15 @@ def verify_inside_image() -> dict: check_file(marker_source, profile["marker"]["source_sha256"]) warmup = verify_warmup(Path("/opt/sparkring/bin/warmup_dflash.py"), source["files"]["warmup_dflash.py"], os.environ) - package_counts = {} - for package in ("vllm", "b12x", "sparkcache"): - manifest = load(BASE_RECEIPTS / f"{package}-source-manifest.json") - if manifest.get("commit") != base[package]["commit"]: - raise ValueError(f"Parent source receipt has the wrong {package} revision") - package_counts[package] = verify_file_map(SITE, manifest["files"]) + compute = verify_compute(profile, base, source, os.environ) + sparkcache_manifest = load(BASE_RECEIPTS / "sparkcache-source-manifest.json") + if sparkcache_manifest.get("commit") != base["sparkcache"]["commit"]: + raise ValueError("Parent source receipt has the wrong sparkcache revision") + package_counts = { + "vllm": compute["vllm_parent_files"], + "b12x": compute["b12x_files"], + "sparkcache": verify_file_map(SITE, sparkcache_manifest["files"]), + } native = load(BASE_RECEIPTS / "native-extension-manifest.json")["files"] native_count = verify_file_map(SITE / "vllm", native) check_file(BUNDLE / "libspark_transport_capi.so", base["sircl"]["native_sha256"]) @@ -117,12 +302,14 @@ def verify_inside_image() -> dict: "source_receipt_sha256": sha256(RECEIPTS / "source-receipt.json"), "bundle_files": verified_bundle, "python_syntax_files": python_files, "parent_package_files": package_counts, "vllm_native_extensions": native_count, - "vllm_commit": base["vllm"]["commit"], "b12x_commit": base["b12x"]["commit"], + "vllm_commit": base["vllm"]["commit"], + "b12x_commit": compute["b12x_revision"], "sparkcache_commit": base["sparkcache"]["commit"], "sircl_native_sha256": base["sircl"]["native_sha256"], "marker_source_sha256": sha256(marker_source), "marker_binary_sha256": sha256(marker), "rocenante_lazy_import": str(roce.__file__), "cuda_initialized": False, "device_access": False, "model_loaded": False, + "compute": compute, "readiness_warmup": warmup, "limitation": "Content and CPU checks do not qualify CUDA graphs, RDMA forwarding, native MTP, cache restoration, or model performance.", } diff --git a/spark_transport/experiments/glm53_rocenante_overlay/rocenante_vllm_overlay.py b/spark_transport/experiments/glm53_rocenante_overlay/rocenante_vllm_overlay.py index 1a0acfd8..068209d2 100644 --- a/spark_transport/experiments/glm53_rocenante_overlay/rocenante_vllm_overlay.py +++ b/spark_transport/experiments/glm53_rocenante_overlay/rocenante_vllm_overlay.py @@ -435,8 +435,7 @@ def all_reduce(self, tensor: Any) -> Any: try: if capturing: stream = torch.cuda.current_stream(self.device) - with self._runtime.capture(stream=stream): - result = self._runtime.all_reduce(tensor, stream=stream) + result = self._runtime.all_reduce(tensor, stream=stream) self._captured_nodes += 1 else: result = self._runtime.all_reduce(tensor) diff --git a/spark_transport/experiments/glm53_rocenante_overlay/test_stream_lifetime.py b/spark_transport/experiments/glm53_rocenante_overlay/test_stream_lifetime.py new file mode 100644 index 00000000..e6b4fb14 --- /dev/null +++ b/spark_transport/experiments/glm53_rocenante_overlay/test_stream_lifetime.py @@ -0,0 +1,81 @@ +"""Exercise stream admission and shared staging order without CUDA hardware.""" +import ast +from contextlib import contextmanager, nullcontext +from pathlib import Path +from types import SimpleNamespace + +import pytest + +ROOT = Path(__file__).resolve().parents[3] +SOURCE = ROOT / 'third_party/b12x_roce/b12x/comm/roce/roce_oneshot.py' + + +def methods(names, namespace): + tree = ast.parse(SOURCE.read_text()) + nodes = [node for cls in tree.body if isinstance(cls, ast.ClassDef) + for node in cls.body if isinstance(node, ast.FunctionDef) and node.name in names] + assert len(nodes) == len(names) + exec(compile(ast.Module(body=nodes, type_ignores=[]), str(SOURCE), 'exec'), namespace) + return namespace + + +def test_per_call_capture_context_cannot_reset_cuda_capture_identity(): + state = SimpleNamespace(stream='A', capture_id=123) + namespace = methods({'capture', '_order_stream'}, { + 'contextmanager':contextmanager, 'Optional':__import__('typing').Optional, + 'torch':SimpleNamespace(cuda=SimpleNamespace(current_stream=lambda device:state.stream)), + '_capture_id':lambda stream:state.capture_id}) + runtime = SimpleNamespace(device=0, _capture_id=0, _capture_stream=None, _last_stream=None) + with namespace['capture'](runtime): + namespace['_order_stream'](runtime, True) + state.stream = 'B' + with pytest.raises(RuntimeError, match='one stream'): + with namespace['capture'](runtime): + namespace['_order_stream'](runtime, True) + state.capture_id = 124 + with namespace['capture'](runtime): + namespace['_order_stream'](runtime, True) + assert runtime._capture_stream == 'B' + + +def test_misaligned_input_waits_before_shared_scratch_copy(): + events = [] + class Tensor: + shape = (4, 4096) + dtype = 'bf16' + device = 0 + def __init__(self, address): self.address = address + def is_contiguous(self): return True + def data_ptr(self): return self.address + def numel(self): return 4 * 4096 + def element_size(self): return 2 + def copy_(self, other): events.append('copy') + cuda = SimpleNamespace(device=lambda _:nullcontext(), is_current_stream_capturing=lambda:False) + namespace = methods({'all_reduce'}, { + 'torch':SimpleNamespace(Tensor=Tensor,cuda=cuda), + 'Optional':__import__('typing').Optional, 'Sequence':__import__('typing').Sequence, + '_nullcontext':nullcontext, 'PACK_BYTES':16, + 'is_launcher_prepared':lambda *args:True, + 'get_launcher':lambda *args:lambda *a:events.append('launch')}) + runtime = SimpleNamespace(_lock=nullcontext(), device=0, check_health=lambda:None, + should_allreduce=lambda inp:True, _launcher_key=lambda dtype:(), + _aligned_scratch=lambda which,like:Tensor(32), + _order_stream=lambda capturing:events.append('wait'), + _mark_stream=lambda capturing:events.append('record'), + _recv_base=0,_flag_base=0,_send_base=0,_ctrl_base=0,_slot_bytes=0, + _epoch_address=0,spin_limit=1,_blocks=1) + namespace['all_reduce'](runtime, Tensor(2), out=Tensor(18)) + assert events == ['wait','copy','launch','copy','record'] + + +def test_padded_gather_orders_staging_and_records_after_output_copy(): + tree = ast.parse(SOURCE.read_text()) + method = next(n for n in ast.walk(tree) if isinstance(n, ast.FunctionDef) and n.name=='all_gather') + calls = [n for n in ast.walk(method) if isinstance(n, ast.Call) and isinstance(n.func, ast.Attribute)] + # Select the padded path after the aligned path's early return. + stage = next(n for n in calls if n.func.attr=='copy_' and isinstance(n.func.value, ast.Subscript)) + waits = [n.lineno for n in calls if n.func.attr=='_order_stream'] + records = [n.lineno for n in calls if n.func.attr=='_mark_stream'] + output = max(n.lineno for n in calls if n.func.attr=='copy_') + assert max(waits) < stage.lineno + assert max(records) > output diff --git a/third_party/b12x_roce/README.md b/third_party/b12x_roce/README.md index f8ae6169..a2120a67 100644 --- a/third_party/b12x_roce/README.md +++ b/third_party/b12x_roce/README.md @@ -10,6 +10,10 @@ unpublished Git ref because the complete selected source is included here. The bundle builder verifies its digest before copying it. `LICENSE` contains the upstream Apache-2.0 license; source copyright notices are retained. +Local stream-safety modifications preserve CUDA capture-ID admission across +Python context boundaries. Eager calls wait for the preceding operation before +writing shared staging buffers, and record completion after output copies. + ## Attribution and design origins RoCEnante originates with Local Inference Lab's contributors, not SparkRing. diff --git a/third_party/b12x_roce/b12x/comm/roce/roce_oneshot.py b/third_party/b12x_roce/b12x/comm/roce/roce_oneshot.py index 84d2c7ec..2ee48e7d 100644 --- a/third_party/b12x_roce/b12x/comm/roce/roce_oneshot.py +++ b/third_party/b12x_roce/b12x/comm/roce/roce_oneshot.py @@ -561,15 +561,17 @@ def all_reduce( "RoCE all-reduce launcher must be prepared before CUDA graph capture" ) launcher = get_launcher(*key) - if out is None: - out = torch.empty_like(inp) - src = inp + if out is None: + out = torch.empty_like(inp) + # Shared staging buffers belong to the preceding operation until + # its completion event, including its output copy, has fired. + self._order_stream(capturing) + src = inp if inp.data_ptr() % PACK_BYTES != 0: src = self._aligned_scratch(0, inp) src.copy_(inp) dst = out if out.data_ptr() % PACK_BYTES == 0 else self._aligned_scratch(1, out) - self._order_stream(capturing) - launcher( + launcher( src.data_ptr(), dst.data_ptr(), nbytes // PACK_BYTES, @@ -760,24 +762,26 @@ def all_gather( # requested layout. nbytes = inp.numel() * inp.element_size() padded = _align_up(nbytes, PACK_BYTES) - staged, gathered = self._gather_scratch(padded) - staged[:nbytes].copy_(inp.reshape(-1).view(torch.uint8)) - self._order_stream(capturing) + staged, gathered = self._gather_scratch(padded) + self._order_stream(capturing) + staged[:nbytes].copy_(inp.reshape(-1).view(torch.uint8)) self._launch_gather( staged.data_ptr(), gathered.data_ptr(), padded, padded // PACK_BYTES ) - self._mark_stream(capturing) - stacked = ( + stacked = ( gathered.view(self.world_size, padded)[:, :nbytes] .reshape(-1) .view(inp.dtype) .reshape(self.world_size, *inp.shape) ) - result = stacked.movedim(0, dim).reshape(shape) - if out is None: - return result.contiguous() - out.copy_(result) - return out + result = stacked.movedim(0, dim).reshape(shape) + if out is None: + result = result.contiguous() + else: + out.copy_(result) + result = out + self._mark_stream(capturing) + return result def _gather_scratch(self, padded: int) -> tuple[torch.Tensor, torch.Tensor]: """Fixed device scratch for the padded all-gather path, allocated once. @@ -838,15 +842,13 @@ def capture(self, stream: object = None, *, channel_id: Optional[str] = None): capture began. Inside, every collective must use one stream. """ - self._last_stream = None - self._capture_stream = None - self._capture_id = 0 - try: - yield self - finally: - self._capture_stream = None - self._capture_id = 0 - self._last_stream = None + self._last_stream = None + # CUDA's capture ID owns stream admission, not this Python context. + # Nested or per-call contexts must retain the same-capture guard. + try: + yield self + finally: + self._last_stream = None # -- diagnostics / lifecycle -------------------------------------------------- diff --git a/third_party/b12x_roce/provenance.json b/third_party/b12x_roce/provenance.json index 70e4e88c..f6deddfb 100644 --- a/third_party/b12x_roce/provenance.json +++ b/third_party/b12x_roce/provenance.json @@ -5,7 +5,8 @@ "upstream_base_commit": "ffb7442d04a9f50b950df1fb17280acad881b7d5", "license": "Apache-2.0", "scope": "Only b12x.comm.roce; no model kernels or other B12X package modules", - "roce_tree_sha256": "3acb9c3cda49931f7988b9b426bd1304df029de95d868a90346df2d09ed8f627", + "roce_tree_sha256": "902a9dfd1a9c8ec379b002b13737701dd6bc58e240abcb5ab9b443455961a3a4", + "local_changes": ["Preserve capture-ID stream admission across Python contexts", "Order shared input staging and output copies across streams"], "git_state": { "commit": "eac260a8257cc6b14e7d4ad674f51e9a09b8790f", "roce_source_dirty": false,