"""Portable native NVFP4 modules for the S2-Pro NVFP4 V1 runtime.""" from __future__ import annotations import hashlib import json import os import sys from dataclasses import asdict, dataclass from functools import lru_cache from pathlib import Path from typing import Any import numpy as np import torch import torch.nn.functional as F from torch import nn from torch.nn.attention import SDPBackend, sdpa_kernel from fish_speech.models.text2semantic.llama import apply_rotary_emb ROOT = Path( os.environ.get("FISH_NVFP4_ROOT", Path(__file__).resolve().parents[3]) ).resolve() NATIVE_SOURCE_ROOT = Path( os.environ.get( "FISH_NVFP4_NATIVE_ROOT", ROOT / "runtime" / "native", ) ).resolve() BUILD_ROOT = Path( os.environ.get( "FISH_NVFP4_BUILD_ROOT", ROOT / "runtime-data" / "torch-extensions", ) ).resolve() COMFY_KITCHEN_ROOT = os.environ.get("COMFY_KITCHEN_ROOT") DIRECT_SOURCE_ROOT = NATIVE_SOURCE_ROOT / "direct_w4a4_m1" SILU_PRODUCT_SOURCE_ROOT = NATIVE_SOURCE_ROOT / "silu_product_nvfp4_m1" RMSNORM_SOURCE_ROOT = NATIVE_SOURCE_ROOT / "rmsnorm_nvfp4_m1" SMALLM_SOURCE_ROOT = NATIVE_SOURCE_ROOT / "smallm_gemv" DIRECT_SOURCE_FILES = ( "direct_w4a4_m1.cpp", "direct_w4a4_m1.cu", "direct_w4a4_m1.h", ) SILU_PRODUCT_SOURCE_FILES = ( "silu_product_nvfp4_m1.cpp", "silu_product_nvfp4_m1.cu", "silu_product_nvfp4_m1.h", ) RMSNORM_SOURCE_FILES = ( "rmsnorm_nvfp4_m1.cpp", "rmsnorm_nvfp4_m1.cu", "rmsnorm_nvfp4_m1.h", ) SMALLM_SOURCE_FILES = ( "smallm_gemv.cpp", "smallm_gemv.cu", "smallm_gemv.h", ) if COMFY_KITCHEN_ROOT and COMFY_KITCHEN_ROOT not in sys.path: sys.path.insert(0, COMFY_KITCHEN_ROOT) from comfy_kitchen.tensor import QuantizedTensor, TensorCoreNVFP4Layout @dataclass class ConversionRecord: name: str in_features: int out_features: int parameters: int probe_cosine: float @dataclass class FusedMLPConversionRecord: name: str in_features: int intermediate_features: int out_features: int parameters: int probe_cosine: float @dataclass class FusedTransformerConversionRecord: name: str parameters: int wqkv_probe_cosine: float wo_probe_cosine: float mlp_probe_cosine: float def quantize_nvfp4( value: torch.Tensor, *, scale: torch.Tensor | float | None = None, ) -> QuantizedTensor: """Apply the pinned tensor-wide dynamic NVFP4 policy.""" return QuantizedTensor.from_float( value, "TensorCoreNVFP4Layout", scale=scale, ) @torch.inference_mode() def apply_nvfp4_rounding_checkpoint( model: nn.Module, checkpoint: Path | str, ) -> dict[str, Any]: """Apply a packed learned-rounding delta after ordinary model conversion.""" from safetensors.torch import load_file checkpoint_path = Path(checkpoint) if checkpoint_path.is_dir(): candidates = sorted(checkpoint_path.glob("*.safetensors")) if len(candidates) != 1: raise ValueError( f"Expected one safetensors file in {checkpoint_path}, found {len(candidates)}" ) checkpoint_path = candidates[0] tensors = load_file(str(checkpoint_path), device=str(next(model.parameters()).device)) qdata_keys = sorted(key for key in tensors if key.endswith(".qdata")) if not qdata_keys: raise ValueError(f"No packed qdata entries in {checkpoint_path}") records = [] for qdata_key in qdata_keys: module_name = qdata_key.removesuffix(".qdata") module = model.get_submodule(module_name) if not isinstance(module, NVFP4Linear): raise TypeError(f"Checkpoint target is not NVFP4Linear: {module_name}") block_key = module_name + ".block_scale" tensor_key = module_name + ".tensor_scale" if block_key not in tensors or tensor_key not in tensors: raise ValueError(f"Checkpoint lacks scales for {module_name}") checkpoint_qdata = tensors[qdata_key] checkpoint_block_scale = tensors[block_key] checkpoint_tensor_scale = tensors[tensor_key] if checkpoint_qdata.shape != module.qdata.shape: raise ValueError(f"qdata shape mismatch for {module_name}") if not torch.equal(checkpoint_block_scale, module.weight_block_scale): raise ValueError(f"Block scale mismatch for {module_name}") if not torch.equal(checkpoint_tensor_scale, module.weight_scale): raise ValueError(f"Tensor scale mismatch for {module_name}") changed_bytes = int((checkpoint_qdata != module.qdata).sum()) old_codes = torch.stack( (module.qdata >> 4, module.qdata & 0x0F), dim=-1, ) new_codes = torch.stack( (checkpoint_qdata >> 4, checkpoint_qdata & 0x0F), dim=-1, ) changed_weights = int((old_codes != new_codes).sum()) module.qdata.copy_(checkpoint_qdata) records.append( { "module": module_name, "changed_packed_bytes": changed_bytes, "changed_weights": changed_weights, "weights": module.out_features * module.in_features, } ) return { "checkpoint": str(checkpoint_path), "modules": len(records), "changed_packed_bytes": sum(row["changed_packed_bytes"] for row in records), "changed_weights": sum(row["changed_weights"] for row in records), "records": records, } def _hadamard_blocks(value: torch.Tensor, block_size: int) -> torch.Tensor: """Materialized orthonormal block Hadamard for quality prototypes.""" if block_size < 2 or block_size & (block_size - 1): raise ValueError(f"Hadamard block size must be a power of two: {block_size}") if value.shape[-1] % block_size: raise ValueError( f"Width {value.shape[-1]} is not divisible by block size {block_size}" ) original_dtype = value.dtype transformed = value.float().reshape(*value.shape[:-1], -1, block_size) stride = 1 while stride < block_size: pairs = transformed.reshape( *transformed.shape[:-1], block_size // (2 * stride), 2, stride, ) left = pairs[..., 0, :] right = pairs[..., 1, :] transformed = torch.stack((left + right, left - right), dim=-2).reshape( *transformed.shape ) stride *= 2 return (transformed.reshape_as(value) / block_size**0.5).to(original_dtype) def _direct_source_hash() -> str: digest = hashlib.sha256() for filename in DIRECT_SOURCE_FILES: digest.update((DIRECT_SOURCE_ROOT / filename).read_bytes()) return digest.hexdigest() def _source_hash(source_root: Path, filenames: tuple[str, ...]) -> str: digest = hashlib.sha256() for filename in filenames: digest.update((source_root / filename).read_bytes()) return digest.hexdigest() @lru_cache(maxsize=1) def load_direct_w4a4_m1_extension(*, verbose: bool = False) -> Any: """Build the pinned direct packed-NVFP4 M=1 primitive locally.""" from torch.utils.cpp_extension import load missing = [ filename for filename in DIRECT_SOURCE_FILES if not (DIRECT_SOURCE_ROOT / filename).is_file() ] if missing: raise RuntimeError(f"Missing pinned NVFP4 native sources: {missing}") build_root = BUILD_ROOT / "direct_w4a4_m1" build_root.mkdir(parents=True, exist_ok=True) os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "12.0") os.environ.setdefault("MAX_JOBS", "2") return load( name=f"s2_pro_direct_w4a4_m1_{_direct_source_hash()[:12]}", sources=[ str(DIRECT_SOURCE_ROOT / "direct_w4a4_m1.cpp"), str(DIRECT_SOURCE_ROOT / "direct_w4a4_m1.cu"), ], extra_cflags=["-O3"], extra_cuda_cflags=["-O3", "--use_fast_math", "-lineinfo"], extra_include_paths=[str(DIRECT_SOURCE_ROOT)], build_directory=str(build_root), with_cuda=True, verbose=verbose, is_python_module=True, ) @lru_cache(maxsize=1) def load_silu_product_nvfp4_m1_extension(*, verbose: bool = False) -> Any: """Build the pinned fused SiLU/product-to-NVFP4 M=1 primitive.""" from torch.utils.cpp_extension import load missing = [ filename for filename in SILU_PRODUCT_SOURCE_FILES if not (SILU_PRODUCT_SOURCE_ROOT / filename).is_file() ] if missing: raise RuntimeError(f"Missing pinned NVFP4 SiLU sources: {missing}") build_root = BUILD_ROOT / "silu_product_nvfp4_m1" build_root.mkdir(parents=True, exist_ok=True) os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "12.0") os.environ.setdefault("MAX_JOBS", "2") return load( name=( "s2_pro_silu_product_nvfp4_m1_" f"{_source_hash(SILU_PRODUCT_SOURCE_ROOT, SILU_PRODUCT_SOURCE_FILES)[:12]}" ), sources=[ str(SILU_PRODUCT_SOURCE_ROOT / "silu_product_nvfp4_m1.cpp"), str(SILU_PRODUCT_SOURCE_ROOT / "silu_product_nvfp4_m1.cu"), ], extra_cflags=["-O3"], extra_cuda_cflags=["-O3", "--use_fast_math", "-lineinfo"], extra_include_paths=[str(SILU_PRODUCT_SOURCE_ROOT)], build_directory=str(build_root), with_cuda=True, verbose=verbose, is_python_module=True, ) @lru_cache(maxsize=1) def load_rmsnorm_nvfp4_m1_extension(*, verbose: bool = False) -> Any: """Build the pinned fused Fish-compatible RMSNorm-to-NVFP4 M=1 primitive.""" from torch.utils.cpp_extension import load missing = [ filename for filename in RMSNORM_SOURCE_FILES if not (RMSNORM_SOURCE_ROOT / filename).is_file() ] if missing: raise RuntimeError(f"Missing pinned NVFP4 RMSNorm sources: {missing}") build_root = BUILD_ROOT / "rmsnorm_nvfp4_m1" build_root.mkdir(parents=True, exist_ok=True) os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "12.0") os.environ.setdefault("MAX_JOBS", "2") return load( name=( "s2_pro_rmsnorm_nvfp4_m1_" f"{_source_hash(RMSNORM_SOURCE_ROOT, RMSNORM_SOURCE_FILES)[:12]}" ), sources=[ str(RMSNORM_SOURCE_ROOT / "rmsnorm_nvfp4_m1.cpp"), str(RMSNORM_SOURCE_ROOT / "rmsnorm_nvfp4_m1.cu"), ], extra_cflags=["-O3"], extra_cuda_cflags=["-O3", "--use_fast_math", "-lineinfo"], extra_include_paths=[str(RMSNORM_SOURCE_ROOT)], build_directory=str(build_root), with_cuda=True, verbose=verbose, is_python_module=True, ) @lru_cache(maxsize=1) def load_smallm_nvfp4_extension(*, verbose: bool = False) -> Any: """Build the pinned packed-weight W4A16 small-M GEMV primitive.""" from torch.utils.cpp_extension import load missing = [ filename for filename in SMALLM_SOURCE_FILES if not (SMALLM_SOURCE_ROOT / filename).is_file() ] if missing: raise RuntimeError(f"Missing pinned NVFP4 small-M sources: {missing}") build_root = BUILD_ROOT / "smallm_gemv" build_root.mkdir(parents=True, exist_ok=True) os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "12.0") os.environ.setdefault("MAX_JOBS", "2") return load( name=( "s2_pro_smallm_nvfp4_" f"{_source_hash(SMALLM_SOURCE_ROOT, SMALLM_SOURCE_FILES)[:12]}" ), sources=[ str(SMALLM_SOURCE_ROOT / "smallm_gemv.cpp"), str(SMALLM_SOURCE_ROOT / "smallm_gemv.cu"), ], extra_cflags=["-O3"], extra_cuda_cflags=["-O3", "--use_fast_math", "-lineinfo"], extra_include_paths=[str(SMALLM_SOURCE_ROOT)], build_directory=str(build_root), with_cuda=True, verbose=verbose, is_python_module=True, ) class NVFP4Linear(nn.Module): """BF16-input linear using a native packed NVFP4 weight with BF16 output.""" def __init__( self, qdata: torch.Tensor, weight_block_scale: torch.Tensor, weight_scale: torch.Tensor, *, in_features: int, out_features: int, direct_m1: bool = False, w4a16_max_m: int = 0, correction_down: torch.Tensor | None = None, correction_up: torch.Tensor | None = None, sparse_correction_indices: torch.Tensor | None = None, sparse_correction_weight: torch.Tensor | None = None, input_hadamard_block_size: int = 0, ) -> None: super().__init__() if w4a16_max_m < 0: raise ValueError("w4a16_max_m must be nonnegative") if direct_m1 and w4a16_max_m: raise ValueError("Choose only one NVFP4 M=1 backend") self.in_features = int(in_features) self.out_features = int(out_features) self.direct_m1 = bool(direct_m1) self.w4a16_max_m = int(w4a16_max_m) if input_hadamard_block_size and ( input_hadamard_block_size < 2 or input_hadamard_block_size & (input_hadamard_block_size - 1) or self.in_features % input_hadamard_block_size ): raise ValueError( "Input Hadamard block size must be a power of two that divides " f"the input width, got {input_hadamard_block_size}" ) self.input_hadamard_block_size = int(input_hadamard_block_size) self.register_buffer("qdata", qdata) self.register_buffer("weight_block_scale", weight_block_scale) self.register_buffer("weight_scale", weight_scale) if (correction_down is None) != (correction_up is None): raise ValueError("Low-rank correction requires both down and up factors") if correction_down is not None: if correction_down.dim() != 2 or correction_up.dim() != 2: raise ValueError("Low-rank correction factors must be 2D") if correction_down.shape[1] != self.in_features: raise ValueError("Low-rank down factor input width does not match") if correction_up.shape != (self.out_features, correction_down.shape[0]): raise ValueError("Low-rank up factor shape does not match") self.register_buffer("correction_down", correction_down) self.register_buffer("correction_up", correction_up) if (sparse_correction_indices is None) != (sparse_correction_weight is None): raise ValueError( "Sparse correction requires both channel indices and a weight" ) if sparse_correction_indices is not None: if ( sparse_correction_indices.dim() != 1 or sparse_correction_indices.dtype != torch.int64 ): raise ValueError("Sparse correction indices must be 1D int64") if sparse_correction_weight.shape != ( self.out_features, sparse_correction_indices.numel(), ): raise ValueError("Sparse correction weight shape does not match") if sparse_correction_weight.dtype != torch.bfloat16: raise ValueError("Sparse correction weight must be BF16") if sparse_correction_indices.numel() and ( int(sparse_correction_indices.min()) < 0 or int(sparse_correction_indices.max()) >= self.in_features ): raise ValueError("Sparse correction channel index is out of range") self.register_buffer("sparse_correction_indices", sparse_correction_indices) self.register_buffer("sparse_correction_weight", sparse_correction_weight) @classmethod @torch.inference_mode() def from_weight( cls, weight: torch.Tensor, *, direct_m1: bool = False, w4a16_max_m: int = 0, quantization_scale: torch.Tensor | float | None = None, correction_down: torch.Tensor | None = None, correction_up: torch.Tensor | None = None, sparse_correction_indices: torch.Tensor | None = None, sparse_correction_weight: torch.Tensor | None = None, input_hadamard_block_size: int = 0, ) -> "NVFP4Linear": if weight.device.type != "cuda": raise ValueError("Quantize S2-Pro weights after moving them to CUDA") if weight.dtype != torch.bfloat16 or weight.dim() != 2: raise ValueError(f"Expected a 2D BF16 source weight, got {weight.dtype} {weight.shape}") packed = quantize_nvfp4( weight.contiguous(), scale=quantization_scale, ) return cls( packed._qdata, packed._params.block_scale, packed._params.scale, in_features=weight.shape[1], out_features=weight.shape[0], direct_m1=direct_m1, w4a16_max_m=w4a16_max_m, correction_down=correction_down, correction_up=correction_up, sparse_correction_indices=sparse_correction_indices, sparse_correction_weight=sparse_correction_weight, input_hadamard_block_size=input_hadamard_block_size, ) @classmethod @torch.inference_mode() def from_linear( cls, linear: nn.Linear, *, direct_m1: bool = False, w4a16_max_m: int = 0, ) -> "NVFP4Linear": if linear.bias is not None: raise ValueError("The initial S2-Pro NVFP4 path supports bias-free linears") if linear.weight.device.type != "cuda": raise ValueError("Quantize S2-Pro linears after moving them to CUDA") if linear.weight.dtype != torch.bfloat16: raise ValueError(f"Expected BF16 source weight, got {linear.weight.dtype}") return cls.from_weight( linear.weight, direct_m1=direct_m1, w4a16_max_m=w4a16_max_m, ) def _weight_quantized_tensor(self) -> QuantizedTensor: params = TensorCoreNVFP4Layout.Params( scale=self.weight_scale, orig_dtype=torch.bfloat16, orig_shape=(self.out_features, self.in_features), block_scale=self.weight_block_scale, ) return QuantizedTensor(self.qdata, "TensorCoreNVFP4Layout", params) def _direct(self, activation: QuantizedTensor) -> torch.Tensor: params = activation._params return load_direct_w4a4_m1_extension().linear( activation._qdata, params.block_scale, params.scale, self.qdata, self.weight_block_scale, self.weight_scale, None, ) def project_packed_m1( self, qdata: torch.Tensor, block_scale: torch.Tensor, tensor_scale: torch.Tensor, ) -> torch.Tensor: """Project an already packed logical M=1 activation without wrappers.""" return load_direct_w4a4_m1_extension().linear( qdata, block_scale, tensor_scale, self.qdata, self.weight_block_scale, self.weight_scale, None, ) def project_quantized(self, activation: QuantizedTensor) -> torch.Tensor: logical_m = int(activation._params.orig_shape[0]) if self.direct_m1 and logical_m == 1: return self._direct(activation) return F.linear(activation, self._weight_quantized_tensor(), None) def forward(self, value: torch.Tensor) -> torch.Tensor: if value.shape[-1] != self.in_features: raise ValueError( f"Expected input width {self.in_features}, got {value.shape[-1]}" ) input_shape = tuple(value.shape) correction_input = value.reshape(-1, self.in_features).contiguous() flattened = correction_input if self.input_hadamard_block_size: flattened = _hadamard_blocks( flattened, self.input_hadamard_block_size, ).contiguous() if 0 < flattened.shape[0] <= self.w4a16_max_m: output = load_smallm_nvfp4_extension().linear( flattened, self.qdata, self.weight_block_scale, self.weight_scale, None, ) else: activation = quantize_nvfp4(flattened) output = self.project_quantized(activation) if self.correction_down is not None: correction = F.linear( F.linear(correction_input, self.correction_down), self.correction_up, ) output = output + correction if self.sparse_correction_indices is not None: selected = correction_input.index_select( 1, self.sparse_correction_indices, ) output = output + F.linear(selected, self.sparse_correction_weight) return output.reshape(*input_shape[:-1], self.out_features) def extra_repr(self) -> str: backend = ( f"w4a16_through_m{self.w4a16_max_m}+w4a4_above_threshold" if self.w4a16_max_m else ("direct_m1+tensorcore" if self.direct_m1 else "tensorcore") ) return ( f"in_features={self.in_features}, out_features={self.out_features}, " f"weight=NVFP4_E2M1, activation=dynamic_NVFP4, output=BF16, backend={backend}, " f"correction_rank={0 if self.correction_down is None else self.correction_down.shape[0]}, " f"sparse_correction_channels=" f"{0 if self.sparse_correction_indices is None else self.sparse_correction_indices.numel()}, " f"input_hadamard_block_size={self.input_hadamard_block_size}" ) class NVFP4FeedForward(nn.Module): """S2 SwiGLU with shared input packing and fused product packing at M=1.""" def __init__( self, w1: NVFP4Linear, w2: NVFP4Linear, w3: NVFP4Linear, ) -> None: super().__init__() self.w1 = w1 self.w2 = w2 self.w3 = w3 self.decode_backend = "w4a16" if w1.w4a16_max_m else "w4a4" @classmethod @torch.inference_mode() def from_module( cls, module: nn.Module, *, decode_backend: str = "w4a4", w4a16_max_m: int = 1, ) -> "NVFP4FeedForward": if decode_backend not in {"w4a4", "w4a16"}: raise ValueError(f"Unsupported NVFP4 MLP decode backend: {decode_backend}") for name in ("w1", "w2", "w3"): if not isinstance(getattr(module, name, None), nn.Linear): raise TypeError(f"Expected BF16 FeedForward.{name} linear") result = cls( NVFP4Linear.from_linear( module.w1, direct_m1=decode_backend == "w4a4", w4a16_max_m=(w4a16_max_m if decode_backend == "w4a16" else 0), ), NVFP4Linear.from_linear( module.w2, direct_m1=decode_backend == "w4a4", w4a16_max_m=(w4a16_max_m if decode_backend == "w4a16" else 0), ), NVFP4Linear.from_linear( module.w3, direct_m1=decode_backend == "w4a4", w4a16_max_m=(w4a16_max_m if decode_backend == "w4a16" else 0), ), ) result.decode_backend = decode_backend return result def _forward_m1(self, flattened: torch.Tensor) -> torch.Tensor: if self.decode_backend == "w4a16": return self.w2(F.silu(self.w1(flattened)) * self.w3(flattened)) packer = load_silu_product_nvfp4_m1_extension() qdata, block_scale, tensor_scale = packer.quantize_input(flattened) return self.forward_packed_m1(qdata, block_scale, tensor_scale) def forward_packed_m1( self, qdata: torch.Tensor, block_scale: torch.Tensor, tensor_scale: torch.Tensor, ) -> torch.Tensor: """Consume an already packed normalized M=1 activation.""" packer = load_silu_product_nvfp4_m1_extension() gate = self.w1.project_packed_m1(qdata, block_scale, tensor_scale) up = self.w3.project_packed_m1(qdata, block_scale, tensor_scale) product_qdata, product_block_scale, product_tensor_scale = packer.quantize( gate, up ) return self.w2.project_packed_m1( product_qdata, product_block_scale, product_tensor_scale, ) def _forward_tensorcore(self, flattened: torch.Tensor) -> torch.Tensor: activation = quantize_nvfp4(flattened) gate = self.w1.project_quantized(activation) up = self.w3.project_quantized(activation) return self.w2(F.silu(gate) * up) def forward(self, value: torch.Tensor) -> torch.Tensor: input_shape = tuple(value.shape) flattened = value.reshape(-1, input_shape[-1]).contiguous() if ( self.decode_backend == "w4a16" and flattened.shape[0] <= self.w1.w4a16_max_m ): output = self.w2(F.silu(self.w1(flattened)) * self.w3(flattened)) elif flattened.shape[0] == 1: output = self._forward_m1(flattened) else: output = self._forward_tensorcore(flattened) return output.reshape(*input_shape[:-1], self.w2.out_features) def extra_repr(self) -> str: return ( f"decode_backend={self.decode_backend}, " f"fused_silu_product_pack_m1={self.decode_backend == 'w4a4'}" ) class NVFP4TransformerBlock(nn.Module): """Slow S2 block with fused norm/MLP packing for autoregressive M=1.""" def __init__(self, source: nn.Module) -> None: super().__init__() source.attention.wqkv = NVFP4Linear.from_linear( source.attention.wqkv, direct_m1=True ) source.attention.wo = NVFP4Linear.from_linear( source.attention.wo, direct_m1=True ) self.attention = source.attention self.feed_forward = NVFP4FeedForward.from_module(source.feed_forward) self.ffn_norm = source.ffn_norm self.attention_norm = source.attention_norm self.train(source.training) def _packed_norm( self, value: torch.Tensor, norm: nn.Module, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: flattened = value.reshape(-1, value.shape[-1]).contiguous() return tuple( load_rmsnorm_nvfp4_m1_extension().quantize( flattened, norm.weight.contiguous(), float(norm.eps), ) ) def _attention_m1( self, x: torch.Tensor, packed_norm: tuple[torch.Tensor, torch.Tensor, torch.Tensor], freqs_cis: torch.Tensor, mask: torch.Tensor | None, input_pos: torch.Tensor | None, ) -> torch.Tensor: attention = self.attention bsz, seqlen, _ = x.shape qkv = attention.wqkv.project_packed_m1(*packed_norm) q_size = attention.n_head * attention.head_dim kv_size = attention.n_local_heads * attention.head_dim q, k, v = qkv.split([q_size, kv_size, kv_size], dim=-1) q = q.view(bsz, seqlen, attention.n_head, attention.head_dim) k = k.view(bsz, seqlen, attention.n_local_heads, attention.head_dim) v = v.view(bsz, seqlen, attention.n_local_heads, attention.head_dim) if attention.attention_qk_norm: q = attention.q_norm(q) k = attention.k_norm(k) q = apply_rotary_emb(q, freqs_cis) k = apply_rotary_emb(k, freqs_cis) q, k, v = (item.transpose(1, 2) for item in (q, k, v)) if attention.kv_cache is not None: k, v = attention.kv_cache.update(input_pos, k, v) repeat = attention.n_head // attention.n_local_heads k = k.repeat_interleave(repeat, dim=1) v = v.repeat_interleave(repeat, dim=1) if attention.use_sdpa: if mask is None: with sdpa_kernel(SDPBackend.FLASH_ATTENTION): y = F.scaled_dot_product_attention( q, k, v, dropout_p=attention.dropout if attention.training else 0.0, is_causal=True, ) else: y = F.scaled_dot_product_attention( q, k, v, attn_mask=mask, dropout_p=attention.dropout if attention.training else 0.0, ) else: y = attention.eq_scaled_dot_product_attention(q, k, v, attn_mask=mask) y = y.transpose(1, 2).contiguous().view(bsz, seqlen, q_size) y_flat = y.reshape(-1, q_size) output_pack = load_silu_product_nvfp4_m1_extension().quantize_input(y_flat) return attention.wo.project_packed_m1(*output_pack).reshape( bsz, seqlen, attention.dim ) def forward( self, x: torch.Tensor, freqs_cis: torch.Tensor, mask: torch.Tensor, input_pos: torch.Tensor | None = None, ) -> torch.Tensor: if x.reshape(-1, x.shape[-1]).shape[0] != 1: h = x + self.attention( self.attention_norm(x), freqs_cis, mask, input_pos ) return h + self.feed_forward(self.ffn_norm(h)) attention_pack = self._packed_norm(x, self.attention_norm) h = x + self._attention_m1(x, attention_pack, freqs_cis, mask, input_pos) ffn_pack = self._packed_norm(h, self.ffn_norm) feed_forward = self.feed_forward.forward_packed_m1(*ffn_pack) return h + feed_forward.reshape_as(h) def extra_repr(self) -> str: return "fused_rmsnorm_pack_m1=True, fused_mlp_pack_m1=True" def _selected_slow_mlp(name: str, module: nn.Module) -> bool: return ( isinstance(module, nn.Linear) and name.startswith("layers.") and ".feed_forward." in name and name.rsplit(".", 1)[-1] in {"w1", "w2", "w3"} ) def _selected_slow_transformer(name: str, module: nn.Module) -> bool: return isinstance(module, nn.Linear) and name.startswith("layers.") def _selected_slow_attention(name: str, module: nn.Module) -> bool: return ( isinstance(module, nn.Linear) and name.startswith("layers.") and ".attention." in name and name.rsplit(".", 1)[-1] in {"wqkv", "wo"} ) def _selected_slow_mlp_block(name: str, module: nn.Module) -> bool: return ( name.startswith("layers.") and name.endswith(".feed_forward") and all(isinstance(getattr(module, role, None), nn.Linear) for role in ("w1", "w2", "w3")) ) def _selected_slow_transformer_block(name: str, module: nn.Module) -> bool: parts = name.split(".") return ( len(parts) == 2 and parts[0] == "layers" and parts[1].isdigit() and hasattr(module, "attention") and hasattr(module, "feed_forward") and hasattr(module, "attention_norm") and hasattr(module, "ffn_norm") ) def _packed_weight_bytes(model: nn.Module) -> int: return sum( int(module.qdata.numel() * module.qdata.element_size()) + int(module.weight_block_scale.numel() * module.weight_block_scale.element_size()) + int(module.weight_scale.numel() * module.weight_scale.element_size()) for module in model.modules() if isinstance(module, NVFP4Linear) ) def _resolve_workspace_path(value: str | Path, *, relative_to: Path) -> Path: path = Path(value) if path.is_absolute(): return path relative_candidate = relative_to / path if relative_candidate.exists(): return relative_candidate return ROOT / path def _clip_weight_blocks(weight: torch.Tensor, ratio: float) -> torch.Tensor: if ratio == 1.0: return weight if not 0 < ratio <= 1: raise ValueError(f"NVFP4 clip ratio must be in (0,1], got {ratio}") if weight.shape[1] % 128: raise ValueError(f"Input width must be divisible by 128, got {weight.shape}") blocks = weight.float().reshape(weight.shape[0], -1, 128) threshold = blocks.abs().amax(dim=-1, keepdim=True) * ratio return torch.clamp(blocks, min=-threshold, max=threshold).reshape_as(weight) @torch.inference_mode() def _prepare_activation_scaling( model: nn.Module, report_path: Path, selected_layers: set[int], *, fold_norms: bool = True, ) -> tuple[ dict[int, torch.Tensor], dict[int, float], dict[int, dict[str, float]], dict[int, dict[str, torch.Tensor]], dict[str, Any], ]: sweep = json.loads(report_path.read_text()) calibration_values = sweep.get("calibration_reports") if calibration_values is None: calibration_values = [sweep["calibration_report"]] calibration_report_paths = [ _resolve_workspace_path(value, relative_to=report_path.parent) for value in calibration_values ] calibrations = [] activation_sets = [] for calibration_report_path in calibration_report_paths: calibration = json.loads(calibration_report_path.read_text()) with np.load(calibration_report_path.parent / calibration["arrays"]) as values: activation_sets.append(values["activations"].copy()) calibrations.append(calibration) activations = np.concatenate(activation_sets, axis=1) rows = {int(row["layer"]): row for row in sweep["results"]} missing = sorted(selected_layers - rows.keys()) if missing: raise ValueError(f"Activation-scaling report lacks layers: {missing}") clamp = float(sweep["scale_clamp"]) activation_statistic = sweep.get("activation_statistic", "absmax") scales: dict[int, torch.Tensor] = {} clip_ratios: dict[int, float] = {} quantization_scales: dict[int, dict[str, float]] = {} corrections: dict[int, dict[str, torch.Tensor]] = {} correction_tensors = None correction_tensors_value = sweep.get("correction_tensors") if correction_tensors_value is not None: from safetensors.torch import load_file correction_path = _resolve_workspace_path( correction_tensors_value, relative_to=report_path.parent, ) correction_tensors = load_file( correction_path, device=str(model.layers[0].ffn_norm.weight.device), ) records = [] for layer_index in sorted(selected_layers): block = model.layers[layer_index] gate = block.feed_forward.w1.weight up = block.feed_forward.w3.weight row = rows[layer_index] improvement = float(row["improvement_fraction_vs_unscaled"]) explicit_channel_scale = row.get("channel_scale") if explicit_channel_scale is not None: if len(explicit_channel_scale) != gate.shape[1]: raise ValueError( f"Layer {layer_index} channel scale has " f"{len(explicit_channel_scale)} values, expected {gate.shape[1]}" ) scale = torch.tensor( explicit_channel_scale, device=gate.device, dtype=torch.float32, ) if not torch.isfinite(scale).all() or (scale <= 0).any(): raise ValueError( f"Layer {layer_index} channel scale must be finite and positive" ) alpha = None clip_ratio = float(row.get("best_clip_ratio", 1.0)) elif improvement <= 0: scale = torch.ones(gate.shape[1], device=gate.device, dtype=torch.float32) alpha = None clip_ratio = 1.0 else: alpha = float(row["best_alpha"]) clip_ratio = float(row.get("best_clip_ratio", 1.0)) activation = torch.from_numpy(activations[layer_index]).to( device=gate.device, dtype=torch.float32, ) if activation_statistic == "absmax": activation_scale = activation.abs().amax(dim=0) elif activation_statistic == "abs_p99": activation_scale = torch.quantile(activation.abs(), 0.99, dim=0) elif activation_statistic == "abs_p999": activation_scale = torch.quantile(activation.abs(), 0.999, dim=0) elif activation_statistic == "mean_abs": activation_scale = activation.abs().mean(dim=0) elif activation_statistic == "rms": activation_scale = activation.square().mean(dim=0).sqrt() else: raise ValueError( f"Unknown activation statistic {activation_statistic!r}" ) activation_scale = activation_scale.clamp_min(1e-6) weight_max = torch.maximum( gate.float().abs().amax(dim=0), up.float().abs().amax(dim=0), ).clamp_min(1e-6) scale = activation_scale.pow(alpha) / weight_max.pow(1.0 - alpha) scale = scale / torch.exp(torch.mean(torch.log(scale))) scale = scale.clamp(min=1.0 / clamp, max=clamp) role_quantization_scales = row.get("weight_tensor_scales") if role_quantization_scales is not None: parsed_role_scales = { role: float(role_quantization_scales[role]) for role in ("w1", "w3") } if any( not np.isfinite(value) or value <= 0 for value in parsed_role_scales.values() ): raise ValueError( f"Layer {layer_index} weight tensor scales must be finite and positive" ) quantization_scales[layer_index] = parsed_role_scales low_rank = row.get("low_rank_correction") if low_rank is not None: if correction_tensors is None: raise ValueError( f"Layer {layer_index} has a low-rank correction without tensors" ) down_key = low_rank["down_key"] up_keys = low_rank["up_keys"] try: correction = { "down": correction_tensors[down_key], "w1_up": correction_tensors[up_keys["w1"]], "w3_up": correction_tensors[up_keys["w3"]], } except KeyError as error: raise ValueError( f"Layer {layer_index} correction tensor is missing: {error}" ) from error if any(value.dtype != torch.bfloat16 for value in correction.values()): raise ValueError( f"Layer {layer_index} correction tensors must be BF16" ) corrections[layer_index] = correction sparse = row.get("sparse_channel_correction") if sparse is not None: if correction_tensors is None: raise ValueError( f"Layer {layer_index} has a sparse correction without tensors" ) try: sparse_correction = { "sparse_indices": correction_tensors[sparse["indices_key"]], "w1_sparse_weight": correction_tensors[ sparse["weight_keys"]["w1"] ], "w3_sparse_weight": correction_tensors[ sparse["weight_keys"]["w3"] ], } except KeyError as error: raise ValueError( f"Layer {layer_index} sparse correction tensor is missing: {error}" ) from error if sparse_correction["sparse_indices"].dtype != torch.int64: raise ValueError( f"Layer {layer_index} sparse correction indices must be int64" ) if any( sparse_correction[key].dtype != torch.bfloat16 for key in ("w1_sparse_weight", "w3_sparse_weight") ): raise ValueError( f"Layer {layer_index} sparse correction weights must be BF16" ) corrections.setdefault(layer_index, {}).update(sparse_correction) input_hadamard_block_size = int(row.get("input_hadamard_block_size", 0)) if input_hadamard_block_size and ( input_hadamard_block_size < 2 or input_hadamard_block_size & (input_hadamard_block_size - 1) or gate.shape[1] % input_hadamard_block_size ): raise ValueError( f"Layer {layer_index} has invalid input Hadamard block size " f"{input_hadamard_block_size}" ) if fold_norms: norm = block.ffn_norm.weight norm.data.copy_((norm.float() / scale).to(dtype=norm.dtype)) scales[layer_index] = scale clip_ratios[layer_index] = clip_ratio records.append( { "layer": layer_index, "alpha": alpha, "calibration_improvement_fraction": improvement, "clip_ratio": clip_ratio, "scale_mode": ( "explicit_channel" if explicit_channel_scale is not None else "activation_formula" ), "correction_rank": ( None if low_rank is None else int(low_rank["rank"]) ), "sparse_correction_channels": ( None if sparse is None else int(sparse["channels"]) ), "input_hadamard_block_size": input_hadamard_block_size, "scale_min": float(scale.min()), "scale_p50": float(torch.quantile(scale, 0.50)), "scale_max": float(scale.max()), } ) return scales, clip_ratios, quantization_scales, corrections, { "sweep_report": str(report_path), "calibration_reports": [str(path) for path in calibration_report_paths], "calibration_history_sha256": [ calibration["history_sha256"] for calibration in calibrations ], "scale_clamp": clamp, "activation_statistic": activation_statistic, "correction_tensors": correction_tensors_value, "layers": records, } @torch.inference_mode() def convert_s2_pro_nvfp4( model: nn.Module, *, policy: str = "slow_mlp", direct_m1: bool = False, w4a16_max_m: int = 1, nvfp4_layers: set[int] | None = None, activation_scaling_report: Path | None = None, probe_seed: int = 20260818, ) -> dict[str, Any]: """Replace selected S2-Pro projections without modifying the FP8 path.""" if w4a16_max_m < 1: raise ValueError("w4a16_max_m must be positive") selectors = { "slow_mlp": (_selected_slow_mlp, 108), "slow_transformer": (_selected_slow_transformer, 180), } if policy == "w4a16_slow_transformer": candidates = [ (name, module) for name, module in model.named_modules() if _selected_slow_transformer(name, module) ] if len(candidates) != 180: raise RuntimeError( f"Expected 180 w4a16_slow_transformer projections, found {len(candidates)}" ) generator = torch.Generator(device=candidates[0][1].weight.device) generator.manual_seed(probe_seed) records = [] for name, linear in candidates: parent_name, attribute = name.rsplit(".", 1) parent = model.get_submodule(parent_name) probe = torch.randn( 1, linear.in_features, dtype=torch.bfloat16, device=linear.weight.device, generator=generator, ) * 0.1 reference = F.linear(probe, linear.weight) replacement = NVFP4Linear.from_linear( linear, w4a16_max_m=w4a16_max_m ) actual = replacement(probe) records.append( ConversionRecord( name=name, in_features=linear.in_features, out_features=linear.out_features, parameters=linear.weight.numel(), probe_cosine=float( F.cosine_similarity( actual.float().flatten(), reference.float().flatten(), dim=0, ).item() ), ) ) setattr(parent, attribute, replacement) torch.cuda.synchronize(candidates[0][1].weight.device) serialized = [asdict(record) for record in records] cosines = [record.probe_cosine for record in records] parameters = sum(record.parameters for record in records) return { "policy": policy, "backend": ( f"w4a16_through_m{w4a16_max_m}+w4a4_above_threshold" ), "w4a16_max_m": w4a16_max_m, "modules": len(records), "projections": len(records), "parameters": parameters, "theoretical_bf16_source_bytes": parameters * 2, "packed_weight_bytes": _packed_weight_bytes(model), "probe_cosine_min": min(cosines), "probe_cosine_mean": sum(cosines) / len(cosines), "probe_cosine_max": max(cosines), "records": serialized, } mixed_role_policies = { "w4a16_mlp_mxfp8_attention": ({"w1", "w2", "w3"}, None), "w4a16_gate_up_mxfp8_rest": ({"w1", "w3"}, None), "w4a16_down_mxfp8_rest": ({"w2"}, None), "w4a16_gate_up_middle6_mxfp8_rest": ({"w1", "w3"}, set(range(15, 21))), "w4a16_gate_up_middle12_mxfp8_rest": ({"w1", "w3"}, set(range(12, 24))), "w4a16_gate_up_middle18_mxfp8_rest": ({"w1", "w3"}, set(range(9, 27))), "w4a16_gate_up_middle24_mxfp8_rest": ({"w1", "w3"}, set(range(6, 30))), "w4a16_gate_up_middle30_mxfp8_rest": ({"w1", "w3"}, set(range(3, 33))), "w4a16_gate_up_custom_mxfp8_rest": ({"w1", "w3"}, "custom"), } if policy in mixed_role_policies: from experimental.fp8 import MXFP8Linear candidates = [ (name, module) for name, module in model.named_modules() if _selected_slow_transformer(name, module) ] if len(candidates) != 180: raise RuntimeError( f"Expected 180 {policy} projections, " f"found {len(candidates)}" ) nvfp4_roles, policy_layers = mixed_role_policies[policy] if policy_layers == "custom": if not nvfp4_layers: raise ValueError(f"{policy} requires at least one nvfp4 layer") invalid_layers = sorted(set(nvfp4_layers) - set(range(36))) if invalid_layers: raise ValueError(f"Invalid slow-transformer layers: {invalid_layers}") selected_nvfp4_layers = set(nvfp4_layers) else: if nvfp4_layers is not None: raise ValueError("nvfp4_layers is only valid with the custom policy") selected_nvfp4_layers = policy_layers activation_scales: dict[int, torch.Tensor] = {} activation_clip_ratios: dict[int, float] = {} activation_quantization_scales: dict[int, dict[str, float]] = {} activation_corrections: dict[int, dict[str, torch.Tensor]] = {} activation_hadamard_blocks: dict[int, int] = {} activation_scaling = None if activation_scaling_report is not None: if nvfp4_roles != {"w1", "w3"}: raise ValueError( "Activation scaling requires NVFP4 gate and up projections together" ) layers_to_scale = ( set(range(36)) if selected_nvfp4_layers is None else set(selected_nvfp4_layers) ) ( activation_scales, activation_clip_ratios, activation_quantization_scales, activation_corrections, activation_scaling, ) = _prepare_activation_scaling( model, Path(activation_scaling_report), layers_to_scale, ) activation_hadamard_blocks = { int(record["layer"]): int(record["input_hadamard_block_size"]) for record in activation_scaling["layers"] if int(record.get("input_hadamard_block_size", 0)) } generator = torch.Generator(device=candidates[0][1].weight.device) generator.manual_seed(probe_seed) records = [] nvfp4_parameters = 0 mxfp8_parameters = 0 for name, linear in candidates: parent_name, attribute = name.rsplit(".", 1) parent = model.get_submodule(parent_name) probe = torch.randn( 1, linear.in_features, dtype=torch.bfloat16, device=linear.weight.device, generator=generator, ) * 0.1 reference = F.linear(probe, linear.weight) role = name.rsplit(".", 1)[-1] layer = int(name.split(".", 2)[1]) selected_layer = ( selected_nvfp4_layers is None or layer in selected_nvfp4_layers ) if ".feed_forward." in name and role in nvfp4_roles and selected_layer: activation_scale = activation_scales.get(layer) if activation_scale is None: replacement = NVFP4Linear.from_linear( linear, w4a16_max_m=w4a16_max_m ) actual_probe = probe else: scaled_weight = ( linear.weight.float() * activation_scale ) scaled_weight = _clip_weight_blocks( scaled_weight, activation_clip_ratios.get(layer, 1.0), ) input_hadamard_block_size = activation_hadamard_blocks.get( layer, 0, ) if input_hadamard_block_size: scaled_weight = _hadamard_blocks( scaled_weight, input_hadamard_block_size, ) scaled_weight = scaled_weight.to(dtype=torch.bfloat16) replacement = NVFP4Linear.from_weight( scaled_weight, w4a16_max_m=w4a16_max_m, quantization_scale=activation_quantization_scales.get( layer, {} ).get(role), correction_down=activation_corrections.get( layer, {} ).get("down"), correction_up=activation_corrections.get( layer, {} ).get(f"{role}_up"), sparse_correction_indices=activation_corrections.get( layer, {} ).get("sparse_indices"), sparse_correction_weight=activation_corrections.get( layer, {} ).get(f"{role}_sparse_weight"), input_hadamard_block_size=input_hadamard_block_size, ) actual_probe = ( probe.float() / activation_scale ).to(dtype=torch.bfloat16) precision = ( f"nvfp4_w4a16_through_m{w4a16_max_m}" + ("_activation_scaled" if activation_scale is not None else "") + ( "_low_rank_corrected" if "down" in activation_corrections.get(layer, {}) else "" ) + ( "_sparse_channel_corrected" if "sparse_indices" in activation_corrections.get(layer, {}) else "" ) + ( f"_hadamard{activation_hadamard_blocks[layer]}" if layer in activation_hadamard_blocks else "" ) ) nvfp4_parameters += linear.weight.numel() else: replacement = MXFP8Linear.from_linear(linear) actual_probe = probe precision = "mxfp8_w8a8" mxfp8_parameters += linear.weight.numel() actual = replacement(actual_probe) record = ConversionRecord( name=name, in_features=linear.in_features, out_features=linear.out_features, parameters=linear.weight.numel(), probe_cosine=float( F.cosine_similarity( actual.float().flatten(), reference.float().flatten(), dim=0, ).item() ), ) records.append({**asdict(record), "precision": precision}) setattr(parent, attribute, replacement) torch.cuda.synchronize(candidates[0][1].weight.device) nvfp4_bytes = _packed_weight_bytes(model) mxfp8_bytes = sum( module.weight_fp8.numel() * module.weight_fp8.element_size() + module.weight_scale_storage.numel() * module.weight_scale_storage.element_size() for module in model.modules() if isinstance(module, MXFP8Linear) ) cosines = [record["probe_cosine"] for record in records] parameters = nvfp4_parameters + mxfp8_parameters correction_tensors = { value.data_ptr(): value for correction in activation_corrections.values() for value in correction.values() } correction_parameters = sum( value.numel() for value in correction_tensors.values() if value.is_floating_point() ) correction_bytes = sum( value.numel() * value.element_size() for value in correction_tensors.values() ) low_rank_tensors = { value.data_ptr(): value for correction in activation_corrections.values() for key, value in correction.items() if key == "down" or key.endswith("_up") } sparse_correction_tensors = { value.data_ptr(): value for correction in activation_corrections.values() for key, value in correction.items() if key == "sparse_indices" or key.endswith("_sparse_weight") } low_rank_correction_parameters = sum( value.numel() for value in low_rank_tensors.values() ) low_rank_correction_bytes = sum( value.numel() * value.element_size() for value in low_rank_tensors.values() ) sparse_correction_parameters = sum( value.numel() for value in sparse_correction_tensors.values() if value.is_floating_point() ) sparse_correction_bytes = sum( value.numel() * value.element_size() for value in sparse_correction_tensors.values() ) return { "policy": policy, "backend": "selective_nvfp4_weight_bf16_activation+mxfp8_rest", "w4a16_max_m": w4a16_max_m, "nvfp4_mlp_roles": sorted(nvfp4_roles), "nvfp4_layers": ( "all" if selected_nvfp4_layers is None else sorted(selected_nvfp4_layers) ), "activation_scaling": activation_scaling, "modules": len(records), "projections": len(records), "parameters": parameters, "nvfp4_parameters": nvfp4_parameters, "mxfp8_parameters": mxfp8_parameters, "low_rank_correction_parameters": low_rank_correction_parameters, "low_rank_correction_bytes": low_rank_correction_bytes, "sparse_correction_parameters": sparse_correction_parameters, "sparse_correction_bytes": sparse_correction_bytes, "correction_parameters": correction_parameters, "correction_bytes": correction_bytes, "theoretical_bf16_source_bytes": parameters * 2, "packed_weight_bytes": nvfp4_bytes + mxfp8_bytes, "nvfp4_packed_weight_bytes": nvfp4_bytes, "mxfp8_packed_weight_bytes": mxfp8_bytes, "probe_cosine_min": min(cosines), "probe_cosine_mean": sum(cosines) / len(cosines), "probe_cosine_max": max(cosines), "records": records, } if policy == "hybrid_slow_transformer": from experimental.fp8 import MXFP8Linear nvfp4 = convert_s2_pro_nvfp4( model, policy="slow_mlp_fused", probe_seed=probe_seed, ) candidates = [ (name, module) for name, module in model.named_modules() if _selected_slow_attention(name, module) ] if len(candidates) != 72: raise RuntimeError( f"Expected 72 hybrid attention projections, found {len(candidates)}" ) generator = torch.Generator(device=candidates[0][1].weight.device) generator.manual_seed(probe_seed + 1) attention_records = [] for name, linear in candidates: parent_name, attribute = name.rsplit(".", 1) parent = model.get_submodule(parent_name) probe = torch.randn( 1, linear.in_features, dtype=torch.bfloat16, device=linear.weight.device, generator=generator, ) * 0.1 reference = F.linear(probe, linear.weight) replacement = MXFP8Linear.from_linear(linear) actual = replacement(probe) attention_records.append( ConversionRecord( name=name, in_features=linear.in_features, out_features=linear.out_features, parameters=linear.weight.numel(), probe_cosine=float( F.cosine_similarity( actual.float().flatten(), reference.float().flatten(), dim=0, ).item() ), ) ) setattr(parent, attribute, replacement) torch.cuda.synchronize(candidates[0][1].weight.device) attention_parameters = sum(record.parameters for record in attention_records) attention_packed_bytes = sum( module.weight_fp8.numel() * module.weight_fp8.element_size() + module.weight_scale_storage.numel() * module.weight_scale_storage.element_size() for module in model.modules() if isinstance(module, MXFP8Linear) ) cosines = [ *[record["probe_cosine"] for record in nvfp4["records"]], *[record.probe_cosine for record in attention_records], ] return { "policy": policy, "backend": "fused_nvfp4_mlp+mxfp8_attention", "modules": nvfp4["modules"] + len(attention_records), "projections": nvfp4["projections"] + len(attention_records), "parameters": nvfp4["parameters"] + attention_parameters, "theoretical_bf16_source_bytes": ( nvfp4["theoretical_bf16_source_bytes"] + attention_parameters * 2 ), "packed_weight_bytes": ( nvfp4["packed_weight_bytes"] + attention_packed_bytes ), "nvfp4_mlp": nvfp4, "mxfp8_attention": { "modules": len(attention_records), "parameters": attention_parameters, "packed_weight_bytes": attention_packed_bytes, "records": [asdict(record) for record in attention_records], }, "probe_cosine_min": min(cosines), "probe_cosine_mean": sum(cosines) / len(cosines), "probe_cosine_max": max(cosines), } if policy == "slow_transformer_fused": candidates = [ (name, module) for name, module in model.named_modules() if _selected_slow_transformer_block(name, module) ] if len(candidates) != 36: raise RuntimeError( f"Expected 36 slow_transformer_fused blocks, found {len(candidates)}" ) generator = torch.Generator(device=candidates[0][1].attention.wqkv.weight.device) generator.manual_seed(probe_seed) records = [] for name, module in candidates: parent_name, attribute = name.rsplit(".", 1) parent = model.get_submodule(parent_name) dim = module.attention.wqkv.in_features probe = torch.randn( 1, dim, dtype=torch.bfloat16, device=module.attention.wqkv.weight.device, generator=generator, ) * 0.1 wo_probe = torch.randn( 1, module.attention.wo.in_features, dtype=torch.bfloat16, device=module.attention.wo.weight.device, generator=generator, ) * 0.1 reference_wqkv = module.attention.wqkv(probe) reference_wo = module.attention.wo(wo_probe) reference_mlp = module.feed_forward(probe) parameters = sum( linear.weight.numel() for linear in ( module.attention.wqkv, module.attention.wo, module.feed_forward.w1, module.feed_forward.w2, module.feed_forward.w3, ) ) replacement = NVFP4TransformerBlock(module) actual_wqkv = replacement.attention.wqkv(probe) actual_wo = replacement.attention.wo(wo_probe) actual_mlp = replacement.feed_forward(probe) records.append( FusedTransformerConversionRecord( name=name, parameters=parameters, wqkv_probe_cosine=float( F.cosine_similarity( actual_wqkv.float().flatten(), reference_wqkv.float().flatten(), dim=0, ).item() ), wo_probe_cosine=float( F.cosine_similarity( actual_wo.float().flatten(), reference_wo.float().flatten(), dim=0, ).item() ), mlp_probe_cosine=float( F.cosine_similarity( actual_mlp.float().flatten(), reference_mlp.float().flatten(), dim=0, ).item() ), ) ) setattr(parent, attribute, replacement) torch.cuda.synchronize(candidates[0][1].attention.wqkv.qdata.device) serialized = [asdict(record) for record in records] cosines = [ cosine for record in records for cosine in ( record.wqkv_probe_cosine, record.wo_probe_cosine, record.mlp_probe_cosine, ) ] parameters = sum(record.parameters for record in records) return { "policy": policy, "backend": "direct_fused_m1+tensorcore_prefill", "modules": len(records), "projections": len(records) * 5, "parameters": parameters, "theoretical_bf16_source_bytes": parameters * 2, "packed_weight_bytes": _packed_weight_bytes(model), "probe_cosine_min": min(cosines), "probe_cosine_mean": sum(cosines) / len(cosines), "probe_cosine_max": max(cosines), "records": serialized, } if policy == "slow_mlp_fused": candidates = [ (name, module) for name, module in model.named_modules() if _selected_slow_mlp_block(name, module) ] if len(candidates) != 36: raise RuntimeError( f"Expected 36 slow_mlp_fused blocks, found {len(candidates)}" ) generator = torch.Generator(device=candidates[0][1].w1.weight.device) generator.manual_seed(probe_seed) records = [] for name, module in candidates: parent_name, attribute = name.rsplit(".", 1) parent = model.get_submodule(parent_name) probe = torch.randn( 1, module.w1.in_features, dtype=torch.bfloat16, device=module.w1.weight.device, generator=generator, ) * 0.1 reference = module(probe) replacement = NVFP4FeedForward.from_module(module) actual = replacement(probe) records.append( FusedMLPConversionRecord( name=name, in_features=module.w1.in_features, intermediate_features=module.w1.out_features, out_features=module.w2.out_features, parameters=( module.w1.weight.numel() + module.w2.weight.numel() + module.w3.weight.numel() ), probe_cosine=float( F.cosine_similarity( actual.float().flatten(), reference.float().flatten(), dim=0, ).item() ), ) ) setattr(parent, attribute, replacement) torch.cuda.synchronize(candidates[0][1].w1.weight.device) serialized = [asdict(record) for record in records] cosines = [record.probe_cosine for record in records] parameters = sum(record.parameters for record in records) return { "policy": policy, "backend": "direct_fused_m1+tensorcore_prefill", "modules": len(records), "projections": len(records) * 3, "parameters": parameters, "theoretical_bf16_source_bytes": parameters * 2, "packed_weight_bytes": _packed_weight_bytes(model), "probe_cosine_min": min(cosines), "probe_cosine_mean": sum(cosines) / len(cosines), "probe_cosine_max": max(cosines), "records": serialized, } if policy not in selectors: raise ValueError(f"Unsupported initial NVFP4 policy: {policy}") selector, expected_modules = selectors[policy] candidates = [ (name, module) for name, module in model.named_modules() if selector(name, module) ] if len(candidates) != expected_modules: raise RuntimeError( f"Expected {expected_modules} {policy} projections, found {len(candidates)}" ) generator = torch.Generator(device=candidates[0][1].weight.device) generator.manual_seed(probe_seed) records = [] for name, linear in candidates: parent_name, attribute = name.rsplit(".", 1) parent = model.get_submodule(parent_name) replacement = NVFP4Linear.from_linear(linear, direct_m1=direct_m1) probe = torch.randn( 1, linear.in_features, dtype=torch.bfloat16, device=linear.weight.device, generator=generator, ) * 0.1 reference = F.linear(probe, linear.weight) actual = replacement(probe) probe_cosine = float( F.cosine_similarity( actual.float().flatten(), reference.float().flatten(), dim=0 ).item() ) records.append( ConversionRecord( name=name, in_features=linear.in_features, out_features=linear.out_features, parameters=linear.weight.numel(), probe_cosine=probe_cosine, ) ) setattr(parent, attribute, replacement) torch.cuda.synchronize(candidates[0][1].weight.device) serialized = [asdict(record) for record in records] cosines = [record.probe_cosine for record in records] return { "policy": policy, "backend": "direct_m1+tensorcore" if direct_m1 else "tensorcore", "modules": len(records), "parameters": sum(record.parameters for record in records), "theoretical_bf16_source_bytes": sum(record.parameters * 2 for record in records), "packed_weight_bytes": _packed_weight_bytes(model), "probe_cosine_min": min(cosines), "probe_cosine_mean": sum(cosines) / len(cosines), "probe_cosine_max": max(cosines), "records": serialized, }