File size: 5,909 Bytes
e1acc15 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | """Torch-free, canonical native drafter arm configuration."""
from __future__ import annotations
import os
from dataclasses import dataclass
from typing import TypedDict
BATCH_INVARIANT_ENV = "DSPARK_NATIVE_BATCH_INVARIANT"
def batch_invariant_from_env() -> int:
"""NB1 batch-invariant dispatch level for native kx2, from one env switch: "0"/unset = off (default); "1" =
every draft-body shape (B x block rows, B = 1..4) uses the dispatch B1 uses; "2" = also window insertion in
fixed 8-row projection sub-batches (the bucket-8 arithmetic for every flush size). Anything else refuses."""
value = os.environ.get(BATCH_INVARIANT_ENV, "0")
if value not in ("0", "1", "2"):
raise ValueError(f"{BATCH_INVARIANT_ENV}={value!r}: expected 0, 1 or 2")
return int(value)
@dataclass(frozen=True)
class ArmConfig:
name: str
experts: str = "k2"
dense: str = "k5"
kv_qdq: str = "engine"
wo_a_qdq: bool = False
combine: str = "fused"
head: str = "k6"
expanded_stages: tuple[int, ...] = ()
expert_kernel: str = "n1" # packed experts: "n1" (ops.ExpertBank Triton pair kernels) | "kx-regfed"
markov: str = "factorized" # head bf16: "factorized" (exact fp32 bias computed per position) | "hot"
markov_rows: int = 4096 # (+ exact fp32 rows cached for the serve's most frequent previous tokens)
dense_kernel: str = "n1" # "kx-selective": measured FP8 shapes, BF16 otherwise
dense_profile: str = "initial" # N2 tuning control: "initial" | "measured-small"
def check(self) -> None:
if self.experts not in ("k2", "packed", "expanded") or self.dense not in ("k5", "fp8", "bf16"):
raise ValueError(f"{self.name}: experts {self.experts} / dense {self.dense}")
if self.kv_qdq not in ("official", "engine") or self.combine not in ("official", "fused"):
raise ValueError(f"{self.name}: kv_qdq {self.kv_qdq} / combine {self.combine}")
if self.expert_kernel not in ("n1", "kx-regfed", "kx-fused"):
raise ValueError(f"{self.name}: expert kernel {self.expert_kernel}")
if self.dense_kernel not in ("n1", "kx-selective") or (self.dense_kernel == "kx-selective" and self.dense != "bf16"):
raise ValueError(f"{self.name}: dense kernel {self.dense_kernel} requires BF16 baseline")
if self.dense_profile not in ("initial", "measured-small"):
raise ValueError(f"{self.name}: dense profile {self.dense_profile}")
if self.head not in ("k6", "bf16") or self.markov not in ("factorized", "hot"):
raise ValueError(f"{self.name}: head {self.head} / markov {self.markov}")
if self.combine == "official" and self.experts == "k2":
raise ValueError(f"{self.name}: the K2 coop kernel sums its experts itself (combine must be fused)")
if self.dense == "k5" and (self.kv_qdq != "engine" or self.wo_a_qdq):
raise ValueError(f"{self.name}: K5 dense runs E's attention (engine KV record, no wo_a QDQ)")
@property
def native(self) -> bool:
return self.experts != "k2" or self.dense != "k5" or self.head != "k6"
@property
def kx2_dispatch(self) -> bool:
"""The measured KX dispatch the batch-invariant mode pins (FP8 dense by M, fused experts by stage/rows)."""
return (self.dense_kernel == "kx-selective" and self.dense_profile == "measured-small"
and self.expert_kernel == "kx-fused")
def stage_expert_kind(self, i: int) -> str:
if self.experts == "packed" and i in self.expanded_stages:
return "expanded"
return self.experts
class OfficialPolicy(TypedDict):
kv_qdq: str
wo_a_qdq: bool
combine: str
OFFICIAL: OfficialPolicy = dict(kv_qdq="official", wo_a_qdq=True, combine="official")
ARMS = {
"k2-tp2": ArmConfig("k2-tp2"), # the deployed control
"native-experts-packed": ArmConfig("native-experts-packed", experts="packed", combine="official"),
"native-experts-expanded": ArmConfig("native-experts-expanded", experts="expanded", combine="official"),
"native-packed": ArmConfig("native-packed", experts="packed", dense="fp8", **OFFICIAL),
"native-expanded": ArmConfig("native-expanded", experts="expanded", dense="bf16", **OFFICIAL),
"native-hybrid": ArmConfig("native-hybrid", experts="packed", dense="bf16", **OFFICIAL),
"native-hybrid-s0x": ArmConfig("native-hybrid-s0x", experts="packed", dense="bf16", expanded_stages=(0,),
**OFFICIAL),
# NDM §6 second step: the draft-only official BF16 head (fp32 logits) with the exact fp32 Markov bias
"native-hybrid-head": ArmConfig("native-hybrid-head", experts="packed", dense="bf16", head="bf16", **OFFICIAL),
"native-expanded-head": ArmConfig("native-expanded-head", experts="expanded", dense="bf16", head="bf16",
**OFFICIAL),
"native-hybrid-head-hot": ArmConfig("native-hybrid-head-hot", experts="packed", dense="bf16", head="bf16",
markov="hot", **OFFICIAL),
"k2-head": ArmConfig("k2-head", head="bf16"), # the head change alone on the control
# KX-MXFP4's register-fed kernel for the packed experts (same math; their schedule, MMA and reuse)
"native-hybrid-kx": ArmConfig("native-hybrid-kx", experts="packed", dense="bf16", expert_kernel="kx-regfed",
**OFFICIAL),
"native-packed-kx": ArmConfig("native-packed-kx", experts="packed", dense="fp8", expert_kernel="kx-regfed",
**OFFICIAL),
"native-hybrid-kx2": ArmConfig("native-hybrid-kx2", experts="packed", dense="bf16", expert_kernel="kx-fused",
dense_kernel="kx-selective", dense_profile="measured-small", **OFFICIAL),
}
|