nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[misc] yapf pass under CI's interpreter (3.12) + pin hook language_version (#1702)

Browse Source
This commit is contained in:
William Lin
2026-08-08 22:24:17 -07:00
committed by GitHub
parent 6eb95693a1
commit 9028953625
582 changed files with 19387 additions and 23526 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo-kernel/python/fastvideo_kernel/triton_kernels/quant_utils.py
+76 -32
View File
@@ -3,33 +3,44 @@ 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,
dst_dtype: tl.constexpr,
mx_tensor_dtype: tl.constexpr = tl.uint8,
use_global_sf: tl.constexpr = True,
two_level_quant_P: tl.constexpr = False):
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,
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,
scale=src_scale,
s_dec=src_s_dec,
BLOCK_SIZE_OUT_DIM=BLOCK_SIZE_OUT_DIM,
BLOCK_SIZE_QUANT_DIM=BLOCK_SIZE_QUANT_DIM,
src_tensor = _compute_dequant(mx_tensor=src_tensor,
scale=src_scale,
s_dec=src_s_dec,
BLOCK_SIZE_OUT_DIM=BLOCK_SIZE_OUT_DIM,
BLOCK_SIZE_QUANT_DIM=BLOCK_SIZE_QUANT_DIM,
dst_dtype=dst_dtype)
return src_tensor, high_prec_src_tensor.to(src_tensor.dtype)
@triton.jit
def fake_quantize_q(Q, fake_Q, stride_z_q, stride_h_q,
stride_tok_q, stride_d_q,
fake_stride_z_q, fake_stride_h_q,
fake_stride_tok_q, fake_stride_d_q,
H, N_CTX_Q,
def fake_quantize_q(Q,
fake_Q,
stride_z_q,
stride_h_q,
stride_tok_q,
stride_d_q,
fake_stride_z_q,
fake_stride_h_q,
fake_stride_tok_q,
fake_stride_d_q,
H,
N_CTX_Q,
BLOCK_M: tl.constexpr,
HEAD_DIM: tl.constexpr,
use_global_sf: tl.constexpr = True):
@@ -46,18 +57,33 @@ def fake_quantize_q(Q, fake_Q, stride_z_q, stride_h_q,
q_valid = offs_m < N_CTX_Q
q = tl.load(Q + offs_m[:, None] * stride_tok_q + offs_k[None, :] * stride_d_q, mask=q_valid[:, None], other=0.0)
q, _ = fake_quantize(src_tensor=q, valid_src_mask=q_valid[:, None], BLOCK_SIZE_OUT_DIM=BLOCK_M, BLOCK_SIZE_QUANT_DIM=HEAD_DIM, dst_dtype=q.dtype, use_global_sf=use_global_sf)
q, _ = fake_quantize(src_tensor=q,
valid_src_mask=q_valid[:, None],
BLOCK_SIZE_OUT_DIM=BLOCK_M,
BLOCK_SIZE_QUANT_DIM=HEAD_DIM,
dst_dtype=q.dtype,
use_global_sf=use_global_sf)
tl.store(fake_Q + offs_m[:, None] * fake_stride_tok_q + offs_k[None, :] * fake_stride_d_q, q, mask=q_valid[:, None])
@triton.jit
def fake_quantize_kv(K, V, fake_K, fake_V, stride_z_kv, stride_h_kv,
stride_tok_kv, stride_d_kv,
fake_stride_z_kv, fake_stride_h_kv,
fake_stride_tok_kv, fake_stride_d_kv,
H, N_CTX_KV,
BLOCK_N: tl.constexpr,
HEAD_DIM: tl.constexpr,
use_global_sf: tl.constexpr = True):
def fake_quantize_kv(K,
V,
fake_K,
fake_V,
stride_z_kv,
stride_h_kv,
stride_tok_kv,
stride_d_kv,
fake_stride_z_kv,
fake_stride_h_kv,
fake_stride_tok_kv,
fake_stride_d_kv,
H,
N_CTX_KV,
BLOCK_N: tl.constexpr,
HEAD_DIM: tl.constexpr,
use_global_sf: tl.constexpr = True):
bhid = tl.program_id(1)
adj_kv = (stride_h_kv * (bhid % H) + stride_z_kv * (bhid // H))
fake_adj_kv = (fake_stride_h_kv * (bhid % H) + fake_stride_z_kv * (bhid // H))
@@ -70,11 +96,29 @@ def fake_quantize_kv(K, V, fake_K, fake_V, stride_z_kv, stride_h_kv,
start_n = pid * BLOCK_N
offs_n = start_n + tl.arange(0, BLOCK_N)
offs_k = tl.arange(0, HEAD_DIM)
kv_valid = offs_n < N_CTX_KV
k_block = tl.load(K + offs_n[:, None] * stride_tok_kv + offs_k[None, :] * stride_d_kv, mask=kv_valid[:, None], other=0.0)
v_block = tl.load(V + offs_n[:, None] * stride_tok_kv + offs_k[None, :] * stride_d_kv, mask=kv_valid[:, None], other=0.0)
k, _ = fake_quantize(src_tensor=k_block, valid_src_mask=kv_valid[:, None], BLOCK_SIZE_OUT_DIM=BLOCK_N, BLOCK_SIZE_QUANT_DIM=HEAD_DIM, dst_dtype=k_block.dtype, use_global_sf=use_global_sf)
v, _ = fake_quantize(src_tensor=v_block, valid_src_mask=kv_valid[:, None], BLOCK_SIZE_OUT_DIM=BLOCK_N, BLOCK_SIZE_QUANT_DIM=HEAD_DIM, dst_dtype=v_block.dtype, use_global_sf=use_global_sf)
tl.store(fake_K + offs_n[:, None] * fake_stride_tok_kv + offs_k[None, :] * fake_stride_d_kv, k, mask=kv_valid[:, None])
tl.store(fake_V + offs_n[:, None] * fake_stride_tok_kv + offs_k[None, :] * fake_stride_d_kv, v, mask=kv_valid[:, None])
k_block = tl.load(K + offs_n[:, None] * stride_tok_kv + offs_k[None, :] * stride_d_kv,
mask=kv_valid[:, None],
other=0.0)
v_block = tl.load(V + offs_n[:, None] * stride_tok_kv + offs_k[None, :] * stride_d_kv,
mask=kv_valid[:, None],
other=0.0)
k, _ = fake_quantize(src_tensor=k_block,
valid_src_mask=kv_valid[:, None],
BLOCK_SIZE_OUT_DIM=BLOCK_N,
BLOCK_SIZE_QUANT_DIM=HEAD_DIM,
dst_dtype=k_block.dtype,
use_global_sf=use_global_sf)
v, _ = fake_quantize(src_tensor=v_block,
valid_src_mask=kv_valid[:, None],
BLOCK_SIZE_OUT_DIM=BLOCK_N,
BLOCK_SIZE_QUANT_DIM=HEAD_DIM,
dst_dtype=v_block.dtype,
use_global_sf=use_global_sf)
tl.store(fake_K + offs_n[:, None] * fake_stride_tok_kv + offs_k[None, :] * fake_stride_d_kv,
k,
mask=kv_valid[:, None])
tl.store(fake_V + offs_n[:, None] * fake_stride_tok_kv + offs_k[None, :] * fake_stride_d_kv,
v,
mask=kv_valid[:, None])