perf(kimi-k3): DMA-ring TP all-reduces and deferred draft ingest for chunked prefill - #564
Open
myshytf wants to merge 68 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
|
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 (69)
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 was referenced Sep 2, 2026
Open
Kimi-K3 DSpark context normalization and per-layer context-KV projection now use the tensor returned by the in-place-capable TP all-reduce interface. The custom all-reduce backend may return distinct storage, while PyNCCL and single-rank paths may return the input. This preserves the reduced squared norm and projected KV values for every backend. The input tensors remain dead after each collective, so no caller-visible aliasing contract changes. Tests: the streamed auxiliary normalization and per-layer context projection tests pass with reducers that return modified out-of-place tensors. Ruff check passed for the implementation and test. Co-authored-by: OpenAI Codex <noreply@openai.com>
B12X DMA all-reduce is now enabled only when the selected PCIe runtime explicitly reports all-peer auxiliary connectivity. A runtime without the capability attribute keeps larger tensors on PyNCCL instead of assuming an unverified topology is safe. Runtimes that report support retain the configured DMA crossover. Runtimes that report false retain the existing fallback. Tests: test_b12x_dispatcher_prepares_single_stable_eager_owner, 4 passed, including a runtime without the capability attribute. Ruff check passed for the implementation and test. Co-authored-by: OpenAI Codex <noreply@openai.com>
Deferred prefill ingest workers append failures to a lock-protected queue. The verifier atomically drains every queued failure and disables all affected request IDs, so concurrent ring jobs cannot overwrite each other and leave a request drafting against incomplete remote KV state. The synchronous RPC lock remains independent from failure delivery, and successful ingest behavior is unchanged. Tests: test_deferred_ingest_failures_disable_every_affected_request passed with two queued failures covering three requests. Ruff check passed for the implementation and test. Co-authored-by: OpenAI Codex <noreply@openai.com>
Author
|
@coderabbitai review Head
Focused results: 2 DSpark out-of-place reduction tests passed, 4 PCIe |
|
|
Feature capture keeps its existing JSON object and records-array format, but appends each record by replacing only the closing byte-count tail. Capture time and in-memory state are therefore constant per proposal instead of growing with the request history. A reset recreates the binary and index files, and freeing a remote request releases its capture bookkeeping. Binary offsets and compact-index fields remain compatible with existing consumers. Tests: test_capture_index_appends_records_without_retaining_the_record_list passed for two records, binary layout, offsets, valid JSON, and bounded state. Ruff check passed for the implementation and test. Co-authored-by: OpenAI Codex <noreply@openai.com>
FlashInfer PCIe IPC workspaces allocate one CUDA slab per rank and open every nonlocal rank handle. Expose all-peer auxiliary support only after a workspace has completed that mapping, and revoke the capability when the pool is empty or closed. This restores the configured DMA crossover for successfully prepared TP2, TP4, and TP8 FlashInfer pools while preserving fail-closed behavior before preparation and after teardown. Tests: tests/distributed/test_flashinfer_pcie_all_reduce.py reports 7 passed. Ruff check and format check pass for both changed files. Co-authored-by: OpenAI Codex <noreply@openai.com>
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.
Purpose and status
Status: the async-ingest and DMA transport is implemented, qualified, and
serving for Kimi-K3 TP8 chunked prefill. The review safeguards are
implemented and unit-qualified, not deployed.
This change removes host serialization from remote draft context ingest and
routes eligible prefill collectives through the b12x PCIe DMA ring. The source
paths define the complete transport and fallback behavior; no deployment-only
code is required.
Resulting behavior
RemoteK3DSparkSpeculatorqueues mid-prefill context ingest when no activerequest sampled a token. A pinned staging ring with
VLLM_K3_DRAFT_ASYNC_PREFILL_INGESTslots owns each multipart request untilits worker receives the reply. Setting the value to
0disables deferral.copy=False; the RPCretains those buffers until the reply arrives.
every failure and disables every affected request until it leaves the batch;
concurrent workers cannot overwrite an earlier failure.
appending one record and the byte-count tail per proposal. Work and in-memory
state stay constant per record, and request release drops its bookkeeping.
VLLM_PCIE_DMA_MIN_BYTESuseb12x.comm.pcie.DmaAllReduceonly when the selected PCIe runtime explicitlyreports all-peer auxiliary connectivity. Missing or false capability keeps
large tensors on PyNCCL.
all_reduce_in_placetries the custom all-reduce before PyNCCL. Callers treatthe input as dead and consume the returned tensor, which permits either
aliased or out-of-place backend results.
The DMA ring uses lossless BF16 reduce-scatter and all-gather hops. Unsupported
topologies, unavailable b12x support, disabled thresholds, and initialization
failures retain PyNCCL. Draft RPC failures retain the no-draft verifier fallback,
so target-model output remains authoritative.
Evidence
Hardware: eight NVIDIA RTX PRO 6000 Blackwell GPUs behind a PEX88000 Gen4
fabric, with a remote DFlash draft on a ninth GPU.
approximately 1.75 ms on PyNCCL.
the DMA route enabled. Decode inter-token latency and greedy output were
unchanged in the qualification runs.
tensors, covering global context normalization and per-layer context-KV
projection.
all-peer capability attribute.
three distinct requests.
AI assistance from Claude Code and OpenAI Codex was used. The submitter must
review and validate every changed line before merge.