perf(mla): fuse K3 DCP verification queries by myshytf · Pull Request #565 · local-inference-lab/vllm

myshytf · 2026-09-01T19:15:26Z

Summary

One commit on top of dev/infernal-invocation (74ec7af): share KV loads across the fixed K=3 DFlash verification rows in the B12X dense MLA decode path, select capacity-specific graph plans, and add guarded q-rep and sparse policies (vllm/envs.py knobs). Kernel side is local-inference-lab/b12x#271 (perf(attention): fuse four-query dense MLA verification); the vLLM side is what production serves (vllm/v1/attention/backends/mla/b12x_mla.py byte-identical to the served overlay).

Tests added: tests/models/kimi_k3/test_mla_padding.py, tests/v1/attention/test_b12x_mla.py extensions.

Rebased from the stacked branch agent/k3-longctx-exact-20260901 so it is independent of the W4A8 opt-in change.

🤖 Generated with Claude Code

https://claude.ai/code/session_01HPWxmKzfikaemyykd3p89D

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>

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>

coderabbitai · 2026-09-01T19:15:34Z

Warning

Review limit reached

Next included review available in 21 minutes.

Check out review usage here.

View limit details

Limit 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.

Learn how review limits work.

Review configuration:

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Team

Run ID: 8b32e920-dc99-4cbb-9241-eb84ef64b2a5

📥 Commits

Reviewing files that changed from the base of the PR and between b5f995e and d461572.

📒 Files selected for processing (66)

csrc/libtorch_stable/attention/merge_attn_states.cu
tests/distributed/test_flashinfer_pcie_all_reduce.py
tests/distributed/test_pynccl.py
tests/kernels/attention/test_merge_attn_states.py
tests/models/kimi_k3/test_aux_attn_res_stream.py
tests/models/kimi_k3/test_eagle3.py
tests/models/kimi_k3/test_mla_padding.py
tests/models/kimi_k3/test_vision_projector.py
tests/models/kimi_k3/test_vision_warmup.py
tests/reasoning/test_kimi_k3_reasoning_parser.py
tests/tool_use/test_kimi_k3_tool_parser.py
tests/v1/attention/test_b12x_mla.py
tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py
tests/v1/core/test_dspark_prefix_cache_policy.py
tests/v1/core/test_kv_cache_utils.py
tests/v1/core/test_scheduler.py
tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py
tests/v1/kv_connector/unit/utils.py
tests/v1/spec_decode/test_acceptance_length_controller.py
tests/v1/spec_decode/test_dflash_causality.py
tests/v1/spec_decode/test_dflash_swa.py
tests/v1/spec_decode/test_dspark_cudagraph_contract.py
tests/v1/spec_decode/test_k3_dspark_remote_speculator.py
tests/v1/spec_decode/test_k3_dspark_standalone.py
tests/v1/spec_decode/test_mtp_structured_output.py
tests/v1/structured_output/test_reasoning_structured_output.py
tests/v1/structured_output/test_utils.py
tests/v1/worker/test_cp_utils.py
tests/v1/worker/test_gpu_structured_outputs.py
tests/v1/worker/test_mamba_hybrid_model_state.py
tests/v1/worker/test_mamba_utils.py
vllm/distributed/communication_op.py
vllm/distributed/device_communicators/flashinfer_pcie_all_reduce.py
vllm/entrypoints/k3_dspark_rpc.py
vllm/entrypoints/k3_dspark_standalone.py
vllm/envs.py
vllm/model_executor/models/kimi_k25_vit.py
vllm/model_executor/models/qwen3_dflash.py
vllm/model_executor/models/vision.py
vllm/models/kimi_k3/nvidia/mla.py
vllm/models/kimi_k3/nvidia/model.py
vllm/parser/kimi_k3.py
vllm/reasoning/kimi_k3_reasoning_parser.py
vllm/tool_parsers/kimi_k3_tool_parser.py
vllm/v1/attention/backends/flash_attn.py
vllm/v1/attention/backends/mla/b12x_mla.py
vllm/v1/core/sched/output.py
vllm/v1/core/sched/scheduler.py
vllm/v1/core/single_type_kv_cache_manager.py
vllm/v1/kv_cache_interface.py
vllm/v1/structured_output/__init__.py
vllm/v1/structured_output/backend_xgrammar.py
vllm/v1/structured_output/utils.py
vllm/v1/worker/cp_utils.py
vllm/v1/worker/gpu/buffer_utils.py
vllm/v1/worker/gpu/input_batch.py
vllm/v1/worker/gpu/model_runner.py
vllm/v1/worker/gpu/model_states/mamba_hybrid.py
vllm/v1/worker/gpu/spec_decode/__init__.py
vllm/v1/worker/gpu/spec_decode/dflash/utils.py
vllm/v1/worker/gpu/spec_decode/dspark/remote_speculator.py
vllm/v1/worker/gpu/spec_decode/dspark/utils.py
vllm/v1/worker/gpu/spec_decode/utils.py
vllm/v1/worker/gpu/structured_outputs.py
vllm/v1/worker/gpu/warmup.py
vllm/v1/worker/mamba_utils.py

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.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands.}

…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

myshytf · 2026-09-02T20:18:10Z

Review follow-up (commit d461572):

Applied

fp8 verify plans are bucketed by power-of-two batch capacity and build selects the smallest covering capacity (previously one plan per batch value, failing past batch 256 at the 1,024-row plan limit).
The builder rejects a plan narrower than the largest local KV shard (a sliding-window spec shrinking max_cache_tokens); the plan-width copy of the worker's block table stays a clamp because KV-block rounding can make the worker table wider while no local sequence references those columns.
_reuse_consumed_query_for_context_output allocates fresh storage when the query holds fewer bytes than the compact output (fp8 query: 192 B per head row vs 256 B bf16 output) instead of raising; the test now uses the production 192-wide query.

Not applied

"Align the B12X API contract" / "remove unsupported dense-MLA keywords": the review compared against the public bhaktatejas922/b12x package; the b12x tree this lineage runs with (local-inference-lab/b12x, shipped in the production image) has uses_query_cache_seqlens, the sparse_* fields and mode="verify".
"Fail fast on a wider block table": a width mismatch alone is legitimate (see above); the hazardous case is a plan that cannot cover the shard, which is now rejected at construction.

Tests: tests/v1/attention/test_b12x_mla.py 38 passed, tests/models/kimi_k3/test_mla_padding.py 14 passed (SM120 image).

voipmonitor and others added 30 commits August 12, 2026 13:46

fix(cache): preserve Mamba CoW after external hits

9972d3c

test(cache): assert external-hit continuation precondition

990a63d

Ignore Kimi tool calls after tools section close

3ebc896

fix(dspark): preserve draft graph capture contract

c805ebd

spec_decode: preserve DFlash-family checkpoint semantics

3938209

Co-authored-by: Codex <codex@openai.com>

attention: execute replicated KV groups without DCP partitioning

8f7f6c6

Co-authored-by: Codex <codex@openai.com>

models: document DFlash rotary layout detection

f85d3d8

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>

Fix invalid block handling for hybrid KV cache groups

c09bed0

Signed-off-by: jungjiyu <libraryofjiyu@gmail.com> Assisted-by: ChatGPT

test(scheduler): preserve hybrid KV recovery edge cases

3c296be

fix(structured-output): map compacted DSpark grammar rows

a18a86e

test(dspark): define graph-tail mode in capture fixture

e7ffc38

fix(mamba): resume prefix hits on checkpoint grid

edb1042

Signed-off-by: Martin Vit <martin@voipmonitor.org>

fix(kimi-k3): preserve final AttnRes block

d19ef45

Signed-off-by: Martin Vit <martin@voipmonitor.org>

voipmonitor and others added 19 commits August 22, 2026 16:00

Merge vLLM PR local-inference-lab#415

bc3dd5a

Merge vLLM PR local-inference-lab#418

afcbf4f

Merge vLLM PR local-inference-lab#419

23e7093

Merge vLLM PR local-inference-lab#459

f52fadc

Merge vLLM PR local-inference-lab#460

94d3833

Merge vLLM PR local-inference-lab#463

dc8f46d

Merge vLLM PR local-inference-lab#464

3c9031e

Merge vLLM PR local-inference-lab#467

7da8bd3

Merge vLLM PR local-inference-lab#468

6ba940c

Merge vLLM PR local-inference-lab#469

89b9db6

Merge upstream vLLM PR vllm-project#50729 backport

1c93ca0

fix(spec_decode): allow scheduler-selected zero draft depth

23f9f27

Signed-off-by: myshytf <9619163+myshytf@users.noreply.github.com>

test(cache): cover DFlash target-cache policy

dbf86de

Signed-off-by: myshytf <9619163+myshytf@users.noreply.github.com>

fix: address CodeRabbit review — document lookahead reservation is ha…

27e71ca

…rmless for single-token requests

fix(comm): prewarm FlashInfer PCIe graph shapes

74ec7af

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>

myshytf requested a review from mgoin as a code owner September 1, 2026 19:15

myshytf mentioned this pull request Sep 2, 2026

perf(mla): launch one dense MLA split per live chunk #587

Open

brandonmmusic-max mentioned this pull request Sep 2, 2026

Support: Jovian r10 TP4 DCP4 MTP3 research image v1 brandonmmusic-max/glm-5.3-exl3-k3-dcp4#1

Open

myshytf mentioned this pull request Sep 7, 2026

perf(mla): stabilize Kimi-K3 long-context verification myshytf/vllm#3

Open

Navigation Menu

Uh oh!

perf(mla): fuse K3 DCP verification queries - #565

perf(mla): fuse K3 DCP verification queries#565
myshytf wants to merge 60 commits into
local-inference-lab:dev/infernal-invocationfrom
myshytf:agent/kimi-k3-fused-dcp-verify-20260901-pr

myshytf commented Sep 1, 2026

Uh oh!

coderabbitai Bot commented Sep 1, 2026 •

edited

Loading

Review limit reached

Uh oh!

myshytf commented Sep 2, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

Uh oh!

Conversation

myshytf commented Sep 1, 2026

Summary

Uh oh!

coderabbitai Bot commented Sep 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Review limit reached

Uh oh!

myshytf commented Sep 2, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

coderabbitai Bot commented Sep 1, 2026 •

edited

Loading