h3-blackwell-runtime/tools/validate_cute_nvfp4_real_tiles.py

"""Compare the CuTe tile producer with every tile of a real H3 activation."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import cutlass
import cutlass.cute as cute
import cutlass.torch as cutlass_torch
import torch
from cutlass.cute.runtime import from_dlpack

from h3_blackwell_runtime.nvfp4_quant import vortex_quantize_nvfp4
from profile_nvfp4_linear import module_for_name, representative_inputs
from validate_cute_nvfp4_tile_producer import (
    BLOCKS_PER_ROW,
    JOBS,
    TILE,
    produce_tile,
    unswizzle_scales,
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument(
        "--linear",
        choices=("attn_qkv_proj", "attn_out_proj", "mlp_fc1"),
        required=True,
    )
    parser.add_argument("--model-path", default="/models/minimax_h3_fl2va_pruned_nvfp4.safetensors")
    parser.add_argument("--block-index", type=int, default=24)
    parser.add_argument("--width", type=int, default=1344)
    parser.add_argument("--height", type=int, default=768)
    parser.add_argument("--frames", type=int, default=124)
    parser.add_argument("--steps", type=int, default=12)
    parser.add_argument("--sampler-step", type=int, default=1)
    parser.add_argument("--seed", type=int, default=440420)
    parser.add_argument("--text-tokens", type=int, default=100)
    parser.add_argument("--attention", default="sage2")
    parser.add_argument("--device", default="cuda")
    return parser.parse_args()


def main() -> None:
    from comfy_kitchen.tensor import TensorCoreNVFP4Layout

    args = parse_args()
    block, inputs, metadata = representative_inputs(args)
    linear = module_for_name(block, args.linear)
    activation = inputs[args.linear].reshape(-1, linear.in_features)[:TILE].contiguous()
    if activation.shape != (TILE, linear.in_features) or activation.shape[1] % TILE:
        raise ValueError(f"Expected a 128-row activation with K divisible by 128, got {activation.shape}")

    packed = vortex_quantize_nvfp4(activation)
    expected_qdata, tensor_scale, expected_physical_scales = (
        TensorCoreNVFP4Layout.get_plain_tensors(packed)
    )
    expected_fp4 = ((expected_qdata & 0x0F) << 4) | ((expected_qdata & 0xF0) >> 4)
    expected_scales = unswizzle_scales(expected_physical_scales.view(torch.uint8))

    source_torch = torch.empty(TILE, TILE, device="cuda", dtype=torch.bfloat16)
    source = from_dlpack(source_torch, assumed_align=16).mark_layout_dynamic(leading_dim=1)
    scale_torch = tensor_scale.float().reshape(1).contiguous()
    scale = from_dlpack(scale_torch, assumed_align=4).mark_layout_dynamic()
    fp4, fp4_torch = cutlass_torch.cute_tensor_like(
        torch.zeros_like(source_torch, dtype=torch.float32),
        cutlass.Float4E2M1FN,
        is_dynamic_layout=True,
        assumed_align=16,
    )
    block_scales_torch = torch.zeros(JOBS, 8, device="cuda", dtype=torch.uint8)
    block_scales = from_dlpack(block_scales_torch.flatten(), assumed_align=16)
    block_scales.element_type = cutlass.Float8E4M3FN
    block_scales = block_scales.mark_layout_dynamic()
    scalar_scales_torch = torch.zeros(JOBS, device="cuda", dtype=torch.uint8)
    scalar_scales = from_dlpack(scalar_scales_torch, assumed_align=16)
    scalar_scales.element_type = cutlass.Float8E4M3FN
    scalar_scales = scalar_scales.mark_layout_dynamic()

    compiled = cute.compile(produce_tile, source, scale, fp4, block_scales, scalar_scales)
    tile_reports = []
    examples = []
    total_fp4_differences = 0
    total_scale_differences = 0
    for k_start in range(0, activation.shape[1], TILE):
        source_torch.copy_(activation[:, k_start : k_start + TILE])
        compiled(source, scale, fp4, block_scales, scalar_scales)
        torch.cuda.synchronize()
        actual_fp4 = fp4_torch.view(torch.uint8).flatten()[: TILE * TILE // 2].reshape(TILE, TILE // 2)
        actual_scales = scalar_scales_torch.reshape(TILE, BLOCKS_PER_ROW)
        fp4_reference = expected_fp4[:, k_start // 2 : (k_start + TILE) // 2]
        scale_reference = expected_scales[:, k_start // 16 : (k_start + TILE) // 16]
        fp4_differences = int((actual_fp4 != fp4_reference).sum().item())
        scale_differences = int((actual_scales != scale_reference).sum().item())
        total_fp4_differences += fp4_differences
        total_scale_differences += scale_differences
        if fp4_differences or scale_differences:
            tile_reports.append({
                "k_start": k_start,
                "fp4_difference_count": fp4_differences,
                "block_scale_difference_count": scale_differences,
            })
            if fp4_differences and len(examples) < 20:
                for row, packed_column in (actual_fp4 != fp4_reference).nonzero().tolist():
                    global_column = k_start + packed_column * 2
                    block_scale_byte = scale_reference[row, packed_column // 8].reshape(1)
                    decoded_scale = block_scale_byte.view(torch.float8_e4m3fn).float()
                    encode_scale = torch.minimum(
                        torch.ones_like(decoded_scale) / (decoded_scale * scale_torch),
                        torch.full_like(decoded_scale, torch.finfo(torch.float32).max),
                    )
                    normalized = activation[row, global_column : global_column + 2].float() * encode_scale
                    examples.append({
                        "row": row,
                        "global_column": global_column,
                        "source": [
                            float(activation[row, global_column].float().item()),
                            float(activation[row, global_column + 1].float().item()),
                        ],
                        "actual_byte": int(actual_fp4[row, packed_column].item()),
                        "expected_byte": int(fp4_reference[row, packed_column].item()),
                        "block_scale_byte": int(block_scale_byte.item()),
                        "decoded_block_scale": float(decoded_scale.item()),
                        "encode_scale": float(encode_scale.item()),
                        "torch_normalized": normalized.tolist(),
                    })
                    if len(examples) == 20:
                        break

    report = {
        "device": torch.cuda.get_device_name(),
        "cutlass_dsl": "4.6.2",
        "metadata": metadata,
        "linear": args.linear,
        "activation_shape": list(activation.shape),
        "tensor_scale": scale_torch.item(),
        "tile_count": activation.shape[1] // TILE,
        "fp4_difference_count": total_fp4_differences,
        "block_scale_difference_count": total_scale_differences,
        "equal": total_fp4_differences == 0 and total_scale_differences == 0,
        "differing_tiles": tile_reports,
        "difference_examples": examples,
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(report, indent=2), flush=True)


if __name__ == "__main__":
    main()