[codex] Fix Kimi DFlash on fathomless - #86
Closed
voipmonitor wants to merge 44 commits into
Closed
Conversation
…decode stages on low-smem GPUs tl.arange requires power-of-2 extents, so cp_lse_ag_out_rs failed to compile for CP/DCP world sizes like 6 or 10; round up and mask with -inf. Also drop to num_stages=1 in the grouped MLA decode kernel on GPUs with <=100KB shared memory per block, where 576-dim MLA with two stages exceeds the limit. Co-Authored-By: Martin Vit <martin@voipmonitor.org>
merge_attn_states derives one source offset from the prefix strides but reads both inputs with it; padded-V attention outputs merged with contiguous chunk outputs read the suffix at wrong offsets once context spans more than one workspace chunk. Force matching strides at all three merge sites. Co-Authored-By: Martin Vit <martin@voipmonitor.org>
…tokens The TP-local argmax masked only the tail org-vocab padding region and mapped all winners through org_vocab_start_index, returning live padding logits and wrong global ids whenever the shard contains added vocab. Mask both padding regions from the shard indices and map added-vocab winners through added_vocab_start_index. Co-Authored-By: Martin Vit <martin@voipmonitor.org>
… sampling kernels Row offsets in the sampling and rejection Triton kernels overflowed int32 for spec-expanded batches times 128k vocab; placeholder (-1) draft tokens were dereferenced before rejection and sampled buffers were left uninitialized; Gumbel noise could produce +inf winners from u==1 and padded CUDA-graph replay rows could clobber cached draft logits. Also mark size-like kernel args do_not_specialize to stop recompile churn. Co-Authored-By: Martin Vit <martin@voipmonitor.org>
…multi-group configs Use the pooled per-group helpers for max-concurrency math, unify same-type page sizes by lcm instead of max, compute engine capacity from the worker view rather than the merged scheduler view, honor per-spec cache_dtype_str when deriving shapes, support per-page-size KV block zeroing, fix uniform-decode misclassification of prefill-shaped batches, and let backends with different declared layouts share one raw KV tensor via as_strided views. Co-Authored-By: Martin Vit <martin@voipmonitor.org>
…r stats every step Prompt throughput was derived from computed tokens at request-output time, so chunked prefills were only counted when a request produced output; track scheduled context tokens per step instead and record scheduler stats even on steps with no outputs. Co-Authored-By: Martin Vit <martin@voipmonitor.org>
Track accepted draft tokens over a configurable window (adaptive_speculative_tokens_window) and adjust the scheduled speculative depth: back off immediately toward mean acceptance + 1, explore upward one token per window, capped by num_speculative_tokens. Works with any model-backed speculative method on both runners; current depth is reported in spec-decode metrics. Co-Authored-By: Martin Vit <martin@voipmonitor.org>
…ading Models can declare checkpoint_weight_name_prefixes; the default loader then downloads and scans only safetensors shards containing matching tensors (via the index weight_map) and skips non-matching tensors inside mixed shards. Draft/MTP-only workers stop pulling the full target checkpoint. Co-Authored-By: Martin Vit <martin@voipmonitor.org>
… weights Dequantize block-FP8 (weight_scale_inv) and MXFP8 (uint8 weight_scale) checkpoint linears to BF16 when the runtime MTP module is unquantized, pad GLM fused kv_a rows to the CUTLASS block shape, remap scale names for compressed-tensors modules, and extend the DeepSeek V3.2 indexer WK loader to MXFP8. Co-Authored-By: Martin Vit <martin@voipmonitor.org>
Add a shared_experts spec to --quantization-config that routes shared-expert projections to online MXFP8, including as an overlay on ModelOpt checkpoints that exclude or omit those layers (e.g. NVFP4 checkpoints with BF16 shared experts). Overlay resolution is fail-closed: only the exact MXFP8-dynamic shared-experts shape is accepted alongside modelopt quantization. Co-Authored-By: Martin Vit <martin@voipmonitor.org>
…uant_algo resolution Support MXFP8 layer configs in ModelOpt mixed-precision checkpoints and fix quant_algo resolution for fused layers: fall back to a default packed-modules mapping when the fused name is missing, expand per-shard prefix candidates, and alias block_sparse_moe and mlp module paths in both directions. Co-Authored-By: Martin Vit <martin@voipmonitor.org>
Stop hardcoding the GPT-OSS clamped-SwiGLU parameters for every MXFP4/NVFP4 model: plumb gemm1_alpha/gemm1_beta through the quant configs, apply GPT-OSS defaults only for model_type gpt_oss, convert natively-loaded MXFP4 checkpoints from [w1,w3] to FlashInfer's [w3,w1] layout, and register moe_forward_shared with mutates_args to match its in-place implementation. Co-Authored-By: Martin Vit <martin@voipmonitor.org>
Zero-pad each gated w13 half (and the w2 K dimension) up to the 128-row scale-swizzle boundary so gated FlashInfer/CUTLASS NVFP4 MoE supports intermediate sizes that are not multiples of 128. Co-Authored-By: Martin Vit <martin@voipmonitor.org>
Detect whether the model actually uses FlashInfer attention or compute kernels before running the multi-minute autotune pass, log the skip reason, and stop every rank from writing the broadcast autotune cache file: only the non-leader local rank 0 persists it. Co-Authored-By: Martin Vit <martin@voipmonitor.org>
…se indexer decode Add a cluster-cooperative top-k op for DSA indexer decode at small batch (TMA row streaming, DSMEM histogram all-reduce, radix tie refinement) plus a single-CTA 4096-bin histogram short path reused by persistent_topk for rows up to 32k. Includes an SM80 shuffle fallback for redux.sync so sm70/75 builds keep compiling, and per-rank tie caps sized to k so k=2048 outputs are always fully written. Co-Authored-By: Martin Vit <martin@voipmonitor.org>
…asoning parsing Add ReasoningParser.is_reasoning_end_for_prompt so multi-turn prompts ending in a previous turn's think-close marker (or containing tool-call examples) no longer flip the stream into reasoning-ended state at start. Rewrite the MiniMax M3 streaming parser to buffer text and withhold partial markers, handling think tags split across deltas by tokenizers that emit them as plain text. Co-Authored-By: Martin Vit <martin@voipmonitor.org>
DeepSeek V3.2/V4 sometimes emit a complete DSML function-call block without closing the think tag, silently losing the tool call in the reasoning channel. Tool parsers can now opt in to recovery: non-streaming re-parses a balanced block out of reasoning content, and streaming intercepts the tool-call start marker inside reasoning with partial-marker buffering. Co-Authored-By: Martin Vit <martin@voipmonitor.org>
…signment into the parser layer Replace the legacy Kimi K2 reasoning and tool parsers with a declarative ParserEngine definition, move kimi-style tool-call ID generation out of the chat serving layer into the parser (extending ID continuity to the Responses API), and make strict tool-choice structural tags model the reasoning prefix so grammar-constrained requests do not reject thinking tokens. Co-Authored-By: Martin Vit <martin@voipmonitor.org>
…tructural tags Lower tool_choice required/named-function into xgrammar structural tags built per tool parser (new structural_tag_model hook on every parser), enforce them through the chat output processor, and model the reasoning prefix in the grammar when a reasoning parser is active so reasoning models do not fail at the reasoning-to-tool boundary. Co-Authored-By: Martin Vit <martin@voipmonitor.org>
Union multiple_of * num_reqs for num_reqs 1..32 into the spec-decode capture sizes so small uniform-decode batches replay exact FULL graphs instead of padded ones, and fall back to plain allreduce + norm in fused_allreduce_gemma_rms_norm under torch.compile where the dispatch is untraceable. Co-Authored-By: Martin Vit <martin@voipmonitor.org>
CommonAttentionBatchTopology caches numpy views of the batch layout, split_decodes_and_prefills moves onto CommonAttentionMetadata with memoization, and builders can declare supports_exact_metadata_reuse plus per-layer metadata group keys to share built metadata across attention groups. Co-Authored-By: Martin Vit <martin@voipmonitor.org>
Decode-context-parallel topk global/local index conversion, DCP-aware indexer and SWA metadata, fp8 KV decode under DCP with LSE return, TRITON_MLA sm120 tuning table and spec-verify decode builder, FlashInfer large-page guards, and DFlash-oriented non-causal SWA semantics in the unified Triton kernel. Co-Authored-By: Martin Vit <martin@voipmonitor.org>
DCP-replicated KV cache specs with per-group hash granularity and block-table cp sizes, sliding-window DFlash draft layers in their own KV groups with null-block slot masking, draft-model config plumbing (kv-cache dtype, attention backend, multi-group init), and per-request temperature for parallel drafting. Co-Authored-By: Martin Vit <martin@voipmonitor.org>
Opt-in B12X_ATTN and B12X_MLA_SPARSE backends backed by the b12x SM12x CuTe DSL kernel package, the b12x paged sparse indexer path with prewarm and DCP merge, CUDA-graph capture guards and prewarm hooks, warmup registry entries, and profiler scope toggling. Co-Authored-By: Martin Vit <martin@voipmonitor.org>
b12x MXFP8/block-FP8 linear kernels, W4A16/W4A8 FP4 MoE experts with planned weights and a replicated-input EP variant, backend registration in the MoE oracles and kernel selection, MXFP4 W2 scale realignment for uneven shards, and MHC custom-op wrappers for torch.compile. Co-Authored-By: Martin Vit <martin@voipmonitor.org>
PCIe oneshot and DMA allreduce backends with NUMA-aware topology detection and autotuned dispatch crossovers, fused allreduce+RMSNorm through the b12x runtime, a cpp custom-allreduce opt-in for PCIe topologies, and b12x DCP LSE reduce-scatter / head all-gather collectives. Co-Authored-By: Martin Vit <martin@voipmonitor.org>
Pad head counts, MoE intermediate sizes, and vocab so allowlisted models run at TP sizes that do not divide them (DeepSeek V4 TP10, GLM DSA TP6, MiniMax M3 TP3), with plan-aware weight loaders zero-filling checkpoint tails past the real data. Co-Authored-By: Martin Vit <martin@voipmonitor.org>
DeepSeek V4 b12x sparse MLA attention, fused WO projection, mHC residual path, compressor DCP support and warmups; MiniMax M3 b12x MSA with torch.compile enablement, pipeline parallelism, and released-checkpoint weight mapping; NextN MoE construction fixes and MTP quant-config handling. Co-Authored-By: Martin Vit <martin@voipmonitor.org>
Local vllm serve wrappers for the lab models, torch 2.12 pins, and cmake fetch/install adjustments for vllm-flash-attention and the SM100 MSA tree. Co-Authored-By: Martin Vit <martin@voipmonitor.org>
Pass vllm_config to _pool_bytes_per_block (signature gained it upstream), update the acceptance-length V2 test for the lifted batch-size dynamic-SD restriction, add kv_transfer_config to the DSv4 test fixture, restore B12xNonCompressedIndexerBackend (dropped when keeping upstream's evolved indexer), and apply cargo fmt to the chat crate.
Enables DSpark speculative decoding with RedHatAI/GLM-5.2-speculator.dspark (override via DSPARK_MODEL). num_speculative_tokens defaults to the speculator checkpoint's n_predict; NUM_SPECULATIVE_TOKENS and GLM52_ADAPTIVE_SPECULATIVE_TOKENS_WINDOW compose as with MTP. Mutually exclusive with --causal-cascade.
Extend triton_filter_and_convert_dcp_index with optional preallocated out/valid_counts buffers so the b12x backend's CUDA-graph-stable views are written directly by the kernel (no per-step allocation or copy), fix the ImportError from the removed pre-rebase helper, and drop a duplicated MXFP8 kernel priority entry left over from the merge.
Restore tip b12x semantics lost in rebase conflict resolution: builder-side b12x prefill chunk planning, schedule metadata and active-width buffers with the exact CPU-shadow decode top-k bound (vllm/v1/attention/backends/mla/ indexer.py); DCP scalars as op args plus aligned fake signature in sparse_attn_indexer; B12X_MLA_SPARSE kv-dtype canonicalization and DCP a2a plumbing in mla_attention; profiling-scope phase restoration and prepare_inputs max_query_len in the v2 runner. Align DFlash/DSpark handling with upstream: per-layer list slot mappings in qwen3_dflash (drop stale Mapping branch), plain use_eagle prefix-cache shift in the scheduler, superseded triton jit_monitor import removed. Restore SpeculativeConfig draft_kv_cache_dtype/draft_attention_backend fields and apply the draft KV dtype override in the dspark/dflash model loaders so dense drafts do not inherit target-only KV layouts (fp8_ds_mla).
Block drafters (DSpark/DFlash) always draft the full block; the adaptive controller caps only verification depth. A reduced depth must verify fewer tokens, so capture a uniform-decode graph for every selectable depth in [1, K] instead of replaying padded max-depth graphs, and gate upstream's batch-size schedule capture on that schedule actually existing. Also add the exact (depth x num_reqs) candidate grid, bounded by min(max_num_reqs, 32) requests, so coverage does not depend on the configured capture-size list. Validated live on GLM-5.2 + RedHatAI dspark speculator at TP8: depth adapts 7->1 under high-entropy load and recovers on predictable content, correct outputs throughout, no errors.
Signed-off-by: Luke Alonso <lalonso@gmail.com>
Signed-off-by: Luke Alonso <lalonso@gmail.com>
…nears Extends the shared-expert MXFP8 overlay (c6902b1) to the rest of the BF16 dense linears a ModelOpt checkpoint excludes. With --quantization-config.linear.weight mxfp8, attention projections, dense MLPs and indexer projections that the checkpoint's exclude list leaves in BF16 are quantized to MXFP8 at load time via mxfp8_e4m3_quantize — bit-identical to an offline-requantized MXFP8-dense checkpoint (verified tensor-for-tensor against one, including TP8 shard-wise == full-tensor grouping, since k stays divisible by 32 under row-parallel sharding). Composition rules: - shared-expert projections are never selected by 'linear'; they remain governed solely by the 'shared_experts' spec, so linear-only gives MXFP8 dense + BF16 shared experts, and linear+shared_experts gives the fully-MXFP8 dense variant - quantization_config.ignore (exact or re: patterns, unfused shard names) keeps individual modules on the BF16 path - routers, lm_head, embeddings and plain nn.Linear modules (GLM MTP eh_proj) are never touched GLM-5.2 recipes (from lukealonso/GLM-5.2-NVFP4, BF16 dense + NVFP4 routed experts): full MXFP8 dense: --quantization modelopt_fp4 --quantization-config '{"linear":{"weight":"mxfp8"},"shared_experts":{"weight":"mxfp8"}}' MXFP8 dense, BF16 shared experts: --quantization modelopt_fp4 --quantization-config.linear.weight mxfp8 Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Luke Alonso <lalonso@gmail.com>
|
Important Review skippedToo many files! This PR contains 251 files, which is 101 over the limit of 150. To get a review, narrow the scope: Upgrade to a paid plan to raise the limit. ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (251)
You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ Finishing Touches🧪 Generate unit tests (beta)
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 |
lukealonso
force-pushed
the
dev/fathomless-firmament
branch
from
July 9, 2026 23:24
3a864ad to
407ace1
Compare
voipmonitor
marked this pull request as ready for review
July 9, 2026 23:35
Author
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
2 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.
What changed
VLLM_DFLASH_FORCE_SINGLE_SWA=1).Why
Kimi-K2.7-Code DFlash on Fathomless either failed DCP4 boot with mixed KV cache specs or ran with poor first-run/JIT-contaminated throughput. The target-only path was healthy after restoring the tuned TRITON_MLA decode path; the remaining DFlash issues were in draft metadata and speculative runner handling.
Validation
python3 -m py_compileon all touched files.max_num_seqs=1, graph cap8, targetTRITON_MLA, draftTRITON_ATTN, KVfp8.139.3 tok/s(/root/bench-results/kimi-k27-v3-ff-dcp4-current-prefill-20260709/decode_cc1_after_prefill.json).8k 7,971 tok/s,64k 4,918 tok/s(/root/bench-results/kimi-k27-v3-ff-dcp4-current-prefill-20260709/prefill_8k_64k.json).python3 /mnt/test.py --port 7801 -L: 7 iterations, 5,807 tokens, avg generation-only205.11 tok/s, CJK0.Note: first decode immediately after boot can be contaminated by Triton JIT (
_topk_topp_kernel,_prepare_dflash_inputs_kernel,_fwd_grouped_kernel_stage1,_fwd_kernel_stage2). Use a warmed decode run for regression comparison.