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.
This commit is contained in:
CalamitousFelicitousness
2026-07-18 01:43:26 +01:00
parent 74e68b58ce
commit ff94565da8
2 changed files with 127 additions and 0 deletions
+23
View File
@@ -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)