From 2e6a7564ecfa6a27bd046b3cb40a6c2c0a48847c Mon Sep 17 00:00:00 2001 From: Yang Date: Sun, 6 Sep 2026 01:53:59 -0700 Subject: [PATCH] Accelerate Wan2.2 inference on NVIDIA Thor Add NVFP4 cuBLASLt projections and fused split-N FFN epilogues, CuTeDSL sparse BF16/FP8 attention, and Thor-specific build and configuration support. Preserve upstream interfaces and include regression coverage. Co-Authored-By: Claude Code --- app/gradio_demo.py | 3 +- ...fficial_stride_qkv_ffn0_gelu_epilogue.json | 40 + ...moe_i2v_distill_nvfp4_fa4_sparse_attn.json | 39 + ...rse_attn_official_stride_qkv_epilogue.json | 40 + ..._nvfp4_fa4_sparse_attn_splitn_compile.json | 40 + .../common/ops/attn/cutedsl_sparse_fmha.py | 168 + .../common/ops/attn/dynamic_sparse_attn.py | 291 +- lightx2v/common/ops/attn/flash_attn.py | 64 +- .../ops/attn/kernels/cutedsl_fmha_helpers.py | 902 ++++ .../kernels/cutedsl_sparse_fmha_kernel.py | 3919 +++++++++++++++++ lightx2v/common/ops/attn/utils/sla_util.py | 78 +- .../common/ops/attn/utils/sla_util_blhd.py | 58 + lightx2v/common/ops/mm/mm_weight.py | 123 +- lightx2v/infer.py | 25 +- .../networks/wan/infer/transformer_infer.py | 58 +- .../wan/weights/transformer_weights.py | 41 +- lightx2v/shot_runner/shot_base.py | 3 +- lightx2v/utils/set_config.py | 16 + lightx2v_kernel/CMakeLists.txt | 56 +- lightx2v_kernel/csrc/common_extension.cc | 40 + .../csrc/gemm/nvfp4_cublaslt_mm.cpp | 398 ++ .../csrc/gemm/nvfp4_quant_kernels_sm120.cu | 600 ++- .../gemm/nvfp4_scaled_mm_kernels_sm120.cu | 937 +++- lightx2v_kernel/include/lightx2v_kernel_ops.h | 50 + .../python/lightx2v_kernel/__init__.py | 3 + .../python/lightx2v_kernel/gemm.py | 101 +- .../test/nvfp4_nvfp4/test_qkv_cublaslt.py | 54 + .../test/nvfp4_nvfp4/test_split_n_stride.py | 133 + test_cases/test_attention_merge.py | 290 ++ test_cases/test_dynamic_sparse_attn_fa4.py | 126 + test_cases/test_infer_merge.py | 69 + test_cases/test_thor_config.py | 96 + test_cases/test_wan_merge.py | 246 ++ test_cases/test_wan_mxfp8_fuse_forwarding.py | 522 +++ test_cases/test_wan_nvfp4_qkv_cublaslt.py | 106 + 35 files changed, 9514 insertions(+), 221 deletions(-) create mode 100644 configs/wan22/thor/wan_moe_i2v_distill_nvfp4_cutedsl_sparse_attn_official_stride_qkv_ffn0_gelu_epilogue.json create mode 100644 configs/wan22/thor/wan_moe_i2v_distill_nvfp4_fa4_sparse_attn.json create mode 100644 configs/wan22/thor/wan_moe_i2v_distill_nvfp4_fa4_sparse_attn_official_stride_qkv_epilogue.json create mode 100644 configs/wan22/thor/wan_moe_i2v_distill_nvfp4_fa4_sparse_attn_splitn_compile.json create mode 100644 lightx2v/common/ops/attn/cutedsl_sparse_fmha.py create mode 100644 lightx2v/common/ops/attn/kernels/cutedsl_fmha_helpers.py create mode 100644 lightx2v/common/ops/attn/kernels/cutedsl_sparse_fmha_kernel.py create mode 100644 lightx2v_kernel/csrc/gemm/nvfp4_cublaslt_mm.cpp create mode 100644 lightx2v_kernel/test/nvfp4_nvfp4/test_qkv_cublaslt.py create mode 100644 lightx2v_kernel/test/nvfp4_nvfp4/test_split_n_stride.py create mode 100644 test_cases/test_attention_merge.py create mode 100644 test_cases/test_dynamic_sparse_attn_fa4.py create mode 100644 test_cases/test_infer_merge.py create mode 100644 test_cases/test_thor_config.py create mode 100644 test_cases/test_wan_merge.py create mode 100644 test_cases/test_wan_mxfp8_fuse_forwarding.py create mode 100644 test_cases/test_wan_nvfp4_qkv_cublaslt.py diff --git a/app/gradio_demo.py b/app/gradio_demo.py index 383a325b2..202b23419 100644 --- a/app/gradio_demo.py +++ b/app/gradio_demo.py @@ -17,7 +17,7 @@ from utils.ui_builder import build_ui, generate_unique_filename, get_auto_config_dict from lightx2v.utils.input_info import init_empty_input_info, update_input_info_from_dict -from lightx2v.utils.set_config import get_default_config +from lightx2v.utils.set_config import get_default_config, validate_thor_config warnings.filterwarnings("ignore", category=UserWarning, module="huggingface_hub") warnings.filterwarnings("ignore", category=UserWarning, module="huggingface_hub.utils") @@ -230,6 +230,7 @@ def run_inference( if config.get("lora_configs"): config["lora_dynamic_apply"] = True + validate_thor_config(config) logger.info(f"Using model: {model_path}") logger.info(f"Inference config:\n{json.dumps(config, indent=4, ensure_ascii=False)}") diff --git a/configs/wan22/thor/wan_moe_i2v_distill_nvfp4_cutedsl_sparse_attn_official_stride_qkv_ffn0_gelu_epilogue.json b/configs/wan22/thor/wan_moe_i2v_distill_nvfp4_cutedsl_sparse_attn_official_stride_qkv_ffn0_gelu_epilogue.json new file mode 100644 index 000000000..52d7359b8 --- /dev/null +++ b/configs/wan22/thor/wan_moe_i2v_distill_nvfp4_cutedsl_sparse_attn_official_stride_qkv_ffn0_gelu_epilogue.json @@ -0,0 +1,40 @@ +{ + "distill_method": "dmd2", + "infer_steps": 4, + "target_video_length": 81, + "text_len": 512, + "target_height": 720, + "target_width": 1280, + "self_attn_1_type": "dynamic_sparse_attn", + "dynamic_sparse_attn_setting": { + "sparsity_ratio": 0.9, + "operator": "cutedsl_fp8" + }, + "cross_attn_1_type": "flash_attn4", + "cross_attn_2_type": "flash_attn4", + "sample_guide_scale": [ + 3.5, + 3.5 + ], + "sample_shift": 7.0, + "enable_cfg": false, + "cpu_offload": false, + "t5_cpu_offload": true, + "vae_cpu_offload": false, + "use_image_encoder": false, + "boundary_step_index": 2, + "denoising_step_list": [ + 1000, + 750, + 500, + 250 + ], + "dit_quantized": true, + "dit_quant_scheme": "nvfp4", + "high_noise_quantized_ckpt": "path to i2v high_nvfp4.safetensors", + "low_noise_quantized_ckpt": "path to i2v low_nvfp4.safetensors", + "high_noise_original_ckpt": null, + "low_noise_original_ckpt": null, + "thor": true, + "use_compile": false +} diff --git a/configs/wan22/thor/wan_moe_i2v_distill_nvfp4_fa4_sparse_attn.json b/configs/wan22/thor/wan_moe_i2v_distill_nvfp4_fa4_sparse_attn.json new file mode 100644 index 000000000..dd99a3f33 --- /dev/null +++ b/configs/wan22/thor/wan_moe_i2v_distill_nvfp4_fa4_sparse_attn.json @@ -0,0 +1,39 @@ +{ + "distill_method": "dmd2", + "infer_steps": 4, + "target_video_length": 81, + "text_len": 512, + "target_height": 480, + "target_width": 832, + "self_attn_1_type": "dynamic_sparse_attn", + "dynamic_sparse_attn_setting": { + "sparsity_ratio": 0.9, + "operator": "fa4" + }, + "cross_attn_1_type": "flash_attn4", + "cross_attn_2_type": "flash_attn4", + "sample_guide_scale": [ + 3.5, + 3.5 + ], + "sample_shift": 7.0, + "enable_cfg": false, + "cpu_offload": false, + "t5_cpu_offload": true, + "vae_cpu_offload": false, + "use_image_encoder": false, + "boundary_step_index": 2, + "denoising_step_list": [ + 1000, + 750, + 500, + 250 + ], + "dit_quantized": true, + "dit_quant_scheme": "nvfp4", + "high_noise_quantized_ckpt": "path to i2v high_nvfp4.safetensors", + "low_noise_quantized_ckpt": "path to i2v low_nvfp4.safetensors", + "high_noise_original_ckpt": null, + "low_noise_original_ckpt": null, + "thor": true +} diff --git a/configs/wan22/thor/wan_moe_i2v_distill_nvfp4_fa4_sparse_attn_official_stride_qkv_epilogue.json b/configs/wan22/thor/wan_moe_i2v_distill_nvfp4_fa4_sparse_attn_official_stride_qkv_epilogue.json new file mode 100644 index 000000000..3dd6de6a9 --- /dev/null +++ b/configs/wan22/thor/wan_moe_i2v_distill_nvfp4_fa4_sparse_attn_official_stride_qkv_epilogue.json @@ -0,0 +1,40 @@ +{ + "distill_method": "dmd2", + "infer_steps": 4, + "target_video_length": 81, + "text_len": 512, + "target_height": 720, + "target_width": 1280, + "self_attn_1_type": "dynamic_sparse_attn", + "dynamic_sparse_attn_setting": { + "sparsity_ratio": 0.9, + "operator": "fa4" + }, + "cross_attn_1_type": "flash_attn4", + "cross_attn_2_type": "flash_attn4", + "sample_guide_scale": [ + 3.5, + 3.5 + ], + "sample_shift": 7.0, + "enable_cfg": false, + "cpu_offload": false, + "t5_cpu_offload": true, + "vae_cpu_offload": false, + "use_image_encoder": false, + "boundary_step_index": 2, + "denoising_step_list": [ + 1000, + 750, + 500, + 250 + ], + "dit_quantized": true, + "dit_quant_scheme": "nvfp4", + "high_noise_quantized_ckpt": "path to i2v high_nvfp4.safetensors", + "low_noise_quantized_ckpt": "path to i2v low_nvfp4.safetensors", + "high_noise_original_ckpt": null, + "low_noise_original_ckpt": null, + "thor": true, + "use_compile": false +} diff --git a/configs/wan22/thor/wan_moe_i2v_distill_nvfp4_fa4_sparse_attn_splitn_compile.json b/configs/wan22/thor/wan_moe_i2v_distill_nvfp4_fa4_sparse_attn_splitn_compile.json new file mode 100644 index 000000000..204a32de4 --- /dev/null +++ b/configs/wan22/thor/wan_moe_i2v_distill_nvfp4_fa4_sparse_attn_splitn_compile.json @@ -0,0 +1,40 @@ +{ + "distill_method": "dmd2", + "infer_steps": 4, + "target_video_length": 81, + "text_len": 512, + "target_height": 480, + "target_width": 832, + "self_attn_1_type": "dynamic_sparse_attn", + "dynamic_sparse_attn_setting": { + "sparsity_ratio": 0.9, + "operator": "fa4" + }, + "cross_attn_1_type": "flash_attn4", + "cross_attn_2_type": "flash_attn4", + "sample_guide_scale": [ + 3.5, + 3.5 + ], + "sample_shift": 7.0, + "enable_cfg": false, + "cpu_offload": false, + "t5_cpu_offload": true, + "vae_cpu_offload": false, + "use_image_encoder": false, + "boundary_step_index": 2, + "denoising_step_list": [ + 1000, + 750, + 500, + 250 + ], + "dit_quantized": true, + "dit_quant_scheme": "nvfp4", + "high_noise_quantized_ckpt": "path to i2v high_nvfp4.safetensors", + "low_noise_quantized_ckpt": "path to i2v low_nvfp4.safetensors", + "high_noise_original_ckpt": null, + "low_noise_original_ckpt": null, + "thor": true, + "use_compile": true +} diff --git a/lightx2v/common/ops/attn/cutedsl_sparse_fmha.py b/lightx2v/common/ops/attn/cutedsl_sparse_fmha.py new file mode 100644 index 000000000..4bc089021 --- /dev/null +++ b/lightx2v/common/ops/attn/cutedsl_sparse_fmha.py @@ -0,0 +1,168 @@ +import math +import threading + +import torch + +_COMPILED = {} +_COMPILE_LOCK = threading.Lock() + + +def _mark_shd_dynamic(tensor): + stride_order = (0, 1, 2) + return tensor.mark_layout_dynamic(leading_dim=2).mark_compact_shape_dynamic(mode=0, stride_order=stride_order).mark_compact_shape_dynamic(mode=1, stride_order=stride_order) + + +def _mark_1d_dynamic(tensor): + return tensor.mark_layout_dynamic(leading_dim=0).mark_compact_shape_dynamic(mode=0, stride_order=(0,)) + + +def _mark_sparse_count_dynamic(tensor): + stride_order = (0, 1, 2) + return ( + tensor.mark_layout_dynamic(leading_dim=2) + .mark_compact_shape_dynamic(mode=0, stride_order=stride_order) + .mark_compact_shape_dynamic(mode=1, stride_order=stride_order) + .mark_compact_shape_dynamic(mode=2, stride_order=stride_order) + ) + + +def _mark_sparse_indices_dynamic(tensor): + stride_order = (0, 1, 2, 3) + return ( + tensor.mark_layout_dynamic(leading_dim=3) + .mark_compact_shape_dynamic(mode=0, stride_order=stride_order) + .mark_compact_shape_dynamic(mode=1, stride_order=stride_order) + .mark_compact_shape_dynamic(mode=2, stride_order=stride_order) + ) + + +def _to_cute(tensor, element_type): + from cutlass.cute.runtime import from_dlpack + + cute_tensor = from_dlpack(tensor, assumed_align=16) + cute_tensor.element_type = element_type + return cute_tensor + + +def _cutlass_dtype(dtype): + import cutlass + + if dtype == torch.float16: + return cutlass.Float16 + if dtype == torch.bfloat16: + return cutlass.BFloat16 + if dtype == torch.float8_e4m3fn: + return cutlass.Float8E4M3FN + raise TypeError(f"unsupported CuTeDSL sparse FMHA dtype: {dtype}") + + +def _validate_inputs(q, k, v, cu_seqlens, max_seqlen, block_count, block_indices): + if q.ndim != 3 or k.ndim != 3 or v.ndim != 3: + raise ValueError("CuTeDSL sparse FMHA expects [total_tokens, heads, dim] Q/K/V") + if k.shape != v.shape or q.shape[0] != k.shape[0]: + raise ValueError("CuTeDSL sparse FMHA requires matching token counts and K/V shapes") + if q.shape[2] != k.shape[2] or not 0 < q.shape[2] <= 128: + raise ValueError("CuTeDSL sparse FMHA requires matching head dimensions in [1, 128]") + if q.shape[1] % k.shape[1] != 0: + raise ValueError("the number of Q heads must be divisible by the number of KV heads") + if q.dtype not in (torch.float16, torch.bfloat16, torch.float8_e4m3fn) or any(t.dtype != q.dtype for t in (k, v)): + raise TypeError("CuTeDSL sparse FMHA requires matching FP16, BF16, or FP8 Q/K/V") + if not all(t.is_cuda and t.is_contiguous() for t in (q, k, v)): + raise ValueError("CuTeDSL sparse FMHA requires contiguous CUDA Q/K/V") + if cu_seqlens.ndim != 1 or cu_seqlens.dtype != torch.int32: + raise TypeError("cu_seqlens must be a one-dimensional int32 tensor") + if max_seqlen <= 0: + raise ValueError("max_seqlen must be positive") + if block_count.dtype != torch.int32 or block_indices.dtype != torch.int32: + raise TypeError("sparse metadata must use int32") + if block_count.ndim != 3 or block_indices.ndim != 4: + raise ValueError("sparse block count/index metadata must have ranks 3 and 4") + if block_indices.shape[:3] != block_count.shape: + raise ValueError("sparse block count/index prefixes must match") + batch_size = cu_seqlens.shape[0] - 1 + q_blocks = (max_seqlen + 255) // 256 + if block_count.shape != (batch_size, q.shape[1], q_blocks): + raise ValueError(f"block_count must have shape [{batch_size}, {q.shape[1]}, {q_blocks}]") + if block_indices.shape[-1] <= 0: + raise ValueError("every sparse metadata row must retain at least one KV block") + + +def _make_args(q, k, v, out, cu_seqlens, max_seqlen, block_count, block_indices): + import cuda.bindings.driver as cuda + import cutlass + from cutlass.cute.typing import Float32, Int32 + + scale = 1.0 / math.sqrt(q.shape[-1]) + stream = cuda.CUstream(torch.cuda.current_stream(q.device).cuda_stream) + input_element_type = _cutlass_dtype(q.dtype) + output_element_type = _cutlass_dtype(out.dtype) + return ( + _mark_shd_dynamic(_to_cute(q, input_element_type)), + _mark_shd_dynamic(_to_cute(k, input_element_type)), + _mark_shd_dynamic(_to_cute(v, input_element_type)), + _mark_shd_dynamic(_to_cute(out, output_element_type)), + _mark_1d_dynamic(_to_cute(cu_seqlens, cutlass.Int32)), + Int32(max_seqlen), + Float32(scale * math.log2(math.e)), + Float32(scale), + Float32(1.0), + _mark_sparse_count_dynamic(_to_cute(block_count, cutlass.Int32)), + _mark_sparse_indices_dynamic(_to_cute(block_indices, cutlass.Int32)), + stream, + ) + + +@torch.compiler.disable +def cutedsl_sparse_fmha( + q, + k, + v, + cu_seqlens, + max_seqlen, + block_count, + block_indices, + output_dtype=None, +): + """Run the standalone Blackwell 256x128 block-sparse ViT FMHA.""" + cu_seqlens = cu_seqlens.to(device=q.device, dtype=torch.int32).contiguous() + block_count = block_count.to(device=q.device, dtype=torch.int32).contiguous() + block_indices = block_indices.to(device=q.device, dtype=torch.int32).contiguous() + _validate_inputs(q, k, v, cu_seqlens, max_seqlen, block_count, block_indices) + + from cutlass import cute + from cutlass.cute.typing import Float32 + + from .kernels import cutedsl_fmha_helpers as fmha_utils + from .kernels.cutedsl_sparse_fmha_kernel import ( + BlackwellFusedMultiHeadAttentionForward, + ) + + if output_dtype is None: + output_dtype = q.dtype + out = torch.empty(q.shape, device=q.device, dtype=output_dtype) + args = _make_args(q, k, v, out, cu_seqlens, max_seqlen, block_count, block_indices) + capacity = block_indices.shape[-1] + key = (q.device.index, q.dtype, out.dtype, q.shape[-1], capacity) + + compiled = _COMPILED.get(key) + if compiled is None: + with _COMPILE_LOCK: + compiled = _COMPILED.get(key) + if compiled is None: + mma_k = ((q.shape[-1] + 15) // 16) * 16 + actual_head_dim = q.shape[-1] if mma_k != q.shape[-1] else None + kernel = BlackwellFusedMultiHeadAttentionForward( + Float32, + Float32, + (128, 128, mma_k), + True, + fmha_utils.MaskEnum.RESIDUAL_MASK, + is_causal=False, + use_sliding_window=False, + actual_head_dim=actual_head_dim, + ) + compiled = cute.compile(kernel.__call_vit_sparse__, *args) + _COMPILED[key] = compiled + + compiled(*args) + return out diff --git a/lightx2v/common/ops/attn/dynamic_sparse_attn.py b/lightx2v/common/ops/attn/dynamic_sparse_attn.py index eb9281e81..9aa4a9548 100644 --- a/lightx2v/common/ops/attn/dynamic_sparse_attn.py +++ b/lightx2v/common/ops/attn/dynamic_sparse_attn.py @@ -1,3 +1,6 @@ +import inspect +from functools import wraps + import torch from loguru import logger @@ -6,16 +9,19 @@ from .kernels.sla_kernel import _attention from .kernels.sla_kernel_ar import _attention_ar from .template import AttnWeightTemplate -from .utils.sla_util import get_block_map, get_cuda_arch -from .utils.sla_util_blhd import get_block_map_blhd +from .utils.sla_util import block_lut_to_ordinal_metadata, get_block_map, get_cuda_arch +from .utils.sla_util_blhd import get_block_lut_blhd, get_block_map_blhd from .utils.sparge_util import block_map_incremental_lut_triton, block_map_ordinal_lut_triton, sage2_block_sparse_attn try: from flash_attn.cute import flash_attn_func as flash_attn_func_v4 - from flash_attn.cute.block_sparsity import BlockSparseTensorsTorch except (ImportError, AttributeError) as exc: logger.info(f"FlashAttention 4 is unavailable: {exc}") flash_attn_func_v4 = None + +try: + from flash_attn.cute.block_sparsity import BlockSparseTensorsTorch +except (ImportError, AttributeError): BlockSparseTensorsTorch = None try: @@ -30,6 +36,93 @@ magi_ffa_func = None +def _detect_fa4_sparse_api(): + if flash_attn_func_v4 is None: + return None + try: + parameters = inspect.signature(flash_attn_func_v4).parameters + except (TypeError, ValueError): + return None + + if "block_sparse_tensors" in parameters: + return "block_sparse_tensors" + + expanded_parameters = { + "mask_block_cnt", + "mask_block_idx", + "full_block_cnt", + "full_block_idx", + "block_size", + } + if expanded_parameters.issubset(parameters): + return "expanded" + return None + + +_FA4_SPARSE_API = _detect_fa4_sparse_api() + + +_FA4_BLOCKSPARSE_OP = None + + +if flash_attn_func_v4 is not None: + _fa4_impl = flash_attn_func_v4 + + # Keep the FA4 kernel eager while exposing a tensor-only boundary to + # Dynamo; upstream blocksparse/callable FA4 is eager-only today. + if hasattr(torch.library, "custom_op"): + try: + + @torch.library.custom_op("lightx2v_internal::fa4_blocksparse", mutates_args=()) + def _fa4_blocksparse_op( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + mask_block_cnt: torch.Tensor, + mask_block_idx: torch.Tensor, + full_block_cnt: torch.Tensor, + full_block_idx: torch.Tensor, + block_q: int, + block_k: int, + ) -> torch.Tensor: + sparse_kwargs = { + "mask_block_cnt": mask_block_cnt, + "mask_block_idx": mask_block_idx, + "full_block_cnt": full_block_cnt, + "full_block_idx": full_block_idx, + "block_size": (block_q, block_k), + } + if _FA4_SPARSE_API == "block_sparse_tensors": + sparse_kwargs = {"block_sparse_tensors": BlockSparseTensorsTorch(**sparse_kwargs)} + out, _ = _fa4_impl(q=q, k=k, v=v, **sparse_kwargs) + return out + + @_fa4_blocksparse_op.register_fake + def _fa4_blocksparse_op_fake( + q, + k, + v, + mask_block_cnt, + mask_block_idx, + full_block_cnt, + full_block_idx, + block_q, + block_k, + ): + return torch.empty_like(q) + + _FA4_BLOCKSPARSE_OP = _fa4_blocksparse_op + except RuntimeError as exc: + # Module reloads can encounter an already registered op. + logger.debug("FA4 blocksparse custom op registration skipped: {}", exc) + + @torch.compiler.disable + @wraps(_fa4_impl) + def flash_attn_func_v4(*args, **kwargs): + """Keep the blocksparse FA4 wrapper and kernel outside Dynamo fake mode.""" + return _fa4_impl(*args, **kwargs) + + @torch.library.custom_op( "lightx2v::dynamic_sparse_sage2", mutates_args=(), @@ -94,6 +187,20 @@ def __init__(self, config=None): elif self.operator == "sage3": self.BLKQ, self.BLKK = 128, 128 self.apply_func = self.apply_sage3 + elif self.operator in ("cutedsl", "cutedsl_vit_fmha"): + self.BLKQ, self.BLKK = 256, 128 + from .cutedsl_sparse_fmha import cutedsl_sparse_fmha + + self.cutedsl_sparse_fmha = cutedsl_sparse_fmha + self.apply_func = self.apply_cutedsl + elif self.operator == "cutedsl_fp8": + # Low-precision Q/K/V on the same CuTeDSL block-sparse FMHA: + # cast BF16 straight to FP8 E4M3 compute operands. + self.BLKQ, self.BLKK = 256, 128 + from .cutedsl_sparse_fmha import cutedsl_sparse_fmha + + self.cutedsl_sparse_fmha = cutedsl_sparse_fmha + self.apply_func = self.apply_cutedsl_fp8 elif self.operator == "fa4": self.BLKQ, self.BLKK = 128, 128 self.apply_func = self.apply_fa4 @@ -210,7 +317,7 @@ def apply_sage3( out = out.transpose(1, 2).reshape(max_seqlen_q, -1) return out - def apply_fa4( + def apply_cutedsl( self, q, k, @@ -221,34 +328,170 @@ def apply_fa4( max_seqlen_kv=None, **kwargs, ): - # (L, H, D) -> (B, L, H, D) - qt = q.unsqueeze(0).transpose(1, 2).contiguous() - kt = k.unsqueeze(0).transpose(1, 2).contiguous() - sparse_map, lut, real_topk = get_block_map(qt, kt, topk_ratio=self.topk, BLKQ=self.BLKQ, BLKK=self.BLKK) + del cu_seqlens_kv, max_seqlen_kv, kwargs + seqlen = q.shape[0] + + if q.shape[0] != k.shape[0] or k.shape != v.shape: + raise ValueError("CuTeDSL sparse attention only supports self-attention") + if cu_seqlens_q is not None and cu_seqlens_q.numel() != 2: + raise ValueError("dynamic CuTeDSL block selection currently requires batch size 1") + + # Pool directly from the layout consumed by CuTeDSL so full Q/K stay + # in their projection output buffers. + lut, _, num_k_blocks = get_block_lut_blhd( + q.unsqueeze(0), + k.unsqueeze(0), + topk_ratio=self.topk, + BLKQ=self.BLKQ, + BLKK=self.BLKK, + ) + if lut.shape[-1] == 0: + raise ValueError("CuTeDSL sparse attention requires at least one selected KV block") + block_indices, block_count = block_lut_to_ordinal_metadata(lut, num_k_blocks) - # (L, H, D) -> (B, L, H, D) + if cu_seqlens_q is None: + cu_seqlens_q = torch.tensor([0, seqlen], dtype=torch.int32, device=q.device) + if max_seqlen_q is None: + max_seqlen_q = seqlen + + out = self.cutedsl_sparse_fmha( + q.contiguous(), + k.contiguous(), + v.contiguous(), + cu_seqlens_q, + max_seqlen_q, + block_count, + block_indices, + ) + return out.reshape(seqlen, -1) + + def apply_cutedsl_fp8( + self, + q, + k, + v, + cu_seqlens_q=None, + cu_seqlens_kv=None, + max_seqlen_q=None, + max_seqlen_kv=None, + **kwargs, + ): + del cu_seqlens_kv, max_seqlen_kv, kwargs + seqlen = q.shape[0] + + if q.shape[0] != k.shape[0] or k.shape != v.shape: + raise ValueError("CuTeDSL FP8 sparse attention only supports self-attention") + if cu_seqlens_q is not None and cu_seqlens_q.numel() != 2: + raise ValueError("dynamic CuTeDSL block selection currently requires batch size 1") + + # Select blocks on the full-precision Q/K before any quantization. + lut, _, num_k_blocks = get_block_lut_blhd( + q.unsqueeze(0), + k.unsqueeze(0), + topk_ratio=self.topk, + BLKQ=self.BLKQ, + BLKK=self.BLKK, + ) + if lut.shape[-1] == 0: + raise ValueError("CuTeDSL sparse attention requires at least one selected KV block") + block_indices, block_count = block_lut_to_ordinal_metadata(lut, num_k_blocks) + + if cu_seqlens_q is None: + cu_seqlens_q = torch.tensor([0, seqlen], dtype=torch.int32, device=q.device) + if max_seqlen_q is None: + max_seqlen_q = seqlen + + # Cast low-precision operands straight to FP8 compute tensors. + q8 = q.to(torch.float8_e4m3fn).contiguous() + k8 = k.to(torch.float8_e4m3fn).contiguous() + v8 = v.to(torch.float8_e4m3fn).contiguous() + + out = self.cutedsl_sparse_fmha( + q8, + k8, + v8, + cu_seqlens_q, + max_seqlen_q, + block_count, + block_indices, + output_dtype=torch.bfloat16, + ) + return out.reshape(seqlen, -1) + + def apply_fa4( + self, + q, + k, + v, + cu_seqlens_q=None, + cu_seqlens_kv=None, + max_seqlen_q=None, + max_seqlen_kv=None, + **kwargs, + ): + # Build the sparse LUT directly from FA4's BLHD input layout. Pooling + # only writes one vector per block instead of cloning full-size Q/K. q = q.unsqueeze(0) k = k.unsqueeze(0) v = v.unsqueeze(0) + lut, _, num_k_blocks = get_block_lut_blhd( + q, + k, + topk_ratio=self.topk, + BLKQ=self.BLKQ, + BLKK=self.BLKK, + ) # (B, H, Q_block_num, K_block_num) - full_block_idx, full_block_cnt = block_map_ordinal_lut_triton(sparse_map) + full_block_idx, full_block_cnt = block_lut_to_ordinal_metadata(lut, num_k_blocks) mask_block_cnt = torch.zeros_like(full_block_cnt) mask_block_idx = torch.zeros_like(full_block_idx) - block_sparse_tensors = BlockSparseTensorsTorch( - mask_block_cnt=mask_block_cnt, - mask_block_idx=mask_block_idx, - full_block_cnt=full_block_cnt, - full_block_idx=full_block_idx, - block_size=(self.BLKQ, self.BLKK), - ) - - out, _ = flash_attn_func_v4( - q=q, - k=k, - v=v, - block_sparse_tensors=block_sparse_tensors, - ) + sparse_kwargs = { + "mask_block_cnt": mask_block_cnt, + "mask_block_idx": mask_block_idx, + "full_block_cnt": full_block_cnt, + "full_block_idx": full_block_idx, + "block_size": (self.BLKQ, self.BLKK), + } + if _FA4_SPARSE_API == "block_sparse_tensors": + if BlockSparseTensorsTorch is None: + raise RuntimeError("FA4 expects block_sparse_tensors, but BlockSparseTensorsTorch is unavailable") + if _FA4_BLOCKSPARSE_OP is not None and torch.compiler.is_compiling(): + out = _FA4_BLOCKSPARSE_OP( + q, + k, + v, + mask_block_cnt, + mask_block_idx, + full_block_cnt, + full_block_idx, + self.BLKQ, + self.BLKK, + ) + else: + out, _ = flash_attn_func_v4( + q=q, + k=k, + v=v, + block_sparse_tensors=BlockSparseTensorsTorch(**sparse_kwargs), + ) + elif _FA4_SPARSE_API == "expanded": + if _FA4_BLOCKSPARSE_OP is not None and torch.compiler.is_compiling(): + out = _FA4_BLOCKSPARSE_OP( + q, + k, + v, + mask_block_cnt, + mask_block_idx, + full_block_cnt, + full_block_idx, + self.BLKQ, + self.BLKK, + ) + else: + out, _ = flash_attn_func_v4(q=q, k=k, v=v, **sparse_kwargs) + else: + raise RuntimeError("Unsupported FA4 sparse attention API: expected block_sparse_tensors or expanded sparse parameters") out = out.reshape(max_seqlen_q, -1) return out diff --git a/lightx2v/common/ops/attn/flash_attn.py b/lightx2v/common/ops/attn/flash_attn.py index 2d19874d6..c6ff8c0a3 100755 --- a/lightx2v/common/ops/attn/flash_attn.py +++ b/lightx2v/common/ops/attn/flash_attn.py @@ -22,12 +22,21 @@ try: from flash_attn.cute import flash_attn_func as flash_attn_func_v4 - from flash_attn.cute.block_sparsity import BlockSparseTensorsTorch except (ImportError, AttributeError) as exc: - logger.info(f"FlashAttention 4 is unavailable: {exc}") + logger.info(f"FlashAttention 4 dense attention is unavailable: {exc}") flash_attn_func_v4 = None + +try: + from flash_attn.cute.block_sparsity import BlockSparseTensorsTorch +except (ImportError, AttributeError): BlockSparseTensorsTorch = None +try: + from flash_attn.cute import flash_attn_varlen_func as flash_attn_varlen_func_v4 +except (ImportError, AttributeError) as exc: + logger.info(f"FlashAttention 4 varlen attention is unavailable: {exc}") + flash_attn_varlen_func_v4 = None + from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER @@ -200,19 +209,42 @@ def apply( max_seqlen_kv=None, **kwargs, ): - if len(q.shape) == 3: - bs = 1 - q, k, v = q.unsqueeze(0), k.unsqueeze(0), v.unsqueeze(0) - elif len(q.shape) == 4: - bs = q.shape[0] - assert bs == 1, "flash_attn4 doesn't support flash_attn_varlen_func now. Just use it for batchsize = 1 for sure." - x, _ = flash_attn_func_v4( - q, - k, - v, - ) - x = x.reshape(bs * max_seqlen_q, -1) - return x + if q.ndim not in (3, 4): + raise ValueError(f"flash_attn4 expects a 3D or 4D tensor, got shape {tuple(q.shape)}") + if (cu_seqlens_q is None) != (cu_seqlens_kv is None): + raise ValueError("flash_attn4 requires both cu_seqlens_q and cu_seqlens_kv") + + total_seqlen = q.shape[0] if q.ndim == 3 else q.shape[0] * q.shape[1] + single_sequence = cu_seqlens_q is None and (q.ndim == 3 or q.shape[0] == 1) + if single_sequence and flash_attn_func_v4 is not None: + if q.ndim == 3: + q, k, v = q.unsqueeze(0), k.unsqueeze(0), v.unsqueeze(0) + x, _ = flash_attn_func_v4(q, k, v, softmax_scale=kwargs.get("softmax_scale"), causal=kwargs.get("causal", False)) + else: + if flash_attn_varlen_func_v4 is None: + raise RuntimeError("flash_attn4 requires flash_attn.cute.flash_attn_varlen_func for this input") + if cu_seqlens_q is None: + raise ValueError("flash_attn4 varlen requires cu_seqlens_q and cu_seqlens_kv") + if q.ndim == 4: + q = q.reshape(-1, q.shape[-2], q.shape[-1]) + k = k.reshape(-1, k.shape[-2], k.shape[-1]) + v = v.reshape(-1, v.shape[-2], v.shape[-1]) + if cu_seqlens_q.is_cpu: + cu_seqlens_q = cu_seqlens_q.to(q.device, non_blocking=True) + if cu_seqlens_kv.is_cpu: + cu_seqlens_kv = cu_seqlens_kv.to(k.device, non_blocking=True) + x, _ = flash_attn_varlen_func_v4( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k=cu_seqlens_kv, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_kv, + softmax_scale=kwargs.get("softmax_scale"), + causal=kwargs.get("causal", False), + ) + return x.reshape(total_seqlen, -1) @ATTN_WEIGHT_REGISTER("spas_flash_attn4") @@ -236,6 +268,8 @@ def apply( max_seqlen_kv=None, **kwargs, ): + if flash_attn_func_v4 is None: + raise RuntimeError("spas_flash_attn4 is not available: could not import flash_attn.cute.flash_attn_func. Install FlashAttention-4 first.") if len(q.shape) == 3: bs = 1 q, k, v = q.unsqueeze(0), k.unsqueeze(0), v.unsqueeze(0) diff --git a/lightx2v/common/ops/attn/kernels/cutedsl_fmha_helpers.py b/lightx2v/common/ops/attn/kernels/cutedsl_fmha_helpers.py new file mode 100644 index 000000000..22824821d --- /dev/null +++ b/lightx2v/common/ops/attn/kernels/cutedsl_fmha_helpers.py @@ -0,0 +1,902 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. + +import enum +from typing import Optional, Tuple + +import cutlass +import cutlass.cute as cute +from cutlass.cute.typing import Boolean +from cutlass.cutlass_dsl import Float32, Int32, extract_mlir_values, min, new_from_mlir_values +from cutlass.utils import WorkTileInfo +from cutlass.utils.hardware_info import HardwareInfo + +############################################################################## +# Fmha static tile scheduler +############################################################################## + + +class FmhaStaticTileSchedulerParams: + """A class to represent parameters for the FMHA (Fused Multi-Head Attention) static tile scheduler. + + This class holds the configuration parameters needed to initialize and configure + the tile scheduler for FMHA operations. + + :ivar is_persistent: Whether to use persistent kernel mode. + :type is_persistent: bool + :ivar problem_shape_mbh: Problem shape in (M, B, H) format. + :type problem_shape_mbh: cute.Shape + """ + + def __init__( + self, + is_persistent: bool, + problem_shape_mbh: cute.Shape, + *, + loc=None, + ip=None, + ): + """ + Initializes the FmhaStaticTileSchedulerParams with the given parameters. + + :param is_persistent: Whether to use persistent kernel mode. + :type is_persistent: bool + :param problem_shape_mbh: Problem shape in (M, B, H) format. + :type problem_shape_mbh: cute.Shape + """ + self.is_persistent = is_persistent + self.problem_shape_mbh = problem_shape_mbh + self._loc = loc + self._ip = ip + + def __extract_mlir_values__(self): + values, self._values_pos = [], [] + for obj in [self.problem_shape_mbh]: + obj_values = extract_mlir_values(obj) + values += obj_values + self._values_pos.append(len(obj_values)) + return values + + def __new_from_mlir_values__(self, values): + obj_list = [] + for obj, n_items in zip([self.problem_shape_mbh], self._values_pos): + obj_list.append(new_from_mlir_values(obj, values[:n_items])) + values = values[n_items:] + return FmhaStaticTileSchedulerParams(self.is_persistent, *(tuple(obj_list)), loc=self._loc) + + +class FmhaStaticTileScheduler: + """A static tile scheduler for FMHA (Fused Multi-Head Attention) operations. + + This class manages the scheduling of work tiles for FMHA kernels, supporting + both persistent and non-persistent kernel modes. It tracks the current work + position and advances through the problem space efficiently. + + :ivar _params: Scheduler parameters. + :type _params: FmhaStaticTileSchedulerParams + :ivar _blk_coord: Block coordinates. + :type _blk_coord: cute.Coord + :ivar _grid_shape: Grid shape for the kernel. + :type _grid_shape: cute.Shape + :ivar _is_persistent: Whether to use persistent kernel mode. + :type _is_persistent: bool + :ivar _current_work_linear_idx: Current linear work index. + :type _current_work_linear_idx: Int32 + :ivar _problem_shape_mbh: Problem shape in (M, B, H) format. + :type _problem_shape_mbh: cute.Layout + :ivar _num_blocks: Number of blocks in the problem. + :type _num_blocks: Int32 + :ivar _is_first_block: Whether this is the first block. + :type _is_first_block: bool + :ivar num_persistent_sm: Number of persistent SMs. + :type num_persistent_sm: Int32 + """ + + def __init__( + self, + params: FmhaStaticTileSchedulerParams, + current_work_linear_idx: Int32, + blk_coord: cute.Coord, + grid_shape: cute.Shape, + *, + loc=None, + ip=None, + ): + """ + Initializes the FmhaStaticTileScheduler with the given parameters. + + :param params: Scheduler parameters. + :type params: FmhaStaticTileSchedulerParams + :param current_work_linear_idx: Current linear work index. + :type current_work_linear_idx: Int32 + :param blk_coord: Block coordinates. + :type blk_coord: cute.Coord + :param grid_shape: Grid shape for the kernel. + :type grid_shape: cute.Shape + """ + self._params = params + self._blk_coord = blk_coord + self._grid_shape = grid_shape + self._is_persistent = params.is_persistent + self._current_work_linear_idx = current_work_linear_idx + self._problem_shape_mbh = cute.make_layout(params.problem_shape_mbh, loc=loc, ip=ip) + self._num_blocks = cute.size(self._problem_shape_mbh, loc=loc, ip=ip) + self._is_first_block = True + self.num_persistent_sm = cute.size(grid_shape, loc=loc, ip=ip) + self._loc = loc + self._ip = ip + + # called by host + @staticmethod + def get_grid_shape( + params: FmhaStaticTileSchedulerParams, + *, + loc=None, + ip=None, + ) -> cute.Shape: + """ + Determine the grid shape for the FMHA kernel. + + For persistent kernels, the grid shape is limited by the number of SMs + (Streaming Multiprocessors) available on the device. For non-persistent + kernels, the grid shape matches the problem shape. + + :param params: Scheduler parameters. + :type params: FmhaStaticTileSchedulerParams + + :return: Grid shape as (M, B, H) tuple. + :rtype: cute.Shape + """ + if params.is_persistent: + hardware_info = HardwareInfo() + sm_count = hardware_info.get_device_multiprocessor_count() + return ( + min(sm_count, cute.size(params.problem_shape_mbh, loc=loc, ip=ip)), + 1, + 1, + ) + else: + return params.problem_shape_mbh + + @staticmethod + def check_valid_work_for_seqlen_q( + q_tiler: int, + current_idx: Int32, + seqlen_q: Int32, + ) -> Boolean: + """ + Check if the current work index is valid for the given query sequence length. + + This method verifies that the current work tile index multiplied by the + query tiler size is within the bounds of the query sequence length. + + :param q_tiler: Query tiler size. + :type q_tiler: int + :param current_idx: Current work index. + :type current_idx: Int32 + :param seqlen_q: Query sequence length. + :type seqlen_q: Int32 + + :return: True if the work is valid, False otherwise. + :rtype: Boolean + """ + return current_idx * q_tiler < seqlen_q + + def get_current_work(self, *, loc=None, ip=None) -> WorkTileInfo: + """ + Get information about the current work tile. + + Determines if the current work is valid and computes the tile coordinates + based on whether the kernel is persistent or non-persistent. + + :return: WorkTileInfo containing tile coordinates and validity flag. + :rtype: WorkTileInfo + """ + is_valid = self._current_work_linear_idx < self._num_blocks if self._is_persistent else self._is_first_block + + blk_coord = (0, 0, 0) + if self._is_persistent: + blk_coord = self._problem_shape_mbh.get_hier_coord(self._current_work_linear_idx, loc=loc, ip=ip) + else: + blk_coord = self._blk_coord + + # cur_tile_coord is (mid, 0, (bid, hid)) + cur_tile_coord = ( + blk_coord[0], + 0, + (blk_coord[1], blk_coord[2]), + ) + + return WorkTileInfo(cur_tile_coord, is_valid) + + def initial_work_tile_info(self, *, loc=None, ip=None): + """ + Get the initial work tile information. + + :return: Initial WorkTileInfo. + :rtype: WorkTileInfo + """ + return self.get_current_work(loc=loc, ip=ip) + + def advance_to_next_work(self, *, advance_count=1, loc=None, ip=None): + """ + Advance to the next work tile. + + For persistent kernels, advances by the number of persistent SMs. + For non-persistent kernels, marks that the first block has been processed. + + :param advance_count: Number of steps to advance (default: 1). + :type advance_count: int + """ + if self._is_persistent: + self._current_work_linear_idx += advance_count * self.num_persistent_sm + self._is_first_block = False + + def __extract_mlir_values__(self): + values = extract_mlir_values(self._params) + values.extend(extract_mlir_values(self._current_work_linear_idx)) + values.extend(extract_mlir_values(self._blk_coord)) + values.extend(extract_mlir_values(self._grid_shape)) + return values + + def __new_from_mlir_values__(self, values): + assert len(values) == 10 + new_params = new_from_mlir_values(self._params, values[0:3]) + new_current_work_linear_idx = new_from_mlir_values(self._current_work_linear_idx, [values[3]]) + new_blk_coord = new_from_mlir_values(self._blk_coord, values[4:7]) + new_grid_shape = new_from_mlir_values(self._grid_shape, values[7:]) + return FmhaStaticTileScheduler(new_params, new_current_work_linear_idx, new_blk_coord, new_grid_shape) + + +def create_fmha_static_tile_scheduler( + params: FmhaStaticTileSchedulerParams, + blk_coord: cute.Coord, + grid_shape: cute.Shape, +) -> FmhaStaticTileScheduler: + """ + Create a new FMHA static tile scheduler. + + :param params: Scheduler parameters. + :type params: FmhaStaticTileSchedulerParams + :param blk_coord: Block coordinates. + :type blk_coord: cute.Coord + :param grid_shape: Grid shape. + :type grid_shape: cute.Shape + + :return: New FmhaStaticTileScheduler instance. + :rtype: FmhaStaticTileScheduler + """ + return FmhaStaticTileScheduler(params, blk_coord[0], blk_coord, grid_shape) + + +def create_fmha_static_tile_scheduler_params( + is_persistent: bool, + problem_shape_mbh: cute.Shape, +) -> FmhaStaticTileSchedulerParams: + """ + Create FMHA static tile scheduler parameters. + + :param is_persistent: Whether to use persistent kernel mode. + :type is_persistent: bool + :param problem_shape_mbh: Problem shape in (M, B, H) format. + :type problem_shape_mbh: cute.Shape + + :return: New FmhaStaticTileSchedulerParams instance. + :rtype: FmhaStaticTileSchedulerParams + """ + return FmhaStaticTileSchedulerParams(is_persistent, problem_shape_mbh) + + +def compute_grid( + o_shape: cute.Shape, + cta_tiler: Tuple[int, int, int], + is_persistent: bool, +) -> Tuple[FmhaStaticTileSchedulerParams, Tuple[int, int, int]]: + """ + Compute grid parameters for FMHA operation. + + This function calculates the appropriate grid shape and scheduler parameters + based on the output tensor shape, CTA (Cooperative Thread Array) tiler, + and whether to use persistent kernel mode. + + The output tensor o has shape (s, d, ((h_r, h_k), b)) where: + - s: sequence length + - d: head dimension + - h_r: number of heads for query + - h_k: number of heads for key + - b: batch size + + :param o_shape: Output tensor shape for grid computation. + :type o_shape: cute.Shape + :param cta_tiler: CTA tiler dimensions (M, N, K). + :type cta_tiler: Tuple[int, int, int] + :param is_persistent: Whether to use persistent kernel mode. + :type is_persistent: bool + + :return: Tuple of (scheduler_params, grid_shape). + :rtype: Tuple[FmhaStaticTileSchedulerParams, Tuple[int, int, int]] + """ + tile_sched_params = create_fmha_static_tile_scheduler_params( + is_persistent, + ( + cute.ceil_div(cute.size(o_shape[0]), cta_tiler[0]), + cute.size(o_shape[2][0]), + cute.size(o_shape[2][1]), + ), + ) + grid = FmhaStaticTileScheduler.get_grid_shape(tile_sched_params) + + return tile_sched_params, grid + + +############################################################################## +# Fused Mask +############################################################################## + + +class MaskEnum(enum.Enum): + """Enumeration of mask types for FMHA operations. + + - RESIDUAL_MASK: Residual mask for handling variable sequence lengths + - WINDOW_MASK: Window mask for attention which also includes causal and no mask + - WINDOW_MASK_INFERENCE: Same as the window mask, but has the limitation that the end of q is aligned with the end of k + - WINDOW_MASK_BWD: Window mask for backward pass + - WINDOW_MASK_BWD_INFERENCE: Same as the window mask for backward pass, but has the limitation that the end of q is aligned with the end of k + """ + + RESIDUAL_MASK = enum.auto() + RESIDUAL_MASK_BWD = enum.auto() + WINDOW_MASK = enum.auto() + WINDOW_MASK_INFERENCE = enum.auto() + WINDOW_MASK_BWD = enum.auto() + WINDOW_MASK_BWD_INFERENCE = enum.auto() + + +class FusedMask: + """A fused mask implementation for FMHA operations. + + This class handles different types of attention masks including no mask, + residual mask for variable sequence lengths, and causal mask for + autoregressive attention patterns. + + The class provides methods to: + - Calculate trip counts for different mask types + - Apply masks to attention scores + - Handle masked and unmasked trip calculations + """ + + def get_trip_count( + mask_type: MaskEnum, + blk_coord: cute.Coord, + tile_shape: cute.Shape, + seqlen_q: Int32, + seqlen_k: Int32, + window_size_left: Optional[Int32] = None, + window_size_right: Optional[Int32] = None, + ) -> Int32: + """ + Calculate the number of trips needed for the current block. + + The trip count depends on the mask type and the block coordinates. + For causal masks, it considers the autoregressive constraint. + + :param mask_type: Type of mask to use + :type mask_type: utils.MaskEnum + :param blk_coord: Block coordinates. + :type blk_coord: cute.Coord + :param tile_shape: Shape of the tile. + :type tile_shape: cute.Shape + :param seqlen_q: Query sequence length for attention computation. + :type seqlen_q: Int32 + :param seqlen_k: Key sequence length for attention computation. + :type seqlen_k: Int32 + :param window_size_left: Left-side sliding window size for attention masking. + :type window_size_left: Optional[Int32] + :param window_size_right: Right-side sliding window size for attention masking. + :type window_size_right: Optional[Int32] + + :return: Number of trips needed. + :rtype: Int32 + """ + result = 0 + offset = 0 + if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_INFERENCE): + offset = seqlen_k - seqlen_q + if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_BWD_INFERENCE): + offset = seqlen_q - seqlen_k + if cutlass.const_expr(mask_type == MaskEnum.RESIDUAL_MASK): + result = cute.ceil_div(seqlen_k, tile_shape[1]) + if cutlass.const_expr(mask_type is MaskEnum.RESIDUAL_MASK_BWD): + result = cute.ceil_div(seqlen_q, tile_shape[0]) + if cutlass.const_expr(mask_type == MaskEnum.WINDOW_MASK or mask_type == MaskEnum.WINDOW_MASK_INFERENCE): + if cutlass.const_expr(window_size_right is None): + result = cute.ceil_div(seqlen_k, tile_shape[1]) + else: + max_idx_q = (blk_coord[0] + 1) * tile_shape[0] + idx_k = max_idx_q + offset + window_size_right + tmp_blocks_k = cute.ceil_div(idx_k, tile_shape[1]) + max_blocks_k = cute.ceil_div(seqlen_k, tile_shape[1]) + result = min(max_blocks_k, tmp_blocks_k) + if cutlass.const_expr(mask_type == MaskEnum.WINDOW_MASK_BWD or mask_type == MaskEnum.WINDOW_MASK_BWD_INFERENCE): + if cutlass.const_expr(window_size_left is None): + result = cute.ceil_div(seqlen_q, tile_shape[0]) + else: + max_idx_k = (blk_coord[1] + 1) * tile_shape[1] + idx_k = max_idx_k + offset + window_size_left + tmp_blocks_q = cute.ceil_div(idx_k, tile_shape[0]) + max_blocks_q = cute.ceil_div(seqlen_q, tile_shape[0]) + result = min(max_blocks_q, tmp_blocks_q) + start_block = FusedMask.get_trip_start( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + result = result - start_block + return result + + @cute.jit + def get_trip_start( + mask_type: MaskEnum, + blk_coord: cute.Coord, + tile_shape: cute.Shape, + seqlen_q: Int32, + seqlen_k: Int32, + window_size_left: Optional[Int32] = None, + window_size_right: Optional[Int32] = None, + ) -> Int32: + """ + Get the start of the trip for the current block. + + :param mask_type: Type of mask to use + :type mask_type: utils.MaskEnum + :param blk_coord: Block coordinates. + :type blk_coord: cute.Coord + :param tile_shape: Shape of the tile. + :type tile_shape: cute.Shape + :param seqlen_q: Query sequence length for attention computation. + :type seqlen_q: Int32 + :param seqlen_k: Key sequence length for attention computation. + :type seqlen_k: Int32 + :param window_size_left: Left-side sliding window size for attention masking. + :type window_size_left: Optional[Int32] + :param window_size_right: Right-side sliding window size for attention masking. + :type window_size_right: Optional[Int32] + """ + result = 0 + offset = 0 + if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_INFERENCE): + offset = seqlen_k - seqlen_q + if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_BWD_INFERENCE): + offset = seqlen_q - seqlen_k + if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK or mask_type is MaskEnum.WINDOW_MASK_INFERENCE): + if cutlass.const_expr(window_size_left is not None): + min_idx_q = blk_coord[0] * tile_shape[0] + idx_k = min_idx_q + offset - window_size_left + tmp_blocks_k = idx_k // tile_shape[1] + result = max(tmp_blocks_k, result) + if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_BWD or mask_type is MaskEnum.WINDOW_MASK_BWD_INFERENCE): + if cutlass.const_expr(window_size_right is not None): + min_idx_k = blk_coord[1] * tile_shape[1] + idx_q = min_idx_k + offset - window_size_right + tmp_blocks_q = idx_q // tile_shape[0] + result = max(tmp_blocks_q, result) + return result + + @cute.jit + def get_leading_mask_id( + mask_type: MaskEnum, + blk_coord: cute.Coord, + tile_shape: cute.Shape, + seqlen_q: Int32, + seqlen_k: Int32, + window_size_left: Optional[Int32] = None, + window_size_right: Optional[Int32] = None, + ) -> Tuple[Int32, Int32]: + """ + Get the begin and end tile idx for the leading mask. + + :param mask_type: Type of mask to use + :type mask_type: utils.MaskEnum + :param blk_coord: Block coordinates. + :type blk_coord: cute.Coord + :param tile_shape: Shape of the tile. + :type tile_shape: cute.Shape + :param seqlen_q: Query sequence length for attention computation. + :type seqlen_q: Int32 + :param seqlen_k: Key sequence length for attention computation. + :type seqlen_k: Int32 + :param window_size_left: Left-side sliding window size for attention masking. + :type window_size_left: Optional[Int32] + :param window_size_right: Right-side sliding window size for attention masking. + :type window_size_right: Optional[Int32] + + :return: Tuple of (begin, end) tile idx for the leading mask. + :rtype: Tuple[Int32, Int32] + """ + offset = 0 + if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_INFERENCE): + offset = seqlen_k - seqlen_q + if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_BWD_INFERENCE): + offset = seqlen_q - seqlen_k + leading_mask_begin = FusedMask.get_trip_start( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + trip_count = FusedMask.get_trip_count( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + + leading_mask_end = leading_mask_begin + if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK or mask_type is MaskEnum.WINDOW_MASK_INFERENCE): + if cutlass.const_expr(window_size_left is not None): + min_idx_q = (blk_coord[0] + 1) * tile_shape[0] + offset - window_size_left + leading_mask_end = min( + cute.ceil_div(min_idx_q, tile_shape[1]) - 1, + trip_count + leading_mask_begin - 1, + ) + else: + leading_mask_end = leading_mask_begin - 1 + elif cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_BWD or mask_type is MaskEnum.WINDOW_MASK_BWD_INFERENCE): + if cutlass.const_expr(window_size_right is not None): + min_idx_k = (blk_coord[1] + 1) * tile_shape[1] + offset - window_size_right + leading_mask_end = cute.ceil_div(min_idx_k, tile_shape[0]) - 1 + else: + leading_mask_end = leading_mask_begin - 1 + return leading_mask_begin, leading_mask_end + + @cute.jit + def get_trailing_mask_id( + mask_type: MaskEnum, + blk_coord: cute.Coord, + tile_shape: cute.Shape, + seqlen_q: Int32, + seqlen_k: Int32, + window_size_left: Optional[Int32] = None, + window_size_right: Optional[Int32] = None, + ) -> Tuple[Optional[Int32], Optional[Int32]]: + """ + Get the begin and end tile idx for the trailing mask. + + :param mask_type: Type of mask to use + :type mask_type: utils.MaskEnum + :param blk_coord: Block coordinates. + :type blk_coord: cute.Coord + :param tile_shape: Shape of the tile. + :type tile_shape: cute.Shape + :param seqlen_q: Query sequence length for attention computation. + :type seqlen_q: Int32 + :param seqlen_k: Key sequence length for attention computation. + :type seqlen_k: Int32 + :param window_size_left: Left-side sliding window size for attention masking. + :type window_size_left: Optional[Int32] + :param window_size_right: Right-side sliding window size for attention masking. + :type window_size_right: Optional[Int32] + + :return: Tuple of (begin, end) tile idx for the trailing mask. + :rtype: Tuple[Int32, Int32] + """ + offset = 0 + if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_INFERENCE): + offset = seqlen_k - seqlen_q + if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_BWD_INFERENCE): + offset = seqlen_q - seqlen_k + trip_start = FusedMask.get_trip_start( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + trip_count = FusedMask.get_trip_count( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + + trailing_mask_begin, trailing_mask_end = None, None + if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK or mask_type is MaskEnum.WINDOW_MASK_INFERENCE): + if cutlass.const_expr(window_size_right is not None): + min_idx_q = blk_coord[0] * tile_shape[0] + offset + window_size_right + trailing_mask_begin = min(min_idx_q // tile_shape[1], trip_count + trip_start - 1) + trailing_mask_end = trip_count + trip_start - 1 + else: + # last tile, we always apply mask on it regardless whether it's a residual tile + trailing_mask_begin = trip_count + trip_start - 1 + trailing_mask_end = trip_count + trip_start - 1 + else: + if cutlass.const_expr(window_size_left is not None): + min_idx_k = blk_coord[1] * tile_shape[1] + offset + window_size_left + 1 + max_idx_k = (blk_coord[1] + 1) * tile_shape[1] + offset + window_size_left + trailing_mask_begin = min( + cute.ceil_div(min_idx_k, tile_shape[0]) - 1, + trip_count + trip_start - 1, + ) + trailing_mask_end = min( + cute.ceil_div(max_idx_k, tile_shape[0]) - 1, + trip_count + trip_start - 1, + ) + else: + # last tile, we always apply mask on it regardless whether it's a residual tile + trailing_mask_begin = trip_count + trip_start - 1 + trailing_mask_end = trip_count + trip_start - 1 + + return trailing_mask_begin, trailing_mask_end + + @cute.jit + def get_masked_leading_count( + mask_type: MaskEnum, + blk_coord: cute.Coord, + tile_shape: cute.Shape, + seqlen_q: Int32, + seqlen_k: Int32, + window_size_left: Optional[Int32] = None, + window_size_right: Optional[Int32] = None, + ) -> Int32: + """ + Calculate the number of masked trips for the leading mask. + + This is used for blocks that need special handling due to masking. + + :param mask_type: Type of mask to use + :type mask_type: utils.MaskEnum + :param blk_coord: Block coordinates. + :type blk_coord: cute.Coord + :param tile_shape: Shape of the tile. + :type tile_shape: cute.Shape + :param seqlen_q: Query sequence length for attention computation. + :type seqlen_q: Int32 + :param seqlen_k: Key sequence length for attention computation. + :type seqlen_k: Int32 + :param window_size_left: Left-side sliding window size for attention masking. + :type window_size_left: Optional[Int32] + :param window_size_right: Right-side sliding window size for attention masking. + :type window_size_right: Optional[Int32] + + :return: Number of masked trips. + :rtype: Int32 + """ + result = 0 + if cutlass.const_expr(mask_type is not MaskEnum.RESIDUAL_MASK and mask_type is not MaskEnum.RESIDUAL_MASK_BWD): + if cutlass.const_expr(window_size_left is not None or window_size_right is not None): + leading_mask_begin, leading_mask_end = FusedMask.get_leading_mask_id( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + result = max(leading_mask_end - leading_mask_begin + 1, 0) + + return result + + @cute.jit + def get_masked_trailing_count( + mask_type: MaskEnum, + blk_coord: cute.Coord, + tile_shape: cute.Shape, + seqlen_q: Int32, + seqlen_k: Int32, + window_size_left: Optional[Int32] = None, + window_size_right: Optional[Int32] = None, + rem_count: Optional[Int32] = 0, + ) -> Int32: + """ + Calculate the number of masked trips for the trailing mask. + + This is used for blocks that need special handling due to masking. + + :param mask_type: Type of mask to use + :type mask_type: utils.MaskEnum + :param blk_coord: Block coordinates. + :type blk_coord: cute.Coord + :param tile_shape: Shape of the tile. + :type tile_shape: cute.Shape + :param seqlen_q: Query sequence length for attention computation. + :type seqlen_q: Int32 + :param seqlen_k: Key sequence length for attention computation. + :type seqlen_k: Int32 + :param window_size_left: Left-side sliding window size for attention masking. + :type window_size_left: Optional[Int32] + :param window_size_right: Right-side sliding window size for attention masking. + :type window_size_right: Optional[Int32] + :param rem_count: Remaining count from previous calculations. + :type rem_count: Int32 + + :return: Number of masked trips. + :rtype: Int32 + """ + result = 0 + + if cutlass.const_expr(mask_type is not MaskEnum.RESIDUAL_MASK and mask_type is not MaskEnum.RESIDUAL_MASK_BWD): + if cutlass.const_expr(window_size_left is not None or window_size_right is not None): + trailing_mask_begin, trailing_mask_end = FusedMask.get_trailing_mask_id( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + leading_mask_begin, leading_mask_end = FusedMask.get_leading_mask_id( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + if cutlass.const_expr(trailing_mask_begin is not None and trailing_mask_end is not None): + if trailing_mask_begin <= leading_mask_end: + result = max(trailing_mask_end - leading_mask_end, 0) + else: + result = max(trailing_mask_end - trailing_mask_begin + 1, 0) + else: + if seqlen_k % tile_shape[1] != 0: + result = 1 + else: + result = 0 + + return result + rem_count + + @cute.jit + def get_unmasked_trip_count( + mask_type: MaskEnum, + blk_coord: cute.Coord, + tile_shape: cute.Shape, + seqlen_q: Int32, + seqlen_k: Int32, + window_size_left: Optional[Int32] = None, + window_size_right: Optional[Int32] = None, + ) -> Int32: + """ + Calculate the number of unmasked trips for the current block. + + This represents the number of trips that don't require special + masking treatment. + + :param mask_type: Type of mask to use + :type mask_type: utils.MaskEnum + :param blk_coord: Block coordinates. + :type blk_coord: cute.Coord + :param tile_shape: Shape of the tile. + :type tile_shape: cute.Shape + :param seqlen_q: Query sequence length for attention computation. + :type seqlen_q: Int32 + :param seqlen_k: Key sequence length for attention computation. + :type seqlen_k: Int32 + :param window_size_left: Left-side sliding window size for attention masking. + :type window_size_left: Optional[Int32] + :param window_size_right: Right-side sliding window size for attention masking. + :type window_size_right: Optional[Int32] + + :return: Number of unmasked trips. + :rtype: Int32 + """ + result = ( + FusedMask.get_trip_count( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + - FusedMask.get_masked_leading_count( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + - FusedMask.get_masked_trailing_count( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + 0, + ) + ) + return result + + @cute.jit + def apply_mask( + mask_type: MaskEnum, + acc_qk: cute.Tensor, + index_qk: cute.Tensor, + seqlen_q: Int32, + seqlen_k: Int32, + window_size_left: Optional[int] = None, + window_size_right: Optional[int] = None, + index_transform: cutlass.Constexpr = lambda index_q, index_k: ( + index_q, + index_k, + ), + ): + """ + Apply the appropriate mask to the attention scores. + + This method modifies the attention scores (acc_qk) based on the mask type + and the positions in the index tensor. + + :param mask_type: Type of mask to use + :type mask_type: utils.MaskEnum + :param acc_qk: Accumulated QK attention scores tensor. + :type acc_qk: cute.Tensor + :param index_qk: Index tensor containing position information. + :type index_qk: cute.Tensor + :param seqlen_k: Key sequence length for attention computation. + :type seqlen_k: Int32 + :param seqlen_q: Query sequence length for attention computation. + :type seqlen_q: Optional[int] + :param window_size_left: Left-side sliding window size for attention masking. + :type window_size_left: Optional[int] + :param window_size_right: Right-side sliding window size for attention masking. + :type window_size_right: Optional[int] + """ + + tidx, tidy, tidx = cute.arch.thread_idx() + offset = 0 + offset = seqlen_k - seqlen_q if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_INFERENCE or mask_type is MaskEnum.WINDOW_MASK_BWD_INFERENCE) else 0 + for i in cutlass.range_constexpr(cute.size(acc_qk)): + index_q, index_k = index_transform(*index_qk[i]) + if cutlass.const_expr(window_size_left is not None or window_size_right is not None): + if cutlass.const_expr(window_size_left is None): + if index_q + offset + window_size_right < index_k: + acc_qk[i] = -Float32.inf + if index_k >= seqlen_k or index_q >= seqlen_q: # residual mask + acc_qk[i] = -Float32.inf + elif cutlass.const_expr(window_size_right is None): + if index_q + offset - window_size_left > index_k: + acc_qk[i] = -Float32.inf + if index_k >= seqlen_k or index_q >= seqlen_q: # residual mask + acc_qk[i] = -Float32.inf + else: + max_K_index = min(index_q + offset + window_size_right, seqlen_k) + min_K_index = max(0, index_q + offset - window_size_left) + if index_k > max_K_index or index_k < min_K_index: + acc_qk[i] = -Float32.inf + if index_k >= seqlen_k or index_q >= seqlen_q: # residual mask + acc_qk[i] = -Float32.inf + + if cutlass.const_expr(mask_type == MaskEnum.RESIDUAL_MASK or mask_type == MaskEnum.RESIDUAL_MASK_BWD): + if index_k >= seqlen_k or index_q >= seqlen_q: + acc_qk[i] = -Float32.inf diff --git a/lightx2v/common/ops/attn/kernels/cutedsl_sparse_fmha_kernel.py b/lightx2v/common/ops/attn/kernels/cutedsl_sparse_fmha_kernel.py new file mode 100644 index 000000000..9a69381de --- /dev/null +++ b/lightx2v/common/ops/attn/kernels/cutedsl_sparse_fmha_kernel.py @@ -0,0 +1,3919 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +import argparse +import math +import os +import sys +import time +from typing import Optional, Tuple, Type, Union + +import cuda.bindings.driver as cuda +import cupy as cp +import cutlass +import cutlass.cute as cute +import cutlass.cute.nvgpu.tcgen05 as tcgen05 +import cutlass.cute.testing as testing +import cutlass.pipeline as pipeline +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +import numpy as np +from cutlass.cute.runtime import from_dlpack +from cutlass.cute.typing import Float32, Int32, Int64 + +if __name__ == "__main__": + current_dir = os.path.dirname(os.path.abspath(__file__)) + sys.path.insert(0, os.path.join(current_dir, "..")) + +from . import cutedsl_fmha_helpers as fmha_utils # isort: skip + +""" +A fused multi-head attention (FMHA) example for the NVIDIA Blackwell SM100 architecture using CUTE DSL + +This example demonstrates an implementation of fused multi-head attention using a TMA + Blackwell SM100 +TensorCore warp-specialized persistent kernel. The implementation integrates the Q*K^T matrix multiplication, +softmax normalization, and softmax(Q*K^T)*V into a single kernel, avoiding intermediate data movement between +global memory and shared memory, thus improving computational efficiency. + +The kernel implements key optimizations including: +- Warp specialization for different computation phases (load, MMA, softmax, correction, epilogue) +- Pipeline stages between different warps for overlapping computation and memory access +- Support for different precision data types +- Optional causal masking for autoregressive models + +To run this example: + +.. code-block:: bash + + python examples/blackwell/fmha.py \ + --qk_acc_dtype Float32 --pv_acc_dtype Float32 \ + --mma_tiler_mn 128,128 \ + --q_shape 4,1024,8,64 --k_shape 4,1024,8,64 \ + --is_persistent + +The above example runs FMHA with batch size 4, sequence length 1024, 8 attention heads, and head +dimension 64. The Blackwell tcgen05 MMA tile shape is (128, 128), and the kernel uses fp16 for input/output +with fp32 for accumulation. + +To collect performance with NCU profiler: + +.. code-block:: bash + + ncu python examples/blackwell/fmha.py \ + --qk_acc_dtype Float32 --pv_acc_dtype Float32 \ + --mma_tiler_mn 128,128 \ + --q_shape 4,1024,8,64 --k_shape 4,1024,8,64 \ + --is_persistent --warmup_iterations 10 \ + --iterations 10 --skip_ref_check + +Constraints for this example: +* Supported head dimensions: 32, 64, and 128 +* Number of heads in Q must be divisible by number of heads in K +* mma_tiler_mn must be 128,128 +* Batch size must be the same for Q, K, and V tensors +* For causal masking, use --is_causal (note: specify without =True/False) +* For persistent scheduling, use --is_persistent (note: specify without =True/False) +""" + + +def make_thread_cooperative_group(size: int): + return pipeline.CooperativeGroup(pipeline.Agent.Thread, size) + + +class BlackwellFusedMultiHeadAttentionForward: + WINDOW_NO_LIMIT = 1 << 30 + + def __init__( + self, + qk_acc_dtype: Type[cutlass.Numeric], + pv_acc_dtype: Type[cutlass.Numeric], + mma_tiler: Tuple[int, int, int], + is_persistent: bool, + mask_type: fmha_utils.MaskEnum, + is_causal: bool = False, + use_sliding_window: bool = False, + actual_head_dim: Optional[int] = None, + ): + """Initializes the configuration for a Blackwell Fused Multi-Head Attention (FMHA) kernel. + + This configuration includes several key aspects: + + 1. Data Type Settings: + - qk_acc_dtype: Data type for Q*K^T matrix multiplication accumulator + - pv_acc_dtype: Data type for P*V matrix multiplication accumulator + + 2. MMA Instruction Settings: + - mma_tiler: The (M, N, K) shape of the MMA instruction unit + - qk_mma_tiler: MMA shape for Q*K^T computation + - pv_mma_tiler: MMA shape for P*V computation + + 3. Kernel Execution Mode: + - is_persistent: Boolean indicating whether to use persistent kernel mode + - mask_type: Specifies the type of mask to use (no mask, residual mask, or causal mask) + - is_causal: Whether to apply causal masking (window_size_right = 0) + - use_sliding_window: Whether to compile with sliding window support + + :param qk_acc_dtype: Data type for Q*K^T matrix multiplication accumulator + :type qk_acc_dtype: Type[cutlass.Numeric] + :param pv_acc_dtype: Data type for P*V matrix multiplication accumulator + :type pv_acc_dtype: Type[cutlass.Numeric] + :param mma_tiler: The (M, N, K) shape of the MMA instruction + :type mma_tiler: Tuple[int, int, int] + :param is_persistent: Whether to use persistent kernel mode + :type is_persistent: bool + :param mask_type: Type of mask to use + :type mask_type: fmha_utils.MaskEnum + :param is_causal: Whether to apply causal masking. When True, window_size_right + is set to Int32(0) as a compile-time constant so tokens cannot attend + to future positions. When False, window_size_right is None (bidirectional). + :type is_causal: bool + :param use_sliding_window: If True, compile with sliding window masking code. + If False, window_size_left is treated as None at compile time, + eliminating left-side window masking code for better performance. + :type use_sliding_window: bool + """ + + self.qk_acc_dtype = qk_acc_dtype + self.pv_acc_dtype = pv_acc_dtype + # head_dim = actual tensor dimension (e.g. 72). + # MMA/SMEM/CTA tilers below use mma_tiler[2] (padded, e.g. 80). + # TMA ZFILL bridges the gap on loads; OOB drop on stores. + self.head_dim = actual_head_dim if actual_head_dim is not None else mma_tiler[2] + self.inv_sqrt_head_dim = 1.0 / math.sqrt(self.head_dim) + self.log2_e = math.log2(math.e) + self.cta_tiler = ( + 2 * mma_tiler[0], # 2 Q tile per CTA + mma_tiler[1], + mma_tiler[2], + ) + self.qk_mma_tiler = mma_tiler + self.pv_mma_tiler = ( + mma_tiler[0], + mma_tiler[2], + mma_tiler[1], + ) + self.cluster_shape_mn = (1, 1) + self.is_persistent = is_persistent + self.mask_type = mask_type + self.is_causal = is_causal + self.use_sliding_window = use_sliding_window + + self.softmax0_warp_ids = (0, 1, 2, 3) + self.softmax1_warp_ids = (4, 5, 6, 7) + self.correction_warp_ids = (8, 9, 10, 11) + self.mma_warp_id = 12 + self.load_warp_id = 13 + self.epilogue_warp_id = 14 + self.empty_warp_id = 15 + self.tmem_alloc_cols = cute.arch.get_max_tmem_alloc_cols("sm_100") + + self.threads_per_warp = 32 + self.threads_per_cta = self.threads_per_warp * len( + ( + *self.softmax0_warp_ids, + *self.softmax1_warp_ids, + *self.correction_warp_ids, + self.mma_warp_id, + self.load_warp_id, + self.epilogue_warp_id, + self.empty_warp_id, + ) + ) + + self.cta_sync_barrier = pipeline.NamedBarrier( + barrier_id=1, + num_threads=self.threads_per_cta, + ) + self.tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=2, + num_threads=self.threads_per_warp, + ) + + self.tmem_s0_offset = 0 + self.tmem_s1_offset = 128 + self.tmem_o0_offset = 256 + self.tmem_o1_offset = 384 + self.tmem_p0_offset = 32 + self.tmem_p1_offset = 160 + + # vec buffer for row_max & row_sum + self.tmem_vec0_offset = 0 + self.tmem_vec1_offset = 128 + + self.num_regs_softmax = int(os.environ.get("SPARSE_FMHA_REGS_SOFTMAX", "192")) + self.num_regs_correction = int(os.environ.get("SPARSE_FMHA_REGS_CORRECTION", "96")) + self.num_regs_other = int(os.environ.get("SPARSE_FMHA_REGS_OTHER", "32")) + self.min_blocks_per_mp = int(os.environ.get("SPARSE_FMHA_MIN_BLOCKS", "1")) + self.fixed_single_batch = os.environ.get("SPARSE_FMHA_FIXED_B1", "0") == "1" + # Delayed/zero-reference max modes are approximate experiments. Keep the + # default path equivalent to the online-softmax reference semantics. + self.max_update_interval = int(os.environ.get("SPARSE_FMHA_MAX_UPDATE_INTERVAL", "1")) + if self.max_update_interval not in (1, 2, 4, 8, 16, 32, 64): + raise ValueError("SPARSE_FMHA_MAX_UPDATE_INTERVAL must be a power of two from 1 to 64") + self.zero_reference_max = os.environ.get("SPARSE_FMHA_ZERO_REFERENCE_MAX", "0") == "1" + + self.buffer_align_bytes = 1024 + + num_warps_per_warpgroup = 4 + self.softmax_warpgroup_count = len((*self.softmax0_warp_ids, *self.softmax1_warp_ids)) // num_warps_per_warpgroup + + def _setup_attributes(self): + """Set up configurations and parameters for the FMHA kernel operation. + + This method initializes and configures various attributes required for the + execution of the fused multi-head attention kernel, mainly about the pipeline stages: + + - Sets up staging parameters for Q, K, V inputs and accumulator data + - Configures pipeline stages for softmax, correction, and epilogue operations + """ + + self.q_stage = int(os.environ.get("SPARSE_FMHA_Q_STAGE", "2")) + # SM110 tuning: BF16 K/V staging is faster at two stages while retaining + # the original four-stage default for the separately tuned FP8 path. + default_kv_stage = 4 if self.q_dtype.width == 8 else 2 + self.kv_stage = int(os.environ.get("SPARSE_FMHA_KV_STAGE", str(default_kv_stage))) + self.acc_stage = 1 + self.softmax_corr_stage = 1 + self.mma_corr_stage = int(os.environ.get("SPARSE_FMHA_MMA_CORR_STAGE", "2")) + self.mma_softmax_stage = 1 + self.epi_stage = int(os.environ.get("SPARSE_FMHA_EPI_STAGE", "2")) + + # Pre-scale softmax output by 2^FP8_E4M3_PRESCALE_LOG2=256 for FP8 to maximize + # E4M3 dynamic range. + # P ∈ [0,1] → P*256 ∈ [0,256], utilizing more of the [0,448] FP8 range. + # Fused into exp2: exp2(x + 8) = exp2(x) * 256, zero extra per-element ops. + FP8_E4M3_PRESCALE_LOG2 = 8.0 # log2(256) — tuned for FP8 E4M3 range [0, 448] + self.softmax_prescale_log2 = FP8_E4M3_PRESCALE_LOG2 if self.q_dtype.width == 8 else 0.0 + self.softmax_prescale_ln = self.softmax_prescale_log2 * 0.6931471805599453 + + @cute.jit + def __call__( + self, + q_tensor: cute.Tensor, # (B, S_q, H_q, D) — B,S_q,H_q dynamic; D static + kv_cache: cute.Tensor, # (B, 2, H_kv, S_k, D) — B,H_kv,S_k dynamic; D static + o_tensor: cute.Tensor, # (B, S_q, H_q, D) — same layout as Q + cum_seqlen_k: cute.Tensor, # (B+1,) Int32 — cumulative KV sequence lengths + window_size_left: Int32, + scale_q: Float32, + scale_k: Float32, + scale_v: Float32, + inv_scale_o: Float32, + stream: cuda.CUstream, + ): + """Execute the Fused Multi-Head Attention operation on the provided tensors. + + This method prepares the input tensors for processing, validates their shapes and types, + configures the computation parameters, and launches the CUDA kernel. + + The method handles: + 1. Tensor layout transformations for specific memory access patterns + 2. Validation of tensor shapes and data types + 3. Initialization of hardware-specific parameters and memory layouts + 4. Configuration of TMA (Tensor Memory Access) operations + 5. Grid and work scheduling computation + 6. Kernel launch with appropriate parameters + + The softmax scale is computed as: + scale_softmax = scale_q * scale_k * (1 / sqrt(head_dim)) + For FP16 (no quantization), pass scale_q = scale_k = scale_v = inv_scale_o = 1.0. + For FP8, pass the dequant scales so the kernel folds them into softmax/output scaling. + + :param q_tensor: The query tensor (B, S_q, H_q, D) with dynamic B, S_q, H_q + :type q_tensor: cute.Tensor + :param kv_cache: The KV cache tensor (B, 2, H_kv, S_k, D) with dynamic B, H_kv, S_k + :type kv_cache: cute.Tensor + :param o_tensor: The output tensor (B, S_q, H_q, D) + :type o_tensor: cute.Tensor + :param window_size_left: Left-side sliding window size for attention masking. + :type window_size_left: Int32 + :param scale_q: Dequantization scale for Q (quant→orig). 1.0 for FP16. + :type scale_q: Float32 + :param scale_k: Dequantization scale for K (quant→orig). 1.0 for FP16. + :type scale_k: Float32 + :param scale_v: Dequantization scale for V (quant→orig). 1.0 for FP16. + :type scale_v: Float32 + :param inv_scale_o: Inverse output quantization scale. 1.0 for FP16 output. + :type inv_scale_o: Float32 + :param stream: The CUDA stream to execute the kernel on + :type stream: cuda.CUstream + :raises TypeError: If tensor data types don't match or aren't supported + :raises RuntimeError: If tensor layouts aren't in supported formats + """ + scale_softmax = scale_q * scale_k * self.inv_sqrt_head_dim + scale_softmax_log2 = scale_softmax * self.log2_e + scale_output = scale_v * inv_scale_o + b = q_tensor.layout.shape[0] + s_q = q_tensor.layout.shape[1] + h_q = q_tensor.layout.shape[2] + h_k = kv_cache.layout.shape[2] + cap = kv_cache.layout.shape[3] + s_lse = s_q + d = self.head_dim + + q_iter = q_tensor.iterator + # KV cache shape[3] = cap (physical capacity, for stride computation). + # Per-batch actual KV lengths come from cum_seqlen_k at runtime. + kv_base = kv_cache.iterator + stride_kv_head = cap * d + stride_kv_select = h_k * stride_kv_head + k_iter = kv_base + v_iter = kv_base + stride_kv_select + o_iter = o_tensor.iterator + + cum_seqlen_q = None + lse_iter = None + h_r = h_q // h_k + # Always use batch-strided layout (not packed varlen). + # cum_seqlen_k is only used inside the kernel for per-batch seqlen_k override. + qo_offset = 0 + kv_offset = 0 + b_qo = b + b_kv = b + stride_b_qo = h_r * h_k * s_q * d + stride_b_kv = 2 * stride_kv_select + b_lse = b + stride_b_lse = h_r * h_k * s_lse + + # (s, d, ((h_r, h_k), b)) + q_layout = cute.make_layout( + (s_q, d, ((h_r, h_k), b_qo)), + stride=(d * h_r * h_k, 1, ((d, d * h_r), stride_b_qo)), + ) + q = cute.make_tensor(q_iter + qo_offset, q_layout) + # (s, d, ((h_r, h_k), b)), 0-stride for h_r to broadcast + k_layout = cute.make_layout( + (cap, d, ((h_r, h_k), b_kv)), + stride=(d, 1, ((0, stride_kv_head), stride_b_kv)), + ) + k = cute.make_tensor(k_iter + kv_offset, k_layout) + # (d, s, ((h_r, h_k), b)), 0-stride for h_r to broadcast + v_layout = cute.make_layout( + (d, cap, ((h_r, h_k), b_kv)), + stride=(1, d, ((0, stride_kv_head), stride_b_kv)), + ) + v = cute.make_tensor(v_iter + kv_offset, v_layout) + # (s, d, ((h_r, h_k), b)) + o_layout = cute.make_layout( + (s_q, d, ((h_r, h_k), b_qo)), + stride=(d * h_r * h_k, 1, ((d, d * h_r), stride_b_qo)), + ) + o = cute.make_tensor(o_iter + qo_offset, o_layout) + if cutlass.const_expr(lse_iter is not None): + # (s, ((h_r, h_k), b)) + lse_layout = cute.make_layout( + (s_lse, ((h_r, h_k), b_lse)), + stride=(1, ((s_lse, h_r * s_lse), stride_b_lse)), + ) + lse = cute.make_tensor(lse_iter, lse_layout) + else: + lse = None + + # setup static attributes before smem/grid/tma computation + self.q_dtype = q.element_type + self.k_dtype = k.element_type + self.v_dtype = v.element_type + self.o_dtype = o.element_type + + self.tile_sched_params, grid = fmha_utils.compute_grid( + cute.shape((s_q, d, ((h_r, h_k), b))), + self.cta_tiler, + self.is_persistent, + ) + + self.q_major_mode = utils.LayoutEnum.from_tensor(q).mma_major_mode() + self.k_major_mode = utils.LayoutEnum.from_tensor(k).mma_major_mode() + self.v_major_mode = utils.LayoutEnum.from_tensor(v).mma_major_mode() + self.o_layout = utils.LayoutEnum.from_tensor(o) + + if cutlass.const_expr(self.q_major_mode != tcgen05.OperandMajorMode.K): + raise RuntimeError("The layout of q is not supported") + if cutlass.const_expr(self.k_major_mode != tcgen05.OperandMajorMode.K): + raise RuntimeError("The layout of k is not supported") + if cutlass.const_expr(self.v_major_mode != tcgen05.OperandMajorMode.MN): + raise RuntimeError("The layout of v is not supported") + + # check type consistency + if cutlass.const_expr(self.q_dtype != self.k_dtype): + raise TypeError(f"Type mismatch: {self.q_dtype} != {self.k_dtype}") + if cutlass.const_expr(self.q_dtype != self.v_dtype): + raise TypeError(f"Type mismatch: {self.q_dtype} != {self.v_dtype}") + self._setup_attributes() + + cta_group = tcgen05.CtaGroup.ONE + # the intermediate tensor p is from tmem & k-major + p_source = tcgen05.OperandSource.TMEM + p_major_mode = tcgen05.OperandMajorMode.K + qk_tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.q_dtype, + self.q_major_mode, + self.k_major_mode, + self.qk_acc_dtype, + cta_group, + self.qk_mma_tiler[:2], + ) + pv_tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.v_dtype, + p_major_mode, + self.v_major_mode, + self.pv_acc_dtype, + cta_group, + self.pv_mma_tiler[:2], + p_source, + ) + + self.cluster_shape_mnk = (*self.cluster_shape_mn, 1) + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout(self.cluster_shape_mnk), + (qk_tiled_mma.thr_id.shape,), + ) + + self.epi_tile = self.pv_mma_tiler[:2] + + q_smem_layout_staged = sm100_utils.make_smem_layout_a( + qk_tiled_mma, + self.qk_mma_tiler, + self.q_dtype, + self.q_stage, + ) + k_smem_layout_staged = sm100_utils.make_smem_layout_b( + qk_tiled_mma, + self.qk_mma_tiler, + self.k_dtype, + self.kv_stage, + ) + p_tmem_layout_staged = sm100_utils.make_smem_layout_a( + pv_tiled_mma, + self.pv_mma_tiler, + self.q_dtype, + self.acc_stage, + ) + v_smem_layout_staged = sm100_utils.make_smem_layout_b( + pv_tiled_mma, + self.pv_mma_tiler, + self.v_dtype, + self.kv_stage, + ) + o_smem_layout_staged = sm100_utils.make_smem_layout_epi( + self.o_dtype, + self.o_layout, + self.epi_tile, + self.epi_stage, + ) + + # TMA load for Q + tma_load_op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp(cta_group) + tma_store_op = cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp() + + q_smem_layout = cute.select(q_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_q, tma_tensor_q = cute.nvgpu.make_tiled_tma_atom_A( + tma_load_op, + q, + q_smem_layout, + self.qk_mma_tiler, + qk_tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + # TMA load for K + k_smem_layout = cute.select(k_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_k, tma_tensor_k = cute.nvgpu.make_tiled_tma_atom_B( + tma_load_op, + k, + k_smem_layout, + self.qk_mma_tiler, + qk_tiled_mma, + self.cluster_layout_vmnk.shape, + ) + # TMA load for V + v_smem_layout = cute.select(v_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_v, tma_tensor_v = cute.nvgpu.make_tiled_tma_atom_B( + tma_load_op, + v, + v_smem_layout, + self.pv_mma_tiler, + pv_tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + o_smem_layout = cute.select(o_smem_layout_staged, mode=[0, 1]) + + tma_atom_o, tma_tensor_o = cute.nvgpu.cpasync.make_tiled_tma_atom( + tma_store_op, + o, + o_smem_layout, + self.epi_tile, + ) + + q_copy_size = cute.size_in_bytes(self.q_dtype, q_smem_layout) + k_copy_size = cute.size_in_bytes(self.k_dtype, k_smem_layout) + self.tma_copy_q_bytes = q_copy_size + self.tma_copy_kv_bytes = k_copy_size + + @cute.struct + class SharedStorage: + # Pipeline barriers + load_q_mbar_ptr: cute.struct.MemRange[Int64, self.q_stage * 2] + load_kv_mbar_ptr: cute.struct.MemRange[Int64, self.kv_stage * 2] + mma_s0_mbar_ptr: cute.struct.MemRange[Int64, self.mma_softmax_stage * 2] + mma_s1_mbar_ptr: cute.struct.MemRange[Int64, self.mma_softmax_stage * 2] + s0_corr_mbar_ptr: cute.struct.MemRange[Int64, self.softmax_corr_stage * 2] + s1_corr_mbar_ptr: cute.struct.MemRange[Int64, self.softmax_corr_stage * 2] + s0_s1_sequence_mbar_ptr: cute.struct.MemRange[Int64, self.softmax_warpgroup_count] + corr_epi_mbar_ptr: cute.struct.MemRange[Int64, self.epi_stage * 2] + mma_corr_mbar_ptr: cute.struct.MemRange[Int64, self.mma_corr_stage * 2] + tmem_dealloc_mbar_ptr: cute.struct.MemRange[Int64, 1] + # Tmem holding buffer + tmem_holding_buf: Int32 + # Smem tensors + sO: cute.struct.Align[ + cute.struct.MemRange[self.o_dtype, cute.cosize(o_smem_layout_staged)], + self.buffer_align_bytes, + ] + sQ: cute.struct.Align[ + cute.struct.MemRange[self.q_dtype, cute.cosize(q_smem_layout_staged)], + self.buffer_align_bytes, + ] + sK: cute.struct.Align[ + cute.struct.MemRange[self.k_dtype, cute.cosize(k_smem_layout_staged)], + self.buffer_align_bytes, + ] + + self.shared_storage = SharedStorage + + # Compile-time dispatch: when use_sliding_window is False, pass None for + # window_size_left to eliminate left-side window masking code in the kernel. + if cutlass.const_expr(self.use_sliding_window): + _wsl = window_size_left + else: + _wsl = None + + _wsr = Int32(0) if cutlass.const_expr(self.is_causal) else None + + # Launch the kernel synchronously + self.kernel( + qk_tiled_mma, + pv_tiled_mma, + tma_atom_q, + tma_tensor_q, + tma_atom_k, + tma_tensor_k, + tma_atom_v, + tma_tensor_v, + tma_atom_o, + tma_tensor_o, + o, + cum_seqlen_q, + cum_seqlen_k, + lse, + scale_softmax_log2, + scale_softmax, + scale_output, + _wsl, + _wsr, + None, + None, + q_smem_layout_staged, + k_smem_layout_staged, + p_tmem_layout_staged, + v_smem_layout_staged, + o_smem_layout_staged, + self.tile_sched_params, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=self.cluster_shape_mnk, + stream=stream, + min_blocks_per_mp=1, + ) + + @cute.jit + def __call_vit__( + self, + q_tensor: cute.Tensor, + k_tensor: cute.Tensor, + v_tensor: cute.Tensor, + o_tensor: cute.Tensor, + cu_seqlens: cute.Tensor, + max_seqlen: Int32, + scale_softmax_log2: Float32, + scale_softmax: Float32, + scale_output: Float32, + stream: cuda.CUstream, + ): + """Compile-time dense entry point for standalone packed ViT attention.""" + self._call_vit_impl( + q_tensor, + k_tensor, + v_tensor, + o_tensor, + cu_seqlens, + max_seqlen, + scale_softmax_log2, + scale_softmax, + scale_output, + None, + None, + stream, + ) + + @cute.jit + def __call_vit_sparse__( + self, + q_tensor: cute.Tensor, + k_tensor: cute.Tensor, + v_tensor: cute.Tensor, + o_tensor: cute.Tensor, + cu_seqlens: cute.Tensor, + max_seqlen: Int32, + scale_softmax_log2: Float32, + scale_softmax: Float32, + scale_output: Float32, + sparse_block_count: cute.Tensor, + sparse_block_indices: cute.Tensor, + stream: cuda.CUstream, + ): + """Block-sparse packed ViT attention. + + Metadata uses one non-empty KV block list per 256-row Q tile: + ``count[B, Hq, ceil(max_seqlen / 256)]`` and + ``indices[B, Hq, ceil(max_seqlen / 256), max_blocks]``. Each index + addresses a 128-token, sequence-local KV tile. + """ + self._call_vit_impl( + q_tensor, + k_tensor, + v_tensor, + o_tensor, + cu_seqlens, + max_seqlen, + scale_softmax_log2, + scale_softmax, + scale_output, + sparse_block_count, + sparse_block_indices, + stream, + ) + + @cute.jit + def _call_vit_impl( + self, + q_tensor: cute.Tensor, # (total_S, H_q, D) — packed varlen + k_tensor: cute.Tensor, # (total_S, H_kv, D) — packed varlen + v_tensor: cute.Tensor, # (total_S, H_kv, D) — packed varlen + o_tensor: cute.Tensor, # (total_S, H_q, D) + cu_seqlens: cute.Tensor, # (B+1,) Int32 + max_seqlen: Int32, + scale_softmax_log2: Float32, + scale_softmax: Float32, + scale_output: Float32, + sparse_block_count: Optional[cute.Tensor], + sparse_block_indices: Optional[cute.Tensor], + stream: cuda.CUstream, + ): + """ViT FMHA: packed varlen, separate Q/K/V, bidirectional (no causal mask). + + All sequences are packed into flat [total_S, H, D] tensors with boundaries + defined by cu_seqlens. Each sequence attends to all tokens in that sequence + (PADDING / RESIDUAL_MASK — no causal ordering). + + max_seqlen is the longest individual sequence length (not total_S). + It controls the grid size and per-sequence tile math. TMA descriptors + stay bounded by the compact total_S tensor extent. + """ + total_s = q_tensor.layout.shape[0] + s_q = max_seqlen + h_q = q_tensor.layout.shape[1] + h_k = k_tensor.layout.shape[1] + d = self.head_dim + + b = cu_seqlens.layout.shape[0] - 1 + h_r = h_q // h_k + + q_iter = q_tensor.iterator + k_iter = k_tensor.iterator + v_iter = v_tensor.iterator + o_iter = o_tensor.iterator + + cum_seqlen_q = None if self.fixed_single_batch else cu_seqlens + cum_seqlen_k = None if self.fixed_single_batch else cu_seqlens + + # The generic packed-varlen path drops the batch coordinate after applying + # cu_seqlens. The B=1 specialization keeps that coordinate in the tensor + # descriptor so the kernel can use its normal non-varlen coordinate path. + if cutlass.const_expr(self.fixed_single_batch): + q_layout = cute.make_layout( + (total_s, d, ((h_r, h_k), 1)), + stride=( + d * h_r * h_k, + 1, + ((d, d * h_r), total_s * d * h_r * h_k), + ), + ) + k_layout = cute.make_layout( + (total_s, d, ((h_r, h_k), 1)), + stride=(d * h_k, 1, ((0, d), total_s * d * h_k)), + ) + v_layout = cute.make_layout( + (d, total_s, ((h_r, h_k), 1)), + stride=(1, d * h_k, ((0, d), total_s * d * h_k)), + ) + o_layout = q_layout + else: + # Descriptor OOB is based on total_s, not on the synthetic + # max_seqlen * batch envelope. + q_layout = cute.make_layout( + (total_s, d, (h_r, h_k)), + stride=(d * h_r * h_k, 1, (d, d * h_r)), + ) + k_layout = cute.make_layout( + (total_s, d, (h_r, h_k)), + stride=(d * h_k, 1, (0, d)), + ) + v_layout = cute.make_layout( + (d, total_s, (h_r, h_k)), + stride=(1, d * h_k, (0, d)), + ) + o_layout = q_layout + q = cute.make_tensor(q_iter, q_layout) + k = cute.make_tensor(k_iter, k_layout) + v = cute.make_tensor(v_iter, v_layout) + o = cute.make_tensor(o_iter, o_layout) + lse = None + + self.q_dtype = q.element_type + self.k_dtype = k.element_type + self.v_dtype = v.element_type + self.o_dtype = o.element_type + + self.tile_sched_params, grid = fmha_utils.compute_grid( + cute.shape((s_q, d, ((h_r, h_k), b))), + self.cta_tiler, + self.is_persistent, + ) + + self.q_major_mode = utils.LayoutEnum.from_tensor(q).mma_major_mode() + self.k_major_mode = utils.LayoutEnum.from_tensor(k).mma_major_mode() + self.v_major_mode = utils.LayoutEnum.from_tensor(v).mma_major_mode() + self.o_layout = utils.LayoutEnum.from_tensor(o) + + if cutlass.const_expr(self.q_major_mode != tcgen05.OperandMajorMode.K): + raise RuntimeError("The layout of q is not supported") + if cutlass.const_expr(self.k_major_mode != tcgen05.OperandMajorMode.K): + raise RuntimeError("The layout of k is not supported") + if cutlass.const_expr(self.v_major_mode != tcgen05.OperandMajorMode.MN): + raise RuntimeError("The layout of v is not supported") + + if cutlass.const_expr(self.q_dtype != self.k_dtype): + raise TypeError(f"Type mismatch: {self.q_dtype} != {self.k_dtype}") + if cutlass.const_expr(self.q_dtype != self.v_dtype): + raise TypeError(f"Type mismatch: {self.q_dtype} != {self.v_dtype}") + self._setup_attributes() + + cta_group = tcgen05.CtaGroup.ONE + p_source = tcgen05.OperandSource.TMEM + p_major_mode = tcgen05.OperandMajorMode.K + qk_tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.q_dtype, + self.q_major_mode, + self.k_major_mode, + self.qk_acc_dtype, + cta_group, + self.qk_mma_tiler[:2], + ) + pv_tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.v_dtype, + p_major_mode, + self.v_major_mode, + self.pv_acc_dtype, + cta_group, + self.pv_mma_tiler[:2], + p_source, + ) + + self.cluster_shape_mnk = (*self.cluster_shape_mn, 1) + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout(self.cluster_shape_mnk), + (qk_tiled_mma.thr_id.shape,), + ) + self.epi_tile = self.pv_mma_tiler[:2] + + q_smem_layout_staged = sm100_utils.make_smem_layout_a(qk_tiled_mma, self.qk_mma_tiler, self.q_dtype, self.q_stage) + k_smem_layout_staged = sm100_utils.make_smem_layout_b(qk_tiled_mma, self.qk_mma_tiler, self.k_dtype, self.kv_stage) + p_tmem_layout_staged = sm100_utils.make_smem_layout_a(pv_tiled_mma, self.pv_mma_tiler, self.q_dtype, self.acc_stage) + v_smem_layout_staged = sm100_utils.make_smem_layout_b(pv_tiled_mma, self.pv_mma_tiler, self.v_dtype, self.kv_stage) + o_smem_layout_staged = sm100_utils.make_smem_layout_epi(self.o_dtype, self.o_layout, self.epi_tile, self.epi_stage) + + tma_load_op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp(cta_group) + tma_store_op = cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp() + + q_smem_layout = cute.select(q_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_q, tma_tensor_q = cute.nvgpu.make_tiled_tma_atom_A(tma_load_op, q, q_smem_layout, self.qk_mma_tiler, qk_tiled_mma, self.cluster_layout_vmnk.shape) + + k_smem_layout = cute.select(k_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_k, tma_tensor_k = cute.nvgpu.make_tiled_tma_atom_B(tma_load_op, k, k_smem_layout, self.qk_mma_tiler, qk_tiled_mma, self.cluster_layout_vmnk.shape) + + v_smem_layout = cute.select(v_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_v, tma_tensor_v = cute.nvgpu.make_tiled_tma_atom_B(tma_load_op, v, v_smem_layout, self.pv_mma_tiler, pv_tiled_mma, self.cluster_layout_vmnk.shape) + + o_smem_layout = cute.select(o_smem_layout_staged, mode=[0, 1]) + tma_atom_o, tma_tensor_o = cute.nvgpu.cpasync.make_tiled_tma_atom(tma_store_op, o, o_smem_layout, self.epi_tile) + + q_copy_size = cute.size_in_bytes(self.q_dtype, q_smem_layout) + k_copy_size = cute.size_in_bytes(self.k_dtype, k_smem_layout) + self.tma_copy_q_bytes = q_copy_size + self.tma_copy_kv_bytes = k_copy_size + + @cute.struct + class SharedStorage: + load_q_mbar_ptr: cute.struct.MemRange[Int64, self.q_stage * 2] + load_kv_mbar_ptr: cute.struct.MemRange[Int64, self.kv_stage * 2] + mma_s0_mbar_ptr: cute.struct.MemRange[Int64, self.mma_softmax_stage * 2] + mma_s1_mbar_ptr: cute.struct.MemRange[Int64, self.mma_softmax_stage * 2] + s0_corr_mbar_ptr: cute.struct.MemRange[Int64, self.softmax_corr_stage * 2] + s1_corr_mbar_ptr: cute.struct.MemRange[Int64, self.softmax_corr_stage * 2] + s0_s1_sequence_mbar_ptr: cute.struct.MemRange[Int64, self.softmax_warpgroup_count] + corr_epi_mbar_ptr: cute.struct.MemRange[Int64, self.epi_stage * 2] + mma_corr_mbar_ptr: cute.struct.MemRange[Int64, self.mma_corr_stage * 2] + tmem_dealloc_mbar_ptr: cute.struct.MemRange[Int64, 1] + tmem_holding_buf: Int32 + sO: cute.struct.Align[cute.struct.MemRange[self.o_dtype, cute.cosize(o_smem_layout_staged)], self.buffer_align_bytes] + sQ: cute.struct.Align[cute.struct.MemRange[self.q_dtype, cute.cosize(q_smem_layout_staged)], self.buffer_align_bytes] + sK: cute.struct.Align[cute.struct.MemRange[self.k_dtype, cute.cosize(k_smem_layout_staged)], self.buffer_align_bytes] + + self.shared_storage = SharedStorage + + # ViT: bidirectional — no sliding window, no causal masking. + # Both window sizes are None (is_causal=False) so apply_mask only + # applies RESIDUAL_MASK for out-of-bounds elements on partial tiles. + _wsl = None + _wsr = Int32(0) if cutlass.const_expr(self.is_causal) else None + + self.kernel( + qk_tiled_mma, + pv_tiled_mma, + tma_atom_q, + tma_tensor_q, + tma_atom_k, + tma_tensor_k, + tma_atom_v, + tma_tensor_v, + tma_atom_o, + tma_tensor_o, + o, + cum_seqlen_q, + cum_seqlen_k, + lse, + scale_softmax_log2, + scale_softmax, + scale_output, + _wsl, + _wsr, + sparse_block_count, + sparse_block_indices, + q_smem_layout_staged, + k_smem_layout_staged, + p_tmem_layout_staged, + v_smem_layout_staged, + o_smem_layout_staged, + self.tile_sched_params, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=self.cluster_shape_mnk, + stream=stream, + min_blocks_per_mp=self.min_blocks_per_mp, + ) + + # GPU device kernel + @cute.kernel + def kernel( + self, + qk_tiled_mma: cute.TiledMma, + pv_tiled_mma: cute.TiledMma, + tma_atom_q: cute.CopyAtom, + mQ_qdl: cute.Tensor, + tma_atom_k: cute.CopyAtom, + mK_kdl: cute.Tensor, + tma_atom_v: cute.CopyAtom, + mV_dkl: cute.Tensor, + tma_atom_o: cute.CopyAtom, + mO_qdl: cute.Tensor, + mO_gmem: cute.Tensor, + cum_seqlen_q: Optional[cute.Tensor], + cum_seqlen_k: Optional[cute.Tensor], + mLSE: Optional[cute.Tensor], + scale_softmax_log2: Float32, + scale_softmax: Float32, + scale_output: Float32, + window_size_left: Optional[Int32], + window_size_right: Optional[Int32], + sparse_block_count: Optional[cute.Tensor], + sparse_block_indices: Optional[cute.Tensor], + q_smem_layout_staged: cute.ComposedLayout, + k_smem_layout_staged: cute.ComposedLayout, + p_tmem_layout_staged: cute.ComposedLayout, + v_smem_layout_staged: cute.ComposedLayout, + o_smem_layout_staged: cute.ComposedLayout, + tile_sched_params: fmha_utils.FmhaStaticTileSchedulerParams, + ): + """The device kernel implementation of the Fused Multi-Head Attention. + + This kernel coordinates multiple specialized warps to perform different phases of the FMHA computation: + 1. Load warp: Loads Q, K, V data from global memory to shared memory using TMA + 2. MMA warp: Performs matrix multiplications (Q*K^T and P*V) + 3. Softmax warps: Compute softmax normalization on attention scores + 4. Correction warps: Apply adjustments to intermediate results + 5. Epilogue warp: Handles final output transformation and storage + + The kernel implements a complex pipeline with overlapping computation and memory operations, + using tensor memory access (TMA) for efficient data loading, warp specialization for different + computation phases, and optional attention masking. + + :param qk_tiled_mma: Tiled MMA for Q*K^T + :type qk_tiled_mma: cute.TiledMma + :param pv_tiled_mma: Tiled MMA for P*V + :type pv_tiled_mma: cute.TiledMma + :param tma_atom_q: TMA copy atom for query tensor + :type tma_atom_q: cute.CopyAtom + :param mQ_qdl: Partitioned query tensor + :type mQ_qdl: cute.Tensor + :param tma_atom_k: TMA copy atom for key tensor + :type tma_atom_k: cute.CopyAtom + :param mK_kdl: Partitioned key tensor + :type mK_kdl: cute.Tensor + :param tma_atom_v: TMA copy atom for value tensor + :type tma_atom_v: cute.CopyAtom + :param mV_dkl: Partitioned value tensor + :type mV_dkl: cute.Tensor + :param tma_atom_o: TMA copy atom for output tensor + :type tma_atom_o: cute.CopyAtom + :param mO_qdl: Partitioned output tensor + :type mO_qdl: cute.Tensor + :param scale_softmax_log2: The log2 scale factor for softmax + :type scale_softmax_log2: Float32 + :param scale_softmax: The scale factor for softmax + :type scale_softmax: Float32 + :param scale_output: The scale factor for the output + :type scale_output: Float32 + :param window_size_left: Left-side sliding window size for attention masking. + :type window_size_left: Optional[Int32] + :param window_size_right: Right-side sliding window size for attention masking. + :type window_size_right: Optional[Int32] + :param q_smem_layout_staged: Shared memory layout for query tensor + :type q_smem_layout_staged: cute.ComposedLayout + :param k_smem_layout_staged: Shared memory layout for key tensor + :type k_smem_layout_staged: cute.ComposedLayout + :param p_tmem_layout_staged: Tensor memory layout for probability matrix + :type p_tmem_layout_staged: cute.ComposedLayout + :param v_smem_layout_staged: Shared memory layout for value tensor + :type v_smem_layout_staged: cute.ComposedLayout + :param o_smem_layout_staged: Shared memory layout for output tensor + :type o_smem_layout_staged: cute.ComposedLayout + :param tile_sched_params: Scheduling parameters for work distribution + :type tile_sched_params: fmha_utils.FmhaStaticTileSchedulerParams + """ + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + # coord inside cta + tidx, _, _ = cute.arch.thread_idx() + + # + # Prefetch tma desc + # + if warp_idx == self.load_warp_id: + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_q) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_k) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_v) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_o) + + # Alloc + smem = utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + + load_q_producer, load_q_consumer = pipeline.PipelineTmaUmma.create( + num_stages=self.q_stage, + producer_group=make_thread_cooperative_group(len([self.load_warp_id])), + consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + tx_count=self.tma_copy_q_bytes, + barrier_storage=storage.load_q_mbar_ptr.data_ptr(), + ).make_participants() + load_kv_producer, load_kv_consumer = pipeline.PipelineTmaUmma.create( + num_stages=self.kv_stage, + producer_group=make_thread_cooperative_group(len([self.load_warp_id])), + consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + tx_count=self.tma_copy_kv_bytes, + barrier_storage=storage.load_kv_mbar_ptr.data_ptr(), + ).make_participants() + mma_s0_producer, mma_s0_consumer = pipeline.PipelineUmmaAsync.create( + num_stages=self.mma_softmax_stage, + producer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + consumer_group=make_thread_cooperative_group(self.threads_per_warp * len(self.softmax0_warp_ids)), + barrier_storage=storage.mma_s0_mbar_ptr.data_ptr(), + ).make_participants() + mma_s1_producer, mma_s1_consumer = pipeline.PipelineUmmaAsync.create( + num_stages=self.mma_softmax_stage, + producer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + consumer_group=make_thread_cooperative_group(self.threads_per_warp * len(self.softmax1_warp_ids)), + barrier_storage=storage.mma_s1_mbar_ptr.data_ptr(), + ).make_participants() + s0_corr_producer, s0_corr_consumer = pipeline.PipelineAsync.create( + num_stages=self.softmax_corr_stage, + producer_group=make_thread_cooperative_group(self.threads_per_warp * len(self.softmax0_warp_ids)), + consumer_group=make_thread_cooperative_group(self.threads_per_warp * len(self.correction_warp_ids)), + barrier_storage=storage.s0_corr_mbar_ptr.data_ptr(), + ).make_participants() + s1_corr_producer, s1_corr_consumer = pipeline.PipelineAsync.create( + num_stages=self.softmax_corr_stage, + producer_group=make_thread_cooperative_group(self.threads_per_warp * len(self.softmax1_warp_ids)), + consumer_group=make_thread_cooperative_group(self.threads_per_warp * len(self.correction_warp_ids)), + barrier_storage=storage.s1_corr_mbar_ptr.data_ptr(), + ).make_participants() + corr_epi_producer, corr_epi_consumer = pipeline.PipelineAsync.create( + num_stages=self.epi_stage, + producer_group=make_thread_cooperative_group(self.threads_per_warp * len(self.correction_warp_ids)), + consumer_group=make_thread_cooperative_group(self.threads_per_warp * len([self.epilogue_warp_id])), + barrier_storage=storage.corr_epi_mbar_ptr.data_ptr(), + ).make_participants() + mma_corr_producer, mma_corr_consumer = pipeline.PipelineUmmaAsync.create( + num_stages=self.mma_corr_stage, + producer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + consumer_group=make_thread_cooperative_group(self.threads_per_warp * len(self.correction_warp_ids)), + barrier_storage=storage.mma_corr_mbar_ptr.data_ptr(), + ).make_participants() + s0_s1_sequence_producer, s0_s1_sequence_consumer = pipeline.PipelineAsync.create( + num_stages=1, + producer_group=make_thread_cooperative_group(self.threads_per_warp * len(self.softmax0_warp_ids)), + consumer_group=make_thread_cooperative_group(self.threads_per_warp * len(self.softmax1_warp_ids)), + barrier_storage=storage.s0_s1_sequence_mbar_ptr.data_ptr(), + ).make_participants() + tmem_dealloc_mbar_ptr = storage.tmem_dealloc_mbar_ptr.data_ptr() + + # Correction & Epilogue & tmem barrier init + if warp_idx == self.empty_warp_id: + cute.arch.mbarrier_init( + tmem_dealloc_mbar_ptr, + self.threads_per_warp + * len( + ( + *self.softmax0_warp_ids, + *self.softmax1_warp_ids, + *self.correction_warp_ids, + ) + ), + ) + cute.arch.mbarrier_init_fence() + + # Generate smem tensor Q/K/V/O + # (MMA, MMA_Q, MMA_D, PIPE) + sQ = storage.sQ.get_tensor(q_smem_layout_staged.outer, swizzle=q_smem_layout_staged.inner) + # (MMA, MMA_K, MMA_D, PIPE) + sK = storage.sK.get_tensor(k_smem_layout_staged.outer, swizzle=k_smem_layout_staged.inner) + # (MMA, MMA_K, MMA_D, PIPE) + # Strip swizzle info to reuse smem + sV_ptr = cute.recast_ptr(sK.iterator, v_smem_layout_staged.inner) + sV = cute.make_tensor(sV_ptr, v_smem_layout_staged.outer) + sO = storage.sO.get_tensor(o_smem_layout_staged.outer, swizzle=o_smem_layout_staged.inner) + qk_thr_mma = qk_tiled_mma.get_slice(0) # default 1sm + pv_thr_mma = pv_tiled_mma.get_slice(0) # default 1sm + tSrQ = qk_thr_mma.make_fragment_A(sQ) + tSrK = qk_thr_mma.make_fragment_B(sK) + tOrV = pv_thr_mma.make_fragment_B(sV) + qk_acc_shape = qk_thr_mma.partition_shape_C((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) + tStS = qk_thr_mma.make_fragment_C(qk_acc_shape) + pv_acc_shape = pv_thr_mma.partition_shape_C((self.pv_mma_tiler[0], self.pv_mma_tiler[1])) + tOtO = pv_thr_mma.make_fragment_C(pv_acc_shape) + + tStS0 = cute.make_tensor(tStS.iterator + self.tmem_s0_offset, tStS.layout) + tStS1 = cute.make_tensor(tStS.iterator + self.tmem_s1_offset, tStS.layout) + tOtO0 = cute.make_tensor(tOtO.iterator + self.tmem_o0_offset, tOtO.layout) + tOtO1 = cute.make_tensor(tOtO.iterator + self.tmem_o1_offset, tOtO.layout) + + tP = cute.make_tensor(tStS.iterator, p_tmem_layout_staged.outer) + tOrP = pv_thr_mma.make_fragment_A(tP)[None, None, None, 0] + tOrP0 = cute.make_tensor( + tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset, + tOrP.layout, + ) + tOrP1 = cute.make_tensor( + tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p1_offset, + tOrP.layout, + ) + self.cta_sync_barrier.arrive_and_wait() + # /////////////////////////////////////////////////////////////////////////////// + # EMPTY + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx == self.empty_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_other) + + # /////////////////////////////////////////////////////////////////////////////// + # LOAD + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx == self.load_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_other) + + tile_sched = fmha_utils.create_fmha_static_tile_scheduler(tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()) + work_tile = tile_sched.initial_work_tile_info() + + while work_tile.is_valid_tile: + curr_block_coord = work_tile.tile_idx + batch_coord = curr_block_coord[2][1] + continue_cond = False + cuseqlen_q = Int32(0) + seqlen_q = mQ_qdl.shape[0] + if cutlass.const_expr(cum_seqlen_q is not None): + cuseqlen_q = cum_seqlen_q[batch_coord] + seqlen_q = cum_seqlen_q[batch_coord + 1] - cuseqlen_q + continue_cond = not fmha_utils.FmhaStaticTileScheduler.check_valid_work_for_seqlen_q( + self.cta_tiler[0], + curr_block_coord[0], + seqlen_q, + ) + if not continue_cond: + mQ_qdl_ = mQ_qdl + mK_kdl_ = mK_kdl + mV_dkl_ = mV_dkl + seqlen_k = mK_kdl.shape[0] + curr_block_coord_q = curr_block_coord + curr_block_coord_kv = curr_block_coord + + if cutlass.const_expr(cum_seqlen_q is not None): + logical_offset_mQ = ( + cuseqlen_q, + 0, + (0, 0), + ) + mQ_qdl_ = cute.domain_offset(logical_offset_mQ, mQ_qdl) + curr_block_coord_q = ( + curr_block_coord[0], + curr_block_coord[1], + curr_block_coord[2][0], + ) + + if cutlass.const_expr(cum_seqlen_k is not None): + cuseqlen_k = cum_seqlen_k[batch_coord] + seqlen_k = cum_seqlen_k[batch_coord + 1] - cuseqlen_k + if cutlass.const_expr(cum_seqlen_q is not None): + logical_offset_mK = ( + cuseqlen_k, + 0, + (0, 0), + ) + logical_offset_mV = ( + 0, + cuseqlen_k, + (0, 0), + ) + mK_kdl_ = cute.domain_offset(logical_offset_mK, mK_kdl) + mV_dkl_ = cute.domain_offset(logical_offset_mV, mV_dkl) + curr_block_coord_kv = ( + curr_block_coord[0], + curr_block_coord[1], + curr_block_coord[2][0], + ) + + # Local tile partition global tensors + # (bM, bK, loopM, loopK, loopL) + gQ_qdl = cute.flat_divide(mQ_qdl_, cute.select(self.qk_mma_tiler, mode=[0, 2])) + tSgQ_qdl = qk_thr_mma.partition_A(gQ_qdl) + tQsQ, tQgQ_qdl = cute.nvgpu.cpasync.tma_partition( + tma_atom_q, + 0, # no multicast + cute.make_layout(1), + cute.group_modes(sQ, 0, 3), + cute.group_modes(tSgQ_qdl, 0, 3), + ) + tQgQ = tQgQ_qdl[None, None, 0, curr_block_coord_q[2]] + + gK_kdl = cute.flat_divide(mK_kdl_, cute.select(self.qk_mma_tiler, mode=[1, 2])) + tSgK_kdl = qk_thr_mma.partition_B(gK_kdl) + tKsK, tKgK_kdl = cute.nvgpu.cpasync.tma_partition( + tma_atom_k, + 0, # no multicast + cute.make_layout(1), + cute.group_modes(sK, 0, 3), + cute.group_modes(tSgK_kdl, 0, 3), + ) + tKgK = tKgK_kdl[None, None, 0, curr_block_coord_kv[2]] + + gV_dkl = cute.flat_divide(mV_dkl_, cute.select(self.pv_mma_tiler, mode=[1, 2])) + tSgV_dkl = pv_thr_mma.partition_B(gV_dkl) + tVsV, tVgV_dkl = cute.nvgpu.cpasync.tma_partition( + tma_atom_v, + 0, # no multicast + cute.make_layout(1), + cute.group_modes(sV, 0, 3), + cute.group_modes(tSgV_dkl, 0, 3), + ) + tVgV = tVgV_dkl[None, 0, None, curr_block_coord_kv[2]] + + # Q0 + q0_coord = 2 * curr_block_coord_q[0] + q0_handle = load_q_producer.acquire_and_advance() + cute.copy( + tma_atom_q, + tQgQ[None, q0_coord], + tQsQ[None, q0_handle.index], + tma_bar_ptr=q0_handle.barrier, + ) + # K0 + if cutlass.const_expr(sparse_block_count is not None): + head_coord = curr_block_coord[2][0] + curr_sparse_block_count = sparse_block_count[batch_coord, head_coord, curr_block_coord[0]] + curr_sparse_block_indices = sparse_block_indices[batch_coord, head_coord, curr_block_coord[0], None] + kv_coord = curr_sparse_block_indices[0] + else: + seqlen_kv_loop_start = fmha_utils.FusedMask.get_trip_start( + self.mask_type, + curr_block_coord, + self.cta_tiler, + seqlen_q, + seqlen_k, + window_size_left, + ) + kv_coord = seqlen_kv_loop_start + k_handle = load_kv_producer.acquire_and_advance() + cute.copy( + tma_atom_k, + tKgK[None, kv_coord], + tKsK[None, k_handle.index], + tma_bar_ptr=k_handle.barrier, + ) + # Q1 + q1_coord = q0_coord + 1 + q1_handle = load_q_producer.acquire_and_advance() + cute.copy( + tma_atom_q, + tQgQ[None, q1_coord], + tQsQ[None, q1_handle.index], + tma_bar_ptr=q1_handle.barrier, + ) + # V0 + v_handle = load_kv_producer.acquire_and_advance() + cute.copy( + tma_atom_v, + tVgV[None, kv_coord], + tVsV[None, v_handle.index], + tma_bar_ptr=v_handle.barrier, + ) + if cutlass.const_expr(sparse_block_count is None): + kv_coord += 1 + + if cutlass.const_expr(sparse_block_count is not None): + seqlen_kv_loop_steps = curr_sparse_block_count - 1 + else: + seqlen_kv_loop_steps = ( + fmha_utils.FusedMask.get_trip_count( + self.mask_type, + curr_block_coord, + self.cta_tiler, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + - 1 + ) + for i in cutlass.range(0, seqlen_kv_loop_steps, 1, unroll=1): + # Ki + if cutlass.const_expr(sparse_block_count is not None): + kv_coord = curr_sparse_block_indices[i + 1] + k_handle = load_kv_producer.acquire_and_advance() + cute.copy( + tma_atom_k, + tKgK[None, kv_coord], + tKsK[None, k_handle.index], + tma_bar_ptr=k_handle.barrier, + ) + # Vi + v_handle = load_kv_producer.acquire_and_advance() + cute.copy( + tma_atom_v, + tVgV[None, kv_coord], + tVsV[None, v_handle.index], + tma_bar_ptr=v_handle.barrier, + ) + if cutlass.const_expr(sparse_block_count is None): + kv_coord += 1 + # End of seqlen_kv loop + + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + # End of persistent scheduler loop + + # /////////////////////////////////////////////////////////////////////////////// + # MMA + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx == self.mma_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_other) + + # Alloc tmem buffer + tmem_alloc_cols = Int32(self.tmem_alloc_cols) + cute.arch.alloc_tmem(tmem_alloc_cols, storage.tmem_holding_buf) + self.tmem_alloc_barrier.arrive_and_wait() + tile_sched = fmha_utils.create_fmha_static_tile_scheduler(tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()) + work_tile = tile_sched.initial_work_tile_info() + + while work_tile.is_valid_tile: + curr_block_coord = work_tile.tile_idx + batch_coord = curr_block_coord[2][1] + continue_cond = False + seqlen_q = mQ_qdl.shape[0] + if cutlass.const_expr(cum_seqlen_q is not None): + cuseqlen_q = cum_seqlen_q[batch_coord] + seqlen_q = cum_seqlen_q[batch_coord + 1] - cuseqlen_q + continue_cond = not fmha_utils.FmhaStaticTileScheduler.check_valid_work_for_seqlen_q( + self.cta_tiler[0], + curr_block_coord[0], + seqlen_q, + ) + + if not continue_cond: + seqlen_k = mK_kdl.shape[0] + if cutlass.const_expr(cum_seqlen_k is not None): + cuseqlen_k = cum_seqlen_k[batch_coord] + seqlen_k = cum_seqlen_k[batch_coord + 1] - cuseqlen_k + + # GEMM_QK00 (Q0 * K0 -> S0) + # 1. wait for Q0 + q0_handle = load_q_consumer.wait_and_advance() + tSrQ0 = tSrQ[None, None, None, q0_handle.index] + # 2. wait for K0 + k_handle = load_kv_consumer.wait_and_advance() + tSrK0 = tSrK[None, None, None, k_handle.index] + # 3. acquire empty S0 buffer + s0_handle = mma_s0_producer.acquire_and_advance() + # 4. gemm + num_kphases = cute.size(tSrQ0, mode=[2]) + for kphase_idx in cutlass.range(num_kphases, unroll_full=True): + kphase_coord = (None, None, kphase_idx) + qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, kphase_idx != 0) + cute.gemm( + qk_tiled_mma, + tStS0, + tSrQ0[kphase_coord], + tSrK0[kphase_coord], + tStS0, + ) + # 5. release S0 + s0_handle.commit() + # End of GEMM (Q0 * K0 -> S0) + + # GEMM_QK10 (Q1 * K0 -> S1), K0 is ready in GEMM_QK00 + # 1. wait for Q1 + q1_handle = load_q_consumer.wait_and_advance() + tSrQ1 = tSrQ[None, None, None, q1_handle.index] + # 2. acquire empty S1 + s1_handle = mma_s1_producer.acquire_and_advance() + # 3. gemm + num_kphases = cute.size(tSrQ1, mode=[2]) + for kphase_idx in cutlass.range(num_kphases, unroll_full=True): + kphase_coord = (None, None, kphase_idx) + qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, kphase_idx != 0) + cute.gemm( + qk_tiled_mma, + tStS1, + tSrQ1[kphase_coord], + tSrK0[kphase_coord], + tStS1, + ) + # 4. release S1 + s1_handle.commit() + # 5. release K0 + k_handle.release() + # End of GEMM (Q1 * K0 -> S1) + # Note: Q0 & Q1 are still needed in the seqlen_kv loop + # so we need to release them after the seqlen_kv loop + + # GEMM_PV00 (P0 * V0 -> O0_partial), O0 needs to be accumulated in the seqlen_kv loop + # 1. wait for V0 + v_handle = load_kv_consumer.wait_and_advance() + tOrVi = tOrV[None, None, None, v_handle.index] + # 2. acquire corrected O0_partial + # Note: acquire corr first to take it out of the critical + # path since softmax takes longer + o0_handle = mma_corr_producer.acquire_and_advance() + # 3. acquire P0 + # this acquire returns the ownership of all of S0 to the mma warp + # including the P0 part (inplaced in S0) + s0_handle = mma_s0_producer.acquire_and_advance() + # 4. gemm + num_kphases = cute.size(tOrP0, mode=[2]) + for kphase_idx in cutlass.range(num_kphases, unroll_full=True): + kphase_coord = (None, None, kphase_idx) + pv_tiled_mma.set(tcgen05.Field.ACCUMULATE, kphase_idx != 0) + cute.gemm( + pv_tiled_mma, + tOtO0, + tOrP0[kphase_coord], + tOrVi[kphase_coord], + tOtO0, + ) + # 5. release accumulated O0_partial + o0_handle.commit() + # End of GEMM_PV00 (P0 * V0 -> O0_partial) + + if cutlass.const_expr(sparse_block_count is not None): + head_coord = curr_block_coord[2][0] + seqlen_kv_loop_steps = sparse_block_count[batch_coord, head_coord, curr_block_coord[0]] - 1 + else: + seqlen_kv_loop_steps = ( + fmha_utils.FusedMask.get_trip_count( + self.mask_type, + curr_block_coord, + self.cta_tiler, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + - 1 + ) + + # O1 hasn't been accumulated yet, its first MMA calculation doesn't need to accumulate + pv_whether_acc = False + for i in cutlass.range(0, seqlen_kv_loop_steps, 1, unroll=1): + # GEMM_QK0i (Q0 * Ki -> S0) + # 1. wait for Ki + k_handle = load_kv_consumer.wait_and_advance() + tSrKi = tSrK[None, None, None, k_handle.index] + # 2. gemm + inner_num_kphases = cute.size(tSrQ0, mode=[2]) + for kphase_idx in cutlass.range(inner_num_kphases, unroll_full=True): + kphase_coord = (None, None, kphase_idx) + qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, kphase_idx != 0) + cute.gemm( + qk_tiled_mma, + tStS0, + tSrQ0[kphase_coord], + tSrKi[kphase_coord], + tStS0, + ) + # 3. release S0 + s0_handle.commit() + # End of GEMM_QK0i (Q0 * Ki -> S0) + + # GEMM_PV1(i-1) (P1 * V(i-1) -> O1_partial), V(i-1) is ready in GEMM_PV0(i-1) + # 1. acquire corrected O1_partial + o1_handle = mma_corr_producer.acquire_and_advance() + # 2. acquire P1 + s1_handle = mma_s1_producer.acquire_and_advance() + # 3. gemm + inner_num_kphases = cute.size(tOrP0, mode=[2]) + for kphase_idx in cutlass.range(inner_num_kphases, unroll_full=True): + kphase_coord = (None, None, kphase_idx) + pv_tiled_mma.set(tcgen05.Field.ACCUMULATE, pv_whether_acc) + cute.gemm( + pv_tiled_mma, + tOtO1, + tOrP1[kphase_coord], + tOrVi[kphase_coord], + tOtO1, + ) + pv_whether_acc = True + # 4. release accumulated O1_partial + o1_handle.commit() + # 5. release V(i-1) + v_handle.release() + # End of GEMM_PV1(i-1) (P1 * V(i-1) -> O1_partial) + + # GEMM_QK1i (Q1 * Ki -> S1), Q1 is ready in GEMM_QK10; Ki is ready in GEMM_QK0i + # 1. gemm + inner_num_kphases = cute.size(tSrQ1, mode=[2]) + for kphase_idx in cutlass.range(inner_num_kphases, unroll_full=True): + kphase_coord = (None, None, kphase_idx) + qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, kphase_idx != 0) + cute.gemm( + qk_tiled_mma, + tStS1, + tSrQ1[kphase_coord], + tSrKi[kphase_coord], + tStS1, + ) + s1_handle.commit() + # 2. release Ki + k_handle.release() + # End of GEMM_QK1i (Q1 * Ki -> S1) + + # GEMM_PV0i (P0 * Vi -> O0_partial) + # 1. wait for Vi + v_handle = load_kv_consumer.wait_and_advance() + tOrVi = tOrV[None, None, None, v_handle.index] + # 2. acquire corrected O0_partial + o0_handle = mma_corr_producer.acquire_and_advance() + # 3. acquire P0 + s0_handle = mma_s0_producer.acquire_and_advance() + # 4. gemm + inner_num_kphases = cute.size(tOrP0, mode=[2]) + for kphase_idx in cutlass.range(inner_num_kphases, unroll_full=True): + kphase_coord = (None, None, kphase_idx) + pv_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + cute.gemm( + pv_tiled_mma, + tOtO0, + tOrP0[kphase_coord], + tOrVi[kphase_coord], + tOtO0, + ) + # 5. release accumulated O0_partial + o0_handle.commit() + # End of GEMM_PV0i (P0 * Vi -> O0_partial) + # End of seqlen_kv loop + + # release Q0 & Q1 + q0_handle.release() + q1_handle.release() + + # GEMM_PV1(i_end) (P1 * Vi_end -> O1) + # 1. acquire corrected O1_partial + o1_handle = mma_corr_producer.acquire_and_advance() + # 2. acquire P1 + s1_handle = mma_s1_producer.acquire_and_advance() + # 3. gemm + num_kphases = cute.size(tOrP1, mode=[2]) + for kphase_idx in cutlass.range(num_kphases, unroll_full=True): + kphase_coord = (None, None, kphase_idx) + pv_tiled_mma.set(tcgen05.Field.ACCUMULATE, pv_whether_acc) + cute.gemm( + pv_tiled_mma, + tOtO1, + tOrP1[kphase_coord], + tOrVi[kphase_coord], + tOtO1, + ) + pv_whether_acc = True + # 4. commit accumulated O1 + o1_handle.commit() + # 5. release Vi_end + v_handle.release() + # End of GEMM_PV1(i_end) (P1 * Vi_end -> O1) + + # Commit S0 and S1 + s0_handle.commit() + s1_handle.commit() + + # Advance to next tile + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + # End of persistent scheduler loop + + # dealloc tmem buffer + cute.arch.relinquish_tmem_alloc_permit() + cute.arch.mbarrier_wait(tmem_dealloc_mbar_ptr, 0) + tmem_alloc_cols = Int32(self.tmem_alloc_cols) + # Retrieving tmem ptr and make acc + tmem_ptr = cute.arch.retrieve_tmem_ptr( + Float32, + alignment=16, + ptr_to_buffer_holding_addr=storage.tmem_holding_buf, + ) + cute.arch.dealloc_tmem(tmem_ptr, tmem_alloc_cols) + + # /////////////////////////////////////////////////////////////////////////////// + # Epilogue + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx == self.epilogue_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_other) + tile_sched = fmha_utils.create_fmha_static_tile_scheduler(tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()) + work_tile = tile_sched.initial_work_tile_info() + + while work_tile.is_valid_tile: + curr_block_coord = work_tile.tile_idx + batch_coord = curr_block_coord[2][1] + continue_cond = False + cuseqlen_q = Int32(0) + seqlen_q = mQ_qdl.shape[0] + + if cutlass.const_expr(cum_seqlen_q is not None): + cuseqlen_q = cum_seqlen_q[batch_coord] + seqlen_q = cum_seqlen_q[batch_coord + 1] - cuseqlen_q + continue_cond = not fmha_utils.FmhaStaticTileScheduler.check_valid_work_for_seqlen_q( + self.cta_tiler[0], + curr_block_coord[0], + seqlen_q, + ) + if not continue_cond: + curr_block_coord_o = curr_block_coord + mO_qdl_ = mO_qdl + mO_gmem_ = mO_gmem + if cutlass.const_expr(cum_seqlen_q is not None): + logical_offset_mO = ( + cuseqlen_q, + 0, + (0, 0), + ) + mO_qdl_ = cute.domain_offset(logical_offset_mO, mO_qdl_) + mO_gmem_ = cute.domain_offset(logical_offset_mO, mO_gmem_) + curr_block_coord_o = ( + curr_block_coord[0], + curr_block_coord[1], + curr_block_coord[2][0], + ) + + o0_coord = 2 * curr_block_coord_o[0] + o1_coord = o0_coord + 1 + gO_qdl = cute.flat_divide(mO_qdl_, cute.select(self.pv_mma_tiler, mode=[0, 1])) + gO = gO_qdl[None, None, None, 0, curr_block_coord_o[2]] + tOsO, tOgO = cute.nvgpu.cpasync.tma_partition( + tma_atom_o, + 0, + cute.make_layout(1), + cute.group_modes(sO, 0, 2), + cute.group_modes(gO, 0, 2), + ) + + # O0 O1 using the same pipeline + # wait from corr, issue tma store on smem + # O0 + # 1. wait for O0 final + o0_handle = corr_epi_consumer.wait_and_advance() + # 2. copy O0 to gmem + o0_row = o0_coord * self.epi_tile[0] + o0_use_tma = o0_row + self.epi_tile[0] <= seqlen_q + if o0_use_tma: + cute.copy(tma_atom_o, tOsO[None, 0], tOgO[None, o0_coord]) + cute.arch.cp_async_bulk_commit_group() + else: + self.store_o_tail(sO, mO_gmem_, o0_row, seqlen_q, curr_block_coord_o[2], 0) + # O1 + # 1. wait for O1 final + o1_handle = corr_epi_consumer.wait_and_advance() + # 2. copy O1 to gmem + o1_row = o1_coord * self.epi_tile[0] + o1_use_tma = o1_row + self.epi_tile[0] <= seqlen_q + if o1_use_tma: + cute.copy(tma_atom_o, tOsO[None, 1], tOgO[None, o1_coord]) + cute.arch.cp_async_bulk_commit_group() + else: + self.store_o_tail(sO, mO_gmem_, o1_row, seqlen_q, curr_block_coord_o[2], 1) + + # Ensure O0 buffer is ready to be released + if o0_use_tma: + if o1_use_tma: + cute.arch.cp_async_bulk_wait_group(1, read=True) + else: + cute.arch.cp_async_bulk_wait_group(0, read=True) + o0_handle.release() + # Ensure O1 buffer is ready to be released + if o1_use_tma: + cute.arch.cp_async_bulk_wait_group(0, read=True) + o1_handle.release() + + # Advance to next tile + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + # End of persistent scheduler loop + + # /////////////////////////////////////////////////////////////////////////////// + # Softmax0 + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx < self.softmax1_warp_ids[0]: + # increase register after decreasing + cute.arch.setmaxregister_increase(self.num_regs_softmax) + + self.softmax( + stage=0, + seqlen_k=mK_kdl.shape[0], + seqlen_q=mQ_qdl.shape[0], + cum_seqlen_q=cum_seqlen_q, + cum_seqlen_k=cum_seqlen_k, + scale_softmax_log2=scale_softmax_log2, + qk_thr_mma=qk_thr_mma, + tStS=tStS, + tStSi=tStS0, + window_size_left=window_size_left, + window_size_right=window_size_right, + sparse_block_count=sparse_block_count, + sparse_block_indices=sparse_block_indices, + mma_si_consumer=mma_s0_consumer, + si_corr_producer=s0_corr_producer, + s0_s1_sequence_consumer=s0_s1_sequence_consumer, + s0_s1_sequence_producer=s0_s1_sequence_producer, + tile_sched_params=tile_sched_params, + ) + cute.arch.mbarrier_arrive(tmem_dealloc_mbar_ptr) + + # /////////////////////////////////////////////////////////////////////////////// + # Softmax1 + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx < self.correction_warp_ids[0] and warp_idx >= self.softmax1_warp_ids[0]: + # increase register after decreasing + cute.arch.setmaxregister_increase(self.num_regs_softmax) + + self.softmax( + stage=1, + seqlen_k=mK_kdl.shape[0], + seqlen_q=mQ_qdl.shape[0], + cum_seqlen_q=cum_seqlen_q, + cum_seqlen_k=cum_seqlen_k, + scale_softmax_log2=scale_softmax_log2, + qk_thr_mma=qk_thr_mma, + tStS=tStS, + tStSi=tStS1, + window_size_left=window_size_left, + window_size_right=window_size_right, + sparse_block_count=sparse_block_count, + sparse_block_indices=sparse_block_indices, + mma_si_consumer=mma_s1_consumer, + si_corr_producer=s1_corr_producer, + s0_s1_sequence_consumer=s0_s1_sequence_consumer, + s0_s1_sequence_producer=s0_s1_sequence_producer, + tile_sched_params=tile_sched_params, + ) + cute.arch.mbarrier_arrive(tmem_dealloc_mbar_ptr) + + # /////////////////////////////////////////////////////////////////////////////// + # Correction + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx >= self.correction_warp_ids[0] and warp_idx < self.mma_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_correction) + + cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) + tScS = qk_thr_mma.partition_C(cS) + + tStS_vec_layout = cute.composition(tStS.layout, cute.make_layout((128, 2))) + + tStS_vec0 = cute.make_tensor(tStS.iterator + self.tmem_vec0_offset, tStS_vec_layout) + tStS_vec1 = cute.make_tensor(tStS.iterator + self.tmem_vec1_offset, tStS_vec_layout) + + tScS_vec_layout = cute.composition(tScS.layout, cute.make_layout((128, 2))) + tScS_vec = cute.make_tensor(tScS.iterator, tScS_vec_layout) + + tmem_load_v_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(2)), + self.qk_acc_dtype, + ) + + tiled_tmem_load_vec = tcgen05.make_tmem_copy(tmem_load_v_atom, tStS_vec0) + thread_idx = tidx % (self.threads_per_warp * len(self.correction_warp_ids)) + thr_tmem_load_vec = tiled_tmem_load_vec.get_slice(thread_idx) + + tTMEM_LOAD_VECtS0 = thr_tmem_load_vec.partition_S(tStS_vec0) + tTMEM_LOAD_VECtS1 = thr_tmem_load_vec.partition_S(tStS_vec1) + tTMEM_LOAD_VECcS = thr_tmem_load_vec.partition_D(tScS_vec) + + tile_sched = fmha_utils.create_fmha_static_tile_scheduler(tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()) + work_tile = tile_sched.initial_work_tile_info() + + while work_tile.is_valid_tile: + curr_block_coord = work_tile.tile_idx + curr_block_coord_lse = curr_block_coord + batch_coord = curr_block_coord[2][1] + seqlen_k = mK_kdl.shape[0] + continue_cond = False + cuseqlen_q = Int32(0) + seqlen_q = mQ_qdl.shape[0] + + if cutlass.const_expr(cum_seqlen_q is not None): + cuseqlen_q = cum_seqlen_q[batch_coord] + seqlen_q = cum_seqlen_q[batch_coord + 1] - cuseqlen_q + # for varlen LSE, batch == 1 + curr_block_coord_lse = ( + curr_block_coord[0], + curr_block_coord[1], + (curr_block_coord[2][0], 0), + ) + continue_cond = not fmha_utils.FmhaStaticTileScheduler.check_valid_work_for_seqlen_q( + self.cta_tiler[0], + curr_block_coord[0], + seqlen_q, + ) + + if not continue_cond: + row_idx = curr_block_coord[0] * self.cta_tiler[0] + tTMEM_LOAD_VECcS[0][0] + if cutlass.const_expr(cum_seqlen_k is not None): + cuseqlen_k = cum_seqlen_k[batch_coord] + seqlen_k = cum_seqlen_k[batch_coord + 1] - cuseqlen_k + # Ignore first signal: the first O partial needs no correction. + vec0_handle = s0_corr_consumer.wait_and_advance() + vec0_handle.release() + vec1_handle = s1_corr_consumer.wait_and_advance() + + if cutlass.const_expr(sparse_block_count is not None): + head_coord = curr_block_coord[2][0] + seqlen_kv_loop_steps = sparse_block_count[batch_coord, head_coord, curr_block_coord[0]] - 1 + else: + seqlen_kv_loop_steps = ( + fmha_utils.FusedMask.get_trip_count( + self.mask_type, + curr_block_coord, + self.cta_tiler, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + - 1 + ) + for i in cutlass.range(0, seqlen_kv_loop_steps, 1, unroll=1): + # wait for vec0 (row_wise current max & previous max) + vec0_handle = s0_corr_consumer.wait_and_advance() + # wait for o0 + o0_handle = mma_corr_consumer.wait_and_advance() + if cutlass.const_expr(not self.zero_reference_max): + tTMEM_LOAD_VECrS = cute.make_rmem_tensor(tTMEM_LOAD_VECcS.shape, self.qk_acc_dtype) + cute.copy( + tiled_tmem_load_vec, + tTMEM_LOAD_VECtS0, + tTMEM_LOAD_VECrS, + ) + scale_ = Float32(0.0) + scale = Float32(1.0) + warp_needs_rescale = cute.arch.vote_any_sync(tTMEM_LOAD_VECrS[0] != tTMEM_LOAD_VECrS[1]) + if warp_needs_rescale: + scale_ = scale_softmax_log2 * (tTMEM_LOAD_VECrS[0] - tTMEM_LOAD_VECrS[1]) + scale = cute.math.exp2(scale_, fastmath=True) + self.correction_rescale(pv_thr_mma, tOtO0, scale) + # release vec1 & o0 + vec1_handle.release() + if cutlass.const_expr(not self.zero_reference_max): + cute.arch.fence_view_async_tmem_store() + o0_handle.release() + + # wait for vec1 (row_wise current max & previous max) + vec1_handle = s1_corr_consumer.wait_and_advance() + o1_handle = mma_corr_consumer.wait_and_advance() + if cutlass.const_expr(not self.zero_reference_max): + cute.copy( + tiled_tmem_load_vec, + tTMEM_LOAD_VECtS1, + tTMEM_LOAD_VECrS, + ) + scale_ = Float32(0.0) + scale = Float32(1.0) + warp_needs_rescale = cute.arch.vote_any_sync(tTMEM_LOAD_VECrS[0] != tTMEM_LOAD_VECrS[1]) + if warp_needs_rescale: + scale_ = scale_softmax_log2 * (tTMEM_LOAD_VECrS[0] - tTMEM_LOAD_VECrS[1]) + scale = cute.math.exp2(scale_, fastmath=True) + self.correction_rescale(pv_thr_mma, tOtO1, scale) + vec0_handle.release() + if cutlass.const_expr(not self.zero_reference_max): + cute.arch.fence_view_async_tmem_store() + o1_handle.release() + # End of seqlen_corr_loop_steps + vec1_handle.release() + + # wait for vec0 (row_wise global sum) + vec0_handle = s0_corr_consumer.wait_and_advance() + tTMEM_LOAD_VECrS = cute.make_rmem_tensor(tTMEM_LOAD_VECcS.shape, self.qk_acc_dtype) + cute.copy(tiled_tmem_load_vec, tTMEM_LOAD_VECtS0, tTMEM_LOAD_VECrS) + cute.arch.fence_view_async_tmem_load() + vec0_handle.release() + # wait for o0 + o0_handle = mma_corr_consumer.wait_and_advance() + o0_final_handle = corr_epi_producer.acquire_and_advance() + self.correction_epilog( + pv_thr_mma, + tOtO0, + mLSE, + tTMEM_LOAD_VECrS, + row_idx, + cuseqlen_q, + seqlen_q, + curr_block_coord_lse, + scale_softmax, + scale_output / tTMEM_LOAD_VECrS[0], + sO[None, None, 0], + ) + o0_handle.release() + o0_final_handle.commit() + + # wait for vec1 (row_wise global sum) + vec1_handle = s1_corr_consumer.wait_and_advance() + cute.copy(tiled_tmem_load_vec, tTMEM_LOAD_VECtS1, tTMEM_LOAD_VECrS) + cute.arch.fence_view_async_tmem_load() + vec1_handle.release() + # wait for o1 + o1_handle = mma_corr_consumer.wait_and_advance() + o1_final_handle = corr_epi_producer.acquire_and_advance() + row_idx += self.qk_mma_tiler[0] + self.correction_epilog( + pv_thr_mma, + tOtO1, + mLSE, + tTMEM_LOAD_VECrS, + row_idx, + cuseqlen_q, + seqlen_q, + curr_block_coord_lse, + scale_softmax, + scale_output / tTMEM_LOAD_VECrS[0], + sO[None, None, 1], + ) + o1_handle.release() + o1_final_handle.commit() + # Advance to next tile + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + # End of persistent scheduler loop + cute.arch.mbarrier_arrive(tmem_dealloc_mbar_ptr) + return + + @cute.jit + def store_o_tail( + self, + sO: cute.Tensor, + mO_gmem: cute.Tensor, + row_start: Int32, + seqlen_q: Int32, + head_coord: cute.Coord, + stage: cutlass.Constexpr, + ): + """Predicated O tile store for packed-varlen sequence tails.""" + tidx, _, _ = cute.arch.thread_idx() + lane_idx = tidx % self.threads_per_warp + valid_rows = seqlen_q - row_start + if valid_rows > self.epi_tile[0]: + valid_rows = self.epi_tile[0] + if valid_rows > 0: + valid_elems = valid_rows * self.head_dim + for elem_idx in cutlass.range(lane_idx, valid_elems, self.threads_per_warp, unroll=1): + row = elem_idx // self.head_dim + col = elem_idx - row * self.head_dim + mO_gmem[row_start + row, col, head_coord] = sO[row, col, stage] + + @cute.jit + def softmax_step( + self, + stage: int, + need_apply_mask: bool, + iter_args: tuple, + value_args: tuple, + pipeline_args: tuple, + atom_args: tuple, + tensor_args: tuple, + ) -> Tuple[ + Float32, + Float32, + Float32, + pipeline.PipelineProducer.ImmutableResourceHandle, + pipeline.PipelineConsumer, + pipeline.PipelineProducer, + pipeline.PipelineConsumer, + pipeline.PipelineProducer, + ]: + """Perform a single step of the softmax computation on a block of attention scores. + + This method processes one block of the attention matrix, computing numerically stable + softmax by first finding the row maximum, subtracting it from all elements, applying + exponential function, and then normalizing by the sum of exponentials. It also handles + optional masking of attention scores. + + The method involves several key operations: + 1. Loading attention scores from tensor memory + 2. Applying optional masking based on position + 3. Computing row-wise maximum values for numerical stability + 4. Transforming scores using exp2(x*scale - max*scale) + 5. Computing row sums for normalization + 6. Coordinating pipeline synchronization between different processing stages + + :param stage: Processing stage (0 for first half, 1 for second half) + :type stage: int + :param need_apply_mask: Whether to apply attention masking + :type need_apply_mask: bool + :param iter_args: Tuple containing the counting tensor, row_max, row_sum, and vector buffer's handle for current iteration + :type iter_args: tuple + :param value_args: Tuple containing seqlen_k, seqlen_q, and scale_softmax_log2 + :type value_args: tuple + :param pipeline_args: Tuple containing pipeline related arguments for MMA, correction, and sequence synchronization + :type pipeline_args: tuple + :param atom_args: Tuple containing mma & copy atoms + :type atom_args: tuple + :param tensor_args: Tuple containing softmax related tensors + :type tensor_args: tuple + :param fused_mask: Compute trip counts and apply masking for attention blocks + :type fused_mask: fmha_utils.FusedMask + :return: Updated state values (row_max, row_sum, and pipeline related arguments) + :rtype: tuple + """ + cS, row_max, observed_row_max, row_sum, update_row_max, vec_i_handle = iter_args + seqlen_k, seqlen_q, scale_softmax_log2, window_size_left, window_size_right = value_args + ( + mma_si_consumer, + si_corr_producer, + s0_s1_sequence_consumer, + s0_s1_sequence_producer, + ) = pipeline_args + ( + qk_thr_mma, + tiled_tmem_load, + tiled_tmem_store, + tiled_tmem_store_vec, + thr_tmem_load, + thr_tmem_store, + thr_tmem_store_vec, + ) = atom_args + ( + tTMEM_LOADtS, + tTMEM_STORE_VECtS, + tTMEM_STOREtS_x4, + ) = tensor_args + + tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width + tScS = qk_thr_mma.partition_C(cS) + tScS_vec_layout = cute.composition(tScS.layout, cute.make_layout((128, 2))) + tScS_vec = cute.make_tensor(tScS.iterator, tScS_vec_layout) + + tScS_P_layout = cute.composition(tScS.layout, cute.make_layout((128, tilePlikeFP32))) + tScS_P = cute.make_tensor(tScS.iterator, tScS_P_layout) + tTMEM_LOADcS = thr_tmem_load.partition_D(tScS) + tTMEM_STORE_VECcS = thr_tmem_store_vec.partition_S(tScS_vec) + tTMEM_STOREcS = thr_tmem_store.partition_S(tScS_P) + + # Wait for Si + si_handle = mma_si_consumer.wait_and_advance() + tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) + cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS) + if need_apply_mask: + fmha_utils.FusedMask.apply_mask( + self.mask_type, + tTMEM_LOADrS, + tTMEM_LOADcS, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + + old_row_max = row_max + if cutlass.const_expr(self.zero_reference_max): + observed_row_max = row_max + row_max_safe = row_max + else: + observed_row_max = tTMEM_LOADrS.load().reduce(cute.ReductionOp.MAX, observed_row_max, 0) + if cutlass.const_expr(self.max_update_interval == 1): + row_max = observed_row_max + else: + if update_row_max or row_max == -cutlass.Float32.inf: + row_max = observed_row_max + row_max_safe = row_max + if row_max == -cutlass.Float32.inf: + row_max_safe = 0.0 + tTMEM_STORE_VECrS = cute.make_rmem_tensor(tTMEM_STORE_VECcS.shape, self.qk_acc_dtype) + tTMEM_STORE_VECrS[0] = old_row_max + tTMEM_STORE_VECrS[1] = row_max_safe + cute.copy(tiled_tmem_store_vec, tTMEM_STORE_VECrS, tTMEM_STORE_VECtS) + cute.arch.fence_view_async_tmem_store() + # Notify correction wg that row_max is ready. + vec_i_handle.commit() + + tTMEM_STORErS_x4 = cute.make_rmem_tensor(tTMEM_STOREcS.shape, self.qk_acc_dtype) + tTMEM_STORErS_x4_e = cute.make_tensor( + cute.recast_ptr(tTMEM_STORErS_x4.iterator, dtype=self.q_dtype), + tTMEM_LOADrS.layout, + ) + + scale = scale_softmax_log2 + minus_row_max_scale = (0.0 - row_max_safe) * scale + self.softmax_prescale_log2 + + # Sequence barrier wait + if cutlass.const_expr(stage == 0): + sequence_producer_handle = s0_s1_sequence_producer.acquire_and_advance() + else: + sequence_consumer_handle = s0_s1_sequence_consumer.wait_and_advance() + frg_cnt = 4 + frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt + tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) + tTMEM_STORErS_x4_e_frg = cute.logical_divide(tTMEM_STORErS_x4_e, cute.make_layout(frg_tile)) + if cutlass.const_expr(not self.zero_reference_max): + acc_scale_ = scale * (old_row_max - row_max_safe) + acc_scale = 1.0 + if old_row_max != row_max_safe: + acc_scale = cute.math.exp2(acc_scale_, fastmath=True) + row_sum *= acc_scale + for j in range(frg_cnt): + for k in cutlass.range(cute.size(tTMEM_LOADrS_frg, mode=[0]), vectorize=True): + tTMEM_LOADrS_frg[k, j] = tTMEM_LOADrS_frg[k, j] * scale + minus_row_max_scale + tTMEM_LOADrS_frg[k, j] = cute.math.exp2(tTMEM_LOADrS_frg[k, j], fastmath=True) + + s_vec = tTMEM_LOADrS_frg[None, j].load() + row_sum = s_vec.reduce(cute.ReductionOp.ADD, row_sum, 0) + tTMEM_STORErS_x4_e_frg[None, j].store(s_vec.to(self.q_dtype)) + # Sequence barrier arrive + if cutlass.const_expr(stage == 0): + sequence_producer_handle.commit() + else: + sequence_consumer_handle.release() + cute.copy(tiled_tmem_store, tTMEM_STORErS_x4, tTMEM_STOREtS_x4) + cute.arch.fence_view_async_tmem_store() + # Notify tensor core warp that softmax(S->P) is ready + si_handle.release() + + vec_i_handle = si_corr_producer.acquire_and_advance() + + return ( + row_max, + observed_row_max, + row_sum, + vec_i_handle, + mma_si_consumer, + si_corr_producer, + s0_s1_sequence_consumer, + s0_s1_sequence_producer, + ) + + # For both softmax0 and softmax1 warp group + @cute.jit + def softmax( + self, + stage: int, + seqlen_k: Int32, + seqlen_q: Int32, + cum_seqlen_q: Optional[cute.Tensor], + cum_seqlen_k: Optional[cute.Tensor], + scale_softmax_log2: Float32, + qk_thr_mma: cute.ThrMma, + tStS: cute.Tensor, + tStSi: cute.Tensor, + window_size_left: Optional[Int32], + window_size_right: Optional[Int32], + sparse_block_count: Optional[cute.Tensor], + sparse_block_indices: Optional[cute.Tensor], + mma_si_consumer: pipeline.PipelineConsumer, + si_corr_producer: pipeline.PipelineProducer, + s0_s1_sequence_consumer: pipeline.PipelineConsumer, + s0_s1_sequence_producer: pipeline.PipelineProducer, + tile_sched_params: fmha_utils.FmhaStaticTileSchedulerParams, + ): + """Compute softmax on attention scores from QK matrix multiplication. + + This method handles the softmax computation for either the first or second half of the + attention matrix, depending on the 'stage' parameter. It calculates row-wise maximum + and sum values needed for stable softmax computation, applies optional masking, and + transforms raw attention scores into probability distributions. + + The implementation uses specialized memory access patterns and efficient math operations + for computing exp(x) using exp2 functions. It also coordinates pipeline + synchronization between MMA, correction, and sequence processing stages. + + :param stage: Processing stage (0 for first half, 1 for second half of attention matrix) + :type stage: int + :param seqlen_k: Length of the key sequence + :type seqlen_k: Int32 + :param seqlen_q: Length of the query sequence + :type seqlen_q: Int32 + :param cum_seqlen_q: Cumulative sequence lengths for queries + :type cum_seqlen_q: cute.Tensor | None + :param cum_seqlen_k: Cumulative sequence lengths for keys + :type cum_seqlen_k: cute.Tensor | None + :param scale_softmax_log2: Log2 scale factor for softmax operation + :type scale_softmax_log2: Float32 + :param qk_thr_mma: Thread MMA operation for QK matrix multiplication + :type qk_thr_mma: cute.ThrMma + :param tStS: Shared tensor for softmax input/output + :type tStS: cute.Tensor + :param tStSi: Input tensor containing attention scores + :type tStSi: cute.Tensor + :param window_size_left: Left-side sliding window size for attention masking. + :type window_size_left: Optional[Int32] + :param window_size_right: Right-side sliding window size for attention masking. + :type window_size_right: Optional[Int32] + :param mma_si_pipeline: Pipeline for synchronizing with MMA operations + :type mma_si_pipeline: pipeline.PipelineAsync + :param si_corr_pipeline: Pipeline for synchronizing with correction operations + :type si_corr_pipeline: pipeline.PipelineAsync + :param s0_s1_sequence_pipeline: Pipeline for synchronizing between stage 0 and 1 + :type s0_s1_sequence_pipeline: pipeline.PipelineAsync + :param tile_sched_params: Parameters for tile scheduling + :type tile_sched_params: fmha_utils.FmhaStaticTileSchedulerParams + :param fused_mask: Compute trip counts and apply masking for attention blocks + :type fused_mask: fmha_utils.FusedMask + """ + tidx, _, _ = cute.arch.thread_idx() + thread_idx = tidx % (self.threads_per_warp * (len(self.softmax0_warp_ids) if stage == 0 else len(self.softmax1_warp_ids))) + + cS_base = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) + tilePlikeFP32 = self.qk_mma_tiler[1] // 32 * self.o_dtype.width + tScS = qk_thr_mma.partition_C(cS_base) + tStS_vec_layout = cute.composition(tStS.layout, cute.make_layout((128, 2))) + tmem_vec_offset = self.tmem_vec0_offset if stage == 0 else self.tmem_vec1_offset + tStS_vec = cute.make_tensor(tStS.iterator + tmem_vec_offset, tStS_vec_layout) + tScS_vec_layout = cute.composition(tScS.layout, cute.make_layout((128, 2))) + tScS_vec = cute.make_tensor(tScS.iterator, tScS_vec_layout) + tStS_P_layout = cute.composition(tStS.layout, cute.make_layout((128, tilePlikeFP32))) + tmem_p_offset = self.tmem_p0_offset if stage == 0 else self.tmem_p1_offset + tStS_P = cute.make_tensor(tStS.iterator + tmem_p_offset, tStS_P_layout) + tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), + self.qk_acc_dtype, + ) + tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStSi) + thread_idx = tidx % (self.threads_per_warp * (len(self.softmax0_warp_ids) if stage == 0 else len(self.softmax1_warp_ids))) + thr_tmem_load = tiled_tmem_load.get_slice(thread_idx) + tTMEM_LOADtS = thr_tmem_load.partition_S(tStSi) + tmem_store_vec_atom = cute.make_copy_atom( + tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(2)), + self.qk_acc_dtype, + ) + tiled_tmem_store_vec = tcgen05.make_tmem_copy(tmem_store_vec_atom, tStS_vec) + thr_tmem_store_vec = tiled_tmem_store_vec.get_slice(thread_idx) + tTMEM_STORE_VECtS = thr_tmem_store_vec.partition_D(tStS_vec) + tTMEM_STORE_VECcS = thr_tmem_store_vec.partition_S(tScS_vec) + tmem_store_atom = cute.make_copy_atom( + tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), + self.qk_acc_dtype, + ) + tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStS_P) + thr_tmem_store = tiled_tmem_store.get_slice(thread_idx) + tTMEM_STOREtS_x4 = thr_tmem_store.partition_D(tStS_P) + + tile_sched = fmha_utils.create_fmha_static_tile_scheduler(tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()) + work_tile = tile_sched.initial_work_tile_info() + + while work_tile.is_valid_tile: + curr_block_coord = work_tile.tile_idx + batch_coord = curr_block_coord[2][1] + seqlen_k_ = seqlen_k + seqlen_q_ = seqlen_q + continue_cond = False + cuseqlen_q = Int32(0) + seqlen_q_ = seqlen_q + if cutlass.const_expr(cum_seqlen_q is not None): + cuseqlen_q = cum_seqlen_q[batch_coord] + seqlen_q_ = cum_seqlen_q[batch_coord + 1] - cuseqlen_q + continue_cond = not fmha_utils.FmhaStaticTileScheduler.check_valid_work_for_seqlen_q( + self.cta_tiler[0], + curr_block_coord[0], + seqlen_q_, + ) + + if not continue_cond: + if cutlass.const_expr(cum_seqlen_k is not None): + cuseqlen_k = cum_seqlen_k[batch_coord] + seqlen_k_ = cum_seqlen_k[batch_coord + 1] - cuseqlen_k + if cutlass.const_expr(self.zero_reference_max): + row_max = Float32(0.0) + observed_row_max = Float32(0.0) + else: + row_max = -Float32.inf + observed_row_max = -Float32.inf + row_sum = 0.0 + value_args = ( + seqlen_k_, + seqlen_q_, + scale_softmax_log2, + window_size_left, + window_size_right, + ) + atom_args = ( + qk_thr_mma, + tiled_tmem_load, + tiled_tmem_store, + tiled_tmem_store_vec, + thr_tmem_load, + thr_tmem_store, + thr_tmem_store_vec, + ) + tensor_args = ( + tTMEM_LOADtS, + tTMEM_STORE_VECtS, + tTMEM_STOREtS_x4, + ) + + logical_offset = ( + curr_block_coord[0] * self.cta_tiler[0] + stage * self.qk_mma_tiler[0], + 0, + ) + cS = cute.domain_offset(logical_offset, cS_base) + vec_i_handle = si_corr_producer.acquire_and_advance() + + sparse_trailing_mask_count = Int32(0) + if cutlass.const_expr(sparse_block_count is not None): + head_coord = curr_block_coord[2][0] + curr_sparse_block_indices = sparse_block_indices[batch_coord, head_coord, curr_block_coord[0], None] + start_count = Int32(0) + sparse_count = sparse_block_count[batch_coord, head_coord, curr_block_coord[0]] + leading_mask_count = Int32(0) + if seqlen_k_ % self.qk_mma_tiler[1] != 0: + last_sparse_block = curr_sparse_block_indices[sparse_count - 1] + if last_sparse_block == (cute.ceil_div(seqlen_k_, self.qk_mma_tiler[1]) - 1): + sparse_trailing_mask_count = Int32(1) + else: + start_count = fmha_utils.FusedMask.get_trip_start( + self.mask_type, + curr_block_coord, + self.cta_tiler, + seqlen_q_, + seqlen_k_, + window_size_left, + ) + + leading_mask_count = fmha_utils.FusedMask.get_masked_leading_count( + self.mask_type, + curr_block_coord, + self.cta_tiler, + seqlen_q_, + seqlen_k_, + window_size_left, + window_size_right, + ) + for i in cutlass.range(start_count, start_count + leading_mask_count, 1, unroll=1): + cS_iter = cute.domain_offset((0, i * self.qk_mma_tiler[1]), cS) + update_row_max = True + if cutlass.const_expr(not self.zero_reference_max and self.max_update_interval > 1): + update_row_max = (i & (self.max_update_interval - 1)) == 0 + iter_args = ( + cS_iter, + row_max, + observed_row_max, + row_sum, + update_row_max, + vec_i_handle, + ) + pipeline_args = ( + mma_si_consumer, + si_corr_producer, + s0_s1_sequence_consumer, + s0_s1_sequence_producer, + ) + ( + row_max, + observed_row_max, + row_sum, + vec_i_handle, + mma_si_consumer, + si_corr_producer, + s0_s1_sequence_consumer, + s0_s1_sequence_producer, + ) = self.softmax_step( + stage, + True, + iter_args, + value_args, + pipeline_args, + atom_args, + tensor_args, + ) + if cutlass.const_expr(sparse_block_count is not None): + unmask_count = sparse_count - sparse_trailing_mask_count + else: + unmask_count = fmha_utils.FusedMask.get_unmasked_trip_count( + self.mask_type, + curr_block_coord, + self.cta_tiler, + seqlen_q_, + seqlen_k_, + window_size_left, + window_size_right, + ) + for i in cutlass.range( + start_count + leading_mask_count, + start_count + leading_mask_count + unmask_count, + 1, + unroll=1, + ): + n_block = i + if cutlass.const_expr(sparse_block_count is not None): + n_block = curr_sparse_block_indices[i] + cS_iter = cute.domain_offset((0, n_block * self.qk_mma_tiler[1]), cS) + update_row_max = True + if cutlass.const_expr(not self.zero_reference_max and self.max_update_interval > 1): + update_row_max = (i & (self.max_update_interval - 1)) == 0 + iter_args = ( + cS_iter, + row_max, + observed_row_max, + row_sum, + update_row_max, + vec_i_handle, + ) + pipeline_args = ( + mma_si_consumer, + si_corr_producer, + s0_s1_sequence_consumer, + s0_s1_sequence_producer, + ) + ( + row_max, + observed_row_max, + row_sum, + vec_i_handle, + mma_si_consumer, + si_corr_producer, + s0_s1_sequence_consumer, + s0_s1_sequence_producer, + ) = self.softmax_step( + stage, + False, + iter_args, + value_args, + pipeline_args, + atom_args, + tensor_args, + ) + if cutlass.const_expr(sparse_block_count is not None): + trailing_mask_count = sparse_trailing_mask_count + else: + trailing_mask_count = fmha_utils.FusedMask.get_masked_trailing_count( + self.mask_type, + curr_block_coord, + self.cta_tiler, + seqlen_q_, + seqlen_k_, + window_size_left, + window_size_right, + ) + + for i in cutlass.range( + start_count + leading_mask_count + unmask_count, + start_count + leading_mask_count + unmask_count + trailing_mask_count, + 1, + unroll=1, + ): + n_block = i + if cutlass.const_expr(sparse_block_count is not None): + n_block = curr_sparse_block_indices[i] + cS_iter = cute.domain_offset((0, n_block * self.qk_mma_tiler[1]), cS) + update_row_max = True + if cutlass.const_expr(not self.zero_reference_max and self.max_update_interval > 1): + update_row_max = (i & (self.max_update_interval - 1)) == 0 + iter_args = ( + cS_iter, + row_max, + observed_row_max, + row_sum, + update_row_max, + vec_i_handle, + ) + pipeline_args = ( + mma_si_consumer, + si_corr_producer, + s0_s1_sequence_consumer, + s0_s1_sequence_producer, + ) + ( + row_max, + observed_row_max, + row_sum, + vec_i_handle, + mma_si_consumer, + si_corr_producer, + s0_s1_sequence_consumer, + s0_s1_sequence_producer, + ) = self.softmax_step( + stage, + True, + iter_args, + value_args, + pipeline_args, + atom_args, + tensor_args, + ) + si_handle = mma_si_consumer.wait_and_advance() + tTMEM_STORE_VECrS = cute.make_rmem_tensor(tTMEM_STORE_VECcS.shape, self.qk_acc_dtype) + tTMEM_STORE_VECrS[0] = row_sum + tTMEM_STORE_VECrS[1] = row_max + cute.copy(tiled_tmem_store_vec, tTMEM_STORE_VECrS, tTMEM_STORE_VECtS) + cute.arch.fence_view_async_tmem_store() + vec_i_handle.commit() + si_corr_producer.acquire() + # Empty step to sync against pipe s + si_handle.release() + + # Advance to next tile + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + # End of persistent scheduler loop + + @cute.jit + def correction_rescale( + self, + thr_mma: cute.ThrMma, + tOtO: cute.Tensor, + scale: Float32, + ): + """Rescale intermediate attention results based on softmax normalization factor. + + This method performs a crucial correction step in the attention computation pipeline. + When processing attention in blocks, the softmax normalization factors may change + as new blocks are processed. This method rescales previously computed partial + output values to account for updated normalization factors. + + The implementation uses efficient tensor memory operations to: + 1. Load existing partial attention output from tensor memory + 2. Apply the scaling factor to all elements + 3. Store the rescaled results back to tensor memory + + :param thr_mma: Thread MMA operation for the computation + :type thr_mma: cute.ThrMma + :param tOtO: Tensor representing partial attention output to be rescaled + :type tOtO: cute.Tensor + :param scale: Scaling factor to apply to the partial results + :type scale: Float32 + """ + pv_tiled_mma_shape = ( + self.pv_mma_tiler[0], + self.pv_mma_tiler[1], + ) + cO = cute.make_identity_tensor(pv_tiled_mma_shape) + tOcO = thr_mma.partition_C(cO) + + corr_tile_size = 16 # tuneable parameter + tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), + self.pv_acc_dtype, + ) + tmem_store_atom = cute.make_copy_atom( + tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), + self.pv_acc_dtype, + ) + + tOtO_i_layout = cute.composition(tOtO.layout, cute.make_layout((128, corr_tile_size))) + tOcO_i_layout = cute.composition(tOcO.layout, cute.make_layout((128, corr_tile_size))) + + tOtO_i = cute.make_tensor(tOtO.iterator, tOtO_i_layout) + tOcO_i = cute.make_tensor(tOcO.iterator, tOcO_i_layout) + + tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tOtO_i) + tiled_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tOtO_i) + tidx, _, _ = cute.arch.thread_idx() + thread_idx = tidx % (self.threads_per_warp * len(self.correction_warp_ids)) + thr_tmem_load = tiled_tmem_load.get_slice(thread_idx) + thr_tmem_store = tiled_tmem_store.get_slice(thread_idx) + + tTMEM_LOADtO = thr_tmem_load.partition_S(tOtO_i) + tTMEM_LOADcO = thr_tmem_load.partition_D(tOcO_i) + + tTMEM_STOREtO = thr_tmem_store.partition_D(tOtO_i) + + tTMrO = cute.make_rmem_tensor((tTMEM_LOADcO.shape, 128 // corr_tile_size), self.pv_acc_dtype) + for i in range(self.cta_tiler[2] // corr_tile_size): + tTMrO_i_ = tTMrO[None, i] + tTMrO_i_layout = cute.composition(tTMrO_i_.layout, cute.make_layout(tTMrO.shape[0])) + tTMrO_i = cute.make_tensor(tTMrO_i_.iterator, tTMrO_i_layout) + tTMEM_LOADtO_i = cute.make_tensor(tTMEM_LOADtO.iterator + i * corr_tile_size, tTMEM_LOADtO.layout) + tTMEM_STOREtO_i = cute.make_tensor(tTMEM_STOREtO.iterator + i * corr_tile_size, tTMEM_STOREtO.layout) + + cute.copy(tiled_tmem_load, tTMEM_LOADtO_i, tTMrO_i) + for j in cutlass.range(cute.size(tTMrO_i), vectorize=True): + tTMrO_i[j] = tTMrO_i[j] * scale + cute.copy(tiled_tmem_store, tTMrO_i, tTMEM_STOREtO_i) + + @cute.jit + def correction_epilog( + self, + thr_mma: cute.ThrMma, + tOtO: cute.Tensor, + mLSE: Optional[cute.Tensor], + tTMEM_LOAD_VECrS: cute.Tensor, + row_idx: Int32, + cuseqlen_q: Int32, + seqlen_q: Int32, + blk_coord: Int32, + scale_softmax: Float32, + scale: Float32, + sO: cute.Tensor, + ): + """Apply final scaling and transformation to attention output before writing to global memory. + + This correction_epilog function handles the final processing step for attention output values. + It applies a scaling factor to the accumulated attention results and prepares the + data for efficient transfer back to global memory. + + The method performs: + 1. Loading of accumulated attention results from tensor memory + 2. Application of the final output scaling factor + 3. Type conversion if necessary (typically from higher precision accumulator to output precision) + 4. Reorganization of data for optimal memory access patterns + 5. Preparation for efficient TMA store operations + + :param thr_mma: Thread MMA operation for the computation + :type thr_mma: cute.ThrMma + :param tOtO: Tensor containing accumulated attention output + :type tOtO: cute.Tensor + :param mLSE: Tensor containing log-sum-exp values for LSE calculation + :type mLSE: cute.Tensor | None + :param tTMEM_LOAD_VECrS: Tensor containing row sum and max values for softmax calculation + :type tTMEM_LOAD_VECrS: cute.Tensor + :param row_idx: Index of the current row being processed + :type row_idx: Int32 + :param cuseqlen_q: Cumulative sequence length of the current query + :type cuseqlen_q: Int32 + :param seqlen_q: Sequence length of the current query + :type seqlen_q: Int32 + :param blk_coord: Coordinate of the current block being processed + :type blk_coord: Int32 + :param scale_softmax: Scaling factor for softmax calculation + :type scale_softmax: Float32 + :param scale: Final scaling factor to apply to the output + :type scale: Float32 + :param sO: Shared memory tensor for the final output + :type sO: cute.Tensor + """ + + pv_tiled_mma_shape = ( + self.pv_mma_tiler[0], + self.pv_mma_tiler[1], + ) + cO = cute.make_identity_tensor(pv_tiled_mma_shape) + + corr_tile_size = 32 * 8 // self.o_dtype.width + tOsO = thr_mma.partition_C(sO) + tOcO = thr_mma.partition_C(cO) + + tOtO_i = cute.logical_divide(tOtO, cute.make_layout((128, corr_tile_size))) + tOcO_i = cute.logical_divide(tOcO, cute.make_layout((128, corr_tile_size))) + tOsO_i = cute.logical_divide(tOsO, cute.make_layout((128, corr_tile_size))) + tidx, _, _ = cute.arch.thread_idx() + thread_idx = tidx % (self.threads_per_warp * len(self.correction_warp_ids)) + + epi_subtile = (self.epi_tile[0], corr_tile_size) + tmem_copy_atom = sm100_utils.get_tmem_load_op( + self.pv_mma_tiler, + self.o_layout, + self.o_dtype, + self.pv_acc_dtype, + epi_subtile, + use_2cta_instrs=False, + ) + + tiled_tmem_load = tcgen05.make_tmem_copy(tmem_copy_atom, tOtO_i[(None, None), 0]) + + thr_tmem_load = tiled_tmem_load.get_slice(thread_idx) + smem_copy_atom = sm100_utils.get_smem_store_op(self.o_layout, self.o_dtype, self.pv_acc_dtype, tiled_tmem_load) + tiled_smem_store = cute.make_tiled_copy_D(smem_copy_atom, tiled_tmem_load) + + tTMEM_LOADtO = thr_tmem_load.partition_S(tOtO_i[(None, None), None]) + tTMEM_LOADsO = thr_tmem_load.partition_D(tOsO_i[(None, None), None]) + tTMEM_LOADoO = thr_tmem_load.partition_D(tOcO_i[(None, None), None]) + + for i in range(self.cta_tiler[2] // corr_tile_size): + tTMEM_LOADtO_i = tTMEM_LOADtO[None, 0, 0, i] + tTMEM_LOADsO_i = tTMEM_LOADsO[None, 0, 0, i] + tTMrO = cute.make_rmem_tensor(tTMEM_LOADoO[None, 0, 0, i].shape, self.pv_acc_dtype) + cute.copy(tiled_tmem_load, tTMEM_LOADtO_i, tTMrO) + for j in range(cute.size(tTMrO), vectorize=True): + tTMrO[j] = tTMrO[j] * scale + tSMrO = cute.make_rmem_tensor(tTMrO.shape, self.o_dtype) + o_vec = tTMrO.load() + tSMrO.store(o_vec.to(self.o_dtype)) + cute.copy(tiled_smem_store, tSMrO, tTMEM_LOADsO_i) + + if cutlass.const_expr(mLSE is not None): + scaled_tmp = scale_softmax * tTMEM_LOAD_VECrS[1] + lse = cute.math.log(tTMEM_LOAD_VECrS[0], fastmath=True) + scaled_tmp - self.softmax_prescale_ln + if row_idx < seqlen_q: + mLSE[row_idx + cuseqlen_q, blk_coord[2]] = lse + + # fence view async shared + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + + +def run( + q_shape: Union[Tuple[int, int, int, int], Tuple[int, Tuple[int, ...], int, int]], + k_shape: Union[Tuple[int, int, int, int], Tuple[int, Tuple[int, ...], int, int]], + in_dtype: Type[cutlass.Numeric], + out_dtype: Type[cutlass.Numeric], + qk_acc_dtype: Type[cutlass.Numeric], + pv_acc_dtype: Type[cutlass.Numeric], + mma_tiler_mn: Tuple[int, int], + is_persistent: bool, + is_causal: bool, + bottom_right_align: bool, + lse_calculation: bool, + window_size: Tuple[int, int], + scale_q: float, + scale_k: float, + scale_v: float, + inv_scale_o: float, + scale_softmax: float, + tolerance: float, + warmup_iterations: int, + iterations: int, + skip_ref_check: bool, + use_cold_l2: bool = False, + output_dir: str = "./fmha_aot_artifacts", + export_only: bool = False, + file_name: str = "fmha", + function_prefix: str = "fmha", + vit_mode: bool = False, + **kwargs, +): + """Execute Fused Multi-Head Attention (FMHA) on Blackwell architecture and validate results. + + This function creates random input tensors for query, key, and value, then performs the + complete FMHA computation pipeline. It supports configurable data types, tiling parameters, + and various attention masking options. Results can be validated against a PyTorch reference + implementation or run multiple times for performance measurement. + + The implementation leverages specialized tensor memory operations and efficient math + operations optimized for Blackwell architecture, including pipelined computation stages + for maximum throughput. + + :param q_shape: Query tensor shape (B, S_q, H, D) where B=batch size, S_q=query sequence length, + H=number of heads, D=head dimension. + If S_q is a tuple, it is the variable sequence length. + :type q_shape: Union[Tuple[int, int, int, int], Tuple[int, Tuple[int, ...], int, int]] + :param k_shape: Key tensor shape (B, S_k, H_k, D) where B=batch size, S_k=key sequence length, + H_k=number of key heads (H must be divisible by H_k), D=head dimension. + If S_k is a tuple, it is the variable sequence length. + :type k_shape: Union[Tuple[int, int, int, int], Tuple[int, Tuple[int, ...], int, int]] + :param in_dtype: Input data type for query, key and value tensors + :type in_dtype: Type[cutlass.Numeric] + :param out_dtype: Output data type for attention output + :type out_dtype: Type[cutlass.Numeric] + :param qk_acc_dtype: Accumulator data type for query-key matrix multiplication + :type qk_acc_dtype: Type[cutlass.Numeric] + :param pv_acc_dtype: Accumulator data type for probability-value matrix multiplication + :type pv_acc_dtype: Type[cutlass.Numeric] + :param mma_tiler_mn: Matrix multiply accumulate tile shape (M, N) + :type mma_tiler_mn: Tuple[int, int] + :param is_persistent: Whether to use persistent kernel optimization + :type is_persistent: bool + :param is_causal: Whether to apply causal masking + :type is_causal: bool + :param lse_calculation: Whether to calculate lse + :type lse_calculation: bool + :param window_size: Sliding window size (left, right) for attention masking. Controls which positions each query can attend to. + :type window_size: Tuple[int, int] + :param scale_q: Scaling factor for query tensor + :type scale_q: float + :param scale_k: Scaling factor for key tensor + :type scale_k: float + :param scale_v: Scaling factor for value tensor + :type scale_v: float + :param inv_scale_o: Inverse scaling factor for output tensor + :type inv_scale_o: float + :param scale_softmax: Attention score scaling factor (defaults to 1/sqrt(D) if set to 0) + :type scale_softmax: float + :param tolerance: Maximum acceptable error for validation + :type tolerance: float + :param warmup_iterations: Number of warmup iterations + :type warmup_iterations: int + :param iterations: Number of iterations to run for performance testing + :type iterations: int + :param skip_ref_check: Skip validation against reference implementation + :type skip_ref_check: bool + :param use_cold_l2: Whether to use circular buffer strategy to ensure cold L2 cache + :type use_cold_l2: bool + + :raises ValueError: If input shapes are incompatible or head dimension is unsupported + :raises RuntimeError: If GPU is unavailable for computation + :return: Execution time of the FMHA kernel in microseconds + :rtype: float + """ + + # Use file_name as tag so parallel process output is identifiable + _tag = f"[{file_name}]" + + if export_only: + print( + f"{_tag} Compiling: head_dim={q_shape[-1]}, in_dtype={in_dtype}, " + f"causal={is_causal}, window={window_size}, " + f"mma_tiler_mn={mma_tiler_mn}, persistent={is_persistent}, " + f"bottom_right_align={bottom_right_align}, " + f"sliding_window={window_size[0] != -1}" + ) + else: + print(f"{_tag} Running Blackwell SM100 FMHA test with:") + print(f"{_tag} q_shape={q_shape}, k_shape={k_shape}") + print(f"{_tag} in_dtype={in_dtype}, out_dtype={out_dtype}") + print(f"{_tag} qk_acc_dtype={qk_acc_dtype}, pv_acc_dtype={pv_acc_dtype}") + print(f"{_tag} mma_tiler_mn={mma_tiler_mn}, is_persistent={is_persistent}") + print(f"{_tag} is_causal={is_causal}, window_size={window_size}") + print(f"{_tag} tolerance={tolerance}, warmup={warmup_iterations}, iterations={iterations}") + + # ---- CuPy/NumPy helpers (replacing cutlass.torch) ---- + def _cutlass_to_cupy_dtype(cutlass_dtype): + if cutlass_dtype == cutlass.Float16: + return cp.float16 + elif cutlass_dtype in (cutlass.Float32, Float32): + return cp.float32 + elif (cutlass_dtype.is_float and cutlass_dtype.width <= 8) or (cutlass_dtype.is_integer and cutlass_dtype.width == 4): + return cp.uint8 # FP8/Int4 stored as bytes + else: + raise ValueError(f"Unsupported dtype for CuPy: {cutlass_dtype}") + + def _get_leading_dim(cp_array): + for i, s in enumerate(cp_array.strides): + if s == cp_array.itemsize: + return i + return len(cp_array.shape) - 1 + + # Unpack parameters + b, s_q, h_q, d = q_shape + b_, s_k, h_k, d_ = k_shape + window_size_left, window_size_right = window_size + if window_size_left == -1: + window_size_left = None + if window_size_right == -1: + window_size_right = None + if is_causal: + window_size_right = 0 + + if b != b_: + raise ValueError("q & k must have the same batch size") + + if d != d_: + raise ValueError("q & k must have the same head dimension") + + # if d not in {32, 64, 128}: + # raise ValueError("head dimension must be 32, 64, or 128") + + if h_q % h_k != 0: + raise ValueError("h_q must be divisible by h_k") + + if isinstance(s_q, tuple) and len(s_q) != b: + raise ValueError("variable_seqlen s_q must have the length of batch size") + if isinstance(s_k, tuple) and len(s_k) != b: + raise ValueError("variable_seqlen s_k must have the length of batch size") + + if in_dtype not in {cutlass.Float8E4M3FN, cutlass.Float16}: + raise ValueError("in_dtype must be Float8E4M3FN or Float16") + + if out_dtype not in {cutlass.Float8E4M3FN, cutlass.Float16}: + raise ValueError("out_dtype must be Float8E4M3FN or Float16") + + if qk_acc_dtype not in {Float32}: + raise ValueError("qk_acc_dtype must be Float32") + + if pv_acc_dtype not in {Float32}: + raise ValueError("pv_acc_dtype must be Float32") + + if iterations < 1: + raise ValueError("iterations must be at least 1") + + h_r = h_q // h_k + + # Prepare GPU tensors: Q, KV cache, O + if cp.cuda.runtime.getDeviceCount() == 0: + raise RuntimeError("GPU is required to run this example!") + + if not export_only: + cp.random.seed(1111) + np.random.seed(1111) + + if isinstance(s_q, tuple) or isinstance(s_k, tuple): + raise NotImplementedError("Variable-length sequences (nested tensors) require PyTorch. Use fmha_runtimeargs_kvcache.py for variable-length support.") + + def create_and_pad_tensor(shape, padding, dtype, is_dynamic_layout=True): + shape_ = tuple(map(lambda x, y: x + y, shape, padding)) + + if export_only: + f32_gpu_full = cp.zeros(shape_, dtype=cp.float32) + else: + min_val = -2 if dtype.is_float or dtype.signed else 0 + f32_gpu_full = cp.random.randint(min_val, 2, shape_).astype(cp.float32) + + # Create dtype GPU buffer and initialize + cp_dtype = _cutlass_to_cupy_dtype(dtype) + is_narrow = (dtype.is_float and dtype.width <= 8) or (dtype.is_integer and dtype.width == 4) + dtype_gpu_full = cp.empty(shape_, dtype=cp_dtype) + + if is_narrow: + # FP8/Int4: use cute.testing.convert + f32_cute = from_dlpack(f32_gpu_full) + if is_dynamic_layout: + f32_cute = f32_cute.mark_layout_dynamic(leading_dim=_get_leading_dim(f32_gpu_full)) + dtype_cute_full = from_dlpack(dtype_gpu_full, assumed_align=16) + dtype_cute_full.element_type = dtype + if is_dynamic_layout: + dtype_cute_full = dtype_cute_full.mark_layout_dynamic(leading_dim=_get_leading_dim(dtype_gpu_full)) + cute.testing.convert(f32_cute, dtype_cute_full) + else: + dtype_gpu_full[:] = f32_gpu_full.astype(cp_dtype) + + # Offset the tensor (slice into padded region) + slices = tuple(slice(s, e) for s, e in zip(padding, shape_)) + dtype_gpu = dtype_gpu_full[slices] + f32_gpu = f32_gpu_full[slices] + + # Create cute tensor from sliced GPU buffer + cute_tensor = from_dlpack(dtype_gpu, assumed_align=16) + cute_tensor.element_type = dtype + + # f32 reference on CPU (numpy) for comparison + f32_ref = f32_gpu.get() + + # Return full buffers too to prevent GC + return (f32_ref, cute_tensor, dtype_gpu, dtype_gpu_full, f32_gpu_full) + + qo_shape = (b, s_q, h_r * h_k, d) + kvcache_shape = (b, 2, h_k, s_k, d) + lse_shape = (b, h_r * h_k, s_q) + qo_padding = (0, 0, 0, 0, 0) + kvcache_padding = (0, 0, 0, 0, 0, 0) + lse_padding = (0, 0, 0, 0) + + q_ref, q_tensor, q_cp, *_q_keep = create_and_pad_tensor( + qo_shape, + qo_padding, + in_dtype, + is_dynamic_layout=True, + ) + kvcache_ref, kvcache_tensor, kvcache_cp, *_kv_keep = create_and_pad_tensor( + kvcache_shape, + kvcache_padding, + in_dtype, + is_dynamic_layout=True, + ) + _, o_tensor, o_cp, *_o_keep = create_and_pad_tensor( + qo_shape, + qo_padding, + out_dtype, + is_dynamic_layout=True, + ) + if lse_calculation: + _, lse_tensor, lse_cp, *_lse_keep = create_and_pad_tensor( + lse_shape, + lse_padding, + cutlass.Float32, + is_dynamic_layout=True, + ) + else: + lse_cp = None + + # SM100 tcgen05.mma atom K = 256 bits / element_bits. For fp16: 16 elems. + # The MMA tiler K must be a multiple of this atom. Non-aligned dims like + # 72 are padded up (→ 80); TMA ZFILL/OOB-drop bridge the gap at zero cost. + _MMA_K_ATOM = 256 // 16 # 16 for fp16 + padded_d = ((d + _MMA_K_ATOM - 1) // _MMA_K_ATOM) * _MMA_K_ATOM + actual_head_dim = d if padded_d != d else None + if actual_head_dim is not None: + print(f"[fmha] head_dim {d} not MMA-aligned; tiler K padded to {padded_d}, tensors stay at {d}") + + mma_tiler = (*mma_tiler_mn, padded_d) + + mask_type = fmha_utils.MaskEnum.WINDOW_MASK + if bottom_right_align: + mask_type = fmha_utils.MaskEnum.WINDOW_MASK_INFERENCE + # Note: window_size_right is always 0 (causal), so window/causal masking is + # always active. RESIDUAL_MASK fallback (no masking) is not reachable. + + s_q_list = s_q if isinstance(s_q, tuple) else [s_q] * b + s_k_list = s_k if isinstance(s_k, tuple) else [s_k] * b + + # To avoid mask out the whole row which results in NaN in softmax + def check_seqlen_valid(s_q, s_k, window_size_left, window_size_right, bottom_right_align): + for i in range(s_q): + offset = 0 if not bottom_right_align else s_k - s_q + + s_q_start = 0 if window_size_left is None else i + offset - window_size_left + s_q_end = s_q if window_size_right is None else i + offset + window_size_right + s_q_min = max(s_q_start, 0) + s_q_max = min(s_q_end, s_k) + + if s_q_max - s_q_min == 0 and (i != 0 and i != s_q - 1): + return False + return True + + need_check_seqlen_valid = window_size_left is not None or window_size_right is not None + for i in range(b): + if need_check_seqlen_valid and not check_seqlen_valid( + s_q_list[i], + s_k_list[i], + window_size_left, + window_size_right, + bottom_right_align, + ): + raise ValueError("sliding window doesn't support current setting") + + use_sliding_window = window_size_left is not None + if vit_mode: + mask_type = fmha_utils.MaskEnum.RESIDUAL_MASK + fmha = BlackwellFusedMultiHeadAttentionForward( + qk_acc_dtype, + pv_acc_dtype, + mma_tiler, + is_persistent, + mask_type, + is_causal=(is_causal and not vit_mode), + use_sliding_window=(use_sliding_window and not vit_mode), + actual_head_dim=actual_head_dim, + ) + + # Initialize Stream + current_stream = cuda.CUstream(cp.cuda.get_current_stream().ptr) + + # Compute folded scales for numpy reference and ViT path. + # The LLM __call__ computes these internally from the raw per-tensor scales. + if scale_softmax == 0.0: # default to 1/sqrt(d) + scale_softmax = 1.0 / math.sqrt(d) + log2_e = math.log2(math.exp(1.0)) + + ref_scale_softmax = scale_q * scale_k * scale_softmax + ref_scale_softmax_log2 = ref_scale_softmax * log2_e + ref_scale_output = scale_v * inv_scale_o + + def mark_bshd_dynamic(tensor): + so = (0, 1, 2, 3) # outermost-to-innermost for contiguous BSHD + return ( + tensor.mark_layout_dynamic(leading_dim=3) + .mark_compact_shape_dynamic(mode=0, stride_order=so) + .mark_compact_shape_dynamic(mode=1, stride_order=so) + .mark_compact_shape_dynamic(mode=2, stride_order=so) + ) + + def mark_shd_dynamic(tensor): + so = (0, 1, 2) # outermost-to-innermost for packed (total_S, H, D) + return tensor.mark_layout_dynamic(leading_dim=2).mark_compact_shape_dynamic(mode=0, stride_order=so).mark_compact_shape_dynamic(mode=1, stride_order=so) + + def mark_kv_cache_dynamic(tensor): + so = (0, 1, 2, 3, 4) # outermost-to-innermost for contiguous (B,2,H,S,D) + return ( + tensor.mark_layout_dynamic(leading_dim=4) + .mark_compact_shape_dynamic(mode=0, stride_order=so) # B + .mark_compact_shape_dynamic(mode=2, stride_order=so) # H_kv + .mark_compact_shape_dynamic(mode=3, stride_order=so) # S + ) + + def mark_1d_dynamic(tensor): + return tensor.mark_layout_dynamic(leading_dim=0).mark_compact_shape_dynamic(mode=0, stride_order=(0,)) + + if vit_mode: + # ViT: packed [total_S, H, D] with cu_seqlens for ragged batching. + # For the reference test, total_S = b * s_q, uniform lengths. + _s = s_q if not isinstance(s_q, tuple) else max(s_q) + total_S = b * _s + cu_seqlens_np = np.arange(b + 1, dtype=np.int32) * _s + + # Reshape Q, K, V from (B, S, H, D) → (total_S, H, D) + q_vit_shape = (total_S, h_r * h_k, d) + k_vit_shape = (total_S, h_k, d) + q_vit_ref, q_vit_tensor, q_vit_cp, *_qv = create_and_pad_tensor(q_vit_shape, (0, 0, 0, 0), in_dtype, is_dynamic_layout=True) + k_vit_ref, k_vit_tensor, k_vit_cp, *_kv = create_and_pad_tensor(k_vit_shape, (0, 0, 0, 0), in_dtype, is_dynamic_layout=True) + v_vit_ref, v_vit_tensor, v_vit_cp, *_vv = create_and_pad_tensor(k_vit_shape, (0, 0, 0, 0), in_dtype, is_dynamic_layout=True) + _, o_vit_tensor, o_vit_cp, *_ov = create_and_pad_tensor(q_vit_shape, (0, 0, 0, 0), out_dtype, is_dynamic_layout=True) + + cu_seqlens_cp = cp.asarray(cu_seqlens_np) + cu_seqlens = from_dlpack(cu_seqlens_cp, assumed_align=16) + + q_dyn = mark_shd_dynamic(q_vit_tensor) + k_dyn = mark_shd_dynamic(k_vit_tensor) + v_dyn = mark_shd_dynamic(v_vit_tensor) + o_dyn = mark_shd_dynamic(o_vit_tensor) + cu_dyn = mark_1d_dynamic(cu_seqlens) + + _max_seqlen = Int32(_s) + + start_time = time.time() + compiled_fmha = cute.compile( + fmha.__call_vit__, + q_dyn, + k_dyn, + v_dyn, + o_dyn, + cu_dyn, + _max_seqlen, + ref_scale_softmax_log2, + ref_scale_softmax, + ref_scale_output, + current_stream, + ) + else: + # LLM: batched Q [B,S,H,D] + combined KV cache [B,2,H,Cap,D] + q_dyn = mark_bshd_dynamic(q_tensor) + kv_dyn = mark_kv_cache_dynamic(kvcache_tensor) + o_dyn = mark_bshd_dynamic(o_tensor) + + _wsl = Int32(window_size_left) if window_size_left is not None else Int32(0) + + _s_k = s_k if not isinstance(s_k, tuple) else max(s_k) + cu_kv_seqlens_np = np.arange(b + 1, dtype=np.int32) * _s_k + cu_kv_seqlens_cp = cp.asarray(cu_kv_seqlens_np) + cu_kv_seqlens = from_dlpack(cu_kv_seqlens_cp, assumed_align=16) + cu_kv_seqlens = mark_1d_dynamic(cu_kv_seqlens) + + start_time = time.time() + compiled_fmha = cute.compile( + fmha, + q_dyn, + kv_dyn, + o_dyn, + cu_kv_seqlens, + _wsl, + scale_q, + scale_k, + scale_v, + inv_scale_o, + current_stream, + ) + + compilation_time = time.time() - start_time + print(f"{_tag} Compilation time: {compilation_time:.4f}s") + + if export_only: + os.makedirs(output_dir, exist_ok=True) + compiled_fmha.export_to_c( + file_path=output_dir, + file_name=file_name, + function_prefix=function_prefix, + ) + print(f"{_tag} Exported to {output_dir}/{file_name}.h and {file_name}.o") + return None + + def _numpy_softmax(x, axis=-1): + x_max = np.max(x, axis=axis, keepdims=True) + # Handle rows that are all -inf + x_max = np.where(np.isfinite(x_max), x_max, 0.0) + e_x = np.exp(x - x_max) + s = np.sum(e_x, axis=axis, keepdims=True) + s = np.where(s == 0, 1.0, s) + return e_x / s + + def _numpy_logsumexp(x, axis=-1): + x_max = np.max(x, axis=axis) + x_max_safe = np.where(np.isfinite(x_max), x_max, 0.0) + return np.log(np.sum(np.exp(x - x_max_safe[..., np.newaxis]), axis=axis)) + x_max + + def run_numpy_single_shot_reference_packed( + q_packed, + k_packed, + v_packed, + cu_seqlens_q, + cu_seqlens_k, + scale_softmax=1.0, + scale_output=1.0, + is_causal=False, + bottom_right_align=False, + lse_calculation=False, + window_size_left=None, + window_size_right=None, + ): + """Packed (ViT-style) numpy reference for single-shot attention. + + q_packed: [total_q, H_q, D] + k_packed/v_packed: [total_k, H_k, D] + cu_seqlens_q/cu_seqlens_k: cumulative offsets of length B+1. + """ + h_q_local = q_packed.shape[1] + h_k_local = k_packed.shape[1] + if h_q_local % h_k_local != 0: + raise ValueError("H_q must be divisible by H_k in packed reference") + repeat_factor = h_q_local // h_k_local + _wsr = 0 if is_causal else window_size_right + + ref_list = [] + lse_list = [] + batch_size = len(cu_seqlens_q) - 1 + for batch_idx in range(batch_size): + q_start = cu_seqlens_q[batch_idx] + q_end = cu_seqlens_q[batch_idx + 1] + k_start = cu_seqlens_k[batch_idx] + k_end = cu_seqlens_k[batch_idx + 1] + + q_i = q_packed[q_start:q_end].transpose(1, 0, 2) # (H_q, S_q, D) + k_i = k_packed[k_start:k_end].transpose(1, 0, 2) # (H_k, S_k, D) + v_i = v_packed[k_start:k_end].transpose(1, 0, 2) # (H_k, S_k, D) + + if repeat_factor > 1: + k_i = np.repeat(k_i, repeat_factor, axis=0) + v_i = np.repeat(v_i, repeat_factor, axis=0) + + s_i = np.einsum("hqd,hkd->hqk", q_i, k_i) * scale_softmax + s_q_local = q_i.shape[1] + s_k_local = k_i.shape[1] + + if window_size_left is not None or _wsr is not None: + q_coords = np.arange(s_q_local).reshape(-1, 1) + k_coords = np.arange(s_k_local).reshape(1, -1) + offset = 0 if not bottom_right_align else s_k_local - s_q_local + if window_size_left is None: + _mask = k_coords > q_coords + offset + _wsr + elif _wsr is None: + _mask = k_coords < q_coords + offset - window_size_left + else: + _mask = (k_coords > q_coords + offset + _wsr) | (k_coords < q_coords + offset - window_size_left) + s_i = np.where(_mask, -np.inf, s_i) + + if lse_calculation: + lse_i = _numpy_logsumexp(s_i, axis=-1) + else: + lse_i = None + + p_i = _numpy_softmax(s_i, axis=-1) + ref_i = np.einsum("hqk,hkd->hqd", p_i, v_i) + ref_i = ref_i.transpose(1, 0, 2) * scale_output + ref_list.append(ref_i) + if lse_calculation: + # (H_q, S_q) -> (S_q, H_q) to align packed output order. + lse_list.append(lse_i.transpose(1, 0)) + + ref = np.concatenate(ref_list, axis=0) + lse = np.concatenate(lse_list, axis=0) if lse_calculation else None + return ref, lse + + def _maybe_quantize_ref_for_narrow_out(o_ref_np): + if not (out_dtype.is_float and out_dtype.width <= 8): + return o_ref_np, tolerance + ref_narrow_cp = cp.empty(o_ref_np.shape, dtype=cp.uint8) + ref_narrow_cute = from_dlpack(ref_narrow_cp, assumed_align=16) + ref_narrow_cute.element_type = out_dtype + ref_narrow_cute = ref_narrow_cute.mark_layout_dynamic(leading_dim=_get_leading_dim(ref_narrow_cp)) + + ref_o_f32_cp = cp.asarray(o_ref_np) + ref_o_f32_cute = from_dlpack(ref_o_f32_cp, assumed_align=16) + ref_o_f32_cute.element_type = cutlass.Float32 + ref_o_f32_cute = ref_o_f32_cute.mark_layout_dynamic(leading_dim=_get_leading_dim(ref_o_f32_cp)) + + cute.testing.convert(ref_o_f32_cute, ref_narrow_cute) + cute.testing.convert(ref_narrow_cute, ref_o_f32_cute) + return ref_o_f32_cp.get(), 0.13 + + if vit_mode: + _vit_test_tag = "[vit_single_shot_test]" + if not skip_ref_check: + print(f"{_vit_test_tag} Running single-shot packed accuracy test:") + print(f"{_vit_test_tag} b={b}, seq_len={_s}, total_s={total_S}, h_q={h_q}, h_k={h_k}, d={d}, is_causal=False") + print(f"{_vit_test_tag} layout=[total_S,H,D], uniform cu_seqlens, max_seqlen={_s}") + compiled_fmha( + q_vit_tensor, + k_vit_tensor, + v_vit_tensor, + o_vit_tensor, + cu_seqlens, + _max_seqlen, + ref_scale_softmax_log2, + ref_scale_softmax, + ref_scale_output, + current_stream, + ) + + o_fp32_cp = cp.empty(o_vit_cp.shape, dtype=cp.float32) + o_fp32_cute = from_dlpack(o_fp32_cp, assumed_align=16) + o_fp32_cute.element_type = Float32 + o_fp32_cute = o_fp32_cute.mark_layout_dynamic(leading_dim=2) + cute.testing.convert(o_vit_tensor, o_fp32_cute) + o_result = o_fp32_cp.get() + + o_ref, _ = run_numpy_single_shot_reference_packed( + q_vit_ref, + k_vit_ref, + v_vit_ref, + cu_seqlens_np, + cu_seqlens_np, + scale_softmax=ref_scale_softmax, + scale_output=ref_scale_output, + is_causal=False, + bottom_right_align=False, + lse_calculation=False, + window_size_left=None, + window_size_right=None, + ) + o_ref, tol_for_check = _maybe_quantize_ref_for_narrow_out(o_ref) + np.testing.assert_allclose(o_result, o_ref, atol=tol_for_check, rtol=1e-05) + print(f"{_vit_test_tag} ViT single-shot accuracy check passed.") + + def generate_vit_tensors(): + _, q_ws, *_gq = create_and_pad_tensor(q_vit_shape, (0, 0, 0, 0), in_dtype, is_dynamic_layout=True) + _, k_ws, *_gk = create_and_pad_tensor(k_vit_shape, (0, 0, 0, 0), in_dtype, is_dynamic_layout=True) + _, v_ws, *_gv = create_and_pad_tensor(k_vit_shape, (0, 0, 0, 0), in_dtype, is_dynamic_layout=True) + _, o_ws, *_go = create_and_pad_tensor(q_vit_shape, (0, 0, 0, 0), out_dtype, is_dynamic_layout=True) + return testing.JitArguments( + mark_shd_dynamic(q_ws), + mark_shd_dynamic(k_ws), + mark_shd_dynamic(v_ws), + mark_shd_dynamic(o_ws), + cu_dyn, + _max_seqlen, + ref_scale_softmax_log2, + ref_scale_softmax, + ref_scale_output, + current_stream, + ) + + exec_time = testing.benchmark( + compiled_fmha, + workspace_generator=generate_vit_tensors, + workspace_count=1, + stream=current_stream, + warmup_iterations=warmup_iterations, + iterations=iterations, + ) + return exec_time + + # LLM path only below: plugin-aligned multi-round prefill regression. + if not skip_ref_check: + # LLM-only regression that mirrors attention plugin unit test + # (`test_plugin_vs_numpy_prefill`) and exercises cap != s_k stride. + # Keep this as the only LLM correctness check. + llm_prefill_tolerance = 0.13 if (out_dtype.is_float and out_dtype.width <= 8) else tolerance + if not isinstance(s_q, tuple) and not isinstance(s_k, tuple): + _num_rounds = 3 + _prefill_seq = min(s_q, s_k) + _cap = _prefill_seq * _num_rounds + print(f"{_tag} Running LLM multi-round prefill test (cap={_cap}, seq_len={_prefill_seq}, rounds={_num_rounds}) ...") + run_llm_multi_round_prefill_test( + batch_size=b, + seq_len=_prefill_seq, + num_rounds=_num_rounds, + h_q=h_q, + h_k=h_k, + d=d, + kv_cache_capacity=_cap, + mma_tiler_mn=mma_tiler_mn, + is_persistent=is_persistent, + is_causal=is_causal, + bottom_right_align=bottom_right_align, + use_sliding_window=use_sliding_window, + window_size_left_val=(window_size_left if window_size_left is not None else -1), + tolerance=llm_prefill_tolerance, + ) + print(f"{_tag} LLM multi-round prefill test passed.") + + def generate_tensors(): + _, q_tensor_workspace, *_gq = create_and_pad_tensor( + qo_shape, + qo_padding, + in_dtype, + is_dynamic_layout=True, + ) + _, kvcache_tensor_workspace, *_gkv = create_and_pad_tensor( + kvcache_shape, + kvcache_padding, + in_dtype, + is_dynamic_layout=True, + ) + _, o_tensor_workspace, *_go = create_and_pad_tensor( + qo_shape, + qo_padding, + out_dtype, + is_dynamic_layout=True, + ) + if lse_calculation: + _, lse_tensor, *_gl = create_and_pad_tensor( + lse_shape, + lse_padding, + cutlass.Float32, + is_dynamic_layout=True, + ) + else: + pass + + q_ws = mark_bshd_dynamic(q_tensor_workspace) + kv_ws = mark_kv_cache_dynamic(kvcache_tensor_workspace) + o_ws = mark_bshd_dynamic(o_tensor_workspace) + + args = testing.JitArguments( + q_ws, + kv_ws, + o_ws, + cu_kv_seqlens, + _wsl, + scale_q, + scale_k, + scale_v, + inv_scale_o, + current_stream, + ) + return args + + workspace_count = 1 + if use_cold_l2: + one_workspace_bytes = q_cp.size * q_cp.itemsize + kvcache_cp.size * kvcache_cp.itemsize + o_cp.size * o_cp.itemsize + (lse_cp.size * lse_cp.itemsize if lse_cp is not None else 0) + workspace_count = testing.get_workspace_count(one_workspace_bytes, warmup_iterations, iterations) + + exec_time = testing.benchmark( + compiled_fmha, + workspace_generator=generate_tensors, + workspace_count=workspace_count, + stream=current_stream, + warmup_iterations=warmup_iterations, + iterations=iterations, + ) + + return exec_time # Return execution time in microseconds + + +def run_llm_multi_round_prefill_test( + batch_size: int = 4, + seq_len: int = 8, + num_rounds: int = 3, + h_q: int = 8, + h_k: int = 8, + d: int = 128, + kv_cache_capacity: int = 64, + mma_tiler_mn: Tuple[int, int] = (128, 128), + is_persistent: bool = True, + is_causal: bool = True, + bottom_right_align: bool = True, + use_sliding_window: bool = False, + window_size_left_val: int = -1, + tolerance: float = 0.1, +): + """LLM FMHA multi-round prefill accuracy test aligned with plugin unit test. + + Each round appends seq_len new tokens to a KV cache with physical capacity + kv_cache_capacity >> effective_kv_len, exercising the cap vs s_k stride + distinction. After each round the CuTe DSL FMHA output is compared against + a numpy reference. + + This helper is intentionally LLM-only: it validates the [B,S,H,D] query path + with KV cache layout [B,2,H,cap,D], matching attention plugin prefill tests. + + :param batch_size: Number of batches. + :param seq_len: Tokens per prefill round (same for Q and new K/V). + :param num_rounds: How many rounds of prefill to run. + :param h_q: Number of query heads. + :param h_k: Number of KV heads. + :param d: Head dimension (64 or 128). + :param kv_cache_capacity: Physical KV cache capacity (cap). + :param mma_tiler_mn: MMA tile shape. + :param is_persistent: Use persistent kernel. + :param is_causal: Enable causal masking (window_size_right = 0). + :param bottom_right_align: bottom-right causal mask alignment. + :param use_sliding_window: Enable sliding window masking. + :param window_size_left_val: Left window size (-1 = disabled). + :param tolerance: Max absolute error tolerance. + """ + _tag = "[llm_prefill_test]" + b = batch_size + cap = kv_cache_capacity + h_r = h_q // h_k + window_size_left = window_size_left_val if use_sliding_window else None + window_size_right = 0 if is_causal else None + + print(f"{_tag} Running multi-round prefill accuracy test:") + print(f"{_tag} b={b}, seq_len={seq_len}, rounds={num_rounds}, cap={cap}, h_q={h_q}, h_k={h_k}, d={d}, is_causal={is_causal}") + + _MMA_K_ATOM = 256 // 16 + padded_d = ((d + _MMA_K_ATOM - 1) // _MMA_K_ATOM) * _MMA_K_ATOM + actual_head_dim = d if padded_d != d else None + + if h_q % h_k != 0: + raise ValueError("h_q must be divisible by h_k") + if num_rounds * seq_len > cap: + raise ValueError(f"total tokens ({num_rounds * seq_len}) exceeds capacity ({cap})") + + cp.random.seed(42) + np.random.seed(42) + + # FP16 test: all per-tensor scales are 1.0. + # The kernel computes softmax_scale = scale_q * scale_k / sqrt(d) internally. + _scale_q = 1.0 + _scale_k = 1.0 + _scale_v = 1.0 + _inv_scale_o = 1.0 + # Reference softmax scale for numpy validation + ref_scale_softmax = 1.0 / math.sqrt(d) + + mask_type = fmha_utils.MaskEnum.WINDOW_MASK + if bottom_right_align: + mask_type = fmha_utils.MaskEnum.WINDOW_MASK_INFERENCE + + fmha_op = BlackwellFusedMultiHeadAttentionForward( + Float32, + Float32, + (*mma_tiler_mn, padded_d), + is_persistent, + mask_type, + use_sliding_window=use_sliding_window, + is_causal=is_causal, + actual_head_dim=actual_head_dim, + ) + current_stream = cuda.CUstream(cp.cuda.get_current_stream().ptr) + _wsl = Int32(window_size_left_val) if use_sliding_window else Int32(0) + + # ---- helpers ---- + def _to_cute(arr, element_type): + t = from_dlpack(arr, assumed_align=16) + t.element_type = element_type + return t + + def mark_bshd_dynamic(tensor): + so = (0, 1, 2, 3) + return ( + tensor.mark_layout_dynamic(leading_dim=3) + .mark_compact_shape_dynamic(mode=0, stride_order=so) + .mark_compact_shape_dynamic(mode=1, stride_order=so) + .mark_compact_shape_dynamic(mode=2, stride_order=so) + ) + + def mark_kv_cache_dynamic(tensor): + so = (0, 1, 2, 3, 4) + return ( + tensor.mark_layout_dynamic(leading_dim=4) + .mark_compact_shape_dynamic(mode=0, stride_order=so) + .mark_compact_shape_dynamic(mode=2, stride_order=so) + .mark_compact_shape_dynamic(mode=3, stride_order=so) + ) + + def _numpy_softmax(x, axis=-1): + x_max = np.max(x, axis=axis, keepdims=True) + x_max = np.where(np.isfinite(x_max), x_max, 0.0) + e_x = np.exp(x - x_max) + s = np.sum(e_x, axis=axis, keepdims=True) + s = np.where(s == 0, 1.0, s) + return e_x / s + + # ---- state: KV cache (zero-initialized, filled progressively) ---- + kv_np = np.zeros((b, 2, h_k, cap, d), dtype=np.float32) + compiled_fmha = None + + all_pass = True + current_pos = 0 + + for round_idx in range(num_rounds): + effective_kv_len = current_pos + seq_len + print(f"\n--- Round {round_idx + 1}/{num_rounds} (pos={current_pos}, s_k={effective_kv_len}, cap={cap}) ---") + + # ---- generate new Q and K/V for this round ---- + q_np = np.random.randint(-2, 2, (b, seq_len, h_q, d)).astype(np.float32) + new_k_np = np.random.randint(-2, 2, (b, h_k, seq_len, d)).astype(np.float32) + new_v_np = np.random.randint(-2, 2, (b, h_k, seq_len, d)).astype(np.float32) + + # write new K/V into cache at [current_pos : current_pos + seq_len] + kv_np[:, 0, :, current_pos : current_pos + seq_len, :] = new_k_np + kv_np[:, 1, :, current_pos : current_pos + seq_len, :] = new_v_np + + # ---- upload to GPU (FP16) ---- + q_cp = cp.asarray(q_np.astype(np.float16)) + kv_cp = cp.asarray(kv_np.astype(np.float16)) + o_cp = cp.zeros((b, seq_len, h_q, d), dtype=cp.float16) + + q_t = mark_bshd_dynamic(_to_cute(q_cp, cutlass.Float16)) + kv_t = mark_kv_cache_dynamic(_to_cute(kv_cp, cutlass.Float16)) + o_t = mark_bshd_dynamic(_to_cute(o_cp, cutlass.Float16)) + + # cum_seqlen_k: uniform effective_kv_len across batches + cu_kv_np = np.arange(b + 1, dtype=np.int32) * effective_kv_len + cu_kv_cp = cp.asarray(cu_kv_np) + cu_kv = from_dlpack(cu_kv_cp, assumed_align=16) + cu_kv = cu_kv.mark_layout_dynamic(leading_dim=0).mark_compact_shape_dynamic(mode=0, stride_order=(0,)) + + # ---- compile on first round ---- + if compiled_fmha is None: + start_time = time.time() + compiled_fmha = cute.compile( + fmha_op, + q_t, + kv_t, + o_t, + cu_kv, + _wsl, + _scale_q, + _scale_k, + _scale_v, + _inv_scale_o, + current_stream, + ) + print(f"{_tag} Compilation time: {time.time() - start_time:.4f}s") + + # ---- run kernel ---- + compiled_fmha( + q_t, + kv_t, + o_t, + cu_kv, + _wsl, + _scale_q, + _scale_k, + _scale_v, + _inv_scale_o, + current_stream, + ) + + # ---- read output ---- + o_f32_cp = cp.empty(o_cp.shape, dtype=cp.float32) + o_f32_cute = from_dlpack(o_f32_cp, assumed_align=16) + o_f32_cute.element_type = Float32 + o_f32_cute = o_f32_cute.mark_layout_dynamic(leading_dim=3) + cute.testing.convert(o_t, o_f32_cute) + o_result = o_f32_cp.get() + + # ---- numpy reference (only attend to valid tokens, not full cap) ---- + for bi in range(b): + q_b = q_np[bi].transpose(1, 0, 2) # (h_q, s_q, d) + k_b = kv_np[bi, 0, :, :effective_kv_len] # (h_k, s_k, d) + v_b = kv_np[bi, 1, :, :effective_kv_len] # (h_k, s_k, d) + + if h_q != h_k: + k_b = np.repeat(k_b, h_r, axis=0) + v_b = np.repeat(v_b, h_r, axis=0) + + scores = np.einsum("hqd,hkd->hqk", q_b, k_b) * ref_scale_softmax + + s_k = effective_kv_len + if window_size_left is not None or window_size_right is not None: + q_coords = np.arange(seq_len).reshape(-1, 1) + k_coords = np.arange(s_k).reshape(1, -1) + offset = (s_k - seq_len) if bottom_right_align else 0 + if window_size_left is None: + mask = k_coords > q_coords + offset + window_size_right + elif window_size_right is None: + mask = k_coords < q_coords + offset - window_size_left + else: + mask = (k_coords > q_coords + offset + window_size_right) | (k_coords < q_coords + offset - window_size_left) + scores = np.where(mask, -np.inf, scores) + probs = _numpy_softmax(scores, axis=-1) + o_ref = np.einsum("hqk,hkd->hqd", probs, v_b) + o_ref = o_ref.transpose(1, 0, 2) * _scale_v * _inv_scale_o + + o_actual = o_result[bi] + max_diff = np.max(np.abs(o_actual - o_ref)) + mean_diff = np.mean(np.abs(o_actual - o_ref)) + + if max_diff > tolerance: + print(f" batch {bi}: FAIL max_diff={max_diff:.6f} mean_diff={mean_diff:.6f}") + all_pass = False + else: + print(f" batch {bi}: PASS max_diff={max_diff:.6f} mean_diff={mean_diff:.6f}") + + current_pos += seq_len + + if all_pass: + print(f"\n{_tag} All {num_rounds} rounds passed.") + else: + raise AssertionError(f"{_tag} Some rounds failed accuracy check!") + + return all_pass + + +if __name__ == "__main__": + + def parse_comma_separated_ints(s: str): + try: + return tuple(int(x.strip()) for x in s.split(",")) + except ValueError: + raise argparse.ArgumentTypeError("Invalid format. Expected comma-separated integers.") + + def parse_nested_comma_separated_ints(s: str): + try: + s = s.strip() + if "(" not in s: + return tuple(int(x.strip()) for x in s.split(",")) + + start = s.find("(") + end = s.find(")") + if start == -1 or end == -1: + raise ValueError("Mismatched parentheses") + + before = s[:start].strip().rstrip(",") + middle = s[start + 1 : end].strip() + after = s[end + 1 :].strip().lstrip(",") + + result = [] + if before: + result.extend(int(x.strip()) for x in before.split(",")) + + if middle: + nested_tuple = tuple(int(x.strip()) for x in middle.split(",")) + result.append(nested_tuple) + + if after: + result.extend(int(x.strip()) for x in after.split(",")) + + return tuple(result) + + except ValueError as e: + if str(e) == "Mismatched parentheses": + raise argparse.ArgumentTypeError("Mismatched parentheses in input") + else: + raise argparse.ArgumentTypeError("Invalid format. Expected comma-separated integers with optional parentheses for nested tuple.") + + parser = argparse.ArgumentParser(description="Example of FMHA on Blackwell.") + + parser.add_argument( + "--in_dtype", + type=cutlass.dtype, + default=cutlass.Float16, + help="Input data type", + ) + + parser.add_argument( + "--out_dtype", + type=cutlass.dtype, + default=cutlass.Float16, + help="Output data type", + ) + + parser.add_argument( + "--qk_acc_dtype", + type=cutlass.dtype, + default=Float32, + help="QK accumulator data type", + ) + + parser.add_argument( + "--pv_acc_dtype", + type=cutlass.dtype, + default=Float32, + help="PV accumulator data type", + ) + + parser.add_argument( + "--mma_tiler_mn", + type=parse_comma_separated_ints, + default=(128, 128), + help="MMA tile shape (M, N)", + ) + + parser.add_argument( + "--is_persistent", + action="store_true", + help="Is persistent", + ) + + parser.add_argument( + "--is_causal", + action="store_true", + help="Whether to use casual mask", + ) + + parser.add_argument( + "--bottom_right_align", + action="store_true", + help="Whether to use bottom right align, under this settion, the end of q is aligned with the end of k.", + ) + + parser.add_argument( + "--lse_calculation", + action="store_true", + help="Whether to calculate lse", + ) + + parser.add_argument( + "--window_size", + type=parse_comma_separated_ints, + default=(-1, -1), + help="Sliding window size (left, right) for attention masking.", + ) + + parser.add_argument( + "--q_shape", + type=parse_nested_comma_separated_ints, + default=(1, 256, 8, 128), + help="Shape of Q (B, S_q, H, D)", + ) + + parser.add_argument( + "--k_shape", + type=parse_nested_comma_separated_ints, + default=(1, 256, 8, 128), + help="Shape of K (B, S_k, H_k, D)", + ) + + parser.add_argument( + "--scale_q", + type=float, + default=1.0, + help="Scaling factors to dequantize Q", + ) + + parser.add_argument( + "--scale_k", + type=float, + default=1.0, + help="Scaling factors to dequantize K", + ) + + parser.add_argument( + "--scale_v", + type=float, + default=1.0, + help="Scaling factors to dequantize V", + ) + + parser.add_argument( + "--inv_scale_o", + type=float, + default=1.0, + help="Scaling factor to quantize O", + ) + + parser.add_argument( + "--scale_softmax", + type=float, + default=0.0, + help="Scaling factor to scale S (i.e. Q*K); if zero, defaults to 1/sqrt(D)", + ) + + parser.add_argument("--tolerance", type=float, default=1e-01, help="Tolerance for validation") + + parser.add_argument( + "--warmup_iterations", + type=int, + default=0, + help="Number of iterations for warmup", + ) + + parser.add_argument( + "--iterations", + type=int, + default=1, + help="Number of iterations after warmup", + ) + + parser.add_argument( + "--skip_ref_check", + action="store_true", + help="Skip reference check", + ) + + parser.add_argument( + "--use_cold_l2", + action="store_true", + default=False, + help="Use circular buffer tensor sets to ensure L2 cold cache", + ) + + parser.add_argument( + "--output_dir", + type=str, + default="./fmha_aot_artifacts", + help="Output directory for AOT compiled artifacts (fmha.h and fmha.o)", + ) + + parser.add_argument( + "--export_only", + action="store_true", + help="Compile and export only; skip reference check and benchmark", + ) + + parser.add_argument( + "--file_name", + type=str, + default="fmha", + help="Base file name for exported artifacts (e.g., fmha_d64 -> fmha_d64.h, fmha_d64.o)", + ) + + parser.add_argument( + "--function_prefix", + type=str, + default="fmha", + help="Function prefix for exported C symbols (avoids conflicts when compiling multiple variants)", + ) + + parser.add_argument( + "--vit_mode", + action="store_true", + help="Compile ViT FMHA variant: packed varlen with separate Q/K/V, bidirectional (no causal mask). Produces a different ABI.", + ) + + args = parser.parse_args() + + if cp.cuda.runtime.getDeviceCount() == 0: + raise RuntimeError("GPU is required to run this example!") + + if len(args.q_shape) != 4: + parser.error("--q_shape must contain exactly 4 values") + + if len(args.k_shape) != 4: + parser.error("--k_shape must contain exactly 4 values") + + if len(args.mma_tiler_mn) != 2: + parser.error("--mma_tiler_mn must contain exactly 2 values") + + if args.vit_mode: + assert args.k_shape == args.q_shape, f"vit_mode requires k_shape == q_shape; got k_shape={args.k_shape}, q_shape={args.q_shape}" + assert not args.is_causal, "vit_mode is bidirectional; --is_causal must not be set" + assert args.window_size == (-1, -1), f"vit_mode does not support sliding window; got --window_size={args.window_size}" + + latency = run( + args.q_shape, + args.k_shape, + args.in_dtype, + args.out_dtype, + args.qk_acc_dtype, + args.pv_acc_dtype, + args.mma_tiler_mn, + args.is_persistent, + args.is_causal, + args.bottom_right_align, + args.lse_calculation, + args.window_size, + args.scale_q, + args.scale_k, + args.scale_v, + args.inv_scale_o, + args.scale_softmax, + args.tolerance, + args.warmup_iterations, + args.iterations, + args.skip_ref_check, + args.use_cold_l2, + output_dir=args.output_dir, + export_only=args.export_only, + file_name=args.file_name, + function_prefix=args.function_prefix, + vit_mode=args.vit_mode, + ) + + if latency is not None: + print(f"Latency: {latency:.4f} us") diff --git a/lightx2v/common/ops/attn/utils/sla_util.py b/lightx2v/common/ops/attn/utils/sla_util.py index d717d9a94..4f7d635b0 100755 --- a/lightx2v/common/ops/attn/utils/sla_util.py +++ b/lightx2v/common/ops/attn/utils/sla_util.py @@ -32,6 +32,35 @@ def compress_kernel( tl.store(XM + xm_offset + idx_l * D + offs_d, x_mean.to(XM.dtype.element_ty)) +@triton.jit +def centered_compress_kernel( + X, + CENTER, + XM, + L: tl.constexpr, + D: tl.constexpr, + BLOCK_L: tl.constexpr, +): + idx_l = tl.program_id(0) + idx_bh = tl.program_id(1) + + offs_l = idx_l * BLOCK_L + tl.arange(0, BLOCK_L) + offs_d = tl.arange(0, D) + valid_l = offs_l[:, None] < L + + x_offset = idx_bh * L * D + xm_offset = idx_bh * ((L + BLOCK_L - 1) // BLOCK_L) * D + center_offset = idx_bh * D + x = tl.load(X + x_offset + offs_l[:, None] * D + offs_d[None, :], mask=valid_l) + center = tl.load(CENTER + center_offset + offs_d) + centered_x = (x - center[None, :]).to(XM.dtype.element_ty) + centered_x = tl.where(valid_l, centered_x, 0.0) + + nx = min(BLOCK_L, L - idx_l * BLOCK_L) + x_mean = tl.sum(centered_x, axis=0, dtype=tl.float32) / nx + tl.store(XM + xm_offset + idx_l * D + offs_d, x_mean.to(XM.dtype.element_ty)) + + def mean_pool(x, BLK): assert x.is_contiguous() @@ -44,14 +73,23 @@ def mean_pool(x, BLK): return x_mean -def get_block_map(q, k, topk_ratio, BLKQ=64, BLKK=64): - arg_k = k - torch.mean(k, dim=-2, keepdim=True) # smooth-k technique in SageAttention - pooled_qblocks = mean_pool(q, BLKQ) - pooled_kblocks = mean_pool(arg_k, BLKK) +def centered_mean_pool(x, center, BLK): + assert x.is_contiguous() + assert center.is_contiguous() + + B, H, L, D = x.shape + L_BLOCKS = (L + BLK - 1) // BLK + x_mean = torch.empty((B, H, L_BLOCKS, D), device=x.device, dtype=x.dtype) + grid = (L_BLOCKS, B * H) + centered_compress_kernel[grid](x, center, x_mean, L, D, BLK) + return x_mean + + +def _get_block_lut(pooled_qblocks, pooled_kblocks, topk_ratio): # GQA - num_q_heads = q.size(1) - num_kv_heads = k.size(1) + num_q_heads = pooled_qblocks.size(1) + num_kv_heads = pooled_kblocks.size(1) if num_q_heads != num_kv_heads: assert num_q_heads % num_kv_heads == 0, f"Number of Q heads ({num_q_heads}) must be divisible by number of KV heads ({num_kv_heads})" repeat_factor = num_q_heads // num_kv_heads @@ -64,7 +102,33 @@ def get_block_map(q, k, topk_ratio, BLKQ=64, BLKK=64): topk = max(1, min(K, int(topk_ratio * K))) lut = torch.topk(pooled_score, topk, dim=-1, sorted=False).indices - sparse_map = torch.zeros_like(pooled_score, dtype=torch.int8) + return lut, topk, K + + +def get_block_lut(q, k, topk_ratio, BLKQ=64, BLKK=64): + pooled_qblocks = mean_pool(q, BLKQ) + k_mean = torch.mean(k, dim=-2, keepdim=True) + pooled_kblocks = centered_mean_pool(k, k_mean, BLKK) + return _get_block_lut(pooled_qblocks, pooled_kblocks, topk_ratio) + + +def block_lut_to_ordinal_metadata(lut, num_k_blocks): + assert lut.dim() == 4 + assert lut.shape[-1] <= num_k_blocks + + full_block_idx = torch.zeros((*lut.shape[:-1], num_k_blocks), dtype=torch.int32, device=lut.device) + full_block_idx[..., : lut.shape[-1]] = torch.sort(lut, dim=-1).values.to(torch.int32) + full_block_cnt = torch.full(lut.shape[:-1], lut.shape[-1], dtype=torch.int32, device=lut.device) + return full_block_idx, full_block_cnt + + +def get_block_map(q, k, topk_ratio, BLKQ=64, BLKK=64): + arg_k = k - torch.mean(k, dim=-2, keepdim=True) # smooth-k technique in SageAttention + pooled_qblocks = mean_pool(q, BLKQ) + pooled_kblocks = mean_pool(arg_k, BLKK) + lut, topk, num_k_blocks = _get_block_lut(pooled_qblocks, pooled_kblocks, topk_ratio) + + sparse_map = torch.zeros((*lut.shape[:-1], num_k_blocks), dtype=torch.int8, device=lut.device) sparse_map.scatter_(-1, lut, 1) return sparse_map, lut, topk diff --git a/lightx2v/common/ops/attn/utils/sla_util_blhd.py b/lightx2v/common/ops/attn/utils/sla_util_blhd.py index 05e9bf4ba..bb360e532 100644 --- a/lightx2v/common/ops/attn/utils/sla_util_blhd.py +++ b/lightx2v/common/ops/attn/utils/sla_util_blhd.py @@ -2,6 +2,8 @@ import triton import triton.language as tl +from .sla_util import _get_block_lut + @triton.jit def compress_kernel( @@ -30,6 +32,42 @@ def compress_kernel( tl.store(XM + xm_offset + idx_l * D + offs_d, x_mean.to(XM.dtype.element_ty)) +@triton.jit +def centered_compress_kernel( + X, + CENTER, + XM, + L: tl.constexpr, + H: tl.constexpr, + D: tl.constexpr, + BLOCK_L: tl.constexpr, +): + idx_l = tl.program_id(0) + idx_bh = tl.program_id(1) + + idx_b = idx_bh // H + idx_h = idx_bh - idx_b * H + + offs_l = idx_l * BLOCK_L + tl.arange(0, BLOCK_L) + offs_d = tl.arange(0, D) + valid_l = offs_l[:, None] < L + + x_offset = idx_b * L * H * D + idx_h * D + xm_offset = idx_bh * ((L + BLOCK_L - 1) // BLOCK_L) * D + center_offset = idx_bh * D + x = tl.load( + X + x_offset + offs_l[:, None] * (H * D) + offs_d[None, :], + mask=valid_l, + ) + center = tl.load(CENTER + center_offset + offs_d) + centered_x = (x - center[None, :]).to(XM.dtype.element_ty) + centered_x = tl.where(valid_l, centered_x, 0.0) + + nx = min(BLOCK_L, L - idx_l * BLOCK_L) + x_mean = tl.sum(centered_x, axis=0, dtype=tl.float32) / nx + tl.store(XM + xm_offset + idx_l * D + offs_d, x_mean.to(XM.dtype.element_ty)) + + def mean_pool(x, BLK): assert x.is_contiguous() B, L, H, D = x.shape @@ -41,6 +79,26 @@ def mean_pool(x, BLK): return x_mean +def centered_mean_pool(x, center, BLK): + assert x.is_contiguous() + assert center.is_contiguous() + + B, L, H, D = x.shape + L_BLOCKS = (L + BLK - 1) // BLK + x_mean = torch.empty((B, H, L_BLOCKS, D), device=x.device, dtype=x.dtype) + + grid = (L_BLOCKS, B * H) + centered_compress_kernel[grid](x, center, x_mean, L, H, D, BLK) + return x_mean + + +def get_block_lut_blhd(q, k, topk_ratio, BLKQ=64, BLKK=64): + pooled_qblocks = mean_pool(q, BLKQ) + k_mean = torch.mean(k, dim=1, keepdim=True) + pooled_kblocks = centered_mean_pool(k, k_mean, BLKK) + return _get_block_lut(pooled_qblocks, pooled_kblocks, topk_ratio) + + def get_block_map_blhd(q, k, topk_ratio, BLKQ=64, BLKK=64): arg_k = k - torch.mean(k, dim=1, keepdim=True) pooled_qblocks = mean_pool(q, BLKQ) diff --git a/lightx2v/common/ops/mm/mm_weight.py b/lightx2v/common/ops/mm/mm_weight.py index 61613dd5a..2880d079f 100755 --- a/lightx2v/common/ops/mm/mm_weight.py +++ b/lightx2v/common/ops/mm/mm_weight.py @@ -31,17 +31,28 @@ try: from lightx2v_kernel.gemm import ( + cublaslt_scaled_nvfp4_mm_bias, cutlass_scaled_mxfp4_mm, cutlass_scaled_mxfp6_mxfp8_mm, cutlass_scaled_mxfp8_mm, cutlass_scaled_nvfp4_mm, + cutlass_scaled_nvfp4_mm_split_n_stride, + cutlass_scaled_nvfp4_mm_split_n_stride_gelu, + cutlass_scaled_nvfp4_mm_split_n_stride_residual_gate, scaled_mxfp4_quant, scaled_mxfp6_quant, scaled_mxfp8_quant, scaled_nvfp4_quant, ) except ImportError: - scaled_nvfp4_quant, cutlass_scaled_nvfp4_mm = None, None + ( + scaled_nvfp4_quant, + cutlass_scaled_nvfp4_mm, + cutlass_scaled_nvfp4_mm_split_n_stride, + cutlass_scaled_nvfp4_mm_split_n_stride_gelu, + cutlass_scaled_nvfp4_mm_split_n_stride_residual_gate, + cublaslt_scaled_nvfp4_mm_bias, + ) = None, None, None, None, None, None scaled_mxfp4_quant, cutlass_scaled_mxfp4_mm = None, None scaled_mxfp6_quant, cutlass_scaled_mxfp6_mxfp8_mm = None, None scaled_mxfp8_quant, cutlass_scaled_mxfp8_mm = None, None @@ -62,6 +73,7 @@ except ImportError: sgl_kernel = None + try: import comfy_kitchen except ImportError: @@ -1247,6 +1259,35 @@ def apply(self, input_tensor): ) return output_tensor + def apply_quantized(self, input_tensor_quant, input_tensor_scale): + return cutlass_scaled_nvfp4_mm( + input_tensor_quant, + self.weight, + input_tensor_scale, + self.weight_scale, + alpha=self.alpha, + bias=self.bias, + ) + + def apply_quantized_cublaslt(self, input_tensor_quant, input_tensor_scale, algorithm_index=-1): + return cublaslt_scaled_nvfp4_mm_bias( + input_tensor_quant, + self.weight, + input_tensor_scale, + self.weight_scale, + alpha=self.alpha, + bias=self.bias, + algorithm_index=algorithm_index, + ) + + def apply_cublaslt(self, input_tensor, algorithm_index=-1): + input_tensor_quant, input_tensor_scale = self.act_quant_func(input_tensor) + return self.apply_quantized_cublaslt( + input_tensor_quant, + input_tensor_scale, + algorithm_index, + ) + def to_cuda(self, non_blocking=False): self.weight = self.pin_weight.to(AI_DEVICE, non_blocking=non_blocking) if hasattr(self, "pin_weight_scale"): @@ -1331,6 +1372,84 @@ def load_state_dict_from_disk(self, block_index, adapter_block_index=None): del weight_scale_tensor +@MM_WEIGHT_REGISTER("nvfp4-split-n-stride-workaround") +class MMWeightWnvfp4Anvfp4dynamicSplitNStrideWorkaround(MMWeightWnvfp4Anvfp4dynamic): + """Represent N shards as batches in one strided CUTLASS GEMM. + + The activation is quantized once and the complete weight and scale tensors + are passed to the backend. A and its scales are broadcast across batches; + weight, weight scales, bias, and output columns advance by batch stride. + """ + + def __init__( + self, + weight_name, + bias_name, + create_cuda_buffer=False, + create_cpu_buffer=False, + lazy_load=False, + lazy_load_file=None, + is_post_adapter=False, + lora_prefix="diffusion_model.blocks", + lora_path="", + split_n_parts=2, + ): + super().__init__( + weight_name, + bias_name, + create_cuda_buffer, + create_cpu_buffer, + lazy_load, + lazy_load_file, + is_post_adapter, + lora_prefix=lora_prefix, + lora_path=lora_path, + ) + if isinstance(split_n_parts, bool) or not isinstance(split_n_parts, int): + raise TypeError("split_n_parts must be an integer") + if split_n_parts < 2: + raise ValueError("split_n_parts must be at least 2 for the split-N weight type") + self.split_n_parts = split_n_parts + + def apply(self, input_tensor): + input_tensor_quant, input_tensor_scale = self.act_quant_func(input_tensor) + return cutlass_scaled_nvfp4_mm_split_n_stride( + input_tensor_quant, + self.weight, + input_tensor_scale, + self.weight_scale, + alpha=self.alpha, + bias=self.bias, + split_n_parts=self.split_n_parts, + ) + + def apply_gelu(self, input_tensor): + input_tensor_quant, input_tensor_scale = self.act_quant_func(input_tensor) + return cutlass_scaled_nvfp4_mm_split_n_stride_gelu( + input_tensor_quant, + self.weight, + input_tensor_scale, + self.weight_scale, + alpha=self.alpha, + bias=self.bias, + split_n_parts=self.split_n_parts, + ) + + def apply_residual_gate(self, input_tensor, residual, gate): + input_tensor_quant, input_tensor_scale = self.act_quant_func(input_tensor) + return cutlass_scaled_nvfp4_mm_split_n_stride_residual_gate( + input_tensor_quant, + self.weight, + input_tensor_scale, + self.weight_scale, + alpha=self.alpha, + residual=residual, + gate=gate, + bias=self.bias, + split_n_parts=self.split_n_parts, + ) + + @MM_WEIGHT_REGISTER("nvfp4-split-n-workaround") class MMWeightWnvfp4Anvfp4dynamicSplitNWorkaround(MMWeightWnvfp4Anvfp4dynamic): """Temporary application-level two-way split-N workaround. @@ -2592,6 +2711,7 @@ def __init__( lora_path="", reduce_output=True, lora_column_chunks=1, + mm_kwargs=None, ): super().__init__( weight_name, @@ -2623,6 +2743,7 @@ def __init__( is_post_adapter=is_post_adapter, lora_prefix=lora_prefix, lora_path=lora_path, + **(mm_kwargs or {}), ) self._row_split_bias = None diff --git a/lightx2v/infer.py b/lightx2v/infer.py index f95818a3e..0e1f53733 100755 --- a/lightx2v/infer.py +++ b/lightx2v/infer.py @@ -349,13 +349,36 @@ def main(): validate_config_paths(config) + use_full_profiler = os.environ.get("LIGHTX2V_TORCH_PROFILER_FULL", "0") == "1" + with ProfilingContext4DebugL1("Total Cost"): # init runner runner = init_runner(config) # start to infer data = args.__dict__ update_input_info_from_dict(input_info, data) - runner.run_pipeline(input_info) + if use_full_profiler: + from torch.profiler import ProfilerActivity, profile, record_function + + with profile( + activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], + record_shapes=True, + ) as profiler: + with record_function("run_pipeline"): + runner.run_pipeline(input_info) + + trace_path = os.environ.get( + "LIGHTX2V_TORCH_PROFILER_FULL_TRACE", + "/tmp/lightx2v_full_trace.json", + ) + trace_dir = os.path.dirname(trace_path) + if trace_dir: + os.makedirs(trace_dir, exist_ok=True) + profiler.export_chrome_trace(trace_path) + logger.info(f"[TorchProfiler] Full pipeline trace exported to {trace_path}") + print(profiler.key_averages().table(sort_by="cuda_time_total", row_limit=40)) + else: + runner.run_pipeline(input_info) # Clean up distributed process group if dist.is_initialized(): diff --git a/lightx2v/models/networks/wan/infer/transformer_infer.py b/lightx2v/models/networks/wan/infer/transformer_infer.py index 3058c4607..c92e1d051 100755 --- a/lightx2v/models/networks/wan/infer/transformer_infer.py +++ b/lightx2v/models/networks/wan/infer/transformer_infer.py @@ -44,6 +44,7 @@ def __init__(self, config): self.modulate_func = modulate self.clean_cuda_cache = self.config.get("clean_cuda_cache", False) self.mxfp8_fuse_enable = self.config.get("mxfp8_fuse_enable", True) + self.thor = self.config.get("thor", False) self.infer_dtype = GET_DTYPE() self.sensitive_layer_dtype = GET_SENSITIVE_DTYPE() @@ -302,7 +303,21 @@ def infer_self_attn(self, phase, x, shift_msa, scale_msa, grid_sizes=None, **sel norm1_out = norm1_out.to(self.infer_dtype) s, n, d = *norm1_out.shape[:1], self.num_heads, self.head_dim - if norm1_quant is not None: + if self.thor: + q_proj = phase.self_attn_q + k_proj = phase.self_attn_k + v_proj = phase.self_attn_v + if not getattr(phase, "_nvfp4_qkv_cublaslt_scale_checked", False): + reference_scale = q_proj.input_global_scale + if not torch.equal(reference_scale, k_proj.input_global_scale) or not torch.equal(reference_scale, v_proj.input_global_scale): + raise ValueError("Thor mode requires identical Q/K/V input_global_scale values") + phase._nvfp4_qkv_cublaslt_scale_checked = True + qkv_quant, qkv_scale = q_proj.act_quant_func(norm1_out) + algorithm_index = -1 + q = phase.self_attn_norm_q.apply(q_proj.apply_quantized_cublaslt(qkv_quant, qkv_scale, algorithm_index)).view(s, n, d) + k = phase.self_attn_norm_k.apply(k_proj.apply_quantized_cublaslt(qkv_quant, qkv_scale, algorithm_index)).view(s, n, d) + v = v_proj.apply_quantized_cublaslt(qkv_quant, qkv_scale, algorithm_index).view(s, n, d) + elif norm1_quant is not None: q = phase.self_attn_norm_q.apply(self._mxfp8_apply_quantized(phase.self_attn_q, norm1_quant, norm1_scale)).view(s, n, d) k = phase.self_attn_norm_k.apply(self._mxfp8_apply_quantized(phase.self_attn_k, norm1_quant, norm1_scale)).view(s, n, d) v = self._mxfp8_apply_quantized(phase.self_attn_v, norm1_quant, norm1_scale).view(s, n, d) @@ -376,7 +391,13 @@ def infer_self_attn(self, phase, x, shift_msa, scale_msa, grid_sizes=None, **sel **attn_running_args, ) - y = phase.self_attn_o.apply(attn_out) + if self.thor: + y = phase.self_attn_o.apply_cublaslt( + attn_out, + -1, + ) + else: + y = phase.self_attn_o.apply(attn_out) if self.clean_cuda_cache: del q, k, v, attn_out @@ -403,7 +424,14 @@ def infer_cross_attn(self, phase, x, context, y_out, gate_msa): context_img = context_img.to(self.infer_dtype) n, d = self.num_heads, self.head_dim - q = phase.cross_attn_norm_q.apply(phase.cross_attn_q.apply(norm3_out)).view(-1, n, d) + if self.thor: + q = phase.cross_attn_q.apply_cublaslt( + norm3_out, + -1, + ) + else: + q = phase.cross_attn_q.apply(norm3_out) + q = phase.cross_attn_norm_q.apply(q).view(-1, n, d) k = phase.cross_attn_norm_k.apply(phase.cross_attn_k.apply(context)).view(-1, n, d) v = phase.cross_attn_v.apply(context).view(-1, n, d) @@ -436,7 +464,13 @@ def infer_cross_attn(self, phase, x, context, y_out, gate_msa): del k_img, v_img, img_attn_out torch_device_module.empty_cache() - attn_out = phase.cross_attn_o.apply(attn_out) + if self.thor: + attn_out = phase.cross_attn_o.apply_cublaslt( + attn_out, + -1, + ) + else: + attn_out = phase.cross_attn_o.apply(attn_out) if self.clean_cuda_cache: del q, k, v, norm3_out, context, context_img @@ -484,13 +518,23 @@ def infer_ffn(self, phase, x, attn_out, c_shift_msa, c_scale_msa, c_gate_msa=Non c_shift_msa=mxfp8_modulate_shift, ) - y = phase.ffn_0.apply(norm2_out) + if self.thor: + y = phase.ffn_0.apply_gelu(norm2_out) + else: + y = phase.ffn_0.apply(norm2_out) if self.clean_cuda_cache: - del norm2_out, x + del norm2_out + if not self.thor: + del x torch_device_module.empty_cache() - y = torch.nn.functional.gelu(y, approximate="tanh") + if not self.thor: + y = torch.nn.functional.gelu(y, approximate="tanh") if self.clean_cuda_cache: torch_device_module.empty_cache() + if self.thor: + phase.ffn_2.apply_residual_gate(y, x, c_gate_msa.squeeze()) + return None + y = phase.ffn_2.apply(y) return y diff --git a/lightx2v/models/networks/wan/weights/transformer_weights.py b/lightx2v/models/networks/wan/weights/transformer_weights.py index 3d16dd6da..e581cd50c 100755 --- a/lightx2v/models/networks/wan/weights/transformer_weights.py +++ b/lightx2v/models/networks/wan/weights/transformer_weights.py @@ -52,6 +52,7 @@ def _mm_weight( lora_prefix="", lora_path="", mm_type_override=None, + mm_kwargs=None, ): mm_type = mm_type_override if mm_type is None: @@ -75,6 +76,7 @@ def _mm_weight( lazy_load_file=lazy_load_file, lora_prefix=lora_prefix, lora_path=lora_path, + mm_kwargs=mm_kwargs, ) return MM_WEIGHT_REGISTER[mm_type]( weight_name, @@ -85,6 +87,7 @@ def _mm_weight( lazy_load_file, lora_prefix=lora_prefix, lora_path=lora_path, + **(mm_kwargs or {}), ) @@ -432,6 +435,20 @@ def __init__( lora_path=lora_path, ), ) + if config.get("thor", False): + if mm_type != "nvfp4": + raise ValueError("thor=true requires dit_quant_scheme='nvfp4'") + if config.get("tensor_parallel", False): + raise NotImplementedError("Thor mode does not support tensor parallelism") + if config.get("cpu_offload", False) or create_cuda_buffer or create_cpu_buffer: + raise NotImplementedError("Thor mode does not support CPU offload") + if lazy_load: + raise NotImplementedError("Thor mode does not support lazy loading") + if lora_path or config.get("lora_configs"): + raise NotImplementedError("Thor mode does not support LoRA") + if config.get("feature_caching", "NoCaching") != "NoCaching": + raise NotImplementedError("Thor mode requires feature_caching='NoCaching'") + self.add_module( "self_attn_o", _mm_weight( @@ -798,12 +815,28 @@ def __init__( LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")](), ) + thor = config.get("thor", False) + if thor: + if self.mm_type != "nvfp4": + raise ValueError("thor=true requires dit_quant_scheme='nvfp4'") + if config.get("tensor_parallel", False): + raise NotImplementedError("Thor mode does not support tensor parallelism") + if config.get("cpu_offload", False) or create_cuda_buffer or create_cpu_buffer: + raise NotImplementedError("Thor mode does not support CPU offload") + if self.lazy_load: + raise NotImplementedError("Thor mode does not support lazy loading") + if lora_path or config.get("lora_configs"): + raise NotImplementedError("Thor mode does not support LoRA") + split_n = config.get("nvfp4_ffn_split_n_workaround", False) if not isinstance(split_n, bool): raise TypeError("nvfp4_ffn_split_n_workaround must be a boolean") - # Temporary and intentionally scoped to Wan FFN. The checkpoint format - # remains ``nvfp4``; only the execution implementation changes. - ffn_mm_type = "nvfp4-split-n-workaround" if self.mm_type == "nvfp4" and split_n else self.mm_type + # The checkpoint format remains ``nvfp4``; only Wan FFN execution changes. + if thor: + ffn_mm_type = "nvfp4-split-n-stride-workaround" + else: + ffn_mm_type = "nvfp4-split-n-workaround" if self.mm_type == "nvfp4" and split_n else self.mm_type + ffn_mm_kwargs = {"split_n_parts": 2} if thor else None fp = f"{block_prefix}.{self.block_index}" self.add_module( @@ -817,6 +850,7 @@ def __init__( create_cpu_buffer=create_cpu_buffer, lazy_load=self.lazy_load, lazy_load_file=self.lazy_load_file, + mm_kwargs=ffn_mm_kwargs, lora_prefix=block_prefix, lora_path=lora_path, mm_type_override=ffn_mm_type, @@ -833,6 +867,7 @@ def __init__( create_cpu_buffer=create_cpu_buffer, lazy_load=self.lazy_load, lazy_load_file=self.lazy_load_file, + mm_kwargs=ffn_mm_kwargs, lora_prefix=block_prefix, lora_path=lora_path, mm_type_override=ffn_mm_type, diff --git a/lightx2v/shot_runner/shot_base.py b/lightx2v/shot_runner/shot_base.py index cc09496a7..902af1184 100755 --- a/lightx2v/shot_runner/shot_base.py +++ b/lightx2v/shot_runner/shot_base.py @@ -10,7 +10,7 @@ from lightx2v.utils.input_info import fill_input_info_from_defaults from lightx2v.utils.profiler import * from lightx2v.utils.registry_factory import RUNNER_REGISTER -from lightx2v.utils.set_config import print_config, set_config, set_parallel_config +from lightx2v.utils.set_config import print_config, set_config, set_parallel_config, validate_thor_config from lightx2v_platform.registry_factory import PLATFORM_DEVICE_REGISTER @@ -59,6 +59,7 @@ def load_clip_configs(main_json_path): config["parallel"] = cfg["parallel"] set_parallel_config(config) + validate_thor_config(config) clip_configs.append(ClipConfig(name=item["name"], config_json=config)) return clip_configs diff --git a/lightx2v/utils/set_config.py b/lightx2v/utils/set_config.py index 86b04fce0..ae8209a50 100755 --- a/lightx2v/utils/set_config.py +++ b/lightx2v/utils/set_config.py @@ -57,6 +57,21 @@ def validate_model_task_args(args): raise ValueError("--omni_vision_subtask is only valid with --task omni_vision_task") +def validate_thor_config(config): + thor = config.get("thor", False) + if not isinstance(thor, bool): + raise TypeError("thor must be a boolean") + if not thor: + return + + if config.get("model_cls") != "wan2.2_moe": + raise ValueError("thor=true only supports model_cls='wan2.2_moe'") + if config.get("task") not in ("i2v", "t2v"): + raise ValueError("thor=true only supports task='i2v' or 't2v'") + if config.get("dit_quant_scheme") != "nvfp4": + raise ValueError("thor=true requires dit_quant_scheme='nvfp4'") + + def set_args2config(args): config = get_default_config() config.update({k: v for k, v in vars(args).items() if k not in ALL_INPUT_INFO_KEYS and v is not None}) @@ -398,6 +413,7 @@ def auto_calc_config(config): config["sound_sampling_rate"] = int(sound_config.get("sampling_rate", 48000)) config["sound_hop_size"] = int(sound_config.get("hop_size", 1920)) + validate_thor_config(config) return config diff --git a/lightx2v_kernel/CMakeLists.txt b/lightx2v_kernel/CMakeLists.txt index 369f4da9c..af153a880 100644 --- a/lightx2v_kernel/CMakeLists.txt +++ b/lightx2v_kernel/CMakeLists.txt @@ -1,6 +1,8 @@ cmake_minimum_required(VERSION 3.22 FATAL_ERROR) project(lightx2v-kernel LANGUAGES CXX CUDA) +option(LIGHTX2V_THOR_NVFP4_ONLY "Build only NVFP4 operators for NVIDIA Thor" OFF) + include(${CMAKE_CURRENT_LIST_DIR}/cmake/utils.cmake) # Python @@ -71,31 +73,41 @@ set(LIGHTX2V_KERNEL_CUDA_FLAGS ) -list(APPEND LIGHTX2V_KERNEL_CUDA_FLAGS - # "-gencode=arch=compute_90,code=sm_90" - # "-gencode=arch=compute_90a,code=sm_90a" - # "-gencode=arch=compute_100,code=sm_100" - # "-gencode=arch=compute_100a,code=sm_100a" - # "-gencode=arch=compute_120,code=sm_120" - "-gencode=arch=compute_120a,code=sm_120a" -) - - -set(SOURCES - "csrc/gemm/nvfp4_scaled_mm_kernels_sm120.cu" - "csrc/gemm/nvfp4_quant_kernels_sm120.cu" - "csrc/gemm/mxfp4_quant_kernels_sm120.cu" - "csrc/gemm/mxfp8_quant_kernels_sm120.cu" - "csrc/gemm/mxfp6_quant_kernels_sm120.cu" - "csrc/gemm/mxfp4_scaled_mm_kernels_sm120.cu" - "csrc/gemm/mxfp6_mxfp8_scaled_mm_kernels_sm120.cu" - "csrc/gemm/mxfp8_scaled_mm_kernels_sm120.cu" - "csrc/kvcache/kv_dequant_cuda.cu" - "csrc/common_extension.cc" -) +if(LIGHTX2V_THOR_NVFP4_ONLY) + list(APPEND LIGHTX2V_KERNEL_CUDA_FLAGS + "-gencode=arch=compute_110a,code=sm_110a" + ) + set(SOURCES + "csrc/gemm/nvfp4_scaled_mm_kernels_sm120.cu" + "csrc/gemm/nvfp4_quant_kernels_sm120.cu" + "csrc/gemm/nvfp4_cublaslt_mm.cpp" + "csrc/common_extension.cc" + ) +else() + list(APPEND LIGHTX2V_KERNEL_CUDA_FLAGS + "-gencode=arch=compute_120a,code=sm_120a" + ) + set(SOURCES + "csrc/gemm/nvfp4_scaled_mm_kernels_sm120.cu" + "csrc/gemm/nvfp4_quant_kernels_sm120.cu" + "csrc/gemm/nvfp4_cublaslt_mm.cpp" + "csrc/gemm/mxfp4_quant_kernels_sm120.cu" + "csrc/gemm/mxfp8_quant_kernels_sm120.cu" + "csrc/gemm/mxfp6_quant_kernels_sm120.cu" + "csrc/gemm/mxfp4_scaled_mm_kernels_sm120.cu" + "csrc/gemm/mxfp6_mxfp8_scaled_mm_kernels_sm120.cu" + "csrc/gemm/mxfp8_scaled_mm_kernels_sm120.cu" + "csrc/kvcache/kv_dequant_cuda.cu" + "csrc/common_extension.cc" + ) +endif() Python_add_library(common_ops MODULE USE_SABI ${SKBUILD_SABI_VERSION} WITH_SOABI ${SOURCES}) +if(LIGHTX2V_THOR_NVFP4_ONLY) + target_compile_definitions(common_ops PRIVATE LIGHTX2V_THOR_NVFP4_ONLY=1) +endif() + message(STATUS "LIGHTX2V_KERNEL_CUDA_FLAGS: ${LIGHTX2V_KERNEL_CUDA_FLAGS}") target_compile_options(common_ops PRIVATE $<$:${LIGHTX2V_KERNEL_CUDA_FLAGS}>) diff --git a/lightx2v_kernel/csrc/common_extension.cc b/lightx2v_kernel/csrc/common_extension.cc index 3bfdc746f..5d409764c 100644 --- a/lightx2v_kernel/csrc/common_extension.cc +++ b/lightx2v_kernel/csrc/common_extension.cc @@ -10,12 +10,51 @@ TORCH_LIBRARY_FRAGMENT(lightx2v_kernel, m) { "cutlass_scaled_nvfp4_mm_sm120(Tensor! out, Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, Tensor " "alpha, Tensor? bias) -> ()"); m.impl("cutlass_scaled_nvfp4_mm_sm120", torch::kCUDA, &cutlass_scaled_nvfp4_mm_sm120); + m.def( + "cublaslt_scaled_nvfp4_mm_bias_sm120(Tensor! out, Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, " + "Tensor alpha, Tensor bias, int algorithm_index) -> ()"); + m.impl( + "cublaslt_scaled_nvfp4_mm_bias_sm120", + torch::kCUDA, + &cublaslt_scaled_nvfp4_mm_bias_sm120); + m.def( + "cublaslt_scaled_nvfp4_mm_bias_algo_count_sm120(Tensor out, Tensor mat_a, Tensor mat_b, Tensor scales_a, " + "Tensor scales_b, Tensor alpha, Tensor bias) -> int"); + m.impl( + "cublaslt_scaled_nvfp4_mm_bias_algo_count_sm120", + torch::kCUDA, + &cublaslt_scaled_nvfp4_mm_bias_algo_count_sm120); + + m.def( + "cutlass_scaled_nvfp4_mm_split_n_stride_sm120(Tensor! out, Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor " + "scales_b, Tensor alpha, Tensor? bias, int split_n_parts) -> ()"); + m.impl( + "cutlass_scaled_nvfp4_mm_split_n_stride_sm120", + torch::kCUDA, + &cutlass_scaled_nvfp4_mm_split_n_stride_sm120); + + m.def( + "cutlass_scaled_nvfp4_mm_split_n_stride_gelu_sm120(Tensor! out, Tensor mat_a, Tensor mat_b, Tensor scales_a, " + "Tensor scales_b, Tensor alpha, Tensor? bias, int split_n_parts) -> ()"); + m.impl( + "cutlass_scaled_nvfp4_mm_split_n_stride_gelu_sm120", + torch::kCUDA, + &cutlass_scaled_nvfp4_mm_split_n_stride_gelu_sm120); + + m.def( + "cutlass_scaled_nvfp4_mm_split_n_stride_residual_gate_sm120(Tensor! residual, Tensor mat_a, Tensor mat_b, " + "Tensor scales_a, Tensor scales_b, Tensor alpha, Tensor? bias, Tensor gate, int split_n_parts) -> ()"); + m.impl( + "cutlass_scaled_nvfp4_mm_split_n_stride_residual_gate_sm120", + torch::kCUDA, + &cutlass_scaled_nvfp4_mm_split_n_stride_residual_gate_sm120); m.def( "scaled_nvfp4_quant_sm120(Tensor! output, Tensor! input," " Tensor! output_scale, Tensor! input_scale) -> ()"); m.impl("scaled_nvfp4_quant_sm120", torch::kCUDA, &scaled_nvfp4_quant_sm120); +#ifndef LIGHTX2V_THOR_NVFP4_ONLY m.def( "scaled_mxfp4_quant_sm120(Tensor! output, Tensor! input," " Tensor! output_scale) -> ()"); @@ -65,6 +104,7 @@ TORCH_LIBRARY_FRAGMENT(lightx2v_kernel, m) { "dequantize_kv_cache_fp4(Tensor[] values, Tensor[] scale_factors, Tensor[] amax, " "int num_heads, int block_token_size, int dtype_code, float e2m1_max, float e4m3_max) -> Tensor"); m.impl("dequantize_kv_cache_fp4", torch::kCUDA, &dequantize_kv_cache_fp4_cuda); +#endif } diff --git a/lightx2v_kernel/csrc/gemm/nvfp4_cublaslt_mm.cpp b/lightx2v_kernel/csrc/gemm/nvfp4_cublaslt_mm.cpp new file mode 100644 index 000000000..848d3cb23 --- /dev/null +++ b/lightx2v_kernel/csrc/gemm/nvfp4_cublaslt_mm.cpp @@ -0,0 +1,398 @@ +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace { + +constexpr size_t kMaxWorkspaceBytes = 128ULL * 1024 * 1024; +constexpr int kMaxHeuristicResults = 64; + +void check_cublas(cublasStatus_t status, char const* expression) { + TORCH_CHECK( + status == CUBLAS_STATUS_SUCCESS, + expression, + " failed: ", + cublasGetStatusString(status)); +} + +#define CUBLAS_CHECK(expression) check_cublas((expression), #expression) + +struct ProblemKey { + int device; + int64_t m; + int64_t n; + int64_t k; + uintptr_t stream; + + bool operator==(ProblemKey const& other) const { + return device == other.device && m == other.m && n == other.n && + k == other.k && stream == other.stream; + } +}; + +struct ProblemKeyHash { + size_t operator()(ProblemKey const& key) const { + size_t result = std::hash{}(key.device); + auto combine = [&result](size_t value) { + result ^= value + 0x9e3779b9 + (result << 6) + (result >> 2); + }; + combine(std::hash{}(key.m)); + combine(std::hash{}(key.n)); + combine(std::hash{}(key.k)); + combine(std::hash{}(key.stream)); + return result; + } +}; + +struct PlanSet { + cublasLtMatmulDesc_t operation_desc = nullptr; + cublasLtMatrixLayout_t weight_layout = nullptr; + cublasLtMatrixLayout_t activation_layout = nullptr; + cublasLtMatrixLayout_t output_layout = nullptr; + std::vector algorithms; + torch::Tensor workspace; + torch::Tensor beta; + std::mutex execution_mutex; + + ~PlanSet() { + if (output_layout != nullptr) { + cublasLtMatrixLayoutDestroy(output_layout); + } + if (activation_layout != nullptr) { + cublasLtMatrixLayoutDestroy(activation_layout); + } + if (weight_layout != nullptr) { + cublasLtMatrixLayoutDestroy(weight_layout); + } + if (operation_desc != nullptr) { + cublasLtMatmulDescDestroy(operation_desc); + } + } +}; + +cublasLtHandle_t cublaslt_handle() { + static cublasLtHandle_t handle = [] { + cublasLtHandle_t value = nullptr; + CUBLAS_CHECK(cublasLtCreate(&value)); + return value; + }(); + return handle; +} + +void check_inputs( + torch::Tensor const& output, + torch::Tensor const& activation, + torch::Tensor const& weight, + torch::Tensor const& activation_scale, + torch::Tensor const& weight_scale, + torch::Tensor const& alpha, + torch::Tensor const& bias) { + TORCH_CHECK( + output.is_cuda() && activation.is_cuda() && weight.is_cuda(), + "output, activation, and weight must be CUDA tensors"); + TORCH_CHECK( + activation_scale.is_cuda() && weight_scale.is_cuda(), + "scale tensors must be CUDA tensors"); + TORCH_CHECK(alpha.is_cuda() && bias.is_cuda(), "alpha and bias must be CUDA tensors"); + + int const device = activation.get_device(); + TORCH_CHECK( + output.get_device() == device && weight.get_device() == device && + activation_scale.get_device() == device && weight_scale.get_device() == device && + alpha.get_device() == device && bias.get_device() == device, + "all tensors must be on the same CUDA device"); + TORCH_CHECK(output.is_contiguous(), "output must be contiguous"); + TORCH_CHECK( + activation.is_contiguous() && weight.is_contiguous(), + "activation and weight must be contiguous"); + TORCH_CHECK( + activation_scale.is_contiguous() && weight_scale.is_contiguous(), + "scale tensors must be contiguous"); + TORCH_CHECK(alpha.is_contiguous() && bias.is_contiguous(), "alpha and bias must be contiguous"); + TORCH_CHECK(output.scalar_type() == at::ScalarType::BFloat16, "output must be BF16"); + TORCH_CHECK(activation.scalar_type() == at::ScalarType::Byte, "packed activation must be uint8"); + TORCH_CHECK(weight.scalar_type() == at::ScalarType::Byte, "packed weight must be uint8"); + TORCH_CHECK( + activation_scale.scalar_type() == at::ScalarType::Float8_e4m3fn && + weight_scale.scalar_type() == at::ScalarType::Float8_e4m3fn, + "scale tensors must be float8_e4m3fn"); + TORCH_CHECK(alpha.scalar_type() == at::ScalarType::Float, "alpha must be FP32"); + TORCH_CHECK(bias.scalar_type() == at::ScalarType::BFloat16, "bias must be BF16"); + TORCH_CHECK( + activation.dim() == 2 && weight.dim() == 2 && output.dim() == 2, + "activation, weight, and output must be matrices"); + TORCH_CHECK( + activation.sizes()[1] == weight.sizes()[1], + "activation and weight packed K dimensions must match"); + + int64_t const m = activation.sizes()[0]; + int64_t const n = weight.sizes()[0]; + int64_t const k = activation.sizes()[1] * 2; + TORCH_CHECK(alpha.numel() == 1, "alpha must contain one value"); + TORCH_CHECK( + output.sizes() == at::IntArrayRef({m, n}), + "output shape must be (", + m, + ", ", + n, + ")"); + TORCH_CHECK(bias.numel() == n, "bias must contain ", n, " values"); + + auto round_up = [](int64_t value, int64_t alignment) { + return (value + alignment - 1) / alignment * alignment; + }; + int64_t const rounded_m = round_up(m, 128); + int64_t const rounded_n = round_up(n, 128); + int64_t const rounded_k_scale = round_up(k / 16, 4); + TORCH_CHECK( + activation_scale.sizes() == at::IntArrayRef({rounded_m, rounded_k_scale}), + "activation scale shape must be (", + rounded_m, + ", ", + rounded_k_scale, + ")"); + TORCH_CHECK( + weight_scale.sizes() == at::IntArrayRef({rounded_n, rounded_k_scale}), + "weight scale shape must be (", + rounded_n, + ", ", + rounded_k_scale, + ")"); +} + +void set_dynamic_pointers( + PlanSet& plan, + torch::Tensor const& activation_scale, + torch::Tensor const& weight_scale, + torch::Tensor const& bias) { + void const* weight_scale_pointer = weight_scale.data_ptr(); + void const* activation_scale_pointer = activation_scale.data_ptr(); + void const* bias_pointer = bias.data_ptr(); + CUBLAS_CHECK(cublasLtMatmulDescSetAttribute( + plan.operation_desc, + CUBLASLT_MATMUL_DESC_A_SCALE_POINTER, + &weight_scale_pointer, + sizeof(weight_scale_pointer))); + CUBLAS_CHECK(cublasLtMatmulDescSetAttribute( + plan.operation_desc, + CUBLASLT_MATMUL_DESC_B_SCALE_POINTER, + &activation_scale_pointer, + sizeof(activation_scale_pointer))); + CUBLAS_CHECK(cublasLtMatmulDescSetAttribute( + plan.operation_desc, + CUBLASLT_MATMUL_DESC_BIAS_POINTER, + &bias_pointer, + sizeof(bias_pointer))); +} + +std::shared_ptr create_plan_set( + torch::Tensor const& activation, + torch::Tensor const& weight, + torch::Tensor const& activation_scale, + torch::Tensor const& weight_scale, + torch::Tensor const& bias) { + auto plan = std::make_shared(); + int64_t const m = activation.sizes()[0]; + int64_t const n = weight.sizes()[0]; + int64_t const k = activation.sizes()[1] * 2; + + CUBLAS_CHECK(cublasLtMatmulDescCreate( + &plan->operation_desc, CUBLAS_COMPUTE_32F, CUDA_R_32F)); + cublasOperation_t transpose_weight = CUBLAS_OP_T; + cublasOperation_t transpose_activation = CUBLAS_OP_N; + cublasLtPointerMode_t pointer_mode = CUBLASLT_POINTER_MODE_DEVICE; + cublasLtMatmulMatrixScale_t block_scale_mode = + CUBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3; + cublasLtEpilogue_t epilogue = CUBLASLT_EPILOGUE_BIAS; + cudaDataType_t bias_type = CUDA_R_16BF; + CUBLAS_CHECK(cublasLtMatmulDescSetAttribute( + plan->operation_desc, + CUBLASLT_MATMUL_DESC_TRANSA, + &transpose_weight, + sizeof(transpose_weight))); + CUBLAS_CHECK(cublasLtMatmulDescSetAttribute( + plan->operation_desc, + CUBLASLT_MATMUL_DESC_TRANSB, + &transpose_activation, + sizeof(transpose_activation))); + CUBLAS_CHECK(cublasLtMatmulDescSetAttribute( + plan->operation_desc, + CUBLASLT_MATMUL_DESC_POINTER_MODE, + &pointer_mode, + sizeof(pointer_mode))); + CUBLAS_CHECK(cublasLtMatmulDescSetAttribute( + plan->operation_desc, + CUBLASLT_MATMUL_DESC_A_SCALE_MODE, + &block_scale_mode, + sizeof(block_scale_mode))); + CUBLAS_CHECK(cublasLtMatmulDescSetAttribute( + plan->operation_desc, + CUBLASLT_MATMUL_DESC_B_SCALE_MODE, + &block_scale_mode, + sizeof(block_scale_mode))); + CUBLAS_CHECK(cublasLtMatmulDescSetAttribute( + plan->operation_desc, + CUBLASLT_MATMUL_DESC_EPILOGUE, + &epilogue, + sizeof(epilogue))); + CUBLAS_CHECK(cublasLtMatmulDescSetAttribute( + plan->operation_desc, + CUBLASLT_MATMUL_DESC_BIAS_DATA_TYPE, + &bias_type, + sizeof(bias_type))); + + // Row-major (M,K) storage is column-major (K,M). Compute output^T. + CUBLAS_CHECK(cublasLtMatrixLayoutCreate( + &plan->weight_layout, CUDA_R_4F_E2M1, k, n, k)); + CUBLAS_CHECK(cublasLtMatrixLayoutCreate( + &plan->activation_layout, CUDA_R_4F_E2M1, k, m, k)); + CUBLAS_CHECK(cublasLtMatrixLayoutCreate( + &plan->output_layout, CUDA_R_16BF, n, m, n)); + set_dynamic_pointers(*plan, activation_scale, weight_scale, bias); + + cublasLtMatmulPreference_t preference = nullptr; + CUBLAS_CHECK(cublasLtMatmulPreferenceCreate(&preference)); + CUBLAS_CHECK(cublasLtMatmulPreferenceSetAttribute( + preference, + CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, + &kMaxWorkspaceBytes, + sizeof(kMaxWorkspaceBytes))); + + std::array results{}; + int returned_results = 0; + cublasStatus_t heuristic_status = cublasLtMatmulAlgoGetHeuristic( + cublaslt_handle(), + plan->operation_desc, + plan->weight_layout, + plan->activation_layout, + plan->output_layout, + plan->output_layout, + preference, + kMaxHeuristicResults, + results.data(), + &returned_results); + cublasLtMatmulPreferenceDestroy(preference); + CUBLAS_CHECK(heuristic_status); + + size_t workspace_size = 0; + for (int index = 0; index < returned_results; ++index) { + if (results[index].state != CUBLAS_STATUS_SUCCESS) { + continue; + } + plan->algorithms.push_back(results[index].algo); + workspace_size = std::max(workspace_size, results[index].workspaceSize); + } + TORCH_CHECK(!plan->algorithms.empty(), "cuBLASLt returned no NVFP4+bias algorithms"); + + auto options = torch::TensorOptions().device(activation.device()); + plan->workspace = torch::empty( + {static_cast(std::max(workspace_size, 1))}, + options.dtype(torch::kUInt8)); + plan->beta = torch::zeros({1}, options.dtype(torch::kFloat32)); + return plan; +} + +std::shared_ptr get_plan_set( + torch::Tensor const& activation, + torch::Tensor const& weight, + torch::Tensor const& activation_scale, + torch::Tensor const& weight_scale, + torch::Tensor const& bias, + cudaStream_t stream) { + static std::mutex cache_mutex; + static std::unordered_map, ProblemKeyHash> cache; + ProblemKey const key{ + activation.get_device(), + activation.sizes()[0], + weight.sizes()[0], + activation.sizes()[1] * 2, + reinterpret_cast(stream)}; + + std::lock_guard lock(cache_mutex); + auto found = cache.find(key); + if (found != cache.end()) { + return found->second; + } + auto plan = create_plan_set( + activation, weight, activation_scale, weight_scale, bias); + cache.emplace(key, plan); + return plan; +} + +} // namespace + +void cublaslt_scaled_nvfp4_mm_bias_sm120( + torch::Tensor& output, + torch::Tensor const& activation, + torch::Tensor const& weight, + torch::Tensor const& activation_scale, + torch::Tensor const& weight_scale, + torch::Tensor const& alpha, + torch::Tensor const& bias, + int64_t algorithm_index) { + check_inputs( + output, activation, weight, activation_scale, weight_scale, alpha, bias); + c10::cuda::CUDAGuard device_guard(activation.device()); + cudaStream_t stream = at::cuda::getCurrentCUDAStream(activation.get_device()); + auto plan = get_plan_set( + activation, weight, activation_scale, weight_scale, bias, stream); + + int64_t selected_algorithm = algorithm_index; + if (selected_algorithm == -1) { + selected_algorithm = plan->algorithms.size() > 1 ? 1 : 0; + } + TORCH_CHECK( + selected_algorithm >= 0 && + selected_algorithm < static_cast(plan->algorithms.size()), + "cuBLASLt algorithm index ", + selected_algorithm, + " is outside [0, ", + plan->algorithms.size(), + ")"); + + std::lock_guard lock(plan->execution_mutex); + set_dynamic_pointers(*plan, activation_scale, weight_scale, bias); + CUBLAS_CHECK(cublasLtMatmul( + cublaslt_handle(), + plan->operation_desc, + alpha.data_ptr(), + weight.data_ptr(), + plan->weight_layout, + activation.data_ptr(), + plan->activation_layout, + plan->beta.data_ptr(), + output.data_ptr(), + plan->output_layout, + output.data_ptr(), + plan->output_layout, + &plan->algorithms[selected_algorithm], + plan->workspace.data_ptr(), + plan->workspace.numel(), + stream)); +} + +int64_t cublaslt_scaled_nvfp4_mm_bias_algo_count_sm120( + torch::Tensor const& output, + torch::Tensor const& activation, + torch::Tensor const& weight, + torch::Tensor const& activation_scale, + torch::Tensor const& weight_scale, + torch::Tensor const& alpha, + torch::Tensor const& bias) { + check_inputs( + output, activation, weight, activation_scale, weight_scale, alpha, bias); + c10::cuda::CUDAGuard device_guard(activation.device()); + cudaStream_t stream = at::cuda::getCurrentCUDAStream(activation.get_device()); + auto plan = get_plan_set( + activation, weight, activation_scale, weight_scale, bias, stream); + return static_cast(plan->algorithms.size()); +} diff --git a/lightx2v_kernel/csrc/gemm/nvfp4_quant_kernels_sm120.cu b/lightx2v_kernel/csrc/gemm/nvfp4_quant_kernels_sm120.cu index ad50950ba..ada1cea66 100644 --- a/lightx2v_kernel/csrc/gemm/nvfp4_quant_kernels_sm120.cu +++ b/lightx2v_kernel/csrc/gemm/nvfp4_quant_kernels_sm120.cu @@ -6,6 +6,8 @@ #include #include +#include + #include "utils.h" // Get type2 from type or vice versa (applied to half and bfloat16) @@ -36,43 +38,50 @@ struct TypeConverter<__nv_bfloat16> { #define ELTS_PER_THREAD 8 +#ifndef EXPERIMENT_BLOCK_THREADS +#define EXPERIMENT_BLOCK_THREADS 256 +#endif +#ifndef EXPERIMENT_LAUNCH_THREADS +#define EXPERIMENT_LAUNCH_THREADS 192 +#endif +#ifndef EXPERIMENT_BLOCKS_PER_SM +#define EXPERIMENT_BLOCKS_PER_SM 4 +#endif +#ifndef EXPERIMENT_TILE4_LOAD_MODE +#define EXPERIMENT_TILE4_LOAD_MODE 3 +#endif +#ifndef EXPERIMENT_TILE4_MAX_REGISTERS +#define EXPERIMENT_TILE4_MAX_REGISTERS 38 +#endif +#ifndef EXPERIMENT_TILE4_OUTER_M +#define EXPERIMENT_TILE4_OUTER_M 1 +#endif +#ifndef EXPERIMENT_PRECOMPUTE_SF +#define EXPERIMENT_PRECOMPUTE_SF 0 +#endif +#ifndef EXPERIMENT_TILE4_PIPELINE +#define EXPERIMENT_TILE4_PIPELINE 0 +#endif +#ifndef EXPERIMENT_TILE4_DIRECT_SCALE +#define EXPERIMENT_TILE4_DIRECT_SCALE 0 +#endif +#ifndef EXPERIMENT_TILE4_DOUBLE_BUFFER +#define EXPERIMENT_TILE4_DOUBLE_BUFFER 0 +#endif +#ifndef EXPERIMENT_SCALE_PACK_MODE +#define EXPERIMENT_SCALE_PACK_MODE 0 +#endif +#ifndef EXPERIMENT_INCREMENTAL_TASK +#define EXPERIMENT_INCREMENTAL_TASK 0 +#endif + +constexpr int TILE4_BLOCK_THREADS = 320; +constexpr int TILE4_MAX_BLOCK_THREADS = 512; +constexpr int TILE4_GRID_BLOCKS_NUMERATOR = 19; +constexpr int TILE4_GRID_BLOCKS_DENOMINATOR = 5; constexpr int CVT_FP4_ELTS_PER_THREAD = 8; constexpr int CVT_FP4_SF_VEC_SIZE = 16; -// Convert 8 float32 values into 8 e2m1 values (represented as one uint32_t). -inline __device__ uint32_t fp32_vec_to_e2m1(float (&array)[8]) { - // PTX instructions used here requires sm100a. -// #if CUDA_VERSION >= 12080 -// #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) && __CUDA_ARCH_HAS_FEATURE__(SM100_ALL) - uint32_t val; - asm volatile( - "{\n" - ".reg .b8 byte0;\n" - ".reg .b8 byte1;\n" - ".reg .b8 byte2;\n" - ".reg .b8 byte3;\n" - "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n" - "cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n" - "cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n" - "cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n" - "mov.b32 %0, {byte0, byte1, byte2, byte3};\n" - "}" - : "=r"(val) - : "f"(array[0]), - "f"(array[1]), - "f"(array[2]), - "f"(array[3]), - "f"(array[4]), - "f"(array[5]), - "f"(array[6]), - "f"(array[7])); - return val; -// #else -// return 0; -// #endif -// #endif -} - // Convert 4 float2 values into 8 e2m1 values (represented as one uint32_t). inline __device__ uint32_t fp32_vec_to_e2m1(float2 (&array)[4]) { // PTX instructions used here requires sm100a. @@ -114,65 +123,140 @@ inline __device__ float reciprocal_approximate_ftz(float a) { return b; } -template -__device__ uint8_t* cvt_quant_to_fp4_get_sf_out_offset(int rowIdx, int colIdx, int numCols, SFType* SFout) { -// #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - static_assert(CVT_FP4_NUM_THREADS_PER_SF == 1 || CVT_FP4_NUM_THREADS_PER_SF == 2); - - // One pair of threads write one SF to global memory. - // TODO: stage through smem for packed STG.32 - // is it better than STG.8 from 4 threads ? - if (threadIdx.x % CVT_FP4_NUM_THREADS_PER_SF == 0) { - // SF vector index (16 elements share one SF in the K dimension). - int32_t kIdx = colIdx / CVT_FP4_NUM_THREADS_PER_SF; - int32_t mIdx = rowIdx; - - // SF layout [numMTiles, numKTiles, 32 (mTile), 4 (mTile), 4(kTile)] - // --> index [mTileIdx, kTileIdx, outerMIdx, innerMIdx, innerKIdx] +// Define a 16 bytes packed data type. +template +struct alignas(16) PackedVec { + typename TypeConverter::Type elts[4]; +}; - int32_t mTileIdx = mIdx / (32 * 4); - // SF vector size 16. - int factor = CVT_FP4_SF_VEC_SIZE * 4; - int32_t numKTiles = (numCols + factor - 1) / factor; - int64_t mTileStride = numKTiles * 32 * 4 * 4; +template <> +struct PackedVec<__nv_fp8_e4m3> { + __nv_fp8x2_e4m3 elts[8]; +}; - int32_t kTileIdx = (kIdx / 4); - int64_t kTileStride = 32 * 4 * 4; +template +inline __device__ PackedVec load_packed_vec(Type const* ptr) { + uint4 raw; + asm volatile( + "ld.global.v4.u32 {%0, %1, %2, %3}, [%4];" + : "=r"(raw.x), "=r"(raw.y), "=r"(raw.z), "=r"(raw.w) + : "l"(ptr)); + union { + uint4 raw; + PackedVec vec; + } value; + value.raw = raw; + return value.vec; +} - // M tile layout [32, 4] is column-major. - int32_t outerMIdx = (mIdx % 32); - int64_t outerMStride = 4 * 4; +template +inline __device__ PackedVec load_packed_vec_v2u64(Type const* ptr) { + unsigned long long lo; + unsigned long long hi; + asm volatile( + "ld.global.v2.u64 {%0, %1}, [%2];" + : "=l"(lo), "=l"(hi) + : "l"(ptr)); + union { + struct { + unsigned long long lo; + unsigned long long hi; + } raw; + PackedVec vec; + } value; + value.raw.lo = lo; + value.raw.hi = hi; + return value.vec; +} - int32_t innerMIdx = (mIdx % (32 * 4)) / 32; - int64_t innerMStride = 4; +template +inline __device__ PackedVec load_packed_vec_streaming(Type const* ptr) { + uint4 raw; + asm volatile( + "ld.global.cs.v4.u32 {%0, %1, %2, %3}, [%4];" + : "=r"(raw.x), "=r"(raw.y), "=r"(raw.z), "=r"(raw.w) + : "l"(ptr)); + union { + uint4 raw; + PackedVec vec; + } value; + value.raw = raw; + return value.vec; +} - int32_t innerKIdx = (kIdx % 4); - int64_t innerKStride = 1; +template +inline __device__ PackedVec load_packed_vec_global(Type const* ptr) { + uint4 raw; + asm volatile( + "ld.global.cg.v4.u32 {%0, %1, %2, %3}, [%4];" + : "=r"(raw.x), "=r"(raw.y), "=r"(raw.z), "=r"(raw.w) + : "l"(ptr)); + union { + uint4 raw; + PackedVec vec; + } value; + value.raw = raw; + return value.vec; +} - // Compute the global offset. - int64_t SFOffset = mTileIdx * mTileStride + kTileIdx * kTileStride + outerMIdx * outerMStride + - innerMIdx * innerMStride + innerKIdx * innerKStride; +template +inline __device__ PackedVec load_tile4_vec(Type const* ptr) { +#if EXPERIMENT_TILE4_LOAD_MODE == 1 + return load_packed_vec(ptr); +#elif EXPERIMENT_TILE4_LOAD_MODE == 2 + return load_packed_vec_v2u64(ptr); +#elif EXPERIMENT_TILE4_LOAD_MODE == 3 + return load_packed_vec_streaming(ptr); +#elif EXPERIMENT_TILE4_LOAD_MODE == 4 + return load_packed_vec_global(ptr); +#else + return *reinterpret_cast const*>(ptr); +#endif +} - return reinterpret_cast(SFout) + SFOffset; +inline __device__ uint32_t pack_scale_byte(uint8_t sfValue) { + uint32_t sf = uint32_t(sfValue); +#if EXPERIMENT_SCALE_PACK_MODE == 1 + uint32_t pair = 0; + uint32_t packed = 0; + if ((threadIdx.x & 1) == 0) { + uint32_t next = __shfl_down_sync(0x55555555, sf, 2, 8); + pair = sf | (next << 8); } -// #endif - return nullptr; + if ((threadIdx.x & 3) == 0) { + uint32_t nextPair = __shfl_down_sync(0x11111111, pair, 4, 8); + packed = pair | (nextPair << 16); + } + return packed; +#elif EXPERIMENT_SCALE_PACK_MODE == 2 + uint32_t pair = sf | (__shfl_down_sync(0xffffffff, sf, 2, 8) << 8); + return pair | (__shfl_down_sync(0xffffffff, pair, 4, 8) << 16); +#else + int32_t groupLane = (threadIdx.x & 31) & ~7; + uint32_t sf1 = __shfl_sync(0xffffffff, sf, groupLane + 2); + uint32_t sf2 = __shfl_sync(0xffffffff, sf, groupLane + 4); + uint32_t sf3 = __shfl_sync(0xffffffff, sf, groupLane + 6); + return sf | (sf1 << 8) | (sf2 << 16) | (sf3 << 24); +#endif } -// Define a 16 bytes packed data type. -template -struct PackedVec { - typename TypeConverter::Type elts[4]; -}; - -template <> -struct PackedVec<__nv_fp8_e4m3> { - __nv_fp8x2_e4m3 elts[8]; -}; +inline __device__ void stage_scale_byte(uint8_t sfValue, int32_t colIdx, uint32_t* sfStageRow) { + uint32_t packedSF = pack_scale_byte(sfValue); + if ((threadIdx.x & 7) == 0) { + sfStageRow[colIdx / 8] = packedSF; + } +} // Quantizes the provided PackedVec into the uint32_t output template -__device__ uint32_t cvt_warp_fp16_to_fp4(PackedVec& vec, float SFScaleVal, uint8_t* SFout) { +__device__ uint32_t cvt_warp_fp16_to_fp4( + PackedVec& vec, + float SFScaleVal, +#if EXPERIMENT_PRECOMPUTE_SF + float SFScaleForMax, +#endif + uint8_t* SFout, + uint8_t* SFValueOut = nullptr) { // #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) // Get absolute maximum values among the local 8 values. auto localMax = __habs2(vec.elts[0]); @@ -188,37 +272,31 @@ __device__ uint32_t cvt_warp_fp16_to_fp4(PackedVec& vec, float SFScaleVal, // Get the final absolute maximum values. float vecMax = float(__hmax(localMax.x, localMax.y)); - // Get the SF (max value of the vector / max value of e2m1). - // maximum value of e2m1 = 6.0. - // TODO: use half as compute data type. +#if EXPERIMENT_PRECOMPUTE_SF + float SFValue = vecMax * SFScaleForMax; +#else float SFValue = SFScaleVal * (vecMax * 0.16666666666666666f); - // 8 bits representation of the SF. +#endif uint8_t fp8SFVal; - // Write the SF to global memory (STG.8). if constexpr (UE8M0_SF) { __nv_fp8_e8m0 tmp; tmp.__x = __nv_cvt_float_to_e8m0(SFValue, __NV_SATFINITE, cudaRoundPosInf); SFValue = static_cast(tmp); fp8SFVal = tmp.__x; } else { - // Here SFValue is always positive, so E4M3 is the same as UE4M3. __nv_fp8_e4m3 tmp = __nv_fp8_e4m3(SFValue); fp8SFVal = tmp.__x; SFValue = static_cast(tmp); } - // Get the output scale. - // Recipe: final_scale = reciprocal(fp32(fp8(SFValue * SFScaleVal))) * - // reciprocal(SFScaleVal)) -// float outputScale = -// SFValue != 0 ? reciprocal_approximate_ftz(SFValue * reciprocal_approximate_ftz(SFScaleVal)) : 0.0f; - - float outputScale = - SFValue != 0 ? SFScaleVal * reciprocal_approximate_ftz(SFValue) : 0.0f; + float outputScale = SFValue != 0 ? SFScaleVal * reciprocal_approximate_ftz(SFValue) : 0.0f; if (SFout) { // Write the SF to global memory (STG.8). *SFout = fp8SFVal; } + if (SFValueOut) { + *SFValueOut = fp8SFVal; + } // Convert the input to float. float2 fp2Vals[CVT_FP4_ELTS_PER_THREAD / 2]; @@ -244,42 +322,271 @@ __device__ uint32_t cvt_warp_fp16_to_fp4(PackedVec& vec, float SFScaleVal, // #endif } -// Use UE4M3 by default. template -__global__ void -// #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) -__launch_bounds__(256, 6) cvt_fp16_to_fp4( -// #else -// cvt_fp16_to_fp4( -// #endif +__global__ __launch_bounds__(EXPERIMENT_BLOCK_THREADS, EXPERIMENT_BLOCKS_PER_SM) +void cvt_fp16_to_fp4_incremental_sf( int32_t numRows, int32_t numCols, Type const* in, float const* SFScale, uint32_t* out, uint32_t* SFout) { -// #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - using PackedVec = PackedVec; - static constexpr int CVT_FP4_NUM_THREADS_PER_SF = (CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD); - static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD, "Vec size is not matched."); - - // Get the global scaling factor, which will be applied to the SF. - // Note SFScale is the same as next GEMM's alpha, which is - // (448.f / (Alpha_A / 6.f)). + using InputVec = PackedVec; float const SFScaleVal = SFScale == nullptr ? 1.0f : SFScale[0]; +#if EXPERIMENT_PRECOMPUTE_SF + float const SFScaleForMax = SFScaleVal * 0.16666666666666666f; +#endif + int32_t numColVecs = numCols / CVT_FP4_ELTS_PER_THREAD; + int32_t numKTiles = (numCols + CVT_FP4_SF_VEC_SIZE * 4 - 1) / (CVT_FP4_SF_VEC_SIZE * 4); + + for (int32_t rowIdx = blockIdx.x; rowIdx < numRows; rowIdx += gridDim.x) { + int64_t rowVecBase = int64_t(rowIdx) * numColVecs; + int64_t sfRowBase = int64_t(rowIdx / 128) * numKTiles * 512 + + int64_t(rowIdx % 32) * 16 + int64_t((rowIdx % 128) / 32) * 4; + int32_t colIdx = threadIdx.x; + int64_t sfOffset = sfRowBase + int64_t(threadIdx.x / 8) * 512 + ((threadIdx.x / 2) & 3); + int32_t colStride = blockDim.x; + int64_t sfStride = int64_t(blockDim.x / 8) * 512; + for (; colIdx < numColVecs; colIdx += colStride, sfOffset += sfStride) { + int64_t vecOffset = rowVecBase + colIdx; + InputVec inVec = reinterpret_cast(in)[vecOffset]; + uint8_t sfValue; + out[vecOffset] = cvt_warp_fp16_to_fp4( + inVec, + SFScaleVal, +#if EXPERIMENT_PRECOMPUTE_SF + SFScaleForMax, +#endif + nullptr, + &sfValue); + + uint32_t sf = uint32_t(sfValue); + int32_t groupLane = threadIdx.x & ~7; + uint32_t sf1 = __shfl_sync(0xffffffff, sf, groupLane + 2); + uint32_t sf2 = __shfl_sync(0xffffffff, sf, groupLane + 4); + uint32_t sf3 = __shfl_sync(0xffffffff, sf, groupLane + 6); + if ((threadIdx.x & 7) == 0) { + uint32_t packedSF = sf | (sf1 << 8) | (sf2 << 16) | (sf3 << 24); + *reinterpret_cast(reinterpret_cast(SFout) + sfOffset) = packedSF; + } + } + } +} - // Input tensor row/col loops. - for (int rowIdx = blockIdx.x; rowIdx < numRows; rowIdx += gridDim.x) { - for (int colIdx = threadIdx.x; colIdx < numCols / CVT_FP4_ELTS_PER_THREAD; colIdx += blockDim.x) { - int64_t inOffset = rowIdx * (numCols / CVT_FP4_ELTS_PER_THREAD) + colIdx; - PackedVec in_vec = reinterpret_cast(in)[inOffset]; - // Get the output tensor offset. - // Same as inOffset because 8 elements are packed into one uint32_t. - int64_t outOffset = inOffset; - auto& out_pos = out[outOffset]; - - auto sf_out = - cvt_quant_to_fp4_get_sf_out_offset(rowIdx, colIdx, numCols, SFout); - - out_pos = cvt_warp_fp16_to_fp4(in_vec, SFScaleVal, sf_out); +template +__global__ __maxnreg__(EXPERIMENT_TILE4_MAX_REGISTERS) +void cvt_fp16_to_fp4_tile4_sf( + int32_t numRows, Type const* in, float const* SFScale, uint32_t* out, uint32_t* SFout) { + using InputVec = PackedVec; + constexpr int32_t numColVecs = 5120 / CVT_FP4_ELTS_PER_THREAD; + constexpr int32_t numKTiles = 5120 / (CVT_FP4_SF_VEC_SIZE * 4); + constexpr int32_t outerMPerTask = EXPERIMENT_TILE4_OUTER_M; + constexpr int32_t scaleBuffers = EXPERIMENT_TILE4_DOUBLE_BUFFER ? 2 : 1; + static_assert(outerMPerTask == 1 || outerMPerTask == 2 || outerMPerTask == 4 || outerMPerTask == 8); + __shared__ uint32_t sfStageStorage[scaleBuffers][outerMPerTask][4][numKTiles]; + float const SFScaleVal = SFScale == nullptr ? 1.0f : SFScale[0]; +#if EXPERIMENT_PRECOMPUTE_SF + float const SFScaleForMax = SFScaleVal * 0.16666666666666666f; +#endif + int32_t numMTiles = (numRows + 127) / 128; + constexpr int32_t tasksPerMTile = 32 / outerMPerTask; + int32_t numTasks = numMTiles * tasksPerMTile; +#if EXPERIMENT_INCREMENTAL_TASK + int32_t taskInMTile = blockIdx.x % tasksPerMTile; + int32_t mTileIdx = blockIdx.x / tasksPerMTile; + int32_t gridMTileStride = gridDim.x / tasksPerMTile; + int32_t gridTaskRemainder = gridDim.x % tasksPerMTile; +#endif + + for (int32_t taskIdx = blockIdx.x, taskIteration = 0; taskIdx < numTasks; + taskIdx += gridDim.x, ++taskIteration) { +#if EXPERIMENT_INCREMENTAL_TASK + int32_t outerMBase = taskInMTile * outerMPerTask; +#else + int32_t mTileIdx = taskIdx / tasksPerMTile; + int32_t outerMBase = (taskIdx % tasksPerMTile) * outerMPerTask; +#endif + uint32_t (*sfStage)[4][numKTiles] = sfStageStorage[taskIteration % scaleBuffers]; +#if EXPERIMENT_TILE4_DIRECT_SCALE + static_assert(outerMPerTask == 1); + int64_t sfTileBase = int64_t(mTileIdx) * numKTiles * 512 + int64_t(outerMBase) * 16; + for (int32_t colIdx = threadIdx.x; colIdx < numColVecs; colIdx += blockDim.x) { + uint32_t packedSF[4]; +#pragma unroll + for (int32_t innerM = 0; innerM < 4; ++innerM) { + int32_t rowIdx = mTileIdx * 128 + innerM * 32 + outerMBase; + if (rowIdx < numRows) { + int64_t vecOffset = int64_t(rowIdx) * numColVecs + colIdx; + InputVec inVec = load_tile4_vec(in + vecOffset * CVT_FP4_ELTS_PER_THREAD); + uint8_t sfValue; + out[vecOffset] = cvt_warp_fp16_to_fp4( + inVec, + SFScaleVal, +#if EXPERIMENT_PRECOMPUTE_SF + SFScaleForMax, +#endif + nullptr, + &sfValue); + packedSF[innerM] = pack_scale_byte(sfValue); + } else { + packedSF[innerM] = 0; + } + } + if ((threadIdx.x & 7) == 0) { + uint4 packed = make_uint4(packedSF[0], packedSF[1], packedSF[2], packedSF[3]); + *reinterpret_cast(reinterpret_cast(SFout) + sfTileBase + + int64_t(colIdx / 8) * 512) = packed; + } + } +#else +#if EXPERIMENT_TILE4_PIPELINE == 2 + static_assert(outerMPerTask == 2); +#pragma unroll + for (int32_t innerM = 0; innerM < 4; ++innerM) { + int32_t rowIdx0 = mTileIdx * 128 + innerM * 32 + outerMBase; + int32_t rowIdx1 = rowIdx0 + 1; + bool valid0 = rowIdx0 < numRows; + bool valid1 = rowIdx1 < numRows; + if (valid0 && valid1) { + int64_t rowVecBase0 = int64_t(rowIdx0) * numColVecs; + int64_t rowVecBase1 = int64_t(rowIdx1) * numColVecs; + for (int32_t colIdx = threadIdx.x; colIdx < numColVecs; colIdx += blockDim.x) { + int64_t vecOffset0 = rowVecBase0 + colIdx; + int64_t vecOffset1 = rowVecBase1 + colIdx; + InputVec inVec0 = load_tile4_vec(in + vecOffset0 * CVT_FP4_ELTS_PER_THREAD); + InputVec inVec1 = load_tile4_vec(in + vecOffset1 * CVT_FP4_ELTS_PER_THREAD); + uint8_t sfValue0; + uint8_t sfValue1; + out[vecOffset0] = cvt_warp_fp16_to_fp4( + inVec0, + SFScaleVal, +#if EXPERIMENT_PRECOMPUTE_SF + SFScaleForMax, +#endif + nullptr, + &sfValue0); + out[vecOffset1] = cvt_warp_fp16_to_fp4( + inVec1, + SFScaleVal, +#if EXPERIMENT_PRECOMPUTE_SF + SFScaleForMax, +#endif + nullptr, + &sfValue1); + stage_scale_byte(sfValue0, colIdx, sfStage[0][innerM]); + stage_scale_byte(sfValue1, colIdx, sfStage[1][innerM]); + } + } else { +#pragma unroll + for (int32_t outerMOffset = 0; outerMOffset < 2; ++outerMOffset) { + int32_t rowIdx = rowIdx0 + outerMOffset; + if (rowIdx < numRows) { + int64_t rowVecBase = int64_t(rowIdx) * numColVecs; + for (int32_t colIdx = threadIdx.x; colIdx < numColVecs; colIdx += blockDim.x) { + int64_t vecOffset = rowVecBase + colIdx; + InputVec inVec = load_tile4_vec(in + vecOffset * CVT_FP4_ELTS_PER_THREAD); + uint8_t sfValue; + out[vecOffset] = cvt_warp_fp16_to_fp4( + inVec, + SFScaleVal, +#if EXPERIMENT_PRECOMPUTE_SF + SFScaleForMax, +#endif + nullptr, + &sfValue); + stage_scale_byte(sfValue, colIdx, sfStage[outerMOffset][innerM]); + } + } else { + for (int32_t kTile = threadIdx.x; kTile < numKTiles; kTile += blockDim.x) { + sfStage[outerMOffset][innerM][kTile] = 0; + } + } + } + } + } +#else +#pragma unroll + for (int32_t outerMOffset = 0; outerMOffset < outerMPerTask; ++outerMOffset) { + int32_t outerM = outerMBase + outerMOffset; +#pragma unroll + for (int32_t innerM = 0; innerM < 4; ++innerM) { + int32_t rowIdx = mTileIdx * 128 + innerM * 32 + outerM; + if (rowIdx < numRows) { + int64_t rowVecBase = int64_t(rowIdx) * numColVecs; + for (int32_t colIdx = threadIdx.x; colIdx < numColVecs; +#if EXPERIMENT_TILE4_PIPELINE == 1 + colIdx += blockDim.x * 2 +#else + colIdx += blockDim.x +#endif + ) { + int64_t vecOffset = rowVecBase + colIdx; + InputVec inVec = load_tile4_vec(in + vecOffset * CVT_FP4_ELTS_PER_THREAD); +#if EXPERIMENT_TILE4_PIPELINE == 1 + int32_t colIdx1 = colIdx + blockDim.x; + int64_t vecOffset1 = rowVecBase + colIdx1; + InputVec inVec1; + if (colIdx1 < numColVecs) { + inVec1 = load_tile4_vec(in + vecOffset1 * CVT_FP4_ELTS_PER_THREAD); + } +#endif + uint8_t sfValue; + out[vecOffset] = cvt_warp_fp16_to_fp4( + inVec, + SFScaleVal, +#if EXPERIMENT_PRECOMPUTE_SF + SFScaleForMax, +#endif + nullptr, + &sfValue); + + stage_scale_byte(sfValue, colIdx, sfStage[outerMOffset][innerM]); +#if EXPERIMENT_TILE4_PIPELINE == 1 + if (colIdx1 < numColVecs) { + uint8_t sfValue1; + out[vecOffset1] = cvt_warp_fp16_to_fp4( + inVec1, + SFScaleVal, +#if EXPERIMENT_PRECOMPUTE_SF + SFScaleForMax, +#endif + nullptr, + &sfValue1); + stage_scale_byte(sfValue1, colIdx1, sfStage[outerMOffset][innerM]); + } +#endif + } + } else { + for (int32_t kTile = threadIdx.x; kTile < numKTiles; kTile += blockDim.x) { + sfStage[outerMOffset][innerM][kTile] = 0; + } + } + } } +#endif + __syncthreads(); + + int64_t sfTileBase = int64_t(mTileIdx) * numKTiles * 512 + int64_t(outerMBase) * 16; + for (int32_t storeIdx = threadIdx.x; storeIdx < numKTiles * outerMPerTask; + storeIdx += blockDim.x) { + int32_t kTile = storeIdx / outerMPerTask; + int32_t outerMOffset = storeIdx % outerMPerTask; + uint4 packed = make_uint4( + sfStage[outerMOffset][0][kTile], + sfStage[outerMOffset][1][kTile], + sfStage[outerMOffset][2][kTile], + sfStage[outerMOffset][3][kTile]); + *reinterpret_cast(reinterpret_cast(SFout) + sfTileBase + + int64_t(kTile) * 512 + outerMOffset * 16) = packed; + } +#if !EXPERIMENT_TILE4_DOUBLE_BUFFER + __syncthreads(); +#endif +#endif +#if EXPERIMENT_INCREMENTAL_TASK + mTileIdx += gridMTileStride; + taskInMTile += gridTaskRemainder; + if (taskInMTile >= tasksPerMTile) { + taskInMTile -= tasksPerMTile; + ++mTileIdx; + } +#endif } -// #endif } template @@ -295,17 +602,46 @@ void invokeFP4Quantization( cudaStream_t stream) { // Grid, Block size. // Each thread converts 8 values. - dim3 block(std::min(int(n / ELTS_PER_THREAD), 256)); - // Get number of blocks per SM (assume we can fully utilize the SM). - int const numBlocksPerSM = 1536 / block.x; - dim3 grid(std::min(int(m), multiProcessorCount * numBlocksPerSM)); + int blockThreads = EXPERIMENT_LAUNCH_THREADS; + if (char const* value = std::getenv("FP4_QUANT_BLOCK_THREADS")) { + blockThreads = std::atoi(value); + } + int gridBlocks = multiProcessorCount * EXPERIMENT_BLOCKS_PER_SM; + if (char const* value = std::getenv("FP4_QUANT_GRID_BLOCKS")) { + gridBlocks = std::atoi(value); + } + TORCH_CHECK(gridBlocks > 0); + if (n == 5120) { + if (std::getenv("FP4_QUANT_BLOCK_THREADS") == nullptr) { + blockThreads = TILE4_BLOCK_THREADS; + } + if (std::getenv("FP4_QUANT_GRID_BLOCKS") == nullptr) { + gridBlocks = + multiProcessorCount * TILE4_GRID_BLOCKS_NUMERATOR / TILE4_GRID_BLOCKS_DENOMINATOR; + } + TORCH_CHECK(blockThreads > 0 && blockThreads <= TILE4_MAX_BLOCK_THREADS && blockThreads % 32 == 0); + dim3 block(blockThreads); + dim3 grid(std::min(int(m), gridBlocks)); + if (useUE8M0) { + cvt_fp16_to_fp4_tile4_sf<<>>( + m, input, SFScale, reinterpret_cast(output), reinterpret_cast(SFOuput)); + } else { + cvt_fp16_to_fp4_tile4_sf<<>>( + m, input, SFScale, reinterpret_cast(output), reinterpret_cast(SFOuput)); + } + return; + } + + TORCH_CHECK(blockThreads > 0 && blockThreads <= EXPERIMENT_BLOCK_THREADS && blockThreads % 32 == 0); + dim3 block(std::min(int(n / ELTS_PER_THREAD), blockThreads)); + dim3 grid(std::min(int(m), gridBlocks)); - // Launch the cvt kernel. + // Launch the generic conversion kernel. if (useUE8M0) { - cvt_fp16_to_fp4<<>>( + cvt_fp16_to_fp4_incremental_sf<<>>( m, n, input, SFScale, reinterpret_cast(output), reinterpret_cast(SFOuput)); } else { - cvt_fp16_to_fp4<<>>( + cvt_fp16_to_fp4_incremental_sf<<>>( m, n, input, SFScale, reinterpret_cast(output), reinterpret_cast(SFOuput)); } } @@ -362,7 +698,7 @@ void scaled_nvfp4_quant_sm120( auto input_sf_ptr = static_cast(input_sf.data_ptr()); auto sf_out = static_cast(output_sf.data_ptr()); auto output_ptr = static_cast(output.data_ptr()); - at::cuda::CUDAGuard device_guard{(char)input.get_device()}; + at::cuda::CUDAGuard device_guard{input.get_device()}; const cudaStream_t stream = at::cuda::getCurrentCUDAStream(input.get_device()); // We don't support e8m0 scales at this moment. diff --git a/lightx2v_kernel/csrc/gemm/nvfp4_scaled_mm_kernels_sm120.cu b/lightx2v_kernel/csrc/gemm/nvfp4_scaled_mm_kernels_sm120.cu index 12d1adebd..bff7bb01f 100644 --- a/lightx2v_kernel/csrc/gemm/nvfp4_scaled_mm_kernels_sm120.cu +++ b/lightx2v_kernel/csrc/gemm/nvfp4_scaled_mm_kernels_sm120.cu @@ -5,6 +5,7 @@ // clang-format off #include "cutlass/cutlass.h" #include "cutlass/epilogue/fusion/operations.hpp" +#include "cutlass/epilogue/fusion/sm90_callbacks_tma_warpspecialized.hpp" #include "cutlass/gemm/collective/collective_builder.hpp" #include "cutlass/epilogue/collective/collective_builder.hpp" #include "cutlass/gemm/device/gemm_universal_adapter.h" @@ -29,8 +30,244 @@ using namespace cute; +namespace cutlass::epilogue::fusion { +template +struct PyTorchTanhGelu { + CUTLASS_HOST_DEVICE + T operator()(T const& x) const { + T const beta = T(0.7978845608028654); + T const kappa = T(0.044715); + return T(0.5) * x * (T(1) + ::tanhf(beta * (x + kappa * x * x * x))); + } +}; + +template +struct PyTorchTanhGelu> { + CUTLASS_HOST_DEVICE + Array operator()(Array const& input) const { + Array output; + PyTorchTanhGelu gelu; + CUTLASS_PRAGMA_UNROLL + for (int index = 0; index < N; ++index) { + output[index] = gelu(input[index]); + } + return output; + } +}; + + +template< + class ElementOutput_, + class ElementCompute_, + class ElementGate_ = ElementOutput_, + class ElementBias_ = ElementOutput_, + class ElementSource_ = ElementOutput_, + class ElementScalar_ = ElementCompute_, + int AlignmentGate_ = 128 / cute::sizeof_bits_v, + int AlignmentBias_ = 128 / cute::sizeof_bits_v, + FloatRoundStyle RoundStyle_ = FloatRoundStyle::round_to_nearest +> +struct LinCombPerColBiasBf16GateResidual + : LinearCombination { + using ElementGate = ElementGate_; + using ElementBias = ElementBias_; + static constexpr int AlignmentGate = AlignmentGate_; + static constexpr int AlignmentBias = AlignmentBias_; + static constexpr bool IsPerColBiasSupported = true; +}; + +template< + class CtaTileShapeMNK, + class ElementOutput, + class ElementCompute, + class ElementGate = ElementOutput, + class ElementBias = ElementOutput, + class ElementSource = ElementOutput, + class ElementScalar = ElementCompute, + int AlignmentGate = 128 / sizeof_bits_v, + int AlignmentBias = 128 / sizeof_bits_v, + FloatRoundStyle RoundStyle = FloatRoundStyle::round_to_nearest +> +using Sm90LinCombPerColBiasBf16GateResidual = + // ElementOutput at each nested compute preserves the original BF16 + // materialization boundaries: GEMM+bias, gate multiply, then residual add. + // Parent nodes convert those BF16 fragments back to ElementCompute. + Sm90EVT, + Sm90SrcFetch, + Sm90EVT, + Sm90EVT, + Sm90ScalarBroadcast>, + Sm90AccFetch, + Sm90RowBroadcast<0, CtaTileShapeMNK, ElementBias, ElementCompute, Stride<_0,_1,int64_t>, AlignmentBias> + >, + Sm90RowBroadcast<0, CtaTileShapeMNK, ElementGate, ElementCompute, Stride<_0,_1,int64_t>, AlignmentGate> + > + >; + +template < + int StagesC, + int StagesD, + int FragmentSize, + bool ReuseSmemC, + bool DelayTmaStore, + class ElementOutput, + class ElementCompute, + class ElementGate, + class ElementBias, + class ElementSource, + class ElementScalar, + int AlignmentGate, + int AlignmentBias, + FloatRoundStyle RoundStyle, + class CtaTileShapeMNK, + class EpilogueTile +> +struct FusionCallbacks< + epilogue::Sm90TmaWarpSpecialized, + fusion::LinCombPerColBiasBf16GateResidual< + ElementOutput, ElementCompute, ElementGate, ElementBias, ElementSource, ElementScalar, + AlignmentGate, AlignmentBias, RoundStyle>, + CtaTileShapeMNK, + EpilogueTile +> : Sm90LinCombPerColBiasBf16GateResidual< + CtaTileShapeMNK, ElementOutput, ElementCompute, ElementGate, ElementBias, ElementSource, + ElementScalar, AlignmentGate, AlignmentBias, RoundStyle> { + using Impl = Sm90LinCombPerColBiasBf16GateResidual< + CtaTileShapeMNK, ElementOutput, ElementCompute, ElementGate, ElementBias, ElementSource, + ElementScalar, AlignmentGate, AlignmentBias, RoundStyle>; + using Operation = fusion::LinCombPerColBiasBf16GateResidual< + ElementOutput, ElementCompute, ElementGate, ElementBias, ElementSource, ElementScalar, + AlignmentGate, AlignmentBias, RoundStyle>; + + struct Arguments { + ElementScalar alpha = ElementScalar(1); + ElementScalar const* alpha_ptr = nullptr; + using StrideAlpha = Stride<_0,_0,int64_t>; + StrideAlpha dAlpha = {_0{}, _0{}, 0}; + + using StrideGate = Stride<_0,_1,int64_t>; + ElementGate const* gate_ptr = nullptr; + StrideGate dGate = {}; + + using StrideBias = Stride<_0,_1,int64_t>; + ElementBias const* bias_ptr = nullptr; + StrideBias dBias = {}; + + operator typename Impl::Arguments() const { + return { + {}, + { + {{{alpha}, {alpha_ptr}, {dAlpha}}, {}, {bias_ptr, ElementBias(0), dBias}, {}}, + {gate_ptr, ElementGate(0), dGate}, + {} + }, + {} + }; + } + }; + + using Impl::Impl; +}; + +template< + class ElementOutput_, + class ElementCompute_, + class ElementBias_ = ElementOutput_, + class ElementSource_ = ElementOutput_, + class ElementScalar_ = ElementCompute_, + int AlignmentBias_ = 128 / cute::sizeof_bits_v, + FloatRoundStyle RoundStyle_ = FloatRoundStyle::round_to_nearest +> +struct LinCombPerColBiasBf16Gelu + : LinearCombination { + using ElementBias = ElementBias_; + static constexpr int AlignmentBias = AlignmentBias_; + static constexpr bool IsPerColBiasSupported = true; +}; + +template< + class CtaTileShapeMNK, + class ElementOutput, + class ElementCompute, + class ElementBias = ElementOutput, + class ElementSource = ElementOutput, + class ElementScalar = ElementCompute, + int AlignmentBias = 128 / sizeof_bits_v, + FloatRoundStyle RoundStyle = FloatRoundStyle::round_to_nearest +> +using Sm90LinCombPerColBiasBf16Gelu = + // Match the unfused path's BF16 materialization before GELU. + Sm90EVT, + Sm90EVT, + Sm90ScalarBroadcast>, + Sm90AccFetch, + Sm90RowBroadcast<0, CtaTileShapeMNK, ElementBias, ElementCompute, Stride<_0,_1,int64_t>, AlignmentBias> + > + >; + +template < + int StagesC, + int StagesD, + int FragmentSize, + bool ReuseSmemC, + bool DelayTmaStore, + class ElementOutput, + class ElementCompute, + class ElementBias, + class ElementSource, + class ElementScalar, + int AlignmentBias, + FloatRoundStyle RoundStyle, + class CtaTileShapeMNK, + class EpilogueTile +> +struct FusionCallbacks< + epilogue::Sm90TmaWarpSpecialized, + fusion::LinCombPerColBiasBf16Gelu< + ElementOutput, ElementCompute, ElementBias, ElementSource, ElementScalar, + AlignmentBias, RoundStyle>, + CtaTileShapeMNK, + EpilogueTile +> : Sm90LinCombPerColBiasBf16Gelu< + CtaTileShapeMNK, ElementOutput, ElementCompute, ElementBias, ElementSource, + ElementScalar, AlignmentBias, RoundStyle> { + using Impl = Sm90LinCombPerColBiasBf16Gelu< + CtaTileShapeMNK, ElementOutput, ElementCompute, ElementBias, ElementSource, + ElementScalar, AlignmentBias, RoundStyle>; + using Operation = fusion::LinCombPerColBiasBf16Gelu< + ElementOutput, ElementCompute, ElementBias, ElementSource, ElementScalar, + AlignmentBias, RoundStyle>; + + struct Arguments { + ElementScalar alpha = ElementScalar(1); + ElementScalar const* alpha_ptr = nullptr; + using StrideAlpha = Stride<_0,_0,int64_t>; + StrideAlpha dAlpha = {_0{}, _0{}, 0}; + + using StrideBias = Stride<_0,_1,int64_t>; + ElementBias const* bias_ptr = nullptr; + StrideBias dBias = {}; + + operator typename Impl::Arguments() const { + return { + {{{alpha}, {alpha_ptr}, {dAlpha}}, {}, {bias_ptr, ElementBias(0), dBias}, {}}, + {} + }; + } + }; + + using Impl::Impl; +}; -struct Fp4GemmSm120 { +} // namespace cutlass::epilogue::fusion + + +template < + class ThreadBlockShape_, + class ClusterShape_, + class MainloopSchedule_, + class EpilogueSchedule_> +struct Fp4GemmSm120Config { ///////////////////////////////////////////////////////////////////////////////////////////////// /// GEMM kernel configurations ///////////////////////////////////////////////////////////////////////////////////////////////// @@ -54,12 +291,16 @@ struct Fp4GemmSm120 { static constexpr int AlignmentC = 128 / cutlass::sizeof_bits::value; // Memory access granularity/alignment of C matrix in units of elements (up to 16 bytes) // Kernel functional config using ElementAccumulator = float; // Element type for internal accumulation - using ArchTag = cutlass::arch::Sm120; // Tag indicating the minimum SM that supports the intended feature +#if defined(LIGHTX2V_THOR_NVFP4_ONLY) + using ArchTag = cutlass::arch::Sm100; +#else + using ArchTag = cutlass::arch::Sm120; +#endif using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp; // Operator class tag // Kernel Perf config - using ThreadBlockShape = Shape<_128,_128,_128>; // Threadblock's tile size - using ClusterShape = Shape<_1,_1,_1>; // Shape of the threadblocks in a cluster + using ThreadBlockShape = ThreadBlockShape_; + using ClusterShape = ClusterShape_; // use per-column bias, i.e. every column has different bias using EVTOp = cutlass::epilogue::fusion::LinCombPerColBias; @@ -71,7 +312,7 @@ struct Fp4GemmSm120 { ElementAccumulator, ElementAccumulator, ElementC, LayoutCTag, AlignmentC, ElementD, LayoutDTag, AlignmentD, - cutlass::epilogue::collective::EpilogueScheduleAuto, // Epilogue schedule policy + EpilogueSchedule_, EVTOp >::CollectiveOp; @@ -82,7 +323,7 @@ struct Fp4GemmSm120 { ElementAccumulator, ThreadBlockShape, ClusterShape, cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, - cutlass::gemm::collective::KernelScheduleAuto // Kernel schedule policy. Auto defaults to cooperative kernel schedule + MainloopSchedule_ >::CollectiveOp; using GemmKernel = cutlass::gemm::kernel::GemmUniversal< @@ -106,6 +347,154 @@ struct Fp4GemmSm120 { using LayoutD = decltype(cute::make_layout(make_shape(0,0,0), StrideD{})); }; +using Fp4GemmSm120 = Fp4GemmSm120Config< + Shape<_128,_128,_128>, Shape<_1,_1,_1>, + cutlass::gemm::collective::KernelScheduleAuto, + cutlass::epilogue::collective::EpilogueScheduleAuto>; + +using Fp4GemmSm120Wan22Ffn2 = Fp4GemmSm120Config< + Shape<_256,_256,_256>, Shape<_2,_1,_1>, + cutlass::gemm::KernelTmaWarpSpecialized2SmNvf4Sm100, + cutlass::epilogue::TmaWarpSpecialized2SmNvf4>; + +template < + class ThreadBlockShape_, + class ClusterShape_, + class MainloopSchedule_, + class EpilogueSchedule_> +struct Fp4GemmResidualGateSm120Config { + using ElementA = cutlass::nv_float4_t; + using LayoutATag = cutlass::layout::RowMajor; + static constexpr int AlignmentA = 32; + using ElementB = cutlass::nv_float4_t; + using LayoutBTag = cutlass::layout::ColumnMajor; + static constexpr int AlignmentB = 32; + + using ElementD = cutlass::bfloat16_t; + using ElementC = cutlass::bfloat16_t; + using LayoutCTag = cutlass::layout::RowMajor; + using LayoutDTag = cutlass::layout::RowMajor; + static constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; + static constexpr int AlignmentC = 128 / cutlass::sizeof_bits::value; + using ElementAccumulator = float; +#if defined(LIGHTX2V_THOR_NVFP4_ONLY) + using ArchTag = cutlass::arch::Sm100; +#else + using ArchTag = cutlass::arch::Sm120; +#endif + using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp; + + using ThreadBlockShape = ThreadBlockShape_; + using ClusterShape = ClusterShape_; + + using EVTOp = cutlass::epilogue::fusion::LinCombPerColBiasBf16GateResidual< + ElementD, ElementAccumulator, ElementD, ElementD, ElementC, ElementAccumulator>; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag, OperatorClass, + ThreadBlockShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementAccumulator, + ElementC, LayoutCTag, AlignmentC, + ElementD, LayoutDTag, AlignmentD, + EpilogueSchedule_, + EVTOp + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, OperatorClass, + ElementA, LayoutATag, AlignmentA, + ElementB, LayoutBTag, AlignmentB, + ElementAccumulator, + ThreadBlockShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule_ + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, CollectiveMainloop, CollectiveEpilogue, void>; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + using StrideA = typename Gemm::GemmKernel::StrideA; + using LayoutSFA = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFA; + using StrideB = typename Gemm::GemmKernel::StrideB; + using LayoutSFB = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFB; + using StrideC = typename Gemm::GemmKernel::StrideC; + using StrideD = typename Gemm::GemmKernel::StrideD; +}; + +using Fp4GemmResidualGateSm120 = Fp4GemmResidualGateSm120Config< + Shape<_128,_128,_128>, Shape<_1,_1,_1>, + cutlass::gemm::collective::KernelScheduleAuto, + cutlass::epilogue::collective::EpilogueScheduleAuto>; + +using Fp4GemmResidualGateSm120Wan22Ffn2 = Fp4GemmResidualGateSm120Config< + Shape<_256,_256,_256>, Shape<_2,_1,_1>, + cutlass::gemm::KernelTmaWarpSpecialized2SmNvf4Sm100, + cutlass::epilogue::TmaWarpSpecialized2SmNvf4>; + +struct Fp4GemmGeluSm120 { + using ElementA = cutlass::nv_float4_t; + using LayoutATag = cutlass::layout::RowMajor; + static constexpr int AlignmentA = 32; + using ElementB = cutlass::nv_float4_t; + using LayoutBTag = cutlass::layout::ColumnMajor; + static constexpr int AlignmentB = 32; + + using ElementD = cutlass::bfloat16_t; + using ElementC = cutlass::bfloat16_t; + using LayoutCTag = cutlass::layout::RowMajor; + using LayoutDTag = cutlass::layout::RowMajor; + static constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; + static constexpr int AlignmentC = 128 / cutlass::sizeof_bits::value; + using ElementAccumulator = float; +#if defined(LIGHTX2V_THOR_NVFP4_ONLY) + using ArchTag = cutlass::arch::Sm100; +#else + using ArchTag = cutlass::arch::Sm120; +#endif + using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp; + + using ThreadBlockShape = Shape<_256,_256,_256>; + using ClusterShape = Shape<_2,_1,_1>; + + using EVTOp = cutlass::epilogue::fusion::LinCombPerColBiasBf16Gelu< + ElementD, ElementAccumulator, ElementD, ElementC, ElementAccumulator>; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag, OperatorClass, + ThreadBlockShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementAccumulator, + ElementC, LayoutCTag, AlignmentC, + ElementD, LayoutDTag, AlignmentD, + cutlass::epilogue::TmaWarpSpecialized2SmNvf4, + EVTOp + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, OperatorClass, + ElementA, LayoutATag, AlignmentA, + ElementB, LayoutBTag, AlignmentB, + ElementAccumulator, + ThreadBlockShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized2SmNvf4Sm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, CollectiveMainloop, CollectiveEpilogue, void>; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + using StrideA = typename Gemm::GemmKernel::StrideA; + using LayoutSFA = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFA; + using StrideB = typename Gemm::GemmKernel::StrideB; + using LayoutSFB = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFB; + using StrideC = typename Gemm::GemmKernel::StrideC; + using StrideD = typename Gemm::GemmKernel::StrideD; +}; // Populates a Gemm::Arguments structure from the given commandline options typename Fp4GemmSm120::Gemm::Arguments args_from_options_nvfp4_nvfp4( @@ -216,6 +605,207 @@ void runGemmNvfp4Sm120( CUTLASS_CHECK(gemm.run(arguments, workspace.data_ptr(), stream)); } +template +typename GemmConfig::Gemm::Arguments +args_from_options_nvfp4_split_n_stride_residual_gate( + at::Tensor& residual, + at::Tensor const& A, + at::Tensor const& B, + at::Tensor const& A_sf, + at::Tensor const& B_sf, + at::Tensor const& alpha, + c10::optional const& bias, + at::Tensor const& gate, + int64_t M, + int64_t N, + int64_t K, + int64_t split_n_parts) { + using Sm1xxBlkScaledConfig = + typename GemmConfig::Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; + + int m = static_cast(M); + int n = static_cast(N); + int k = static_cast(K); + int batch_count = static_cast(split_n_parts); + int shard_n = n / batch_count; + + auto stride_A = cutlass::make_cute_packed_stride( + typename GemmConfig::StrideA{}, {m, k, batch_count}); + auto stride_B = cutlass::make_cute_packed_stride( + typename GemmConfig::StrideB{}, {shard_n, k, batch_count}); + auto stride_D = cutlass::make_cute_packed_stride( + typename GemmConfig::StrideD{}, {m, shard_n, batch_count}); + cute::get<2>(stride_A) = 0; + cute::get<0>(stride_D) = n; + cute::get<2>(stride_D) = shard_n; + + auto problem_shape = cute::make_shape(m, shard_n, k, batch_count); + auto layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(problem_shape); + auto layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(problem_shape); + cute::get<2, 1>(cute::stride(layout_SFA)) = 0; + + typename GemmConfig::Gemm::Arguments arguments{ + cutlass::gemm::GemmUniversalMode::kBatched, + problem_shape, + { + static_cast(A.data_ptr()), + stride_A, + static_cast(B.data_ptr()), + stride_B, + static_cast(A_sf.data_ptr()), + layout_SFA, + static_cast(B_sf.data_ptr()), + layout_SFB + }, + { + {}, + static_cast(residual.data_ptr()), + stride_D, + static_cast(residual.data_ptr()), + stride_D + } + }; + + auto& fusion_args = arguments.epilogue.thread; + fusion_args.alpha_ptr = static_cast(alpha.data_ptr()); + fusion_args.gate_ptr = static_cast(gate.data_ptr()); + using StrideGate = Stride; + auto gate_stride = StrideGate{}; + cute::get<2>(gate_stride) = shard_n; + fusion_args.dGate = gate_stride; + if (bias) { + fusion_args.bias_ptr = + static_cast(bias->data_ptr()); + using StrideBias = Stride; + auto bias_stride = StrideBias{}; + cute::get<2>(bias_stride) = shard_n; + fusion_args.dBias = bias_stride; + } + return arguments; +} + +template +void runGemmNvfp4SplitNStrideResidualGateSm120( + at::Tensor& residual, + at::Tensor const& A, + at::Tensor const& B, + at::Tensor const& A_sf, + at::Tensor const& B_sf, + at::Tensor const& alpha, + c10::optional const& bias, + at::Tensor const& gate, + int64_t m, + int64_t n, + int64_t k, + int64_t split_n_parts, + cudaStream_t stream) { + typename GemmConfig::Gemm gemm; + auto arguments = args_from_options_nvfp4_split_n_stride_residual_gate( + residual, A, B, A_sf, B_sf, alpha, bias, gate, m, n, k, split_n_parts); + size_t workspace_size = GemmConfig::Gemm::get_workspace_size(arguments); + auto workspace = torch::empty( + workspace_size, torch::TensorOptions().dtype(torch::kUInt8).device(A.device())); + + CUTLASS_CHECK(gemm.can_implement(arguments)); + CUTLASS_CHECK(gemm.initialize(arguments, workspace.data_ptr(), stream)); + CUTLASS_CHECK(gemm.run(arguments, workspace.data_ptr(), stream)); +} + +typename Fp4GemmGeluSm120::Gemm::Arguments +args_from_options_nvfp4_split_n_stride_gelu( + at::Tensor& output, + at::Tensor const& A, + at::Tensor const& B, + at::Tensor const& A_sf, + at::Tensor const& B_sf, + at::Tensor const& alpha, + c10::optional const& bias, + int64_t M, + int64_t N, + int64_t K, + int64_t split_n_parts) { + using Sm1xxBlkScaledConfig = + typename Fp4GemmGeluSm120::Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; + + int m = static_cast(M); + int n = static_cast(N); + int k = static_cast(K); + int batch_count = static_cast(split_n_parts); + int shard_n = n / batch_count; + + auto stride_A = cutlass::make_cute_packed_stride( + Fp4GemmGeluSm120::StrideA{}, {m, k, batch_count}); + auto stride_B = cutlass::make_cute_packed_stride( + Fp4GemmGeluSm120::StrideB{}, {shard_n, k, batch_count}); + auto stride_D = cutlass::make_cute_packed_stride( + Fp4GemmGeluSm120::StrideD{}, {m, shard_n, batch_count}); + cute::get<2>(stride_A) = 0; + cute::get<0>(stride_D) = n; + cute::get<2>(stride_D) = shard_n; + + auto problem_shape = cute::make_shape(m, shard_n, k, batch_count); + auto layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(problem_shape); + auto layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(problem_shape); + cute::get<2, 1>(cute::stride(layout_SFA)) = 0; + + typename Fp4GemmGeluSm120::Gemm::Arguments arguments{ + cutlass::gemm::GemmUniversalMode::kBatched, + problem_shape, + { + static_cast(A.data_ptr()), + stride_A, + static_cast(B.data_ptr()), + stride_B, + static_cast(A_sf.data_ptr()), + layout_SFA, + static_cast(B_sf.data_ptr()), + layout_SFB + }, + { + {}, + static_cast(output.data_ptr()), + stride_D, + static_cast(output.data_ptr()), + stride_D + } + }; + + auto& fusion_args = arguments.epilogue.thread; + fusion_args.alpha_ptr = static_cast(alpha.data_ptr()); + if (bias) { + fusion_args.bias_ptr = static_cast(bias->data_ptr()); + using StrideBias = Stride; + auto bias_stride = StrideBias{}; + cute::get<2>(bias_stride) = shard_n; + fusion_args.dBias = bias_stride; + } + return arguments; +} + +void runGemmNvfp4SplitNStrideGeluSm120( + at::Tensor& output, + at::Tensor const& A, + at::Tensor const& B, + at::Tensor const& A_sf, + at::Tensor const& B_sf, + at::Tensor const& alpha, + c10::optional const& bias, + int64_t m, + int64_t n, + int64_t k, + int64_t split_n_parts, + cudaStream_t stream) { + typename Fp4GemmGeluSm120::Gemm gemm; + auto arguments = args_from_options_nvfp4_split_n_stride_gelu( + output, A, B, A_sf, B_sf, alpha, bias, m, n, k, split_n_parts); + size_t workspace_size = Fp4GemmGeluSm120::Gemm::get_workspace_size(arguments); + auto workspace = torch::empty( + workspace_size, torch::TensorOptions().dtype(torch::kUInt8).device(A.device())); + + CUTLASS_CHECK(gemm.can_implement(arguments)); + CUTLASS_CHECK(gemm.initialize(arguments, workspace.data_ptr(), stream)); + CUTLASS_CHECK(gemm.run(arguments, workspace.data_ptr(), stream)); +} constexpr auto FLOAT4_E2M1X2 = at::ScalarType::Byte; constexpr auto SF_DTYPE = at::ScalarType::Float8_e4m3fn; @@ -326,3 +916,338 @@ void cutlass_scaled_nvfp4_mm_sm120( runGemmNvfp4Sm120(D, A, B, A_sf, B_sf, alpha, bias, m, n, k, stream); } + + +// Keep split-N stride argument construction and execution isolated from the +// regular NVFP4 operator so changes here cannot alter its behavior. +// prepare the calculation parameters for the gemm +template +typename GemmConfig::Gemm::Arguments args_from_options_nvfp4_nvfp4_split_n_stride( + at::Tensor& D, + at::Tensor const& A, + at::Tensor const& B, + at::Tensor const& A_sf, + at::Tensor const& B_sf, + at::Tensor const& alpha, + c10::optional const& bias, + int64_t M, + int64_t N, + int64_t K, + int64_t split_n_parts) { + using Sm1xxBlkScaledConfig = typename GemmConfig::Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; + + int m = static_cast(M); + int n = static_cast(N); + int k = static_cast(K); + int batch_count = static_cast(split_n_parts); + int shard_n = n / batch_count; + + auto stride_A = cutlass::make_cute_packed_stride(typename GemmConfig::StrideA{}, {m, k, batch_count}); + auto stride_B = cutlass::make_cute_packed_stride(typename GemmConfig::StrideB{}, {shard_n, k, batch_count}); + auto stride_D = cutlass::make_cute_packed_stride(typename GemmConfig::StrideD{}, {m, shard_n, batch_count}); + + // Broadcast A across batches. B batches are consecutive N shards. + cute::get<2>(stride_A) = 0; + // D is one [M, N] tensor: batch l starts at column l * shard_n. + cute::get<0>(stride_D) = n; + cute::get<2>(stride_D) = shard_n; + + auto problem_shape = cute::make_shape(m, shard_n, k, batch_count); + auto layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(problem_shape); + auto layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(problem_shape); + // SFB remains packed by batch. Broadcast only SFA's nested batch mode. + cute::get<2, 1>(cute::stride(layout_SFA)) = 0; + + if (bias) { + using StrideBias = Stride; + + typename GemmConfig::Gemm::Arguments arguments{ + cutlass::gemm::GemmUniversalMode::kBatched, + problem_shape, + {// Mainloop arguments + static_cast(A.data_ptr()), + stride_A, + static_cast(B.data_ptr()), + stride_B, + static_cast(A_sf.data_ptr()), + layout_SFA, + static_cast(B_sf.data_ptr()), + layout_SFB}, + { // Epilogue arguments + {}, // epilogue.thread + static_cast(D.data_ptr()), + stride_D, + static_cast(D.data_ptr()), + stride_D}}; + auto& fusion_args = arguments.epilogue.thread; + fusion_args.alpha_ptr = static_cast(alpha.data_ptr()); + fusion_args.bias_ptr = static_cast(bias->data_ptr()); + auto stride_bias = StrideBias{}; + cute::get<2>(stride_bias) = shard_n; + fusion_args.dBias = stride_bias; + return arguments; + } + else + { + typename GemmConfig::Gemm::Arguments arguments{ + cutlass::gemm::GemmUniversalMode::kBatched, + problem_shape, + {// Mainloop arguments + static_cast(A.data_ptr()), + stride_A, + static_cast(B.data_ptr()), + stride_B, + static_cast(A_sf.data_ptr()), + layout_SFA, + static_cast(B_sf.data_ptr()), + layout_SFB}, + { // Epilogue arguments + {}, // epilogue.thread + static_cast(D.data_ptr()), + stride_D, + static_cast(D.data_ptr()), + stride_D}}; + auto& fusion_args = arguments.epilogue.thread; + fusion_args.alpha_ptr = static_cast(alpha.data_ptr()); + return arguments; + } +} + +// implement the gemm for nvfp4splitnstride +template +void runGemmNvfp4SplitNStrideSm120( + at::Tensor& D, + at::Tensor const& A, + at::Tensor const& B, + at::Tensor const& A_sf, + at::Tensor const& B_sf, + at::Tensor const& alpha, + c10::optional const& bias, + int64_t m, + int64_t n, + int64_t k, + int64_t split_n_parts, + cudaStream_t stream) { + typename GemmConfig::Gemm gemm; + + auto arguments = args_from_options_nvfp4_nvfp4_split_n_stride( + D, A, B, A_sf, B_sf, alpha, bias, m, n, k, split_n_parts); + auto beta_dev = torch::zeros({1}, torch::TensorOptions() + .dtype(torch::kFloat32) + .device(A.device())); + arguments.epilogue.thread.beta_ptr = + static_cast(beta_dev.data_ptr()); + size_t workspace_size = GemmConfig::Gemm::get_workspace_size(arguments); + auto const workspace_options = torch::TensorOptions().dtype(torch::kUInt8).device(A.device()); + auto workspace = torch::empty(workspace_size, workspace_options); + + CUTLASS_CHECK(gemm.can_implement(arguments)); + CUTLASS_CHECK(gemm.initialize(arguments, workspace.data_ptr(), stream)); + CUTLASS_CHECK(gemm.run(arguments, workspace.data_ptr(), stream)); +} + +// check the inputs and run the NVFP4 split-N stride GEMM kernel +void check_nvfp4_split_n_stride_inputs( + torch::Tensor& D, + torch::Tensor const& A, + torch::Tensor const& B, + torch::Tensor const& A_sf, + torch::Tensor const& B_sf, + torch::Tensor const& alpha, + c10::optional const& bias, + int64_t split_n_parts) { + CHECK_INPUT(D, at::ScalarType::BFloat16, "out"); + CHECK_INPUT(A, FLOAT4_E2M1X2, "a"); + CHECK_INPUT(B, FLOAT4_E2M1X2, "b"); + CHECK_INPUT(A_sf, SF_DTYPE, "scale_a"); + CHECK_INPUT(B_sf, SF_DTYPE, "scale_b"); + CHECK_INPUT(alpha, at::ScalarType::Float, "alpha"); + + TORCH_CHECK(D.dim() == 2, "out must be a matrix"); + TORCH_CHECK(A.dim() == 2, "a must be a matrix"); + TORCH_CHECK(B.dim() == 2, "b must be a matrix"); + TORCH_CHECK( + A.sizes()[1] == B.sizes()[1], + "a and b shapes cannot be multiplied (", + A.sizes()[0], + "x", + A.sizes()[1], + " and ", + B.sizes()[0], + "x", + B.sizes()[1], + ")"); + + auto const m = A.sizes()[0]; + auto const n = B.sizes()[0]; + auto const k = A.sizes()[1] * 2; + + TORCH_CHECK( + D.sizes()[0] == m && D.sizes()[1] == n, + "out must have shape (", + m, + "x", + n, + "), but got (", + D.sizes()[0], + "x", + D.sizes()[1], + ")"); + if (bias) { + auto const& bias_tensor = bias.value(); + CHECK_INPUT(bias_tensor, at::ScalarType::BFloat16, "bias"); + TORCH_CHECK(bias_tensor.numel() == n, "bias must contain ", n, " elements, but got ", bias_tensor.numel()); + } + + constexpr int alignment = 32; + TORCH_CHECK( + k % alignment == 0, + "Expected k to be divisible by ", + alignment, + ", but got a shape: (", + A.sizes()[0], + "x", + A.sizes()[1], + "), k: ", + k, + "."); + TORCH_CHECK( + n % alignment == 0, + "Expected n to be divisible by ", + alignment, + ", but got b shape: (", + B.sizes()[0], + "x", + B.sizes()[1], + ")."); + + auto round_up = [](int x, int y) { return (x + y - 1) / y * y; }; + int rounded_m = round_up(m, 128); + int rounded_n = round_up(n, 128); + int rounded_k = round_up(k / 16, 4); + + TORCH_CHECK(A_sf.dim() == 2, "scale_a must be a matrix"); + TORCH_CHECK(B_sf.dim() == 2, "scale_b must be a matrix"); + TORCH_CHECK( + A_sf.sizes()[1] == B_sf.sizes()[1], + "scale_a and scale_b shapes cannot be multiplied (", + A_sf.sizes()[0], + "x", + A_sf.sizes()[1], + " and ", + B_sf.sizes()[0], + "x", + B_sf.sizes()[1], + ")"); + TORCH_CHECK( + A_sf.sizes()[0] == rounded_m && A_sf.sizes()[1] == rounded_k, + "scale_a must be padded and swizzled to a shape (", + rounded_m, + "x", + rounded_k, + "), but got a shape (", + A_sf.sizes()[0], + "x", + A_sf.sizes()[1], + ")"); + TORCH_CHECK( + B_sf.sizes()[0] == rounded_n && B_sf.sizes()[1] == rounded_k, + "scale_b must be padded and swizzled to a shape (", + rounded_n, + "x", + rounded_k, + "), but got a shape (", + B_sf.sizes()[0], + "x", + B_sf.sizes()[1], + ")"); + + TORCH_CHECK(split_n_parts >= 2, "split_n_parts must be at least 2, but got ", split_n_parts); + TORCH_CHECK( + n % split_n_parts == 0, + "Expected n to be divisible by split_n_parts, but got n=", + n, + " and split_n_parts=", + split_n_parts); + TORCH_CHECK( + (n / split_n_parts) % 128 == 0, + "Each NVFP4 split-N shard must be divisible by 128 for the swizzled scale layout, but got shard n=", + n / split_n_parts); + +} + +void cutlass_scaled_nvfp4_mm_split_n_stride_sm120( + torch::Tensor& D, + torch::Tensor const& A, + torch::Tensor const& B, + torch::Tensor const& A_sf, + torch::Tensor const& B_sf, + torch::Tensor const& alpha, + c10::optional const& bias, + int64_t split_n_parts) { + check_nvfp4_split_n_stride_inputs(D, A, B, A_sf, B_sf, alpha, bias, split_n_parts); + + auto const m = A.sizes()[0]; + auto const n = B.sizes()[0]; + auto const k = A.sizes()[1] * 2; + at::cuda::CUDAGuard device_guard{(char)A.get_device()}; + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(A.get_device()); + if (m == 75348 && n == 5120 && k == 13824 && split_n_parts == 2) { + runGemmNvfp4SplitNStrideSm120( + D, A, B, A_sf, B_sf, alpha, bias, m, n, k, split_n_parts, stream); + } else { + runGemmNvfp4SplitNStrideSm120( + D, A, B, A_sf, B_sf, alpha, bias, m, n, k, split_n_parts, stream); + } +} + +void cutlass_scaled_nvfp4_mm_split_n_stride_gelu_sm120( + torch::Tensor& D, + torch::Tensor const& A, + torch::Tensor const& B, + torch::Tensor const& A_sf, + torch::Tensor const& B_sf, + torch::Tensor const& alpha, + c10::optional const& bias, + int64_t split_n_parts) { + check_nvfp4_split_n_stride_inputs(D, A, B, A_sf, B_sf, alpha, bias, split_n_parts); + + auto const m = A.sizes()[0]; + auto const n = B.sizes()[0]; + auto const k = A.sizes()[1] * 2; + at::cuda::CUDAGuard device_guard{(char)A.get_device()}; + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(A.get_device()); + runGemmNvfp4SplitNStrideGeluSm120( + D, A, B, A_sf, B_sf, alpha, bias, m, n, k, split_n_parts, stream); +} + +void cutlass_scaled_nvfp4_mm_split_n_stride_residual_gate_sm120( + torch::Tensor& residual, + torch::Tensor const& A, + torch::Tensor const& B, + torch::Tensor const& A_sf, + torch::Tensor const& B_sf, + torch::Tensor const& alpha, + c10::optional const& bias, + torch::Tensor const& gate, + int64_t split_n_parts) { + check_nvfp4_split_n_stride_inputs( + residual, A, B, A_sf, B_sf, alpha, bias, split_n_parts); + CHECK_INPUT(gate, at::ScalarType::BFloat16, "gate"); + TORCH_CHECK(gate.device() == residual.device(), "gate and residual must be on the same CUDA device"); + TORCH_CHECK(gate.dim() == 1, "gate must be a 1D per-column tensor"); + TORCH_CHECK(gate.sizes()[0] == residual.sizes()[1], "gate size must match residual columns"); + + auto const m = A.sizes()[0]; + auto const n = B.sizes()[0]; + auto const k = A.sizes()[1] * 2; + at::cuda::CUDAGuard device_guard{(char)A.get_device()}; + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(A.get_device()); + if (m == 75348 && n == 5120 && k == 13824 && split_n_parts == 2) { + runGemmNvfp4SplitNStrideResidualGateSm120( + residual, A, B, A_sf, B_sf, alpha, bias, gate, m, n, k, split_n_parts, stream); + } else { + runGemmNvfp4SplitNStrideResidualGateSm120( + residual, A, B, A_sf, B_sf, alpha, bias, gate, m, n, k, split_n_parts, stream); + } +} diff --git a/lightx2v_kernel/include/lightx2v_kernel_ops.h b/lightx2v_kernel/include/lightx2v_kernel_ops.h index 04b380596..15453b294 100644 --- a/lightx2v_kernel/include/lightx2v_kernel_ops.h +++ b/lightx2v_kernel/include/lightx2v_kernel_ops.h @@ -73,6 +73,56 @@ void cutlass_scaled_nvfp4_mm_sm120( torch::Tensor const& alpha, c10::optional const& bias); +void cublaslt_scaled_nvfp4_mm_bias_sm120( + torch::Tensor& output, + torch::Tensor const& activation, + torch::Tensor const& weight, + torch::Tensor const& activation_scale, + torch::Tensor const& weight_scale, + torch::Tensor const& alpha, + torch::Tensor const& bias, + int64_t algorithm_index); + +int64_t cublaslt_scaled_nvfp4_mm_bias_algo_count_sm120( + torch::Tensor const& output, + torch::Tensor const& activation, + torch::Tensor const& weight, + torch::Tensor const& activation_scale, + torch::Tensor const& weight_scale, + torch::Tensor const& alpha, + torch::Tensor const& bias); + +void cutlass_scaled_nvfp4_mm_split_n_stride_sm120( + torch::Tensor& D, + torch::Tensor const& A, + torch::Tensor const& B, + torch::Tensor const& A_sf, + torch::Tensor const& B_sf, + torch::Tensor const& alpha, + c10::optional const& bias, + int64_t split_n_parts); + +void cutlass_scaled_nvfp4_mm_split_n_stride_gelu_sm120( + torch::Tensor& D, + torch::Tensor const& A, + torch::Tensor const& B, + torch::Tensor const& A_sf, + torch::Tensor const& B_sf, + torch::Tensor const& alpha, + c10::optional const& bias, + int64_t split_n_parts); + +void cutlass_scaled_nvfp4_mm_split_n_stride_residual_gate_sm120( + torch::Tensor& residual, + torch::Tensor const& A, + torch::Tensor const& B, + torch::Tensor const& A_sf, + torch::Tensor const& B_sf, + torch::Tensor const& alpha, + c10::optional const& bias, + torch::Tensor const& gate, + int64_t split_n_parts); + void cutlass_scaled_mxfp4_mm_sm120( torch::Tensor& D, torch::Tensor const& A, diff --git a/lightx2v_kernel/python/lightx2v_kernel/__init__.py b/lightx2v_kernel/python/lightx2v_kernel/__init__.py index 3f78cf886..565477ce5 100644 --- a/lightx2v_kernel/python/lightx2v_kernel/__init__.py +++ b/lightx2v_kernel/python/lightx2v_kernel/__init__.py @@ -1,6 +1,9 @@ import ctypes import os import platform + +import torch # noqa: F401 + from lightx2v_kernel import common_ops # noqa: F401 from lightx2v_kernel.version import __version__ diff --git a/lightx2v_kernel/python/lightx2v_kernel/gemm.py b/lightx2v_kernel/python/lightx2v_kernel/gemm.py index 8ae4b956e..2ceb05e4f 100644 --- a/lightx2v_kernel/python/lightx2v_kernel/gemm.py +++ b/lightx2v_kernel/python/lightx2v_kernel/gemm.py @@ -8,6 +8,101 @@ def cutlass_scaled_nvfp4_mm(mat_a, mat_b, scales_a, scales_b, alpha, bias=None): return out +def cublaslt_scaled_nvfp4_mm_bias( + mat_a, + mat_b, + scales_a, + scales_b, + alpha, + bias, + algorithm_index=-1, +): + if bias is None: + raise ValueError("cuBLASLt NVFP4 requires a bias tensor") + m, n = mat_a.shape[0], mat_b.shape[0] + out = torch.empty((m, n), dtype=torch.bfloat16, device=mat_a.device) + torch.ops.lightx2v_kernel.cublaslt_scaled_nvfp4_mm_bias_sm120.default( + out, + mat_a, + mat_b, + scales_a, + scales_b, + alpha, + bias, + algorithm_index, + ) + return out + + +def cublaslt_scaled_nvfp4_mm_bias_algo_count( + mat_a, mat_b, scales_a, scales_b, alpha, bias +): + """Return the cached cuBLASLt heuristic count for this exact M/N/K shape.""" + m, n = mat_a.shape[0], mat_b.shape[0] + out = torch.empty((m, n), dtype=torch.bfloat16, device=mat_a.device) + return torch.ops.lightx2v_kernel.cublaslt_scaled_nvfp4_mm_bias_algo_count_sm120.default( + out, mat_a, mat_b, scales_a, scales_b, alpha, bias + ) + + +def cutlass_scaled_nvfp4_mm_split_n_stride(mat_a, mat_b, scales_a, scales_b, alpha, bias=None, split_n_parts=2): + m, n = mat_a.shape[0], mat_b.shape[0] + out = torch.empty((m, n), dtype=torch.bfloat16, device=mat_a.device) + torch.ops.lightx2v_kernel.cutlass_scaled_nvfp4_mm_split_n_stride_sm120.default( + out, + mat_a, + mat_b, + scales_a, + scales_b, + alpha, + bias, + split_n_parts, + ) + return out + + +def cutlass_scaled_nvfp4_mm_split_n_stride_gelu( + mat_a, + mat_b, + scales_a, + scales_b, + alpha, + bias=None, + split_n_parts=2, +): + m, n = mat_a.shape[0], mat_b.shape[0] + out = torch.empty((m, n), dtype=torch.bfloat16, device=mat_a.device) + torch.ops.lightx2v_kernel.cutlass_scaled_nvfp4_mm_split_n_stride_gelu_sm120.default( + out, mat_a, mat_b, scales_a, scales_b, alpha, bias, split_n_parts + ) + return out + + +def cutlass_scaled_nvfp4_mm_split_n_stride_residual_gate( + mat_a, + mat_b, + scales_a, + scales_b, + alpha, + residual, + gate, + bias=None, + split_n_parts=2, +): + torch.ops.lightx2v_kernel.cutlass_scaled_nvfp4_mm_split_n_stride_residual_gate_sm120.default( + residual, + mat_a, + mat_b, + scales_a, + scales_b, + alpha, + bias, + gate.contiguous(), + split_n_parts, + ) + return residual + + def scaled_nvfp4_quant(input: torch.Tensor, input_global_scale: torch.Tensor): """ Quantize input tensor to FP4 and return quantized tensor and scale. @@ -48,7 +143,11 @@ def scaled_nvfp4_quant(input: torch.Tensor, input_global_scale: torch.Tensor): # rounded_m = ((m + 128 - 1) // 128) * 128 # scale_n = n // block_size # rounded_n = ((scale_n + 4 - 1) // 4) * 4 - output_scale = torch.zeros((((m + 128 - 1) // 128) * 128, (n // block_size + 4 - 1) // 4), device=device, dtype=torch.int32) + rounded_m = ((m + 128 - 1) // 128) * 128 + scale_shape = (rounded_m, (n // block_size + 4 - 1) // 4) + # The kernel writes every logical SF element for complete 128-row tiles. + # Keep zeros only for a partial tile, where swizzled padding must remain 0. + output_scale = torch.empty(scale_shape, device=device, dtype=torch.int32) if rounded_m == m else torch.zeros(scale_shape, device=device, dtype=torch.int32) torch.ops.lightx2v_kernel.scaled_nvfp4_quant_sm120.default(output, input, output_scale, input_global_scale) output_scale = output_scale.view(torch.float8_e4m3fn) diff --git a/lightx2v_kernel/test/nvfp4_nvfp4/test_qkv_cublaslt.py b/lightx2v_kernel/test/nvfp4_nvfp4/test_qkv_cublaslt.py new file mode 100644 index 000000000..f995f6ae6 --- /dev/null +++ b/lightx2v_kernel/test/nvfp4_nvfp4/test_qkv_cublaslt.py @@ -0,0 +1,54 @@ +import torch +from lightx2v_kernel.gemm import ( + cublaslt_scaled_nvfp4_mm_bias, + cutlass_scaled_nvfp4_mm, + scaled_nvfp4_quant, +) + +FLOAT4_E2M1_MAX = 6.0 +FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max + + +@torch.inference_mode() +def test_cublaslt_qkv_projection_matches_cutlass_with_auto_algorithm(): + torch.manual_seed(0) + m, n, k = 257, 512, 512 + activation = torch.randn((m, k), dtype=torch.bfloat16, device="cuda") + weight = torch.randn((n, k), dtype=torch.bfloat16, device="cuda") + bias = torch.randn((n,), dtype=torch.bfloat16, device="cuda") + + activation_global_scale = ( + (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / activation.abs().max() + ).float() + weight_global_scale = ( + (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / weight.abs().max() + ).float() + activation_fp4, activation_scale = scaled_nvfp4_quant( + activation, + activation_global_scale, + ) + weight_fp4, weight_scale = scaled_nvfp4_quant( + weight, + weight_global_scale, + ) + alpha = 1.0 / (activation_global_scale * weight_global_scale) + + expected = cutlass_scaled_nvfp4_mm( + activation_fp4, + weight_fp4, + activation_scale, + weight_scale, + alpha, + bias, + ) + actual = cublaslt_scaled_nvfp4_mm_bias( + activation_fp4, + weight_fp4, + activation_scale, + weight_scale, + alpha, + bias, + algorithm_index=-1, + ) + + torch.testing.assert_close(actual, expected, atol=1e-1, rtol=1e-1) diff --git a/lightx2v_kernel/test/nvfp4_nvfp4/test_split_n_stride.py b/lightx2v_kernel/test/nvfp4_nvfp4/test_split_n_stride.py new file mode 100644 index 000000000..fb02232cf --- /dev/null +++ b/lightx2v_kernel/test/nvfp4_nvfp4/test_split_n_stride.py @@ -0,0 +1,133 @@ +import pytest +import torch + +from lightx2v_kernel.gemm import ( + cutlass_scaled_nvfp4_mm, + cutlass_scaled_nvfp4_mm_split_n_stride, + cutlass_scaled_nvfp4_mm_split_n_stride_gelu, + cutlass_scaled_nvfp4_mm_split_n_stride_residual_gate, + scaled_nvfp4_quant, +) + + +FLOAT4_E2M1_MAX = 6.0 +FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max + + +@pytest.mark.parametrize("split_n_parts,bias_enabled", [(2, False), (4, True)]) +@torch.inference_mode() +def test_split_n_stride_matches_full_gemm(split_n_parts, bias_enabled): + m, n, k = 129, 512, 256 + activation = torch.randn((m, k), dtype=torch.bfloat16, device="cuda") + weight = torch.randn((n, k), dtype=torch.bfloat16, device="cuda") + bias = torch.randn((n,), dtype=torch.bfloat16, device="cuda") if bias_enabled else None + + activation_global_scale = ((FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / activation.abs().max()).float() + weight_global_scale = ((FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / weight.abs().max()).float() + alpha = 1.0 / (activation_global_scale * weight_global_scale) + activation_fp4, activation_scale = scaled_nvfp4_quant(activation, activation_global_scale) + weight_fp4, weight_scale = scaled_nvfp4_quant(weight, weight_global_scale) + + expected = cutlass_scaled_nvfp4_mm( + activation_fp4, + weight_fp4, + activation_scale, + weight_scale, + alpha, + bias, + ) + actual = cutlass_scaled_nvfp4_mm_split_n_stride( + activation_fp4, + weight_fp4, + activation_scale, + weight_scale, + alpha, + bias, + split_n_parts=split_n_parts, + ) + + torch.testing.assert_close(actual, expected, atol=1e-1, rtol=1e-1) + + +@pytest.mark.parametrize("split_n_parts,bias_enabled", [(2, False), (4, True)]) +@torch.inference_mode() +def test_split_n_stride_gelu_matches_unfused_exactly(split_n_parts, bias_enabled): + torch.manual_seed(11) + m, n, k = 129, 512, 256 + activation = torch.randn((m, k), dtype=torch.bfloat16, device="cuda") + weight = torch.randn((n, k), dtype=torch.bfloat16, device="cuda") + bias = torch.randn((n,), dtype=torch.bfloat16, device="cuda") if bias_enabled else None + + activation_global_scale = ((FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / activation.abs().max()).float() + weight_global_scale = ((FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / weight.abs().max()).float() + alpha = 1.0 / (activation_global_scale * weight_global_scale) + activation_fp4, activation_scale = scaled_nvfp4_quant(activation, activation_global_scale) + weight_fp4, weight_scale = scaled_nvfp4_quant(weight, weight_global_scale) + + gemm_output = cutlass_scaled_nvfp4_mm_split_n_stride( + activation_fp4, + weight_fp4, + activation_scale, + weight_scale, + alpha, + bias, + split_n_parts=split_n_parts, + ) + expected = torch.nn.functional.gelu(gemm_output, approximate="tanh") + actual = cutlass_scaled_nvfp4_mm_split_n_stride_gelu( + activation_fp4, + weight_fp4, + activation_scale, + weight_scale, + alpha, + bias, + split_n_parts=split_n_parts, + ) + + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +@pytest.mark.parametrize("split_n_parts,bias_enabled", [(2, False), (4, True)]) +@torch.inference_mode() +def test_split_n_stride_residual_gate_is_exact(split_n_parts, bias_enabled): + torch.manual_seed(7) + m, n, k = 129, 512, 256 + activation = torch.randn((m, k), dtype=torch.bfloat16, device="cuda") + weight = torch.randn((n, k), dtype=torch.bfloat16, device="cuda") + bias = torch.randn((n,), dtype=torch.bfloat16, device="cuda") if bias_enabled else None + gate = torch.randn((n,), dtype=torch.bfloat16, device="cuda") + residual = torch.randn((m, n), dtype=torch.bfloat16, device="cuda") + + activation_global_scale = ((FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / activation.abs().max()).float() + weight_global_scale = ((FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / weight.abs().max()).float() + alpha = 1.0 / (activation_global_scale * weight_global_scale) + activation_fp4, activation_scale = scaled_nvfp4_quant(activation, activation_global_scale) + weight_fp4, weight_scale = scaled_nvfp4_quant(weight, weight_global_scale) + + ffn_out = cutlass_scaled_nvfp4_mm_split_n_stride( + activation_fp4, + weight_fp4, + activation_scale, + weight_scale, + alpha, + bias, + split_n_parts=split_n_parts, + ) + expected = residual.clone() + expected.add_(ffn_out * gate) + + actual = residual.clone() + returned = cutlass_scaled_nvfp4_mm_split_n_stride_residual_gate( + activation_fp4, + weight_fp4, + activation_scale, + weight_scale, + alpha, + actual, + gate, + bias, + split_n_parts=split_n_parts, + ) + + assert returned.data_ptr() == actual.data_ptr() + torch.testing.assert_close(actual, expected, atol=0, rtol=0) diff --git a/test_cases/test_attention_merge.py b/test_cases/test_attention_merge.py new file mode 100644 index 000000000..ebf977d79 --- /dev/null +++ b/test_cases/test_attention_merge.py @@ -0,0 +1,290 @@ +"""CPU-only merge regression tests; do not import the LightX2V package.""" + +import ast +import builtins +import inspect +import types +import unittest +from pathlib import Path +from unittest.mock import Mock, patch + +import torch + +ATTN = Path(__file__).resolve().parents[1] / "lightx2v/common/ops/attn" + + +def load_attention(filename, dense=None, varlen=None, metadata=None, import_error=ImportError): + registrations = {} + ops = {} + + def register(name): + def decorate(cls): + if name in registrations: + raise AssertionError(f"Duplicate attention registration: {name}") + registrations[name] = cls + return cls + + return decorate + + def custom_op(name, **options): + def decorate(fn): + if name in ops: + raise AssertionError(f"Duplicate custom op: {name}") + ops[name] = fn + fn.options = options + + def register_fake(fake): + fn.fake = fake + return fake + + fn.register_fake = register_fake + return fn + + return decorate + + dependencies = { + "lightx2v.utils.registry_factory": types.SimpleNamespace(ATTN_WEIGHT_REGISTER=register), + "template": types.SimpleNamespace(AttnWeightTemplate=object), + "utils.sla_util": types.SimpleNamespace( + get_block_map=Mock(), + get_cuda_arch=lambda _: "sm110", + block_lut_to_ordinal_metadata=Mock(), + ), + "utils.sla_util_blhd": types.SimpleNamespace(get_block_lut_blhd=Mock(), get_block_map_blhd=Mock()), + "utils.sparge_util": types.SimpleNamespace( + block_map_incremental_lut_triton=Mock(), + block_map_ordinal_lut_triton=Mock(), + sage2_block_sparse_attn=Mock(), + get_block_map_meansim=Mock(), + ), + "kernels.sla_kernel": types.SimpleNamespace(_attention=Mock()), + "kernels.sla_kernel_ar": types.SimpleNamespace(_attention_ar=Mock()), + } + real_import = builtins.__import__ + + def isolated_import(name, globals=None, locals=None, fromlist=(), level=0): + if name in dependencies: + return dependencies[name] + if name == "flash_attn.cute": + value = dense if fromlist[0] == "flash_attn_func" else varlen + if value is None: + raise import_error(fromlist[0]) + return types.SimpleNamespace(**{fromlist[0]: value}) + if name == "flash_attn.cute.block_sparsity": + if metadata is None: + raise import_error("BlockSparseTensorsTorch") + return types.SimpleNamespace(BlockSparseTensorsTorch=metadata) + if name.startswith(("flash_attn", "sageattn3_sparse", "magi_attention")): + raise ImportError(name) + return real_import(name, globals, locals, fromlist, level) + + namespace = {"__name__": "attention_merge_test", "__builtins__": dict(vars(builtins), __import__=isolated_import)} + with patch.object(torch.library, "custom_op", custom_op), patch.object(torch.compiler, "disable", lambda f: f): + exec(compile((ATTN / filename).read_text(), str(ATTN / filename), "exec"), namespace) # noqa: S102 - trusted local source + return types.SimpleNamespace(**namespace), registrations, ops + + +def load_helpers(): + tree = ast.parse((ATTN / "utils/sla_util.py").read_text()) + names = {"_get_block_lut", "block_lut_to_ordinal_metadata", "get_block_lut", "get_block_map"} + tree.body = [node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name in names] + namespace = {"torch": torch} + exec(compile(tree, str(ATTN / "utils/sla_util.py"), "exec"), namespace) # noqa: S102 - trusted local helpers + return namespace + + +class FlashAttentionMergeTests(unittest.TestCase): + def setUp(self): + self.q = torch.randn(5, 2, 4) + self.k = torch.randn(7, 2, 4) + self.dense = Mock(side_effect=lambda q, k, v, **kw: (q.clone(), None)) + self.varlen = Mock(side_effect=lambda q, k, v, cu, **kw: (q.clone(), None)) + + def test_dense_single_sequence_without_varlen_or_metadata(self): + for shape4 in (False, True): + for error in (ImportError, AttributeError): + with self.subTest(shape4=shape4, error=error): + mod, registered, _ = load_attention("flash_attn.py", dense=self.dense, import_error=error) + q = self.q.unsqueeze(0) if shape4 else self.q + k = self.k.unsqueeze(0) if shape4 else self.k + out = mod.FlashAttn4Weight().apply(q, k, k, causal=True, softmax_scale=0.3) + self.assertEqual(out.shape, (5, 8)) + self.assertEqual(self.dense.call_args.args[0].shape, (1, 5, 2, 4)) + self.assertEqual(self.dense.call_args.kwargs, {"causal": True, "softmax_scale": 0.3}) + self.assertEqual(len(registered), 4) + self.assertIsNone(mod.BlockSparseTensorsTorch) + + def test_single_sequence_with_cumulative_lengths_preserves_varlen(self): + for q_end, k_end in ((5, 7), (3, 4)): + with self.subTest(q_end=q_end, k_end=k_end): + mod, _, _ = load_attention("flash_attn.py", self.dense, self.varlen) + cuq, cuk = torch.tensor([0, q_end]), torch.tensor([0, k_end]) + out = mod.FlashAttn4Weight().apply(self.q, self.k, self.k, cuq, cuk, q_end, k_end) + self.assertEqual(out.shape, (5, 8)) + self.assertIs(self.varlen.call_args.args[3], cuq) + self.assertIs(self.varlen.call_args.kwargs["cu_seqlens_k"], cuk) + self.dense.assert_not_called() + + def test_single_sequence_with_lengths_requires_varlen(self): + mod, _, _ = load_attention("flash_attn.py", self.dense) + with self.assertRaisesRegex(RuntimeError, "varlen"): + mod.FlashAttn4Weight().apply(self.q, self.k, self.k, torch.tensor([0, 3]), torch.tensor([0, 4]), 3, 4) + self.dense.assert_not_called() + + def test_packed_unequal_sequences_use_varlen_and_total_tokens(self): + mod, _, _ = load_attention("flash_attn.py", self.dense, self.varlen) + cuq, cuk = torch.tensor([0, 2, 5]), torch.tensor([0, 3, 7]) + out = mod.FlashAttn4Weight().apply(self.q, self.k, self.k, cuq, cuk, 3, 4, causal=True, softmax_scale=0.2) + self.assertEqual(out.shape, (5, 8)) + self.dense.assert_not_called() + args, kw = self.varlen.call_args + self.assertIs(args[3], cuq) + self.assertIs(kw["cu_seqlens_k"], cuk) + self.assertEqual((kw["max_seqlen_q"], kw["max_seqlen_k"]), (3, 4)) + self.assertEqual((kw["causal"], kw["softmax_scale"]), (True, 0.2)) + + def test_packed_input_never_falls_back_to_dense(self): + mod, _, _ = load_attention("flash_attn.py", self.dense) + with self.assertRaisesRegex(RuntimeError, "varlen"): + mod.FlashAttn4Weight().apply(self.q, self.k, self.k, torch.tensor([0, 2, 5]), torch.tensor([0, 3, 7]), 3, 4) + self.dense.assert_not_called() + + def test_batched_4d_varlen_and_missing_lengths(self): + mod, _, _ = load_attention("flash_attn.py", self.dense, self.varlen) + q = torch.randn(2, 3, 2, 4) + cu = torch.tensor([0, 3, 6]) + out = mod.FlashAttn4Weight().apply(q, q, q, cu, cu, 3, 3) + self.assertEqual(out.shape, (6, 8)) + self.assertEqual(self.varlen.call_args.args[0].shape, (6, 2, 4)) + with self.assertRaises(ValueError): + mod.FlashAttn4Weight().apply(q, q, q) + with self.assertRaises(ValueError): + mod.FlashAttn4Weight().apply(self.q, self.k, self.k, cu_seqlens_q=cu) + + def test_varlen_remains_available_without_dense(self): + mod, _, _ = load_attention("flash_attn.py", varlen=self.varlen, import_error=AttributeError) + out = mod.FlashAttn4Weight().apply(self.q, self.k, self.k, torch.tensor([0, 5]), torch.tensor([0, 7]), 5, 7) + self.assertEqual(out.shape, (5, 8)) + self.varlen.assert_called_once() + + def test_fa3_lse_stays_on_fa3_class(self): + mod, _, _ = load_attention("flash_attn.py") + lse = torch.arange(10).reshape(1, 2, 5) + fn = Mock(return_value=(self.q.unsqueeze(0), lse, None)) + mod.FlashAttn3Weight.apply_with_lse.__globals__["flash_attn_func_v3"] = fn + out, actual_lse = mod.FlashAttn3Weight().apply_with_lse(self.q, self.q, self.q, 0.5) + self.assertEqual(out.shape, (5, 8)) + torch.testing.assert_close(actual_lse, lse.transpose(1, 2).reshape(5, 2)) + self.assertTrue(fn.call_args.kwargs["return_attn_probs"]) + self.assertFalse(hasattr(mod.FlashAttn4Weight, "apply_with_lse")) + + +class DynamicSparseMergeTests(unittest.TestCase): + def test_sparse_apis_eager_and_compile_boundary(self): + for api in ("expanded", "block_sparse_tensors"): + for error in (ImportError, AttributeError): + with self.subTest(api=api, error=error): + calls = [] + + def expanded(q, k, v, mask_block_cnt, mask_block_idx, full_block_cnt, full_block_idx, block_size, calls=calls): + calls.append((full_block_cnt, full_block_idx, block_size)) + return q.clone(), None + + def bundled(q, k, v, block_sparse_tensors, calls=calls): + calls.append((block_sparse_tensors.full_block_cnt, block_sparse_tensors.full_block_idx, block_sparse_tensors.block_size)) + return q.clone(), None + + fn = expanded if api == "expanded" else bundled + metadata = None if api == "expanded" else types.SimpleNamespace + mod, _, ops = load_attention("dynamic_sparse_attn.py", dense=fn, metadata=metadata, import_error=error) + self.assertEqual(mod._FA4_SPARSE_API, api) + self.assertEqual(inspect.signature(mod.flash_attn_func_v4), inspect.signature(fn)) + helpers = load_helpers() + ns = mod.DynamicSparseAttnWeight.apply_fa4.__globals__ + ns["get_block_lut_blhd"] = Mock(return_value=(torch.zeros(1, 2, 1, 1, dtype=torch.long), 1, 1)) + ns["block_lut_to_ordinal_metadata"] = helpers["block_lut_to_ordinal_metadata"] + weight = object.__new__(mod.DynamicSparseAttnWeight) + weight.topk, weight.BLKQ, weight.BLKK = 0.2, 128, 128 + q = torch.randn(5, 2, 4) + for compiling in (False, True): + with patch.object(torch.compiler, "is_compiling", return_value=compiling): + out = weight.apply_fa4(q, q, q, max_seqlen_q=5) + torch.testing.assert_close(out, q.reshape(5, 8)) + self.assertEqual(len(calls), 2) + self.assertEqual(calls[0][2], (128, 128)) + torch.testing.assert_close(calls[0][0], calls[1][0]) + op = ops["lightx2v_internal::fa4_blocksparse"] + fake = op.fake(q, q, q, None, None, None, None, 128, 128) + self.assertEqual(fake.shape, q.shape) + self.assertEqual(fake.dtype, q.dtype) + + def test_sage2_custom_op_and_apply(self): + mod, registered, ops = load_attention("dynamic_sparse_attn.py") + self.assertEqual(list(registered), ["dynamic_sparse_attn"]) + op = ops["lightx2v::dynamic_sparse_sage2"] + self.assertEqual(op.options, {"mutates_args": (), "device_types": "cuda"}) + q = torch.randn(5, 2, 4) + ns = op.__globals__ + sparse_map, lut, counts = object(), object(), object() + ns["get_block_map"] = Mock(return_value=(sparse_map, None, None)) + ns["block_map_incremental_lut_triton"] = Mock(return_value=(lut, counts)) + ns["sage2_block_sparse_attn"] = Mock(side_effect=lambda q, *args: q.clone()) + weight = object.__new__(mod.DynamicSparseAttnWeight) + weight.topk, weight.BLKQ, weight.BLKK, weight.arch = 0.2, 128, 64, "sm110" + out = weight.apply_sage2(q, q, q, max_seqlen_q=5) + torch.testing.assert_close(out, q.reshape(5, 8)) + ns["block_map_incremental_lut_triton"].assert_called_once_with(sparse_map) + args = ns["sage2_block_sparse_attn"].call_args.args + self.assertEqual(args[0].shape, (1, 2, 5, 4)) + self.assertTrue(args[0].is_contiguous()) + self.assertEqual(args[3:], (lut, counts, 128, 64, "sm110")) + fake = op.fake(q, q, q, 0.2, 128, 64, "sm110") + self.assertEqual((fake.shape, fake.dtype, fake.device), (q.shape, q.dtype, q.device)) + + def test_cutedsl_fp8_retains_blhd_selection(self): + mod, _, _ = load_attention("dynamic_sparse_attn.py") + weight = object.__new__(mod.DynamicSparseAttnWeight) + weight.topk, weight.BLKQ, weight.BLKK = 0.2, 256, 128 + weight.cutedsl_sparse_fmha = Mock(side_effect=lambda q, *args, **kw: q.to(torch.bfloat16)) + ns = mod.DynamicSparseAttnWeight.apply_cutedsl_fp8.__globals__ + pool = Mock(return_value=(torch.zeros(1, 2, 1, 1, dtype=torch.long), 1, 1)) + ns["get_block_lut_blhd"] = pool + ns["block_lut_to_ordinal_metadata"] = load_helpers()["block_lut_to_ordinal_metadata"] + q = torch.ones(5, 2, 4, dtype=torch.bfloat16) + out = weight.apply_cutedsl_fp8(q, q, q) + self.assertEqual(pool.call_args.args[0].shape, (1, 5, 2, 4)) + self.assertEqual(pool.call_args.args[0].dtype, torch.bfloat16) + args, kw = weight.cutedsl_sparse_fmha.call_args + self.assertTrue(all(t.dtype == torch.float8_e4m3fn for t in args[:3])) + self.assertEqual(kw["output_dtype"], torch.bfloat16) + self.assertEqual(out.shape, (5, 8)) + + +class SparseHelperMergeTests(unittest.TestCase): + def test_short_sequences_keep_one_block_and_gqa(self): + helpers = load_helpers() + q, k = torch.ones(1, 4, 2, 4), torch.ones(1, 2, 1, 4) + for ratio in (0.0, 0.01, 0.2, 1.0): + lut, topk, count = helpers["_get_block_lut"](q, k, ratio) + self.assertEqual((topk, count), (1, 1)) + self.assertEqual(lut.shape, (1, 4, 2, 1)) + indices, counts = helpers["block_lut_to_ordinal_metadata"](lut, count) + self.assertEqual(indices.dtype, torch.int32) + self.assertTrue(torch.all(counts == 1)) + helpers["mean_pool"] = lambda x, block: x.mean(-2, keepdim=True) + sparse_map, _, topk = helpers["get_block_map"](q, k, 0.01) + self.assertEqual(topk, 1) + self.assertTrue(torch.all(sparse_map == 1)) + + def test_syntax_and_unique_definitions(self): + for path in (ATTN / "dynamic_sparse_attn.py", ATTN / "flash_attn.py", ATTN / "utils/sla_util.py"): + tree = ast.parse(path.read_text()) + definitions = [node.name for node in tree.body if isinstance(node, (ast.ClassDef, ast.FunctionDef))] + self.assertEqual(len(definitions), len(set(definitions)), str(path)) + for cls in (node for node in tree.body if isinstance(node, ast.ClassDef)): + methods = [node.name for node in cls.body if isinstance(node, ast.FunctionDef)] + self.assertEqual(len(methods), len(set(methods)), cls.name) + + +if __name__ == "__main__": + unittest.main() diff --git a/test_cases/test_dynamic_sparse_attn_fa4.py b/test_cases/test_dynamic_sparse_attn_fa4.py new file mode 100644 index 000000000..273ea3dc5 --- /dev/null +++ b/test_cases/test_dynamic_sparse_attn_fa4.py @@ -0,0 +1,126 @@ +import pytest +import torch + +from lightx2v.common.ops.attn import dynamic_sparse_attn +from lightx2v.common.ops.attn.utils.sla_util import ( + block_lut_to_ordinal_metadata, + centered_mean_pool, + get_block_lut, + get_block_map, + mean_pool, +) +from lightx2v.common.ops.attn.utils.sla_util_blhd import ( + centered_mean_pool as centered_mean_pool_blhd, +) +from lightx2v.common.ops.attn.utils.sla_util_blhd import ( + get_block_lut_blhd, +) +from lightx2v.common.ops.attn.utils.sparge_util import block_map_ordinal_lut_triton + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_cutedsl_builds_lut_and_preserves_input_dtype(monkeypatch, dtype): + q = torch.randn((256, 2, 128), dtype=dtype) + k = torch.randn_like(q) + v = torch.randn_like(q) + expected_lut = torch.zeros((1, 2, 1, 1), dtype=torch.int64) + + def fake_get_block_lut_blhd(qt, kt, **kwargs): + assert qt.shape == (1, 256, 2, 128) + assert kt.shape == (1, 256, 2, 128) + assert qt.dtype == dtype + assert kt.dtype == dtype + assert kwargs["BLKQ"] == 256 + assert kwargs["BLKK"] == 128 + return expected_lut, 1, 2 + + def fake_cutedsl_sparse_fmha(qt, kt, vt, cu_seqlens, max_seqlen, block_count, block_indices): + assert qt.dtype == dtype + assert kt.dtype == dtype + assert vt.dtype == dtype + assert qt.is_contiguous() and kt.is_contiguous() and vt.is_contiguous() + assert cu_seqlens.tolist() == [0, 256] + assert max_seqlen == 256 + torch.testing.assert_close(block_count, torch.ones((1, 2, 1), dtype=torch.int32)) + assert block_indices.shape == (1, 2, 1, 2) + return torch.zeros_like(qt) + + monkeypatch.setattr(dynamic_sparse_attn, "get_block_lut_blhd", fake_get_block_lut_blhd) + attn = object.__new__(dynamic_sparse_attn.DynamicSparseAttnWeight) + attn.BLKQ, attn.BLKK = 256, 128 + attn.topk = 0.5 + attn.cutedsl_sparse_fmha = fake_cutedsl_sparse_fmha + + out = attn.apply_cutedsl(q, k, v) + + assert out.shape == (256, 2 * 128) + assert out.dtype == dtype + + +def test_block_lut_to_ordinal_metadata_sorts_and_pads(): + lut = torch.tensor([[[[5, 1, 3], [4, 0, 2]]]]) + + full_block_idx, full_block_cnt = block_lut_to_ordinal_metadata(lut, num_k_blocks=6) + + expected_idx = torch.tensor([[[[1, 3, 5, 0, 0, 0], [0, 2, 4, 0, 0, 0]]]], dtype=torch.int32) + expected_cnt = torch.tensor([[[3, 3]]], dtype=torch.int32) + torch.testing.assert_close(full_block_idx, expected_idx, atol=0, rtol=0) + torch.testing.assert_close(full_block_cnt, expected_cnt, atol=0, rtol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton mean pooling requires CUDA") +def test_fa4_lut_matches_existing_map_path_for_gqa_and_partial_blocks(): + torch.manual_seed(42) + q = torch.randn((1, 4, 257, 64), device="cuda", dtype=torch.bfloat16) + k = torch.randn((1, 2, 259, 64), device="cuda", dtype=torch.bfloat16) + + k_mean = torch.mean(k, dim=-2, keepdim=True) + old_pooled_k = mean_pool(k - k_mean, 128) + new_pooled_k = centered_mean_pool(k, k_mean, 128) + torch.testing.assert_close(new_pooled_k, old_pooled_k, atol=0, rtol=0) + + sparse_map, old_lut, old_topk = get_block_map(q, k, topk_ratio=0.67, BLKQ=128, BLKK=128) + new_lut, new_topk, num_k_blocks = get_block_lut(q, k, topk_ratio=0.67, BLKQ=128, BLKK=128) + + torch.testing.assert_close( + torch.sort(new_lut, dim=-1).values, + torch.sort(old_lut, dim=-1).values, + atol=0, + rtol=0, + ) + assert new_topk == old_topk + assert num_k_blocks == 3 + + old_idx, old_cnt = block_map_ordinal_lut_triton(sparse_map) + new_idx, new_cnt = block_lut_to_ordinal_metadata(new_lut, num_k_blocks) + torch.testing.assert_close(new_idx, old_idx, atol=0, rtol=0) + torch.testing.assert_close(new_cnt, old_cnt, atol=0, rtol=0) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("block_q,block_k", [(128, 128), (256, 128)]) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton mean pooling requires CUDA") +def test_blhd_lut_matches_bhld_without_full_tensor_clones(dtype, block_q, block_k): + torch.manual_seed(123) + q = torch.randn((1, 257, 4, 64), device="cuda", dtype=dtype) + k = torch.randn((1, 259, 2, 64), device="cuda", dtype=dtype) + + q_bhld = q.transpose(1, 2).contiguous() + k_bhld = k.transpose(1, 2).contiguous() + expected_lut, expected_topk, expected_k_blocks = get_block_lut(q_bhld, k_bhld, topk_ratio=0.67, BLKQ=block_q, BLKK=block_k) + actual_lut, actual_topk, actual_k_blocks = get_block_lut_blhd(q, k, topk_ratio=0.67, BLKQ=block_q, BLKK=block_k) + + torch.testing.assert_close( + torch.sort(actual_lut, dim=-1).values, + torch.sort(expected_lut, dim=-1).values, + atol=0, + rtol=0, + ) + assert actual_topk == expected_topk + assert actual_k_blocks == expected_k_blocks + + k_center_blhd = torch.mean(k, dim=1, keepdim=True) + pooled_k_blhd = centered_mean_pool_blhd(k, k_center_blhd, block_k) + k_center_bhld = torch.mean(k_bhld, dim=-2, keepdim=True) + pooled_k_bhld = centered_mean_pool(k_bhld, k_center_bhld, block_k) + torch.testing.assert_close(pooled_k_blhd, pooled_k_bhld, atol=0, rtol=0) diff --git a/test_cases/test_infer_merge.py b/test_cases/test_infer_merge.py new file mode 100644 index 000000000..e869e0b58 --- /dev/null +++ b/test_cases/test_infer_merge.py @@ -0,0 +1,69 @@ +import argparse +import ast +import os +import sys +import types +import unittest +from contextlib import nullcontext +from pathlib import Path +from unittest.mock import MagicMock, patch + + +class InferMergeTest(unittest.TestCase): + def test_profiler_is_opt_in_and_preserves_json_warmup(self): + path = Path(__file__).resolve().parents[1] / "lightx2v" / "infer.py" + tree = ast.parse(path.read_text()) + main = next(node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == "main") + code = compile(ast.Module(body=[main], type_ignores=[]), str(path), "exec") + for value in (None, "0", "1"): + with self.subTest(profiler=value): + runner = MagicMock() + config = {"parallel": False, "warmup": True} + profiler_module = types.ModuleType("torch.profiler") + profiler_module.ProfilerActivity = types.SimpleNamespace(CPU="cpu", CUDA="cuda") + profiler_module.profile = MagicMock() + profiler_module.record_function = MagicMock(side_effect=lambda _: nullcontext()) + profiler = profiler_module.profile.return_value.__enter__.return_value + namespace = { + "argparse": argparse, + "os": os, + "OMNI_VISION_SUBTASK_CHOICES": (), + "WAN_ANIMATE2_MODEL_ID": "wan2.2_animate2", + "seed_all": MagicMock(), + "set_config": MagicMock(return_value=config), + "init_empty_input_info": MagicMock(return_value=object()), + "print_config": MagicMock(), + "validate_config_paths": MagicMock(), + "ProfilingContext4DebugL1": lambda _: nullcontext(), + "init_runner": MagicMock(return_value=runner), + "update_input_info_from_dict": MagicMock(), + "dist": types.SimpleNamespace(is_initialized=lambda: False), + "logger": MagicMock(), + } + env = {"LIGHTX2V_TORCH_PROFILER_FULL_TRACE": "/tmp/lightx2v-test/trace.json"} + if value is not None: + env["LIGHTX2V_TORCH_PROFILER_FULL"] = value + argv = ["infer", "--model_cls", "wan2.2_moe", "--model_path", "/models/wan", "--config_json", "thor.json"] + with ( + patch.dict(os.environ, env, clear=True), + patch.dict(sys.modules, {"torch.profiler": profiler_module}), + patch.object(sys, "argv", argv), + patch("os.makedirs") as mkdir, + patch("builtins.print"), + ): + exec(code, namespace) # noqa: S102 - Execute the local CLI function without loading GPU runners. + namespace["main"]() + runner.run_pipeline.assert_called_once() + self.assertIs(config["warmup"], True) + if value == "1": + profiler_module.profile.assert_called_once() + profiler.export_chrome_trace.assert_called_once_with(env["LIGHTX2V_TORCH_PROFILER_FULL_TRACE"]) + mkdir.assert_called_once_with("/tmp/lightx2v-test", exist_ok=True) + else: + profiler_module.profile.assert_not_called() + profiler.export_chrome_trace.assert_not_called() + mkdir.assert_not_called() + + +if __name__ == "__main__": + unittest.main() diff --git a/test_cases/test_thor_config.py b/test_cases/test_thor_config.py new file mode 100644 index 000000000..438faec1f --- /dev/null +++ b/test_cases/test_thor_config.py @@ -0,0 +1,96 @@ +import json +import os +import unittest +from pathlib import Path + +os.environ.setdefault("SKIP_PLATFORM_CHECK", "1") + +from lightx2v.models.schedulers.wan.scheduler_factory import get_wan_distill_method +from lightx2v.utils.set_config import validate_thor_config + + +class ThorConfigValidationTest(unittest.TestCase): + @staticmethod + def make_config(**overrides): + config = { + "model_cls": "wan2.2_moe", + "distill_method": "dmd2", + "task": "i2v", + "dit_quant_scheme": "nvfp4", + } + config.update(overrides) + return config + + def test_supported_model_task_combinations(self): + for distill_method in (None, "dmd2"): + for task in ("i2v", "t2v"): + with self.subTest(distill_method=distill_method, task=task): + config = self.make_config(thor=True, task=task) + if distill_method is None: + config.pop("distill_method") + validate_thor_config(config) + self.assertEqual(get_wan_distill_method(config), distill_method) + + def test_retired_distill_model_name_is_rejected(self): + with self.assertRaisesRegex(ValueError, "model_cls"): + validate_thor_config(self.make_config(thor=True, model_cls="wan2.2_moe_distill")) + + def test_thor_presets_select_dmd2(self): + config_dir = Path(__file__).resolve().parents[1] / "configs" / "wan22" / "thor" + paths = sorted(config_dir.glob("*.json")) + self.assertTrue(paths) + for path in paths: + with self.subTest(config=path.name): + config = json.loads(path.read_text()) + config.update(model_cls="wan2.2_moe", task="i2v") + self.assertIs(config["thor"], True) + validate_thor_config(config) + self.assertEqual(get_wan_distill_method(config), "dmd2") + + def test_disabled_or_missing_thor_is_ignored(self): + validate_thor_config({}) + validate_thor_config({"thor": False, "model_cls": "wan2.1", "dit_quant_scheme": "fp8"}) + + def test_thor_must_be_boolean(self): + for value in (None, 0, 1, "true", [], {}): + with self.subTest(value=value): + with self.assertRaisesRegex(TypeError, "thor must be a boolean"): + validate_thor_config(self.make_config(thor=value)) + + def test_unsupported_configurations_are_rejected(self): + for overrides, message in ( + ({"model_cls": "wan2.1"}, "model_cls"), + ({"task": "flf2v"}, "task"), + ({"dit_quant_scheme": "fp8"}, "dit_quant_scheme"), + ): + with self.subTest(overrides=overrides): + with self.assertRaisesRegex(ValueError, message): + validate_thor_config(self.make_config(thor=True, **overrides)) + + def test_legacy_fields_are_accepted_without_mutating_config(self): + legacy_options = { + "nvfp4_ffn0_gelu_fusion": True, + "nvfp4_ffn2_residual_gate_fusion": False, + "nvfp4_ffn_split_n_parts": 4, + "nvfp4_ffn_split_n_stride_workaround": True, + "nvfp4_large_gemm_cublaslt": False, + "nvfp4_large_gemm_cublaslt_algorithm": 3, + "nvfp4_qkv_cublaslt": True, + "nvfp4_qkv_cublaslt_algorithm": 2, + } + for options in ({key: value} for key, value in legacy_options.items()): + with self.subTest(options=options): + config = dict(options) + validate_thor_config(config) + self.assertEqual(config, options) + self.assertNotIn("thor", config) + for thor in (False, True): + with self.subTest(thor=thor): + config = self.make_config(thor=thor, **legacy_options) + expected = dict(config) + validate_thor_config(config) + self.assertEqual(config, expected) + + +if __name__ == "__main__": + unittest.main() diff --git a/test_cases/test_wan_merge.py b/test_cases/test_wan_merge.py new file mode 100644 index 000000000..709ed748d --- /dev/null +++ b/test_cases/test_wan_merge.py @@ -0,0 +1,246 @@ +"""Dependency-free routing tests for the Wan/Thor merge (no torch imports).""" + +import ast +import inspect +import unittest +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock + +ROOT = Path(__file__).resolve().parents[1] +MM = ROOT / "lightx2v/common/ops/mm/mm_weight.py" +WEIGHTS = ROOT / "lightx2v/models/networks/wan/weights/transformer_weights.py" +INFER = ROOT / "lightx2v/models/networks/wan/infer/transformer_infer.py" + + +def load_definitions(path, names, namespace): + tree = ast.parse(path.read_text()) + tree.body = [node for node in tree.body if isinstance(node, (ast.FunctionDef, ast.ClassDef)) and node.name in names] + # Execute only selected definitions from local source files, without package imports. + exec(compile(tree, str(path), "exec"), namespace) # noqa: S102 + return namespace + + +class Registry(dict): + def __call__(self, name): + def register(cls): + self[name] = cls + return cls + + return register + + +class WeightModule: + def __init__(self, *args, **kwargs): + pass + + def add_module(self, name, module): + setattr(self, name, module) + + +class Tensor: + shape = (4, 2, 8) + + def squeeze(self): + return self + + def view(self, *shape): + return self + + +class WanMergeTests(unittest.TestCase): + def setUp(self): + self.registry = Registry({name: Mock(name=name) for name in ("Default", "Calib", "nvfp4", "nvfp4-split-n-workaround", "nvfp4-split-n-stride-workaround", "TensorParallel")}) + self.namespace = load_definitions( + WEIGHTS, + {"_mm_weight", "WanFFN"}, + { + "MM_WEIGHT_REGISTER": self.registry, + "LN_WEIGHT_REGISTER": {"torch": Mock()}, + "WeightModule": WeightModule, + "dist": SimpleNamespace(get_rank=lambda group: 1, get_world_size=lambda group: 2), + }, + ) + + def test_override_precedes_calibration_and_forwards_kwargs(self): + mm_weight = self.namespace["_mm_weight"] + config = {"dit_quant_scheme": "nvfp4", "do_mm_calib": True} + mm_weight(config, "w", "b") + self.registry["Calib"].assert_called_once() + options = {"split_n_parts": 2} + mm_weight(config, "w", "b", mm_type_override="nvfp4-split-n-stride-workaround", mm_kwargs=options) + self.assertEqual(self.registry["nvfp4-split-n-stride-workaround"].call_args.kwargs["split_n_parts"], 2) + self.assertEqual(options, {"split_n_parts": 2}) + + def test_tp_helper_forwards_override_and_kwargs(self): + group = object() + config = {"tensor_parallel": True, "device_mesh": SimpleNamespace(get_group=lambda **kwargs: group), "do_mm_calib": True} + options = {"split_n_parts": 2} + self.namespace["_mm_weight"](config, "w", "b", split_dim="col", mm_type_override="nvfp4", mm_kwargs=options) + kwargs = self.registry["TensorParallel"].call_args.kwargs + self.assertEqual(kwargs["mm_type"], "nvfp4") + self.assertIs(kwargs["mm_kwargs"], options) + self.assertIs(kwargs["tp_group"], group) + self.assertEqual((kwargs["tp_rank"], kwargs["tp_size"]), (1, 2)) + + def test_ffn_registry_routing(self): + for thor, split_n, mm_type, expected in ( + (True, False, "nvfp4", "nvfp4-split-n-stride-workaround"), + (True, True, "nvfp4", "nvfp4-split-n-stride-workaround"), + (False, True, "nvfp4", "nvfp4-split-n-workaround"), + (False, False, "nvfp4", "nvfp4"), + (False, True, "Default", "Default"), + (False, True, "Calib", "Calib"), + ): + with self.subTest(thor=thor, split_n=split_n, mm_type=mm_type): + for factory in self.registry.values(): + factory.reset_mock() + config = {"thor": thor, "nvfp4_ffn_split_n_workaround": split_n, "dit_quant_scheme": mm_type} + self.namespace["WanFFN"](0, "blocks", "t2v", mm_type, config) + factory = self.registry[expected] + self.assertEqual(factory.call_count, 2) + for call in factory.call_args_list: + self.assertEqual(call.kwargs.get("split_n_parts"), 2 if thor else None) + self.assertEqual(sum(factory.call_count for factory in self.registry.values()), 2) + + def test_split_n_requires_boolean(self): + with self.assertRaises(TypeError): + self.namespace["WanFFN"](0, "blocks", "t2v", "nvfp4", {"nvfp4_ffn_split_n_workaround": 1}) + + def test_tp_positional_lora_chunks_and_mm_kwargs_coexist(self): + inner = Mock() + registry = Registry({"nvfp4": inner}) + namespace = load_definitions( + MM, + {"MMWeightTP"}, + { + "MMWeightTemplate": WeightModule, + "MMWeight": Mock(), + "MM_WEIGHT_REGISTER": registry, + }, + ) + cls = namespace["MMWeightTP"] + params = list(inspect.signature(cls).parameters) + self.assertEqual(params[-3:], ["reduce_output", "lora_column_chunks", "mm_kwargs"]) + weight = cls("w", "b", "nvfp4", None, 0, 2, "col", False, False, False, None, False, "prefix", "lora", False, 2, mm_kwargs={"split_n_parts": 2}) + self.assertEqual(weight.lora_column_chunks, 2) + self.assertFalse(weight.reduce_output) + self.assertEqual(inner.call_args.kwargs["split_n_parts"], 2) + self.assertEqual(inner.call_args.kwargs["lora_path"], "lora") + self.assertNotIn("lora_column_chunks", inner.call_args.kwargs) + + def test_split_n_registries_and_single_fp8_registration(self): + registry = Registry() + load_definitions( + MM, + {"MMWeightWnvfp4Anvfp4dynamicSplitNWorkaround", "MMWeightWnvfp4Anvfp4dynamicSplitNStrideWorkaround"}, + { + "MM_WEIGHT_REGISTER": registry, + "MMWeightWnvfp4Anvfp4dynamic": WeightModule, + }, + ) + self.assertEqual(set(registry), {"nvfp4-split-n-workaround", "nvfp4-split-n-stride-workaround"}) + self.assertEqual(registry["nvfp4-split-n-stride-workaround"]("w", "b", split_n_parts=2).split_n_parts, 2) + tree = ast.parse(MM.read_text()) + wrappers = [node for node in ast.walk(tree) if isinstance(node, ast.FunctionDef) and node.name == "_fp8_scaled_mm"] + self.assertEqual(len(wrappers), 1) + self.assertIn("sgl_fp8_scaled_mm_meta", ast.unparse(wrappers[0])) + self.assertIn("return sgl_fp8_scaled_mm(", ast.unparse(wrappers[0])) + + def test_attention_kwargs_reach_all_routes_with_thor_qkv(self): + namespace = load_definitions( + INFER, + {"WanTransformerInfer"}, + { + "WanMxfp8FuseMixin": type("Mixin", (), {}), + "BaseTransformerInfer": object, + "torch": SimpleNamespace(no_grad=lambda: lambda method: method, equal=lambda a, b: a == b), + }, + ) + cls = namespace["WanTransformerInfer"] + tensor = Tensor() + for thor in (False, True): + for route in ("local", "new", "legacy"): + with self.subTest(thor=thor, route=route): + infer = cls.__new__(cls) + infer.__dict__.update( + thor=thor, + cos_sin=None, + rope_positions=None, + sensitive_layer_dtype="bf16", + infer_dtype="bf16", + num_heads=2, + head_dim=8, + clean_cuda_cache=False, + block_idx=3, + scheduler=object(), + _sol_morton_preordered=False, + seq_parallel=route != "local", + use_new_seq_p_interface=route == "new", + self_attn_cu_seqlens_qkv=object(), + seq_p_group=None, + seq_p_prepost_backend="torch", + seq_p_a2a_backend="torch", + seq_p_quant_scheme=None, + seq_p_tensor_fusion=False, + seq_p_head_parallel=False, + seq_p_fp8_comm=False, + seq_p_fp4_comm=False, + seq_p_configured_quant_scheme=None, + has_post_adapter=False, + ) + infer._use_mxfp8_quant_fuse = lambda: False + infer._can_reuse_self_attn_mxfp8_quant = lambda *args: False + infer.modulate_func = lambda *args, **kwargs: tensor + + def projection(): + return SimpleNamespace( + apply=Mock(return_value=tensor), + apply_cublaslt=Mock(return_value=tensor), + apply_quantized_cublaslt=Mock(return_value=tensor), + act_quant_func=Mock(return_value=(tensor, tensor)), + input_global_scale=1, + ) + + phase = SimpleNamespace( + modulation=object(), + norm1=projection(), + self_attn_q=projection(), + self_attn_k=projection(), + self_attn_v=projection(), + self_attn_o=projection(), + self_attn_norm_q=projection(), + self_attn_norm_k=projection(), + rope=SimpleNamespace(apply=Mock(return_value=(tensor, tensor))), + self_attn_1=projection(), + self_attn_1_parallel=SimpleNamespace(apply=Mock(return_value=tensor), apply_new=Mock(return_value=(tensor, None))), + ) + infer.pre_process = lambda *args: (tensor,) * 6 + infer.infer_cross_attn = lambda *args: (tensor, tensor) + infer.infer_ffn = lambda *args: None + pre = SimpleNamespace(embed0=tensor, context=tensor, grid_sizes=SimpleNamespace(tuple=(1, 2, 2))) + marker = object() + options = {"cache_state": marker, "block_idx": 99} + self.assertIs(infer.infer_block(SimpleNamespace(compute_phases=[phase, object(), object()]), tensor, pre, options), tensor) + if route == "new": + kwargs = phase.self_attn_1_parallel.apply_new.call_args.kwargs["attention_kwargs"] + else: + attention = phase.self_attn_1 if route == "local" else phase.self_attn_1_parallel + kwargs = attention.apply.call_args.kwargs + self.assertIs(kwargs["cache_state"], marker) + self.assertEqual(kwargs["block_idx"], 99) + self.assertEqual(kwargs["grid_sizes"], (1, 2, 2)) + self.assertEqual(options, {"cache_state": marker, "block_idx": 99}) + if thor: + phase.self_attn_q.act_quant_func.assert_called_once() + phase.self_attn_k.act_quant_func.assert_not_called() + phase.self_attn_v.act_quant_func.assert_not_called() + for proj in (phase.self_attn_q, phase.self_attn_k, phase.self_attn_v): + proj.apply_quantized_cublaslt.assert_called_once_with(tensor, tensor, -1) + phase.self_attn_o.apply_cublaslt.assert_called_once_with(tensor, -1) + else: + phase.self_attn_o.apply.assert_called_once_with(tensor) + + +if __name__ == "__main__": + unittest.main() diff --git a/test_cases/test_wan_mxfp8_fuse_forwarding.py b/test_cases/test_wan_mxfp8_fuse_forwarding.py new file mode 100644 index 000000000..ebb21a3d7 --- /dev/null +++ b/test_cases/test_wan_mxfp8_fuse_forwarding.py @@ -0,0 +1,522 @@ +import os +import sys +import types +import unittest +from importlib import import_module +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +os.environ.setdefault("SKIP_PLATFORM_CHECK", "1") + + +def ensure_lightx2v_pipeline_stub(): + if "lightx2v.pipeline" not in sys.modules: + pipeline_stub = types.ModuleType("lightx2v.pipeline") + pipeline_stub.LightX2VPipeline = object + sys.modules["lightx2v.pipeline"] = pipeline_stub + + +def ensure_local_lightx2v_kernel(): + kernel_python_root = Path(__file__).resolve().parents[1] / "lightx2v_kernel" / "python" + kernel_python_root_str = str(kernel_python_root) + if kernel_python_root_str in sys.path: + sys.path.remove(kernel_python_root_str) + sys.path.insert(0, kernel_python_root_str) + for module_name in list(sys.modules): + if module_name == "lightx2v_kernel" or module_name.startswith("lightx2v_kernel."): + del sys.modules[module_name] + + +def make_config(**overrides): + config = { + "task": "i2v", + "num_layers": 1, + "num_heads": 1, + "dim": 8, + "seq_parallel": False, + "cpu_offload": False, + "modulate_type": "torch", + "rope_type": "torch", + "dit_quant_scheme": "mxfp8", + "mxfp8_fuse_enable": True, + "infer_steps": 4, + "teacache_thresh": 0.1, + "use_ret_steps": False, + "coefficients": ([1.0], [1.0]), + } + config.update(overrides) + return config + + +def make_linear_phase(): + return SimpleNamespace( + norm2=SimpleNamespace(apply=lambda x: x.clone()), + ffn_0=SimpleNamespace(apply=lambda x: x + 1), + ffn_2=SimpleNamespace(apply=lambda x: x + 2), + ) + + +class WanNvfp4SplitNStrideTest(unittest.TestCase): + @staticmethod + def make_ffn(**overrides): + ensure_lightx2v_pipeline_stub() + ensure_local_lightx2v_kernel() + lazy_load = overrides.pop("lazy_load", False) + transformer_weights = import_module("lightx2v.models.networks.wan.weights.transformer_weights") + config = { + "layer_norm_type": "torch", + "tensor_parallel": False, + } + config.update(overrides) + return transformer_weights.WanFFN( + block_index=0, + block_prefix="blocks", + task="i2v", + mm_type="nvfp4", + config=config, + lazy_load=lazy_load, + ) + + def test_stride_workaround_selects_full_weight_operator_for_both_ffn_layers(self): + ffn = self.make_ffn(thor=True) + mm_weight = import_module("lightx2v.common.ops.mm.mm_weight") + + self.assertIsInstance(ffn.ffn_0, mm_weight.MMWeightWnvfp4Anvfp4dynamicSplitNStrideWorkaround) + self.assertIsInstance(ffn.ffn_2, mm_weight.MMWeightWnvfp4Anvfp4dynamicSplitNStrideWorkaround) + self.assertEqual(ffn.ffn_0.split_n_parts, 2) + self.assertEqual(ffn.ffn_2.split_n_parts, 2) + + def test_stride_operator_passes_complete_weight_and_scale_tensors(self): + ensure_lightx2v_pipeline_stub() + ensure_local_lightx2v_kernel() + mm_weight = import_module("lightx2v.common.ops.mm.mm_weight") + operator = mm_weight.MMWeightWnvfp4Anvfp4dynamicSplitNStrideWorkaround( + "blocks.0.ffn.0.weight", + "blocks.0.ffn.0.bias", + split_n_parts=2, + ) + input_tensor = object() + input_quant = object() + input_scale = object() + weight = object() + weight_scale = object() + alpha = object() + bias = object() + output = object() + operator.act_quant_func = lambda value: (input_quant, input_scale) + operator.weight = weight + operator.weight_scale = weight_scale + operator.alpha = alpha + operator.bias = bias + + with patch.object(mm_weight, "cutlass_scaled_nvfp4_mm_split_n_stride", return_value=output) as kernel: + actual = operator.apply(input_tensor) + + self.assertIs(actual, output) + kernel.assert_called_once_with( + input_quant, + weight, + input_scale, + weight_scale, + alpha=alpha, + bias=bias, + split_n_parts=2, + ) + + def test_residual_gate_operator_forwards_complete_tensors(self): + ensure_lightx2v_pipeline_stub() + ensure_local_lightx2v_kernel() + mm_weight = import_module("lightx2v.common.ops.mm.mm_weight") + operator = mm_weight.MMWeightWnvfp4Anvfp4dynamicSplitNStrideWorkaround( + "blocks.0.ffn.2.weight", + "blocks.0.ffn.2.bias", + split_n_parts=2, + ) + values = {name: object() for name in ("input", "input_quant", "input_scale", "weight", "weight_scale", "alpha", "bias", "residual", "gate")} + operator.act_quant_func = lambda value: (values["input_quant"], values["input_scale"]) + operator.weight = values["weight"] + operator.weight_scale = values["weight_scale"] + operator.alpha = values["alpha"] + operator.bias = values["bias"] + + with patch.object( + mm_weight, + "cutlass_scaled_nvfp4_mm_split_n_stride_residual_gate", + return_value=values["residual"], + ) as kernel: + actual = operator.apply_residual_gate(values["input"], values["residual"], values["gate"]) + + self.assertIs(actual, values["residual"]) + kernel.assert_called_once_with( + values["input_quant"], + values["weight"], + values["input_scale"], + values["weight_scale"], + alpha=values["alpha"], + residual=values["residual"], + gate=values["gate"], + bias=values["bias"], + split_n_parts=2, + ) + + def test_gelu_operator_forwards_complete_tensors(self): + ensure_lightx2v_pipeline_stub() + ensure_local_lightx2v_kernel() + mm_weight = import_module("lightx2v.common.ops.mm.mm_weight") + operator = mm_weight.MMWeightWnvfp4Anvfp4dynamicSplitNStrideWorkaround( + "blocks.0.ffn.0.weight", + "blocks.0.ffn.0.bias", + split_n_parts=2, + ) + values = {name: object() for name in ("input", "input_quant", "input_scale", "weight", "weight_scale", "alpha", "bias", "output")} + operator.act_quant_func = lambda value: (values["input_quant"], values["input_scale"]) + operator.weight = values["weight"] + operator.weight_scale = values["weight_scale"] + operator.alpha = values["alpha"] + operator.bias = values["bias"] + + with patch.object( + mm_weight, + "cutlass_scaled_nvfp4_mm_split_n_stride_gelu", + return_value=values["output"], + ) as kernel: + actual = operator.apply_gelu(values["input"]) + + self.assertIs(actual, values["output"]) + kernel.assert_called_once_with( + values["input_quant"], + values["weight"], + values["input_scale"], + values["weight_scale"], + alpha=values["alpha"], + bias=values["bias"], + split_n_parts=2, + ) + + def test_thor_incompatibilities(self): + invalid_cases = [ + ({"tensor_parallel": True}, "tensor parallelism"), + ({"cpu_offload": True}, "CPU offload"), + ({"lazy_load": True}, "lazy loading"), + ({"lora_configs": [{}]}, "LoRA"), + ] + for extra_config, message in invalid_cases: + with self.subTest(extra_config=extra_config): + with self.assertRaisesRegex(NotImplementedError, message): + self.make_ffn(thor=True, **extra_config) + + +class WanMxfp8FuseForwardingTest(unittest.TestCase): + def test_base_infer_ffn_respects_mxfp8_fuse_enable(self): + ensure_lightx2v_pipeline_stub() + ensure_local_lightx2v_kernel() + transformer_infer = import_module("lightx2v.models.networks.wan.infer.transformer_infer") + + phase = make_linear_phase() + x = torch.zeros(1, 8) + attn_out = torch.zeros(1, 8) + c_shift = torch.zeros(1, 8) + c_scale = torch.zeros(1, 8) + + disabled = transformer_infer.WanTransformerInfer(make_config(mxfp8_fuse_enable=False)) + disabled._mxfp8_fuse_available = True + disabled._ensure_mxfp8_quant_ffn_ready = lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError("should not be called")) + y = disabled.infer_ffn(phase, x.clone(), attn_out.clone(), c_shift, c_scale, c_gate_msa=None) + self.assertIsInstance(y, torch.Tensor) + + enabled = transformer_infer.WanTransformerInfer(make_config(mxfp8_fuse_enable=True)) + enabled._mxfp8_fuse_available = True + enabled._ensure_mxfp8_quant_ffn_ready = lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError("fuse gate check reached")) + with self.assertRaisesRegex(RuntimeError, "fuse gate check reached"): + enabled.infer_ffn(phase, x.clone(), attn_out.clone(), c_shift, c_scale, c_gate_msa=None) + + def test_nvfp4_ffn0_gelu_fusion_dispatches_once(self): + ensure_lightx2v_pipeline_stub() + ensure_local_lightx2v_kernel() + transformer_infer = import_module("lightx2v.models.networks.wan.infer.transformer_infer") + + seen = {} + + def apply_gelu(value): + seen["input"] = value.clone() + return torch.full_like(value, 3) + + phase = make_linear_phase() + phase.ffn_0 = SimpleNamespace( + apply=lambda value: (_ for _ in ()).throw(RuntimeError("ordinary FFN0 must be skipped")), + apply_gelu=apply_gelu, + ) + phase.ffn_2 = SimpleNamespace( + apply=lambda value: (_ for _ in ()).throw(RuntimeError("ordinary FFN2 must be skipped")), + apply_residual_gate=lambda y, residual, gate: residual.copy_(y), + ) + infer = transformer_infer.WanTransformerInfer( + make_config( + dit_quant_scheme="nvfp4", + mxfp8_fuse_enable=False, + thor=True, + ) + ) + x = torch.zeros(1, 8) + gate = torch.ones_like(x) + output = infer.infer_ffn( + phase, + x, + torch.zeros_like(x), + torch.zeros_like(x), + torch.zeros_like(x), + gate, + ) + + self.assertTrue(torch.equal(seen["input"], torch.zeros(8))) + self.assertIsNone(output) + self.assertTrue(torch.equal(x, torch.full_like(x, 3.0))) + + def test_nvfp4_residual_gate_fusion_mutates_x_and_skips_post_process(self): + ensure_lightx2v_pipeline_stub() + ensure_local_lightx2v_kernel() + transformer_infer = import_module("lightx2v.models.networks.wan.infer.transformer_infer") + + gate = torch.full((1, 8), 0.5) + phase = make_linear_phase() + phase.ffn_0 = SimpleNamespace( + apply=lambda value: (_ for _ in ()).throw(RuntimeError("ordinary FFN0 must be skipped")), + apply_gelu=lambda value: value, + ) + seen = {} + + def apply_residual_gate(y, residual, actual_gate): + seen["gate"] = actual_gate + residual.add_(3) + return residual + + phase.ffn_2 = SimpleNamespace( + apply=lambda y: (_ for _ in ()).throw(RuntimeError("ordinary FFN2 must be skipped")), + apply_residual_gate=apply_residual_gate, + ) + infer = transformer_infer.WanTransformerInfer( + make_config( + dit_quant_scheme="nvfp4", + mxfp8_fuse_enable=False, + thor=True, + clean_cuda_cache=True, + ) + ) + x = torch.zeros(1, 8) + y = infer.infer_ffn( + phase, + x, + torch.zeros_like(x), + torch.zeros_like(x), + torch.zeros_like(x), + gate, + ) + + self.assertIsNone(y) + self.assertTrue(torch.equal(x, torch.full_like(x, 3))) + self.assertTrue(torch.equal(seen["gate"], gate.squeeze())) + self.assertIs(infer.post_process(x, y, gate), x) + + def test_offload_phase_two_forwards_c_gate_msa(self): + ensure_lightx2v_pipeline_stub() + ensure_local_lightx2v_kernel() + offload_infer = import_module("lightx2v.models.networks.wan.infer.offload.transformer_infer") + + infer = offload_infer.WanOffloadTransformerInfer(make_config()) + gate = torch.ones(1, 8) + infer.phase_params = { + "attn_out": torch.zeros(1, 8), + "c_shift_msa": torch.zeros(1, 8), + "c_scale_msa": torch.zeros(1, 8), + "c_gate_msa": gate, + "y": None, + } + seen = {} + + def fake_infer_ffn(phase, x, attn_out, c_shift, c_scale, c_gate=None): + seen["gate"] = c_gate + return torch.zeros_like(x) + + infer.infer_ffn = fake_infer_ffn + infer.post_process = lambda x, y, c_gate, pre_infer_out=None: x + + x = infer.infer_phase(2, SimpleNamespace(), torch.zeros(1, 8), SimpleNamespace(adapter_args={"hints": []})) + self.assertIsInstance(x, torch.Tensor) + self.assertIs(seen["gate"], gate) + + def test_self_forcing_forwards_c_gate_msa(self): + ensure_lightx2v_pipeline_stub() + ensure_local_lightx2v_kernel() + self_forcing = import_module("lightx2v.models.networks.wan.infer.self_forcing.transformer_infer") + + infer = self_forcing.WanSFTransformerInfer(make_config()) + gate = torch.ones(2, 1, 8) + infer.pre_process = lambda modulation, embed0: (gate, gate, gate, gate, gate, gate) + infer.infer_self_attn_with_kvcache = lambda *args, **kwargs: torch.zeros(4, 8) + infer.infer_cross_attn_with_kvcache = lambda *args, **kwargs: (torch.zeros(4, 8), torch.zeros(4, 8)) + seen = {} + + def fake_infer_ffn(phase, x, attn_out, c_shift, c_scale, c_gate=None): + seen["gate"] = c_gate + return torch.zeros_like(x) + + infer.infer_ffn = fake_infer_ffn + infer.post_process = lambda x, y, c_gate, pre_infer_out=None: x + + block = SimpleNamespace(compute_phases=[SimpleNamespace(modulation=None), SimpleNamespace(), SimpleNamespace()]) + pre_infer_out = SimpleNamespace( + x=torch.zeros(4, 8), + embed0=torch.zeros(2, 6, 8), + grid_sizes=SimpleNamespace(tensor=torch.ones(1, dtype=torch.int32)), + seq_lens=torch.ones(1, dtype=torch.int32), + freqs=torch.zeros(1), + context=torch.zeros(1, 8), + ) + infer.infer_block_with_kvcache(block, torch.zeros(4, 8), pre_infer_out) + self.assertIs(seen["gate"], gate) + + def test_lingbot_forwards_c_gate_msa(self): + ensure_lightx2v_pipeline_stub() + ensure_local_lightx2v_kernel() + lingbot = import_module("lightx2v.models.networks.wan.infer.lingbot.transformer_infer") + + infer = lingbot.WanLingbotTransformerInfer(make_config()) + gate = torch.ones(1, 8) + infer.pre_process = lambda modulation, embed0: (gate, gate, gate, gate, gate, gate) + infer.infer_self_attn = lambda *args, **kwargs: torch.zeros(1, 8) + infer.infer_cross_attn = lambda *args, **kwargs: (torch.zeros(1, 8), torch.zeros(1, 8)) + seen = {} + + def fake_infer_ffn(phase, x, attn_out, c_shift, c_scale, c_gate=None): + seen["gate"] = c_gate + return torch.zeros_like(x) + + infer.infer_ffn = fake_infer_ffn + infer.post_process = lambda x, y, c_gate, pre_infer_out=None: x + + block = SimpleNamespace(compute_phases=[SimpleNamespace(modulation=None), SimpleNamespace(), SimpleNamespace()]) + pre_infer_out = SimpleNamespace( + x=torch.zeros(1, 8), + embed0=torch.zeros(1, 6, 8), + context=torch.zeros(1, 8), + conditional_dict={}, + adapter_args={"hints": []}, + ) + infer.infer_block(block, torch.zeros(1, 8), pre_infer_out) + self.assertIs(seen["gate"], gate) + + def test_audio_forwards_c_gate_msa(self): + ensure_lightx2v_pipeline_stub() + ensure_local_lightx2v_kernel() + audio = import_module("lightx2v.models.networks.wan.infer.audio.transformer_infer") + + infer = audio.WanAudioARTransformerInfer(make_config()) + gate = torch.ones(1, 8) + infer.pre_process = lambda modulation, embed0: (gate, gate, gate, gate, gate, gate) + infer.infer_self_attn_with_kvcache = lambda *args, **kwargs: torch.zeros(1, 8) + infer.infer_cross_attn_with_kvcache = lambda *args, **kwargs: (torch.zeros(1, 8), torch.zeros(1, 8)) + seen = {} + + def fake_infer_ffn(phase, x, attn_out, c_shift, c_scale, c_gate=None): + seen["gate"] = c_gate + return torch.zeros_like(x) + + infer.infer_ffn = fake_infer_ffn + infer.post_process = lambda x, y, c_gate, pre_infer_out=None: x + + block = SimpleNamespace(compute_phases=[SimpleNamespace(modulation=None), SimpleNamespace(), SimpleNamespace(), SimpleNamespace()]) + pre_infer_out = SimpleNamespace( + x=torch.zeros(1, 8), + embed0=torch.zeros(1, 6, 8), + grid_sizes=SimpleNamespace(tensor=torch.ones(1, 3, dtype=torch.int32)), + seq_lens=torch.ones(1, dtype=torch.int32), + freqs=torch.zeros(1), + context=torch.zeros(1, 8), + adapter_args={"audio_encoder_output": None}, + ) + infer.infer_block_with_kvcache(block, torch.zeros(1, 8), pre_infer_out) + self.assertIs(seen["gate"], gate) + + def test_audio_ar_rope_modes(self): + ensure_lightx2v_pipeline_stub() + ensure_local_lightx2v_kernel() + audio = import_module("lightx2v.models.networks.wan.infer.audio.transformer_infer") + + local = audio.WanAudioARTransformerInfer(make_config()) + self.assertEqual(local.rope_position_mode, "local") + self.assertIsNone(local.rope_max_frames) + self.assertTrue(torch.equal(local._cache_positions_for_range(2, 5, "cpu"), torch.tensor([2, 3, 4]))) + + global_rope = audio.WanAudioARTransformerInfer(make_config(ar_config={"num_frame_per_chunk": 1, "rope_position_mode": "global", "rope_max_frames": 40})) + self.assertEqual(global_rope.rope_position_mode, "global") + self.assertEqual(global_rope._global_rope_start_frame(80, 1, 1), 40) + self.assertTrue( + torch.equal( + global_rope._cache_positions_for_range(2, 5, "cpu", global_end=51, sink_tokens=2), + torch.tensor([48, 49, 50]), + ) + ) + + def test_audio_ar_rope_mode_validation(self): + ensure_lightx2v_pipeline_stub() + ensure_local_lightx2v_kernel() + audio = import_module("lightx2v.models.networks.wan.infer.audio.transformer_infer") + + with self.assertRaisesRegex(ValueError, "rope_position_mode"): + audio.WanAudioARTransformerInfer(make_config(ar_config={"rope_position_mode": "invalid"})) + with self.assertRaisesRegex(ValueError, "requires"): + audio.WanAudioARTransformerInfer(make_config(ar_config={"rope_position_mode": "local", "rope_max_frames": 40})) + with self.assertRaisesRegex(ValueError, "at least"): + audio.WanAudioARTransformerInfer(make_config(ar_config={"num_frame_per_chunk": 2, "rope_position_mode": "global", "rope_max_frames": 1})) + + def test_feature_caching_variants_forward_c_gate_msa(self): + ensure_lightx2v_pipeline_stub() + ensure_local_lightx2v_kernel() + feature_caching = import_module("lightx2v.models.networks.wan.infer.feature_caching.transformer_infer") + + cases = [ + ( + feature_caching.WanTransformerInferTaylorCaching, + make_config(), + lambda infer, weights, x, embed0: infer.infer_calculating(weights, None, None, x, embed0, None, None, None), + ), + ( + feature_caching.WanTransformerInferAdaCaching, + make_config(), + lambda infer, weights, x, embed0: infer.infer_calculating(weights, None, None, x, embed0, None, None, None), + ), + ( + feature_caching.WanTransformerInferCustomCaching, + make_config(), + lambda infer, weights, x, embed0: infer.infer_calculating(weights, None, None, x, embed0, None, None, None), + ), + ] + + for cls, config, runner in cases: + with self.subTest(cls=cls.__name__): + infer = cls(config) + infer.scheduler = SimpleNamespace(infer_condition=True) + infer.derivative_approximation = lambda *args, **kwargs: None + gate = torch.ones(1, 8) + infer.infer_modulation = lambda phase, embed0: (gate, gate, gate, gate, gate, gate) + infer.infer_self_attn = lambda *args, **kwargs: torch.zeros(1, 8) + infer.infer_cross_attn = lambda *args, **kwargs: (torch.zeros(1, 8), torch.zeros(1, 8)) + seen = {} + + def fake_infer_ffn(phase, x, attn_out, c_shift, c_scale, c_gate=None): + seen["gate"] = c_gate + return torch.zeros_like(x) + + infer.infer_ffn = fake_infer_ffn + infer.post_process = lambda x, y, c_gate, pre_infer_out=None: x + + weights = SimpleNamespace(blocks=[SimpleNamespace(compute_phases=[SimpleNamespace(), SimpleNamespace(), SimpleNamespace(), SimpleNamespace()])]) + runner(infer, weights, torch.zeros(1, 8), torch.zeros(1, 6, 8)) + self.assertIs(seen["gate"], gate) + + +if __name__ == "__main__": + unittest.main() diff --git a/test_cases/test_wan_nvfp4_qkv_cublaslt.py b/test_cases/test_wan_nvfp4_qkv_cublaslt.py new file mode 100644 index 000000000..e7978d3c2 --- /dev/null +++ b/test_cases/test_wan_nvfp4_qkv_cublaslt.py @@ -0,0 +1,106 @@ +import os +import sys +import types +import unittest +from importlib import import_module +from pathlib import Path +from unittest.mock import patch + +os.environ.setdefault("SKIP_PLATFORM_CHECK", "1") + + +def ensure_lightx2v_pipeline_stub(): + if "lightx2v.pipeline" not in sys.modules: + pipeline_stub = types.ModuleType("lightx2v.pipeline") + pipeline_stub.LightX2VPipeline = object + sys.modules["lightx2v.pipeline"] = pipeline_stub + + +def ensure_local_lightx2v_kernel(): + kernel_python_root = Path(__file__).resolve().parents[1] / "lightx2v_kernel" / "python" + kernel_python_root_str = str(kernel_python_root) + if kernel_python_root_str in sys.path: + sys.path.remove(kernel_python_root_str) + sys.path.insert(0, kernel_python_root_str) + + +class WanNvfp4QkvCublasltTest(unittest.TestCase): + @classmethod + def setUpClass(cls): + ensure_lightx2v_pipeline_stub() + ensure_local_lightx2v_kernel() + + def test_quantized_projection_forwards_to_cublaslt(self): + mm_weight = import_module("lightx2v.common.ops.mm.mm_weight") + operator = mm_weight.MMWeightWnvfp4Anvfp4dynamic( + "blocks.0.self_attn.q.weight", + "blocks.0.self_attn.q.bias", + ) + input_quant = object() + input_scale = object() + weight = object() + weight_scale = object() + alpha = object() + bias = object() + output = object() + operator.weight = weight + operator.weight_scale = weight_scale + operator.alpha = alpha + operator.bias = bias + + with patch.object(mm_weight, "cublaslt_scaled_nvfp4_mm_bias", return_value=output) as kernel: + actual = operator.apply_quantized_cublaslt(input_quant, input_scale) + + self.assertIs(actual, output) + kernel.assert_called_once_with( + input_quant, + weight, + input_scale, + weight_scale, + alpha=alpha, + bias=bias, + algorithm_index=-1, + ) + + def test_projection_quantizes_then_forwards_to_cublaslt(self): + mm_weight = import_module("lightx2v.common.ops.mm.mm_weight") + operator = mm_weight.MMWeightWnvfp4Anvfp4dynamic( + "blocks.0.self_attn.o.weight", + "blocks.0.self_attn.o.bias", + ) + input_tensor = object() + input_quant = object() + input_scale = object() + output = object() + operator.act_quant_func = unittest.mock.Mock(return_value=(input_quant, input_scale)) + operator.apply_quantized_cublaslt = unittest.mock.Mock(return_value=output) + + actual = operator.apply_cublaslt(input_tensor, algorithm_index=2) + + self.assertIs(actual, output) + operator.act_quant_func.assert_called_once_with(input_tensor) + operator.apply_quantized_cublaslt.assert_called_once_with( + input_quant, + input_scale, + 2, + ) + + def test_thor_requires_nvfp4_weights(self): + transformer_weights = import_module("lightx2v.models.networks.wan.weights.transformer_weights") + config = { + "rope_type": "torch_real_rope", + "layer_norm_type": "torch", + "rms_norm_type": "torch", + "tensor_parallel": False, + "seq_parallel": False, + "thor": True, + } + + with self.assertRaisesRegex(ValueError, "requires dit_quant_scheme='nvfp4'"): + transformer_weights.WanSelfAttention( + block_index=0, + block_prefix="blocks", + task="i2v", + mm_type="Default", + config=config, + )