diff --git a/src/h3_blackwell_runtime/nvfp4.py b/src/h3_blackwell_runtime/nvfp4.py index a479c87..2c2d0ef 100644 --- a/src/h3_blackwell_runtime/nvfp4.py +++ b/src/h3_blackwell_runtime/nvfp4.py @@ -1,6 +1,7 @@ """Standalone Blackwell NVFP4 linear adapter for Comfy-format checkpoints.""" import json +import os from dataclasses import dataclass import torch @@ -33,7 +34,7 @@ def parse_quant_sidecar(sidecar: torch.Tensor) -> dict: 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): + 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.") @@ -45,6 +46,7 @@ class Nvfp4Linear(DynamicLoraMixin, nn.Module): 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) @@ -88,13 +90,91 @@ class Nvfp4Linear(DynamicLoraMixin, nn.Module): else: if x.dtype == torch.float32: raise ValueError("Quantized NVFP4 activation GEMM requires FP16 or BF16 activations.") - packed_x = QuantizedTensor.from_float(flat_x, "TensorCoreNVFP4Layout") - output = functional.linear(packed_x, packed_weight, bias)[:flat_x.shape[0], :self.out_features] + 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 load_nvfp4_linear(tensors: dict[str, torch.Tensor], prefix: str, *, output_dtype=torch.bfloat16) -> Nvfp4Linear: + 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]) @@ -110,4 +190,4 @@ def load_nvfp4_linear(tensors: dict[str, torch.Tensor], prefix: str, *, output_d in_features=in_features, out_features=weight.shape[0], ) - return Nvfp4Linear(packed, output_dtype=output_dtype) + return Nvfp4Linear(packed, output_dtype=output_dtype, role=role)