From ff94565da879aa185274a0906ad03ddc103dc3e7 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sat, 18 Jul 2026 01:43:26 +0100 Subject: [PATCH] perf(lora): pad side-channel factor ranks to fixed buckets The compiled dequant specializes per factor rank, so each distinct lora set shape paid a compile stall on switch. Pad appended factors to a power-of-two rank ladder (multiples of 64 past the hosted cap) with zero fill: switches inside a bucket reuse the compiled graph, and common trained ranks land on their bucket exactly so padding is usually a no-op. Regression tests pin the factor add inside the single compiled graph and the bucket reuse. --- modules/lora/lora_sdnq.py | 23 ++++++++ test/test-sdnq-lora-factors.py | 104 +++++++++++++++++++++++++++++++++ 2 files changed, 127 insertions(+) diff --git a/modules/lora/lora_sdnq.py b/modules/lora/lora_sdnq.py index b2a91de6c..5635f8495 100644 --- a/modules/lora/lora_sdnq.py +++ b/modules/lora/lora_sdnq.py @@ -46,6 +46,22 @@ from modules.logger import log fallback_layers: list[str] = [] hosted_layers: list[tuple[str, float, bool]] = [] +def rank_bucket(r): + """Fixed rank ladder for compiled-graph reuse: powers of two up to 256, multiples of 64 above (hosted rank plus exact members).""" + if r <= 8: + return 8 + if r <= 256: + return 1 << (r - 1).bit_length() + return -(-r // 64) * 64 + + +def pad_rank(t, dim, bucket): + if t.shape[dim] >= bucket: + return t + shape = list(t.shape) + shape[dim] = bucket - t.shape[dim] + return torch.cat([t, t.new_zeros(shape)], dim=dim) + def enabled(): """True while the exact svd-channel machinery may take quantized layers; the requantize choice routes every layer to the legacy weight-rewrite path.""" @@ -187,6 +203,13 @@ def append_factors(self, ups, downs): parts_down = ([orig_down.to(device=devices.device, dtype=dtype)] if orig_down is not None else []) + downs new_up = torch.cat(parts_up, dim=1).contiguous() new_down = torch.cat(parts_down, dim=0).contiguous() + from sdnq.common import use_torch_compile + if use_torch_compile: + # the compiled dequant specializes per factor rank; pad to a fixed bucket so set switches inside a bucket reuse the graph (zero columns contribute exactly nothing) + dim_up, dim_down = (0, 1) if deq.use_quantized_matmul else (1, 0) + bucket = rank_bucket(new_up.shape[dim_up]) + new_up = pad_rank(new_up, dim_up, bucket) + new_down = pad_rank(new_down, dim_down, bucket) self.sdnq_lora_svd_stash = (orig_up, orig_down) self.svd_up = torch.nn.Parameter(new_up.to(device=device), requires_grad=False) self.svd_down = torch.nn.Parameter(new_down.to(device=device), requires_grad=False) diff --git a/test/test-sdnq-lora-factors.py b/test/test-sdnq-lora-factors.py index 58028801c..4c2c17ce6 100644 --- a/test/test-sdnq-lora-factors.py +++ b/test/test-sdnq-lora-factors.py @@ -37,6 +37,11 @@ for the per-model analyzer): - Factor cache: hosted factors replay bit-identically from the disk cache without re-running the svd, a configuration change (multiplier) misses and writes a separate entry, and budget 0 writes nothing. +- Compile: the factor add runs inside the single compiled dequant graph + (fullgraph, no breaks) and matches the eager result; factor ranks pad to + a fixed bucket ladder so set switches inside a bucket reuse the compiled + graph while a novel bucket compiles exactly once, and padding changes + the dequantized weight by nothing beyond reduction-order ulp. All tensors are synthetic; no model files or running server required. @@ -1013,6 +1018,102 @@ def test_factor_cache_disabled_at_zero(): return True +CAT_COMPILE = category('compile') + + +def dq_compiled(layer): + # the production entry: skip_compile left at its default so the shared compiled dequant runs + return layer.sdnq_dequantizer(layer.weight, layer.scale, zero_point=layer.zero_point, + svd_up=layer.svd_up, svd_down=layer.svd_down, + skip_quantized_matmul=layer.sdnq_dequantizer.use_quantized_matmul, + dtype=torch.float32) + + +def graph_stats(): + from torch._dynamo.utils import counters + return int(counters['stats']['unique_graphs']), sum(counters['graph_break'].values()) + + +def test_factor_add_inside_compiled_graph(): + from sdnq.common import use_torch_compile + if not use_torch_compile: + return True # compile disabled at sdnq import (no triton); nothing to pin + import torch._dynamo + from torch._dynamo.utils import counters + layer = build_layer('uint4') + A, B, _D = make_delta() + dtype = layer.sdnq_dequantizer.result_dtype + torch._dynamo.reset() + counters.clear() + lora_sdnq.append_factors(layer, [B.to(dtype)], [A.to(dtype)]) + W_c = dq_compiled(layer) + graphs, breaks = graph_stats() + assert breaks == 0, f'graph breaks in the compiled dequant: {breaks}' + assert graphs == 1, f'factor-bearing dequant must be one compiled region, got {graphs} graphs' + W_e = dq(layer) + assert torch.allclose(W_c, W_e, rtol=1e-3, atol=1e-4), f'compiled vs eager dequant diverged, max {float((W_c - W_e).abs().max()):.3e}' + lora_sdnq.remove_factors(layer) + return True + + +def test_rank_bucket_graph_reuse(): + from sdnq.common import use_torch_compile + if not use_torch_compile: + return True + import torch._dynamo + from torch._dynamo.utils import counters + import sdnq.common as sdnq_common + layer = build_layer('uint4', use_hadamard=False) + dtype = layer.sdnq_dequantizer.result_dtype + torch.manual_seed(13) + mk = lambda r: (torch.randn(OUT_F, r, device=DEVICE, dtype=dtype) * 0.01, torch.randn(r, IN_F, device=DEVICE, dtype=dtype) * 0.01) + B8, A8 = mk(8) + B6, A6 = mk(6) + B24, A24 = mk(24) + + torch._dynamo.reset() + counters.clear() + lora_sdnq.append_factors(layer, [B8], [A8]) + assert layer.svd_up.shape[1] == 8, f'rank 8 must bucket to 8, got {layer.svd_up.shape[1]}' + dq_compiled(layer) + g_first, _ = graph_stats() + + lora_sdnq.remove_factors(layer) + lora_sdnq.append_factors(layer, [B6], [A6]) + assert layer.svd_up.shape[1] == 8, f'rank 6 must pad to bucket 8, got {layer.svd_up.shape[1]}' + assert float(layer.svd_up[:, 6:].abs().sum()) == 0.0, 'pad columns must be exact zeros' + dq_compiled(layer) + g_same, _ = graph_stats() + assert g_same == g_first, f'same bucket must reuse the graph: {g_first} -> {g_same}' + + lora_sdnq.remove_factors(layer) + lora_sdnq.append_factors(layer, [B24], [A24]) + assert layer.svd_up.shape[1] == 32, f'rank 24 must pad to bucket 32, got {layer.svd_up.shape[1]}' + dq_compiled(layer) + g_novel, _ = graph_stats() + assert g_novel == g_first + 1, f'novel bucket must compile exactly one new graph: {g_first} -> {g_novel}' + + lora_sdnq.remove_factors(layer) + lora_sdnq.append_factors(layer, [B8], [A8]) + dq_compiled(layer) + g_back, _ = graph_stats() + assert g_back == g_novel, f'returning to a seen bucket must be free: {g_novel} -> {g_back}' + + W_padded = dq(layer) + lora_sdnq.remove_factors(layer) + old_flag = sdnq_common.use_torch_compile + sdnq_common.use_torch_compile = False + try: + lora_sdnq.append_factors(layer, [B8], [A8]) + assert layer.svd_up.shape[1] == 8 + W_unpadded = dq(layer) + finally: + sdnq_common.use_torch_compile = old_flag + lora_sdnq.remove_factors(layer) + assert torch.allclose(W_padded, W_unpadded, rtol=0.0, atol=1e-6), f'padding must be inert beyond reduction-order ulp, max {float((W_padded - W_unpadded).abs().max()):.3e}' + return True + + CAT_ROBUST = category('robustness') @@ -1090,6 +1191,9 @@ def run_tests(): for fn in [test_factor_cache_roundtrip_bitexact, test_factor_cache_invalidates_on_multiplier, test_factor_cache_int8_quantization, test_factor_cache_disabled_at_zero]: run_test(CAT_FCACHE, fn) + log.warning('=== Compile ===') + for fn in [test_factor_add_inside_compiled_graph, test_rank_bucket_graph_reuse]: + run_test(CAT_COMPILE, fn) log.warning('=== Robustness ===') for fn in [test_remove_factors_after_device_move, test_stacked_shape_mismatch_falls_back]: run_test(CAT_ROBUST, fn)