diff --git a/src/h3_blackwell_runtime/checkpoint.py b/src/h3_blackwell_runtime/checkpoint.py index 0019b96..b5aee5e 100644 --- a/src/h3_blackwell_runtime/checkpoint.py +++ b/src/h3_blackwell_runtime/checkpoint.py @@ -22,7 +22,7 @@ class H3Checkpoint: return value.to(dtype=dtype) if dtype is not None else value def nvfp4_linear(self, prefix: str, *, output_dtype=torch.bfloat16) -> Nvfp4Linear: - names = ("comfy_quant", "weight", "weight_scale", "weight_scale_2", "bias") + names = ("comfy_quant", "weight", "weight_scale", "weight_scale_2", "bias", "pre_quant_scale") tensors = {} from safetensors import safe_open diff --git a/src/h3_blackwell_runtime/nvfp4.py b/src/h3_blackwell_runtime/nvfp4.py index 5c7d458..4587cd6 100644 --- a/src/h3_blackwell_runtime/nvfp4.py +++ b/src/h3_blackwell_runtime/nvfp4.py @@ -14,6 +14,8 @@ class Nvfp4LinearTensors: weight_scale: torch.Tensor weight_scale_2: torch.Tensor bias: torch.Tensor | None + pre_quant_scale: torch.Tensor | None + full_precision_matrix_mult: bool in_features: int out_features: int @@ -41,23 +43,17 @@ class Nvfp4Linear(nn.Module): self.in_features = tensors.in_features self.out_features = tensors.out_features self.output_dtype = output_dtype + self.full_precision_matrix_mult = tensors.full_precision_matrix_mult self.register_buffer("weight", tensors.weight.contiguous(), persistent=False) self.register_buffer("weight_scale", tensors.weight_scale.view(torch.float8_e4m3fn).contiguous(), persistent=False) self.register_buffer("weight_scale_2", tensors.weight_scale_2.to(torch.float32).contiguous(), persistent=False) self.register_buffer("bias", tensors.bias.contiguous() if tensors.bias is not None else None, persistent=False) + self.register_buffer("pre_quant_scale", tensors.pre_quant_scale.contiguous() if tensors.pre_quant_scale is not None else None, persistent=False) - def forward(self, x: torch.Tensor) -> torch.Tensor: - if x.shape[-1] != self.in_features: - raise ValueError(f"Expected feature width {self.in_features}, received {x.shape[-1]}.") - if x.dtype not in (torch.float16, torch.bfloat16): - raise ValueError("NVFP4 linear accepts FP16 or BF16 activations.") - + def _packed_weight(self): from comfy_kitchen.tensor import QuantizedTensor, TensorCoreNVFP4Layout - original_shape = x.shape[:-1] - flat_x = x.reshape(-1, self.in_features).contiguous() - packed_x = QuantizedTensor.from_float(flat_x, "TensorCoreNVFP4Layout") - packed_weight = QuantizedTensor( + return QuantizedTensor( self.weight, "TensorCoreNVFP4Layout", TensorCoreNVFP4Layout.Params( @@ -67,14 +63,34 @@ class Nvfp4Linear(nn.Module): orig_shape=(self.out_features, self.in_features), ), ) - output = functional.linear(packed_x, packed_weight, self.bias) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if x.shape[-1] != self.in_features: + raise ValueError(f"Expected feature width {self.in_features}, received {x.shape[-1]}.") + if x.dtype not in (torch.float16, torch.bfloat16): + raise ValueError("NVFP4 linear accepts FP16 or BF16 activations.") + + from comfy_kitchen.tensor import QuantizedTensor + + original_shape = x.shape[:-1] + flat_x = x.reshape(-1, self.in_features).contiguous() + if self.pre_quant_scale is not None: + flat_x = flat_x * self.pre_quant_scale.to(flat_x) + packed_weight = self._packed_weight() + bias = self.bias.to(flat_x) if self.bias is not None else None + if self.full_precision_matrix_mult: + weight = packed_weight.dequantize().to(flat_x) + output = functional.linear(flat_x, weight, bias) + return output.reshape(*original_shape, self.out_features) + packed_x = QuantizedTensor.from_float(flat_x, "TensorCoreNVFP4Layout") + output = functional.linear(packed_x, packed_weight, bias) return output[:flat_x.shape[0], :self.out_features].reshape(*original_shape, self.out_features) def load_nvfp4_linear(tensors: dict[str, torch.Tensor], prefix: str, *, output_dtype=torch.bfloat16) -> Nvfp4Linear: """Load one Comfy-format NVFP4 linear from a safetensors tensor mapping.""" sidecar_key = f"{prefix}.comfy_quant" - parse_quant_sidecar(tensors[sidecar_key]) + metadata = parse_quant_sidecar(tensors[sidecar_key]) weight = tensors[f"{prefix}.weight"] in_features = weight.shape[1] * 2 packed = Nvfp4LinearTensors( @@ -82,6 +98,8 @@ def load_nvfp4_linear(tensors: dict[str, torch.Tensor], prefix: str, *, output_d weight_scale=tensors[f"{prefix}.weight_scale"], weight_scale_2=tensors[f"{prefix}.weight_scale_2"], bias=tensors.get(f"{prefix}.bias"), + pre_quant_scale=tensors.get(f"{prefix}.pre_quant_scale"), + full_precision_matrix_mult=metadata.get("full_precision_matrix_mult", False), in_features=in_features, out_features=weight.shape[0], )