feat: DeepSeek-V4-Flash integration + Blackwell (sm120) bring-up - #153
Draft
drunkcoding wants to merge 96 commits into
Draft
feat: DeepSeek-V4-Flash integration + Blackwell (sm120) bring-up#153drunkcoding wants to merge 96 commits into
drunkcoding wants to merge 96 commits into
Conversation
drunkcoding
force-pushed
the
feature/deepseek-v4-kernel-integration
branch
from
May 12, 2026 17:11
948efa6 to
1a584bc
Compare
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
Fixes two bugs in the original v4_fused_qk_rmsnorm: (1) variance was computed over the entire flattened row instead of per-head, and (2) only 2 CTAs per token regardless of n_heads. New kernel uses grid (T, n_heads+1) with 3D Q input [T, n_heads, head_dim] for correct per-head normalization. 37 tests including per-head independence gate, parametrized n_heads, and non-contiguous stride coverage. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
Handwritten CUDA kernel replacing Triton dispatch overhead for the fused per-head Q + global KV RMSNorm. Key design: 4 warps/CTA, 16-byte vectorized loads, warp-shuffle reduction, __launch_bounds__(128,16) for 32 regs/thread. Microbench results vs Triton: geomean 1.52x on Blackwell, 2.00x on H20. Vs sglang: wins 24/32 cells (12/16 per arch), losses only at large T where sglang's tile::Memory abstraction extracts 4% more DRAM bandwidth. NCU verified: 86-92% DRAM peak utilization, 32 regs/thread. Also includes minor formatting cleanup in setup.py and _jit_registry.py. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
Move all V4 kernel modules from batchgen/attention/mla/ to batchgen_kernels/{triton,attention,common,moe,src}/. Includes @triton.autotune for cache_utils (5.6x speedup on k7) and CUDA fused_silu_mul_quant port.
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)
Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
- fast_topk_cuda: add valid-m WGMMA variant for variable-length batches - fused_indexer_score: support 4D aux_blocked_k, expand topk configs - fused_kv_norm_rope_cache: extend CUDA kernel for V4 compress pipeline Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
20 test files covering all V4 kernel families: cache_utils, compress_quant, compressor, fp4_dequant, fp4_kv, fused_silu_quant, hash_routing, hyper_connections, indexer, inv_rope_fp8, mxfp4_marlin, qnorm_rope_kv, routing, topk, tilelang_score, c128_online, and integration/attn_backend. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
- DeepSeekV4Flash model: decoder layer, attention, MoE, compressor modules - Wire hc_pre/hc_post from batchgen_kernels.common.v4_hyper_connections (eliminates inline _hc_split/_hc_pre/_hc_post duplicates) - Add v4_backend, v4_indexer_metadata, v4_mxfp4_marlin_moe, v4_fp4_kv_cache Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
Update raw_bytes_per_token from 583 to 584 for both V4-Flash and V4-Pro profiles, matching TOKEN_BYTES (NOPE_DIM=448 + ROPE_DIM*2=128 + SCALE_DIM=8) from v4_cache_utils.py. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
- Gate: replace inline PyTorch with hash_routing/sqrtsoftplus_topk from batchgen_kernels - Expert: fused_silu_mul_quant_cuda for SiLU*mul+FP8 quant (PyTorch fallback) - Wrapper: _forward_decode_optimized routes decode to DeepseekV4AttnBackend (Q/KV projection via module ops, attention via FlashMLA sparse/dense/compressed) Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
- Fix _ref_pre/_ref_post/_hc_split to standalone PyTorch (deleted model methods) - New test_v4_wiring_correctness.py: gate routing (T1), expert activation (T2), attention decode shapes (T3), KV cache bytes (T4), decoder layer HC (T5), performance regression gate (T6) — 54/54 passed on H20 - New test_v4_e2e_logits.py: scaffolding for full model logit validation Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
DeepSeekV4FlashMoE ran each rank's expert shard over only its own local decode tokens, then dense all_reduce'd a [local_tokens, H] buffer. Under DP attention (each rank owns different sequences) this both (a) silently dropped contributions for tokens routed to experts on other ranks and (b) deadlocked NCCL whenever ranks had different token counts. Adopt the established EP-decode contract used by minimax/kimi/glm5: pad local tokens to a global-max (num_tokens_per_rank), all-gather hidden states and input_ids, run the owned expert shard over all global tokens, all_reduce(SUM), then slice back this rank's rows. Wire padding_bsz through configure_decoding (previously discarded) via _init_decoding_padding_bsz/set_num_tokens_per_rank, which the worker's existing _sync_decode_moe_rank_counts already drives. Add a 2-rank torchrun test suite covering remote-rank routing, uneven token counts, hash routing + topk=2 + empty rank, dynamic resize, and guard errors.
…ecode Three fixes that let DeepSeek-V4-Flash decode execute end-to-end: 1. c4 indexer DP-attention gather: the indexer wq_b/weights_proj were left TP-sharded (n_local_heads) while the decode wrapper reshaped to the full n_heads, crashing on c4 layers. Gather them across ranks alongside the main attention tensors (conditional on c4 layers) and use the full-head weights in DP mode, mirroring the main attention path. 2. Compressor bridge init: DeepSeekV4Compressor ran xavier_uniform_ at construction even though the runtime bridge immediately rebinds the weights; on degenerate freshly-created params this raised. Add init_weights=False so the bridge skips random init. 3. Compressor inference-tensor safety: assigning runtime-loaded inference tensors into nn.Linear/nn.Parameter .data made F.linear raise "Inference tensors do not track version counter". Convert the kernel compressor to functional _runtime_linear with buffer-backed ape/norm/wkv_weight/wgate_weight (mirroring _linear_from_weight), and have the bridge detach + assign plain tensors.
… IMA) The c4 sparse path called the fused indexer with the RAW sequence length (meta.seq_lens_casual) while its cached_k buffer is sized for the COMPRESSED length (seq_len // 4). The scoring kernel masks K loads with `s_offs < cache_seqlens` over a buffer whose stride is seq_len//4, so once the raw length exceeded the compressed buffer it read ~4x out of bounds, causing a CUDA illegal memory access partway through decode. Pass meta.c4_topk_lengths_clamp1 (= max(seq_len // 4, 1)) instead. Verified on 4xH20: 2 sequences decode the full 64 tokens (128 total) with zero illegal memory errors; decode now runs end-to-end.
The continuous-decode path pops completed sequences out of
self._local_to_uuid_map / self.query_book before generate()'s result
finalize runs, so the gather loop found them empty and returned no results
("Detokenization complete: 0 sequences" -> empty HTTP response). Iterate
global_batch.get_sequences_for_rank(rank) and read each sequence's own
decoded_tokens buffer-pool view + decoded_length instead, which persist on
the SequenceEntry across the decode lifecycle.
Verified on 4xH20: HTTP /v1/inference now returns
{"status":"success","results":[...]} with 2 sequences detokenized.
NOTE: this commit also carries prior in-progress V4 worker changes that were
already uncommitted in the working tree on this WIP branch (the file could not
be split cleanly). Adds an env-gated V4_RESULT_DEBUG diagnostic.
…A decode Pure-Triton (no wgmma/TMA) slot-based grouped FP4 MoE and sparse MLA decode kernels for Blackwell sm120, where the wgmma-based FlashMLA path cannot run. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
584-byte/token paged FP8 KV pool (FP8 nope + BF16 rope + UE8M0 scales) and the decode-time KV coordinator for V4 sparse MLA. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
…efill populate FlashMLA decode adapter dispatching across torch-ref / sm120-Triton / wgmma backends, the eager torch reference (correctness oracle), and prefill KV population. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
…perts, and PyNCCL collectives Adds decode-timing hooks, env-gated grouped FP4 MoE (BATCHGEN_V4_GROUPED_MOE) with shared-scratch weight staging, persistent local experts (GLM5-style get_tensor load), PyNCCL EP collectives (BATCHGEN_V4_PYNCCL_COMM), and the sm120 Triton MLA decode branch (BATCHGEN_V4_MLA_SM120_TRITON). All flag-gated; default paths preserved. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
Metadata-driven dtype/key resolution for the V4 checkpoint converter and assets package init. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
…hanges Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
…ted) Investigated last near-tie residual: MoE router is float32-faithful to official (refuted); SM120 attn_sink logaddexp is algebraically identical to official/torch (agent's 'LSE mismatch' claim refuted by direct math). Real residual = inherent FP rounding between equivalent kernels (SM120 log2-domain softmax + different GEMM tiling vs official natural-log), flipping only genuine near-ties. Not a discrete bug; documents the diminishing-returns options. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
drunkcoding
changed the title
added 18 commits
June 16, 2026 13:36
Add ARG GPU_ARCH (blackwell default | hopper) conditional branching to the single Dockerfile instead of a separate file, keeping the ~90% shared base and gating only the arch-specific steps: - batchgen_kernels: sm90a (Hopper WGMMA) vs sm120 (Blackwell) - FlashMLA: built from public upstream source at an arch-pinned commit — deepseek-ai/FlashMLA@c741387 for Hopper (the deferred-scheduling / DeepSeek-V3.2 sparse API with zero-arg get_mla_metadata + attn_sink that V4-Flash requires), 1408756a for Blackwell (never calls FlashMLA at runtime) - tilelang + fast_hadamard_transform added for BOTH archs (V4-Flash sparse prefill needs them; the image previously lacked them). fht is built from GitHub source because its PyPI sdist is broken (missing csrc). Everything is built from public source — no prebuilt binaries vendored. Fix sm90a WGMMA build in batchgen_kernels/setup.py: replace bare -arch=sm_90a with explicit -gencode arch=compute_90a,code=sm_90a (6 sites). Bare -arch=sm_90a also emits a plain compute_90 PTX pass that ptxas rejects for wgmma.*. Re-pin apache-tvm-ffi==0.1.5 after flashinfer install: flashinfer pulls in 0.1.12 which breaks tilelang import on torch 2.9. Both GPU_ARCH=hopper and GPU_ARCH=blackwell images build clean from source and import all V4 deps (verified locally; Hopper runtime validated at 82.5% MMLU-Pro on H20).
Remove HANDOFF.md, VLLM_ENTRY_POINTS.md, and the tools/ diagnostic scripts that are no longer referenced. Ignore agent/session docs (AGENTS.md, HANDOFF*.md, opencode.json, etc.) and clean the dangling HANDOFF.md reference in model.py.
Replace the per-sequence Python loops (with .item() host syncs) in the full-prefix, SWA-window, and single-token slot builders with vectorized page-table gathers, gated to the sm120 triton backend via BATCHGEN_V4_FAST_PREFIX_INDICES (default on, slow-path fallback retained). Measured on RTX PRO 6000: full-prefix index build 1.29ms->0.23ms (5.8x), SWA-window 0.22ms->0.12ms (1.8x), resolve_swa_token_slots 104us->63us. The shared _physicalize_tail_padded_contiguous helper is sound only for tail-padded contiguous positions (no interior holes), so sparse c4/c128 selection paths are intentionally untouched. Bit-exact vs the slow path: unit equivalence across page boundaries + extension, plus full c128/dense e2e decode parity (atol=0).
…20 sparse decode Two ~374ms Triton compilations otherwise landed on live decode steps 0 (first-call JIT) and 64 (topk_rounded 64->128 transition), inflating attn_sm120_main mean to 1666us vs 170us p50. Add warmup_sm120_sparse_decode covering both topk buckets (called lazily before serving / graph capture) and pin the autotuned config (BLOCK_T=32, num_warps=8) since the topk is SWA-capped at 128. Result: attn_sm120_main 1666us->240us, step-0 spike 377ms->25us. Also add a geometry guard in _run_triton_sparse_decode: the kernel infers page_size from k_cache.shape[1] and addresses via stride(0), so a flat [num_pages, page_bytes] cache silently misreads page_size and causes OOB scale/rope reads. Require stride(0) >= page_size*584 (allows allocator padding) and raise a clear error instead of a CUDA illegal access.
_pack_model1_rows ran a 7-iteration Python loop (~70 tiny ops) for the per-tile UE8M0 scale + e4m3 quantization, flat at 0.40ms from B=1 to 256 (launch-bound). Reshape nope to [N,7,64] and compute all tiles in one batched op. Byte-exact output; pack 0.40ms->0.075ms (5.3x), cutting attn_kv_store from 16.8% to ~8.6% of the decode step.
…arity Update test references from the removed wkv/wgate nn.Linear submodules to the current wkv_weight/wgate_weight buffers (F.linear), restoring the c128 eager-decode canonical-parity coverage that had silently broken. Add eager<->fused-kernel bridge parity tests: drive the eager compressor and the fused compress+norm+rope+quant+store kernels from identical projections and assert the emitted compressed-KV agrees within fp8/mxfp4 tolerance (overlap=0). Closes the previously-uncovered cross-path equivalence gap before any swap of eager for fused write.
MXFP4 cvt.e2m1x2 PTX is rejected by ptxas on sm_90a (Hopper). Make the DeepSeek-V4 indexer quantization arch-aware so Hopper falls back to FP8: - fused_indexer_q: default use_fp4 now derived from device capability (cc>=12) instead of always-False; explicit override still honored. - _maybe_rotate (eager indexer compressor): the runtime sm90 blocker. Its fp4_act_quant (tilelang cvt.e2m1x2) now falls back to an fp8 e4m3 block-32 fake-quant approximation on sm90. - New _fused_kv_compress_norm_rope_insert_indexer_fp8_attn kernel matching the indexer pool's [128B fp8 | 4B fp32 scale] layout, plus a capability/env-gated dispatcher fused_kv_compress_norm_rope_insert_indexer (BATCHGEN_V4_INDEXER_QUANT=auto|mxfp4|fp8). Validated: sm120 parity unchanged; H20/sm90 runs the fp8 paths with no PTXASError.
The grouped MoE path uses MXFP4 WGMMA kernels (cvt.e2m1x2), rejected by ptxas on sm_90a. Add _v4_grouped_moe_enabled() to AND the BATCHGEN_V4_GROUPED_MOE flag with a cc>=12 check, so Hopper uses the per-expert loop fallback (_dequant_fp4_e2m1_weight is pure-torch, correct on all archs).
…atch coverage Add requires_mxfp4 marker (skipif cc<12) to the 6 MXFP4-only tests so they SKIP cleanly on sm90 instead of crashing ptxas. Add coverage for the new arch-aware paths: fp8 indexer-compress parity, indexer-quant selector, and the fused_indexer_q capability dispatch default. Validated on H20: MXFP4 tests skip, fp8/dispatch tests pass, no PTXASError.
…ode-loop e2e The DeepSeekV4Compressor kernel class moved from wkv/wgate nn.Module submodules to wkv_weight/wgate_weight buffers; update the two stale assertions. Add world_size=1 to the SimpleNamespace module mocks the current wrapper reads. Skip the indexer fp4_act_quant tests on sm120 when tilelang is unavailable (sm90 exercises the fp8 fallback and runs them). Validated: sm120 11 passed / 2 skipped; H20/sm90 13 passed, 0 PTXASError.
One-stop script (build/launch/wait/smoke/mmlu/stop) to rebuild the Hopper image from current source and run DeepSeek-V4-Flash on 4x H20. Encodes the launch gotchas found during sm-aware validation: 512G --shm-size (100G host-KV + 320G weight region), stale-shm cleanup between launches, and the sm-aware env flags (BATCHGEN_V4_INDEXER_QUANT=auto, grouped-MoE off). Context: the sm-aware kernel switch is validated to load + serve on H20 with zero PTXASError; this runbook closes the loop so a fresh-from-HEAD image runs inference (fixing the stale-image FlashMLA/tilelang drift).
Building the image from inside China was failing/crawling on three upstream sources. Parameterize them (defaults unchanged, so non-CN builds are unaffected): - TORCH_FIND_LINKS: torch==2.9.0+cu129 via uv --find-links to a flat wheel mirror (download.pytorch.org was ~127KB/s in CN; Aliyun is multi-MB/s). - UV_DEFAULT_INDEX: PyPI mirror for the remaining uv installs (uv ignores PIP_INDEX_URL); the flashinfer-python step otherwise crawled on files.pythonhosted.org. - GitHub recursive clones (FlashMLA/DeepGEMM + cutlass): force git HTTP/1.1 + bigger postBuffer, and wrap DeepGEMM in a 5x retry loop to survive the intermittent 'HTTP2 framing layer' / GnuTLS clone failures from CN. Verified by a full hopper rebuild on a CN H20 node (batchgen:v4flash-hopper-current, 27.2GB): sm90a kernels + FlashMLA c741387 + tilelang all build clean. docker/README.md documents the CN build invocation.
A rank with an empty local decode_uuids returned early from _sync_decode_uuids_tensor / _sync_completion_status_tensor BEFORE their dist.all_reduce, skipping a decode-entry collective the other ranks ran. The idle rank then slept on dist.barrier() (futex) while the rest blocked forever in the skipped all_reduce / the per-layer MoE all-gather -> the multi-GPU decode hang (7 spinning + 1 futex on MP8). Derive the tensor size and the empty-decision from an all_reduce(MAX) of a local max-index so every rank runs the identical collective sequence even when its local list is empty. Add a global all_reduce(MAX) break-validation before the decode-entry break so ranks leave the loop together, with a fail-fast RuntimeError on residual desync. Validated on H20 MP8: decode advances through all 43 layers with all 8 ranks in lockstep (was: hung at iteration 0).
Per-rank, immediately-flushed fd-2 markers ([DDL] pid=... rank=... ...) at the decode-entry syncs, the decoding_continuous loop, and the per-layer MoE collectives (states/ids all-gather, all-reduce). fd 2 is written directly so markers survive a hung/buffered logging pipeline. Off by default; enable on multi-GPU H20 with BATCHGEN_DECODE_DEADLOCK_TRACE=1. Used to confirm rank lockstep during the decode-deadlock validation; paired with docker/v4_h20_validate_decode_fix.sh diagnose.
_run_owned_experts called counts[expert_idx].item() once per owned expert to skip empty experts -- ~32 device-to-host syncs/layer, ~1.4k/token over 43 layers, the dominant decode-step cost. Replace with a single .tolist() of the owned-counts slice, then index the Python list. Numerically identical; removes the per-expert sync from the decode hot path.
Mount a persistent host dir at /root/.cache/torch_extensions (TORCH_EXT_CACHE) so core_engine + the runtime load() kernels are not recompiled (~15-20min, 4 extensions) on every fresh container; later launches reuse the cached .so. Add warmup / cache-status subcommands to populate and inspect it. Add --ulimit memlock=-1 --ulimit stack=67108864: pinning the 100GB host-KV region via cudaHostRegister fails with the default 64KB memlock (workers crash with 'cudaHostRegister failed: invalid argument'). Matches the known-good sibling container's limits. Add RUN_ENV_EXTRA hook and the v4_h20_validate_decode_fix.sh harness (launch + decode smoke + wchan/marker capture) used to validate the decode-deadlock fix on H20 MP8.
Patch/minor bumps clearing ~17 Dependabot alerts (no API breakage): - filelock 3.20.0 -> 3.20.3 (TOCTOU symlink) - idna 3.11 -> 3.15 (CVE-2024-3651 bypass) - python-dotenv 1.2.1 -> 1.2.2 (symlink overwrite) - python-multipart 0.0.21 -> 0.0.31 (DoS / file write) - requests 2.32.5 -> 2.33.0 (temp file reuse) - sentencepiece 0.2.0 -> 0.2.1 (heap overflow) - urllib3 2.5.0 -> 2.7.0 (redirect header leak / decompression bomb)
Collaborator
Author
sm120 (consumer Blackwell) decode path — implemented & mergedThe sm120 (RTX PRO 6000 / 5090, compute 12.x, no wgmma/TMA) decode path is implemented and merged via #253 (merged 2026-06-09, Kernels & code locations
Arch-aware MoE gates (below sm120 fall back / skip)
Status
|
added 8 commits
July 4, 2026 13:57
…) + FP4 utils MXFP4 grouped-expert decode/prefill kernels for RTX 6000 Pro (sm120), mega3 as the production decode path (int64-indexed), plus JIT registry and setup wiring. Includes microbenchmark harness.
Route V4-Flash MoE through the sm120 grouped kernels as the only path (fallbacks raise), add BATCHGEN_V4_RESIDENT_EXPERTS env propagation, EP offloading config plumbing, and DP-prefill/EP-decode phase strategy.
The prefill/compress RoPE caches were capped at max_position_embeddings (default 8192; V4-Flash config only sets original_seq_len=65536), so any position >= 8192 indexed out of bounds -> device-side assert during long sequence prefill and decode. Floor at original_seq_len and rebuild when a longer sequence arrives. Verified: 8192-token prefill 390 tok/s, 0 asserts.
Prefill admission only checked host-KV capacity and decode selection used the two-page-buffer working-set estimate, while actual allocation needs full-context KV: oversized batches raised 'Insufficient free pages' mid-forward, silently dropping all but ~2 sequences from the response (prefill) or failing the whole request (decode). Add _truncate_batch_to_gpu_kv_fit: cumulative per-pool preflight via can_allocate_pages_for_sequences, MIN-all-reduced so the SPMD batch stays rank-identical; excess sequences stay queued for the next wave. Also clear pending KV-append tasks on batch reset (post-OOM host leak path). Verified e2e: 192x8192 request now returns 192/192 results with wave-cycling ([PREFILL] 27/189..., [DECODE] 33/192...) and zero allocation failures (was 2/192 silent).
…ch harness Setup doc records verified launch config, tuning-lever findings (decode best c20_f06 22.3 tok/s; recommended prefill 128x8192 ~2.4K tok/s aggregate), bug root causes (RoPE cap, silent drop, decode over-admission) and fix verification. Autoresearch harness (force-added past benchmarks/ gitignore) is the reproducible experiment runner with the full results ledger.
Expose 7 planner-internal knobs (BATCHGEN_{ATTN,MOE}_{PREFILL,DECODE}_MB,
BATCHGEN_EXPERT_{PREFILL,DECODE}_CAP, BATCHGEN_PREFILL_TOKEN_CAP) applied
after model-specific planning; unset keeps planner defaults.
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.
Summary
End-to-end integration of DeepSeek-V4-Flash (284B) into BatchGen, from kernel consolidation through full production model wiring and Blackwell (sm120) bring-up. The branch now runs V4-Flash end-to-end on 4× RTX PRO 6000 Blackwell: stable under sustained load, 72–75% MMLU-Pro, single-token character-exact vs the official reference, and dozens-to-hundreds of characters exact on long greedy generations.
129 files, ~50k insertions. Three broad phases:
Phase 1 — Kernels (original scope)
batchgen_kernels/{triton,attention,common,moe,src}/(cache_utils, compress_quant, indexer_q, inv_rope_fp8, hyper_connections, fp4_dequant, silu_mul_quant, hash_routing, sqrtsoftplus_topk, …);@triton.autotuneon cache_utils (5.6× on k7).Phase 2 — V4-Flash model wiring + Blackwell bring-up
Correctness fixes found during bring-up (all verified)
configure_decodingmutated the shared expert-streaming task, leaking into the next prefill ("expert weights not loaded"); restore pristine task inconfigure_prefill.mgr._sequenceswhich the V4 coordinator lacks (AttributeError→SIGSEGV); (b) coordinator not re-initialized on watermark re-prefill (destroyed-but-not-None). Both fixed; verified through a 500-prompt run that exercises the full on-hold→re-prefill→resume cycle.Phase 3 — Character-exact decode (QAT numerics)
The model is QAT-trained (activations quantized to fp8/fp4 before each GEMM). The default paths dequantized weights to bf16 instead, causing greedy decode to drift from the reference. Fixes (all env-gated, default off; defaults preserve the fast throughput path):
BATCHGEN_V4_QAT_LINEAR=1— QAT-faithful dense/attention linears (act_quant + fp8/fp4 GEMM). Bit-exact vs official (cos=1.0).BATCHGEN_V4_QAT_MOE=1— new QAT-faithful grouped MoE decode kernel (v4_grouped_mxfp4_moe_forward_qat), per-expert act_quant + fp4_gemm. Bit-exact vs per-expert reference; the prior bf16-weight-dequant grouped kernel (cos~0.9988) is annotated non-QAT and kept as the fast path.index_score.relu_()before the gate-weighted sum, changing ~30% of the topk KV selection; added relu (3 score sites) + fp4-quant on indexer q/k to match official.Char-exact progress (greedy A/B vs golden)
identityfirst divergence: ~char 10 (no QAT) → char 122;haiku: char 2 → char 71;tiny-math: exact. The remaining residual is inherent FP rounding between batchgen's optimized sm120 kernels (log2-domain softmax, different GEMM tiling) and the reference's natural-log kernels — i.e. near-tie token flips, not a discrete bug (verified: fp8 KV is bit-exact, router is float32-faithful, attn_sink is algebraically equivalent).Tests
tools/.Notes / follow-ups
HANDOFF.mdfor the complete bring-up log, root-cause analyses, and the working launch configuration.