89 lines
3.8 KiB
Python
89 lines
3.8 KiB
Python
|
"""Standalone Blackwell NVFP4 linear adapter for Comfy-format checkpoints."""
|
||
|
|
|
||
|
|
import json
|
||
|
|
from dataclasses import dataclass
|
||
|
|
|
||
|
|
import torch
|
||
|
|
import torch.nn.functional as functional
|
||
|
|
from torch import nn
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass(frozen=True)
|
||
|
|
class Nvfp4LinearTensors:
|
||
|
|
weight: torch.Tensor
|
||
|
|
weight_scale: torch.Tensor
|
||
|
|
weight_scale_2: torch.Tensor
|
||
|
|
bias: torch.Tensor | None
|
||
|
|
in_features: int
|
||
|
|
out_features: int
|
||
|
|
|
||
|
|
|
||
|
|
def parse_quant_sidecar(sidecar: torch.Tensor) -> dict:
|
||
|
|
"""Validate the small JSON descriptor stored beside each packed weight."""
|
||
|
|
metadata = json.loads(bytes(sidecar.cpu().tolist()))
|
||
|
|
if metadata.get("format") != "nvfp4":
|
||
|
|
raise ValueError(f"Unsupported quantization metadata: {metadata!r}")
|
||
|
|
return metadata
|
||
|
|
|
||
|
|
|
||
|
|
class Nvfp4Linear(nn.Module):
|
||
|
|
"""Execute a packed Comfy NVFP4 linear with Comfy Kitchen's CUDA 13 kernel."""
|
||
|
|
|
||
|
|
def __init__(self, tensors: Nvfp4LinearTensors, output_dtype=torch.bfloat16):
|
||
|
|
super().__init__()
|
||
|
|
if tensors.weight.dtype != torch.uint8 or tensors.weight.ndim != 2:
|
||
|
|
raise ValueError("NVFP4 weights must be a rank-2 packed uint8 tensor.")
|
||
|
|
if tensors.weight.shape != (tensors.out_features, tensors.in_features // 2):
|
||
|
|
raise ValueError("Packed NVFP4 dimensions do not match logical linear dimensions.")
|
||
|
|
if tensors.in_features % 32:
|
||
|
|
raise ValueError("Blackwell NVFP4 GEMM requires input width divisible by 32.")
|
||
|
|
|
||
|
|
self.in_features = tensors.in_features
|
||
|
|
self.out_features = tensors.out_features
|
||
|
|
self.output_dtype = output_dtype
|
||
|
|
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)
|
||
|
|
|
||
|
|
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, 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(
|
||
|
|
self.weight,
|
||
|
|
"TensorCoreNVFP4Layout",
|
||
|
|
TensorCoreNVFP4Layout.Params(
|
||
|
|
scale=self.weight_scale_2,
|
||
|
|
block_scale=self.weight_scale,
|
||
|
|
orig_dtype=self.output_dtype,
|
||
|
|
orig_shape=(self.out_features, self.in_features),
|
||
|
|
),
|
||
|
|
)
|
||
|
|
output = functional.linear(packed_x, packed_weight, self.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])
|
||
|
|
weight = tensors[f"{prefix}.weight"]
|
||
|
|
in_features = weight.shape[1] * 2
|
||
|
|
packed = Nvfp4LinearTensors(
|
||
|
|
weight=weight,
|
||
|
|
weight_scale=tensors[f"{prefix}.weight_scale"],
|
||
|
|
weight_scale_2=tensors[f"{prefix}.weight_scale_2"],
|
||
|
|
bias=tensors.get(f"{prefix}.bias"),
|
||
|
|
in_features=in_features,
|
||
|
|
out_features=weight.shape[0],
|
||
|
|
)
|
||
|
|
return Nvfp4Linear(packed, output_dtype=output_dtype)
|