nvfp4-megamoe-kernel/src/nvfp4_megamoe_kernel/weight_transform.py

"""
NVFP4 Weight Transformation for CUTLASS mega_moe kernel.

Converts raw NVFP4 checkpoint weights (uint8 E2M1 + float8_e4m3fn UE4M3 + float32 global scale)
into the format expected by the CUTLASS block-scaled GEMM kernel:
- Packed FP4 weights (int8, K-major)
- UE4M3 block scales (float8_e4m3fn, row-major — CUTLASS SF remap handles interleaving)
- float32 global scales (NOT folded into block scales — passed separately for per-expert alpha)

Previous versions folded weight_scale_2 into block scales via float8 round-trip, which caused
25% relative error (product of ~56-448 block_sf × ~4.65e-05 global_sf lands in the low-precision
zone of float8_e4m3fn where step size is 25%). The global scale is now applied as a per-expert
multiplier to the GEMM alpha, preserving full float32 precision.

Call signature matches the nightly vLLM deepseek_v4.py finalize_weights:
    transform_nvfp4_weights_for_mega_moe(
        (l1_weight, l1_weight_scale),
        (l2_weight, l2_weight_scale),
        l1_weight_scale_2=...,
        l2_weight_scale_2=...,
    )
"""

import torch


def transform_nvfp4_weights_for_mega_moe(
    l1_tuple: tuple[torch.Tensor, torch.Tensor],  # (weight, weight_scale)
    l2_tuple: tuple[torch.Tensor, torch.Tensor],  # (weight, weight_scale)
    l1_weight_scale_2: torch.Tensor = None,       # float32 global scale for L1
    l2_weight_scale_2: torch.Tensor = None,       # float32 global scale for L2
) -> tuple[tuple[torch.Tensor, torch.Tensor, torch.Tensor],
           tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
    """Transform NVFP4 weights for the CUTLASS block-scaled GEMM.

    NO LONGER FOLDS GLOBAL SCALES INTO BLOCK SCALES.
    Folding block_sf (float8) × global_sf (float32) → float8 loses ~25% precision
    because the product lands in the low-precision zone of float8_e4m3fn.
    Instead, global scales are returned separately and applied as per-expert GEMM alpha.

    Args:
        l1_tuple: (w13_weight, w13_weight_scale) — gate_up proj
        l2_tuple: (w2_weight, w2_weight_scale) — down proj
        l1_weight_scale_2: global scale for L1 (float32)
            Shape (E, 2) for gate+up, or (E,) per-expert, or scalar
        l2_weight_scale_2: global scale for L2 (float32)
            Shape (E,) per-expert, or scalar

    Returns:
        ((l1_weight, l1_sf, l1_global_sf), (l2_weight, l2_sf, l2_global_sf))
        where global_sf is (E,) float32 — the geometric mean of gate/up for L1,
        or the per-expert global scale for L2.
        The caller must apply global_sf as a per-expert multiplier to the GEMM alpha.
    """
    l1_weight, l1_weight_scale = l1_tuple
    l2_weight, l2_weight_scale = l2_tuple

    # Extract global scales as per-expert float32 vectors
    # L1: gate/up have separate global scales — store both
    # The caller (nvfp4_mega_moe_full) will apply the right one per-expert
    if l1_weight_scale_2 is not None:
        l1_gs = l1_weight_scale_2.to(torch.float32)
        if l1_gs.dim() == 2 and l1_gs.shape[1] == 2:
            # (E, 2) — gate_gs and up_gs separate
            # For L1 alpha, use the geometric mean (close enough since gate and up
            # global scales are typically similar). Actually, we need BOTH because
            # the GEMM produces gate and up in one shot.
            # Better: just store (E, 2) and let the caller apply post-GEMM scaling.
            l1_global_sf = l1_gs  # (E, 2) float32
        else:
            l1_global_sf = l1_gs  # (E,) float32
    else:
        l1_global_sf = torch.ones(l1_weight.shape[0], dtype=torch.float32, device=l1_weight.device)

    if l2_weight_scale_2 is not None:
        l2_gs = l2_weight_scale_2.to(torch.float32)
        l2_global_sf = l2_gs  # (E,) or scalar → broadcast to (E,)
        if l2_global_sf.dim() == 0:
            l2_global_sf = l2_global_sf.expand(l2_weight.shape[0])
    else:
        l2_global_sf = torch.ones(l2_weight.shape[0], dtype=torch.float32, device=l2_weight.device)

    # Debug: one-time diagnostic
    if not getattr(transform_nvfp4_weights_for_mega_moe, '_diag', False):
        transform_nvfp4_weights_for_mega_moe._diag = True
        print(f"[WT-XFORM] L1 block_sf range=[{l1_weight_scale.float().min():.4e}, "
              f"{l1_weight_scale.float().max():.4e}] unique={torch.unique(l1_weight_scale.view(torch.uint8)).numel()}")
        print(f"[WT-XFORM] L1 global_sf: shape={tuple(l1_global_sf.shape)} "
              f"range=[{l1_global_sf.min():.4e}, {l1_global_sf.max():.4e}]")
        print(f"[WT-XFORM] L2 block_sf range=[{l2_weight_scale.float().min():.4e}, "
              f"{l2_weight_scale.float().max():.4e}] unique={torch.unique(l2_weight_scale.view(torch.uint8)).numel()}")
        print(f"[WT-XFORM] L2 global_sf: shape={tuple(l2_global_sf.shape)} "
              f"range=[{l2_global_sf.min():.4e}, {l2_global_sf.max():.4e}]")

    # Block scales stay as original float8 — NO FOLDING
    l1_sf_out = l1_weight_scale.contiguous()
    l2_sf_out = l2_weight_scale.contiguous()

    # CUTLASS B is declared ColumnMajor — it expects (K, N) in memory.
    # Checkpoint weights are (N, K_half) row-major, so we transpose to (K_half, N)
    l1_weight_out = l1_weight.transpose(-2, -1).contiguous()
    l2_weight_out = l2_weight.transpose(-2, -1).contiguous()

    # Same for scale factors: (N, sf_k) row-major → (sf_k, N) column-major
    l1_sf_out = l1_sf_out.transpose(-2, -1).contiguous()
    l2_sf_out = l2_sf_out.transpose(-2, -1).contiguous()

    return (l1_weight_out, l1_sf_out, l1_global_sf), (l2_weight_out, l2_sf_out, l2_global_sf)