perf(lora): chunk select scoring and cache select scores for replay

Select scoring staged full fp32 copies, an abs copy and top-k workspace per
layer (hundreds of MB of transients that collide with block swapping on
offloaded denoisers) and recomputed scores from freshly assembled deltas on
every apply, which kept select modes out of the factor-cache fast path.

- score_pair: row-chunked fp32 interiors, fp64 accumulators, one device sync
- select scores persist in the factor cache as additive per-layer records
  under the existing configuration signature
- apply_select_cached and register_weight_pair_cached replay a pair without
  assembling deltas; the weight-kind winner is still computed at schedule
  time

modules/lora/lora_factor_cache.py

@@ -157,6 +157,30 @@ def fetch(network_layer_name):
     return entry
 
 
+def lookup_scores(network_layer_name):
+    """Cached select scores for a layer as ((s0, s1), (a0, a1)), or None.
+
+    Score records ride the same signature-keyed entry as factors, and the
+    signature already pins everything the scores depend on (pair, multipliers,
+    stack mode and params). No hit/miss accounting: a record saves scoring and
+    delta assembly, not a sketch.
+    """
+    if state['sig'] is None:
+        return None
+    t = state['store'].get(f'{network_layer_name}.sel')
+    if t is None:
+        return None
+    return (float(t[0]), float(t[1])), (float(t[2]), float(t[3]))
+
+
+def store_scores(network_layer_name, scores, abs_sums):
+    """Persist a select-mode score record; additive to the entry, older files upgrade on their next pass."""
+    if state['sig'] is None:
+        return
+    state['store'][f'{network_layer_name}.sel'] = torch.tensor([scores[0], scores[1], abs_sums[0], abs_sums[1]], dtype=torch.float64)
+    state['dirty'] = True
+
+
 def store(network_layer_name, up, down, energy, calibrated, rms):
     """Quantize-before-use: returns the dequantized round-trip the caller must apply.
 

modules/lora/lora_sdnq.py

@@ -460,6 +460,54 @@ def truncate_delta(self, D, dtype):
     return up_h, down_h, energy, rms is not None
 
 
+def apply_select_cached(self, network_layer_name, wanted_names):
+    """Serve a select pair from cache and live factors before the walk assembles deltas.
+
+    A cached score record plus a factor pair per network (exact factors for
+    factorable members, cached truncations otherwise) rebuild the segments and
+    the selection registration without any ``calc_updown``. Returns None when
+    any piece is missing; the caller assembles and ``apply_select`` recomputes
+    and stores.
+    """
+    from sdnq.quant_utils import rotate_hadamard
+    deq = self.sdnq_dequantizer
+    changed = remove_factors(self)
+    if wanted_names == ():
+        return changed
+    if len(l.loaded_networks) != 2:
+        return None
+    dtype = deq.result_dtype
+    lora_factor_cache.begin_pass(wanted_names)
+    rec = lora_factor_cache.lookup_scores(network_layer_name)
+    if rec is None:
+        return None
+    pairs, notes = [], []
+    for i, net in enumerate(l.loaded_networks):
+        module = net.modules.get(network_layer_name, None)
+        if module is None:
+            return None
+        factors = get_module_factors(module, devices.device, dtype, original_shape=deq.original_shape)
+        if factors is not None:
+            up_i, down_i = factors
+            if deq.use_hadamard:
+                down_i = rotate_hadamard(down_i.to(dtype=torch.float32), group_size=deq.hadamard_group_size).to(dtype=dtype)
+        else:
+            cached = lora_factor_cache.lookup(f'{network_layer_name}#{i}')
+            if cached is None:
+                return None
+            up_i, down_i = cached[0].to(device=devices.device, dtype=dtype), cached[1].to(device=devices.device, dtype=dtype)
+            notes.append((f'{network_layer_name}#{i}', cached[2], cached[3]))
+        pairs.append((up_i, down_i))
+    scores, abs_sums = rec
+    segments, transposed = append_factors(self, [pairs[0][0], pairs[1][0]], [pairs[0][1], pairs[1][1]])
+    lora_stack.register(network_layer_name, self, 'factor', scores, segments=(segments[0], segments[1], transposed), abs_sums=abs_sums)
+    for note in notes:
+        lora_factor_cache.note_hit()
+        hosted_layers.append(note)
+    select_layers.append(network_layer_name)
+    return True
+
+
 def apply_select(self, network_layer_name, per_net, wanted_names):
     """Attach two networks' contributions as separate side-channel segments for per-layer selection.
 
@@ -502,10 +550,8 @@ def apply_select(self, network_layer_name, per_net, wanted_names):
                 up_i, down_i = lora_factor_cache.store(key, up_i, down_i, energy, calibrated, float(D.detach().float().square().mean().sqrt()))
                 hosted_layers.append((key, energy, calibrated))
         pairs.append((up_i, down_i))
-    d0 = per_net[0][1].detach().to(devices.device, torch.float32)
-    d1 = per_net[1][1].detach().to(devices.device, torch.float32)
-    scores, abs_sums = lora_stack.score_pair(d0, d1, ranks[0], ranks[1])
-    del d0, d1
+    scores, abs_sums = lora_stack.score_pair(per_net[0][1].detach(), per_net[1][1].detach(), ranks[0], ranks[1])
+    lora_factor_cache.store_scores(network_layer_name, scores, abs_sums)
     segments, transposed = append_factors(self, [pairs[0][0], pairs[1][0]], [pairs[0][1], pairs[1][1]])
     lora_stack.register(network_layer_name, self, 'factor', scores, segments=(segments[0], segments[1], transposed), abs_sums=abs_sums)
     select_layers.append(network_layer_name) # counted apart from the plain concat: both ride the svd channel but only one is a summed set

modules/lora/lora_stack.py

@@ -145,20 +145,43 @@ def combine(named_deltas, layer_name):
 
 
 def score_pair(d0, d1, rank0, rank1):
-    """Selection scores for a dense delta pair: klora top-K sums (K = rank product) or est energies; plus abs-sums for the global balance."""
-    abs_sums = (float(d0.abs().sum()), float(d1.abs().sum()))
-    if mode() == 'klora':
-        k = max(1, int(rank0) * int(rank1))
-        s0 = float(torch.topk(d0.abs().flatten(), min(k, d0.numel()), sorted=False).values.sum())
-        s1 = float(torch.topk(d1.abs().flatten(), min(k, d1.numel()), sorted=False).values.sum())
-    else:
-        s0 = float(d0.float().square().sum())
-        s1 = float(d1.float().square().sum())
-    return (s0, s1), abs_sums
+    """Selection scores for a dense delta pair: klora top-K sums (K = rank product) or est energies; plus abs-sums for the global balance.
+
+    Row-chunked fp32 interiors with fp64 accumulators and one device sync for
+    all four reductions. Full-tensor staging (fp32 copy, abs copy, top-k
+    workspace) peaks hundreds of MB per large layer, which collides with block
+    swapping on offloaded denoisers; chunking bounds the transient to the
+    chunk. The global top-K over per-chunk top-K candidates selects the same
+    element set as a whole-tensor top-K.
+    """
+    k = max(1, int(rank0) * int(rank1)) if mode() == 'klora' else 0
+    accs = []
+    for d in (d0, d1):
+        score = torch.zeros((), device=d.device, dtype=torch.float64)
+        abs_sum = torch.zeros((), device=d.device, dtype=torch.float64)
+        cands = []
+        for start in range(0, d.shape[0], ROW_CHUNK):
+            c = d[start:start + ROW_CHUNK].to(torch.float32).abs() # out-of-place abs: to() may alias a caller-owned fp32 tensor
+            abs_sum += c.sum(dtype=torch.float64)
+            if k:
+                flat = c.flatten()
+                cands.append(torch.topk(flat, min(k, flat.numel()), sorted=False).values)
+            else:
+                score += c.square().sum(dtype=torch.float64)
+        if k and cands:
+            allc = torch.cat(cands) if len(cands) > 1 else cands[0]
+            score = torch.topk(allc, min(k, allc.numel()), sorted=False).values.sum(dtype=torch.float64)
+        accs.append((score, abs_sum))
+    packed = torch.stack([accs[0][0], accs[0][1], accs[1][0], accs[1][1]]).cpu()
+    return (float(packed[0]), float(packed[2])), (float(packed[1]), float(packed[3]))
 
 
-def register_weight_pair(layer_name, module, per_net):
-    """Score and register a weight-kind selection pair; True when the layer is scheduled."""
+def register_weight_pair(layer_name, module, per_net, wanted_names=None):
+    """Score and register a weight-kind selection pair; True when the layer is scheduled.
+
+    The scores persist in the factor cache when a pass identity is given, so a
+    later apply of the same configuration registers from the record alone.
+    """
     from modules.lora import lora_common as l
     if per_net is None or len(per_net) != 2:
         return False
@@ -172,11 +195,35 @@ def register_weight_pair(layer_name, module, per_net):
             return False
         names.append(net_name)
         ranks.append(int(getattr(net_module, 'dim', 0) or 0) or 64)
-    scores, abs_sums = score_pair(per_net[0][1].float(), per_net[1][1].float(), ranks[0], ranks[1])
+    scores, abs_sums = score_pair(per_net[0][1], per_net[1][1], ranks[0], ranks[1])
+    if wanted_names is not None:
+        from modules.lora import lora_factor_cache
+        lora_factor_cache.begin_pass(wanted_names)
+        lora_factor_cache.store_scores(layer_name, scores, abs_sums)
     register(layer_name, module, 'weight', scores, nets=tuple(names), abs_sums=abs_sums)
     return True
 
 
+def register_weight_pair_cached(layer_name, module, wanted_names):
+    """Register a weight-kind pair from its cached score record; True when served.
+
+    The record was stored under the same configuration signature, which pins
+    the loaded pair, multipliers and stack settings, so both networks are known
+    to target the layer and the prompt-order roles are unchanged.
+    """
+    from modules.lora import lora_common as l
+    from modules.lora import lora_factor_cache
+    if len(l.loaded_networks) != 2:
+        return False
+    lora_factor_cache.begin_pass(wanted_names)
+    rec = lora_factor_cache.lookup_scores(layer_name)
+    if rec is None:
+        return False
+    scores, abs_sums = rec
+    register(layer_name, module, 'weight', scores, nets=tuple(n.name for n in l.loaded_networks), abs_sums=abs_sums)
+    return True
+
+
 def drop(layer_name):
     """Forget a layer's selection entry (its factors were removed or restored)."""
     if layer_name is not None and state['entries'].pop(layer_name, None) is not None:

modules/lora/networks.py

@@ -117,18 +117,20 @@ def network_activate(include=None, exclude=None):
                             weights_backup = getattr(module, "network_weights_backup", None)
                             if weights_backup is not None and not isinstance(weights_backup, bool):
                                 network_apply_weights(module, None, None, device=device)
-                            per_net, sel_bias = network_calc_weights(module, network_layer_name, elimit=elimit, per_net=True)
-                            if sel_bias is None:
-                                applied = lora_sdnq.apply_select(module, network_layer_name, per_net, component_wanted)
-                                if applied is not None:
-                                    if applied and component_wanted:
-                                        applied_layers.append(network_layer_name)
-                                        applied_weight += 1
-                                    module.network_current_names = component_wanted
-                                    module.network_current_stack = stack_sig
-                                    if task is not None:
-                                        pbar.update(task, advance=1)
-                                    continue
+                            applied = lora_sdnq.apply_select_cached(module, network_layer_name, component_wanted) # a stored score record and factor pair serve before the deltas are assembled
+                            if applied is None:
+                                per_net, sel_bias = network_calc_weights(module, network_layer_name, elimit=elimit, per_net=True)
+                                if sel_bias is None:
+                                    applied = lora_sdnq.apply_select(module, network_layer_name, per_net, component_wanted)
+                            if applied is not None:
+                                if applied and component_wanted:
+                                    applied_layers.append(network_layer_name)
+                                    applied_weight += 1
+                                module.network_current_names = component_wanted
+                                module.network_current_stack = stack_sig
+                                if task is not None:
+                                    pbar.update(task, advance=1)
+                                continue
                             lora_stack.warn_once('select-unridable', f'Network stack: mode={lora_stack.mode()} layer="{network_layer_name}" fallback=sum') # a pair the channel cannot carry (bias delta or malformed member) sums like any unsupported set
                         elif getattr(module, 'sdnq_dequantizer', None) is not None: # hosting disabled: quantized layers have no side-channel to carry segments and packed backups cannot flip, so the sum paths below take the layer
                             if any(net.modules.get(network_layer_name, None) is not None for net in l.loaded_networks):
@@ -137,8 +139,18 @@ def network_activate(include=None, exclude=None):
                             sel_backup = network_backup_weights(module, network_layer_name, component_wanted, fuse)
                             weights_backup = getattr(module, "network_weights_backup", None)
                             if weights_backup is not None and not isinstance(weights_backup, bool):
+                                if lora_stack.register_weight_pair_cached(network_layer_name, module, component_wanted): # a stored score record registers without assembling the pair
+                                    backup_size += sel_backup
+                                    network_apply_weights(module, None, None, device=device) # pristine until the schedule applies the winner
+                                    applied_layers.append(network_layer_name)
+                                    applied_weight += 1
+                                    module.network_current_names = component_wanted
+                                    module.network_current_stack = stack_sig
+                                    if task is not None:
+                                        pbar.update(task, advance=1)
+                                    continue
                                 per_net, sel_bias = network_calc_weights(module, network_layer_name, elimit=elimit, per_net=True)
-                                if sel_bias is None and lora_stack.register_weight_pair(network_layer_name, module, per_net):
+                                if sel_bias is None and lora_stack.register_weight_pair(network_layer_name, module, per_net, component_wanted):
                                     backup_size += sel_backup # counted only when this branch keeps the layer; the fallthrough re-enters the shared backup call below, which counts it then
                                     network_apply_weights(module, None, None, device=device) # pristine until the schedule applies the winner
                                     applied_layers.append(network_layer_name)

test/test-sdnq-lora-factors.py

@@ -2003,6 +2003,122 @@ def test_flip_lands_before_crossover_step():
     return True
 
 
+def test_score_pair_chunked_precision():
+    torch.manual_seed(71)
+    shapes = [(700, 460), (OUT_F, IN_F), (64,)] # off-chunk rows, square, and a 1-D norm delta
+    for shape in shapes:
+        d0 = (torch.randn(*shape, device=DEVICE) * 1e-2).to(torch.bfloat16)
+        d1 = (torch.randn(*shape, device=DEVICE) * 1e-2).to(torch.bfloat16)
+        for mode_name in ('klora', 'estlora'):
+            with select_mode(mode_name):
+                (s0, s1), (a0, a1) = lora_stack.score_pair(d0, d1, 8, 8)
+            f0, f1 = d0.to(torch.float64), d1.to(torch.float64)
+            ra0, ra1 = float(f0.abs().sum()), float(f1.abs().sum())
+            if mode_name == 'klora':
+                k = 64
+                rs0 = float(torch.topk(f0.abs().flatten(), min(k, f0.numel()), sorted=False).values.sum())
+                rs1 = float(torch.topk(f1.abs().flatten(), min(k, f1.numel()), sorted=False).values.sum())
+            else:
+                rs0, rs1 = float(f0.square().sum()), float(f1.square().sum())
+            for got, ref, label in ((s0, rs0, 'score0'), (s1, rs1, 'score1'), (a0, ra0, 'abs0'), (a1, ra1, 'abs1')):
+                assert abs(got - ref) <= 1e-9 * max(abs(ref), 1e-12), f'{mode_name} {label} shape={shape}: got {got!r} ref {ref!r}'
+    frozen = torch.randn(300, 200, device=DEVICE, dtype=torch.float32) * 1e-2 # fp32 input aliases through to(); abs must stay out-of-place
+    pristine = frozen.clone()
+    with select_mode('klora'):
+        lora_stack.score_pair(frozen, frozen, 4, 4)
+    assert torch.equal(frozen, pristine), 'score_pair must not mutate a caller-owned fp32 delta'
+    return True
+
+
+def test_select_replay_from_cache_skips_calc():
+    import tempfile
+    layer = build_layer('uint4')
+    torch.manual_seed(73)
+    Dd1 = (torch.randn(OUT_F, 24, device=DEVICE) @ torch.randn(24, IN_F, device=DEVICE)) * 1e-3
+    Dd2 = (torch.randn(OUT_F, 24, device=DEVICE) @ torch.randn(24, IN_F, device=DEVICE)) * 1e-3
+    with tempfile.TemporaryDirectory() as tmp:
+        with host_rank(32), host_cache(10, os.path.join(tmp, 'cache')), mock_model(lin=layer), select_mode('klora'):
+            n1 = make_dense_net('selk1', layer, Dd1)
+            n2 = make_dense_net('selk2', layer, Dd2)
+            for net in (n1, n2):
+                lora_file = os.path.join(tmp, f'{net.name}.safetensors')
+                with open(lora_file, 'wb') as f:
+                    f.write(b'0' * 64)
+                net.network_on_disk.filename = lora_file
+            from modules.modeldata import model_data
+            model_data.sd_model.sd_checkpoint_info = MockCheckpointInfo('test/cache-model')
+            Wdq0 = dq(layer)
+            with counting_calc() as calls:
+                activate(n1, n2)
+                assert calls['n'] > 0, 'a fresh select apply must assemble both deltas'
+                entry = lora_stack.state['entries'].get('lora_transformer_test')
+                assert entry is not None and entry['kind'] == 'factor'
+                fresh_scores, fresh_abs = entry['scores'], entry['abs_sums']
+                first = dq(layer)
+                activate()
+                calls['n'] = 0
+                real_svd = torch.svd_lowrank
+                torch.svd_lowrank = raise_no_svd
+                try:
+                    activate(n1, n2)
+                finally:
+                    torch.svd_lowrank = real_svd
+                assert calls['n'] == 0, f'a select replay must not assemble deltas: calc_updown ran {calls["n"]} times'
+                entry = lora_stack.state['entries'].get('lora_transformer_test')
+                assert entry is not None and entry['kind'] == 'factor', 'the replay must re-register the selection'
+                assert entry['scores'] == fresh_scores and entry['abs_sums'] == fresh_abs, 'cached scores must replay exactly'
+                assert torch.equal(dq(layer), first), 'select replay must be bit-identical to the fresh apply'
+                activate()
+                assert torch.equal(dq(layer), Wdq0)
+    return True
+
+
+def test_select_weight_replay_from_cache_skips_calc():
+    import tempfile
+    lin = torch.nn.Linear(IN_F, OUT_F, bias=False, dtype=torch.bfloat16, device=DEVICE)
+    with torch.no_grad():
+        lin.weight.copy_(torch.randn(OUT_F, IN_F, device=DEVICE) * 0.02)
+    lin.network_layer_name = 'lora_transformer_plainsel'
+    lin.network_current_names = ()
+    torch.manual_seed(75)
+    Dd1 = (torch.randn(OUT_F, 24, device=DEVICE) @ torch.randn(24, IN_F, device=DEVICE)) * 1e-3
+    Dd2 = (torch.randn(OUT_F, 24, device=DEVICE) @ torch.randn(24, IN_F, device=DEVICE)) * 1e-3
+    W0 = lin.weight.detach().float().clone()
+    with tempfile.TemporaryDirectory() as tmp:
+        with host_cache(10, os.path.join(tmp, 'cache')), mock_model(lin=lin), select_mode('klora'):
+            n1 = make_dense_net('selw1', lin, Dd1)
+            n2 = make_dense_net('selw2', lin, Dd2)
+            for net in (n1, n2):
+                lora_file = os.path.join(tmp, f'{net.name}.safetensors')
+                with open(lora_file, 'wb') as f:
+                    f.write(b'0' * 64)
+                net.network_on_disk.filename = lora_file
+            from modules.modeldata import model_data
+            model_data.sd_model.sd_checkpoint_info = MockCheckpointInfo('test/cache-model')
+            with counting_calc() as calls:
+                activate(n1, n2)
+                assert calls['n'] > 0, 'a fresh weight-kind select apply must assemble the pair'
+                entry = lora_stack.state['entries'].get('lora_transformer_plainsel')
+                assert entry is not None and entry['kind'] == 'weight'
+                fresh_scores = entry['scores']
+                assert torch.equal(lin.weight.detach().float(), W0), 'weights stay pristine until the schedule applies a winner'
+                lora_stack.reset(20)
+                fresh_selected = lin.weight.detach().float().clone()
+                activate()
+                assert torch.equal(lin.weight.detach().float(), W0)
+                calls['n'] = 0
+                activate(n1, n2)
+                assert calls['n'] == 0, f'a weight-kind select replay must not assemble the pair: calc_updown ran {calls["n"]} times'
+                entry = lora_stack.state['entries'].get('lora_transformer_plainsel')
+                assert entry is not None and entry['kind'] == 'weight', 'the replay must register from the score record'
+                assert entry['scores'] == fresh_scores, 'cached scores must replay exactly'
+                lora_stack.reset(20) # the winner recompute at schedule time still assembles its own delta, by design
+                assert torch.equal(lin.weight.detach().float(), fresh_selected), 'the replayed schedule must select the same winner'
+                activate()
+                assert torch.equal(lin.weight.detach().float(), W0)
+    return True
+
+
 CAT_COMPILE = category('compile')
 
 
@@ -2194,7 +2310,8 @@ def run_tests():
                test_select_deactivate_from_midflip, test_select_requires_exactly_two_nets, test_select_gated_off_when_compiled,
                test_select_finalize_drops_dead_module, test_select_int8_pair_rides_segments, test_select_gate_dormant_without_pair, test_stale_schedule_dropped_on_reapply,
                test_est_energy_matches_full_frobenius, test_select_weight_kind_plain_layer,
-               test_select_gamma_tracks_live_entries, test_select_host_disabled_falls_back_to_sum, test_flip_lands_before_crossover_step]:
+               test_select_gamma_tracks_live_entries, test_select_host_disabled_falls_back_to_sum, test_flip_lands_before_crossover_step,
+               test_score_pair_chunked_precision, test_select_replay_from_cache_skips_calc, test_select_weight_replay_from_cache_skips_calc]:
         run_test(CAT_SELECT, fn)
     log.warning('=== Compile ===')
     for fn in [test_factor_add_inside_compiled_graph, test_rank_bucket_graph_reuse]: