"""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()