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

.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/

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).
 

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]

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
@@ -209,12 +210,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
@@ -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]
 
@@ -295,12 +351,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