diff --git a/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B_sage3.sh b/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B_sage3.sh index b4da98b1c..74bba9b97 100644 --- a/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B_sage3.sh +++ b/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 diff --git a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_sage3.sh b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_sage3.sh index 7ca00def6..bcc1a99d8 100644 --- a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_sage3.sh +++ b/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" ) diff --git a/fastvideo/attention/backends/sage_attn3.py b/fastvideo/attention/backends/sage_attn3.py index 2bfb5a283..495e50228 100644 --- a/fastvideo/attention/backends/sage_attn3.py +++ b/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 diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py index 75140222e..37af9a8bf 100644 --- a/fastvideo/training/training_pipeline.py +++ b/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 diff --git a/test_quantization_precision.py b/test_quantization_precision.py index d3da670c6..023d40a19 100644 --- a/test_quantization_precision.py +++ b/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"), ] diff --git a/test_quantization_simple.py b/test_quantization_simple.py deleted file mode 100644 index 0e2f92ca6..000000000 --- a/test_quantization_simple.py +++ /dev/null @@ -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() diff --git a/test_triton_fp4_vs_cutlass.py b/test_triton_fp4_vs_cutlass.py new file mode 100644 index 000000000..c028425fb --- /dev/null +++ b/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() \ No newline at end of file diff --git a/tinyllama.py b/tinyllama.py new file mode 100644 index 000000000..48eee9a8c --- /dev/null +++ b/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() \ No newline at end of file diff --git a/triton_fp4.py b/triton_fp4.py index c74dc6271..782838132 100644 --- a/triton_fp4.py +++ b/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") \ No newline at end of file diff --git a/triton_original.py b/triton_original.py new file mode 100644 index 000000000..b52103e6e --- /dev/null +++ b/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") \ No newline at end of file