From e061384aa41df3b6b0ab68a1df415a8133cd17d2 Mon Sep 17 00:00:00 2001 From: heroarmor <162866837+heroarmor@users.noreply.github.com> Date: Sun, 5 Jul 2026 14:18:05 -0400 Subject: [PATCH 1/7] fix(mp): FLOP-weighted calibration budget, per-row qk MP, drop alpha/beta - calibrate_mp_thresholds.py: budget is now sum(rows * macs * stoc_len) instead of sum(rows * stoc_len), so an "avgX" target means X in real compute, not row count. macs_per_row (out*in for linears, N*head_dim for qk/av) is threaded through add() -> the global lambda solve + iso-budget check. qk recording is switched to per-row. Validated: realized FLOP-weighted avg = 47.97 for target 48. - sc_attention.py: _sc_qk MP is now per-query-row (||Q_row||_inf), matching the per_row quant scale; removed the per-head classifier and static per-head path. - quant_sc_main.py: removed all alpha/beta remnants -- --mp_alpha/--mp_beta and the 14 per-op variants, the crashing AdaptiveMPConfig(alpha=,beta=, operator_params=) construction, and the run-name tag. --adaptive_mp now requires --adaptive_mp_table. - sc_attention.py + sc_mlp.py: repaired the adaptive_classify_rows signature drift vs kernel 9358239 (6 call sites passed the old arg order -> TypeError, so the adaptive path could not run at all). Invalidates existing calibration threshold tables (new cost model -> new allocations); re-calibration required. Co-Authored-By: Claude Opus 4.8 (1M context) --- qdit/sc_integration/sc_attention.py | 134 +++++++++++++--------------- qdit/sc_integration/sc_mlp.py | 8 +- scripts/calibrate_mp_thresholds.py | 97 +++++++++++++++----- scripts/quant_sc_main.py | 126 +++----------------------- 4 files changed, 156 insertions(+), 209 deletions(-) diff --git a/qdit/sc_integration/sc_attention.py b/qdit/sc_integration/sc_attention.py index e12c914..65f31b5 100644 --- a/qdit/sc_integration/sc_attention.py +++ b/qdit/sc_integration/sc_attention.py @@ -321,12 +321,12 @@ def _sc_linear_dynamic_mp(self, x, weight, bias, operator, chunk_d=0, grouped=Fa if self.sc_controller.adaptive_mp_config is not None: assignment = adaptive_classify_rows( row_metric, - self.sc_controller.current_timestep, - self.sc_controller.total_timesteps, self.sc_controller.adaptive_mp_config, operator=operator, block_idx=self.block_idx, total_blocks=self.sc_controller.total_blocks, + timestep=self.sc_controller.current_timestep, + total_timesteps=self.sc_controller.total_timesteps, ) else: mp_config = self.sc_controller.mp_config @@ -426,12 +426,12 @@ def _sc_linear_combined_mp(self, x, weight, bias, operator, dispatch, if self.sc_controller.adaptive_mp_config is not None: assignment = adaptive_classify_rows( row_metric, - self.sc_controller.current_timestep, - self.sc_controller.total_timesteps, self.sc_controller.adaptive_mp_config, operator=operator, block_idx=self.block_idx, total_blocks=self.sc_controller.total_blocks, + timestep=self.sc_controller.current_timestep, + total_timesteps=self.sc_controller.total_timesteps, ) else: mp_config = self.sc_controller.mp_config @@ -645,95 +645,85 @@ def _sc_linear(self, x, weight, bias, operator=None, chunk_d=0, # ================================================================= def _sc_qk(self, q, k): - """SC q@k^T matmul — supports per-head mixed precision.""" + """SC q@k^T matmul — per-query-row mixed precision. + + MP granularity is per query-row (per (batch, head)), consistent with the + per_row quantization scale. Rows are classified by ``||Q_row||_inf``. The + old per-head stoc_len path was removed so MP and quant granularity match. + """ q_scaled = q * self.scale B, H, N, D = q_scaled.shape - head_stoc_lens = self.sc_controller.get_group_stoc_lens( - self.block_idx, "qk") + has_dynamic_mp = (self.sc_controller.adaptive_mp_config is not None + or self.sc_controller.mp_config is not None) + + if not has_dynamic_mp: + # Uniform precision — existing fast path (batched kernel). + stoc_len = self.sc_controller.get_stoc_len(self.block_idx, "qk") + sc_prec = self.sc_controller.resolve_sc_prec(stoc_len) + return self._sc_qk_uniform(q_scaled, k, sc_prec, stoc_len) - # Dynamic MP: compute per-head stoc_lens from Q magnitude - if head_stoc_lens is None and (self.sc_controller.adaptive_mp_config is not None - or self.sc_controller.mp_config is not None): - q_metric = q_scaled.float().abs().amax(dim=(0, 2, 3)) # [H] - MetricProfiler.record(q_metric, self.sc_controller.current_timestep, - self.block_idx, "qk") + # Dynamic per-query-row MP (adaptive or fixed): one stoc_len per query + # row within each (batch, head). Mirrors _sc_av's per-row path and keeps + # MP granularity aligned with the per_row quant scale. + BH = B * H + q_flat = q_scaled.reshape(BH, N, D).float() + k_flat = k.reshape(BH, N, D).float() + output = torch.zeros(BH, N, N, device=q.device, dtype=torch.float32) + + matmul_fn = self._get_matmul_fn() + baseline_stoc_len = self.sc_controller.stoc_len + compute_baseline = 0 + compute_actual = 0.0 + + for i in range(BH): + row_metric = q_flat[i].abs().amax(dim=-1) # [N] = ||Q_row||_inf + if i == 0: # profile / log once per (t, block) + MetricProfiler.record(row_metric, self.sc_controller.current_timestep, + self.block_idx, "qk") if self.sc_controller.adaptive_mp_config is not None: assignment = adaptive_classify_rows( - q_metric, - self.sc_controller.current_timestep, - self.sc_controller.total_timesteps, + row_metric, self.sc_controller.adaptive_mp_config, operator="qk", block_idx=self.block_idx, total_blocks=self.sc_controller.total_blocks, + timestep=self.sc_controller.current_timestep, + total_timesteps=self.sc_controller.total_timesteps, ) else: mp_config = self.sc_controller.mp_config assignment = classify_rows_by_metric( - q_metric, mp_config.stoc_len_levels, mp_config.level_fractions) - - MPDistributionLogger.log( - self.sc_controller.current_timestep, self.block_idx, - "qk", assignment, H) - - head_stoc_lens = [0] * H - for sl, heads in assignment.level_row_indices.items(): - for h_idx in heads: - head_stoc_lens[h_idx.item()] = sl + row_metric, mp_config.stoc_len_levels, mp_config.level_fractions) - if head_stoc_lens is None: - # Uniform precision — existing fast path - stoc_len = self.sc_controller.get_stoc_len(self.block_idx, "qk") - sc_prec = self.sc_controller.resolve_sc_prec(stoc_len) - return self._sc_qk_uniform(q_scaled, k, sc_prec, stoc_len) - - # Mixed: group heads by stoc_len, compute within each group - output = torch.zeros(B, H, N, N, device=q.device, dtype=torch.float32) - stoc_len_to_heads: dict[int, list[int]] = defaultdict(list) - for h, sl in enumerate(head_stoc_lens): - stoc_len_to_heads[sl].append(h) - - baseline_stoc_len = self.sc_controller.stoc_len - compute_baseline = B * H * N * N * D * baseline_stoc_len - compute_actual = 0.0 + if i == 0: + MPDistributionLogger.log( + self.sc_controller.current_timestep, self.block_idx, + "qk", assignment, N) - for sl, heads in stoc_len_to_heads.items(): - # Pruned heads (stoc_len=0): output already zeroed - if sl == 0: - continue - compute_actual += B * len(heads) * N * N * D * sl + k_i = k_flat[i] # [N, D] + for sl, rows in assignment.level_row_indices.items(): + if len(rows) == 0 or sl == 0: + continue # pruned rows: output already zeroed + n_rows = len(rows) + compute_baseline += n_rows * N * D * baseline_stoc_len + compute_actual += n_rows * N * D * sl - sp = self.sc_controller.resolve_sc_prec(sl) - config = self._get_sc_config(D, sp) - - for h in heads: - # Extract single head: [B, 1, N, D] -> [B, N, D] - q_h = q_scaled[:, h].float() - k_h = k[:, h].float() - - matmul_fn = self._get_matmul_fn() - if self.sc_mode == "bipolar": - output[:, h] = matmul_fn( - q_h, k_h, granularity="per_head", mode="bipolar", - sc_prec=sp, config=config, stoc_len=sl, - rng_levels=self._rng_levels(sl), - ) - else: - for b_idx in range(B): - output[b_idx, h] = matmul_fn( - q_h[b_idx], k_h[b_idx], - granularity="per_tensor", mode=self.sc_mode, - sc_prec=sp, config=config, stoc_len=sl, - rng_levels=self._rng_levels(sl), - ) + sp = self.sc_controller.resolve_sc_prec(sl) + config = self._get_sc_config(D, sp) + q_sub = q_flat[i][rows] # [n_rows, D] + output[i, rows] = matmul_fn( + q_sub, k_i, + granularity="per_row", group_a=1, group_b=1, + mode=self.sc_mode, sc_prec=sp, config=config, + stoc_len=sl, rng_levels=self._rng_levels(sl)) MPDistributionLogger.log_compute( self.sc_controller.current_timestep, self.block_idx, "qk", compute_baseline, compute_actual) - return output + return output.reshape(B, H, N, N) def _sc_qk_uniform(self, q_scaled, k, sc_prec, stoc_len): """Uniform precision QK — fully batched kernel or per-head loop.""" @@ -852,12 +842,12 @@ def _sc_av(self, attn, v): if self.sc_controller.adaptive_mp_config is not None: assignment = adaptive_classify_rows( row_max, - self.sc_controller.current_timestep, - self.sc_controller.total_timesteps, self.sc_controller.adaptive_mp_config, operator="av", block_idx=self.block_idx, total_blocks=self.sc_controller.total_blocks, + timestep=self.sc_controller.current_timestep, + total_timesteps=self.sc_controller.total_timesteps, ) else: mp_config = self.sc_controller.mp_config diff --git a/qdit/sc_integration/sc_mlp.py b/qdit/sc_integration/sc_mlp.py index 031198b..215936a 100644 --- a/qdit/sc_integration/sc_mlp.py +++ b/qdit/sc_integration/sc_mlp.py @@ -209,12 +209,12 @@ def _sc_linear_dynamic_mp(self, x, weight, bias, operator, chunk_d=0): if self.sc_controller.adaptive_mp_config is not None: assignment = adaptive_classify_rows( row_metric, - self.sc_controller.current_timestep, - self.sc_controller.total_timesteps, self.sc_controller.adaptive_mp_config, operator=operator, block_idx=self.block_idx, total_blocks=self.sc_controller.total_blocks, + timestep=self.sc_controller.current_timestep, + total_timesteps=self.sc_controller.total_timesteps, ) else: mp_config = self.sc_controller.mp_config @@ -295,12 +295,12 @@ def _sc_linear_combined_mp(self, x, weight, bias, operator, dispatch, if self.sc_controller.adaptive_mp_config is not None: assignment = adaptive_classify_rows( row_metric, - self.sc_controller.current_timestep, - self.sc_controller.total_timesteps, self.sc_controller.adaptive_mp_config, operator=operator, block_idx=self.block_idx, total_blocks=self.sc_controller.total_blocks, + timestep=self.sc_controller.current_timestep, + total_timesteps=self.sc_controller.total_timesteps, ) else: mp_config = self.sc_controller.mp_config diff --git a/scripts/calibrate_mp_thresholds.py b/scripts/calibrate_mp_thresholds.py index c301dfe..4b70425 100644 --- a/scripts/calibrate_mp_thresholds.py +++ b/scripts/calibrate_mp_thresholds.py @@ -145,6 +145,7 @@ def _cosine_dist_heads(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor _METRIC_ROWS = _relative_l2_rows _METRIC_HEADS = _relative_l2_heads _USE_FP_TEACHER = False +_SC_HALVE = False # bipolar sign-magnitude halving during calibration probes def _save_fp_weights(model) -> None: @@ -337,6 +338,12 @@ def __init__( # group's pre-subsample row count R_g, used for the R_g-weighted global lambda. self.budget_scope = budget_scope self.true_counts: dict[tuple[str, int, int], int] = defaultdict(int) + # Per-operator MACs-per-row (out_features x in_features for linears, + # N x head_dim for qk/av). The GLOBAL budget prices each row by + # macs_per_row, so the budget is FLOP-weighted rather than row-weighted: + # a target avg stoc_len then means that value in real compute, not per + # row. Set once per operator by add(); constant across its rows. + self.op_macs: dict[str, float] = {} self._rng = np.random.default_rng(rng_seed) # Optional full-resolution log keyed by raw (operator, block_idx, timestep) so the # bucketing can be re-derived offline (any timestep/layer_buckets) without re-probing. @@ -365,10 +372,15 @@ def add( timestep: int, metric_norm: torch.Tensor, errors_by_level: list[torch.Tensor], + macs_per_row: float = 1.0, ): if not self.should_record(operator, block_idx, timestep): return + # Record this operator's per-row MAC cost (constant across its rows) so + # the global budget can weight rows by FLOPs, not just row count. + self.op_macs[operator] = float(macs_per_row) + metrics = metric_norm.detach().float().reshape(-1).cpu().numpy() err = torch.stack([e.detach().float().reshape(-1) for e in errors_by_level], dim=-1) errors = err.cpu().numpy() @@ -475,14 +487,20 @@ def export(self) -> tuple[dict, list[dict]]: # independently pinned to budget_ratio*ref -- original behaviour). global_lam = None if self.budget_scope == "global": - werrs, row_counts = [], [] + werrs, flop_counts = [], [] for key, rec in self.records.items(): e = np.concatenate(rec["errors"], axis=0) werrs.append(e) - row_counts.append(float(self.true_counts.get(key, e.shape[0]))) + R = float(self.true_counts.get(key, e.shape[0])) + # FLOP-weighted budget: price each true row by macs_per_row so + # cheap-but-numerous ops (per-head-per-row av/qk) can't dominate + # the budget over few-but-expensive linear rows. _global_lambda + # treats the passed counts as the budget/cost weights, so passing + # R*macs makes the realized average FLOP-weighted. + flop_counts.append(R * self.op_macs.get(key[0], 1.0)) global_lam = _global_lambda( werrs, self.costs, self.budget_ratio, - float(self.budget_ref_stoc_len), row_counts) + float(self.budget_ref_stoc_len), flop_counts) payload["global_lambda"] = float(global_lam) for operator in sorted(self.operators): @@ -497,7 +515,8 @@ def export(self) -> tuple[dict, list[dict]]: fitted = self._fit_group( np.concatenate(op_metrics, axis=0), np.concatenate(op_errors, axis=0), - lam=global_lam, weight=1.0, rep=self._rep_for_op(operator), + lam=global_lam, weight=1.0, + rep=self._rep_for_op(operator) * self.op_macs.get(operator, 1.0), ) payload["operator_defaults"][operator] = fitted summary_rows.append( @@ -518,8 +537,9 @@ def export(self) -> tuple[dict, list[dict]]: if metrics.size < self.min_bucket_units: continue operator, t_bucket, l_bucket = key - fitted = self._fit_group(metrics, errors, lam=global_lam, - weight=1.0, rep=self._rep_for_key(key)) + fitted = self._fit_group( + metrics, errors, lam=global_lam, weight=1.0, + rep=self._rep_for_key(key) * self.op_macs.get(key[0], 1.0)) bucket_key = f"{operator}:t{t_bucket}:l{l_bucket}" payload["buckets"][bucket_key] = fitted summary_rows.append( @@ -532,16 +552,19 @@ def export(self) -> tuple[dict, list[dict]]: } ) - # Honest iso-budget check: row-weighted (R_g) average stoc_len across all - # buckets. For per_bucket this is ~budget_ratio*ref by construction; for - # global it's the realized average -- should also land near the target. + # Honest iso-budget check: FLOP-weighted (R_g * macs_g) average stoc_len + # across all buckets. For global scope this is the realized compute + # average -- should land near the target (unlike the old row-weighted + # average, which under-counted the FLOP-heavy linears). + payload["op_macs"] = dict(self.op_macs) num = den = 0.0 for bkey, fitted in payload["buckets"].items(): op, tpart, lpart = bkey.split(":") k = (op, int(tpart[1:]), int(lpart[1:])) R = float(self.true_counts.get(k, fitted["num_units"])) - num += R * fitted["avg_stoc_len"] - den += R + w = R * self.op_macs.get(op, 1.0) + num += w * fitted["avg_stoc_len"] + den += w payload["expected_avg_stoc_len"] = (num / den) if den else 0.0 return payload, summary_rows @@ -605,6 +628,7 @@ def _run_attention_linear_level( chunk_d=chunk_d, stoc_len=stoc_len, rng_levels=rng_levels, + halve_bipolar_stoc_len=_SC_HALVE, ) result = chunk_result if result is None else result + chunk_result else: @@ -620,6 +644,7 @@ def _run_attention_linear_level( group_b=1, stoc_len=stoc_len, rng_levels=rng_levels, + halve_bipolar_stoc_len=_SC_HALVE, ) if bias is not None: @@ -674,6 +699,7 @@ def _run_mlp_linear_level( chunk_d=chunk_d, stoc_len=stoc_len, rng_levels=rng_levels, + halve_bipolar_stoc_len=_SC_HALVE, ) result = chunk_result if result is None else result + chunk_result else: @@ -690,6 +716,7 @@ def _run_mlp_linear_level( chunk_d=chunk_d, stoc_len=stoc_len, rng_levels=rng_levels, + halve_bipolar_stoc_len=_SC_HALVE, ) if bias is not None: result = result + bias @@ -758,7 +785,9 @@ def _attn_hook(self, module, inputs, _output): teacher_qkv.float().reshape(-1, teacher_qkv.shape[-1]), ) ) - self.calibrator.add("input_proj", block_idx, timestep, row_metric, level_errors) + self.calibrator.add( + "input_proj", block_idx, timestep, row_metric, level_errors, + macs_per_row=float(module.qkv.weight.shape[0] * module.qkv.weight.shape[1])) else: if _USE_FP_TEACHER: x_for_teacher = (torch.index_select(x, 2, module.reorder_index_qkv) @@ -784,16 +813,32 @@ def _attn_hook(self, module, inputs, _output): teacher_attn = (q_scaled @ k.transpose(-2, -1)) if self.calibrator.should_record("qk", block_idx, timestep): - q_metric = _normalize_metric(q_scaled.float().abs().amax(dim=(0, 2, 3))) - level_errors = [] + # Per-query-row recording — matches the runtime per-row qk MP + # (rows classified by ||Q_row||_inf), mirroring the av path. + q_flat = q_scaled.reshape(-1, q_scaled.shape[-2], q_scaled.shape[-1]) + teacher_flat = teacher_attn.reshape( + -1, teacher_attn.shape[-2], teacher_attn.shape[-1]) + qk_errors_by_level: list[list[torch.Tensor]] = [[] for _ in self.calibrator.levels] for sl in self.calibrator.levels: if sl == 0: sc_attn = torch.zeros_like(teacher_attn, dtype=torch.float32) else: sc_prec = _resolve_level_sc_prec(module, sl) sc_attn = module._sc_qk_uniform(q, k, sc_prec, sl).float() - level_errors.append(_METRIC_HEADS(sc_attn, teacher_attn.float())) - self.calibrator.add("qk", block_idx, timestep, q_metric, level_errors) + sc_flat = sc_attn.reshape(-1, sc_attn.shape[-2], sc_attn.shape[-1]) + for i in range(sc_flat.shape[0]): + qk_errors_by_level[self.calibrator.levels.index(sl)].append( + _METRIC_ROWS(sc_flat[i], teacher_flat[i].float()) + ) + qk_macs = float(n_tokens * module.head_dim) + for i in range(q_flat.shape[0]): + row_metric = _normalize_metric(q_flat[i].abs().amax(dim=-1)) + level_errors = [ + qk_errors_by_level[level_idx][i] + for level_idx in range(len(self.calibrator.levels)) + ] + self.calibrator.add("qk", block_idx, timestep, row_metric, + level_errors, macs_per_row=qk_macs) attn = module.attn_drop(teacher_attn.softmax(dim=-1)) teacher_av = (attn @ v) @@ -820,7 +865,9 @@ def _attn_hook(self, module, inputs, _output): level_errors_by_level[level_idx][i] for level_idx in range(len(self.calibrator.levels)) ] - self.calibrator.add("av", block_idx, timestep, row_metric, level_errors) + self.calibrator.add( + "av", block_idx, timestep, row_metric, level_errors, + macs_per_row=float(module.head_dim * n_tokens)) proj_hidden = module.num_heads * module.head_dim proj_in = teacher_av.transpose(1, 2).reshape(bsz, n_tokens, proj_hidden) @@ -850,7 +897,9 @@ def _attn_hook(self, module, inputs, _output): teacher_proj.float().reshape(-1, teacher_proj.shape[-1]), ) ) - self.calibrator.add("proj", block_idx, timestep, row_metric, level_errors) + self.calibrator.add( + "proj", block_idx, timestep, row_metric, level_errors, + macs_per_row=float(module.proj.weight.shape[0] * module.proj.weight.shape[1])) def _mlp_hook(self, module, inputs, _output): timestep = module.sc_controller.current_timestep @@ -887,7 +936,9 @@ def _mlp_hook(self, module, inputs, _output): teacher_fc1.float().reshape(-1, teacher_fc1.shape[-1]), ) ) - self.calibrator.add("mlp_fc1", block_idx, timestep, row_metric, level_errors) + self.calibrator.add( + "mlp_fc1", block_idx, timestep, row_metric, level_errors, + macs_per_row=float(module.fc1.weight.shape[0] * module.fc1.weight.shape[1])) else: if _USE_FP_TEACHER: x_for_teacher = (torch.index_select(x, 2, module.reorder_index_fc1) @@ -919,7 +970,9 @@ def _mlp_hook(self, module, inputs, _output): teacher_fc2.float().reshape(-1, teacher_fc2.shape[-1]), ) ) - self.calibrator.add("mlp_fc2", block_idx, timestep, row_metric, level_errors) + self.calibrator.add( + "mlp_fc2", block_idx, timestep, row_metric, level_errors, + macs_per_row=float(module.fc2.weight.shape[0] * module.fc2.weight.shape[1])) def register_hooks(self): hooks = [] @@ -1111,11 +1164,13 @@ def main(): model = quantize_sc_model(model, device, args, sc_controller=sc_controller) # Wire metric + teacher dispatch globals. - global _METRIC_ROWS, _METRIC_HEADS, _USE_FP_TEACHER + global _METRIC_ROWS, _METRIC_HEADS, _USE_FP_TEACHER, _SC_HALVE if args.metric == "cosine": _METRIC_ROWS = _cosine_dist_rows _METRIC_HEADS = _cosine_dist_heads _USE_FP_TEACHER = (args.teacher == "fp") + _SC_HALVE = bool(args.sc_halve) + print(f"SC halve during calibration: {_SC_HALVE}") print(f"Calibration metric: {args.metric}, teacher: {args.teacher}") model.to(device) model.eval().half() diff --git a/scripts/quant_sc_main.py b/scripts/quant_sc_main.py index 7e07f3d..9159685 100644 --- a/scripts/quant_sc_main.py +++ b/scripts/quant_sc_main.py @@ -424,20 +424,12 @@ def main(): quant_string_name += "_fixlvlprec" if args.sc_noise_model: quant_string_name += "_noisemodel" - # Append per-operator MP alpha/beta to folder name when adaptive_mp is used + # Tag the folder name with the calibrated MP table when adaptive_mp is used if getattr(args, 'adaptive_mp', False) or getattr(args, 'adaptive_mp_table', None): - _mp_parts = [] - for _op in ["qk", "av", "proj", "input_proj", "mlp_fc1", "mlp_fc2"]: - _a = getattr(args, f"mp_alpha_{_op}", None) - _b = getattr(args, f"mp_beta_{_op}", None) - if _a is not None and _b is not None: - _mp_parts.append(f"{_op}_a{_a}_b{_b}") if getattr(args, 'adaptive_mp_table', None): quant_string_name += f"_mptbl_{Path(args.adaptive_mp_table).stem}" - elif _mp_parts: - quant_string_name += "_mp_" + "_".join(_mp_parts) else: - quant_string_name += f"_mp_a{args.mp_alpha}_b{args.mp_beta}" + quant_string_name += "_adaptivemp" if getattr(args, 'range_mp', False): _rmp_parts = [] for _op in ["qk", "av", "proj", "input_proj", "mlp"]: @@ -545,44 +537,20 @@ def main(): # Initialize mixed precision if requested if args.adaptive_mp or args.adaptive_mp_table: + if not args.adaptive_mp_table: + raise ValueError( + "--adaptive_mp requires --adaptive_mp_table (a calibrated " + "threshold table from calibrate_mp_thresholds.py). The old " + "closed-form alpha/beta fallback has been removed.") levels = [int(x) for x in args.mp_levels.split(',')] - # Build per-operator overrides (None means use global default) - operator_params = {} - # Lookup order: exact op name → group key → global - for op, group_key in [("qk", "qk"), ("av", "av"), - ("mlp_fc1", "mlp_fc1"), ("mlp_fc2", "mlp_fc2"), - ("input_proj", "input_proj"), ("proj", "proj")]: - # Try exact op-level arg first, then group-level fallback - a = getattr(args, f"mp_alpha_{op}", None) - b = getattr(args, f"mp_beta_{op}", None) - if a is None: - # Group fallback: mlp_fc1/mlp_fc2 → mlp - group_fallback = {"mlp_fc1": "mlp", "mlp_fc2": "mlp"}.get(op) - if group_fallback: - a = getattr(args, f"mp_alpha_{group_fallback}", None) - if b is None: - group_fallback = {"mlp_fc1": "mlp", "mlp_fc2": "mlp"}.get(op) - if group_fallback: - b = getattr(args, f"mp_beta_{group_fallback}", None) - if a is not None or b is not None: - operator_params[op] = ( - a if a is not None else args.mp_alpha, - b if b is not None else args.mp_beta, - ) - adaptive_config = AdaptiveMPConfig( stoc_len_levels=levels, - alpha=args.mp_alpha, - beta=args.mp_beta, enable_pruning=args.mp_enable_pruning, - operator_params=operator_params, threshold_table_path=args.adaptive_mp_table, ) sc_controller.init_adaptive_mp(adaptive_config) - logging.info(f"Adaptive mixed precision V2 enabled: levels={levels}, " - f"alpha={args.mp_alpha}, beta={args.mp_beta}, " + logging.info(f"Adaptive mixed precision enabled: levels={levels}, " f"pruning={args.mp_enable_pruning}, " - f"operator_params={operator_params}, " f"threshold_table={args.adaptive_mp_table}") elif args.mp: levels = [int(x) for x in args.mp_levels.split(',')] @@ -744,25 +712,17 @@ def create_argparser(): help='Comma-separated fractions per level (default: equal). E.g. "0.1,0.2,0.3,0.4".' ) - # Adaptive mixed precision (timestep-aware, inspired by APT) + # Adaptive mixed precision (calibrated per-row thresholds) parser.add_argument( '--adaptive_mp', action='store_true', - help='Enable adaptive mixed precision with timestep-aware thresholds.' + help='Enable adaptive mixed precision. Requires --adaptive_mp_table ' + '(a calibrated threshold table); there is no closed-form fallback.' ) parser.add_argument( '--adaptive_mp_table', type=str, default=None, - help='Path to a calibrated adaptive-MP threshold JSON table. ' - 'When provided, runtime classification uses operator/timestep/layer ' - 'bucket thresholds from the table and falls back to alpha/beta only ' - 'for operators missing from the table.' - ) - parser.add_argument( - '--mp_alpha', type=float, default=0.3, - help='Adaptive MP: threshold sensitivity to timestep progress.' - ) - parser.add_argument( - '--mp_beta', type=float, default=0.05, - help='Adaptive MP: base threshold offset.' + help='Path to a calibrated adaptive-MP threshold JSON table from ' + 'calibrate_mp_thresholds.py. Runtime classification uses the ' + 'per-(operator, timestep, layer) bucket thresholds from the table.' ) parser.add_argument( '--mp_enable_pruning', action='store_true', default=True, @@ -772,64 +732,6 @@ def create_argparser(): '--no_mp_pruning', dest='mp_enable_pruning', action='store_false', help='Adaptive MP: disable row pruning.' ) - # Per-operator alpha/beta overrides - parser.add_argument( - '--mp_alpha_qk', type=float, default=None, - help='Per-operator alpha for QK (default: use --mp_alpha).' - ) - parser.add_argument( - '--mp_alpha_av', type=float, default=None, - help='Per-operator alpha for AV (default: use --mp_alpha).' - ) - parser.add_argument( - '--mp_alpha_mlp', type=float, default=None, - help='Per-operator alpha for MLP fc1/fc2 (default: use --mp_alpha).' - ) - parser.add_argument( - '--mp_alpha_proj', type=float, default=None, - help='Per-operator alpha for proj (default: use --mp_alpha).' - ) - parser.add_argument( - '--mp_alpha_input_proj', type=float, default=None, - help='Per-operator alpha for input_proj (default: fallback to --mp_alpha_proj, then --mp_alpha).' - ) - parser.add_argument( - '--mp_alpha_mlp_fc1', type=float, default=None, - help='Per-operator alpha for mlp_fc1 (default: fallback to --mp_alpha_mlp, then --mp_alpha).' - ) - parser.add_argument( - '--mp_alpha_mlp_fc2', type=float, default=None, - help='Per-operator alpha for mlp_fc2 (default: fallback to --mp_alpha_mlp, then --mp_alpha).' - ) - parser.add_argument( - '--mp_beta_qk', type=float, default=None, - help='Per-operator beta for QK (default: use --mp_beta).' - ) - parser.add_argument( - '--mp_beta_av', type=float, default=None, - help='Per-operator beta for AV (default: use --mp_beta).' - ) - parser.add_argument( - '--mp_beta_mlp', type=float, default=None, - help='Per-operator beta for MLP fc1/fc2 (default: use --mp_beta).' - ) - parser.add_argument( - '--mp_beta_proj', type=float, default=None, - help='Per-operator beta for proj (default: use --mp_beta).' - ) - parser.add_argument( - '--mp_beta_input_proj', type=float, default=None, - help='Per-operator beta for input_proj (default: fallback to --mp_beta_proj, then --mp_beta).' - ) - parser.add_argument( - '--mp_beta_mlp_fc1', type=float, default=None, - help='Per-operator beta for mlp_fc1 (default: fallback to --mp_beta_mlp, then --mp_beta).' - ) - parser.add_argument( - '--mp_beta_mlp_fc2', type=float, default=None, - help='Per-operator beta for mlp_fc2 (default: fallback to --mp_beta_mlp, then --mp_beta).' - ) - # Range-based mixed precision (weight min/max range) parser.add_argument( '--range_mp', action='store_true', From 7a326953ba9f3155c304beeefaea4143c3fb6713 Mon Sep 17 00:00:00 2001 From: heroarmor <162866837+heroarmor@users.noreply.github.com> Date: Sat, 11 Jul 2026 13:55:25 -0400 Subject: [PATCH 2/7] chore: gitignore SLURM job logs and generated figures Run artifacts (slurm-*.out, figures/) are regenerated, not source. Co-Authored-By: Claude Opus 4.8 (1M context) --- .gitignore | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/.gitignore b/.gitignore index c6277fb..8120ce8 100644 --- a/.gitignore +++ b/.gitignore @@ -16,3 +16,7 @@ logs_mp_sweep/ evaluation/sample_grids/ evaluation/imagenet_ref/images/ evaluation/imagenet_ref/*.npz + +# SLURM job logs and generated figures (run artifacts — regenerated, not source) +slurm-*.out +figures/ From 0b66c15d0897b4da71b82a8319619105d7ef64fc Mon Sep 17 00:00:00 2001 From: heroarmor <162866837+heroarmor@users.noreply.github.com> Date: Sat, 11 Jul 2026 13:55:47 -0400 Subject: [PATCH 3/7] tools+scripts: add SC comparison tools and MP/calib sbatch scripts tools/: compare_old_vs_new{,_isolated,_thorough}.py, compare_vs_cpu_reference.py, sc_call_counter{,_full}.py, smoke_test_e2e.py scripts/: global-PR calib + MP sbatch (calib_globalpr{,_flop}{,_halve}, mp_globalpr_{auto,fill,flop_halve_auto}, uniform{,_halve}_auto, opfreeze_probe) + eval/ CPU eval Co-Authored-By: Claude Opus 4.8 (1M context) --- scripts/eval/eval_mp_flop_halve_cpu.sh | 37 +++ scripts/eval/sbatch_eval_mp_flop_halve.sb | 42 ++++ scripts/sbatch_calib_globalpr.sb | 31 +++ scripts/sbatch_calib_globalpr_flop.sb | 32 +++ scripts/sbatch_calib_globalpr_flop_halve.sb | 33 +++ scripts/sbatch_calib_globalpr_halve.sb | 31 +++ scripts/sbatch_mp_globalpr_auto.sb | 56 +++++ scripts/sbatch_mp_globalpr_fill.sb | 56 +++++ scripts/sbatch_mp_globalpr_flop_halve_auto.sb | 56 +++++ scripts/sbatch_opfreeze_probe.sb | 45 ++++ scripts/sbatch_uniform_auto.sb | 49 ++++ scripts/sbatch_uniform_halve_auto.sb | 49 ++++ tools/compare_old_vs_new.py | 207 ++++++++++++++++ tools/compare_old_vs_new_isolated.py | 119 ++++++++++ tools/compare_old_vs_new_thorough.py | 167 +++++++++++++ tools/compare_vs_cpu_reference.py | 156 ++++++++++++ tools/sc_call_counter.py | 61 +++++ tools/sc_call_counter_full.py | 59 +++++ tools/smoke_test_e2e.py | 223 ++++++++++++++++++ 19 files changed, 1509 insertions(+) create mode 100644 scripts/eval/eval_mp_flop_halve_cpu.sh create mode 100644 scripts/eval/sbatch_eval_mp_flop_halve.sb create mode 100644 scripts/sbatch_calib_globalpr.sb create mode 100644 scripts/sbatch_calib_globalpr_flop.sb create mode 100644 scripts/sbatch_calib_globalpr_flop_halve.sb create mode 100755 scripts/sbatch_calib_globalpr_halve.sb create mode 100644 scripts/sbatch_mp_globalpr_auto.sb create mode 100755 scripts/sbatch_mp_globalpr_fill.sb create mode 100644 scripts/sbatch_mp_globalpr_flop_halve_auto.sb create mode 100644 scripts/sbatch_opfreeze_probe.sb create mode 100644 scripts/sbatch_uniform_auto.sb create mode 100755 scripts/sbatch_uniform_halve_auto.sb create mode 100644 tools/compare_old_vs_new.py create mode 100644 tools/compare_old_vs_new_isolated.py create mode 100644 tools/compare_old_vs_new_thorough.py create mode 100644 tools/compare_vs_cpu_reference.py create mode 100644 tools/sc_call_counter.py create mode 100644 tools/sc_call_counter_full.py create mode 100644 tools/smoke_test_e2e.py diff --git a/scripts/eval/eval_mp_flop_halve_cpu.sh b/scripts/eval/eval_mp_flop_halve_cpu.sh new file mode 100644 index 0000000..0135f64 --- /dev/null +++ b/scripts/eval/eval_mp_flop_halve_cpu.sh @@ -0,0 +1,37 @@ +#!/bin/bash +# CPU-only eval for one FLOP+halve adaptive-MP config (no GPU -> respects the +# 6-GPU cap). KID + FID/IS/Precision/Recall vs the ImageNet-256 reference. +# Usage: bash eval_mp_flop_halve_cpu.sh +set -uo pipefail +AVG="${1:?usage: eval_mp_flop_halve_cpu.sh }" +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +PREV=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion_prev +PACKER=$PREV/imagenet256_ref/parallel_npz.py +REF=$PREV/imagenet256_ref/VIRTUAL_imagenet256_labeled.npz +BASE=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion_fid_mp_globalpr_flop_halve_cfg15/adaptive_avg${AVG} +AD=$BASE/samples +WORK=$BASE/_eval; mkdir -p "$WORK" +OUT=$WORK/eval_avg${AVG}.txt +export CUDA_VISIBLE_DEVICES="" # force CPU +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh +exec > >(tee "$OUT") 2>&1 +echo "############ FLOP+HALVE MP EVAL avg${AVG} (CPU) $(date) ############" + +IDX=$WORK/idx.txt; NPZ=$WORK/avg${AVG}.npz +(cd "$AD" && ls [0-9][0-9][0-9][0-9][0-9][0-9].png | awk '($1+0)<2000' | sort) > "$IDX" +SEL=$WORK/sel; rm -rf "$SEL"; mkdir -p "$SEL" +while read -r f; do [[ -n "$f" ]] && ln -sf "$AD/$f" "$SEL/$f"; done < "$IDX" +echo "=== matched idx 0-1999: $(wc -l < "$IDX") ===" + +conda activate qdit +python -u "$PACKER" "$SEL" "$NPZ" +conda deactivate + +conda activate tfeval +export TF_CPP_MIN_LOG_LEVEL=2 EVALUATOR=$PREV/Q-DiT/models/evaluations/evaluator.py +echo ""; echo "===== [1] KID vs ref (idx 0-1999) =====" +bash "$REPO/scripts/eval/kid_openai.sh" "$WORK/kid_avg${AVG}.txt" "$NPZ" +echo ""; echo "===== [2] FID / IS / sFID / Precision / Recall vs ref =====" +python -u "$EVALUATOR" "$REF" "$NPZ" 2>&1 | grep -E "^(Inception Score|FID|sFID|Precision|Recall):" | sed "s/^/[avg${AVG}] /" +conda deactivate +echo ""; echo "############ DONE avg${AVG} $(date) ############" diff --git a/scripts/eval/sbatch_eval_mp_flop_halve.sb b/scripts/eval/sbatch_eval_mp_flop_halve.sb new file mode 100644 index 0000000..0a4cfa0 --- /dev/null +++ b/scripts/eval/sbatch_eval_mp_flop_halve.sb @@ -0,0 +1,42 @@ +#!/bin/bash +#SBATCH --job-name=eval_mp +#SBATCH --partition=gpu-rtx6000 +#SBATCH --account=nbleier_owned1 +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-gpu=4 +#SBATCH --mem=48G +#SBATCH --time=00:45:00 +# GPU eval for one FLOP+halve adaptive-MP config: KID + FID/IS/sFID/Prec/Recall. +# Usage: sbatch sbatch_eval_mp_flop_halve.sb +set -uo pipefail +AVG="${1:?usage: sbatch sbatch_eval_mp_flop_halve.sb }" +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +PREV=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion_prev +PACKER=$PREV/imagenet256_ref/parallel_npz.py +REF=$PREV/imagenet256_ref/VIRTUAL_imagenet256_labeled.npz +BASE=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion_fid_mp_globalpr_flop_halve_cfg15/adaptive_avg${AVG} +AD=$BASE/samples +WORK=$BASE/_eval; mkdir -p "$WORK" +OUT=$WORK/eval_avg${AVG}.txt +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh +exec > >(tee "$OUT") 2>&1 +echo "############ FLOP+HALVE MP EVAL avg${AVG} (GPU) $(date) ############" + +IDX=$WORK/idx.txt; NPZ=$WORK/avg${AVG}.npz +(cd "$AD" && ls [0-9][0-9][0-9][0-9][0-9][0-9].png | awk '($1+0)<2000' | sort) > "$IDX" +SEL=$WORK/sel; rm -rf "$SEL"; mkdir -p "$SEL" +while read -r f; do [[ -n "$f" ]] && ln -sf "$AD/$f" "$SEL/$f"; done < "$IDX" +echo "=== matched idx 0-1999: $(wc -l < "$IDX") ===" + +conda activate qdit +python -u "$PACKER" "$SEL" "$NPZ" +conda deactivate + +conda activate tfeval +export TF_CPP_MIN_LOG_LEVEL=2 EVALUATOR=$PREV/Q-DiT/models/evaluations/evaluator.py +echo ""; echo "===== [1] KID vs ref (idx 0-1999) =====" +bash "$REPO/scripts/eval/kid_openai.sh" "$WORK/kid_avg${AVG}.txt" "$NPZ" +echo ""; echo "===== [2] FID / IS / sFID / Precision / Recall vs ref =====" +python -u "$EVALUATOR" "$REF" "$NPZ" 2>&1 | grep -E "^(Inception Score|FID|sFID|Precision|Recall):" | sed "s/^/[avg${AVG}] /" +conda deactivate +echo ""; echo "############ DONE avg${AVG} $(date) ############" diff --git a/scripts/sbatch_calib_globalpr.sb b/scripts/sbatch_calib_globalpr.sb new file mode 100644 index 0000000..f32bc9c --- /dev/null +++ b/scripts/sbatch_calib_globalpr.sb @@ -0,0 +1,31 @@ +#!/bin/bash +#SBATCH --job-name=calib_gpr +#SBATCH --partition=gpu-rtx6000 +#SBATCH --account=nbleier_owned1 +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-gpu=8 +#SBATCH --mem=110G +#SBATCH --time=03:00:00 +# per-row GLOBAL calibration for one budget. Usage: sbatch sbatch_calib_globalpr.sb +set -uo pipefail +AVG="${1:?AVG}"; BR="${2:?BUDGET_RATIO}" +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +SCRATCH=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi +CKPT=$SCRATCH/pretrained_models/DiT-XL-2-256x256.pt +OUTDIR=$SCRATCH/scmp_diffusion_fid_mp_cfg15/calib_global_perrow +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh; conda activate qdit +export PYTHONUNBUFFERED=1 SC_OWEN_MODE=bitrev; cd "$REPO"; mkdir -p "$OUTDIR" +echo "=== calib GLOBAL per-row avg${AVG} br=${BR} $(date) ===" +python -u scripts/calibrate_mp_thresholds.py \ + --mp_levels 256,192,128,96,64,48,32,16 \ + --budget_ratio "$BR" --budget_ref_stoc_len 256 --budget_scope global \ + --metric cosine --teacher fp \ + --sc_prec 8 --sc_fixed_level_prec --sc_qk_granularity per_row \ + --wbits 8 --abits 8 --w_sym --a_sym \ + --image-size 256 --num-sampling-steps 50 \ + --num_calib_batches 1 --num_calib_timesteps 50 \ + --timestep_buckets 10 --layer_buckets 4 \ + --teacher_cfg_scale 0.0 --ckpt "$CKPT" \ + --calib_output_json "$OUTDIR/calib_fix_avg${AVG}_l256_ref192.json" \ + --calib_summary_csv "$OUTDIR/calib_fix_avg${AVG}_summary.csv" 2>&1 | tee "$OUTDIR/calib_avg${AVG}.log" +echo "DONE avg${AVG} $(date)" diff --git a/scripts/sbatch_calib_globalpr_flop.sb b/scripts/sbatch_calib_globalpr_flop.sb new file mode 100644 index 0000000..6e5113a --- /dev/null +++ b/scripts/sbatch_calib_globalpr_flop.sb @@ -0,0 +1,32 @@ +#!/bin/bash +#SBATCH --job-name=calib_flop +#SBATCH --partition=gpu-rtx6000 +#SBATCH --account=nbleier_owned1 +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-gpu=8 +#SBATCH --mem=110G +#SBATCH --time=03:00:00 +# FLOP-weighted per-row GLOBAL calibration for one budget (fix branch: +# fix/flop-cost-perrow-qk-drop-alphabeta). Usage: sbatch sbatch_calib_globalpr_flop.sb +set -uo pipefail +AVG="${1:?AVG}"; BR="${2:?BUDGET_RATIO}" +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +SCRATCH=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi +CKPT=$SCRATCH/pretrained_models/DiT-XL-2-256x256.pt +OUTDIR=$SCRATCH/scmp_diffusion_fid_mp_cfg15/calib_global_perrow_flop +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh; conda activate qdit +export PYTHONUNBUFFERED=1 SC_OWEN_MODE=bitrev; cd "$REPO"; mkdir -p "$OUTDIR" +echo "=== FLOP-weighted calib GLOBAL per-row avg${AVG} br=${BR} branch=$(git branch --show-current) $(date) ===" +python -u scripts/calibrate_mp_thresholds.py \ + --mp_levels 256,192,128,96,64,48,32,16 \ + --budget_ratio "$BR" --budget_ref_stoc_len 256 --budget_scope global \ + --metric cosine --teacher fp \ + --sc_prec 8 --sc_fixed_level_prec --sc_qk_granularity per_row \ + --wbits 8 --abits 8 --w_sym --a_sym \ + --image-size 256 --num-sampling-steps 50 \ + --num_calib_batches 1 --num_calib_timesteps 50 \ + --timestep_buckets 10 --layer_buckets 4 \ + --teacher_cfg_scale 0.0 --ckpt "$CKPT" \ + --calib_output_json "$OUTDIR/calib_fix_avg${AVG}_l256_ref256.json" \ + --calib_summary_csv "$OUTDIR/calib_fix_avg${AVG}_summary.csv" 2>&1 | tee "$OUTDIR/calib_avg${AVG}.log" +echo "DONE avg${AVG} $(date)" diff --git a/scripts/sbatch_calib_globalpr_flop_halve.sb b/scripts/sbatch_calib_globalpr_flop_halve.sb new file mode 100644 index 0000000..06d1984 --- /dev/null +++ b/scripts/sbatch_calib_globalpr_flop_halve.sb @@ -0,0 +1,33 @@ +#!/bin/bash +#SBATCH --job-name=calib_flopH +#SBATCH --partition=gpu-rtx6000 +#SBATCH --account=nbleier_owned1 +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-gpu=8 +#SBATCH --mem=110G +#SBATCH --time=03:00:00 +# FLOP-weighted + HALVE per-row GLOBAL calibration for one budget (fix branch: +# fix/flop-cost-perrow-qk-drop-alphabeta). Matches the halve MP deployment. +# Usage: sbatch sbatch_calib_globalpr_flop_halve.sb +set -uo pipefail +AVG="${1:?AVG}"; BR="${2:?BUDGET_RATIO}" +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +SCRATCH=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi +CKPT=$SCRATCH/pretrained_models/DiT-XL-2-256x256.pt +OUTDIR=$SCRATCH/scmp_diffusion_fid_mp_cfg15/calib_global_perrow_flop_halve +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh; conda activate qdit +export PYTHONUNBUFFERED=1 SC_OWEN_MODE=bitrev; cd "$REPO"; mkdir -p "$OUTDIR" +echo "=== FLOP+HALVE calib GLOBAL per-row avg${AVG} br=${BR} branch=$(git branch --show-current) $(date) ===" +python -u scripts/calibrate_mp_thresholds.py \ + --mp_levels 256,192,128,96,64,48,32,16 \ + --budget_ratio "$BR" --budget_ref_stoc_len 256 --budget_scope global \ + --metric cosine --teacher fp \ + --sc_prec 8 --sc_fixed_level_prec --sc_qk_granularity per_row \ + --wbits 8 --abits 8 --w_sym --a_sym \ + --image-size 256 --num-sampling-steps 50 \ + --num_calib_batches 1 --num_calib_timesteps 50 \ + --timestep_buckets 10 --layer_buckets 4 \ + --sc_halve --teacher_cfg_scale 0.0 --ckpt "$CKPT" \ + --calib_output_json "$OUTDIR/calib_fix_avg${AVG}.json" \ + --calib_summary_csv "$OUTDIR/calib_fix_avg${AVG}_summary.csv" 2>&1 | tee "$OUTDIR/calib_avg${AVG}.log" +echo "DONE avg${AVG} $(date)" diff --git a/scripts/sbatch_calib_globalpr_halve.sb b/scripts/sbatch_calib_globalpr_halve.sb new file mode 100755 index 0000000..d2f0403 --- /dev/null +++ b/scripts/sbatch_calib_globalpr_halve.sb @@ -0,0 +1,31 @@ +#!/bin/bash +#SBATCH --job-name=calib_gpr +#SBATCH --partition=gpu-rtx6000 +#SBATCH --account=nbleier_owned1 +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-gpu=8 +#SBATCH --mem=110G +#SBATCH --time=03:00:00 +# per-row GLOBAL calibration for one budget. Usage: sbatch sbatch_calib_globalpr.sb +set -uo pipefail +AVG="${1:?AVG}"; BR="${2:?BUDGET_RATIO}" +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +SCRATCH=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi +CKPT=$SCRATCH/pretrained_models/DiT-XL-2-256x256.pt +OUTDIR=$SCRATCH/scmp_diffusion_fid_mp_cfg15/calib_global_perrow_halve +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh; conda activate qdit +export PYTHONUNBUFFERED=1 SC_OWEN_MODE=bitrev; cd "$REPO"; mkdir -p "$OUTDIR" +echo "=== calib GLOBAL per-row HALVE avg${AVG} br=${BR} $(date) ===" +python -u scripts/calibrate_mp_thresholds.py \ + --mp_levels 256,192,128,96,64,48,32,16 \ + --budget_ratio "$BR" --budget_ref_stoc_len 256 --budget_scope global \ + --metric cosine --teacher fp \ + --sc_prec 8 --sc_fixed_level_prec --sc_qk_granularity per_row \ + --wbits 8 --abits 8 --w_sym --a_sym \ + --image-size 256 --num-sampling-steps 50 \ + --num_calib_batches 1 --num_calib_timesteps 50 \ + --timestep_buckets 10 --layer_buckets 4 \ + --sc_halve --teacher_cfg_scale 0.0 --ckpt "$CKPT" \ + --calib_output_json "$OUTDIR/calib_fix_avg${AVG}_l256_ref192.json" \ + --calib_summary_csv "$OUTDIR/calib_fix_avg${AVG}_summary.csv" 2>&1 | tee "$OUTDIR/calib_avg${AVG}.log" +echo "DONE avg${AVG} $(date)" diff --git a/scripts/sbatch_mp_globalpr_auto.sb b/scripts/sbatch_mp_globalpr_auto.sb new file mode 100644 index 0000000..e358c2a --- /dev/null +++ b/scripts/sbatch_mp_globalpr_auto.sb @@ -0,0 +1,56 @@ +#!/bin/bash +#SBATCH --job-name=mp_gpr_auto +#SBATCH --partition=gpu-rtx6000 +#SBATCH --account=nbleier_owned1 +#SBATCH --gres=gpu:3 +#SBATCH --cpus-per-gpu=2 +#SBATCH --mem-per-gpu=12G +#SBATCH --time=1-12:00:00 +# +# global+per-row generation that AUTO-RESUMES to NUM_FID=10000, TIMEOUT-safe: +# it queues its successor (afterany, chain-capped) at the START, before it can +# time out. Usage: sbatch sbatch_mp_globalpr_auto.sb [CHAIN] +set -uo pipefail +AVG="${1:?usage: sbatch sbatch_mp_globalpr_auto.sb [CHAIN]}"; CHAIN="${2:-0}"; NUM_GPUS=3; MAXCHAIN=15 +SELF=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion/scripts/sbatch_mp_globalpr_auto.sb +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +SCRATCH=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi +CKPT=$SCRATCH/pretrained_models/DiT-XL-2-256x256.pt +CALIB=$SCRATCH/scmp_diffusion_fid_mp_cfg15/calib_global_perrow/calib_fix_avg${AVG}_l256_ref192.json +OUT=$SCRATCH/scmp_diffusion_fid_mp_globalpr_cfg15/adaptive_avg${AVG} +SAMPLES=$OUT/samples; IDX=$OUT/_indices; LOG=$OUT/_logs +NUM_FID=2000; BALANCED=10000; BATCH=64; STEPS=50; CFG=1.5; SEED=0; NUM_CLASSES=1000 +MP_LEVELS=256,192,128,96,64,48,32,16 +mkdir -p "$SAMPLES" "$IDX" "$LOG" +[[ -f "$CALIB" ]] || { echo "ERROR: missing calib $CALIB" >&2; exit 1; } +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh; conda activate qdit +export SC_OWEN_MODE=bitrev PYTHONUNBUFFERED=1 OMP_NUM_THREADS=2; cd "$REPO" +cnt(){ find "$SAMPLES" -maxdepth 1 -name '[0-9][0-9][0-9][0-9][0-9][0-9].png' | wc -l; } +START=$(cnt) +echo "=== avg${AVG} auto: start=$START/$NUM_FID chain=$CHAIN job=${SLURM_JOB_ID:-?} $(date) ===" +if [[ "$START" -ge "$NUM_FID" ]]; then echo "[complete] avg${AVG} already at $NUM_FID"; exit 0; fi +# Queue successor NOW (afterany -> runs after me regardless of timeout/complete); chain-capped. +if [[ "$CHAIN" -lt "$MAXCHAIN" ]]; then + sbatch -o "$OUT/slurm-%j.out" -e "$OUT/slurm-%j.err" \ + --dependency=afterany:${SLURM_JOB_ID} "$SELF" "$AVG" $((CHAIN+1)) \ + && echo "[chain] queued successor (chain=$((CHAIN+1)))" +fi +python -u scripts/_plan_missing_indices.py "$SAMPLES" "$NUM_FID" "$NUM_GPUS" "$IDX" $((BALANCED/NUM_CLASSES)) +pids=() +for ((g=0; g "$LOG/gpu_${g}.log" 2>&1 & + pids+=($!) +done +for p in "${pids[@]}"; do wait "$p" || true; done +echo "=== avg${AVG}: $START -> $(cnt) / $NUM_FID at $(date) ===" diff --git a/scripts/sbatch_mp_globalpr_fill.sb b/scripts/sbatch_mp_globalpr_fill.sb new file mode 100755 index 0000000..b88939a --- /dev/null +++ b/scripts/sbatch_mp_globalpr_fill.sb @@ -0,0 +1,56 @@ +#!/bin/bash +#SBATCH --job-name=mp_gpr_fill +#SBATCH --partition=gpu-rtx6000 +#SBATCH --account=nbleier_owned1 +#SBATCH --gres=gpu:3 +#SBATCH --cpus-per-gpu=2 +#SBATCH --mem-per-gpu=12G +#SBATCH --time=1-12:00:00 +# +# global+per-row generation that AUTO-RESUMES to NUM_FID=10000, TIMEOUT-safe: +# it queues its successor (afterany, chain-capped) at the START, before it can +# time out. Usage: sbatch sbatch_mp_globalpr_auto.sb [CHAIN] +set -uo pipefail +AVG="${1:?usage: sbatch sbatch_mp_globalpr_auto.sb [CHAIN]}"; CHAIN="${2:-0}"; NUM_GPUS=3; MAXCHAIN=15 +SELF=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion/scripts/sbatch_mp_globalpr_fill.sb +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +SCRATCH=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi +CKPT=$SCRATCH/pretrained_models/DiT-XL-2-256x256.pt +CALIB=$SCRATCH/scmp_diffusion_fid_mp_cfg15/calib_global_perrow/calib_fix_avg${AVG}_l256_ref192.json +OUT=$SCRATCH/scmp_diffusion_fid_mp_globalpr_cfg15/adaptive_avg${AVG} +SAMPLES=$OUT/samples; IDX=$OUT/_indices; LOG=$OUT/_logs +NUM_FID=2000; BALANCED=10000; BATCH=64; STEPS=50; CFG=1.5; SEED=0; NUM_CLASSES=1000 +MP_LEVELS=256,192,128,96,64,48,32,16 +mkdir -p "$SAMPLES" "$IDX" "$LOG" +[[ -f "$CALIB" ]] || { echo "ERROR: missing calib $CALIB" >&2; exit 1; } +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh; conda activate qdit +export SC_OWEN_MODE=bitrev PYTHONUNBUFFERED=1 OMP_NUM_THREADS=2; cd "$REPO" +cnt(){ find "$SAMPLES" -maxdepth 1 -name '[0-9][0-9][0-9][0-9][0-9][0-9].png' -printf '%f\n' | awk -v N=$NUM_FID '{n=$1+0; if(n runs after me regardless of timeout/complete); chain-capped. +if [[ "$CHAIN" -lt "$MAXCHAIN" ]]; then + sbatch -o "$OUT/slurm-%j.out" -e "$OUT/slurm-%j.err" \ + --dependency=afterany:${SLURM_JOB_ID} "$SELF" "$AVG" $((CHAIN+1)) \ + && echo "[chain] queued successor (chain=$((CHAIN+1)))" +fi +python -u scripts/_plan_missing_indices.py "$SAMPLES" "$NUM_FID" "$NUM_GPUS" "$IDX" $((BALANCED/NUM_CLASSES)) +pids=() +for ((g=0; g "$LOG/gpu_${g}.log" 2>&1 & + pids+=($!) +done +for p in "${pids[@]}"; do wait "$p" || true; done +echo "=== avg${AVG}: $START -> $(cnt) / $NUM_FID at $(date) ===" diff --git a/scripts/sbatch_mp_globalpr_flop_halve_auto.sb b/scripts/sbatch_mp_globalpr_flop_halve_auto.sb new file mode 100644 index 0000000..7debbc4 --- /dev/null +++ b/scripts/sbatch_mp_globalpr_flop_halve_auto.sb @@ -0,0 +1,56 @@ +#!/bin/bash +#SBATCH --job-name=mp_flopH +#SBATCH --partition=gpu-rtx6000 +#SBATCH --account=nbleier_owned1 +#SBATCH --gres=gpu:3 +#SBATCH --cpus-per-gpu=2 +#SBATCH --mem-per-gpu=12G +#SBATCH --time=1-12:00:00 +# +# FLOP-weighted + HALVE adaptive-MP generation, auto-resumes to NUM_FID, TIMEOUT-safe: +# queues its successor (afterany, chain-capped) at the START. Uses the FLOP+halve +# calibrated tables. Usage: sbatch sbatch_mp_globalpr_flop_halve_auto.sb [CHAIN] +set -uo pipefail +AVG="${1:?usage: sbatch sbatch_mp_globalpr_flop_halve_auto.sb [CHAIN]}"; CHAIN="${2:-0}"; NUM_GPUS=3; MAXCHAIN=15 +SELF=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion/scripts/sbatch_mp_globalpr_flop_halve_auto.sb +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +SCRATCH=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi +CKPT=$SCRATCH/pretrained_models/DiT-XL-2-256x256.pt +CALIB=$SCRATCH/scmp_diffusion_fid_mp_cfg15/calib_global_perrow_flop_halve/calib_fix_avg${AVG}.json +OUT=$SCRATCH/scmp_diffusion_fid_mp_globalpr_flop_halve_cfg15/adaptive_avg${AVG} +SAMPLES=$OUT/samples; IDX=$OUT/_indices; LOG=$OUT/_logs +NUM_FID=2000; BALANCED=10000; BATCH=64; STEPS=50; CFG=1.5; SEED=0; NUM_CLASSES=1000 +MP_LEVELS=256,192,128,96,64,48,32,16 +mkdir -p "$SAMPLES" "$IDX" "$LOG" +[[ -f "$CALIB" ]] || { echo "ERROR: missing calib $CALIB" >&2; exit 1; } +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh; conda activate qdit +export SC_OWEN_MODE=bitrev PYTHONUNBUFFERED=1 OMP_NUM_THREADS=2; cd "$REPO" +cnt(){ find "$SAMPLES" -maxdepth 1 -name '[0-9][0-9][0-9][0-9][0-9][0-9].png' | wc -l; } +START=$(cnt) +echo "=== FLOP+HALVE avg${AVG} auto: start=$START/$NUM_FID chain=$CHAIN job=${SLURM_JOB_ID:-?} $(date) ===" +if [[ "$START" -ge "$NUM_FID" ]]; then echo "[complete] avg${AVG} already at $NUM_FID"; exit 0; fi +# Queue successor NOW (afterany -> runs after me regardless of timeout/complete); chain-capped. +if [[ "$CHAIN" -lt "$MAXCHAIN" ]]; then + sbatch -o "$OUT/slurm-%j.out" -e "$OUT/slurm-%j.err" \ + --dependency=afterany:${SLURM_JOB_ID} "$SELF" "$AVG" $((CHAIN+1)) \ + && echo "[chain] queued successor (chain=$((CHAIN+1)))" +fi +python -u scripts/_plan_missing_indices.py "$SAMPLES" "$NUM_FID" "$NUM_GPUS" "$IDX" $((BALANCED/NUM_CLASSES)) +pids=() +for ((g=0; g "$LOG/gpu_${g}.log" 2>&1 & + pids+=($!) +done +for p in "${pids[@]}"; do wait "$p" || true; done +echo "=== FLOP+HALVE avg${AVG}: $START -> $(cnt) / $NUM_FID at $(date) ===" diff --git a/scripts/sbatch_opfreeze_probe.sb b/scripts/sbatch_opfreeze_probe.sb new file mode 100644 index 0000000..3342bfb --- /dev/null +++ b/scripts/sbatch_opfreeze_probe.sb @@ -0,0 +1,45 @@ +#!/bin/bash +#SBATCH --job-name=opfreeze +#SBATCH --partition=gpu-rtx6000 +#SBATCH --account=nbleier_owned1 +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-gpu=4 +#SBATCH --mem-per-gpu=16G +#SBATCH --time=04:00:00 +# +# Op-freeze sensitivity probe: all ops at the avg48 FLOP-halve budget EXCEPT one +# op forced to stoc_len 256 (frozen-to-FP-precision table). Generates the same +# indices 0-63 (seed-matched to FP & the avg48 baseline) so PSNR-vs-FP is paired. +# Usage: sbatch sbatch_opfreeze_probe.sb (OP in qk av input_proj proj mlp_fc1 mlp_fc2) +set -uo pipefail +OP="${1:?usage: sbatch sbatch_opfreeze_probe.sb }" +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +SCRATCH=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi +CKPT=$SCRATCH/pretrained_models/DiT-XL-2-256x256.pt +CALIB=$SCRATCH/scmp_diffusion_fid_mp_cfg15/calib_opfreeze_avg48/freeze_${OP}.json +PROBE=$SCRATCH/scmp_diffusion_opfreeze_avg48 +OUT=$PROBE/freeze_${OP}; SAMPLES=$OUT/samples; LOG=$OUT/_logs +IDXF="${2:-$PROBE/idx_0_23.txt}" +N=$(grep -c . "$IDXF") +BALANCED=10000; BATCH=$N; STEPS=50; CFG=1.5; SEED=0; NUM_CLASSES=1000 +MP_LEVELS=256,192,128,96,64,48,32,16 +mkdir -p "$SAMPLES" "$LOG" +[[ -f "$CALIB" ]] || { echo "ERROR: missing calib $CALIB" >&2; exit 1; } +[[ -s "$IDXF" ]] || { echo "ERROR: missing index file $IDXF" >&2; exit 1; } +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh; conda activate qdit +export SC_OWEN_MODE=bitrev PYTHONUNBUFFERED=1 OMP_NUM_THREADS=4; cd "$REPO" +have(){ find "$SAMPLES" -maxdepth 1 -name '[0-9][0-9][0-9][0-9][0-9][0-9].png' | wc -l; } +echo "=== opfreeze OP=$OP start=$(have)/$N job=${SLURM_JOB_ID:-?} $(date) ===" +if [[ "$(have)" -ge "$N" ]]; then echo "[complete] $OP already at $N"; exit 0; fi +CUDA_VISIBLE_DEVICES=0 python -u scripts/quant_sc_main.py \ + --wbits 8 --abits 8 --w_sym --a_sym \ + --timewise 1 --qklayerwise 1.0 --avlayerwise 1.0 --projlayerwise 1.0 --mlplayerwise 1.0 --inputprojlayerwise 1.0 \ + --sc_prec 8 --sc_fixed_level_prec --sc_qk_granularity per_row --sc_halve \ + --adaptive_mp --adaptive_mp_table "$CALIB" --mp_levels "$MP_LEVELS" \ + --image-size 256 --num-sampling-steps "$STEPS" --cfg-scale "$CFG" --batch-size "$BATCH" \ + --generate-fid-samples --balanced_classes --num-classes "$NUM_CLASSES" \ + --balanced_total_samples "$BALANCED" --num-fid-samples "$BALANCED" \ + --target_indices_path "$IDXF" --samples_dir_override "$SAMPLES" \ + --seed "$SEED" --results-dir "$LOG" --ckpt "$CKPT" \ + 2>&1 | tee "$LOG/run.log" +echo "=== opfreeze OP=$OP done=$(have)/$N $(date) ===" diff --git a/scripts/sbatch_uniform_auto.sb b/scripts/sbatch_uniform_auto.sb new file mode 100644 index 0000000..77e1ab5 --- /dev/null +++ b/scripts/sbatch_uniform_auto.sb @@ -0,0 +1,49 @@ +#!/bin/bash +#SBATCH --job-name=uni_auto +#SBATCH --partition=gpu-rtx6000 +#SBATCH --account=nbleier_owned1 +#SBATCH --gres=gpu:3 +#SBATCH --cpus-per-gpu=2 +#SBATCH --mem-per-gpu=12G +#SBATCH --time=1-12:00:00 +# uniform generation, TIMEOUT-safe auto-resume to 10000. Usage: sbatch sbatch_uniform_auto.sb [CHAIN] +set -uo pipefail +AVG="${1:?AVG}"; CHAIN="${2:-0}"; NUM_GPUS=3; MAXCHAIN=15 +SELF=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion/scripts/sbatch_uniform_auto.sb +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +SCRATCH=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi +CKPT=$SCRATCH/pretrained_models/DiT-XL-2-256x256.pt +SC_JSON=$SCRATCH/scmp_diffusion_fid_halve_bitrev/configs/sc_cfg_uniform${AVG}_all.json +OUT=$SCRATCH/scmp_diffusion_fid_uniform_bitrev_cfg15/uniform${AVG} +SAMPLES=$OUT/samples; IDX=$OUT/_indices; LOG=$OUT/_logs +NUM_FID=2000; BALANCED=10000; BATCH=64; STEPS=50; CFG=1.5; SEED=0; NUM_CLASSES=1000 +mkdir -p "$SAMPLES" "$IDX" "$LOG" +[[ -f "$SC_JSON" ]] || { echo "ERROR: missing $SC_JSON" >&2; exit 1; } +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh; conda activate qdit +export SC_OWEN_MODE=bitrev PYTHONUNBUFFERED=1 OMP_NUM_THREADS=2; cd "$REPO" +cnt(){ find "$SAMPLES" -maxdepth 1 -name '[0-9][0-9][0-9][0-9][0-9][0-9].png' | wc -l; } +START=$(cnt) +echo "=== uniform${AVG} auto: start=$START/$NUM_FID chain=$CHAIN job=${SLURM_JOB_ID:-?} $(date) ===" +if [[ "$START" -ge "$NUM_FID" ]]; then echo "[complete]"; exit 0; fi +if [[ "$CHAIN" -lt "$MAXCHAIN" ]]; then + sbatch -o "$OUT/slurm-%j.out" -e "$OUT/slurm-%j.err" --dependency=afterany:${SLURM_JOB_ID} "$SELF" "$AVG" $((CHAIN+1)) && echo "[chain] queued successor" +fi +python -u scripts/_plan_missing_indices.py "$SAMPLES" "$NUM_FID" "$NUM_GPUS" "$IDX" $((BALANCED/NUM_CLASSES)) +pids=() +for ((g=0; g "$LOG/gpu_${g}.log" 2>&1 & + pids+=($!) +done +for p in "${pids[@]}"; do wait "$p" || true; done +echo "=== uniform${AVG}: $START -> $(cnt)/$NUM_FID $(date) ===" diff --git a/scripts/sbatch_uniform_halve_auto.sb b/scripts/sbatch_uniform_halve_auto.sb new file mode 100755 index 0000000..f3fa9eb --- /dev/null +++ b/scripts/sbatch_uniform_halve_auto.sb @@ -0,0 +1,49 @@ +#!/bin/bash +#SBATCH --job-name=uni_halve +#SBATCH --partition=gpu-rtx6000 +#SBATCH --account=nbleier_owned1 +#SBATCH --gres=gpu:3 +#SBATCH --cpus-per-gpu=2 +#SBATCH --mem-per-gpu=12G +#SBATCH --time=1-12:00:00 +# uniform generation, TIMEOUT-safe auto-resume to 10000. Usage: sbatch sbatch_uniform_auto.sb [CHAIN] +set -uo pipefail +AVG="${1:?AVG}"; CHAIN="${2:-0}"; NUM_GPUS=3; MAXCHAIN=15 +SELF=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion/scripts/sbatch_uniform_halve_auto.sb +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +SCRATCH=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi +CKPT=$SCRATCH/pretrained_models/DiT-XL-2-256x256.pt +SC_JSON=$SCRATCH/scmp_diffusion_fid_halve_bitrev/configs/sc_cfg_uniform${AVG}_all.json +OUT=$SCRATCH/scmp_diffusion_fid_uniform_halveON_cfg15/uniform${AVG} +SAMPLES=$OUT/samples; IDX=$OUT/_indices; LOG=$OUT/_logs +NUM_FID=2000; BALANCED=10000; BATCH=64; STEPS=50; CFG=1.5; SEED=0; NUM_CLASSES=1000 +mkdir -p "$SAMPLES" "$IDX" "$LOG" +[[ -f "$SC_JSON" ]] || { echo "ERROR: missing $SC_JSON" >&2; exit 1; } +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh; conda activate qdit +export SC_OWEN_MODE=bitrev PYTHONUNBUFFERED=1 OMP_NUM_THREADS=2; cd "$REPO" +cnt(){ find "$SAMPLES" -maxdepth 1 -name '[0-9][0-9][0-9][0-9][0-9][0-9].png' | wc -l; } +START=$(cnt) +echo "=== uniform${AVG} auto: start=$START/$NUM_FID chain=$CHAIN job=${SLURM_JOB_ID:-?} $(date) ===" +if [[ "$START" -ge "$NUM_FID" ]]; then echo "[complete]"; exit 0; fi +if [[ "$CHAIN" -lt "$MAXCHAIN" ]]; then + sbatch -o "$OUT/slurm-%j.out" -e "$OUT/slurm-%j.err" --dependency=afterany:${SLURM_JOB_ID} "$SELF" "$AVG" $((CHAIN+1)) && echo "[chain] queued successor" +fi +python -u scripts/_plan_missing_indices.py "$SAMPLES" "$NUM_FID" "$NUM_GPUS" "$IDX" $((BALANCED/NUM_CLASSES)) +pids=() +for ((g=0; g "$LOG/gpu_${g}.log" 2>&1 & + pids+=($!) +done +for p in "${pids[@]}"; do wait "$p" || true; done +echo "=== uniform${AVG}: $START -> $(cnt)/$NUM_FID $(date) ===" diff --git a/tools/compare_old_vs_new.py b/tools/compare_old_vs_new.py new file mode 100644 index 0000000..fe5483d --- /dev/null +++ b/tools/compare_old_vs_new.py @@ -0,0 +1,207 @@ +"""Numerical comparison: scmp_llm/SC/sc_triton.py (old) vs scmp_kernels.sc.sc_matmul (new). + +For each shape pattern and mode, runs both implementations on identical inputs +with identical Sobol config + identical stoc_len, then reports: + - max abs diff between the two SC outputs + - max rel diff between the two SC outputs + - rel_err of each vs torch.matmul (the fp baseline) + +Expected outcome: outputs differ slightly due to the clipping margin removal +(bipolar ±125→±127, unipolar [2,253]→[0,255]). Both should be within sane +SC noise distance from the fp baseline. +""" +from __future__ import annotations +import sys, os +from pathlib import Path + +# Old impl — bare imports relative to SC/ folder +SCMP_LLM_SC = Path("/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_llm/SC") +sys.path.insert(0, str(SCMP_LLM_SC)) +import sc_triton as old_sc +from config_helpers import make_sobol_simple_config + +# New impl — installed as scmp_kernels package +from scmp_kernels.sc import sc_matmul as new_sc_matmul, clear_rng_cache +import scmp_kernels.sc.kernels as new_sc_kernels + +import torch + + +DEVICE = "cuda" + + +def diff(a, b): + """Return (max_abs, max_rel).""" + d = (a - b).abs() + rel = d / b.abs().clamp_min(1e-6) + return float(d.max().item()), float(rel.max().item()) + + +def rel_err(pred, target): + num = (pred - target).pow(2).mean().sqrt() + den = target.pow(2).mean().sqrt().clamp_min(1e-8) + return float((num / den).item()) + + +def line(label, vals): + print(f" {label:<60} " + " ".join(f"{v:>10}" for v in vals), flush=True) + + +def header(title): + print(f"\n{title}\n{'='*120}", flush=True) + + +def reset(): + old_sc.clear_rng_cache() + clear_rng_cache() + + +def run_case(name, a, b, *, mode, granularity_new, old_call, sc_prec=8, stoc_len=256): + """Run both impls on (a, b), print diff vs each other + each vs fp.""" + D = a.shape[-1] + cfg = make_sobol_simple_config(D, D, sc_prec) + fp = a @ b.transpose(-1, -2) if b.dim() == a.dim() else a @ b.t() + + reset() + out_old = old_call(a, b, cfg, sc_prec, stoc_len, mode) + + reset() + out_new = new_sc_matmul(a, b, granularity=granularity_new, mode=mode, + sc_prec=sc_prec, stoc_len=stoc_len, config=cfg) + + assert out_old.shape == out_new.shape == fp.shape, \ + f"{name}: shape mismatch old={out_old.shape} new={out_new.shape} fp={fp.shape}" + + max_abs, max_rel = diff(out_new, out_old) + re_old = rel_err(out_old, fp) + re_new = rel_err(out_new, fp) + bitwise = "BIT-IDENTICAL" if max_abs == 0 else ("≈identical" if max_abs < 1e-5 else "differ") + line(name, [f"{max_abs:.3e}", f"{max_rel:.3e}", + f"{re_old:.4f}", f"{re_new:.4f}", bitwise]) + + +# -------------------------------------------------------------------------- +# Adapters — each wraps an old-impl call to a uniform signature +# -------------------------------------------------------------------------- +def _old_enable_triton(a, b, cfg, sc_prec, stoc_len, mode): + return old_sc.sc_matmul_enable_triton( + a, b, a.max().item(), a.min().item(), b.max().item(), b.min().item(), + mode=mode, sc_prec=sc_prec, config=cfg, stoc_len=stoc_len, + ) + + +def _old_enable_batched_bipolar(q, k, cfg, sc_prec, stoc_len, mode): + q_maxs = q.amax(dim=(1, 2)) + q_mins = q.amin(dim=(1, 2)) + k_maxs = k.amax(dim=(1, 2)) + k_mins = k.amin(dim=(1, 2)) + return old_sc.sc_matmul_enable_batched_bipolar( + q, k, q_maxs, q_mins, k_maxs, k_mins, + sc_prec, cfg, stoc_len=stoc_len, + ) + + +def _old_grouped_enable(a, b, cfg, sc_prec, stoc_len, mode): + return old_sc.sc_matmul_grouped_enable_triton( + a, b, group_a=a.shape[0], group_b=b.shape[0], + mode=mode, sc_prec=sc_prec, config=cfg, stoc_len=stoc_len, + ) + + +def main(): + print(f"scmp_llm SC source: {SCMP_LLM_SC}", flush=True) + print(f"new scmp_kernels: {Path(new_sc_kernels.__file__).parent}", flush=True) + print(f"device: {torch.cuda.get_device_name(0)}", flush=True) + + line("CASE", ["max|Δ|", "max rel-Δ", "rel_err old", "rel_err new", "vs old"]) + print("-" * 120, flush=True) + + torch.manual_seed(0) + + # 2D matmul: per_tensor bipolar — sc_matmul_enable_triton vs sc_matmul(granularity="per_tensor") + header("BIPOLAR — 2D matmul A(N×D) @ B(M×D).T ─ enable-signal table-lookup") + line("CASE", ["max|Δ|", "max rel-Δ", "rel_err old", "rel_err new", "vs old"]) + print("-" * 120, flush=True) + + for shape in [(32, 64, 16), (32, 128, 64), (128, 1152, 1152)]: + N, D, M = shape + a = torch.randn(N, D, device=DEVICE) + b = torch.randn(M, D, device=DEVICE) * 0.1 + run_case(f"matmul N={N:<4} D={D:<5} M={M:<5} bipolar per_tensor", + a, b, mode="bipolar", granularity_new="per_tensor", + old_call=_old_enable_triton) + + # 2D unipolar + header("UNIPOLAR — 2D matmul A(N×D) @ B(M×D).T") + line("CASE", ["max|Δ|", "max rel-Δ", "rel_err old", "rel_err new", "vs old"]) + print("-" * 120, flush=True) + + for shape in [(32, 64, 16), (32, 128, 64)]: + N, D, M = shape + a = torch.rand(N, D, device=DEVICE) # unipolar wants non-neg-ish + b = torch.randn(M, D, device=DEVICE) * 0.1 + run_case(f"matmul N={N:<4} D={D:<5} M={M:<5} unipolar per_tensor", + a, b, mode="unipolar", granularity_new="per_tensor", + old_call=_old_enable_triton) + + # GEMV — rank-1 a (matrix-vector): A(1×D) @ B(M×D).T → (1, M) + header("GEMV — A(1×D) @ B(M×D).T ─ same path, just N=1") + line("CASE", ["max|Δ|", "max rel-Δ", "rel_err old", "rel_err new", "vs old"]) + print("-" * 120, flush=True) + + for D, M in [(64, 16), (128, 64), (1152, 1152)]: + a = torch.randn(1, D, device=DEVICE) + b = torch.randn(M, D, device=DEVICE) * 0.1 + run_case(f"GEMV D={D:<5} M={M:<5} bipolar per_tensor", + a, b, mode="bipolar", granularity_new="per_tensor", + old_call=_old_enable_triton) + + # 3D per-head bipolar — QK pattern in attention + header("3D BATCHED BIPOLAR — Q(BH×N×D) @ K(BH×N×D).T ─ per-head SC") + line("CASE", ["max|Δ|", "max rel-Δ", "rel_err old", "rel_err new", "vs old"]) + print("-" * 120, flush=True) + + for BH, N, D in [(8, 32, 64), (16, 256, 72)]: + q = torch.randn(BH, N, D, device=DEVICE) + k = torch.randn(BH, N, D, device=DEVICE) + run_case(f"QK BH={BH:<3} N={N:<5} D={D:<5} bipolar per_head", + q, k, mode="bipolar", granularity_new="per_head", + old_call=_old_enable_batched_bipolar) + + # AV-style grouped matmul (per-row groups) + header("GROUPED — A(N×D) @ B(M×D).T ─ per-row groups, group_a=N, group_b=M") + line("CASE", ["max|Δ|", "max rel-Δ", "rel_err old", "rel_err new", "vs old"]) + print("-" * 120, flush=True) + + for N, D, M in [(64, 32, 128), (256, 72, 256)]: + a = torch.softmax(torch.randn(N, M, device=DEVICE), dim=-1) # softmax-like + b = torch.randn(D, M, device=DEVICE) * 0.1 + N2 = a.shape[0]; M2 = b.shape[0] + # new dispatcher: granularity=per_row + explicit group_a/group_b + cfg = make_sobol_simple_config(M, M, 8) + fp = a @ b.t() + reset() + out_old = old_sc.sc_matmul_grouped_enable_triton( + a, b, group_a=N, group_b=D, mode="bipolar", sc_prec=8, config=cfg, stoc_len=256) + reset() + out_new = new_sc_matmul( + a, b, granularity="per_row", group_a=N, group_b=D, + mode="bipolar", sc_prec=8, stoc_len=256, config=cfg) + ma, mr = diff(out_new, out_old) + re_o = rel_err(out_old, fp); re_n = rel_err(out_new, fp) + bit = "BIT-IDENTICAL" if ma == 0 else ("≈identical" if ma < 1e-5 else "differ") + line(f"grouped N={N:<4} D={D:<4} M={M:<4} bipolar per_row", + [f"{ma:.3e}", f"{mr:.3e}", f"{re_o:.4f}", f"{re_n:.4f}", bit]) + + print("\nLegend:") + print(" max|Δ| = max |out_new - out_old| (0 ⇒ bit-identical)") + print(" max rel-Δ = max |out_new - out_old| / |out_old|") + print(" rel_err = SC vs fp baseline (lower = closer to fp)") + print() + print("Note: any non-zero Δ is explained entirely by the clipping-margin removal") + print(" (bipolar: ±125 → ±127, unipolar: [2,253] → [0,255]) which uses more") + print(" quantization levels — so the new version's rel_err should be ≤ the old.") + + +if __name__ == "__main__": + main() diff --git a/tools/compare_old_vs_new_isolated.py b/tools/compare_old_vs_new_isolated.py new file mode 100644 index 0000000..e805feb --- /dev/null +++ b/tools/compare_old_vs_new_isolated.py @@ -0,0 +1,119 @@ +"""Isolate the cause of the bipolar per_tensor diff. + +Monkey-patches scmp_llm's ``fused_quantize_bipolar`` to use ``q_clip = q_norm`` +(no margin). If this is the sole source of the diff, the output must become +bit-identical to scmp_kernels. +""" +from __future__ import annotations +import sys +from pathlib import Path + +SCMP_LLM_SC = Path("/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_llm/SC") +sys.path.insert(0, str(SCMP_LLM_SC)) +import sc_triton as old_sc +from config_helpers import make_sobol_simple_config + +from scmp_kernels.sc import sc_matmul as new_sc_matmul, clear_rng_cache + +import torch + + +# --- Save originals + install patched fused_quantize_bipolar --------------- +_orig_fqb = old_sc.fused_quantize_bipolar + + +def patched_fused_quantize_bipolar(fp_tensor, abs_max, sc_prec, rng_levels=None): + """scmp_llm's body, but with q_clip = q_norm (no margin) — matches scmp_kernels.""" + import triton + rows, cols = fp_tensor.shape + q_norm = 2 ** (sc_prec - 1) - 1 # 127 + q_clip = q_norm # 127 ← was q_norm - 2 + max_rng_val = old_sc._resolve_rng_levels(sc_prec, rng_levels) + abs_max = max(abs_max, 1e-5) + scale = abs_max / q_clip + inv_scale = 1.0 / scale + + boundary = torch.empty(rows, cols, dtype=torch.int16, device=fp_tensor.device) + sign = torch.empty(rows, cols, dtype=torch.int8, device=fp_tensor.device) + total = rows * cols + BLOCK = 1024 + grid = (triton.cdiv(total, BLOCK),) + old_sc.fused_quant_bipolar_kernel[grid]( + fp_tensor, boundary, sign, + inv_scale, q_clip, -q_clip, q_clip, max_rng_val, + rows, cols, BLOCK, + ) + return boundary, sign, scale + + +def diff(a, b): + d = (a - b).abs() + return float(d.max().item()), float((d / b.abs().clamp_min(1e-6)).max().item()) + + +def rel_err(pred, target): + num = (pred - target).pow(2).mean().sqrt() + den = target.pow(2).mean().sqrt().clamp_min(1e-8) + return float((num / den).item()) + + +def run_pair(name, a, b, *, patch): + """If patch=True, use the no-margin version of fused_quantize_bipolar.""" + D = a.shape[-1] + cfg = make_sobol_simple_config(D, D, 8) + fp = a @ b.t() + + if patch: + old_sc.fused_quantize_bipolar = patched_fused_quantize_bipolar + else: + old_sc.fused_quantize_bipolar = _orig_fqb + + old_sc.clear_rng_cache(); clear_rng_cache() + out_old = old_sc.sc_matmul_enable_triton( + a, b, a.max().item(), a.min().item(), b.max().item(), b.min().item(), + mode="bipolar", sc_prec=8, config=cfg, stoc_len=256, + ) + + old_sc.clear_rng_cache(); clear_rng_cache() + out_new = new_sc_matmul(a, b, granularity="per_tensor", mode="bipolar", + sc_prec=8, stoc_len=256, config=cfg) + + ma, mr = diff(out_new, out_old) + re_o = rel_err(out_old, fp) + re_n = rel_err(out_new, fp) + bit = "BIT-IDENTICAL" if ma == 0.0 else "differ" + flag = "[PATCHED]" if patch else "[ORIGINAL]" + print(f" {flag:<11} {name:<46} max|Δ|={ma:.3e} rel_err: old={re_o:.4f} new={re_n:.4f} {bit}", flush=True) + + +def main(): + torch.manual_seed(0) + print(f"device: {torch.cuda.get_device_name(0)}\n", flush=True) + + shapes = [ + ("matmul N=32 D=64 M=16 ", (32, 64, 16)), + ("matmul N=32 D=128 M=64 ", (32, 128, 64)), + ("matmul N=128 D=1152 M=1152", (128, 1152, 1152)), + ("GEMV N=1 D=64 M=16 ", (1, 64, 16)), + ("GEMV N=1 D=128 M=64 ", (1, 128, 64)), + ("GEMV N=1 D=1152 M=1152", (1, 1152, 1152)), + ] + + print("ORIGINAL scmp_llm (q_clip = q_norm - 2 = 125 for 8-bit)") + print("-" * 110, flush=True) + inputs = [] + for name, (N, D, M) in shapes: + a = torch.randn(N, D, device="cuda") + b = torch.randn(M, D, device="cuda") * 0.1 + inputs.append((name, a, b)) + run_pair(name, a, b, patch=False) + + print() + print("PATCHED scmp_llm (q_clip = q_norm = 127 ⇒ should match scmp_kernels)") + print("-" * 110, flush=True) + for name, a, b in inputs: + run_pair(name, a, b, patch=True) + + +if __name__ == "__main__": + main() diff --git a/tools/compare_old_vs_new_thorough.py b/tools/compare_old_vs_new_thorough.py new file mode 100644 index 0000000..faa3e5b --- /dev/null +++ b/tools/compare_old_vs_new_thorough.py @@ -0,0 +1,167 @@ +"""Exhaustive re-verification: scmp_llm/SC vs scmp_kernels.sc.sc_matmul. + +For each (granularity, mode) combination, sweeps: + - 3 different shapes + - 3 different random seeds + - 3 different stoc_len values (64, 128, 256) + - 2 different sc_prec values (6 and 8) where applicable + +Every case must report BIT-IDENTICAL (max|Δ|=0). Anything else means a real +divergence remains. +""" +from __future__ import annotations +import sys +from pathlib import Path + +# Old impl +SCMP_LLM_SC = Path("/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_llm/SC") +sys.path.insert(0, str(SCMP_LLM_SC)) +import sc_triton as old_sc +from config_helpers import make_sobol_simple_config + +# New impl +from scmp_kernels.sc import sc_matmul as new_sc_matmul, clear_rng_cache as new_clear + +import torch + + +# Sanity: prove the two impls are genuinely separate modules +print(f"old module: {old_sc.__file__}", flush=True) +print(f"old enable fn: id={id(old_sc.sc_matmul_enable_triton)}", flush=True) + +import scmp_kernels.sc.kernels as new_k +print(f"new module: {new_k.__file__}", flush=True) +print(f"new dispatch: id={id(new_sc_matmul)}", flush=True) + +assert old_sc.__file__ != new_k.__file__, \ + "scmp_llm and scmp_kernels share a file — comparison would be vacuous" +print() + + +def reset(): + old_sc.clear_rng_cache() + new_clear() + + +fail_count = 0 +case_count = 0 + + +def case(label, a, b, *, granularity, mode, sc_prec, stoc_len, + group_a=1, group_b=1, old_kind="enable_triton"): + """Run both, compare, increment counters.""" + global fail_count, case_count + case_count += 1 + D = a.shape[-1] + cfg = make_sobol_simple_config(D, D, sc_prec) + + reset() + if old_kind == "enable_triton": + out_old = old_sc.sc_matmul_enable_triton( + a, b, a.max().item(), a.min().item(), b.max().item(), b.min().item(), + mode=mode, sc_prec=sc_prec, config=cfg, stoc_len=stoc_len, + ) + elif old_kind == "enable_batched_bipolar": + out_old = old_sc.sc_matmul_enable_batched_bipolar( + a, b, a.amax(dim=(1, 2)), a.amin(dim=(1, 2)), + b.amax(dim=(1, 2)), b.amin(dim=(1, 2)), + sc_prec, cfg, stoc_len=stoc_len, + ) + elif old_kind == "grouped_enable": + out_old = old_sc.sc_matmul_grouped_enable_triton( + a, b, group_a=group_a, group_b=group_b, + mode=mode, sc_prec=sc_prec, config=cfg, stoc_len=stoc_len, + ) + else: + raise ValueError(old_kind) + + reset() + kw = dict(granularity=granularity, mode=mode, + sc_prec=sc_prec, stoc_len=stoc_len, config=cfg) + if granularity == "per_row": + kw.update(group_a=group_a, group_b=group_b) + out_new = new_sc_matmul(a, b, **kw) + + max_abs = float((out_old - out_new).abs().max().item()) + ok = (max_abs == 0.0) + if not ok: + fail_count += 1 + flag = "✓" if ok else "✗" + print(f" {flag} {label:<70} max|Δ|={max_abs:.3e}", flush=True) + + +def main(): + print(f"device: {torch.cuda.get_device_name(0)}", flush=True) + print(f"\n{'─'*100}\nBIPOLAR 2D per_tensor ─ the path where the clipping fix matters most", flush=True) + print("─" * 100, flush=True) + for seed in [0, 1, 42]: + torch.manual_seed(seed) + for (N, D, M) in [(8, 32, 8), (32, 128, 64), (128, 1152, 1152)]: + for sc_prec, stoc_len in [(8, 64), (8, 128), (8, 256), (6, 64)]: + a = torch.randn(N, D, device="cuda") + b = torch.randn(M, D, device="cuda") * 0.1 + label = f"seed={seed} N={N:<4} D={D:<5} M={M:<5} sc_prec={sc_prec} stoc_len={stoc_len:<3}" + case(label, a, b, granularity="per_tensor", mode="bipolar", + sc_prec=sc_prec, stoc_len=stoc_len) + + print(f"\n{'─'*100}\nBIPOLAR GEMV per_tensor (N=1)", flush=True) + print("─" * 100, flush=True) + for seed in [0, 1, 42]: + torch.manual_seed(seed) + for (D, M) in [(32, 8), (128, 64), (1152, 1152)]: + for stoc_len in [64, 128, 256]: + a = torch.randn(1, D, device="cuda") + b = torch.randn(M, D, device="cuda") * 0.1 + label = f"seed={seed} GEMV D={D:<5} M={M:<5} sc_prec=8 stoc_len={stoc_len:<3}" + case(label, a, b, granularity="per_tensor", mode="bipolar", + sc_prec=8, stoc_len=stoc_len) + + print(f"\n{'─'*100}\nUNIPOLAR 2D per_tensor", flush=True) + print("─" * 100, flush=True) + for seed in [0, 1, 42]: + torch.manual_seed(seed) + for (N, D, M) in [(8, 32, 8), (32, 128, 64)]: + for stoc_len in [128, 256]: + a = torch.rand(N, D, device="cuda") + b = torch.randn(M, D, device="cuda") * 0.1 + label = f"seed={seed} unipolar N={N:<4} D={D:<5} M={M:<5} stoc_len={stoc_len:<3}" + case(label, a, b, granularity="per_tensor", mode="unipolar", + sc_prec=8, stoc_len=stoc_len) + + print(f"\n{'─'*100}\nQK per_head bipolar (3D batched)", flush=True) + print("─" * 100, flush=True) + for seed in [0, 1, 42]: + torch.manual_seed(seed) + for (BH, N, D) in [(4, 16, 32), (8, 64, 64), (16, 256, 72)]: + for stoc_len in [128, 256]: + q = torch.randn(BH, N, D, device="cuda") + k = torch.randn(BH, N, D, device="cuda") + label = f"seed={seed} QK BH={BH:<3} N={N:<5} D={D:<5} stoc_len={stoc_len:<3}" + case(label, q, k, granularity="per_head", mode="bipolar", + sc_prec=8, stoc_len=stoc_len, old_kind="enable_batched_bipolar") + + print(f"\n{'─'*100}\nAV grouped per_row bipolar", flush=True) + print("─" * 100, flush=True) + for seed in [0, 1, 42]: + torch.manual_seed(seed) + for (N, D, M) in [(16, 8, 32), (64, 32, 128), (256, 72, 256)]: + for stoc_len in [128, 256]: + attn = torch.softmax(torch.randn(N, M, device="cuda"), dim=-1) + v = torch.randn(D, M, device="cuda") * 0.1 + label = f"seed={seed} AV N={N:<4} D={D:<4} M={M:<4} stoc_len={stoc_len:<3}" + case(label, attn, v, granularity="per_row", mode="bipolar", + sc_prec=8, stoc_len=stoc_len, + group_a=N, group_b=D, old_kind="grouped_enable") + + print(f"\n{'='*100}", flush=True) + print(f"Total cases: {case_count} Bit-identical: {case_count - fail_count} Diverged: {fail_count}", flush=True) + print("=" * 100, flush=True) + if fail_count == 0: + print("\n ✓ ALL CASES BIT-IDENTICAL — scmp_llm and scmp_kernels produce equivalent SC outputs.", flush=True) + else: + print(f"\n ✗ {fail_count} cases diverged — investigate.", flush=True) + sys.exit(0 if fail_count == 0 else 1) + + +if __name__ == "__main__": + main() diff --git a/tools/compare_vs_cpu_reference.py b/tools/compare_vs_cpu_reference.py new file mode 100644 index 0000000..0330a96 --- /dev/null +++ b/tools/compare_vs_cpu_reference.py @@ -0,0 +1,156 @@ +"""Compare scmp_kernels.sc.sc_matmul (GPU Triton) against the CPU reference +in scmp_llm/SC/sc_enable.py (sc_matmul_enable). + +The CPU reference is the gold-standard semantic check: it uses Python/PyTorch +ops, no Triton, two algorithms: + - cycle_by_cycle: exact UnarySim FSUMul simulator (very slow, small shapes) + - k_shortcut: vectorized prefix-sum equivalent (faster) + +For each shape × mode × stoc_len, runs: + GPU Triton (new scmp_kernels) + CPU k_shortcut (scmp_llm reference) + CPU cycle_by_cycle (scmp_llm reference, on small shapes only) + +Compares each pair. Expectations: + - CPU k_shortcut vs cycle_by_cycle: should be bit-identical (same math). + - GPU vs CPU: should be bit-identical OR differ only by float32 rounding + (≤ 1e-4 absolute, dominated by accumulation order). +""" +from __future__ import annotations +import sys, time +from pathlib import Path + +SCMP_LLM_SC = Path("/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_llm/SC") +sys.path.insert(0, str(SCMP_LLM_SC)) +import sc_enable as cpu_ref +from config_helpers import make_sobol_simple_config + +from scmp_kernels.sc import sc_matmul as gpu_matmul, clear_rng_cache + +import torch + + +def diff(a, b): + d = (a - b).abs() + rel = d / b.abs().clamp_min(1e-6) + return float(d.max().item()), float(d.mean().item()), float(rel.max().item()) + + +def rel_err(pred, target): + num = (pred - target).pow(2).mean().sqrt() + den = target.pow(2).mean().sqrt().clamp_min(1e-8) + return float((num / den).item()) + + +def reset(): + clear_rng_cache() + + +def section(title): + print(f"\n{title}\n{'─'*120}", flush=True) + print(f" {'CASE':<55} {'max|Δ|':>10} {'mean|Δ|':>10} {'max rel':>10} {'verdict':>20}", flush=True) + + +def verdict(ma, threshold=1e-4): + if ma == 0.0: return "BIT-IDENTICAL" + if ma < threshold: return f"≈identical ({ma:.1e})" + return f"differ ({ma:.2e})" + + +def main(): + print(f"device: {torch.cuda.get_device_name(0)}", flush=True) + print(f"CPU reference: {cpu_ref.__file__}", flush=True) + print(f"GPU impl: scmp_kernels.sc.sc_matmul", flush=True) + + # ---------------------------------------------------------------- + # 1. CPU k_shortcut vs CPU cycle_by_cycle (both reference) + # ---------------------------------------------------------------- + section("1. CPU references self-consistent? — k_shortcut vs cycle_by_cycle") + torch.manual_seed(0) + for (N, D, M) in [(4, 16, 4), (8, 32, 8)]: # cycle_by_cycle is slow + for mode in ["bipolar", "unipolar"]: + if mode == "unipolar": + a = torch.rand(N, D, device="cuda") + else: + a = torch.randn(N, D, device="cuda") + b = torch.randn(M, D, device="cuda") * 0.1 + cfg = make_sobol_simple_config(D, D, 8) + out_cyc = cpu_ref.sc_matmul_enable( + a, b, a.max().item(), a.min().item(), b.max().item(), b.min().item(), + mode=mode, sc_prec=8, config=cfg, method="cycle_by_cycle") + out_ks = cpu_ref.sc_matmul_enable( + a, b, a.max().item(), a.min().item(), b.max().item(), b.min().item(), + mode=mode, sc_prec=8, config=cfg, method="k_shortcut") + ma, me, mr = diff(out_ks, out_cyc) + print(f" {f'{mode:<8} N={N:<3} D={D:<4} M={M:<3} sc_prec=8 stoc_len=256':<55} " + f"{ma:>10.3e} {me:>10.3e} {mr:>10.3e} {verdict(ma):>20}", flush=True) + + # ---------------------------------------------------------------- + # 2. GPU Triton vs CPU k_shortcut — bipolar per_tensor + # ---------------------------------------------------------------- + section("2. GPU Triton vs CPU k_shortcut — bipolar per_tensor") + for (N, D, M) in [(8, 32, 8), (32, 128, 64), (64, 256, 32)]: + for stoc_len in [64, 256]: + torch.manual_seed(0) + a = torch.randn(N, D, device="cuda") + b = torch.randn(M, D, device="cuda") * 0.1 + cfg = make_sobol_simple_config(D, D, 8) + reset() + out_cpu = cpu_ref.sc_matmul_enable( + a, b, a.max().item(), a.min().item(), b.max().item(), b.min().item(), + mode="bipolar", sc_prec=8, config=cfg, method="k_shortcut") + # CPU ref uses stoc_len = 2**sc_prec = 256 always (no stoc_len arg!) + # So we only compare at stoc_len=256 for full equivalence + if stoc_len != 256: + continue + reset() + out_gpu = gpu_matmul(a, b, granularity="per_tensor", mode="bipolar", + sc_prec=8, stoc_len=stoc_len, config=cfg) + ma, me, mr = diff(out_gpu, out_cpu) + print(f" {f'bipolar N={N:<3} D={D:<4} M={M:<3} stoc_len={stoc_len:<3}':<55} " + f"{ma:>10.3e} {me:>10.3e} {mr:>10.3e} {verdict(ma):>20}", flush=True) + + # ---------------------------------------------------------------- + # 3. GPU vs CPU — unipolar + # ---------------------------------------------------------------- + section("3. GPU Triton vs CPU k_shortcut — unipolar per_tensor") + for (N, D, M) in [(8, 32, 8), (16, 128, 32)]: + torch.manual_seed(0) + a = torch.rand(N, D, device="cuda") + b = torch.randn(M, D, device="cuda") * 0.1 + cfg = make_sobol_simple_config(D, D, 8) + reset() + out_cpu = cpu_ref.sc_matmul_enable( + a, b, a.max().item(), a.min().item(), b.max().item(), b.min().item(), + mode="unipolar", sc_prec=8, config=cfg, method="k_shortcut") + reset() + out_gpu = gpu_matmul(a, b, granularity="per_tensor", mode="unipolar", + sc_prec=8, stoc_len=256, config=cfg) + ma, me, mr = diff(out_gpu, out_cpu) + print(f" {f'unipolar N={N:<3} D={D:<4} M={M:<3} stoc_len=256':<55} " + f"{ma:>10.3e} {me:>10.3e} {mr:>10.3e} {verdict(ma):>20}", flush=True) + + # ---------------------------------------------------------------- + # 4. Numerical accuracy: both vs torch.matmul fp baseline + # ---------------------------------------------------------------- + section("4. SC vs fp baseline — both impls should match torch.matmul within SC noise band") + torch.manual_seed(0) + N, D, M = 32, 128, 64 + a = torch.randn(N, D, device="cuda") + b = torch.randn(M, D, device="cuda") * 0.1 + fp = a @ b.t() + cfg = make_sobol_simple_config(D, D, 8) + reset() + out_gpu = gpu_matmul(a, b, granularity="per_tensor", mode="bipolar", + sc_prec=8, stoc_len=256, config=cfg) + out_cpu = cpu_ref.sc_matmul_enable( + a, b, a.max().item(), a.min().item(), b.max().item(), b.min().item(), + mode="bipolar", sc_prec=8, config=cfg, method="k_shortcut") + print(f" bipolar N=32 D=128 M=64 stoc_len=256", flush=True) + print(f" rel_err GPU vs fp: {rel_err(out_gpu, fp):.4f}", flush=True) + print(f" rel_err CPU vs fp: {rel_err(out_cpu, fp):.4f}", flush=True) + print(f" rel_err GPU vs CPU: {rel_err(out_gpu, out_cpu):.4e}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/tools/sc_call_counter.py b/tools/sc_call_counter.py new file mode 100644 index 0000000..ab59822 --- /dev/null +++ b/tools/sc_call_counter.py @@ -0,0 +1,61 @@ +"""Wrap sc_matmul to count calls per granularity, then run quant_sc_main.py. + +Confirms SC kernels are actually invoked during DiT sampling, not skipped. +Prints a count summary at process exit. +""" +import sys, os, atexit +from pathlib import Path +from collections import Counter + +ROOT = Path('/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion') +sys.path.insert(0, str(ROOT)) + +import scmp_kernels.sc.matmul as mm +_orig = mm.sc_matmul +counts = Counter() + +def _wrapped(a, b, granularity="per_row", **kw): + key = ( + granularity, + kw.get('mode', 'bipolar'), + 'chunk' if kw.get('chunk_d', 0) > 0 else 'nochunk', + 'group' if (kw.get('group_a', 1) > 1 or kw.get('group_b', 1) > 1) else 'nogroup', + ) + counts[key] += 1 + return _orig(a, b, granularity=granularity, **kw) + +mm.sc_matmul = _wrapped + +# Patch every site that did `from scmp_kernels.sc import sc_matmul` +import scmp_kernels.sc as sc_pkg +sc_pkg.sc_matmul = _wrapped +import qdit.sc_integration.sc_attention as sa +import qdit.sc_integration.sc_mlp as sm +sa.sc_matmul = _wrapped +sm.sc_matmul = _wrapped + +def _summary(): + total = sum(counts.values()) + print(f"\n========== sc_matmul call summary ==========", flush=True) + print(f"Total sc_matmul invocations: {total}", flush=True) + for k, n in sorted(counts.items(), key=lambda kv: -kv[1]): + gran, mode, chunk, group = k + print(f" {n:>6} granularity={gran:<10} mode={mode:<8} {chunk:<8} {group}", flush=True) + print(f"============================================\n", flush=True) +atexit.register(_summary) + +# Now run quant_sc_main as if invoked directly +sys.argv = [ + 'quant_sc_main.py', + '--ckpt', '/nfs/turbo/coe-nbleier/zhkangqi/pretrained_models/DiT-XL-2-256x256.pt', + '--wbits', '8', '--abits', '8', '--w_sym', '--a_sym', + '--timewise', '0.5', '--qklayerwise', '0.8', + '--avlayerwise', '0.0', '--projlayerwise', '0.0', + '--mlplayerwise', '0.0', '--inputprojlayerwise', '0.0', + '--sc_prec', '8', + '--image-size', '256', '--num-sampling-steps', '50', + '--cfg-scale', '4', '--batch-size', '8', + '--results-dir', 'results/smoke_e2e_count', +] +os.chdir(str(ROOT)) +exec(compile(open('scripts/quant_sc_main.py').read(), 'scripts/quant_sc_main.py', 'exec')) diff --git a/tools/sc_call_counter_full.py b/tools/sc_call_counter_full.py new file mode 100644 index 0000000..197738c --- /dev/null +++ b/tools/sc_call_counter_full.py @@ -0,0 +1,59 @@ +"""Full SC: all operators enabled, uniform stoc_len=128. Count sc_matmul calls.""" +import sys, os, atexit +from pathlib import Path +from collections import Counter + +ROOT = Path('/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion') +sys.path.insert(0, str(ROOT)) + +import scmp_kernels.sc.matmul as mm +_orig = mm.sc_matmul +counts = Counter() + +def _wrapped(a, b, granularity="per_row", **kw): + key = ( + granularity, + kw.get('mode', 'bipolar'), + 'chunk' if kw.get('chunk_d', 0) > 0 else 'nochunk', + 'group' if (kw.get('group_a', 1) > 1 or kw.get('group_b', 1) > 1) else 'nogroup', + ) + counts[key] += 1 + return _orig(a, b, granularity=granularity, **kw) + +mm.sc_matmul = _wrapped +import scmp_kernels.sc as sc_pkg +sc_pkg.sc_matmul = _wrapped +import qdit.sc_integration.sc_attention as sa +import qdit.sc_integration.sc_mlp as sm +sa.sc_matmul = _wrapped +sm.sc_matmul = _wrapped + +def _summary(): + total = sum(counts.values()) + print(f"\n========== sc_matmul call summary (uniform stoc_len=128, ALL OPS) ==========", flush=True) + print(f"Total sc_matmul invocations: {total}", flush=True) + for k, n in sorted(counts.items(), key=lambda kv: -kv[1]): + gran, mode, chunk, group = k + print(f" {n:>6} granularity={gran:<10} mode={mode:<8} {chunk:<8} {group}", flush=True) + print(f"============================================================================\n", flush=True) +atexit.register(_summary) + +sys.argv = [ + 'quant_sc_main.py', + '--ckpt', '/nfs/turbo/coe-nbleier/zhkangqi/pretrained_models/DiT-XL-2-256x256.pt', + '--wbits', '8', '--abits', '8', '--w_sym', '--a_sym', + '--timewise', '1.0', + '--qklayerwise', '1.0', + '--avlayerwise', '1.0', + '--projlayerwise', '1.0', + '--mlplayerwise', '1.0', + '--inputprojlayerwise', '1.0', + '--sc_prec', '8', + '--sc_fixed_level_prec', + '--sc_config', 'results/sc_cfg_uniform128_all.json', + '--image-size', '256', '--num-sampling-steps', '50', + '--cfg-scale', '4', '--batch-size', '8', + '--results-dir', 'results/smoke_e2e_all_L128', +] +os.chdir(str(ROOT)) +exec(compile(open('scripts/quant_sc_main.py').read(), 'scripts/quant_sc_main.py', 'exec')) diff --git a/tools/smoke_test_e2e.py b/tools/smoke_test_e2e.py new file mode 100644 index 0000000..3ee3700 --- /dev/null +++ b/tools/smoke_test_e2e.py @@ -0,0 +1,223 @@ +"""End-to-end smoke test for the scmp_diffusion bootstrap. + +Run on a GPU node. Verifies: + 1. submodule imports resolve (scmp_kernels.sc / scmp_kernels.mp) + 2. sc_matmul reaches every granularity + mode combination + 3. SC results stay within sane rel-err of torch.matmul + 4. Q-DiT integration imports resolve (qdit.sc_integration.{sc_attention, sc_mlp, noise_matmul}) + 5. SCMlp forward pass produces a finite, non-zero, sensibly-scaled output + 6. SCAttention forward pass produces a finite output + 7. No remaining references to deprecated names + +Exits non-zero on any failure with a clear message. +""" +from __future__ import annotations + +import importlib +import sys +import time +import traceback +from pathlib import Path + +# Make local Q-DiT package importable +ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(ROOT)) + + +def section(name): + print(f"\n{'='*80}\n{name}\n{'='*80}", flush=True) + + +def check(label, fn): + t0 = time.time() + try: + result = fn() + dt = (time.time() - t0) * 1000 + msg = f" PASS {label:<60} ({dt:>6.1f} ms)" + if result: + msg += f" {result}" + print(msg, flush=True) + return True + except Exception: + print(f" FAIL {label}", flush=True) + traceback.print_exc() + return False + + +def main(): + failures = [] + + section("1. Imports") + import torch + print(f" torch {torch.__version__}, cuda available: {torch.cuda.is_available()}", flush=True) + if not torch.cuda.is_available(): + print(" no CUDA — skipping kernel execution"); sys.exit(2) + print(f" device: {torch.cuda.get_device_name(0)}", flush=True) + try: + import triton + print(f" triton {triton.__version__}", flush=True) + except Exception as e: + print(f" triton import: {e}", flush=True) + sys.exit(2) + + if not check("import scmp_kernels.sc", + lambda: importlib.import_module("scmp_kernels.sc").__name__): + failures.append("scmp_kernels.sc") + + if not check("import scmp_kernels.mp", + lambda: importlib.import_module("scmp_kernels.mp").__name__): + failures.append("scmp_kernels.mp") + + if not check("from scmp_kernels.sc import sc_matmul, clear_rng_cache, det_kernel_tuning", + lambda: __import__("scmp_kernels.sc", fromlist=["sc_matmul","clear_rng_cache","det_kernel_tuning"])): + failures.append("sc public API") + + if not check("from scmp_kernels.mp import (9 names)", + lambda: [getattr(__import__("scmp_kernels.mp", fromlist=[n]), n) + for n in ("MPConfig","AdaptiveMPConfig","RangeMPConfig","RowAssignment", + "classify_rows_by_metric","adaptive_classify_rows", + "classify_groups_by_range","MPDistributionLogger","MetricProfiler")] + and "OK"): + failures.append("mp public API") + + section("2. sc_matmul granularity sweep — basic correctness vs torch.matmul") + from scmp_kernels.sc import sc_matmul + from scmp_kernels.sc.config_helpers import make_sobol_simple_config + + def rel_err(pred, target): + num = (pred - target).pow(2).mean().sqrt() + den = target.pow(2).mean().sqrt().clamp_min(1e-8) + return float((num / den).item()) + + torch.manual_seed(0) + device = "cuda" + + # per_tensor 2D bipolar + def per_tensor_2d_bipolar(): + a = torch.randn(32, 128, device=device, dtype=torch.float32) + b = torch.randn(64, 128, device=device, dtype=torch.float32) * 0.1 + fp = a @ b.t() + sc = sc_matmul(a, b, granularity="per_tensor", mode="bipolar", sc_prec=8, stoc_len=256) + e = rel_err(sc, fp) + return f"rel_err={e:.4f} shape={tuple(sc.shape)}" + + if not check("sc_matmul per_tensor 2D bipolar (32×128)@(64×128).T", per_tensor_2d_bipolar): + failures.append("per_tensor 2D bipolar") + + def per_row_2d_bipolar(): + a = torch.randn(32, 128, device=device, dtype=torch.float32) + b = torch.randn(64, 128, device=device, dtype=torch.float32) * 0.1 + fp = a @ b.t() + sc = sc_matmul(a, b, granularity="per_row", mode="bipolar", sc_prec=8, stoc_len=256) + e = rel_err(sc, fp) + return f"rel_err={e:.4f} shape={tuple(sc.shape)}" + + if not check("sc_matmul per_row 2D bipolar", per_row_2d_bipolar): + failures.append("per_row 2D bipolar") + + def per_row_mlp_chunked(): + a = torch.randn(32, 1152, device=device, dtype=torch.float32) + b = torch.randn(64, 1152, device=device, dtype=torch.float32) * 0.05 + fp = a @ b.t() + sc = sc_matmul(a, b, granularity="per_row", mode="bipolar", + chunk_d=72, sc_prec=8, stoc_len=256) + e = rel_err(sc, fp) + return f"rel_err={e:.4f} shape={tuple(sc.shape)} (chunked D=1152→72)" + + if not check("sc_matmul per_row MLP chunked", per_row_mlp_chunked): + failures.append("per_row MLP chunked") + + def per_row_grouped(): + N, M, D = 64, 32, 128 + a = torch.softmax(torch.randn(N, M, device=device), dim=-1) + b = torch.randn(D, M, device=device, dtype=torch.float32) * 0.1 + fp = a @ b.t() + sc = sc_matmul(a, b, granularity="per_row", mode="bipolar", + group_a=N, group_b=D, sc_prec=8, stoc_len=256) + e = rel_err(sc, fp) + return f"rel_err={e:.4f} shape={tuple(sc.shape)} (group_a={N},group_b={D})" + + if not check("sc_matmul per_row grouped (AV pattern)", per_row_grouped): + failures.append("per_row grouped") + + def per_head_bipolar(): + BH, N, D = 8, 32, 64 + q = torch.randn(BH, N, D, device=device, dtype=torch.float32) + k = torch.randn(BH, N, D, device=device, dtype=torch.float32) + fp = q @ k.transpose(-1, -2) + sc = sc_matmul(q, k, granularity="per_head", mode="bipolar", sc_prec=8, stoc_len=256) + e = rel_err(sc, fp) + return f"rel_err={e:.4f} shape={tuple(sc.shape)}" + + if not check("sc_matmul per_head bipolar (8×32×64)@(8×32×64).T", per_head_bipolar): + failures.append("per_head bipolar") + + def per_row_unipolar(): + a = torch.rand(32, 64, device=device, dtype=torch.float32) + b = torch.randn(16, 64, device=device, dtype=torch.float32) * 0.1 + fp = a @ b.t() + sc = sc_matmul(a, b, granularity="per_row", mode="unipolar", sc_prec=8, stoc_len=256) + e = rel_err(sc, fp) + return f"rel_err={e:.4f} shape={tuple(sc.shape)}" + + if not check("sc_matmul per_row 2D unipolar", per_row_unipolar): + failures.append("per_row unipolar") + + section("3. Q-DiT integration imports") + + qdit_imports = [ + ("qdit.sc_integration", ["SCController", "SCAttention", "SCMlp", "SCDiTBlock", "MPConfig", "add_sc_wrapper", "create_sc_controller_from_args"]), + ("qdit.sc_integration.sc_attention", ["SCAttention"]), + ("qdit.sc_integration.sc_mlp", ["SCMlp"]), + ("qdit.sc_integration.noise_matmul", ["noisy_sc_matmul"]), + ("qdit.sc_integration.sc_controller", ["SCController"]), + ("qdit.sc_integration.mp_config", ["MPConfig", "AdaptiveMPConfig"]), + ] + + for mod_name, names in qdit_imports: + def _doit(m=mod_name, ns=names): + mod = importlib.import_module(m) + missing = [n for n in ns if not hasattr(mod, n)] + if missing: + raise AttributeError(f"missing names in {m}: {missing}") + return f"resolved: {', '.join(ns)}" + if not check(f"import {mod_name}", _doit): + failures.append(mod_name) + + section("4. Deprecated-name surveillance") + + deprecated = ["sc_matmul_per_tensor", "sc_matmul_mlp", "sc_matmul_grouped", + "sc_matmul_enable_triton", "sc_matmul_enable_triton_mlp", + "sc_matmul_grouped_enable_triton", "sc_matmul_enable_batched_bipolar", + "bin_to_stoc_packed", "xnor_matmul"] + def deprecated_absent(): + import scmp_kernels.sc as sck_sc + present = [n for n in deprecated if hasattr(sck_sc, n)] + if present: + raise AssertionError(f"deprecated names still exported: {present}") + return f"none of {len(deprecated)} deprecated names re-emerged" + if not check("scmp_kernels.sc has no deprecated public names", deprecated_absent): + failures.append("deprecated surveillance") + + def sc_enable_absent(): + import qdit.sc_integration.sc_controller as ctrl + # constructor should not accept sc_enable + from inspect import signature + params = signature(ctrl.SCController.__init__).parameters + if "sc_enable" in params: + raise AssertionError("SCController.__init__ still accepts sc_enable") + return "SCController has no sc_enable parameter" + if not check("SCController no longer accepts sc_enable", sc_enable_absent): + failures.append("sc_enable removed") + + section("Summary") + if failures: + print(f"\n {len(failures)} FAILURES:") + for f in failures: print(f" - {f}") + sys.exit(1) + print("\n ALL CHECKS PASSED") + sys.exit(0) + + +if __name__ == "__main__": + main() From ca5eab07f2d74f9a6f998b0e5571f151b0efa77c Mon Sep 17 00:00:00 2001 From: heroarmor <162866837+heroarmor@users.noreply.github.com> Date: Sat, 11 Jul 2026 13:55:47 -0400 Subject: [PATCH 4/7] kernels: bump submodule to trace(#25)+#28 scramble-cache fix + doc cleanup Submodule now at local branch local/trace+scramble-cache (6c4539a): PR#25 trace subsystem + cherry-picked PR#28 enable-table scramble cache-key fix + mp/config doc cleanup. NOTE: 6c4539a is LOCAL-ONLY (not yet pushed to the heroarmor fork); a fresh 'git submodule update' cannot resolve it until pushed. Co-Authored-By: Claude Opus 4.8 (1M context) --- scmp_kernels | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scmp_kernels b/scmp_kernels index 9358239..6c4539a 160000 --- a/scmp_kernels +++ b/scmp_kernels @@ -1 +1 @@ -Subproject commit 9358239af8d65f54a08e4a47e090933947626a41 +Subproject commit 6c4539ae35a60ce95f3e7747edc81444d823cc55 From 1b74fab82b16ba3ebf0f3b887744d6c823cbead2 Mon Sep 17 00:00:00 2001 From: heroarmor <162866837+heroarmor@users.noreply.github.com> Date: Wed, 15 Jul 2026 04:58:22 -0400 Subject: [PATCH 5/7] fix(halve): enforce max stoc_len 128 end-to-end; cfg-matched calibration; mse_rel metric MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The halve MP line calibrated and ran levels above the halve-mode maximum (2^(sc_prec-1)=128): the 256/192 levels are unrealizable on halve hardware and silently inflated attention accuracy (up to 99% of qk rows ran 256/192- cycle streams). Fixes, one per layer: - scmp_kernels bump -> c967d3a: sc_matmul now raises on explicit stoc_len/rng_levels above 2^(sc_prec-1) under halve+bipolar (upstream kernels PR #30), so stale 256-max tables crash instead of overspending. - calibrate_mp_thresholds.py: hard-validate mp_levels and budget_ref_stoc_len <= 2^(sc_prec-1) under --sc_halve; add --metric mse_rel (squared relative L2 = cosine's quadratic curvature + gain-error sensitivity; a pure 10% gain error scores 0.0 under cosine, 0.01 under mse_rel). - sbatch_calib_globalpr_flop_halve.sb / sbatch_calib_globalpr_halve.sb: 128-capped level grid, budget_ref 128, BUDGET_RATIO computed from AVG in-script (avg/128 — passing the old avg/256 ratios silently halved every budget), optional METRIC arg with suffixed OUTDIR, and calibration now runs at the deployment CFG (teacher_cfg_scale 1.5, batch 16 so the CFG-doubled batch keeps the old 32-row memory footprint) instead of cfg=0. - sbatch_mp_globalpr_flop_halve_auto.sb: [METRIC] and [NUM_FID] args (metric-suffixed calib/output dirs; NUM_FID for fast probes), points at the 128-capped tables; old 256-grid tables/arms preserved in the un-suffixed dirs for comparison. Validated: avg48 recalibration hits expected_avg_stoc_len=48.00; runtime realized per-op distribution matches the table within 1-3pp per level; on the noise-aligned index prefix the fixed adaptive arms (cosine and mse_rel) beat uniform48 on fidelity-to-FP (2.5% vs 5% trajectory divergence). Co-Authored-By: Claude Fable 5 --- scmp_kernels | 2 +- scripts/calibrate_mp_thresholds.py | 41 ++++++++++++++++++- scripts/sbatch_calib_globalpr_flop_halve.sb | 30 ++++++++++---- scripts/sbatch_calib_globalpr_halve.sb | 26 ++++++++---- scripts/sbatch_mp_globalpr_flop_halve_auto.sb | 22 ++++++---- 5 files changed, 95 insertions(+), 26 deletions(-) diff --git a/scmp_kernels b/scmp_kernels index 6c4539a..c967d3a 160000 --- a/scmp_kernels +++ b/scmp_kernels @@ -1 +1 @@ -Subproject commit 6c4539ae35a60ce95f3e7747edc81444d823cc55 +Subproject commit c967d3aa9014f37a2f6823048dd02993e2e1dd21 diff --git a/scripts/calibrate_mp_thresholds.py b/scripts/calibrate_mp_thresholds.py index 4b70425..92f5596 100644 --- a/scripts/calibrate_mp_thresholds.py +++ b/scripts/calibrate_mp_thresholds.py @@ -141,6 +141,22 @@ def _cosine_dist_heads(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor return 1.0 - cos +def _relative_mse_rows(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor: + """Squared relative L2 per row: ||d||^2 / ||t||^2. + + Small-error expansion = gain_err^2 + orthogonal_r^2: keeps cosine's + quadratic (second-order-loss) curvature and relative row weighting, but + also sees pure gain/bias errors that cosine is blind to. + """ + r = _relative_l2_rows(pred, target) + return r * r + + +def _relative_mse_heads(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor: + r = _relative_l2_heads(pred, target) + return r * r + + # Globals dispatched from CLI flags --metric / --teacher. _METRIC_ROWS = _relative_l2_rows _METRIC_HEADS = _relative_l2_heads @@ -1098,9 +1114,10 @@ def _build_parser(): parser.add_argument( "--metric", type=str, - choices=["l2", "cosine"], + choices=["l2", "cosine", "mse_rel"], default="cosine", - help="Per-unit error metric. cosine = 1 - cos_sim (recommended for fix mode).", + help="Per-unit error metric. cosine = 1 - cos_sim (recommended for fix mode); " + "mse_rel = ||d||^2/||t||^2 (cosine's curvature + gain-error sensitivity).", ) parser.add_argument( "--teacher", @@ -1127,6 +1144,23 @@ def main(): explicit_timesteps = _parse_csv_ints(args.calib_timesteps) budget_ref_stoc_len = args.budget_ref_stoc_len or max(levels) + # Halve mode: the sign-magnitude grid (and hence the longest realizable + # stream) is 2^(sc_prec-1). Levels above that would probe accuracy the + # halve hardware cannot deliver, and a larger budget_ref silently rescales + # every budget_ratio, so reject both loudly. + if args.sc_halve: + halve_max = 2 ** (args.sc_prec - 1) + bad = [lv for lv in levels if lv > halve_max] + if bad: + raise ValueError( + f"--sc_halve: mp_levels {bad} exceed the halve-mode maximum " + f"2^(sc_prec-1)={halve_max}; use levels <= {halve_max}.") + if budget_ref_stoc_len > halve_max: + raise ValueError( + f"--sc_halve: budget_ref_stoc_len={budget_ref_stoc_len} exceeds " + f"the halve-mode maximum {halve_max}. Express budget_ratio " + f"against {halve_max} (avg = budget_ratio * {halve_max}).") + latent_size = args.image_size // 8 model = DiT_models[args.model]( input_size=latent_size, @@ -1168,6 +1202,9 @@ def main(): if args.metric == "cosine": _METRIC_ROWS = _cosine_dist_rows _METRIC_HEADS = _cosine_dist_heads + elif args.metric == "mse_rel": + _METRIC_ROWS = _relative_mse_rows + _METRIC_HEADS = _relative_mse_heads _USE_FP_TEACHER = (args.teacher == "fp") _SC_HALVE = bool(args.sc_halve) print(f"SC halve during calibration: {_SC_HALVE}") diff --git a/scripts/sbatch_calib_globalpr_flop_halve.sb b/scripts/sbatch_calib_globalpr_flop_halve.sb index 06d1984..4239749 100644 --- a/scripts/sbatch_calib_globalpr_flop_halve.sb +++ b/scripts/sbatch_calib_globalpr_flop_halve.sb @@ -8,26 +8,40 @@ #SBATCH --time=03:00:00 # FLOP-weighted + HALVE per-row GLOBAL calibration for one budget (fix branch: # fix/flop-cost-perrow-qk-drop-alphabeta). Matches the halve MP deployment. -# Usage: sbatch sbatch_calib_globalpr_flop_halve.sb +# Halve mode caps stream length at 2^(sc_prec-1)=128: levels above 128 are +# unrealizable on halve hardware, so the level grid tops out at 128 and the +# budget is expressed against a 128 reference (BUDGET_RATIO = AVG/128, +# computed here — do NOT pass the old AVG/256 ratios). +# Usage: sbatch sbatch_calib_globalpr_flop_halve.sb [METRIC] +# METRIC: cosine (default) | mse_rel | l2 — non-cosine tables land in a +# metric-suffixed OUTDIR so A/B arms never overwrite each other. set -uo pipefail -AVG="${1:?AVG}"; BR="${2:?BUDGET_RATIO}" +AVG="${1:?AVG}" +METRIC="${2:-cosine}" +HALVE_REF=128 +if (( AVG > HALVE_REF )); then + echo "ERROR: AVG=${AVG} exceeds halve-mode max stoc_len ${HALVE_REF}" >&2 + exit 1 +fi +BR=$(awk "BEGIN{print ${AVG}/${HALVE_REF}}") REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion SCRATCH=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi CKPT=$SCRATCH/pretrained_models/DiT-XL-2-256x256.pt -OUTDIR=$SCRATCH/scmp_diffusion_fid_mp_cfg15/calib_global_perrow_flop_halve +OUTDIR=$SCRATCH/scmp_diffusion_fid_mp_cfg15/calib_global_perrow_flop_halve_max128 +[[ "$METRIC" != "cosine" ]] && OUTDIR="${OUTDIR}_${METRIC}" source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh; conda activate qdit export PYTHONUNBUFFERED=1 SC_OWEN_MODE=bitrev; cd "$REPO"; mkdir -p "$OUTDIR" -echo "=== FLOP+HALVE calib GLOBAL per-row avg${AVG} br=${BR} branch=$(git branch --show-current) $(date) ===" +echo "=== FLOP+HALVE calib GLOBAL per-row avg${AVG} br=${BR} (ref=${HALVE_REF}) metric=${METRIC} branch=$(git branch --show-current) $(date) ===" python -u scripts/calibrate_mp_thresholds.py \ - --mp_levels 256,192,128,96,64,48,32,16 \ - --budget_ratio "$BR" --budget_ref_stoc_len 256 --budget_scope global \ - --metric cosine --teacher fp \ + --mp_levels 128,96,64,48,32,16 \ + --budget_ratio "$BR" --budget_ref_stoc_len "$HALVE_REF" --budget_scope global \ + --metric "$METRIC" --teacher fp \ --sc_prec 8 --sc_fixed_level_prec --sc_qk_granularity per_row \ --wbits 8 --abits 8 --w_sym --a_sym \ --image-size 256 --num-sampling-steps 50 \ --num_calib_batches 1 --num_calib_timesteps 50 \ --timestep_buckets 10 --layer_buckets 4 \ - --sc_halve --teacher_cfg_scale 0.0 --ckpt "$CKPT" \ + --sc_halve --teacher_cfg_scale 1.5 --batch-size 16 --ckpt "$CKPT" \ --calib_output_json "$OUTDIR/calib_fix_avg${AVG}.json" \ --calib_summary_csv "$OUTDIR/calib_fix_avg${AVG}_summary.csv" 2>&1 | tee "$OUTDIR/calib_avg${AVG}.log" echo "DONE avg${AVG} $(date)" diff --git a/scripts/sbatch_calib_globalpr_halve.sb b/scripts/sbatch_calib_globalpr_halve.sb index d2f0403..139eca2 100755 --- a/scripts/sbatch_calib_globalpr_halve.sb +++ b/scripts/sbatch_calib_globalpr_halve.sb @@ -6,26 +6,36 @@ #SBATCH --cpus-per-gpu=8 #SBATCH --mem=110G #SBATCH --time=03:00:00 -# per-row GLOBAL calibration for one budget. Usage: sbatch sbatch_calib_globalpr.sb +# per-row GLOBAL calibration (row-weighted) for one budget, HALVE mode. +# Halve mode caps stream length at 2^(sc_prec-1)=128, so the level grid tops +# out at 128 and the budget is expressed against a 128 reference +# (BUDGET_RATIO = AVG/128, computed here — do NOT pass the old AVG/256 ratios). +# Usage: sbatch sbatch_calib_globalpr_halve.sb set -uo pipefail -AVG="${1:?AVG}"; BR="${2:?BUDGET_RATIO}" +AVG="${1:?AVG}" +HALVE_REF=128 +if (( AVG > HALVE_REF )); then + echo "ERROR: AVG=${AVG} exceeds halve-mode max stoc_len ${HALVE_REF}" >&2 + exit 1 +fi +BR=$(awk "BEGIN{print ${AVG}/${HALVE_REF}}") REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion SCRATCH=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi CKPT=$SCRATCH/pretrained_models/DiT-XL-2-256x256.pt -OUTDIR=$SCRATCH/scmp_diffusion_fid_mp_cfg15/calib_global_perrow_halve +OUTDIR=$SCRATCH/scmp_diffusion_fid_mp_cfg15/calib_global_perrow_halve_max128 source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh; conda activate qdit export PYTHONUNBUFFERED=1 SC_OWEN_MODE=bitrev; cd "$REPO"; mkdir -p "$OUTDIR" -echo "=== calib GLOBAL per-row HALVE avg${AVG} br=${BR} $(date) ===" +echo "=== calib GLOBAL per-row HALVE avg${AVG} br=${BR} (ref=${HALVE_REF}) $(date) ===" python -u scripts/calibrate_mp_thresholds.py \ - --mp_levels 256,192,128,96,64,48,32,16 \ - --budget_ratio "$BR" --budget_ref_stoc_len 256 --budget_scope global \ + --mp_levels 128,96,64,48,32,16 \ + --budget_ratio "$BR" --budget_ref_stoc_len "$HALVE_REF" --budget_scope global \ --metric cosine --teacher fp \ --sc_prec 8 --sc_fixed_level_prec --sc_qk_granularity per_row \ --wbits 8 --abits 8 --w_sym --a_sym \ --image-size 256 --num-sampling-steps 50 \ --num_calib_batches 1 --num_calib_timesteps 50 \ --timestep_buckets 10 --layer_buckets 4 \ - --sc_halve --teacher_cfg_scale 0.0 --ckpt "$CKPT" \ - --calib_output_json "$OUTDIR/calib_fix_avg${AVG}_l256_ref192.json" \ + --sc_halve --teacher_cfg_scale 1.5 --batch-size 16 --ckpt "$CKPT" \ + --calib_output_json "$OUTDIR/calib_fix_avg${AVG}.json" \ --calib_summary_csv "$OUTDIR/calib_fix_avg${AVG}_summary.csv" 2>&1 | tee "$OUTDIR/calib_avg${AVG}.log" echo "DONE avg${AVG} $(date)" diff --git a/scripts/sbatch_mp_globalpr_flop_halve_auto.sb b/scripts/sbatch_mp_globalpr_flop_halve_auto.sb index 7debbc4..d3f5b9a 100644 --- a/scripts/sbatch_mp_globalpr_flop_halve_auto.sb +++ b/scripts/sbatch_mp_globalpr_flop_halve_auto.sb @@ -9,18 +9,26 @@ # # FLOP-weighted + HALVE adaptive-MP generation, auto-resumes to NUM_FID, TIMEOUT-safe: # queues its successor (afterany, chain-capped) at the START. Uses the FLOP+halve -# calibrated tables. Usage: sbatch sbatch_mp_globalpr_flop_halve_auto.sb [CHAIN] +# calibrated tables with the level grid capped at 128 (halve-mode max stoc_len = +# 2^(sc_prec-1); the old 256-max tables probed unrealizable levels — kept in +# calib_global_perrow_flop_halve/ + ..._flop_halve_cfg15/ for comparison only). +# Usage: sbatch sbatch_mp_globalpr_flop_halve_auto.sb [CHAIN] [METRIC] [NUM_FID] +# METRIC: cosine (default) | mse_rel | l2 — must match the calib arm; picks the +# metric-suffixed calib table and output dir. +# NUM_FID: target image count (default 2000); e.g. 100 for a fast PSNR probe. set -uo pipefail -AVG="${1:?usage: sbatch sbatch_mp_globalpr_flop_halve_auto.sb [CHAIN]}"; CHAIN="${2:-0}"; NUM_GPUS=3; MAXCHAIN=15 +AVG="${1:?usage: sbatch sbatch_mp_globalpr_flop_halve_auto.sb [CHAIN] [METRIC] [NUM_FID]}"; CHAIN="${2:-0}"; METRIC="${3:-cosine}"; NUM_FID_ARG="${4:-2000}"; NUM_GPUS=3; MAXCHAIN=15 SELF=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion/scripts/sbatch_mp_globalpr_flop_halve_auto.sb REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion SCRATCH=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi CKPT=$SCRATCH/pretrained_models/DiT-XL-2-256x256.pt -CALIB=$SCRATCH/scmp_diffusion_fid_mp_cfg15/calib_global_perrow_flop_halve/calib_fix_avg${AVG}.json -OUT=$SCRATCH/scmp_diffusion_fid_mp_globalpr_flop_halve_cfg15/adaptive_avg${AVG} +CALIBDIR=$SCRATCH/scmp_diffusion_fid_mp_cfg15/calib_global_perrow_flop_halve_max128 +OUT=$SCRATCH/scmp_diffusion_fid_mp_globalpr_flop_halve128_cfg15/adaptive_avg${AVG} +if [[ "$METRIC" != "cosine" ]]; then CALIBDIR="${CALIBDIR}_${METRIC}"; OUT="${OUT}_${METRIC}"; fi +CALIB=$CALIBDIR/calib_fix_avg${AVG}.json SAMPLES=$OUT/samples; IDX=$OUT/_indices; LOG=$OUT/_logs -NUM_FID=2000; BALANCED=10000; BATCH=64; STEPS=50; CFG=1.5; SEED=0; NUM_CLASSES=1000 -MP_LEVELS=256,192,128,96,64,48,32,16 +NUM_FID=$NUM_FID_ARG; BALANCED=10000; BATCH=64; STEPS=50; CFG=1.5; SEED=0; NUM_CLASSES=1000 +MP_LEVELS=128,96,64,48,32,16 mkdir -p "$SAMPLES" "$IDX" "$LOG" [[ -f "$CALIB" ]] || { echo "ERROR: missing calib $CALIB" >&2; exit 1; } source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh; conda activate qdit @@ -32,7 +40,7 @@ if [[ "$START" -ge "$NUM_FID" ]]; then echo "[complete] avg${AVG} already at $NU # Queue successor NOW (afterany -> runs after me regardless of timeout/complete); chain-capped. if [[ "$CHAIN" -lt "$MAXCHAIN" ]]; then sbatch -o "$OUT/slurm-%j.out" -e "$OUT/slurm-%j.err" \ - --dependency=afterany:${SLURM_JOB_ID} "$SELF" "$AVG" $((CHAIN+1)) \ + --dependency=afterany:${SLURM_JOB_ID} "$SELF" "$AVG" $((CHAIN+1)) "$METRIC" "$NUM_FID_ARG" \ && echo "[chain] queued successor (chain=$((CHAIN+1)))" fi python -u scripts/_plan_missing_indices.py "$SAMPLES" "$NUM_FID" "$NUM_GPUS" "$IDX" $((BALANCED/NUM_CLASSES)) From c328d36521448b54d0640701ff206b865bee110d Mon Sep 17 00:00:00 2001 From: heroarmor <162866837+heroarmor@users.noreply.github.com> Date: Wed, 15 Jul 2026 04:58:39 -0400 Subject: [PATCH 6/7] tools(probe): verification sbatch scripts for the halve MP line MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Small single-purpose probes developed while root-causing the halve-line regressions; kept because each verifies an invariant worth re-checking: - eval/sbatch_eval_uniform_halve.sb: KID + FID/IS/Prec/Recall for one uniform-halve arm (companion to sbatch_eval_mp_flop_halve.sb). - eval/sbatch_eval_ab_kid100.sb: matched-index KID across probe arms + an FP anchor on the same indices (the anchor exposes the class-composition floor that dominates absolute KID at small n — FP scores the same ~30x1e-3 as SC arms on 10 classes). - eval/sbatch_probe_realized_dist{,_old}.sb: validate-path run that dumps debug_mp_distribution.csv — compares the RUNTIME-realized per-op level fractions against the calibration table's intent (A4 check). - eval/sbatch_kernel_level_err.sb: kernel-level rel-err vs stoc_len under halve on qk/av/mlp shapes (catches level-grid pathologies without running the diffusion model). - sbatch_uniform48_rerun100.sb: regenerate uniform48 idx 0-99 into a fresh dir — reproducibility check. NOTE: latent noise pairs by position in each worker's index list, not by global index, so PSNR comparisons against runs with a different GPU split are only valid on the aligned worker-0 prefix. - sbatch_only1op_sc32.sb: single-op SC isolation (everything FP except one op at stoc_len 32 + halve) via per-op config JSONs; the legacy layerwise flags are ignored whenever --sc_config supplies a precision map. Co-Authored-By: Claude Fable 5 --- scripts/eval/sbatch_eval_ab_kid100.sb | 47 ++++++++++++++ scripts/eval/sbatch_eval_uniform_halve.sb | 43 +++++++++++++ scripts/eval/sbatch_kernel_level_err.sb | 61 ++++++++++++++++++ scripts/eval/sbatch_probe_realized_dist.sb | 62 +++++++++++++++++++ .../eval/sbatch_probe_realized_dist_old.sb | 51 +++++++++++++++ scripts/sbatch_only1op_sc32.sb | 58 +++++++++++++++++ scripts/sbatch_uniform48_rerun100.sb | 45 ++++++++++++++ 7 files changed, 367 insertions(+) create mode 100644 scripts/eval/sbatch_eval_ab_kid100.sb create mode 100644 scripts/eval/sbatch_eval_uniform_halve.sb create mode 100644 scripts/eval/sbatch_kernel_level_err.sb create mode 100644 scripts/eval/sbatch_probe_realized_dist.sb create mode 100644 scripts/eval/sbatch_probe_realized_dist_old.sb create mode 100644 scripts/sbatch_only1op_sc32.sb create mode 100644 scripts/sbatch_uniform48_rerun100.sb diff --git a/scripts/eval/sbatch_eval_ab_kid100.sb b/scripts/eval/sbatch_eval_ab_kid100.sb new file mode 100644 index 0000000..12e87a2 --- /dev/null +++ b/scripts/eval/sbatch_eval_ab_kid100.sb @@ -0,0 +1,47 @@ +#!/bin/bash +#SBATCH --job-name=kid_ab100 +#SBATCH --partition=gpu-rtx6000 +#SBATCH --account=nbleier_owned1 +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-gpu=4 +#SBATCH --mem=48G +#SBATCH --time=00:30:00 +# One-off: KID on idx 0-99 for the metric A/B probe arms + uniform48 anchor. +# All three use the same indices (classes 0-9, 10 imgs each) vs the full +# ImageNet-256 ref -> absolute KID inflated by class mismatch, but the +# three-way RANKING is apples-to-apples. +set -uo pipefail +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +PREV=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion_prev +REF=$PREV/imagenet256_ref/VIRTUAL_imagenet256_labeled.npz +S=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi +A=$S/scmp_diffusion_fid_mp_globalpr_flop_halve128_cfg15/adaptive_avg48/samples +B=$S/scmp_diffusion_fid_mp_globalpr_flop_halve128_cfg15/adaptive_avg48_mse_rel/samples +U=$S/scmp_diffusion_fid_uniform_halveON_cfg15/uniform48/samples +F=$S/scmp_diffusion_fid_fp_6shard_cfg15/samples +WORK=$S/scmp_diffusion_fid_mp_globalpr_flop_halve128_cfg15/_ab_eval; mkdir -p "$WORK" +OUT=$WORK/kid_ab100.txt +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh +exec > >(tee "$WORK/eval_ab100.log") 2>&1 +echo "############ A/B KID idx0-99 (GPU) $(date) ############" + +conda activate qdit +for arm in A B U F; do + src_var=$arm; src=${!src_var} + SEL=$WORK/sel_$arm; rm -rf "$SEL"; mkdir -p "$SEL" + n=0 + for i in $(seq 0 99); do + f=$(printf '%06d.png' "$i") + [[ -f "$src/$f" ]] && ln -sf "$src/$f" "$SEL/$f" && n=$((n+1)) + done + echo "[$arm] linked $n/100 from $src" + python -u "$REPO/scripts/eval/pngs_to_npz.py" "$SEL" "$WORK/arm_$arm.npz" +done +conda deactivate + +conda activate tfeval +export TF_CPP_MIN_LOG_LEVEL=2 EVALUATOR=$PREV/Q-DiT/models/evaluations/evaluator.py +python -u "$REPO/scripts/eval/kid_openai.py" "$REF" "$OUT" \ + "$WORK/arm_A.npz" "$WORK/arm_B.npz" "$WORK/arm_U.npz" "$WORK/arm_F.npz" +conda deactivate +echo "############ DONE $(date) ############" diff --git a/scripts/eval/sbatch_eval_uniform_halve.sb b/scripts/eval/sbatch_eval_uniform_halve.sb new file mode 100644 index 0000000..26e9e4e --- /dev/null +++ b/scripts/eval/sbatch_eval_uniform_halve.sb @@ -0,0 +1,43 @@ +#!/bin/bash +#SBATCH --job-name=eval_uni +#SBATCH --partition=gpu-rtx6000 +#SBATCH --account=nbleier_owned1 +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-gpu=4 +#SBATCH --mem=48G +#SBATCH --time=00:45:00 +# GPU eval for one uniform (halve) config: KID + FID/IS/sFID/Prec/Recall. +# Same methodology as sbatch_eval_mp_flop_halve.sb, repointed to the uniform dir. +# Usage: sbatch sbatch_eval_uniform_halve.sb +set -uo pipefail +AVG="${1:?usage: sbatch sbatch_eval_uniform_halve.sb }" +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +PREV=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion_prev +PACKER=$PREV/imagenet256_ref/parallel_npz.py +REF=$PREV/imagenet256_ref/VIRTUAL_imagenet256_labeled.npz +BASE=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion_fid_uniform_halveON_cfg15/uniform${AVG} +AD=$BASE/samples +WORK=$BASE/_eval; mkdir -p "$WORK" +OUT=$WORK/eval_uniform${AVG}.txt +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh +exec > >(tee "$OUT") 2>&1 +echo "############ UNIFORM(halve) EVAL uniform${AVG} (GPU) $(date) ############" + +IDX=$WORK/idx.txt; NPZ=$WORK/uniform${AVG}.npz +(cd "$AD" && ls [0-9][0-9][0-9][0-9][0-9][0-9].png | awk '($1+0)<2000' | sort) > "$IDX" +SEL=$WORK/sel; rm -rf "$SEL"; mkdir -p "$SEL" +while read -r f; do [[ -n "$f" ]] && ln -sf "$AD/$f" "$SEL/$f"; done < "$IDX" +echo "=== matched idx 0-1999: $(wc -l < "$IDX") ===" + +conda activate qdit +python -u "$PACKER" "$SEL" "$NPZ" +conda deactivate + +conda activate tfeval +export TF_CPP_MIN_LOG_LEVEL=2 EVALUATOR=$PREV/Q-DiT/models/evaluations/evaluator.py +echo ""; echo "===== [1] KID vs ref (idx 0-1999) =====" +bash "$REPO/scripts/eval/kid_openai.sh" "$WORK/kid_uniform${AVG}.txt" "$NPZ" +echo ""; echo "===== [2] FID / IS / sFID / Precision / Recall vs ref =====" +python -u "$EVALUATOR" "$REF" "$NPZ" 2>&1 | grep -E "^(Inception Score|FID|sFID|Precision|Recall):" | sed "s/^/[uniform${AVG}] /" +conda deactivate +echo ""; echo "############ DONE uniform${AVG} $(date) ############" diff --git a/scripts/eval/sbatch_kernel_level_err.sb b/scripts/eval/sbatch_kernel_level_err.sb new file mode 100644 index 0000000..70b9ca0 --- /dev/null +++ b/scripts/eval/sbatch_kernel_level_err.sb @@ -0,0 +1,61 @@ +#!/bin/bash +#SBATCH --job-name=kern_lvl_err +#SBATCH --partition=gpu-rtx6000 +#SBATCH --account=nbleier_owned1 +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-gpu=4 +#SBATCH --mem=24G +#SBATCH --time=00:15:00 +# Kernel-level SC error vs stoc_len under halve: is L=128 anomalously bad? +# Random matrices, per-row bipolar, qk-like and mlp-like shapes, mixed-level +# calls in one process (cache interactions included, like MP dispatch does). +set -uo pipefail +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +S=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi +WORK=$S/scmp_diffusion_fid_mp_globalpr_flop_halve128_cfg15/_dist_probe +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh; conda activate qdit +export SC_OWEN_MODE=bitrev PYTHONUNBUFFERED=1; cd "$REPO" +exec > >(tee "$WORK/kernel_level_err.log") 2>&1 +python - << 'EOF' +import torch +from scmp_kernels.sc.matmul import sc_matmul + +torch.manual_seed(0) +dev = "cuda" + +def rel_err(out, exact): + return ((out - exact).norm() / exact.norm()).item() + +# (name, N, M, D): a is (N, D), b is (M, D), product a @ b.T +SHAPES = [ + ("qk-like (256x256, D=72)", 256, 256, 72), + ("av-like (256x72, D=256)", 256, 72, 256), + ("mlp-like (512x4608, D=1152, chunk144)", 512, 4608, 1152), +] +LEVELS = [32, 48, 64, 96, 128] + +for name, N, M, D in SHAPES: + a = torch.randn(N, D, device=dev) + b = torch.randn(M, D, device=dev) + exact = a @ b.T + chunk = 144 if D > 1024 else 0 + print(f"=== {name} ===") + # halve ON, ascending then repeated (cache warm) order + for tag, levels in (("halve asc ", LEVELS), ("halve rpt ", LEVELS)): + errs = [] + for sl in levels: + out = sc_matmul(a, b, granularity="per_row", mode="bipolar", + sc_prec=8, stoc_len=sl, group_a=1, group_b=1, + chunk_d=chunk, halve_bipolar_stoc_len=True) + errs.append(f"{sl}:{rel_err(out, exact):.4f}") + print(f" {tag}: " + " ".join(errs)) + # halve OFF reference at same lengths + 256 + errs = [] + for sl in LEVELS + [256]: + out = sc_matmul(a, b, granularity="per_row", mode="bipolar", + sc_prec=8, stoc_len=sl, group_a=1, group_b=1, + chunk_d=chunk, halve_bipolar_stoc_len=False) + errs.append(f"{sl}:{rel_err(out, exact):.4f}") + print(f" no-halve : " + " ".join(errs)) +EOF +echo "DONE $(date)" diff --git a/scripts/eval/sbatch_probe_realized_dist.sb b/scripts/eval/sbatch_probe_realized_dist.sb new file mode 100644 index 0000000..cc28e41 --- /dev/null +++ b/scripts/eval/sbatch_probe_realized_dist.sb @@ -0,0 +1,62 @@ +#!/bin/bash +#SBATCH --job-name=mp_dist_probe +#SBATCH --partition=gpu-rtx6000 +#SBATCH --account=nbleier_owned1 +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-gpu=4 +#SBATCH --mem=48G +#SBATCH --time=01:00:00 +# A4 instrumented probe: run the VALIDATE path (8 imgs, dumps +# debug_mp_distribution.csv) once with the OLD 256-grid flop-halve table and +# once with the NEW 128-cap table, then print realized per-op level fractions +# side by side. 10 sampling steps -> one visit per t-bucket, enough for the +# realized-distribution comparison at ~7 min/run. +set -uo pipefail +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +S=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi +CKPT=$S/pretrained_models/DiT-XL-2-256x256.pt +OLD_CALIB=$S/scmp_diffusion_fid_mp_cfg15/calib_global_perrow_flop_halve/calib_fix_avg48.json +NEW_CALIB=$S/scmp_diffusion_fid_mp_cfg15/calib_global_perrow_flop_halve_max128/calib_fix_avg48.json +WORK=$S/scmp_diffusion_fid_mp_globalpr_flop_halve128_cfg15/_dist_probe +mkdir -p "$WORK" +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh; conda activate qdit +export SC_OWEN_MODE=bitrev PYTHONUNBUFFERED=1 OMP_NUM_THREADS=2; cd "$REPO" +exec > >(tee "$WORK/probe.log") 2>&1 + +run_probe () { + local tag="$1" calib="$2" levels="$3" + echo "===== probe $tag levels=$levels =====" + rm -rf "$WORK/$tag"; mkdir -p "$WORK/$tag" + python -u scripts/quant_sc_main.py \ + --wbits 8 --abits 8 --w_sym --a_sym \ + --timewise 1 --qklayerwise 1.0 --avlayerwise 1.0 --projlayerwise 1.0 --mlplayerwise 1.0 --inputprojlayerwise 1.0 \ + --sc_prec 8 --sc_fixed_level_prec --sc_qk_granularity per_row --sc_halve \ + --adaptive_mp --adaptive_mp_table "$calib" --mp_levels "$levels" \ + --image-size 256 --num-sampling-steps 10 --cfg-scale 1.5 \ + --seed 0 --results-dir "$WORK/$tag" --ckpt "$CKPT" + find "$WORK/$tag" -name 'debug_mp_distribution.csv' -exec cp {} "$WORK/dist_$tag.csv" \; + find "$WORK/$tag" -name 'mp_savings_summary.txt' -exec cp {} "$WORK/savings_$tag.txt" \; +} + +run_probe old "$OLD_CALIB" "256,192,128,96,64,48,32,16" +run_probe new "$NEW_CALIB" "128,96,64,48,32,16" + +python - << 'EOF' +import csv, collections +for tag in ("old", "new"): + per_op = collections.defaultdict(lambda: collections.Counter()) + tot = collections.Counter() + with open(f"/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion_fid_mp_globalpr_flop_halve128_cfg15/_dist_probe/dist_{tag}.csv") as f: + for row in csv.DictReader(f): + op = row["operator"] + for k, v in row.items(): + if k.startswith("sl_") and k.endswith("_count") and v: + sl = int(k.split("_")[1]); per_op[op][sl] += int(v); tot[op] += int(v) + print(f"===== REALIZED distribution ({tag} table) =====") + for op in sorted(per_op): + c = per_op[op]; n = tot[op] + avg = sum(sl * cnt for sl, cnt in c.items()) / max(n, 1) + fr = " ".join(f"{sl}:{100*cnt/n:.1f}%" for sl, cnt in sorted(c.items(), reverse=True) if cnt) + print(f" {op:12s} avg={avg:7.1f} {fr}") +EOF +echo "DONE $(date)" diff --git a/scripts/eval/sbatch_probe_realized_dist_old.sb b/scripts/eval/sbatch_probe_realized_dist_old.sb new file mode 100644 index 0000000..55ae58d --- /dev/null +++ b/scripts/eval/sbatch_probe_realized_dist_old.sb @@ -0,0 +1,51 @@ +#!/bin/bash +#SBATCH --job-name=mp_dist_old +#SBATCH --partition=gpu-rtx6000 +#SBATCH --account=nbleier_owned1 +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-gpu=4 +#SBATCH --mem=48G +#SBATCH --time=00:40:00 +# Control probe: realized distribution under the OLD 256-grid cfg0-calib table. +# Runs WITHOUT --sc_halve because the fixed kernels (correctly) refuse +# stoc_len>128 under halve — halve only changes SC noise, not the +# threshold classification, so realized level fractions remain valid. +set -uo pipefail +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +S=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi +CKPT=$S/pretrained_models/DiT-XL-2-256x256.pt +OLD_CALIB=$S/scmp_diffusion_fid_mp_cfg15/calib_global_perrow_flop_halve/calib_fix_avg48.json +WORK=$S/scmp_diffusion_fid_mp_globalpr_flop_halve128_cfg15/_dist_probe +mkdir -p "$WORK" +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh; conda activate qdit +export SC_OWEN_MODE=bitrev PYTHONUNBUFFERED=1 OMP_NUM_THREADS=2; cd "$REPO" +exec > >(tee "$WORK/probe_old.log") 2>&1 + +echo "===== probe old(halve-off) levels=256,192,128,96,64,48,32,16 =====" +rm -rf "$WORK/old"; mkdir -p "$WORK/old" +python -u scripts/quant_sc_main.py \ + --wbits 8 --abits 8 --w_sym --a_sym \ + --timewise 1 --qklayerwise 1.0 --avlayerwise 1.0 --projlayerwise 1.0 --mlplayerwise 1.0 --inputprojlayerwise 1.0 \ + --sc_prec 8 --sc_fixed_level_prec --sc_qk_granularity per_row \ + --adaptive_mp --adaptive_mp_table "$OLD_CALIB" --mp_levels "256,192,128,96,64,48,32,16" \ + --image-size 256 --num-sampling-steps 10 --cfg-scale 1.5 \ + --seed 0 --results-dir "$WORK/old" --ckpt "$CKPT" +find "$WORK/old" -name 'debug_mp_distribution.csv' -exec cp {} "$WORK/dist_old.csv" \; + +python - << 'EOF' +import csv, collections +W = "/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion_fid_mp_globalpr_flop_halve128_cfg15/_dist_probe" +per_op = collections.defaultdict(collections.Counter); tot = collections.Counter() +for row in csv.DictReader(open(f"{W}/dist_old.csv")): + op = row["operator"] + for k, v in row.items(): + if k.startswith("sl_") and k.endswith("_count") and v: + per_op[op][int(k.split("_")[1])] += int(v); tot[op] += int(v) +print("===== REALIZED (old table, runtime cfg1.5, halve-off) =====") +for op in sorted(per_op): + c = per_op[op]; n = tot[op] + avg = sum(sl * cnt for sl, cnt in c.items()) / max(n, 1) + fr = " ".join(f"{sl}:{100*cnt/n:.1f}%" for sl, cnt in sorted(c.items(), reverse=True)) + print(f" {op:12s} avg={avg:7.1f} {fr}") +EOF +echo "DONE $(date)" diff --git a/scripts/sbatch_only1op_sc32.sb b/scripts/sbatch_only1op_sc32.sb new file mode 100644 index 0000000..fe2cdfa --- /dev/null +++ b/scripts/sbatch_only1op_sc32.sb @@ -0,0 +1,58 @@ +#!/bin/bash +#SBATCH --job-name=only1op32 +#SBATCH --partition=gpu-rtx6000 +#SBATCH --account=nbleier_owned1 +#SBATCH --gres=gpu:3 +#SBATCH --cpus-per-gpu=2 +#SBATCH --mem-per-gpu=12G +#SBATCH --time=06:00:00 +# Single-op SC isolation probe (era-2 code): EVERYTHING runs FP except the ONE +# selected op-group, which runs SC at uniform stoc_len=32 + halve. Generates +# idx 0-99. The bifurcation rate vs FP directly measures that op's solo +# trajectory-branching power at a stressed length. +# v2: isolation via per-op config JSON (enabled=false for others) — the +# layerwise flags are IGNORED when --sc_config provides a precision map +# (sc_controller.py:115), which silently turned v1 into uniform32 x5. +# Usage: sbatch sbatch_only1op_sc32.sb OP in: qk av proj inproj mlp +set -uo pipefail +OP="${1:?usage: sbatch sbatch_only1op_sc32.sb }" +QK=0; AV=0; PROJ=0; INPROJ=0; MLP=0 +case "$OP" in + qk) QK=1.0 ;; + av) AV=1.0 ;; + proj) PROJ=1.0 ;; + inproj) INPROJ=1.0 ;; + mlp) MLP=1.0 ;; + *) echo "unknown OP $OP" >&2; exit 1 ;; +esac +NUM_GPUS=3 +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +SCRATCH=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi +CKPT=$SCRATCH/pretrained_models/DiT-XL-2-256x256.pt +SC_JSON=$SCRATCH/scmp_diffusion_fid_halve_bitrev/configs/sc_cfg_only32_${OP}.json +OUT=$SCRATCH/scmp_diffusion_fid_uniform_halveON_cfg15/only1op32b_${OP} +SAMPLES=$OUT/samples; IDX=$OUT/_indices; LOG=$OUT/_logs +NUM_FID=100; BALANCED=10000; BATCH=64; STEPS=50; CFG=1.5; SEED=0; NUM_CLASSES=1000 +mkdir -p "$SAMPLES" "$IDX" "$LOG" +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh; conda activate qdit +export SC_OWEN_MODE=bitrev PYTHONUNBUFFERED=1 OMP_NUM_THREADS=2; cd "$REPO" +echo "=== only1op SC32 op=$OP (qk=$QK av=$AV proj=$PROJ inproj=$INPROJ mlp=$MLP) job=${SLURM_JOB_ID:-?} $(date) ===" +python -u scripts/_plan_missing_indices.py "$SAMPLES" "$NUM_FID" "$NUM_GPUS" "$IDX" $((BALANCED/NUM_CLASSES)) +pids=() +for ((g=0; g "$LOG/gpu_${g}.log" 2>&1 & + pids+=($!) +done +for p in "${pids[@]}"; do wait "$p" || true; done +echo "=== only1op $OP done: $(find "$SAMPLES" -name '*.png' | wc -l)/$NUM_FID $(date) ===" diff --git a/scripts/sbatch_uniform48_rerun100.sb b/scripts/sbatch_uniform48_rerun100.sb new file mode 100644 index 0000000..0f512be --- /dev/null +++ b/scripts/sbatch_uniform48_rerun100.sb @@ -0,0 +1,45 @@ +#!/bin/bash +#SBATCH --job-name=uni48_rerun +#SBATCH --partition=gpu-rtx6000 +#SBATCH --account=nbleier_owned1 +#SBATCH --gres=gpu:3 +#SBATCH --cpus-per-gpu=2 +#SBATCH --mem-per-gpu=12G +#SBATCH --time=06:00:00 +# One-off reproducibility probe: regenerate uniform48 (halve) idx 0-99 under +# the CURRENT code state into a fresh dir, to verify the uniform48 anchor +# still reproduces (~24 dB vs FP, ~5% bifurcation) after the kernel guard / +# recent changes. Identical flags/config/seed to sbatch_uniform_halve_auto.sb. +set -uo pipefail +NUM_GPUS=3 +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +SCRATCH=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi +CKPT=$SCRATCH/pretrained_models/DiT-XL-2-256x256.pt +SC_JSON=$SCRATCH/scmp_diffusion_fid_halve_bitrev/configs/sc_cfg_uniform48_all.json +OUT=$SCRATCH/scmp_diffusion_fid_uniform_halveON_cfg15/uniform48_rerun +SAMPLES=$OUT/samples; IDX=$OUT/_indices; LOG=$OUT/_logs +NUM_FID=100; BALANCED=10000; BATCH=64; STEPS=50; CFG=1.5; SEED=0; NUM_CLASSES=1000 +mkdir -p "$SAMPLES" "$IDX" "$LOG" +[[ -f "$SC_JSON" ]] || { echo "ERROR: missing $SC_JSON" >&2; exit 1; } +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh; conda activate qdit +export SC_OWEN_MODE=bitrev PYTHONUNBUFFERED=1 OMP_NUM_THREADS=2; cd "$REPO" +echo "=== uniform48 RERUN probe: target=$NUM_FID job=${SLURM_JOB_ID:-?} $(date) ===" +python -u scripts/_plan_missing_indices.py "$SAMPLES" "$NUM_FID" "$NUM_GPUS" "$IDX" $((BALANCED/NUM_CLASSES)) +pids=() +for ((g=0; g "$LOG/gpu_${g}.log" 2>&1 & + pids+=($!) +done +for p in "${pids[@]}"; do wait "$p" || true; done +echo "=== uniform48 RERUN done: $(find "$SAMPLES" -name '*.png' | wc -l)/$NUM_FID $(date) ===" From dfb061fcbca975fe8113d6a8c2efcb6b03da02d4 Mon Sep 17 00:00:00 2001 From: heroarmor <162866837+heroarmor@users.noreply.github.com> Date: Tue, 18 Aug 2026 15:29:43 -0400 Subject: [PATCH 7/7] feat(mp): per-group (K-band) stream lengths on the contraction axis MP was per-row on the token axis only: every chunk of D ran the same stoc_len. This lets each band of quantization chunks run its own length at each rung. RUNTIME (sc_kbands.py, sc_mlp.py, sc_attention.py) The kernel still takes a SCALAR stoc_len -- _sc_matmul_per_row_mlp builds cum_indicator once outside the chunk loop and every chunk shares it -- so the banding happens above it: rows are grouped by RUNG index (not by parent stoc_len value, since the rung is what indexes every band's ladder) and each (band, rung) is a separate call whose partial products are summed. The two paths band differently because they chunk differently: - attention already walks D one chunk at a time in Python (it must: chunk_d requires granularity="per_row" and this path can run per_tensor), so a band is just a label on those chunks and only the stoc_len lookup changes. That is why it comes out BIT-IDENTICAL to the parent under the parent ladder. - mlp hands chunking to the kernel, so bands gather their columns and each band re-chunks at the same boundaries. Measured 6.8e-08 against the parent -- fp32 summation order, nothing else. GUARDS (all three would otherwise surface only as a wrong number) - A table naming qk/av is rejected at init_adaptive_mp: they contract over head_dim, have no band dispatch, and would run the parent while the run carries the k-bands label. - A band length above the halve ceiling is rejected at startup rather than mid-run. Bands RAISE the rungs they favour, so a band ladder can exceed 2^(sc_prec-1) even when the parent ladder does not. - k-bands + --range_mp raises: the combined path has no contraction-axis dispatch and would drop the allocation on the floor. - A table covering a module at a different chunk_d or contraction width raises instead of falling back to the parent. TOOLING build_kband_table.py adds a k_bands section to an existing calibrated table under a straight-line tilt, rounding DOWN so the identity is never overspent and handing single units back largest-lost-fraction-first. --tilt 0 must reproduce the parent and is the sanity arm for any sweep. This is a sweepable starting point, NOT the solved allocation -- the LLM side solves it by Lagrangian water-fill on measured per-band error curves. Note for the halve line: at the top rung the parent already sits at the 128 ceiling, so no band can be lengthened and that rung cannot be tilted at all. Bands only have freedom on rungs below the ceiling. TESTS 24 tests: table validation (overspend, underspend, whole-chunk rule, band-id contiguity, per-operator band counts, width agreement), column partitioning incl. the tail-chunk ordering, the controller guards, and on GPU both the parent-equivalence and non-degeneracy gates. verify_kband_e2e.py drives the real _sc_linear_dynamic_mp methods rather than reimplementing the loop. Verified on gl1802 (job 58184021): 24 passed, e2e PASS. Co-Authored-By: Claude Opus 5 (1M context) --- qdit/sc_integration/sc_attention.py | 109 ++++++--- qdit/sc_integration/sc_controller.py | 44 ++++ qdit/sc_integration/sc_kbands.py | 145 ++++++++++++ qdit/sc_integration/sc_mlp.py | 110 ++++++--- scmp_kernels | 2 +- scripts/build_kband_table.py | 230 ++++++++++++++++++ scripts/eval/sbatch_kband_equivalence.sb | 43 ++++ scripts/verify_kband_e2e.py | 178 ++++++++++++++ tests/test_kbands.py | 290 +++++++++++++++++++++++ 9 files changed, 1094 insertions(+), 57 deletions(-) create mode 100644 qdit/sc_integration/sc_kbands.py create mode 100755 scripts/build_kband_table.py create mode 100755 scripts/eval/sbatch_kband_equivalence.sb create mode 100755 scripts/verify_kband_e2e.py create mode 100644 tests/test_kbands.py diff --git a/qdit/sc_integration/sc_attention.py b/qdit/sc_integration/sc_attention.py index 65f31b5..218ee36 100644 --- a/qdit/sc_integration/sc_attention.py +++ b/qdit/sc_integration/sc_attention.py @@ -25,6 +25,7 @@ from ..qLinearLayer import QLinearLayer from .sc_controller import SCController from .mp_config import classify_rows_by_metric, adaptive_classify_rows, MPDistributionLogger, MetricProfiler +from .sc_kbands import resolve_k_bands, band_of_chunk_index, _reject_k_bands_on_combined_mp from scmp_kernels.sc import sc_matmul from scmp_kernels.sc.config_helpers import make_sobol_simple_config @@ -348,48 +349,96 @@ def _sc_linear_dynamic_mp(self, x, weight, bias, operator, chunk_d=0, grouped=Fa compute_baseline = 0 compute_actual = 0.0 - for sl, rows in assignment.level_row_indices.items(): - if len(rows) == 0 or sl == 0: - continue - n_rows = len(rows) - compute_baseline += n_rows * out_features * D * baseline_stoc_len - compute_actual += n_rows * out_features * D * sl + k_bands = resolve_k_bands( + self.sc_controller, operator, self.block_idx, D, chunk_d) + + if k_bands is not None: + # ---- per-group (K-band) stream lengths ----------------------- + # The app-level chunk loop already walks the contraction axis one + # quantization chunk at a time; a band is just a label on those + # chunks, so the only change is which stoc_len each chunk runs at. + # Rows are grouped by RUNG index (not by parent stoc_len value), + # because the rung is what indexes every band's ladder. + band_of_chunk, band_ladders = k_bands + n_rungs = len(band_ladders[0]) + parent_levels = self.sc_controller.adaptive_mp_config.stoc_len_levels + granularity = "per_row" if grouped else "per_tensor" - sp = self.sc_controller.resolve_sc_prec(sl) - x_sub = x_flat[rows].contiguous() + for k in range(n_rungs): + rows = (assignment.row_levels == k).nonzero(as_tuple=True)[0] + if rows.numel() == 0 or parent_levels[k] == 0: + continue + n_rows = rows.numel() + compute_baseline += n_rows * out_features * D * baseline_stoc_len - granularity = "per_row" if grouped else "per_tensor" - if chunk_d > 0 and D > chunk_d: + x_sub = x_flat[rows].contiguous() sub_result = None - for start in range(0, D, chunk_d): + for chunk_idx, start in enumerate(range(0, D, chunk_d)): end = min(start + chunk_d, D) - x_chunk = x_sub[:, start:end].contiguous() - w_chunk = weight[:, start:end].contiguous() - config = self._get_sc_config(end - start, sp) + band = band_of_chunk_index(band_of_chunk, chunk_idx) + sl = int(band_ladders[band][k]) + if sl <= 0: + continue # this band contributes nothing at this rung + compute_actual += n_rows * out_features * (end - start) * sl + sp = self.sc_controller.resolve_sc_prec(sl) + config = self._get_sc_config(end - start, sp) kwargs = dict(granularity=granularity, mode=self.sc_mode, sc_prec=sp, config=config, stoc_len=sl, rng_levels=self._rng_levels(sl)) if grouped: kwargs.update(group_a=1, group_b=1) - chunk_out = matmul_fn(x_chunk, w_chunk, **kwargs) + chunk_out = matmul_fn(x_sub[:, start:end].contiguous(), + weight[:, start:end].contiguous(), + **kwargs) + sub_result = (chunk_out if sub_result is None + else sub_result + chunk_out) + if sub_result is not None: + result[rows] = sub_result + else: + for sl, rows in assignment.level_row_indices.items(): + if len(rows) == 0 or sl == 0: + continue + n_rows = len(rows) + compute_baseline += n_rows * out_features * D * baseline_stoc_len + compute_actual += n_rows * out_features * D * sl - if sub_result is None: - sub_result = chunk_out - else: - sub_result = sub_result + chunk_out - result[rows] = sub_result - else: - config = self._get_sc_config(D, sp) - kwargs = dict(granularity=granularity, - mode=self.sc_mode, sc_prec=sp, config=config, - stoc_len=sl, rng_levels=self._rng_levels(sl)) - if grouped: - kwargs.update(group_a=1, group_b=1) + sp = self.sc_controller.resolve_sc_prec(sl) + x_sub = x_flat[rows].contiguous() + + granularity = "per_row" if grouped else "per_tensor" + if chunk_d > 0 and D > chunk_d: + sub_result = None + for start in range(0, D, chunk_d): + end = min(start + chunk_d, D) + x_chunk = x_sub[:, start:end].contiguous() + w_chunk = weight[:, start:end].contiguous() + config = self._get_sc_config(end - start, sp) + + kwargs = dict(granularity=granularity, + mode=self.sc_mode, sc_prec=sp, config=config, + stoc_len=sl, rng_levels=self._rng_levels(sl)) + if grouped: + kwargs.update(group_a=1, group_b=1) + + chunk_out = matmul_fn(x_chunk, w_chunk, **kwargs) + + if sub_result is None: + sub_result = chunk_out + else: + sub_result = sub_result + chunk_out + result[rows] = sub_result + else: + config = self._get_sc_config(D, sp) + kwargs = dict(granularity=granularity, + mode=self.sc_mode, sc_prec=sp, config=config, + stoc_len=sl, rng_levels=self._rng_levels(sl)) + if grouped: + kwargs.update(group_a=1, group_b=1) - sub = matmul_fn(x_sub, weight, **kwargs) - result[rows] = sub + sub = matmul_fn(x_sub, weight, **kwargs) + result[rows] = sub MPDistributionLogger.log_compute( self.sc_controller.current_timestep, self.block_idx, @@ -415,6 +464,8 @@ def _sc_linear_combined_mp(self, x, weight, bias, operator, dispatch, """ orig_shape = x.shape D = x.shape[-1] + _reject_k_bands_on_combined_mp( + self.sc_controller, operator, self.block_idx, D, chunk_d) x_flat = x.reshape(-1, D) M = x_flat.shape[0] diff --git a/qdit/sc_integration/sc_controller.py b/qdit/sc_integration/sc_controller.py index de8d52e..55b6138 100644 --- a/qdit/sc_integration/sc_controller.py +++ b/qdit/sc_integration/sc_controller.py @@ -273,8 +273,52 @@ def init_adaptive_mp(self, adaptive_mp_config: AdaptiveMPConfig): When set, this takes priority over mp_config for dynamic MP paths. """ + if getattr(adaptive_mp_config, "k_band_count", 0): + self._validate_k_bands(adaptive_mp_config) self.adaptive_mp_config = adaptive_mp_config + # Operators whose contraction axis is an input-channel dim wide enough to + # partition into quantization chunks. qk/av contract over head_dim, which + # is smaller than one chunk, so they have no band dispatch at all -- a + # table naming them would be silently ignored, which is why this raises. + K_BAND_OPERATORS = ("mlp_fc1", "mlp_fc2", "proj", "input_proj") + + def _validate_k_bands(self, adaptive_mp_config) -> None: + """Reject band allocations this controller cannot actually execute. + + Both failures below would otherwise surface only as a wrong number: + an unsupported operator runs the per-row parent while the run is + labelled k-bands, and an over-long band length is a stream the halve + hardware cannot realize. + """ + named = {op for op, _ in adaptive_mp_config.k_band_chunks} + unsupported = sorted(named - set(self.K_BAND_OPERATORS)) + if unsupported: + raise ValueError( + f"K-band table names operators with no band dispatch: " + f"{unsupported}. Only {list(self.K_BAND_OPERATORS)} contract " + f"over an input-channel axis; qk/av contract over head_dim and " + f"would run the per-row parent while the run is labelled " + f"k-bands.") + + # Bands raise the length of the rungs they favour, so a band ladder can + # exceed the halve-mode ceiling even when the parent ladder does not. + # sc_matmul enforces this per call for bipolar; checking it here turns + # a mid-run crash on some later timestep into a startup error. + if not self.halve: + return + max_len = 2 ** (self.sc_prec - 1) + for key, band_ladders in adaptive_mp_config.k_band_ladders.items(): + for b, rungs in enumerate(band_ladders): + for k, sl in enumerate(rungs): + if sl > max_len: + raise ValueError( + f"K-band ladder {key} band {b} rung {k} is " + f"stoc_len={sl}, above the halve-mode maximum " + f"{max_len} for sc_prec={self.sc_prec}. Such a " + f"stream is unrealizable on halve hardware and " + f"would silently inflate simulated accuracy.") + def init_range_mp(self, range_mp_config: RangeMPConfig): """Initialize range-based mixed precision (weight min/max range). diff --git a/qdit/sc_integration/sc_kbands.py b/qdit/sc_integration/sc_kbands.py new file mode 100644 index 0000000..6f8fcc4 --- /dev/null +++ b/qdit/sc_integration/sc_kbands.py @@ -0,0 +1,145 @@ +"""K-bands: per-group (contraction-axis) stream lengths. + +Row dispatch is unchanged -- one metric, one rung index per row -- but the +contraction axis is partitioned into bands of whole quantization chunks and +band ``b`` runs rung ``k`` at its own stream length ``L[b][k]``. + +Why this cannot lose: setting every band's ladder equal to the parent ladder +(``L[b][k] == stoc_len_levels[k]``) reproduces the per-row parent, so the +parent sits inside the search space. The budget is an exact per-rung identity +rather than a tolerance, because MACs are linear in the contraction dim and a +band's column fraction therefore IS its MAC fraction:: + + sum_b (w_b / R) * L[b][k] == L_parent[k] + +Why whole chunks: the kernel builds its RNG tables over ``chunk_d`` dims and +reuses them for every chunk, so a chunk relocated to another band keeps the +exact scale and Owen mask it would have had in the unbanded call. Relocating +arbitrary channels does not. + +The allocation itself (which chunk goes in which band, and each band's ladder) +is solved offline and shipped in the MP table's ``k_bands`` section; see +``AdaptiveMPConfig._load_k_bands`` for the schema and its validation. +""" + +import torch + + +def resolve_k_bands(sc_controller, operator, block_idx, D, chunk_d): + """Band map + ladders for one module, or ``None`` to run the per-row parent. + + Returns ``(band_of_chunk, band_ladders)``. ``None`` is returned whenever + the adaptive config carries no allocation for this module, which is the + normal path for operators the solver left alone. + + A table that DOES cover this module but was priced against a different + contraction width or chunk size raises instead of falling back: the band + map would no longer describe this matmul, and silently running the parent + while the run is labelled "k-bands" is the one failure mode that looks + exactly like a result. + """ + mp_config = getattr(sc_controller, "adaptive_mp_config", None) + if mp_config is None or not getattr(mp_config, "k_band_count", 0): + return None + + k_bands = mp_config.get_k_bands( + operator, block_idx, sc_controller.total_blocks, + timestep=sc_controller.current_timestep, + total_timesteps=sc_controller.total_timesteps, + ) + if k_bands is None: + return None + + table_chunk_d = int(mp_config.k_band_chunk_d) + if chunk_d != table_chunk_d: + raise ValueError( + f"K-band table for {operator} block {block_idx} was solved at " + f"chunk_d={table_chunk_d} but this call runs chunk_d={chunk_d}. " + f"Bands own whole quantization chunks, so the band map does not " + f"describe this matmul.") + table_width = int(mp_config.k_band_residual_width[operator]) + if D != table_width: + raise ValueError( + f"K-band table for {operator} block {block_idx} was priced against " + f"a contraction width of {table_width} but this call has D={D}. " + f"Band widths set the iso-compute budget, so the allocation is " + f"not transferable.") + if D <= chunk_d: + raise ValueError( + f"K-band table covers {operator} block {block_idx} but D={D} is " + f"within a single chunk of {chunk_d}; there is nothing to band.") + return k_bands + + +def band_columns(module, band_of_chunk, n_bands, chunk_d, D, device): + """Column indices owned by each band, plus each band's channel width. + + Chunks stay in ASCENDING order inside a band, which puts the short tail + chunk last in whichever band owns it. That is the only ordering under + which a band's gathered columns re-chunk to the same boundaries the + unbanded call would have used, and therefore the only one under which the + per-chunk scales and RNG tables are unchanged. + + Cached on the module: the band map is fixed for a module's lifetime, while + this runs per forward. + """ + cache = getattr(module, "_sc_k_band_cache", None) + key = (tuple(band_of_chunk), n_bands, chunk_d, D, str(device)) + if cache is not None and cache[0] == key: + return cache[1], cache[2] + + cols_per_band, width_per_band = [], [] + for b in range(n_bands): + cols = [] + for chunk_idx, band in enumerate(band_of_chunk): + if band != b: + continue + start = chunk_idx * chunk_d + cols.extend(range(start, min(start + chunk_d, D))) + width_per_band.append(len(cols)) + cols_per_band.append( + torch.tensor(cols, dtype=torch.long, device=device)) + + if sum(width_per_band) != D: + raise ValueError( + f"K-band columns cover {sum(width_per_band)} of {D} channels; the " + f"band map does not partition the contraction axis.") + + module._sc_k_band_cache = (key, cols_per_band, width_per_band) + return cols_per_band, width_per_band + + +def _reject_k_bands_on_combined_mp(sc_controller, operator, block_idx, D, + chunk_d): + """Refuse to run the combined range+dynamic path on a banded operator. + + That path iterates range-based weight groups and takes + ``min(range_stoc_len, dynamic_stoc_len)`` per (group, row). It has no + contraction-axis dispatch, so a band allocation for this operator would be + dropped on the floor while the run still carries the k-bands label. + """ + mp_config = getattr(sc_controller, "adaptive_mp_config", None) + if mp_config is None or not getattr(mp_config, "k_band_count", 0): + return + if resolve_k_bands(sc_controller, operator, block_idx, D, chunk_d) is None: + return + raise NotImplementedError( + f"K-bands and range-based per-group weight dispatch are both active " + f"for '{operator}' (block {block_idx}). The combined path has no " + f"contraction-axis dispatch and would silently discard the band " + f"allocation. Drop --range_mp, or drop '{operator}' from the table's " + f"k_bands section.") + + +def band_of_chunk_index(band_of_chunk, chunk_idx): + """Band owning chunk ``chunk_idx``, bounds-checked. + + Used by the app-level chunk loops (attention), which walk chunks directly + instead of gathering per-band columns. + """ + if not 0 <= chunk_idx < len(band_of_chunk): + raise ValueError( + f"chunk index {chunk_idx} outside the band map of length " + f"{len(band_of_chunk)}; the table's chunk count and the runtime " + f"chunking disagree.") + return band_of_chunk[chunk_idx] diff --git a/qdit/sc_integration/sc_mlp.py b/qdit/sc_integration/sc_mlp.py index 215936a..67194ee 100644 --- a/qdit/sc_integration/sc_mlp.py +++ b/qdit/sc_integration/sc_mlp.py @@ -25,6 +25,7 @@ from ..qLinearLayer import QLinearLayer from .sc_controller import SCController from .mp_config import classify_rows_by_metric, adaptive_classify_rows, MPDistributionLogger, MetricProfiler +from .sc_kbands import resolve_k_bands, band_columns, _reject_k_bands_on_combined_mp from scmp_kernels.sc import sc_matmul from scmp_kernels.sc.config_helpers import make_sobol_simple_config @@ -234,35 +235,88 @@ def _sc_linear_dynamic_mp(self, x, weight, bias, operator, chunk_d=0): compute_baseline = 0 compute_actual = 0.0 - for sl, rows in assignment.level_row_indices.items(): - if len(rows) == 0 or sl == 0: - continue # pruned rows: result already zeroed - n_rows = len(rows) - compute_baseline += n_rows * out_features * D * baseline_stoc_len - compute_actual += n_rows * out_features * D * sl + k_bands = resolve_k_bands( + self.sc_controller, operator, self.block_idx, D, chunk_d) + + if k_bands is not None: + # ---- per-group (K-band) stream lengths ----------------------- + # Same rung index per row as the per-row parent; band b runs that + # rung at its own length. Bands partition the contraction axis + # into whole quantization chunks, so each chunk keeps the exact + # scale and RNG table it would have had in the unbanded call, and + # the bands' partial products sum to the same matmul. + band_of_chunk, band_ladders = k_bands + n_bands = len(band_ladders) + n_rungs = len(band_ladders[0]) + parent_levels = self.sc_controller.adaptive_mp_config.stoc_len_levels + cols_per_band, width_per_band = band_columns( + self, band_of_chunk, n_bands, chunk_d, D, x_flat.device) + + rows_by_rung = [ + (assignment.row_levels == k).nonzero(as_tuple=True)[0] + for k in range(n_rungs) + ] + for k, rows in enumerate(rows_by_rung): + if rows.numel() == 0 or parent_levels[k] == 0: + continue + compute_baseline += ( + rows.numel() * out_features * D * baseline_stoc_len) - sp = self.sc_controller.resolve_sc_prec(sl) - x_sub = x_flat[rows].contiguous() # [len(rows), D] + for b in range(n_bands): + cols = cols_per_band[b] + if cols.numel() == 0: + continue + x_band = x_flat.index_select(1, cols).contiguous() + w_band = weight.index_select(1, cols).contiguous() + band_width = width_per_band[b] + for k, rows in enumerate(rows_by_rung): + sl = int(band_ladders[b][k]) + if sl <= 0 or rows.numel() == 0: + continue # this band contributes nothing for these rows + compute_actual += ( + rows.numel() * out_features * band_width * sl) + sp = self.sc_controller.resolve_sc_prec(sl) + # Same config as the unbanded call: a band is >= 2 chunks + # wide, so the kernel re-chunks it at the same boundaries. + config = self._get_sc_config(chunk_d, sp) + # bands accumulate into the shared row buffer + result[rows] += matmul_fn( + x_band.index_select(0, rows).contiguous(), w_band, + granularity="per_row", + mode=self.sc_mode, sc_prec=sp, config=config, + group_a=1, group_b=1, chunk_d=chunk_d, + stoc_len=sl, + rng_levels=self._rng_levels(sl)) + else: + for sl, rows in assignment.level_row_indices.items(): + if len(rows) == 0 or sl == 0: + continue # pruned rows: result already zeroed + n_rows = len(rows) + compute_baseline += n_rows * out_features * D * baseline_stoc_len + compute_actual += n_rows * out_features * D * sl - if chunk_d > 0 and D > chunk_d: - config = self._get_sc_config(chunk_d, sp) - sub = matmul_fn( - x_sub, weight, - granularity="per_row", - mode=self.sc_mode, sc_prec=sp, config=config, - group_a=1, group_b=1, chunk_d=chunk_d, - stoc_len=sl, - rng_levels=self._rng_levels(sl)) - else: - config = self._get_sc_config(D, sp) - sub = matmul_fn( - x_sub, weight, - granularity="per_row", - mode=self.sc_mode, sc_prec=sp, config=config, - group_a=1, group_b=1, - stoc_len=sl, - rng_levels=self._rng_levels(sl)) - result[rows] = sub + sp = self.sc_controller.resolve_sc_prec(sl) + x_sub = x_flat[rows].contiguous() # [len(rows), D] + + if chunk_d > 0 and D > chunk_d: + config = self._get_sc_config(chunk_d, sp) + sub = matmul_fn( + x_sub, weight, + granularity="per_row", + mode=self.sc_mode, sc_prec=sp, config=config, + group_a=1, group_b=1, chunk_d=chunk_d, + stoc_len=sl, + rng_levels=self._rng_levels(sl)) + else: + config = self._get_sc_config(D, sp) + sub = matmul_fn( + x_sub, weight, + granularity="per_row", + mode=self.sc_mode, sc_prec=sp, config=config, + group_a=1, group_b=1, + stoc_len=sl, + rng_levels=self._rng_levels(sl)) + result[rows] = sub MPDistributionLogger.log_compute( self.sc_controller.current_timestep, self.block_idx, @@ -284,6 +338,8 @@ def _sc_linear_combined_mp(self, x, weight, bias, operator, dispatch, """ orig_shape = x.shape D = x.shape[-1] + _reject_k_bands_on_combined_mp( + self.sc_controller, operator, self.block_idx, D, chunk_d) x_flat = x.reshape(-1, D) M = x_flat.shape[0] diff --git a/scmp_kernels b/scmp_kernels index c967d3a..acf8eca 160000 --- a/scmp_kernels +++ b/scmp_kernels @@ -1 +1 @@ -Subproject commit c967d3aa9014f37a2f6823048dd02993e2e1dd21 +Subproject commit acf8ecaeef0cbfa3763577f3fc5b7073e8f586f7 diff --git a/scripts/build_kband_table.py b/scripts/build_kband_table.py new file mode 100755 index 0000000..68e00a7 --- /dev/null +++ b/scripts/build_kband_table.py @@ -0,0 +1,230 @@ +#!/usr/bin/env python +"""Add a ``k_bands`` section to an existing adaptive MP threshold table. + +Row dispatch is untouched: the same calibrated thresholds put each row on the +same rung. What this adds is a partition of the contraction axis into bands of +whole quantization chunks, and a per-band ladder that runs each rung at its own +stream length under an exact per-rung iso-compute identity:: + + sum_b (w_b / R) * L[b][k] == L_parent[k] + +Because the parent sits inside that space (``L[b][k] == L_parent[k]`` for every +band), the refinement cannot lose -- a tilt that does not help degenerates back +to the parent rather than costing accuracy. + +The tilt here is a straight-line shape across bands, which is a sweepable +starting point, not a solved allocation. ``--tilt 0`` reproduces the parent +exactly and is the sanity arm every sweep should include. + +Usage:: + + # sanity arm: must reproduce the per-row parent + python scripts/build_kband_table.py --in mp_best.json --out mp_kb0.json \\ + --tilt 0 --n-bands 4 --max-stoc-len 128 + + # front-loaded: early chunks longer, late chunks shorter + python scripts/build_kband_table.py --in mp_best.json --out mp_kb.json \\ + --tilt 0.25 --n-bands 4 --max-stoc-len 128 +""" +import argparse +import json +import sys + +# DiT-XL/2 contraction widths per operator. Override with --width op=N. +DEFAULT_WIDTHS = { + "mlp_fc1": 1152, + "mlp_fc2": 4608, + "proj": 1152, + "input_proj": 1152, +} + + +def chunk_widths(residual_width, chunk_d): + n = (residual_width + chunk_d - 1) // chunk_d + return [min(chunk_d, residual_width - c * chunk_d) for c in range(n)] + + +def contiguous_bands(n_chunks, n_bands): + """Assign chunks to bands in contiguous, as-equal-as-possible runs. + + Contiguous because a tilt is a statement about position along the + contraction axis; scattered bands would make the shape meaningless. + Every band needs >= 2 chunks to stay on the chunked kernel path, so the + caller must cap n_bands at n_chunks // 2. + """ + if n_bands * 2 > n_chunks: + raise ValueError( + f"{n_chunks} chunks cannot be split into {n_bands} bands of >= 2 " + f"chunks; the widest allocation here is {n_chunks // 2} bands.") + base, extra = divmod(n_chunks, n_bands) + bands, chunk = [], 0 + for b in range(n_bands): + size = base + (1 if b < extra else 0) + bands.extend([b] * size) + chunk += size + return bands + + +def solve_ladder(parent_levels, widths, tilt, max_stoc_len, snap): + """Per-band lengths that spend at most the parent's per-rung budget. + + Rounds DOWN off the tilt shape so the identity is never overspent, then + hands single units back to the bands with the largest lost fraction while + they still fit. The leftover is bounded by max(w_b) / (R * L_k), which is + far inside the loader's underspend tolerance -- an exact integer solution + does not exist for arbitrary widths, and overspending to reach one would + price the cell as cheaper than it runs. + """ + n_bands = len(widths) + total_width = sum(widths) + # Straight line from (1 + tilt) at band 0 down to (1 - tilt) at the last. + if n_bands == 1: + shape = [1.0] + else: + shape = [1.0 + tilt - 2.0 * tilt * b / (n_bands - 1) + for b in range(n_bands)] + # Renormalize so the shape itself is width-weighted iso-compute; otherwise + # unequal band widths bias the mean and every rung starts overspent. + mean = sum(widths[b] * shape[b] for b in range(n_bands)) / total_width + shape = [s / mean for s in shape] + + ladder = [] + for parent in parent_levels: + if parent == 0: + ladder.append([0] * n_bands) + continue + budget = total_width * parent # sum_b w_b * L[b] must be <= + raw = [parent * s for s in shape] + lengths = [] + for b in range(n_bands): + v = int(raw[b] // snap) * snap + v = max(snap, min(v, max_stoc_len)) + lengths.append(v) + spent = sum(widths[b] * lengths[b] for b in range(n_bands)) + if spent > budget: + # Only reachable via the max_stoc_len clamp lifting a band. + order = sorted(range(n_bands), key=lambda b: -lengths[b]) + for b in order: + while spent > budget and lengths[b] > snap: + lengths[b] -= snap + spent -= widths[b] * snap + # Hand back what is left, largest lost fraction first. + order = sorted(range(n_bands), key=lambda b: -(raw[b] - lengths[b])) + progressed = True + while progressed: + progressed = False + for b in order: + if lengths[b] + snap > max_stoc_len: + continue + if spent + widths[b] * snap <= budget: + lengths[b] += snap + spent += widths[b] * snap + progressed = True + ladder.append(lengths) + + # ladder is [n_rungs][n_bands]; the table stores [n_bands][n_rungs]. + return [[ladder[k][b] for k in range(len(parent_levels))] + for b in range(n_bands)] + + +def main(): + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--in", dest="src", required=True, + help="existing adaptive MP table (JSON)") + p.add_argument("--out", dest="dst", required=True) + p.add_argument("--n-bands", type=int, default=4, + help="maximum bands; clamped per operator by its chunk count") + p.add_argument("--chunk-d", type=int, default=128) + p.add_argument("--tilt", type=float, default=0.0, + help="0 reproduces the per-row parent exactly; 0.25 runs " + "the first band ~25%% longer and the last ~25%% shorter") + p.add_argument("--max-stoc-len", type=int, default=0, + help="halve-mode ceiling, e.g. 128 for sc_prec=8 (0 = none)") + p.add_argument("--snap", type=int, default=1, + help="round stream lengths to a multiple of this") + p.add_argument("--operators", default="mlp_fc1,mlp_fc2,proj,input_proj", + help="comma-separated operators to band") + p.add_argument("--blocks", type=int, default=28, + help="number of transformer blocks (28 for DiT-XL/2); the " + "band map is emitted per (operator, block)") + p.add_argument("--width", action="append", default=[], metavar="OP=N", + help="override an operator's contraction width") + args = p.parse_args() + + with open(args.src) as f: + payload = json.load(f) + + levels = [int(x) for x in payload["stoc_len_levels"]] + max_stoc_len = args.max_stoc_len or max(levels) + widths_by_op = dict(DEFAULT_WIDTHS) + for spec in args.width: + op, _, n = spec.partition("=") + widths_by_op[op] = int(n) + + operators = [op for op in args.operators.split(",") if op] + # Only band operators the table actually has thresholds for; a band map for + # an operator the run never dispatches is dead weight that still validates. + seen_ops = {key.split(":")[0] for key in payload.get("buckets", {})} + seen_ops |= set(payload.get("operator_defaults", {})) + + n_blocks = args.blocks + residual_width, chunk_bands, ladders = {}, {}, {} + for op in operators: + if op not in seen_ops: + print(f"[build_kband_table] skipping '{op}': no thresholds in the " + f"source table", file=sys.stderr) + continue + if op not in widths_by_op: + raise SystemExit( + f"no contraction width known for '{op}'; pass --width {op}=N") + R = widths_by_op[op] + cw = chunk_widths(R, args.chunk_d) + n_bands = min(args.n_bands, len(cw) // 2) + if n_bands < 2: + print(f"[build_kband_table] skipping '{op}': {len(cw)} chunks at " + f"chunk_d={args.chunk_d} leaves room for {len(cw)//2} bands", + file=sys.stderr) + continue + if n_bands < args.n_bands: + print(f"[build_kband_table] '{op}': capped to {n_bands} bands " + f"({len(cw)} chunks)", file=sys.stderr) + bands = contiguous_bands(len(cw), n_bands) + bw = [0] * n_bands + for c, b in enumerate(bands): + bw[b] += cw[c] + + residual_width[op] = R + for blk in range(n_blocks): + chunk_bands[f"{op}:{blk}"] = bands + band_ladder = solve_ladder(levels, bw, args.tilt, max_stoc_len, + args.snap) + for t in range(int(payload.get("timestep_buckets", 1))): + for l in range(int(payload.get("layer_buckets", 1))): + ladders[f"{op}:t{t}:l{l}"] = band_ladder + + spend = [sum(bw[b] * band_ladder[b][k] for b in range(n_bands)) / + (sum(bw) * levels[k]) if levels[k] else 1.0 + for k in range(len(levels))] + print(f"[build_kband_table] {op}: {n_bands} bands, widths {bw}, " + f"budget used per rung {[f'{s:.4f}' for s in spend]}") + + if not residual_width: + raise SystemExit("no operator could be banded; nothing written") + + # n_bands in the section is the max across operators; the loader checks + # band ids against it and the per-operator ladder length carries the rest. + payload["k_bands"] = { + "n_bands": max(len(v) for v in ladders.values()), + "chunk_d": args.chunk_d, + "residual_width": residual_width, + "chunk_bands": chunk_bands, + "ladders": ladders, + } + with open(args.dst, "w") as f: + json.dump(payload, f, indent=2) + print(f"[build_kband_table] wrote {args.dst}") + + +if __name__ == "__main__": + main() diff --git a/scripts/eval/sbatch_kband_equivalence.sb b/scripts/eval/sbatch_kband_equivalence.sb new file mode 100755 index 0000000..c547b77 --- /dev/null +++ b/scripts/eval/sbatch_kband_equivalence.sb @@ -0,0 +1,43 @@ +#!/bin/bash +#SBATCH --job-name=kband_eq +#SBATCH --partition=gpu-rtx6000 +#SBATCH --account=nbleier_owned1 +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-gpu=4 +#SBATCH --mem-per-gpu=32G +#SBATCH --time=00:30:00 +#SBATCH --output=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion/slurm/kband_eq_%j.out +# +# K-band correctness gate. Two things must hold or the whole allocation story +# is void: +# 1. With every band's ladder equal to the parent ladder, banded dispatch +# reproduces the per-row parent (fp32 summation order only). +# 2. A non-uniform ladder actually changes the result -- otherwise the +# dispatch is a no-op dressed as a refinement. +set -uo pipefail +# Slurm stages the script into /var/spool, so $0 is not the repo copy. +# Override to verify a worktree / clone; the default is the main checkout. +REPO="${KBAND_REPO:-/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion}" +cd "$REPO" || exit 1 +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh; conda activate qdit + +echo "=== node $(hostname) ===" +nvidia-smi --query-gpu=name,memory.total --format=csv,noheader + +# scmp_kernels is an EDITABLE install, so it does not follow $REPO and cannot be +# redirected with PYTHONPATH. Print what actually got imported: a gate whose log +# cannot prove which trees it ran against proves nothing. +echo "=== trees under test ===" +python -c " +import inspect, scmp_kernels.mp.config as c, qdit.sc_integration.sc_kbands as k +print(' scmp_kernels:', inspect.getfile(c)) +print(' qdit :', inspect.getfile(k)) +" + +echo +echo "=== k-band unit + equivalence tests ===" +python -m pytest tests/test_kbands.py -v 2>&1 | tail -40 + +echo +echo "=== end-to-end: banded SCMlp vs per-row parent on one block ===" +python scripts/verify_kband_e2e.py diff --git a/scripts/verify_kband_e2e.py b/scripts/verify_kband_e2e.py new file mode 100755 index 0000000..0424b58 --- /dev/null +++ b/scripts/verify_kband_e2e.py @@ -0,0 +1,178 @@ +#!/usr/bin/env python +"""Drive the real SCMlp / SCAttention band dispatch and check it against the parent. + +The unit tests reimplement the band loop against raw ``sc_matmul``; this runs +the actual ``_sc_linear_dynamic_mp`` methods the model calls, so a mistake in +the dispatch itself (row indexing, accumulation, config selection) shows up +here and not only in production. + +Two gates: + +1. **Parent equivalence.** With ``L[b][k] == stoc_len_levels[k]`` for every + band, the banded path must reproduce the unbanded per-row path. Not bit + identity -- splitting the contraction axis changes only the fp32 summation + order, since each chunk keeps its own scale and RNG table. + +2. **Non-degeneracy.** A tilted ladder must actually move the output. A band + dispatch that silently collapses back to the parent would pass gate 1 while + measuring nothing. +""" +import json +import os +import sys +import tempfile + +import torch + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from scmp_kernels.mp import AdaptiveMPConfig +from scmp_kernels.sc.config_helpers import make_sobol_simple_config +from qdit.sc_integration.sc_mlp import SCMlp +from qdit.sc_integration.sc_attention import SCAttention + +LEVELS = [128, 96, 64, 48, 32, 16] +CHUNK_D = 128 +WIDTH = 1152 # DiT-XL/2 hidden -> 9 chunks +CHUNK_BANDS = [0, 0, 0, 1, 1, 2, 2, 3, 3] # widths 384/256/256/256 +TOTAL_BLOCKS = 28 + + +class StubController: + """Only the surface ``_sc_linear_dynamic_mp`` actually reads.""" + + def __init__(self, adaptive_mp_config): + self.adaptive_mp_config = adaptive_mp_config + self.mp_config = None + self.current_timestep = 10 + self.total_timesteps = 250 + self.total_blocks = TOTAL_BLOCKS + self.sc_prec = 8 + self.stoc_len = 128 + self.noise_model = False + self.halve = True + self.fixed_level_sc_prec = True + + def resolve_sc_prec(self, stoc_len): + return self.sc_prec + + +class StubModule: + """Duck-types the parts of SCMlp / SCAttention the dispatch touches.""" + + def __init__(self, controller, block_idx=3): + self.sc_controller = controller + self.block_idx = block_idx + self.sc_mode = "bipolar" + self._sc_configs = {} + + _get_sc_config = SCMlp._get_sc_config + _get_matmul_fn = SCMlp._get_matmul_fn + _rng_levels = SCMlp._rng_levels + + +def make_config(band_ladders): + """An AdaptiveMPConfig whose thresholds spread rows over every rung.""" + n = len(LEVELS) + thresholds = [round(1.0 - (i + 1) / n, 4) for i in range(n - 1)] + payload = { + "stoc_len_levels": LEVELS, + "timestep_buckets": 1, + "layer_buckets": 1, + "operator_defaults": { + op: {"thresholds": thresholds} + for op in ("mlp_fc1", "proj") + }, + "k_bands": { + "n_bands": 4, + "chunk_d": CHUNK_D, + "residual_width": {"mlp_fc1": WIDTH, "proj": WIDTH}, + "chunk_bands": {f"{op}:{b}": CHUNK_BANDS + for op in ("mlp_fc1", "proj") + for b in range(TOTAL_BLOCKS)}, + "ladders": {f"{op}:t0:l0": band_ladders + for op in ("mlp_fc1", "proj")}, + }, + } + fd, path = tempfile.mkstemp(suffix=".json") + with os.fdopen(fd, "w") as f: + json.dump(payload, f) + try: + return AdaptiveMPConfig(stoc_len_levels=list(LEVELS), + threshold_table_path=path) + finally: + os.unlink(path) + + +def run(method, cfg, x, w, operator, banded, **kwargs): + saved = cfg.k_band_count + if not banded: + cfg.k_band_count = 0 + try: + mod = StubModule(StubController(cfg)) + return method(mod, x, w, None, operator, chunk_d=CHUNK_D, **kwargs) + finally: + cfg.k_band_count = saved + + +def rel(a, b): + return ((a - b).norm() / b.norm()).item() + + +def check(name, banded, parent, *, expect_close): + r = rel(banded, parent) + if expect_close: + ok = r < 1e-5 + verdict = "OK" if ok else "FAIL" + print(f"[{verdict}] {name}: rel err vs parent {r:.3e} (want < 1e-5)") + else: + ok = r > 1e-4 + verdict = "OK" if ok else "FAIL" + print(f"[{verdict}] {name}: rel err vs parent {r:.3e} (want > 1e-4)") + return ok + + +def main(): + if not torch.cuda.is_available(): + print("no CUDA; SC kernels cannot run") + return 1 + + torch.manual_seed(0) + M, N = 128, 1152 + x = torch.randn(M, WIDTH, device="cuda") * 0.05 + w_mlp = torch.randn(4608, WIDTH, device="cuda") * 0.02 + w_proj = torch.randn(N, WIDTH, device="cuda") * 0.02 + + parent_ladders = [list(LEVELS) for _ in range(4)] + # 2a + b == 3 * parent per rung for widths 384/256/256/256 is + # (1/3)a0 + (2/9)(a1+a2+a3) == parent; keep it simple and provably legal by + # trading the widest band down against the narrow ones. + tilted = [ + [128, 72, 48, 36, 24, 12], + [128, 108, 72, 54, 36, 18], + [128, 108, 72, 54, 36, 18], + [128, 108, 72, 54, 36, 18], + ] + + cfg_parent = make_config(parent_ladders) + cfg_tilt = make_config(tilted) + + ok = True + for label, method, weight, operator, kwargs in [ + ("SCMlp.mlp_fc1", SCMlp._sc_linear_dynamic_mp, w_mlp, "mlp_fc1", {}), + ("SCAttention.proj", SCAttention._sc_linear_dynamic_mp, w_proj, "proj", + {"grouped": True}), + ]: + parent = run(method, cfg_parent, x, weight, operator, False, **kwargs) + banded = run(method, cfg_parent, x, weight, operator, True, **kwargs) + ok &= check(f"{label} parent-ladder", banded, parent, expect_close=True) + + tilt = run(method, cfg_tilt, x, weight, operator, True, **kwargs) + ok &= check(f"{label} tilted-ladder", tilt, parent, expect_close=False) + + print("\nRESULT:", "PASS" if ok else "FAIL") + return 0 if ok else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_kbands.py b/tests/test_kbands.py new file mode 100644 index 0000000..3dbc8e3 --- /dev/null +++ b/tests/test_kbands.py @@ -0,0 +1,290 @@ +"""K-bands: table validation, column partitioning, and numeric equivalence. + +The equivalence test is the one that matters: with every band's ladder equal to +the parent ladder, the banded dispatch must reproduce the per-row parent. That +is what makes the refinement unable to lose, so if it ever stops holding the +whole allocation story is void. +""" +import json +import os +import sys +import tempfile +import unittest + +import torch + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from scmp_kernels.mp import AdaptiveMPConfig # noqa: E402 +from qdit.sc_integration.sc_kbands import band_columns # noqa: E402 + + +LEVELS = [128, 96, 64, 32] +# DiT-XL/2: hidden 1152 -> 9 chunks of 128 for proj / mlp_fc1. +FC1_WIDTH = 1152 +CHUNK_D = 128 +# 4 bands over 9 chunks: widths 256/256/256/384 channels. +CHUNK_BANDS = [0, 0, 1, 1, 2, 2, 3, 3, 3] +# Solved so 2a + b == 3 * parent for every rung, which is the iso-compute +# identity for these widths: (2/9)*3a + (1/3)b == parent. +WIDE = [144, 112, 80, 40] +NARROW = [96, 64, 32, 16] + + +def make_table(*, ladders=None, chunk_bands=None, n_bands=4, + residual_width=None, chunk_d=CHUNK_D, levels=None): + levels = LEVELS if levels is None else levels + payload = { + "stoc_len_levels": levels, + "timestep_buckets": 1, + "layer_buckets": 1, + "buckets": { + "mlp_fc1:t0:l0": {"thresholds": [0.75, 0.5, 0.25][: len(levels) - 1]}, + }, + "k_bands": { + "n_bands": n_bands, + "chunk_d": chunk_d, + "residual_width": ({"mlp_fc1": FC1_WIDTH} + if residual_width is None else residual_width), + "chunk_bands": ({"mlp_fc1:0": CHUNK_BANDS} + if chunk_bands is None else chunk_bands), + "ladders": ({"mlp_fc1:t0:l0": [WIDE, WIDE, WIDE, NARROW]} + if ladders is None else ladders), + }, + } + fd, path = tempfile.mkstemp(suffix=".json") + with os.fdopen(fd, "w") as f: + json.dump(payload, f) + return path + + +def load(**kwargs): + levels = kwargs.pop("levels", None) or LEVELS + path = make_table(levels=levels, **kwargs) + try: + return AdaptiveMPConfig(stoc_len_levels=list(levels), + threshold_table_path=path) + finally: + os.unlink(path) + + +class TestKBandTableValidation(unittest.TestCase): + + def test_loads_and_exposes_bands(self): + cfg = load() + self.assertEqual(cfg.k_band_count, 4) + self.assertEqual(cfg.k_band_chunk_d, CHUNK_D) + bands, ladders = cfg.get_k_bands("mlp_fc1", 0, 28) + self.assertEqual(bands, CHUNK_BANDS) + self.assertEqual(ladders, [WIDE, WIDE, WIDE, NARROW]) + + def test_absent_section_leaves_path_disabled(self): + cfg = AdaptiveMPConfig(stoc_len_levels=list(LEVELS)) + self.assertEqual(cfg.k_band_count, 0) + self.assertIsNone(cfg.get_k_bands("mlp_fc1", 0, 28)) + + def test_uncovered_operator_returns_none(self): + cfg = load() + self.assertIsNone(cfg.get_k_bands("mlp_fc2", 0, 28)) + + def test_parent_ladder_in_every_band_is_iso_compute(self): + """L[b][k] == parent[k] must pass the identity exactly.""" + parent_everywhere = [list(LEVELS) for _ in range(4)] + cfg = load(ladders={"mlp_fc1:t0:l0": parent_everywhere}) + _, ladders = cfg.get_k_bands("mlp_fc1", 0, 28) + self.assertEqual(ladders, parent_everywhere) + + def test_overspend_is_rejected(self): + # +16 on one wide band with no compensation elsewhere. + over = [[160, 112, 80, 40], WIDE, WIDE, NARROW] + with self.assertRaisesRegex(ValueError, "overspends rung 0"): + load(ladders={"mlp_fc1:t0:l0": over}) + + def test_float_noise_is_not_overspend(self): + """The identity is exact in reals; only true overspend may raise.""" + cfg = load() # WIDE/NARROW hit the identity dead on + self.assertIsNotNone(cfg.get_k_bands("mlp_fc1", 0, 28)) + + def test_large_underspend_is_rejected(self): + half = [[x // 2 for x in WIDE] for _ in range(3)] + [ + [x // 2 for x in NARROW]] + with self.assertRaisesRegex(ValueError, "underspends rung 0"): + load(ladders={"mlp_fc1:t0:l0": half}) + + def test_band_with_one_chunk_is_rejected(self): + lonely = [0, 0, 1, 1, 1, 2, 2, 2, 3] # band 3 owns a single chunk + with self.assertRaisesRegex(ValueError, "band 3 with 1 chunk"): + load(chunk_bands={"mlp_fc1:0": lonely}) + + def test_wrong_chunk_count_is_rejected(self): + with self.assertRaisesRegex(ValueError, "chunk entries"): + load(chunk_bands={"mlp_fc1:0": CHUNK_BANDS[:-1]}) + + def test_ladder_rung_count_must_match_levels(self): + short = [WIDE[:-1], WIDE[:-1], WIDE[:-1], NARROW[:-1]] + with self.assertRaisesRegex(ValueError, "rungs, expected 4"): + load(ladders={"mlp_fc1:t0:l0": short}) + + def test_chunk_bands_without_ladders_is_rejected(self): + """Silently falling back to the parent is the failure to prevent.""" + with self.assertRaisesRegex(ValueError, "no.*ladders entry"): + load(ladders={}) + + def test_missing_residual_width_is_rejected(self): + with self.assertRaisesRegex(ValueError, "residual_width is required"): + load(residual_width={}) + + def test_band_widths_must_agree_across_blocks(self): + other = [0, 0, 0, 1, 1, 2, 2, 3, 3] # same op, different widths + with self.assertRaisesRegex(ValueError, "band widths"): + load(chunk_bands={"mlp_fc1:0": CHUNK_BANDS, "mlp_fc1:1": other}) + + def test_n_bands_below_two_is_rejected(self): + with self.assertRaisesRegex(ValueError, "n_bands must be >= 2"): + load(n_bands=1) + + +class TestBandColumns(unittest.TestCase): + + class _Module: + pass + + def test_partitions_the_contraction_axis(self): + m = self._Module() + cols, widths = band_columns(m, CHUNK_BANDS, 4, CHUNK_D, FC1_WIDTH, + torch.device("cpu")) + self.assertEqual(widths, [256, 256, 256, 384]) + self.assertEqual(sum(widths), FC1_WIDTH) + seen = torch.cat(cols).sort().values + self.assertTrue(torch.equal(seen, torch.arange(FC1_WIDTH))) + + def test_chunks_stay_ascending_inside_a_band(self): + """The tail chunk must land last, or a band re-chunks differently.""" + # 1160 channels -> 10 chunks: nine 128-wide plus an 8-wide tail. + bands = [0, 0, 1, 1, 2, 2, 3, 3, 3, 3] + m = self._Module() + cols, widths = band_columns(m, bands, 4, CHUNK_D, 1160, + torch.device("cpu")) + self.assertEqual(widths[3], 128 + 128 + 128 + 8) + tail = cols[3] + self.assertTrue(bool((tail[1:] > tail[:-1]).all())) + self.assertEqual(int(tail[-1]), 1159) + + def test_cache_hits_on_repeat(self): + m = self._Module() + a, _ = band_columns(m, CHUNK_BANDS, 4, CHUNK_D, FC1_WIDTH, + torch.device("cpu")) + b, _ = band_columns(m, CHUNK_BANDS, 4, CHUNK_D, FC1_WIDTH, + torch.device("cpu")) + self.assertIs(a, b) + + def test_non_partition_is_rejected(self): + m = self._Module() + bands = [0, 0, 1, 1, 2, 2, 3, 3, 3] + with self.assertRaisesRegex(ValueError, "does not partition"): + band_columns(m, bands, 3, CHUNK_D, FC1_WIDTH, + torch.device("cpu")) + + +class TestControllerGuards(unittest.TestCase): + """The controller rejects allocations it cannot actually execute.""" + + def _controller(self, halve=True, sc_prec=8): + from qdit.sc_integration.sc_controller import SCController + c = SCController.__new__(SCController) + c.halve = halve + c.sc_prec = sc_prec + c.adaptive_mp_config = None + return c + + def test_qk_in_the_table_is_rejected(self): + cfg = load() + cfg.k_band_chunks[("qk", 0)] = CHUNK_BANDS + with self.assertRaisesRegex(ValueError, r"no band dispatch.*\['qk'\]"): + self._controller().init_adaptive_mp(cfg) + + def test_band_above_halve_ceiling_is_rejected(self): + cfg = load() + # 144 > 2**(8-1); legal by iso-compute, unrealizable on halve hardware. + with self.assertRaisesRegex(ValueError, "above the halve-mode maximum"): + self._controller(halve=True, sc_prec=8).init_adaptive_mp(cfg) + + def test_same_table_passes_without_halve(self): + cfg = load() + self._controller(halve=False).init_adaptive_mp(cfg) # must not raise + + def test_clean_table_passes(self): + cfg = load(ladders={"mlp_fc1:t0:l0": [list(LEVELS) for _ in range(4)]}) + self._controller(halve=True, sc_prec=8).init_adaptive_mp(cfg) + + +@unittest.skipUnless(torch.cuda.is_available(), "SC kernels need CUDA") +class TestKBandEquivalence(unittest.TestCase): + """Banded dispatch with the parent ladder must reproduce the parent. + + Not bit-identity: splitting the contraction axis changes only the order in + which per-chunk partial products are summed in fp32. Every chunk keeps its + own scale and RNG table, so the partials themselves are unchanged. + """ + + def _run(self, band_ladders): + from scmp_kernels.sc import sc_matmul + from scmp_kernels.sc.config_helpers import make_sobol_simple_config + + torch.manual_seed(0) + M, D, N = 64, FC1_WIDTH, 256 + x = torch.randn(M, D, device="cuda") + w = torch.randn(N, D, device="cuda") + rungs = torch.randint(0, len(LEVELS), (M,), device="cuda") + + out = torch.zeros(M, N, device="cuda", dtype=torch.float32) + m = TestBandColumns._Module() + cols, widths = band_columns(m, CHUNK_BANDS, 4, CHUNK_D, D, x.device) + for b in range(4): + xb = x.index_select(1, cols[b]).contiguous() + wb = w.index_select(1, cols[b]).contiguous() + for k in range(len(LEVELS)): + sl = band_ladders[b][k] + rows = (rungs == k).nonzero(as_tuple=True)[0] + if rows.numel() == 0 or sl <= 0: + continue + cfg = make_sobol_simple_config(CHUNK_D, CHUNK_D, 8) + out[rows] += sc_matmul( + xb.index_select(0, rows).contiguous(), wb, + granularity="per_row", mode="bipolar", sc_prec=8, + config=cfg, group_a=1, group_b=1, chunk_d=CHUNK_D, + stoc_len=sl) + return x, w, rungs, out + + def test_parent_ladder_reproduces_unbanded(self): + from scmp_kernels.sc import sc_matmul + from scmp_kernels.sc.config_helpers import make_sobol_simple_config + + parent_everywhere = [list(LEVELS) for _ in range(4)] + x, w, rungs, banded = self._run(parent_everywhere) + + unbanded = torch.zeros_like(banded) + cfg = make_sobol_simple_config(CHUNK_D, CHUNK_D, 8) + for k, sl in enumerate(LEVELS): + rows = (rungs == k).nonzero(as_tuple=True)[0] + if rows.numel() == 0: + continue + unbanded[rows] = sc_matmul( + x.index_select(0, rows).contiguous(), w, + granularity="per_row", mode="bipolar", sc_prec=8, + config=cfg, group_a=1, group_b=1, chunk_d=CHUNK_D, + stoc_len=sl) + + rel = ((banded - unbanded).norm() / unbanded.norm()).item() + self.assertLess(rel, 1e-5, f"banded vs parent rel err {rel:.3e}") + + def test_nonuniform_ladder_changes_the_result(self): + """A real allocation must not be a no-op dressed as one.""" + _, _, _, parent = self._run([list(LEVELS) for _ in range(4)]) + _, _, _, tuned = self._run([WIDE, WIDE, WIDE, NARROW]) + rel = ((tuned - parent).norm() / parent.norm()).item() + self.assertGreater(rel, 1e-4, "band ladders had no effect") + + +if __name__ == "__main__": + unittest.main(verbosity=2)