# TP4 adapters

The adapters of the TP4 production line ([docs/tp4.md](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 w_i = 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`.