w4a4: pad svdq linear to multiple of 128

src/cache_dit/quantization/svdquant/fused.py

@@ -53,29 +53,30 @@ 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,
     )
   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,25 +165,25 @@ 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,
     )
   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)
 
 

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,58 +62,66 @@ 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
     self.precision = precision
     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"]

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,