perf(lora): materialize select winners on the accelerator and time the reset

The weight-kind schedule reset ran each winner's calc_updown on the target
weight's device, which on a block-swapped denoiser is the cpu; at hundreds
of layers per pass the cpu matmuls dominated every select generation. The
delta now computes on the accelerator and moves back, matching the activate
walk's convention.

- reset and flip execution log a debug timing line (materialize, select
  loop, move/calc/apply split); the reset runs outside the activate walk,
  so its cost was invisible to the load timers

modules/lora/lora_stack.py

@@ -15,6 +15,7 @@ EST-LoRA arXiv:2508.02165 (its measured style-discrepancy estimate is
 exposed as an option instead of being derived from probe generations).
 """
 
+import time
 import weakref
 import hashlib
 
@@ -305,9 +306,16 @@ def finalize(total_steps):
     e_den = sum(e['scores'][1] for e in state['entries'].values())
     state['gamma_e'] = (e_num / e_den) if e_den > 0 else 1.0
     state['flips'] = {}
-    if any(e['kind'] == 'weight' for e in state['entries'].values()):
+    stats = {'weight_n': 0, 'factor_n': 0, 'materialize': 0.0, 'select': 0.0, 'w_move': 0.0, 'w_calc': 0.0, 'w_apply': 0.0}
+    state['stats'] = stats
+    stats['weight_n'] = sum(1 for e in state['entries'].values() if e['kind'] == 'weight')
+    stats['factor_n'] = len(state['entries']) - stats['weight_n']
+    if stats['weight_n'] > 0:
+        t0 = time.time()
         materialize_model()
+        stats['materialize'] = time.time() - t0
     style_first = 0
+    t0 = time.time()
     for layer_name, entry in list(state['entries'].items()): # snapshot: apply_selection drops entries whose module died
         flip_at = layer_flip_step(entry['scores'], state['total_steps'])
         initial = 1 if flip_at == 0 else 0
@@ -315,6 +323,7 @@ def finalize(total_steps):
         apply_selection(layer_name, entry, initial)
         if 0 < flip_at < state['total_steps']:
             state['flips'].setdefault(flip_at - 1, []).append(layer_name) # step callbacks fire after the denoise, so the flip runs one step early to be live during the crossover step's forward
+    stats['select'] = time.time() - t0
     state['finalized'] = True
     if len(state['entries']) > 0: # only a built schedule can carry a flip count, so this is the line that shows selection is live rather than requested
         gamma = state['gamma_e'] if mode() == 'estlora' else state['gamma']
@@ -322,6 +331,8 @@ def finalize(total_steps):
         if report != state['reported']: # rebuilt every pass, so a batch would otherwise repeat one line per image
             state['reported'] = report
             log.info(f'Network load: type=LoRA stack={report[0]} layers={report[1]} style={report[2]} flips={report[3]} steps={report[4]} gamma={report[5]:.3f}')
+        # logged every pass: the reset runs outside the activate walk, so its cost is invisible to the load timers
+        log.debug(f'Network select: type=LoRA reset weight={stats["weight_n"]} factor={stats["factor_n"]} time={{materialize: {stats["materialize"]:.2f}, select: {stats["select"]:.2f}, move: {stats["w_move"]:.2f}, calc: {stats["w_calc"]:.2f}, apply: {stats["w_apply"]:.2f}}}')
 
 
 def reset(total_steps):
@@ -335,10 +346,15 @@ def on_step(step):
     """Flip the layers whose crossover is this step; non-flip steps are a dict miss."""
     if not state['finalized']:
         return
-    for layer_name in state['flips'].get(int(step), ()):
+    layers = state['flips'].get(int(step), ())
+    if not layers:
+        return
+    t0 = time.time()
+    for layer_name in layers:
         entry = state['entries'].get(layer_name)
         if entry is not None:
             apply_selection(layer_name, entry, 1)
+    log.debug(f'Network select: type=LoRA flip step={int(step)} layers={len(layers)} time={time.time() - t0:.2f}')
 
 
 def apply_selection(layer_name, entry, winner):
@@ -375,6 +391,15 @@ def weight_selection(module, entry, winner):
     if weight is None or weight.is_meta:
         warn_once('select-offloaded', 'Network stack: flip=skipped weight=offloaded')
         return
+    from modules import devices
+    stats = state.get('stats') or {}
     device = weight.device
-    updown = net_module.calc_updown(backup.to(device))[0]
+    t0 = time.time()
+    base = backup.to(devices.device) # a swapped-out layer keeps its weight on cpu; the delta matmul belongs on the accelerator regardless
+    t1 = time.time()
+    updown = net_module.calc_updown(base)[0].to(device)
+    t2 = time.time()
     network_apply_weights(module, updown, None, device=device) # recomputes from the pristine backup, requantizing where the layer needs it
+    stats['w_move'] = stats.get('w_move', 0.0) + (t1 - t0)
+    stats['w_calc'] = stats.get('w_calc', 0.0) + (t2 - t1)
+    stats['w_apply'] = stats.get('w_apply', 0.0) + (time.time() - t2)

test/test-sdnq-lora-factors.py

@@ -2119,6 +2119,66 @@ def test_select_weight_replay_from_cache_skips_calc():
     return True
 
 
+def test_select_reset_reports_timing():
+    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_timed'
+    lin.network_current_names = ()
+    A1, B1, _D1 = make_delta(seed=77, sigma=1e-2)
+    A2, B2, _D2 = make_delta(seed=78, sigma=1e-2)
+    n1 = make_net('t1', lin, A1, B1)
+    n2 = make_net('t2', lin, A2, B2)
+    with mock_model(lin=lin), select_mode('klora'):
+        activate(n1, n2)
+        lora_stack.reset(20)
+        stats = lora_stack.state.get('stats')
+        assert stats is not None, 'a reset must publish its timing stats'
+        assert stats['weight_n'] == 1 and stats['factor_n'] == 0, f'weight-kind counts wrong: {stats}'
+        assert stats['w_calc'] > 0.0, 'the weight-kind winner apply must account its calc time'
+        assert stats['select'] > 0.0
+        activate()
+    layer = build_layer('uint4')
+    f1, f2, _Df1, _Df2 = select_pair(layer, seed0=79, seed1=80)
+    with mock_model(lin=layer), select_mode('klora'):
+        activate(f1, f2)
+        lora_stack.reset(20)
+        stats = lora_stack.state.get('stats')
+        assert stats is not None and stats['factor_n'] == 1 and stats['weight_n'] == 0, f'factor-kind counts wrong: {stats}'
+        assert stats['w_calc'] == 0.0, 'factor-kind resets flip segments and must not touch the weight path'
+        activate()
+    return True
+
+
+def test_select_weight_flip_calcs_on_accelerator():
+    from modules.lora import network_lora
+    lin = torch.nn.Linear(IN_F, OUT_F, bias=False, dtype=torch.bfloat16, device='cpu') # a swapped-out layer: weight lives on cpu
+    with torch.no_grad():
+        lin.weight.copy_(torch.randn(OUT_F, IN_F) * 0.02)
+    lin.network_layer_name = 'lora_transformer_swapped'
+    lin.network_current_names = ()
+    A1, B1, _D1 = make_delta(seed=81, sigma=1e-2)
+    A2, B2, _D2 = make_delta(seed=82, sigma=1e-2)
+    n1 = make_net('s1', lin, A1, B1)
+    n2 = make_net('s2', lin, A2, B2)
+    seen = []
+    real = network_lora.NetworkModuleLora.calc_updown
+    def spy(self, target, *args, **kwargs):
+        seen.append(target.device.type)
+        return real(self, target, *args, **kwargs)
+    with mock_model(lin=lin), select_mode('klora'):
+        activate(n1, n2)
+        lin.to('cpu') # the offload dispatch swaps blocks back out after the walk; the reset must not follow the weight onto the cpu
+        network_lora.NetworkModuleLora.calc_updown = spy
+        try:
+            lora_stack.reset(20)
+        finally:
+            network_lora.NetworkModuleLora.calc_updown = real
+        assert seen and all(d == DEVICE.type for d in seen), f'winner materialization must calc on the accelerator, saw {seen}'
+        activate()
+    return True
+
+
 CAT_COMPILE = category('compile')
 
 
@@ -2311,7 +2371,8 @@ def run_tests():
                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_score_pair_chunked_precision, test_select_replay_from_cache_skips_calc, test_select_weight_replay_from_cache_skips_calc]:
+               test_score_pair_chunked_precision, test_select_replay_from_cache_skips_calc, test_select_weight_replay_from_cache_skips_calc,
+               test_select_reset_reports_timing, test_select_weight_flip_calcs_on_accelerator]:
         run_test(CAT_SELECT, fn)
     log.warning('=== Compile ===')
     for fn in [test_factor_add_inside_compiled_graph, test_rank_bucket_graph_reuse]: