h3-blackwell-runtime/src/h3_blackwell_runtime/nvfp4.py

"""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)