daniel/h3-blackwell-runtime
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions
h3-blackwell-runtime/research/cute_nvfp4_ring/patches/0006-h3-gemm-validator.patch
Daniel Maddern ea14b98544 Preserve Vortex performance research history
2026-08-25 20:30:22 +07:00

543 lines
27 KiB
Diff
Raw Blame History

diff --git a/tools/validate_cute_nvfp4_h3.py b/tools/validate_cute_nvfp4_h3.py
new file mode 100644
index 0000000..f7a3aa3
--- /dev/null
+++ b/tools/validate_cute_nvfp4_h3.py
@@ -0,0 +1,537 @@
+"""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()
Reference in a new issue View git blame Copy permalink