131 lines
6.4 KiB
Diff
131 lines
6.4 KiB
Diff
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)
|