"""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()