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>
Describe TP-owned KV heads retaining window positions under target DCP and fixed-address context buffers captured by supported draft configurations. Co-authored-by: Codex <codex@openai.com> Signed-off-by: Codex <codex@openai.com>
The draft receives concatenated target features before its row-sharded projection slices the input. Check the global source width for TP9 while preserving local matrix dimensions. Exercise combine_hidden_states for batched and single-token inputs and reject incompatible feature counts. Validation: both TP9 projection tests pass through the model combination method, including all nine aux extents. Co-authored-by: Codex <codex@openai.com> Signed-off-by: Codex <codex@openai.com>
`_B12X_DCP_WORLD_SIZES` gains 9 so the decode query-head gather, LSE reduce-scatter and Kimi projection pair gather use the B12X PCIe IPC pool at DCP9 instead of the NCCL fallbacks. The B12X kernel accepts nine ranks (`b12x.comm.pcie.pcie_dcp_a2a.SUPPORTED_WORLD_SIZES`) and the Kimi-K3 TP9 head padding (99 MLA heads, eleven per rank) satisfies its divisibility contract; pool initialization still falls back to NCCL on any rank failure. Scope: with VLLM_USE_B12X_DCP_A2A=1 every MLA layer's DCP exchange and every MoE layer's router/latent gather at DCP9 leave the NCCL path. The fused Kimi top-k selection stays unavailable at nine ranks and uses the ordinary paired gather plus grouped top-k. Validation: tests/distributed/test_dcp_a2a.py::test_b12x_dcp_world_size_gate covers 8 and 9 (pool reached) and 6 (declined); it needs a CUDA device and has not run yet. Nine-GPU serving effect is unmeasured. Co-Authored-By: Claude Code <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011aPCVZBsteYgs4PTAQYuE5 Signed-off-by: g0san <hhg698pknv@privaterelay.appleid.com>
`VLLM_PCIE_TWOSHOT_ALLREDUCE_MODE` (`pull`, the default, or `push`) is applied to the b12x `PCIeTwoShotBF16` runtime before its graph launchers are prepared, and the two-shot window log line reports the mode. A b12x runtime without `all_reduce_mode` rejects a non-default request instead of silently using the pull kernel. Scope: with `VLLM_PCIE_ONESHOT_ALLREDUCE_MAX_SIZE` and the fused add-RMSNorm limit lowered below the decode payload sizes, every decode all-reduce takes the two-shot push kernel; the fused add-RMSNorm call falls back to its unfused form above that limit. Validation: syntax and import structure checked; serving effect at TP9 is measured by the TP9 qualification boot. Co-Authored-By: Claude Code <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011aPCVZBsteYgs4PTAQYuE5 Signed-off-by: g0san <hhg698pknv@privaterelay.appleid.com>
…ment `try_dcp_b12x_all_gather_pair_kimi_topk` returns None when the per-rank latent or router width does not form whole 16-byte packs, which is the case at nine ranks (398 bf16 and 99 fp32 elements). The B12X pool raised `paired row widths must be multiples of 16 bytes` from the kernel warmup and aborted the boot; callers now take the ordinary paired gather and router as they do for other unsupported shapes. Validation: the TP9 boot that hit the warmup failure is rerun by the qualification script; the eight-rank path is unchanged. Co-Authored-By: Claude Code <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011aPCVZBsteYgs4PTAQYuE5 Signed-off-by: g0san <hhg698pknv@privaterelay.appleid.com>
`kimi_projection_shard_width` rounds each rank's slice of a padded Kimi column-parallel projection up to a multiple of eight elements, and the B12X projection-gather warmup builds its sample rows with the same rule. The B12X paired gather moves 16-byte packs, so at nine ranks the latent (3584 -> 400 per rank) and router (896 -> 104 per rank) projections now use the PCIe IPC pool instead of two NCCL all-gathers per MoE layer; the gathered tensors are still sliced to their logical widths. The fused gather-plus-top-k variant keeps its exact-width contract and stays on the ordinary paired gather at nine ranks. Compatibility: TP8 shards (448 and 112) are already aligned, so weights, gathers and outputs are unchanged there. mypy findings at model.py lines 2303-2350 predate this change and are unrelated. Validation: tests/models/kimi_k3/test_tp_projection.py (18 passed, CPU) covers the width rule at TP8 and TP9; the nine-rank serving effect is measured by the TP9 qualification boot. Co-Authored-By: Claude Code <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011aPCVZBsteYgs4PTAQYuE5 Signed-off-by: g0san <hhg698pknv@privaterelay.appleid.com>
`_try_b12x_dcp_all_gather_heads` gathers into a packed intermediate and copies into `out` when the caller-owned output is not contiguous. The B12X gather kernel writes a packed [batch, heads, head_dim] result, and the Kimi-K3 dense MLA decode path at nine DCP ranks hands it a 99-head view of its 104-head query tile, which the kernel rejected with `output must be contiguous` during CUDA graph capture. Compatibility: contiguous outputs (eight DCP ranks: 96 heads in a 96-head tile) keep the direct in-place gather. Validation: exercised by the TP9 qualification boot; no unit test yet because the path needs a CUDA device and a B12X pool. Co-Authored-By: Claude Code <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011aPCVZBsteYgs4PTAQYuE5 Signed-off-by: g0san <hhg698pknv@privaterelay.appleid.com>
`KimiK3PrefillProjectionWorkspace.get` returns None, with a one-time warning, when a context chunk has more rows than the retained buffer, so `project_context` falls back to the allocating projection instead of raising. The buffer is reserved when weights are loaded with the cache block size known then; the attention backend aligns its chunked-context workspace to the final block size and DCP geometry, which at DCP9 rounds `VLLM_MLA_INTERNAL_CONTEXT_WORKSPACE_SIZE=24576` up to 27,648 rows while the buffer holds 24,624. Four concurrent 60k-token prefills killed the TP9 server on that mismatch (`context projection needs 27648 rows, but the retained workspace has 24624`). Compatibility: TP8/DCP8 sizes coincide (24,576 both ways), so the retained path is unchanged there. The `import re` lint finding at line 33 predates this change. Validation: the TP9 serve script also sets the workspace to 27,648 so the retained path stays in use; the fallback is exercised by the TP9 qualification runs. Co-Authored-By: Claude Code <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011aPCVZBsteYgs4PTAQYuE5 Signed-off-by: g0san <hhg698pknv@privaterelay.appleid.com>
The hybrid coordinator enables fine-grained prefix hits whenever an aligned Mamba group exists and lifts the hit alignment to the least common multiple of the hash unit and every aligned Mamba block size. The full-attention manager probes interior boundaries at hash granularity and reports the longest boundary that is a multiple of that alignment. With a 1,536-token hash unit and a 1,536-token Mamba block (TP8) this is the previous behaviour; with a 4,608-token Mamba block (TP9/DCP9) local hits now land on 4,608-token boundaries instead of being disabled, which had forced every repeated prompt shorter than the 13,824-token scheduler block through an external LMCache retrieve (0 local hits, 100% external hits on the TP9 validation server). Validation: tests/v1/core/test_prefix_caching.py, test_mamba_align_chunk_split.py and prefix_cache/ test_partial_prefix_cache_primitives.py: 123 passed (CPU). The TP9 serving effect (local hit on the second 5,048-token probe) is checked by the qualification run. Co-Authored-By: Claude Code <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011aPCVZBsteYgs4PTAQYuE5 Signed-off-by: g0san <hhg698pknv@privaterelay.appleid.com>
Each single-type manager caches its reachable tails at `cache_alignment_tokens`, which the hybrid coordinator sets to its partial-hit alignment (the hash unit lifted to the Mamba block cadence). Before, tails were aligned to the scheduler block, so a replicated sliding-window draft group with 1,536-token blocks under a 13,824-token scheduler block cached only the last blocks of each scheduler segment and no lookup shorter than a scheduler block could hit in that group, which made the hybrid hit zero even though the attention and Mamba groups had cached the prefix. Validation: tests/v1/core/test_prefix_caching.py gains test_hybrid_partial_hits_align_to_mamba_block_under_dcp (16-token hash unit, 48-token aligned Mamba block, 64x3-token DCP attention block, replicated 16-token sliding-window group: a 96-token replay hits 96 in all three groups); test_prefix_caching.py, test_mamba_align_chunk_split.py and prefix_cache/test_partial_prefix_cache_primitives.py pass (CPU). The test_kv_cache_utils.py cases that download Hugging Face configs are not runnable offline and were not evaluated. Co-Authored-By: Claude Code <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011aPCVZBsteYgs4PTAQYuE5 Signed-off-by: g0san <hhg698pknv@privaterelay.appleid.com>
The varlen tile scheduler sizes the grid from total_q + num_batch * (tile_m - 1) query rows, so a batch whose sequence lengths do not fill whole query tiles (a padded zero-length request, a prefix-hit tail next to a fresh request, two chunked requests with tile remainders) gets extra CTAs whose work tile maps to batch index num_batch with is_valid_tile False. The SM80/SM120 kernel derived those CTAs' sequence lengths from cu_seqlens[num_batch + 1], one element past the tensor, and then loaded K/V tiles and stored O rows for the garbage length: out-of-bounds reads past K/V (a CUDA illegal address when the offset leaves mapped memory) and writes past the end of the attention output. The SM100 kernel already skips invalid tiles; this kernel now does the same before touching global memory. Valid tiles compute exactly as before. Scope: vllm/vllm_flash_attn/cute/flash_fwd.py is the copy of vllm-project/flash-attention 7484d475 flash_attn/cute/flash_fwd.py that the vLLM build installs (untracked in this tree, force-added as an override for the serving overlay); the same defect exists upstream. Kimi-K3 dense MLA prefill on SM120 (VLLM_MLA_SM120_FA4_PREFILL=1) runs this kernel for every prefill batch, at TP8 and TP9. Validation: tests/kernels/attention/test_fa4_sm120_varlen_padding.py plants 2^30 after cu_seqlens and a NaN canary after out for batches [3615, 0], [1055, 744], [1057, 111], [744, 1055] (11 heads, qk 192, v 128, bf16, causal): 4 passed with this kernel; the image copy fails all 4 by writing past the output. TP9/DCP9 serving reproduced the fault deterministically (GPU core dump: FlashAttentionForwardSm120 grid 330, faulting CTAs 319-329, LDGSTS at the K/V load) with a 12,831-token prompt replayed twice from a 9,216-token prefix hit next to a 927-token request. Co-Authored-By: Claude Code <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011aPCVZBsteYgs4PTAQYuE5 Signed-off-by: g0san <hhg698pknv@privaterelay.appleid.com>
…ranks
VLLM_PCIE_TWOSHOT_ALLREDUCE_MODE accepts island9 in addition to pull and
push. At nine ranks it constructs b12x's PCIeIsland9AllReduce (ranks 0-3 and
4-7 reduce quarters inside their PCIe switch clusters, rank 8 on the CPU root
contributes through island 0, fp32 island partials, one bf16 rounding) for
the payload window above the one-shot ceiling and below the DMA ring, with
the same capacity vocabulary (row_elems x max_rows) and dispatch as the
two-shot. Other world sizes reject the mode at initialization.
Validation: the runtime is exercised by b12x tests/comm/test_pcie_tp9_physical.py
and benchmarks/{benchmark,precision}_pcie_tp9_allreduce.py on nine idle
GPUs (pending the next maintenance window). Serving selection follows the
measured latency and the precision comparison against the two-shot.
Co-Authored-By: Claude Code <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011aPCVZBsteYgs4PTAQYuE5
Signed-off-by: g0san <hhg698pknv@privaterelay.appleid.com>
…se format
The kquant serialized-MXFP8 dense format (dense_format "mxfp8" with the
checkpoint's ignored_layers) describes the language model only. The vision
tower and multimodal projector are stored in bf16 without MXFP8 scales, so
selecting the quantized linear method for them made the loader fail with
"serialized quantization parameters were not initialized from checkpoint:
mm_projector.linear_1.weight_scale, vision_tower.encoder.blocks.0.mlp.fc0
.weight_scale, ...". The multimodal wrapper now drops the quantization
config for those modules when the dense format is mxfp8, as it already does
for compressed-tensors configs; the language model is unchanged.
Validation: Kimi-K3 TP9/DCP9 boot with the image encoder enabled
(--limit-mm-per-prompt {"image":5}) on the QSRT-K2 checkpoint, whose
safetensors index carries 168 vision keys and no scale tensor; before the
change the boot failed at weight loading, after it the encoder loads and the
image probe (research/tp9-colocated-qsrt-20260905/vision_probe.py) runs
(result recorded in the maintenance window evidence).
Co-Authored-By: Claude Code <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011aPCVZBsteYgs4PTAQYuE5
Signed-off-by: g0san <hhg698pknv@privaterelay.appleid.com>
When the vision tower's head count does not divide the tensor-parallel size, vLLM replicates the tower on every rank and shards a batch by image, so a request with one large image runs the whole encoder on a single rank while the others wait: on the Kimi-K3 TP9 target a 40,000-patch image costs 1.09 s of encoder time on rank 0, 0.90 s of it in the 27 self-attention kernels (FlashAttention 2, 12 heads x 128). vllm/model_executor/models/vit_query_shard.py plans a query split of that replicated forward: patch embedding, norms, projections and MLPs stay replicated (every rank computes the identical full tensors) and each rank runs the varlen attention for its share of every image's 128-row query tiles against the image's complete keys and values; the shares are all-gathered in the TP group and restored to packed row order (one query segment per image, empty where a rank has no rows, so segment b pairs with key segment b). FlashAttention computes every query row independently over the key blocks in a fixed order, so the gathered output equals the unsharded computation bit for bit; only the row-to-GPU assignment changes. MoonViT's encoder layer takes the plan when it is active and the FLASH_ATTN backend is in use; vision_tower_forward activates it on the data-parallel path for batches of at least VLLM_VIT_QUERY_SHARD_MIN_PATCHES patches (default 8192; below that the gather cost exceeds the attention saved) unless VLLM_VIT_QUERY_SHARD=0. MMEncoderAttention exposes its FlashAttention version as fa_version. Every rank now materializes the full encoder activations of a batch (previously only the rank that owned the image), which raises the other ranks' transient to the owning rank's level. Validation: tests/models/multimodal/test_vit_query_shard.py (29 cases: tile-aligned single coverage of every row, balanced tiles, gather order for one and several images, ranks without rows) passes in the serving image; a nine-rank CPU run with stand-in attention and gathers matches the unsharded path to bf16 precision for one and three images (research/tp9-colocated-qsrt-20260905 scratch check); the nine-GPU bit-identity and timing check is research/tp9-colocated-qsrt-20260905/ vit_query_shard_check.py, run at the maintenance window before the option is served. Co-Authored-By: Claude Code <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011aPCVZBsteYgs4PTAQYuE5 Signed-off-by: g0san <hhg698pknv@privaterelay.appleid.com>
…dings
The QSRT trellis prefill path binds the caller-owned scratch plan with
empty launch slots, so b12x's run_w4a16_moe resolves the fused MoE launch
on every layer call through compile_w4a16_fused_moe: tile selection,
W4A16FusedMoeKernel construction (two W4A16GemmKernel objects) and a
cache-key lookup, about 0.10 ms per call on an idle core and 4.5x that
under the with_stack profiler that produced the tail-chunk traces. The
compiled kernel is cached under a key that does not contain the row
count (only M == 1 and the small-M split-K band are specialized), so the
launch the plan compiled for the scheduler capacity at boot is the same
object every prefill row count resolves to.
_ensure_runtime now takes the prefill plan's prewarmed capacity fused
launch and the int32/mapped top-k sum launch once per layer state, and
_apply_once binds them into the binding (dataclasses.replace) for every
launch above the prefill route threshold. run_w4a16_moe then skips its
own resolution and validates the launch against the live contract, as
it does for the kept tier's preplanned launches. The decode plan keeps
lazy resolution, where the direct-route selection is dynamic.
Numerics: bit-identical. The kernel binary, the live row count, the
persistent grid (route-pack path: sms x blocks_per_sm) and the buffers
are unchanged; only the host-side derivation of the launch object is
skipped. VLLM_KQUANT_W4A16_PREPLANNED_LAUNCH=0 restores the lazy path
for an A/B comparison. Plans without prewarmed launches keep the lazy
path with a one-time warning.
Validation: CPU tests in the serving image (tests/quantization/
test_kquant_hybrid.py: launch selection, fallback, switch, and the
b12x cache key being identical for M in {2, 257, 830, 1536, 4608} and
different only for M = 1). GPU: logits digest with the switch on/off at
the next maintenance window.
Co-Authored-By: Claude Code <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011aPCVZBsteYgs4PTAQYuE5
Signed-off-by: g0san <hhg698pknv@privaterelay.appleid.com>
(cherry picked from commit 5b89012)
KimiPaddedRowParallelLinear splits its input axis into per-rank shards of kimi_projection_shard_width columns (a multiple of eight elements; 400 for the 3,584-wide routed latent at TP9) instead of the ceiling division (399), and its prefill entry points forward_into / accumulate_into read the rank's K window as a strided view of the full-width latent, skipping the last rank's zero-filled weight tail (384 real columns at TP9) rather than padding the input. Reason: with 399-column shards the bf16 operands of the 4,608-token prefill up-projection (K=399, N=7,168 per rank) have a 798-byte row pitch and, on odd ranks, 2-byte-aligned window starts, so cuBLAS falls back to an align-1 SM75 tensor-op kernel: 202 us and 130 TFLOP/s per MoE layer on the served TP9 stack (rank-0 trace, host frame model.py forward_into), against 380-450 TFLOP/s for the other dense projections of the same chunk. 400-column shards give every rank a 16-byte-aligned window start and row pitch and a K that is a multiple of eight, which admits the vectorized SM120 kernels. The strided view also removes the per-layer pad and shard-copy kernels (three elementwise launches, ~39 us) that materialized the contiguous shard. Compatibility: TP sizes that divide the latent width into 8-aligned shards (TP8: 448 columns) keep their partition, operands and results; their only change is the copy-free view. Decode keeps forward(), which zero-pads the input to the same 400-column split. Weight loading needs no change: the loader already zero-fills the last rank's tail from the parameter shape. Numerics: not bit-identical where the shard width changes (TP9). Each rank's fp32 partial sum covers a different K range (400-column instead of 399-column blocks) and a different kernel may order the reduction differently, so the bf16 partials and the all-reduced sum can differ in the last bit. Same precision class (bf16 operands, fp32 accumulation, one bf16 rounding per rank); deployable only after the window measurement against an fp64 reference shows equal or lower error (research/prefill-campaign-20260906/evidence/item9/up_proj_alignment_probe.py). Validation: CPU tests in the serving image (tests/models/kimi_k3/ test_mlp_tensor_lifetime.py: 16-byte-aligned copy-free windows on all nine ranks, nine partials summing to the full projection with a zero-filled last-rank tail, accumulate_into adding the same partial); existing routed-output-transform and loader tests unchanged. The mypy hook is skipped: it reports five pre-existing union-attr errors in the unrelated auxiliary-stream code of this file on the base revision. Co-Authored-By: Claude Code <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011aPCVZBsteYgs4PTAQYuE5 Signed-off-by: g0san <hhg698pknv@privaterelay.appleid.com> (cherry picked from commit d8c5e5d)
VLLM_K3_DUMP_TOPK_IDS=<dir> makes the hybrid QSRT expert path (TP rank 0
only) save the routed expert ids of the first VLLM_K3_DUMP_TOPK_IDS_CHUNKS
(default 16) prefill launches of every MoE layer as
<dir>/layer{L}_chunk{c}_m{tokens}.pt holding {"topk_ids": int32
[tokens, top_k], "layer": int, "num_tokens": int}. Decode launches
(m <= 8) are not recorded. The variable is unset by default and the hook
then costs one environment lookup per layer call.
Reason: the W4A16 prefill route-block choice (48 vs 64 rows) depends on
the per-expert row histogram of real traffic, which synthetic uniform
routing cannot reproduce; the b12x extent harness replays these files
(--topk-ids) to measure the block-size trade-off on the served
distribution.
Compatibility: no behavior change while the variable is unset. When set,
the device-to-host copy synchronizes the stream once per recorded
launch; diagnostic runs only.
Validation: CPU test in the serving image (tests/quantization/
test_kquant_hybrid.py::test_topk_id_dump_writes_prefill_chunks_on_rank_zero:
file names, payload dtype/shape, chunk cap, rank filter).
Co-Authored-By: Claude Code <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011aPCVZBsteYgs4PTAQYuE5
Signed-off-by: g0san <hhg698pknv@privaterelay.appleid.com>
(cherry picked from commit f17a7bf)
With VLLM_K3_RING_STATIC_IO=1 (default off) the in-place prefill all-reduces of a Kimi-K3 layer return the B12X DMA ring's static output instead of a copy: the attention-output reduction (tp_projection.reduce_kimi_full_width_projection) and the MoE runner's two in-place reductions (routed latent before the output transform, final output) pass borrow_output=True through tensor_model_parallel_all_reduce_in_place -> GroupCoordinator -> CudaCommunicator -> CustomAllreduce -> PCIeDmaAllReduce.all_reduce. The ring's lossless replay entry is in place (one buffer is static input and output), and the layer already chains one buffer through pre-attention AttnRes output -> attention caller output -> o_proj out= -> in-place reduction -> post-attention AttnRes output -> MoE up_proj out= -> in-place reduction, so once a reduction is borrowed every later reduction of that shape reads its input from the static buffer and returns it: the two 66 MB device copies per replayed reduction (about 18 ms per 4,608-token chunk on rank 0) disappear. A borrowed tensor is overwritten by the next reduction of the same numel/dtype, so the two places that retain a reduction longer copy it out: the AttnRes block-write layers, whose attention-output reduction becomes the retained prefix (materialize_kimi_reduction), and the model output, where the final AttnRes no longer reuses a borrowed last-layer MoE output as its storage. Every other consumer (post-attention AttnRes, the MoE front end, the latent RMSNorm and up-projection, the next layer's pre-attention AttnRes, the auxiliary-state capture) reads the tensor before that reduction. The keyword is forwarded only when set, so communicators and call sites without the gate keep their exact calls. Compatibility: no behavior change with the gate off; other models and the functional all-reduce never borrow. Numerics: bit-identical (the ring replays the same kernels; only staging copies are removed). Validation: CPU tests in the serving image with --noconftest (tests/models/kimi_k3/test_ring_static_io.py: the keyword reaches the ring and the runner only when enabled, materialization copies only borrowed storage, block-write layers copy the prefix while other layers consume the reduction in place; tests/models/kimi_k3/test_tp_projection.py and test_mlp_tensor_lifetime.py unchanged). tests/distributed/ test_inplace_allreduce.py's two CudaCommunicator cases fail on the base revision as well (they predate the ca_comm dispatch). The mypy hook is skipped: it reports five pre-existing union-attr errors in the unrelated auxiliary-stream code of model.py. GPU: same-boot greedy fingerprint with the gate on and off at the next maintenance window. Co-Authored-By: Claude Code <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011aPCVZBsteYgs4PTAQYuE5 Signed-off-by: g0san <hhg698pknv@privaterelay.appleid.com> (cherry picked from commit 64017bc)
With VLLM_K3_PROJECTION_GATHER=dma_pair (default nccl) a prefill MoE layer (>= 1,024 tokens, TP-sharded auxiliary projections, latent MoE, no fused tail) gathers its fp32 router-logit shard ([rows, 104] at TP9) and bf16 routed-latent shard ([rows, 400]) with one PCIeDmaAllReduce. all_gather_pair issued on the ring's side stream (CustomAllreduce.pcie_dma_all_gather_pair, reached through tensor_model_parallel_pcie_all_gather_pair -> GroupCoordinator -> CudaCommunicator). The copy engine carries both blocks as one 8/9-loaded ring pass while the main stream runs the shared experts, using the same layer call the runner would make (writing into the consumed input when the runner would); the runner then adopts that result (SharedExperts.install_precomputed_output, MoERunner.forward( shared_output=...)), waits for the gather, and assembles the logical [rows, 896] and [rows, 3584] tensors from the rank-major blocks in one pass each. The served path serializes two PyNCCL all-gathers (2.2 ms per layer at 4,608 tokens, SM-occupying) before the shared experts and top-k (0.85 ms); the ring pass takes about 1.6 ms and the shared experts hide under it. When the ring declines a call (decode-sized, graph capture, no ring) the layer falls back to the NCCL gathers of the same local shards. Compatibility: no behavior change with the gate off; decode keeps its paired B12X projection gather. Numerics: bit-identical (an all-gather only copies; the GEMMs and the shared experts see identical operands, only the issue order changes). Validation: CPU tests in the serving image with --noconftest (tests/models/kimi_k3/test_projection_gather_pair.py: block assembly equals concatenation byte for byte for the 104/896 fp32 and 400/3584 bf16 shapes, the gate and the decode-size fallback, the pending gather's wait assembles the ring outputs, the NCCL fallback, the shared experts adopt a precomputed output and require aliasing under input reuse, and KimiMoE.forward issues the gather, runs the shared experts, then waits; test_ring_static_io.py, test_tp_projection.py and test_mlp_tensor_lifetime.py unchanged). The mypy hook is skipped for the five pre-existing union-attr errors in model.py. GPU: nine-rank physical gather test, graph-replay mix with the ring all-reduce, and the same-boot greedy fingerprint at the next maintenance window. Co-Authored-By: Claude Code <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011aPCVZBsteYgs4PTAQYuE5 Signed-off-by: g0san <hhg698pknv@privaterelay.appleid.com> (cherry picked from commit 3df9d9a)
…on shard With VLLM_K3_LATENT_REDUCE=rs_fp32 or rs_bf16 (default allreduce) a prefill MoE layer (>= 1,024 tokens, TP-sharded up-projection, no latent capture) reduces its TP-partial routed latent ([rows, 3584] bf16 at TP9) with PCIeDmaAllReduce.reduce_scatter_columns instead of the 33 MB all-reduce: the MoE runner asks the output transform for reduce_scatter_tp_partial, KimiRoutedOutputTransform hands the partial to the ring with the up-projection's shard width (400 columns; the last rank owns 384) through tensor_model_parallel_pcie_reduce_scatter_columns -> GroupCoordinator -> CudaCommunicator -> CustomAllreduce, and the transform continues on the returned [rows, 400] column block. The RMSNorm variance comes from each rank's fp64 sum of squares of its block combined by an fp64 all-reduce ([rows, 1], 37 KB over PyNCCL), the block is scaled by rsqrt(variance + eps) and its slice of the norm weight in fp32 and rounded to bf16 once, and KimiPaddedRowParallelLinear consumes it as its input shard (local_block_shard: the last rank skips its 16 zero columns as it skips the zero-filled weight tail) through forward_local_block, forward_into(x_is_local_block=True) or accumulate_into(x_is_local_block=True). KimiMoE compiles the ring's reduce-scatter kernels at construction, ahead of the kernel-resolution freeze and graph capture. When the ring declines a call the layer keeps the all-reduce. Bytes per layer: 1.67 N link bytes (rs_fp32) or 0.89 N (rs_bf16) instead of the all-reduce's 1.78 N, and the row-parallel up-projection reads only its shard. Numerics: not bit-identical. rs_fp32 rounds each latent element once (the served ring rounds after each of eight hops), and the column-block RMSNorm uses fp64 partial sums and a single rounding where the full-width kernel uses an fp32 block reduction and two roundings; the up-projection GEMM reads a contiguous [rows, 400] block (row pitch 800 B) where the all-reduce path reads a strided window of the full-width latent, so the kernel selection may differ within the same precision class. rs_bf16 keeps eight roundings in column-block order (candidate only). Deployment follows the campaign rule: only after the precision harness (b12x benchmarks/precision_pcie_tp9_allreduce.py --rows 4608 --width 3584 --norm) and the fp64 up-projection probe show equal or lower error than the served path. Compatibility: no behavior change with the gate off; the runner's reduce hook returns a third value (column block) and passes column_block to the transform only when set, so other transforms keep their calls. Validation: CPU tests in the serving image with --noconftest (tests/models/kimi_k3/test_latent_reduce_scatter.py: the gate and the decode-size fallback, the transform hook's capture and reduce_results guards, nine emulated ranks' column-block RMSNorm within one bf16 ulp of the fp64 reference and with no more mismatches and no larger relative L2 error than the served arithmetic, the up-projection's block entry points producing the full-width window's projection on both ranks of a two-way split with a padded last shard, and the runner preferring the block and flagging it; the ring-side, gather-pair, in-place and tensor-lifetime tests unchanged, 57 passing). The mypy hook is skipped for the five pre-existing union-attr errors in model.py. GPU: the nine-rank physical reduce-scatter test, the precision harness, the up-projection probe and the same-boot fingerprint at the next maintenance window. Co-Authored-By: Claude Code <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011aPCVZBsteYgs4PTAQYuE5 Signed-off-by: g0san <hhg698pknv@privaterelay.appleid.com> (cherry picked from commit 42d24bd)
VLLM_K3_DECODE_SHARD_PACK=1 avoids allocating, zero-filling and copying the full routed latent on every TP rank before extracting a contiguous shard. Interior ranks copy only their own columns. The tail rank packs valid columns and zeros into the same contiguous padded K width in one Triton launch. The normal GEMM dispatcher, weights, dimensions and arithmetic remain unchanged. Prefill and unsupported shapes retain the existing path. The tail pack is warmed for decode row counts before graph capture. CUDA graph replay reuses captured addresses and needs no replay allocation. The pack uses 64-bit row addressing and preserves payload bits, including negative zero and NaN payloads. It never writes the source tensor. Validation on RTX PRO 6000 Max-Q and Workstation Edition, SM120: 37 tests passed per GPU across test_mlp_tensor_lifetime.py, test_vision_warmup.py and test_kimi_k3_triton_warmup.py. Actual projection outputs match bitwise at TP9 ranks 0-8, rows 1/2/4/8/9/16; interleaved CUDA graph diagnostics include changed inputs and stable output addresses. At rows=4 projection graph latency is 0.749-0.750x control on Max-Q and 0.601-0.800x on Workstation Edition. Whole-model validation is pending. Lint and pre-commit passed except five pre-existing optional-projector union-attr errors in model.py; the mypy-3.10 hook is skipped for that unchanged baseline defect. The added test's type annotation is fixed. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01KxvNwugeU8RJFd7WRYwNLG Signed-off-by: g0san <hhg698pknv@privaterelay.appleid.com>
Model runner v2 matched full speculative-decode graphs using only request and token counts. A prompt ending four tokens past a 4608-token boundary therefore replayed the verifier graph even though KDA metadata classified those four prompt tokens as prefill. The older runner checks prefill progress, but v2 did not. Derive has_prefill from scheduled requests' existing CPU prefill progress before graph selection. Exclude uniform and bounded-verifier FULL candidates when a request is prefilling, retaining piecewise, generic mixed FULL, and eager fallbacks. Propagate the flag through DP agreement so its second dispatch cannot promote a peer's prefill to a decode graph. True decode graph selection, model arithmetic, and precision are unchanged. Validation: 13 CPU graph-dispatch and dynamic-speculation tests pass with an inert graph-pool stub in the CUDA serving image. They cover an exact four-token tail, varlen decode candidates, mixed FULL compatibility, eager fallback and DP agreement. Before the fix, the serving target corrupts cold and cached prompts of 4612 and 9220 tokens; neighboring tails 1/2/3/5/8/9/16/53 return the exact expected JSON. The captured incident had 115204 = 25*4608+4 prompt tokens and repeated one sentence 21747 times. The same short-tail prompt succeeds during eager profiling. Whole-model post-fix qualification is pending. Pre-commit including mypy passes; the CUDA API policy hook is skipped only for an existing diagnostic synchronize call in model_runner.py, outside the changed path. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01KxvNwugeU8RJFd7WRYwNLG Signed-off-by: g0san <hhg698pknv@privaterelay.appleid.com>
|
Important Review skippedToo many files! This PR contains 136 files, which is 36 over the limit of 100. To get a review, reduce the PR to 100 files or fewer by splitting it into smaller PRs or changing its base branch. Upgrade to a paid plan to raise the limit. This review couldn't start because sufficient usage credits or metered capacity aren't available. Add credits or update usage-based reviews in the billing tab, then retry. ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Team Run ID: 📒 Files selected for processing (136)
You can disable this status message by setting the 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 |
added 2 commits
September 7, 2026 12:39
An opt-in test targets an explicitly supplied idle test server. It calibrates prompt lengths with the tokenizer endpoint, checks completion usage for the same count, and requires an exact JSON answer on cold and cached requests. It covers prefill chunk boundaries plus the observed 115204-token incident length. Requests use fresh salts and never execute tools or reset caches. Validation: the equivalent local serving harness reproduced failures only at four-token tails before the graph-dispatch fix. After the fix all 36 short-tail checks and six 115200-token-boundary neighbor checks pass, including cold and cached prompts of exactly 115204 tokens. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01KxvNwugeU8RJFd7WRYwNLG Signed-off-by: g0san <hhg698pknv@privaterelay.appleid.com>
Execute the real runner until graph dispatch with minimal request state. Assert that four-token tails at 4612 and 115204 tokens are marked prefill, completed prompts remain decode, reordered mixed batches preserve the guard, inactive slots cannot trigger it, and dummy capture keeps prior behavior. Validation: 19 graph-manager, runner-entry and dynamic-speculation tests pass in the serving image with an inert CPU graph-pool handle. The tests introduce no runtime or serving source change. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01KxvNwugeU8RJFd7WRYwNLG Signed-off-by: g0san <hhg698pknv@privaterelay.appleid.com>
myshytf
marked this pull request as ready for review
September 7, 2026 07:08
This was referenced Sep 7, 2026
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
8 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.
Prefill-aware graph dispatch
Status: qualified and serving. The reproduced boundary failure, broader serving regressions and final paired serving benchmark have completed. This lab PR preserves an operator-authorized diagnosis and fix. Human review is not claimed.
Model runner v2 used token counts alone to select a full speculative-decode graph. A prefill tail with exactly the verifier's query length was therefore eligible for decode graph replay even though its KDA metadata described a prefill. The graph contains capture-time decode kernels and cannot substitute for the prefill path.
The failure reproduces on Kimi-K3 TP9/DCP9 with colocated DFlash, speculative depth 3, 4,608-token prefill chunks, FULL_AND_PIECEWISE graphs, and the v2 runner:
The affected user request had 115,204 input tokens, exactly
25 * 4608 + 4. Generation was already incoherent at its first tokens, then repeated one sentence until the output limit. An eager-profiled request with the failing 4,612-token shape returns the expected JSON. No CUDA Xid, finite-output alarm, or HTTP error accompanied the bad generation.Change
Commit
d4c3c7578fderiveshas_prefillfrom the already-available CPU prefill progress of scheduled requests. The graph manager excludes uniform and bounded-verifier FULL candidates when that flag is set. It retains piecewise graphs, generic FULL graphs captured for mixed batches, and eager fallback. Dummy capture and true decode retain their existing behavior.The prefill flag participates in data-parallel graph agreement so a second dispatch cannot promote another replica's prefill into a decode graph. No GPU synchronization or new device kernel is introduced. Model arithmetic, precision, and the true-decode graph remain unchanged.
This branch is stacked on the served integration and lab #686, which only packs the local padded projection input. The graph-classification fix is independent of that packing optimization.
Validation
The tests cover four-token prefill versus identical-shaped decode, varlen FULL candidates, generic mixed FULL compatibility, eager fallback, and data-parallel agreement. Pre-commit and mypy passed. The CUDA API hook was skipped only for an unchanged diagnostic
torch.cuda.synchronize()call outside the changed path.Post-fix serving at the unchanged nine-GPU configuration passes all 36 short-boundary checks plus six incident-length neighbor checks. In particular, 4,612, 9,220 and 115,204-token prompts now return the exact JSON on both cold and cached runs. Another 16 semantic checks pass across cold, GPU-prefix, external LMCache restore and concurrent requests. Existing functional, vision, cache-resume and 220k prefill regressions pass. No repetition penalty or output truncation guard is used as a substitute for the correction.
The fixed source boots in a fresh cache namespace. Historical entries are retained without being reused for qualification. Twelve further checks pass for exact four-token tails with reasoning enabled, external restoration and simultaneous requests. At the incident length, the external path restored 115,200 tokens and correctly processed the final four prompt tokens.
The final
llm_decode_bench.pycc=1 run over 0/8k/16k/32k/64k/128k completed at 2026-09-07 04:03 UTC. Compared with the same-day pre-change baseline, the combined fixed build with #686 gives decode steps/s +0.6–1.1%, and prefill -0.5–0.7%. No competing or queued requests were recorded. These are whole-build measurements, not isolated speed claims for this correctness fix. An opt-in serving regression is preserved intests/models/kimi_k3/test_prefill_tail_serving.py. Additional CPU tests execute the real v2 runner through graph dispatch to cover active-slot mapping, inactive slots, dummy capture and completed prompts.Duplicate-work check
Open-PR searches covered short-prefill graph dispatch in both lab and upstream repositories. Upstream vllm-project#51483 changes Kimi KDA metadata for stateless first chunks. It does not change model-runner-v2 graph selection, which happens before that metadata is consumed. The v1 runner already checks prefill progress, but the active v2 runner did not.
Written, investigated, tested and reviewed with Claude assistance. Private request contents are not included in this PR.
Publication status
Published as a non-draft on 2026-09-07 at the operator’s explicit request to expose all prepared PRs in local-inference-lab. Existing qualification evidence and limitations above are unchanged. No code, deployment configuration, merge approval or automatic merge is changed by this status update.
🤖 Generated with Claude Code
https://claude.ai/code/session_01KxvNwugeU8RJFd7WRYwNLG