# Lucifer-CUTLASS TP3 virtual-padding audit ## Scope - Audited image: `ai01/vllm-v9-b12x-tp3:20260710-v4-init` - Image ID: `sha256:7413040323dff257965bb30153de56c0720ce699408a35be749dc4fc8446fbb5` - vLLM line: v9 / `45c1582e9` - Model: local `DeepSeek-V4-Flash` - Requested backend pair: - attention: `FLASHINFER_MLA_SPARSE_DSV4` - MoE: `flashinfer_cutlass` - Requested topology: tensor parallel size 3, expert parallel disabled This is a source and shape audit only. No service, production configuration, container, image, or vLLM source was changed. ## Conclusion Lucifer TP3 is feasible, but removing the B12X backend guard alone would load corrupt routed-expert weights. Most of virtual TP is already backend-neutral in practice. DeepSeek attention, the output projection, shared experts, vocabulary, and ordinary linear weight loaders can represent the TP3 padded geometry. The blocking implementation gap is the MXFP4 `w2` loader after FlashInfer-CUTLASS rounds the logical local MoE width from 704 to 768. The work is moderate and localized: 1. Replace the B12X-only activation/validation guard with an explicit, capability-checked Lucifer combination. 2. Make MXFP4 TP slicing aware of storage units: unpacked elements, two FP4 elements per packed byte, and 32 elements per E8M0 scale. 3. Prove odd-world-size CUTLASS execution and the padded attention/output path on hardware before model-level benchmarking. ## Exact guards ### Activation and validation `vllm/config/virtual_tp.py` contains the direct policy gates: - Lines 22-27 define plan kind `b12x-padded` and the alignments. - Lines 32-61 activate virtual TP only for supported models when either `_uses_b12x_attention()` or `_uses_native_b12x_moe()` is true. - Lines 442-446 reject expert parallelism. This restriction should remain. - Lines 448-451 reject DeepGEMM MegaMoE. This does not block Lucifer. - Lines 453-457 require native B12X MoE. This is the immediate Lucifer blocker. - Lines 468-472 require B12X sparse MLA attention. - Lines 532-535 define native B12X MoE as backend `b12x`, or `auto` with `VLLM_USE_B12X_MOE=1`. There is a subtle second issue in `_uses_b12x_attention()` at lines 537-554. An explicit `B12X_MLA_SPARSE` backend returns true, but lines 549-554 also return true for any supported virtual-TP model on CUDA SM120. On ai01, the Lucifer attention backend therefore passes the B12X attention check accidentally. That fallback is not evidence that FlashInfer attention was intentionally validated. The correct policy change is not to delete lines 453-472. It should explicitly allow a vetted capability tuple: - model type `deepseek_v4` - CUDA capability 120 - attention backend `FLASHINFER_MLA_SPARSE_DSV4` - MoE backend `flashinfer_cutlass` - tensor parallel only; expert parallel remains disabled The `b12x-padded` plan name is misleading but not itself a runtime blocker. Renaming it can be deferred if compatibility matters; backend validation must not rely on the broad SM120 fallback. ## TP3 shape plan DeepSeek-V4-Flash starts with: - hidden size 4096 - 64 attention heads, head dimension 512, one shared KV head - 8 output groups, output LoRA rank 1024 - routed MoE intermediate size 2048, 256 experts, top-k 6 - one shared expert with intermediate size 2048 - vocabulary size 129280 - 64 replicated index heads of dimension 128 - MXFP4 experts and `swiglu_limit=10.0` The existing plan builder produces: | Axis | Original | Virtual global | Logical per rank | Kernel-local | | --- | ---: | ---: | ---: | ---: | | Attention heads | 64 | 72 | 24 | FlashInfer pads 24 to 32 | | Output groups | 8 | 9 | 3 | 3 | | Routed MoE intermediate | 2048 | 2112 | 704 | CUTLASS rounds 704 to 768 | | Shared-expert intermediate | 2048 | 2304 | 768 | 768 | | Vocabulary | 129280 | 129408 | 43136 | 43136 | The formulas are: - attention: `align_up(ceil(64 / 3), 8) = 24`, global 72 - output groups: preserve 8 heads per group, so `72 / 8 = 9` - routed MoE: `align_up(ceil(2048 / 3), 32) = 704`, global 2112 - shared expert: `align_up(ceil(2048 / 3), 128) = 768`, global 2304 - vocabulary: align globally to `lcm(64, 3) = 192`, yielding 129408 The KV head, index heads, hidden size, expert count, and routing top-k are not padded. The indexer remains replicated. ## Attention trace The Lucifer attention path is already structurally compatible with the plan: - `vllm/models/deepseek_v4/attention.py:176-205` reads the mutated 72 heads and 9 output groups, producing 24 heads and 3 groups per rank. - `attention.py:204-211` allocates the sink at the backend-padded head count and initializes unused entries to negative infinity. - `attention.py:222-255` builds `wq_b`, `wo_a`, and `wo_b` from the mutated dimensions. - `attention.py:485-556` allocates backend-padded output and slices it back to the 24 logical local heads before output projection. - `nvidia/flashinfer_sparse.py:528-540` supports SM120 local head counts by padding to 16, 32, 64, or 128. Local 24 therefore runs as 32. - `nvidia/flashinfer_sparse.py:542-555` passes dynamic local group counts to the fused inverse-RoPE/DeepGEMM output projection; no `groups == 2/4/8` guard was found. The resulting projection shapes are: | Weight | Virtual global shape change | Local TP3 width | Rank 2 tail | | --- | --- | ---: | ---: | | `wq_b` output | 32768 to 36864 | 12288 | 4096 zero rows, 8 heads | | `wo_a` output | 8192 to 9216 | 3072 | 1024 zero rows, 1 group | | `wo_b` input | 8192 to 9216 | 3072 | 1024 zero columns, 1 group | `wo_a` input remains 4096 because `72 * 512 / 9` preserves the original eight heads per output group. Ranks 0 and 1 each own 24 real heads. Rank 2 owns 16 real heads followed by 8 virtual heads. FlashInfer adds another 8 backend-only heads on every rank to reach 32; those backend-only heads are sliced away. One correctness condition still needs a hardware test: zero queries do not guarantee zero attention values. The eight rank-2 virtual logical heads are safe only if padded output group 8 has exactly zero `wo_a` rows and `wo_b` columns, making its final contribution zero. ## Routed-expert loader failure ### Where the two sizes diverge `vllm/model_executor/layers/fused_moe/config.py:1313-1332` divides the virtual global intermediate size by TP and preserves the result as the unpadded local size: `2112 / 3 = 704`. `RoutedExperts.__init__` at `routed_experts.py:125-137` then asks the selected MXFP4 backend to round kernel dimensions. In `fused_moe/oracle/mxfp4.py:652-686`, FlashInfer-CUTLASS aligns intermediate and hidden dimensions to 128. Hidden 4096 is unchanged; local intermediate 704 becomes 768. `vllm/model_executor/virtual_tp.py:43-46` currently returns `min(virtual_local_size, param_axis_size)`. It assumes both values use the same unit. That is true for unpacked W13 rows, but false for packed W2 bytes and W2 scale groups. ### W13 happens to load correctly W13 stores the intermediate dimension unpacked. Its parameter has 768 rows per half and the virtual plan says 704 logical elements per rank. The loader at `routed_experts.py:451-512` therefore selects checkpoint offsets 0, 704, and 1408 and zero-fills the kernel tails: | Rank | Checkpoint elements per W1/W3 half | Zero tail in 768 rows | | ---: | --- | ---: | | 0 | `[0, 704)` | 64 | | 1 | `[704, 1408)` | 64 | | 2 | `[1408, 2048)` | 128 | Rank 2's 128 zeros comprise 64 virtual-plan elements plus 64 CUTLASS-only elements. ### Packed W2 weight is wrong MXFP4 W2 stores two FP4 elements per byte. A 768-element kernel shard therefore has a parameter axis of 384 bytes, while the correct logical checkpoint shard is `704 / 2 = 352` bytes. The generic fallback at `routed_experts.py:551-555` asks for the virtual shard size using `param_axis_size=384`. The unit-blind helper returns 384 and chooses byte offsets 0, 384, and 768. Correct offsets are 0, 352, and 704. | Rank | Correct source bytes | Destination bytes | Required zero tail | | ---: | --- | --- | ---: | | 0 | `[0, 352)` | `[0, 352)` | 32 | | 1 | `[352, 704)` | `[0, 352)` | 32 | | 2 | `[704, 1024)` | `[0, 320)` | 64 | Without a change, data is assigned to the wrong TP-local columns and CUTLASS-only tail columns contain checkpoint data instead of zeros. ### W2 E8M0 scales are wrong W2 has one scale per 32 logical elements. A kernel shard has `768 / 32 = 24` scale groups, while a full logical rank has `704 / 32 = 22` groups. The Lucifer parameter does not receive the B12X-only metadata attached in `mxfp4.py:815-821`, so the same generic fallback selects group offsets 0, 24, and 48. Correct offsets are 0, 22, and 44. | Rank | Correct source groups | Destination groups | Required zero tail | | ---: | --- | --- | ---: | | 0 | `[0, 22)` | `[0, 22)` | 2 | | 1 | `[22, 44)` | `[0, 22)` | 2 | | 2 | `[44, 64)` | `[0, 20)` | 4 | The existing B12X attributes are read at `routed_experts.py:741-750` and routed through the scale-aware branch at lines 532-550. They are backend-specific and cover scale groups, not the FlashInfer packed-weight case. ## Required loader changes No implementation should enable the Lucifer guard until all of these are in place: 1. Preserve the logical local width 704 independently from the CUTLASS kernel width 768. `intermediate_size_per_partition_unpadded` already carries it. 2. Replace the `b12x_mxfp4_*` loader metadata with backend-neutral MXFP4 shard metadata. 3. Annotate W2 packed weight with logical elements per rank and packing factor 2. Compute source starts from `704 * tp_rank`, then convert to bytes. 4. Annotate W2 scales with logical elements per rank and group size 32. Compute source starts in elements, then convert to scale-group offsets. 5. Zero the full destination shard before copying so virtual and kernel-only tails are deterministic. 6. Make W13 use the same explicit logical-width metadata even though its current unpacked path happens to be correct. 7. Verify the post-load CUTLASS permutation/interleave preserves the intended logical slices and keeps all padded tails zero. Likely source surface: - `vllm/config/virtual_tp.py`: capability-based activation and validation - `vllm/model_executor/virtual_tp.py`: unit-aware shard conversion helper - `vllm/model_executor/layers/fused_moe/routed_experts.py`: metadata-driven W13/W2 source slicing - `vllm/model_executor/layers/quantization/mxfp4.py`: attach logical-width, packing, and scale-group metadata before loading No attention implementation change is currently indicated. ## Paths that already appear compatible - Generic column/row-parallel linear loaders call `pad_or_narrow_weight()` and can zero-fill attention projection tails. - `nvidia/model.py:1759+` pads shared-expert block-FP8 weights and scales. For TP3, ranks 0/1 load 768 real elements; rank 2 loads 512 real plus 256 zeros. - Vocabulary-parallel embedding uses the virtual vocabulary padding helper, yielding 43136 rows per rank and 128 global zero rows. - The MTP draft configuration receives the same virtual plan through `vllm/config/speculative.py`. This must still be tested after base-model correctness is established. ## Remaining kernel assumptions `flashinfer_cutlass_moe.py:235-240` has no odd-TP rejection and lines 427-428 pass integer `tp_size` and `tp_rank` to FlashInfer. The Python API therefore permits TP3, but this does not prove the selected CUTLASS tactic and its prepare/finalize path are correct for world size 3. The first hardware MoE test must verify: - partial-output/finalization semantics for `tp_size=3` - BF16 activation with MXFP4 weights - hidden 4096, logical local intermediate 704, kernel local intermediate 768 - 256 experts, top-k 6, and `swiglu_limit=10.0` - eager execution and CUDA graph replay Communication is a separate concern. Lucifer can remain on NCCL for TP3; B12X one-shot all-reduce support is not required to prove Lucifer virtual padding. ## Correctness test plan ### 1. Configuration tests - Construct the exact DeepSeek V4 TP3 Lucifer configuration. - Assert heads 72, output groups 9, routed intermediate 2112, shared intermediate 2304, and vocabulary 129408. - Assert the explicit FlashInfer attention/CUTLASS MoE pair is accepted. - Reject unsupported attention/MoE combinations and non-SM120 platforms. - Keep expert parallel rejected. - Assert the MTP draft config receives the identical plan. ### 2. CPU loader tests Use monotonic synthetic checkpoint tensors so every source position is visible. For all three ranks, assert exactly: - W13 offsets 0/704/1408 and zero tails 64/64/128 - packed W2 byte offsets 0/352/704 and zero tails 32/32/64 - W2 scale offsets 0/22/44 and zero tails 2/2/4 - shared-expert FP8 weights/scales, `wq_b`, `wo_a`, `wo_b`, attention sink, and vocabulary tails match their planned slices - CUTLASS post-load permutation preserves source values and zero tails Also concatenate each rank's logical slice and prove byte-for-byte equality to the original checkpoint tensor. This catches correct-looking shapes with wrong rank boundaries. ### 3. Attention hardware tests - Exercise local 24-head input padded to FlashInfer 32 in prefill and decode. - Compare the first 24 logical outputs with an unpadded reference. - Prove rank-2 virtual heads and output group 8 contribute exactly zero after `wo_a`/`wo_b`. - Exercise three local output groups through inverse RoPE and DeepGEMM. - Cover each DeepSeek compression mode used by the model. ### 4. MoE hardware tests - Compare TP3 summed output with a TP1 or dequantized reference. - Test token counts `M=1, 8, 16, 64`. - Inspect every padded weight/scale tail before kernel conversion. - Test eager, graph warmup, capture, and replay. - Reject NaN/Inf and verify zero E8M0 tails do not create visible output. ### 5. End-to-end model tests 1. Start with MTP disabled and deterministic greedy prompts. 2. Compare TP2 Lucifer and TP3 Lucifer tokens, logits, and log probabilities within an agreed numerical tolerance. 3. Run short prefill/decode tests and repeated CUDA graph replay. 4. Enable MTP2 and validate draft acceptance and final-token correctness. 5. Run the full benchmark only after all earlier checks pass. ## Go/no-go criterion Do not launch a full TP3 Lucifer model merely because it loads. Proceed only after the synthetic loader test proves the 0/352/704 packed-W2 offsets and the 0/22/44 scale offsets, and a small CUTLASS hardware test matches a reference at world size 3. Those two checks directly cover the currently unproven behavior.