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>
Share KV loads across fixed K=3 verification rows, select capacity-specific graph plans, and add guarded q-rep and sparse policies. Assisted-by: OpenAI Codex Signed-off-by: myshytf <9619163+myshytf@users.noreply.github.com>
The dense MLA kernel (b12x) now shares each request's live 64-token chunks evenly over the launched splits, so an eager launch needs min(num_splits, live chunks) splits rather than the plan-prefix ceil(live chunks / chunks_per_split). Both launches partition the chunks exactly as the full-plan CUDA-graph launch does; the removed formula left most CTAs idle on sequences shorter than the plan. Co-Authored-By: Claude Code <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01HPWxmKzfikaemyykd3p89D
The b12x dense MLA plan accepts `partial_dtype` (element type of the split partials the merge reads) and `single_split_chunks` (largest live chunk count one split scans alone; balanced ranges above). Two environment variables select them for every K3 dense MLA plan: - VLLM_K3_DENSE_MLA_PARTIAL_DTYPE: "bf16" (default, unchanged) or "fp32" (partials kept exact, merged results rounded once). - VLLM_K3_DENSE_MLA_SINGLE_SPLIT_CHUNKS: -1 (default: the plan's chunks per split, i.e. the fixed-range association for requests that fit one run) or a chunk count; 0 balances every request. Eager launches now use one split for requests within the threshold, so those write the output directly without a merge; longer requests keep one split per live chunk. Validation: tests/v1/attention/test_b12x_mla.py (45 passed in the production image) covers the launch count with and without a threshold and the environment parsing; the served-lineage test expectations for the balanced launch count are updated to the balanced rule. Co-Authored-By: Claude Code <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01HPWxmKzfikaemyykd3p89D
…ext output Three dense-MLA metadata and output-storage fixes for the fused DCP verification path: - Verify plans (fp8 KV, four-query tiles) are created per power-of-two batch capacity (`_dense_mla_plan_row_caps`) and `build` selects the smallest covering capacity, like the decode plans; the batch range is bounded by the flattened row capacity (four rows per request). One plan per batch value grew linearly with max_num_seqs and exceeded the 1,024-row plan limit from batch 257. - The plan's page table must cover the largest local KV shard: `build` copies the worker's block table into the plan-width flattened table and drops columns past that width (KV-block rounding can make the worker table wider while no local sequence references those columns); a plan narrower than the shard would drop referenced pages, so the builder now rejects it (a sliding-window spec shrinking the plan) instead of clamping. - `_reuse_consumed_query_for_context_output` allocates fresh storage when the consumed query holds fewer bytes than the compact bf16 context output (an fp8 Kimi-K3 query row is 192 bytes, the output row 256), instead of raising on every fp8 prefill with chunked context. Validation: tests/v1/attention/test_b12x_mla.py (38 passed, new covering- bucket test) and tests/models/kimi_k3/test_mla_padding.py (14 passed; the fp8 case now uses the production 192-wide query) in the SM120 image. Co-Authored-By: Claude Code <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01HPWxmKzfikaemyykd3p89D
…kimi-k3-dense-mla-balanced-splits-20260902-pr
…to agent/kimi-k3-dense-mla-split-precision-20260903-pr
Copy cache-on-write source pages through each tensor view and its actual block axis, preserving storage offsets and padded block strides in heterogeneous DCP layouts. Keep the raw-storage path only for layouts without an identifiable block dimension. This preserves 1,536-token partial prefix reuse for Kimi-K3 fp8_ds_mla instead of coarsening hits to the 12,288-token DCP page boundary. Validated on CUDA with strided uint8 views, nonzero block-axis layouts, and the production 94x1536x656 packed MLA shape. Signed-off-by: myshytf <9619163+myshytf@users.noreply.github.com>
Carry the originating KV cache group on each scheduler CoW pair and apply it only to that group’s logical cache views. Shared allocator block IDs can otherwise overwrite live recurrent state in another heterogeneous cache group. Keep untagged copies as a compatibility broadcast and preserve 1,536-token DCP partial-prefix reuse. Assisted-by: OpenAI Codex Signed-off-by: myshytf <9619163+myshytf@users.noreply.github.com>
Defer only the first draft proposal after an fp8_ds_mla cache hit that ends inside a DCP attention page. The target advances one token without speculative verification, then the configured draft depth resumes on the following step. This preserves 1,536-token prefix reuse, packed KV cache, DFlash, CUDA graphs, and asynchronous LMCache while avoiding non-finite target state in the immediate multi-token verification block. Assisted-by: OpenAI Codex Signed-off-by: myshytf <9619163+myshytf@users.noreply.github.com>
Emit zero draft placeholders on the final prefill step of an fp8_ds_mla cache resume that lands inside a DCP page. This makes the first resumed target step single-token under async scheduling, then restores the configured speculative depth on the following step. Remove the worker-side draft trimming because async scheduling had already created the next step placeholders before worker output arrived. Assisted-by: OpenAI Codex <codex@openai.com> Signed-off-by: myshytf <9619163+myshytf@users.noreply.github.com>
Treat a completed remote-KV request as already admitted when it resumes with its reserved HMA blocks. Per-step allocation still enforces free-block limits, while avoiding a second full-sequence reservation check. This makes the committed scheduler safe to mount read-only in the production runtime, where the same guard was previously injected at boot. Assisted-by: OpenAI Codex <codex@openai.com> Signed-off-by: myshytf <9619163+myshytf@users.noreply.github.com>
Block 0 is the placeholder every block table pads with, and its pages are shared by all KV cache groups of one layer index, so attention readers and recurrent kernels treat it as all zeros. A KV connector that receives a request's block table with placeholder slots can scatter loaded state into it; the LMCache MP connector did so for the recurrent checkpoint slots of a resumed hybrid request, and the packed fp8_ds_mla MLA reader then loaded non-finite records from slot 0 for its masked candidates. The V2 runner now zeroes block 0 through the existing KV block zeroer whenever a step reports completed KV receives, including no-forward steps. The loaded request is scheduled only after that report, so it never reads a page a load overwrote. Steps without completed receives are unchanged. Validation: tests/v1/worker/test_gpu_model_runner_v2_null_block_guard.py (3 passed) with the production image interpreter.
Adds an opt-in decode window to the B12X packed fp8_ds_mla adapter. With VLLM_K3_MLA_DECODE_WINDOW_TOKENS=N (or an integer in the file named by VLLM_K3_MLA_DECODE_WINDOW_FILE, re-read on every metadata build so a research boot can switch without restarting), every decode query attends to the N most recent global positions instead of the whole visible context. Under DCP the window start is projected to each rank's local slot (token p lives on rank p % dcp at slot p // dcp); the physical-slot table is compacted from that start and the candidate count passed to the reader shrinks to match. Prefill stays exact-dense and every token is still written to the cache, so this changes the attention semantics of the 24 MLA layers at decode time only. It is a measurement aid for the sparse-attention question and is not a production feature. With the window disabled the slot table and the kernel arguments carry the same values as before; the candidate count now travels through a persistent builder buffer so a CUDA graph captured in either mode replays with the current step's values. Validation: tests/v1/attention/test_b12x_mla.py, 66 passed on an SM120 GPU with the production image interpreter (one pre-existing failure needs the split-precision env of PR local-inference-lab#595); window start arithmetic covered for DCP 8 and DCP 1, the Triton compaction, the builder attach in both modes, and the file override.
|
Warning Review limit reachedNext included review available in 41 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 (74)
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.
Packed MLA recent-window research control
Status: research-only; unsupported for production. This PR publishes an already-prepared experiment at the operator's request. It does not enable the feature in serving and must not be interpreted as a precision-preserving optimization.
Commit
d619dd51e573adds an opt-in recent-token window to the packed Kimi-K3 MLA decode adapter. A positiveVLLM_K3_MLA_DECODE_WINDOW_TOKENSrestricts decode attention to that many recent global positions. A file override allows a research process to change the value between metadata builds. With the window disabled, exact-dense slot selection remains the default.Under decode context parallelism the global window start is projected into each rank's local cache slots. The physical-slot table is compacted, and a persistent per-row candidate-count buffer supplies updated visibility to graph replay. Prefill remains dense and every token is still stored, but a positive window changes decode attention semantics and therefore model outputs.
Compatibility and scope
The branch is based on a historical packed-KV Kimi integration, with the incremental window change confined to
vllm/v1/attention/backends/mla/b12x_mla.pyand its tests. Lab #644 supplies the broader packed-KV path. This PR is separate research infrastructure and does not replace #644 or its exact behavior.The configured serving policy forbids precision-lowering optimization. That policy remains intact: this branch is being published for inspection only. No quality or throughput claim is made for a nonzero window. The current target retains exact context attention.
Recorded validation
The original commit records
tests/v1/attention/test_b12x_mla.py: 66 passed on SM120 in the production image, with one unrelated existing test requiring the split-precision configuration of #595. Coverage includes DCP1/DCP8 window-start arithmetic, Triton slot compaction, metadata attachment with the control disabled/enabled, and the file override.Those GPU tests were not rerun during this publication-only audit.
git diff --checkpasses. The branch requires a current-base port and quality evaluation before any deployment consideration.Duplicate searches found no lab PR for this recent-token-window control. Source and its recorded evidence were reviewed before publication. AI assistance was used; no human merge approval is claimed.
🤖 Generated with Claude Code
https://claude.ai/code/session_01KxvNwugeU8RJFd7WRYwNLG