Upload tools/fp8_to_nvfp4_streaming.py with huggingface_hub

tools/fp8_to_nvfp4_streaming.py

@@ -0,0 +1,1290 @@
+#!/usr/bin/env python3
+"""
+Streaming FP8 to NVFP4 Conversion for DeepSeek V3.2
+
+Converts FP8 e4m3 quantized weights (128x128 block scales) to NVFP4 e2m1 format
+(16-element block scales) via FP32 intermediates.
+
+Target: vLLM-compatible checkpoint with compressed-tensors format.
+"""
+
+import os
+import json
+import torch
+import gc
+import re
+import shutil
+import time
+import logging
+from typing import Dict, Any, Optional, Tuple, List, Set
+from pathlib import Path
+from dataclasses import dataclass, field
+from safetensors.torch import save_file as st_save_file
+from safetensors import safe_open
+
+logger = logging.getLogger(__name__)
+
+# ============================================================================
+# NVFP4 E2M1 Constants (from TensorRT-Model-Optimizer nvfp4_tensor.py)
+# ============================================================================
+
+# E2M1 quantization boundaries for searchsorted
+E2M1_BOUNDS = torch.tensor([0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0])
+
+# E2M1 representable values (index 0-7 = positive, 8-15 = negative with sign bit)
+E2M1_VALUES = torch.tensor([0, 0.5, 1, 1.5, 2, 3, 4, 6, 0, -0.5, -1, -1.5, -2, -3, -4, -6])
+
+# Maximum representable FP4 value
+FP4_MAX = 6.0
+
+# Maximum FP8 E4M3 value (for scale normalization)
+FP8_E4M3_MAX = 448.0
+
+# ============================================================================
+# Tensor Classification Patterns
+# ============================================================================
+
+# Patterns for tensors that should NOT be quantized (preserve in original dtype)
+PRESERVE_PATTERNS = [
+    r"embed_tokens",              # Embeddings
+    r"lm_head",                   # Output head
+    r"\.mlp\.gate\.",             # MoE router gate (NOT gate_proj) - note: .gate. not .gate$
+    r"shared_experts\.gate\.",    # Shared expert routing
+    r"shared_expert_gate",        # Alternative naming
+    r"layernorm",                 # LayerNorm weights
+    r"_norm\.",                   # RMSNorm weights (input_layernorm, etc.)
+    r"\.norm\.",                  # Norm weights
+    r"\.bias$",                   # Bias terms
+    # V3.2 DSA-specific (CRITICAL):
+    r"indexer\.weights_proj",     # Sparse pattern selector - MUST preserve!
+    r"indexer\.k_norm",           # Indexer normalization
+    # Scale tensors (handled separately)
+    r"_scale_inv$",               # FP8 scale_inv tensors
+    r"_scale$",                   # Scale tensors
+    r"_scale_2$",                 # Global scale tensors
+]
+
+# Compile patterns for efficiency
+PRESERVE_PATTERNS_COMPILED = [re.compile(p) for p in PRESERVE_PATTERNS]
+
+
+# ============================================================================
+# ShardedSafeTensorWriter (adapted from fp8_fp4_llmcompressor_streaming.py)
+# ============================================================================
+
+class ShardedSafeTensorWriter:
+    """
+    Stream tensors into numbered .safetensors shards and build a HF-style index JSON.
+    """
+    def __init__(self, out_dir: str, max_shard_size: str = "5GB"):
+        self.out_dir = os.path.abspath(out_dir)
+        os.makedirs(self.out_dir, exist_ok=True)
+        self.max_bytes = self._parse_size_to_bytes(max_shard_size)
+        self.curr_tensors: Dict[str, torch.Tensor] = {}
+        self.curr_bytes = 0
+        self.shard_idx = 1
+        self.weight_map: Dict[str, str] = {}
+        self.total_bytes = 0
+
+    def _parse_size_to_bytes(self, size_str: str) -> int:
+        size_str = size_str.upper().strip()
+        if size_str.endswith('GB'):
+            return int(float(size_str[:-2]) * 1024 * 1024 * 1024)
+        elif size_str.endswith('MB'):
+            return int(float(size_str[:-2]) * 1024 * 1024)
+        elif size_str.endswith('KB'):
+            return int(float(size_str[:-2]) * 1024)
+        else:
+            return int(size_str)
+
+    def _next_shard_name(self) -> str:
+        return f"model-{self.shard_idx:05d}.safetensors"
+
+    def _flush(self):
+        if not self.curr_tensors:
+            return
+        fname = self._next_shard_name()
+        path = os.path.join(self.out_dir, fname)
+        st_save_file(self.curr_tensors, path, metadata={"format": "nvfp4"})
+        logger.info(f"  Saved shard {fname}: {len(self.curr_tensors)} tensors, {self.curr_bytes / 1e9:.2f} GB")
+        for k in self.curr_tensors.keys():
+            self.weight_map[k] = fname
+        self.total_bytes += self.curr_bytes
+        self.curr_tensors.clear()
+        self.curr_bytes = 0
+        self.shard_idx += 1
+
+    def add_tensor(self, name: str, tensor: torch.Tensor):
+        if tensor.device.type != "cpu":
+            tensor = tensor.to("cpu")
+        if not tensor.is_contiguous():
+            tensor = tensor.contiguous()
+        tbytes = tensor.element_size() * tensor.numel()
+        if self.curr_bytes > 0 and self.curr_bytes + tbytes > self.max_bytes:
+            self._flush()
+        self.curr_tensors[name] = tensor
+        self.curr_bytes += tbytes
+
+    def finalize(self) -> int:
+        self._flush()
+        index_path = os.path.join(self.out_dir, "model.safetensors.index.json")
+        index = {"metadata": {"total_size": self.total_bytes}, "weight_map": self.weight_map}
+        with open(index_path, "w") as f:
+            json.dump(index, f, indent=2)
+        logger.info(f"Finalized: {self.shard_idx - 1} shards, {self.total_bytes / 1e9:.2f} GB total")
+        return self.shard_idx - 1
+
+
+# ============================================================================
+# Conversion Statistics
+# ============================================================================
+
+@dataclass
+class ConversionStats:
+    """Track conversion statistics."""
+    total_tensors: int = 0
+    fp8_tensors: int = 0
+    # Primary conversions: FP8 tensors where we ran the full conversion logic
+    primary_conversions: int = 0
+    # MoE partner conversions: FP8 tensors converted as partners during joint scale computation
+    # These are cached during primary conversion and written when encountered in stream
+    moe_partner_conversions: int = 0
+    preserved_sensitive: int = 0
+    copied_unchanged: int = 0
+    total_params: int = 0
+    layers_processed: Set[str] = field(default_factory=set)
+    warnings: List[Dict] = field(default_factory=list)
+    errors: List[Dict] = field(default_factory=list)
+    start_time: float = 0
+    end_time: float = 0
+
+    @property
+    def total_nvfp4_tensors(self) -> int:
+        """Total FP8 tensors converted to NVFP4 (primary + partner)."""
+        return self.primary_conversions + self.moe_partner_conversions
+
+    def log_warning(self, key: str, reason: str):
+        self.warnings.append({"tensor": key, "reason": reason})
+
+    def log_error(self, key: str, error: str):
+        self.errors.append({"tensor": key, "error": error})
+
+
+# ============================================================================
+# FP8 Block Dequantization
+# ============================================================================
+
+def dequantize_fp8_block_to_fp32(
+    fp8_weight: torch.Tensor,
+    scale_inv: torch.Tensor,
+    block_size: int = 128,
+    device: Optional[torch.device] = None
+) -> torch.Tensor:
+    """
+    Dequantize FP8 e4m3 weight using block-wise scale_inv.
+
+    The DeepSeek FP8 format uses 128x128 blocks where each block
+    shares a single inverse scale factor.
+
+    Formula: fp32_weight = fp8_weight.to(float32) * scale_inv[block_i, block_j]
+
+    Reference: TensorRT-Model-Optimizer/examples/deepseek/ds_kernel.py:89-110
+
+    Args:
+        fp8_weight: FP8 e4m3 weight tensor [M, N]
+        scale_inv: Inverse scale tensor [M/block_size, N/block_size]
+        block_size: Block size (default 128)
+        device: Device to compute on (None = same as input)
+
+    Returns:
+        FP32 dequantized weight tensor [M, N]
+    """
+    if device is not None:
+        fp8_weight = fp8_weight.to(device)
+        scale_inv = scale_inv.to(device)
+
+    M, N = fp8_weight.shape
+
+    # Handle case where dimensions aren't divisible by block_size
+    M_blocks = (M + block_size - 1) // block_size
+    N_blocks = (N + block_size - 1) // block_size
+
+    # Validate scale_inv shape
+    expected_scale_shape = (M_blocks, N_blocks)
+    if scale_inv.shape != expected_scale_shape:
+        # Some weights have different scale shapes (e.g., per-row scaling)
+        if scale_inv.numel() == 1:
+            # Scalar scale
+            return fp8_weight.to(torch.float32) * scale_inv.item()
+        elif scale_inv.shape[0] == 1 or scale_inv.shape[1] == 1:
+            # Per-row or per-column scaling
+            return fp8_weight.to(torch.float32) * scale_inv.to(torch.float32)
+        else:
+            logger.warning(f"Unexpected scale_inv shape {scale_inv.shape} for weight {fp8_weight.shape}, expected {expected_scale_shape}")
+            # Try to broadcast
+            return fp8_weight.to(torch.float32) * scale_inv.to(torch.float32)
+
+    # Convert FP8 to FP32
+    fp32_weight = fp8_weight.to(torch.float32)
+
+    # If dimensions match exactly, use efficient block multiplication
+    if M % block_size == 0 and N % block_size == 0:
+        # Reshape to blocks: [M/bs, bs, N/bs, bs]
+        weight_blocks = fp32_weight.view(M_blocks, block_size, N_blocks, block_size)
+
+        # Apply scale: scale_inv[i, j] applies to weight_blocks[i, :, j, :]
+        # scale_inv shape: [M_blocks, N_blocks] -> [M_blocks, 1, N_blocks, 1]
+        scaled = weight_blocks * scale_inv[:, None, :, None].to(torch.float32)
+
+        # Reshape back
+        return scaled.view(M, N)
+    else:
+        # Handle non-divisible dimensions with padding
+        M_pad = M_blocks * block_size
+        N_pad = N_blocks * block_size
+
+        padded_weight = torch.zeros(M_pad, N_pad, dtype=torch.float32, device=fp32_weight.device)
+        padded_weight[:M, :N] = fp32_weight
+
+        weight_blocks = padded_weight.view(M_blocks, block_size, N_blocks, block_size)
+        scaled = weight_blocks * scale_inv[:, None, :, None].to(torch.float32)
+
+        return scaled.view(M_pad, N_pad)[:M, :N]
+
+
+# ============================================================================
+# NVFP4 Scale Computation
+# ============================================================================
+
+def compute_nvfp4_scales(
+    fp32_weight: torch.Tensor,
+    block_size: int = 16
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Compute two-level NVFP4 scaling factors.
+
+    NVFP4 uses dual-level scaling:
+    1. Per-tensor global scale (scale_2): amax / (6.0 * 448.0)
+    2. Per-block scale: per_block_amax / (6.0 * scale_2)
+
+    Reference: TensorRT-Model-Optimizer nvfp4_tensor.py:94-97, 63-92
+
+    Args:
+        fp32_weight: FP32 weight tensor
+        block_size: Block size for per-block scaling (default 16)
+
+    Returns:
+        Tuple of:
+            - weight_scale: Per-block FP8 E4M3 scale [M, N/block_size]
+            - weight_scale_2: Per-tensor FP32 global scale (scalar tensor)
+    """
+    # Step 1: Compute per-tensor global scale (scale_2)
+    global_amax = fp32_weight.abs().max()
+    weight_scale_2 = global_amax / (FP4_MAX * FP8_E4M3_MAX)
+
+    # Ensure non-zero scale (use abs comparison to avoid float precision issues)
+    if weight_scale_2.abs() < 1e-10:
+        weight_scale_2 = torch.tensor(1e-8, dtype=torch.float32, device=fp32_weight.device)
+
+    # Step 2: Compute per-block scale
+    original_shape = fp32_weight.shape
+
+    # Handle N dimension for block quantization
+    M = fp32_weight.shape[0] if fp32_weight.dim() > 1 else 1
+    N = fp32_weight.shape[-1]
+
+    # Pad N if not divisible by block_size
+    N_padded = ((N + block_size - 1) // block_size) * block_size
+    if N_padded != N:
+        if fp32_weight.dim() == 1:
+            padded = torch.zeros(N_padded, dtype=fp32_weight.dtype, device=fp32_weight.device)
+            padded[:N] = fp32_weight
+            fp32_weight = padded
+        else:
+            padded = torch.zeros(*original_shape[:-1], N_padded, dtype=fp32_weight.dtype, device=fp32_weight.device)
+            padded[..., :N] = fp32_weight
+            fp32_weight = padded
+
+    # Reshape to blocks along last dimension
+    if fp32_weight.dim() == 1:
+        weight_blocks = fp32_weight.view(-1, block_size)
+    else:
+        weight_blocks = fp32_weight.view(*original_shape[:-1], -1, block_size)
+
+    # Compute per-block amax
+    per_block_amax = weight_blocks.abs().amax(dim=-1)  # [..., N/block_size]
+
+    # Per-block scale = per_block_amax / (6.0 * scale_2)
+    per_block_scale = per_block_amax / (FP4_MAX * weight_scale_2)
+
+    # Clamp to avoid division by zero, set zeros to 1.0
+    per_block_scale = per_block_scale.clamp(min=1e-8)
+    per_block_scale[per_block_scale < 1e-7] = 1.0
+
+    # Convert to FP8 E4M3 (if available, otherwise keep as float32)
+    try:
+        weight_scale = per_block_scale.to(torch.float8_e4m3fn)
+    except (RuntimeError, TypeError):
+        # FP8 not supported on this device/PyTorch version
+        weight_scale = per_block_scale.to(torch.float32)
+
+    return weight_scale, weight_scale_2
+
+
+# ============================================================================
+# NVFP4 Quantization and Packing
+# ============================================================================
+
+def quantize_to_nvfp4_packed(
+    fp32_weight: torch.Tensor,
+    weight_scale: torch.Tensor,
+    weight_scale_2: torch.Tensor,
+    block_size: int = 16
+) -> torch.Tensor:
+    """
+    Quantize FP32 weight to NVFP4 packed uint8 format.
+
+    E2M1 values: {0, 0.5, 1, 1.5, 2, 3, 4, 6} with sign (16 total values)
+    Packing: (code[..., 1::2] << 4) | code[..., 0::2]
+
+    Reference: TensorRT-Model-Optimizer nvfp4_tensor.py:119-140, 224-227
+
+    Args:
+        fp32_weight: FP32 weight tensor
+        weight_scale: Per-block FP8 E4M3 scale
+        weight_scale_2: Per-tensor FP32 global scale
+        block_size: Block size (default 16)
+
+    Returns:
+        Packed uint8 tensor [M, N/2]
+    """
+    device = fp32_weight.device
+    original_shape = fp32_weight.shape
+    N = original_shape[-1]
+
+    # Pad N if not divisible by block_size
+    N_padded = ((N + block_size - 1) // block_size) * block_size
+    if N_padded != N:
+        if fp32_weight.dim() == 1:
+            padded = torch.zeros(N_padded, dtype=fp32_weight.dtype, device=device)
+            padded[:N] = fp32_weight
+            fp32_weight = padded
+        else:
+            padded = torch.zeros(*original_shape[:-1], N_padded, dtype=fp32_weight.dtype, device=device)
+            padded[..., :N] = fp32_weight
+            fp32_weight = padded
+
+    # Reshape for block-wise processing
+    if fp32_weight.dim() == 1:
+        weight_blocks = fp32_weight.view(-1, block_size)
+    else:
+        weight_blocks = fp32_weight.view(*original_shape[:-1], -1, block_size)
+
+    # Compute combined scale and apply
+    # scaled_weight = weight / (scale * scale_2)
+    combined_scale = weight_scale.to(torch.float32) * weight_scale_2
+    scaled_weight = weight_blocks / combined_scale.unsqueeze(-1)
+
+    # Flatten back to original shape (with padding)
+    if fp32_weight.dim() == 1:
+        scaled_weight = scaled_weight.view(-1)
+    else:
+        scaled_weight = scaled_weight.view(*original_shape[:-1], -1)
+
+    # Get E2M1 bounds on device
+    e2m1_bounds = E2M1_BOUNDS.to(device)
+
+    # Extract sign bit and compute absolute values
+    sign_bit = (scaled_weight < 0).to(torch.uint8)
+    weight_abs = scaled_weight.abs()
+
+    # Find nearest E2M1 magnitude index (0-7) using searchsorted
+    # searchsorted returns index where value should be inserted
+    ord_idx = torch.searchsorted(e2m1_bounds, weight_abs, out_int32=True).to(torch.uint8)
+
+    # Handle rounding at boundary values (odd indices need special treatment)
+    # For values exactly at odd boundaries [0.75, 1.75, 2.5], round up
+    odd_bounds = e2m1_bounds[[1, 3, 5]]  # [0.75, 1.75, 2.5]
+    equals_odd = torch.any(weight_abs.unsqueeze(-1) == odd_bounds, dim=-1).to(torch.uint8)
+
+    # Combine sign and ordinal: code = (sign << 3) | (ord + round_adjust)
+    fp4_codes = (sign_bit << 3) | (ord_idx + equals_odd)
+
+    # Ensure codes are in valid range [0, 15]
+    fp4_codes = fp4_codes.clamp(0, 15)
+
+    # Pack pairs of FP4 values into uint8
+    # Even indices in low nibble, odd indices in high nibble
+    packed = (fp4_codes[..., 1::2] << 4) | fp4_codes[..., 0::2]
+    packed = packed.to(torch.uint8)
+
+    return packed
+
+
+# ============================================================================
+# Tensor Classification
+# ============================================================================
+
+def should_preserve_tensor(key: str) -> bool:
+    """
+    Check if a tensor should be preserved (not quantized).
+
+    Args:
+        key: Tensor name/key
+
+    Returns:
+        True if tensor should be preserved in original dtype
+    """
+    for pattern in PRESERVE_PATTERNS_COMPILED:
+        if pattern.search(key):
+            return True
+    return False
+
+
+def is_fp8_weight(key: str, tensor: torch.Tensor) -> bool:
+    """
+    Check if a tensor is an FP8 quantized weight.
+
+    Args:
+        key: Tensor name
+        tensor: The tensor to check
+
+    Returns:
+        True if this is an FP8 weight that should be converted
+    """
+    # Check dtype
+    if tensor.dtype != torch.float8_e4m3fn:
+        return False
+
+    # Check it's a weight (not a scale or bias)
+    if not key.endswith('.weight'):
+        return False
+
+    # Check it's not a preserved tensor
+    if should_preserve_tensor(key):
+        return False
+
+    return True
+
+
+# ============================================================================
+# MoE Expert Pair Helper Functions
+# ============================================================================
+
+def get_moe_expert_pair_key(weight_key: str) -> Optional[str]:
+    """
+    Get the expert pair identifier for MoE gate_proj/up_proj weights.
+
+    For vLLM's fused MoE kernels, gate_proj (w1) and up_proj (w3) must share
+    the same weight_scale_2 because they're fused together.
+
+    Args:
+        weight_key: Tensor name (e.g., "model.layers.0.mlp.experts.5.gate_proj.weight")
+
+    Returns:
+        Expert pair key (e.g., "model.layers.0.mlp.experts.5") or None if not MoE weight
+    """
+    # Match MoE expert gate_proj or up_proj patterns
+    # Pattern: model.layers.{L}.mlp.experts.{E}.gate_proj.weight
+    # Pattern: model.layers.{L}.mlp.experts.{E}.up_proj.weight
+    moe_pattern = re.match(r'(model\.layers\.\d+\.mlp\.experts\.\d+)\.(gate_proj|up_proj)\.weight$', weight_key)
+    if moe_pattern:
+        return moe_pattern.group(1)
+
+    # Also match shared_experts pattern if present
+    shared_pattern = re.match(r'(model\.layers\.\d+\.mlp\.shared_experts)\.(gate_proj|up_proj)\.weight$', weight_key)
+    if shared_pattern:
+        return shared_pattern.group(1)
+
+    return None
+
+
+# ============================================================================
+# Main Converter Class
+# ============================================================================
+
+class FP8ToNVFP4StreamingConverter:
+    """
+    Streaming FP8 to NVFP4 converter for DeepSeek V3.2.
+
+    Processes safetensor shards sequentially with GPU acceleration,
+    converting FP8 e4m3 weights to NVFP4 e2m1 format.
+    """
+
+    def __init__(
+        self,
+        model_path: str,
+        output_dir: str,
+        device: str = "cuda",
+        max_shard_size: str = "5GB",
+        fp8_block_size: int = 128,
+        nvfp4_block_size: int = 16
+    ):
+        """
+        Initialize the converter.
+
+        Args:
+            model_path: Path to source FP8 model
+            output_dir: Output directory for NVFP4 model
+            device: Device for computation (cuda or cpu)
+            max_shard_size: Maximum output shard size
+            fp8_block_size: FP8 quantization block size (default 128)
+            nvfp4_block_size: NVFP4 quantization block size (default 16)
+        """
+        self.model_path = Path(model_path)
+        self.output_dir = Path(output_dir)
+        self.device = torch.device(device if torch.cuda.is_available() else "cpu")
+        self.max_shard_size = max_shard_size
+        self.fp8_block_size = fp8_block_size
+        self.nvfp4_block_size = nvfp4_block_size
+
+        # Load model index
+        self.weight_map, self.shard_to_keys = self._load_index()
+
+        # Initialize statistics
+        self.stats = ConversionStats()
+
+        # Cache for cross-shard scale_inv tensors
+        self.scale_cache: Dict[str, torch.Tensor] = {}
+
+        # Cache for processed MoE weights (for streaming partner handling)
+        # When we process gate_proj, we also load up_proj, process both with joint scale,
+        # and cache up_proj's result here so we can skip it when we encounter it later
+        # Key: weight_key (e.g., "model.layers.0.mlp.experts.5.up_proj.weight")
+        # Value: Dict of converted tensors
+        self.moe_processed_cache: Dict[str, Dict[str, torch.Tensor]] = {}
+
+        # Build MoE pair mapping from index for efficient lookup
+        self.moe_pairs: Dict[str, Dict[str, str]] = self._build_moe_pair_map()
+
+        # Initialize writer
+        self.writer = ShardedSafeTensorWriter(str(self.output_dir), max_shard_size)
+
+        logger.info(f"Initialized FP8→NVFP4 converter")
+        logger.info(f"  Source: {self.model_path}")
+        logger.info(f"  Output: {self.output_dir}")
+        logger.info(f"  Device: {self.device}")
+        logger.info(f"  FP8 block size: {self.fp8_block_size}")
+        logger.info(f"  NVFP4 block size: {self.nvfp4_block_size}")
+
+    def _load_index(self) -> Tuple[Dict[str, str], Dict[str, List[str]]]:
+        """Load model index and build shard-to-keys mapping."""
+        index_path = self.model_path / "model.safetensors.index.json"
+
+        if not index_path.exists():
+            raise FileNotFoundError(f"Model index not found: {index_path}")
+
+        with open(index_path) as f:
+            index = json.load(f)
+
+        weight_map = index.get("weight_map", {})
+
+        # Build reverse mapping: shard -> list of keys
+        shard_to_keys: Dict[str, List[str]] = {}
+        for key, shard in weight_map.items():
+            if shard not in shard_to_keys:
+                shard_to_keys[shard] = []
+            shard_to_keys[shard].append(key)
+
+        logger.info(f"Loaded index: {len(weight_map)} tensors across {len(shard_to_keys)} shards")
+
+        return weight_map, shard_to_keys
+
+    def _build_moe_pair_map(self) -> Dict[str, Dict[str, str]]:
+        """
+        Build mapping of MoE gate_proj/up_proj pairs from the index file.
+
+        This is a lightweight operation that just scans tensor names without
+        loading any weights, enabling efficient streaming processing.
+
+        Returns:
+            Dict mapping pair_key -> {"gate_proj": full_key, "up_proj": full_key}
+        """
+        moe_pairs: Dict[str, Dict[str, str]] = {}
+
+        for weight_key in self.weight_map.keys():
+            pair_key = get_moe_expert_pair_key(weight_key)
+            if pair_key:
+                if pair_key not in moe_pairs:
+                    moe_pairs[pair_key] = {}
+                if "gate_proj" in weight_key:
+                    moe_pairs[pair_key]["gate_proj"] = weight_key
+                elif "up_proj" in weight_key:
+                    moe_pairs[pair_key]["up_proj"] = weight_key
+
+        # Filter to complete pairs only
+        complete_pairs = {k: v for k, v in moe_pairs.items()
+                        if "gate_proj" in v and "up_proj" in v}
+
+        logger.info(f"Found {len(complete_pairs)} MoE expert pairs (gate_proj + up_proj)")
+        return complete_pairs
+
+    def _load_weight_from_shard(
+        self,
+        weight_key: str
+    ) -> Optional[Tuple[torch.Tensor, torch.Tensor]]:
+        """
+        Load an FP8 weight and its scale_inv from the appropriate shard.
+
+        Uses the index to locate which shard contains the weight.
+
+        Args:
+            weight_key: Full tensor key (e.g., "model.layers.0.mlp.experts.5.up_proj.weight")
+
+        Returns:
+            Tuple of (fp8_weight, scale_inv) or None if not found
+        """
+        if weight_key not in self.weight_map:
+            return None
+
+        shard_name = self.weight_map[weight_key]
+        shard_path = self.model_path / shard_name
+
+        if not shard_path.exists():
+            logger.warning(f"Shard not found: {shard_path}")
+            return None
+
+        try:
+            with safe_open(shard_path, framework="pt", device="cpu") as f:
+                shard_keys = list(f.keys())
+
+                if weight_key not in shard_keys:
+                    return None
+
+                fp8_weight = f.get_tensor(weight_key)
+
+                # Get scale_inv (may be in this shard or another)
+                scale_inv = self._get_scale_inv(weight_key, shard_keys, f)
+                if scale_inv is None:
+                    logger.warning(f"Missing scale_inv for {weight_key}")
+                    return None
+
+                return fp8_weight, scale_inv
+        except Exception as e:
+            logger.warning(f"Failed to load {weight_key}: {e}")
+            return None
+
+    def _get_partner_key(self, weight_key: str) -> Optional[str]:
+        """
+        Get the partner key for an MoE gate_proj/up_proj weight.
+
+        Args:
+            weight_key: Full tensor key
+
+        Returns:
+            Partner weight key or None if not an MoE pair weight
+        """
+        pair_key = get_moe_expert_pair_key(weight_key)
+        if not pair_key or pair_key not in self.moe_pairs:
+            return None
+
+        pair = self.moe_pairs[pair_key]
+        if "gate_proj" in weight_key:
+            return pair.get("up_proj")
+        elif "up_proj" in weight_key:
+            return pair.get("gate_proj")
+        return None
+
+    def _get_scale_inv(
+        self,
+        weight_key: str,
+        current_shard_keys: List[str],
+        current_shard_file: Any  # safetensors file handle from safe_open()
+    ) -> Optional[torch.Tensor]:
+        """
+        Get scale_inv tensor, loading from other shard if needed.
+
+        Uses the model index to find which shard contains the scale_inv
+        and loads it on demand. Caches loaded scales for efficiency.
+
+        Args:
+            weight_key: The weight tensor key (e.g., "model.layers.X.mlp.gate_proj.weight")
+            current_shard_keys: List of keys in the current shard
+            current_shard_file: Open safetensors file handle for current shard
+
+        Returns:
+            scale_inv tensor or None if not found
+        """
+        scale_key = weight_key.replace('.weight', '.weight_scale_inv')
+
+        # Fast path: check current shard first
+        if scale_key in current_shard_keys:
+            return current_shard_file.get_tensor(scale_key)
+
+        # Check cache
+        if scale_key in self.scale_cache:
+            return self.scale_cache[scale_key]
+
+        # Look up in index and load from correct shard
+        if scale_key in self.weight_map:
+            scale_shard = self.weight_map[scale_key]
+            scale_path = self.model_path / scale_shard
+
+            try:
+                with safe_open(scale_path, framework="pt", device="cpu") as f:
+                    scale_inv = f.get_tensor(scale_key)
+                    # Cache for future use (scales are small ~32KB each)
+                    self.scale_cache[scale_key] = scale_inv
+                    logger.debug(f"Loaded cross-shard scale_inv from {scale_shard}: {scale_key}")
+                    return scale_inv
+            except Exception as e:
+                logger.warning(f"Failed to load scale_inv from {scale_shard}: {e}")
+                return None
+
+        return None
+
+    def _convert_fp8_to_nvfp4(
+        self,
+        key: str,
+        fp8_weight: torch.Tensor,
+        scale_inv: torch.Tensor
+    ) -> Dict[str, torch.Tensor]:
+        """
+        Convert a single FP8 weight to NVFP4 format.
+
+        For MoE gate_proj/up_proj weights, loads the partner weight on-demand
+        to compute a joint scale_2, ensuring vLLM's fused MoE kernels work correctly.
+        The partner's result is cached to avoid reprocessing.
+
+        Args:
+            key: Tensor name
+            fp8_weight: FP8 e4m3 weight tensor
+            scale_inv: FP8 inverse scale tensor
+
+        Returns:
+            Dict with converted tensors:
+                - key: packed NVFP4 weight
+                - key.replace('.weight', '.weight_scale'): per-block scale
+                - key.replace('.weight', '.weight_scale_2'): global scale
+        """
+        # Move to processing device
+        fp8_weight = fp8_weight.to(self.device)
+        scale_inv = scale_inv.to(self.device)
+
+        # Step 1: Dequantize FP8 to FP32
+        fp32_weight = dequantize_fp8_block_to_fp32(
+            fp8_weight, scale_inv, block_size=self.fp8_block_size
+        )
+
+        # Step 2: Compute NVFP4 scales
+        # Check if this is an MoE weight that needs shared scale_2 with partner
+        partner_key = self._get_partner_key(key)
+
+        if partner_key:
+            # MoE gate_proj/up_proj - need joint scale with partner
+            # Load partner weight on-demand
+            partner_data = self._load_weight_from_shard(partner_key)
+
+            if partner_data:
+                partner_fp8, partner_scale_inv = partner_data
+                partner_fp8 = partner_fp8.to(self.device)
+                partner_scale_inv = partner_scale_inv.to(self.device)
+
+                # Dequantize partner
+                partner_fp32 = dequantize_fp8_block_to_fp32(
+                    partner_fp8, partner_scale_inv, block_size=self.fp8_block_size
+                )
+
+                # Compute joint amax and scale_2
+                my_amax = fp32_weight.abs().max()
+                partner_amax = partner_fp32.abs().max()
+                joint_amax = torch.max(my_amax, partner_amax)
+                joint_scale_2 = joint_amax / (FP4_MAX * FP8_E4M3_MAX)
+
+                # Ensure non-zero (use abs comparison to avoid float precision issues)
+                if joint_scale_2.abs() < 1e-10:
+                    joint_scale_2 = torch.tensor(1e-8, dtype=torch.float32, device=self.device)
+
+                # Compute per-block scale for this weight using joint scale_2
+                weight_scale = self._compute_per_block_scale(fp32_weight, joint_scale_2)
+                weight_scale_2 = joint_scale_2
+
+                # Also convert partner and cache its result
+                partner_scale = self._compute_per_block_scale(partner_fp32, joint_scale_2)
+                partner_packed = quantize_to_nvfp4_packed(
+                    partner_fp32, partner_scale, joint_scale_2, block_size=self.nvfp4_block_size
+                )
+
+                partner_base = partner_key.replace('.weight', '')
+                self.moe_processed_cache[partner_key] = {
+                    f"{partner_base}.weight": partner_packed.cpu(),
+                    f"{partner_base}.weight_scale": partner_scale.cpu(),
+                    f"{partner_base}.weight_scale_2": joint_scale_2.cpu().view(1),
+                }
+
+                logger.debug(f"Computed joint scale_2 for {key} + {partner_key}: {joint_scale_2.item():.6e}")
+
+                # Cleanup partner tensors
+                del partner_fp32, partner_fp8, partner_scale_inv
+            else:
+                # Partner not found - use standard per-tensor scale
+                logger.warning(f"Partner {partner_key} not found for {key}, using independent scale")
+                weight_scale, weight_scale_2 = compute_nvfp4_scales(
+                    fp32_weight, block_size=self.nvfp4_block_size
+                )
+        else:
+            # Non-MoE weight - standard per-tensor scale computation
+            weight_scale, weight_scale_2 = compute_nvfp4_scales(
+                fp32_weight, block_size=self.nvfp4_block_size
+            )
+
+        # Step 3: Quantize to NVFP4 packed format
+        packed_weight = quantize_to_nvfp4_packed(
+            fp32_weight, weight_scale, weight_scale_2, block_size=self.nvfp4_block_size
+        )
+
+        # Build output tensor names
+        base_name = key.replace('.weight', '')
+        result = {
+            f"{base_name}.weight": packed_weight.cpu(),
+            f"{base_name}.weight_scale": weight_scale.cpu(),
+            f"{base_name}.weight_scale_2": weight_scale_2.cpu().view(1),
+        }
+
+        # Update statistics - this is a "primary" conversion (not from MoE partner cache)
+        self.stats.primary_conversions += 1
+
+        # Free GPU memory
+        del fp32_weight
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+        return result
+
+    def _compute_per_block_scale(
+        self,
+        fp32_weight: torch.Tensor,
+        weight_scale_2: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Compute per-block scale given a fixed weight_scale_2.
+
+        Args:
+            fp32_weight: FP32 weight tensor
+            weight_scale_2: Global scale (FP32 scalar)
+
+        Returns:
+            Per-block FP8 E4M3 scale tensor
+        """
+        original_shape = fp32_weight.shape
+        N = fp32_weight.shape[-1]
+        block_size = self.nvfp4_block_size
+
+        # Pad N if not divisible by block_size
+        N_padded = ((N + block_size - 1) // block_size) * block_size
+        if N_padded != N:
+            if fp32_weight.dim() == 1:
+                padded = torch.zeros(N_padded, dtype=fp32_weight.dtype, device=fp32_weight.device)
+                padded[:N] = fp32_weight
+                fp32_padded = padded
+            else:
+                padded = torch.zeros(*original_shape[:-1], N_padded, dtype=fp32_weight.dtype, device=fp32_weight.device)
+                padded[..., :N] = fp32_weight
+                fp32_padded = padded
+        else:
+            fp32_padded = fp32_weight
+
+        # Reshape to blocks
+        if fp32_padded.dim() == 1:
+            weight_blocks = fp32_padded.view(-1, block_size)
+        else:
+            weight_blocks = fp32_padded.view(*original_shape[:-1], -1, block_size)
+
+        # Per-block amax
+        per_block_amax = weight_blocks.abs().amax(dim=-1)
+
+        # Per-block scale with the given scale_2
+        per_block_scale = per_block_amax / (FP4_MAX * weight_scale_2)
+        per_block_scale = per_block_scale.clamp(min=1e-8)
+        per_block_scale[per_block_scale < 1e-7] = 1.0
+
+        # Convert to FP8 E4M3
+        try:
+            return per_block_scale.to(torch.float8_e4m3fn)
+        except (RuntimeError, TypeError):
+            return per_block_scale.to(torch.float32)
+
+    def process_shard(self, shard_name: str) -> int:
+        """
+        Process a single shard, converting FP8 weights to NVFP4.
+
+        Args:
+            shard_name: Name of the shard file
+
+        Returns:
+            Number of tensors processed
+        """
+        shard_path = self.model_path / shard_name
+
+        if not shard_path.exists():
+            logger.error(f"Shard not found: {shard_path}")
+            return 0
+
+        tensors_processed = 0
+
+        with safe_open(shard_path, framework="pt", device="cpu") as f:
+            keys = list(f.keys())
+
+            # Process each tensor
+            for key in keys:
+                tensor = f.get_tensor(key)
+                self.stats.total_tensors += 1
+                self.stats.total_params += tensor.numel()
+
+                # Track layer (safely handle edge cases)
+                if '.layers.' in key:
+                    parts = key.split('.layers.')
+                    if len(parts) > 1 and '.' in parts[1]:
+                        layer_num = parts[1].split('.')[0]
+                        self.stats.layers_processed.add(layer_num)
+
+                # Skip scale_inv tensors (handled with weights)
+                if key.endswith('_scale_inv'):
+                    continue
+
+                # Check if this is an FP8 weight to convert
+                if is_fp8_weight(key, tensor):
+                    self.stats.fp8_tensors += 1
+
+                    # Check if this weight was already processed as a partner
+                    if key in self.moe_processed_cache:
+                        # Use cached result from partner processing
+                        # This tensor was converted when its MoE partner was processed
+                        # (gate_proj and up_proj share weight_scale_2 for vLLM fused kernels)
+                        cached = self.moe_processed_cache.pop(key)  # Pop to free memory
+                        for name, t in cached.items():
+                            self.writer.add_tensor(name, t)
+                        self.stats.moe_partner_conversions += 1
+                        tensors_processed += 1
+                        logger.debug(f"Using cached result for MoE partner: {key}")
+                        continue
+
+                    # Find corresponding scale_inv (with cross-shard lookup)
+                    scale_inv = self._get_scale_inv(key, keys, f)
+
+                    if scale_inv is not None:
+                        try:
+                            # Convert FP8 → NVFP4
+                            converted = self._convert_fp8_to_nvfp4(key, tensor, scale_inv)
+
+                            # Add to writer
+                            for name, t in converted.items():
+                                self.writer.add_tensor(name, t)
+
+                            tensors_processed += 1
+
+                        except Exception as e:
+                            logger.error(f"Error converting {key}: {e}")
+                            self.stats.log_error(key, str(e))
+                            # Skip this tensor - preserving FP8 would create corrupt checkpoint
+                            # vLLM expects NVFP4 format for all quantized weights
+                            logger.warning(f"Skipping {key} due to conversion error - checkpoint may be incomplete")
+                    else:
+                        # Missing scale_inv - skip this tensor
+                        # Preserving FP8 would create corrupt checkpoint
+                        logger.warning(f"Missing scale_inv for {key} (not found in any shard) - skipping")
+                        self.stats.log_warning(key, "missing_scale_inv")
+
+                elif should_preserve_tensor(key):
+                    # Preserve sensitive tensors
+                    self.writer.add_tensor(key, tensor)
+                    self.stats.preserved_sensitive += 1
+                    tensors_processed += 1
+
+                else:
+                    # Copy other tensors unchanged (norms, biases, etc.)
+                    self.writer.add_tensor(key, tensor)
+                    self.stats.copied_unchanged += 1
+                    tensors_processed += 1
+
+                # Free memory
+                del tensor
+
+        # Clear scale cache - scales from this shard won't be needed again
+        # This prevents unbounded memory growth for large models
+        self.scale_cache.clear()
+
+        # Garbage collection
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+        return tensors_processed
+
+    def generate_config(self) -> Dict[str, Any]:
+        """Generate vLLM-compatible config.json with modelopt NVFP4 format."""
+        # Load original config
+        config_path = self.model_path / "config.json"
+        with open(config_path) as f:
+            config = json.load(f)
+
+        # Update quantization config for NVFP4 using modelopt format
+        # This format is compatible with vLLM's modelopt_fp4 quantization handler
+        # Reference: https://huggingface.co/nvidia/Llama-3.1-8B-Instruct-NVFP4/blob/main/config.json
+        config["quantization_config"] = {
+            "quant_method": "modelopt",
+            "quant_algo": "NVFP4",
+            "config_groups": {
+                "group_0": {
+                    "targets": ["Linear"],
+                    "weights": {
+                        "num_bits": 4,
+                        "type": "float",
+                        "group_size": self.nvfp4_block_size,
+                        "dynamic": False
+                    },
+                    "input_activations": None
+                }
+            },
+            "ignore": [
+                "lm_head",
+                "model.embed_tokens",
+                "re:.*\\.mlp\\.gate$",
+                "re:.*layernorm.*",
+                "re:.*_norm.*",
+                "re:.*indexer\\.weights_proj.*",
+                "re:.*indexer\\.k_norm.*"
+            ],
+            "kv_cache_scheme": None,
+            "original_format": {
+                "quant_method": "fp8",
+                "fmt": "e4m3",
+                "scale_fmt": "ue8m0",
+                "weight_block_size": [self.fp8_block_size, self.fp8_block_size]
+            },
+            "conversion_info": {
+                "source": "fp8_e4m3",
+                "target": "nvfp4_e2m1",
+                "intermediate": "fp32",
+                "timestamp": time.strftime("%Y-%m-%d %H:%M:%S")
+            }
+        }
+
+        return config
+
+    def copy_auxiliary_files(self):
+        """Copy tokenizer and other auxiliary files."""
+        aux_files = [
+            "tokenizer.json",
+            "tokenizer_config.json",
+            "special_tokens_map.json",
+            "vocab.json",
+            "merges.txt",
+            "tokenizer.model",
+            "generation_config.json"
+        ]
+
+        for filename in aux_files:
+            src = self.model_path / filename
+            if src.exists():
+                dst = self.output_dir / filename
+                shutil.copy2(src, dst)
+                logger.info(f"Copied {filename}")
+
+        # Copy encoding folder if exists (V3.2 specific)
+        encoding_src = self.model_path / "encoding"
+        if encoding_src.exists() and encoding_src.is_dir():
+            encoding_dst = self.output_dir / "encoding"
+            shutil.copytree(encoding_src, encoding_dst, dirs_exist_ok=True)
+            logger.info("Copied encoding folder")
+
+    def generate_report(self) -> Dict[str, Any]:
+        """Generate conversion report."""
+        elapsed = self.stats.end_time - self.stats.start_time
+
+        report = {
+            "conversion_summary": {
+                "source_format": "FP8 E4M3 (DeepSeek block-quantized)",
+                "target_format": "NVFP4 E2M1 (16-element blocks)",
+                "intermediate_format": "FP32",
+                "model": str(self.model_path),
+                "output": str(self.output_dir),
+                "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
+                "elapsed_seconds": round(elapsed, 2),
+                "elapsed_minutes": round(elapsed / 60, 2)
+            },
+            "tensor_statistics": {
+                "total_tensors": self.stats.total_tensors,
+                "fp8_tensors_found": self.stats.fp8_tensors,
+                "primary_conversions": self.stats.primary_conversions,
+                "moe_partner_conversions": self.stats.moe_partner_conversions,
+                "total_nvfp4_tensors": self.stats.total_nvfp4_tensors,
+                "preserved_sensitive": self.stats.preserved_sensitive,
+                "copied_unchanged": self.stats.copied_unchanged,
+                "total_parameters": self.stats.total_params
+            },
+            "layer_statistics": {
+                "layers_processed": len(self.stats.layers_processed),
+                "layer_ids": sorted(self.stats.layers_processed, key=lambda x: int(x) if x.isdigit() else 0)
+            },
+            "output_statistics": {
+                "output_shards": self.writer.shard_idx - 1,
+                "output_size_gb": round(self.writer.total_bytes / 1e9, 2)
+            },
+            "issues": {
+                "warnings": self.stats.warnings[:20],
+                "errors": self.stats.errors[:20],
+                "total_warnings": len(self.stats.warnings),
+                "total_errors": len(self.stats.errors)
+            }
+        }
+
+        # Log truncation if applicable
+        if len(self.stats.warnings) > 20:
+            logger.info(f"Report truncated: showing 20 of {len(self.stats.warnings)} warnings")
+        if len(self.stats.errors) > 20:
+            logger.info(f"Report truncated: showing 20 of {len(self.stats.errors)} errors")
+
+        return report
+
+    def run(self) -> Dict[str, Any]:
+        """
+        Run the full conversion process.
+
+        Returns:
+            Conversion report dictionary
+        """
+        logger.info("=" * 70)
+        logger.info("Starting FP8 to NVFP4 Streaming Conversion")
+        logger.info("=" * 70)
+
+        self.stats.start_time = time.time()
+
+        # Get sorted list of shards
+        shard_names = sorted(self.shard_to_keys.keys())
+        total_shards = len(shard_names)
+
+        logger.info(f"Processing {total_shards} shards...")
+
+        # Process each shard
+        for idx, shard_name in enumerate(shard_names, 1):
+            logger.info(f"\n[{idx}/{total_shards}] Processing {shard_name}")
+            tensors = self.process_shard(shard_name)
+            logger.info(f"  Processed {tensors} tensors")
+
+        # Check for orphaned MoE cache entries (partner never encountered)
+        if self.moe_processed_cache:
+            orphan_count = len(self.moe_processed_cache)
+            logger.warning(f"Found {orphan_count} orphaned MoE cache entries (partner weight never processed):")
+            for key in list(self.moe_processed_cache.keys())[:5]:
+                logger.warning(f"  - {key}")
+            if orphan_count > 5:
+                logger.warning(f"  ... and {orphan_count - 5} more")
+            self.moe_processed_cache.clear()
+
+        # Finalize output
+        logger.info("\nFinalizing output...")
+        self.writer.finalize()
+
+        # Generate and save config
+        logger.info("Generating config.json...")
+        config = self.generate_config()
+        config_path = self.output_dir / "config.json"
+        with open(config_path, 'w') as f:
+            json.dump(config, f, indent=2)
+
+        # Copy auxiliary files
+        logger.info("Copying auxiliary files...")
+        self.copy_auxiliary_files()
+
+        self.stats.end_time = time.time()
+
+        # Generate report
+        report = self.generate_report()
+
+        # Save report
+        report_path = self.output_dir / "conversion_report.json"
+        with open(report_path, 'w') as f:
+            json.dump(report, f, indent=2)
+        logger.info(f"Saved conversion report: {report_path}")
+
+        # Print summary
+        elapsed = self.stats.end_time - self.stats.start_time
+        logger.info("\n" + "=" * 70)
+        logger.info("Conversion Complete!")
+        logger.info(f"  Time: {elapsed / 60:.1f} minutes")
+        logger.info(f"  FP8 tensors found: {self.stats.fp8_tensors}")
+        logger.info(f"  Primary conversions: {self.stats.primary_conversions}")
+        logger.info(f"  MoE partner conversions: {self.stats.moe_partner_conversions}")
+        logger.info(f"  Total NVFP4 tensors: {self.stats.total_nvfp4_tensors}")
+        logger.info(f"  Tensors preserved: {self.stats.preserved_sensitive}")
+        logger.info(f"  Output shards: {self.writer.shard_idx - 1}")
+        logger.info(f"  Output size: {self.writer.total_bytes / 1e9:.2f} GB")
+        logger.info(f"  Output: {self.output_dir}")
+        logger.info("=" * 70)
+
+        return report
+
+
+# ============================================================================
+# Main Entry Point
+# ============================================================================
+
+def main():
+    import argparse
+
+    parser = argparse.ArgumentParser(
+        description="Streaming FP8 to NVFP4 converter for DeepSeek V3.2"
+    )
+    parser.add_argument(
+        "model_path",
+        help="Path to FP8 model (e.g., /mnt/models/deepseek-v3.2)"
+    )
+    parser.add_argument(
+        "--output_dir",
+        default=None,
+        help="Output directory (default: {model_path}-nvfp4)"
+    )
+    parser.add_argument(
+        "--device",
+        default="cuda",
+        choices=["cuda", "cpu"],
+        help="Device for computation (default: cuda)"
+    )
+    parser.add_argument(
+        "--max_shard_size",
+        default="5GB",
+        help="Maximum output shard size (default: 5GB)"
+    )
+    parser.add_argument(
+        "--fp8_block_size",
+        type=int,
+        default=128,
+        help="FP8 quantization block size (default: 128)"
+    )
+    parser.add_argument(
+        "--nvfp4_block_size",
+        type=int,
+        default=16,
+        help="NVFP4 quantization block size (default: 16)"
+    )
+
+    args = parser.parse_args()
+
+    # Default output directory
+    if args.output_dir is None:
+        args.output_dir = f"{args.model_path.rstrip('/')}-nvfp4"
+
+    # Set up logging
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s - %(levelname)s - %(message)s"
+    )
+
+    # Create and run converter
+    converter = FP8ToNVFP4StreamingConverter(
+        model_path=args.model_path,
+        output_dir=args.output_dir,
+        device=args.device,
+        max_shard_size=args.max_shard_size,
+        fp8_block_size=args.fp8_block_size,
+        nvfp4_block_size=args.nvfp4_block_size
+    )
+
+    report = converter.run()
+
+    return report
+
+
+if __name__ == "__main__":
+    main()