save

examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B_sage3.sh

@@ -4,8 +4,8 @@ set -euo pipefail
 ############################################
 # Single node, 6 GPUs
 ############################################
-NUM_GPUS=6
-export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5
+NUM_GPUS=4
+export CUDA_VISIBLE_DEVICES=0,1,2,3
 
 
 export WANDB_MODE="online"
@@ -24,7 +24,7 @@ echo "CUDA_VISIBLE_DEVICES: $CUDA_VISIBLE_DEVICES"
 ############################################
 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 DATA_DIR=data/Wan-Syn_77x448x832_600k/train
-VALIDATION_DATASET_FILE=data/Wan-Syn_77x448x832_600k/validation_6.json
+VALIDATION_DATASET_FILE=data/Wan-Syn_77x448x832_600k/validation_4.json
 OUTPUT_DIR="checkpoints/sage3_distill"
 
 # Training arguments
@@ -35,21 +35,21 @@ training_args=(
   --train_batch_size 1
   --train_sp_batch_size 1
   --gradient_accumulation_steps 1
-  --num_latent_t 21
+  --num_latent_t 12
   --num_height 480
   --num_width 832
-  --num_frames 81
+  --num_frames 45
   --enable_gradient_checkpointing_type "full"
   --generator_4bit_attn True
 )
 
 # Parallel arguments (adjusted to 6 GPUs)
 parallel_args=(
-  --num_gpus 6
+  --num_gpus 4
   --sp_size 1
   --tp_size 1
   --hsdp_replicate_dim 1
-  --hsdp_shard_dim 6
+  --hsdp_shard_dim 4
 )
 
 # Model arguments

examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_sage3.sh

@@ -7,8 +7,8 @@ export TOKENIZERS_PARALLELISM=false
 
 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 DATA_DIR="data/Wan-Syn_77x448x832_600k/train/"
-VALIDATION_DATASET_FILE="data/Wan-Syn_77x448x832_600k/validation_6.json"
-NUM_GPUS=6
+VALIDATION_DATASET_FILE="data/Wan-Syn_77x448x832_600k/validation_4.json"
+NUM_GPUS=4
 # export CUDA_VISIBLE_DEVICES=4,5
 
 
@@ -20,10 +20,10 @@ training_args=(
   --train_batch_size 1
   --train_sp_batch_size 1
   --gradient_accumulation_steps 1
-  --num_latent_t 6
+  --num_latent_t 12
   --num_height 480
   --num_width 832
-  --num_frames 21
+  --num_frames 45
   --enable_gradient_checkpointing_type "full"
   --generator_4bit_attn True
 )
@@ -53,7 +53,7 @@ dataset_args=(
 validation_args=(
   --log_validation 
   --validation_dataset_file $VALIDATION_DATASET_FILE
-  --validation_steps 200
+  --validation_steps 50
   --validation_sampling_steps "50" 
   --validation_guidance_scale "6.0"
 )

fastvideo/attention/backends/sage_attn3.py

@@ -2,7 +2,10 @@
 
 from curses import KEY_B2
 import torch
-from sageattn import sageattn_blackwell
+try:
+    from sageattn import sageattn_blackwell
+except ImportError:
+    sageattn_blackwell = None
 
 from fastvideo.attention.backends.abstract import (AttentionBackend,
                                                    AttentionImpl,
@@ -13,8 +16,12 @@ from fastvideo.logger import init_logger
 logger = init_logger(__name__)
 
 from math import sqrt
+try:
+    from flash_attn.flash_attn_interface import _wrapped_flash_attn_forward, _wrapped_flash_attn_backward
+except ImportError:
+    _wrapped_flash_attn_forward = None
+    _wrapped_flash_attn_backward = None
 
-from flash_attn.flash_attn_interface import _wrapped_flash_attn_forward, _wrapped_flash_attn_backward
 logger = init_logger(__name__)
 from flashinfer import SfLayout, mm_fp4, nvfp4_quantize
 
@@ -186,7 +193,7 @@ def flash_attn_backward(q_BLHD, k_BLHD, v_BLHD, out_BLHD, grad_out_BLHD, is_caus
     )
     return dq_BLHD, dk_BLHD, dv_BLHD, None, None
 
-USE_TORCH_4BIT_FWD = False
+USE_TORCH_4BIT_FWD = True
 USE_TORCH_16BIT_BWD = True
 class _SageAttnBlackwellWith16bitBwd(torch.autograd.Function):
     @staticmethod

fastvideo/training/training_pipeline.py

@@ -928,85 +928,85 @@ class TrainingPipeline(LoRAPipeline, ABC):
         num_sp_groups = world_group.world_size // self.sp_group.world_size
 
         # Process each validation prompt for each validation step
-        # for num_inference_steps in validation_steps:
-        #     logger.info("rank: %s: num_inference_steps: %s",
-        #                 self.global_rank,
-        #                 num_inference_steps,
-        #                 local_main_process_only=False)
-        #     step_videos: list[np.ndarray] = []
-        #     step_captions: list[str] = []
+        for num_inference_steps in validation_steps:
+            logger.info("rank: %s: num_inference_steps: %s",
+                        self.global_rank,
+                        num_inference_steps,
+                        local_main_process_only=False)
+            step_videos: list[np.ndarray] = []
+            step_captions: list[str] = []
 
-        #     for validation_batch in validation_dataloader:
-        #         batch = self._prepare_validation_batch(sampling_param,
-        #                                                training_args,
-        #                                                validation_batch,
-        #                                                num_inference_steps)
-        #         logger.info("rank: %s: rank_in_sp_group: %s, batch.prompt: %s",
-        #                     self.global_rank,
-        #                     self.rank_in_sp_group,
-        #                     batch.prompt,
-        #                     local_main_process_only=False)
+            for validation_batch in validation_dataloader:
+                batch = self._prepare_validation_batch(sampling_param,
+                                                       training_args,
+                                                       validation_batch,
+                                                       num_inference_steps)
+                logger.info("rank: %s: rank_in_sp_group: %s, batch.prompt: %s",
+                            self.global_rank,
+                            self.rank_in_sp_group,
+                            batch.prompt,
+                            local_main_process_only=False)
 
-        #         assert batch.prompt is not None and isinstance(
-        #             batch.prompt, str)
-        #         step_captions.append(batch.prompt)
+                assert batch.prompt is not None and isinstance(
+                    batch.prompt, str)
+                step_captions.append(batch.prompt)
 
-        #         # Run validation inference
-        #         output_batch = self.validation_pipeline.forward(
-        #             batch, training_args)
-        #         samples = output_batch.output
+                # Run validation inference
+                output_batch = self.validation_pipeline.forward(
+                    batch, training_args)
+                samples = output_batch.output
 
-        #         if self.rank_in_sp_group != 0:
-        #             continue
+                if self.rank_in_sp_group != 0:
+                    continue
 
-        #         # Process outputs
-        #         video = rearrange(samples, "b c t h w -> t b c h w")
-        #         frames = []
-        #         for x in video:
-        #             x = torchvision.utils.make_grid(x, nrow=6)
-        #             x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
-        #             frames.append((x * 255).numpy().astype(np.uint8))
-        #         step_videos.append(frames)
+                # Process outputs
+                video = rearrange(samples, "b c t h w -> t b c h w")
+                frames = []
+                for x in video:
+                    x = torchvision.utils.make_grid(x, nrow=6)
+                    x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
+                    frames.append((x * 255).numpy().astype(np.uint8))
+                step_videos.append(frames)
 
-        #     # Only sp_group leaders (rank_in_sp_group == 0) need to send their
-        #     # results to global rank 0
-        #     if self.rank_in_sp_group == 0:
-        #         if self.global_rank == 0:
-        #             # Global rank 0 collects results from all sp_group leaders
-        #             all_videos = step_videos  # Start with own results
-        #             all_captions = step_captions
+            # Only sp_group leaders (rank_in_sp_group == 0) need to send their
+            # results to global rank 0
+            if self.rank_in_sp_group == 0:
+                if self.global_rank == 0:
+                    # Global rank 0 collects results from all sp_group leaders
+                    all_videos = step_videos  # Start with own results
+                    all_captions = step_captions
 
-        #             # Receive from other sp_group leaders
-        #             for sp_group_idx in range(1, num_sp_groups):
-        #                 src_rank = sp_group_idx * self.sp_world_size  # Global rank of other sp_group leaders
-        #                 recv_videos = world_group.recv_object(src=src_rank)
-        #                 recv_captions = world_group.recv_object(src=src_rank)
-        #                 all_videos.extend(recv_videos)
-        #                 all_captions.extend(recv_captions)
+                    # Receive from other sp_group leaders
+                    for sp_group_idx in range(1, num_sp_groups):
+                        src_rank = sp_group_idx * self.sp_world_size  # Global rank of other sp_group leaders
+                        recv_videos = world_group.recv_object(src=src_rank)
+                        recv_captions = world_group.recv_object(src=src_rank)
+                        all_videos.extend(recv_videos)
+                        all_captions.extend(recv_captions)
 
-        #             video_filenames = []
-        #             for i, (video, caption) in enumerate(
-        #                     zip(all_videos, all_captions, strict=True)):
-        #                 os.makedirs(training_args.output_dir, exist_ok=True)
-        #                 filename = os.path.join(
-        #                     training_args.output_dir,
-        #                     f"validation_step_{global_step}_inference_steps_{num_inference_steps}_video_{i}.mp4"
-        #                 )
-        #                 imageio.mimsave(filename, video, fps=sampling_param.fps)
-        #                 video_filenames.append(filename)
+                    video_filenames = []
+                    for i, (video, caption) in enumerate(
+                            zip(all_videos, all_captions, strict=True)):
+                        os.makedirs(training_args.output_dir, exist_ok=True)
+                        filename = os.path.join(
+                            training_args.output_dir,
+                            f"validation_step_{global_step}_inference_steps_{num_inference_steps}_video_{i}.mp4"
+                        )
+                        imageio.mimsave(filename, video, fps=sampling_param.fps)
+                        video_filenames.append(filename)
 
-        #             logs = {
-        #                 f"validation_videos_{num_inference_steps}_steps": [
-        #                     wandb.Video(filename, caption=caption)
-        #                     for filename, caption in zip(
-        #                         video_filenames, all_captions, strict=True)
-        #                 ]
-        #             }
-        #             wandb.log(logs, step=global_step)
-        #         else:
-        #             # Other sp_group leaders send their results to global rank 0
-        #             world_group.send_object(step_videos, dst=0)
-        #             world_group.send_object(step_captions, dst=0)
+                    logs = {
+                        f"validation_videos_{num_inference_steps}_steps": [
+                            wandb.Video(filename, caption=caption)
+                            for filename, caption in zip(
+                                video_filenames, all_captions, strict=True)
+                        ]
+                    }
+                    wandb.log(logs, step=global_step)
+                else:
+                    # Other sp_group leaders send their results to global rank 0
+                    world_group.send_object(step_videos, dst=0)
+                    world_group.send_object(step_captions, dst=0)
 
         # Re-enable gradients for training
         training_args.inference_mode = False

test_quantization_precision.py

@@ -12,7 +12,90 @@ import numpy as np
 from flashinfer import SfLayout, mm_fp4, nvfp4_quantize
 import unittest
 from typing import Tuple, Dict, Any
-from fake_quant_fp4 import fake_quant_fp4
+
+import torch
+import triton
+import triton.language as tl
+import triton.profiler as proton
+from triton.tools.tensor_descriptor import TensorDescriptor
+from triton.tools.mxfp import MXFP4Tensor, MXScaleTensor
+
+
+
+def generate_fake_quant(a_f32: torch.Tensor,
+                        b_f32: torch.Tensor,
+                        vec_size: int = 16,
+                        max_mxf4: float = 6.0,
+                        eps: float = 1e-8):
+    """
+    Fake-quantize fp32 A (M,K) and B (K,N) with nvfp4 (E2M1) using FP8 E4M3 (finite-only) block scales.
+    Returns dequantized fp32 tensors (A_qdq, B_qdq). Assumes:
+      - M % 128 == 0, N % 128 == 0, K % (vec_size * 4) == 0
+      - A is (M, K), B is (K, N)
+      - Scales are uniform per 128x64 (A) and 64x128 (B) tiles (matching your kernel’s layout)
+    """
+    assert a_f32.is_cuda and b_f32.is_cuda
+    M, K = a_f32.shape
+    Kb, N = b_f32.shape
+    assert Kb == K, "A and B inner dims must match (A:(M,K), B:(K,N))"
+    assert M % 128 == 0 and N % 128 == 0, "M,N must be multiples of 128"
+    assert K % (vec_size * 4) == 0, f"K must be multiple of {vec_size*4}"
+
+    device = a_f32.device
+
+    # ---- tile-scale helpers (produce FP8 E4M3 scales with shape [*, *, 32, 16]) ----
+    def _tile_scales_A(a_fp32):
+        m_chunks = M // 128
+        k_chunks = K // (vec_size * 4)              # 64-wide tiles along K
+        x = a_fp32.view(m_chunks, 128, k_chunks, 4, vec_size)
+        tile_max = x.abs().amax(dim=(1, 3, 4))      # [m_chunks, k_chunks]
+        scale = torch.clamp(tile_max / max_mxf4, min=eps)  # fp32
+        s = scale[:, :, None, None].expand(m_chunks, k_chunks, 32, 16)
+        return s.to(torch.float8_e4m3fn)            # [m_chunks, k_chunks, 32, 16]
+
+    def _tile_scales_B(b_fp32):
+        k_chunks = K // (vec_size * 4)
+        n_chunks = N // 128
+        x = b_fp32.view(k_chunks, 4, vec_size, n_chunks, 128)
+        tile_max = x.abs().amax(dim=(1, 2, 4))      # [k_chunks, n_chunks]
+        scale = torch.clamp(tile_max / max_mxf4, min=eps).t()  # [n_chunks, k_chunks]
+        s = scale[:, :, None, None].expand(n_chunks, k_chunks, 32, 16)
+        return s.to(torch.float8_e4m3fn)            # [n_chunks, k_chunks, 32, 16]
+
+    # ---- expand tile scales to per-element fp32 ----
+    def _expand_A_scale_full(a_scale_fp8):
+        m_chunks, k_chunks, _, _ = a_scale_fp8.shape
+        base = a_scale_fp8[..., 0, 0].to(a_f32.dtype)          # [m_chunks, k_chunks]
+        tile = base[:, :, None, None].expand(m_chunks, k_chunks, 128, 64)
+        return tile.permute(0, 2, 1, 3).reshape(M, k_chunks * 64).contiguous()  # [M,K]
+
+    def _expand_B_scale_full(b_scale_fp8):
+        n_chunks, k_chunks, _, _ = b_scale_fp8.shape
+        base = b_scale_fp8[..., 0, 0].to(b_f32.dtype)          # [n_chunks, k_chunks]
+        tile = base[:, :, None, None].expand(n_chunks, k_chunks, 64, 128)
+        return tile.permute(1, 2, 0, 3).reshape(k_chunks * 64, n_chunks * 128).contiguous()  # [K,N]
+
+    # ---- 1) compute FP8 E4M3 tile scales from fp32 ----
+    a_scale_fp8 = _tile_scales_A(a_f32)   # [M//128, K//64, 32, 16]
+    b_scale_fp8 = _tile_scales_B(b_f32)   # [N//128, K//64, 32, 16]
+
+    # ---- 2) expand scales to full fp32 for quant/dequant math ----
+    a_scale_full = _expand_A_scale_full(a_scale_fp8)   # [M, K] fp32
+    b_scale_full = _expand_B_scale_full(b_scale_fp8)   # [K, N] fp32
+
+    # ---- 3) quantize to nvfp4 codes (E2M1) ----
+    a_codes = MXFP4Tensor(data=(a_f32 / a_scale_full)).data                   # [M, K] uint8 (low nibble)
+    b_codes = MXFP4Tensor(data=(b_f32 / b_scale_full)).data                   # [K, N] uint8
+
+    # ---- 4) dequantize back to fp32 (fake-quant output) ----
+    a_qdq = MXFP4Tensor(size=(M, K), device=device); a_qdq.data = a_codes
+    b_qdq = MXFP4Tensor(size=(K, N), device=device); b_qdq.data = b_codes
+
+    a_qdq = a_qdq.to(torch.float32).to(a_f32.dtype) * a_scale_full
+    b_qdq = b_qdq.to(torch.float32).to(b_f32.dtype) * b_scale_full
+
+    return a_qdq, b_qdq
+
 
 class QuantizationPrecisionTest:
     """Test class for comparing quantization precision methods."""
@@ -38,7 +121,8 @@ class QuantizationPrecisionTest:
         Returns:
             Tuple of (quantized_matrix, inverse_scale_factor)
         """
-        scale_factor = self._compute_scale_factor(matrix)
+        # scale_factor = self._compute_scale_factor(matrix)
+        scale_factor = torch.tensor(1.0, device=matrix.device, dtype=torch.float32)
         quantized, inv_scale = nvfp4_quantize(
             matrix, 
             scale_factor, 
@@ -127,18 +211,18 @@ class QuantizationPrecisionTest:
         """
         # Quantize both matrices
         
-        # X_quant, X_inv_scale, X_global_scale_factor = self.quantize_matrix(X)
-        # W_quant, W_inv_scale, W_global_scale_factor = self.quantize_matrix(W)
-        # # Dequantize both matrices
-        # X_dequant = self.dequantize_matrix(X.shape, X_quant, X_inv_scale, X_global_scale_factor)
-        # W_dequant = self.dequantize_matrix(W.shape, W_quant, W_inv_scale, W_global_scale_factor)
+        X_quant, X_inv_scale, X_global_scale_factor = self.quantize_matrix(X)
+        W_quant, W_inv_scale, W_global_scale_factor = self.quantize_matrix(W)
+        # Dequantize both matrices
+        X_dequant = self.dequantize_matrix(X.shape, X_quant, X_inv_scale, X_global_scale_factor)
+        W_dequant = self.dequantize_matrix(W.shape, W_quant, W_inv_scale, W_global_scale_factor)
         
-        # result = torch.matmul(X_dequant, W_dequant.T)
+        result = torch.matmul(X_dequant, W_dequant.T)
         
         
-        X_fake_quant = fake_quant_fp4(X, stochastic_rounding=True, block_size=16, scale_format='e4m3')
-        W_fake_quant = fake_quant_fp4(W, stochastic_rounding=True, block_size=16, scale_format='e4m3')
-        result = torch.matmul(X_fake_quant, W_fake_quant.T)
+        # X_fake_quant, W_fake_quant_T = generate_fake_quant(X, W.T.contiguous())
+        # import pdb; pdb.set_trace()
+        # result = torch.matmul(X_fake_quant, W_fake_quant_T)
         
         return result
     
@@ -161,7 +245,6 @@ class QuantizationPrecisionTest:
         result1 = self.real_quant_matmul(X, W)
         result2 = self.fake_quant_matmul(X, W)
         import pdb; pdb.set_trace()
-        
         # Compute reference result (no quantization)
         reference = torch.matmul(X, W.T)
         # Calculate differences
@@ -206,7 +289,7 @@ class QuantizationPrecisionTest:
         # Test cases with different sizes and distributions
         test_cases = [
             # (X_shape, W_shape, distribution_name)
-            ((64, 128), (256, 128), "small_matrices"),
+            ((128, 128), (256, 128), "small_matrices"),
             ((256, 512), (1024, 512), "medium_matrices"),
             ((1024, 2048), (4096, 2048), "large_matrices"),
         ]

test_quantization_simple.py

@@ -1,235 +0,0 @@
-#!/usr/bin/env python3
-"""
-Simplified precision test for comparing quantization methods using FlashInfer API.
-
-This test compares:
-1. dequant(quant(X) @ quant(W)) - quantized matrix multiplication
-2. dequant(quant(X)) @ dequant(quant(W)) - dequantized matrix multiplication
-
-Both methods use 4-bit quantization and produce bf16 matrices.
-"""
-
-import torch
-import numpy as np
-from flashinfer import SfLayout, mm_fp4, nvfp4_quantize
-
-
-def compute_scale_factor(tensor: torch.Tensor) -> torch.Tensor:
-    """Compute scale factor for quantization using FastVideo's method."""
-    maxabs = tensor.float().abs().nan_to_num().max()
-    maxabs = torch.maximum(maxabs, torch.tensor(1e-12, device=tensor.device, dtype=maxabs.dtype))
-    return (448.0 * 6.0) / maxabs
-
-
-def quantize_matrix(matrix: torch.Tensor, sf_layout: SfLayout = SfLayout.layout_128x4):
-    """Quantize matrix to 4-bit using FlashInfer."""
-    scale_factor = compute_scale_factor(matrix)
-    quantized, inv_scale = nvfp4_quantize(
-        matrix, 
-        scale_factor, 
-        sfLayout=sf_layout, 
-        do_shuffle=False
-    )
-    return quantized, inv_scale
-
-
-def method1_quantized_matmul(X: torch.Tensor, W: torch.Tensor, device: str = "cuda", dtype: torch.dtype = torch.bfloat16):
-    """
-    Method 1: dequant(quant(X) @ quant(W))
-    Perform matrix multiplication in quantized space.
-    """
-    # Quantize both matrices
-    X_quant, X_inv_scale = quantize_matrix(X)
-    W_quant, W_inv_scale = quantize_matrix(W)
-    
-    # Compute alpha scaling factor
-    alpha = 1.0 / (X_inv_scale * W_inv_scale)
-    
-    # Perform quantized matrix multiplication
-    result = torch.empty((X.shape[0], W.shape[0]), device=device, dtype=dtype)
-    mm_fp4(
-        X_quant, W_quant.T, X_inv_scale, W_inv_scale.T, alpha,
-        dtype, result,
-        block_size=16,
-        use_8x4_sf_layout=False,
-        backend="cutlass"
-    )
-    
-    return result
-
-
-def method2_dequantized_matmul(X: torch.Tensor, W: torch.Tensor, device: str = "cuda", dtype: torch.dtype = torch.bfloat16):
-    """
-    Method 2: dequant(quant(X)) @ dequant(quant(W))
-    Dequantize both matrices first, then perform matrix multiplication.
-    
-    Note: Since FlashInfer doesn't provide direct dequantization, we approximate
-    this by using the quantized matrices with identity matrices to extract
-    the dequantized values.
-    """
-    # Quantize both matrices
-    X_quant, X_inv_scale = quantize_matrix(X)
-    W_quant, W_inv_scale = quantize_matrix(W)
-    
-    # Approximate dequantization by multiplying with identity matrices
-    # This extracts the dequantized values
-    X_dequant = torch.empty_like(X)
-    W_dequant = torch.empty_like(W)
-    
-    # Dequantize X
-    identity_X = torch.eye(X.shape[1], device=device, dtype=dtype)
-    mm_fp4(
-        X_quant, identity_X, X_inv_scale, X_inv_scale,
-        torch.tensor(1.0, device=device, dtype=torch.float32),
-        dtype, X_dequant,
-        block_size=16,
-        use_8x4_sf_layout=False,
-        backend="cutlass"
-    )
-    
-    # Dequantize W
-    identity_W = torch.eye(W.shape[1], device=device, dtype=dtype)
-    mm_fp4(
-        W_quant, identity_W, W_inv_scale, W_inv_scale,
-        torch.tensor(1.0, device=device, dtype=torch.float32),
-        dtype, W_dequant,
-        block_size=16,
-        use_8x4_sf_layout=False,
-        backend="cutlass"
-    )
-    
-    # Perform standard matrix multiplication
-    result = torch.matmul(X_dequant, W_dequant)
-    
-    return result
-
-
-def compare_quantization_methods(X: torch.Tensor, W: torch.Tensor, device: str = "cuda", dtype: torch.dtype = torch.bfloat16):
-    """
-    Compare the two quantization methods and return detailed analysis.
-    """
-    # Ensure matrices are on correct device and dtype
-    X = X.to(device=device, dtype=dtype)
-    W = W.to(device=device, dtype=dtype)
-    
-    # Compute results using both methods
-    result1 = method1_quantized_matmul(X, W, device, dtype)
-    result2 = method2_dequantized_matmul(X, W, device, dtype)
-    
-    # Compute reference result (no quantization)
-    reference = torch.matmul(X, W)
-    
-    # Calculate differences
-    diff_methods = torch.abs(result1 - result2)
-    diff1_ref = torch.abs(result1 - reference)
-    diff2_ref = torch.abs(result2 - reference)
-    
-    # Statistical analysis
-    analysis = {
-        'method1_result': result1,
-        'method2_result': result2,
-        'reference_result': reference,
-        'method_difference': {
-            'max': torch.max(diff_methods).item(),
-            'mean': torch.mean(diff_methods).item(),
-            'std': torch.std(diff_methods).item(),
-            'relative_max': torch.max(diff_methods / (torch.abs(reference) + 1e-8)).item(),
-            'relative_mean': torch.mean(diff_methods / (torch.abs(reference) + 1e-8)).item(),
-        },
-        'method1_vs_reference': {
-            'max': torch.max(diff1_ref).item(),
-            'mean': torch.mean(diff1_ref).item(),
-            'relative_max': torch.max(diff1_ref / (torch.abs(reference) + 1e-8)).item(),
-        },
-        'method2_vs_reference': {
-            'max': torch.max(diff2_ref).item(),
-            'mean': torch.mean(diff2_ref).item(),
-            'relative_max': torch.max(diff2_ref / (torch.abs(reference) + 1e-8)).item(),
-        },
-        'input_shapes': {
-            'X': list(X.shape),
-            'W': list(W.shape),
-        }
-    }
-    
-    return analysis
-
-
-def run_test_suite():
-    """Run comprehensive test suite."""
-    print("Quantization Precision Test")
-    print("=" * 50)
-    
-    # Check device availability
-    if torch.cuda.is_available():
-        device = "cuda"
-        print(f"Using CUDA device: {torch.cuda.get_device_name()}")
-    else:
-        device = "cpu"
-        print("CUDA not available, using CPU")
-    
-    dtype = torch.bfloat16
-    
-    # Test cases
-    test_cases = [
-        # (X_shape, W_shape, description)
-        ((32, 64), (64, 128), "Small matrices"),
-        ((128, 256), (256, 512), "Medium matrices"),
-        ((256, 512), (512, 1024), "Large matrices"),
-        ((64, 128), (128, 64), "Rectangular matrices"),
-    ]
-    
-    distributions = [
-        ("Normal", lambda shape: torch.randn(shape, device=device)),
-        ("Uniform", lambda shape: torch.rand(shape, device=device) * 2 - 1),
-        ("Small values", lambda shape: torch.randn(shape, device=device) * 0.1),
-        ("Large values", lambda shape: torch.randn(shape, device=device) * 10),
-    ]
-    
-    all_max_diffs = []
-    all_mean_diffs = []
-    
-    for x_shape, w_shape, size_desc in test_cases:
-        print(f"\n{size_desc} ({x_shape} x {w_shape}):")
-        print("-" * 40)
-        
-        for dist_name, dist_func in distributions:
-            # Generate test matrices
-            X = dist_func(x_shape).to(dtype=dtype)
-            W = dist_func(w_shape).to(dtype=dtype)
-            
-            # Run comparison
-            analysis = compare_quantization_methods(X, W, device, dtype)
-            
-            max_diff = analysis['method_difference']['max']
-            mean_diff = analysis['method_difference']['mean']
-            rel_max_diff = analysis['method_difference']['relative_max']
-            
-            all_max_diffs.append(max_diff)
-            all_mean_diffs.append(mean_diff)
-            
-            print(f"  {dist_name:12} | Max diff: {max_diff:.6f} | Mean diff: {mean_diff:.6f} | Rel max: {rel_max_diff:.6f}")
-    
-    # Summary
-    print(f"\nSummary:")
-    print(f"  Overall max difference: {max(all_max_diffs):.6f}")
-    print(f"  Overall mean difference: {np.mean(all_mean_diffs):.6f}")
-    print(f"  Overall std difference: {np.std(all_mean_diffs):.6f}")
-    
-    # Test assertions
-    print(f"\nTest Results:")
-    print(f"  ✓ Max difference < 0.01: {max(all_max_diffs) < 0.01}")
-    print(f"  ✓ Mean difference < 0.001: {np.mean(all_mean_diffs) < 0.001}")
-    
-    return all_max_diffs, all_mean_diffs
-
-
-if __name__ == "__main__":
-    try:
-        max_diffs, mean_diffs = run_test_suite()
-        print(f"\n✅ Test completed successfully!")
-        print(f"   Both quantization methods produce similar results within acceptable precision.")
-    except Exception as e:
-        print(f"\n❌ Test failed with error: {e}")
-        import traceback
-        traceback.print_exc()

test_triton_fp4_vs_cutlass.py

@@ -0,0 +1,276 @@
+import torch
+from triton.tools.tensor_descriptor import TensorDescriptor
+from triton.tools.mxfp import MXFP4Tensor
+from triton_fp4 import block_scaled_matmul
+from flashinfer import SfLayout, mm_fp4, nvfp4_quantize
+
+
+
+def triton_nvfp4(A: torch.Tensor,
+                 B: torch.Tensor,
+                 out_dtype: torch.dtype = torch.float16):
+    """
+    Block-scaled nvfp4 GEMM: (A @ B) with A:(M,K), B:(K,N) given in fp32/bf16.
+    Returns the GEMM output in `out_dtype`.
+
+    Requirements (match your kernel):
+      - M % 128 == 0, N % 128 == 0, K % 64 == 0 (since vec_size=16 and 4 groups → 64)
+      - Device: CUDA, Blackwell (SM=10)
+    """
+    assert A.is_cuda and B.is_cuda, "Inputs must be on CUDA"
+    assert A.dtype in (torch.float32, torch.bfloat16), "A must be fp32/bf16"
+    assert B.dtype in (torch.float32, torch.bfloat16), "B must be fp32/bf16"
+    M, K = A.shape
+    Kb, N = B.shape
+    assert Kb == K, "Inner dims must match"
+    assert M % 128 == 0 and N % 128 == 0 and K % 64 == 0, "Shape must respect tiling"
+
+    # ---- kernel constants / config ----
+    BLOCK_M = 128
+    BLOCK_N = 256
+    BLOCK_K = 256
+    VEC_SIZE = 16
+    ELEM_PER_BYTE_A = 2
+    ELEM_PER_BYTE_B = 2
+    MAX_MXF4 = 6.0    # nvfp4 (E2M1) max magnitude
+    EPS = 1e-8
+
+    device = A.device
+
+    # ---- tile-scale helpers (FP8 E4M3, finite-only) ----
+    # A tiles are 128 x 64 along (M,K)
+    def _tile_scales_A(a_fp32):
+        m_chunks = M // 128
+        k_chunks = K // (VEC_SIZE * 4)  # 64 wide
+        x = a_fp32.view(m_chunks, 128, k_chunks, 4, VEC_SIZE)
+        tile_max = x.abs().amax(dim=(1, 3, 4))                       # [m_chunks, k_chunks]
+        scale = torch.clamp(tile_max / MAX_MXF4, min=EPS)            # fp32
+        s = scale[:, :, None, None].expand(m_chunks, k_chunks, 32, 16)
+        return s.to(torch.float8_e4m3fn)                              # [m_chunks, k_chunks, 32, 16]
+
+    # B tiles are 64 x 128 over (K,N)
+    def _tile_scales_B(b_fp32):
+        k_chunks = K // (VEC_SIZE * 4)
+        n_chunks = N // 128
+        x = b_fp32.view(k_chunks, 4, VEC_SIZE, n_chunks, 128)
+        tile_max = x.abs().amax(dim=(1, 2, 4))                        # [k_chunks, n_chunks]
+        scale = torch.clamp(tile_max / MAX_MXF4, min=EPS).t()         # [n_chunks, k_chunks]
+        s = scale[:, :, None, None].expand(n_chunks, k_chunks, 32, 16)
+        return s.to(torch.float8_e4m3fn)                              # [n_chunks, k_chunks, 32, 16]
+
+    def _expand_A_scale_full(a_scale_fp8):
+        m_chunks, k_chunks, _, _ = a_scale_fp8.shape
+        base = a_scale_fp8[..., 0, 0].to(A.dtype)               # [m_chunks, k_chunks]
+        tile = base[:, :, None, None].expand(m_chunks, k_chunks, 128, 64)
+        return tile.permute(0, 2, 1, 3).reshape(M, k_chunks * 64).contiguous()  # [M,K]
+
+    def _expand_B_scale_full(b_scale_fp8):
+        n_chunks, k_chunks, _, _ = b_scale_fp8.shape
+        base = b_scale_fp8[..., 0, 0].to(B.dtype)               # [n_chunks, k_chunks]
+        tile = base[:, :, None, None].expand(n_chunks, k_chunks, 64, 128)
+        return tile.permute(1, 2, 0, 3).reshape(k_chunks * 64, n_chunks * 128).contiguous()  # [K,N]
+
+    # ---- 1) scales from fp32 ----
+    a_scale_fp8 = _tile_scales_A(A)             # [M//128, K//64, 32, 16]
+    b_scale_fp8 = _tile_scales_B(B)             # [N//128, K//64, 32, 16]
+    a_scale_full = _expand_A_scale_full(a_scale_fp8)   # [M,K] fp32
+    b_scale_full = _expand_B_scale_full(b_scale_fp8)   # [K,N] fp32
+
+    # ---- 2) quantize to nvfp4 codes (E2M1) ----
+    a_codes = MXFP4Tensor(data=(A / a_scale_full)).data            # [M,K] uint8 (low nibble used)
+    b_codes_KN = MXFP4Tensor(data=(B / b_scale_full)).data         # [K,N] uint8
+
+    # Kernel expects B laid out as (N, K) packed along K (like the original example)
+    b_codes_NK = b_codes_KN.transpose(0, 1).contiguous()             # [N,K]
+
+    # ---- 3) pack to 2-per-byte along K and build descriptors ----
+    a_fp4 = MXFP4Tensor(size=(M, K), device=device); a_fp4.data = a_codes
+    b_fp4_desc = MXFP4Tensor(size=(N, K), device=device); b_fp4_desc.data = b_codes_NK
+
+    a_packed = a_fp4.to_packed_tensor(dim=1)      # (M, K/2)
+    b_packed = b_fp4_desc.to_packed_tensor(dim=1) # (N, K/2)    
+    
+    a_desc = TensorDescriptor.from_tensor(a_packed, [BLOCK_M, BLOCK_K // ELEM_PER_BYTE_A])
+    b_desc = TensorDescriptor.from_tensor(b_packed, [BLOCK_N, BLOCK_K // ELEM_PER_BYTE_B])
+
+    # reps for scale swizzle
+    rep_m = BLOCK_M // 128
+    rep_n = BLOCK_N // 128
+    rep_k = BLOCK_K // VEC_SIZE // 4
+
+    # reshape FP8 scales into 5D TMA blocks [1, rep_*, rep_k, 2, 256]
+    a_scale_5d = a_scale_fp8.reshape(1, a_scale_fp8.shape[0], a_scale_fp8.shape[1], 2, 256)
+    b_scale_5d = b_scale_fp8.reshape(1, b_scale_fp8.shape[0], b_scale_fp8.shape[1], 2, 256)
+    a_scale_desc = TensorDescriptor.from_tensor(a_scale_5d, [1, rep_m, rep_k, 2, 256])
+    b_scale_desc = TensorDescriptor.from_tensor(b_scale_5d, [1, rep_n, rep_k, 2, 256])
+
+    a_packed_cutlass, a_scale_fp8_cutlass = nvfp4_quantize(
+        A, torch.tensor(1.0, device=device, dtype=torch.float32), sfLayout=SfLayout.layout_128x4, do_shuffle=False
+    )
+    import pdb; pdb.set_trace()
+    # ---- 4) launch GEMM ----
+    configs = {
+        "BLOCK_SIZE_M": BLOCK_M,
+        "BLOCK_SIZE_N": BLOCK_N,
+        "BLOCK_SIZE_K": BLOCK_K,
+        "num_stages": 4,
+        "ELEM_PER_BYTE_A": ELEM_PER_BYTE_A,
+        "ELEM_PER_BYTE_B": ELEM_PER_BYTE_B,
+        "VEC_SIZE": VEC_SIZE,
+    }
+    # Uses your already-defined block_scaled_matmul(...)
+    C = block_scaled_matmul(
+        a_desc, a_scale_desc,
+        b_desc, b_scale_desc,
+        out_dtype, M, N, K,
+        rep_m, rep_n, rep_k, configs
+    )
+    return C
+
+
+
+def cutlass_nvfp4(A: torch.Tensor,
+                  B: torch.Tensor,
+                  out_dtype: torch.dtype = torch.bfloat16,
+                  sf_layout: SfLayout = SfLayout.layout_128x4,
+                  block_size: int = 16,
+                  do_shuffle: bool = False,
+                  scale_factor_A: torch.Tensor | float | None = None,
+                  scale_factor_B: torch.Tensor | float | None = None) -> torch.Tensor:
+    """
+    FlashInfer/CUTLASS nvfp4 GEMM: returns (A @ B) using 4-bit E2M1 operands with per-tile scaling.
+
+    Args:
+      A: (M, K) tensor, fp32 or bf16, CUDA.
+      B: (K, N) tensor, fp32 or bf16, CUDA.
+      out_dtype: torch.float32, torch.float16, or torch.bfloat16 (recommended).
+      sf_layout: FlashInfer scale-factor layout (default SfLayout.layout_128x4).
+      block_size: internal tiling, default 16.
+      do_shuffle: pass-through to nvfp4_quantize.
+      scale_factor_A, scale_factor_B: optional global scale factors (float or 0-dim tensor).
+        If None, defaults to 1.0 (FlashInfer still applies its internal per-tile scales).
+
+    Returns:
+      C: (M, N) tensor in out_dtype
+    """
+    assert A.is_cuda and B.is_cuda, "Inputs must be CUDA tensors"
+    assert A.dim() == 2 and B.dim() == 2, "A and B must be 2D"
+    M, K = A.shape
+    Kb, N = B.shape
+    assert K == Kb, f"Incompatible shapes: A is (M={M},K={K}), B is (K={Kb},N={N})"
+
+    # Cast inputs to supported accumulation input types for FlashInfer quantizer.
+    # (fp32/bf16 are fine; we keep user dtype if already bf16 to save bandwidth.)
+    A_in = A if A.dtype in (torch.bfloat16, torch.float32) else A.float()
+    B_in = B if B.dtype in (torch.bfloat16, torch.float32) else B.float()
+    B_in = B_in.T.contiguous()
+
+    # Choose global scaling (FlashInfer also uses per-tile scales internally).
+    if scale_factor_A is None:
+        scale_factor_A = torch.tensor(1.0, device=A.device, dtype=torch.float32)
+    elif not torch.is_tensor(scale_factor_A):
+        scale_factor_A = torch.tensor(float(scale_factor_A), device=A.device, dtype=torch.float32)
+
+    if scale_factor_B is None:
+        scale_factor_B = torch.tensor(1.0, device=B.device, dtype=torch.float32)
+    elif not torch.is_tensor(scale_factor_B):
+        scale_factor_B = torch.tensor(float(scale_factor_B), device=B.device, dtype=torch.float32)
+
+    # Quantize A and B to nvfp4; inv scales are FlashInfer’s per-tile inverse scales.
+    # Returned quantized matrices are FP4-packed formats understood by mm_fp4.
+    A_q, A_inv_sf = nvfp4_quantize(
+        A_in, scale_factor_A, sfLayout=sf_layout, do_shuffle=do_shuffle
+    )
+    B_q, B_inv_sf = nvfp4_quantize(
+        B_in, scale_factor_B, sfLayout=sf_layout, do_shuffle=do_shuffle
+    )
+
+    # Alpha rescales the product of the two global scale factors.
+    alpha = 1.0 / (scale_factor_A * scale_factor_B)
+
+    # Pick FlashInfer flag for the scale-factor layout
+    use_8x4_sf_layout = (sf_layout == SfLayout.layout_8x4)
+
+    # Allocate output and call CUTLASS-backed FP4 GEMM.
+    C = torch.empty((M, N), device=A.device, dtype=out_dtype)
+    mm_fp4(
+        A_q, B_q.T,             # note the transpose on B_q (expecting KxN → NxK in kernel)
+        A_inv_sf, B_inv_sf.T,   # matching transposition of per-tile inverse scales
+        alpha,                  # global alpha
+        out_dtype, C,           # output dtype & buffer
+        block_size=block_size,
+        use_8x4_sf_layout=use_8x4_sf_layout,
+        backend="cutlass"
+    )
+    return C
+
+
+def main():
+    import argparse, math
+    torch.manual_seed(0)
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--M", type=int, default=256, help="Rows of A / C (multiple of 128)")
+    parser.add_argument("--N", type=int, default=256, help="Cols of B / C (multiple of 128)")
+    parser.add_argument("--K", type=int, default=512, help="Inner dim (multiple of 64)")
+    parser.add_argument("--dtype", type=str, choices=["fp32", "bf16"], default="bf16")
+    parser.add_argument("--trials", type=int, default=1)
+    parser.add_argument("--verbose", action="store_true")
+    args = parser.parse_args()
+
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA required")
+
+    M, N, K = args.M, args.N, args.K
+    assert M % 128 == 0 and N % 128 == 0 and K % 64 == 0, "Shape must satisfy: M%128==0, N%128==0, K%64==0."
+
+    dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float32
+    device = "cuda"
+
+    def stats(name, X):
+        X = X.float()
+        return f"{name}: max={X.abs().max().item():.6e} mean={X.abs().mean().item():.6e} std={X.std().item():.6e}"
+
+    for t in range(args.trials):
+        # Inputs
+        A = torch.randn((M, K), device=device, dtype=dtype)
+        B = torch.randn((K, N), device=device, dtype=dtype)
+
+        # GEMMs
+        C_triton  = triton_nvfp4(A, B, out_dtype=torch.bfloat16 if dtype is torch.bfloat16 else torch.float32)
+        C_cutlass = cutlass_nvfp4(A, B, out_dtype=torch.bfloat16 if dtype is torch.bfloat16 else torch.float32)
+        # Bring to fp32 for error analysis
+        Ct = C_triton.float()
+        Cc = C_cutlass.float()
+        Cref = (A.float() @ B.float())
+        # Triton vs Cutlass
+        diff_tc = (Ct - Cc)
+        denom_tc = Cc.abs().clamp_min(1e-8)
+        rel_tc = (diff_tc.abs() / denom_tc)
+
+        # Each vs reference (optional sanity)
+        diff_t_ref = (Ct - Cref)
+        diff_c_ref = (Cc - Cref)
+        rel_t_ref = diff_t_ref.abs() / Cref.abs().clamp_min(1e-8)
+        rel_c_ref = diff_c_ref.abs() / Cref.abs().clamp_min(1e-8)
+
+
+        print("\nTriton vs Cutlass:")
+        print(f"  abs diff   -> max={diff_tc.abs().max().item():.6e}  mean={diff_tc.abs().mean().item():.6e}")
+        print(f"  rel diff   -> max={rel_tc.max().item():.6e}       mean={rel_tc.mean().item():.6e}")
+
+        print("\nTriton vs Reference:")
+        print(f"  abs diff   -> max={diff_t_ref.abs().max().item():.6e}  mean={diff_t_ref.abs().mean().item():.6e}")
+        print(f"  rel diff   -> max={rel_t_ref.max().item():.6e}         mean={rel_t_ref.mean().item():.6e}")
+
+        print("\nCutlass vs Reference:")
+        print(f"  abs diff   -> max={diff_c_ref.abs().max().item():.6e}  mean={diff_c_ref.abs().mean().item():.6e}")
+        print(f"  rel diff   -> max={rel_c_ref.max().item():.6e}         mean={rel_c_ref.mean().item():.6e}")
+
+        if args.verbose:
+            # a tiny sample print to eyeball values
+            i, j = 0, 0
+            print(f"\nSample C[0,0]: triton={Ct[i,j].item():.6e}, cutlass={Cc[i,j].item():.6e}, ref={Cref[i,j].item():.6e}")
+
+if __name__ == "__main__":
+    main()

tinyllama.py

@@ -0,0 +1,33 @@
+from tensorrt_llm import LLM, SamplingParams
+
+
+def main():
+
+    # Model could accept HF model name, a path to local HF model,
+    # or TensorRT Model Optimizer's quantized checkpoints like nvidia/Llama-3.1-8B-Instruct-FP8 on HF.
+    llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0")
+
+    # Sample prompts.
+    prompts = [
+        "Hello, my name is",
+        "The capital of France is",
+        "The future of AI is",
+    ]
+
+    # Create a sampling params.
+    sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
+
+    for output in llm.generate(prompts, sampling_params):
+        print(
+            f"Prompt: {output.prompt!r}, Generated text: {output.outputs[0].text!r}"
+        )
+
+    # Got output like
+    # Prompt: 'Hello, my name is', Generated text: '\n\nJane Smith. I am a student pursuing my degree in Computer Science at [university]. I enjoy learning new things, especially technology and programming'
+    # Prompt: 'The president of the United States is', Generated text: 'likely to nominate a new Supreme Court justice to fill the seat vacated by the death of Antonin Scalia. The Senate should vote to confirm the'
+    # Prompt: 'The capital of France is', Generated text: 'Paris.'
+    # Prompt: 'The future of AI is', Generated text: 'an exciting time for us. We are constantly researching, developing, and improving our platform to create the most advanced and efficient model available. We are'
+
+
+if __name__ == '__main__':
+    main()

triton_fp4.py

@@ -6,7 +6,7 @@ import triton.language as tl
 import triton.profiler as proton
 from triton.tools.tensor_descriptor import TensorDescriptor
 from triton.tools.mxfp import MXFP4Tensor, MXScaleTensor
-
+from test_quantization_precision import generate_fake_quant
 
 VEC_SIZE = 16
 MAX_MXF4 = 6.0      # max magnitude for E2M1 (nvfp4) with M∈{0,1}, E∈{0..3}
@@ -142,7 +142,7 @@ def block_scaled_matmul(a_desc, a_scale_desc, b_desc, b_scale_desc, dtype_dst, M
         configs["num_stages"],
     )
     return output
-def initialize_block_scaled(M, N, K, block_scale_type="nvfp4", compute_reference=False):
+def initialize_block_scaled(M, N, K, block_scale_type="nvfp4", compute_reference=False, input_output_dtype=torch.bfloat16):
     assert block_scale_type == "nvfp4", f"Only nvfp4 supported here (got {block_scale_type})"
 
     # --- constants ---
@@ -182,20 +182,20 @@ def initialize_block_scaled(M, N, K, block_scale_type="nvfp4", compute_reference
     # --- expand tile scales to per-element fp32 (for quant/dequant math) ---
     def _expand_A_scale_full(a_scale_fp8):
         m_chunks, k_chunks, _, _ = a_scale_fp8.shape
-        base = a_scale_fp8[..., 0, 0].to(torch.float32)            # [m_chunks, k_chunks]
+        base = a_scale_fp8[..., 0, 0].to(input_output_dtype)            # [m_chunks, k_chunks]
         tile = base[:, :, None, None].expand(m_chunks, k_chunks, 128, 64)
         return tile.permute(0, 2, 1, 3).reshape(M, k_chunks * 64).contiguous()  # [M, K]
 
     def _expand_B_scale_full(b_scale_fp8):
         n_chunks, k_chunks, _, _ = b_scale_fp8.shape
-        base = b_scale_fp8[..., 0, 0].to(torch.float32)            # [n_chunks, k_chunks]
+        base = b_scale_fp8[..., 0, 0].to(input_output_dtype)            # [n_chunks, k_chunks]
         tile = base[:, :, None, None].expand(n_chunks, k_chunks, 64, 128)
         return tile.permute(1, 2, 0, 3).reshape(k_chunks * 64, n_chunks * 128).contiguous()  # [K, N]
 
     # --- 1) fp32 sources (pre-quant) ---
     torch.manual_seed(0)
-    a_f32_pre = torch.randn((M, K), device=device, dtype=torch.float32) * 0.5
-    b_f32_pre = torch.randn((K, N), device=device, dtype=torch.float32) * 0.5
+    a_f32_pre = torch.randn((M, K), device=device, dtype=input_output_dtype)
+    b_f32_pre = torch.randn((K, N), device=device, dtype=input_output_dtype)
 
     # --- 2) tile scales (FP8 E4M3, finite-only) ---
     a_scale_fp8 = _tile_scales_A(a_f32_pre)    # [M//128, K//64, 32, 16]
@@ -237,8 +237,7 @@ def initialize_block_scaled(M, N, K, block_scale_type="nvfp4", compute_reference
     a_ref_fp4 = MXFP4Tensor(size=(M, K), device=device); a_ref_fp4.data = a_codes
     b_ref_fp4 = MXFP4Tensor(size=(K, N), device=device); b_ref_fp4.data = b_codes_KN
 
-    a_ref_f32 = a_ref_fp4.to(torch.float32) * a_scale_full          # [M, K]
-    b_ref_f32 = b_ref_fp4.to(torch.float32) * b_scale_full          # [K, N]
+    a_ref_f32, b_ref_f32 = generate_fake_quant(a_f32_pre, b_f32_pre)
 
     reference = a_ref_f32 @ b_ref_f32 if compute_reference else None
 
@@ -255,11 +254,11 @@ def initialize_block_scaled(M, N, K, block_scale_type="nvfp4", compute_reference
     return a_desc, a_scale_desc, b_desc, b_scale_desc, rep_m, rep_n, rep_k, configs, reference
 
 
-def validate_block_scaled(M, N, K, block_scale_type="nvfp4"):
+def validate_block_scaled(M, N, K, block_scale_type="nvfp4", input_output_dtype=torch.bfloat16):
     a_desc, a_scale, b_desc, b_scale, rep_m, rep_n, rep_k, configs, reference = initialize_block_scaled(
         M, N, K, block_scale_type, compute_reference=True)
-    output = block_scaled_matmul(a_desc, a_scale, b_desc, b_scale, torch.float32, M, N, K, rep_m, rep_n, rep_k, configs)
-    torch.testing.assert_close(reference, output.to(torch.float32), atol=1e-3, rtol=1e-3)
+    output = block_scaled_matmul(a_desc, a_scale, b_desc, b_scale, input_output_dtype, M, N, K, rep_m, rep_n, rep_k, configs)
+    torch.testing.assert_close(reference, output.to(input_output_dtype), atol=1e-3, rtol=1e-3)
     print(f"✅ (pass {block_scale_type})")
 
 
@@ -297,6 +296,7 @@ if __name__ == "__main__":
     parser.add_argument("--K_step", type=int, default=512)
     parser.add_argument("--bench", action="store_true", default=True)
     parser.add_argument("--format", type=str, choices=["nvfp4"], default="nvfp4")
+    parser.add_argument("--input_output_dtype", type=str, choices=["float32", "float16", "bfloat16"], default="bfloat16")
     args = parser.parse_args()
 
     if not supports_block_scaling():
@@ -307,6 +307,14 @@ if __name__ == "__main__":
             args.K_step = 1  # doesn't matter as long as it's not 0
 
         torch.manual_seed(42)
+        dtype =torch.float32 if args.input_output_dtype == "float32" else torch.float16 if args.input_output_dtype == "float16" else torch.bfloat16
+        validate_block_scaled(8192, 8192, 8192, block_scale_type=args.format, input_output_dtype=dtype)
 
-        validate_block_scaled(8192, 8192, 8192, block_scale_type=args.format)
 
+        if args.bench:
+            proton.start("block_scaled_matmul", hook="triton")
+            proton.deactivate(0)  # Skip argument creation
+            for K in range(args.K_range[0], args.K_range[1] + 1, args.K_step):
+                bench_block_scaled(K, reps=10000, block_scale_type=args.format)
+            proton.finalize()
+            show_profile("block_scaled_matmul")

triton_original.py

@@ -0,0 +1,302 @@
+import argparse
+
+import torch
+import triton
+import triton.language as tl
+import triton.profiler as proton
+from triton.tools.tensor_descriptor import TensorDescriptor
+from triton.tools.mxfp import MXFP4Tensor, MXScaleTensor
+
+
+def is_cuda():
+    return triton.runtime.driver.active.get_current_target().backend == "cuda"
+
+
+def supports_block_scaling():
+    return is_cuda() and torch.cuda.get_device_capability()[0] == 10
+
+
+def _matmul_launch_metadata(grid, kernel, args):
+    ret = {}
+    M, N, K = args["M"], args["N"], args["K"]
+    kernel_name = kernel.name
+    if "ELEM_PER_BYTE_A" and "ELEM_PER_BYTE_B" and "VEC_SIZE" in args:
+        if args["ELEM_PER_BYTE_A"] == 1 and args["ELEM_PER_BYTE_B"] == 1:
+            kernel_name += "_mxfp8"
+        elif args["ELEM_PER_BYTE_A"] == 1 and args["ELEM_PER_BYTE_B"] == 2:
+            kernel_name += "_mixed"
+        elif args["ELEM_PER_BYTE_A"] == 2 and args["ELEM_PER_BYTE_B"] == 2:
+            if args["VEC_SIZE"] == 16:
+                kernel_name += "_nvfp4"
+            elif args["VEC_SIZE"] == 32:
+                kernel_name += "_mxfp4"
+    ret["name"] = f"{kernel_name} [M={M}, N={N}, K={K}]"
+    ret["flops"] = 2.0 * M * N * K
+    return ret
+
+
+@triton.jit(launch_metadata=_matmul_launch_metadata)
+def block_scaled_matmul_kernel(  #
+        a_desc,  #
+        a_scale_desc,  #
+        b_desc,  #
+        b_scale_desc,  #
+        c_desc,  #
+        M: tl.constexpr,  #
+        N: tl.constexpr,  #
+        K: tl.constexpr,  #
+        output_type: tl.constexpr,  #
+        ELEM_PER_BYTE_A: tl.constexpr,  #
+        ELEM_PER_BYTE_B: tl.constexpr,  #
+        VEC_SIZE: tl.constexpr,  #
+        BLOCK_M: tl.constexpr,  #
+        BLOCK_N: tl.constexpr,  #
+        BLOCK_K: tl.constexpr,  #
+        rep_m: tl.constexpr,  #
+        rep_n: tl.constexpr,  #
+        rep_k: tl.constexpr,  #
+        NUM_STAGES: tl.constexpr,  #
+):  #
+    if output_type == 0:
+        output_dtype = tl.float32
+    elif output_type == 1:
+        output_dtype = tl.float16
+    elif output_type == 2:
+        output_dtype = tl.float8e4nv
+
+    pid = tl.program_id(axis=0)
+    num_pid_m = tl.cdiv(M, BLOCK_M)
+    pid_m = pid % num_pid_m
+    pid_n = pid // num_pid_m
+    offs_am = pid_m * BLOCK_M
+    offs_bn = pid_n * BLOCK_N
+    offs_k_a = 0
+    offs_k_b = 0
+    offs_scale_m = pid_m * rep_m
+    offs_scale_n = pid_n * rep_n
+    offs_scale_k = 0
+
+    MIXED_PREC: tl.constexpr = ELEM_PER_BYTE_A == 1 and ELEM_PER_BYTE_B == 2
+
+    accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    for k in tl.range(0, tl.cdiv(K, BLOCK_K), num_stages=NUM_STAGES):
+        a = a_desc.load([offs_am, offs_k_a])
+        b = b_desc.load([offs_bn, offs_k_b])
+        scale_a = a_scale_desc.load([0, offs_scale_m, offs_scale_k, 0, 0])
+        scale_b = b_scale_desc.load([0, offs_scale_n, offs_scale_k, 0, 0])
+
+        scale_a = scale_a.reshape(rep_m, rep_k, 32, 4, 4).trans(0, 3, 2, 1, 4).reshape(BLOCK_M, BLOCK_K // VEC_SIZE)
+        scale_b = scale_b.reshape(rep_n, rep_k, 32, 4, 4).trans(0, 3, 2, 1, 4).reshape(BLOCK_N, BLOCK_K // VEC_SIZE)
+
+        if MIXED_PREC:
+            accumulator = tl.dot_scaled(a, scale_a, "e4m3", b.T, scale_b, "e2m1", accumulator)
+        elif ELEM_PER_BYTE_A == 2 and ELEM_PER_BYTE_B == 2:
+            accumulator = tl.dot_scaled(a, scale_a, "e2m1", b.T, scale_b, "e2m1", accumulator)
+        else:
+            accumulator = tl.dot_scaled(a, scale_a, "e4m3", b.T, scale_b, "e4m3", accumulator)
+
+        offs_k_a += BLOCK_K // ELEM_PER_BYTE_A
+        offs_k_b += BLOCK_K // ELEM_PER_BYTE_B
+        offs_scale_k += rep_k
+
+    c_desc.store([offs_am, offs_bn], accumulator.to(output_dtype))
+
+
+def block_scaled_matmul(a_desc, a_scale_desc, b_desc, b_scale_desc, dtype_dst, M, N, K, rep_m, rep_n, rep_k, configs):
+    output = torch.empty((M, N), dtype=dtype_dst, device="cuda")
+    if dtype_dst == torch.float32:
+        dtype_dst = 0
+    elif dtype_dst == torch.float16:
+        dtype_dst = 1
+    elif dtype_dst == torch.float8_e4m3fn:
+        dtype_dst = 2
+    else:
+        raise ValueError(f"Unsupported dtype: {dtype_dst}")
+
+    BLOCK_M = configs["BLOCK_SIZE_M"]
+    BLOCK_N = configs["BLOCK_SIZE_N"]
+    c_desc = TensorDescriptor.from_tensor(output, [BLOCK_M, BLOCK_N])
+
+    grid = (triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N), 1)
+    block_scaled_matmul_kernel[grid](
+        a_desc,
+        a_scale_desc,
+        b_desc,
+        b_scale_desc,
+        c_desc,
+        M,
+        N,
+        K,
+        dtype_dst,
+        configs["ELEM_PER_BYTE_A"],
+        configs["ELEM_PER_BYTE_B"],
+        configs["VEC_SIZE"],
+        configs["BLOCK_SIZE_M"],
+        configs["BLOCK_SIZE_N"],
+        configs["BLOCK_SIZE_K"],
+        rep_m,
+        rep_n,
+        rep_k,
+        configs["num_stages"],
+    )
+    return output
+
+
+def initialize_block_scaled(M, N, K, block_scale_type="nvfp4", compute_reference=False):
+    BLOCK_M = 128
+    BLOCK_N = 256
+    BLOCK_K = 256 if "fp4" in block_scale_type else 128
+    VEC_SIZE = 16 if block_scale_type == "nvfp4" else 32
+    assert block_scale_type in ["nvfp4", "mxfp4", "mxfp8", "mixed"], f"Invalid block scale type: {block_scale_type}"
+    ELEM_PER_BYTE_A = 2 if "fp4" in block_scale_type else 1
+    ELEM_PER_BYTE_B = 1 if block_scale_type == "mxfp8" else 2
+
+    device = "cuda"
+    a_ref = MXFP4Tensor(size=(M, K), device=device).random()
+    # Similar to Hopper's wgmma symmetric fp8 instruction, the RHS is expected
+    # to be in col-major layout for Blackwell's tcgen05.mma when using fp4 operands.
+    # To conform to the expected semantics of tl.dot_scaled, (M, K) x (K, N),
+    # the data is generated in col-major layout, packed along K for fp4, and then
+    # logically transposed. Note that if one operand is of fp8 precision, unlike Hopper,
+    # Blackwell supports both row-major and col-major layouts for the RHS matrix.
+    # For the mixed-precision case, the fp4 RHS can be either in row or col-major layout.
+    # But for performance reason, it is recommended to use col-major layout. If TMA is used
+    # for the fp4 RHS operand load in mixed-precision dot, as in this tutorial, it must be
+    # in col-major layout.
+    b_ref = MXFP4Tensor(size=(N, K), device=device).random()
+    if block_scale_type in ["mxfp8", "mixed"]:
+        a_ref = a_ref.to(torch.float32)
+        a = a_ref.to(torch.float8_e4m3fn)
+    else:
+        # Pack two fp4 elements per byte along K
+        a = a_ref.to_packed_tensor(dim=1)
+
+    if block_scale_type == "mxfp8":
+        b_ref = b_ref.to(torch.float32)
+        b = b_ref.to(torch.float8_e4m3fn)
+    else:
+        b = b_ref.to_packed_tensor(dim=1)
+
+    b_ref = b_ref.to(torch.float32).T
+
+    a_desc = TensorDescriptor.from_tensor(a, [BLOCK_M, BLOCK_K // ELEM_PER_BYTE_A])
+    b_desc = TensorDescriptor.from_tensor(b, [BLOCK_N, BLOCK_K // ELEM_PER_BYTE_B])
+
+    a_scale_shape = [M // 128, K // VEC_SIZE // 4, 32, 16]
+    b_scale_shape = [N // 128, K // VEC_SIZE // 4, 32, 16]
+    epsilon = 1e-8
+    a_scale = torch.rand(a_scale_shape, device=device) + epsilon
+    b_scale = torch.rand(b_scale_shape, device=device) + epsilon
+    if block_scale_type == "nvfp4":
+        a_scale = a_scale.to(torch.float8_e4m3fn)
+        b_scale = b_scale.to(torch.float8_e4m3fn)
+        a_scale_ref = a_scale
+        b_scale_ref = b_scale
+    elif block_scale_type in ["mxfp4", "mxfp8", "mixed"]:
+        a_scale_ref = MXScaleTensor(a_scale)
+        b_scale_ref = MXScaleTensor(b_scale)
+        a_scale = a_scale_ref.data
+        b_scale = b_scale_ref.data
+
+    rep_m = BLOCK_M // 128
+    rep_n = BLOCK_N // 128
+    rep_k = BLOCK_K // VEC_SIZE // 4
+
+    # Use 5D TMA descriptor [1, rep_m, rep_k, 2, 256] with uint8 elements.
+    # With 256 elements we better utilize the L2 and don't require the TMA
+    # engine to emit many small messages (16B) messages as with 32x16xu8.
+    a_scale_block_shape = [1, rep_m, rep_k, 2, 256]
+    b_scale_block_shape = [1, rep_n, rep_k, 2, 256]
+    a_scale = a_scale.reshape(1, a_scale_shape[0], a_scale.shape[1], 2, 256)
+    b_scale = b_scale.reshape(1, b_scale_shape[0], b_scale.shape[1], 2, 256)
+    a_scale_desc = TensorDescriptor.from_tensor(a_scale, block_shape=a_scale_block_shape)
+    b_scale_desc = TensorDescriptor.from_tensor(b_scale, block_shape=b_scale_block_shape)
+
+    reference = None
+    if compute_reference:
+        a_scale_ref = a_scale_ref.to(torch.float32)
+        b_scale_ref = b_scale_ref.to(torch.float32)
+
+        def unpack_scale(packed):
+            packed = packed.reshape(*packed.shape[:-2], 32, 4, 4)
+            num_chunk_m, num_chunk_k, _, _, _ = packed.shape
+            return packed.permute(0, 3, 2, 1, 4).reshape(num_chunk_m * 128, num_chunk_k * 4).contiguous()
+
+        a_scale_ref = unpack_scale(a_scale_ref).repeat_interleave(VEC_SIZE, dim=1)[:M, :K]
+        b_scale_ref = unpack_scale(b_scale_ref).repeat_interleave(VEC_SIZE, dim=1).T.contiguous()[:K, :N]
+        reference = torch.matmul(a_ref.to(torch.float32) * a_scale_ref, b_ref * b_scale_ref)
+
+    configs = {
+        "BLOCK_SIZE_M": BLOCK_M,
+        "BLOCK_SIZE_N": BLOCK_N,
+        "BLOCK_SIZE_K": BLOCK_K,
+        "num_stages": 4,
+        "ELEM_PER_BYTE_A": ELEM_PER_BYTE_A,
+        "ELEM_PER_BYTE_B": ELEM_PER_BYTE_B,
+        "VEC_SIZE": VEC_SIZE,
+    }
+    return a_desc, a_scale_desc, b_desc, b_scale_desc, rep_m, rep_n, rep_k, configs, reference
+
+
+def validate_block_scaled(M, N, K, block_scale_type="nvfp4"):
+    a_desc, a_scale, b_desc, b_scale, rep_m, rep_n, rep_k, configs, reference = initialize_block_scaled(
+        M, N, K, block_scale_type, compute_reference=True)
+    output = block_scaled_matmul(a_desc, a_scale, b_desc, b_scale, torch.float16, M, N, K, rep_m, rep_n, rep_k, configs)
+    torch.testing.assert_close(reference, output.to(torch.float32), atol=1e-3, rtol=1e-3)
+    print(f"✅ (pass {block_scale_type})")
+
+
+def bench_block_scaled(K, block_scale_type="nvfp4", reps=10):
+    assert K % 128 == 0
+    M = 8192
+    N = 8192
+    print(f"Problem Shape = {M}x{N}x{K}")
+
+    a_desc, a_scale, b_desc, b_scale, rep_m, rep_n, rep_k, configs, _ = initialize_block_scaled(
+        M, N, K, block_scale_type, compute_reference=False)
+    _ = block_scaled_matmul(a_desc, a_scale, b_desc, b_scale, torch.float16, M, N, K, rep_m, rep_n, rep_k, configs)
+
+    proton.activate(0)
+    for _ in range(reps):
+        _ = block_scaled_matmul(a_desc, a_scale, b_desc, b_scale, torch.float16, M, N, K, rep_m, rep_n, rep_k, configs)
+    proton.deactivate(0)
+    print("Done benchmarking")
+
+
+def show_profile(profile_name):
+    import triton.profiler.viewer as proton_viewer
+
+    metric_names = ["time/ms"]
+    metric_names = ["tflop/s"] + metric_names
+    file_name = f"{profile_name}.hatchet"
+    tree, metrics = proton_viewer.parse(metric_names, file_name)
+    proton_viewer.print_tree(tree, metrics)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-K", type=int, required=False, default=8192)
+    parser.add_argument("--K_range", type=int, nargs=2)
+    parser.add_argument("--K_step", type=int, default=512)
+    parser.add_argument("--bench", action="store_true", default=True)
+    parser.add_argument("--format", type=str, choices=["mxfp4", "nvfp4", "mxfp8", "mixed"], default="nvfp4")
+    args = parser.parse_args()
+
+    if not supports_block_scaling():
+        print("⛔ This example requires GPU support for block scaled matmul")
+    else:
+        if args.K and args.K_range is None:
+            args.K_range = [args.K, args.K]
+            args.K_step = 1  # doesn't matter as long as it's not 0
+
+        torch.manual_seed(42)
+
+        validate_block_scaled(8192, 8192, 8192, block_scale_type=args.format)
+
+        if args.bench:
+            proton.start("block_scaled_matmul", hook="triton")
+            proton.deactivate(0)  # Skip argument creation
+            for K in range(args.K_range[0], args.K_range[1] + 1, args.K_step):
+                bench_block_scaled(K, reps=10000, block_scale_type=args.format)
+            proton.finalize()
+            show_profile("block_scaled_matmul")