diff --git a/scmp_kernels/mp/__init__.py b/scmp_kernels/mp/__init__.py index 9ba1383..10dcdeb 100644 --- a/scmp_kernels/mp/__init__.py +++ b/scmp_kernels/mp/__init__.py @@ -6,6 +6,8 @@ from .config import ( MPConfig, AdaptiveMPConfig, + compute_row_metric, + ROW_METRIC_NAMES, FreeBoundaryMPConfig, RangeMPConfig, RowAssignment, @@ -28,6 +30,8 @@ __all__ = [ "MPConfig", "AdaptiveMPConfig", + "compute_row_metric", + "ROW_METRIC_NAMES", "FreeBoundaryMPConfig", "RangeMPConfig", "RowAssignment", diff --git a/scmp_kernels/mp/config.py b/scmp_kernels/mp/config.py index 4a988a3..602849e 100644 --- a/scmp_kernels/mp/config.py +++ b/scmp_kernels/mp/config.py @@ -91,6 +91,87 @@ def _parse_bucket_key(bucket_key: str) -> tuple[str, int, int]: ) from exc +def _parse_protected_channel_key(key: str) -> tuple[str, int, Optional[int]]: + """Parse protected-channel keys like 'q_proj:b12' or 'up_proj:b3:u7'.""" + try: + parts = key.split(":") + if len(parts) not in (2, 3): + raise ValueError + operator, b_part = parts[0], parts[1] + if not b_part.startswith("b"): + raise ValueError + unit = None + if len(parts) == 3: + u_part = parts[2] + if not u_part.startswith("u"): + raise ValueError + unit = int(u_part[1:]) + return operator, int(b_part[1:]), unit + except Exception as exc: # pragma: no cover - defensive parsing + raise ValueError( + f"Invalid protected-channel key '{key}'. " + "Expected ':b' or ':b:u'." + ) from exc + + +def _extract_group_levels(payload, source: str) -> Optional[list[int]]: + """Extract a per-group ladder from a bucket payload, or None if absent. + + Returning None (not the global list) is what keeps a table without + per-group ladders on the byte-identical path. + + STRICTLY descending is required, not merely non-increasing: RowAssignment + keys level_row_indices by the stoc_len VALUE, so two rungs sharing a value + would silently collapse into one entry and the rows of the lower rung + would be evaluated at the higher rung's length with no error anywhere. + """ + if not isinstance(payload, dict): + return None + raw = payload.get("stoc_len_levels") + if raw is None: + return None + levels = [int(x) for x in raw] + if len(levels) < 2: + raise ValueError( + f"Per-group stoc_len_levels for {source} must have >= 2 rungs, " + f"got {levels}.") + for idx in range(1, len(levels)): + if levels[idx] >= levels[idx - 1]: + raise ValueError( + f"Per-group stoc_len_levels for {source} must be strictly " + f"descending (duplicate or ascending rungs silently merge in " + f"level_row_indices), got {levels}.") + if levels[-1] < 0: + raise ValueError( + f"Per-group stoc_len_levels for {source} has a negative rung: " + f"{levels}.") + return levels + + +def residual_chunk_widths(residual_width: int, chunk_d: int) -> list[int]: + """Widths of the residual's quantization chunks, tail last. + + Mirrors the kernel's chunk loop (`for d_start in range(0, D, chunk_d)`, + sc/kernels.py) so calibration and runtime agree on what a "group" is. The + last chunk is short whenever chunk_d does not divide the residual width, + which is the common case once protected channels are carved out (e.g. + down_proj 9728 - 584 = 9144 -> 71 x 128 + 56). + """ + if residual_width <= 0 or chunk_d <= 0: + return [] + full, tail = divmod(residual_width, chunk_d) + return [chunk_d] * full + ([tail] if tail else []) + + +def _band_widths(band_of_chunk: list[int], chunk_widths: list[int], + n_bands: int) -> list[int]: + """Total columns per band. Bands need not be contiguous in chunk index.""" + widths = [0] * n_bands + for chunk_idx, band in enumerate(band_of_chunk): + widths[band] += chunk_widths[chunk_idx] + return widths + + def _extract_thresholds(payload, n_levels: int, source: str) -> list[float]: """Extract a threshold list of length n_levels-1 from a table payload.""" raw_thresholds = payload.get("thresholds") if isinstance(payload, dict) else payload @@ -116,6 +197,35 @@ def _extract_thresholds(payload, n_levels: int, source: str) -> list[float]: return thresholds +_ROW_METRIC_EPS = 1e-12 +# Candidate per-row dispatch metrics (act_global_v2 ρ-selection). All are O(D) +# reductions over the last dim — same runtime cost class as the original amax. +ROW_METRIC_NAMES = ("amax", "l2", "crest") + + +def compute_row_metric(x: torch.Tensor, name: str) -> torch.Tensor: + """Per-row dispatch metric over the LAST dim of ``x``. + + ``amax`` = ‖row‖_inf (the original metric), + ``l2`` = ‖row‖_2, + ``crest`` = ‖row‖_inf / ‖row‖_2 (scale-free peakedness). + + The caller multiplies by the calibrated sign (−1 inverts the ranking; the + min–max normalization inside ``adaptive_classify_rows`` maps sign-flipped + values to exactly ``1 − normalized(raw)``, matching the calibration-side + transform in calibrate_mp_thresholds.py). + """ + if name == "amax": + return x.abs().amax(dim=-1) + if name == "l2": + return x.float().norm(dim=-1) + if name == "crest": + xf = x.float() + return xf.abs().amax(dim=-1) / (xf.norm(dim=-1) + _ROW_METRIC_EPS) + raise ValueError(f"unknown dispatch metric '{name}' " + f"(expected one of {ROW_METRIC_NAMES})") + + def _classify_rows_by_thresholds( metric_norm: torch.Tensor, stoc_len_levels: list[int], @@ -228,10 +338,61 @@ class AdaptiveMPConfig: layer_buckets: int = 1 operator_default_thresholds: dict[str, list[float]] = field(default_factory=dict) bucket_thresholds: dict[tuple[str, int, int], list[float]] = field(default_factory=dict) + # Per-(op, timestep-bucket, layer-bucket) LADDER, parallel to + # bucket_thresholds. Empty => every bucket uses stoc_len_levels, which is + # the historical behaviour. See get_levels(). + bucket_stoc_len_levels: dict[tuple[str, int, int], list[int]] = field( + default_factory=dict) + operator_default_stoc_len_levels: dict[str, list[int]] = field( + default_factory=dict) + protected_channel_stoc_len: Optional[int] = None + protected_channel_indices: dict[tuple[str, int, Optional[int]], list[int]] = field(default_factory=dict) + # ---- K-bands (Phase 3: per-group stream lengths) ---------------------- + # The residual (non-protected) contraction axis is partitioned into + # ``k_band_count`` bands of WHOLE quantization chunks. Row dispatch is + # unchanged -- one metric, one rung index k per row -- but band b executes + # rung k at its own length ``k_band_ladders[(op,t,l)][b][k]``. Setting + # every band's ladder equal to the bucket ladder reproduces the per-row + # parent exactly, which is what makes this refinement unable to lose. + # k_band_count == 0 (default) disables the whole path. + k_band_count: int = 0 + k_band_chunk_d: int = 128 + # (operator, block_idx) -> band id per residual chunk, ascending chunk order + k_band_chunks: dict[tuple[str, int], list[int]] = field(default_factory=dict) + # (operator, t_bucket, l_bucket) -> [n_bands][n_rungs] stream lengths + k_band_ladders: dict[tuple[str, int, int], list[list[int]]] = field( + default_factory=dict) + # operator -> residual width R (constant per op: |protected| is per-op + # constant even though WHICH channels are protected varies per block) + k_band_residual_width: dict[str, int] = field(default_factory=dict) + # act_global_v2 ρ-selected dispatch metric: operator -> (name, sign). + # Absent operator => ("amax", +1.0), byte-identical to the original + # dispatch. Populated from the table's "dispatch_metrics" payload. + dispatch_metrics: dict[str, tuple[str, float]] = field(default_factory=dict) # When set, bypass the linear-threshold classifier and use these fractions # as quantile targets per level (top frac[0] rows -> levels[0], etc.). # Length must match stoc_len_levels; sums to 1. target_fractions: Optional[list[float]] = None + # ---- Absolute escape gate (R7) ---------------------------------------- + # When ``escape_gate_k`` is set, calibrated-table classification adds one + # extra compare after the per-call min–max normalization: rows whose + # NORMALIZED metric exceeds t_esc_b = metric_mean_b + k * metric_std_b + # escape to ``escape_stoc_len``, regardless of the band thresholds. + # mu_b / sigma_b are the per-bucket ``metric_mean`` / ``metric_std`` + # already stored in the calibration table — statistics of the pooled + # per-call-normalized SIGNED dispatch metric, i.e. the SAME post-sign + # normalized space the band thresholds live in (sign −1 metrics were + # stored as 1 − normalized(raw)), so the compare direction never flips. + # ``escape_gate_k=None`` (default) disables the gate and is byte-identical + # to the pre-gate classifier. + escape_gate_k: Optional[float] = None + escape_stoc_len: int = 128 + # Precomputed at table load (one float per bucket / operator default). + # NOT clamped to <= 1.0 on purpose: the normalized metric lives in [0, 1] + # (max EXACTLY 1.0) and the compare is strict, so a t_esc >= 1.0 simply + # never fires for that bucket. + bucket_escape_thresholds: dict[tuple[str, int, int], float] = field(default_factory=dict) + operator_default_escape_thresholds: dict[str, float] = field(default_factory=dict) def __post_init__(self): assert len(self.stoc_len_levels) >= 2, ( @@ -242,6 +403,18 @@ def __post_init__(self): f"got {self.stoc_len_levels}") if not self.enable_pruning and 0 in self.stoc_len_levels: self.stoc_len_levels = [s for s in self.stoc_len_levels if s > 0] + if self.escape_gate_k is not None: + self.escape_gate_k = float(self.escape_gate_k) + self.escape_stoc_len = int(self.escape_stoc_len) + if self.escape_stoc_len <= 0: + raise ValueError( + f"escape_stoc_len must be a positive cycle count, " + f"got {self.escape_stoc_len}") + if not self.threshold_table_path: + raise ValueError( + "escape_gate_k requires a calibrated threshold table " + "(threshold_table_path): the gate constants are the " + "table's per-bucket metric_mean/metric_std.") if self.threshold_table_path: self.load_threshold_table(self.threshold_table_path) if self.target_fractions is not None: @@ -269,21 +442,84 @@ def load_threshold_table(self, path: str): self.layer_buckets = int(payload.get("layer_buckets", 1)) self.operator_default_thresholds = {} self.bucket_thresholds = {} + self.operator_default_escape_thresholds = {} + self.bucket_escape_thresholds = {} + self.bucket_stoc_len_levels = {} + self.operator_default_stoc_len_levels = {} + self.protected_channel_stoc_len = None + self.protected_channel_indices = {} + self.k_band_count = 0 + self.k_band_chunks = {} + self.k_band_ladders = {} + self.k_band_residual_width = {} for operator, operator_payload in payload.get("operator_defaults", {}).items(): + op_levels = _extract_group_levels( + operator_payload, f"operator_default:{operator}") + if op_levels is not None: + self.operator_default_stoc_len_levels[operator] = op_levels self.operator_default_thresholds[operator] = _extract_thresholds( operator_payload, - len(self.stoc_len_levels), + len(op_levels if op_levels is not None + else self.stoc_len_levels), f"operator_default:{operator}", ) + t_esc = self._escape_threshold_from_payload(operator_payload) + if t_esc is not None: + self.operator_default_escape_thresholds[operator] = t_esc for bucket_key, bucket_payload in payload.get("buckets", {}).items(): operator, t_bucket, l_bucket = _parse_bucket_key(bucket_key) + # Per-group ladder, if this table declares one. Thresholds are + # then validated against THIS bucket's rung count, not the global + # one -- a per-group ladder may legitimately have a different + # number of rungs from the table-level default. + grp_levels = _extract_group_levels(bucket_payload, bucket_key) + if grp_levels is not None: + self.bucket_stoc_len_levels[ + (operator, t_bucket, l_bucket)] = grp_levels self.bucket_thresholds[(operator, t_bucket, l_bucket)] = _extract_thresholds( bucket_payload, - len(self.stoc_len_levels), + len(grp_levels if grp_levels is not None + else self.stoc_len_levels), bucket_key, ) + t_esc = self._escape_threshold_from_payload(bucket_payload) + if t_esc is not None: + self.bucket_escape_thresholds[(operator, t_bucket, l_bucket)] = t_esc + + if self.escape_gate_k is not None and not ( + self.bucket_escape_thresholds + or self.operator_default_escape_thresholds): + raise ValueError( + f"escape_gate_k={self.escape_gate_k} is set but the threshold " + f"table {path} carries no metric_mean/metric_std in any bucket " + "or operator default — the gate constants cannot be derived. " + "Recalibrate with a calibrator that exports metric stats, or " + "unset escape_gate_k.") + + self._load_k_bands( + payload.get("k_bands"), + sc_prec=int(payload.get("sc_prec", 8)), + halve=bool(payload.get("halve_bipolar_stoc_len", True))) + + protected = payload.get("protected_channels") or {} + indices = protected.get("indices") or {} + if indices: + self.protected_channel_stoc_len = int(protected.get("stoc_len", 128)) + for key, vals in indices.items(): + self.protected_channel_indices[_parse_protected_channel_key(key)] = [ + int(v) for v in vals + ] + + self.dispatch_metrics = {} + for op, spec in (payload.get("dispatch_metrics") or {}).items(): + name = str(spec.get("metric", "amax")) + if name not in ROW_METRIC_NAMES: + raise ValueError( + f"Adaptive MP table dispatch_metrics[{op}] names unknown " + f"metric '{name}' (expected one of {ROW_METRIC_NAMES}).") + self.dispatch_metrics[op] = (name, float(spec.get("sign", 1.0))) def get_thresholds( self, @@ -304,6 +540,353 @@ def get_thresholds( return self.operator_default_thresholds[operator] return None + # ---- K-bands --------------------------------------------------------- + # Tolerance on the per-rung iso-cost identity, in halved cycles, MAC + # weighted. Band ladders are integers so exact equality is generally + # unreachable; 0.25 is well inside the smallest effect worth chasing (a + # t48 budget moves ~1 cycle for a 2% change) and far tighter than the + # realized-trace band check used downstream. + K_BAND_ISO_COST_TOL = 0.25 + # Underspend is SAFE for the iso-compute claim (a win at lower cost is + # strictly stronger), so it gets a looser bound than overspend -- but it is + # still bounded, because a solver leaving many cycles unspent is a solver + # bug, not a conservative choice. + K_BAND_MAX_UNDERSPEND = 2.0 + + def _load_k_bands(self, section, sc_prec: int = 8, + halve: bool = True) -> None: + """Parse and VALIDATE the k_bands section (Phase 3). + + Everything here is a hard failure rather than a fallback: a malformed + band spec that silently degrades to per-row would produce a cell that + looks like a Phase-3 result but is not one, and realized_flop_avg_sl + would not reveal it. + """ + if not section: + return + n_bands = int(section.get("n_bands", 0)) + if n_bands < 2: + raise ValueError( + f"k_bands.n_bands must be >= 2 (got {n_bands}); omit the " + f"section entirely to run the per-row parent.") + chunk_d = int(section.get("chunk_d", 128)) + cap = 2 ** (sc_prec - 1) if halve else 2 ** sc_prec + residual_width = {str(k): int(v) + for k, v in (section.get("residual_width") or {}).items()} + if not residual_width: + raise ValueError( + "k_bands.residual_width is required: band widths price the " + "iso-cost identity and cannot be inferred at load time.") + + chunk_bands: dict[tuple[str, int], list[int]] = {} + # band -> width, per operator; must agree across every block of an op + op_band_widths: dict[str, list[int]] = {} + for key_str, band_ids in (section.get("chunk_bands") or {}).items(): + operator, block_idx, _ = _parse_protected_channel_key(key_str) + if operator not in residual_width: + raise ValueError( + f"k_bands.chunk_bands has '{key_str}' but no " + f"residual_width['{operator}'].") + widths = residual_chunk_widths(residual_width[operator], chunk_d) + bands = [int(b) for b in band_ids] + if len(bands) != len(widths): + raise ValueError( + f"k_bands.chunk_bands['{key_str}'] has {len(bands)} " + f"entries but the residual has {len(widths)} chunks " + f"(residual_width={residual_width[operator]}, " + f"chunk_d={chunk_d}).") + if any(b < 0 or b >= n_bands for b in bands): + raise ValueError( + f"k_bands.chunk_bands['{key_str}'] has a band id outside " + f"[0, {n_bands}).") + n_op = max(bands) + 1 + counts = [bands.count(b) for b in range(n_op)] + if min(counts) < 2: + # A band of a single chunk is <= chunk_d wide and would fall + # off the chunked kernel path onto a different quantization + # implementation (sc/kernels.py gates on `D > chunk_d`). + raise ValueError( + f"k_bands.chunk_bands['{key_str}'] leaves a band with " + f"{min(counts)} chunk(s); every band needs >= 2 so its " + f"width exceeds chunk_d={chunk_d}.") + bw = _band_widths(bands, widths, n_op) + prior = op_band_widths.setdefault(operator, bw) + if prior != bw: + # Ladders are per (op, layer-bucket) but membership is per + # (op, block); if widths drifted across blocks of a bucket the + # per-rung identity could not hold for all of them at once. + raise ValueError( + f"k_bands.chunk_bands['{key_str}'] gives band widths {bw}, " + f"but another block of '{operator}' gives {prior}. Band " + f"widths must be constant per operator so one ladder set " + f"can satisfy the iso-cost identity for every block.") + chunk_bands[(operator, block_idx)] = bands + + ladders: dict[tuple[str, int, int], list[list[int]]] = {} + for bucket_key, per_band in (section.get("ladders") or {}).items(): + operator, t_bucket, l_bucket = _parse_bucket_key(bucket_key) + # Band count is PER OPERATOR. `n_bands` is the MAXIMUM: an operator + # whose residual has too few chunks to give every band >= 2 uses + # fewer. Requiring exactly n_bands everywhere silently dropped every + # narrow projection (~19 chunks) out of the K-band path whenever + # n_bands was raised for down_proj (71 chunks), which is what made + # `--n-bands 16` cover 4 of 28 buckets. + n_op = len(per_band) + if n_op < 2 or n_op > n_bands: + raise ValueError( + f"k_bands.ladders['{bucket_key}'] has {n_op} ladders; " + f"expected between 2 and n_bands={n_bands}.") + op_bands = {b for (op, _), bm in chunk_bands.items() if op == operator + for b in bm} + if op_bands and max(op_bands) + 1 != n_op: + raise ValueError( + f"k_bands.ladders['{bucket_key}'] has {n_op} ladders but " + f"'{operator}' chunk_bands use {max(op_bands) + 1} bands.") + parent = self.get_levels( + operator=operator, + block_idx=l_bucket, total_blocks=self.layer_buckets) + band_ladders = [] + for b, rungs in enumerate(per_band): + rungs = [int(x) for x in rungs] + if len(rungs) != len(parent): + raise ValueError( + f"k_bands.ladders['{bucket_key}'] band {b} has " + f"{len(rungs)} rungs but the row ladder has " + f"{len(parent)}; the row's rung index indexes EVERY " + f"band ladder, so they must agree.") + if any(r <= 0 for r in rungs): + raise ValueError( + f"k_bands.ladders['{bucket_key}'] band {b} has a " + f"non-positive rung: {rungs}.") + # Above the halved cap the stream WRAPS -- the result is + # meaningless rather than merely worse -- and a band that + # "wins" by overrunning it would look like a Phase-3 gain. + if any(r > cap for r in rungs): + raise ValueError( + f"k_bands.ladders['{bucket_key}'] band {b} exceeds the " + f"stream-length cap {cap}: {rungs}.") + band_ladders.append(rungs) + widths = op_band_widths.get(operator) + if widths is None: + raise ValueError( + f"k_bands.ladders['{bucket_key}'] has no matching " + f"chunk_bands entry for operator '{operator}'.") + if len(widths) != n_op: + raise ValueError( + f"k_bands.ladders['{bucket_key}'] has {n_op} ladders but " + f"'{operator}' has {len(widths)} band widths.") + total = float(sum(widths)) + for k, parent_len in enumerate(parent): + # per-operator band count, NOT the global maximum + realized = sum(widths[b] * band_ladders[b][k] + for b in range(n_op)) / total + # ONE-SIDED. Overspending breaks the iso-compute claim and is a + # hard error. UNDERspending cannot: it only means the cell used + # less compute than the parent, so a quality win there is + # strictly stronger, not weaker. A two-sided check was correct + # only while the solver was two-sided; discrete water-filling + # legitimately leaves a fraction of a cycle unspent when the + # measurement grid is coarse, and rejecting that threw away + # valid allocations (v_proj:t0:l1 at 47.70 vs parent 48). + if realized - parent_len > self.K_BAND_ISO_COST_TOL: + raise ValueError( + f"k_bands.ladders['{bucket_key}'] OVERSPENDS at rung " + f"{k}: MAC-weighted band mean {realized:.4f} vs parent " + f"rung {parent_len} (tolerance " + f"{self.K_BAND_ISO_COST_TOL}). Phase 3 redistributes " + f"stream length, it does not spend more.") + if parent_len - realized > self.K_BAND_MAX_UNDERSPEND: + raise ValueError( + f"k_bands.ladders['{bucket_key}'] UNDERSPENDS at rung " + f"{k} by {parent_len - realized:.4f} cycles (limit " + f"{self.K_BAND_MAX_UNDERSPEND}). That is safe for the " + f"iso-compute claim but means the solver is leaving " + f"budget on the table — investigate the grid, do not " + f"just widen this.") + ladders[(operator, t_bucket, l_bucket)] = band_ladders + + missing = {op for (op, _) in chunk_bands} - {op for (op, _, _) in ladders} + if missing: + raise ValueError( + f"k_bands: operators {sorted(missing)} have chunk_bands but no " + f"ladders; they would silently run per-row.") + + self.k_band_count = n_bands + self.k_band_chunk_d = chunk_d + self.k_band_chunks = chunk_bands + self.k_band_ladders = ladders + self.k_band_residual_width = residual_width + + def get_k_bands( + self, + operator: Optional[str], + block_idx: Optional[int], + total_blocks: Optional[int], + *, + timestep: int = 0, + total_timesteps: int = 1, + ): + """(band_of_chunk, band_ladders) for one module, or None to run per-row. + + ``band_of_chunk[c]`` is the band owning residual chunk ``c`` (ascending + chunk order, tail last); ``band_ladders[b][k]`` is the stream length + band ``b`` runs when the row landed on rung ``k``. + """ + if not self.k_band_count or operator is None or block_idx is None: + return None + bands = self.k_band_chunks.get((operator, int(block_idx))) + if bands is None: + return None + t_bucket = _bucket_index(timestep, total_timesteps, self.timestep_buckets) + l_bucket = _bucket_index(block_idx, total_blocks, self.layer_buckets) \ + if total_blocks is not None else 0 + ladders = self.k_band_ladders.get((operator, t_bucket, l_bucket)) + if ladders is None: + return None + return bands, ladders + + def get_levels( + self, + timestep: int = 0, + total_timesteps: int = 1, + operator: Optional[str] = None, + block_idx: Optional[int] = None, + total_blocks: Optional[int] = None, + ) -> list[int]: + """Per-(operator, layer-bucket) ladder, falling back to the global one. + + The calibrated table historically carried ONE ladder shared by every + one of its 36 (op x layer-quartile) buckets, with only the thresholds + varying per bucket. Measured occupancy shows that wastes roughly half + the rungs: on 14B t32 the MLP buckets put ZERO MAC on rungs 0-1 (97, + 64) while qk puts 94-97% on rung 0 and nothing below rung 2 -- the two + populations live at opposite ends of a shared ladder, so each gets + about half the available resolution. + + When ``bucket_stoc_len_levels`` is empty this returns the global list + UNCHANGED (identity is preserved, not just equality), so a table + without per-group ladders behaves exactly as before. + + Bucket resolution deliberately mirrors ``get_thresholds`` so a bucket's + ladder and its thresholds can never disagree about which bucket it is. + """ + if (self.bucket_stoc_len_levels and operator + and block_idx is not None and total_blocks is not None): + t_bucket = _bucket_index(timestep, total_timesteps, + self.timestep_buckets) + l_bucket = _bucket_index(block_idx, total_blocks, + self.layer_buckets) + levels = self.bucket_stoc_len_levels.get( + (operator, t_bucket, l_bucket)) + if levels is not None: + return levels + if operator and operator in self.operator_default_stoc_len_levels: + return self.operator_default_stoc_len_levels[operator] + return self.stoc_len_levels + + def _escape_threshold_from_payload(self, bucket_payload) -> Optional[float]: + """t_esc = metric_mean + k * metric_std for one table payload. + + Returns None when the gate is off or the payload has no stats. + Deliberately NOT clamped to <= 1.0 — a t_esc >= 1.0 never fires + because the normalized metric support is [0, 1].""" + if self.escape_gate_k is None or not isinstance(bucket_payload, dict): + return None + mu = bucket_payload.get("metric_mean") + sigma = bucket_payload.get("metric_std") + if mu is None or sigma is None: + return None + return float(mu) + float(self.escape_gate_k) * float(sigma) + + def get_escape_threshold( + self, + timestep: int, + total_timesteps: int, + operator: Optional[str] = None, + block_idx: Optional[int] = None, + total_blocks: Optional[int] = None, + ) -> Optional[float]: + """Escape-gate threshold for one operator/timestep/block bucket. + + Mirrors :meth:`get_thresholds` lookup order (bucket, then operator + default). None = the gate cannot fire for this call (gate off, or no + stats for this bucket).""" + if self.escape_gate_k is None: + return None + if (self.bucket_escape_thresholds and operator and block_idx is not None + and total_blocks is not None): + t_bucket = _bucket_index(timestep, total_timesteps, self.timestep_buckets) + l_bucket = _bucket_index(block_idx, total_blocks, self.layer_buckets) + t_esc = self.bucket_escape_thresholds.get((operator, t_bucket, l_bucket)) + if t_esc is not None: + return t_esc + if operator and operator in self.operator_default_escape_thresholds: + return self.operator_default_escape_thresholds[operator] + return None + + def classify_level_values( + self, + timestep: int = 0, + total_timesteps: int = 1, + operator: Optional[str] = None, + block_idx: Optional[int] = None, + total_blocks: Optional[int] = None, + ) -> list[int]: + """Level-index -> stoc_len map for consumers of ``row_levels``. + + MUST resolve the SAME bucket that produced ``row_levels``. The SC + attention path classifies rows with adaptive_classify_rows (which uses + the bucket's own ladder via get_levels) and then maps index -> stream + length through here. When this ignored the bucket and returned the + global list, indices stayed in range (the ladders have equal length) + but every row ran at the WRONG stream length -- silently, since + nothing downstream cross-checks the two. It surfaced only as a + realized-cost reconciliation failure ("SC lengths outside the + adaptive/protected set") once per-group ladders diverged. + + With the escape gate ON and ``escape_stoc_len`` not already a rung, + escaped rows carry level index ``len(levels)``; this returns the + ladder plus that appended escape entry so index-driven dispatch loops + cover it. Gate off -- or escape length already a rung -- returns the + ladder itself, so the dispatch loop is byte-identical to the pre-gate + code. The appended entry intentionally breaks the descending-order + convention: this is an index map for dispatch, not a ladder.""" + levels = self.get_levels( + timestep=timestep, + total_timesteps=total_timesteps, + operator=operator, + block_idx=block_idx, + total_blocks=total_blocks, + ) + if (self.escape_gate_k is None + or int(self.escape_stoc_len) in levels): + return levels + return list(levels) + [int(self.escape_stoc_len)] + + def get_dispatch_metric(self, operator: Optional[str]) -> tuple[str, float]: + """(metric_name, sign) for one operator's per-row dispatch. + + Default ("amax", +1.0) — the original dispatch — for operators the + table did not switch (or when the table predates dispatch_metrics).""" + if operator and self.dispatch_metrics: + return self.dispatch_metrics.get(operator, ("amax", 1.0)) + return ("amax", 1.0) + + def get_protected_channels( + self, + operator: Optional[str] = None, + block_idx: Optional[int] = None, + unit_idx: Optional[int] = None, + ) -> Optional[list[int]]: + """Return protected input-channel indices for one linear module.""" + if not operator or block_idx is None: + return None + key = (operator, int(block_idx), unit_idx) + vals = self.protected_channel_indices.get(key) + if vals is not None: + return vals + return self.protected_channel_indices.get((operator, int(block_idx), None)) + def adaptive_classify_rows( metric: torch.Tensor, @@ -333,7 +916,15 @@ def adaptive_classify_rows( RowAssignment compatible with existing dispatch code. """ N = metric.shape[0] - levels = config.stoc_len_levels + # Per-group ladder when the table declares one; otherwise this IS + # config.stoc_len_levels, so the no-groups path is byte-identical. + levels = config.get_levels( + timestep=timestep, + total_timesteps=total_timesteps, + operator=operator, + block_idx=block_idx, + total_blocks=total_blocks, + ) n_levels = len(levels) # Empty row batch — e.g. a MoE expert that received ZERO tokens this forward @@ -402,7 +993,15 @@ def adaptive_classify_rows( total_blocks=total_blocks, ) if calibrated_thresholds is not None: - return _classify_rows_by_thresholds(metric_norm, levels, calibrated_thresholds) + assignment = _classify_rows_by_thresholds( + metric_norm, levels, calibrated_thresholds) + _apply_escape_gate( + assignment, metric_norm, config, + operator=operator, block_idx=block_idx, + total_blocks=total_blocks, + timestep=timestep, total_timesteps=total_timesteps, + ) + return assignment # No path matched (not free-boundary, no target_fractions, and no # calibrated thresholds for this operator/bucket). There is no closed-form @@ -416,6 +1015,73 @@ def adaptive_classify_rows( ) +def _apply_escape_gate( + assignment: RowAssignment, + metric_norm: torch.Tensor, + config: AdaptiveMPConfig, + *, + operator: Optional[str], + block_idx: Optional[int], + total_blocks: Optional[int], + timestep: int, + total_timesteps: int, +) -> None: + """In-place absolute escape gate (R7) on a threshold classification. + + Rows with normalized metric STRICTLY above t_esc_b = mu_b + k * sigma_b + (per-bucket constants precomputed at table load) are reassigned to + ``config.escape_stoc_len`` regardless of the band thresholds. The compare + happens in the same post-sign normalized space the band thresholds live + in, so no extra sign handling is needed (sign −1 metrics arrive already + negated; min–max normalization maps them to 1 − normalized(raw), exactly + as calibration stored them before computing mu/sigma). + + Escaped rows get level index ``len(stoc_len_levels)`` (a NEW index one + past the ladder) and a ``level_row_indices[escape_stoc_len]`` entry — + unless the escape length already is a ladder rung, in which case they are + folded into that rung's existing index. Gate off (``escape_gate_k`` None + ⇒ ``get_escape_threshold`` returns None) leaves the assignment untouched. + """ + t_esc = config.get_escape_threshold( + timestep=timestep, + total_timesteps=total_timesteps, + operator=operator, + block_idx=block_idx, + total_blocks=total_blocks, + ) + if t_esc is None or t_esc >= 1.0: + # t_esc >= 1.0 can never fire: metric_norm lives in [0, 1] with max + # EXACTLY 1.0 and the compare is strict. Skipping the tensor compare + # here is a shortcut, not a clamp — behavior is identical. + return + esc_mask = metric_norm > metric_norm.new_tensor(t_esc) + if not bool(esc_mask.any().item()): + return + # MUST be the SAME ladder the classification used (adaptive_classify_rows + # resolves it via get_levels). Reading config.stoc_len_levels here instead + # was latent-correct only while every bucket shared the global ladder: the + # rebuild below re-keys level_row_indices by stoc_len VALUE, so with a + # per-bucket ladder every row would be re-keyed to the GLOBAL rung value at + # its index and the bucket's ladder would be silently discarded. That is + # unobservable in realized_flop_avg_sl (the tracker prices the assignment it + # is handed), so the cell would look in-budget while running wrong lengths. + levels = config.get_levels( + timestep=timestep, + total_timesteps=total_timesteps, + operator=operator, + block_idx=block_idx, + total_blocks=total_blocks, + ) + esc_sl = int(config.escape_stoc_len) + esc_idx = levels.index(esc_sl) if esc_sl in levels else len(levels) + assignment.row_levels[esc_mask] = esc_idx + for level_idx, sl in enumerate(levels): + assignment.level_row_indices[sl] = torch.where( + assignment.row_levels == level_idx)[0] + if esc_idx == len(levels): + assignment.level_row_indices[esc_sl] = torch.where(esc_mask)[0] + + # ===================================================================== # Free-boundary MP (zero hyperparameter; offline oracle-search populated) # ===================================================================== diff --git a/scmp_kernels/sc/kernels.py b/scmp_kernels/sc/kernels.py index 86bdeb6..9f74d49 100644 --- a/scmp_kernels/sc/kernels.py +++ b/scmp_kernels/sc/kernels.py @@ -92,6 +92,7 @@ def clear_rng_cache(): _rng_seq_cache.clear() _enable_table_cache.clear() _k_table_cache.clear() + _cum_indicator_cache.clear() @@ -589,8 +590,44 @@ def _sc_matmul_per_head_bipolar( # Enable-Signal Host Functions # ============================================================================= -_enable_table_cache: dict[str, tuple[torch.Tensor, torch.Tensor]] = {} -_k_table_cache: dict[str, torch.Tensor] = {} +from collections import OrderedDict as _OrderedDict + +# Bounded LRU: autoregressive DECODE grows the attention inner dim every step, +# so per-(D, stoc_len) tables otherwise accumulate without bound (MP's 5 levels +# x ~50 decode steps x ~270MB @ ctx-4096 = tens of GB -> OOM). The cap is far +# above any fixed-ctx (PPL) working set, so eviction never occurs there and +# results are unaffected; on decode, stale-D entries are evicted (rebuild cost +# only). Tables are pure functions of the key, so eviction is always safe. +_ENABLE_TABLE_CACHE_MAX = int(os.environ.get("SC_ENABLE_TABLE_CACHE_MAX", "32")) +_enable_table_cache: "_OrderedDict[str, tuple[torch.Tensor, torch.Tensor]]" = _OrderedDict() +_k_table_cache: "_OrderedDict[str, torch.Tensor]" = _OrderedDict() +# Chunked-MLP cum_indicator, same key discipline as _enable_table_cache. +_cum_indicator_cache: "_OrderedDict[str, torch.Tensor]" = _OrderedDict() + + +def _cum_cache_enabled() -> bool: + """SC_CUM_INDICATOR_CACHE=0 restores the per-call rebuild. + + Exists so the bit-identity control run can be produced from the SAME + binary: the cache must not change a single output bit, only how often the + table is built. + """ + return os.environ.get("SC_CUM_INDICATOR_CACHE", "1") != "0" + + +def _lru_get(cache, key): + if key in cache: + cache.move_to_end(key) + return cache[key] + return None + + +def _lru_put(cache, key, value): + cache[key] = value + cache.move_to_end(key) + while len(cache) > _ENABLE_TABLE_CACHE_MAX: + cache.popitem(last=False) + return value def _resolve_rng_levels(sc_prec: int, rng_levels: Optional[int]) -> int: @@ -830,11 +867,45 @@ def _get_cached_enable_tables( stoc_len = 2 ** sc_prec grid_levels = _resolve_rng_levels(sc_prec, rng_levels) key = _enable_table_cache_key(config, sc_prec, device) + f"|sl={stoc_len}|rng={grid_levels}" - if key not in _enable_table_cache: - _enable_table_cache[key] = build_enable_tables( - rng_a, rng_b, sc_prec, stoc_len, rng_levels=grid_levels - ) - return _enable_table_cache[key] + cached = _lru_get(_enable_table_cache, key) + if cached is not None: + return cached + return _lru_put(_enable_table_cache, key, build_enable_tables( + rng_a, rng_b, sc_prec, stoc_len, rng_levels=grid_levels + )) + + +def _get_cached_cum_indicator( + key: str, + rng_b: torch.Tensor, + chunk_d: int, + stoc_len: int, + V: int, + device: torch.device, +) -> torch.Tensor: + """Cache the chunked-MLP cum_indicator across calls. + + The table is a pure function of (rng_b, chunk_d, stoc_len, V) and rng_b is + itself served from _get_cached_sequences, so every SCLinear in the model + that shares (config, sc_prec, stoc_len, rng_levels) builds a byte-identical + table. The non-chunked (attention) path has cached this since day one via + _get_cached_enable_tables; the chunked MLP path rebuilt it on EVERY call -- + measured at 37,621 rebuilds per 2048-token window on Qwen3-30B-A3B (128 + experts x 3 projections x 48 layers) and 983 on the 14B dense model, each + one a 4.3 MB int16 allocation plus a Triton launch. + + Returned tensors are READ-ONLY by contract: the matmul kernels only load + from cum_indicator, and nothing writes to it after build_cum_indicator_kernel. + """ + cached = _lru_get(_cum_indicator_cache, key) + if cached is not None: + return cached + cum_indicator = torch.zeros( + chunk_d, stoc_len + 1, V, dtype=torch.int16, device=device) + build_cum_indicator_kernel[(chunk_d,)]( + rng_b, cum_indicator, chunk_d, stoc_len, V, + ) + return _lru_put(_cum_indicator_cache, key, cum_indicator) def _get_cached_k_table( @@ -848,11 +919,12 @@ def _get_cached_k_table( stoc_len = 2 ** sc_prec grid_levels = _resolve_rng_levels(sc_prec, rng_levels) key = _enable_table_cache_key(config, sc_prec, device) + f"|k_only|sl={stoc_len}|rng={grid_levels}" - if key not in _k_table_cache: - _k_table_cache[key] = build_k_table_only( - rng_a, sc_prec, stoc_len, rng_levels=grid_levels - ) - return _k_table_cache[key] + cached = _lru_get(_k_table_cache, key) + if cached is not None: + return cached + return _lru_put(_k_table_cache, key, build_k_table_only( + rng_a, sc_prec, stoc_len, rng_levels=grid_levels + )) def enable_matmul_triton( @@ -1335,6 +1407,7 @@ def _sc_matmul_bipolar_mlp_chunked( a, b, sc_prec, k_table, rng_b, chunk_d, stoc_len=None, rng_levels: Optional[int] = None, + cum_cache_key: Optional[str] = None, ): """ Bipolar SC matmul for MLP with internal chunk_d loop. @@ -1361,12 +1434,20 @@ def _sc_matmul_bipolar_mlp_chunked( device = a.device output = torch.zeros(N, M, dtype=torch.float32, device=device) - # Build cum_indicator ONCE — all chunks share the same RNG sequences - cum_indicator = torch.zeros(chunk_d, stoc_len + 1, V, dtype=torch.int16, device=device) - build_cum_indicator_kernel[(chunk_d,)]( - rng_b, cum_indicator, - chunk_d, stoc_len, V, - ) + # Build cum_indicator ONCE — all chunks share the same RNG sequences. + # With a cache key, reuse it across CALLS too (see + # _get_cached_cum_indicator); without one, keep the per-call build so any + # future caller is bit-identical to the old behaviour by default. + if cum_cache_key is not None and _cum_cache_enabled(): + cum_indicator = _get_cached_cum_indicator( + cum_cache_key, rng_b, chunk_d, stoc_len, V, device) + else: + cum_indicator = torch.zeros( + chunk_d, stoc_len + 1, V, dtype=torch.int16, device=device) + build_cum_indicator_kernel[(chunk_d,)]( + rng_b, cum_indicator, + chunk_d, stoc_len, V, + ) # Tiled matmul params — adaptive tile size for small N/M if N <= 64 or M <= 64: @@ -1573,6 +1654,9 @@ def _sc_matmul_per_row_mlp( result = _sc_matmul_bipolar_mlp_chunked( a, b, sc_prec, k_table, rng_b, chunk_d, stoc_len=stoc_len, rng_levels=grid_levels, + cum_cache_key=( + _enable_table_cache_key(config, sc_prec, a.device) + + f"|cum|cd={chunk_d}|sl={stoc_len}|rng={grid_levels}"), ) if device.type != 'cuda': diff --git a/scmp_kernels/sc/matmul.py b/scmp_kernels/sc/matmul.py index ca11139..b4e70ae 100644 --- a/scmp_kernels/sc/matmul.py +++ b/scmp_kernels/sc/matmul.py @@ -28,6 +28,7 @@ _sc_matmul_per_head_bipolar, ) from ..quant.smoothquant import apply_smoothing +from .. import trace as _trace _VALID_GRANULARITIES = ("per_tensor", "per_row", "per_head") @@ -220,6 +221,34 @@ def _sc_matmul_impl( f"sc_matmul: granularity='per_head' currently only supports " f"mode='bipolar', got '{mode}'.") + # ---- precision trace (see scmp_kernels/trace.py) ------------------------- + # Host-side shape metadata only — no tensor reads, no device sync. The + # bool guard keeps the off-path cost at one attribute read. Placed after + # the validation gates so rejected calls are never recorded, and after the + # halving block so stoc_len is the TRUE cycle count. + if _trace._ENABLED: + if a.dim() == 3: + _batch, _rows, _d_in = a.shape + else: + _rows, _d_in = a.shape + _batch = 1 + _trace.record_matmul( + rows=int(_rows), d_in=int(_d_in), d_out=int(b.shape[-2]), + batch=int(_batch), + stoc_len=int(stoc_len) if stoc_len is not None else 2 ** sc_prec, + sc_prec=sc_prec, mode=mode, granularity=granularity, + # halving only takes effect in bipolar mode (guard above) — record + # whether it was APPLIED, not merely requested, to keep the field + # consistent with the effective stoc_len/rng_levels. + halve=bool(halve_bipolar_stoc_len and mode == "bipolar"), + # RESOLVED enable-grid size (mirrors _resolve_rng_levels), so the + # field has one meaning across halved and non-halved runs. + rng_levels=(int(rng_levels) if rng_levels is not None + else 2 ** sc_prec), + chunk_d=chunk_d, + smoothed=smooth_scales is not None, + ) + # ---- dispatch ------------------------------------------------------------ if granularity == "per_tensor": # Compute the per-tensor range on host (one .max / .min sync). diff --git a/scmp_kernels/trace.py b/scmp_kernels/trace.py new file mode 100644 index 0000000..a1961f4 --- /dev/null +++ b/scmp_kernels/trace.py @@ -0,0 +1,406 @@ +"""SC computation trace — precision/shape log for external HW simulators. + +Records, for every ``sc_matmul`` call, the *effective* precision +(``stoc_len`` = SC cycle count, post-halving; ``rng_levels`` = resolved +enable-grid size) together with the matmul shape and the caller-supplied +identity (operator, decoder block, sub-unit such as a MoE expert or +attention head). An energy/latency simulator can then compute e.g. +``MACs x energy_per_MAC(stoc_len)`` per group, or replay the ordered +per-call timeline. + +Design constraints (why it looks the way it does): + +* **Zero overhead when off.** The only cost on the hot path is reading one + module-level bool (``_ENABLED``). Callers guard with + ``if trace._ENABLED:`` before doing any work. +* **No device synchronization.** Only host-side metadata is recorded + (tensor *shapes* and plain ints) — never tensor values, ``.item()``, or + anything that would stall the async CUDA pipeline. +* **No heavy imports.** Pure stdlib; importable everywhere the package is. +* **Bounded memory.** Trace mode spills to disk every ``_SPILL`` records, + so week-long decode runs cannot OOM the host. +* **Thread-safe.** Context is thread-local; accumulation is locked + (relevant for e.g. RULER's ``--threads > 1`` prediction driver). + +Usage +----- +Enable via env (picked up at import time) or API:: + + SC_MP_TRACE=/path/out.json # enables; summary mode (default) + SC_MP_TRACE_MODE=trace # per-call JSONL instead + + from scmp_kernels import trace + trace.enable("/path/out.json", mode="summary") + +Applications tag identity before their matmuls (cheap; overwrite per call):: + + if trace._ENABLED: + trace.set_context(op="down_proj", block=12, unit=expert_idx) + +``flush()`` writes the file and resets by default. An ``atexit`` hook +flushes automatically (header carries ``"atexit": true`` so unattended +files are recognizable). Multi-config sweeps must call ``reset()`` before +and ``flush(path)`` after each config, or configs merge into one file. + +Coverage and known semantics +---------------------------- +* Only SC matmuls are recorded. FP16 ops — ``lm_head``, embeddings, the + MoE router ``gate``, norms, softmax — do NOT appear; total model work is + strictly larger than the trace's MAC sum (header carries a ``coverage`` + note). +* Rows dispatched to a drop level (``stoc_len <= 0``) are zero-filled by + the callers without an ``sc_matmul`` call and are NOT recorded; with + such configs (none shipped today) per-op rows no longer sum to tokens. +* Uniform (non-MP) attention runs record one 3D call with + ``batch = B*H`` and ``unit = null``; MP attention runs record per-head + 2D calls (``unit`` = head index). Consumers must accept both flavors. +* ``d_in``/``d_out`` are part of the summary key, so + ``rows * d_in * d_out == macs`` holds for every group with + ``dims_vary: false``. The exception is bounded memory: at most + ``_MAX_SHAPES`` (env ``SC_MP_TRACE_MAX_SHAPES``, default 32) distinct + shapes per base key get their own group, because autoregressive decode + grows attention's KV length every step and would otherwise mint one + group per step. Shapes past the cap fold into one catch-all group with + *representative* dims and ``dims_vary: true``; there and only there, + ``rows * d_in * d_out != macs``. ``macs``/``rows``/``row_cycles`` are + always exact — they accumulate true per-call values. + Traces written before 2026-07-29 keyed without dims and can collapse + two static shapes (e.g. an operator's protected-channel slice sharing a + ``stoc_len`` with its main slice) into a fictitious shape; see + ``hpca_results/llm/ppl/mp_best/repair_traces.py`` for the offline fix. + +Output schemas +-------------- +summary (JSON):: + + {"schema": "scmp-trace-summary-v1", + "header": {...static config + coverage...}, + "groups": [{"block": 12, "op": "down_proj", "unit": null, + "stoc_len": 128, "sc_prec": 8, "halve": true, + "granularity": "per_row", "mode": "bipolar", + "rng_levels": 128, "chunk_d": 128, "smoothed": true, + "d_in": 9728, "d_out": 2560, "dims_vary": false, + "calls": 64, "rows": 81920, "macs": 5497558138880, + "row_cycles": 10485760}, ...]} + +trace (JSONL; header line first, then one object per call in execution +order — ``seq``):: + + {"seq": 18211, "block": 12, "op": "qk", "unit": 7, "rows": 1024, + "d_in": 128, "d_out": 1024, "batch": 1, "stoc_len": 64, ...} + +``rows`` counts a-operand rows across the batch (``batch * N`` for 3D); +``macs = batch * rows_per_slice * d_out * d_in``; ``row_cycles = rows * +stoc_len`` (per-row stream cost — the first-order SC latency/energy proxy). +""" +from __future__ import annotations + +import atexit +import json +import os +import sys +import threading +from typing import Optional + +__all__ = [ + "enable", "disable", "reset", "flush", + "set_context", "clear_context", "is_enabled", +] + +# --------------------------------------------------------------------------- +# Module state. _ENABLED is read directly by hot-path guards. +# --------------------------------------------------------------------------- +_ENABLED: bool = False +_MODE: str = "summary" # "summary" | "trace" +_PATH: Optional[str] = None +_SEQ: int = 0 +_LOCK = threading.Lock() +_TLS = threading.local() # per-thread (op, block, unit) +# summary accumulator: +# key -> [calls, rows, macs, row_cycles, d_in, d_out, dims_vary] +# where key = (block, op, unit, stoc_len, sc_prec, halve, mode, granularity, +# rng_levels, chunk_d, smoothed, d_in, d_out) +_SUMMARY: dict = {} +# base key (the above minus d_in/d_out) -> set of shapes that own a group. +# Bounds the group count when dims genuinely vary per call (KV growth). +_SHAPES: dict = {} +_TRACE: list = [] +_SPILL = 100_000 # trace-mode records per disk spill +_FH = None # open spill file handle (trace mode) + +# Distinct (d_in, d_out) per base key that get their own exact group before +# the rest fold into a dims_vary catch-all. Static-shape operators use 1-2. +try: + _MAX_SHAPES = int(os.environ.get("SC_MP_TRACE_MAX_SHAPES", "32")) +except ValueError: + _MAX_SHAPES = 32 +if _MAX_SHAPES < 1: + _MAX_SHAPES = 1 + + +def is_enabled() -> bool: + return _ENABLED + + +def enable(path: str, mode: str = "summary") -> None: + """Turn tracing on. ``mode``: 'summary' (aggregate) or 'trace' (per-call). + + Refuses to switch mode while unflushed records exist — call ``flush()`` + or ``reset()`` first (otherwise the buffered records of the old mode + would be silently stranded).""" + global _ENABLED, _MODE, _PATH + if mode not in ("summary", "trace"): + raise ValueError(f"trace mode must be 'summary' or 'trace', got {mode!r}") + if _SEQ and mode != _MODE: + raise RuntimeError( + f"trace.enable: {_SEQ} unflushed records in mode {_MODE!r}; " + f"flush() or reset() before switching to {mode!r}.") + _MODE = mode + _PATH = path + _ENABLED = True + + +def disable() -> None: + global _ENABLED + _ENABLED = False + + +def reset() -> None: + """Drop all accumulated records (keeps enabled state and context).""" + global _SEQ, _FH + with _LOCK: + _SUMMARY.clear() + _SHAPES.clear() + _TRACE.clear() + _SEQ = 0 + if _FH is not None: + _FH.close() + _FH = None + + +def set_context(op: Optional[str], block: Optional[int], + unit: Optional[int] = None) -> None: + """Tag subsequent sc_matmul calls with (operator, block, sub-unit). + + ``unit`` disambiguates calls sharing (op, block): MoE expert index, + attention head index. Thread-local, so concurrent forwards (e.g. RULER + --threads > 1) do not cross-tag. Overwrite per call site.""" + _TLS.ctx = (op, block, unit) + + +def clear_context() -> None: + set_context(None, None, None) + + +def _open_spill(): + """Open the trace-mode spill file and write the header line.""" + global _FH + out = _PATH + os.makedirs(os.path.dirname(os.path.abspath(out)) or ".", exist_ok=True) + _FH = open(out, "w") + hdr = _header() + hdr["n_records"] = None # unknown until the run ends + _FH.write(json.dumps({"schema": "scmp-trace-v1", "header": hdr}) + "\n") + + +def record_matmul(*, rows: int, d_in: int, d_out: int, batch: int, + stoc_len: int, sc_prec: int, mode: str, granularity: str, + halve: bool, rng_levels: Optional[int], chunk_d: int, + smoothed: bool = False) -> None: + """Record one sc_matmul call. Called from the sc_matmul entry when enabled. + + All arguments are host-side ints/strs (shape metadata) — never tensors. + ``rows`` is a-rows per slice; total a-rows = batch * rows. + ``rng_levels`` must be the RESOLVED enable-grid size (never None).""" + global _SEQ + op, block, unit = getattr(_TLS, "ctx", (None, None, None)) + rows_total = batch * rows + macs = batch * rows * d_out * d_in + row_cycles = rows_total * stoc_len + with _LOCK: + if _MODE == "trace": + _TRACE.append({ + "seq": _SEQ, + "op": op, "block": block, "unit": unit, + "rows": rows_total, "d_in": d_in, "d_out": d_out, + "batch": batch, + "stoc_len": stoc_len, "sc_prec": sc_prec, "halve": halve, + "mode": mode, "granularity": granularity, + "rng_levels": rng_levels, "chunk_d": chunk_d, + "smoothed": smoothed, + "macs": macs, "row_cycles": row_cycles, + }) + # Bounded memory: spill to disk periodically so long decode runs + # cannot OOM the host (an OOM kill would also skip atexit and + # lose everything). + if len(_TRACE) >= _SPILL: + if _FH is None: + _open_spill() + for rec in _TRACE: + _FH.write(json.dumps(rec) + "\n") + _TRACE.clear() + else: + # d_in/d_out ARE part of the key, so rows*d_in*d_out == macs holds + # per group and consumers get real shapes. The reason they were + # once excluded is bounded memory: attention's KV length grows per + # decode step, which would mint one group per step. That is handled + # by a per-base-key shape cap instead — the first _MAX_SHAPES + # distinct (d_in, d_out) get exact groups, anything beyond folds + # into one catch-all group carrying representative dims and + # dims_vary=True. Static-shape operators (linears, and the + # protected-channel slice that shares their stoc_len) never reach + # the cap, so they no longer collapse into a fictitious shape. + base = (block, op, unit, stoc_len, sc_prec, halve, mode, + granularity, rng_levels, chunk_d, smoothed) + shapes = _SHAPES.get(base) + if shapes is None: + shapes = _SHAPES[base] = set() + if (d_in, d_out) in shapes: + key = base + (d_in, d_out) + elif len(shapes) < _MAX_SHAPES: + shapes.add((d_in, d_out)) + key = base + (d_in, d_out) + else: + key = base + (None, None) # over cap: dims genuinely vary + acc = _SUMMARY.get(key) + if acc is None: + _SUMMARY[key] = [1, rows_total, macs, row_cycles, + d_in, d_out, key[-1] is None] + else: + acc[0] += 1 + acc[1] += rows_total + acc[2] += macs + acc[3] += row_cycles + if acc[4] != d_in or acc[5] != d_out: + acc[6] = True + _SEQ += 1 + + +def _header() -> dict: + """Static config snapshot written once per output file. + + ``scramble_masks`` is the env-knob ECHO (int when parseable), not the + per-call resolved mask count — that is ``min(scramble_masks, + 2**sc_prec)`` using each record's own ``sc_prec``, and it only applies + when ``owen_mode`` is ``bitrev``.""" + try: + masks = int(os.environ.get("SC_SCRAMBLE_MASKS", "64")) + except ValueError: + # Garbage env value: echo it raw rather than crash flush() — the + # kernel's own validation raises at matmul time, not here. + masks = os.environ.get("SC_SCRAMBLE_MASKS") + return { + "owen_mode": os.environ.get("SC_OWEN_MODE", "bitrev"), + "scramble_masks": masks, + "mode": _MODE, + "n_records": _SEQ, + "coverage": "sc_matmul only — FP16 ops (lm_head, embeddings, MoE " + "router gate, norms, softmax) are NOT recorded", + } + + +def flush(path: Optional[str] = None, header_extra: Optional[dict] = None, + reset_after: bool = True) -> Optional[str]: + """Write accumulated records; returns the output path (None if nothing). + + ``path`` overrides the enable-time path (e.g. one file per sweep + config) — EXCEPT in trace mode once spilling has begun: the spill file + (enable-time path) is already on disk, so the remainder is appended + there and the override is ignored with a warning. ``header_extra`` + merges caller metadata (model id, eval tag) into the header.""" + global _FH, _SEQ + with _LOCK: + if _SEQ == 0: + return None + if _MODE == "trace" and _FH is not None: + # Already spilling: finish the spill file. + if path is not None and path != _PATH: + print(f"[trace] WARN: flush path override {path!r} ignored — " + f"trace already spilled to {_PATH!r}.", file=sys.stderr) + for rec in _TRACE: + _FH.write(json.dumps(rec) + "\n") + # The header line (written at first spill) has n_records=null — + # the final count is only known now. ALWAYS append the trailer so + # the file carries its count even for a bare flush(); merge any + # caller metadata into it. + trailer = {"trailer": True, "n_records": _SEQ} + if header_extra: + trailer.update(header_extra) + _FH.write(json.dumps(trailer) + "\n") + _FH.close() + _FH = None + out = _PATH + else: + out = path or _PATH + if out is None: + return None + os.makedirs(os.path.dirname(os.path.abspath(out)) or ".", + exist_ok=True) + header = _header() + if header_extra: + header.update(header_extra) + if _MODE == "trace": + with open(out, "w") as f: + f.write(json.dumps({"schema": "scmp-trace-v1", + "header": header}) + "\n") + for rec in _TRACE: + f.write(json.dumps(rec) + "\n") + else: + groups = [] + for (block, op, unit, sl, prec, halve, mode, gran, rngl, + ckd, smoothed, _kd_in, _kd_out), ( + calls, rows, macs, cyc, d_in, d_out, + vary) in sorted( + _SUMMARY.items(), + key=lambda kv: (kv[0][0] if kv[0][0] is not None + else -1, str(kv[0][1]), + str(kv[0][2]), -kv[0][3])): + groups.append({ + "block": block, "op": op, "unit": unit, + "stoc_len": sl, "sc_prec": prec, "halve": halve, + "mode": mode, "granularity": gran, + "rng_levels": rngl, "chunk_d": ckd, + "smoothed": smoothed, + "d_in": d_in, "d_out": d_out, "dims_vary": vary, + "calls": calls, "rows": rows, "macs": macs, + "row_cycles": cyc, + }) + with open(out, "w") as f: + json.dump({"schema": "scmp-trace-summary-v1", + "header": header, "groups": groups}, f, + indent=1) + if reset_after: + _SUMMARY.clear() + _SHAPES.clear() + _TRACE.clear() + _SEQ = 0 + return out + + +def _atexit_flush() -> None: + # Safety net: apps that enable via env but never call flush() still get + # their file at interpreter exit. The header/trailer is marked so + # consumers can tell an unattended flush (which may span multiple + # configs) from a deliberate per-config one. + if _ENABLED and _SEQ: + try: + flush(header_extra={"atexit": True}) + except Exception: + pass + + +atexit.register(_atexit_flush) + +# Env-driven enable (picked up when the package is first imported). A bad +# SC_MP_TRACE_MODE must not take down the whole package import (the callers' +# try/except guards catch ImportError, not ValueError) — warn and fall back. +_env_path = os.environ.get("SC_MP_TRACE", "").strip() +if _env_path: + _env_mode = (os.environ.get("SC_MP_TRACE_MODE", "summary").strip() + or "summary") + if _env_mode not in ("summary", "trace"): + print(f"[trace] WARN: SC_MP_TRACE_MODE={_env_mode!r} is not " + f"'summary'|'trace'; falling back to 'summary'.", + file=sys.stderr) + _env_mode = "summary" + enable(_env_path, _env_mode) diff --git a/tests/sweep_rows.sbatch b/tests/sweep_rows.sbatch new file mode 100644 index 0000000..271fc77 --- /dev/null +++ b/tests/sweep_rows.sbatch @@ -0,0 +1,21 @@ +#!/bin/bash +#SBATCH --job-name=sweep_rows +#SBATCH --account=nbleier_owned1 +#SBATCH --reservation=rtx6000_arph_nodes +#SBATCH --partition=gpu-rtx6000 +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=12 +#SBATCH --mem=180G +#SBATCH --time=00:30:00 +#SBATCH --output=/scratch/nbleier_owned_root/nbleier_owned1/shared_data/allenjin/hpca/logs/sweep_rows_%j.out +set +u +source ~/.bashrc +conda activate annstention +set -u +export HF_HOME=/nfs/turbo/coe-nbleier/allenjin/hf_cache +export TMPDIR=/scratch/nbleier_owned_root/nbleier_owned1/shared_data/allenjin/hpca/tmp +export SC_OWEN_MODE=bitrev +export SC_SCRAMBLE_MASKS=64 +cd /home/allenjin/Projects/scmp_llm/kernels +echo "=== row-throughput sweep ===" +python -u tests/sweep_rows_throughput.py diff --git a/tests/sweep_rows_throughput.py b/tests/sweep_rows_throughput.py new file mode 100644 index 0000000..1b3b79b --- /dev/null +++ b/tests/sweep_rows_throughput.py @@ -0,0 +1,97 @@ +"""Throughput vs row count for the chunked-MLP SC path. + +Measured B1 result: caching the cum_indicator is bit-identical but worth 0.4%. +The two shapes measured there instead showed a UTILIZATION problem -- + + 30B expert (N=128, D=2048, M=768 ) 2.01e8 MAC in 2.748 ms = 7.3e10 MAC/s + 14B dense (N=2048, D=5120, M=2048) 2.15e10 MAC in 53.6 ms = 4.0e11 MAC/s + +87x the arithmetic for 19.5x the time. Qwen3-30B-A3B gives each expert only +2048 tokens * top_k 8 / 128 experts = 128 rows per window, so if throughput is +row-starved then batching windows (B*128 rows per expert) is the real lever and +caching tables never was. + +This sweeps N at FIXED D, M to get the actual curve, so the batching speedup is +predicted BEFORE any batching code is written. It also checks bit-identity of +the per-row outputs across N, which is the property window batching must not +break: rows are independent, so row i's result must not depend on how many +other rows were in the call. +""" + +import os +import sys +import time + +import torch + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from scmp_kernels.sc.kernels import ( # noqa: E402 + _sc_matmul_per_row_mlp, + clear_rng_cache, +) + +D, M = 2048, 768 # Qwen3-30B-A3B expert projection +ROW_COUNTS = [128, 256, 512, 1024, 2048, 4096] +STOC_LEN = 64 +CHUNK_D = 128 +SC_PREC = 8 + + +def call(a, b): + return _sc_matmul_per_row_mlp( + a, b, mode="bipolar", sc_prec=SC_PREC, chunk_d=CHUNK_D, + stoc_len=STOC_LEN) + + +def main(): + if not torch.cuda.is_available(): + print("FAIL: no CUDA device") + return 1 + print(f"device: {torch.cuda.get_device_name(0)}") + print(f"D={D} M={M} stoc_len={STOC_LEN} chunk_d={CHUNK_D}\n") + os.environ["SC_CUM_INDICATOR_CACHE"] = "1" + clear_rng_cache() + + torch.manual_seed(0) + b = torch.randn(M, D, device="cuda") + a_full = torch.randn(max(ROW_COUNTS), D, device="cuda") + + base = None + print(f"{'N':>6} {'ms/call':>9} {'MAC/s':>11} {'vs N=128':>9} " + f"{'ms per 128 rows':>16}") + for n in ROW_COUNTS: + a = a_full[:n].contiguous() + call(a, b) # warm + torch.cuda.synchronize() + reps = 20 + t0 = time.perf_counter() + for _ in range(reps): + call(a, b) + torch.cuda.synchronize() + ms = (time.perf_counter() - t0) / reps * 1e3 + macs = n * D * M / (ms * 1e-3) + base = base or macs + per128 = ms / (n / 128) + print(f"{n:6d} {ms:9.3f} {macs:11.3e} {macs/base:8.2f}x " + f"{per128:16.3f}") + + # Row independence: row i's output must not depend on batch size. This is + # exactly what window batching relies on, and what would silently break if + # tile selection changed the reduction order. + print("\nrow-independence across N (bit-exact):") + a_small = a_full[:128].contiguous() + ref = call(a_small, b) + for n in (256, 1024, 4096): + big = call(a_full[:n].contiguous(), b) + same = torch.equal(ref, big[:128]) + if not same: + d = (ref - big[:128]).abs().max().item() + print(f" N={n:5d}: *** DIFFERS *** max|d|={d:.3e}") + else: + print(f" N={n:5d}: first 128 rows bit-identical to the N=128 call") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/test_group_ladders.py b/tests/test_group_ladders.py new file mode 100644 index 0000000..1fa70bd --- /dev/null +++ b/tests/test_group_ladders.py @@ -0,0 +1,351 @@ +"""Per-(op x layer-bucket) ladders in the calibrated MP table. + +The table historically carried ONE ladder shared by all 36 (op x layer) +buckets, with only thresholds varying per bucket. Measured rung occupancy +(14B t32, ladder [97,64,49,32,24,20]) shows the MLP buckets put ZERO MAC on +rungs 0-1 while qk puts 94-97% on rung 0 and nothing below rung 2, so a shared +ladder gives each population roughly half its resolution. + +These tests pin the two properties that matter: + 1. a table WITHOUT per-group ladders is byte-identical to the old behaviour + (get_levels returns the very same list object), and + 2. a table WITH them routes each bucket to its own ladder. +""" + +import json +import tempfile +import unittest +from pathlib import Path + +import torch + +from scmp_kernels.mp.config import ( + AdaptiveMPConfig, + _extract_group_levels, + adaptive_classify_rows, +) + +GLOBAL = [96, 64, 48, 32] +ATTN = [112, 96, 80, 64] +MLP = [48, 32, 24, 16] + + +def _bucket(thresholds, levels=None): + payload = {"thresholds": list(thresholds), + "metric_mean": 0.5, "metric_std": 0.1} + if levels is not None: + payload["stoc_len_levels"] = list(levels) + return payload + + +def _write(payload): + payload = {"stoc_len_levels": list(GLOBAL), **payload} + path = Path(tempfile.mkdtemp()) / "table.json" + path.write_text(json.dumps(payload)) + return str(path) + + +def _config(path): + return AdaptiveMPConfig( + stoc_len_levels=list(GLOBAL), + threshold_table_path=path, + timestep_buckets=1, + layer_buckets=4, + ) + + +class ExtractGroupLevelsTest(unittest.TestCase): + + def test_absent_returns_none_not_the_global_list(self): + self.assertIsNone(_extract_group_levels({"thresholds": []}, "k")) + + def test_strictly_descending_required(self): + # duplicates would silently merge in level_row_indices, which is keyed + # by the stoc_len VALUE -- rows of the lower rung would run at the + # higher rung's length with no error raised anywhere + with self.assertRaisesRegex(ValueError, "strictly"): + _extract_group_levels({"stoc_len_levels": [64, 32, 32, 16]}, "k") + with self.assertRaisesRegex(ValueError, "strictly"): + _extract_group_levels({"stoc_len_levels": [16, 32, 48, 64]}, "k") + + def test_needs_at_least_two_rungs(self): + with self.assertRaisesRegex(ValueError, ">= 2 rungs"): + _extract_group_levels({"stoc_len_levels": [64]}, "k") + + +class NoGroupLaddersIsUnchangedTest(unittest.TestCase): + + def setUp(self): + self.cfg = _config(_write({ + "buckets": {f"qk:t0:l{i}": _bucket([0.7, 0.5, 0.3]) + for i in range(4)}, + })) + + def test_get_levels_returns_the_global_list_itself(self): + got = self.cfg.get_levels( + operator="qk", block_idx=0, total_blocks=4) + self.assertIs(got, self.cfg.stoc_len_levels) + + def test_unknown_operator_also_falls_back(self): + self.assertIs( + self.cfg.get_levels(operator="nope", block_idx=0, total_blocks=4), + self.cfg.stoc_len_levels) + + def test_missing_block_context_falls_back(self): + self.assertIs(self.cfg.get_levels(operator="qk"), + self.cfg.stoc_len_levels) + + +class GroupLaddersTest(unittest.TestCase): + + def setUp(self): + buckets = {} + for i in range(4): + buckets[f"qk:t0:l{i}"] = _bucket([0.7, 0.5, 0.3], ATTN) + buckets[f"down_proj:t0:l{i}"] = _bucket([0.7, 0.5, 0.3], MLP) + buckets[f"o_proj:t0:l{i}"] = _bucket([0.7, 0.5, 0.3]) # no group + self.cfg = _config(_write({"buckets": buckets})) + + def test_each_bucket_gets_its_own_ladder(self): + self.assertEqual( + self.cfg.get_levels(operator="qk", block_idx=0, total_blocks=4), + ATTN) + self.assertEqual( + self.cfg.get_levels(operator="down_proj", block_idx=0, + total_blocks=4), + MLP) + + def test_bucket_without_a_group_ladder_still_falls_back(self): + self.assertIs( + self.cfg.get_levels(operator="o_proj", block_idx=0, + total_blocks=4), + self.cfg.stoc_len_levels) + + def test_layer_bucket_is_resolved_like_thresholds(self): + # same bucket arithmetic as get_thresholds: block 3 of 4 -> l3 + for block in range(4): + self.assertEqual( + self.cfg.get_levels(operator="qk", block_idx=block, + total_blocks=4), + ATTN) + + def test_classification_dispatches_on_the_group_ladder(self): + metric = torch.linspace(0.0, 1.0, 32) + attn = adaptive_classify_rows( + metric, self.cfg, operator="qk", block_idx=0, total_blocks=4) + mlp = adaptive_classify_rows( + metric, self.cfg, operator="down_proj", block_idx=0, + total_blocks=4) + self.assertEqual(sorted(attn.level_row_indices), sorted(ATTN)) + self.assertEqual(sorted(mlp.level_row_indices), sorted(MLP)) + # every row is assigned exactly once, in both groups + for assignment, levels in ((attn, ATTN), (mlp, MLP)): + total = sum(v.numel() for v in assignment.level_row_indices.values()) + self.assertEqual(total, metric.numel()) + self.assertEqual(len(levels), len(assignment.level_row_indices)) + + def test_empty_row_batch_uses_the_group_ladder(self): + # MoE experts routinely receive zero tokens in a forward + out = adaptive_classify_rows( + torch.empty(0), self.cfg, operator="down_proj", block_idx=0, + total_blocks=4) + self.assertEqual(sorted(out.level_row_indices), sorted(MLP)) + + +class GroupLadderRungCountTest(unittest.TestCase): + + def test_group_may_have_a_different_number_of_rungs(self): + cfg = _config(_write({ + "buckets": { + # 3 rungs -> 2 thresholds, against a 4-rung global default + "qk:t0:l0": _bucket([0.6, 0.3], [112, 96, 64]), + }, + })) + self.assertEqual( + cfg.get_levels(operator="qk", block_idx=0, total_blocks=4), + [112, 96, 64]) + + def test_threshold_count_is_validated_against_the_group_not_the_global(self): + with self.assertRaisesRegex(ValueError, "length 3, expected 2"): + _config(_write({ + "buckets": {"qk:t0:l0": _bucket([0.7, 0.5, 0.3], + [112, 96, 64])}, + })) + + +class ClassifyLevelValuesTest(unittest.TestCase): + """index -> stream-length map MUST resolve the same bucket as the rows. + + The SC attention path classifies rows with adaptive_classify_rows (which + uses the bucket ladder) and then maps index -> stoc_len via + classify_level_values. When the latter ignored the bucket and returned + the global list, indices stayed in range -- the ladders have equal length + -- but every attention row ran at the WRONG stream length, with nothing + downstream cross-checking the two. It only surfaced as a realized-cost + reconciliation failure once per-group ladders diverged from the global + one, i.e. it would have been silent in every prior wave. + """ + + def _cfg(self, grouped): + buckets = {} + for i in range(4): + buckets[f"qk:t0:l{i}"] = _bucket( + [0.7, 0.5, 0.3], ATTN if grouped else None) + buckets[f"down_proj:t0:l{i}"] = _bucket( + [0.7, 0.5, 0.3], MLP if grouped else None) + return _config(_write({"buckets": buckets})) + + def test_resolves_the_bucket_ladder(self): + cfg = self._cfg(grouped=True) + self.assertEqual( + cfg.classify_level_values(operator="qk", block_idx=0, + total_blocks=4), ATTN) + self.assertEqual( + cfg.classify_level_values(operator="down_proj", block_idx=0, + total_blocks=4), MLP) + + def test_matches_what_classification_actually_used(self): + # the invariant that was violated: dispatch map == classification map + cfg = self._cfg(grouped=True) + for op in ("qk", "down_proj"): + rows = adaptive_classify_rows( + torch.linspace(0.0, 1.0, 16), cfg, operator=op, + block_idx=0, total_blocks=4) + values = cfg.classify_level_values( + operator=op, block_idx=0, total_blocks=4) + self.assertEqual(sorted(rows.level_row_indices), sorted(values)) + + def test_no_groups_is_unchanged(self): + cfg = self._cfg(grouped=False) + self.assertEqual(cfg.classify_level_values(), cfg.stoc_len_levels) + self.assertEqual( + cfg.classify_level_values(operator="qk", block_idx=0, + total_blocks=4), + cfg.stoc_len_levels) + + def test_no_arg_call_still_works(self): + # existing callers (test_mp_escape_gate) invoke it with no arguments + self.assertEqual(self._cfg(grouped=True).classify_level_values(), + GLOBAL) + + +class EscapeGateUsesTheBucketLadderTest(unittest.TestCase): + """The escape gate must rebuild against the ladder classification USED. + + Third instance of the same family as ClassifyLevelValuesTest above. + ``_apply_escape_gate`` re-keys ``level_row_indices`` by stoc_len VALUE for + every rung, so reading the GLOBAL ladder there while + ``adaptive_classify_rows`` classified against the per-bucket ladder makes + every row run at the global rung value for its index -- the bucket ladder is + discarded outright. It is invisible downstream: SCLinear just iterates + whatever keys it is handed (sc_common.py:581) and the MP tracker prices the + assignment rather than measuring it, so realized_flop_avg_sl still lands in + band and the cell looks valid. + + Dormant in every shipped config -- all 8 deployed t32/t48 cells run + escape_gate_k=2.0 with ZERO per-bucket ladders -- so the fix is a no-op on + them (pinned by NoBucketLadderIsUnaffected below) and fires only for + per-band/per-group work. + """ + + ESC = 128 # not a rung of GLOBAL, ATTN or MLP -> gets its own index + + def _cfg(self, grouped, k=2.0): + # _bucket sets metric_mean=0.5, metric_std=0.1 -> t_esc = 0.5 + k*0.1 + buckets = {} + for i in range(4): + buckets[f"down_proj:t0:l{i}"] = _bucket( + [0.7, 0.5, 0.3], MLP if grouped else None) + buckets[f"qk:t0:l{i}"] = _bucket( + [0.7, 0.5, 0.3], ATTN if grouped else None) + return AdaptiveMPConfig( + stoc_len_levels=list(GLOBAL), + threshold_table_path=_write({"buckets": buckets}), + timestep_buckets=1, + layer_buckets=4, + escape_gate_k=k, + escape_stoc_len=self.ESC, + ) + + def _classify(self, cfg, op): + # linspace spans [0,1] so metric_norm == metric; t_esc = 0.7 fires. + return adaptive_classify_rows( + torch.linspace(0.0, 1.0, 64), cfg, operator=op, + block_idx=0, total_blocks=4) + + def test_escaped_assignment_is_keyed_by_the_bucket_ladder(self): + cfg = self._cfg(grouped=True) + got = self._classify(cfg, "down_proj") + self.assertEqual(sorted(got.level_row_indices), + sorted(MLP + [self.ESC]), + "gate re-keyed the assignment against the wrong " + "ladder; rows would run at global rung values") + # and the other bucket keeps ITS ladder, not down_proj's and not global + self.assertEqual(sorted(self._classify(cfg, "qk").level_row_indices), + sorted(ATTN + [self.ESC])) + + def test_no_global_rung_value_leaks_in(self): + # the pre-fix failure mode, stated positively: GLOBAL-only values + # (96, 64) must not appear for a bucket whose ladder is MLP + keys = set(self._classify(self._cfg(grouped=True), + "down_proj").level_row_indices) + self.assertFalse(keys & (set(GLOBAL) - set(MLP)), + f"global rung values leaked into the assignment: " + f"{sorted(keys)}") + + def test_every_row_still_assigned_exactly_once(self): + got = self._classify(self._cfg(grouped=True), "down_proj") + total = sum(v.numel() for v in got.level_row_indices.values()) + self.assertEqual(total, 64) + # escaped rows are exactly those strictly above t_esc = 0.7 + metric = torch.linspace(0.0, 1.0, 64) + expected = torch.where(metric > 0.7)[0] + self.assertTrue( + torch.equal(got.level_row_indices[self.ESC].sort().values, + expected)) + + def test_gate_that_cannot_fire_leaves_the_bucket_ladder_alone(self): + # k=6 -> t_esc = 1.1 >= 1.0, early return before any rebuild + got = self._classify(self._cfg(grouped=True, k=6.0), "down_proj") + self.assertEqual(sorted(got.level_row_indices), sorted(MLP)) + + +class NoBucketLadderIsUnaffectedTest(unittest.TestCase): + """Regression guard: the deployed shape (gate on, no per-bucket ladder). + + All 8 deployed t32/t48 cells are exactly this. The gate fix must not move + them at all, so their archived PPLs stay reproducible. + """ + + ESC = 128 + + def _cfg(self): + buckets = {f"down_proj:t0:l{i}": _bucket([0.7, 0.5, 0.3]) + for i in range(4)} + return AdaptiveMPConfig( + stoc_len_levels=list(GLOBAL), + threshold_table_path=_write({"buckets": buckets}), + timestep_buckets=1, + layer_buckets=4, + escape_gate_k=2.0, + escape_stoc_len=self.ESC, + ) + + def test_keys_are_the_global_ladder_plus_escape(self): + cfg = self._cfg() + got = adaptive_classify_rows( + torch.linspace(0.0, 1.0, 64), cfg, operator="down_proj", + block_idx=0, total_blocks=4) + self.assertEqual(sorted(got.level_row_indices), + sorted(GLOBAL + [self.ESC])) + + def test_get_levels_still_returns_the_global_list_itself(self): + # identity, not equality -- what the fix reads must BE the global list + cfg = self._cfg() + self.assertIs( + cfg.get_levels(operator="down_proj", block_idx=0, total_blocks=4), + cfg.stoc_len_levels) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_k_bands.py b/tests/test_k_bands.py new file mode 100644 index 0000000..d11be61 --- /dev/null +++ b/tests/test_k_bands.py @@ -0,0 +1,293 @@ +"""K-band (Phase 3) table schema: per-group stream lengths inside a row. + +Row dispatch is unchanged -- one metric, one rung index per row. What changes +is that the residual contraction axis is partitioned into bands of whole +quantization chunks, and band b runs rung k at its own length. Setting every +band's ladder equal to the row ladder reproduces the per-row parent, which is +what makes the refinement unable to lose. + +The load-time checks are the safety net for the whole experiment: a malformed +band spec that quietly degraded to per-row would produce a cell that LOOKS like +a Phase-3 result, and realized_flop_avg_sl would not reveal it (the tracker +prices the assignment it is handed rather than measuring the kernel). +""" + +import json +import tempfile +import unittest +from pathlib import Path + +from scmp_kernels.mp.config import ( + AdaptiveMPConfig, + residual_chunk_widths, +) + +GLOBAL = [96, 64, 48, 32] +# down_proj-shaped: 9728 - 584 protected = 9144 = 71 x 128 + 56 (ragged tail) +R_DOWN = 9144 +N_CHUNKS = len(residual_chunk_widths(R_DOWN, 128)) # 72 + + +def _bucket(thresholds): + return {"thresholds": list(thresholds), + "metric_mean": 0.5, "metric_std": 0.1} + + +def _halves(): + """Band map splitting the 72 chunks into two equal-count halves.""" + return [0] * (N_CHUNKS // 2) + [1] * (N_CHUNKS - N_CHUNKS // 2) + + +def _write(k_bands, buckets=None): + payload = { + "stoc_len_levels": list(GLOBAL), + "buckets": buckets if buckets is not None else { + f"down_proj:t0:l{i}": _bucket([0.7, 0.5, 0.3]) for i in range(4)}, + } + if k_bands is not None: + payload["k_bands"] = k_bands + path = Path(tempfile.mkdtemp()) / "table.json" + path.write_text(json.dumps(payload)) + return str(path) + + +def _cfg(k_bands, buckets=None): + return AdaptiveMPConfig( + stoc_len_levels=list(GLOBAL), + threshold_table_path=_write(k_bands, buckets), + timestep_buckets=1, + layer_buckets=4, + ) + + +def _spec(ladders, bands=None, n_bands=2, n_blocks=4): + return { + "n_bands": n_bands, + "chunk_d": 128, + "residual_width": {"down_proj": R_DOWN}, + "chunk_bands": {f"down_proj:b{b}": (bands if bands is not None + else _halves()) + for b in range(n_blocks)}, + "ladders": {f"down_proj:t0:l{i}": ladders for i in range(4)}, + } + + +class ChunkGeometryTest(unittest.TestCase): + + def test_ragged_tail_is_its_own_chunk_and_comes_last(self): + w = residual_chunk_widths(R_DOWN, 128) + self.assertEqual(len(w), 72) + self.assertEqual(w[:-1], [128] * 71) + self.assertEqual(w[-1], 56) + self.assertEqual(sum(w), R_DOWN) + + def test_exact_multiple_has_no_tail(self): + self.assertEqual(residual_chunk_widths(512, 128), [128] * 4) + + +class ParentIsInTheSpaceTest(unittest.TestCase): + """All bands = the row ladder must load, and must be the identity.""" + + def test_parent_ladder_in_every_band_loads(self): + cfg = _cfg(_spec([list(GLOBAL), list(GLOBAL)])) + self.assertEqual(cfg.k_band_count, 2) + bands, ladders = cfg.get_k_bands("down_proj", 0, 4) + self.assertEqual(ladders, [GLOBAL, GLOBAL]) + self.assertEqual(len(bands), N_CHUNKS) + + def test_absent_section_disables_the_path(self): + cfg = _cfg(None) + self.assertEqual(cfg.k_band_count, 0) + self.assertIsNone(cfg.get_k_bands("down_proj", 0, 4)) + + def test_operator_without_bands_runs_per_row(self): + cfg = _cfg(_spec([list(GLOBAL), list(GLOBAL)])) + self.assertIsNone(cfg.get_k_bands("q_proj", 0, 4)) + + +class IsoCostIdentityTest(unittest.TestCase): + """Phase 3 REDISTRIBUTES stream length; it must not spend more.""" + + def test_overspending_band_is_rejected(self): + hot = [x + 16 for x in GLOBAL] # every rung longer, no payback + with self.assertRaisesRegex(ValueError, "OVERSPENDS"): + _cfg(_spec([hot, list(GLOBAL)])) + + def test_gross_underspend_is_rejected_as_a_solver_bug(self): + # Underspending is SAFE for the iso-compute claim, but a solver leaving + # 8 cycles unspent is broken, not conservative. + cold = [max(x - 16, 1) for x in GLOBAL] + with self.assertRaisesRegex(ValueError, "UNDERSPENDS"): + _cfg(_spec([cold, list(GLOBAL)])) + + def test_small_underspend_is_ALLOWED(self): + # Discrete water-filling legitimately leaves a fraction of a cycle + # unspent on a coarse grid. Rejecting that threw away valid + # allocations (v_proj:t0:l1 realized 47.70 against a parent rung of 48). + # A cheaper cell that still wins is a STRONGER result, not a failure. + w0, w1 = 36 * 128, 35 * 128 + 56 + cold, warm = [], [] + for L in GLOBAL: + # shave ~0.5 cycle off the MAC-weighted mean + cold.append(L - 1) + warm.append(L) + cfg = _cfg(_spec([cold, warm])) + _, ladders = cfg.get_k_bands("down_proj", 0, 4) + self.assertEqual(ladders, [cold, warm]) + for k, L in enumerate(GLOBAL): + realized = (w0 * cold[k] + w1 * warm[k]) / R_DOWN + self.assertLess(realized, L, "should be under the parent rung") + self.assertLess(L - realized, cfg.K_BAND_MAX_UNDERSPEND) + + def test_legal_redistribution_within_tolerance_loads(self): + # widths: band0 = 36*128 = 4608, band1 = 35*128 + 56 = 4536; R = 9144. + # Move d cycles onto band0 and pay it back from band1 exactly: + # 4608*(L+d) + 4536*(L-e) = 9144*L -> e = d * 4608/4536 + w0, w1 = 36 * 128, 35 * 128 + 56 + self.assertEqual(w0 + w1, R_DOWN) + hot, cold = [], [] + for L in GLOBAL: + d = 8 + e = round(d * w0 / w1) + hot.append(L + d) + cold.append(L - e) + cfg = _cfg(_spec([hot, cold])) + _, ladders = cfg.get_k_bands("down_proj", 0, 4) + self.assertEqual(ladders, [hot, cold]) + # and the realized MAC-weighted mean really is the parent + for k, L in enumerate(GLOBAL): + realized = (w0 * hot[k] + w1 * cold[k]) / R_DOWN + self.assertLess(abs(realized - L), cfg.K_BAND_ISO_COST_TOL) + + +class RungIndexResolverTest(unittest.TestCase): + """Band ladders are indexed by the ROW LADDER's rung index. + + Regression: the runtime originally sized its dispatch loop with + ``classify_level_values``, which APPENDS the escape length as an extra + dispatch index when the gate is on and escape_stoc_len is not already a + rung. Band ladders carry one entry per real rung, so the loop ran one index + past every band ladder and died with IndexError inside o_proj's forward. + ``get_levels`` is the correct resolver; the escape slot is handled + explicitly. These pin the two lists apart so the distinction cannot be + quietly lost again. + """ + + ESC = 128 # deliberately NOT a rung of GLOBAL + + def _cfg(self, gate=True): + return AdaptiveMPConfig( + stoc_len_levels=list(GLOBAL), + threshold_table_path=_write(_spec([list(GLOBAL), list(GLOBAL)])), + timestep_buckets=1, layer_buckets=4, + escape_gate_k=2.0 if gate else None, + escape_stoc_len=self.ESC, + ) + + def test_classify_level_values_is_longer_than_the_ladder_with_the_gate_on(self): + cfg = self._cfg(gate=True) + ladder = cfg.get_levels(operator="down_proj", block_idx=0, + total_blocks=36) + dispatch = cfg.classify_level_values(operator="down_proj", block_idx=0, + total_blocks=36) + self.assertEqual(len(ladder), len(GLOBAL)) + self.assertEqual(len(dispatch), len(GLOBAL) + 1, + "escape entry should be appended for dispatch") + self.assertEqual(dispatch[-1], self.ESC) + + def test_band_ladders_match_get_levels_not_the_dispatch_map(self): + cfg = self._cfg(gate=True) + _, band_ladders = cfg.get_k_bands("down_proj", 0, 36) + ladder = cfg.get_levels(operator="down_proj", block_idx=0, + total_blocks=36) + for b, lad in enumerate(band_ladders): + self.assertEqual( + len(lad), len(ladder), + f"band {b} must have one entry per REAL rung; sizing the " + f"dispatch loop off classify_level_values overruns it") + + def test_escape_length_already_a_rung_appends_nothing(self): + # gate on, escape length IS a rung -> escaped rows fold into it, so + # dispatch map and ladder have equal length and the band loop's + # explicit escape slot simply finds no rows + cfg = AdaptiveMPConfig( + stoc_len_levels=list(GLOBAL), + threshold_table_path=_write(_spec([list(GLOBAL), list(GLOBAL)])), + timestep_buckets=1, layer_buckets=4, + escape_gate_k=2.0, escape_stoc_len=GLOBAL[0], + ) + self.assertEqual( + len(cfg.classify_level_values(operator="down_proj", block_idx=0, + total_blocks=36)), + len(cfg.get_levels(operator="down_proj", block_idx=0, + total_blocks=36))) + + +class MalformedSpecIsRejectedTest(unittest.TestCase): + + def test_single_band_is_not_a_band_split(self): + with self.assertRaisesRegex(ValueError, "n_bands must be >= 2"): + _cfg(_spec([list(GLOBAL)], n_bands=1)) + + def test_band_with_one_chunk_falls_off_the_chunked_kernel_path(self): + lonely = [0] + [1] * (N_CHUNKS - 1) + with self.assertRaisesRegex(ValueError, r"needs >= 2"): + _cfg(_spec([list(GLOBAL), list(GLOBAL)], bands=lonely)) + + def test_chunk_count_mismatch_is_rejected(self): + with self.assertRaisesRegex(ValueError, "entries but the residual has"): + _cfg(_spec([list(GLOBAL), list(GLOBAL)], bands=[0, 1] * 4)) + + def test_rung_count_must_match_the_row_ladder(self): + short = GLOBAL[:3] + with self.assertRaisesRegex(ValueError, "rungs but the row ladder has"): + _cfg(_spec([short, short])) + + def test_band_id_out_of_range_is_rejected(self): + bad = _halves() + bad[0] = 5 + with self.assertRaisesRegex(ValueError, "outside"): + _cfg(_spec([list(GLOBAL), list(GLOBAL)], bands=bad)) + + def test_missing_residual_width_is_rejected(self): + spec = _spec([list(GLOBAL), list(GLOBAL)]) + spec.pop("residual_width") + with self.assertRaisesRegex(ValueError, "residual_width is required"): + _cfg(spec) + + def test_ladders_without_chunk_bands_is_rejected(self): + spec = _spec([list(GLOBAL), list(GLOBAL)]) + spec["chunk_bands"] = {} + with self.assertRaisesRegex(ValueError, "no matching chunk_bands"): + _cfg(spec) + + def test_chunk_bands_without_ladders_would_silently_run_per_row(self): + spec = _spec([list(GLOBAL), list(GLOBAL)]) + spec["ladders"] = {} + with self.assertRaisesRegex(ValueError, "no ladders"): + _cfg(spec) + + def test_band_widths_must_agree_across_blocks_of_an_operator(self): + # ladders are per (op, layer-bucket) but membership is per (op, block); + # drifting widths mean no single ladder set can hold the identity + spec = _spec([list(GLOBAL), list(GLOBAL)]) + skewed = [0] * 40 + [1] * (N_CHUNKS - 40) + spec["chunk_bands"]["down_proj:b2"] = skewed + with self.assertRaisesRegex(ValueError, "must be constant per operator"): + _cfg(spec) + + def test_rung_above_the_halved_cap_is_rejected(self): + # Above 2**(sc_prec-1) the stream WRAPS -- meaningless, not merely + # worse -- so a band that "won" by overrunning the cap would read as a + # Phase-3 gain. The allocator's search grid is clamped too; this is the + # backstop for a hand-written or mis-solved table. + with self.assertRaisesRegex(ValueError, "exceeds the stream-length cap"): + _cfg(_spec([[136, 64, 48, 32], [56, 64, 48, 32]])) + + def test_non_positive_rung_is_rejected(self): + with self.assertRaisesRegex(ValueError, "non-positive rung"): + _cfg(_spec([[96, 64, 48, 0], [96, 64, 48, 64]])) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_trace_summary.py b/tests/test_trace_summary.py new file mode 100644 index 0000000..093f676 --- /dev/null +++ b/tests/test_trace_summary.py @@ -0,0 +1,108 @@ +"""Summary-trace grouping invariants (no GPU, no torch). + +The regression these guard: `d_in`/`d_out` used to be excluded from the +summary key so KV growth during decode could not mint a group per step. That +also merged *static* shapes whenever they shared a stoc_len — in particular an +operator's protected-channel slice landing on the same rung as its main slice — +producing groups whose reported dims describe a matmul that was never run. +""" +import importlib.util +import json +import os + +# Load trace.py by path: it is pure stdlib, while scmp_kernels/__init__.py +# pulls in torch. Keeps these invariants testable on a login node. +_PATH = os.path.join(os.path.dirname(__file__), "..", "scmp_kernels", "trace.py") +_spec = importlib.util.spec_from_file_location("_scmp_trace_under_test", _PATH) +trace = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(trace) + + +def _rec(op, block, stoc_len, d_in, d_out, rows=8, unit=None): + trace.set_context(op, block, unit) + trace.record_matmul(rows=rows, d_in=d_in, d_out=d_out, batch=1, + stoc_len=stoc_len, sc_prec=8, mode="bipolar", + granularity="per_row", halve=True, rng_levels=128, + chunk_d=128, smoothed=True) + + +def _flush(tmp_path, name="t.json"): + out = str(tmp_path / name) + trace.flush(path=out) + with open(out) as f: + return json.load(f)["groups"] + + +def _fresh(tmp_path): + trace.reset() + trace.enable(str(tmp_path / "unused.json"), mode="summary") + + +def test_pc_slice_sharing_a_rung_stays_a_separate_group(tmp_path): + """The exact mp_best t96 failure: 9144-wide main + 584-wide protected + slice, both at stoc_len 128, must NOT merge into one 584-wide group.""" + _fresh(tmp_path) + for _ in range(3): + _rec("down_proj", 0, 128, 9144, 2560, rows=100) # main slice + _rec("down_proj", 0, 128, 584, 2560, rows=100) # protected slice + groups = _flush(tmp_path) + + assert len(groups) == 2, f"expected 2 shape groups, got {len(groups)}" + assert {g["d_in"] for g in groups} == {9144, 584} + for g in groups: + assert g["dims_vary"] is False + assert g["rows"] * g["d_in"] * g["d_out"] == g["macs"] + assert sum(g["calls"] for g in groups) == 6 + + +def test_macs_identity_holds_for_every_non_varying_group(tmp_path): + _fresh(tmp_path) + for sl in (32, 48, 96, 128): + _rec("up_proj", 1, sl, 2483, 9728, rows=7) + _rec("up_proj", 1, sl, 77, 9728, rows=5) + _rec("qk", 2, 64, 128, 2048, rows=64) + groups = _flush(tmp_path) + + assert groups, "no groups recorded" + for g in groups: + assert not g["dims_vary"] + assert g["rows"] * g["d_in"] * g["d_out"] == g["macs"], g + + +def test_kv_growth_is_bounded_by_the_shape_cap(tmp_path): + """Decode-style growth must not mint an unbounded number of groups.""" + _fresh(tmp_path) + cap = trace._MAX_SHAPES + n = cap + 40 + total_macs = 0 + for step in range(n): + kv = 128 + step + _rec("av", 3, 64, kv, 128, rows=4) + total_macs += 4 * kv * 128 + groups = _flush(tmp_path) + + assert len(groups) == cap + 1, ( + f"expected {cap} exact + 1 catch-all, got {len(groups)}") + varying = [g for g in groups if g["dims_vary"]] + assert len(varying) == 1 + assert varying[0]["calls"] == n - cap + # Totals stay exact even where the shape is only representative. + assert sum(g["macs"] for g in groups) == total_macs + assert sum(g["calls"] for g in groups) == n + for g in groups: + if not g["dims_vary"]: + assert g["rows"] * g["d_in"] * g["d_out"] == g["macs"] + + +def test_reset_clears_the_shape_registry(tmp_path): + _fresh(tmp_path) + _rec("o_proj", 0, 128, 4055, 2560) + trace.reset() + trace.enable(str(tmp_path / "unused.json"), mode="summary") + _rec("o_proj", 0, 128, 4055, 2560) + groups = _flush(tmp_path, "t2.json") + + assert len(groups) == 1 + assert groups[0]["calls"] == 1, "reset() leaked state into the next run" + trace.reset() + trace.disable() diff --git a/tests/verify_cum_cache.py b/tests/verify_cum_cache.py new file mode 100644 index 0000000..d532f71 --- /dev/null +++ b/tests/verify_cum_cache.py @@ -0,0 +1,94 @@ +"""B1 verification: the chunked-MLP cum_indicator cache must be BIT-IDENTICAL. + +Runs the same chunked-MLP SC matmul with the cache on and off (same binary, +switched by SC_CUM_INDICATOR_CACHE) and requires torch.equal -- not allclose. +Bit-identity is the property that lets 30B results stay comparable with the +4B/8B/14B results produced before the cache existed. + +Also reports the speedup, which is the whole point of the change: the table was +rebuilt on EVERY sc_matmul call (37,621 times per 2048-token window on +Qwen3-30B-A3B, 983 on 14B), each a ~4.3 MB int16 alloc plus a Triton launch. +""" + +import os +import sys +import time + +import torch + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from scmp_kernels.sc.kernels import ( # noqa: E402 + _sc_matmul_per_row_mlp, + _cum_indicator_cache, + clear_rng_cache, +) + +# Qwen3-30B-A3B expert-ish shape: 128 rows/expert, moe_intermediate 768. +SHAPES = [ + ("30B expert (N=128, D=2048, M=768)", 128, 2048, 768), + ("14B dense (N=2048, D=5120, M=2048)", 2048, 5120, 2048), +] +STOC_LENS = [96, 64, 48, 32, 24, 16] +CHUNK_D = 128 +SC_PREC = 8 + + +def run_once(a, b, stoc_len, cache_on, clear=True): + os.environ["SC_CUM_INDICATOR_CACHE"] = "1" if cache_on else "0" + if clear: + clear_rng_cache() + return _sc_matmul_per_row_mlp( + a, b, mode="bipolar", sc_prec=SC_PREC, chunk_d=CHUNK_D, + stoc_len=stoc_len) + + +def main(): + if not torch.cuda.is_available(): + print("FAIL: no CUDA device") + return 1 + print(f"device: {torch.cuda.get_device_name(0)}") + print(f"torch: {torch.__version__}\n") + + ok = True + for label, N, D, M in SHAPES: + torch.manual_seed(0) + a = torch.randn(N, D, device="cuda") + b = torch.randn(M, D, device="cuda") + print(f"--- {label}") + for stoc_len in STOC_LENS: + ref = run_once(a, b, stoc_len, cache_on=False) + got = run_once(a, b, stoc_len, cache_on=True) + # second cached call exercises the HIT path, not just the fill + hit = run_once(a, b, stoc_len, cache_on=True, clear=False) + torch.cuda.synchronize() + exact = torch.equal(ref, got) and torch.equal(ref, hit) + if not exact: + d1 = (ref - got).abs().max().item() + d2 = (ref - hit).abs().max().item() + print(f" sl={stoc_len:3d} *** NOT BIT-IDENTICAL *** " + f"max|d| fill={d1:.3e} hit={d2:.3e}") + ok = False + else: + print(f" sl={stoc_len:3d} bit-identical (fill and hit)") + + # timing: repeated calls, which is what the model actually does + reps = 30 + for cache_on in (False, True): + run_once(a, b, 64, cache_on) # warm triton + fill + torch.cuda.synchronize() + t0 = time.perf_counter() + for _ in range(reps): + run_once(a, b, 64, cache_on, clear=False) + torch.cuda.synchronize() + dt = (time.perf_counter() - t0) / reps * 1e3 + print(f" cache={'ON ' if cache_on else 'OFF'} " + f"{dt:8.3f} ms/call") + print(f" cache entries held: {len(_cum_indicator_cache)}\n") + + print("RESULT:", "PASS - bit-identical everywhere" if ok else "FAIL") + return 0 if ok else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/verify_cum_cache.sbatch b/tests/verify_cum_cache.sbatch new file mode 100644 index 0000000..887b9d5 --- /dev/null +++ b/tests/verify_cum_cache.sbatch @@ -0,0 +1,23 @@ +#!/bin/bash +#SBATCH --job-name=verify_cum_cache +#SBATCH --account=nbleier_owned1 +#SBATCH --reservation=rtx6000_arph_nodes +#SBATCH --partition=gpu-rtx6000 +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=12 +#SBATCH --mem=180G +#SBATCH --time=00:30:00 +#SBATCH --output=/scratch/nbleier_owned_root/nbleier_owned1/shared_data/allenjin/hpca/logs/verify_cum_cache_%j.out +set +u +source ~/.bashrc +conda activate annstention +set -u +export HF_HOME=/nfs/turbo/coe-nbleier/allenjin/hf_cache +export TMPDIR=/scratch/nbleier_owned_root/nbleier_owned1/shared_data/allenjin/hpca/tmp +export SC_OWEN_MODE=bitrev +export SC_SCRAMBLE_MASKS=64 +cd /home/allenjin/Projects/scmp_llm/kernels +echo "=== B1 bit-identity + speedup ===" +python -u tests/verify_cum_cache.py +echo "=== per-group ladder runtime tests (need torch) ===" +python -u -m unittest tests.test_group_ladders -v 2>&1 | tail -20