daniel/h3-blackwell-runtime
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Honor NVFP4 full precision metadata

Browse source
This commit is contained in:
Daniel Maddern 2026-08-12 21:08:25 +07:00
parent 054b8d58a0
commit 0cad1dbc31
2 changed files with 31 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
src/h3_blackwell_runtime/checkpoint.py
View file

@ -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

42
src/h3_blackwell_runtime/nvfp4.py
View file

@ -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],
)