recipes/docker-images/mia-tp4/docs/tp4-adapters.md

TP4 adapters

The adapters of the TP4 production line (docs/tp4.md). All are import hooks in adapter/sitecustomize.py, gated by an environment variable, off unless the variable is set, and each refuses to boot if the engine symbol it patches has drifted. engram_prefetch, wo_a_w8, draft_tau, draft_head_fp8, block_verify, folded_result_fence, verify_cap and autotune_keep were already used by the TP3 profile; they are listed for completeness.

File	Gate	What
adapter/shared_pad_k.py	DSV41_SHARED_PAD_K	Shared-expert down_proj K padded 576 → 640 with re-blocked scales, so the shape stays on the b12x MXFP8 kernel instead of the CUTLASS fallback
adapter/indexer_chunked.py	DSV41_INDEXER_CHUNKED	sglang#39187 adapted to the dev-dsv41 image backend (self.candidate_masks)
adapter/indexer_chunked_v3.py + spark_prefill_dense.py	DSV41_INDEXER_CHUNKED, SPARK_PREFILL_TP_SPLIT	sglang#39187 for the candidate_metadata backend of the dsv4.1 branch, plus the SG18 prefill TP split inside the chunked path; sitecustomize picks v1 or v3 from the stock source
adapter/engram_prefetch.py	DSV41_ENGRAM_PREFETCH	Forks a side stream after EngramHasher.forward, runs the ids copy, the row-store host callback and the row copies for every Engram layer there, and joins at the layer's gather; DSV41_ENGRAM_PREFETCH_CHECK=1 re-runs the synchronous path and counts differing gathers (0 over benches and needles)
adapter/wo_a_w8.py	DSV41_WO_A_W8	After load, builds an fp8 twin (e4m3 + one exponent per 32x32 block) of every bf16 wo_a and keeps it only when it reproduces the weight exactly; the 2–8-row verify/draft path then runs a copy of the stock _wo_a_partial that loads the twin and rebuilds the same bf16 tile in registers. Other shapes stay on the stock kernel
adapter/draft_tau.py	DSV41_DRAFT_TAU	Scales the per-request temperature of the folded DSpark draft sampler; the same tensor drives draft sampling and the draft probabilities in the rejection sampler
adapter/draft_head_fp8.py, adapter/draft_head_fp8_tp4.py	DSV41_DRAFT_HEAD_FP8, DSV41_DRAFT_HEAD_FP8_IMPL=tp4 selects the TP4 variant	When the draft attaches the target's vocab-parallel LM head, keeps an fp8 copy (e4m3 + one power-of-two scale per 32x32 block) and computes the draft's block logits from it with a Triton kernel (up to 64 rows in 16-row tiles; the TP4 variant sizes the row tile to the batch, 16/64/128 up to DSV41_DRAFT_HEAD_FP8_MAX_M, default 64, so the head is streamed once per call, with the same fp32 logits per row); the target keeps the bf16 head
adapter/block_verify.py	DSV41_BLOCK_VERIFY	Replaces AcceptSampling (sampled rows) with block verification: weights wi = min(w{i-1} p/q, 1), joint accept probabilities h_i, correction token from max(w p − q, 0); same target probabilities (temperature/top-k/top-p) and draft probabilities the stock sampler uses; ragged verify keeps the stock rule
adapter/folded_result_fence.py	DSV41_FOLDED_FENCE	Clones the six AcceptOuts tensors on the folded path so the overlapped D2H copy never reads the verify graph's persistent buffers
adapter/verify_cap.py	DSV41_VERIFY_CAP	Per-request verify length inside the fixed [bs, 6] verify: an in-graph Triton kernel copies the anchor row's router ids and weights over the rows past each request's live length (one launch per layer, 38 us per step for 43 layers), acceptance is capped through cutoff_verify_lens, and the live length comes from the draft confidence head (built in static mode, which the engine otherwise skips) as a stopping rule over earlier drafts only. 1..5 fix the verified drafts instead (for measurement); conf:a,b,c,d sets one threshold per draft 2..5; DSV41_VERIFY_CAP_LOG=<file> records the confidences, live length and accepted drafts per step (rank 0, costs a sync per step). A fit on 6,074 uncut steps found per-position thresholds worth only +0.4 % over the single 0.1, and a perfect cut worth +15.5 %: what is left is in the confidence head, not the rule
adapter/wo_a_w8.py (mid, drop)	DSV41_WO_A_W8_MID, DSV41_WO_A_W8_DROP	MID: a second Triton kernel for 9-192 verify/draft rows from the fp8 twin (split-K partials summed in a fixed order). DROP: after a dequantized copy of each twin is checked bit-identical, the bf16 weight becomes a zero-stride view of a one-element tag and every call is served from the twin; prefill dequantizes into a fresh buffer per call (a single shared buffer raced between the target and draft streams)
adapter/router_live.py	DSV41_ROUTER_LIVE	Reads sglang.kernels.ops.moe.moe_fused_gate's source, adds a live_ptr/LIVE_STRIDE argument to _router_triton_kernel so each row loads its request's anchor row when its position is past the live length, writes the patched module to a temp file (Triton needs a real file) and imports it; verify_cap calls it for target-verify gates. Any drift in the engine source fails the build at import
adapter/replicated_split.py	DSV41_REPLICATED_SPLIT	Per-instance ReplicatedLinear.forward override for layers whose prefix ends with a listed suffix, at M <= DSV41_REPLICATED_SPLIT_MAX_M (96). Tiles of 128 output rows are spread over the TP ranks (4/4/3/3 for 14), the weight and its scales (row-indexed, or the flat 128x4-swizzled MXFP8 buffer whose 128-row tiles are contiguous) are sliced, and the same quantized linear runs on the slice; outputs are all-gathered (padded to the widest slice). At the first eager call each rank compares its slice against the stock layer on the real input and three random ones, bit for bit, and the ranks agree with a MIN all-reduce: a layer is split on all ranks or on none
adapter/draft_main_proj.py	DSV41_DRAFT_MAIN_PROJ_SPLIT	The draft's project_target_hidden for up to 16 rows: the MXFP8 main_proj is read out of the layer itself (applied to identity rows, which is exact), each rank keeps its 1280 output columns as e4m3 with per-(row, 32-column) power-of-two scales (checked to reproduce the weight exactly), a Triton GEMV computes them from the bf16 activation and the columns are all-gathered
adapter/roce_gather.py	DSV41_ROCE_GATHER	Sets the RoCEnante runtime's max_gather_bytes (the SG17 overlay builds it with 0) up to the slot size, prepares the gather launcher and its padded scratch before graph capture, and routes PyNcclCommunicator.all_gather calls whose shard fits through roce.all_gather
adapter/moe_b12x_next.py	DSV41_MOE_B12X_NEXT	Routed MoE on b12x main (package b12x_next, fetched and renamed at build by scripts/fetch_runtime.sh) instead of FlashInfer's CUTLASS MXFP4 path, which cannot run the 576-wide per-rank experts of EP_SIZE=1. Accepts the width in create_weights, repacks the checkpoint bytes in place (no extra weight memory), swaps the fused MoE function, prepares one execution plan per decode/verify row count of the CUDA-graph list plus a prefill ladder up to the chunk size, and pins two plans b12x's heuristic gets wrong for 576-wide experts: dynamic at <= 8 rows (the micro plan re-reads an expert slice per token and slot) and 64-row tiles from 2048 rows (16-row tiles re-stage the weights per 16 rows). At EP1 the router's ids go straight to b12x (it skips the -1 padding ids); at EP2 an in-graph kernel maps them to local experts. Refuses to boot on any drift of the pinned b12x commit or the engine symbols it patches
adapter/hc_fused.py	DSV41_HC_FUSED	For prefill-size inputs (>= 512 rows) the hyper-connection mix statistics are accumulated in one pass over K instead of 80 partial slices plus a reduce, with the same partial sums in the same order (bit-identical); decode sizes keep the stock kernels, which run hidden on a side stream. Compares itself with the stock kernels on the first live call and falls back for good on any differing bit
adapter/prefill_sp.py	DSV41_PREFILL_SP	Prefill sequence parallel for chunks of >= 2048 rows that split evenly over TP (eager prefill only; decode, verify, the draft and CUDA graphs never enter it): attention and MoE take all-gathered inputs and return reduce-scattered outputs, and the per-row work between them (hc post/combine/norm/mix stats, the Engram wkv GEMM and gate) runs on the rank's quarter of the rows. Every row-count-dependent kernel choice keys on the full chunk (sp_logical_rows()), so each per-row kernel gives the same bits on a quarter as on the full chunk; the wkv shard GEMM is self-checked per chunk size and rank-agreed. DSV41_PREFILL_SP_EXACT=1 keeps the all-reduce (bit-identical to stock), =comm swaps only the collectives, DSV41_PREFILL_SP_DEBUG=compare runs the stock op next to each sharded one and logs any differing bit. Refuses to boot on engine drift
adapter/l2_prefetch.py	DSV41_L2_PREFETCH	Learns, per decode/verify graph shape, which MXFP8 linears and router GEMM follow each RoCE collective, and during capture forks a side-stream kernel at every collective that issues cp.async.bulk.prefetch.L2 for the first DSV41_L2_PREFETCH_MB (6) of those weights, joined at the next collective. Forwards on the NCCL path (the pre-capture autotune dummy run) are never learned; a captured shape whose plan does not match logs MISMATCH and runs without prefetch
adapter/autotune_keep.py	DSV41_AUTOTUNE_KEEP	Replaces the per-rank byte digest of the FlashInfer autotune cache gate with the load decision (as in sglang#40420) plus a per-rank launch fingerprint written next to the cache after tuning, so a configuration change never leaves one rank with a partial hit
adapter/fast_load.py	DSV41_FAST_LOAD	Checkpoint tensors this rank will copy (owned experts, no Engram tables) read eagerly by a 16-thread pread pool into 256 MiB pinned slabs and returned from safe_open as views; at EP_SIZE=1 (DSV41_FAST_LOAD_TP_SLICE=auto) only the rank's slice of each expert is kept, behind a full-shape stand-in whose narrow returns it and which refuses any other access; the model's async copies paced to a byte budget so the reads stay just ahead; the DSpark draft load opens only the mtp.* shards. Loader-only: the model still does every narrow and copy

Tests: tests/test_indexer_chunked.py and tests/test_indexer_chunked_v3.py lift the stock function out of the engine's source, drive it and the backport with deterministic fake kernels, and require bitwise-equal page_indices, raw_indices and candidate masks over six scenarios (ragged batches, empty requests, full and tail-only mask publishing, mask consumption, one row per chunk up to a single chunk). Both run inside the image build (Dockerfile runs v1, the canary images v3), together with the thinking-alias, output-cap and loop-abort tests; the canary-roce build also runs the fast-load, prefill-SP, b12x MoE (CPU) and ring RoCE tests. tests/test_fast_load_pacing.py (all images) checks the copy pacing; the in-image checkpoint test for the eager reads (bitwise equality against the stock safe_open for every dtype, draft shard filtering, memory release) needs the weights mounted and is described in docs/fast-load.md.