200 lines
9 KiB
Python
200 lines
9 KiB
Python
"""Standalone Blackwell NVFP4 linear adapter for Comfy-format checkpoints."""
|
|
|
|
import json
|
|
import os
|
|
from dataclasses import dataclass
|
|
|
|
import torch
|
|
import torch.nn.functional as functional
|
|
from torch import nn
|
|
|
|
from .lora import DynamicLoraMixin
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Nvfp4LinearTensors:
|
|
weight: torch.Tensor
|
|
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
|
|
|
|
|
|
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(DynamicLoraMixin, nn.Module):
|
|
"""Execute a packed Comfy NVFP4 linear with Comfy Kitchen's CUDA 13 kernel."""
|
|
|
|
def __init__(self, tensors: Nvfp4LinearTensors, output_dtype=torch.bfloat16, *, role: str = "other"):
|
|
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.role = role
|
|
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)
|
|
self._init_dynamic_lora()
|
|
|
|
def _packed_weight(self):
|
|
from comfy_kitchen.tensor import QuantizedTensor, TensorCoreNVFP4Layout
|
|
|
|
return 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),
|
|
),
|
|
)
|
|
|
|
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, torch.float32):
|
|
raise ValueError("NVFP4 linear accepts FP16, BF16, or FP32 activations.")
|
|
|
|
from comfy_kitchen.tensor import QuantizedTensor
|
|
|
|
original_shape = x.shape[:-1]
|
|
original_x = x
|
|
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)
|
|
else:
|
|
if x.dtype == torch.float32:
|
|
raise ValueError("Quantized NVFP4 activation GEMM requires FP16 or BF16 activations.")
|
|
ring_output = None
|
|
if self.role == "h3_attn_qkv" and os.getenv("H3_CUTE_QKV_RING", "").lower() in {
|
|
"1", "true", "yes", "on",
|
|
}:
|
|
from .cute_qkv_ring import qkv_ring_linear
|
|
|
|
ring_output = qkv_ring_linear(self, flat_x)
|
|
if ring_output is not None:
|
|
output = ring_output
|
|
elif os.getenv("H3_NVFP4_SCALE_BACKEND", "torch").lower() == "torch":
|
|
packed_x = QuantizedTensor.from_float(flat_x, "TensorCoreNVFP4Layout")
|
|
output = functional.linear(packed_x, packed_weight, bias)[:flat_x.shape[0], :self.out_features]
|
|
else:
|
|
from .nvfp4_quant import vortex_quantize_nvfp4
|
|
|
|
packed_x = vortex_quantize_nvfp4(flat_x)
|
|
output = functional.linear(packed_x, packed_weight, bias)[:flat_x.shape[0], :self.out_features]
|
|
base = output.reshape(*original_shape, self.out_features)
|
|
return self._apply_lora(original_x, base)
|
|
|
|
def forward_modulated(
|
|
self,
|
|
x: torch.Tensor,
|
|
shift: torch.Tensor,
|
|
scale: torch.Tensor,
|
|
row_index: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
"""Fuse exact H3 modulation into activation packing for QKV or FC1."""
|
|
if self.role not in {"h3_attn_qkv", "h3_mlp_fc1"}:
|
|
raise ValueError("modulated NVFP4 dispatch is restricted to H3 QKV and FC1")
|
|
if (
|
|
self.full_precision_matrix_mult
|
|
or self.pre_quant_scale is not None
|
|
or self.active_lora is not None
|
|
or self.lora_strength != 0.0
|
|
or torch.is_grad_enabled()
|
|
or x.requires_grad
|
|
):
|
|
from .h3_fusion import fused_modulate_
|
|
|
|
return self(fused_modulate_(x, shift, scale, row_index))
|
|
if x.shape != (row_index.numel(), self.in_features):
|
|
raise ValueError("modulated NVFP4 input and row index have incompatible shapes")
|
|
if x.dtype != torch.bfloat16 or not x.is_cuda or not x.is_contiguous():
|
|
raise ValueError("modulated NVFP4 dispatch requires contiguous CUDA BF16 input")
|
|
|
|
from .nvfp4_quant import vortex_quantize_modulated_nvfp4
|
|
|
|
packed_x = vortex_quantize_modulated_nvfp4(x, shift, scale, row_index)
|
|
bias = self.bias.to(x) if self.bias is not None else None
|
|
output = functional.linear(packed_x, self._packed_weight(), bias)[
|
|
: x.shape[0], : self.out_features
|
|
]
|
|
return output
|
|
|
|
def forward_swiglu(self, gate_up: torch.Tensor) -> torch.Tensor:
|
|
"""Fuse exact BF16 SwiGLU into activation packing for H3 FC2."""
|
|
if self.role != "h3_mlp_fc2":
|
|
raise ValueError("SwiGLU NVFP4 dispatch is restricted to H3 FC2")
|
|
if gate_up.shape[-1] != self.in_features * 2:
|
|
raise ValueError("SwiGLU input width must be twice the FC2 input width")
|
|
if (
|
|
self.full_precision_matrix_mult
|
|
or self.pre_quant_scale is not None
|
|
or self.active_lora is not None
|
|
or self.lora_strength != 0.0
|
|
or torch.is_grad_enabled()
|
|
or gate_up.requires_grad
|
|
):
|
|
gate, up = gate_up.chunk(2, dim=-1)
|
|
return self(torch.nn.functional.silu(gate).mul_(up))
|
|
if gate_up.dtype != torch.bfloat16 or not gate_up.is_cuda or not gate_up.is_contiguous():
|
|
raise ValueError("SwiGLU NVFP4 dispatch requires contiguous CUDA BF16 input")
|
|
|
|
from .nvfp4_quant import vortex_native_quantize_swiglu_nvfp4, wrap_native_swiglu_nvfp4
|
|
|
|
packed = vortex_native_quantize_swiglu_nvfp4(gate_up)
|
|
if os.getenv("H3_NVFP4_FC2_LT_SPLITK1", "").lower() in {"1", "true", "yes", "on"}:
|
|
from .fc2_lt import fc2_lt_linear
|
|
|
|
scheduled = fc2_lt_linear(self, gate_up, *packed)
|
|
if scheduled is not None:
|
|
return scheduled
|
|
packed_x = wrap_native_swiglu_nvfp4(gate_up, packed)
|
|
bias = self.bias.to(gate_up) if self.bias is not None else None
|
|
output = functional.linear(packed_x, self._packed_weight(), bias)[
|
|
: gate_up.shape[0], : self.out_features
|
|
]
|
|
return output
|
|
|
|
|
|
def load_nvfp4_linear(tensors: dict[str, torch.Tensor], prefix: str, *, output_dtype=torch.bfloat16, role: str = "other") -> Nvfp4Linear:
|
|
"""Load one Comfy-format NVFP4 linear from a safetensors tensor mapping."""
|
|
sidecar_key = f"{prefix}.comfy_quant"
|
|
metadata = 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"),
|
|
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],
|
|
)
|
|
return Nvfp4Linear(packed, output_dtype=output_dtype, role=role)
|