From 52e3484f6ead29521e48e831d5f2305115327030 Mon Sep 17 00:00:00 2001 From: RandNMR73 Date: Fri, 10 Apr 2026 05:59:36 +0000 Subject: [PATCH] fix --- .../triton_kernels/nvfp4_utils.py | 22 ++++++++++++++++--- .../triton_kernels/quant_utils.py | 8 ++++--- 2 files changed, 24 insertions(+), 6 deletions(-) diff --git a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/nvfp4_utils.py b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/nvfp4_utils.py index afefe3742..6fdbb36e9 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/nvfp4_utils.py +++ b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/nvfp4_utils.py @@ -9,10 +9,24 @@ from triton.language.target_info import cuda_capability_geq MXFP_BLOCK_SIZE = tl.constexpr(16) @triton.jit -def _compute_quant_and_scale(src_tensor, valid_src_mask, use_global_sf=True, two_level_quant_P=False): +def _compute_quant_and_scale( + src_tensor, + valid_src_mask, + mx_tensor_dtype: tl.constexpr = tl.uint8, + use_global_sf=True, + two_level_quant_P=False, +): BLOCK_SIZE_OUT_DIM: tl.constexpr = src_tensor.shape[0] BLOCK_SIZE_QUANT_DIM: tl.constexpr = src_tensor.shape[1] BLOCK_SIZE_QUANT_MX_SCALE: tl.constexpr = src_tensor.shape[1] // MXFP_BLOCK_SIZE + is_fp4: tl.constexpr = mx_tensor_dtype == tl.uint8 + + tl.static_assert( + is_fp4 + or mx_tensor_dtype == tl.float8e4nv + or mx_tensor_dtype == tl.float8e5, + "mx_tensor_dtype must be uint8, float8e4nv, or float8e5", + ) # Explicit cast to fp32 since most ops are not supported on bfloat16. We avoid needless conversions to and from bf16 f32_tensor = src_tensor.to(tl.float32) @@ -52,7 +66,7 @@ def _compute_quant_and_scale(src_tensor, valid_src_mask, use_global_sf=True, two quant_tensor = tl.where(valid_src_mask, quant_tensor, 0.0) dequant_scale = s_dec_b_e4m3.reshape([BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE]) - if cuda_capability_geq(10, 0): + if is_fp4 and cuda_capability_geq(10, 0): # Convert scaled values to two f32 lanes and use PTX cvt to e2m1x2 with two f32 operands. pairs = tl.reshape(quant_tensor, [BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_DIM // 2, 2]) lo_f, hi_f = tl.split(pairs) @@ -74,7 +88,7 @@ def _compute_quant_and_scale(src_tensor, valid_src_mask, use_global_sf=True, two is_pure=True, pack=1, ) - else: + elif is_fp4: quant_tensor = quant_tensor.to(tl.uint32, bitcast=True) signs = quant_tensor & 0x80000000 exponents = (quant_tensor >> 23) & 0xFF @@ -113,6 +127,8 @@ def _compute_quant_and_scale(src_tensor, valid_src_mask, use_global_sf=True, two e2m1_value = tl.reshape(e2m1_value, [BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_DIM // 2, 2]) evens, odds = tl.split(e2m1_value) out_tensor = evens | (odds << 4) + else: + out_tensor = quant_tensor.to(mx_tensor_dtype) return out_tensor, dequant_scale, s_dec diff --git a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/quant_utils.py b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/quant_utils.py index 4355f70ec..31d4acc79 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/quant_utils.py +++ b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/quant_utils.py @@ -4,14 +4,16 @@ import triton.language as tl from .nvfp4_utils import _compute_quant_and_scale, _compute_dequant @triton.jit -def fake_quantize(src_tensor, valid_src_mask, BLOCK_SIZE_OUT_DIM: tl.constexpr, - BLOCK_SIZE_QUANT_DIM: tl.constexpr, +def fake_quantize(src_tensor, valid_src_mask, BLOCK_SIZE_OUT_DIM: tl.constexpr, + BLOCK_SIZE_QUANT_DIM: tl.constexpr, dst_dtype: tl.constexpr, + mx_tensor_dtype: tl.constexpr = tl.uint8, use_global_sf: tl.constexpr = True, two_level_quant_P: tl.constexpr = False): high_prec_src_tensor = src_tensor - src_tensor, src_scale, src_s_dec = _compute_quant_and_scale(src_tensor=src_tensor, + src_tensor, src_scale, src_s_dec = _compute_quant_and_scale(src_tensor=src_tensor, valid_src_mask=valid_src_mask, + mx_tensor_dtype=mx_tensor_dtype, use_global_sf=use_global_sf, two_level_quant_P=two_level_quant_P) src_tensor = _compute_dequant(mx_tensor=src_tensor,