160 lines
7.4 KiB
Diff
160 lines
7.4 KiB
Diff
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()
|