diff --git a/tools/validate_cute_nvfp4_real_tiles.py b/tools/validate_cute_nvfp4_real_tiles.py new file mode 100644 index 0000000..be72bd2 --- /dev/null +++ b/tools/validate_cute_nvfp4_real_tiles.py @@ -0,0 +1,154 @@ +"""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()