272 lines
10 KiB
Python
272 lines
10 KiB
Python
|
"""Validate and time a bounded 128-row NVFP4 packed-tile ring prototype."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import json
|
||
|
|
import math
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
import cutlass
|
||
|
|
import cutlass.cute as cute
|
||
|
|
import cutlass.torch as cutlass_torch
|
||
|
|
import torch
|
||
|
|
import torch.nn.functional as functional
|
||
|
|
from cutlass.cute.runtime import from_dlpack
|
||
|
|
|
||
|
|
from h3_blackwell_runtime.nvfp4_quant import (
|
||
|
|
nvfp4_activation_scale,
|
||
|
|
vortex_native_quantize_nvfp4,
|
||
|
|
vortex_native_quantize_nvfp4_into,
|
||
|
|
vortex_quantize_nvfp4,
|
||
|
|
)
|
||
|
|
from profile_nvfp4_linear import module_for_name, representative_inputs
|
||
|
|
from validate_cute_nvfp4_h3 import (
|
||
|
|
fp4_tensor,
|
||
|
|
load_cutlass_example,
|
||
|
|
output_tensor,
|
||
|
|
scale_tensor,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def parse_args() -> argparse.Namespace:
|
||
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
||
|
|
parser.add_argument("--cutlass-example", type=Path, required=True)
|
||
|
|
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("--warmup", type=int, default=5)
|
||
|
|
parser.add_argument("--iterations", type=int, default=20)
|
||
|
|
parser.add_argument("--rows", type=int, default=128)
|
||
|
|
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 measure_cuda(fn, *, warmup: int, iterations: int) -> tuple[object, float]:
|
||
|
|
result = None
|
||
|
|
for _ in range(warmup):
|
||
|
|
result = fn()
|
||
|
|
torch.cuda.synchronize()
|
||
|
|
started = torch.cuda.Event(enable_timing=True)
|
||
|
|
finished = torch.cuda.Event(enable_timing=True)
|
||
|
|
started.record()
|
||
|
|
for _ in range(iterations):
|
||
|
|
result = fn()
|
||
|
|
finished.record()
|
||
|
|
finished.synchronize()
|
||
|
|
return result, started.elapsed_time(finished) / iterations
|
||
|
|
|
||
|
|
|
||
|
|
def main() -> None:
|
||
|
|
from comfy_kitchen.tensor import TensorCoreNVFP4Layout
|
||
|
|
import comfy_kitchen as ck
|
||
|
|
|
||
|
|
args = parse_args()
|
||
|
|
if args.warmup < 0 or args.iterations <= 0:
|
||
|
|
raise ValueError("--warmup must be non-negative and --iterations must be positive")
|
||
|
|
if args.rows <= 0 or args.rows % 128:
|
||
|
|
raise ValueError("--rows must be a positive multiple of 128")
|
||
|
|
|
||
|
|
example = load_cutlass_example(args.cutlass_example, fuse_alpha=True)
|
||
|
|
block, inputs, metadata = representative_inputs(args)
|
||
|
|
linear = module_for_name(block, args.linear)
|
||
|
|
full_activation = inputs[args.linear].reshape(-1, linear.in_features).contiguous()
|
||
|
|
if args.rows > full_activation.shape[0]:
|
||
|
|
raise ValueError(f"--rows exceeds the available {full_activation.shape[0]} rows")
|
||
|
|
activation = full_activation[:args.rows].contiguous()
|
||
|
|
if activation.shape[1] % 128:
|
||
|
|
raise ValueError(f"Ring prototype requires K divisible by 128, got {activation.shape}")
|
||
|
|
global_scale = nvfp4_activation_scale(full_activation).float()
|
||
|
|
|
||
|
|
with torch.inference_mode():
|
||
|
|
expected_packed = vortex_quantize_nvfp4(activation, scale=global_scale)
|
||
|
|
ring_packed = vortex_native_quantize_nvfp4(activation, scale=global_scale)
|
||
|
|
packed_weight = linear._packed_weight()
|
||
|
|
expected_qdata, expected_tensor_scale, expected_block_scales = (
|
||
|
|
TensorCoreNVFP4Layout.get_plain_tensors(expected_packed)
|
||
|
|
)
|
||
|
|
ring_qdata, ring_tensor_scale, ring_block_scales = (
|
||
|
|
TensorCoreNVFP4Layout.get_plain_tensors(ring_packed)
|
||
|
|
)
|
||
|
|
b_qdata, tensor_scale_b, b_block_scales = (
|
||
|
|
TensorCoreNVFP4Layout.get_plain_tensors(packed_weight)
|
||
|
|
)
|
||
|
|
reference = functional.linear(expected_packed, packed_weight, None)
|
||
|
|
|
||
|
|
a, _ = fp4_tensor(ring_qdata, swap_nibbles=False, reencode=True)
|
||
|
|
b, _ = fp4_tensor(b_qdata, swap_nibbles=False, reencode=True)
|
||
|
|
sfa = scale_tensor(ring_block_scales)
|
||
|
|
sfb = scale_tensor(b_block_scales)
|
||
|
|
output_bf16 = torch.zeros(
|
||
|
|
args.rows, b_qdata.shape[0], device="cuda", dtype=torch.bfloat16,
|
||
|
|
)
|
||
|
|
c = output_tensor(output_bf16)
|
||
|
|
alpha = ring_tensor_scale.float() * tensor_scale_b.float()
|
||
|
|
alpha_argument = from_dlpack(alpha.reshape(1).contiguous(), assumed_align=4)
|
||
|
|
gemm = example.Sm120BlockScaledGemmKernel(
|
||
|
|
cutlass.Float32, 16, (128, 128, 128), (128, 128),
|
||
|
|
)
|
||
|
|
max_active_clusters = cutlass.utils.HardwareInfo().get_max_active_clusters(1)
|
||
|
|
stream = cutlass_torch.default_stream()
|
||
|
|
compiled_gemm = cute.compile(
|
||
|
|
gemm, a, b, sfa, sfb, c, alpha_argument, max_active_clusters, stream,
|
||
|
|
)
|
||
|
|
|
||
|
|
def run_consumer():
|
||
|
|
return compiled_gemm(a, b, sfa, sfb, c, alpha_argument, stream)
|
||
|
|
|
||
|
|
def run_pack_into():
|
||
|
|
return vortex_native_quantize_nvfp4_into(
|
||
|
|
activation, global_scale, ring_qdata, ring_block_scales,
|
||
|
|
)
|
||
|
|
|
||
|
|
def run_cute_ring():
|
||
|
|
run_pack_into()
|
||
|
|
return run_consumer()
|
||
|
|
|
||
|
|
def run_comfy_consumer():
|
||
|
|
return ck.scaled_mm_nvfp4(
|
||
|
|
ring_qdata,
|
||
|
|
b_qdata,
|
||
|
|
tensor_scale_a=ring_tensor_scale,
|
||
|
|
tensor_scale_b=tensor_scale_b,
|
||
|
|
block_scale_a=ring_block_scales,
|
||
|
|
block_scale_b=b_block_scales,
|
||
|
|
out_dtype=torch.bfloat16,
|
||
|
|
alpha=ring_tensor_scale.float() * tensor_scale_b.float(),
|
||
|
|
)
|
||
|
|
|
||
|
|
def run_comfy_ring():
|
||
|
|
run_pack_into()
|
||
|
|
return run_comfy_consumer()
|
||
|
|
|
||
|
|
run_consumer()
|
||
|
|
torch.cuda.synchronize()
|
||
|
|
candidate = output_bf16[:, : linear.out_features]
|
||
|
|
delta = candidate.float() - reference.float()
|
||
|
|
|
||
|
|
_, scale_ms = measure_cuda(
|
||
|
|
lambda: nvfp4_activation_scale(full_activation),
|
||
|
|
warmup=args.warmup,
|
||
|
|
iterations=args.iterations,
|
||
|
|
)
|
||
|
|
_, producer_ms = measure_cuda(
|
||
|
|
run_pack_into,
|
||
|
|
warmup=args.warmup,
|
||
|
|
iterations=args.iterations,
|
||
|
|
)
|
||
|
|
_, consumer_ms = measure_cuda(
|
||
|
|
run_consumer,
|
||
|
|
warmup=args.warmup,
|
||
|
|
iterations=args.iterations,
|
||
|
|
)
|
||
|
|
_, actual_ring_ms = measure_cuda(
|
||
|
|
run_cute_ring,
|
||
|
|
warmup=args.warmup,
|
||
|
|
iterations=args.iterations,
|
||
|
|
)
|
||
|
|
comfy_candidate, comfy_consumer_ms = measure_cuda(
|
||
|
|
run_comfy_consumer,
|
||
|
|
warmup=args.warmup,
|
||
|
|
iterations=args.iterations,
|
||
|
|
)
|
||
|
|
_, comfy_ring_ms = measure_cuda(
|
||
|
|
run_comfy_ring,
|
||
|
|
warmup=args.warmup,
|
||
|
|
iterations=args.iterations,
|
||
|
|
)
|
||
|
|
_, reference_ms = measure_cuda(
|
||
|
|
lambda: functional.linear(
|
||
|
|
vortex_quantize_nvfp4(activation, scale=global_scale),
|
||
|
|
packed_weight,
|
||
|
|
None,
|
||
|
|
),
|
||
|
|
warmup=args.warmup,
|
||
|
|
iterations=args.iterations,
|
||
|
|
)
|
||
|
|
|
||
|
|
qdata_bytes = ring_qdata.numel() * ring_qdata.element_size()
|
||
|
|
sfa_bytes = ring_block_scales.numel() * ring_block_scales.element_size()
|
||
|
|
chunk_count = math.ceil(full_activation.shape[0] / args.rows)
|
||
|
|
modeled_cute_chunk_ms = producer_ms + consumer_ms
|
||
|
|
modeled_comfy_chunk_ms = producer_ms + comfy_consumer_ms
|
||
|
|
modeled_reference_canonical_ms = scale_ms + chunk_count * reference_ms
|
||
|
|
report = {
|
||
|
|
"device": torch.cuda.get_device_name(),
|
||
|
|
"cutlass_dsl": "4.6.2",
|
||
|
|
"metadata": metadata,
|
||
|
|
"linear": args.linear,
|
||
|
|
"mnk": [args.rows, linear.out_features, activation.shape[1]],
|
||
|
|
"full_activation_rows": full_activation.shape[0],
|
||
|
|
"modeled_chunk_count": chunk_count,
|
||
|
|
"ring": {
|
||
|
|
"row_capacity": args.rows,
|
||
|
|
"producer": "vortex_native_quantize_nvfp4",
|
||
|
|
"qdata_bytes": qdata_bytes,
|
||
|
|
"sfa_bytes": sfa_bytes,
|
||
|
|
"logical_bytes": qdata_bytes + sfa_bytes,
|
||
|
|
},
|
||
|
|
"parity": {
|
||
|
|
"tensor_scale_equal": torch.equal(
|
||
|
|
ring_tensor_scale, expected_tensor_scale,
|
||
|
|
),
|
||
|
|
"fp4_difference_count": int(
|
||
|
|
(ring_qdata != expected_qdata).sum().item()
|
||
|
|
),
|
||
|
|
"block_scale_difference_count": int(
|
||
|
|
(ring_block_scales.view(torch.uint8)
|
||
|
|
!= expected_block_scales.view(torch.uint8)).sum().item()
|
||
|
|
),
|
||
|
|
"output_equal": torch.equal(candidate, reference),
|
||
|
|
"max_abs": delta.abs().max().item(),
|
||
|
|
"mean_abs": delta.abs().mean().item(),
|
||
|
|
"comfy_output_equal": torch.equal(
|
||
|
|
comfy_candidate[: args.rows, : linear.out_features], reference,
|
||
|
|
),
|
||
|
|
},
|
||
|
|
"timing": {
|
||
|
|
"warmup": args.warmup,
|
||
|
|
"iterations": args.iterations,
|
||
|
|
"producer_ms": producer_ms,
|
||
|
|
"global_scale_ms": scale_ms,
|
||
|
|
"cute_consumer_ms": consumer_ms,
|
||
|
|
"modeled_cute_chunk_ms": modeled_cute_chunk_ms,
|
||
|
|
"modeled_comfy_chunk_ms": modeled_comfy_chunk_ms,
|
||
|
|
"actual_into_ring_cute_gemm_ms": actual_ring_ms,
|
||
|
|
"comfy_consumer_ms": comfy_consumer_ms,
|
||
|
|
"actual_into_ring_comfy_gemm_ms": comfy_ring_ms,
|
||
|
|
"reference_vortex_scale_comfy_pack_gemm_ms": reference_ms,
|
||
|
|
"modeled_canonical_reference_ms": modeled_reference_canonical_ms,
|
||
|
|
"modeled_canonical_cute_ring_ms": scale_ms + chunk_count * actual_ring_ms,
|
||
|
|
"modeled_canonical_comfy_ring_ms": scale_ms + chunk_count * comfy_ring_ms,
|
||
|
|
"actual_ring_vs_reference": actual_ring_ms / reference_ms,
|
||
|
|
"comfy_ring_vs_reference": comfy_ring_ms / reference_ms,
|
||
|
|
"modeled_canonical_cute_ring_vs_reference": (
|
||
|
|
scale_ms + chunk_count * actual_ring_ms
|
||
|
|
) / modeled_reference_canonical_ms,
|
||
|
|
"modeled_canonical_comfy_ring_vs_reference": (
|
||
|
|
scale_ms + chunk_count * comfy_ring_ms
|
||
|
|
) / modeled_reference_canonical_ms,
|
||
|
|
},
|
||
|
|
}
|
||
|
|
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()
|