Conversation
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>
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c7fd9c8d20
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| device, | ||
| page_size=self.page_size, | ||
| num_q_heads=self._kernel_heads, | ||
| max_total_q=batch * 4, |
There was a problem hiding this comment.
With FP8 KV and max_num_seqs > 256, even a non-speculative configuration fails during metadata-builder initialization: this loop eventually requests max_total_q=1028, while _create_dense_mla_plan rejects anything above _MAX_B12X_QUERY_ROWS (1024). Such configurations otherwise pass supports_combination, which explicitly permits up to 1024 sequences, and previously initialized successfully. Only create Q4 plans for batches whose total rows fit the limit—and ideally only when four-token verification is enabled.
Useful? React with 👍 / 👎.
| or parallel_config.prefill_context_parallel_size > 1 | ||
| ): | ||
| return False | ||
| layer_spec = envs.VLLM_K3_DCP_Q_REPLICATE_LAYERS |
There was a problem hiding this comment.
When query replication is enabled for an allowlisted non-causal draft layer, this predicate enables DCPGroupColumnParallelLinear even though get_kv_cache_spec marks that layer's cache dcp_replicated by default. Its metadata consequently reports one DCP KV shard, but the projection still emits the full DCP group's heads; B12xMLAImpl.forward_mqa then expects only self.num_heads and rejects the tensor. This makes configurations such as an all allowlist fail on the draft layer unless draft-cache sharding is separately enabled, so q-rep must account for whether the layer's KV cache is actually DCP-sharded (or slice the decode query for replicated caches).
Useful? React with 👍 / 👎.
1 participant
Summary
1,2,4,8,16,28and select the smallest covering plan at metadata-build timeK=3verification as one four-query request, with a bounded request page table and per-query DCP-local causal visibility, instead of flattening it into four independent decode rowsDependencies and default behavior
myshytf/vllm#1so the review diff is against the exact production integration snapshot rather than 59 unrelated composition commitslocal-inference-lab/b12x#271for fused verification and per-query visibilityVLLM_DCP_Q_REPLICATE=0andVLLM_K3_DYNAMIC_SPARSE_STRIDE=1Full Kimi-K3 q-rep duplicates query-projection and absorbed-weight storage. The layer allowlist is mandatory when q-rep is enabled; an explicit
allis required to accept the full persistent VRAM cost.Validation
The matching B12X SM120 suite passed 20 tests, including CUDA Graph replay and FP32-oracle coverage for multi-request four-query verification and DCP-visible causal prefixes.
The exact production-generation backport loaded TP8/DCP8 Kimi-K3 with FP8 KV and a 1,048,576-token limit, logged all decode and verify plans, captured all PIECEWISE/FULL graphs without eager fallback, passed API and LMCache health checks, and served text canaries. The operator's long-context benchmark reported that decode speed no longer degraded with context length.
No model-quality claim is made for dynamic sparsity. It remains disabled and requires a separate model evaluation before any stride above one can be enabled. Query replication is arithmetic-equivalent but also remains disabled in the production profile pending its VRAM/throughput A/B.
Duplicate check
Searches in
vllm-project/vllmandlocal-inference-lab/vllmfound no open PR for fused four-query B12X MLA verification or the shape plan bank.vllm-project/vllm#45964supplies the generic DCP query-replication primitive reused herevllm-project/vllm#54494wires that primitive into the ROCm Kimi-K3 wrapper; this PR covers the separate NVIDIA standalone Kimi-K3 MLA/B12X path and adds a per-layer VRAM safety policyReview status
AI assistance from OpenAI Codex was used to prepare this change. This PR is intentionally a draft. The submitting human must review every changed line, rerun the relevant tests and final benchmark, and understand and defend the implementation before marking it ready.
Local inference lab publication
This prepared implementation is available for review in local-inference-lab#565. That lab PR is the maintained review destination. This fork-only PR is a historical record; its non-draft status does not change validation limitations or authorize deployment.
🤖 Generated with Claude Code
https://claude.ai/code/session_01KxvNwugeU8RJFd7WRYwNLG