diff --git a/src/cache_dit/quantization/svdquant/fused.py b/src/cache_dit/quantization/svdquant/fused.py index 6aa2b228..8b2b48d7 100644 --- a/src/cache_dit/quantization/svdquant/fused.py +++ b/src/cache_dit/quantization/svdquant/fused.py @@ -53,16 +53,17 @@ def fused_gelu_mlp( quantized_x, ascales, lora_act = fc1.quantize(x, pad_size=pad_size) batch_size_pad = quantized_x.shape[0] + fc1_out_padded = getattr(fc1, '_padded_out_features', fc1.out_features) if fc2.precision == "nvfp4": qout_act = torch.empty( batch_size_pad, - fc1.out_features // 2, + fc1_out_padded // 2, dtype=torch.uint8, device=x.device, ) qout_ascales = torch.empty( - fc1.out_features // 16, + fc1_out_padded // 16, batch_size_pad, dtype=torch.float8_e4m3fn, device=x.device, @@ -70,12 +71,12 @@ def fused_gelu_mlp( else: qout_act = torch.empty( batch_size_pad, - fc1.out_features // 2, + fc1_out_padded // 2, dtype=torch.uint8, device=x.device, ) qout_ascales = torch.empty( - fc1.out_features // 64, + fc1_out_padded // 64, batch_size_pad, dtype=x.dtype, device=x.device, @@ -107,6 +108,8 @@ def fused_gelu_mlp( output = fc2.forward_quant(qout_act, qout_ascales, qout_lora_act) output = output[:token_count] + if hasattr(fc2, '_needs_n_strip') and fc2._needs_n_strip: + output = output[:, :fc2.out_features] return output.reshape(*leading_shape, fc2.out_features) @@ -149,9 +152,10 @@ def fused_gelu_proj( quantized_x, ascales, lora_act = fc1.quantize(x, pad_size=pad_size) batch_size_pad = quantized_x.shape[0] + fc1_out_padded = getattr(fc1, '_padded_out_features', fc1.out_features) out = torch.empty( batch_size_pad, - fc1.out_features, + fc1_out_padded, dtype=x.dtype, device=x.device, ) @@ -161,12 +165,12 @@ def fused_gelu_proj( if fc1.precision == "nvfp4": qout_act = torch.empty( batch_size_pad, - fc1.out_features // 2, + fc1_out_padded // 2, dtype=torch.uint8, device=x.device, ) qout_ascales = torch.empty( - fc1.out_features // 16, + fc1_out_padded // 16, batch_size_pad, dtype=torch.float8_e4m3fn, device=x.device, @@ -174,12 +178,12 @@ def fused_gelu_proj( else: qout_act = torch.empty( batch_size_pad, - fc1.out_features // 2, + fc1_out_padded // 2, dtype=torch.uint8, device=x.device, ) qout_ascales = torch.empty( - fc1.out_features // 64, + fc1_out_padded // 64, batch_size_pad, dtype=x.dtype, device=x.device, @@ -188,7 +192,7 @@ def fused_gelu_proj( # smooth_factor must match fc1.out_features (output dim N), # not fc1.in_features (fc1.smooth_factor is for input quant). smooth_ones = torch.ones( - fc1.out_features, + fc1_out_padded, dtype=x.dtype, device=x.device, ) @@ -211,6 +215,8 @@ def fused_gelu_proj( ) out = out[:token_count] + if hasattr(fc1, '_needs_n_strip') and fc1._needs_n_strip: + out = out[:, :fc1.out_features] return out.reshape(*leading_shape, fc1.out_features) diff --git a/src/cache_dit/quantization/svdquant/linear.py b/src/cache_dit/quantization/svdquant/linear.py index 7486586a..15b43253 100644 --- a/src/cache_dit/quantization/svdquant/linear.py +++ b/src/cache_dit/quantization/svdquant/linear.py @@ -1,5 +1,6 @@ # Adapted from nunchaku's implementation of SVDQW4A4Linear. import torch +import torch.nn.functional as F from torch import nn from ...kernels import svdq_gemm_w4a4 @@ -61,11 +62,16 @@ def __init__( device = torch.device("cpu") if rank < 0: raise ValueError(f"rank must be non-negative, got {rank}.") - if in_features % 2 != 0: - raise ValueError(f"in_features must be divisible by 2, got {in_features}.") if runtime_kernel not in {"v1", "v2"}: raise ValueError(f"runtime_kernel must be 'v1' or 'v2', got {runtime_kernel!r}.") + if precision == "nvfp4": + self.group_size = 16 + elif precision == "int4": + self.group_size = 64 + else: + raise ValueError(f"Invalid precision: {precision}") + self.in_features = in_features self.out_features = out_features self.rank = rank @@ -73,46 +79,49 @@ def __init__( self.torch_dtype = torch_dtype self.runtime_kernel = runtime_kernel - if precision == "nvfp4": - self.group_size = 16 - elif precision == "int4": - self.group_size = 64 - else: - raise ValueError(f"Invalid precision: {precision}") - if in_features % self.group_size != 0: + padded_in = ((in_features + 127) // 128) * 128 + padded_out = ((out_features + 127) // 128) * 128 + self._padded_in_features = padded_in + self._padded_out_features = padded_out + self._needs_k_pad = padded_in != in_features + self._needs_n_strip = padded_out != out_features + + if in_features % 2 != 0: + raise ValueError(f"in_features must be divisible by 2, got {in_features}.") + if padded_in % self.group_size != 0: raise ValueError( - f"in_features must be divisible by group_size={self.group_size} for {precision}, got {in_features}." - ) + f"Padded in_features ({padded_in}) must be divisible by group_size={self.group_size} " + f"for {precision}, got logical in_features={in_features}.") self.qweight = nn.Parameter( - torch.empty(out_features, in_features // 2, dtype=torch.int8, device=device), + torch.empty(padded_out, padded_in // 2, dtype=torch.int8, device=device), requires_grad=False, ) - self.bias = (nn.Parameter(torch.empty(out_features, dtype=torch_dtype, device=device), + self.bias = (nn.Parameter(torch.empty(padded_out, dtype=torch_dtype, device=device), requires_grad=True) if bias else None) self.wscales = nn.Parameter( torch.empty( - in_features // self.group_size, - out_features, + padded_in // self.group_size, + padded_out, dtype=torch_dtype if precision == "int4" else torch.float8_e4m3fn, device=device, ), requires_grad=False, ) self.smooth_factor = nn.Parameter( - torch.empty(in_features, dtype=torch_dtype, device=device), + torch.empty(padded_in, dtype=torch_dtype, device=device), requires_grad=False, ) self.smooth_factor_orig = nn.Parameter( - torch.empty(in_features, dtype=torch_dtype, device=device), + torch.empty(padded_in, dtype=torch_dtype, device=device), requires_grad=False, ) - self.proj_down = nn.Parameter(torch.empty(in_features, rank, dtype=torch_dtype, device=device)) - self.proj_up = nn.Parameter(torch.empty(out_features, rank, dtype=torch_dtype, device=device)) + self.proj_down = nn.Parameter(torch.empty(padded_in, rank, dtype=torch_dtype, device=device)) + self.proj_up = nn.Parameter(torch.empty(padded_out, rank, dtype=torch_dtype, device=device)) if precision == "nvfp4": self.wcscales = nn.Parameter( - torch.ones(out_features, dtype=torch_dtype, device=device), + torch.ones(padded_out, dtype=torch_dtype, device=device), requires_grad=False, ) self.wtscale = 1.0 @@ -173,10 +182,14 @@ def _forward_plain(self, x: torch.Tensor, output: torch.Tensor | None = None) -> for extent in leading_shape: token_count *= extent x = x.reshape(token_count, channels) + + if self._needs_k_pad: + x = F.pad(x, (0, self._padded_in_features - self.in_features)) + quantized_x, ascales, lora_act_out = self.quantize(x) use_direct_output = output is not None and output.shape == ( quantized_x.shape[0], - self.out_features, + self._padded_out_features, ) padded_output = self.forward_quant( quantized_x, @@ -186,6 +199,9 @@ def _forward_plain(self, x: torch.Tensor, output: torch.Tensor | None = None) -> ) logical_output = padded_output[:token_count] + if self._needs_n_strip: + logical_output = logical_output[:, :self.out_features] + if output is not None and not use_direct_output: expected_shape = (*leading_shape, self.out_features) if output.shape != expected_shape: @@ -263,9 +279,13 @@ def forward_quant( return result def __repr__(self) -> str: + extra = "" + if self._needs_k_pad or self._needs_n_strip: + extra = (f", padded_in={self._padded_in_features}" + f", padded_out={self._padded_out_features}") return (f"SVDQW4A4Linear(in_features={self.in_features}, out_features={self.out_features}, " f"rank={self.rank}, precision={self.precision}, act_unsigned={self.act_unsigned}, " - f"runtime_kernel={self.runtime_kernel})") + f"runtime_kernel={self.runtime_kernel}{extra})") __all__ = ["SVDQW4A4Linear"] diff --git a/src/cache_dit/quantization/svdquant/quantizer.py b/src/cache_dit/quantization/svdquant/quantizer.py index f07ba6a4..0a38a456 100644 --- a/src/cache_dit/quantization/svdquant/quantizer.py +++ b/src/cache_dit/quantization/svdquant/quantizer.py @@ -7,6 +7,7 @@ import torch from torch import nn +import torch.nn.functional as F from .dtensor import SVDQShardSpec from .dtensor import SVDQW4A4ShardLinear @@ -93,19 +94,17 @@ def validate_linear_geometry( raise ValueError( f"{precision.upper()} SVDQuant requires in_features divisible by {group_size}, got {in_features}." ) - if in_features % 128 != 0: - raise ValueError( - f"The migrated W4A4 packer/runtime requires in_features divisible by 128, got {in_features}.") - if out_features % 128 != 0: - raise ValueError( - f"The migrated W4A4 packer/runtime requires out_features divisible by 128, got {out_features}." - ) if rank < 0: raise ValueError(f"rank must be non-negative, got {rank}.") if rank != 0 and rank % 16 != 0: raise ValueError(f"The migrated W4A4 runtime requires rank 0 or a multiple of 16, got {rank}.") +def _pad_to_multiple_128(x: int) -> int: + """Round up to the nearest multiple of 128 for the SVDQ W4A4 packer/runtime tile.""" + return ((x + 127) // 128) * 128 + + def _normalize_dtype(torch_dtype: torch.dtype | None, device: torch.device | str) -> torch.dtype: if torch_dtype is not None: if torch_dtype not in (torch.float16, torch.bfloat16): @@ -972,6 +971,9 @@ def _quantize_from_smooth_scale( smooth_scale_orig = _repair_invalid_scale(smooth_scale_orig.to(device=device, dtype=torch_dtype)) smoothed_weight = weight.to(dtype=math_dtype) * smooth_scale.to(dtype=math_dtype).unsqueeze(0) + # NOTE: The SVD decomposition is performed before the padding step, so the low-rank factors + # are computed for the original weight shape and make it strictly match the math precision. + # The padding below is applied afterwards to ensure compatibility with hardware constraints. lowrank_down, lowrank_up, residual = decompose_lowrank_residual( smoothed_weight, rank, @@ -998,6 +1000,36 @@ def _quantize_from_smooth_scale( local_bias, _ = SVDQW4A4ShardLinear.resolve_local(linear.bias) bias = None if local_bias is None else local_bias.detach().to(device=device, dtype=torch_dtype) + + padded_in = _pad_to_multiple_128(local_in_features) + padded_out = _pad_to_multiple_128(local_out_features) + needs_pad = padded_in != local_in_features or padded_out != local_out_features + + if needs_pad: + residual = F.pad(residual, + (0, padded_in - local_in_features, 0, padded_out - local_out_features)) + smooth_scale = F.pad(smooth_scale, (0, padded_in - local_in_features), value=1.0) + smooth_scale_orig = F.pad(smooth_scale_orig, (0, padded_in - local_in_features), value=1.0) + if bias is not None: + bias = F.pad(bias, (0, padded_out - local_out_features)) + if rank > 0: + lowrank_down = F.pad(lowrank_down, (0, padded_in - local_in_features)) + lowrank_up = F.pad(lowrank_up, (0, 0, 0, padded_out - local_out_features)) + if channel_scales is not None: + channel_scales = F.pad( + channel_scales, + (0, 0, 0, 0, 0, 0, 0, padded_out - local_out_features), + value=1.0, + ) + group_scales_pad_out = padded_out - local_out_features + group_scales_pad_in = (padded_in - local_in_features) // _resolve_svdq_group_size(precision) + if group_scales_pad_out > 0 or group_scales_pad_in > 0: + group_scales = F.pad( + group_scales, + (0, 0, 0, group_scales_pad_in, 0, 0, 0, group_scales_pad_out), + value=1.0, + ) + raw_state_dict = export_raw_svdq_w4a4_state_dict( residual, scale=channel_scales if channel_scales is not None else group_scales, @@ -1010,8 +1042,8 @@ def _quantize_from_smooth_scale( ) module_state_dict = adapt_svdq_module_state_dict( raw_state_dict, - in_features=local_in_features, - out_features=local_out_features, + in_features=padded_in, + out_features=padded_out, rank=rank, precision=_normalize_svdq_precision(precision), torch_dtype=torch_dtype, diff --git a/tests/quantization/test_svdquant_padding.py b/tests/quantization/test_svdquant_padding.py new file mode 100644 index 00000000..51980c73 --- /dev/null +++ b/tests/quantization/test_svdquant_padding.py @@ -0,0 +1,262 @@ +"""Operator-level correctness tests for SVDQ NVFP4 padding support. + +Tests SVDQW4A4Linear with non-128-multiple in_features / out_features to verify that the padding/de- +padding logic produces correct results (within quantization tolerance) compared to the original FP16 +linear layer. +""" + +from __future__ import annotations + +import pytest +import torch + +from cache_dit.quantization.svdquant.linear import SVDQW4A4Linear +from cache_dit.quantization.svdquant.quantizer import _quantize_from_smooth_scale +from cache_dit.quantization.svdquant.quantizer import _pad_to_multiple_128 +from cache_dit.quantization.svdquant.quantizer import validate_linear_geometry + +DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu" +TORCH_DTYPE = torch.bfloat16 + + +def _create_linear(in_features, out_features, bias=True): + return torch.nn.Linear(in_features, out_features, bias=bias, dtype=TORCH_DTYPE, device=DEVICE) + + +def _quantize_linear(linear, rank=0, precision="nvfp4"): + """Run a minimal smooth-scale quantize for testing, returns SVDQW4A4Linear.""" + smooth = torch.ones(linear.in_features, dtype=TORCH_DTYPE, device=DEVICE) + smooth_orig = smooth.clone() + result = _quantize_from_smooth_scale( + linear, + smooth, + smooth_scale_orig=smooth_orig, + rank=rank, + precision=precision, + act_unsigned=False, + torch_dtype=TORCH_DTYPE, + device=DEVICE, + calibrate_precision="low", + runtime_kernel="v1", + ) + return result + + +def _assert_output_sane(quant_output: torch.Tensor): + """Sanity checks: no NaN/Inf, output is finite, max abs is reasonable (<100).""" + assert not torch.isnan(quant_output).any(), "Quantized output contains NaN" + assert not torch.isinf(quant_output).any(), "Quantized output contains Inf" + max_abs = quant_output.abs().max().item() + assert max_abs < 100.0, f"Output max abs {max_abs:.2f} exceeds sanity limit" + + +class TestSVDQPadding: + """Test SVDQ NVFP4 with non-128-multiple dimensions.""" + + @pytest.mark.parametrize( + "in_features,out_features", + [ + (3360, 3360), # Boogu-Image typical dims + (3360, 1280), # mixed alignment + (1280, 3360), # mixed alignment (reverse) + (512, 512), # already aligned (regression test) + (320, 640), # small non-aligned + (640, 320), # small non-aligned (reverse) + ]) + def test_padded_quantize_output_shape(self, in_features, out_features): + """SVDQW4A4Linear forward output shape matches logical dims.""" + linear = _create_linear(in_features, out_features, bias=False) + quantized = _quantize_linear(linear, rank=0) + assert quantized.in_features == in_features + assert quantized.out_features == out_features + + x = torch.randn(2, 3, in_features, dtype=TORCH_DTYPE, device=DEVICE) + output = quantized(x) + assert output.shape == (2, 3, out_features) + _assert_output_sane(output) + + @pytest.mark.parametrize("in_features,out_features,rank", [ + (3360, 3360, 0), + (320, 640, 0), + (3360, 3360, 32), + (320, 640, 32), + ]) + def test_padded_quantize_no_nan_inf(self, in_features, out_features, rank): + """Quantized output from padded layers contains no NaN/Inf.""" + linear = _create_linear(in_features, out_features, bias=False) + quantized = _quantize_linear(linear, rank=rank) + + x = torch.randn(1, 17, in_features, dtype=TORCH_DTYPE, device=DEVICE) + with torch.inference_mode(): + output = quantized(x) + _assert_output_sane(output) + + @pytest.mark.parametrize("in_features,out_features", [ + (3360, 3360), + (320, 640), + ]) + def test_padded_quantize_output_reproducible(self, in_features, out_features): + """Same input twice produces identical quantized output.""" + linear = _create_linear(in_features, out_features, bias=False) + quantized = _quantize_linear(linear, rank=0) + + x = torch.randn(1, 17, in_features, dtype=TORCH_DTYPE, device=DEVICE) + with torch.inference_mode(): + out1 = quantized(x) + out2 = quantized(x) + assert torch.equal(out1, out2) + + @pytest.mark.parametrize("in_features,out_features", [ + (512, 512), + (512, 1024), + ]) + def test_already_aligned_no_overhead(self, in_features, out_features): + """Already-128-aligned layers have zero-overhead padding path.""" + linear = _create_linear(in_features, out_features, bias=False) + quantized = _quantize_linear(linear, rank=0) + + assert not quantized._needs_k_pad + assert not quantized._needs_n_strip + assert quantized._padded_in_features == in_features + assert quantized._padded_out_features == out_features + + @pytest.mark.parametrize("in_features,out_features", [ + (512, 512), + (512, 1024), + ]) + def test_aligned_output_with_and_without_pad_same(self, in_features, out_features): + """Same already-aligned layer gives identical output via forward or manual path.""" + linear = _create_linear(in_features, out_features, bias=False) + quantized = _quantize_linear(linear, rank=0) + + x = torch.randn(1, 17, in_features, dtype=TORCH_DTYPE, device=DEVICE).reshape(17, in_features) + with torch.inference_mode(): + # Through normal forward (with potential K-pad + N-strip which are no-ops) + forward_out = quantized(x.unsqueeze(0).unsqueeze(0)) + + # Through raw quantize + forward_quant + strip + quantized_x, ascales, lora_act = quantized.quantize(x) + padded_out = quantized.forward_quant(quantized_x, ascales, lora_act) + raw_out = padded_out[:x.shape[0]] + raw_out = raw_out[:, :out_features] + + assert torch.equal(forward_out.squeeze(0).squeeze(0), raw_out) + + def test_pad_multiple_computation(self): + """_pad_to_multiple_128 round-up logic.""" + assert _pad_to_multiple_128(1) == 128 + assert _pad_to_multiple_128(127) == 128 + assert _pad_to_multiple_128(128) == 128 + assert _pad_to_multiple_128(129) == 256 + assert _pad_to_multiple_128(3360) == 3456 + assert _pad_to_multiple_128(0) == 0 + + def test_validate_linear_geometry_allows_non_128_multiple(self): + """validate_linear_geometry no longer rejects non-128-multiple dims.""" + validate_linear_geometry(3360, 3360, rank=0, precision="nvfp4") + validate_linear_geometry(3360, 3360, rank=32, precision="nvfp4") + validate_linear_geometry(320, 640, rank=0, precision="int4") + + def test_validate_linear_geometry_still_rejects_bad_group_size(self): + """validate_linear_geometry still rejects in_features not divisible by group_size.""" + with pytest.raises(ValueError, match="in_features divisible by 16"): + validate_linear_geometry(15, 128, rank=0, precision="nvfp4") + + def test_validate_linear_geometry_still_rejects_bad_rank(self): + """validate_linear_geometry still rejects rank not divisible by 16.""" + with pytest.raises(ValueError, match="multiple of 16"): + validate_linear_geometry(128, 128, rank=5, precision="nvfp4") + + def test_svdq_linear_repr_shows_logical_dims(self): + """SVDQW4A4Linear.__repr__ shows logical dims with padded info.""" + linear = SVDQW4A4Linear(3360, 3360, rank=0, precision="nvfp4", torch_dtype=TORCH_DTYPE) + rep = repr(linear) + assert "in_features=3360" in rep + assert "out_features=3360" in rep + assert "padded_in=3456" in rep + assert "padded_out=3456" in rep + + def test_svdq_linear_repr_aligned_no_padded_info(self): + """SVDQW4A4Linear.__repr__ omits padded info when already aligned.""" + linear = SVDQW4A4Linear(512, 512, rank=0, precision="int4", torch_dtype=TORCH_DTYPE) + rep = repr(linear) + assert "in_features=512" in rep + assert "padded_in" not in rep + + def test_padded_weights_produce_zero_output_on_zero_input(self): + """Zero input to padded SVDQW4A4Linear produces zero output.""" + in_features, out_features = 3360, 3360 + linear = _create_linear(in_features, out_features, bias=False) + quantized = _quantize_linear(linear, rank=0) + + x = torch.zeros(1, 1, in_features, dtype=TORCH_DTYPE, device=DEVICE) + with torch.inference_mode(): + output = quantized(x) + assert output.abs().max().item() < 1e-3, "Zero input should produce near-zero output" + + +class TestSVDQPaddingINT4: + """Test SVDQ INT4 with non-128-multiple dimensions (group_size=64).""" + + @pytest.mark.parametrize( + "in_features,out_features", + [ + (320, 640), # non-128-multiple but 64-multiple + (192, 384), # smaller non-aligned + (512, 512), # already aligned (regression) + ]) + def test_int4_quantize_output_shape(self, in_features, out_features): + """INT4 SVDQW4A4Linear forward output shape matches logical dims.""" + linear = _create_linear(in_features, out_features, bias=False) + quantized = _quantize_linear(linear, rank=0, precision="int4") + assert quantized.in_features == in_features + assert quantized.out_features == out_features + assert quantized.precision == "int4" + + x = torch.randn(2, 3, in_features, dtype=TORCH_DTYPE, device=DEVICE) + output = quantized(x) + assert output.shape == (2, 3, out_features) + _assert_output_sane(output) + + @pytest.mark.parametrize("in_features,out_features,rank", [ + (320, 640, 0), + (320, 640, 32), + (192, 384, 0), + ]) + def test_int4_quantize_no_nan_inf(self, in_features, out_features, rank): + """INT4 quantized output from padded layers contains no NaN/Inf.""" + linear = _create_linear(in_features, out_features, bias=False) + quantized = _quantize_linear(linear, rank=rank, precision="int4") + + x = torch.randn(1, 17, in_features, dtype=TORCH_DTYPE, device=DEVICE) + with torch.inference_mode(): + output = quantized(x) + _assert_output_sane(output) + + def test_int4_already_aligned_no_overhead(self): + """Already-128-aligned INT4 layers have zero-overhead padding path.""" + linear = _create_linear(512, 512, bias=False) + quantized = _quantize_linear(linear, rank=0, precision="int4") + + assert not quantized._needs_k_pad + assert not quantized._needs_n_strip + + def test_int4_zero_input_produces_zero_output(self): + """Zero input to padded INT4 SVDQW4A4Linear produces zero output.""" + linear = _create_linear(320, 640, bias=False) + quantized = _quantize_linear(linear, rank=0, precision="int4") + + x = torch.zeros(1, 1, 320, dtype=TORCH_DTYPE, device=DEVICE) + with torch.inference_mode(): + output = quantized(x) + assert output.abs().max().item() < 1e-3 + + def test_int4_validate_geometry_allows_non_128_multiple(self): + """INT4 validate_linear_geometry allows 64-multiples that aren't 128-multiples.""" + validate_linear_geometry(320, 640, rank=0, precision="int4") + validate_linear_geometry(192, 384, rank=16, precision="int4") + + def test_int4_validate_geometry_rejects_non_64_multiple(self): + """INT4 validate_linear_geometry still rejects in_features not divisible by 64.""" + with pytest.raises(ValueError, match="in_features divisible by 64"): + validate_linear_geometry(100, 256, rank=0, precision="int4")