recipes/docker-images/mia-tp4/adapter/prefill_sp.py

"""Prefill sequence parallel for DeepSeek-V4.1 at TP4 (DSV41_PREFILL_SP=comm | 1).

Today every TP rank runs every row of a prefill chunk through the row-local work (the
hyper-connection post / combine / norm / stats on the [M, 4, 5120] residual, the Engram `wkv`
GEMM and gate), and each sublayer ends with an all-reduce of [M, 5120]. Here, on the full-row
layers (0 .. late_layer_start-1 under decoder SWA bounded replay), the residual stays split by
rows: rank r owns rows [r*M/W, (r+1)*M/W).

  attention  : all-gather x (bf16) -> stock attention on all M rows, wo_b without its
               all-reduce -> reduce-scatter along rows
  MoE        : all-gather x (bf16) -> stock DeepseekV2MoE under mlp_reduce_scatter=True
               (skips the post-experts all-reduce) -> reduce-scatter along rows
  Engram     : full-row lookup (unchanged, so the prefetch still matches) -> reduce-scatter
               -> wkv GEMM + gate on the shard (wkv per (layer, M) bit-exact self-check,
               agreed over TP with MIN, full-M fallback)
  hc post / combine / norm / stats: stock code on the shard; every kernel choice that the
               engine keys on the row count is keyed on the FULL chunk row count
               (sp_logical_rows(), used by _hc_combine / hc_post here, by hc_fused and by
               hc_prefill_fused)
  layer 21   : the tail rows of R and prev_pre are gathered (byte copy) and the late layers
               run stock
  no tail    : R and the last pre are gathered at exit

Modes
  DSV41_PREFILL_SP=comm  P0, transport only: nothing is sharded; the same three reduction sites
                         (wo_b, MoE output, Engram lookup) run reduce-scatter + all-gather
                         instead of the all-reduce. Same numerics as stage 1 fast by
                         construction (same reduce-scatters on the same tensors).
  DSV41_PREFILL_SP=1     stage 1, bf16 gathers.
  DSV41_PREFILL_SP_EXACT=1  (stage 1) every reduce-scatter is all-reduce + keep the quarter:
                         must reproduce production bit for bit.
  DSV41_PREFILL_SP_MIN_ROWS  default 2048; chunks below it or not divisible by the TP size run
                         the stock path in every mode (same predicate).
  DSV41_PREFILL_SP_WKV=full  never run wkv on the shard (gather the lookup, full-M GEMM).

Scope: extend without speculation only (plain prefill chunks), eager (never under graph
capture or torch.compile), TP only (no DP attention, no CP, no a2a MoE, PP 1). Decode,
verify, draft and CUDA graphs never enter this code. Every decision depends only on shape,
forward mode and static config, so it is identical on every rank.

Drift guard: sha256[:16] of the engine sources this relies on are checked at install; any
difference raises (refuses to boot). hc_prefill_fused, when enabled, must key its gates on
sp_logical_rows() (it declares SP_LOGICAL_ROWS = True); otherwise install raises.
"""
import hashlib
import inspect
import logging
import os
import threading
from contextlib import nullcontext

import torch

logger = logging.getLogger(__name__)

_RAW = os.environ.get("DSV41_PREFILL_SP", "0").strip().lower()
MODE = "comm" if _RAW == "comm" else ("shard" if _RAW in ("1", "on", "true", "shard") else "off")
EXACT = os.environ.get("DSV41_PREFILL_SP_EXACT", "0").strip().lower() in ("1", "on", "true")
MIN_ROWS = int(os.environ.get("DSV41_PREFILL_SP_MIN_ROWS", "2048"))
WKV_SHARD = os.environ.get("DSV41_PREFILL_SP_WKV", "shard").strip().lower() != "full"
# Stage 2 (DSV41_PREFILL_SP_FP8=1, with DSV41_PREFILL_SP=1): on full-row layers whose attention
# reads its input only through wqkv_a (no compressor / indexer: all but the kv/index sources 2, 8,
# 14, 20), each rank MXFP8-quantizes its shard (per row, per 32 columns; linear scale layout),
# all-gathers the fp8 bytes + ue8m0 scales, swizzles the scales to 128x4 and hands attention the
# pre-quantized input (x_quant) instead of the bf16 rows. Exact by construction if the quantizer is
# row-local; checked, not assumed: the first chunk of every size compares the gathered MXFP8 input
# with the stock quantization of the gathered bf16 rows AND wqkv_a's output on both, agrees over
# TP (MIN), and keeps the bf16 gather for that size on any difference.
FP8 = os.environ.get("DSV41_PREFILL_SP_FP8", "0").strip().lower() in ("1", "on", "true")
# Debug (never in production): DSV41_PREFILL_SP_DEBUG=compare,fp (either or both).
#   fp       per traced prefill chunk, rank 0 logs the chunk layout and a checksum of every
#            layer's input / attention in+out / MoE in+out / output residual and pre, over the full
#            rows. Works with DSV41_PREFILL_SP unset too (then only the tracing is installed), so an
#            off boot and an SP boot can be diffed chunk by chunk.
#   compare  (stage 1) next to every sharded op, the stock op runs on the same (gathered) inputs
#            and every rank logs each op whose bits differ: layer, op, differing elements, max abs.
#            wo_b and the MoE also run their stock all-reduce path on the same partials.
# DSV41_PREFILL_SP_DEBUG_CHUNKS  traced chunks (default 64, every extend forward counts)
# DSV41_PREFILL_SP_DEBUG_DUMP    directory: torch.save of the first mismatch per op and rank
DEBUG = {t.strip() for t in os.environ.get("DSV41_PREFILL_SP_DEBUG", "").lower().split(",") if t.strip()}
if "compare" in DEBUG:
    DEBUG.add("fp")
DEBUG_CHUNKS = int(os.environ.get("DSV41_PREFILL_SP_DEBUG_CHUNKS", "64"))
DEBUG_DUMP = os.environ.get("DSV41_PREFILL_SP_DEBUG_DUMP", "").strip()

# Engine sources (sglang dsv4.1 f80c91a4b as in dsv41-4x-spark:ds41-verify-0924d), sha256[:16]
# of inspect.getsource(inspect.unwrap(obj)). The loop below is a copy of the first; the others
# are the call sites and kernel gates whose behaviour the exactness argument depends on.
_EXPECTED = {
    "v4.DeepseekV4Model._forward_layers_hc_pre_from_prev": "bc059d80f038e7f8",
    "v4.DeepseekV4DecoderLayer.forward_hc_pre_from_prev": "49dc092d369923cf",
    "v4.DeepseekV4DecoderLayer._hc_combine": "da5688eff1cf6fab",
    "v4.DeepseekV4DecoderLayer._hc_mix_stats": "0a6edf4b45cd160b",
    "v4.DeepseekV4DecoderLayer.hc_post": "8adb8062e8e66c31",
    "v4.DeepseekV4DecoderLayer._run_moe_ffn_dp_sync": "065622659105f3fe",
    "v4.DeepseekV4Model.forward": "3ab2a4745377b83b",
    "v4.MQALayer.forward": "d21dfc6533e96207",
    "v2.DeepseekV2MoE.forward": "d931dfa4d2acacd8",
    "v2.DeepseekV2MoE.forward_normal": "9da361af16dafbec",
    "engram.Engram.forward": "841b358015e044b8",
    "engram.EngramEmbedding.forward": "b2f3ea98879b9437",
    "engram.EngramEmbedding._lookup": "66a0f98a49fcd266",
    "hcn.hc_combine_norm": "afa73fe49eb1d4b0",
    "hcn._hc_combine_norm_prefill": "74f71ebf85b010ee",
    "linear.RowParallelLinear.forward": "30be12725b32806b",
}


# sglang.srt.layers.attention.deepseek_v4_backend: the tail-row semantics _tail_rows reproduces.
# Stage 2 additionally relies on these (checked at install when DSV41_PREFILL_SP_FP8=1): which
# inputs attention's prefill path reads, and when it takes a pre-quantized input.
_EXPECTED_FP8 = {
    "v4.MQALayer._forward_prepare": "b3adfe1a07ee6127",
    "v4.MQALayer.accepts_mxfp8_swizzled_input": "5bb4fddfc5b329ca",
    "v4.MQALayer._compute_kv_to_cache": "f52f600c4986428c",
    "v4.MQALayer._compute_kv_bf16": "00c4ab90b403d1e9",
}
_EXPECTED_TAIL = {
    "LateLayerTail.rows": "c58392d3ef651408",
    "LateLayerTail.real_rows": "c78962c2e01e0e3d",
    "_tail_rows": "54a3cbff5e674e84",
}


class _Ctx(threading.local):
    def __init__(self):
        self.plan = None        # _Plan while the layer loop is inside the full-row region
        self.inner = 0          # > 0 inside a gathered attention / MoE call
        self.skip_wo_b = None   # the wo_b whose all-reduce the current attention call skips
        self.dbg = None         # _Dbg while a traced chunk runs


_ctx = _Ctx()
_M = {}                         # engine modules, filled by install
_STATIC = {"checked": False}
_WKV_OK = {}                    # (layer_id, M) -> bool (agreed over TP)
_SEEN_M = set()
_REQUIRE_CUDA = True            # CPU tests switch this off
_DBG = {"chunks": 0, "dumped": 0}


def _agree_min(p, local):
    """Rank-agreed AND of a local boolean (MIN over the TP group)."""
    flag = torch.tensor([1 if local else 0], dtype=torch.int32, device=torch.cuda.current_device())
    torch.distributed.all_reduce(flag, op=torch.distributed.ReduceOp.MIN, group=p.group.device_group)
    return bool(flag.item())


def _mega_moe(moe, x):
    try:
        from sglang.srt.layers.moe.mega_moe import should_use_mega_moe
    except ImportError:
        return False
    return should_use_mega_moe(moe, x)


class _Dbg:
    def __init__(self, idx, rows, plan, rank):
        self.idx, self.rows, self.plan, self.rank = idx, rows, plan, rank
        self.compare = "compare" in DEBUG and plan is not None and plan.sharded
        self.layer = -1
        self.fps = []           # (layer, key, fp)
        self.checked = 0
        self.bad = []           # (layer, op, ndiff, numel, maxabs)
        self.dumped = set()


def _fp(t):
    """Order-sensitive checksum of a tensor's bytes (GPU, chunked int64 sums)."""
    if t is None:
        return "none"
    t = t.contiguous().reshape(-1)
    nbytes = t.numel() * t.element_size()
    v = t.view(torch.int32) if nbytes % 4 == 0 else t.view(torch.uint8)
    s1 = torch.zeros((), dtype=torch.int64, device=v.device)
    s2 = torch.zeros((), dtype=torch.int64, device=v.device)
    step = 1 << 22
    for i in range(0, v.numel(), step):
        c = v[i:i + step].to(torch.int64)
        w = torch.arange(i, i + c.numel(), device=v.device, dtype=torch.int64) % 65521 + 1
        s1 += c.sum()
        s2 += (c * w).sum()
    return f"{(int(s1) * 1000003 + int(s2)) & ((1 << 48) - 1):012x}"


def _full_rows(t):
    """t over every row of the chunk: gathered when it is this rank's shard (collective)."""
    p = _ctx.plan
    if (t is not None and p is not None and p.sharded and isinstance(t, torch.Tensor)
            and t.dim() >= 1 and t.shape[0] == p.shard and p.shard != p.rows):
        return _gather_rows(p, t)
    return t


def _rec(key, t, full=False):
    d = _ctx.dbg
    if d is None:
        return
    d.fps.append((d.layer, key, _fp(t if full else _full_rows(t))))


class _Inner:
    def __enter__(self):
        _ctx.inner += 1

    def __exit__(self, *a):
        _ctx.inner -= 1


def _comparing():
    d = _ctx.dbg
    return d is not None and d.compare and _ctx.plan is not None and _ctx.plan.sharded


def _cmp(op, ours, ref, inputs=None):
    """Bits of our result vs the stock op's on the same inputs; log and dump differences."""
    d = _ctx.dbg
    if d is None:
        return
    outs = ours if isinstance(ours, (tuple, list)) else (ours,)
    refs = ref if isinstance(ref, (tuple, list)) else (ref,)
    d.checked += 1
    for k, (a, b) in enumerate(zip(outs, refs)):
        name = op if len(outs) == 1 else f"{op}[{k}]"
        if a.shape == b.shape and a.dtype == b.dtype and torch.equal(a, b):
            continue
        if a.shape != b.shape:
            nd, mx = -1, float("nan")
        else:
            diff = (a.float() - b.float()).abs()
            nd, mx = int((a != b).sum()), float(diff.max())
        d.bad.append((d.layer, name, nd, a.numel(), mx))
        if len(d.bad) <= 24:
            print(f"DSV41_PREFILL_SP compare MISMATCH chunk={d.idx} M={d.rows} rank={d.rank} "
                  f"layer={d.layer} op={name} shape={tuple(a.shape)}/{tuple(b.shape)} "
                  f"differing={nd}/{a.numel()} max_abs={mx:.4g}", flush=True)
        if DEBUG_DUMP and op not in d.dumped and _DBG["dumped"] < 16:
            d.dumped.add(op)
            _DBG["dumped"] += 1
            try:
                os.makedirs(DEBUG_DUMP, exist_ok=True)
                path = os.path.join(DEBUG_DUMP, f"psp_c{d.idx}_r{d.rank}_L{d.layer}_{op.replace(' ', '_')}.pt")
                blob = {"op": name, "layer": d.layer, "chunk": d.idx, "rows": d.rows, "rank": d.rank,
                        "ours": a.detach().cpu(), "ref": b.detach().cpu()}
                if inputs:
                    blob["inputs"] = {n: (v.detach().cpu() if isinstance(v, torch.Tensor) else v)
                                      for n, v in inputs.items()}
                torch.save(blob, path)
                print(f"DSV41_PREFILL_SP compare dumped {path}", flush=True)
            except Exception as exc:  # noqa: BLE001
                print(f"DSV41_PREFILL_SP compare dump failed: {exc!r}", flush=True)


def _debug_begin(model, hidden_states, forward_batch, input_ids, p):
    if not DEBUG or _DBG["chunks"] >= DEBUG_CHUNKS:
        return None
    if torch.compiler.is_compiling() or (torch.cuda.is_available()
                                         and torch.cuda.is_current_stream_capturing()):
        return None
    if not forward_batch.forward_mode.is_extend_without_speculative():
        return None
    _DBG["chunks"] += 1
    rank = p.rank if p is not None else _M["v4"].get_tp_group().rank_in_group
    d = _Dbg(_DBG["chunks"], hidden_states.shape[0], p, rank)
    if rank == 0:
        lens = getattr(forward_batch, "extend_seq_lens_cpu", None)
        pref = getattr(forward_batch, "extend_prefix_lens_cpu", None)
        mode = "off" if p is None else (MODE + ("-exact" if p.exact else "") if p.sharded else MODE)
        print(f"DSV41_PREFILL_SP fp chunk={d.idx} M={d.rows} mode={mode} reqs={len(lens) if lens is not None else '?'} "
              f"extend={list(lens)[:16] if lens is not None else '?'} prefix={list(pref)[:16] if pref is not None else '?'} "
              f"ids={_fp(input_ids)} R0={_fp(hidden_states)}", flush=True)
    return d


def _debug_end(d, out):
    try:
        hs, pre = out[0], out[1]
        fh, fp_ = _fp(hs), _fp(pre)
    except Exception as exc:  # noqa: BLE001
        fh = fp_ = f"err {exc!r}"
    if d.rank == 0:
        by_layer = {}
        for layer, key, f in d.fps:
            by_layer.setdefault(layer, []).append(f"{key}={f}")
        for layer, items in by_layer.items():
            print(f"DSV41_PREFILL_SP fp chunk={d.idx} L={layer} " + " ".join(items), flush=True)
        print(f"DSV41_PREFILL_SP fp chunk={d.idx} out R={fh} pre={fp_}", flush=True)
    if d.compare:
        first = f" first: layer {d.bad[0][0]} {d.bad[0][1]}" if d.bad else ""
        print(f"DSV41_PREFILL_SP compare chunk={d.idx} M={d.rows} rank={d.rank}: {d.checked} ops checked, "
              f"{len(d.bad)} mismatched{first}", flush=True)


def _make_layer_forward(orig):
    """Tracing only: record each layer's input / output residual and pre."""
    def forward_hc_pre_from_prev(self, *args, **kwargs):
        d = _ctx.dbg
        if d is None or _ctx.inner:
            return orig(self, *args, **kwargs)
        hs = kwargs["hidden_states"] if "hidden_states" in kwargs else args[1]
        d.layer = getattr(self, "layer_id", d.layer + 1)
        _rec("Rin", hs)
        out = orig(self, *args, **kwargs)
        _rec("Rout", out[0])
        _rec("pre", out[1])
        return out

    forward_hc_pre_from_prev.__wrapped__ = orig
    return forward_hc_pre_from_prev


class _Plan:
    __slots__ = ("rows", "world", "rank", "shard", "lo", "hi", "group", "sharded", "exact")

    def __init__(self, rows, world, rank, group, sharded, exact):
        self.rows, self.world, self.rank, self.group = rows, world, rank, group
        self.shard = rows // world
        self.lo, self.hi = rank * self.shard, (rank + 1) * self.shard
        self.sharded, self.exact = sharded, exact

    @property
    def in_rows(self):
        return self.shard if self.sharded else self.rows


# ------------------------------------------------------------------------------------------
# the logical-row helper (kernel gates key on the full chunk row count)
# ------------------------------------------------------------------------------------------
def sp_logical_rows(x):
    """Row count a kernel gate must use for `x`: the full chunk row count M when `x` is this
    rank's shard of a sequence-parallel chunk, else x.shape[0]. Cheap and safe to call when
    the adapter is off."""
    n = x.shape[0]
    p = _ctx.plan
    if p is None or _ctx.inner or not p.sharded:
        return n
    return p.rows if n == p.shard else n


def enabled():
    return MODE != "off"


# ------------------------------------------------------------------------------------------
# row partition (pure; unit-tested on CPU)
# ------------------------------------------------------------------------------------------
def eligible_rows(rows, world, min_rows=None):
    """Shape predicate shared by every mode: rows >= MIN_ROWS and divisible by the TP size."""
    min_rows = MIN_ROWS if min_rows is None else min_rows
    return world > 1 and rows >= min_rows and rows % world == 0


def shard_range(rows, world, rank):
    s = rows // world
    return rank * s, (rank + 1) * s


def tail_plan(rows, world, tail_len):
    """'full' when gathering all rows is no more traffic than gathering world * T rows."""
    return "full" if world * tail_len >= rows else "owners"


def owner_local(idx, shard):
    owner = torch.div(idx, shard, rounding_mode="floor")
    return owner, idx - owner * shard


# ------------------------------------------------------------------------------------------
# collectives (the group is a sglang GroupCoordinator; tests pass a fake one)
# ------------------------------------------------------------------------------------------
def _gather_rows(p, x):
    """[S, ...] shard on every rank -> [M, ...] (byte copy)."""
    x = x.contiguous()
    out = x.new_empty((p.rows, *x.shape[1:]))
    p.group.all_gather_into_tensor(out, x)
    return out


def _reduce_rows(p, x):
    """[M, ...] per-rank partial -> this rank's reduced [S, ...] rows."""
    x = x.contiguous()
    if p.exact:
        full = p.group.all_reduce(x)
        return full[p.lo:p.hi]
    out = x.new_empty((p.shard, *x.shape[1:]))
    p.group.reduce_scatter_tensor(out, x)
    return out


def _reduce_out(p, x):
    """Output of a reduction site: this rank's rows (stage 1) or every row (P0: RS + AG)."""
    if p.sharded:
        return _reduce_rows(p, x)
    x = x.contiguous()
    part = x.new_empty((p.shard, *x.shape[1:]))
    p.group.reduce_scatter_tensor(part, x)
    return _gather_rows(p, part)


def _tail_real_rows(p, tail, t):
    """tail.real_rows(concat of every rank's shard t), as a byte copy."""
    if tail.contiguous_start is not None:
        start, idx = int(tail.contiguous_start), None
        n_tail = p.rows - start
    else:
        start, idx = None, tail.token_indices
        n_tail = int(idx.shape[0])
    if tail_plan(p.rows, p.world, n_tail) == "full" or n_tail == 0:
        full = _gather_rows(p, t)
        return full[start:] if idx is None else full[idx]
    if idx is None:
        idx = torch.arange(start, p.rows, device=t.device)
    owner, local = owner_local(idx.to(torch.int64), p.shard)
    mine = t.contiguous()[local.clamp(0, p.shard - 1)].contiguous()
    buf = t.new_empty((p.world * n_tail, *t.shape[1:]))
    p.group.all_gather_into_tensor(buf, mine)
    buf = buf.view(p.world, n_tail, *t.shape[1:])
    return buf[owner, torch.arange(n_tail, device=t.device)]


def _tail_rows(p, tail, t):
    rows = _tail_real_rows(p, tail, t)
    if tail.pad_rows:
        rows = torch.cat([rows, rows.new_zeros((tail.pad_rows, *rows.shape[1:]))])
    return rows


# ------------------------------------------------------------------------------------------
# drift guard
# ------------------------------------------------------------------------------------------
def _src_hash(obj):
    obj = getattr(obj, "fn", obj)                       # triton JITFunction
    obj = inspect.unwrap(obj)
    return hashlib.sha256(inspect.getsource(obj).encode()).hexdigest()[:16]


def _resolve(name):
    mod, attr = name.split(".", 1)
    obj = _M[mod]
    for part in attr.split("."):
        obj = getattr(obj, part)
    return obj


def engine_hashes():
    return {k: _src_hash(_resolve(k)) for k in _EXPECTED}


def check_engine():
    exp = dict(_EXPECTED)
    if FP8:
        exp.update(_EXPECTED_FP8)
    got = {k: _src_hash(_resolve(k)) for k in exp}
    bad = [f"{k} {got[k]} != {exp[k]}" for k in got if exp[k] != got[k]]
    if bad:
        raise RuntimeError("DSV41_PREFILL_SP: engine source drifted, refusing to run: " + "; ".join(bad))


def _load_modules(model_mod):
    from sglang.kernels.ops.layernorm import hc_combine_norm as hcn
    from sglang.kernels.ops.layernorm import mhc
    from sglang.srt.layers import engram, linear
    from sglang.srt.models import deepseek_v2 as v2

    _M.update(v4=model_mod, v2=v2, engram=engram, hcn=hcn, mhc=mhc, linear=linear)


def _check_hc_prefill_fused():
    if os.environ.get("DSV41_HC_PREFILL_FUSED", "0").strip() in ("0", "off", "false", ""):
        return
    import hc_prefill_fused as hpf
    if not getattr(hpf, "SP_LOGICAL_ROWS", False):
        raise RuntimeError(
            "DSV41_PREFILL_SP=1 with DSV41_HC_PREFILL_FUSED=1 needs hc_prefill_fused to key its "
            "row gates on prefill_sp.sp_logical_rows() (and declare SP_LOGICAL_ROWS = True); "
            "refusing to run. See diagnostics/dsv41-prefill-sp/IMPL.md.")


# ------------------------------------------------------------------------------------------
# plan (rank-invariant: shape, forward mode, static config)
# ------------------------------------------------------------------------------------------
def _check_tail_engine():
    from sglang.srt.layers.attention import deepseek_v4_backend as b
    bad = []
    for k, want in _EXPECTED_TAIL.items():
        obj = b
        for part in k.split("."):
            obj = getattr(obj, part)
        got = _src_hash(obj)
        if got != want:
            bad.append(f"deepseek_v4_backend.{k} {got} != {want}")
    if bad:
        raise RuntimeError("DSV41_PREFILL_SP: engine source drifted, refusing to run: " + "; ".join(bad))


def _static_check(model, world):
    v4 = _M["v4"]
    if not _STATIC.get("tail_checked"):
        _check_tail_engine()
        _STATIC["tail_checked"] = True
    par = v4.get_parallel()
    errs = []
    if model.pp_group.world_size != 1:
        errs.append("PP > 1")
    if par.attn_dp_size != 1:
        errs.append(f"attn_dp_size {par.attn_dp_size}")
    if par.tp_size != world or par.attn_tp_size != world:
        errs.append(f"tp {par.tp_size} / attn_tp {par.attn_tp_size} / group {world}")
    if not v4.get_moe_a2a_backend().is_none():
        errs.append("a2a MoE backend")
    if not model.hc_pre_from_prev_sublayer or model.hc_mult != 4:
        errs.append("hc layout")
    if v4.get_platform().is_sm100:
        errs.append("sm100 (_hc_mix_stats bf16x3/deepgemm gates are row-count keyed)")
    if MODE == "comm" and v4.envs.SGLANG_ENABLE_DETERMINISTIC_INFERENCE.get():
        errs.append("deterministic inference (reduce_scatter_tensor is all_reduce there)")
    if MIN_ROWS < 256:
        errs.append("DSV41_PREFILL_SP_MIN_ROWS below 256 would reach decode/verify sizes")
    last = model.late_layer_start if model.late_layer_start is not None else model.end_layer
    for i in range(model.start_layer, last):
        layer = model.layers[i]
        attn, mlp = layer.self_attn, layer.mlp
        if not isinstance(attn, v4.MQALayer) or not isinstance(mlp, _M["v2"].DeepseekV2MoE):
            errs.append(f"layer {i}: {type(attn).__name__}/{type(mlp).__name__}")
            continue
        rpl = getattr(_M.get("linear"), "RowParallelLinear", None)
        if rpl is not None and not isinstance(attn.wo_b, rpl):
            errs.append(f"layer {i}: wo_b is {type(attn.wo_b).__name__}")
        if not (attn.wo_b.reduce_results and attn.wo_b.tp_size == world and attn.attn_tp_size == world):
            errs.append(f"layer {i}: wo_b reduce/tp")
        if getattr(mlp, "_shared_expert_tp1", False) or getattr(mlp, "_enable_a2a_moe", False):
            errs.append(f"layer {i}: TP1 shared expert or a2a MoE")
        if mlp.tp_size != world:
            errs.append(f"layer {i}: moe tp {mlp.tp_size}")
        if getattr(layer, "dsa_enable_prefill_cp", False) or layer.use_fused_mhc_post_pre:
            errs.append(f"layer {i}: prefill CP / fused post-pre")
        if layer.engram is not None and layer.engram.embed.tp_size != world:
            errs.append(f"layer {i}: engram tp")
    if errs:
        raise RuntimeError("DSV41_PREFILL_SP: unsupported configuration, refusing to run: " + "; ".join(errs))
    if FP8 and MODE == "shard" and world > 1:
        fp8_layers = [i for i in range(model.start_layer, last)
                      if fp8_eligible_attn(model.layers[i].self_attn)]
        if _M["v4"].get_tp_group().rank_in_group == 0:
            print(f"DSV41_PREFILL_SP fp8 attention gather on layers {fp8_layers} "
                  f"({len(fp8_layers)} of {last - model.start_layer}; the rest gather bf16)", flush=True)
    _STATIC["checked"] = True


def _plan(model, hidden_states, forward_batch):
    if MODE == "off":
        return None
    v4 = _M["v4"]
    if torch.compiler.is_compiling() or (torch.cuda.is_available()
                                         and torch.cuda.is_current_stream_capturing()):
        return None
    if not forward_batch.forward_mode.is_extend_without_speculative():
        return None
    if v4.is_in_breakable_cuda_graph() or v4.is_cp_active(forward_batch):
        return None
    if v4.get_attn_tp_context().input_scattered:
        return None
    if not (hidden_states.dim() == 3 and hidden_states.dtype == torch.bfloat16
            and (hidden_states.is_cuda or not _REQUIRE_CUDA) and hidden_states.is_contiguous()):
        return None
    group = v4.get_tp_group()
    rows, world = hidden_states.shape[0], group.world_size
    if not eligible_rows(rows, world):
        return None
    if not _STATIC["checked"]:
        _static_check(model, world)
    sharded = MODE == "shard"
    p = _Plan(rows, world, group.rank_in_group, group, sharded, EXACT and sharded)
    if rows not in _SEEN_M:
        _SEEN_M.add(rows)
        if p.rank == 0 and len(_SEEN_M) <= 8:
            print(f"DSV41_PREFILL_SP {MODE}{' exact' if p.exact else ''}: chunk of {rows} rows -> "
                  f"{'shards of ' + str(p.shard) if sharded else 'reduce-scatter + all-gather'}",
                  flush=True)
    return p


# ------------------------------------------------------------------------------------------
# Engram on the shard
# ------------------------------------------------------------------------------------------
def _wkv_shard(p, engram, emb):
    """kv rows [lo, hi) from the shard's lookup rows, bit-identical to the full-M GEMM."""
    if not WKV_SHARD:
        return engram.wkv(_gather_rows(p, emb).flatten(-2))[0][p.lo:p.hi]
    key = (engram.layer_hash_index, p.rows)
    ok = _WKV_OK.get(key)
    if ok is None:
        full = engram.wkv(_gather_rows(p, emb).flatten(-2))[0]
        mine = full[p.lo:p.hi].clone()
        del full                # the [M, 25600] output is not held across the agreement
        part = engram.wkv(emb.flatten(-2))[0]
        local = bool(torch.equal(mine, part))
        ok = _WKV_OK[key] = _agree_min(p, local)
        if p.rank == 0 and (len(_WKV_OK) <= 4 or not ok):
            print(f"DSV41_PREFILL_SP wkv (engram {engram.layer_hash_index}, M={p.rows}): shard GEMM "
                  f"{'bit-exact on every rank -> ON' if ok else 'differs -> full-M GEMM'} "
                  f"(rank 0 local {local})", flush=True)
        return part if ok else mine
    if ok:
        return engram.wkv(emb.flatten(-2))[0]
    return engram.wkv(_gather_rows(p, emb).flatten(-2))[0][p.lo:p.hi]


def _engram_lookup_partial(engram, ids):
    embed = engram.embed
    if embed._shared:           # host-shared table: every rank already has every row
        return embed(ids), False
    return embed._owned_rows(ids), True


def _sp_engram(p, engram, x, ids, forward_batch=None):
    """Stock Engram.forward on a shard of x: lookup on all rows, reduce, wkv + gate on rows."""
    emb, partial = _engram_lookup_partial(engram, ids)
    if p.sharded:
        emb = _reduce_rows(p, emb) if partial else emb[p.lo:p.hi]
        kv = _wkv_shard(p, engram, emb)
        out = engram.apply_gate(x, kv.contiguous())
        if _comparing():
            with _Inner():
                xg = _gather_rows(p, x)
                # slices are cloned at once: the [M, ...] references are not held across collectives
                ref_kv = engram.wkv(_gather_rows(p, emb).flatten(-2))[0][p.lo:p.hi].clone()
                ref = engram(xg, ids, forward_batch, cp_all_tokens=False)[p.lo:p.hi].clone()
                del xg
            _cmp("engram wkv", kv, ref_kv)
            _cmp("engram out", out, ref, {"x": x, "kv": kv})
        return out
    if partial:                 # P0: reduce-scatter + all-gather in place of the all-reduce
        emb = _reduce_out(p, emb)
    kv, _ = engram.wkv(emb.flatten(-2))
    return engram.apply_gate(x, kv)


# ------------------------------------------------------------------------------------------
# the layer loop (copy of DeepseekV4Model._forward_layers_hc_pre_from_prev, hash-pinned)
# ------------------------------------------------------------------------------------------
def _sp_forward_layers(self, p, positions, hidden_states, forward_batch, input_ids,
                       input_ids_global, capture_dspark, dspark_aux_hidden_states):
    v4 = _M["v4"]
    assert self.pp_group.world_size == 1, "pre-mix hand-off across PP is not wired"
    hash_ids = None
    if self.engram_hasher is not None:
        hash_ids = self.engram_hasher(input_ids, forward_batch)
    tail = None
    attn_backend = None
    if self.late_layer_start is not None and forward_batch.forward_mode.is_extend_without_speculative():
        self._check_late_layer_tail_readers(forward_batch)
        attn_backend = v4.get_attn_backend()
        tail = attn_backend.tail_forward_metadata.late_layer_tail
    in_region = True
    R = hidden_states[p.lo:p.hi] if p.sharded else hidden_states
    ids_rows = None
    saved_full = None
    prev_pre = None
    precomputed_attn = None
    _ctx.plan = p
    try:
        for i in range(self.start_layer, self.end_layer):
            if tail is not None and i == self.late_layer_start:
                saved_full = attn_backend.enter_late_layer_tail(forward_batch)
                _ctx.plan = None
                in_region = False
                if p.sharded:
                    R_sh, pre_sh = R, prev_pre
                    R, prev_pre = _tail_rows(p, tail, R), _tail_rows(p, tail, prev_pre)
                    if _ctx.dbg is not None and _ctx.dbg.compare:
                        _cmp("tail rows R", R, tail.rows(_gather_rows(p, R_sh)))
                        _cmp("tail rows pre", prev_pre, tail.rows(_gather_rows(p, pre_sh)))
                    del R_sh, pre_sh
                else:
                    R, prev_pre = tail.rows(R), tail.rows(prev_pre)
                input_ids, input_ids_global = tail.rows(input_ids), tail.rows(input_ids_global)
                positions = tail.positions
                if hash_ids is not None:
                    hash_ids = tail.rows(hash_ids)
            engram = self.layers[i].engram
            if engram is not None:
                precomputed_attn = None
                before_engram = R
                ids = hash_ids[:, engram.layer_hash_index]
                if in_region:
                    R = _sp_engram(p, engram, R, ids, forward_batch)
                else:
                    R = engram(R, ids, forward_batch, cp_all_tokens=False)
                if self.config.model_type == "deepseek_v41" and self.config.vision_n_layers > 0:
                    if in_region and p.sharded:
                        if ids_rows is None:
                            ids_rows = input_ids[p.lo:p.hi]
                        vis = ids_rows
                    else:
                        vis = input_ids
                    R = torch.where((vis == self.config.image_token_id)[:, None, None],
                                    before_engram, R)
            if capture_dspark and i in self.dspark_layers_to_capture:
                aux = R
                if in_region and p.sharded:
                    aux = _tail_rows(p, tail, R) if tail is not None else _gather_rows(p, R)
                elif tail is not None and i < self.late_layer_start:
                    aux = tail.rows(aux)
                dspark_aux_hidden_states.append(aux.mean(dim=1))
            ctx = (
                nullcontext()
                if v4.check_cuda_graph_backend(v4.Phase.PREFILL, v4.Backend.TC_PIECEWISE)
                else v4.get_global_expert_distribution_recorder().with_current_layer(i)
            )
            # next_norm / next_input only serve decode and verify (<= 8 rows): never here.
            next_input = []
            with ctx:
                R, prev_pre = self.layers[i].forward_hc_pre_from_prev(
                    positions=positions,
                    hidden_states=R,
                    input_ids=input_ids,
                    forward_batch=forward_batch,
                    input_ids_global=input_ids_global,
                    prev_pre=prev_pre,
                    precomputed_attn=precomputed_attn,
                    next_norm=None,
                    next_input=next_input,
                )
            precomputed_attn = next_input[0] if next_input else None
    finally:
        _ctx.plan = None
    if saved_full is not None:
        attn_backend.exit_late_layer_tail(saved_full, forward_batch)
        return R, prev_pre, tail
    if p.sharded:
        R, prev_pre = _gather_rows(p, R), _gather_rows(p, prev_pre)
    return R, prev_pre, None


# ------------------------------------------------------------------------------------------
# stage 2: MXFP8 attention-input gather
# ------------------------------------------------------------------------------------------
_FP8_M = {}                     # M -> (ok, scale shape, scale dtype) agreed over TP


def fp8_eligible_attn(attn):
    """Attention whose prefill path reads x only through wqkv_a (plus shape / dtype)."""
    v = getattr(attn, "_dsv41_sp_fp8", None)
    if v is None:
        try:
            v = bool(getattr(attn, "compressor", None) is None
                     and getattr(attn, "indexer", None) is None
                     and getattr(attn, "fuse_wqa_wkv", False)
                     and getattr(attn, "wqkv_a", None) is not None
                     and attn.accepts_mxfp8_swizzled_input())
        except Exception:  # noqa: BLE001
            v = False
        attn._dsv41_sp_fp8 = v
    return v


def _mxfp8_quantize(x, swizzled):
    from sglang.srt.layers.quantization.fp8_utils import flashinfer_mxfp8_quantize
    return flashinfer_mxfp8_quantize(x, is_sf_swizzled_layout=swizzled, alignment=32)


def swizzle_128x4(s_lin):
    """[M, G] linear ue8m0 scales -> the 128x4 tiled layout (rows padded to 128 with zeros):
    tile (m // 128, g // 4) in row-tile-major order, inside a tile [m % 32][(m % 128) // 32][g % 4]."""
    m, g = s_lin.shape
    assert g % 4 == 0
    mt = (m + 127) // 128
    if mt * 128 != m:
        s_lin = torch.cat([s_lin, s_lin.new_zeros((mt * 128 - m, g))])
    return s_lin.view(mt, 4, 32, g // 4, 4).permute(0, 3, 2, 1, 4).contiguous().view(-1)


def _mxfp8_input_cls():
    from sglang.srt.layers.quantization.mxfp8_input import Mxfp8SwizzledInput
    return Mxfp8SwizzledInput


def mxfp8_gather(p, x):
    """This rank's bf16 shard [S, K] -> the whole chunk's (fp8 [M, K], linear scales [M, K/32])."""
    s_rows, k = x.shape
    q, sf = _mxfp8_quantize(x.contiguous(), False)
    sf = sf.view(torch.uint8).reshape(-1)[: s_rows * (k // 32)].view(s_rows, k // 32)
    qg = _gather_rows(p, q.view(torch.uint8)).view(q.dtype)
    sg = _gather_rows(p, sf)
    return qg, sg


def _fp8_input(p, attn, x):
    """Pre-quantized attention input for this chunk, or None (bf16 gather) with the bf16 rows
    when the first-chunk check had to gather them anyway."""
    state = _FP8_M.get(p.rows)
    if state is not None and not state[0]:
        return None, None
    qg, sg = mxfp8_gather(p, x)
    sw = swizzle_128x4(sg)
    if state is None:
        xb = _gather_rows(p, x)
        rq, rs = _mxfp8_quantize(xb, True)
        rs_u8 = rs.view(torch.uint8).reshape(-1)
        local = bool(rq.shape == qg.shape and torch.equal(rq.view(torch.uint8), qg.view(torch.uint8))
                     and rs_u8.numel() == sw.numel() and torch.equal(rs_u8, sw))
        if local:
            xq = _mxfp8_input_cls()(qg, sw.view(rs.dtype).view(rs.shape))
            with _Inner():
                local = bool(torch.equal(attn.wqkv_a(xb)[0], attn.wqkv_a(xq)[0]))
        ok = _agree_min(p, local)
        _FP8_M[p.rows] = state = (ok, tuple(rs.shape), rs.dtype)
        if p.rank == 0 and (len(_FP8_M) <= 8 or not ok):
            print(f"DSV41_PREFILL_SP fp8 attention input (M={p.rows}): MXFP8 bytes + scales and wqkv_a "
                  f"output {'equal the stock bf16 path on every rank -> ON' if ok else 'differ -> bf16 gather'}"
                  f" (rank 0 local {local})", flush=True)
        if not ok:
            return None, xb
    return _mxfp8_input_cls()(qg, sw.view(state[2]).view(state[1])), None


# ------------------------------------------------------------------------------------------
# attention / MoE: gather in, reduce-scatter out
# ------------------------------------------------------------------------------------------
def _hook_wo_b(attn):
    wo_b = attn.wo_b
    if getattr(wo_b, "_dsv41_sp_hooked", False):
        return
    orig = wo_b.forward

    def forward(input_, skip_all_reduce=False, *a, **kw):
        if _ctx.skip_wo_b is not wo_b:
            return orig(input_, skip_all_reduce, *a, **kw)
        out = orig(input_, True, *a, **kw)
        if _comparing():
            # stock: the same GEMM with the engine's own all-reduce, vs our all-reduce of the partial
            ref = orig(input_, False, *a, **kw)[0]
            ours = _ctx.plan.group.all_reduce(out[0].clone())
            _cmp("wo_b all-reduce", ours, ref)
        return out

    wo_b.forward = forward
    wo_b._dsv41_sp_hooked = True


def _make_attn_forward(orig):
    def forward(self, x, positions, forward_batch, x_quant=None):
        p = _ctx.plan
        if (p is None or _ctx.inner or x_quant is not None or not isinstance(x, torch.Tensor)
                or x.dim() != 2 or x.shape[0] != p.in_rows):
            if _ctx.dbg is None or _ctx.inner:
                return orig(self, x, positions, forward_batch, x_quant)
            _rec("ain", x, True)
            o = orig(self, x, positions, forward_batch, x_quant)
            _rec("aout", o, True)
            return o
        _hook_wo_b(self)
        xq = xb = None
        if p.sharded and FP8 and fp8_eligible_attn(self):
            xq, xb = _fp8_input(p, self, x)
        if xq is not None:
            # shape / dtype only: this layer's attention reads x solely through wqkv_a(x_quant).
            # NaN, so any other read would poison the output instead of passing silently.
            xf = x.new_full((1, 1), float("nan")).expand(p.rows, x.shape[1])
            if _ctx.dbg is not None:
                xb = _gather_rows(p, x)
                _rec("ain", xb, True)
                if _comparing():
                    rq, rs = _mxfp8_quantize(xb, True)
                    _cmp("attn x mxfp8 data", xq[0].view(torch.uint8), rq.view(torch.uint8))
                    _cmp("attn x mxfp8 scales", xq[1].view(torch.uint8).reshape(-1),
                         rs.view(torch.uint8).reshape(-1))
                    with _Inner():
                        _cmp("attn wqkv_a", self.wqkv_a(xq)[0], self.wqkv_a(xb)[0])
                    del rq, rs
                del xb
        else:
            xf = xb if xb is not None else (_gather_rows(p, x) if p.sharded else x)
            if _ctx.dbg is not None:
                _rec("ain", xf, True)
        _ctx.inner += 1
        _ctx.skip_wo_b = self.wo_b
        try:
            o = orig(self, xf, positions, forward_batch, xq)
        finally:
            _ctx.inner -= 1
            _ctx.skip_wo_b = None
        if o.shape[0] != p.rows:
            raise RuntimeError(f"DSV41_PREFILL_SP: attention returned {tuple(o.shape)} for {p.rows} rows")
        out = _reduce_out(p, o)
        if _ctx.dbg is not None:
            _rec("aout", out)
        return out

    forward.__wrapped__ = orig
    return forward


def _make_moe_forward(orig):
    def forward(self, hidden_states, *args, **kwargs):
        p = _ctx.plan
        if (p is None or _ctx.inner or not isinstance(hidden_states, torch.Tensor)
                or hidden_states.dim() != 2 or hidden_states.shape[0] != p.in_rows
                or kwargs.get("skip_shared_experts", False)):
            if _ctx.dbg is None or _ctx.inner or not isinstance(hidden_states, torch.Tensor):
                return orig(self, hidden_states, *args, **kwargs)
            _rec("min", hidden_states, True)
            o = orig(self, hidden_states, *args, **kwargs)
            _rec("mout", o, True)
            return o
        v4 = _M["v4"]
        xf = _gather_rows(p, hidden_states) if p.sharded else hidden_states
        if _mega_moe(self, xf):
            raise RuntimeError("DSV41_PREFILL_SP: mega MoE path is not wired")
        if _ctx.dbg is not None:
            _rec("min", xf, True)
        x_ref = xf.clone() if _comparing() else None     # the MoE may write its input in place
        _ctx.inner += 1
        try:
            with v4.get_forward().scoped(mlp_reduce_scatter=True):
                out = orig(self, xf, *args, **kwargs)
        finally:
            _ctx.inner -= 1
        if out.shape[0] != p.rows:
            raise RuntimeError(f"DSV41_PREFILL_SP: MoE returned {tuple(out.shape)} for {p.rows} rows")
        if x_ref is not None:
            # stock: the same MoE on the same rows with its own post-experts all-reduce
            ours = p.group.all_reduce(out.clone())
            with _Inner():
                ref = orig(self, x_ref, *args, **kwargs)
            _cmp("moe all-reduce", ours, ref)
            del x_ref, ref, ours
        res = _reduce_out(p, out)
        if _ctx.dbg is not None:
            _rec("mout", res)
        return res

    forward.__wrapped__ = orig
    return forward


# ------------------------------------------------------------------------------------------
# hc gates keyed on the logical row count (stage 1 only)
# ------------------------------------------------------------------------------------------
def _combine_norm_fused(layer, x, apply_pre, norm, m):
    """The engine's _hc_combine fused-norm gate, evaluated at row count m."""
    v4 = _M["v4"]
    from sglang.srt.batch_invariant_ops import is_batch_invariant_mode_enabled
    return bool(
        x.is_cuda
        and v4.get_platform().is_blackwell
        and (0 < m <= 8 or (layer.config.model_type == "deepseek_v41" and 4096 <= m <= 65536))
        and layer.hc_mult == 4
        and x.flatten(1).shape[1] == 20480
        and x.dtype == norm.weight.dtype == torch.bfloat16
        and apply_pre.stride(1) == 1
        and not norm.cast_x_before_out_mul
        and norm.variance_size_override is None
        and not is_batch_invariant_mode_enabled()
    )


def hc_combine_norm_rows(x_flat, pre, weight, eps, logical_m):
    """hc_combine_norm on a row slice with the kernel the full chunk (logical_m rows) uses."""
    hcn = _M["hcn"]
    m = x_flat.shape[0]
    assert 4096 <= logical_m <= 65536 and x_flat.shape == (m, 20480)
    assert pre.shape == (m, 4) and pre.stride(1) == 1
    assert weight.shape == (5120,) and weight.is_contiguous()
    assert x_flat.dtype == weight.dtype == torch.bfloat16 and x_flat.stride(1) == 1
    y = torch.empty((m, 5120), dtype=x_flat.dtype, device=x_flat.device)
    hcn._hc_combine_norm_prefill[(m,)](
        x_flat, pre, weight, y, x_flat.stride(0), pre.stride(0), eps, num_warps=4
    )
    return y


def _make_hc_combine(orig):
    def _hc_combine(self, x, apply_pre, norm, stats_stream=None, quantized=None, normalized=None,
                    precomputed=None):
        p = _ctx.plan
        if (p is None or _ctx.inner or not p.sharded or x.shape[0] != p.shard
                or precomputed is not None or normalized is not None or stats_stream is not None):
            return orig(self, x, apply_pre, norm, stats_stream, quantized, normalized, precomputed)
        y = _sp_combine(self, orig, p, x, apply_pre, norm, quantized)
        if _comparing():
            with _Inner():
                xg = _gather_rows(p, x)
                pg = _gather_rows(p, apply_pre) if apply_pre is not None else None
                ref = orig(self, xg, pg, norm, None, None, None, None)[p.lo:p.hi].clone()
                del xg, pg
            which = "attn" if norm is getattr(self, "input_layernorm", None) else "ffn"
            _cmp(f"hc_combine {which}{'' if apply_pre is not None else ' (stream 0)'}", y, ref,
                 {"x": x, "pre": apply_pre})
        return y

    def _sp_combine(self, orig, p, x, apply_pre, norm, quantized):
        if apply_pre is None:
            return orig(self, x, apply_pre, norm, None, quantized, None, None)
        m = p.rows
        full = _combine_norm_fused(self, x, apply_pre, norm, m)
        here = _combine_norm_fused(self, x, apply_pre, norm, x.shape[0])
        if full == here:
            return orig(self, x, apply_pre, norm, None, quantized, None, None)
        x_flat = x.flatten(1)
        if full:
            return hc_combine_norm_rows(x_flat, apply_pre, norm.weight, norm.variance_epsilon, m)
        return norm(_M["mhc"].hc_combine(x_flat, apply_pre, self.hc_mult, x.dtype))

    _hc_combine.__wrapped__ = orig
    return _hc_combine


def _post_split_h(layer, x, residual, post, comb, m):
    v4 = _M["v4"]
    return bool(
        v4._is_cuda
        and v4.get_platform().is_blackwell
        and layer.hc_pre_from_prev_sublayer
        and layer.hc_mult == 4
        and x.shape[1] == 5120
        and m <= 384
        and x.dtype == residual.dtype == torch.bfloat16
        and post.dtype == comb.dtype == torch.float32
        and all(t.is_contiguous() for t in (x, residual, post, comb))
    )


def _make_hc_post(orig):
    def hc_post(self, x, residual, post, comb):
        p = _ctx.plan
        if p is None or _ctx.inner or not p.sharded or x.shape[0] != p.shard:
            return orig(self, x, residual, post, comb)
        out = _sp_post(self, x, residual, post, comb, p)
        if _comparing():
            with _Inner():
                ref = orig(self, _gather_rows(p, x), _gather_rows(p, residual), _gather_rows(p, post),
                           _gather_rows(p, comb))[p.lo:p.hi].clone()
            _cmp("hc_post", out, ref, {"x": x, "post": post, "comb": comb})
        return out

    def _sp_post(self, x, residual, post, comb, p):
        v4 = _M["v4"]
        full = _post_split_h(self, x, residual, post, comb, p.rows)
        if full == _post_split_h(self, x, residual, post, comb, x.shape[0]):
            return orig(self, x, residual, post, comb)
        if full:
            return v4.mhc_post_split_h(x, residual, post, comb)
        # the full chunk skips split-h: the rest of the engine's CUDA chain
        if v4.envs.SGLANG_OPT_USE_FLASHINFER_MHC.get():
            from flashinfer.mhc import mhc_post
            return mhc_post(x, residual, post, comb)
        if v4.envs.SGLANG_OPT_USE_TILELANG_MHC_POST.get():
            if v4.get_platform().is_sm90 and 1 <= p.rows <= 64:
                return v4.mhc_post_split_h(x, residual, post, comb)
            return _M["mhc"].mhc_post(x, residual, post, comb)
        raise RuntimeError("DSV41_PREFILL_SP: hc_post fallback chain not wired for this platform")

    hc_post.__wrapped__ = orig
    return hc_post


def _make_hc_mix_stats(orig):
    """Debug compare only: the stats on the shard vs the stock stats on the gathered rows."""
    def _hc_mix_stats(self, x, hc_fn, hc_scale, hc_base, stats_stream=None):
        out = orig(self, x, hc_fn, hc_scale, hc_base, stats_stream)
        p = _ctx.plan
        if (_comparing() and not _ctx.inner and stats_stream is None and x.shape[0] == p.shard
                and p.shard != p.rows):
            with _Inner():
                ref = orig(self, _gather_rows(p, x), hc_fn, hc_scale, hc_base, None)
            which = "attn" if hc_fn is getattr(self, "hc_attn_fn", None) else "ffn"
            _cmp(f"hc_mix_stats {which}", tuple(out), tuple(r[p.lo:p.hi].clone() for r in ref))
        return out

    _hc_mix_stats.__wrapped__ = orig
    return _hc_mix_stats


# ------------------------------------------------------------------------------------------
# install
# ------------------------------------------------------------------------------------------
def install(model_mod):
    """sglang.srt.models.deepseek_v4. Must run BEFORE hc_prefill_fused.install, so that its
    'stock' _hc_combine / hc_post are these logical-row-aware ones."""
    if MODE == "off" and not DEBUG:
        return
    if getattr(model_mod, "_dsv41_prefill_sp", False):
        return
    _load_modules(model_mod)
    check_engine()
    if MODE == "shard":
        _check_hc_prefill_fused()
    Model, Layer = model_mod.DeepseekV4Model, model_mod.DeepseekV4DecoderLayer
    MoE = _M["v2"].DeepseekV2MoE
    stock_layers = Model._forward_layers_hc_pre_from_prev

    def _forward_layers_hc_pre_from_prev(self, positions, hidden_states, forward_batch, input_ids,
                                         input_ids_global, capture_dspark, dspark_aux_hidden_states):
        p = _plan(self, hidden_states, forward_batch)
        d = _debug_begin(self, hidden_states, forward_batch, input_ids, p) if DEBUG else None
        _ctx.dbg = d
        try:
            if p is None:
                out = stock_layers(self, positions, hidden_states, forward_batch, input_ids,
                                   input_ids_global, capture_dspark, dspark_aux_hidden_states)
            else:
                out = _sp_forward_layers(self, p, positions, hidden_states, forward_batch, input_ids,
                                         input_ids_global, capture_dspark, dspark_aux_hidden_states)
        finally:
            _ctx.dbg = None
        if d is not None:
            _debug_end(d, out)
        return out

    _forward_layers_hc_pre_from_prev.__wrapped__ = stock_layers
    Model._forward_layers_hc_pre_from_prev = _forward_layers_hc_pre_from_prev
    model_mod.MQALayer.forward = _make_attn_forward(model_mod.MQALayer.forward)
    MoE.forward = _make_moe_forward(MoE.forward)
    if MODE == "shard":
        Layer._hc_combine = _make_hc_combine(Layer._hc_combine)
        Layer.hc_post = _make_hc_post(Layer.hc_post)
        if "compare" in DEBUG:
            Layer._hc_mix_stats = _make_hc_mix_stats(Layer._hc_mix_stats)
    if DEBUG:
        Layer.forward_hc_pre_from_prev = _make_layer_forward(Layer.forward_hc_pre_from_prev)
        print(f"DSV41_PREFILL_SP DEBUG {sorted(DEBUG)} ARMED for the first {DEBUG_CHUNKS} extend "
              f"forwards{' (dumps in ' + DEBUG_DUMP + ')' if DEBUG_DUMP else ''}: NOT FOR PRODUCTION",
              flush=True)
    model_mod._dsv41_prefill_sp = True
    if MODE == "off":
        return
    what = ("reduce-scatter + all-gather at wo_b / MoE / Engram (comm only, nothing sharded)"
            if MODE == "comm" else
            f"rows sharded over TP on the full-row layers, bf16 gathers"
            f"{', EXACT (all-reduce + keep the quarter)' if EXACT else ''}"
            f"{', MXFP8 attention-input gathers (checked per chunk size)' if FP8 else ''}"
            f", wkv {'on the shard (self-checked)' if WKV_SHARD else 'full-M'}")
    print(f"DSV41_PREFILL_SP={_RAW} ARMED: prefill chunks >= {MIN_ROWS} rows: {what}", flush=True)
    if EXACT and MODE == "comm":
        print("DSV41_PREFILL_SP: DSV41_PREFILL_SP_EXACT has no effect in comm mode", flush=True)