154 lines
7.1 KiB
Python
154 lines
7.1 KiB
Python
"""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()
|