"""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_quantize_swiglu_nvfp4 packed_x = vortex_quantize_swiglu_nvfp4(gate_up) 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)