|
|
@@ -0,0 +1,1137 @@
|
|
|
+"""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)
|