[voipmonitor]

Summary

DeepSeek-V4.1-Flash on the lossless MXFP4-CSF checkpoint (TP4, MXFP4 routed experts, MXFP8 activations, w4a8_mx) ran
experts.mxfp4_csf.decode(...) before every MoE call: it expanded the routed experts' E8M0 scales into the shared native
SFB scratch, and the W4A8 kernels then read that scratch. In serving this expansion was 2.73% of GPU kernel time at C16
(DSpark K7) and CSF decode was 4.3% (C8) / 1.7% (C16) slower per step than the uncompressed checkpoint.

This PR makes the compact W4A8 kernels read the compressed scales inline. No scale expansion runs before the MoE for
compact geometries, and outputs are bit-identical to both the expansion path and uncompressed scales.

Plans above W4A8_CSF_INLINE_MAX_TOKENS (1536, env B12X_W4A8_CSF_INLINE_MAX_TOKENS) are compute bound: they expand every
expert's scales once from the inline storage (a Triton kernel, 61 us per layer) and run the native kernels.

Serving (DeepSeek-V4.1-Flash TP4, DSpark K7, 4x RTX PRO 6000 Max-Q, lossless CSF checkpoint)

Beta code (b12x d4b73caf, vLLM ff292a9c63) with this branch; llm_decode_bench decode at context 0 (two 30 s runs), standalone prefill
(two runs). Same day and GPUs for all three arms.

arm	C8 tok/s	C16 tok/s	C32 tok/s	prefill 8K / 32K	KV cache
uncompressed checkpoint	807	1108 / 1101	1624 / 1600	17425 / 17609	9.73M
CSF, per-call expansion (B12X_W4A8_CSF_INLINE=0)	783 / 790	1068 / 1070	1584 / 1577	17980, 17746 / 17584, 17494	13.11M
CSF, this branch	823 / 810	1101 / 1115	1648 / 1642	17970, 17413 / 17604, 17584	12.79M

Decode +3.6% (C8, C16) and +4.1% (C32) over the expansion, at or above the uncompressed checkpoint; prefill unchanged. Greedy smoke
output is coherent; server-level greedy text differs between restarts for every arm (also uncompressed vs expansion), so it is not an
identity check here. The kernel-level checks below are.

Fixes found by the serving run

• KV cache: the first serving run had 11.96M KV tokens, 8.8% below the expansion. vLLM loads weights with expandable segments and
  max_split_size_mb:20; the inline storage, allocated after the encoder's large intermediates, kept their freed pages mapped
  (torch.cuda.memory_reserved() - memory_allocated() after load: 1.01 GiB vs 0.04 GiB). Encoding and verification now run in a private
  torch.cuda.MemPool and only the storage is copied out (0.06 GiB; KV 12.79-12.84M). A subprocess test with vLLM's allocator settings
  reproduces it: 115 MiB stranded over eight DS4.1 planes before, 15 MiB after.
• Prefill: pure inline reads lost at 2048 and 4096 tokens (layer 20: +2.5% / +4.4% against the expansion, layer 39: +5.6% / +6.3%),
  hence the 1536-token limit above. With it, layer 20 against the expansion: 8 tokens -9.0%, 128 -10.1%, 1024 -7.5%, 1536 -3.5%,
  2048 and 4096 0.0%; layer 39: -9.1%, -9.9%, -6.9%, +0.7%, +0.5%, 0.0%. Bit-identical throughout.
• Frozen serving: the expansion is a Triton program, which serving resolves without lowering. Fused-MoE preparation compiles it
  (compile_inline_scale_expansion, a geometry-only payload) for dynamic plans above the limit; the launch goes through
  compile_plan.launch_triton. Without this the m4096 plan failed to prepare (no-compilation phase encountered an uncached Triton program: _expand_tiles).

Design

b12x/_lib/quant/mxfp4_csf_inline.py (new):

• Storage per projection, built once at weight preparation on the GPU from the native compact scale plane
  (the scratch already holds every expert's native scales at that point):
  • one 96-byte block per native scale tile (128 rows x 4 K32 columns, 64 rows in an N64 group tail, 2 columns in a
    K64 tail): 64 bytes of selector nibbles (one bit per scale: base or base + 1), a header word and up to 7
    exception records (slot, column, byte);
  • 128 row bases per row block (one byte per row);
  • a raw 512-byte copy of the rare tiles with more than 7 exceptions (header marks them).
    Rows are stored in consumer slot order (slot = 32 warp + 4 quad + nt), so every MMA thread reads its four row
    bases with one 32-bit load per task and its four selector nibbles with one shared 32-bit load per stage.
    Bases are chosen per row from the native bytes (most bytes equal to base or base + 1), independently of the
    checkpoint's unsliced row bases. Preparation verifies that the storage decodes to the native plane byte for byte.
• Kernels (behind a new compile-time csf_inline flag; every existing path is unchanged when it is off):
  • the stage copies the 96-byte tile block with six 16-byte cp.async in the same commit group as A/B (replacing the
    512-byte SFB copy, or the byte-wise synchronous SFB loads of phase 2);
  • after the stage barrier each thread rebuilds exactly the four scale words it feeds to its MMAs, in registers:
    base * 0x01010101 + spread(nibble), then patches exception records in registers (or, for a raw tile, loads its
    four native words); K-tail tiles clear the two padding columns as the native stage does. No extra barrier.
  • Inline kernels: compact micro FC1 (W4A8CompactMicroProjectionKernel), compact micro FC2 (W4A8MaterializedPhase2Kernel,
    direct routes), and the split-materialized M16 FC1/FC2 (W4A8MaterializedPhase1Kernel / W4A8MaterializedPhase2Kernel).
    These are all kernels that compact (N % 128 == 64) W4A8 runs, so compact CSF experts no longer keep the expansion
    decoder at all.
• Plumbing: MoEWeightPreparationPlan.w4a8_csf_inline (in the weight payload, so compile jobs and program keys
  distinguish it), B12XFP4ExpertWeights.mxfp4_csf_inline, w4a8_csf_inline compile key / launch-policy bit for the
  dynamic kernel and csf_inline key for the compact micro launch. b12x_moe_fp4 skips the decode and passes the inline
  storage in every SFB operand slot, so no launch argument references the caller's expansion scratch.
• B12X_W4A8_CSF_INLINE=0 at preparation restores the per-call expansion. N % 128 == 0 geometries (N256 repacked)
  keep the expansion path unchanged.

Measurements

benchmarks/benchmark_w4a8_csf_inline.py: one DS4.1 TP4 MoE layer (E=384, H=5120, N=576 per rank, top-6, MXFP8
activations), three arms sharing weights, scales and routes, planned/prepared/graph-captured outside timing; median of
60 interleaved replays of a 20-call CUDA graph, per call. All arms bit-identical (checked by the benchmark).
RTX PRO 6000 Blackwell Workstation (600 W, GPU-a2f726b2-60db-8102-abef-abf76c457fe0), image
ghcr.io/local-inference-lab/vllm:karmic-kraken-beta-20261004-be47852127fafda5 (CUTLASS DSL 4.7.1), dynamic clocks.

data	tokens	backend	native (us)	CSF + expansion (us)	CSF inline (us)	inline vs native	inline vs expansion
real layer 0, rank 0	8	micro	111.9	118.6	107.9	-3.6%	-9.0%
real layer 0, rank 0	64	dynamic	603.1	645.8	587.5	-2.6%	-9.0%
real layer 0, rank 0	128	dynamic	813.2	871.0	783.9	-3.6%	-10.0%
real layer 20, rank 0	8	micro	111.8	118.8	108.4	-3.1%	-8.8%
real layer 20, rank 0	64	dynamic	603.0	645.7	589.2	-2.3%	-8.7%
real layer 20, rank 0	128	dynamic	813.0	870.9	783.2	-3.7%	-10.1%
real layer 39, rank 0 (densest exceptions)	8	micro	111.7	121.0	109.9	-1.7%	-9.2%
real layer 39, rank 0	64	dynamic	602.8	648.9	591.3	-1.9%	-8.9%
real layer 39, rank 0	128	dynamic	815.5	873.9	788.8	-3.3%	-9.7%
synthetic	8	micro	110.8	118.0	109.6	-1.1%	-7.1%
synthetic	64	dynamic	608.9	645.9	593.1	-2.6%	-8.2%
synthetic	128	dynamic	811.6	877.0	787.0	-3.0%	-10.3%

Inline is faster than uncompressed scales because a stage reads 96 instead of 512 scale bytes per tile.
Scale storage of one layer and rank: native 101.2 MiB, inline 23.4 MiB (layer 0) to 27.3 MiB (layer 39).

Command (real layers): python benchmarks/benchmark_w4a8_csf_inline.py --checkpoint <DeepSeek-V4.1-Flash-MXFP4-CSF> --layer {0,20,39} --tokens 8,64,128 --rounds 60.

Tests

Added with the fixes: test_build_does_not_strand_allocator_pages, test_expansion_reproduces_native_bytes (all inline geometries),
test_expansion_launches_the_program_preparation_compiles (empty Triton cache, no_compilation(); fails without the compile job),
test_inline_plans_above_the_limit_expand_the_inline_storage (limit lowered to 16; 17 and 512 tokens under CUDA graph replay with a
poisoned scratch). Final run of the CSF/inline/x4t/NVFP4-CSF suites plus test_fused_moe_variant_selection.py and
test_prepared_plan_compile.py: 252 passed, 1 failed (test_mhc_runtime_projection_splits_share_code_across_prepared_plans,
pre-existing on the base).

All runs on one RTX PRO 6000 (GPU 14/15) in the image above, PYTHONPATH=<worktree> python -m pytest -q -p no:cacheprovider ....

• New tests/quantization/test_mxfp4_csf_inline.py: storage reproduces native compact planes (DS4.1 W13 w13/w31 with
  N64 group tails, DS4.1 W2 with K64 tails, Kimi TP16, 3584x6) with sparse and dense exceptions, clamped and zero rows,
  raw tiles; a probe kernel stages every tile like the compact kernels (thread offsets 0 and 32) and its rebuilt words
  equal the native staged words exactly.
• New in tests/moe/test_mxfp4_csf.py:
  • test_inline_scales_match_native_and_expansion_under_poisoned_replay: native vs per-call expansion
    (B12X_W4A8_CSF_INLINE=0) vs inline, H/N = 5120/576 (DS4.1 TP4), 1024/320 and 512/256 (N128: expansion kept),
    tokens 1, 8, 17, 64, 512 (micro and split M16 dynamic), w13 and w31 layouts, int32/int64 routes, CUDA graph capture
    and three replays with mutated inputs, the shared expansion scratch filled with 0xFF (NaN E8M0) before every replay
    and still 0xFF afterwards, zero replay allocations, outputs equal bit for bit. Fixture scales put exceptions in most
    tiles and dense rows that overflow the records (raw tiles).
  • test_inline_capacity_plan_reuses_its_launches_for_live_counts: capacities 8 (micro) and 512 (dynamic), live
    counts 1, 3, capacity/2 + 1 and capacity after session.freeze() (frozen kernel resolution), bit-exact.
  • test_inline_requires_compact_split_kernels.
• Existing suites, base dfe61efc vs this branch (same command, files: tests/moe/test_mxfp4_csf.py,
  tests/quantization/test_mxfp4_csf.py, tests/quantization/test_x4t_scales.py, tests/moe/test_w4a8_mx_tp_moe.py,
  tests/moe/test_w4a8_tp_moe.py, tests/moe/test_w4a8_dynamic_kernel.py, tests/moe/test_nvfp4_csf.py,
  tests/moe/test_nvfp4_csf_dispatch.py, tests/moe/test_fused_moe_variant_selection.py,
  tests/preparation/test_moe_corpus.py, tests/preparation/test_tuning_predicates.py,
  tests/preparation/test_prepared_plan_compile.py, tests/preparation/test_precision_choices.py, plus the new
  tests/quantization/test_mxfp4_csf_inline.py):
  • base: 41 failed, 485 passed, 5 skipped;
  • this branch: 41 failed, 560 passed, 5 skipped (the 75 new tests pass). The 41 failures are identical and
    pre-existing: 40 W4A8 TP-MoE tests call stale helpers (plan_tp_moe_execution()/build_tp_moe_fp4_binding()
    without decode_config, non-canonical experts) and test_mhc_runtime_projection_splits_share_code_across_prepared_plans.
  • After rebasing onto the current beta tip d4b73caf (NVFP4-CSF: expand_scales() for ahead-of-call scale expansion #477): tests/quantization/test_mxfp4_csf_inline.py,
    tests/moe/test_mxfp4_csf.py, tests/quantization/test_mxfp4_csf.py, tests/quantization/test_x4t_scales.py
    and tests/moe/test_nvfp4_csf.py: 185 passed.
• Real checkpoint (validation/csf/check_mxfp4_csf_inline.py, vLLM read_mxfp4_csf_layer):
  DeepSeek-V4.1-Flash-MXFP4-CSF layers 0, 20, 38, 39, TP4 ranks 0 and 3: inline storage decodes to the native planes
  exactly (raw tiles: 0 to 7367 per plane), and native / expansion / inline MoE outputs are bit-identical at 1, 8
  (micro), 64 and 128 (dynamic) tokens under graph replay with the poisoned scratch.
• Compile identity: with the flag off, the cubins of phase 1, phase 2 (grouped and direct) and the micro projection are
  byte-identical to the base revision. Flag on, no spills or local memory: registers phase 1 168 -> 228, phase 2
  102 -> 126 (direct 100 -> 120), micro projection 84 -> 124; the persistent grids stay at two CTAs per SM.

Risks and open items

• Only compact (N % 128 == 64) geometries are inline. N256-repacked W4A8 (for example DS4-Flash TP2, N=1024) keeps the
  per-call expansion; inline staging there needs the monolithic dynamic kernel's repacked SFB consumers.
• Inline storage is about 1.5x the bytes of the CSF batch the expansion decoder kept (about 25 MiB vs 16 MiB per DS4.1
  TP4 layer, 0.33 GiB per GPU in total). The caller's shared expansion scratch is filled during preparation and serves
  the calls above the inline limit (about 101 MiB per rank).
• The MoE tuning query does not distinguish inline from expanded scales; cached tunings measured with the expansion
  are reused.

AI assistance (Claude Code) was used for the implementation, tests, measurements and this description.

🤖 Generated with Claude Code

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

• Configuration used: Repository: local-inference-lab/b12x/.coderabbit.yaml
• Review profile: CHILL
• Plan: Advanced
• Run ID: 5625af79-6e54-4ae2-9cff-f0105e1865d3

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

• Autopilot · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

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