feat(mp): per-group (K-band) stream lengths on the contraction axis (stacks on #6)

[heroarmor]

Summary

MP was per-row on the token axis only: every chunk of the contraction axis ran
the same stoc_len. This adds per-group stream lengths — the contraction
axis is partitioned into bands of whole quantization chunks, and band b runs
rung k at its own length L[b][k], under the exact per-rung iso-compute
identity that scmp_kernels#31
validates at table load.

Row dispatch is unchanged. Because L[b][k] == stoc_len_levels[k] reproduces
the per-row parent exactly, the parent is a legal point in the space and the
refinement cannot lose: a tilt that does not help degenerates back to the
parent rather than costing accuracy.

⚠️ Stacks on #6 — and why it has to

main does not currently run adaptive MP at all. It calls the pre-refactor
convention

adaptive_classify_rows(row_metric, current_timestep, total_timesteps, config, operator=..., ...)

while the submodule main pins (9358239) takes (metric, config, ...). Any
adaptive-MP forward off main raises
TypeError: adaptive_classify_rows() got multiple values for argument 'operator'.
#6's e061384 is what fixes the call sites.

I first built this as a standalone branch off main — the Python cherry-picked
cleanly, but it cannot run there. So this branch is #6 + one commit. Please
merge #6 first; review this PR's single commit dfb061f.

How the banding works

sc_matmul takes a scalar stoc_len — the chunked MLP path builds
cum_indicator once outside the chunk loop and every chunk shares it — so the
banding happens above the kernel: 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:

path	how	vs. the parent under the parent ladder
sc_attention.py	already walks the axis 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	0.000e+00 — bit-identical
sc_mlp.py	hands chunking to the kernel, so bands gather their columns and each band re-chunks at the same boundaries	6.8e-08 — fp32 summation order, nothing else

Guards

All four failures below would otherwise surface only as a wrong number:

• qk/av in the table → rejected at init_adaptive_mp. They contract
  over head_dim, have no band dispatch, and would run the per-row parent
  while the run carries the k-bands label.
• A band length above the halve ceiling → rejected at startup rather than
  mid-run on some later timestep. 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 range+dynamic path has no
  contraction-axis dispatch and would drop the allocation on the floor.
• A table priced at a different chunk_d or contraction width → raises
  instead of falling back to the parent.

Tooling

scripts/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 a solved allocation — the LLM side
solves it by Lagrangian water-fill on measured per-band error curves.

Verified against the real calib_fix_avg32.json: budget used per rung
0.9931–1.0000, zero overspend.

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.

Testing

tests/test_kbands.py — 24 tests: table validation, column partitioning
including tail-chunk ordering, the controller guards, and on GPU both the
parent-equivalence and non-degeneracy gates.

scripts/verify_kband_e2e.py drives the real _sc_linear_dynamic_mp methods
rather than reimplementing the band loop, so a mistake in row indexing or
accumulation shows up there and not only in production. It also prints which
trees were imported — scmp_kernels is an editable install and does not follow
$REPO, so a gate whose log cannot prove what it ran against proves nothing.

Verified on gl1802 (job 58187543), qdit @ dfb061f + scmp_kernels @ acf8eca:

24 passed
[OK] SCMlp.mlp_fc1 parent-ladder: rel err vs parent 6.776e-08 (want < 1e-5)
[OK] SCMlp.mlp_fc1 tilted-ladder: rel err vs parent 1.119e-01 (want > 1e-4)
[OK] SCAttention.proj parent-ladder: rel err vs parent 0.000e+00 (want < 1e-5)
[OK] SCAttention.proj tilted-ladder: rel err vs parent 1.113e-01 (want > 1e-4)
RESULT: PASS

Submodule

Bumps scmp_kernels to heroarmor:feat/k-bands (acf8eca) = #6's pin
c967d3a + the k_bands config layer, one commit and nothing else. The clean
version of that layer against kernels main is
scmp_kernels#31;
this pin exists because #6's line carries trace (#25), which never reached
kernels main.

🤖 Generated with Claude Code

[Copilot]

Pull request overview

This PR extends adaptive mixed precision (MP) to support per-group (K-band) stream lengths along the contraction axis, enabling each contraction-axis band to run each rung at its own stream length while preserving a per-rung iso-compute identity (so the per-row parent remains a valid degenerate point). It also adds validation/guards, calibration improvements (FLOP-weighted budgeting + halve-aware checks + new metric), and extensive tooling/tests to verify equivalence and non-degeneracy.

Changes:

• Add K-band dispatch + validation plumbing across SCController, SCMlp, and SCAttention via new sc_kbands helpers.
• Update calibration + runtime wiring around adaptive MP (table-required), including FLOP-weighted budget accounting, halve-mode guards, and mse_rel metric support.
• Add tests and multiple end-to-end / cluster scripts to validate correctness and evaluate configurations.

Reviewed changes

Copilot reviewed 37 out of 38 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
tools/smoke_test_e2e.py	GPU smoke test for imports + basic sc_matmul correctness + deprecated-name surveillance.
tools/sc_call_counter.py	Wrapper to count sc_matmul invocations during quant_sc_main.py runs.
tools/sc_call_counter_full.py	Variant of call counter configured for “all ops on” uniform settings.
tools/compare_vs_cpu_reference.py	Compares GPU Triton implementation vs CPU reference implementation.
tools/compare_old_vs_new.py	Numerical comparison harness for old SC implementation vs new scmp_kernels.
tools/compare_old_vs_new_thorough.py	Exhaustive (seed/shape/level) bit-identity verification harness.
tools/compare_old_vs_new_isolated.py	Isolation tool to attribute observed diffs to quantization margin behavior.
tests/test_kbands.py	Unit tests for K-band schema validation, column banding, controller guards, and equivalence.
scripts/verify_kband_e2e.py	End-to-end driver exercising real _sc_linear_dynamic_mp band dispatch vs parent.
scripts/sbatch_uniform48_rerun100.sb	Cluster job script for uniform48 reproducibility rerun.
scripts/sbatch_uniform_halve_auto.sb	Cluster uniform generation with halve-mode + timeout-safe auto-resume.
scripts/sbatch_uniform_auto.sb	Cluster uniform generation with timeout-safe auto-resume (non-halve arm).
scripts/sbatch_opfreeze_probe.sb	Cluster probe forcing one op to FP-like precision to measure sensitivity.
scripts/sbatch_only1op_sc32.sb	Cluster single-op isolation probe at stressed SC length.
scripts/sbatch_mp_globalpr_flop_halve_auto.sb	Cluster FLOP-weighted + halve adaptive-MP generation with auto-resume.
scripts/sbatch_mp_globalpr_fill.sb	Cluster global+per-row MP generation “fill” script with chaining.
scripts/sbatch_mp_globalpr_auto.sb	Cluster global+per-row MP generation auto-resume script.
scripts/sbatch_calib_globalpr.sb	Cluster per-row global calibration job script (row-weighted arm).
scripts/sbatch_calib_globalpr_halve.sb	Cluster calibration job script with halve-aware level/budget handling.
scripts/sbatch_calib_globalpr_flop.sb	Cluster FLOP-weighted calibration job script.
scripts/sbatch_calib_globalpr_flop_halve.sb	Cluster FLOP-weighted + halve calibration job script.
scripts/quant_sc_main.py	Requires calibrated table for adaptive MP; removes alpha/beta fallback + updates naming/logging.
scripts/eval/sbatch_probe_realized_dist.sb	Probe script to compare realized MP level distributions (old vs new tables).
scripts/eval/sbatch_probe_realized_dist_old.sb	Control probe for realized distribution using old table without halve.
scripts/eval/sbatch_kernel_level_err.sb	Kernel-level error probe vs stoc_len under halve/no-halve.
scripts/eval/sbatch_kband_equivalence.sb	Cluster gate running K-band unit tests + end-to-end equivalence check.
scripts/eval/sbatch_eval_uniform_halve.sb	GPU evaluation script for a uniform (halve) config (KID/FID/etc).
scripts/eval/sbatch_eval_mp_flop_halve.sb	GPU evaluation script for FLOP+halve adaptive MP configs.
scripts/eval/sbatch_eval_ab_kid100.sb	One-off A/B KID evaluation on a fixed 100-image index set.
scripts/eval/eval_mp_flop_halve_cpu.sh	CPU-only eval script for FLOP+halve adaptive MP configs.
scripts/calibrate_mp_thresholds.py	Calibration updates: FLOP-weighted budgeting, halve guards, mse_rel, per-row qk recording, and halve propagation into probes.
scripts/build_kband_table.py	Tool to add a k_bands section to a calibrated MP table (tilt-based ladder builder).
qdit/sc_integration/sc_mlp.py	Implements K-band dispatch for MLP contraction-axis banding + guards for combined MP path.
qdit/sc_integration/sc_kbands.py	New helper module: resolve/validate band configs, build band column indices, reject unsupported combined modes.
qdit/sc_integration/sc_controller.py	Validates K-band allocations at init (supported ops + halve ceiling).
qdit/sc_integration/sc_attention.py	Implements K-band dispatch for attention linears + switches QK MP to per-query-row path; adds combined-path guard.
.gitignore	Ignores SLURM logs and generated figures.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.