537 lines
26 KiB
Python
537 lines
26 KiB
Python
"""Compare the CUTLASS DSL SM121 block-scaled GEMM with real H3 NVFP4 tensors."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import importlib.util
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
import torch.nn.functional as functional
|
|
|
|
from h3_blackwell_runtime.nvfp4_quant import vortex_quantize_nvfp4
|
|
from profile_nvfp4_linear import module_for_name, representative_inputs
|
|
|
|
|
|
def _replace_once(source: str, old: str, new: str, name: str) -> str:
|
|
if source.count(old) != 1:
|
|
raise RuntimeError(
|
|
f"CUTLASS {name} patch expected one occurrence, found {source.count(old)}: {old!r}"
|
|
)
|
|
return source.replace(old, new)
|
|
|
|
|
|
def _patch_streaming_a(source: str) -> str:
|
|
source = _replace_once(
|
|
source,
|
|
"import cutlass\nimport cutlass.cute as cute",
|
|
"""import cutlass
|
|
import cutlass.cute as cute
|
|
from cutlass import Float32
|
|
from cutlass.cutlass_dsl import dsl_user_op
|
|
from cutlass._mlir import ir
|
|
from cutlass._mlir.dialects import llvm
|
|
|
|
|
|
@dsl_user_op
|
|
def vortex_rcp_approx_ftz_f32(
|
|
x: Float32,
|
|
*,
|
|
loc: Optional[ir.Location] = None,
|
|
ip: Optional[ir.InsertionPoint] = None,
|
|
) -> Float32:
|
|
result = llvm.inline_asm(
|
|
Float32.mlir_type,
|
|
[x.ir_value(loc=loc, ip=ip)],
|
|
"rcp.approx.ftz.f32 $0, $1;",
|
|
"=f,f",
|
|
has_side_effects=False,
|
|
asm_dialect=0,
|
|
loc=loc,
|
|
ip=ip,
|
|
)
|
|
return Float32(result)""",
|
|
"streaming-A reciprocal",
|
|
)
|
|
replacements = (
|
|
(
|
|
" self.a_dtype = a.element_type\n self.b_dtype = b.element_type\n self.c_dtype = c.element_type\n self.sf_dtype = sfa.element_type",
|
|
" self.a_source_dtype = a.element_type\n self.a_dtype = cutlass.Float4E2M1FN\n self.b_dtype = b.element_type\n self.c_dtype = c.element_type\n self.sf_dtype = cutlass.Float8E4M3FN",
|
|
),
|
|
(
|
|
" self.sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(\n a.shape, self.sf_vec_size\n )\n sfa_tensor = cute.make_tensor(sfa.iterator, self.sfa_layout)\n",
|
|
"",
|
|
),
|
|
(
|
|
" tma_atom_a, tma_tensor_a = self._make_tma_atoms_and_tensors(\n a,\n self.a_smem_layout_staged,\n (self.tile_shape_mnk[0], self.tile_shape_mnk[2]),\n 1,\n internal_type=self.tma_internal_a_dtype,\n )\n\n",
|
|
"",
|
|
),
|
|
(
|
|
" tma_atom_sfa, tma_tensor_sfa = self._make_tma_atoms_and_tensors(\n sfa_tensor,\n self.sfa_smem_layout_staged,\n (self.tile_shape_mnk[0], self.tile_shape_mnk[2]),\n 1,\n internal_type=cutlass.Int16,\n )\n\n",
|
|
"",
|
|
),
|
|
(
|
|
" self.kernel(\n tma_atom_a,\n tma_tensor_a,\n tma_atom_b,\n tma_tensor_b,\n tma_atom_sfa,\n tma_tensor_sfa,\n tma_atom_sfb,",
|
|
" self.kernel(\n a,\n sfa,\n tma_atom_b,\n tma_tensor_b,\n tma_atom_sfb,",
|
|
),
|
|
(
|
|
" tma_atom_a: cute.CopyAtom,\n mA_mkl: cute.Tensor,\n tma_atom_b: cute.CopyAtom,\n mB_nkl: cute.Tensor,\n tma_atom_sfa: cute.CopyAtom,\n mSFA_mkl: cute.Tensor,\n tma_atom_sfb: cute.CopyAtom,",
|
|
" mA_mkl: cute.Tensor,\n tensor_scale_a: cute.Tensor,\n tma_atom_b: cute.CopyAtom,\n mB_nkl: cute.Tensor,\n tma_atom_sfb: cute.CopyAtom,",
|
|
),
|
|
(
|
|
" cpasync.prefetch_descriptor(tma_atom_a)\n cpasync.prefetch_descriptor(tma_atom_b)\n cpasync.prefetch_descriptor(tma_atom_sfa)\n cpasync.prefetch_descriptor(tma_atom_sfb)",
|
|
" cpasync.prefetch_descriptor(tma_atom_b)\n cpasync.prefetch_descriptor(tma_atom_sfb)",
|
|
),
|
|
(
|
|
" tma_copy_bytes = (\n cute.size_in_bytes(self.a_dtype, a_smem_layout)\n + cute.size_in_bytes(self.b_dtype, b_smem_layout)\n + cute.size_in_bytes(self.sf_dtype, sfa_smem_layout)\n + cute.size_in_bytes(self.sf_dtype, sfb_smem_layout)\n )",
|
|
" tma_copy_bytes = (\n cute.size_in_bytes(self.b_dtype, b_smem_layout)\n + cute.size_in_bytes(self.sf_dtype, sfb_smem_layout)\n )",
|
|
),
|
|
(
|
|
" # (tM, tK, loopM, loopK, loopL)\n gSFA_mkl = cute.local_tile(\n mSFA_mkl,\n cute.slice_(self.tile_shape_mnk, (None, 0, None)),\n (None, None, None),\n )\n",
|
|
"",
|
|
),
|
|
(
|
|
" # TMA load A partition_S/D\n a_cta_layout = cute.make_layout(cute.slice_(cta_layout_mnk, (0, None, 0)).shape)\n a_cta_crd = cluster_coord_mnk[1]\n tAsA, tAgA = cpasync.tma_partition(\n tma_atom_a,\n a_cta_crd,\n a_cta_layout,\n cute.group_modes(sA, 0, 2),\n cute.group_modes(gA_mkl, 0, 2),\n )\n\n",
|
|
"",
|
|
),
|
|
(
|
|
" tAsSFA, tAgSFA = cpasync.tma_partition(\n tma_atom_sfa,\n a_cta_crd,\n a_cta_layout,\n cute.group_modes(sSFA, 0, 2),\n cute.group_modes(gSFA_mkl, 0, 2),\n )\n tAsSFA = cute.filter_zeros(tAsSFA)\n tAgSFA = cute.filter_zeros(tAgSFA)\n\n",
|
|
"",
|
|
),
|
|
(
|
|
" cute.arch.setmaxregister_decrease(self.load_register_requirement)\n\n while work_tile.is_valid_tile:",
|
|
""" cute.arch.setmaxregister_decrease(self.load_register_requirement)
|
|
producer_lane = tidx - self.tma_load_warp_id * self.num_threads_per_warp
|
|
fp4_store = cute.make_copy_atom(
|
|
cute.nvgpu.CopyUniversalOp(), cutlass.Float4E2M1FN
|
|
)
|
|
source_fragment = cute.make_rmem_tensor((16,), cutlass.Float32)
|
|
normalized_fragment = cute.make_rmem_tensor((8,), cutlass.Float32)
|
|
fp4_fragment = cute.make_rmem_tensor((8,), cutlass.Float4E2M1FN)
|
|
scale_source = cute.make_rmem_tensor((8,), cutlass.Float32)
|
|
scale_fragment = cute.make_rmem_tensor((8,), cutlass.Float8E4M3FN)
|
|
decoded_scale_fragment = cute.make_rmem_tensor((8,), cutlass.Float32)
|
|
while work_tile.is_valid_tile:""",
|
|
),
|
|
(
|
|
" tAgA_mkl = tAgA[(None, tile_coord_mnl[0], None, tile_coord_mnl[2])]\n tBgB_nkl = tBgB[(None, tile_coord_mnl[1], None, tile_coord_mnl[2])]\n tAgSFA_mkl = tAgSFA[(None, tile_coord_mnl[0], None, tile_coord_mnl[2])]\n tBgSFB_nkl = tBgSFB[(None, tile_coord_mnl[1], None, tile_coord_mnl[2])]",
|
|
" tBgB_nkl = tBgB[(None, tile_coord_mnl[1], None, tile_coord_mnl[2])]\n tBgSFB_nkl = tBgSFB[(None, tile_coord_mnl[1], None, tile_coord_mnl[2])]",
|
|
),
|
|
(
|
|
" tAgA_k = tAgA_mkl[(None, mainloop_producer_state.count)]\n tAsA_pipe = tAsA[(None, mainloop_producer_state.index)]\n\n tBgB_k = tBgB_nkl[(None, mainloop_producer_state.count)]\n tBsB_pipe = tBsB[(None, mainloop_producer_state.index)]\n\n tAgSFA_k = tAgSFA_mkl[(None, mainloop_producer_state.count)]\n tAsSFA_pipe = tAsSFA[(None, mainloop_producer_state.index)]\n\n tBgSFB_k = tBgSFB_nkl[(None, mainloop_producer_state.count)]",
|
|
" tBgB_k = tBgB_nkl[(None, mainloop_producer_state.count)]\n tBsB_pipe = tBsB[(None, mainloop_producer_state.index)]\n\n tBgSFB_k = tBgSFB_nkl[(None, mainloop_producer_state.count)]",
|
|
),
|
|
(
|
|
""" cute.copy(
|
|
tma_atom_a,
|
|
tAgA_k,
|
|
tAsA_pipe,
|
|
tma_bar_ptr=mainloop_pipeline.producer_get_barrier(
|
|
mainloop_producer_state
|
|
),
|
|
)
|
|
""",
|
|
""" scale = tensor_scale_a[0]
|
|
stage = mainloop_producer_state.index
|
|
source_k_base = mainloop_producer_state.count * self.tile_shape_mnk[2]
|
|
source_m_base = tile_coord_mnl[0] * self.tile_shape_mnk[0]
|
|
for row_group in cutlass.range_constexpr(4):
|
|
row = producer_lane + row_group * self.num_threads_per_warp
|
|
source_row = source_m_base + row
|
|
sA_row = sA[(row, None, stage)]
|
|
sA_tiles = cute.zipped_divide(sA_row, (8,))
|
|
for block_column in cutlass.range_constexpr(8):
|
|
source_column = source_k_base + block_column * 16
|
|
maximum = cutlass.Float32(0.0)
|
|
for element in cutlass.range_constexpr(16):
|
|
value = mA_mkl[
|
|
source_row,
|
|
source_column + element,
|
|
tile_coord_mnl[2],
|
|
]
|
|
source_fragment[element] = value
|
|
maximum = cutlass.max(
|
|
cutlass.max(value, -value), maximum
|
|
)
|
|
|
|
raw_block_scale = (
|
|
maximum / cutlass.Float32(6.0)
|
|
) / scale
|
|
for element in cutlass.range_constexpr(8):
|
|
scale_source[element] = raw_block_scale
|
|
scale_values = scale_source.load()
|
|
scale_values = cute.where(
|
|
scale_values <= cutlass.Float32(448.0),
|
|
scale_values,
|
|
cutlass.Float32(448.0),
|
|
)
|
|
scale_fragment.store(
|
|
scale_values.to(cutlass.Float8E4M3FN)
|
|
)
|
|
sSFA[row, block_column * 16, stage] = scale_fragment[0]
|
|
decoded_scale_fragment.store(
|
|
scale_fragment.load().to(cutlass.Float32)
|
|
)
|
|
raw_encode_scale = vortex_rcp_approx_ftz_f32(
|
|
decoded_scale_fragment[0] * scale
|
|
)
|
|
for element in cutlass.range_constexpr(8):
|
|
scale_source[element] = raw_encode_scale
|
|
encode_scale_values = scale_source.load()
|
|
encode_scale_values = cute.where(
|
|
encode_scale_values
|
|
<= cutlass.Float32(3.402823466e38),
|
|
encode_scale_values,
|
|
cutlass.Float32(3.402823466e38),
|
|
)
|
|
scale_source.store(encode_scale_values)
|
|
encode_scale = scale_source[0]
|
|
|
|
for half in cutlass.range_constexpr(2):
|
|
for element in cutlass.range_constexpr(8):
|
|
normalized_fragment[element] = source_fragment[
|
|
half * 8 + element
|
|
] * encode_scale
|
|
fp4_fragment.store(
|
|
normalized_fragment.load().to(
|
|
cutlass.Float4E2M1FN
|
|
)
|
|
)
|
|
cute.copy(
|
|
fp4_store,
|
|
fp4_fragment,
|
|
sA_tiles[(None, block_column * 2 + half)],
|
|
)
|
|
cute.arch.fence_proxy("async.shared", space="cta")
|
|
""",
|
|
),
|
|
(
|
|
""" cute.copy(
|
|
tma_atom_sfa,
|
|
tAgSFA_k,
|
|
tAsSFA_pipe,
|
|
tma_bar_ptr=mainloop_pipeline.producer_get_barrier(
|
|
mainloop_producer_state
|
|
),
|
|
)
|
|
""",
|
|
"",
|
|
),
|
|
)
|
|
for index, (old, new) in enumerate(replacements):
|
|
source = _replace_once(source, old, new, f"streaming-A[{index}]")
|
|
return source
|
|
|
|
|
|
def load_cutlass_example(path: Path, *, fuse_alpha: bool, stream_a: bool = False):
|
|
if not path.is_file():
|
|
raise FileNotFoundError(f"CUTLASS DSL example not found: {path}")
|
|
sys.path.insert(0, str(path.parent))
|
|
source = path.read_text(encoding="utf-8")
|
|
if stream_a:
|
|
source = _patch_streaming_a(source)
|
|
if fuse_alpha:
|
|
replacements = (
|
|
(" c: cute.Tensor,\n max_active_clusters: cutlass.Constexpr,", " c: cute.Tensor,\n alpha: cute.Tensor,\n max_active_clusters: cutlass.Constexpr,"),
|
|
(" mC_mnl: cute.Tensor,\n tiled_mma: cute.TiledMma,", " mC_mnl: cute.Tensor,\n alpha: cute.Tensor,\n tiled_mma: cute.TiledMma,"),
|
|
(" tma_tensor_c,\n self.tiled_mma,", " tma_tensor_c,\n alpha,\n self.tiled_mma,"),
|
|
(" tRS_rD_out.store(acc_vec.to(self.c_dtype))", " tRS_rD_out.store((acc_vec * alpha[0]).to(self.c_dtype))"),
|
|
)
|
|
for old, new in replacements:
|
|
source = _replace_once(source, old, new, "alpha")
|
|
if stream_a or fuse_alpha:
|
|
suffix = "_vortex_stream_a" if stream_a else "_vortex"
|
|
suffix += "_alpha" if fuse_alpha else ""
|
|
load_path = path.with_name(f"{path.stem}{suffix}.py")
|
|
load_path.write_text(source, encoding="utf-8")
|
|
else:
|
|
load_path = path
|
|
spec = importlib.util.spec_from_file_location("vortex_cutlass_blockscaled", load_path)
|
|
if spec is None or spec.loader is None:
|
|
raise ImportError(f"Cannot load CUTLASS DSL example: {load_path}")
|
|
module = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
def decode_comfy_fp4(storage: torch.Tensor) -> torch.Tensor:
|
|
lookup = torch.tensor(
|
|
[0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0],
|
|
device=storage.device,
|
|
dtype=torch.float32,
|
|
)
|
|
codes = torch.stack((storage >> 4, storage & 0x0F), dim=-1).reshape(storage.shape[0], -1)
|
|
return lookup[codes.long()]
|
|
|
|
|
|
def fp4_tensor(storage: torch.Tensor, *, swap_nibbles: bool, reencode: bool):
|
|
import cutlass
|
|
import cutlass.torch as cutlass_torch
|
|
from cutlass.cute.runtime import from_dlpack
|
|
|
|
rows, packed_columns = storage.shape
|
|
if reencode:
|
|
logical = decode_comfy_fp4(storage).unsqueeze(-1)
|
|
return cutlass_torch.cute_tensor_like(
|
|
logical, cutlass.Float4E2M1FN, is_dynamic_layout=True, assumed_align=16,
|
|
)
|
|
if swap_nibbles:
|
|
storage = ((storage & 0x0F) << 4) | ((storage & 0xF0) >> 4)
|
|
# DLPack cannot export Torch's packed FP4 dtype yet. Build the desired
|
|
# logical layout from an oversized uint8 allocation, then reinterpret its
|
|
# iterator as FP4 and populate only the packed storage that the layout uses.
|
|
backing = torch.empty(
|
|
(rows, packed_columns * 2, 1), device=storage.device, dtype=torch.uint8,
|
|
)
|
|
backing.zero_()
|
|
backing[:, :packed_columns, 0].copy_(storage)
|
|
tensor = from_dlpack(backing, assumed_align=16)
|
|
tensor.element_type = cutlass.Float4E2M1FN
|
|
tensor = tensor.mark_layout_dynamic(leading_dim=1)
|
|
return tensor, backing
|
|
|
|
|
|
def output_tensor(storage: torch.Tensor):
|
|
from cutlass.cute.runtime import from_dlpack
|
|
|
|
tensor = from_dlpack(storage.unsqueeze(-1), assumed_align=16)
|
|
tensor = tensor.mark_compact_shape_dynamic(mode=1, stride_order=(2, 0, 1), divisibility=1)
|
|
return tensor
|
|
|
|
|
|
def scale_tensor(storage: torch.Tensor):
|
|
import cutlass
|
|
from cutlass.cute.runtime import from_dlpack
|
|
|
|
tensor = from_dlpack(storage.view(torch.uint8).unsqueeze(-1), assumed_align=16)
|
|
tensor.element_type = cutlass.Float8E4M3FN
|
|
return tensor.mark_layout_dynamic(leading_dim=1)
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--cutlass-example", type=Path, required=True)
|
|
parser.add_argument("--model-path", default="/models/minimax_h3_fl2va_pruned_nvfp4.safetensors")
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument("--linear", choices=("attn_qkv_proj", "attn_out_proj", "mlp_fc1", "mlp_fc2"), default="attn_qkv_proj")
|
|
parser.add_argument("--rows", type=int, default=128)
|
|
parser.add_argument("--tile-k", type=int, choices=(128, 256), default=128)
|
|
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")
|
|
parser.add_argument("--swap-nibbles", action="store_true")
|
|
parser.add_argument("--reencode-fp4", action="store_true")
|
|
parser.add_argument("--zero-a", action="store_true")
|
|
parser.add_argument("--fuse-alpha", action="store_true")
|
|
parser.add_argument("--stream-a", action="store_true")
|
|
parser.add_argument("--warmup", type=int, default=0)
|
|
parser.add_argument("--iterations", type=int, default=0)
|
|
parser.add_argument("--benchmark-reference", action="store_true")
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
if args.rows <= 0 or args.rows % 128:
|
|
raise ValueError("--rows must be a positive multiple of 128")
|
|
if args.stream_a and args.linear == "mlp_fc2":
|
|
raise ValueError("--stream-a deliberately excludes mlp_fc2; retain the reference cuBLAS path")
|
|
if args.warmup < 0 or args.iterations < 0:
|
|
raise ValueError("--warmup and --iterations must be non-negative")
|
|
|
|
import cutlass
|
|
import cutlass.cute as cute
|
|
import cutlass.torch as cutlass_torch
|
|
import comfy_kitchen as ck
|
|
from cutlass.cute.runtime import from_dlpack
|
|
from comfy_kitchen.tensor import TensorCoreNVFP4Layout
|
|
|
|
example = load_cutlass_example(
|
|
args.cutlass_example, fuse_alpha=args.fuse_alpha, stream_a=args.stream_a,
|
|
)
|
|
block, inputs, metadata = representative_inputs(args)
|
|
linear = module_for_name(block, args.linear)
|
|
activation = inputs[args.linear].reshape(-1, linear.in_features)[:args.rows].contiguous()
|
|
|
|
with torch.inference_mode():
|
|
packed_activation = vortex_quantize_nvfp4(activation)
|
|
packed_weight = linear._packed_weight()
|
|
a_qdata, a_tensor_scale, a_block_scale = TensorCoreNVFP4Layout.get_plain_tensors(packed_activation)
|
|
b_qdata, b_tensor_scale, b_block_scale = TensorCoreNVFP4Layout.get_plain_tensors(packed_weight)
|
|
reference = functional.linear(packed_activation, packed_weight, None)[:args.rows, :linear.out_features]
|
|
raw_reference = ck.scaled_mm_nvfp4(
|
|
a_qdata,
|
|
b_qdata,
|
|
tensor_scale_a=a_tensor_scale,
|
|
tensor_scale_b=b_tensor_scale,
|
|
block_scale_a=a_block_scale,
|
|
block_scale_b=b_block_scale,
|
|
out_dtype=torch.bfloat16,
|
|
alpha=torch.ones(1, device=activation.device, dtype=torch.float32),
|
|
)[:args.rows, :linear.out_features]
|
|
if args.zero_a:
|
|
a_qdata = torch.zeros_like(a_qdata)
|
|
|
|
output_bf16 = torch.zeros((a_qdata.shape[0], b_qdata.shape[0]), device=activation.device, dtype=torch.bfloat16)
|
|
if args.stream_a:
|
|
if args.rows != 128 or activation.shape[1] % 128:
|
|
raise ValueError("--stream-a currently requires exactly 128 rows and K divisible by 128")
|
|
a = output_tensor(activation)
|
|
a_backing = activation
|
|
else:
|
|
a, a_backing = fp4_tensor(a_qdata, swap_nibbles=args.swap_nibbles, reencode=args.reencode_fp4)
|
|
b, b_backing = fp4_tensor(b_qdata, swap_nibbles=args.swap_nibbles, reencode=args.reencode_fp4)
|
|
sfa = (
|
|
from_dlpack(a_tensor_scale.float().reshape(1).contiguous(), assumed_align=4)
|
|
if args.stream_a
|
|
else scale_tensor(a_block_scale)
|
|
)
|
|
sfb = scale_tensor(b_block_scale)
|
|
c = output_tensor(output_bf16)
|
|
|
|
gemm = example.Sm120BlockScaledGemmKernel(
|
|
cutlass.Float32,
|
|
16,
|
|
(128, 128, args.tile_k),
|
|
(128, 128),
|
|
)
|
|
hardware_info = cutlass.utils.HardwareInfo()
|
|
max_active_clusters = hardware_info.get_max_active_clusters(1)
|
|
stream = cutlass_torch.default_stream()
|
|
alpha = a_tensor_scale.float() * b_tensor_scale.float()
|
|
alpha_argument = from_dlpack(alpha.reshape(1).contiguous(), assumed_align=4)
|
|
if args.fuse_alpha:
|
|
compiled = cute.compile(gemm, a, b, sfa, sfb, c, alpha_argument, max_active_clusters, stream)
|
|
arguments = (a, b, sfa, sfb, c, alpha_argument, stream)
|
|
else:
|
|
compiled = cute.compile(gemm, a, b, sfa, sfb, c, max_active_clusters, stream)
|
|
arguments = (a, b, sfa, sfb, c, stream)
|
|
compiled(*arguments)
|
|
torch.cuda.synchronize()
|
|
|
|
timing = None
|
|
if args.iterations:
|
|
for _ in range(args.warmup):
|
|
compiled(*arguments)
|
|
torch.cuda.synchronize()
|
|
started = torch.cuda.Event(enable_timing=True)
|
|
finished = torch.cuda.Event(enable_timing=True)
|
|
started.record()
|
|
for _ in range(args.iterations):
|
|
compiled(*arguments)
|
|
finished.record()
|
|
finished.synchronize()
|
|
total_ms = started.elapsed_time(finished)
|
|
timing = {
|
|
"warmup": args.warmup,
|
|
"iterations": args.iterations,
|
|
"total_ms": total_ms,
|
|
"mean_ms": total_ms / args.iterations,
|
|
}
|
|
|
|
reference_timing = None
|
|
if args.benchmark_reference:
|
|
if not args.iterations:
|
|
raise ValueError("--benchmark-reference requires --iterations")
|
|
|
|
def measure_cuda(fn):
|
|
result = None
|
|
for _ in range(args.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(args.iterations):
|
|
result = fn()
|
|
finished.record()
|
|
finished.synchronize()
|
|
total = started.elapsed_time(finished)
|
|
return result, total / args.iterations
|
|
|
|
_, quantize_ms = measure_cuda(lambda: vortex_quantize_nvfp4(activation))
|
|
_, complete_ms = measure_cuda(
|
|
lambda: functional.linear(
|
|
vortex_quantize_nvfp4(activation), packed_weight, None,
|
|
)
|
|
)
|
|
reference_timing = {
|
|
"backend": "vortex_scale_plus_comfy_pack_gemm",
|
|
"activation_quantize_mean_ms": quantize_ms,
|
|
"complete_projection_mean_ms": complete_ms,
|
|
}
|
|
|
|
candidate = (
|
|
output_bf16[:args.rows, :linear.out_features]
|
|
if args.fuse_alpha
|
|
else (output_bf16[:args.rows, :linear.out_features].float() * alpha).to(reference.dtype)
|
|
)
|
|
raw_candidate = output_bf16[:args.rows, :linear.out_features]
|
|
raw_delta = raw_candidate.float() - raw_reference.float()
|
|
delta = candidate.float() - reference.float()
|
|
report = {
|
|
"device": torch.cuda.get_device_name(),
|
|
"torch": torch.__version__,
|
|
"cutlass_dsl": "4.6.2",
|
|
"metadata": metadata,
|
|
"linear": args.linear,
|
|
"rows": args.rows,
|
|
"mnk": [args.rows, linear.out_features, linear.in_features],
|
|
"tile_shape_mnk": [128, 128, args.tile_k],
|
|
"swap_nibbles": args.swap_nibbles,
|
|
"reencode_fp4": args.reencode_fp4,
|
|
"zero_a": args.zero_a,
|
|
"fuse_alpha": args.fuse_alpha,
|
|
"stream_a": args.stream_a,
|
|
"streamed_activation_materialization": (
|
|
{"global_qdata": False, "global_sfa": False}
|
|
if args.stream_a
|
|
else None
|
|
),
|
|
"timing": timing,
|
|
"reference_timing": reference_timing,
|
|
"cute_shapes": {"a": str(a.shape), "b": str(b.shape), "sfa": str(sfa.shape), "sfb": str(sfb.shape), "c": str(c.shape)},
|
|
"tensor_scales": {"a": a_tensor_scale.float().item(), "b": b_tensor_scale.float().item(), "alpha": alpha.item()},
|
|
"raw_output": {
|
|
"dtype": str(output_bf16.dtype),
|
|
"checksum": output_bf16.float().sum().item(),
|
|
"max_abs": output_bf16.float().abs().max().item(),
|
|
"nonzero": int(torch.count_nonzero(output_bf16).item()),
|
|
"finite": bool(torch.isfinite(output_bf16).all().item()),
|
|
},
|
|
"raw_blockscaled_parity": {
|
|
"applicable": not args.fuse_alpha,
|
|
"reference_checksum": raw_reference.float().sum().item(),
|
|
"candidate_checksum": raw_candidate.float().sum().item(),
|
|
"equal": torch.equal(raw_candidate, raw_reference),
|
|
"max_abs": raw_delta.abs().max().item(),
|
|
"mean_abs": raw_delta.abs().mean().item(),
|
|
},
|
|
"reference_checksum": reference.float().sum().item(),
|
|
"candidate_checksum": candidate.float().sum().item(),
|
|
"equal": torch.equal(candidate, reference),
|
|
"max_abs": delta.abs().max().item(),
|
|
"mean_abs": delta.abs().mean().item(),
|
|
"relative_l2": (delta.norm() / reference.float().norm().clamp_min(1e-12)).item(),
|
|
"numerical_note": (
|
|
"The experimental epilogue applies the FP32 global-scale product before BF16 conversion."
|
|
if args.fuse_alpha
|
|
else "The stock SM121 kernel rounds before the external global-scale product; exact H3 integration requires the fused-alpha epilogue."
|
|
),
|
|
}
|
|
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()
|