dpskv3-logistic-review

name

dpskv3-logistic-review

description

Audit the DeepSeek V3 MPK demo + builder chain end-to-end and confirm logical equivalence with vLLM's reference implementation. Use after structural changes to `python/mirage/mpk/models/deepseek_v3/builder.py`, `demo/deepseek_v3/demo.py`, or any MLA / MoE / MTP task in `python/mirage/mpk/persistent_kernel.py`, `src/kernel/task_register.cc`, and `include/mirage/persistent_kernel/tasks/blackwell/mla_*.cuh` / `moe_*.cuh`. Produces a structured drift report so the change reviewer can confirm the math/topology is still equivalent.

DeepSeek V3 Logistic Review

This skill produces a drift report that compares the current MPK DeepSeek V3 demo chain against vLLM's reference implementation, flagging any places where the two diverge in math, KV-cache topology, weight layout, or scheduling. Use it whenever someone (Claude, Codex, human) edits the demo / builder / MLA / MoE / MTP path and you need to confirm the change preserves correctness.

The skill is read-only — it audits and reports, it does not edit. Apply fixes separately, then re-run the skill.

---

When to invoke

• After any commit that touches python/mirage/mpk/models/deepseek_v3/builder.py, demo/deepseek_v3/demo.py, or python/mirage/mpk/persistent_kernel.py's MLA / MoE / MTP / KV-gather wrappers.
• After kernel changes under include/mirage/persistent_kernel/tasks/blackwell/mla_*.cuh, moe_*.cuh, or linear_fp8_*.cuh.
• After updating task registrations in src/kernel/graph.cc or src/kernel/task_register.cc for MLA / MoE / MTP.
• Before merging a feature branch back to mpk.
• When a regression-test PASS leaves you uncertain whether the behavior is right (PASS only proves "ran without crash and produced a latency line"; it does not prove correctness).

This skill is not a substitute for scripts/regression_test.sh — they are complementary. Regression catches "does it run and produce output?"; this skill catches "is what it produces actually equivalent to the reference implementation?".

---

The four-pillar review structure

Every audit pass should cover four pillars in order. Skip a pillar only when the change being reviewed obviously cannot affect it.

1. Weight pipeline — demo.py weight load, FP8 absorption, per-rank shard, MTP-layer-included absorption, weight-cache load.
2. Builder topology — builder.py graph construction, tensor lifetime, KV cache layout, dual-dispatch gates.
3. Task / kernel wiring — Python wrapper → graph.cc dispatch → task_register.cc codegen → device kernel header.
4. Runtime + scheduler — task metadata setup in runtime.cc, per-iteration prepare_next_batch, AllReduce/sync-counter ordering.

Each pillar has a checklist below. Cite file:line for every claim in the report.

---

Pillar 1 — Weight pipeline

Reference: vLLM's DeepSeek V2 model loader

• vllm/model_executor/models/deepseek_v2.py — DeepseekV2ForCausalLM load + weight absorption.
• vllm/v1/spec_decode/eagle/speculator.py — MTP-specific weight loading (eh_proj, enorm, hnorm, shared lm_head).

MPK chain to audit:

• demo/deepseek_v3/demo.py lines ~600-760 (selective load + absorption):
  • absorb_layers includes num_layers (the MTP layer at idx 61) when args.mtp > 0. Double-check on every audit.
  • For each attn key in absorb_layers:
    • Original q_b_proj.weight becomes the absorbed fused 576-d Q (per absorb_kv_into_q); written back to state_dict[q_key].
    • Original Q is also split into per-head q_b_nope.weight ([H*128, q_lora], unabsorbed nope) and q_b_pe.weight ([H*64, q_lora]).
    • kv_b_proj.weight is split into kv_b_k.weight ([H*128, kv_lora]) and kv_b_v.weight ([H*128, kv_lora]).
    • The original o_proj is preserved as o_proj_original for the unabsorbed prefill output; the absorbed o_proj (post-W_UV fold-in) is what stays as o_proj.weight.
  • pad_token_id = eos_token_id fallback in qwen3 demo (orthogonal).

Audit checklist:

• Are all three forms (fused absorbed + split absorbed + split unabsorbed) produced for every layer including idx 61?
• Do q_b_nope, q_b_pe, kv_b_k, kv_b_v shapes match what the chunked-prefill kernel expects ([H * d_per_head, q_lora])? See mla_prefill_tp8_chunked_sm100.cuh for the consumer.
• Does the absorbed q_b_proj shape match what the absorbed decode + absorbed prefill kernels expect ([H*576, q_lora])?
• When args.mtp > 0, does layer 61 also get q_b_nope / q_b_pe / kv_b_k / kv_b_v? Per demo.py:644 absorb_layers.append(num_layers).
• TP sharding rules in demo.py's SHARD_RULES regex list match the corresponding rules in builder.py for any new parallel linear?
• Pre-converted model{rank}-mp{world_size}.safetensors (the DeepSeek-V3-Demo-mp8 path) is consistent with the absorb path — meaning the same weight set is produced regardless of which load path is taken?

Common drift:

• Adding a new sharded linear without updating BOTH demo.py shard rules and builder.py's SHARD_RULES.
• Forgetting to extend absorb_layers when adding a new layer that needs the absorbed split (e.g., a second MTP layer).
• Changing the absorbed Q layout in the kernel without updating the load-time absorption math.

---

Pillar 2 — Builder topology

Reference: vLLM's DeepseekV2DecoderLayer (per-layer pattern) and EagleSpeculator (MTP forward path), specifically:

• vllm/model_executor/models/deepseek_v2.py:1100-1200 — decoder layer forward.
• vllm/v1/worker/gpu/spec_decode/eagle/speculator.py:374-420 — MTP first-pass + draft loop.

MPK chain to audit (builder.py):

• __init__ and _new_intermediate_tensors (line 811): tensor allocation. Critical buffers:
  • q_a_out, q_nope_pe (fused 576-d), q_nope, q_pe (split), c_latent_out, k_pe_out, ckv_sep, kpe_sep, prefill_k_nope, prefill_v, attn_out, attn_unabsorbed, mla_partial_o, mla_partial_lse, contiguous_kv.
  • _use_prefill = mbt > 8 — chunked-prefill gate.
  • _direct_paged_decode_kv — eligibility for skipping the dense KV gather copy (TP=1 single-request, TP=2, TP=4; TP=8 disabled).
• _build_mla_attention_layer (line 1190) — main model MLA.
• _build_mla_attention_layer_with_prefix (line 2114) — MTP MLA. Note: use_mtp_prefill_attention flag at line 2144 should be True (post-fix f526c7ab). If you find False, that's a regression.
• _build_moe_mlp (line 1660) and _build_moe_mlp_with_prefix (line 2463) — routed-experts group GEMM + topk + scatter + mul_sum_add. Watch w13 (dim=1) + w2 (dim=2) shard directions.
• _build_dense_mlp (line 1604) — gate/up + down with residual fusion.
• _build_mtp_layer (line 2654) — MTP draft loop. Step-0 input is mtp_step0_input_tokens (shifted prompt + main argmax tail); steps 1+ are autoregressive draft tokens.
• build_from_dict (line 3154) — top-level orchestration.

Audit checklist:

• _use_prefill = mbt > 8. Demo readme says >=32; trust the code.
• use_mtp_prefill_attention = True at builder.py:2144. If regressed to False, MTP's prefill attention is silently skipped → garbage hidden states.
• _direct_paged_decode_kv eligibility includes the TP variant under test. TP=8 stays on dense gather until the hang is fixed.
• Each mla_kv_gather* writes to the cache the subsequent attention task reads. Mismatched buffers = empty KV history.
• Main model uses mla_prefill_tp8_chunked with unabsorbed Q/K/V; MTP uses mla_prefill_absorbed with fused 576-d Q.
• Decode wrappers (mla_mtp_decode_tp{1,2,4,8}_layer) are registered alongside prefill so the dual-dispatch runtime gate can route by Q_LEN. Decode kernel returns early on Q_LEN > 8.
• Reduce stage (mla_mtp_decode_tp*_reduce_layer or mla_reduce_layer) is registered iff num_splits > 1. Single split uses attn_out directly.
• AllReduce + residual fusion (linear_fp8_with_residual or explicit allreduce_layer(residual=...)) — does TP > 1 external AllReduce only when fuse_residual is False?
• LM head: vocab-parallel split at TP > 1 with MTP off; full vocab when MTP on (build_from_dict:3206-3231).
• MTP draft loop: step-0 hidden_input is main_hidden_states (target's final hidden state), step-1+ is self.mtp_x from previous draft.
• mtp_build_embed_input_layer runs once per MPK iteration before the draft loop, populating shifted prompt tokens.

Common drift:

• Forgetting the runtime Q_LEN gate when adding a new prefill or decode kernel — both will run at all Q_LEN.
• Allocating a buffer for one path (e.g., split q_nope/q_pe) but the consumer kernel reads a different one (e.g., q_nope_pe fused).
• Re-using a shared partial buffer across MTP draft steps without resetting it.
• Changing _use_prefill threshold in builder without updating the runtime Q_LEN gate in the kernel wrappers.

---

Pillar 3 — Task / kernel wiring

MPK paths:

• Python wrapper: python/mirage/mpk/persistent_kernel.py. Each layer-method registers a tb_graph and calls kn_graph.register_task(tb_graph, "<task_name>", params). The wrapper signature decides what gets bound to input_ptrs[].
• Graph dispatch: src/kernel/graph.cc::Graph::register_task matches the task-name string to a TASK_* enum + register_*_task function.
• Codegen: src/kernel/task_register.cc::register_<task>_task emits the C++ snippet that runs in the per-worker dispatch switch. The snippet:
  • Reads task_desc->task_metadata.{request_id, kv_idx, merge_task_offset} set by the runtime scheduler (per runtime.cc::register_mugraph).
  • Reads runtime_config.{qo_indptr_buffer, paged_kv_indptr_buffer, paged_kv_last_page_len_buffer, request_ids, step, prompt_length} to compute current Q_LEN, S, KV pointer offsets.
  • Calls into the device function in include/mirage/persistent_kernel/tasks/blackwell/<task>.cuh.
• Device kernel: *.cuh headers under tasks/blackwell/. SM100a only — Hopper variants live under tasks/hopper/ and are separate.

Audit checklist:

• Python wrapper's tb_graph.new_input(...) count matches the tuple in graph.cc ((num_inputs, num_outputs, TASK_*, variant_id)). Mismatch → "Invalid global read" at runtime.
• If the task uses store_in_dmem=True for an output (the "MPK convention"), the tuple says (N+1, 0, ...) not (N, 1, ...).
• task_register.cc codegen reads input pointers from the right indices and computes per-request offsets correctly: qo_fp_ * (H * D) for Q, bi_ * MPK_MAX_SEQ_LENGTH * D for KV (per-request stride), etc.
• Runtime Q_LEN gate matches the kernel's intended dispatch: prefill returns when Q_LEN <= 8; decode returns when Q_LEN > 8. Off-by-one is a silent miscompare.
• runtime.cc::register_mugraph's metadata setup matches the kernel's expected mapping. e.g., TASK_MLA_KV_GATHER_SM100 uses request_id = bid.x; TASK_MLA_PREFILL_TP8_CHUNKED_SM100 uses request_id = bid.x (head), kv_idx = bid.y (q_block), merge_task_offset = bid.z (batch).
• Newly-added tasks have entries in runtime_header.h (enum) + task_register.h (declaration) + task_register.cc (impl) + graph.cc (dispatch) + runtime.cc (task_type_to_name + metadata handler if non-default) + task_header.cuh (include) + Python wrapper. Check ALL of these.
• grid_dim in the Python wrapper matches what the kernel and the metadata setup expect. Wrong grid_dim silently produces wrong pointer offsets per task.
• Dynamic SMEM per worker stays under ~205 KB on B200. New tile shape that exceeds this hits Invalid __shared__ write at runtime, not compile time.

Common drift:

• Adding a kernel that reads q_nope_pe (fused) when the wrapper passes q_nope (split), or vice versa.
• Forgetting to register the task in runtime.cc::register_mugraph's metadata switch — request_id defaults to 0 silently.
• Editing a .cuh file's TMA descriptor signature without updating the corresponding task_register.cc snippet.
• Adding a new TP variant without mirroring the warp-0 tcgen05.alloc + bar.sync 1, 128 pattern from the existing TP{2,4,8} kernels.

---

Pillar 4 — Runtime + scheduler

MPK paths:

• include/mirage/persistent_kernel/persistent_kernel.cuh — worker loop, scheduler, NVSHMEM team setup, prepare_next_batch.
• src/kernel/runtime.cc::print_task_graph — emits the megakernel test.cu, schedules tasks via dependency graph.
• python/mirage/mpk/persistent_kernel.py::compile — invokes nvcc, loads the built .so via ctypes / Python ext.

Audit checklist:

• MPK_MAX_NUM_BATCHED_TOKENS, MPK_MAX_NUM_BATCHED_REQUESTS, MPK_MAX_SEQ_LENGTH, MPK_PAGE_SIZE, MPK_MAX_NUM_PAGES compile-time constants reflect the demo CLI args.
• -rdc=true is the default for NVSHMEM builds. MPK_RDC_FALSE=1 is an escape hatch only.