perf(distributed): route mid-size TP all-reduces to the bf16 PCIe two-shot (Kimi-K3 wrapper) - #593
Open
myshytf wants to merge 64 commits into
Conversation
Record scheduler-side speculative widths in GrammarOutput so worker-side draft trimming cannot shift flattened grammar masks onto later requests. Destination logits continue to use the worker-visible width, while source offsets use the serialized scheduler width. Validated with focused unit coverage and a 160-request concurrent DeepSeek V4 structured-output workload.
KimiK3ToolParser.extract_tool_calls_streaming matched calls with _call_re, which requires the closing <|close|>call<|sep|> marker. Until that marker arrived nothing was emitted for the call, so a long tool call produced no SSE deltas for the whole generation and then dumped the entire arguments JSON in one delta. Track the call from its <|open|>call ...<|sep|> marker instead. The name goes out immediately, and _partial_arguments serializes the arguments seen so far as a prefix of the final JSON, so each step can stream the difference against what it already sent. String argument bodies are raw text, so they are forwarded as they arrive with a trailing partial close marker held back; other types still need the whole literal to decode and are held until their block closes. The concatenated deltas are byte-identical to the non-streaming extract_tool_calls output. Signed-off-by: guptaishaan <guptaishaan@users.noreply.github.com>
Withhold whitespace-tolerant argument-close fragments until they form a complete XTML marker. This keeps streamed JSON argument deltas prefix-stable for every marker form accepted by the parser. Co-authored-by: OpenAI Codex <noreply@openai.com>
Co-authored-by: Codex <codex@openai.com>
Co-authored-by: Codex <codex@openai.com>
Document the target model input and optional NeoX layout result using the repository's Google-style docstring contract. This is documentation-only and does not change runtime behavior. Co-authored-by: OpenAI Codex <codex@openai.com>
Initialize fresh assistant generations in the reasoning channel when Kimi thinking is enabled, while preserving rendered marker state for continued assistant messages. Filter complete and split XTML control markers at the composed parser boundary so malformed model transitions cannot expose protocol syntax as API content. The thinking-disabled path and continuation semantics remain unchanged. Validation: 72 Kimi K3 reasoning and tool-parser tests; Ruff format and lint; git diff whitespace validation.
Signed-off-by: jungjiyu <libraryofjiyu@gmail.com> Assisted-by: ChatGPT
Model a 17-group hybrid KV layout and report a load failure from the final group. The test requires failure_policy=fail to finish only the affected request, emit an error result, and schedule a subsequent healthy request.\n\nValidation: 20 KV load-failure tests and 7 hybrid/Mamba scheduler tests pass in the CUDA 13.3 PyTorch 2.13 runtime.
Stop accepting speculative token batches when the grammar matcher reaches its terminal state. Preserve terminal-state tracking across validation and acceptance calls so tokens after a complete structured value cannot be committed. This is the Infernal Invocation backport of vllm-project#52805 commits d8cde608cf1f3de406c75f081a76a0e6eb55a9cb, 1cf6f25351357354cf8c520c0b2976b029429668, and 1856abd22452c3da67364986ece7245fce52c950. Signed-off-by: Martin Vit <martin@voipmonitor.org>
Structured-output masks are prepared before speculative verification. An accepted block can cross reasoning activation or grammar termination, so its suffix may have been sampled under a grammar state that no longer applies at commit time. Validate the accepted block without advancing the matcher, commit only its valid prefix, and roll scheduler accounting back for resampling. Preserve the unstructured and single-token fast paths, and report only committed draft tokens in speculative metrics. Co-authored-by: Adam Moisa <adammoisa@gmail.com> Assisted-by: OpenAI Codex Signed-off-by: Martin Vit <martin@voipmonitor.org> (cherry picked from commit fa0777f) Signed-off-by: Martin Vit <martin@voipmonitor.org>
Infernal Invocation exposes prompt inspection through is_reasoning_end_for_prompt. Make the upstream structured-output regression fixture implement the branch contract so it exercises the production method instead of a stale mock interface. Signed-off-by: Martin Vit <martin@voipmonitor.org>
Type the conditional Kimi compact-RoPE protection scope through the shared context-manager interface. Both the Kimi protection context and the no-op context retain their existing runtime behavior. Signed-off-by: Martin Vit <martin@voipmonitor.org>
The debug branch initializes the event list before every sweep point. Assert that invariant after detaching the list from the model runner so static analysis can verify indexed event access. Profiling and warmup behavior are unchanged. Signed-off-by: Martin Vit <martin@voipmonitor.org>
…DFlash aux state (vllm-project#50487) Signed-off-by: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com> Co-authored-by: Janelle Cai <janelle.cai@modal.com> (cherry picked from commit 03a8d0b)
Verify that disabled AttnRes capture returns before reading unavailable weights and that enabled capture selects both normalization and projection weights from the correct consumer. Document the capture interface parameters and return value.
Compute MoonViT rotary frequencies only for the image grid sizes present in each request instead of materializing the configured 512x512 ceiling. This reduces the measured first-image CUDA allocation peak from 340,018,176 bytes to 1,990,656 bytes for a 36x36 grid while preserving bit-identical CPU and CUDA output. Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Martin Vit <martin@voipmonitor.org>
Project independent Kimi vision features separately so MXFP8/Marlin workspace scales with the largest image instead of the sum of all scheduled images. Preserve output order, shape, activation dtype, and numerical results while reducing the measured TP16 three-image transient peak by 32.52 MiB. Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Martin Vit <martin@voipmonitor.org>
Define token-position DCP shard count on each cache specification and use max_num_blocks_per_req as the worker block-table width contract. Attention caches retain full, partial, or replicated DCP layouts; recurrent caches report one token-position shard and preserve their mode-specific table width. This removes the model runner's cache-type special case while retaining the 1,310-column Mamba align table required by a 1,000,000-token model length with 768-token blocks and seven speculative blocks. Assisted-by: OpenAI Codex <noreply@openai.com> Signed-off-by: Martin Vit <martin@voipmonitor.org>
Signed-off-by: Martin Vit <martin@voipmonitor.org>
Gather each tensor-parallel vision shard at its produced row count instead of padding every rank to the largest shard. This preserves embedding order and the uniform-size fast path while preventing the transient allocation from scaling with TP size when a request contains fewer images than ranks. Validate zero-length PyNccl inputs, single-image output parity, empty inputs, uneven four-GPU assignments, and multi-image assignments. A TP16 Kimi-K3-shaped harness reduces the collective output from 224 MiB to 14 MiB per GPU with bit-exact gathered content. Signed-off-by: Martin Vit <martin@voipmonitor.org>
Signed-off-by: Martin Vit <martin@voipmonitor.org>
Cache each head's prefix and suffix log-sum-exp values before any output write when the thread group fits inside a CUDA block. This preserves chunked-attention accumulators that pass the running LSE tensor as both prefix input and output destination, while retaining the direct-load path for head groups that cross block boundaries. Index all cached values through the declared tensor strides.\n\nAdd exact in-place versus disjoint-output coverage for the six-head, 128-element MLA geometry at 256 and 4096 tokens.\n\nThe shared-memory loading structure adapts vLLM PR vllm-project#45778 (commit c71576f) to the strided-LSE kernel contract.\n\nCo-authored-by: nicole-lihui <nicole.li@daocloud.io> Signed-off-by: Martin Vit <martin@voipmonitor.org>
… GPU Adds a verifier-side proxy (RemoteK3DSparkSpeculator) and a standalone draft server (vllm.entrypoints.k3_dspark_standalone + k3_dspark_rpc) so the DSpark draft model executes on its own single GPU while the target runs TP/DCP on separate GPUs. Draft weights, KV, Markov head, and CUDA graphs live entirely on the draft process; the target exchanges context and proposals over a versioned ZMQ/TCP protocol (PROTOCOL_VERSION=2). Behavior and invariants: - VLLM_K3_DRAFT_REMOTE_ADDRESS selects the remote path at speculator construction; unset preserves the existing local DSpark/DFlash path. - propose() matches BaseSpeculator's signature; rank 0 performs RPC and all ranks consume the broadcast result. - Fail closed: any RPC failure fills draft tokens with -1 (no speculation for the step) and disables affected requests until they leave the batch; FREE remains safe for never-created remote state. - Retained-prefix reconnection validates a target prefix-cache hit against retained draft state via a host-visible view of the request token table (InputBatch.all_token_ids_cpu, backed by StagedWriteTensor.cpu). - CUDA-graph capture interface preserved: init_cudagraph_manager and capture(capture_phase=...) conform to BaseSpeculator. Compatibility: no change when the remote address is unset; draft side supports DSpark and DFlash checkpoints on a single GPU including Ampere-class cards. Validation: 19 new CPU unit tests pass (test_k3_dspark_remote_speculator.py, test_k3_dspark_standalone.py); production-qualified serving lukealonso/Kimi-K3-QSRT-K2 TP8/DCP8 with an Inferact BF16 DSpark draft on a dedicated RTX 3090. Limitations: one remote draft process (draft TP1); TCP transport; greedy draft sampling with block rejection sampling on the verifier. AI assistance was used in the preparation of this change; every line was reviewed and the listed tests were run by the submitter. Signed-off-by: myshytf <9619163+myshytf@users.noreply.github.com>
Signed-off-by: myshytf <9619163+myshytf@users.noreply.github.com>
Keep DSpark and DFlash scheduling lookahead semantics while applying EAGLE's last-hash target-cache drop only when an actual target KV group is marked as EAGLE. This preserves fine target APC tails for remote/disaggregated drafts and retains the legacy fallback for classic EAGLE. Signed-off-by: myshytf <9619163+myshytf@users.noreply.github.com>
Signed-off-by: myshytf <9619163+myshytf@users.noreply.github.com>
… <=1 output token When a batch contains only new requests (no running ones) and every one has max_tokens <= 1, set num_spec_tokens_to_schedule = 0. Speculative decoding cannot help a 1-token output, so the draft pass and verification are pure overhead. This is the shape of every max_tokens=1 API call, every prefill-throughput benchmark, and every embedding/classification-style request. Measured on RTX 5090 (31.4 GiB), Qwen3.8-27B EXL3, MTP=6: 1-token request latency 141 ms -> 127 ms 2051-token prefill bench 7445 -> 7635 tok/s (+2.5%) TG on normal requests 189.8 tok/s (unchanged) The guard is conservative: it requires scheduled_running_reqs to be empty, so an in-flight multi-token generation can never lose its draft tokens. Signed-off-by: Michel Belleau <michel.belleau@malaiwah.com>
…rmless for single-token requests
Call the finalized FlashInfer workspace prepare API during vLLM graph warmup so autotune and cache lookup complete before CUDA graph capture. Signed-off-by: myshytf <9619163+myshytf@users.noreply.github.com>
Records the remote DFlash/DSpark speculator as served on 2026-09-01: the probabilistic logits transport (multipart RPC frames), the prefix-reconnect partial-window gate and its diagnostics, the per-phase timing detail, and the feature-capture hook. These run in production as marker-guarded overlay patches on top of this branch's file; committing the snapshot gives the following performance commits a reviewable base. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Bt7bK1Ru7Ywq4s7QBwunsa
For a scheduler step in which no active request sampled a token (every request is still inside its prompt) the draft server's proposal is never consumed, but rank 0 still copied the aux hidden rows to host, serialized 132 MiB per 1,536-token chunk and blocked on the ZMQ round trip while the other TP ranks waited at the next collective: 0.2-0.3 s of device idle at every chunk boundary. Such steps now hand the frames to a worker thread that owns the socket in FIFO order and return after the pinned-memory copy. A ring of pinned staging buffers (depth `VLLM_K3_DRAFT_ASYNC_PREFILL_INGEST`, default 2, 0 disables) keeps each slot reserved until its reply arrives, every synchronous RPC drains the queue first so request ordering on the draft server is unchanged, and a failed deferred ingest disables drafting for its requests exactly like a failed synchronous proposal. Steps that sampled a token keep the synchronous path. Validation: 8K cold prefill 1,105-1,128 -> 1,348-1,375 tok/s on the production target with decode, acceptance length and outputs unchanged; no deferred-ingest failures in production logs. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Bt7bK1Ru7Ywq4s7QBwunsa
The synchronous proposal serialized the aux rows with `.numpy().tobytes()`, a Python-level copy of up to 132 MiB (~40 ms) before the ZMQ send. The pinned staging buffers are not rewritten until the next proposal and `_rpc` waits for the reply, so the frames now reference the staging memory directly and are sent with `copy=False`. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Bt7bK1Ru7Ywq4s7QBwunsa
With the FlashInfer PCIe IPC backend selected for decode-size one-shot shapes, every all-reduce above the one-shot limit fell back to PyNCCL. The b12x `DmaAllReduce` (a CE-driven reduce-scatter + all-gather ring, lossless bf16 per hop like the NCCL ring) is now initialized alongside the FlashInfer one-shot when `VLLM_PCIE_DMA_MIN_BYTES` is set, and tensors at or above that size take it. The backend name reports `FLASHINFER_PCIE_IPC_B12X_DMA` when the ring is active. `VLLM_PCIE_DMA_MIN_BYTES=off` keeps PyNCCL. Measured on the production target (TP8 over PCIe Gen4, 1,536-token chunks): a 22 MiB reduce takes ~1.0 ms on the ring versus ~1.75 ms on NCCL; a rank-0 trace of a 7,680-token prefill had NCCL all-reduce at 57% of GPU time before this change. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Bt7bK1Ru7Ywq4s7QBwunsa
`all_reduce_in_place` went straight to PyNCCL. Kimi's row-parallel prefill projections (attention output and MoE final output for >= 1,024 rows) use that entry, so 805 of the 1,395 all-reduces of a 7,680-token prefill (all of the 22 MiB ones) stayed on the NCCL ring while the functional entry already used the b12x DMA ring. Eligible tensors now take `ca_comm.custom_all_reduce` first; the caller consumes the returned tensor and treats the source as dead, so returning the collective's output is equivalent to the in-place result. Validation: production 8K cold prefill 1,259-1,415 -> 1,476-1,486 tok/s, decode ITL unchanged, greedy outputs unchanged. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Bt7bK1Ru7Ywq4s7QBwunsa
…-shot Between the PCIe one-shot ceiling (112 KB on the Kimi-K3 rig) and the DMA ring floor (6 MB) the custom all-reduce wrapper handed TP all-reduces to the PyNCCL ring. At TP8 with the K=3 DFlash draft every decode all-reduce is 4 tokens x 7168 x 2 B = 57 KB per request, so the ring served all 283 all-reduces of a step from concurrency 2 upwards. This routes that window to b12x's lossless bf16 two-shot (local-inference-lab/b12x#290, PCIeTwoShotBF16; the vLLM side of local-inference-lab#580 for the Kimi-K3 wrapper): payloads above the one-shot ceiling and up to VLLM_PCIE_TWOSHOT_ALLREDUCE_MAX_SIZE (default 768 KB, 0/off disables) go to the two-shot when it accepts the tensor; the one-shot keeps priority below its ceiling and the DMA ring keeps everything above the window. VLLM_PCIE_TWOSHOT_ROW_ELEMS (default 896 = 7168 / 8) is the payload row width: the two-shot serves element counts that are multiples of row_elems x world_size, and 896 keeps every token count of the K3 decode payload eligible at TP8. The two-shot is initialized on all ranks or on none (rank-consensus on failure), captured inside the wrapper's graph capture, and closed with the wrapper. Numerics: the two-shot reduces each shard in FP32 in a fixed rank order and rounds once, where the bf16 ring rounds after every hop. Validation: tests/distributed/test_k3_pcie_twoshot_route.py (device-free, mocked runtimes; Kimi-K3 production image): limit parsing, one-shot priority at and below its ceiling, two-shot from 9 tokens (129 KB) to 48 tokens (688 KB), PyNCCL between the two-shot limit and the DMA floor, DMA above, dtype acceptance and the disabled route: 10 passed. Runtime correctness on the fabric is the four-rank test of b12x#290; production measurement at concurrency 4-8 pending. Co-Authored-By: Claude Code <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01HPWxmKzfikaemyykd3p89D
|
Warning Review limit reachedNext included review available in 59 minutes. View limit detailsLimit details: You’ve used the included review currently available. You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository. Review configuration: ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Team Run ID: 📒 Files selected for processing (67)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
7 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Kimi-K3 wrapper routing for the lossless bf16 PCIe two-shot all-reduce (local-inference-lab/b12x#290,
PCIeTwoShotBF16; the design of #580 applied tocustom_all_reduce.py, which the Kimi-K3 stack uses with the FlashInfer one-shot and the B12X DMA ring). Between the one-shot ceiling (112 KB on the production rig) and the DMA floor (6 MB) the wrapper handed TP all-reduces to the PyNCCL ring; at TP8 with the K=3 DFlash draft every decode all-reduce is 57 KB per request, so the ring served all 283 all-reduces of a step from concurrency 2 upwards.VLLM_PCIE_TWOSHOT_ALLREDUCE_MAX_SIZE(default 768 KB;0/offdisables) go to the two-shot when it accepts the tensor; the one-shot keeps priority below its ceiling, the DMA ring keeps everything above the window.VLLM_PCIE_TWOSHOT_ROW_ELEMS(default 896 = 7168 / 8) is the payload row width: the two-shot serves element counts that are multiples ofrow_elems × world_size, and 896 keeps every token count of the Kimi-K3 decode payload eligible at TP8.Numerics: the two-shot reduces each shard in FP32 in a fixed rank order and rounds once; the bf16 ring it replaces rounds after every hop.
Stacked on #564 (
agent/kimi-k3-prefill-comm-20260901-pr).Validation
tests/distributed/test_k3_pcie_twoshot_route.py(device-free, mocked runtimes; Kimi-K3 production image): limit parsing, one-shot priority at and below its ceiling, two-shot from 9 tokens (129 KB) to 48 tokens (688 KB), PyNCCL between the two-shot limit and the DMA floor, DMA above, dtype acceptance and the disabled route — 10 passed. The two-shot runtime imports on the served b12x snapshot (itspcie_oneshot/pcie_twoshot/_cuda_ipcdependencies are byte-identical there). Fabric correctness runs b12x#290's four-rank torchrun test withrow_elems=896in the deployment window (candidates/k3-decode2-20260903/deploy.sh); production measurement at concurrency 4–8 pending.🤖 Generated with Claude Code
https://claude.ai/code/session_01HPWxmKzfikaemyykd3p89D