From 058f61236d8cee3a19a9e474a0c598ecd7a74776 Mon Sep 17 00:00:00 2001 From: jiahanc <173873397+jiahanc@users.noreply.github.com> Date: Wed, 8 Jul 2026 21:29:24 -0700 Subject: [PATCH 01/10] add m3 config MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 中文：添加 MiniMax-M3 配置。 --- .../fixed_seq_len/minimaxm3_fp4_b300.sh | 42 +- .../flashinfer-cutedsl-splitk-gemm.patch | 3115 +++++++++++++++++ configs/nvidia-master.yaml | 2 +- perf-changelog.yaml | 8 + 4 files changed, 3165 insertions(+), 2 deletions(-) create mode 100644 benchmarks/single_node/fixed_seq_len/patches/flashinfer-cutedsl-splitk-gemm.patch diff --git a/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh b/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh index f91419edb7..ad955e031f 100755 --- a/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh +++ b/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh @@ -3,7 +3,11 @@ # MiniMax-M3 NVFP4 B300 single-node vLLM recipe. # Same shape as minimaxm3_fp8_b300.sh but uses the nvidia/MiniMax-M3-NVFP4 # checkpoint. MiniMax-M3 modelopt NVFP4 support (vllm-project/vllm PR #46380) is -# baked into the perf container image, so no runtime patch is needed. +# baked into the perf container image. +# +# At runtime the recipe swaps the image's FlashInfer for the pinned nightly +# release and applies the CuTeDSL split-K gemm patch on top, mirroring vLLM +# commits 82a00090a (nightly install) and 0a16bea6b (patch). source "$(dirname "$0")/../../benchmark_lib.sh" @@ -19,6 +23,40 @@ check_env_vars \ RANDOM_RANGE_RATIO \ RESULT_FILENAME +# # --- FlashInfer nightly + CuTeDSL split-K gemm patch ------------------------- +# # Pinned to nightly release nightly-v0.6.14-20260708. The wheels are not yet on +# # the flashinfer.ai/whl/nightly flat index, so install by direct release URL. +# FLASHINFER_VERSION=0.6.14.dev20260708 +# FLASHINFER_NIGHTLY_TAG=nightly-v0.6.14-20260708 +# FLASHINFER_RELEASE_URL="https://github.com/flashinfer-ai/flashinfer/releases/download/${FLASHINFER_NIGHTLY_TAG}" + +# # The flashinfer-jit-cache wheel is CUDA-version-specific; derive cuXYZ from +# # the torch build inside the container (e.g. 13.0 -> cu130). +# CUDA_MAJOR_MINOR="cu$(python3 -c 'import torch; print(torch.version.cuda.replace(".", ""))')" \ +# || { echo "Failed to determine CUDA version from torch" >&2; exit 1; } + +# python3 -m pip uninstall -y flashinfer-python flashinfer-cubin flashinfer-jit-cache + +# python3 -m pip install --pre \ +# "${FLASHINFER_RELEASE_URL}/flashinfer_python-${FLASHINFER_VERSION}-py3-none-any.whl" \ +# "${FLASHINFER_RELEASE_URL}/flashinfer_cubin-${FLASHINFER_VERSION}-py3-none-any.whl" \ +# "${FLASHINFER_RELEASE_URL}/flashinfer_jit_cache-${FLASHINFER_VERSION}+${CUDA_MAJOR_MINOR}-cp39-abi3-manylinux_2_28_$(uname -m).whl" \ +# || { echo "FlashInfer nightly install failed" >&2; exit 1; } + +# # Apply CuTeDSL split-K gemm patch (jiahanc/flashinfer branch cutedsl-splitK, +# # flashinfer/gemm changes only) on top of the pinned nightly. The reinstall +# # above guarantees pristine files, so the patch always applies cleanly. +# FLASHINFER_PATCH="$(dirname "$0")/patches/flashinfer-cutedsl-splitk-gemm.patch" +# if ! command -v patch >/dev/null 2>&1; then +# apt-get update -y && apt-get install -y --no-install-recommends patch \ +# || { echo "Failed to install patch(1)" >&2; exit 1; } +# fi +# SITE_PACKAGES=$(dirname "$(python3 -c "import importlib.util; print(importlib.util.find_spec('flashinfer').submodule_search_locations[0])")") \ +# || { echo "Could not locate the installed flashinfer package" >&2; exit 1; } +# patch -p1 -d "${SITE_PACKAGES}" < "$FLASHINFER_PATCH" \ +# || { echo "FlashInfer CuTeDSL split-K gemm patch failed" >&2; exit 1; } +# # ----------------------------------------------------------------------------- + if [[ -n "${MODEL_PATH:-}" ]]; then if [[ ! -d "$MODEL_PATH" || -z "$(ls -A "$MODEL_PATH" 2>/dev/null)" ]]; then hf download "$MODEL" --local-dir "$MODEL_PATH" @@ -39,6 +77,7 @@ SERVER_LOG=/workspace/server.log export VLLM_ENGINE_READY_TIMEOUT_S=3600 export VLLM_FLOAT32_MATMUL_PRECISION=high export VLLM_FLASHINFER_ALLREDUCE_BACKEND=trtllm +export VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=1800 if [ "${DP_ATTENTION}" = "true" ]; then PARALLEL_ARGS="--tensor-parallel-size=1 --data-parallel-size=$TP --enable-expert-parallel" @@ -57,6 +96,7 @@ start_gpu_monitor set -x vllm serve "$MODEL_PATH" --served-model-name "$MODEL" --host 0.0.0.0 --port $PORT \ $PARALLEL_ARGS \ +--attention_config.indexer_kv_dtype fp8 \ --gpu-memory-utilization 0.95 \ --max-model-len $MAX_MODEL_LEN \ --kv-cache-dtype fp8 \ diff --git a/benchmarks/single_node/fixed_seq_len/patches/flashinfer-cutedsl-splitk-gemm.patch b/benchmarks/single_node/fixed_seq_len/patches/flashinfer-cutedsl-splitk-gemm.patch new file mode 100644 index 0000000000..965114c94e --- /dev/null +++ b/benchmarks/single_node/fixed_seq_len/patches/flashinfer-cutedsl-splitk-gemm.patch @@ -0,0 +1,3115 @@ +diff --git a/flashinfer/gemm/gemm_base.py b/flashinfer/gemm/gemm_base.py +index 1a665cfd..c95bfaa4 100755 +--- a/flashinfer/gemm/gemm_base.py ++++ b/flashinfer/gemm/gemm_base.py +@@ -15,9 +15,10 @@ limitations under the License. + """ + + import functools ++from dataclasses import replace + from enum import Enum + from types import SimpleNamespace +-from typing import List, Literal, Optional, Tuple ++from typing import Callable, List, Literal, Optional, Tuple + + from flashinfer.trtllm_low_latency_gemm import trtllm_low_latency_gemm + import torch +@@ -4606,6 +4607,7 @@ def _compile_block_scaled_gemm( + sf_n, + sf_k, + batch_size, ++ cluster_shape_k=1, + ): + """Compile a block-scaled GEMM kernel via CuTe DSL and cache it. + +@@ -4671,9 +4673,8 @@ def _compile_block_scaled_gemm( + cutlass.Float32, (1,), assumed_align=4 + ) + +- max_active_clusters = get_max_active_clusters( +- cluster_shape_mn[0] * cluster_shape_mn[1] +- ) ++ launch_cluster_size = cluster_shape_mn[0] * cluster_shape_mn[1] * cluster_shape_k ++ max_active_clusters = get_max_active_clusters(launch_cluster_size) + stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + + compiled_gemm = cute.compile( +@@ -4721,6 +4722,35 @@ def _prepare_alpha_for_launch(alpha_tensor, device): + + _CUTE_DSL_MM_MXFP8_KERNEL_CACHE: dict[tuple, tuple] = {} + ++_SM100_MXFP8_SPLIT_K_MAX_M = 32 ++_SM100_MXFP8_SPLIT_K_MMA_INST_BITS_K = 256 ++_SM100_MXFP8_SPLIT_K_MMA_INST_TILE_K = 4 ++_SM100_MXFP8_SPLIT_K_CANDIDATES = (2, 4) ++ ++ ++def _is_sm100_split_k_valid( ++ m: int, ++ k: int, ++ ab_dtype, ++ split_k_slices: int, ++) -> bool: ++ """Check the specialized low-M split-K shape constraints.""" ++ if ab_dtype.width != 8: ++ return False ++ tile_k = ( ++ _SM100_MXFP8_SPLIT_K_MMA_INST_BITS_K * _SM100_MXFP8_SPLIT_K_MMA_INST_TILE_K ++ ) // ab_dtype.width ++ return ( ++ m <= _SM100_MXFP8_SPLIT_K_MAX_M ++ and split_k_slices in _SM100_MXFP8_SPLIT_K_CANDIDATES ++ and k % (tile_k * split_k_slices) == 0 ++ ) ++ ++ ++def _sm100_mxfp8_low_m_tile(m: int) -> tuple[int, int]: ++ tile_n = 8 if m <= 8 else 16 if m <= 16 else 32 ++ return (128, tile_n) ++ + + def _check_cute_dsl_availability(): + try: +@@ -4743,6 +4773,9 @@ def _cute_dsl_gemm_mxfp8_runner( + from .kernels.dense_blockscaled_gemm_sm100 import ( + Sm100BlockScaledPersistentDenseGemmKernel, + ) ++ from .kernels.dense_blockscaled_gemm_sm100_splitk import ( ++ Sm100BlockScaledSplitKGemmKernel, ++ ) + + if out_dtype not in (torch.bfloat16, torch.float16): + raise ValueError( +@@ -4760,22 +4793,51 @@ def _cute_dsl_gemm_mxfp8_runner( + _ = sm_major, sm_minor + + class CuteDSLMxfp8GemmRunner(TunableRunner): ++ def get_cache_key_extras(self, inputs: List[torch.Tensor]) -> tuple: ++ _, _, _, _, _, out, _ = inputs ++ return (str(out.dtype), enable_pdl) ++ + def get_valid_tactics( + self, + inputs: List[torch.Tensor], + profile: OptimizationProfile, + ) -> list: + (a, b, a_descale, b_descale, _, out, _) = inputs +- return _get_sm100_block_scaled_tactics( +- m=a.shape[0], +- n=b.shape[1], +- real_k=a.shape[1], +- ab_dtype=cutlass.Float8E4M3FN, ++ m = a.shape[0] ++ n = b.shape[1] ++ real_k = a.shape[1] ++ ab_dtype = cutlass.Float8E4M3FN ++ base_tactics = _get_sm100_block_scaled_tactics( ++ m=m, ++ n=n, ++ real_k=real_k, ++ ab_dtype=ab_dtype, + sf_dtype=cutlass.Float8E8M0FNU, + sf_vec_size=32, + c_cutlass_dtype=c_cutlass_dtype, + device=a.device, + ) ++ valid_tactics = [(*tactic, 1) for tactic in base_tactics] ++ if not out.is_contiguous(): ++ return valid_tactics ++ ++ for split_k_slices in _SM100_MXFP8_SPLIT_K_CANDIDATES: ++ if _is_sm100_split_k_valid( ++ m, ++ real_k, ++ ab_dtype, ++ split_k_slices, ++ ): ++ valid_tactics.append( ++ ( ++ _sm100_mxfp8_low_m_tile(m), ++ (1, 1), ++ True, ++ False, ++ split_k_slices, ++ ) ++ ) ++ return valid_tactics + + def forward( + self, +@@ -4794,14 +4856,50 @@ def _cute_dsl_gemm_mxfp8_runner( + batch_size = 1 + + if tactic is None or tactic == -1: +- tactic = ( +- _SM100_DEFAULT_MMA_TILER_MN, +- _SM100_DEFAULT_CLUSTER_SHAPE_MN, +- False, +- False, +- ) ++ if m <= _SM100_MXFP8_SPLIT_K_MAX_M and out.is_contiguous(): ++ # Untuned low-M execution uses the corresponding base tactic. ++ tactic = ( ++ _sm100_mxfp8_low_m_tile(m), ++ (1, 1), ++ True, ++ False, ++ 1, ++ ) ++ else: ++ tactic = ( ++ _SM100_DEFAULT_MMA_TILER_MN, ++ _SM100_DEFAULT_CLUSTER_SHAPE_MN, ++ False, ++ False, ++ 1, ++ ) + +- (mma_tiler_mn, cluster_shape_mn, swap_ab, use_prefetch) = tactic ++ ( ++ mma_tiler_mn, ++ cluster_shape_mn, ++ swap_ab, ++ use_prefetch, ++ split_k_slices, ++ ) = tactic ++ if split_k_slices < 1: ++ raise ValueError(f"Invalid MXFP8 split-K tactic: {tactic}") ++ is_split_k = split_k_slices > 1 ++ ++ if is_split_k: ++ if ( ++ cluster_shape_mn != (1, 1) ++ or not swap_ab ++ or use_prefetch ++ or not out.is_contiguous() ++ or mma_tiler_mn != _sm100_mxfp8_low_m_tile(m) ++ or not _is_sm100_split_k_valid( ++ m, ++ real_k, ++ cutlass.Float8E4M3FN, ++ split_k_slices, ++ ) ++ ): ++ raise ValueError(f"Invalid MXFP8 split-K tactic: {tactic}") + + if swap_ab: + kernel_m, kernel_n = n, m +@@ -4827,18 +4925,31 @@ def _cute_dsl_gemm_mxfp8_runner( + use_prefetch, + enable_pdl, + out_dtype, ++ split_k_slices, + ) + +- compiled_gemm, _ = _compile_block_scaled_gemm( +- _CUTE_DSL_MM_MXFP8_KERNEL_CACHE, +- cache_key, +- lambda: Sm100BlockScaledPersistentDenseGemmKernel( ++ make_kernel: Callable[[], object] ++ kernel_cache = _CUTE_DSL_MM_MXFP8_KERNEL_CACHE ++ if is_split_k: ++ make_kernel = lambda: Sm100BlockScaledSplitKGemmKernel( ++ sf_vec_size, ++ mma_tiler_mn, ++ split_k_slices, ++ enable_pdl, ++ ) ++ else: ++ make_kernel = lambda: Sm100BlockScaledPersistentDenseGemmKernel( + sf_vec_size, + mma_tiler_mn, + cluster_shape_mn, + use_prefetch, + enable_pdl, +- ), ++ ) ++ ++ compiled_gemm, _ = _compile_block_scaled_gemm( ++ kernel_cache, ++ cache_key, ++ make_kernel, + ab_cutlass_dtype=cutlass.Float8E4M3FN, + sf_dtype=sf_dtype, + c_cutlass_dtype=c_cutlass_dtype, +@@ -4849,6 +4960,7 @@ def _cute_dsl_gemm_mxfp8_runner( + sf_n=sf_n, + sf_k=sf_k, + batch_size=batch_size, ++ cluster_shape_k=split_k_slices, + ) + + alpha_for_launch = _prepare_alpha_for_launch(None, a.device) +@@ -4856,6 +4968,7 @@ def _cute_dsl_gemm_mxfp8_runner( + launch_out = ( + out.as_strided(out.shape, (1, out.shape[0])) if swap_ab else out + ) ++ + compiled_gemm( + kernel_a, + kernel_b, +@@ -5163,7 +5276,11 @@ def mm_mxfp8( + + tuner = AutoTuner.get() + +- tuning_config = _MM_MXFP8_TUNING_CONFIG ++ tuning_config = ( ++ _MM_MXFP8_CUTE_DSL_TUNING_CONFIG ++ if backends == ["cute-dsl"] ++ else _MM_MXFP8_TUNING_CONFIG ++ ) + + inputs = [ + a, +@@ -6323,6 +6440,12 @@ _MM_MXFP8_TUNING_CONFIG = TuningConfig( + ), + ) + ++_MM_MXFP8_CUTE_DSL_TUNING_CONFIG = replace( ++ _MM_MXFP8_TUNING_CONFIG, ++ use_cuda_graph=True, ++ use_cold_l2_cache=True, ++) ++ + + @backend_requirement( + { +diff --git a/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm100.py b/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm100.py +index 85a006d2..f653ea8a 100644 +--- a/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm100.py ++++ b/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm100.py +@@ -41,11 +41,13 @@ import cutlass.utils.blackwell_helpers as sm100_utils + import cutlass.utils.blockscaled_layout as blockscaled_utils + from cutlass.cute.nvgpu import cpasync, tcgen05 + ++from .dense_blockscaled_gemm_sm100_common import _Sm100BlockScaledGemmCommon ++ + from cutlass.cute.arch import griddepcontrol_launch_dependents, griddepcontrol_wait + from cutlass.pipeline import PipelineTmaUmma, PipelineUmmaAsync + + +-class Sm100BlockScaledPersistentDenseGemmKernel: ++class Sm100BlockScaledPersistentDenseGemmKernel(_Sm100BlockScaledGemmCommon): + """Implements batched matrix multiplication (C = A x SFA x B x SFB) with support for various data types + and Blackwell GPU architectural features, including persistent tile scheduling and warp specialization. + +@@ -254,6 +256,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel: + self.sf_vec_size, + self.smem_capacity, + self.occupancy, ++ 1024, + ) + + # Compute A/B/SFA/SFB/C shared memory layout +@@ -344,155 +347,25 @@ class Sm100BlockScaledPersistentDenseGemmKernel: + Raises: + TypeError: If input data types are incompatible with the MMA instruction. + """ +- # Setup static attributes before smem/grid/tma computation +- self.a_dtype: Type[cutlass.Numeric] = a_tensor.element_type +- self.b_dtype: Type[cutlass.Numeric] = b_tensor.element_type +- self.sf_dtype: Type[cutlass.Numeric] = sfa_tensor.element_type +- self.c_dtype: Type[cutlass.Numeric] = c_tensor.element_type +- self.a_major_mode = utils.LayoutEnum.from_tensor(a_tensor).mma_major_mode() +- self.b_major_mode = utils.LayoutEnum.from_tensor(b_tensor).mma_major_mode() +- self.c_layout = utils.LayoutEnum.from_tensor(c_tensor) +- +- # Check if input data types are compatible with MMA instruction +- if cutlass.const_expr(self.a_dtype != self.b_dtype): +- raise TypeError(f"Type must match: {self.a_dtype} != {self.b_dtype}") +- +- # Setup attributes that dependent on gemm inputs +- self._setup_attributes() +- +- # Setup sfa/sfb tensor by filling A/B tensor to scale factor atom layout +- # ((Atom_M, Rest_M),(Atom_K, Rest_K),RestL) +- sfa_layout = blockscaled_utils.tile_atom_to_shape_SF( +- a_tensor.shape, self.sf_vec_size +- ) +- sfa_tensor = cute.make_tensor(sfa_tensor.iterator, sfa_layout) +- +- # ((Atom_N, Rest_N),(Atom_K, Rest_K),RestL) +- sfb_layout = blockscaled_utils.tile_atom_to_shape_SF( +- b_tensor.shape, self.sf_vec_size +- ) +- sfb_tensor = cute.make_tensor(sfb_tensor.iterator, sfb_layout) +- +- tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma( +- self.a_dtype, +- self.b_dtype, +- self.a_major_mode, +- self.b_major_mode, +- self.sf_dtype, +- self.sf_vec_size, +- self.cta_group, +- self.mma_inst_shape_mnk[:2], +- ) +- +- tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma( +- self.a_dtype, +- self.b_dtype, +- self.a_major_mode, +- self.b_major_mode, +- self.sf_dtype, +- self.sf_vec_size, +- cute.nvgpu.tcgen05.CtaGroup.ONE, +- self.mma_inst_shape_mnk_sfb[:2], +- ) +- atom_thr_size = cute.size(tiled_mma.thr_id.shape) +- +- # Setup TMA load for A +- a_op = sm100_utils.cluster_shape_to_tma_atom_A( +- self.cluster_shape_mn, tiled_mma.thr_id +- ) +- a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) +- tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( +- a_op, +- a_tensor, +- a_smem_layout, +- self.mma_tiler, ++ ( + tiled_mma, +- self.cluster_layout_vmnk.shape, +- ) +- +- # Setup TMA load for B +- b_op = sm100_utils.cluster_shape_to_tma_atom_B( +- self.cluster_shape_mn, tiled_mma.thr_id +- ) +- b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) +- tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( +- b_op, ++ tiled_mma_sfb, ++ tma_atom_a, ++ tma_tensor_a, ++ tma_atom_b, ++ tma_tensor_b, ++ tma_atom_sfa, ++ tma_tensor_sfa, ++ tma_atom_sfb, ++ tma_tensor_sfb, ++ tma_atom_c, ++ tma_tensor_c, ++ ) = self._prepare_tma( ++ a_tensor, + b_tensor, +- b_smem_layout, +- self.mma_tiler, +- tiled_mma, +- self.cluster_layout_vmnk.shape, +- ) +- +- # Setup TMA load for SFA +- sfa_op = sm100_utils.cluster_shape_to_tma_atom_A( +- self.cluster_shape_mn, tiled_mma.thr_id +- ) +- sfa_smem_layout = cute.slice_( +- self.sfa_smem_layout_staged, (None, None, None, 0) +- ) +- tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A( +- sfa_op, + sfa_tensor, +- sfa_smem_layout, +- self.mma_tiler, +- tiled_mma, +- self.cluster_layout_vmnk.shape, +- internal_type=cutlass.Int16, +- ) +- +- # Setup TMA load for SFB +- sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB( +- self.cluster_shape_mn, tiled_mma.thr_id +- ) +- sfb_smem_layout = cute.slice_( +- self.sfb_smem_layout_staged, (None, None, None, 0) +- ) +- tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B( +- sfb_op, + sfb_tensor, +- sfb_smem_layout, +- self.mma_tiler_sfb, +- tiled_mma_sfb, +- self.cluster_layout_sfb_vmnk.shape, +- internal_type=cutlass.Int16, +- ) +- if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 192): +- x = tma_tensor_sfb.stride[0][1] +- y = cute.ceil_div(tma_tensor_sfb.shape[0][1], 4) +- +- new_shape = ( +- (tma_tensor_sfb.shape[0][0], ((2, 2), y)), +- tma_tensor_sfb.shape[1], +- tma_tensor_sfb.shape[2], +- ) +- # Use right multiplication for ScaledBasis (3 * x instead of x * 3) +- x_times_3 = 3 * x +- new_stride = ( +- (tma_tensor_sfb.stride[0][0], ((x, x), x_times_3)), +- tma_tensor_sfb.stride[1], +- tma_tensor_sfb.stride[2], +- ) +- tma_tensor_sfb_new_layout = cute.make_layout(new_shape, stride=new_stride) +- tma_tensor_sfb = cute.make_tensor( +- tma_tensor_sfb.iterator, tma_tensor_sfb_new_layout +- ) +- +- a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout) +- b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout) +- sfa_copy_size = cute.size_in_bytes(self.sf_dtype, sfa_smem_layout) +- sfb_copy_size = cute.size_in_bytes(self.sf_dtype, sfb_smem_layout) +- self.num_tma_load_bytes = ( +- a_copy_size + b_copy_size + sfa_copy_size + sfb_copy_size +- ) * atom_thr_size +- +- # Setup TMA store for C +- epi_smem_layout = cute.slice_(self.c_smem_layout_staged, (None, None, 0)) +- tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( +- cpasync.CopyBulkTensorTileS2GOp(), + c_tensor, +- epi_smem_layout, +- self.epi_tile, + ) + + # Compute grid size +@@ -1537,295 +1410,6 @@ class Sm100BlockScaledPersistentDenseGemmKernel: + + griddepcontrol_launch_dependents() + +- def mainloop_s2t_copy_and_partition( +- self, +- sSF: cute.Tensor, +- tSF: cute.Tensor, +- ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: +- """ +- Make tiledCopy for smem to tmem load for scale factor tensor, then use it to partition smem memory (source) and tensor memory (destination). +- +- Args: +- sSF (cute.Tensor): The scale factor tensor in smem +- tSF (cute.Tensor): The scale factor tensor in tmem +- +- Returns: +- A tuple containing (tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t) where: +- - tiled_copy_s2t: The tiled copy operation for smem to tmem load for scale factor tensor(s2t) +- - tCsSF_compact_s2t: The partitioned scale factor tensor in smem +- - tSF_compact_s2t: The partitioned scale factor tensor in tmem +- """ +- # (MMA, MMA_MN, MMA_K, STAGE) +- tCsSF_compact = cute.filter_zeros(sSF) +- # (MMA, MMA_MN, MMA_K) +- tCtSF_compact = cute.filter_zeros(tSF) +- +- # Make S2T CopyAtom and tiledCopy +- copy_atom_s2t = cute.make_copy_atom( +- tcgen05.Cp4x32x128bOp(self.cta_group), +- self.sf_dtype, +- ) +- tiled_copy_s2t = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSF_compact) +- thr_copy_s2t = tiled_copy_s2t.get_slice(0) +- +- # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE) +- tCsSF_compact_s2t_ = thr_copy_s2t.partition_S(tCsSF_compact) +- # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE) +- tCsSF_compact_s2t = tcgen05.get_s2t_smem_desc_tensor( +- tiled_copy_s2t, tCsSF_compact_s2t_ +- ) +- # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K) +- tCtSF_compact_s2t = thr_copy_s2t.partition_D(tCtSF_compact) +- +- return tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t +- +- def epilog_tmem_copy_and_partition( +- self, +- tidx: cutlass.Int32, +- tAcc: cute.Tensor, +- gC_mnl: cute.Tensor, +- epi_tile: cute.Tile, +- use_2cta_instrs: Union[cutlass.Boolean, bool], +- ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: +- """ +- Make tiledCopy for tensor memory load, then use it to partition tensor memory (source) and register array (destination). +- +- Args: +- tidx (cutlass.Int32): The thread index in epilogue warp groups +- tAcc (cute.Tensor): The accumulator tensor to be copied and partitioned +- gC_mnl (cute.Tensor): The global tensor C +- epi_tile (cute.Tile): The epilogue tiler +- use_2cta_instrs (bool): Whether use_2cta_instrs is enabled +- +- Returns: +- A tuple containing (tiled_copy_t2r, tTR_tAcc, tTR_rAcc) where: +- - tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r) +- - tTR_tAcc: The partitioned accumulator tensor +- - tTR_rAcc: The accumulated tensor in register used to hold t2r results +- """ +- # Make tiledCopy for tensor memory load +- copy_atom_t2r = sm100_utils.get_tmem_load_op( +- self.cta_tile_shape_mnk, +- self.c_layout, +- self.c_dtype, +- self.acc_dtype, +- epi_tile, +- use_2cta_instrs, +- ) +- # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, STAGE) +- tAcc_epi = cute.flat_divide( +- tAcc[((None, None), 0, 0, None)], +- epi_tile, +- ) +- # (EPI_TILE_M, EPI_TILE_N) +- tiled_copy_t2r = tcgen05.make_tmem_copy( +- copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)] +- ) +- +- thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) +- # (T2R, T2R_M, T2R_N, EPI_M, EPI_M, STAGE) +- tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi) +- +- # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL) +- gC_mnl_epi = cute.flat_divide( +- gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile +- ) +- # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL) +- tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi) +- # (T2R, T2R_M, T2R_N) +- tTR_rAcc = cute.make_rmem_tensor( +- tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype +- ) +- return tiled_copy_t2r, tTR_tAcc, tTR_rAcc +- +- def epilog_smem_copy_and_partition( +- self, +- tiled_copy_t2r: cute.TiledCopy, +- tTR_rC: cute.Tensor, +- tidx: cutlass.Int32, +- sC: cute.Tensor, +- ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: +- """ +- Make tiledCopy for shared memory store, then use it to partition register array (source) and shared memory (destination). +- +- Args: +- tiled_copy_t2r (cute.TiledCopy): The tiled copy operation for tmem to register copy(t2r) +- tTR_rC (cute.Tensor): The partitioned accumulator tensor +- tidx (cutlass.Int32): The thread index in epilogue warp groups +- sC (cute.Tensor): The shared memory tensor to be copied and partitioned +- sepi (cute.Tensor): +- +- Returns: +- A tuple containing (tiled_copy_r2s, tRS_rC, tRS_sC) where: +- - tiled_copy_r2s: The tiled copy operation for register to smem copy(r2s) +- - tRS_rC: The partitioned tensor C (register source) +- - tRS_sC: The partitioned tensor C (smem destination) +- """ +- copy_atom_r2s = sm100_utils.get_smem_store_op( +- self.c_layout, self.c_dtype, self.acc_dtype, tiled_copy_t2r +- ) +- tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r) +- # (R2S, R2S_M, R2S_N, PIPE_D) +- thr_copy_r2s = tiled_copy_r2s.get_slice(tidx) +- tRS_sC = thr_copy_r2s.partition_D(sC) +- # (R2S, R2S_M, R2S_N) +- tRS_rC = tiled_copy_r2s.retile(tTR_rC) +- return tiled_copy_r2s, tRS_rC, tRS_sC +- +- def epilog_gmem_copy_and_partition( +- self, +- tidx: cutlass.Int32, +- atom: Union[cute.CopyAtom, cute.TiledCopy], +- gC_mnl: cute.Tensor, +- epi_tile: cute.Tile, +- sC: cute.Tensor, +- ) -> Tuple[cute.CopyAtom, cute.Tensor, cute.Tensor]: +- """ +- Make tiledCopy for global memory store, then use it to: +- partition shared memory (source) and global memory (destination) for TMA store version. +- +- Args: +- tidx (cutlass.Int32): The thread index in epilogue warp groups +- atom (Union[cute.CopyAtom, cute.TiledCopy]): The copy_atom_c to be used for TMA store version, or tiled_copy_t2r for none TMA store version +- gC_mnl (cute.Tensor): The global tensor C +- epi_tile (cute.Tile): The epilogue tiler +- sC (cute.Tensor): The shared memory tensor to be copied and partitioned +- +- Returns: +- A tuple containing (tma_atom_c, bSG_sC, bSG_gC) where: +- - tma_atom_c: The TMA copy atom +- - bSG_sC: The partitioned shared memory tensor C +- - bSG_gC: The partitioned global tensor C +- """ +- # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL) +- gC_epi = cute.flat_divide( +- gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile +- ) +- +- tma_atom_c = atom +- sC_for_tma_partition = cute.group_modes(sC, 0, 2) +- gC_for_tma_partition = cute.group_modes(gC_epi, 0, 2) +- # ((ATOM_V, REST_V), EPI_M, EPI_N) +- # ((ATOM_V, REST_V), EPI_M, EPI_N, RestM, RestN, RestL) +- bSG_sC, bSG_gC = cpasync.tma_partition( +- tma_atom_c, +- 0, +- cute.make_layout(1), +- sC_for_tma_partition, +- gC_for_tma_partition, +- ) +- return tma_atom_c, bSG_sC, bSG_gC +- +- @staticmethod +- def _compute_stages( +- tiled_mma: cute.TiledMma, +- mma_tiler_mnk: Tuple[int, int, int], +- a_dtype: Type[cutlass.Numeric], +- a_major_mode: cute.nvgpu.OperandMajorMode, +- b_dtype: Type[cutlass.Numeric], +- b_major_mode: cute.nvgpu.OperandMajorMode, +- epi_tile: cute.Tile, +- c_dtype: Type[cutlass.Numeric], +- c_layout: utils.LayoutEnum, +- sf_dtype: Type[cutlass.Numeric], +- sf_vec_size: int, +- smem_capacity: int, +- occupancy: int, +- ) -> Tuple[int, int, int]: +- """Computes the optimal number of pipeline stages for operands. +- +- This method uses heuristics to determine the number of stages for the +- accumulator (ACC), A/B operands, and C operand based on the kernel +- configuration and available shared memory. +- +- Args: +- tiled_mma (cute.TiledMma): The tiled MMA object. +- mma_tiler_mnk (Tuple[int, int, int]): The shape (M, N, K) of the +- MMA tiler. +- a_dtype (Type[cutlass.Numeric]): Data type of operand A. +- a_major_mode (cute.nvgpu.OperandMajorMode): Layout of operand A. +- b_dtype (Type[cutlass.Numeric]): Data type of operand B. +- b_major_mode (cute.nvgpu.OperandMajorMode): Layout of operand B. +- epi_tile (cute.Tile): The epilogue tile shape. +- c_dtype (Type[cutlass.Numeric]): Data type of operand C. +- c_layout (utils.LayoutEnum): Layout of operand C. +- sf_dtype (Type[cutlass.Numeric]): Data type of the scale factors. +- sf_vec_size (int): Vector size of the scale factors. +- smem_capacity (int): Total available shared memory in bytes. +- occupancy (int): Target number of CTAs per SM. +- +- Returns: +- Tuple[int, int, int]: A tuple containing the number of stages for +- (ACC, A/B, C). +- """ +- # ACC stages +- num_acc_stage = 1 if mma_tiler_mnk[1] == 256 else 2 +- +- # Default C stages +- num_c_stage = 2 +- +- # Calculate smem layout and size for one stage of A, B, SFA, SFB and C +- a_smem_layout_stage_one = sm100_utils.make_smem_layout_a( +- tiled_mma, +- mma_tiler_mnk, +- a_dtype, +- 1, # a tmp 1 stage is provided +- ) +- b_smem_layout_staged_one = sm100_utils.make_smem_layout_b( +- tiled_mma, +- mma_tiler_mnk, +- b_dtype, +- 1, # a tmp 1 stage is provided +- ) +- sfa_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfa( +- tiled_mma, +- mma_tiler_mnk, +- sf_vec_size, +- 1, # a tmp 1 stage is provided +- ) +- sfb_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfb( +- tiled_mma, +- mma_tiler_mnk, +- sf_vec_size, +- 1, # a tmp 1 stage is provided +- ) +- +- c_smem_layout_staged_one = sm100_utils.make_smem_layout_epi( +- c_dtype, +- c_layout, +- epi_tile, +- 1, +- ) +- +- ab_bytes_per_stage = ( +- cute.size_in_bytes(a_dtype, a_smem_layout_stage_one) +- + cute.size_in_bytes(b_dtype, b_smem_layout_staged_one) +- + cute.size_in_bytes(sf_dtype, sfa_smem_layout_staged_one) +- + cute.size_in_bytes(sf_dtype, sfb_smem_layout_staged_one) +- ) +- mbar_helpers_bytes = 1024 +- c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout_staged_one) +- c_bytes = c_bytes_per_stage * num_c_stage +- +- # Calculate A/B/SFA/SFB stages: +- # Start with total smem per CTA (capacity / occupancy) +- # Subtract reserved bytes and initial C stages bytes +- # Divide remaining by bytes needed per A/B/SFA/SFB stage +- num_ab_stage = ( +- smem_capacity // occupancy - (mbar_helpers_bytes + c_bytes) +- ) // ab_bytes_per_stage +- +- # Refine epilogue stages: +- # Calculate remaining smem after allocating for A/B/SFA/SFB stages and reserved bytes +- # Add remaining unused smem to epilogue +- num_c_stage += ( +- smem_capacity +- - occupancy * ab_bytes_per_stage * num_ab_stage +- - occupancy * (mbar_helpers_bytes + c_bytes) +- ) // (occupancy * c_bytes_per_stage) +- +- return num_acc_stage, num_ab_stage, num_c_stage +- + @staticmethod + def _compute_grid( + c: cute.Tensor, +@@ -2089,125 +1673,3 @@ class Sm100BlockScaledPersistentDenseGemmKernel: + if mma_tiler_mn[1] < 64 and (n > mma_tiler_mn[1] or cluster_shape_mn[1] > 1): + can_implement = False + return can_implement +- +- # fully dynamic shape +- @cute.jit +- def wrapper( +- self, +- mA: cute.Tensor, +- mB: cute.Tensor, +- mC: cute.Tensor, +- sf_m: cutlass.Int64, +- sf_n: cutlass.Int64, +- sf_k: cutlass.Int64, +- l: cutlass.Constexpr, +- a_sf_ptr: cute.Pointer, +- b_sf_ptr: cute.Pointer, +- alpha_tensor: cute.Tensor, +- max_active_clusters: cutlass.Constexpr, +- current_stream, +- swap_ab: cutlass.Constexpr = False, +- epilogue_op: cutlass.Constexpr = lambda x: x, +- ): +- """Executes the wrapped GEMM kernel with dynamically shaped tensors. +- +- Uses TVM-FFI for efficient tensor passing: A, B, C, and alpha are passed +- as cute.Tensor directly (torch tensors at runtime via TVM-FFI's C-level +- dlpack, with negligible conversion cost). Scale factor tensors remain as +- pointers because their 6D BlockScaledBasicChunk layout cannot be expressed +- as a torch tensor. +- +- Args: +- mA (cute.Tensor): Input tensor A. +- FP4 path: shape (m, k_packed), dtype Uint8 (2xFP4 packed per byte). +- MXFP8 path: shape (m, k), dtype Float8. +- mB (cute.Tensor): Input tensor B. +- FP4 path: shape (n, k_packed), dtype Uint8. +- MXFP8 path: shape (n, k), dtype Float8. +- mC (cute.Tensor): Output tensor C, shape (m, n). +- sf_m (cutlass.Int64): Scale factor M dim (ceil(m/128)). +- sf_n (cutlass.Int64): Scale factor N dim (ceil(n/128)). +- sf_k (cutlass.Int64): Scale factor K dim (ceil(k/sf_vec_size/4)). +- l (cutlass.Constexpr): Batch dimension (L). +- a_sf_ptr (cute.Pointer): Pointer to scale factor tensor for A. +- b_sf_ptr (cute.Pointer): Pointer to scale factor tensor for B. +- alpha_tensor (cute.Tensor): Alpha scaling factor, shape (1,), float32. +- max_active_clusters (cutlass.Constexpr): Max active clusters. +- current_stream: CUDA stream (managed by TVM-FFI fake stream). +- swap_ab (cutlass.Constexpr): Whether A/B are swapped (controls C layout). +- epilogue_op (cutlass.Constexpr): Elementwise epilogue function. +- """ +- # A, B, C are passed as cute.Tensor via TVM-FFI. +- # Support two input encodings: +- # 1) FP4 packed as uint8 (recast to Float4E2M1FN, k = k_packed * 2) +- # 2) MXFP8 as float8 (k = k_raw, no recast) +- m = cute.size(mA, mode=[0]) +- k_raw = cute.size(mA, mode=[1]) +- n = cute.size(mB, mode=[0]) +- +- if cutlass.const_expr( +- mA.element_type == cutlass.Uint8 and mB.element_type == cutlass.Uint8 +- ): +- # FP4 packed path: 2 FP4 values per uint8 byte +- k = k_raw * 2 +- a_ptr = cute.recast_ptr(mA.iterator, dtype=cutlass.Float4E2M1FN) +- b_ptr = cute.recast_ptr(mB.iterator, dtype=cutlass.Float4E2M1FN) +- elif cutlass.const_expr(mA.element_type != mB.element_type): +- raise TypeError( +- "Unsupported mixed input dtypes for block-scaled GEMM: " +- "mA and mB must have matching element_type " +- "(both Uint8 for FP4 path, or both FP8 for MXFP8 path)." +- ) +- else: +- # MXFP8 path: input tensors are already FP8. +- k = k_raw +- a_ptr = mA.iterator +- b_ptr = mB.iterator +- +- a_tensor = cute.make_tensor( +- a_ptr, +- layout=cute.make_ordered_layout((m, k, l), order=(1, 0, 2)), +- ) +- b_tensor = cute.make_tensor( +- b_ptr, +- layout=cute.make_ordered_layout((n, k, l), order=(1, 0, 2)), +- ) +- # Reshape C to (m, n, l) -- swap_ab is constexpr, determines layout at compile time +- if cutlass.const_expr(swap_ab): +- c_tensor = cute.make_tensor( +- mC.iterator, +- layout=cute.make_ordered_layout((m, n, l), order=(0, 1, 2)), +- ) +- else: +- c_tensor = cute.make_tensor( +- mC.iterator, +- layout=cute.make_ordered_layout((m, n, l), order=(1, 0, 2)), +- ) +- # Scale factor tensors: 6D BlockScaledBasicChunk layout from pointers +- # (32, 4, sf_m, 4, sf_k, l) with order (2, 1, 4, 0, 3, 5) +- sfa_tensor = cute.make_tensor( +- a_sf_ptr, +- layout=cute.make_ordered_layout( +- (32, 4, sf_m, 4, sf_k, l), +- order=(2, 1, 4, 0, 3, 5), +- ), +- ) +- sfb_tensor = cute.make_tensor( +- b_sf_ptr, +- layout=cute.make_ordered_layout( +- (32, 4, sf_n, 4, sf_k, l), +- order=(2, 1, 4, 0, 3, 5), +- ), +- ) +- +- self( +- a_tensor, +- b_tensor, +- sfa_tensor, +- sfb_tensor, +- c_tensor, +- alpha_tensor, +- max_active_clusters, +- current_stream, +- epilogue_op, +- ) +diff --git a/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm100_common.py b/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm100_common.py +new file mode 100644 +index 00000000..0d1642b1 +--- /dev/null ++++ b/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm100_common.py +@@ -0,0 +1,669 @@ ++# Copyright (c) 2025 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. ++ ++ ++from typing import Any, Tuple, Type, Union ++ ++import cutlass ++import cutlass.cute as cute ++import cutlass.utils as utils ++import cutlass.utils.blackwell_helpers as sm100_utils ++import cutlass.utils.blockscaled_layout as blockscaled_utils ++from cutlass.cute.nvgpu import cpasync, tcgen05 ++ ++ ++class _Sm100BlockScaledGemmCommon: ++ """Shared launch and copy-partition helpers for SM100 block-scaled GEMM.""" ++ ++ acc_dtype: Any ++ c_dtype: Any ++ c_layout: Any ++ cta_group: Any ++ cta_tile_shape_mnk: Any ++ sf_dtype: Any ++ # Concrete kernels initialize these in __init__ or _setup_attributes. ++ # CuTe layout/configuration objects are generated dynamically, so Any is ++ # intentional at this shared interface boundary. ++ a_smem_layout_staged: Any ++ b_smem_layout_staged: Any ++ c_smem_layout_staged: Any ++ sfa_smem_layout_staged: Any ++ sfb_smem_layout_staged: Any ++ cluster_layout_sfb_vmnk: Any ++ cluster_layout_vmnk: Any ++ cluster_shape_mn: Any ++ epi_tile: Any ++ mma_inst_shape_mnk: Any ++ mma_inst_shape_mnk_sfb: Any ++ mma_tiler: Any ++ mma_tiler_sfb: Any ++ sf_vec_size: int ++ ++ def _setup_attributes(self) -> None: ++ raise NotImplementedError ++ ++ def __call__( ++ self, ++ a_tensor: cute.Tensor, ++ b_tensor: cute.Tensor, ++ sfa_tensor: cute.Tensor, ++ sfb_tensor: cute.Tensor, ++ c_tensor: cute.Tensor, ++ alpha: cute.Tensor, ++ max_active_clusters: cutlass.Constexpr, ++ stream: Any, ++ epilogue_op: cutlass.Constexpr = lambda x: x, ++ ) -> Any: ++ raise NotImplementedError ++ ++ def _prepare_tma( ++ self, ++ a_tensor: cute.Tensor, ++ b_tensor: cute.Tensor, ++ sfa_tensor: cute.Tensor, ++ sfb_tensor: cute.Tensor, ++ c_tensor: cute.Tensor, ++ ): ++ """Build the shared block-scaled MMA and TMA descriptors.""" ++ # Setup static attributes before smem/grid/tma computation ++ self.a_dtype: Type[cutlass.Numeric] = a_tensor.element_type ++ self.b_dtype: Type[cutlass.Numeric] = b_tensor.element_type ++ self.sf_dtype: Type[cutlass.Numeric] = sfa_tensor.element_type ++ self.c_dtype: Type[cutlass.Numeric] = c_tensor.element_type ++ self.a_major_mode = utils.LayoutEnum.from_tensor(a_tensor).mma_major_mode() ++ self.b_major_mode = utils.LayoutEnum.from_tensor(b_tensor).mma_major_mode() ++ self.c_layout = utils.LayoutEnum.from_tensor(c_tensor) ++ ++ # Check if input data types are compatible with MMA instruction ++ if cutlass.const_expr(self.a_dtype != self.b_dtype): ++ raise TypeError(f"Type must match: {self.a_dtype} != {self.b_dtype}") ++ ++ # Setup attributes that dependent on gemm inputs ++ self._setup_attributes() ++ ++ # Setup sfa/sfb tensor by filling A/B tensor to scale factor atom layout ++ # ((Atom_M, Rest_M),(Atom_K, Rest_K),RestL) ++ sfa_layout = blockscaled_utils.tile_atom_to_shape_SF( ++ a_tensor.shape, self.sf_vec_size ++ ) ++ sfa_tensor = cute.make_tensor(sfa_tensor.iterator, sfa_layout) ++ ++ # ((Atom_N, Rest_N),(Atom_K, Rest_K),RestL) ++ sfb_layout = blockscaled_utils.tile_atom_to_shape_SF( ++ b_tensor.shape, self.sf_vec_size ++ ) ++ sfb_tensor = cute.make_tensor(sfb_tensor.iterator, sfb_layout) ++ ++ tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma( ++ self.a_dtype, ++ self.b_dtype, ++ self.a_major_mode, ++ self.b_major_mode, ++ self.sf_dtype, ++ self.sf_vec_size, ++ self.cta_group, ++ self.mma_inst_shape_mnk[:2], ++ ) ++ ++ tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma( ++ self.a_dtype, ++ self.b_dtype, ++ self.a_major_mode, ++ self.b_major_mode, ++ self.sf_dtype, ++ self.sf_vec_size, ++ cute.nvgpu.tcgen05.CtaGroup.ONE, ++ self.mma_inst_shape_mnk_sfb[:2], ++ ) ++ atom_thr_size = cute.size(tiled_mma.thr_id.shape) ++ ++ # Setup TMA load for A ++ a_op = sm100_utils.cluster_shape_to_tma_atom_A( ++ self.cluster_shape_mn, tiled_mma.thr_id ++ ) ++ a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) ++ tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( ++ a_op, ++ a_tensor, ++ a_smem_layout, ++ self.mma_tiler, ++ tiled_mma, ++ self.cluster_layout_vmnk.shape, ++ ) ++ ++ # Setup TMA load for B ++ b_op = sm100_utils.cluster_shape_to_tma_atom_B( ++ self.cluster_shape_mn, tiled_mma.thr_id ++ ) ++ b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) ++ tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( ++ b_op, ++ b_tensor, ++ b_smem_layout, ++ self.mma_tiler, ++ tiled_mma, ++ self.cluster_layout_vmnk.shape, ++ ) ++ ++ # Setup TMA load for SFA ++ sfa_op = sm100_utils.cluster_shape_to_tma_atom_A( ++ self.cluster_shape_mn, tiled_mma.thr_id ++ ) ++ sfa_smem_layout = cute.slice_( ++ self.sfa_smem_layout_staged, (None, None, None, 0) ++ ) ++ tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A( ++ sfa_op, ++ sfa_tensor, ++ sfa_smem_layout, ++ self.mma_tiler, ++ tiled_mma, ++ self.cluster_layout_vmnk.shape, ++ internal_type=cutlass.Int16, ++ ) ++ ++ # Setup TMA load for SFB ++ sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB( ++ self.cluster_shape_mn, tiled_mma.thr_id ++ ) ++ sfb_smem_layout = cute.slice_( ++ self.sfb_smem_layout_staged, (None, None, None, 0) ++ ) ++ tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B( ++ sfb_op, ++ sfb_tensor, ++ sfb_smem_layout, ++ self.mma_tiler_sfb, ++ tiled_mma_sfb, ++ self.cluster_layout_sfb_vmnk.shape, ++ internal_type=cutlass.Int16, ++ ) ++ if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 192): ++ x = tma_tensor_sfb.stride[0][1] ++ y = cute.ceil_div(tma_tensor_sfb.shape[0][1], 4) ++ ++ new_shape = ( ++ (tma_tensor_sfb.shape[0][0], ((2, 2), y)), ++ tma_tensor_sfb.shape[1], ++ tma_tensor_sfb.shape[2], ++ ) ++ # Use right multiplication for ScaledBasis (3 * x instead of x * 3) ++ x_times_3 = 3 * x ++ new_stride = ( ++ (tma_tensor_sfb.stride[0][0], ((x, x), x_times_3)), ++ tma_tensor_sfb.stride[1], ++ tma_tensor_sfb.stride[2], ++ ) ++ tma_tensor_sfb_new_layout = cute.make_layout(new_shape, stride=new_stride) ++ tma_tensor_sfb = cute.make_tensor( ++ tma_tensor_sfb.iterator, tma_tensor_sfb_new_layout ++ ) ++ ++ a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout) ++ b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout) ++ sfa_copy_size = cute.size_in_bytes(self.sf_dtype, sfa_smem_layout) ++ sfb_copy_size = cute.size_in_bytes(self.sf_dtype, sfb_smem_layout) ++ self.num_tma_load_bytes = ( ++ a_copy_size + b_copy_size + sfa_copy_size + sfb_copy_size ++ ) * atom_thr_size ++ ++ # Setup TMA store for C ++ epi_smem_layout = cute.slice_(self.c_smem_layout_staged, (None, None, 0)) ++ tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( ++ cpasync.CopyBulkTensorTileS2GOp(), ++ c_tensor, ++ epi_smem_layout, ++ self.epi_tile, ++ ) ++ ++ return ( ++ tiled_mma, ++ tiled_mma_sfb, ++ tma_atom_a, ++ tma_tensor_a, ++ tma_atom_b, ++ tma_tensor_b, ++ tma_atom_sfa, ++ tma_tensor_sfa, ++ tma_atom_sfb, ++ tma_tensor_sfb, ++ tma_atom_c, ++ tma_tensor_c, ++ ) ++ ++ @staticmethod ++ def _compute_stages( ++ tiled_mma: cute.TiledMma, ++ mma_tiler_mnk: Tuple[int, int, int], ++ a_dtype: Type[cutlass.Numeric], ++ a_major_mode: cute.nvgpu.OperandMajorMode, ++ b_dtype: Type[cutlass.Numeric], ++ b_major_mode: cute.nvgpu.OperandMajorMode, ++ epi_tile: cute.Tile, ++ c_dtype: Type[cutlass.Numeric], ++ c_layout: utils.LayoutEnum, ++ sf_dtype: Type[cutlass.Numeric], ++ sf_vec_size: int, ++ smem_capacity: int, ++ occupancy: int, ++ reserved_smem_bytes: int, ++ ) -> Tuple[int, int, int]: ++ """Computes the optimal number of pipeline stages for operands. ++ ++ This method uses heuristics to determine the number of stages for the ++ accumulator (ACC), A/B operands, and C operand based on the kernel ++ configuration and available shared memory. ++ ++ Args: ++ tiled_mma (cute.TiledMma): The tiled MMA object. ++ mma_tiler_mnk (Tuple[int, int, int]): The shape (M, N, K) of the ++ MMA tiler. ++ a_dtype (Type[cutlass.Numeric]): Data type of operand A. ++ a_major_mode (cute.nvgpu.OperandMajorMode): Layout of operand A. ++ b_dtype (Type[cutlass.Numeric]): Data type of operand B. ++ b_major_mode (cute.nvgpu.OperandMajorMode): Layout of operand B. ++ epi_tile (cute.Tile): The epilogue tile shape. ++ c_dtype (Type[cutlass.Numeric]): Data type of operand C. ++ c_layout (utils.LayoutEnum): Layout of operand C. ++ sf_dtype (Type[cutlass.Numeric]): Data type of the scale factors. ++ sf_vec_size (int): Vector size of the scale factors. ++ smem_capacity (int): Total available shared memory in bytes. ++ occupancy (int): Target number of CTAs per SM. ++ reserved_smem_bytes (int): Per-CTA bytes reserved outside the staged operands and output. ++ ++ Returns: ++ Tuple[int, int, int]: A tuple containing the number of stages for ++ (ACC, A/B, C). ++ """ ++ # ACC stages ++ num_acc_stage = 1 if mma_tiler_mnk[1] == 256 else 2 ++ ++ # Default C stages ++ num_c_stage = 2 ++ ++ # Calculate smem layout and size for one stage of A, B, SFA, SFB and C ++ a_smem_layout_stage_one = sm100_utils.make_smem_layout_a( ++ tiled_mma, ++ mma_tiler_mnk, ++ a_dtype, ++ 1, # a tmp 1 stage is provided ++ ) ++ b_smem_layout_staged_one = sm100_utils.make_smem_layout_b( ++ tiled_mma, ++ mma_tiler_mnk, ++ b_dtype, ++ 1, # a tmp 1 stage is provided ++ ) ++ sfa_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfa( ++ tiled_mma, ++ mma_tiler_mnk, ++ sf_vec_size, ++ 1, # a tmp 1 stage is provided ++ ) ++ sfb_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfb( ++ tiled_mma, ++ mma_tiler_mnk, ++ sf_vec_size, ++ 1, # a tmp 1 stage is provided ++ ) ++ ++ c_smem_layout_staged_one = sm100_utils.make_smem_layout_epi( ++ c_dtype, ++ c_layout, ++ epi_tile, ++ 1, ++ ) ++ ++ ab_bytes_per_stage = ( ++ cute.size_in_bytes(a_dtype, a_smem_layout_stage_one) ++ + cute.size_in_bytes(b_dtype, b_smem_layout_staged_one) ++ + cute.size_in_bytes(sf_dtype, sfa_smem_layout_staged_one) ++ + cute.size_in_bytes(sf_dtype, sfb_smem_layout_staged_one) ++ ) ++ c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout_staged_one) ++ c_bytes = c_bytes_per_stage * num_c_stage ++ ++ # Calculate A/B/SFA/SFB stages: ++ # Start with total smem per CTA (capacity / occupancy) ++ # Subtract reserved bytes and initial C stages bytes ++ # Divide remaining by bytes needed per A/B/SFA/SFB stage ++ num_ab_stage = ( ++ smem_capacity // occupancy - (reserved_smem_bytes + c_bytes) ++ ) // ab_bytes_per_stage ++ ++ # Refine epilogue stages: ++ # Calculate remaining smem after allocating for A/B/SFA/SFB stages and reserved bytes ++ # Add remaining unused smem to epilogue ++ num_c_stage += ( ++ smem_capacity ++ - occupancy * ab_bytes_per_stage * num_ab_stage ++ - occupancy * (reserved_smem_bytes + c_bytes) ++ ) // (occupancy * c_bytes_per_stage) ++ ++ return num_acc_stage, num_ab_stage, num_c_stage ++ ++ # fully dynamic shape ++ @cute.jit ++ def wrapper( ++ self, ++ mA: cute.Tensor, ++ mB: cute.Tensor, ++ mC: cute.Tensor, ++ sf_m: cutlass.Int64, ++ sf_n: cutlass.Int64, ++ sf_k: cutlass.Int64, ++ l: cutlass.Constexpr, ++ a_sf_ptr: cute.Pointer, ++ b_sf_ptr: cute.Pointer, ++ alpha_tensor: cute.Tensor, ++ max_active_clusters: cutlass.Constexpr, ++ current_stream, ++ swap_ab: cutlass.Constexpr = False, ++ epilogue_op: cutlass.Constexpr = lambda x: x, ++ ): ++ """Executes the wrapped GEMM kernel with dynamically shaped tensors. ++ ++ Uses TVM-FFI for efficient tensor passing: A, B, C, and alpha are passed ++ as cute.Tensor directly (torch tensors at runtime via TVM-FFI's C-level ++ dlpack, with negligible conversion cost). Scale factor tensors remain as ++ pointers because their 6D BlockScaledBasicChunk layout cannot be expressed ++ as a torch tensor. ++ ++ Args: ++ mA (cute.Tensor): Input tensor A. ++ FP4 path: shape (m, k_packed), dtype Uint8 (2xFP4 packed per byte). ++ MXFP8 path: shape (m, k), dtype Float8. ++ mB (cute.Tensor): Input tensor B. ++ FP4 path: shape (n, k_packed), dtype Uint8. ++ MXFP8 path: shape (n, k), dtype Float8. ++ mC (cute.Tensor): Output tensor C, shape (m, n). ++ sf_m (cutlass.Int64): Scale factor M dim (ceil(m/128)). ++ sf_n (cutlass.Int64): Scale factor N dim (ceil(n/128)). ++ sf_k (cutlass.Int64): Scale factor K dim (ceil(k/sf_vec_size/4)). ++ l (cutlass.Constexpr): Batch dimension (L). ++ a_sf_ptr (cute.Pointer): Pointer to scale factor tensor for A. ++ b_sf_ptr (cute.Pointer): Pointer to scale factor tensor for B. ++ alpha_tensor (cute.Tensor): Alpha scaling factor, shape (1,), float32. ++ max_active_clusters (cutlass.Constexpr): Max active clusters. ++ current_stream: CUDA stream (managed by TVM-FFI fake stream). ++ swap_ab (cutlass.Constexpr): Whether A/B are swapped (controls C layout). ++ epilogue_op (cutlass.Constexpr): Elementwise epilogue function. ++ """ ++ # A, B, C are passed as cute.Tensor via TVM-FFI. ++ # Support two input encodings: ++ # 1) FP4 packed as uint8 (recast to Float4E2M1FN, k = k_packed * 2) ++ # 2) MXFP8 as float8 (k = k_raw, no recast) ++ m = cute.size(mA, mode=[0]) ++ k_raw = cute.size(mA, mode=[1]) ++ n = cute.size(mB, mode=[0]) ++ ++ if cutlass.const_expr( ++ mA.element_type == cutlass.Uint8 and mB.element_type == cutlass.Uint8 ++ ): ++ # FP4 packed path: 2 FP4 values per uint8 byte ++ k = k_raw * 2 ++ a_ptr = cute.recast_ptr(mA.iterator, dtype=cutlass.Float4E2M1FN) ++ b_ptr = cute.recast_ptr(mB.iterator, dtype=cutlass.Float4E2M1FN) ++ elif cutlass.const_expr(mA.element_type != mB.element_type): ++ raise TypeError( ++ "Unsupported mixed input dtypes for block-scaled GEMM: " ++ "mA and mB must have matching element_type " ++ "(both Uint8 for FP4 path, or both FP8 for MXFP8 path)." ++ ) ++ else: ++ # MXFP8 path: input tensors are already FP8. ++ k = k_raw ++ a_ptr = mA.iterator ++ b_ptr = mB.iterator ++ ++ a_tensor = cute.make_tensor( ++ a_ptr, ++ layout=cute.make_ordered_layout((m, k, l), order=(1, 0, 2)), ++ ) ++ b_tensor = cute.make_tensor( ++ b_ptr, ++ layout=cute.make_ordered_layout((n, k, l), order=(1, 0, 2)), ++ ) ++ # Reshape C to (m, n, l) -- swap_ab is constexpr, determines layout at compile time ++ if cutlass.const_expr(swap_ab): ++ c_tensor = cute.make_tensor( ++ mC.iterator, ++ layout=cute.make_ordered_layout((m, n, l), order=(0, 1, 2)), ++ ) ++ else: ++ c_tensor = cute.make_tensor( ++ mC.iterator, ++ layout=cute.make_ordered_layout((m, n, l), order=(1, 0, 2)), ++ ) ++ # Scale factor tensors: 6D BlockScaledBasicChunk layout from pointers ++ # (32, 4, sf_m, 4, sf_k, l) with order (2, 1, 4, 0, 3, 5) ++ sfa_tensor = cute.make_tensor( ++ a_sf_ptr, ++ layout=cute.make_ordered_layout( ++ (32, 4, sf_m, 4, sf_k, l), ++ order=(2, 1, 4, 0, 3, 5), ++ ), ++ ) ++ sfb_tensor = cute.make_tensor( ++ b_sf_ptr, ++ layout=cute.make_ordered_layout( ++ (32, 4, sf_n, 4, sf_k, l), ++ order=(2, 1, 4, 0, 3, 5), ++ ), ++ ) ++ ++ self( ++ a_tensor, ++ b_tensor, ++ sfa_tensor, ++ sfb_tensor, ++ c_tensor, ++ alpha_tensor, ++ max_active_clusters, ++ current_stream, ++ epilogue_op, ++ ) ++ ++ def mainloop_s2t_copy_and_partition( ++ self, ++ sSF: cute.Tensor, ++ tSF: cute.Tensor, ++ ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: ++ """ ++ Make tiledCopy for smem to tmem load for scale factor tensor, then use it to partition smem memory (source) and tensor memory (destination). ++ ++ Args: ++ sSF (cute.Tensor): The scale factor tensor in smem ++ tSF (cute.Tensor): The scale factor tensor in tmem ++ ++ Returns: ++ A tuple containing (tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t) where: ++ - tiled_copy_s2t: The tiled copy operation for smem to tmem load for scale factor tensor(s2t) ++ - tCsSF_compact_s2t: The partitioned scale factor tensor in smem ++ - tSF_compact_s2t: The partitioned scale factor tensor in tmem ++ """ ++ # (MMA, MMA_MN, MMA_K, STAGE) ++ tCsSF_compact = cute.filter_zeros(sSF) ++ # (MMA, MMA_MN, MMA_K) ++ tCtSF_compact = cute.filter_zeros(tSF) ++ ++ # Make S2T CopyAtom and tiledCopy ++ copy_atom_s2t = cute.make_copy_atom( ++ tcgen05.Cp4x32x128bOp(self.cta_group), ++ self.sf_dtype, ++ ) ++ tiled_copy_s2t = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSF_compact) ++ thr_copy_s2t = tiled_copy_s2t.get_slice(0) ++ ++ # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE) ++ tCsSF_compact_s2t_ = thr_copy_s2t.partition_S(tCsSF_compact) ++ # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE) ++ tCsSF_compact_s2t = tcgen05.get_s2t_smem_desc_tensor( ++ tiled_copy_s2t, tCsSF_compact_s2t_ ++ ) ++ # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K) ++ tCtSF_compact_s2t = thr_copy_s2t.partition_D(tCtSF_compact) ++ ++ return tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t ++ ++ def epilog_tmem_copy_and_partition( ++ self, ++ tidx: cutlass.Int32, ++ tAcc: cute.Tensor, ++ gC_mnl: cute.Tensor, ++ epi_tile: cute.Tile, ++ use_2cta_instrs: Union[cutlass.Boolean, bool], ++ ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: ++ """ ++ Make tiledCopy for tensor memory load, then use it to partition tensor memory (source) and register array (destination). ++ ++ Args: ++ tidx (cutlass.Int32): The thread index in epilogue warp groups ++ tAcc (cute.Tensor): The accumulator tensor to be copied and partitioned ++ gC_mnl (cute.Tensor): The global tensor C ++ epi_tile (cute.Tile): The epilogue tiler ++ use_2cta_instrs (bool): Whether use_2cta_instrs is enabled ++ ++ Returns: ++ A tuple containing (tiled_copy_t2r, tTR_tAcc, tTR_rAcc) where: ++ - tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r) ++ - tTR_tAcc: The partitioned accumulator tensor ++ - tTR_rAcc: The accumulated tensor in register used to hold t2r results ++ """ ++ # Make tiledCopy for tensor memory load ++ copy_atom_t2r = sm100_utils.get_tmem_load_op( ++ self.cta_tile_shape_mnk, ++ self.c_layout, ++ self.c_dtype, ++ self.acc_dtype, ++ epi_tile, ++ use_2cta_instrs, ++ ) ++ # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, STAGE) ++ tAcc_epi = cute.flat_divide( ++ tAcc[((None, None), 0, 0, None)], ++ epi_tile, ++ ) ++ # (EPI_TILE_M, EPI_TILE_N) ++ tiled_copy_t2r = tcgen05.make_tmem_copy( ++ copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)] ++ ) ++ ++ thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) ++ # (T2R, T2R_M, T2R_N, EPI_M, EPI_M, STAGE) ++ tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi) ++ ++ # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL) ++ gC_mnl_epi = cute.flat_divide( ++ gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile ++ ) ++ # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL) ++ tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi) ++ # (T2R, T2R_M, T2R_N) ++ tTR_rAcc = cute.make_rmem_tensor( ++ tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype ++ ) ++ return tiled_copy_t2r, tTR_tAcc, tTR_rAcc ++ ++ def epilog_smem_copy_and_partition( ++ self, ++ tiled_copy_t2r: cute.TiledCopy, ++ tTR_rC: cute.Tensor, ++ tidx: cutlass.Int32, ++ sC: cute.Tensor, ++ ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: ++ """ ++ Make tiledCopy for shared memory store, then use it to partition register array (source) and shared memory (destination). ++ ++ Args: ++ tiled_copy_t2r (cute.TiledCopy): The tiled copy operation for tmem to register copy(t2r) ++ tTR_rC (cute.Tensor): The partitioned accumulator tensor ++ tidx (cutlass.Int32): The thread index in epilogue warp groups ++ sC (cute.Tensor): The shared memory tensor to be copied and partitioned ++ sepi (cute.Tensor): ++ ++ Returns: ++ A tuple containing (tiled_copy_r2s, tRS_rC, tRS_sC) where: ++ - tiled_copy_r2s: The tiled copy operation for register to smem copy(r2s) ++ - tRS_rC: The partitioned tensor C (register source) ++ - tRS_sC: The partitioned tensor C (smem destination) ++ """ ++ copy_atom_r2s = sm100_utils.get_smem_store_op( ++ self.c_layout, self.c_dtype, self.acc_dtype, tiled_copy_t2r ++ ) ++ tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r) ++ # (R2S, R2S_M, R2S_N, PIPE_D) ++ thr_copy_r2s = tiled_copy_r2s.get_slice(tidx) ++ tRS_sC = thr_copy_r2s.partition_D(sC) ++ # (R2S, R2S_M, R2S_N) ++ tRS_rC = tiled_copy_r2s.retile(tTR_rC) ++ return tiled_copy_r2s, tRS_rC, tRS_sC ++ ++ def epilog_gmem_copy_and_partition( ++ self, ++ tidx: cutlass.Int32, ++ atom: Union[cute.CopyAtom, cute.TiledCopy], ++ gC_mnl: cute.Tensor, ++ epi_tile: cute.Tile, ++ sC: cute.Tensor, ++ ) -> Tuple[cute.CopyAtom, cute.Tensor, cute.Tensor]: ++ """ ++ Make tiledCopy for global memory store, then use it to: ++ partition shared memory (source) and global memory (destination) for TMA store version. ++ ++ Args: ++ tidx (cutlass.Int32): The thread index in epilogue warp groups ++ atom (Union[cute.CopyAtom, cute.TiledCopy]): The copy_atom_c to be used for TMA store version, or tiled_copy_t2r for none TMA store version ++ gC_mnl (cute.Tensor): The global tensor C ++ epi_tile (cute.Tile): The epilogue tiler ++ sC (cute.Tensor): The shared memory tensor to be copied and partitioned ++ ++ Returns: ++ A tuple containing (tma_atom_c, bSG_sC, bSG_gC) where: ++ - tma_atom_c: The TMA copy atom ++ - bSG_sC: The partitioned shared memory tensor C ++ - bSG_gC: The partitioned global tensor C ++ """ ++ # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL) ++ gC_epi = cute.flat_divide( ++ gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile ++ ) ++ ++ tma_atom_c = atom ++ sC_for_tma_partition = cute.group_modes(sC, 0, 2) ++ gC_for_tma_partition = cute.group_modes(gC_epi, 0, 2) ++ # ((ATOM_V, REST_V), EPI_M, EPI_N) ++ # ((ATOM_V, REST_V), EPI_M, EPI_N, RestM, RestN, RestL) ++ bSG_sC, bSG_gC = cpasync.tma_partition( ++ tma_atom_c, ++ 0, ++ cute.make_layout(1), ++ sC_for_tma_partition, ++ gC_for_tma_partition, ++ ) ++ return tma_atom_c, bSG_sC, bSG_gC +diff --git a/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm100_splitk.py b/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm100_splitk.py +new file mode 100644 +index 00000000..f78dd8b3 +--- /dev/null ++++ b/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm100_splitk.py +@@ -0,0 +1,1539 @@ ++# Copyright (c) 2025 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. ++ ++# This kernel keeps FlashInfer's block-scaled MXFP8 mainloop and adapts the ++# cluster-local reduction protocol from NVIDIA DKG's FP8 split-K tutorial. ++ ++from typing import Optional, Tuple, Union ++ ++import cuda.bindings.driver as cuda ++import cutlass ++import cutlass.cute as cute ++import cutlass.pipeline as pipeline ++import cutlass.utils as utils ++import cutlass.utils.blackwell_helpers as sm100_utils ++import cutlass.utils.blockscaled_layout as blockscaled_utils ++from cutlass.cute.nvgpu import cpasync, tcgen05 ++ ++from cutlass.cute.arch import griddepcontrol_launch_dependents, griddepcontrol_wait ++from cutlass.pipeline import PipelineTmaUmma, PipelineUmmaAsync ++from cutlass import Float32, Int32 ++from cutlass._mlir.dialects import llvm ++from .dense_blockscaled_gemm_sm100_common import _Sm100BlockScaledGemmCommon ++ ++from cutlass.cutlass_dsl import ( ++ T, ++ dsl_user_op, ++ extract_mlir_values, ++ new_from_mlir_values, ++) ++ ++ ++@dsl_user_op ++def _map_shared_rank( ++ smem_ptr: cute.Pointer, ++ peer_cta_rank_in_cluster: Int32, ++ *, ++ loc=None, ++ ip=None, ++) -> Int32: ++ """Map a local shared-memory pointer to a peer CTA in the cluster.""" ++ smem_ptr_i32 = smem_ptr.toint(loc=loc, ip=ip).ir_value() ++ return Int32( ++ llvm.inline_asm( ++ T.i32(), ++ [smem_ptr_i32, peer_cta_rank_in_cluster.ir_value()], ++ "mapa.shared::cluster.u32 $0, $1, $2;", ++ "=r,r,r", ++ has_side_effects=False, ++ is_align_stack=False, ++ asm_dialect=llvm.AsmDialect.AD_ATT, ++ ) ++ ) ++ ++ ++@dsl_user_op ++def _store_shared_remote_v4( ++ value0: Float32, ++ value1: Float32, ++ value2: Float32, ++ value3: Float32, ++ smem_ptr: cute.Pointer, ++ mbar_ptr: cute.Pointer, ++ peer_cta_rank_in_cluster: Int32, ++ *, ++ loc=None, ++ ip=None, ++) -> None: ++ """Store four FP32 values to peer SMEM and complete 16 mbarrier bytes.""" ++ remote_smem_ptr = _map_shared_rank( ++ smem_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip ++ ).ir_value() ++ remote_mbar_ptr = _map_shared_rank( ++ mbar_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip ++ ).ir_value() ++ values = [ ++ value0.bitcast(Int32).ir_value(loc=loc, ip=ip), ++ value1.bitcast(Int32).ir_value(loc=loc, ip=ip), ++ value2.bitcast(Int32).ir_value(loc=loc, ip=ip), ++ value3.bitcast(Int32).ir_value(loc=loc, ip=ip), ++ ] ++ llvm.inline_asm( ++ None, ++ [remote_smem_ptr, *values, remote_mbar_ptr], ++ "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 " ++ "[$0], {$1, $2, $3, $4}, [$5];", ++ "r,r,r,r,r,r", ++ has_side_effects=True, ++ is_align_stack=False, ++ asm_dialect=llvm.AsmDialect.AD_ATT, ++ ) ++ ++ ++class _SplitKOneTileScheduler: ++ """One-work-item scheduler for a physical K-cluster. ++ ++ Grid X/Y identify the output tile. Every CTA in grid Z belongs to the same ++ physical cluster and uses its cluster rank as the logical K-slice index. ++ """ ++ ++ def __init__(self, tile_idx, num_tiles_executed): ++ self._tile_idx = tile_idx ++ self._num_tiles_executed = num_tiles_executed ++ self._tile_idx_num_values = None ++ ++ @staticmethod ++ @cute.jit ++ def create(block_idx, cta_rank_in_cluster): ++ tile_idx = ( ++ cutlass.Int32(block_idx[0]), ++ cutlass.Int32(block_idx[1]), ++ cutlass.Int32(cta_rank_in_cluster), ++ ) ++ return _SplitKOneTileScheduler(tile_idx, cutlass.Int32(0)) ++ ++ def __extract_mlir_values__(self): ++ tile_idx_values = extract_mlir_values(self._tile_idx) ++ self._tile_idx_num_values = len(tile_idx_values) ++ return tile_idx_values + extract_mlir_values(self._num_tiles_executed) ++ ++ def __new_from_mlir_values__(self, values): ++ if self._tile_idx_num_values is None: ++ raise ValueError("Split-K scheduler tile width was not recorded") ++ width = self._tile_idx_num_values ++ tile_idx = new_from_mlir_values(self._tile_idx, values[:width]) ++ num_tiles_executed = new_from_mlir_values( ++ self._num_tiles_executed, values[width:] ++ ) ++ return _SplitKOneTileScheduler(tile_idx, num_tiles_executed) ++ ++ @cute.jit ++ def get_current_work(self): ++ return utils.WorkTileInfo( ++ self._tile_idx, ++ self._num_tiles_executed < cutlass.Int32(1), ++ ) ++ ++ @cute.jit ++ def initial_work_tile_info(self): ++ return self.get_current_work() ++ ++ @cute.jit ++ def advance_to_next_work(self): ++ self._num_tiles_executed += cutlass.Int32(1) ++ ++ @property ++ @cute.jit ++ def num_tiles_executed(self): ++ return self._num_tiles_executed ++ ++ ++class Sm100BlockScaledSplitKGemmKernel(_Sm100BlockScaledGemmCommon): ++ """Low-M MXFP8 GEMM with an in-kernel split-K reduction. ++ ++ The kernel is specialized for the swap_ab=True dense MXFP8 path with ++ R128C4 scale-factor layout. A physical (1, 1, split_k_slices) cluster ++ computes one output tile. Each CTA accumulates a disjoint K slice in FP32; ++ peer CTAs send their partials through distributed shared memory, and the ++ owner CTA writes the final BF16/FP16 output without a global workspace or ++ second reduction kernel. ++ ++ The supported tactics use an MMA-M tile of 128, MMA-N tiles of 8/16/32, ++ and two or four K slices. ++ """ ++ ++ def __init__( ++ self, ++ sf_vec_size: int, ++ mma_tiler_mn: Tuple[int, int], ++ split_k_slices: int, ++ enable_pdl: bool = True, ++ ): ++ """Configure one low-M split-K tactic. ++ ++ Args: ++ sf_vec_size: MXFP8 scale-factor vector size; must be 32. ++ mma_tiler_mn: MMA tile shape; supported values are (128, 8), ++ (128, 16), and (128, 32). ++ split_k_slices: Physical cluster-K size; must be 2 or 4. ++ enable_pdl: Whether to enable programmatic dependent launch. ++ """ ++ ++ self.acc_dtype = cutlass.Float32 ++ if sf_vec_size != 32: ++ raise ValueError( ++ f"MXFP8 split-K requires sf_vec_size=32, got {sf_vec_size}" ++ ) ++ self.sf_vec_size = sf_vec_size ++ if mma_tiler_mn not in ((128, 8), (128, 16), (128, 32)): ++ raise ValueError( ++ "MXFP8 split-K requires mma_tiler_mn=(128, 8/16/32), " ++ f"got {mma_tiler_mn}" ++ ) ++ if split_k_slices not in (2, 4): ++ raise ValueError(f"split_k_slices must be 2 or 4, got {split_k_slices}") ++ self.cluster_shape_mn = (1, 1) ++ self.cluster_shape_mnk = (1, 1, split_k_slices) ++ # K dimension is deferred in _setup_attributes ++ self.mma_tiler = (*mma_tiler_mn, 1) ++ self.enable_pdl = enable_pdl ++ self.split_k_slices = split_k_slices ++ self.cta_group = tcgen05.CtaGroup.ONE ++ ++ self.occupancy = 1 ++ # Set specialized warp ids ++ self.epilog_warp_id = ( ++ 0, ++ 1, ++ 2, ++ 3, ++ ) ++ self.mma_warp_id = 4 ++ self.tma_warp_id = 5 ++ self.threads_per_cta = 32 * len( ++ (self.mma_warp_id, self.tma_warp_id, *self.epilog_warp_id) ++ ) ++ # Set barrier id for cta sync, epilogue sync and tmem ptr sync ++ self.cta_sync_bar_id = 0 ++ self.epilog_sync_bar_id = 1 ++ self.tmem_ptr_sync_bar_id = 2 ++ self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") ++ SM100_TMEM_CAPACITY_COLUMNS = 512 ++ self.num_tmem_alloc_cols = SM100_TMEM_CAPACITY_COLUMNS ++ ++ def _setup_attributes(self): ++ """Set up configurations that are dependent on GEMM inputs ++ ++ This method configures various attributes based on the input tensor properties ++ (data types, leading dimensions) and kernel settings: ++ - Configuring tiled MMA ++ - Computing MMA/cluster/tile shapes ++ - Computing cluster layout ++ - Computing multicast CTAs for A/B/SFA/SFB ++ - Computing epilogue subtile ++ - Setting up A/B/SFA/SFB/C stage counts in shared memory ++ - Computing A/B/SFA/SFB/C shared memory layout ++ - Computing tensor memory allocation columns ++ """ ++ # Compute mma instruction shapes ++ mma_inst_bits_k = 256 ++ # (MMA_Tile_Shape_M, MMA_Tile_Shape_N, MMA_Inst_Shape_K) ++ self.mma_inst_shape_mnk = ( ++ self.mma_tiler[0], ++ self.mma_tiler[1], ++ mma_inst_bits_k // self.a_dtype.width, ++ ) ++ # (CTA_Tile_Shape_M, Round_Up(MMA_Tile_Shape_N, 128), MMA_Inst_Shape_K) ++ self.mma_inst_shape_mnk_sfb = ( ++ self.mma_inst_shape_mnk[0], ++ cute.round_up(self.mma_inst_shape_mnk[1], 128), ++ self.mma_inst_shape_mnk[2], ++ ) ++ ++ tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma( ++ self.a_dtype, ++ self.b_dtype, ++ self.a_major_mode, ++ self.b_major_mode, ++ self.sf_dtype, ++ self.sf_vec_size, ++ self.cta_group, ++ self.mma_inst_shape_mnk[:2], ++ ) ++ ++ tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma( ++ self.a_dtype, ++ self.b_dtype, ++ self.a_major_mode, ++ self.b_major_mode, ++ self.sf_dtype, ++ self.sf_vec_size, ++ cute.nvgpu.tcgen05.CtaGroup.ONE, ++ self.mma_inst_shape_mnk_sfb[:2], ++ ) ++ ++ # Compute mma/cluster/tile shapes ++ mma_inst_tile_k = 4 ++ self.mma_tiler = ( ++ self.mma_inst_shape_mnk[0], ++ self.mma_inst_shape_mnk[1], ++ self.mma_inst_shape_mnk[2] * mma_inst_tile_k, ++ ) ++ self.mma_tiler_sfb = ( ++ self.mma_inst_shape_mnk_sfb[0], ++ self.mma_inst_shape_mnk_sfb[1], ++ self.mma_inst_shape_mnk_sfb[2] * mma_inst_tile_k, ++ ) ++ self.cta_tile_shape_mnk = ( ++ self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), ++ self.mma_tiler[1], ++ self.mma_tiler[2], ++ ) ++ self.cta_tile_shape_mnk_sfb = ( ++ self.mma_tiler_sfb[0] // cute.size(tiled_mma.thr_id.shape), ++ self.mma_tiler_sfb[1], ++ self.mma_tiler_sfb[2], ++ ) ++ ++ # Compute cluster layout ++ self.cluster_layout_vmnk = cute.tiled_divide( ++ cute.make_layout(self.cluster_shape_mnk), ++ (tiled_mma.thr_id.shape,), ++ ) ++ self.cluster_layout_sfb_vmnk = cute.tiled_divide( ++ cute.make_layout(self.cluster_shape_mnk), ++ (tiled_mma_sfb.thr_id.shape,), ++ ) ++ ++ # Compute number of multicast CTAs for A/B ++ self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2]) ++ self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1]) ++ self.num_mcast_ctas_sfb = cute.size(self.cluster_layout_sfb_vmnk.shape[1]) ++ self.is_a_mcast = self.num_mcast_ctas_a > 1 ++ self.is_b_mcast = self.num_mcast_ctas_b > 1 ++ self.is_sfb_mcast = self.num_mcast_ctas_sfb > 1 ++ ++ # Compute epilogue subtile ++ self.epi_tile = sm100_utils.compute_epilogue_tile_shape( ++ self.cta_tile_shape_mnk, ++ False, ++ self.c_layout, ++ self.c_dtype, ++ ) ++ ++ self.epi_tile_n = cute.size(self.epi_tile[1]) ++ self.num_epilogue_subtiles = ( ++ self.cta_tile_shape_mnk[0] ++ * self.cta_tile_shape_mnk[1] ++ // (cute.size(self.epi_tile[0]) * cute.size(self.epi_tile[1])) ++ ) ++ self.num_epilogue_threads = 32 * len(self.epilog_warp_id) ++ self.reduce_values_per_thread = ( ++ cute.size(self.epi_tile[0]) * cute.size(self.epi_tile[1]) ++ ) // (self.num_epilogue_threads) ++ self.reduce_mailbox_elements = ( ++ (self.split_k_slices - 1) ++ * self.cta_tile_shape_mnk[0] ++ * self.cta_tile_shape_mnk[1] ++ ) ++ self.reduce_smem_bytes = ( ++ self.reduce_mailbox_elements * cutlass.Float32.width // 8 ++ + self.num_epilogue_subtiles * 8 ++ ) ++ ++ # Setup A/B/C stage count in shared memory and ACC stage count in tensor memory ++ self.num_acc_stage, self.num_ab_stage, self.num_c_stage = self._compute_stages( ++ tiled_mma, ++ self.mma_tiler, ++ self.a_dtype, ++ self.a_major_mode, ++ self.b_dtype, ++ self.b_major_mode, ++ self.epi_tile, ++ self.c_dtype, ++ self.c_layout, ++ self.sf_dtype, ++ self.sf_vec_size, ++ self.smem_capacity, ++ self.occupancy, ++ 1024 + self.reduce_smem_bytes, ++ ) ++ ++ # Compute A/B/SFA/SFB/C shared memory layout ++ self.a_smem_layout_staged = sm100_utils.make_smem_layout_a( ++ tiled_mma, ++ self.mma_tiler, ++ self.a_dtype, ++ self.num_ab_stage, ++ ) ++ self.b_smem_layout_staged = sm100_utils.make_smem_layout_b( ++ tiled_mma, ++ self.mma_tiler, ++ self.b_dtype, ++ self.num_ab_stage, ++ ) ++ self.sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa( ++ tiled_mma, ++ self.mma_tiler, ++ self.sf_vec_size, ++ self.num_ab_stage, ++ ) ++ self.sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb( ++ tiled_mma, ++ self.mma_tiler, ++ self.sf_vec_size, ++ self.num_ab_stage, ++ ) ++ self.c_smem_layout_staged = sm100_utils.make_smem_layout_epi( ++ self.c_dtype, ++ self.c_layout, ++ self.epi_tile, ++ self.num_c_stage, ++ ) ++ ++ self.overlapping_accum = self.num_acc_stage == 1 ++ sf_atom_mn = 32 ++ self.num_sfa_tmem_cols = ( ++ self.cta_tile_shape_mnk[0] // sf_atom_mn ++ ) * mma_inst_tile_k ++ self.num_sfb_tmem_cols = ( ++ self.cta_tile_shape_mnk_sfb[1] // sf_atom_mn ++ ) * mma_inst_tile_k ++ self.num_sf_tmem_cols = self.num_sfa_tmem_cols + self.num_sfb_tmem_cols ++ self.num_accumulator_tmem_cols = ( ++ self.cta_tile_shape_mnk[1] * self.num_acc_stage ++ if not self.overlapping_accum ++ else self.cta_tile_shape_mnk[1] * 2 - self.num_sf_tmem_cols ++ ) ++ ++ # Only when overlapping_accum is enabled, we need to release accumulator buffer early in epilogue ++ self.iter_acc_early_release_in_epilogue = ( ++ self.num_sf_tmem_cols // self.epi_tile_n ++ ) ++ ++ @cute.jit ++ def __call__( ++ self, ++ a_tensor: cute.Tensor, ++ b_tensor: cute.Tensor, ++ sfa_tensor: cute.Tensor, ++ sfb_tensor: cute.Tensor, ++ c_tensor: cute.Tensor, ++ alpha: cute.Tensor, # Single-element tensor containing alpha value ++ max_active_clusters: cutlass.Constexpr, ++ stream: cuda.CUstream, ++ epilogue_op: cutlass.Constexpr = lambda x: x, ++ ): ++ """Execute the GEMM operation in steps: ++ - Setup static attributes before smem/grid/tma computation ++ - Setup TMA load/store atoms and tensors ++ - Compute grid size with regard to hardware constraints ++ - Define shared storage for kernel ++ - Launch the kernel synchronously ++ ++ Args: ++ a_tensor (cute.Tensor): Input tensor A ++ b_tensor (cute.Tensor): Input tensor B ++ sfa_tensor (cute.Tensor): Scale factor tensor A ++ sfb_tensor (cute.Tensor): Scale factor tensor B ++ c_tensor (cute.Tensor): Output tensor C ++ max_active_clusters (cutlass.Constexpr): Maximum number of active clusters ++ stream (cuda.CUstream): CUDA stream for asynchronous execution ++ epilogue_op (cutlass.Constexpr): Optional elementwise lambda function to apply to the output tensor ++ ++ Raises: ++ TypeError: If input data types are incompatible with the MMA instruction. ++ """ ++ ( ++ tiled_mma, ++ tiled_mma_sfb, ++ tma_atom_a, ++ tma_tensor_a, ++ tma_atom_b, ++ tma_tensor_b, ++ tma_atom_sfa, ++ tma_tensor_sfa, ++ tma_atom_sfb, ++ tma_tensor_sfb, ++ tma_atom_c, ++ tma_tensor_c, ++ ) = self._prepare_tma( ++ a_tensor, ++ b_tensor, ++ sfa_tensor, ++ sfb_tensor, ++ c_tensor, ++ ) ++ ++ # Compute grid size ++ grid = self._compute_grid(c_tensor, self.cta_tile_shape_mnk) ++ ++ self.buffer_align_bytes = 1024 ++ ++ # Define shared storage for kernel ++ @cute.struct ++ class SharedStorage: ++ ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage] ++ ab_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage] ++ acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage] ++ acc_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage] ++ tmem_holding_buf: cutlass.Int32 ++ reduce_full_mbar_ptr: cute.struct.MemRange[ ++ cutlass.Int64, self.num_epilogue_subtiles ++ ] ++ reduce_mailbox: cute.struct.Align[ ++ cute.struct.MemRange[cutlass.Float32, self.reduce_mailbox_elements], ++ 128, ++ ] ++ # (EPI_TILE_M, EPI_TILE_N, STAGE) ++ sC: cute.struct.Align[ ++ cute.struct.MemRange[ ++ self.c_dtype, ++ cute.cosize(self.c_smem_layout_staged.outer), ++ ], ++ self.buffer_align_bytes, ++ ] ++ # (MMA, MMA_M, MMA_K, STAGE) ++ sA: cute.struct.Align[ ++ cute.struct.MemRange[ ++ self.a_dtype, cute.cosize(self.a_smem_layout_staged.outer) ++ ], ++ self.buffer_align_bytes, ++ ] ++ # (MMA, MMA_N, MMA_K, STAGE) ++ sB: cute.struct.Align[ ++ cute.struct.MemRange[ ++ self.b_dtype, cute.cosize(self.b_smem_layout_staged.outer) ++ ], ++ self.buffer_align_bytes, ++ ] ++ # (MMA, MMA_M, MMA_K, STAGE) ++ sSFA: cute.struct.Align[ ++ cute.struct.MemRange[ ++ self.sf_dtype, cute.cosize(self.sfa_smem_layout_staged) ++ ], ++ self.buffer_align_bytes, ++ ] ++ # (MMA, MMA_N, MMA_K, STAGE) ++ sSFB: cute.struct.Align[ ++ cute.struct.MemRange[ ++ self.sf_dtype, cute.cosize(self.sfb_smem_layout_staged) ++ ], ++ self.buffer_align_bytes, ++ ] ++ ++ self.shared_storage = SharedStorage ++ ++ # Launch the kernel synchronously ++ self.kernel( ++ tiled_mma, ++ tiled_mma_sfb, ++ tma_atom_a, ++ tma_tensor_a, ++ tma_atom_b, ++ tma_tensor_b, ++ tma_atom_sfa, ++ tma_tensor_sfa, ++ tma_atom_sfb, ++ tma_tensor_sfb, ++ tma_atom_c, ++ tma_tensor_c, ++ self.cluster_layout_vmnk, ++ self.cluster_layout_sfb_vmnk, ++ self.a_smem_layout_staged, ++ self.b_smem_layout_staged, ++ self.sfa_smem_layout_staged, ++ self.sfb_smem_layout_staged, ++ self.c_smem_layout_staged, ++ self.epi_tile, ++ epilogue_op, ++ alpha, ++ ).launch( ++ grid=grid, ++ block=[self.threads_per_cta, 1, 1], ++ cluster=self.cluster_shape_mnk, ++ smem=self.shared_storage.size_in_bytes(), # type: ignore[attr-defined] ++ min_blocks_per_mp=1, ++ stream=stream, ++ use_pdl=self.enable_pdl, ++ ) ++ return ++ ++ # GPU device kernel ++ @cute.kernel ++ def kernel( ++ self, ++ tiled_mma: cute.TiledMma, ++ tiled_mma_sfb: cute.TiledMma, ++ tma_atom_a: cute.CopyAtom, ++ mA_mkl: cute.Tensor, ++ tma_atom_b: cute.CopyAtom, ++ mB_nkl: cute.Tensor, ++ tma_atom_sfa: cute.CopyAtom, ++ mSFA_mkl: cute.Tensor, ++ tma_atom_sfb: cute.CopyAtom, ++ mSFB_nkl: cute.Tensor, ++ tma_atom_c: Optional[cute.CopyAtom], ++ mC_mnl: cute.Tensor, ++ cluster_layout_vmnk: cute.Layout, ++ cluster_layout_sfb_vmnk: cute.Layout, ++ a_smem_layout_staged: cute.ComposedLayout, ++ b_smem_layout_staged: cute.ComposedLayout, ++ sfa_smem_layout_staged: cute.Layout, ++ sfb_smem_layout_staged: cute.Layout, ++ c_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout, None], ++ epi_tile: cute.Tile, ++ epilogue_op: cutlass.Constexpr, ++ alpha: cute.Tensor, ++ ): ++ """ ++ GPU device kernel performing the Persistent batched GEMM computation. ++ """ ++ # Keep alpha in FP32 for precision: the accumulator is in FP32 and alpha ++ # may be a very small scaling factor. Converting to c_dtype (e.g., FP16) ++ # before multiplication could cause overflow when acc values are large. ++ alpha_value = alpha[0].to(cutlass.Float32) ++ ++ warp_idx = cute.arch.warp_idx() ++ warp_idx = cute.arch.make_warp_uniform(warp_idx) ++ ++ # ++ # Prefetch tma desc ++ # ++ if warp_idx == self.tma_warp_id: ++ cpasync.prefetch_descriptor(tma_atom_a) ++ cpasync.prefetch_descriptor(tma_atom_b) ++ cpasync.prefetch_descriptor(tma_atom_sfa) ++ cpasync.prefetch_descriptor(tma_atom_sfb) ++ cpasync.prefetch_descriptor(tma_atom_c) ++ ++ # ++ # Setup cta/thread coordinates ++ # ++ # Coords inside cluster ++ bidx, bidy, bidz = cute.arch.block_idx() ++ mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape) ++ is_leader_cta = mma_tile_coord_v == 0 ++ cta_rank_in_cluster = cute.arch.make_warp_uniform( ++ cute.arch.block_idx_in_cluster() ++ ) ++ block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord( ++ cta_rank_in_cluster ++ ) ++ block_in_cluster_coord_sfb_vmnk = cluster_layout_sfb_vmnk.get_flat_coord( ++ cta_rank_in_cluster ++ ) ++ # Coord inside cta ++ tidx, _, _ = cute.arch.thread_idx() ++ ++ # ++ # Alloc and init: a+b full/empty, accumulator full/empty, tensor memory dealloc barrier ++ # ++ smem = utils.SmemAllocator() ++ storage = smem.allocate(self.shared_storage) ++ ++ tmem_holding_buf = storage.tmem_holding_buf.ptr ++ reduce_full_mbar_ptr = storage.reduce_full_mbar_ptr.data_ptr() ++ ++ # Initialize mainloop ab_pipeline (barrier) and states ++ ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) ++ num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 ++ ab_pipeline_consumer_group = pipeline.CooperativeGroup( ++ pipeline.Agent.Thread, num_tma_producer ++ ) ++ ab_pipeline = PipelineTmaUmma.create( ++ barrier_storage=storage.ab_full_mbar_ptr.data_ptr(), ++ num_stages=self.num_ab_stage, ++ producer_group=ab_pipeline_producer_group, ++ consumer_group=ab_pipeline_consumer_group, ++ tx_count=self.num_tma_load_bytes, ++ cta_layout_vmnk=cluster_layout_vmnk, ++ ) ++ ++ # Initialize acc_pipeline (barrier) and states ++ acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) ++ num_acc_consumer_threads = len(self.epilog_warp_id) ++ acc_pipeline_consumer_group = pipeline.CooperativeGroup( ++ pipeline.Agent.Thread, num_acc_consumer_threads ++ ) ++ acc_pipeline = PipelineUmmaAsync.create( ++ barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), ++ num_stages=self.num_acc_stage, ++ producer_group=acc_pipeline_producer_group, ++ consumer_group=acc_pipeline_consumer_group, ++ cta_layout_vmnk=cluster_layout_vmnk, ++ ) ++ ++ if warp_idx == self.tma_warp_id: ++ with cute.arch.elect_one(): ++ for subtile_idx in cutlass.range_constexpr(self.num_epilogue_subtiles): ++ cute.arch.mbarrier_init(reduce_full_mbar_ptr + subtile_idx, 1) ++ cute.arch.mbarrier_init_fence() ++ ++ # Cluster arrive after barrier init ++ if cute.size(self.cluster_shape_mnk) > 1: ++ cute.arch.cluster_arrive_relaxed(aligned=True) ++ ++ # ++ # Setup smem tensor A/B/SFA/SFB/C ++ reduce_mailbox = storage.reduce_mailbox.get_tensor( ++ cute.make_layout((self.reduce_mailbox_elements,)) ++ ) ++ # ++ # (EPI_TILE_M, EPI_TILE_N, STAGE) ++ sC = storage.sC.get_tensor( ++ c_smem_layout_staged.outer, swizzle=c_smem_layout_staged.inner ++ ) ++ # (MMA, MMA_M, MMA_K, STAGE) ++ sA = storage.sA.get_tensor( ++ a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner ++ ) ++ # (MMA, MMA_N, MMA_K, STAGE) ++ sB = storage.sB.get_tensor( ++ b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner ++ ) ++ # (MMA, MMA_M, MMA_K, STAGE) ++ sSFA = storage.sSFA.get_tensor(sfa_smem_layout_staged) ++ # (MMA, MMA_N, MMA_K, STAGE) ++ sSFB = storage.sSFB.get_tensor(sfb_smem_layout_staged) ++ ++ # ++ # Compute multicast mask for A/B/SFA/SFB buffer full ++ # ++ a_full_mcast_mask = None ++ b_full_mcast_mask = None ++ sfa_full_mcast_mask = None ++ sfb_full_mcast_mask = None ++ if cutlass.const_expr(self.is_a_mcast or self.is_b_mcast): ++ a_full_mcast_mask = cpasync.create_tma_multicast_mask( ++ cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 ++ ) ++ b_full_mcast_mask = cpasync.create_tma_multicast_mask( ++ cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1 ++ ) ++ sfa_full_mcast_mask = cpasync.create_tma_multicast_mask( ++ cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 ++ ) ++ sfb_full_mcast_mask = cpasync.create_tma_multicast_mask( ++ cluster_layout_sfb_vmnk, block_in_cluster_coord_sfb_vmnk, mcast_mode=1 ++ ) ++ ++ # ++ # Local_tile partition global tensors ++ # ++ # (bM, bK, RestM, RestK, RestL) ++ gA_mkl = cute.local_tile( ++ mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) ++ ) ++ # (bN, bK, RestN, RestK, RestL) ++ gB_nkl = cute.local_tile( ++ mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None) ++ ) ++ # (bM, bK, RestM, RestK, RestL) ++ gSFA_mkl = cute.local_tile( ++ mSFA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) ++ ) ++ # (bN, bK, RestN, RestK, RestL) ++ gSFB_nkl = cute.local_tile( ++ mSFB_nkl, ++ cute.slice_(self.mma_tiler_sfb, (0, None, None)), ++ (None, None, None), ++ ) ++ # (bM, bN, RestM, RestN, RestL) ++ gC_mnl = cute.local_tile( ++ mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) ++ ) ++ total_k_block_cnt = cutlass.Int32(cute.size(gA_mkl, mode=[3])) ++ k_block_cnt = total_k_block_cnt ++ if cutlass.const_expr(self.split_k_slices > 1): ++ k_block_cnt = total_k_block_cnt // self.split_k_slices ++ ++ # ++ # Partition global tensor for TiledMMA_A/B/C ++ # ++ thr_mma = tiled_mma.get_slice(mma_tile_coord_v) ++ thr_mma_sfb = tiled_mma_sfb.get_slice(mma_tile_coord_v) ++ # (MMA, MMA_M, MMA_K, RestM, RestK, RestL) ++ tCgA = thr_mma.partition_A(gA_mkl) ++ # (MMA, MMA_N, MMA_K, RestN, RestK, RestL) ++ tCgB = thr_mma.partition_B(gB_nkl) ++ # (MMA, MMA_M, MMA_K, RestM, RestK, RestL) ++ tCgSFA = thr_mma.partition_A(gSFA_mkl) ++ # (MMA, MMA_N, MMA_K, RestN, RestK, RestL) ++ tCgSFB = thr_mma_sfb.partition_B(gSFB_nkl) ++ # (MMA, MMA_M, MMA_N, RestM, RestN, RestL) ++ tCgC = thr_mma.partition_C(gC_mnl) ++ ++ # ++ # Partition global/shared tensor for TMA load A/B ++ # ++ # TMA load A partition_S/D ++ a_cta_layout = cute.make_layout( ++ cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape ++ ) ++ # ((atom_v, rest_v), STAGE) ++ # ((atom_v, rest_v), RestM, RestK, RestL) ++ tAsA, tAgA = cpasync.tma_partition( ++ tma_atom_a, ++ block_in_cluster_coord_vmnk[2], ++ a_cta_layout, ++ cute.group_modes(sA, 0, 3), ++ cute.group_modes(tCgA, 0, 3), ++ ) ++ # TMA load B partition_S/D ++ b_cta_layout = cute.make_layout( ++ cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape ++ ) ++ # ((atom_v, rest_v), STAGE) ++ # ((atom_v, rest_v), RestN, RestK, RestL) ++ tBsB, tBgB = cpasync.tma_partition( ++ tma_atom_b, ++ block_in_cluster_coord_vmnk[1], ++ b_cta_layout, ++ cute.group_modes(sB, 0, 3), ++ cute.group_modes(tCgB, 0, 3), ++ ) ++ ++ # TMA load SFA partition_S/D ++ sfa_cta_layout = a_cta_layout ++ # ((atom_v, rest_v), STAGE) ++ # ((atom_v, rest_v), RestM, RestK, RestL) ++ tAsSFA, tAgSFA = cute.nvgpu.cpasync.tma_partition( ++ tma_atom_sfa, ++ block_in_cluster_coord_vmnk[2], ++ sfa_cta_layout, ++ cute.group_modes(sSFA, 0, 3), ++ cute.group_modes(tCgSFA, 0, 3), ++ ) ++ tAsSFA = cute.filter_zeros(tAsSFA) ++ tAgSFA = cute.filter_zeros(tAgSFA) ++ ++ # TMA load SFB partition_S/D ++ sfb_cta_layout = cute.make_layout( ++ cute.slice_(cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape ++ ) ++ # ((atom_v, rest_v), STAGE) ++ # ((atom_v, rest_v), RestN, RestK, RestL) ++ tBsSFB, tBgSFB = cute.nvgpu.cpasync.tma_partition( ++ tma_atom_sfb, ++ block_in_cluster_coord_sfb_vmnk[1], ++ sfb_cta_layout, ++ cute.group_modes(sSFB, 0, 3), ++ cute.group_modes(tCgSFB, 0, 3), ++ ) ++ tBsSFB = cute.filter_zeros(tBsSFB) ++ tBgSFB = cute.filter_zeros(tBgSFB) ++ ++ # ++ # Partition shared/tensor memory tensor for TiledMMA_A/B/C ++ # ++ # (MMA, MMA_M, MMA_K, STAGE) ++ tCrA = tiled_mma.make_fragment_A(sA) ++ # (MMA, MMA_N, MMA_K, STAGE) ++ tCrB = tiled_mma.make_fragment_B(sB) ++ # (MMA, MMA_M, MMA_N) ++ acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) ++ # (MMA, MMA_M, MMA_N, STAGE) ++ if cutlass.const_expr(self.overlapping_accum): ++ num_acc_stage_overlapped = 2 ++ tCtAcc_fake = tiled_mma.make_fragment_C( ++ cute.append(acc_shape, num_acc_stage_overlapped) ++ ) ++ # (MMA, MMA_M, MMA_N, STAGE) ++ tCtAcc_fake = cute.make_tensor( ++ tCtAcc_fake.iterator, ++ cute.make_layout( ++ tCtAcc_fake.shape, ++ stride=( ++ tCtAcc_fake.stride[0], ++ tCtAcc_fake.stride[1], ++ tCtAcc_fake.stride[2], ++ (256 - self.num_sf_tmem_cols) * tCtAcc_fake.stride[0][1], ++ ), ++ ), ++ ) ++ else: ++ # (MMA, MMA_M, MMA_N, STAGE) ++ tCtAcc_fake = tiled_mma.make_fragment_C( ++ cute.append(acc_shape, self.num_acc_stage) ++ ) ++ ++ # ++ # Cluster wait before tensor memory alloc ++ # ++ if cute.size(self.cluster_shape_mnk) > 1: ++ cute.arch.cluster_wait() ++ else: ++ cute.arch.barrier( ++ barrier_id=self.cta_sync_bar_id, number_of_threads=self.threads_per_cta ++ ) ++ ++ # PDL bookend: griddepcontrol_wait is always emitted; the actual PDL ++ # behavior is controlled by use_pdl= in .launch(). The griddepcontrol ++ # instructions are effectively no-ops when PDL is not enabled at ++ # the driver level. ++ griddepcontrol_wait() ++ ++ # ++ # Specialized TMA load warp ++ # ++ if warp_idx == self.tma_warp_id: ++ # ++ # Persistent tile scheduling loop ++ # ++ tile_sched = _SplitKOneTileScheduler.create( ++ cute.arch.block_idx(), cta_rank_in_cluster ++ ) ++ work_tile = tile_sched.initial_work_tile_info() ++ ++ ab_producer_state = pipeline.make_pipeline_state( ++ pipeline.PipelineUserType.Producer, self.num_ab_stage ++ ) ++ ++ while work_tile.is_valid_tile: ++ # Get tile coord from tile scheduler ++ cur_tile_coord = work_tile.tile_idx ++ mma_tile_coord_mnl = ( ++ cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), ++ cur_tile_coord[1], ++ cur_tile_coord[2], ++ ) ++ ++ input_l = mma_tile_coord_mnl[2] ++ k_block_start = cutlass.Int32(0) ++ if cutlass.const_expr(self.split_k_slices > 1): ++ # C L is the split-K index; dense inputs always use L=0. ++ input_l = cutlass.Int32(0) ++ k_block_start = cutlass.Int32(mma_tile_coord_mnl[2]) * k_block_cnt ++ ++ # ++ # Slice to per mma tile index ++ # ++ # ((atom_v, rest_v), RestK) ++ tAgA_slice = tAgA[(None, mma_tile_coord_mnl[0], None, input_l)] ++ # ((atom_v, rest_v), RestK) ++ tBgB_slice = tBgB[(None, mma_tile_coord_mnl[1], None, input_l)] ++ ++ # ((atom_v, rest_v), RestK) ++ tAgSFA_slice = tAgSFA[(None, mma_tile_coord_mnl[0], None, input_l)] ++ slice_n = mma_tile_coord_mnl[1] ++ if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64): ++ slice_n = mma_tile_coord_mnl[1] // 2 ++ # ((atom_v, rest_v), RestK) ++ tBgSFB_slice = tBgSFB[(None, slice_n, None, input_l)] ++ ++ # Peek (try_wait) AB buffer empty for k_block = prefetch_k_block_cnt ++ ab_producer_state.reset_count() ++ peek_ab_empty_status = cutlass.Boolean(1) ++ if ab_producer_state.count < k_block_cnt: ++ peek_ab_empty_status = ab_pipeline.producer_try_acquire( ++ ab_producer_state ++ ) ++ # ++ # Tma load loop ++ # ++ for _k_block in cutlass.range(0, k_block_cnt, 1, unroll=1): ++ # Conditionally wait for AB buffer empty ++ ab_pipeline.producer_acquire( ++ ab_producer_state, peek_ab_empty_status ++ ) ++ ++ # TMA load A/B/SFA/SFB ++ cute.copy( ++ tma_atom_a, ++ tAgA_slice[(None, k_block_start + ab_producer_state.count)], ++ tAsA[(None, ab_producer_state.index)], ++ tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), ++ mcast_mask=a_full_mcast_mask, ++ ) ++ cute.copy( ++ tma_atom_b, ++ tBgB_slice[(None, k_block_start + ab_producer_state.count)], ++ tBsB[(None, ab_producer_state.index)], ++ tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), ++ mcast_mask=b_full_mcast_mask, ++ ) ++ cute.copy( ++ tma_atom_sfa, ++ tAgSFA_slice[(None, k_block_start + ab_producer_state.count)], ++ tAsSFA[(None, ab_producer_state.index)], ++ tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), ++ mcast_mask=sfa_full_mcast_mask, ++ ) ++ cute.copy( ++ tma_atom_sfb, ++ tBgSFB_slice[(None, k_block_start + ab_producer_state.count)], ++ tBsSFB[(None, ab_producer_state.index)], ++ tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), ++ mcast_mask=sfb_full_mcast_mask, ++ ) ++ ++ # Peek (try_wait) AB buffer empty for k_block = prefetch_k_block_cnt + k_block + 1 ++ ab_producer_state.advance() ++ peek_ab_empty_status = cutlass.Boolean(1) ++ if ab_producer_state.count < k_block_cnt: ++ peek_ab_empty_status = ab_pipeline.producer_try_acquire( ++ ab_producer_state ++ ) ++ ++ # ++ # Advance to next tile ++ # ++ tile_sched.advance_to_next_work() ++ work_tile = tile_sched.get_current_work() ++ ++ # ++ # Wait A/B buffer empty ++ # ++ ab_pipeline.producer_tail(ab_producer_state) ++ ++ # ++ # Specialized MMA warp ++ # ++ if warp_idx == self.mma_warp_id: ++ # ++ # Bar sync for retrieve tensor memory ptr from shared mem ++ # ++ tmem_ptr_read_threads = 32 * len((self.mma_warp_id, *self.epilog_warp_id)) ++ cute.arch.barrier( ++ barrier_id=self.tmem_ptr_sync_bar_id, ++ number_of_threads=tmem_ptr_read_threads, ++ ) ++ ++ # ++ # Retrieving tensor memory ptr and make accumulator/SFA/SFB tensor ++ # ++ # Make accumulator tmem tensor ++ acc_tmem_ptr = cute.arch.retrieve_tmem_ptr( ++ self.acc_dtype, ++ alignment=16, ++ ptr_to_buffer_holding_addr=tmem_holding_buf, ++ ) ++ # (MMA, MMA_M, MMA_N, STAGE) ++ tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) ++ ++ # Make SFA tmem tensor ++ sfa_tmem_ptr = cute.recast_ptr( ++ acc_tmem_ptr + self.num_accumulator_tmem_cols, ++ dtype=self.sf_dtype, ++ ) ++ # (MMA, MMA_M, MMA_K) ++ tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa( ++ tiled_mma, ++ self.mma_tiler, ++ self.sf_vec_size, ++ cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)), ++ ) ++ tCtSFA = cute.make_tensor(sfa_tmem_ptr, tCtSFA_layout) ++ ++ # Make SFB tmem tensor ++ sfb_tmem_ptr = cute.recast_ptr( ++ acc_tmem_ptr + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols, ++ dtype=self.sf_dtype, ++ ) ++ # (MMA, MMA_N, MMA_K) ++ tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb( ++ tiled_mma, ++ self.mma_tiler, ++ self.sf_vec_size, ++ cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)), ++ ) ++ tCtSFB = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout) ++ # ++ # Partition for S2T copy of SFA/SFB ++ # ++ tiled_copy_s2t_sfa, tCsSFA_compact_s2t, tCtSFA_compact_s2t = ( ++ self.mainloop_s2t_copy_and_partition(sSFA, tCtSFA) ++ ) ++ tiled_copy_s2t_sfb, tCsSFB_compact_s2t, tCtSFB_compact_s2t = ( ++ self.mainloop_s2t_copy_and_partition(sSFB, tCtSFB) ++ ) ++ ++ # ++ # Persistent tile scheduling loop ++ # ++ tile_sched = _SplitKOneTileScheduler.create( ++ cute.arch.block_idx(), cta_rank_in_cluster ++ ) ++ work_tile = tile_sched.initial_work_tile_info() ++ ++ ab_consumer_state = pipeline.make_pipeline_state( ++ pipeline.PipelineUserType.Consumer, self.num_ab_stage ++ ) ++ acc_producer_state = pipeline.make_pipeline_state( ++ pipeline.PipelineUserType.Producer, self.num_acc_stage ++ ) ++ ++ while work_tile.is_valid_tile: ++ # Get tile coord from tile scheduler ++ cur_tile_coord = work_tile.tile_idx ++ mma_tile_coord_mnl = ( ++ cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), ++ cur_tile_coord[1], ++ cur_tile_coord[2], ++ ) ++ ++ if cutlass.const_expr(self.overlapping_accum): ++ acc_stage_index = acc_producer_state.phase ^ 1 ++ else: ++ acc_stage_index = acc_producer_state.index ++ ++ # Set tensor memory buffer for current tile ++ # (MMA, MMA_M, MMA_N) ++ tCtAcc = tCtAcc_base[(None, None, None, acc_stage_index)] ++ ++ # Peek (try_wait) AB buffer full for k_block = 0 ++ ab_consumer_state.reset_count() ++ peek_ab_full_status = cutlass.Boolean(1) ++ if ab_consumer_state.count < k_block_cnt and is_leader_cta: ++ peek_ab_full_status = ab_pipeline.consumer_try_wait( ++ ab_consumer_state ++ ) ++ ++ # ++ # Wait for accumulator buffer empty ++ # ++ if is_leader_cta: ++ acc_pipeline.producer_acquire(acc_producer_state) ++ ++ tCtSFB_mma = tCtSFB ++ if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 192): ++ # If this is an ODD tile, shift the TMEM start address for cta_tile_shape_n=192 case by two words (ignores first 64 columns of SFB) ++ offset = ( ++ cutlass.Int32(2) ++ if mma_tile_coord_mnl[1] % 2 == 1 ++ else cutlass.Int32(0) ++ ) ++ shifted_ptr = cute.recast_ptr( ++ acc_tmem_ptr ++ + self.num_accumulator_tmem_cols ++ + self.num_sfa_tmem_cols ++ + offset, ++ dtype=self.sf_dtype, ++ ) ++ tCtSFB_mma = cute.make_tensor(shifted_ptr, tCtSFB_layout) ++ elif cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64): ++ # Move in increments of 64 columns of SFB ++ offset = cutlass.Int32((mma_tile_coord_mnl[1] % 2) * 2) ++ shifted_ptr = cute.recast_ptr( ++ acc_tmem_ptr ++ + self.num_accumulator_tmem_cols ++ + self.num_sfa_tmem_cols ++ + offset, ++ dtype=self.sf_dtype, ++ ) ++ tCtSFB_mma = cute.make_tensor(shifted_ptr, tCtSFB_layout) ++ ++ # ++ # Reset the ACCUMULATE field for each tile ++ # ++ tiled_mma.set(tcgen05.Field.ACCUMULATE, False) ++ ++ # ++ # Mma mainloop ++ # ++ for _k_block in range(k_block_cnt): ++ if is_leader_cta: ++ # Conditionally wait for AB buffer full ++ ab_pipeline.consumer_wait( ++ ab_consumer_state, peek_ab_full_status ++ ) ++ ++ # Copy SFA/SFB from smem to tmem ++ s2t_stage_coord = ( ++ None, ++ None, ++ None, ++ None, ++ ab_consumer_state.index, ++ ) ++ tCsSFA_compact_s2t_staged = tCsSFA_compact_s2t[s2t_stage_coord] ++ tCsSFB_compact_s2t_staged = tCsSFB_compact_s2t[s2t_stage_coord] ++ cute.copy( ++ tiled_copy_s2t_sfa, ++ tCsSFA_compact_s2t_staged, ++ tCtSFA_compact_s2t, ++ ) ++ cute.copy( ++ tiled_copy_s2t_sfb, ++ tCsSFB_compact_s2t_staged, ++ tCtSFB_compact_s2t, ++ ) ++ ++ # tCtAcc += tCrA * tCrSFA * tCrB * tCrSFB ++ num_kphases = cute.size(tCrA, mode=[2]) ++ for kphase_idx in cutlass.range(num_kphases, unroll_full=True): ++ kphase_coord = ( ++ None, ++ None, ++ kphase_idx, ++ ab_consumer_state.index, ++ ) ++ ++ # Set SFA/SFB tensor to tiled_mma ++ sf_kphase_coord = (None, None, kphase_idx) ++ tiled_mma.set( ++ tcgen05.Field.SFA, ++ tCtSFA[sf_kphase_coord].iterator, ++ ) ++ tiled_mma.set( ++ tcgen05.Field.SFB, ++ tCtSFB_mma[sf_kphase_coord].iterator, ++ ) ++ ++ cute.gemm( ++ tiled_mma, ++ tCtAcc, ++ tCrA[kphase_coord], ++ tCrB[kphase_coord], ++ tCtAcc, ++ ) ++ ++ # Enable accumulate on tCtAcc after first kphase ++ tiled_mma.set(tcgen05.Field.ACCUMULATE, True) ++ ++ # Async arrive AB buffer empty ++ ab_pipeline.consumer_release(ab_consumer_state) ++ ++ # Peek (try_wait) AB buffer full for k_block = k_block + 1 ++ ab_consumer_state.advance() ++ peek_ab_full_status = cutlass.Boolean(1) ++ if ab_consumer_state.count < k_block_cnt: ++ if is_leader_cta: ++ peek_ab_full_status = ab_pipeline.consumer_try_wait( ++ ab_consumer_state ++ ) ++ ++ # ++ # Async arrive accumulator buffer full ++ # ++ if is_leader_cta: ++ acc_pipeline.producer_commit(acc_producer_state) ++ acc_producer_state.advance() ++ ++ # ++ # Advance to next tile ++ # ++ tile_sched.advance_to_next_work() ++ work_tile = tile_sched.get_current_work() ++ ++ # ++ # Wait for accumulator buffer empty ++ # ++ acc_pipeline.producer_tail(acc_producer_state) ++ # ++ # Specialized epilogue warps ++ # ++ if warp_idx < self.mma_warp_id: ++ # ++ # Alloc tensor memory buffer ++ # ++ if warp_idx == self.epilog_warp_id[0]: ++ cute.arch.alloc_tmem( ++ self.num_tmem_alloc_cols, ++ tmem_holding_buf, ++ is_two_cta=False, ++ ) ++ ++ # ++ # Bar sync for retrieve tensor memory ptr from shared memory ++ # ++ tmem_ptr_read_threads = 32 * len((self.mma_warp_id, *self.epilog_warp_id)) ++ cute.arch.barrier( ++ barrier_id=self.tmem_ptr_sync_bar_id, ++ number_of_threads=tmem_ptr_read_threads, ++ ) ++ ++ # ++ # Retrieving tensor memory ptr and make accumulator tensor ++ # ++ acc_tmem_ptr = cute.arch.retrieve_tmem_ptr( ++ self.acc_dtype, ++ alignment=16, ++ ptr_to_buffer_holding_addr=tmem_holding_buf, ++ ) ++ # (MMA, MMA_M, MMA_N, STAGE) ++ tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) ++ ++ # ++ # Partition for epilogue ++ # ++ epi_tidx = tidx ++ tiled_copy_t2r, tTR_tAcc_base, tTR_rAcc = ( ++ self.epilog_tmem_copy_and_partition( ++ epi_tidx, tCtAcc_base, tCgC, epi_tile, False ++ ) ++ ) ++ ++ tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype) ++ tiled_copy_r2s, tRS_rC, tRS_sC = self.epilog_smem_copy_and_partition( ++ tiled_copy_t2r, tTR_rC, epi_tidx, sC ++ ) ++ tma_atom_c, bSG_sC, bSG_gC_partitioned = ( ++ self.epilog_gmem_copy_and_partition( ++ epi_tidx, tma_atom_c, tCgC, epi_tile, sC ++ ) ++ ) ++ ++ # ++ # Persistent tile scheduling loop ++ # ++ tile_sched = _SplitKOneTileScheduler.create( ++ cute.arch.block_idx(), cta_rank_in_cluster ++ ) ++ work_tile = tile_sched.initial_work_tile_info() ++ ++ acc_consumer_state = pipeline.make_pipeline_state( ++ pipeline.PipelineUserType.Consumer, self.num_acc_stage ++ ) ++ ++ # Threads/warps participating in tma store pipeline ++ c_producer_group = pipeline.CooperativeGroup( ++ pipeline.Agent.Thread, ++ 32 * len(self.epilog_warp_id), ++ ) ++ c_pipeline = pipeline.PipelineTmaStore.create( ++ num_stages=self.num_c_stage, ++ producer_group=c_producer_group, ++ ) ++ ++ while work_tile.is_valid_tile: ++ # Get tile coord from tile scheduler ++ cur_tile_coord = work_tile.tile_idx ++ mma_tile_coord_mnl = ( ++ cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), ++ cur_tile_coord[1], ++ cur_tile_coord[2], ++ ) ++ ++ # ++ # Slice to per mma tile index ++ # ++ # ((ATOM_V, REST_V), EPI_M, EPI_N) ++ bSG_gC = bSG_gC_partitioned[ ++ ( ++ None, ++ None, ++ None, ++ mma_tile_coord_mnl[0], ++ mma_tile_coord_mnl[1], ++ 0, ++ ) ++ ] ++ ++ if cutlass.const_expr(self.overlapping_accum): ++ acc_stage_index = acc_consumer_state.phase ++ reverse_subtile = ( ++ cutlass.Boolean(True) ++ if acc_stage_index == 0 ++ else cutlass.Boolean(False) ++ ) ++ else: ++ acc_stage_index = acc_consumer_state.index ++ ++ # Set tensor memory buffer for current tile ++ # (T2R, T2R_M, T2R_N, EPI_M, EPI_M) ++ tTR_tAcc = tTR_tAcc_base[ ++ (None, None, None, None, None, acc_stage_index) ++ ] ++ ++ # ++ # Wait for accumulator buffer full ++ # ++ acc_pipeline.consumer_wait(acc_consumer_state) ++ ++ tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) ++ bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC)) ++ ++ # ++ # Store accumulator to global memory in sub-tiles ++ # ++ subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) ++ ++ for subtile_idx in cutlass.range(subtile_cnt): ++ real_subtile_idx = subtile_idx ++ if cutlass.const_expr(self.overlapping_accum): ++ if reverse_subtile: ++ real_subtile_idx = ( ++ self.cta_tile_shape_mnk[1] // self.epi_tile_n ++ - 1 ++ - subtile_idx ++ ) ++ # ++ # Load accumulator from tensor memory buffer to register ++ # ++ tTR_tAcc_mn = tTR_tAcc[(None, None, None, real_subtile_idx)] ++ cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) ++ ++ # ++ # Async arrive accumulator buffer empty earlier when overlapping_accum is enabled ++ # ++ if cutlass.const_expr(self.overlapping_accum): ++ if subtile_idx == self.iter_acc_early_release_in_epilogue: ++ # Fence for TMEM load ++ cute.arch.fence_view_async_tmem_load() ++ with cute.arch.elect_one(): ++ acc_pipeline.consumer_release(acc_consumer_state) ++ acc_consumer_state.advance() ++ # Retile the FP32 accumulator fragment into the same register ++ # order used by the existing shared-memory epilogue. ++ tRS_rAcc = tiled_copy_r2s.retile(tTR_rAcc) ++ owner_rank = cutlass.Int32(real_subtile_idx % self.split_k_slices) ++ peer_count = self.split_k_slices - 1 ++ values_per_thread = self.reduce_values_per_thread ++ values_per_peer = self.num_epilogue_threads * values_per_thread ++ subtile_mailbox_base = ( ++ real_subtile_idx * peer_count * values_per_peer ++ ) ++ reduce_mbar_ptr = reduce_full_mbar_ptr + real_subtile_idx ++ ++ if cta_rank_in_cluster != owner_rank: ++ peer_slot = cta_rank_in_cluster ++ if cta_rank_in_cluster > owner_rank: ++ peer_slot = cta_rank_in_cluster - cutlass.Int32(1) ++ remote_thread_base = ( ++ subtile_mailbox_base ++ + peer_slot * values_per_peer ++ + epi_tidx * values_per_thread ++ ) ++ for value_idx in cutlass.range_constexpr( ++ 0, values_per_thread, 4 ++ ): ++ _store_shared_remote_v4( ++ tRS_rAcc[value_idx], ++ tRS_rAcc[value_idx + 1], ++ tRS_rAcc[value_idx + 2], ++ tRS_rAcc[value_idx + 3], ++ reduce_mailbox.iterator ++ + remote_thread_base ++ + value_idx, ++ reduce_mbar_ptr, ++ owner_rank, ++ ) ++ else: ++ if warp_idx == self.epilog_warp_id[0]: ++ with cute.arch.elect_one(): ++ cute.arch.mbarrier_arrive_and_expect_tx( ++ reduce_mbar_ptr, ++ peer_count ++ * (cute.size(epi_tile[0]) * cute.size(epi_tile[1])) ++ * 4, ++ ) ++ cute.arch.mbarrier_wait(reduce_mbar_ptr, phase=0) ++ ++ local_thread_base = ( ++ subtile_mailbox_base + epi_tidx * values_per_thread ++ ) ++ for peer_slot in cutlass.range_constexpr(peer_count): ++ peer_thread_base = ( ++ local_thread_base + peer_slot * values_per_peer ++ ) ++ for value_idx in cutlass.range_constexpr(values_per_thread): ++ tRS_rAcc[value_idx] = ( ++ tRS_rAcc[value_idx] ++ + reduce_mailbox[peer_thread_base + value_idx] ++ ) ++ ++ # Apply alpha and cast only after the FP32 cluster sum. ++ acc_vec = tRS_rAcc.load() ++ acc_vec = epilogue_op((alpha_value * acc_vec).to(self.c_dtype)) ++ tRS_rC.store(acc_vec) ++ ++ c_buffer = subtile_idx % self.num_c_stage ++ cute.copy( ++ tiled_copy_r2s, ++ tRS_rC, ++ tRS_sC[(None, None, None, c_buffer)], ++ ) ++ cute.arch.fence_proxy( ++ "async.shared", ++ space="cta", ++ ) ++ epilog_threads = 32 * len(self.epilog_warp_id) ++ cute.arch.barrier( ++ barrier_id=self.epilog_sync_bar_id, ++ number_of_threads=epilog_threads, ++ ) ++ ++ if warp_idx == self.epilog_warp_id[0]: ++ cute.copy( ++ tma_atom_c, ++ bSG_sC[(None, c_buffer)], ++ bSG_gC[(None, real_subtile_idx)], ++ ) ++ c_pipeline.producer_commit() ++ c_pipeline.producer_acquire() ++ cute.arch.barrier( ++ barrier_id=self.epilog_sync_bar_id, ++ number_of_threads=epilog_threads, ++ ) ++ ++ # ++ # Async arrive accumulator buffer empty ++ # ++ if cutlass.const_expr(not self.overlapping_accum): ++ with cute.arch.elect_one(): ++ acc_pipeline.consumer_release(acc_consumer_state) ++ acc_consumer_state.advance() ++ ++ # ++ # Advance to next tile ++ # ++ tile_sched.advance_to_next_work() ++ work_tile = tile_sched.get_current_work() ++ ++ # ++ # Dealloc the tensor memory buffer ++ # ++ if warp_idx == self.epilog_warp_id[0]: ++ cute.arch.relinquish_tmem_alloc_permit(is_two_cta=False) ++ epilog_threads = 32 * len(self.epilog_warp_id) ++ cute.arch.barrier( ++ barrier_id=self.epilog_sync_bar_id, number_of_threads=epilog_threads ++ ) ++ if warp_idx == self.epilog_warp_id[0]: ++ cute.arch.dealloc_tmem( ++ acc_tmem_ptr, self.num_tmem_alloc_cols, is_two_cta=False ++ ) ++ # ++ # Wait for C store complete ++ # ++ c_pipeline.producer_tail() ++ ++ griddepcontrol_launch_dependents() ++ ++ def _compute_grid( ++ self, ++ c: cute.Tensor, ++ cta_tile_shape_mnk: Tuple[int, int, int], ++ ) -> Tuple[int, int, int]: ++ """Return one physical K-cluster for every output tile.""" ++ c_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0)) ++ gc = cute.zipped_divide(c, tiler=c_shape) ++ num_ctas_mnl = gc[(0, (None, None, None))].shape ++ return (num_ctas_mnl[0], num_ctas_mnl[1], self.split_k_slices) diff --git a/configs/nvidia-master.yaml b/configs/nvidia-master.yaml index 6be8e006f9..3f3fc6c4c5 100644 --- a/configs/nvidia-master.yaml +++ b/configs/nvidia-master.yaml @@ -13559,7 +13559,7 @@ minimaxm3-fp8-b300-vllm: # weights are pre-staged read-only at /scratch/models/MiniMax-M3-NVFP4 (added to # the STAGED_MODELS allow-list in launch_b300-nv.sh). minimaxm3-fp4-b300-vllm: - image: vllm/vllm-openai:nightly-93d8f834dd8acf33eb0e2a75b2711b628cb6e226 + image: vllm/vllm-openai:nightly-2afa3f7e950264bb179d030c23a1ed1f46558fd9 model: nvidia/MiniMax-M3-NVFP4 model-prefix: minimaxm3 runner: b300 diff --git a/perf-changelog.yaml b/perf-changelog.yaml index be28208ab4..2cd5c553ea 100644 --- a/perf-changelog.yaml +++ b/perf-changelog.yaml @@ -4716,3 +4716,11 @@ - "Clean the export envs" - "Enable two batch overlap" pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2093 + +- config-keys: + - minimaxm3-fp4-b300-vllm + description: + - "Update the vLLM image to nightly-2afa3f7e950264bb179d030c23a1ed1f46558fd9" + - "Install FlashInfer 0.6.15.dev20260710 with upstream SM100 low-M MXFP8 split-K support" + - "Remove the explicit FlashInfer TRT-LLM all-reduce backend override" + pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2124 From 0119a7681acd355d0e48e1a3abb9a2646fc8b443 Mon Sep 17 00:00:00 2001 From: jiahanc <173873397+jiahanc@users.noreply.github.com> Date: Wed, 8 Jul 2026 23:55:17 -0700 Subject: [PATCH 02/10] use bf16 route gemm MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 中文：路由 GEMM 使用 BF16 权重。 --- .../fixed_seq_len/minimaxm3_fp4_b300.sh | 15 +++++++++++++++ .../patches/vllm-minimaxm3-gate-bf16.patch | 11 +++++++++++ 2 files changed, 26 insertions(+) create mode 100644 benchmarks/single_node/fixed_seq_len/patches/vllm-minimaxm3-gate-bf16.patch diff --git a/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh b/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh index ad955e031f..67bb5314a1 100755 --- a/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh +++ b/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh @@ -57,6 +57,21 @@ check_env_vars \ # || { echo "FlashInfer CuTeDSL split-K gemm patch failed" >&2; exit 1; } # # ----------------------------------------------------------------------------- +# --- vLLM MiniMax-M3 GateLinear bf16 patch (vllm-project/vllm#47393) ---------- +VLLM_GATE_PATCH="$(dirname "$0")/patches/vllm-minimaxm3-gate-bf16.patch" +if ! command -v patch >/dev/null 2>&1; then + apt-get update -y && apt-get install -y --no-install-recommends patch \ + || { echo "Failed to install patch(1)" >&2; exit 1; } +fi +VLLM_PACKAGE_DIR=$(python3 -c "import importlib.util; print(importlib.util.find_spec(\"vllm\").submodule_search_locations[0])") \ + || { echo "Could not locate the installed vllm package" >&2; exit 1; } +VLLM_SITE_PACKAGES=$(dirname "$VLLM_PACKAGE_DIR") +patch --dry-run -p1 -d "$VLLM_SITE_PACKAGES" < "$VLLM_GATE_PATCH" >/dev/null \ + || { echo "vLLM MiniMax-M3 GateLinear bf16 patch does not apply" >&2; exit 1; } +patch -p1 -d "$VLLM_SITE_PACKAGES" < "$VLLM_GATE_PATCH" \ + || { echo "vLLM MiniMax-M3 GateLinear bf16 patch failed" >&2; exit 1; } +# ----------------------------------------------------------------------------- + if [[ -n "${MODEL_PATH:-}" ]]; then if [[ ! -d "$MODEL_PATH" || -z "$(ls -A "$MODEL_PATH" 2>/dev/null)" ]]; then hf download "$MODEL" --local-dir "$MODEL_PATH" diff --git a/benchmarks/single_node/fixed_seq_len/patches/vllm-minimaxm3-gate-bf16.patch b/benchmarks/single_node/fixed_seq_len/patches/vllm-minimaxm3-gate-bf16.patch new file mode 100644 index 0000000000..831480e370 --- /dev/null +++ b/benchmarks/single_node/fixed_seq_len/patches/vllm-minimaxm3-gate-bf16.patch @@ -0,0 +1,11 @@ +--- a/vllm/models/minimax_m3/nvidia/model.py ++++ b/vllm/models/minimax_m3/nvidia/model.py +@@ -228,7 +228,7 @@ + config.hidden_size, + config.num_local_experts, + bias=False, +- params_dtype=torch.float32, ++ params_dtype=torch.bfloat16, + out_dtype=torch.float32, + prefix=f"{prefix}.gate", + ) From 437c606b0c21b0e0210d55e5caea4472724cf73e Mon Sep 17 00:00:00 2001 From: jiahanc <173873397+jiahanc@users.noreply.github.com> Date: Thu, 9 Jul 2026 05:12:06 -0700 Subject: [PATCH 03/10] install nightly fi and remove bf16 route MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 中文：安装固定版本的 FlashInfer nightly，启用 CuTeDSL split-K 补丁，并移除 BF16 路由 GEMM 补丁。 --- .../fixed_seq_len/minimaxm3_fp4_b300.sh | 73 ++++++++----------- .../patches/vllm-minimaxm3-gate-bf16.patch | 11 --- 2 files changed, 30 insertions(+), 54 deletions(-) delete mode 100644 benchmarks/single_node/fixed_seq_len/patches/vllm-minimaxm3-gate-bf16.patch diff --git a/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh b/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh index 67bb5314a1..35e7cc486a 100755 --- a/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh +++ b/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh @@ -23,53 +23,40 @@ check_env_vars \ RANDOM_RANGE_RATIO \ RESULT_FILENAME -# # --- FlashInfer nightly + CuTeDSL split-K gemm patch ------------------------- -# # Pinned to nightly release nightly-v0.6.14-20260708. The wheels are not yet on -# # the flashinfer.ai/whl/nightly flat index, so install by direct release URL. -# FLASHINFER_VERSION=0.6.14.dev20260708 -# FLASHINFER_NIGHTLY_TAG=nightly-v0.6.14-20260708 -# FLASHINFER_RELEASE_URL="https://github.com/flashinfer-ai/flashinfer/releases/download/${FLASHINFER_NIGHTLY_TAG}" - -# # The flashinfer-jit-cache wheel is CUDA-version-specific; derive cuXYZ from -# # the torch build inside the container (e.g. 13.0 -> cu130). -# CUDA_MAJOR_MINOR="cu$(python3 -c 'import torch; print(torch.version.cuda.replace(".", ""))')" \ -# || { echo "Failed to determine CUDA version from torch" >&2; exit 1; } - -# python3 -m pip uninstall -y flashinfer-python flashinfer-cubin flashinfer-jit-cache - -# python3 -m pip install --pre \ -# "${FLASHINFER_RELEASE_URL}/flashinfer_python-${FLASHINFER_VERSION}-py3-none-any.whl" \ -# "${FLASHINFER_RELEASE_URL}/flashinfer_cubin-${FLASHINFER_VERSION}-py3-none-any.whl" \ -# "${FLASHINFER_RELEASE_URL}/flashinfer_jit_cache-${FLASHINFER_VERSION}+${CUDA_MAJOR_MINOR}-cp39-abi3-manylinux_2_28_$(uname -m).whl" \ -# || { echo "FlashInfer nightly install failed" >&2; exit 1; } - -# # Apply CuTeDSL split-K gemm patch (jiahanc/flashinfer branch cutedsl-splitK, -# # flashinfer/gemm changes only) on top of the pinned nightly. The reinstall -# # above guarantees pristine files, so the patch always applies cleanly. -# FLASHINFER_PATCH="$(dirname "$0")/patches/flashinfer-cutedsl-splitk-gemm.patch" -# if ! command -v patch >/dev/null 2>&1; then -# apt-get update -y && apt-get install -y --no-install-recommends patch \ -# || { echo "Failed to install patch(1)" >&2; exit 1; } -# fi -# SITE_PACKAGES=$(dirname "$(python3 -c "import importlib.util; print(importlib.util.find_spec('flashinfer').submodule_search_locations[0])")") \ -# || { echo "Could not locate the installed flashinfer package" >&2; exit 1; } -# patch -p1 -d "${SITE_PACKAGES}" < "$FLASHINFER_PATCH" \ -# || { echo "FlashInfer CuTeDSL split-K gemm patch failed" >&2; exit 1; } -# # ----------------------------------------------------------------------------- - -# --- vLLM MiniMax-M3 GateLinear bf16 patch (vllm-project/vllm#47393) ---------- -VLLM_GATE_PATCH="$(dirname "$0")/patches/vllm-minimaxm3-gate-bf16.patch" +# --- FlashInfer nightly + CuTeDSL split-K gemm patch ------------------------- +# Pinned to nightly release nightly-v0.6.14-20260708. The wheels are not yet on +# the flashinfer.ai/whl/nightly flat index, so install by direct release URL. +FLASHINFER_VERSION=0.6.14.dev20260708 +FLASHINFER_NIGHTLY_TAG=nightly-v0.6.14-20260708 +FLASHINFER_RELEASE_URL="https://github.com/flashinfer-ai/flashinfer/releases/download/${FLASHINFER_NIGHTLY_TAG}" + +# The flashinfer-jit-cache wheel is CUDA-version-specific; derive cuXYZ from +# the torch build inside the container (e.g. 13.0 -> cu130). +CUDA_MAJOR_MINOR="cu$(python3 -c 'import torch; print(torch.version.cuda.replace(".", ""))')" \ + || { echo "Failed to determine CUDA version from torch" >&2; exit 1; } + +python3 -m pip uninstall -y flashinfer-python flashinfer-cubin flashinfer-jit-cache + +python3 -m pip install --pre \ + "${FLASHINFER_RELEASE_URL}/flashinfer_python-${FLASHINFER_VERSION}-py3-none-any.whl" \ + "${FLASHINFER_RELEASE_URL}/flashinfer_cubin-${FLASHINFER_VERSION}-py3-none-any.whl" \ + "${FLASHINFER_RELEASE_URL}/flashinfer_jit_cache-${FLASHINFER_VERSION}+${CUDA_MAJOR_MINOR}-cp39-abi3-manylinux_2_28_$(uname -m).whl" \ + || { echo "FlashInfer nightly install failed" >&2; exit 1; } + +python -m pip install "nvidia-cutlass-dsl[cu13]>=4.5.2" + +# Apply CuTeDSL split-K gemm patch (jiahanc/flashinfer branch cutedsl-splitK, +# flashinfer/gemm changes only) on top of the pinned nightly. The reinstall +# above guarantees pristine files, so the patch always applies cleanly. +FLASHINFER_PATCH="$(dirname "$0")/patches/flashinfer-cutedsl-splitk-gemm.patch" if ! command -v patch >/dev/null 2>&1; then apt-get update -y && apt-get install -y --no-install-recommends patch \ || { echo "Failed to install patch(1)" >&2; exit 1; } fi -VLLM_PACKAGE_DIR=$(python3 -c "import importlib.util; print(importlib.util.find_spec(\"vllm\").submodule_search_locations[0])") \ - || { echo "Could not locate the installed vllm package" >&2; exit 1; } -VLLM_SITE_PACKAGES=$(dirname "$VLLM_PACKAGE_DIR") -patch --dry-run -p1 -d "$VLLM_SITE_PACKAGES" < "$VLLM_GATE_PATCH" >/dev/null \ - || { echo "vLLM MiniMax-M3 GateLinear bf16 patch does not apply" >&2; exit 1; } -patch -p1 -d "$VLLM_SITE_PACKAGES" < "$VLLM_GATE_PATCH" \ - || { echo "vLLM MiniMax-M3 GateLinear bf16 patch failed" >&2; exit 1; } +SITE_PACKAGES=$(dirname "$(python3 -c "import importlib.util; print(importlib.util.find_spec('flashinfer').submodule_search_locations[0])")") \ + || { echo "Could not locate the installed flashinfer package" >&2; exit 1; } +patch -p1 -d "${SITE_PACKAGES}" < "$FLASHINFER_PATCH" \ + || { echo "FlashInfer CuTeDSL split-K gemm patch failed" >&2; exit 1; } # ----------------------------------------------------------------------------- if [[ -n "${MODEL_PATH:-}" ]]; then diff --git a/benchmarks/single_node/fixed_seq_len/patches/vllm-minimaxm3-gate-bf16.patch b/benchmarks/single_node/fixed_seq_len/patches/vllm-minimaxm3-gate-bf16.patch deleted file mode 100644 index 831480e370..0000000000 --- a/benchmarks/single_node/fixed_seq_len/patches/vllm-minimaxm3-gate-bf16.patch +++ /dev/null @@ -1,11 +0,0 @@ ---- a/vllm/models/minimax_m3/nvidia/model.py -+++ b/vllm/models/minimax_m3/nvidia/model.py -@@ -228,7 +228,7 @@ - config.hidden_size, - config.num_local_experts, - bias=False, -- params_dtype=torch.float32, -+ params_dtype=torch.bfloat16, - out_dtype=torch.float32, - prefix=f"{prefix}.gate", - ) From dbd30885de4458584a877e70f85eccc4c9fb231d Mon Sep 17 00:00:00 2001 From: jiahanc <173873397+jiahanc@users.noreply.github.com> Date: Thu, 9 Jul 2026 07:42:44 -0700 Subject: [PATCH 04/10] use stable flashinfer 0.6.14 cu13 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 中文：改用 FlashInfer 0.6.14 的 CUDA 13 extra，并删除单独的 CUTLASS DSL 安装步骤。 --- .../fixed_seq_len/minimaxm3_fp4_b300.sh | 30 +++++-------------- 1 file changed, 8 insertions(+), 22 deletions(-) diff --git a/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh b/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh index 35e7cc486a..7d850788a8 100755 --- a/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh +++ b/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh @@ -5,9 +5,9 @@ # checkpoint. MiniMax-M3 modelopt NVFP4 support (vllm-project/vllm PR #46380) is # baked into the perf container image. # -# At runtime the recipe swaps the image's FlashInfer for the pinned nightly -# release and applies the CuTeDSL split-K gemm patch on top, mirroring vLLM -# commits 82a00090a (nightly install) and 0a16bea6b (patch). +# At runtime the recipe swaps the image's FlashInfer for the pinned stable +# 0.6.14 release with CUDA 13 dependencies and applies the CuTeDSL split-K gemm +# patch on top, mirroring vLLM commit 0a16bea6b. source "$(dirname "$0")/../../benchmark_lib.sh" @@ -23,30 +23,16 @@ check_env_vars \ RANDOM_RANGE_RATIO \ RESULT_FILENAME -# --- FlashInfer nightly + CuTeDSL split-K gemm patch ------------------------- -# Pinned to nightly release nightly-v0.6.14-20260708. The wheels are not yet on -# the flashinfer.ai/whl/nightly flat index, so install by direct release URL. -FLASHINFER_VERSION=0.6.14.dev20260708 -FLASHINFER_NIGHTLY_TAG=nightly-v0.6.14-20260708 -FLASHINFER_RELEASE_URL="https://github.com/flashinfer-ai/flashinfer/releases/download/${FLASHINFER_NIGHTLY_TAG}" - -# The flashinfer-jit-cache wheel is CUDA-version-specific; derive cuXYZ from -# the torch build inside the container (e.g. 13.0 -> cu130). -CUDA_MAJOR_MINOR="cu$(python3 -c 'import torch; print(torch.version.cuda.replace(".", ""))')" \ - || { echo "Failed to determine CUDA version from torch" >&2; exit 1; } +# --- FlashInfer 0.6.14 + CuTeDSL split-K gemm patch -------------------------- +FLASHINFER_VERSION=0.6.14 python3 -m pip uninstall -y flashinfer-python flashinfer-cubin flashinfer-jit-cache -python3 -m pip install --pre \ - "${FLASHINFER_RELEASE_URL}/flashinfer_python-${FLASHINFER_VERSION}-py3-none-any.whl" \ - "${FLASHINFER_RELEASE_URL}/flashinfer_cubin-${FLASHINFER_VERSION}-py3-none-any.whl" \ - "${FLASHINFER_RELEASE_URL}/flashinfer_jit_cache-${FLASHINFER_VERSION}+${CUDA_MAJOR_MINOR}-cp39-abi3-manylinux_2_28_$(uname -m).whl" \ - || { echo "FlashInfer nightly install failed" >&2; exit 1; } - -python -m pip install "nvidia-cutlass-dsl[cu13]>=4.5.2" +python -m pip install "flashinfer-python[cu13]==${FLASHINFER_VERSION}" \ + || { echo "FlashInfer ${FLASHINFER_VERSION} install failed" >&2; exit 1; } # Apply CuTeDSL split-K gemm patch (jiahanc/flashinfer branch cutedsl-splitK, -# flashinfer/gemm changes only) on top of the pinned nightly. The reinstall +# flashinfer/gemm changes only) on top of the pinned release. The reinstall # above guarantees pristine files, so the patch always applies cleanly. FLASHINFER_PATCH="$(dirname "$0")/patches/flashinfer-cutedsl-splitk-gemm.patch" if ! command -v patch >/dev/null 2>&1; then From bfae28ae5a1dbbd59a4d06f0001c7e56d1a0c4fe Mon Sep 17 00:00:00 2001 From: jiahanc <173873397+jiahanc@users.noreply.github.com> Date: Thu, 9 Jul 2026 07:47:51 -0700 Subject: [PATCH 05/10] fix: install matching FlashInfer 0.6.14 runtime packages MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 中文：改用 python3，并安装相互匹配的 FlashInfer 0.6.14 core、cubin 和 CUDA 13.0 JIT cache。 --- .../single_node/fixed_seq_len/minimaxm3_fp4_b300.sh | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh b/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh index 7d850788a8..fed69ecd55 100755 --- a/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh +++ b/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh @@ -28,9 +28,17 @@ FLASHINFER_VERSION=0.6.14 python3 -m pip uninstall -y flashinfer-python flashinfer-cubin flashinfer-jit-cache -python -m pip install "flashinfer-python[cu13]==${FLASHINFER_VERSION}" \ +python3 -m pip install "flashinfer-python[cu13]==${FLASHINFER_VERSION}" \ || { echo "FlashInfer ${FLASHINFER_VERSION} install failed" >&2; exit 1; } +python3 -m pip install "flashinfer-cubin==${FLASHINFER_VERSION}" \ + --index-url https://flashinfer.ai/whl \ + || { echo "FlashInfer cubin ${FLASHINFER_VERSION} install failed" >&2; exit 1; } + +python3 -m pip install "flashinfer-jit-cache==${FLASHINFER_VERSION}+cu130" \ + --index-url https://flashinfer.ai/whl/cu130 \ + || { echo "FlashInfer JIT cache ${FLASHINFER_VERSION}+cu130 install failed" >&2; exit 1; } + # Apply CuTeDSL split-K gemm patch (jiahanc/flashinfer branch cutedsl-splitK, # flashinfer/gemm changes only) on top of the pinned release. The reinstall # above guarantees pristine files, so the patch always applies cleanly. From 756a7d937911b9fecaec325d681ac372fe428d23 Mon Sep 17 00:00:00 2001 From: jiahanc <173873397+jiahanc@users.noreply.github.com> Date: Thu, 9 Jul 2026 07:57:51 -0700 Subject: [PATCH 06/10] no-deps MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 中文：为 FlashInfer core、cubin 和 JIT cache 安装统一添加 --no-deps，避免重新解析镜像依赖。 --- benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh b/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh index fed69ecd55..6b3b9bcf0c 100755 --- a/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh +++ b/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh @@ -28,14 +28,14 @@ FLASHINFER_VERSION=0.6.14 python3 -m pip uninstall -y flashinfer-python flashinfer-cubin flashinfer-jit-cache -python3 -m pip install "flashinfer-python[cu13]==${FLASHINFER_VERSION}" \ +python3 -m pip install --no-deps "flashinfer-python[cu13]==${FLASHINFER_VERSION}" \ || { echo "FlashInfer ${FLASHINFER_VERSION} install failed" >&2; exit 1; } -python3 -m pip install "flashinfer-cubin==${FLASHINFER_VERSION}" \ +python3 -m pip install --no-deps "flashinfer-cubin==${FLASHINFER_VERSION}" \ --index-url https://flashinfer.ai/whl \ || { echo "FlashInfer cubin ${FLASHINFER_VERSION} install failed" >&2; exit 1; } -python3 -m pip install "flashinfer-jit-cache==${FLASHINFER_VERSION}+cu130" \ +python3 -m pip install --no-deps "flashinfer-jit-cache==${FLASHINFER_VERSION}+cu130" \ --index-url https://flashinfer.ai/whl/cu130 \ || { echo "FlashInfer JIT cache ${FLASHINFER_VERSION}+cu130 install failed" >&2; exit 1; } From db31e46a621341ab7134c404a4bfa5b4adba55e5 Mon Sep 17 00:00:00 2001 From: jiahanc <173873397+jiahanc@users.noreply.github.com> Date: Thu, 9 Jul 2026 17:40:20 -0700 Subject: [PATCH 07/10] revert fi upgrade Signed-off-by: jiahanc <173873397+jiahanc@users.noreply.github.com> --- .../fixed_seq_len/minimaxm3_fp4_b300.sh | 22 +++++++++---------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh b/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh index 6b3b9bcf0c..ce6ba5f43f 100755 --- a/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh +++ b/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh @@ -23,21 +23,21 @@ check_env_vars \ RANDOM_RANGE_RATIO \ RESULT_FILENAME -# --- FlashInfer 0.6.14 + CuTeDSL split-K gemm patch -------------------------- -FLASHINFER_VERSION=0.6.14 +# # --- FlashInfer 0.6.14 + CuTeDSL split-K gemm patch -------------------------- +# FLASHINFER_VERSION=0.6.14 -python3 -m pip uninstall -y flashinfer-python flashinfer-cubin flashinfer-jit-cache +# python3 -m pip uninstall -y flashinfer-python flashinfer-cubin flashinfer-jit-cache -python3 -m pip install --no-deps "flashinfer-python[cu13]==${FLASHINFER_VERSION}" \ - || { echo "FlashInfer ${FLASHINFER_VERSION} install failed" >&2; exit 1; } +# python3 -m pip install --no-deps "flashinfer-python[cu13]==${FLASHINFER_VERSION}" \ +# || { echo "FlashInfer ${FLASHINFER_VERSION} install failed" >&2; exit 1; } -python3 -m pip install --no-deps "flashinfer-cubin==${FLASHINFER_VERSION}" \ - --index-url https://flashinfer.ai/whl \ - || { echo "FlashInfer cubin ${FLASHINFER_VERSION} install failed" >&2; exit 1; } +# python3 -m pip install --no-deps "flashinfer-cubin==${FLASHINFER_VERSION}" \ +# --index-url https://flashinfer.ai/whl \ +# || { echo "FlashInfer cubin ${FLASHINFER_VERSION} install failed" >&2; exit 1; } -python3 -m pip install --no-deps "flashinfer-jit-cache==${FLASHINFER_VERSION}+cu130" \ - --index-url https://flashinfer.ai/whl/cu130 \ - || { echo "FlashInfer JIT cache ${FLASHINFER_VERSION}+cu130 install failed" >&2; exit 1; } +# python3 -m pip install --no-deps "flashinfer-jit-cache==${FLASHINFER_VERSION}+cu130" \ +# --index-url https://flashinfer.ai/whl/cu130 \ +# || { echo "FlashInfer JIT cache ${FLASHINFER_VERSION}+cu130 install failed" >&2; exit 1; } # Apply CuTeDSL split-K gemm patch (jiahanc/flashinfer branch cutedsl-splitK, # flashinfer/gemm changes only) on top of the pinned release. The reinstall From 81f3cbf0dd622bb2f86d9e7d4630fb4d602add4b Mon Sep 17 00:00:00 2001 From: jiahanc <173873397+jiahanc@users.noreply.github.com> Date: Fri, 10 Jul 2026 06:55:36 -0700 Subject: [PATCH 08/10] fix: use upstream FlashInfer split-K nightly MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 中文：升级到包含上游 SM100 low-M MXFP8 split-K 的 FlashInfer nightly，删除本地 split-K 补丁，并移除显式的 TRT-LLM all-reduce backend 配置。 --- .../fixed_seq_len/minimaxm3_fp4_b300.sh | 42 +- .../flashinfer-cutedsl-splitk-gemm.patch | 3115 ----------------- 2 files changed, 13 insertions(+), 3144 deletions(-) delete mode 100644 benchmarks/single_node/fixed_seq_len/patches/flashinfer-cutedsl-splitk-gemm.patch diff --git a/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh b/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh index ce6ba5f43f..5b732b7f5e 100755 --- a/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh +++ b/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh @@ -5,9 +5,9 @@ # checkpoint. MiniMax-M3 modelopt NVFP4 support (vllm-project/vllm PR #46380) is # baked into the perf container image. # -# At runtime the recipe swaps the image's FlashInfer for the pinned stable -# 0.6.14 release with CUDA 13 dependencies and applies the CuTeDSL split-K gemm -# patch on top, mirroring vLLM commit 0a16bea6b. +# At runtime the recipe swaps the image's FlashInfer for the first pinned +# nightly containing the upstream SM100 low-M MXFP8 split-K kernel +# (flashinfer-ai/flashinfer#3847). source "$(dirname "$0")/../../benchmark_lib.sh" @@ -23,34 +23,19 @@ check_env_vars \ RANDOM_RANGE_RATIO \ RESULT_FILENAME -# # --- FlashInfer 0.6.14 + CuTeDSL split-K gemm patch -------------------------- -# FLASHINFER_VERSION=0.6.14 +# --- FlashInfer nightly ------------------------------------------------------ +FLASHINFER_VERSION=0.6.15.dev20260710 +FLASHINFER_NIGHTLY_TAG=nightly-v0.6.15-20260710 +FLASHINFER_RELEASE_URL="https://github.com/flashinfer-ai/flashinfer/releases/download/${FLASHINFER_NIGHTLY_TAG}" -# python3 -m pip uninstall -y flashinfer-python flashinfer-cubin flashinfer-jit-cache +python3 -m pip uninstall -y flashinfer-python flashinfer-cubin flashinfer-jit-cache -# python3 -m pip install --no-deps "flashinfer-python[cu13]==${FLASHINFER_VERSION}" \ -# || { echo "FlashInfer ${FLASHINFER_VERSION} install failed" >&2; exit 1; } +python3 -m pip install \ + "${FLASHINFER_RELEASE_URL}/flashinfer_python-${FLASHINFER_VERSION}-py3-none-any.whl" \ + "${FLASHINFER_RELEASE_URL}/flashinfer_cubin-${FLASHINFER_VERSION}-py3-none-any.whl" \ + "${FLASHINFER_RELEASE_URL}/flashinfer_jit_cache-${FLASHINFER_VERSION}+cu130-cp39-abi3-manylinux_2_28_$(uname -m).whl" \ + || { echo "FlashInfer nightly install failed" >&2; exit 1; } -# python3 -m pip install --no-deps "flashinfer-cubin==${FLASHINFER_VERSION}" \ -# --index-url https://flashinfer.ai/whl \ -# || { echo "FlashInfer cubin ${FLASHINFER_VERSION} install failed" >&2; exit 1; } - -# python3 -m pip install --no-deps "flashinfer-jit-cache==${FLASHINFER_VERSION}+cu130" \ -# --index-url https://flashinfer.ai/whl/cu130 \ -# || { echo "FlashInfer JIT cache ${FLASHINFER_VERSION}+cu130 install failed" >&2; exit 1; } - -# Apply CuTeDSL split-K gemm patch (jiahanc/flashinfer branch cutedsl-splitK, -# flashinfer/gemm changes only) on top of the pinned release. The reinstall -# above guarantees pristine files, so the patch always applies cleanly. -FLASHINFER_PATCH="$(dirname "$0")/patches/flashinfer-cutedsl-splitk-gemm.patch" -if ! command -v patch >/dev/null 2>&1; then - apt-get update -y && apt-get install -y --no-install-recommends patch \ - || { echo "Failed to install patch(1)" >&2; exit 1; } -fi -SITE_PACKAGES=$(dirname "$(python3 -c "import importlib.util; print(importlib.util.find_spec('flashinfer').submodule_search_locations[0])")") \ - || { echo "Could not locate the installed flashinfer package" >&2; exit 1; } -patch -p1 -d "${SITE_PACKAGES}" < "$FLASHINFER_PATCH" \ - || { echo "FlashInfer CuTeDSL split-K gemm patch failed" >&2; exit 1; } # ----------------------------------------------------------------------------- if [[ -n "${MODEL_PATH:-}" ]]; then @@ -72,7 +57,6 @@ SERVER_LOG=/workspace/server.log export VLLM_ENGINE_READY_TIMEOUT_S=3600 export VLLM_FLOAT32_MATMUL_PRECISION=high -export VLLM_FLASHINFER_ALLREDUCE_BACKEND=trtllm export VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=1800 if [ "${DP_ATTENTION}" = "true" ]; then diff --git a/benchmarks/single_node/fixed_seq_len/patches/flashinfer-cutedsl-splitk-gemm.patch b/benchmarks/single_node/fixed_seq_len/patches/flashinfer-cutedsl-splitk-gemm.patch deleted file mode 100644 index 965114c94e..0000000000 --- a/benchmarks/single_node/fixed_seq_len/patches/flashinfer-cutedsl-splitk-gemm.patch +++ /dev/null @@ -1,3115 +0,0 @@ -diff --git a/flashinfer/gemm/gemm_base.py b/flashinfer/gemm/gemm_base.py -index 1a665cfd..c95bfaa4 100755 ---- a/flashinfer/gemm/gemm_base.py -+++ b/flashinfer/gemm/gemm_base.py -@@ -15,9 +15,10 @@ limitations under the License. - """ - - import functools -+from dataclasses import replace - from enum import Enum - from types import SimpleNamespace --from typing import List, Literal, Optional, Tuple -+from typing import Callable, List, Literal, Optional, Tuple - - from flashinfer.trtllm_low_latency_gemm import trtllm_low_latency_gemm - import torch -@@ -4606,6 +4607,7 @@ def _compile_block_scaled_gemm( - sf_n, - sf_k, - batch_size, -+ cluster_shape_k=1, - ): - """Compile a block-scaled GEMM kernel via CuTe DSL and cache it. - -@@ -4671,9 +4673,8 @@ def _compile_block_scaled_gemm( - cutlass.Float32, (1,), assumed_align=4 - ) - -- max_active_clusters = get_max_active_clusters( -- cluster_shape_mn[0] * cluster_shape_mn[1] -- ) -+ launch_cluster_size = cluster_shape_mn[0] * cluster_shape_mn[1] * cluster_shape_k -+ max_active_clusters = get_max_active_clusters(launch_cluster_size) - stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) - - compiled_gemm = cute.compile( -@@ -4721,6 +4722,35 @@ def _prepare_alpha_for_launch(alpha_tensor, device): - - _CUTE_DSL_MM_MXFP8_KERNEL_CACHE: dict[tuple, tuple] = {} - -+_SM100_MXFP8_SPLIT_K_MAX_M = 32 -+_SM100_MXFP8_SPLIT_K_MMA_INST_BITS_K = 256 -+_SM100_MXFP8_SPLIT_K_MMA_INST_TILE_K = 4 -+_SM100_MXFP8_SPLIT_K_CANDIDATES = (2, 4) -+ -+ -+def _is_sm100_split_k_valid( -+ m: int, -+ k: int, -+ ab_dtype, -+ split_k_slices: int, -+) -> bool: -+ """Check the specialized low-M split-K shape constraints.""" -+ if ab_dtype.width != 8: -+ return False -+ tile_k = ( -+ _SM100_MXFP8_SPLIT_K_MMA_INST_BITS_K * _SM100_MXFP8_SPLIT_K_MMA_INST_TILE_K -+ ) // ab_dtype.width -+ return ( -+ m <= _SM100_MXFP8_SPLIT_K_MAX_M -+ and split_k_slices in _SM100_MXFP8_SPLIT_K_CANDIDATES -+ and k % (tile_k * split_k_slices) == 0 -+ ) -+ -+ -+def _sm100_mxfp8_low_m_tile(m: int) -> tuple[int, int]: -+ tile_n = 8 if m <= 8 else 16 if m <= 16 else 32 -+ return (128, tile_n) -+ - - def _check_cute_dsl_availability(): - try: -@@ -4743,6 +4773,9 @@ def _cute_dsl_gemm_mxfp8_runner( - from .kernels.dense_blockscaled_gemm_sm100 import ( - Sm100BlockScaledPersistentDenseGemmKernel, - ) -+ from .kernels.dense_blockscaled_gemm_sm100_splitk import ( -+ Sm100BlockScaledSplitKGemmKernel, -+ ) - - if out_dtype not in (torch.bfloat16, torch.float16): - raise ValueError( -@@ -4760,22 +4793,51 @@ def _cute_dsl_gemm_mxfp8_runner( - _ = sm_major, sm_minor - - class CuteDSLMxfp8GemmRunner(TunableRunner): -+ def get_cache_key_extras(self, inputs: List[torch.Tensor]) -> tuple: -+ _, _, _, _, _, out, _ = inputs -+ return (str(out.dtype), enable_pdl) -+ - def get_valid_tactics( - self, - inputs: List[torch.Tensor], - profile: OptimizationProfile, - ) -> list: - (a, b, a_descale, b_descale, _, out, _) = inputs -- return _get_sm100_block_scaled_tactics( -- m=a.shape[0], -- n=b.shape[1], -- real_k=a.shape[1], -- ab_dtype=cutlass.Float8E4M3FN, -+ m = a.shape[0] -+ n = b.shape[1] -+ real_k = a.shape[1] -+ ab_dtype = cutlass.Float8E4M3FN -+ base_tactics = _get_sm100_block_scaled_tactics( -+ m=m, -+ n=n, -+ real_k=real_k, -+ ab_dtype=ab_dtype, - sf_dtype=cutlass.Float8E8M0FNU, - sf_vec_size=32, - c_cutlass_dtype=c_cutlass_dtype, - device=a.device, - ) -+ valid_tactics = [(*tactic, 1) for tactic in base_tactics] -+ if not out.is_contiguous(): -+ return valid_tactics -+ -+ for split_k_slices in _SM100_MXFP8_SPLIT_K_CANDIDATES: -+ if _is_sm100_split_k_valid( -+ m, -+ real_k, -+ ab_dtype, -+ split_k_slices, -+ ): -+ valid_tactics.append( -+ ( -+ _sm100_mxfp8_low_m_tile(m), -+ (1, 1), -+ True, -+ False, -+ split_k_slices, -+ ) -+ ) -+ return valid_tactics - - def forward( - self, -@@ -4794,14 +4856,50 @@ def _cute_dsl_gemm_mxfp8_runner( - batch_size = 1 - - if tactic is None or tactic == -1: -- tactic = ( -- _SM100_DEFAULT_MMA_TILER_MN, -- _SM100_DEFAULT_CLUSTER_SHAPE_MN, -- False, -- False, -- ) -+ if m <= _SM100_MXFP8_SPLIT_K_MAX_M and out.is_contiguous(): -+ # Untuned low-M execution uses the corresponding base tactic. -+ tactic = ( -+ _sm100_mxfp8_low_m_tile(m), -+ (1, 1), -+ True, -+ False, -+ 1, -+ ) -+ else: -+ tactic = ( -+ _SM100_DEFAULT_MMA_TILER_MN, -+ _SM100_DEFAULT_CLUSTER_SHAPE_MN, -+ False, -+ False, -+ 1, -+ ) - -- (mma_tiler_mn, cluster_shape_mn, swap_ab, use_prefetch) = tactic -+ ( -+ mma_tiler_mn, -+ cluster_shape_mn, -+ swap_ab, -+ use_prefetch, -+ split_k_slices, -+ ) = tactic -+ if split_k_slices < 1: -+ raise ValueError(f"Invalid MXFP8 split-K tactic: {tactic}") -+ is_split_k = split_k_slices > 1 -+ -+ if is_split_k: -+ if ( -+ cluster_shape_mn != (1, 1) -+ or not swap_ab -+ or use_prefetch -+ or not out.is_contiguous() -+ or mma_tiler_mn != _sm100_mxfp8_low_m_tile(m) -+ or not _is_sm100_split_k_valid( -+ m, -+ real_k, -+ cutlass.Float8E4M3FN, -+ split_k_slices, -+ ) -+ ): -+ raise ValueError(f"Invalid MXFP8 split-K tactic: {tactic}") - - if swap_ab: - kernel_m, kernel_n = n, m -@@ -4827,18 +4925,31 @@ def _cute_dsl_gemm_mxfp8_runner( - use_prefetch, - enable_pdl, - out_dtype, -+ split_k_slices, - ) - -- compiled_gemm, _ = _compile_block_scaled_gemm( -- _CUTE_DSL_MM_MXFP8_KERNEL_CACHE, -- cache_key, -- lambda: Sm100BlockScaledPersistentDenseGemmKernel( -+ make_kernel: Callable[[], object] -+ kernel_cache = _CUTE_DSL_MM_MXFP8_KERNEL_CACHE -+ if is_split_k: -+ make_kernel = lambda: Sm100BlockScaledSplitKGemmKernel( -+ sf_vec_size, -+ mma_tiler_mn, -+ split_k_slices, -+ enable_pdl, -+ ) -+ else: -+ make_kernel = lambda: Sm100BlockScaledPersistentDenseGemmKernel( - sf_vec_size, - mma_tiler_mn, - cluster_shape_mn, - use_prefetch, - enable_pdl, -- ), -+ ) -+ -+ compiled_gemm, _ = _compile_block_scaled_gemm( -+ kernel_cache, -+ cache_key, -+ make_kernel, - ab_cutlass_dtype=cutlass.Float8E4M3FN, - sf_dtype=sf_dtype, - c_cutlass_dtype=c_cutlass_dtype, -@@ -4849,6 +4960,7 @@ def _cute_dsl_gemm_mxfp8_runner( - sf_n=sf_n, - sf_k=sf_k, - batch_size=batch_size, -+ cluster_shape_k=split_k_slices, - ) - - alpha_for_launch = _prepare_alpha_for_launch(None, a.device) -@@ -4856,6 +4968,7 @@ def _cute_dsl_gemm_mxfp8_runner( - launch_out = ( - out.as_strided(out.shape, (1, out.shape[0])) if swap_ab else out - ) -+ - compiled_gemm( - kernel_a, - kernel_b, -@@ -5163,7 +5276,11 @@ def mm_mxfp8( - - tuner = AutoTuner.get() - -- tuning_config = _MM_MXFP8_TUNING_CONFIG -+ tuning_config = ( -+ _MM_MXFP8_CUTE_DSL_TUNING_CONFIG -+ if backends == ["cute-dsl"] -+ else _MM_MXFP8_TUNING_CONFIG -+ ) - - inputs = [ - a, -@@ -6323,6 +6440,12 @@ _MM_MXFP8_TUNING_CONFIG = TuningConfig( - ), - ) - -+_MM_MXFP8_CUTE_DSL_TUNING_CONFIG = replace( -+ _MM_MXFP8_TUNING_CONFIG, -+ use_cuda_graph=True, -+ use_cold_l2_cache=True, -+) -+ - - @backend_requirement( - { -diff --git a/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm100.py b/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm100.py -index 85a006d2..f653ea8a 100644 ---- a/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm100.py -+++ b/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm100.py -@@ -41,11 +41,13 @@ import cutlass.utils.blackwell_helpers as sm100_utils - import cutlass.utils.blockscaled_layout as blockscaled_utils - from cutlass.cute.nvgpu import cpasync, tcgen05 - -+from .dense_blockscaled_gemm_sm100_common import _Sm100BlockScaledGemmCommon -+ - from cutlass.cute.arch import griddepcontrol_launch_dependents, griddepcontrol_wait - from cutlass.pipeline import PipelineTmaUmma, PipelineUmmaAsync - - --class Sm100BlockScaledPersistentDenseGemmKernel: -+class Sm100BlockScaledPersistentDenseGemmKernel(_Sm100BlockScaledGemmCommon): - """Implements batched matrix multiplication (C = A x SFA x B x SFB) with support for various data types - and Blackwell GPU architectural features, including persistent tile scheduling and warp specialization. - -@@ -254,6 +256,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel: - self.sf_vec_size, - self.smem_capacity, - self.occupancy, -+ 1024, - ) - - # Compute A/B/SFA/SFB/C shared memory layout -@@ -344,155 +347,25 @@ class Sm100BlockScaledPersistentDenseGemmKernel: - Raises: - TypeError: If input data types are incompatible with the MMA instruction. - """ -- # Setup static attributes before smem/grid/tma computation -- self.a_dtype: Type[cutlass.Numeric] = a_tensor.element_type -- self.b_dtype: Type[cutlass.Numeric] = b_tensor.element_type -- self.sf_dtype: Type[cutlass.Numeric] = sfa_tensor.element_type -- self.c_dtype: Type[cutlass.Numeric] = c_tensor.element_type -- self.a_major_mode = utils.LayoutEnum.from_tensor(a_tensor).mma_major_mode() -- self.b_major_mode = utils.LayoutEnum.from_tensor(b_tensor).mma_major_mode() -- self.c_layout = utils.LayoutEnum.from_tensor(c_tensor) -- -- # Check if input data types are compatible with MMA instruction -- if cutlass.const_expr(self.a_dtype != self.b_dtype): -- raise TypeError(f"Type must match: {self.a_dtype} != {self.b_dtype}") -- -- # Setup attributes that dependent on gemm inputs -- self._setup_attributes() -- -- # Setup sfa/sfb tensor by filling A/B tensor to scale factor atom layout -- # ((Atom_M, Rest_M),(Atom_K, Rest_K),RestL) -- sfa_layout = blockscaled_utils.tile_atom_to_shape_SF( -- a_tensor.shape, self.sf_vec_size -- ) -- sfa_tensor = cute.make_tensor(sfa_tensor.iterator, sfa_layout) -- -- # ((Atom_N, Rest_N),(Atom_K, Rest_K),RestL) -- sfb_layout = blockscaled_utils.tile_atom_to_shape_SF( -- b_tensor.shape, self.sf_vec_size -- ) -- sfb_tensor = cute.make_tensor(sfb_tensor.iterator, sfb_layout) -- -- tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma( -- self.a_dtype, -- self.b_dtype, -- self.a_major_mode, -- self.b_major_mode, -- self.sf_dtype, -- self.sf_vec_size, -- self.cta_group, -- self.mma_inst_shape_mnk[:2], -- ) -- -- tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma( -- self.a_dtype, -- self.b_dtype, -- self.a_major_mode, -- self.b_major_mode, -- self.sf_dtype, -- self.sf_vec_size, -- cute.nvgpu.tcgen05.CtaGroup.ONE, -- self.mma_inst_shape_mnk_sfb[:2], -- ) -- atom_thr_size = cute.size(tiled_mma.thr_id.shape) -- -- # Setup TMA load for A -- a_op = sm100_utils.cluster_shape_to_tma_atom_A( -- self.cluster_shape_mn, tiled_mma.thr_id -- ) -- a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) -- tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( -- a_op, -- a_tensor, -- a_smem_layout, -- self.mma_tiler, -+ ( - tiled_mma, -- self.cluster_layout_vmnk.shape, -- ) -- -- # Setup TMA load for B -- b_op = sm100_utils.cluster_shape_to_tma_atom_B( -- self.cluster_shape_mn, tiled_mma.thr_id -- ) -- b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) -- tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( -- b_op, -+ tiled_mma_sfb, -+ tma_atom_a, -+ tma_tensor_a, -+ tma_atom_b, -+ tma_tensor_b, -+ tma_atom_sfa, -+ tma_tensor_sfa, -+ tma_atom_sfb, -+ tma_tensor_sfb, -+ tma_atom_c, -+ tma_tensor_c, -+ ) = self._prepare_tma( -+ a_tensor, - b_tensor, -- b_smem_layout, -- self.mma_tiler, -- tiled_mma, -- self.cluster_layout_vmnk.shape, -- ) -- -- # Setup TMA load for SFA -- sfa_op = sm100_utils.cluster_shape_to_tma_atom_A( -- self.cluster_shape_mn, tiled_mma.thr_id -- ) -- sfa_smem_layout = cute.slice_( -- self.sfa_smem_layout_staged, (None, None, None, 0) -- ) -- tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A( -- sfa_op, - sfa_tensor, -- sfa_smem_layout, -- self.mma_tiler, -- tiled_mma, -- self.cluster_layout_vmnk.shape, -- internal_type=cutlass.Int16, -- ) -- -- # Setup TMA load for SFB -- sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB( -- self.cluster_shape_mn, tiled_mma.thr_id -- ) -- sfb_smem_layout = cute.slice_( -- self.sfb_smem_layout_staged, (None, None, None, 0) -- ) -- tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B( -- sfb_op, - sfb_tensor, -- sfb_smem_layout, -- self.mma_tiler_sfb, -- tiled_mma_sfb, -- self.cluster_layout_sfb_vmnk.shape, -- internal_type=cutlass.Int16, -- ) -- if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 192): -- x = tma_tensor_sfb.stride[0][1] -- y = cute.ceil_div(tma_tensor_sfb.shape[0][1], 4) -- -- new_shape = ( -- (tma_tensor_sfb.shape[0][0], ((2, 2), y)), -- tma_tensor_sfb.shape[1], -- tma_tensor_sfb.shape[2], -- ) -- # Use right multiplication for ScaledBasis (3 * x instead of x * 3) -- x_times_3 = 3 * x -- new_stride = ( -- (tma_tensor_sfb.stride[0][0], ((x, x), x_times_3)), -- tma_tensor_sfb.stride[1], -- tma_tensor_sfb.stride[2], -- ) -- tma_tensor_sfb_new_layout = cute.make_layout(new_shape, stride=new_stride) -- tma_tensor_sfb = cute.make_tensor( -- tma_tensor_sfb.iterator, tma_tensor_sfb_new_layout -- ) -- -- a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout) -- b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout) -- sfa_copy_size = cute.size_in_bytes(self.sf_dtype, sfa_smem_layout) -- sfb_copy_size = cute.size_in_bytes(self.sf_dtype, sfb_smem_layout) -- self.num_tma_load_bytes = ( -- a_copy_size + b_copy_size + sfa_copy_size + sfb_copy_size -- ) * atom_thr_size -- -- # Setup TMA store for C -- epi_smem_layout = cute.slice_(self.c_smem_layout_staged, (None, None, 0)) -- tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( -- cpasync.CopyBulkTensorTileS2GOp(), - c_tensor, -- epi_smem_layout, -- self.epi_tile, - ) - - # Compute grid size -@@ -1537,295 +1410,6 @@ class Sm100BlockScaledPersistentDenseGemmKernel: - - griddepcontrol_launch_dependents() - -- def mainloop_s2t_copy_and_partition( -- self, -- sSF: cute.Tensor, -- tSF: cute.Tensor, -- ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: -- """ -- Make tiledCopy for smem to tmem load for scale factor tensor, then use it to partition smem memory (source) and tensor memory (destination). -- -- Args: -- sSF (cute.Tensor): The scale factor tensor in smem -- tSF (cute.Tensor): The scale factor tensor in tmem -- -- Returns: -- A tuple containing (tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t) where: -- - tiled_copy_s2t: The tiled copy operation for smem to tmem load for scale factor tensor(s2t) -- - tCsSF_compact_s2t: The partitioned scale factor tensor in smem -- - tSF_compact_s2t: The partitioned scale factor tensor in tmem -- """ -- # (MMA, MMA_MN, MMA_K, STAGE) -- tCsSF_compact = cute.filter_zeros(sSF) -- # (MMA, MMA_MN, MMA_K) -- tCtSF_compact = cute.filter_zeros(tSF) -- -- # Make S2T CopyAtom and tiledCopy -- copy_atom_s2t = cute.make_copy_atom( -- tcgen05.Cp4x32x128bOp(self.cta_group), -- self.sf_dtype, -- ) -- tiled_copy_s2t = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSF_compact) -- thr_copy_s2t = tiled_copy_s2t.get_slice(0) -- -- # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE) -- tCsSF_compact_s2t_ = thr_copy_s2t.partition_S(tCsSF_compact) -- # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE) -- tCsSF_compact_s2t = tcgen05.get_s2t_smem_desc_tensor( -- tiled_copy_s2t, tCsSF_compact_s2t_ -- ) -- # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K) -- tCtSF_compact_s2t = thr_copy_s2t.partition_D(tCtSF_compact) -- -- return tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t -- -- def epilog_tmem_copy_and_partition( -- self, -- tidx: cutlass.Int32, -- tAcc: cute.Tensor, -- gC_mnl: cute.Tensor, -- epi_tile: cute.Tile, -- use_2cta_instrs: Union[cutlass.Boolean, bool], -- ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: -- """ -- Make tiledCopy for tensor memory load, then use it to partition tensor memory (source) and register array (destination). -- -- Args: -- tidx (cutlass.Int32): The thread index in epilogue warp groups -- tAcc (cute.Tensor): The accumulator tensor to be copied and partitioned -- gC_mnl (cute.Tensor): The global tensor C -- epi_tile (cute.Tile): The epilogue tiler -- use_2cta_instrs (bool): Whether use_2cta_instrs is enabled -- -- Returns: -- A tuple containing (tiled_copy_t2r, tTR_tAcc, tTR_rAcc) where: -- - tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r) -- - tTR_tAcc: The partitioned accumulator tensor -- - tTR_rAcc: The accumulated tensor in register used to hold t2r results -- """ -- # Make tiledCopy for tensor memory load -- copy_atom_t2r = sm100_utils.get_tmem_load_op( -- self.cta_tile_shape_mnk, -- self.c_layout, -- self.c_dtype, -- self.acc_dtype, -- epi_tile, -- use_2cta_instrs, -- ) -- # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, STAGE) -- tAcc_epi = cute.flat_divide( -- tAcc[((None, None), 0, 0, None)], -- epi_tile, -- ) -- # (EPI_TILE_M, EPI_TILE_N) -- tiled_copy_t2r = tcgen05.make_tmem_copy( -- copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)] -- ) -- -- thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) -- # (T2R, T2R_M, T2R_N, EPI_M, EPI_M, STAGE) -- tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi) -- -- # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL) -- gC_mnl_epi = cute.flat_divide( -- gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile -- ) -- # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL) -- tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi) -- # (T2R, T2R_M, T2R_N) -- tTR_rAcc = cute.make_rmem_tensor( -- tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype -- ) -- return tiled_copy_t2r, tTR_tAcc, tTR_rAcc -- -- def epilog_smem_copy_and_partition( -- self, -- tiled_copy_t2r: cute.TiledCopy, -- tTR_rC: cute.Tensor, -- tidx: cutlass.Int32, -- sC: cute.Tensor, -- ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: -- """ -- Make tiledCopy for shared memory store, then use it to partition register array (source) and shared memory (destination). -- -- Args: -- tiled_copy_t2r (cute.TiledCopy): The tiled copy operation for tmem to register copy(t2r) -- tTR_rC (cute.Tensor): The partitioned accumulator tensor -- tidx (cutlass.Int32): The thread index in epilogue warp groups -- sC (cute.Tensor): The shared memory tensor to be copied and partitioned -- sepi (cute.Tensor): -- -- Returns: -- A tuple containing (tiled_copy_r2s, tRS_rC, tRS_sC) where: -- - tiled_copy_r2s: The tiled copy operation for register to smem copy(r2s) -- - tRS_rC: The partitioned tensor C (register source) -- - tRS_sC: The partitioned tensor C (smem destination) -- """ -- copy_atom_r2s = sm100_utils.get_smem_store_op( -- self.c_layout, self.c_dtype, self.acc_dtype, tiled_copy_t2r -- ) -- tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r) -- # (R2S, R2S_M, R2S_N, PIPE_D) -- thr_copy_r2s = tiled_copy_r2s.get_slice(tidx) -- tRS_sC = thr_copy_r2s.partition_D(sC) -- # (R2S, R2S_M, R2S_N) -- tRS_rC = tiled_copy_r2s.retile(tTR_rC) -- return tiled_copy_r2s, tRS_rC, tRS_sC -- -- def epilog_gmem_copy_and_partition( -- self, -- tidx: cutlass.Int32, -- atom: Union[cute.CopyAtom, cute.TiledCopy], -- gC_mnl: cute.Tensor, -- epi_tile: cute.Tile, -- sC: cute.Tensor, -- ) -> Tuple[cute.CopyAtom, cute.Tensor, cute.Tensor]: -- """ -- Make tiledCopy for global memory store, then use it to: -- partition shared memory (source) and global memory (destination) for TMA store version. -- -- Args: -- tidx (cutlass.Int32): The thread index in epilogue warp groups -- atom (Union[cute.CopyAtom, cute.TiledCopy]): The copy_atom_c to be used for TMA store version, or tiled_copy_t2r for none TMA store version -- gC_mnl (cute.Tensor): The global tensor C -- epi_tile (cute.Tile): The epilogue tiler -- sC (cute.Tensor): The shared memory tensor to be copied and partitioned -- -- Returns: -- A tuple containing (tma_atom_c, bSG_sC, bSG_gC) where: -- - tma_atom_c: The TMA copy atom -- - bSG_sC: The partitioned shared memory tensor C -- - bSG_gC: The partitioned global tensor C -- """ -- # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL) -- gC_epi = cute.flat_divide( -- gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile -- ) -- -- tma_atom_c = atom -- sC_for_tma_partition = cute.group_modes(sC, 0, 2) -- gC_for_tma_partition = cute.group_modes(gC_epi, 0, 2) -- # ((ATOM_V, REST_V), EPI_M, EPI_N) -- # ((ATOM_V, REST_V), EPI_M, EPI_N, RestM, RestN, RestL) -- bSG_sC, bSG_gC = cpasync.tma_partition( -- tma_atom_c, -- 0, -- cute.make_layout(1), -- sC_for_tma_partition, -- gC_for_tma_partition, -- ) -- return tma_atom_c, bSG_sC, bSG_gC -- -- @staticmethod -- def _compute_stages( -- tiled_mma: cute.TiledMma, -- mma_tiler_mnk: Tuple[int, int, int], -- a_dtype: Type[cutlass.Numeric], -- a_major_mode: cute.nvgpu.OperandMajorMode, -- b_dtype: Type[cutlass.Numeric], -- b_major_mode: cute.nvgpu.OperandMajorMode, -- epi_tile: cute.Tile, -- c_dtype: Type[cutlass.Numeric], -- c_layout: utils.LayoutEnum, -- sf_dtype: Type[cutlass.Numeric], -- sf_vec_size: int, -- smem_capacity: int, -- occupancy: int, -- ) -> Tuple[int, int, int]: -- """Computes the optimal number of pipeline stages for operands. -- -- This method uses heuristics to determine the number of stages for the -- accumulator (ACC), A/B operands, and C operand based on the kernel -- configuration and available shared memory. -- -- Args: -- tiled_mma (cute.TiledMma): The tiled MMA object. -- mma_tiler_mnk (Tuple[int, int, int]): The shape (M, N, K) of the -- MMA tiler. -- a_dtype (Type[cutlass.Numeric]): Data type of operand A. -- a_major_mode (cute.nvgpu.OperandMajorMode): Layout of operand A. -- b_dtype (Type[cutlass.Numeric]): Data type of operand B. -- b_major_mode (cute.nvgpu.OperandMajorMode): Layout of operand B. -- epi_tile (cute.Tile): The epilogue tile shape. -- c_dtype (Type[cutlass.Numeric]): Data type of operand C. -- c_layout (utils.LayoutEnum): Layout of operand C. -- sf_dtype (Type[cutlass.Numeric]): Data type of the scale factors. -- sf_vec_size (int): Vector size of the scale factors. -- smem_capacity (int): Total available shared memory in bytes. -- occupancy (int): Target number of CTAs per SM. -- -- Returns: -- Tuple[int, int, int]: A tuple containing the number of stages for -- (ACC, A/B, C). -- """ -- # ACC stages -- num_acc_stage = 1 if mma_tiler_mnk[1] == 256 else 2 -- -- # Default C stages -- num_c_stage = 2 -- -- # Calculate smem layout and size for one stage of A, B, SFA, SFB and C -- a_smem_layout_stage_one = sm100_utils.make_smem_layout_a( -- tiled_mma, -- mma_tiler_mnk, -- a_dtype, -- 1, # a tmp 1 stage is provided -- ) -- b_smem_layout_staged_one = sm100_utils.make_smem_layout_b( -- tiled_mma, -- mma_tiler_mnk, -- b_dtype, -- 1, # a tmp 1 stage is provided -- ) -- sfa_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfa( -- tiled_mma, -- mma_tiler_mnk, -- sf_vec_size, -- 1, # a tmp 1 stage is provided -- ) -- sfb_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfb( -- tiled_mma, -- mma_tiler_mnk, -- sf_vec_size, -- 1, # a tmp 1 stage is provided -- ) -- -- c_smem_layout_staged_one = sm100_utils.make_smem_layout_epi( -- c_dtype, -- c_layout, -- epi_tile, -- 1, -- ) -- -- ab_bytes_per_stage = ( -- cute.size_in_bytes(a_dtype, a_smem_layout_stage_one) -- + cute.size_in_bytes(b_dtype, b_smem_layout_staged_one) -- + cute.size_in_bytes(sf_dtype, sfa_smem_layout_staged_one) -- + cute.size_in_bytes(sf_dtype, sfb_smem_layout_staged_one) -- ) -- mbar_helpers_bytes = 1024 -- c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout_staged_one) -- c_bytes = c_bytes_per_stage * num_c_stage -- -- # Calculate A/B/SFA/SFB stages: -- # Start with total smem per CTA (capacity / occupancy) -- # Subtract reserved bytes and initial C stages bytes -- # Divide remaining by bytes needed per A/B/SFA/SFB stage -- num_ab_stage = ( -- smem_capacity // occupancy - (mbar_helpers_bytes + c_bytes) -- ) // ab_bytes_per_stage -- -- # Refine epilogue stages: -- # Calculate remaining smem after allocating for A/B/SFA/SFB stages and reserved bytes -- # Add remaining unused smem to epilogue -- num_c_stage += ( -- smem_capacity -- - occupancy * ab_bytes_per_stage * num_ab_stage -- - occupancy * (mbar_helpers_bytes + c_bytes) -- ) // (occupancy * c_bytes_per_stage) -- -- return num_acc_stage, num_ab_stage, num_c_stage -- - @staticmethod - def _compute_grid( - c: cute.Tensor, -@@ -2089,125 +1673,3 @@ class Sm100BlockScaledPersistentDenseGemmKernel: - if mma_tiler_mn[1] < 64 and (n > mma_tiler_mn[1] or cluster_shape_mn[1] > 1): - can_implement = False - return can_implement -- -- # fully dynamic shape -- @cute.jit -- def wrapper( -- self, -- mA: cute.Tensor, -- mB: cute.Tensor, -- mC: cute.Tensor, -- sf_m: cutlass.Int64, -- sf_n: cutlass.Int64, -- sf_k: cutlass.Int64, -- l: cutlass.Constexpr, -- a_sf_ptr: cute.Pointer, -- b_sf_ptr: cute.Pointer, -- alpha_tensor: cute.Tensor, -- max_active_clusters: cutlass.Constexpr, -- current_stream, -- swap_ab: cutlass.Constexpr = False, -- epilogue_op: cutlass.Constexpr = lambda x: x, -- ): -- """Executes the wrapped GEMM kernel with dynamically shaped tensors. -- -- Uses TVM-FFI for efficient tensor passing: A, B, C, and alpha are passed -- as cute.Tensor directly (torch tensors at runtime via TVM-FFI's C-level -- dlpack, with negligible conversion cost). Scale factor tensors remain as -- pointers because their 6D BlockScaledBasicChunk layout cannot be expressed -- as a torch tensor. -- -- Args: -- mA (cute.Tensor): Input tensor A. -- FP4 path: shape (m, k_packed), dtype Uint8 (2xFP4 packed per byte). -- MXFP8 path: shape (m, k), dtype Float8. -- mB (cute.Tensor): Input tensor B. -- FP4 path: shape (n, k_packed), dtype Uint8. -- MXFP8 path: shape (n, k), dtype Float8. -- mC (cute.Tensor): Output tensor C, shape (m, n). -- sf_m (cutlass.Int64): Scale factor M dim (ceil(m/128)). -- sf_n (cutlass.Int64): Scale factor N dim (ceil(n/128)). -- sf_k (cutlass.Int64): Scale factor K dim (ceil(k/sf_vec_size/4)). -- l (cutlass.Constexpr): Batch dimension (L). -- a_sf_ptr (cute.Pointer): Pointer to scale factor tensor for A. -- b_sf_ptr (cute.Pointer): Pointer to scale factor tensor for B. -- alpha_tensor (cute.Tensor): Alpha scaling factor, shape (1,), float32. -- max_active_clusters (cutlass.Constexpr): Max active clusters. -- current_stream: CUDA stream (managed by TVM-FFI fake stream). -- swap_ab (cutlass.Constexpr): Whether A/B are swapped (controls C layout). -- epilogue_op (cutlass.Constexpr): Elementwise epilogue function. -- """ -- # A, B, C are passed as cute.Tensor via TVM-FFI. -- # Support two input encodings: -- # 1) FP4 packed as uint8 (recast to Float4E2M1FN, k = k_packed * 2) -- # 2) MXFP8 as float8 (k = k_raw, no recast) -- m = cute.size(mA, mode=[0]) -- k_raw = cute.size(mA, mode=[1]) -- n = cute.size(mB, mode=[0]) -- -- if cutlass.const_expr( -- mA.element_type == cutlass.Uint8 and mB.element_type == cutlass.Uint8 -- ): -- # FP4 packed path: 2 FP4 values per uint8 byte -- k = k_raw * 2 -- a_ptr = cute.recast_ptr(mA.iterator, dtype=cutlass.Float4E2M1FN) -- b_ptr = cute.recast_ptr(mB.iterator, dtype=cutlass.Float4E2M1FN) -- elif cutlass.const_expr(mA.element_type != mB.element_type): -- raise TypeError( -- "Unsupported mixed input dtypes for block-scaled GEMM: " -- "mA and mB must have matching element_type " -- "(both Uint8 for FP4 path, or both FP8 for MXFP8 path)." -- ) -- else: -- # MXFP8 path: input tensors are already FP8. -- k = k_raw -- a_ptr = mA.iterator -- b_ptr = mB.iterator -- -- a_tensor = cute.make_tensor( -- a_ptr, -- layout=cute.make_ordered_layout((m, k, l), order=(1, 0, 2)), -- ) -- b_tensor = cute.make_tensor( -- b_ptr, -- layout=cute.make_ordered_layout((n, k, l), order=(1, 0, 2)), -- ) -- # Reshape C to (m, n, l) -- swap_ab is constexpr, determines layout at compile time -- if cutlass.const_expr(swap_ab): -- c_tensor = cute.make_tensor( -- mC.iterator, -- layout=cute.make_ordered_layout((m, n, l), order=(0, 1, 2)), -- ) -- else: -- c_tensor = cute.make_tensor( -- mC.iterator, -- layout=cute.make_ordered_layout((m, n, l), order=(1, 0, 2)), -- ) -- # Scale factor tensors: 6D BlockScaledBasicChunk layout from pointers -- # (32, 4, sf_m, 4, sf_k, l) with order (2, 1, 4, 0, 3, 5) -- sfa_tensor = cute.make_tensor( -- a_sf_ptr, -- layout=cute.make_ordered_layout( -- (32, 4, sf_m, 4, sf_k, l), -- order=(2, 1, 4, 0, 3, 5), -- ), -- ) -- sfb_tensor = cute.make_tensor( -- b_sf_ptr, -- layout=cute.make_ordered_layout( -- (32, 4, sf_n, 4, sf_k, l), -- order=(2, 1, 4, 0, 3, 5), -- ), -- ) -- -- self( -- a_tensor, -- b_tensor, -- sfa_tensor, -- sfb_tensor, -- c_tensor, -- alpha_tensor, -- max_active_clusters, -- current_stream, -- epilogue_op, -- ) -diff --git a/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm100_common.py b/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm100_common.py -new file mode 100644 -index 00000000..0d1642b1 ---- /dev/null -+++ b/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm100_common.py -@@ -0,0 +1,669 @@ -+# Copyright (c) 2025 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. -+ -+ -+from typing import Any, Tuple, Type, Union -+ -+import cutlass -+import cutlass.cute as cute -+import cutlass.utils as utils -+import cutlass.utils.blackwell_helpers as sm100_utils -+import cutlass.utils.blockscaled_layout as blockscaled_utils -+from cutlass.cute.nvgpu import cpasync, tcgen05 -+ -+ -+class _Sm100BlockScaledGemmCommon: -+ """Shared launch and copy-partition helpers for SM100 block-scaled GEMM.""" -+ -+ acc_dtype: Any -+ c_dtype: Any -+ c_layout: Any -+ cta_group: Any -+ cta_tile_shape_mnk: Any -+ sf_dtype: Any -+ # Concrete kernels initialize these in __init__ or _setup_attributes. -+ # CuTe layout/configuration objects are generated dynamically, so Any is -+ # intentional at this shared interface boundary. -+ a_smem_layout_staged: Any -+ b_smem_layout_staged: Any -+ c_smem_layout_staged: Any -+ sfa_smem_layout_staged: Any -+ sfb_smem_layout_staged: Any -+ cluster_layout_sfb_vmnk: Any -+ cluster_layout_vmnk: Any -+ cluster_shape_mn: Any -+ epi_tile: Any -+ mma_inst_shape_mnk: Any -+ mma_inst_shape_mnk_sfb: Any -+ mma_tiler: Any -+ mma_tiler_sfb: Any -+ sf_vec_size: int -+ -+ def _setup_attributes(self) -> None: -+ raise NotImplementedError -+ -+ def __call__( -+ self, -+ a_tensor: cute.Tensor, -+ b_tensor: cute.Tensor, -+ sfa_tensor: cute.Tensor, -+ sfb_tensor: cute.Tensor, -+ c_tensor: cute.Tensor, -+ alpha: cute.Tensor, -+ max_active_clusters: cutlass.Constexpr, -+ stream: Any, -+ epilogue_op: cutlass.Constexpr = lambda x: x, -+ ) -> Any: -+ raise NotImplementedError -+ -+ def _prepare_tma( -+ self, -+ a_tensor: cute.Tensor, -+ b_tensor: cute.Tensor, -+ sfa_tensor: cute.Tensor, -+ sfb_tensor: cute.Tensor, -+ c_tensor: cute.Tensor, -+ ): -+ """Build the shared block-scaled MMA and TMA descriptors.""" -+ # Setup static attributes before smem/grid/tma computation -+ self.a_dtype: Type[cutlass.Numeric] = a_tensor.element_type -+ self.b_dtype: Type[cutlass.Numeric] = b_tensor.element_type -+ self.sf_dtype: Type[cutlass.Numeric] = sfa_tensor.element_type -+ self.c_dtype: Type[cutlass.Numeric] = c_tensor.element_type -+ self.a_major_mode = utils.LayoutEnum.from_tensor(a_tensor).mma_major_mode() -+ self.b_major_mode = utils.LayoutEnum.from_tensor(b_tensor).mma_major_mode() -+ self.c_layout = utils.LayoutEnum.from_tensor(c_tensor) -+ -+ # Check if input data types are compatible with MMA instruction -+ if cutlass.const_expr(self.a_dtype != self.b_dtype): -+ raise TypeError(f"Type must match: {self.a_dtype} != {self.b_dtype}") -+ -+ # Setup attributes that dependent on gemm inputs -+ self._setup_attributes() -+ -+ # Setup sfa/sfb tensor by filling A/B tensor to scale factor atom layout -+ # ((Atom_M, Rest_M),(Atom_K, Rest_K),RestL) -+ sfa_layout = blockscaled_utils.tile_atom_to_shape_SF( -+ a_tensor.shape, self.sf_vec_size -+ ) -+ sfa_tensor = cute.make_tensor(sfa_tensor.iterator, sfa_layout) -+ -+ # ((Atom_N, Rest_N),(Atom_K, Rest_K),RestL) -+ sfb_layout = blockscaled_utils.tile_atom_to_shape_SF( -+ b_tensor.shape, self.sf_vec_size -+ ) -+ sfb_tensor = cute.make_tensor(sfb_tensor.iterator, sfb_layout) -+ -+ tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma( -+ self.a_dtype, -+ self.b_dtype, -+ self.a_major_mode, -+ self.b_major_mode, -+ self.sf_dtype, -+ self.sf_vec_size, -+ self.cta_group, -+ self.mma_inst_shape_mnk[:2], -+ ) -+ -+ tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma( -+ self.a_dtype, -+ self.b_dtype, -+ self.a_major_mode, -+ self.b_major_mode, -+ self.sf_dtype, -+ self.sf_vec_size, -+ cute.nvgpu.tcgen05.CtaGroup.ONE, -+ self.mma_inst_shape_mnk_sfb[:2], -+ ) -+ atom_thr_size = cute.size(tiled_mma.thr_id.shape) -+ -+ # Setup TMA load for A -+ a_op = sm100_utils.cluster_shape_to_tma_atom_A( -+ self.cluster_shape_mn, tiled_mma.thr_id -+ ) -+ a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) -+ tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( -+ a_op, -+ a_tensor, -+ a_smem_layout, -+ self.mma_tiler, -+ tiled_mma, -+ self.cluster_layout_vmnk.shape, -+ ) -+ -+ # Setup TMA load for B -+ b_op = sm100_utils.cluster_shape_to_tma_atom_B( -+ self.cluster_shape_mn, tiled_mma.thr_id -+ ) -+ b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) -+ tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( -+ b_op, -+ b_tensor, -+ b_smem_layout, -+ self.mma_tiler, -+ tiled_mma, -+ self.cluster_layout_vmnk.shape, -+ ) -+ -+ # Setup TMA load for SFA -+ sfa_op = sm100_utils.cluster_shape_to_tma_atom_A( -+ self.cluster_shape_mn, tiled_mma.thr_id -+ ) -+ sfa_smem_layout = cute.slice_( -+ self.sfa_smem_layout_staged, (None, None, None, 0) -+ ) -+ tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A( -+ sfa_op, -+ sfa_tensor, -+ sfa_smem_layout, -+ self.mma_tiler, -+ tiled_mma, -+ self.cluster_layout_vmnk.shape, -+ internal_type=cutlass.Int16, -+ ) -+ -+ # Setup TMA load for SFB -+ sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB( -+ self.cluster_shape_mn, tiled_mma.thr_id -+ ) -+ sfb_smem_layout = cute.slice_( -+ self.sfb_smem_layout_staged, (None, None, None, 0) -+ ) -+ tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B( -+ sfb_op, -+ sfb_tensor, -+ sfb_smem_layout, -+ self.mma_tiler_sfb, -+ tiled_mma_sfb, -+ self.cluster_layout_sfb_vmnk.shape, -+ internal_type=cutlass.Int16, -+ ) -+ if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 192): -+ x = tma_tensor_sfb.stride[0][1] -+ y = cute.ceil_div(tma_tensor_sfb.shape[0][1], 4) -+ -+ new_shape = ( -+ (tma_tensor_sfb.shape[0][0], ((2, 2), y)), -+ tma_tensor_sfb.shape[1], -+ tma_tensor_sfb.shape[2], -+ ) -+ # Use right multiplication for ScaledBasis (3 * x instead of x * 3) -+ x_times_3 = 3 * x -+ new_stride = ( -+ (tma_tensor_sfb.stride[0][0], ((x, x), x_times_3)), -+ tma_tensor_sfb.stride[1], -+ tma_tensor_sfb.stride[2], -+ ) -+ tma_tensor_sfb_new_layout = cute.make_layout(new_shape, stride=new_stride) -+ tma_tensor_sfb = cute.make_tensor( -+ tma_tensor_sfb.iterator, tma_tensor_sfb_new_layout -+ ) -+ -+ a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout) -+ b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout) -+ sfa_copy_size = cute.size_in_bytes(self.sf_dtype, sfa_smem_layout) -+ sfb_copy_size = cute.size_in_bytes(self.sf_dtype, sfb_smem_layout) -+ self.num_tma_load_bytes = ( -+ a_copy_size + b_copy_size + sfa_copy_size + sfb_copy_size -+ ) * atom_thr_size -+ -+ # Setup TMA store for C -+ epi_smem_layout = cute.slice_(self.c_smem_layout_staged, (None, None, 0)) -+ tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( -+ cpasync.CopyBulkTensorTileS2GOp(), -+ c_tensor, -+ epi_smem_layout, -+ self.epi_tile, -+ ) -+ -+ return ( -+ tiled_mma, -+ tiled_mma_sfb, -+ tma_atom_a, -+ tma_tensor_a, -+ tma_atom_b, -+ tma_tensor_b, -+ tma_atom_sfa, -+ tma_tensor_sfa, -+ tma_atom_sfb, -+ tma_tensor_sfb, -+ tma_atom_c, -+ tma_tensor_c, -+ ) -+ -+ @staticmethod -+ def _compute_stages( -+ tiled_mma: cute.TiledMma, -+ mma_tiler_mnk: Tuple[int, int, int], -+ a_dtype: Type[cutlass.Numeric], -+ a_major_mode: cute.nvgpu.OperandMajorMode, -+ b_dtype: Type[cutlass.Numeric], -+ b_major_mode: cute.nvgpu.OperandMajorMode, -+ epi_tile: cute.Tile, -+ c_dtype: Type[cutlass.Numeric], -+ c_layout: utils.LayoutEnum, -+ sf_dtype: Type[cutlass.Numeric], -+ sf_vec_size: int, -+ smem_capacity: int, -+ occupancy: int, -+ reserved_smem_bytes: int, -+ ) -> Tuple[int, int, int]: -+ """Computes the optimal number of pipeline stages for operands. -+ -+ This method uses heuristics to determine the number of stages for the -+ accumulator (ACC), A/B operands, and C operand based on the kernel -+ configuration and available shared memory. -+ -+ Args: -+ tiled_mma (cute.TiledMma): The tiled MMA object. -+ mma_tiler_mnk (Tuple[int, int, int]): The shape (M, N, K) of the -+ MMA tiler. -+ a_dtype (Type[cutlass.Numeric]): Data type of operand A. -+ a_major_mode (cute.nvgpu.OperandMajorMode): Layout of operand A. -+ b_dtype (Type[cutlass.Numeric]): Data type of operand B. -+ b_major_mode (cute.nvgpu.OperandMajorMode): Layout of operand B. -+ epi_tile (cute.Tile): The epilogue tile shape. -+ c_dtype (Type[cutlass.Numeric]): Data type of operand C. -+ c_layout (utils.LayoutEnum): Layout of operand C. -+ sf_dtype (Type[cutlass.Numeric]): Data type of the scale factors. -+ sf_vec_size (int): Vector size of the scale factors. -+ smem_capacity (int): Total available shared memory in bytes. -+ occupancy (int): Target number of CTAs per SM. -+ reserved_smem_bytes (int): Per-CTA bytes reserved outside the staged operands and output. -+ -+ Returns: -+ Tuple[int, int, int]: A tuple containing the number of stages for -+ (ACC, A/B, C). -+ """ -+ # ACC stages -+ num_acc_stage = 1 if mma_tiler_mnk[1] == 256 else 2 -+ -+ # Default C stages -+ num_c_stage = 2 -+ -+ # Calculate smem layout and size for one stage of A, B, SFA, SFB and C -+ a_smem_layout_stage_one = sm100_utils.make_smem_layout_a( -+ tiled_mma, -+ mma_tiler_mnk, -+ a_dtype, -+ 1, # a tmp 1 stage is provided -+ ) -+ b_smem_layout_staged_one = sm100_utils.make_smem_layout_b( -+ tiled_mma, -+ mma_tiler_mnk, -+ b_dtype, -+ 1, # a tmp 1 stage is provided -+ ) -+ sfa_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfa( -+ tiled_mma, -+ mma_tiler_mnk, -+ sf_vec_size, -+ 1, # a tmp 1 stage is provided -+ ) -+ sfb_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfb( -+ tiled_mma, -+ mma_tiler_mnk, -+ sf_vec_size, -+ 1, # a tmp 1 stage is provided -+ ) -+ -+ c_smem_layout_staged_one = sm100_utils.make_smem_layout_epi( -+ c_dtype, -+ c_layout, -+ epi_tile, -+ 1, -+ ) -+ -+ ab_bytes_per_stage = ( -+ cute.size_in_bytes(a_dtype, a_smem_layout_stage_one) -+ + cute.size_in_bytes(b_dtype, b_smem_layout_staged_one) -+ + cute.size_in_bytes(sf_dtype, sfa_smem_layout_staged_one) -+ + cute.size_in_bytes(sf_dtype, sfb_smem_layout_staged_one) -+ ) -+ c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout_staged_one) -+ c_bytes = c_bytes_per_stage * num_c_stage -+ -+ # Calculate A/B/SFA/SFB stages: -+ # Start with total smem per CTA (capacity / occupancy) -+ # Subtract reserved bytes and initial C stages bytes -+ # Divide remaining by bytes needed per A/B/SFA/SFB stage -+ num_ab_stage = ( -+ smem_capacity // occupancy - (reserved_smem_bytes + c_bytes) -+ ) // ab_bytes_per_stage -+ -+ # Refine epilogue stages: -+ # Calculate remaining smem after allocating for A/B/SFA/SFB stages and reserved bytes -+ # Add remaining unused smem to epilogue -+ num_c_stage += ( -+ smem_capacity -+ - occupancy * ab_bytes_per_stage * num_ab_stage -+ - occupancy * (reserved_smem_bytes + c_bytes) -+ ) // (occupancy * c_bytes_per_stage) -+ -+ return num_acc_stage, num_ab_stage, num_c_stage -+ -+ # fully dynamic shape -+ @cute.jit -+ def wrapper( -+ self, -+ mA: cute.Tensor, -+ mB: cute.Tensor, -+ mC: cute.Tensor, -+ sf_m: cutlass.Int64, -+ sf_n: cutlass.Int64, -+ sf_k: cutlass.Int64, -+ l: cutlass.Constexpr, -+ a_sf_ptr: cute.Pointer, -+ b_sf_ptr: cute.Pointer, -+ alpha_tensor: cute.Tensor, -+ max_active_clusters: cutlass.Constexpr, -+ current_stream, -+ swap_ab: cutlass.Constexpr = False, -+ epilogue_op: cutlass.Constexpr = lambda x: x, -+ ): -+ """Executes the wrapped GEMM kernel with dynamically shaped tensors. -+ -+ Uses TVM-FFI for efficient tensor passing: A, B, C, and alpha are passed -+ as cute.Tensor directly (torch tensors at runtime via TVM-FFI's C-level -+ dlpack, with negligible conversion cost). Scale factor tensors remain as -+ pointers because their 6D BlockScaledBasicChunk layout cannot be expressed -+ as a torch tensor. -+ -+ Args: -+ mA (cute.Tensor): Input tensor A. -+ FP4 path: shape (m, k_packed), dtype Uint8 (2xFP4 packed per byte). -+ MXFP8 path: shape (m, k), dtype Float8. -+ mB (cute.Tensor): Input tensor B. -+ FP4 path: shape (n, k_packed), dtype Uint8. -+ MXFP8 path: shape (n, k), dtype Float8. -+ mC (cute.Tensor): Output tensor C, shape (m, n). -+ sf_m (cutlass.Int64): Scale factor M dim (ceil(m/128)). -+ sf_n (cutlass.Int64): Scale factor N dim (ceil(n/128)). -+ sf_k (cutlass.Int64): Scale factor K dim (ceil(k/sf_vec_size/4)). -+ l (cutlass.Constexpr): Batch dimension (L). -+ a_sf_ptr (cute.Pointer): Pointer to scale factor tensor for A. -+ b_sf_ptr (cute.Pointer): Pointer to scale factor tensor for B. -+ alpha_tensor (cute.Tensor): Alpha scaling factor, shape (1,), float32. -+ max_active_clusters (cutlass.Constexpr): Max active clusters. -+ current_stream: CUDA stream (managed by TVM-FFI fake stream). -+ swap_ab (cutlass.Constexpr): Whether A/B are swapped (controls C layout). -+ epilogue_op (cutlass.Constexpr): Elementwise epilogue function. -+ """ -+ # A, B, C are passed as cute.Tensor via TVM-FFI. -+ # Support two input encodings: -+ # 1) FP4 packed as uint8 (recast to Float4E2M1FN, k = k_packed * 2) -+ # 2) MXFP8 as float8 (k = k_raw, no recast) -+ m = cute.size(mA, mode=[0]) -+ k_raw = cute.size(mA, mode=[1]) -+ n = cute.size(mB, mode=[0]) -+ -+ if cutlass.const_expr( -+ mA.element_type == cutlass.Uint8 and mB.element_type == cutlass.Uint8 -+ ): -+ # FP4 packed path: 2 FP4 values per uint8 byte -+ k = k_raw * 2 -+ a_ptr = cute.recast_ptr(mA.iterator, dtype=cutlass.Float4E2M1FN) -+ b_ptr = cute.recast_ptr(mB.iterator, dtype=cutlass.Float4E2M1FN) -+ elif cutlass.const_expr(mA.element_type != mB.element_type): -+ raise TypeError( -+ "Unsupported mixed input dtypes for block-scaled GEMM: " -+ "mA and mB must have matching element_type " -+ "(both Uint8 for FP4 path, or both FP8 for MXFP8 path)." -+ ) -+ else: -+ # MXFP8 path: input tensors are already FP8. -+ k = k_raw -+ a_ptr = mA.iterator -+ b_ptr = mB.iterator -+ -+ a_tensor = cute.make_tensor( -+ a_ptr, -+ layout=cute.make_ordered_layout((m, k, l), order=(1, 0, 2)), -+ ) -+ b_tensor = cute.make_tensor( -+ b_ptr, -+ layout=cute.make_ordered_layout((n, k, l), order=(1, 0, 2)), -+ ) -+ # Reshape C to (m, n, l) -- swap_ab is constexpr, determines layout at compile time -+ if cutlass.const_expr(swap_ab): -+ c_tensor = cute.make_tensor( -+ mC.iterator, -+ layout=cute.make_ordered_layout((m, n, l), order=(0, 1, 2)), -+ ) -+ else: -+ c_tensor = cute.make_tensor( -+ mC.iterator, -+ layout=cute.make_ordered_layout((m, n, l), order=(1, 0, 2)), -+ ) -+ # Scale factor tensors: 6D BlockScaledBasicChunk layout from pointers -+ # (32, 4, sf_m, 4, sf_k, l) with order (2, 1, 4, 0, 3, 5) -+ sfa_tensor = cute.make_tensor( -+ a_sf_ptr, -+ layout=cute.make_ordered_layout( -+ (32, 4, sf_m, 4, sf_k, l), -+ order=(2, 1, 4, 0, 3, 5), -+ ), -+ ) -+ sfb_tensor = cute.make_tensor( -+ b_sf_ptr, -+ layout=cute.make_ordered_layout( -+ (32, 4, sf_n, 4, sf_k, l), -+ order=(2, 1, 4, 0, 3, 5), -+ ), -+ ) -+ -+ self( -+ a_tensor, -+ b_tensor, -+ sfa_tensor, -+ sfb_tensor, -+ c_tensor, -+ alpha_tensor, -+ max_active_clusters, -+ current_stream, -+ epilogue_op, -+ ) -+ -+ def mainloop_s2t_copy_and_partition( -+ self, -+ sSF: cute.Tensor, -+ tSF: cute.Tensor, -+ ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: -+ """ -+ Make tiledCopy for smem to tmem load for scale factor tensor, then use it to partition smem memory (source) and tensor memory (destination). -+ -+ Args: -+ sSF (cute.Tensor): The scale factor tensor in smem -+ tSF (cute.Tensor): The scale factor tensor in tmem -+ -+ Returns: -+ A tuple containing (tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t) where: -+ - tiled_copy_s2t: The tiled copy operation for smem to tmem load for scale factor tensor(s2t) -+ - tCsSF_compact_s2t: The partitioned scale factor tensor in smem -+ - tSF_compact_s2t: The partitioned scale factor tensor in tmem -+ """ -+ # (MMA, MMA_MN, MMA_K, STAGE) -+ tCsSF_compact = cute.filter_zeros(sSF) -+ # (MMA, MMA_MN, MMA_K) -+ tCtSF_compact = cute.filter_zeros(tSF) -+ -+ # Make S2T CopyAtom and tiledCopy -+ copy_atom_s2t = cute.make_copy_atom( -+ tcgen05.Cp4x32x128bOp(self.cta_group), -+ self.sf_dtype, -+ ) -+ tiled_copy_s2t = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSF_compact) -+ thr_copy_s2t = tiled_copy_s2t.get_slice(0) -+ -+ # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE) -+ tCsSF_compact_s2t_ = thr_copy_s2t.partition_S(tCsSF_compact) -+ # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE) -+ tCsSF_compact_s2t = tcgen05.get_s2t_smem_desc_tensor( -+ tiled_copy_s2t, tCsSF_compact_s2t_ -+ ) -+ # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K) -+ tCtSF_compact_s2t = thr_copy_s2t.partition_D(tCtSF_compact) -+ -+ return tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t -+ -+ def epilog_tmem_copy_and_partition( -+ self, -+ tidx: cutlass.Int32, -+ tAcc: cute.Tensor, -+ gC_mnl: cute.Tensor, -+ epi_tile: cute.Tile, -+ use_2cta_instrs: Union[cutlass.Boolean, bool], -+ ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: -+ """ -+ Make tiledCopy for tensor memory load, then use it to partition tensor memory (source) and register array (destination). -+ -+ Args: -+ tidx (cutlass.Int32): The thread index in epilogue warp groups -+ tAcc (cute.Tensor): The accumulator tensor to be copied and partitioned -+ gC_mnl (cute.Tensor): The global tensor C -+ epi_tile (cute.Tile): The epilogue tiler -+ use_2cta_instrs (bool): Whether use_2cta_instrs is enabled -+ -+ Returns: -+ A tuple containing (tiled_copy_t2r, tTR_tAcc, tTR_rAcc) where: -+ - tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r) -+ - tTR_tAcc: The partitioned accumulator tensor -+ - tTR_rAcc: The accumulated tensor in register used to hold t2r results -+ """ -+ # Make tiledCopy for tensor memory load -+ copy_atom_t2r = sm100_utils.get_tmem_load_op( -+ self.cta_tile_shape_mnk, -+ self.c_layout, -+ self.c_dtype, -+ self.acc_dtype, -+ epi_tile, -+ use_2cta_instrs, -+ ) -+ # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, STAGE) -+ tAcc_epi = cute.flat_divide( -+ tAcc[((None, None), 0, 0, None)], -+ epi_tile, -+ ) -+ # (EPI_TILE_M, EPI_TILE_N) -+ tiled_copy_t2r = tcgen05.make_tmem_copy( -+ copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)] -+ ) -+ -+ thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) -+ # (T2R, T2R_M, T2R_N, EPI_M, EPI_M, STAGE) -+ tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi) -+ -+ # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL) -+ gC_mnl_epi = cute.flat_divide( -+ gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile -+ ) -+ # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL) -+ tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi) -+ # (T2R, T2R_M, T2R_N) -+ tTR_rAcc = cute.make_rmem_tensor( -+ tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype -+ ) -+ return tiled_copy_t2r, tTR_tAcc, tTR_rAcc -+ -+ def epilog_smem_copy_and_partition( -+ self, -+ tiled_copy_t2r: cute.TiledCopy, -+ tTR_rC: cute.Tensor, -+ tidx: cutlass.Int32, -+ sC: cute.Tensor, -+ ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: -+ """ -+ Make tiledCopy for shared memory store, then use it to partition register array (source) and shared memory (destination). -+ -+ Args: -+ tiled_copy_t2r (cute.TiledCopy): The tiled copy operation for tmem to register copy(t2r) -+ tTR_rC (cute.Tensor): The partitioned accumulator tensor -+ tidx (cutlass.Int32): The thread index in epilogue warp groups -+ sC (cute.Tensor): The shared memory tensor to be copied and partitioned -+ sepi (cute.Tensor): -+ -+ Returns: -+ A tuple containing (tiled_copy_r2s, tRS_rC, tRS_sC) where: -+ - tiled_copy_r2s: The tiled copy operation for register to smem copy(r2s) -+ - tRS_rC: The partitioned tensor C (register source) -+ - tRS_sC: The partitioned tensor C (smem destination) -+ """ -+ copy_atom_r2s = sm100_utils.get_smem_store_op( -+ self.c_layout, self.c_dtype, self.acc_dtype, tiled_copy_t2r -+ ) -+ tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r) -+ # (R2S, R2S_M, R2S_N, PIPE_D) -+ thr_copy_r2s = tiled_copy_r2s.get_slice(tidx) -+ tRS_sC = thr_copy_r2s.partition_D(sC) -+ # (R2S, R2S_M, R2S_N) -+ tRS_rC = tiled_copy_r2s.retile(tTR_rC) -+ return tiled_copy_r2s, tRS_rC, tRS_sC -+ -+ def epilog_gmem_copy_and_partition( -+ self, -+ tidx: cutlass.Int32, -+ atom: Union[cute.CopyAtom, cute.TiledCopy], -+ gC_mnl: cute.Tensor, -+ epi_tile: cute.Tile, -+ sC: cute.Tensor, -+ ) -> Tuple[cute.CopyAtom, cute.Tensor, cute.Tensor]: -+ """ -+ Make tiledCopy for global memory store, then use it to: -+ partition shared memory (source) and global memory (destination) for TMA store version. -+ -+ Args: -+ tidx (cutlass.Int32): The thread index in epilogue warp groups -+ atom (Union[cute.CopyAtom, cute.TiledCopy]): The copy_atom_c to be used for TMA store version, or tiled_copy_t2r for none TMA store version -+ gC_mnl (cute.Tensor): The global tensor C -+ epi_tile (cute.Tile): The epilogue tiler -+ sC (cute.Tensor): The shared memory tensor to be copied and partitioned -+ -+ Returns: -+ A tuple containing (tma_atom_c, bSG_sC, bSG_gC) where: -+ - tma_atom_c: The TMA copy atom -+ - bSG_sC: The partitioned shared memory tensor C -+ - bSG_gC: The partitioned global tensor C -+ """ -+ # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL) -+ gC_epi = cute.flat_divide( -+ gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile -+ ) -+ -+ tma_atom_c = atom -+ sC_for_tma_partition = cute.group_modes(sC, 0, 2) -+ gC_for_tma_partition = cute.group_modes(gC_epi, 0, 2) -+ # ((ATOM_V, REST_V), EPI_M, EPI_N) -+ # ((ATOM_V, REST_V), EPI_M, EPI_N, RestM, RestN, RestL) -+ bSG_sC, bSG_gC = cpasync.tma_partition( -+ tma_atom_c, -+ 0, -+ cute.make_layout(1), -+ sC_for_tma_partition, -+ gC_for_tma_partition, -+ ) -+ return tma_atom_c, bSG_sC, bSG_gC -diff --git a/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm100_splitk.py b/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm100_splitk.py -new file mode 100644 -index 00000000..f78dd8b3 ---- /dev/null -+++ b/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm100_splitk.py -@@ -0,0 +1,1539 @@ -+# Copyright (c) 2025 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. -+ -+# This kernel keeps FlashInfer's block-scaled MXFP8 mainloop and adapts the -+# cluster-local reduction protocol from NVIDIA DKG's FP8 split-K tutorial. -+ -+from typing import Optional, Tuple, Union -+ -+import cuda.bindings.driver as cuda -+import cutlass -+import cutlass.cute as cute -+import cutlass.pipeline as pipeline -+import cutlass.utils as utils -+import cutlass.utils.blackwell_helpers as sm100_utils -+import cutlass.utils.blockscaled_layout as blockscaled_utils -+from cutlass.cute.nvgpu import cpasync, tcgen05 -+ -+from cutlass.cute.arch import griddepcontrol_launch_dependents, griddepcontrol_wait -+from cutlass.pipeline import PipelineTmaUmma, PipelineUmmaAsync -+from cutlass import Float32, Int32 -+from cutlass._mlir.dialects import llvm -+from .dense_blockscaled_gemm_sm100_common import _Sm100BlockScaledGemmCommon -+ -+from cutlass.cutlass_dsl import ( -+ T, -+ dsl_user_op, -+ extract_mlir_values, -+ new_from_mlir_values, -+) -+ -+ -+@dsl_user_op -+def _map_shared_rank( -+ smem_ptr: cute.Pointer, -+ peer_cta_rank_in_cluster: Int32, -+ *, -+ loc=None, -+ ip=None, -+) -> Int32: -+ """Map a local shared-memory pointer to a peer CTA in the cluster.""" -+ smem_ptr_i32 = smem_ptr.toint(loc=loc, ip=ip).ir_value() -+ return Int32( -+ llvm.inline_asm( -+ T.i32(), -+ [smem_ptr_i32, peer_cta_rank_in_cluster.ir_value()], -+ "mapa.shared::cluster.u32 $0, $1, $2;", -+ "=r,r,r", -+ has_side_effects=False, -+ is_align_stack=False, -+ asm_dialect=llvm.AsmDialect.AD_ATT, -+ ) -+ ) -+ -+ -+@dsl_user_op -+def _store_shared_remote_v4( -+ value0: Float32, -+ value1: Float32, -+ value2: Float32, -+ value3: Float32, -+ smem_ptr: cute.Pointer, -+ mbar_ptr: cute.Pointer, -+ peer_cta_rank_in_cluster: Int32, -+ *, -+ loc=None, -+ ip=None, -+) -> None: -+ """Store four FP32 values to peer SMEM and complete 16 mbarrier bytes.""" -+ remote_smem_ptr = _map_shared_rank( -+ smem_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip -+ ).ir_value() -+ remote_mbar_ptr = _map_shared_rank( -+ mbar_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip -+ ).ir_value() -+ values = [ -+ value0.bitcast(Int32).ir_value(loc=loc, ip=ip), -+ value1.bitcast(Int32).ir_value(loc=loc, ip=ip), -+ value2.bitcast(Int32).ir_value(loc=loc, ip=ip), -+ value3.bitcast(Int32).ir_value(loc=loc, ip=ip), -+ ] -+ llvm.inline_asm( -+ None, -+ [remote_smem_ptr, *values, remote_mbar_ptr], -+ "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 " -+ "[$0], {$1, $2, $3, $4}, [$5];", -+ "r,r,r,r,r,r", -+ has_side_effects=True, -+ is_align_stack=False, -+ asm_dialect=llvm.AsmDialect.AD_ATT, -+ ) -+ -+ -+class _SplitKOneTileScheduler: -+ """One-work-item scheduler for a physical K-cluster. -+ -+ Grid X/Y identify the output tile. Every CTA in grid Z belongs to the same -+ physical cluster and uses its cluster rank as the logical K-slice index. -+ """ -+ -+ def __init__(self, tile_idx, num_tiles_executed): -+ self._tile_idx = tile_idx -+ self._num_tiles_executed = num_tiles_executed -+ self._tile_idx_num_values = None -+ -+ @staticmethod -+ @cute.jit -+ def create(block_idx, cta_rank_in_cluster): -+ tile_idx = ( -+ cutlass.Int32(block_idx[0]), -+ cutlass.Int32(block_idx[1]), -+ cutlass.Int32(cta_rank_in_cluster), -+ ) -+ return _SplitKOneTileScheduler(tile_idx, cutlass.Int32(0)) -+ -+ def __extract_mlir_values__(self): -+ tile_idx_values = extract_mlir_values(self._tile_idx) -+ self._tile_idx_num_values = len(tile_idx_values) -+ return tile_idx_values + extract_mlir_values(self._num_tiles_executed) -+ -+ def __new_from_mlir_values__(self, values): -+ if self._tile_idx_num_values is None: -+ raise ValueError("Split-K scheduler tile width was not recorded") -+ width = self._tile_idx_num_values -+ tile_idx = new_from_mlir_values(self._tile_idx, values[:width]) -+ num_tiles_executed = new_from_mlir_values( -+ self._num_tiles_executed, values[width:] -+ ) -+ return _SplitKOneTileScheduler(tile_idx, num_tiles_executed) -+ -+ @cute.jit -+ def get_current_work(self): -+ return utils.WorkTileInfo( -+ self._tile_idx, -+ self._num_tiles_executed < cutlass.Int32(1), -+ ) -+ -+ @cute.jit -+ def initial_work_tile_info(self): -+ return self.get_current_work() -+ -+ @cute.jit -+ def advance_to_next_work(self): -+ self._num_tiles_executed += cutlass.Int32(1) -+ -+ @property -+ @cute.jit -+ def num_tiles_executed(self): -+ return self._num_tiles_executed -+ -+ -+class Sm100BlockScaledSplitKGemmKernel(_Sm100BlockScaledGemmCommon): -+ """Low-M MXFP8 GEMM with an in-kernel split-K reduction. -+ -+ The kernel is specialized for the swap_ab=True dense MXFP8 path with -+ R128C4 scale-factor layout. A physical (1, 1, split_k_slices) cluster -+ computes one output tile. Each CTA accumulates a disjoint K slice in FP32; -+ peer CTAs send their partials through distributed shared memory, and the -+ owner CTA writes the final BF16/FP16 output without a global workspace or -+ second reduction kernel. -+ -+ The supported tactics use an MMA-M tile of 128, MMA-N tiles of 8/16/32, -+ and two or four K slices. -+ """ -+ -+ def __init__( -+ self, -+ sf_vec_size: int, -+ mma_tiler_mn: Tuple[int, int], -+ split_k_slices: int, -+ enable_pdl: bool = True, -+ ): -+ """Configure one low-M split-K tactic. -+ -+ Args: -+ sf_vec_size: MXFP8 scale-factor vector size; must be 32. -+ mma_tiler_mn: MMA tile shape; supported values are (128, 8), -+ (128, 16), and (128, 32). -+ split_k_slices: Physical cluster-K size; must be 2 or 4. -+ enable_pdl: Whether to enable programmatic dependent launch. -+ """ -+ -+ self.acc_dtype = cutlass.Float32 -+ if sf_vec_size != 32: -+ raise ValueError( -+ f"MXFP8 split-K requires sf_vec_size=32, got {sf_vec_size}" -+ ) -+ self.sf_vec_size = sf_vec_size -+ if mma_tiler_mn not in ((128, 8), (128, 16), (128, 32)): -+ raise ValueError( -+ "MXFP8 split-K requires mma_tiler_mn=(128, 8/16/32), " -+ f"got {mma_tiler_mn}" -+ ) -+ if split_k_slices not in (2, 4): -+ raise ValueError(f"split_k_slices must be 2 or 4, got {split_k_slices}") -+ self.cluster_shape_mn = (1, 1) -+ self.cluster_shape_mnk = (1, 1, split_k_slices) -+ # K dimension is deferred in _setup_attributes -+ self.mma_tiler = (*mma_tiler_mn, 1) -+ self.enable_pdl = enable_pdl -+ self.split_k_slices = split_k_slices -+ self.cta_group = tcgen05.CtaGroup.ONE -+ -+ self.occupancy = 1 -+ # Set specialized warp ids -+ self.epilog_warp_id = ( -+ 0, -+ 1, -+ 2, -+ 3, -+ ) -+ self.mma_warp_id = 4 -+ self.tma_warp_id = 5 -+ self.threads_per_cta = 32 * len( -+ (self.mma_warp_id, self.tma_warp_id, *self.epilog_warp_id) -+ ) -+ # Set barrier id for cta sync, epilogue sync and tmem ptr sync -+ self.cta_sync_bar_id = 0 -+ self.epilog_sync_bar_id = 1 -+ self.tmem_ptr_sync_bar_id = 2 -+ self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") -+ SM100_TMEM_CAPACITY_COLUMNS = 512 -+ self.num_tmem_alloc_cols = SM100_TMEM_CAPACITY_COLUMNS -+ -+ def _setup_attributes(self): -+ """Set up configurations that are dependent on GEMM inputs -+ -+ This method configures various attributes based on the input tensor properties -+ (data types, leading dimensions) and kernel settings: -+ - Configuring tiled MMA -+ - Computing MMA/cluster/tile shapes -+ - Computing cluster layout -+ - Computing multicast CTAs for A/B/SFA/SFB -+ - Computing epilogue subtile -+ - Setting up A/B/SFA/SFB/C stage counts in shared memory -+ - Computing A/B/SFA/SFB/C shared memory layout -+ - Computing tensor memory allocation columns -+ """ -+ # Compute mma instruction shapes -+ mma_inst_bits_k = 256 -+ # (MMA_Tile_Shape_M, MMA_Tile_Shape_N, MMA_Inst_Shape_K) -+ self.mma_inst_shape_mnk = ( -+ self.mma_tiler[0], -+ self.mma_tiler[1], -+ mma_inst_bits_k // self.a_dtype.width, -+ ) -+ # (CTA_Tile_Shape_M, Round_Up(MMA_Tile_Shape_N, 128), MMA_Inst_Shape_K) -+ self.mma_inst_shape_mnk_sfb = ( -+ self.mma_inst_shape_mnk[0], -+ cute.round_up(self.mma_inst_shape_mnk[1], 128), -+ self.mma_inst_shape_mnk[2], -+ ) -+ -+ tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma( -+ self.a_dtype, -+ self.b_dtype, -+ self.a_major_mode, -+ self.b_major_mode, -+ self.sf_dtype, -+ self.sf_vec_size, -+ self.cta_group, -+ self.mma_inst_shape_mnk[:2], -+ ) -+ -+ tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma( -+ self.a_dtype, -+ self.b_dtype, -+ self.a_major_mode, -+ self.b_major_mode, -+ self.sf_dtype, -+ self.sf_vec_size, -+ cute.nvgpu.tcgen05.CtaGroup.ONE, -+ self.mma_inst_shape_mnk_sfb[:2], -+ ) -+ -+ # Compute mma/cluster/tile shapes -+ mma_inst_tile_k = 4 -+ self.mma_tiler = ( -+ self.mma_inst_shape_mnk[0], -+ self.mma_inst_shape_mnk[1], -+ self.mma_inst_shape_mnk[2] * mma_inst_tile_k, -+ ) -+ self.mma_tiler_sfb = ( -+ self.mma_inst_shape_mnk_sfb[0], -+ self.mma_inst_shape_mnk_sfb[1], -+ self.mma_inst_shape_mnk_sfb[2] * mma_inst_tile_k, -+ ) -+ self.cta_tile_shape_mnk = ( -+ self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), -+ self.mma_tiler[1], -+ self.mma_tiler[2], -+ ) -+ self.cta_tile_shape_mnk_sfb = ( -+ self.mma_tiler_sfb[0] // cute.size(tiled_mma.thr_id.shape), -+ self.mma_tiler_sfb[1], -+ self.mma_tiler_sfb[2], -+ ) -+ -+ # Compute cluster layout -+ self.cluster_layout_vmnk = cute.tiled_divide( -+ cute.make_layout(self.cluster_shape_mnk), -+ (tiled_mma.thr_id.shape,), -+ ) -+ self.cluster_layout_sfb_vmnk = cute.tiled_divide( -+ cute.make_layout(self.cluster_shape_mnk), -+ (tiled_mma_sfb.thr_id.shape,), -+ ) -+ -+ # Compute number of multicast CTAs for A/B -+ self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2]) -+ self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1]) -+ self.num_mcast_ctas_sfb = cute.size(self.cluster_layout_sfb_vmnk.shape[1]) -+ self.is_a_mcast = self.num_mcast_ctas_a > 1 -+ self.is_b_mcast = self.num_mcast_ctas_b > 1 -+ self.is_sfb_mcast = self.num_mcast_ctas_sfb > 1 -+ -+ # Compute epilogue subtile -+ self.epi_tile = sm100_utils.compute_epilogue_tile_shape( -+ self.cta_tile_shape_mnk, -+ False, -+ self.c_layout, -+ self.c_dtype, -+ ) -+ -+ self.epi_tile_n = cute.size(self.epi_tile[1]) -+ self.num_epilogue_subtiles = ( -+ self.cta_tile_shape_mnk[0] -+ * self.cta_tile_shape_mnk[1] -+ // (cute.size(self.epi_tile[0]) * cute.size(self.epi_tile[1])) -+ ) -+ self.num_epilogue_threads = 32 * len(self.epilog_warp_id) -+ self.reduce_values_per_thread = ( -+ cute.size(self.epi_tile[0]) * cute.size(self.epi_tile[1]) -+ ) // (self.num_epilogue_threads) -+ self.reduce_mailbox_elements = ( -+ (self.split_k_slices - 1) -+ * self.cta_tile_shape_mnk[0] -+ * self.cta_tile_shape_mnk[1] -+ ) -+ self.reduce_smem_bytes = ( -+ self.reduce_mailbox_elements * cutlass.Float32.width // 8 -+ + self.num_epilogue_subtiles * 8 -+ ) -+ -+ # Setup A/B/C stage count in shared memory and ACC stage count in tensor memory -+ self.num_acc_stage, self.num_ab_stage, self.num_c_stage = self._compute_stages( -+ tiled_mma, -+ self.mma_tiler, -+ self.a_dtype, -+ self.a_major_mode, -+ self.b_dtype, -+ self.b_major_mode, -+ self.epi_tile, -+ self.c_dtype, -+ self.c_layout, -+ self.sf_dtype, -+ self.sf_vec_size, -+ self.smem_capacity, -+ self.occupancy, -+ 1024 + self.reduce_smem_bytes, -+ ) -+ -+ # Compute A/B/SFA/SFB/C shared memory layout -+ self.a_smem_layout_staged = sm100_utils.make_smem_layout_a( -+ tiled_mma, -+ self.mma_tiler, -+ self.a_dtype, -+ self.num_ab_stage, -+ ) -+ self.b_smem_layout_staged = sm100_utils.make_smem_layout_b( -+ tiled_mma, -+ self.mma_tiler, -+ self.b_dtype, -+ self.num_ab_stage, -+ ) -+ self.sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa( -+ tiled_mma, -+ self.mma_tiler, -+ self.sf_vec_size, -+ self.num_ab_stage, -+ ) -+ self.sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb( -+ tiled_mma, -+ self.mma_tiler, -+ self.sf_vec_size, -+ self.num_ab_stage, -+ ) -+ self.c_smem_layout_staged = sm100_utils.make_smem_layout_epi( -+ self.c_dtype, -+ self.c_layout, -+ self.epi_tile, -+ self.num_c_stage, -+ ) -+ -+ self.overlapping_accum = self.num_acc_stage == 1 -+ sf_atom_mn = 32 -+ self.num_sfa_tmem_cols = ( -+ self.cta_tile_shape_mnk[0] // sf_atom_mn -+ ) * mma_inst_tile_k -+ self.num_sfb_tmem_cols = ( -+ self.cta_tile_shape_mnk_sfb[1] // sf_atom_mn -+ ) * mma_inst_tile_k -+ self.num_sf_tmem_cols = self.num_sfa_tmem_cols + self.num_sfb_tmem_cols -+ self.num_accumulator_tmem_cols = ( -+ self.cta_tile_shape_mnk[1] * self.num_acc_stage -+ if not self.overlapping_accum -+ else self.cta_tile_shape_mnk[1] * 2 - self.num_sf_tmem_cols -+ ) -+ -+ # Only when overlapping_accum is enabled, we need to release accumulator buffer early in epilogue -+ self.iter_acc_early_release_in_epilogue = ( -+ self.num_sf_tmem_cols // self.epi_tile_n -+ ) -+ -+ @cute.jit -+ def __call__( -+ self, -+ a_tensor: cute.Tensor, -+ b_tensor: cute.Tensor, -+ sfa_tensor: cute.Tensor, -+ sfb_tensor: cute.Tensor, -+ c_tensor: cute.Tensor, -+ alpha: cute.Tensor, # Single-element tensor containing alpha value -+ max_active_clusters: cutlass.Constexpr, -+ stream: cuda.CUstream, -+ epilogue_op: cutlass.Constexpr = lambda x: x, -+ ): -+ """Execute the GEMM operation in steps: -+ - Setup static attributes before smem/grid/tma computation -+ - Setup TMA load/store atoms and tensors -+ - Compute grid size with regard to hardware constraints -+ - Define shared storage for kernel -+ - Launch the kernel synchronously -+ -+ Args: -+ a_tensor (cute.Tensor): Input tensor A -+ b_tensor (cute.Tensor): Input tensor B -+ sfa_tensor (cute.Tensor): Scale factor tensor A -+ sfb_tensor (cute.Tensor): Scale factor tensor B -+ c_tensor (cute.Tensor): Output tensor C -+ max_active_clusters (cutlass.Constexpr): Maximum number of active clusters -+ stream (cuda.CUstream): CUDA stream for asynchronous execution -+ epilogue_op (cutlass.Constexpr): Optional elementwise lambda function to apply to the output tensor -+ -+ Raises: -+ TypeError: If input data types are incompatible with the MMA instruction. -+ """ -+ ( -+ tiled_mma, -+ tiled_mma_sfb, -+ tma_atom_a, -+ tma_tensor_a, -+ tma_atom_b, -+ tma_tensor_b, -+ tma_atom_sfa, -+ tma_tensor_sfa, -+ tma_atom_sfb, -+ tma_tensor_sfb, -+ tma_atom_c, -+ tma_tensor_c, -+ ) = self._prepare_tma( -+ a_tensor, -+ b_tensor, -+ sfa_tensor, -+ sfb_tensor, -+ c_tensor, -+ ) -+ -+ # Compute grid size -+ grid = self._compute_grid(c_tensor, self.cta_tile_shape_mnk) -+ -+ self.buffer_align_bytes = 1024 -+ -+ # Define shared storage for kernel -+ @cute.struct -+ class SharedStorage: -+ ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage] -+ ab_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage] -+ acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage] -+ acc_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage] -+ tmem_holding_buf: cutlass.Int32 -+ reduce_full_mbar_ptr: cute.struct.MemRange[ -+ cutlass.Int64, self.num_epilogue_subtiles -+ ] -+ reduce_mailbox: cute.struct.Align[ -+ cute.struct.MemRange[cutlass.Float32, self.reduce_mailbox_elements], -+ 128, -+ ] -+ # (EPI_TILE_M, EPI_TILE_N, STAGE) -+ sC: cute.struct.Align[ -+ cute.struct.MemRange[ -+ self.c_dtype, -+ cute.cosize(self.c_smem_layout_staged.outer), -+ ], -+ self.buffer_align_bytes, -+ ] -+ # (MMA, MMA_M, MMA_K, STAGE) -+ sA: cute.struct.Align[ -+ cute.struct.MemRange[ -+ self.a_dtype, cute.cosize(self.a_smem_layout_staged.outer) -+ ], -+ self.buffer_align_bytes, -+ ] -+ # (MMA, MMA_N, MMA_K, STAGE) -+ sB: cute.struct.Align[ -+ cute.struct.MemRange[ -+ self.b_dtype, cute.cosize(self.b_smem_layout_staged.outer) -+ ], -+ self.buffer_align_bytes, -+ ] -+ # (MMA, MMA_M, MMA_K, STAGE) -+ sSFA: cute.struct.Align[ -+ cute.struct.MemRange[ -+ self.sf_dtype, cute.cosize(self.sfa_smem_layout_staged) -+ ], -+ self.buffer_align_bytes, -+ ] -+ # (MMA, MMA_N, MMA_K, STAGE) -+ sSFB: cute.struct.Align[ -+ cute.struct.MemRange[ -+ self.sf_dtype, cute.cosize(self.sfb_smem_layout_staged) -+ ], -+ self.buffer_align_bytes, -+ ] -+ -+ self.shared_storage = SharedStorage -+ -+ # Launch the kernel synchronously -+ self.kernel( -+ tiled_mma, -+ tiled_mma_sfb, -+ tma_atom_a, -+ tma_tensor_a, -+ tma_atom_b, -+ tma_tensor_b, -+ tma_atom_sfa, -+ tma_tensor_sfa, -+ tma_atom_sfb, -+ tma_tensor_sfb, -+ tma_atom_c, -+ tma_tensor_c, -+ self.cluster_layout_vmnk, -+ self.cluster_layout_sfb_vmnk, -+ self.a_smem_layout_staged, -+ self.b_smem_layout_staged, -+ self.sfa_smem_layout_staged, -+ self.sfb_smem_layout_staged, -+ self.c_smem_layout_staged, -+ self.epi_tile, -+ epilogue_op, -+ alpha, -+ ).launch( -+ grid=grid, -+ block=[self.threads_per_cta, 1, 1], -+ cluster=self.cluster_shape_mnk, -+ smem=self.shared_storage.size_in_bytes(), # type: ignore[attr-defined] -+ min_blocks_per_mp=1, -+ stream=stream, -+ use_pdl=self.enable_pdl, -+ ) -+ return -+ -+ # GPU device kernel -+ @cute.kernel -+ def kernel( -+ self, -+ tiled_mma: cute.TiledMma, -+ tiled_mma_sfb: cute.TiledMma, -+ tma_atom_a: cute.CopyAtom, -+ mA_mkl: cute.Tensor, -+ tma_atom_b: cute.CopyAtom, -+ mB_nkl: cute.Tensor, -+ tma_atom_sfa: cute.CopyAtom, -+ mSFA_mkl: cute.Tensor, -+ tma_atom_sfb: cute.CopyAtom, -+ mSFB_nkl: cute.Tensor, -+ tma_atom_c: Optional[cute.CopyAtom], -+ mC_mnl: cute.Tensor, -+ cluster_layout_vmnk: cute.Layout, -+ cluster_layout_sfb_vmnk: cute.Layout, -+ a_smem_layout_staged: cute.ComposedLayout, -+ b_smem_layout_staged: cute.ComposedLayout, -+ sfa_smem_layout_staged: cute.Layout, -+ sfb_smem_layout_staged: cute.Layout, -+ c_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout, None], -+ epi_tile: cute.Tile, -+ epilogue_op: cutlass.Constexpr, -+ alpha: cute.Tensor, -+ ): -+ """ -+ GPU device kernel performing the Persistent batched GEMM computation. -+ """ -+ # Keep alpha in FP32 for precision: the accumulator is in FP32 and alpha -+ # may be a very small scaling factor. Converting to c_dtype (e.g., FP16) -+ # before multiplication could cause overflow when acc values are large. -+ alpha_value = alpha[0].to(cutlass.Float32) -+ -+ warp_idx = cute.arch.warp_idx() -+ warp_idx = cute.arch.make_warp_uniform(warp_idx) -+ -+ # -+ # Prefetch tma desc -+ # -+ if warp_idx == self.tma_warp_id: -+ cpasync.prefetch_descriptor(tma_atom_a) -+ cpasync.prefetch_descriptor(tma_atom_b) -+ cpasync.prefetch_descriptor(tma_atom_sfa) -+ cpasync.prefetch_descriptor(tma_atom_sfb) -+ cpasync.prefetch_descriptor(tma_atom_c) -+ -+ # -+ # Setup cta/thread coordinates -+ # -+ # Coords inside cluster -+ bidx, bidy, bidz = cute.arch.block_idx() -+ mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape) -+ is_leader_cta = mma_tile_coord_v == 0 -+ cta_rank_in_cluster = cute.arch.make_warp_uniform( -+ cute.arch.block_idx_in_cluster() -+ ) -+ block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord( -+ cta_rank_in_cluster -+ ) -+ block_in_cluster_coord_sfb_vmnk = cluster_layout_sfb_vmnk.get_flat_coord( -+ cta_rank_in_cluster -+ ) -+ # Coord inside cta -+ tidx, _, _ = cute.arch.thread_idx() -+ -+ # -+ # Alloc and init: a+b full/empty, accumulator full/empty, tensor memory dealloc barrier -+ # -+ smem = utils.SmemAllocator() -+ storage = smem.allocate(self.shared_storage) -+ -+ tmem_holding_buf = storage.tmem_holding_buf.ptr -+ reduce_full_mbar_ptr = storage.reduce_full_mbar_ptr.data_ptr() -+ -+ # Initialize mainloop ab_pipeline (barrier) and states -+ ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) -+ num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 -+ ab_pipeline_consumer_group = pipeline.CooperativeGroup( -+ pipeline.Agent.Thread, num_tma_producer -+ ) -+ ab_pipeline = PipelineTmaUmma.create( -+ barrier_storage=storage.ab_full_mbar_ptr.data_ptr(), -+ num_stages=self.num_ab_stage, -+ producer_group=ab_pipeline_producer_group, -+ consumer_group=ab_pipeline_consumer_group, -+ tx_count=self.num_tma_load_bytes, -+ cta_layout_vmnk=cluster_layout_vmnk, -+ ) -+ -+ # Initialize acc_pipeline (barrier) and states -+ acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) -+ num_acc_consumer_threads = len(self.epilog_warp_id) -+ acc_pipeline_consumer_group = pipeline.CooperativeGroup( -+ pipeline.Agent.Thread, num_acc_consumer_threads -+ ) -+ acc_pipeline = PipelineUmmaAsync.create( -+ barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), -+ num_stages=self.num_acc_stage, -+ producer_group=acc_pipeline_producer_group, -+ consumer_group=acc_pipeline_consumer_group, -+ cta_layout_vmnk=cluster_layout_vmnk, -+ ) -+ -+ if warp_idx == self.tma_warp_id: -+ with cute.arch.elect_one(): -+ for subtile_idx in cutlass.range_constexpr(self.num_epilogue_subtiles): -+ cute.arch.mbarrier_init(reduce_full_mbar_ptr + subtile_idx, 1) -+ cute.arch.mbarrier_init_fence() -+ -+ # Cluster arrive after barrier init -+ if cute.size(self.cluster_shape_mnk) > 1: -+ cute.arch.cluster_arrive_relaxed(aligned=True) -+ -+ # -+ # Setup smem tensor A/B/SFA/SFB/C -+ reduce_mailbox = storage.reduce_mailbox.get_tensor( -+ cute.make_layout((self.reduce_mailbox_elements,)) -+ ) -+ # -+ # (EPI_TILE_M, EPI_TILE_N, STAGE) -+ sC = storage.sC.get_tensor( -+ c_smem_layout_staged.outer, swizzle=c_smem_layout_staged.inner -+ ) -+ # (MMA, MMA_M, MMA_K, STAGE) -+ sA = storage.sA.get_tensor( -+ a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner -+ ) -+ # (MMA, MMA_N, MMA_K, STAGE) -+ sB = storage.sB.get_tensor( -+ b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner -+ ) -+ # (MMA, MMA_M, MMA_K, STAGE) -+ sSFA = storage.sSFA.get_tensor(sfa_smem_layout_staged) -+ # (MMA, MMA_N, MMA_K, STAGE) -+ sSFB = storage.sSFB.get_tensor(sfb_smem_layout_staged) -+ -+ # -+ # Compute multicast mask for A/B/SFA/SFB buffer full -+ # -+ a_full_mcast_mask = None -+ b_full_mcast_mask = None -+ sfa_full_mcast_mask = None -+ sfb_full_mcast_mask = None -+ if cutlass.const_expr(self.is_a_mcast or self.is_b_mcast): -+ a_full_mcast_mask = cpasync.create_tma_multicast_mask( -+ cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 -+ ) -+ b_full_mcast_mask = cpasync.create_tma_multicast_mask( -+ cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1 -+ ) -+ sfa_full_mcast_mask = cpasync.create_tma_multicast_mask( -+ cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 -+ ) -+ sfb_full_mcast_mask = cpasync.create_tma_multicast_mask( -+ cluster_layout_sfb_vmnk, block_in_cluster_coord_sfb_vmnk, mcast_mode=1 -+ ) -+ -+ # -+ # Local_tile partition global tensors -+ # -+ # (bM, bK, RestM, RestK, RestL) -+ gA_mkl = cute.local_tile( -+ mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) -+ ) -+ # (bN, bK, RestN, RestK, RestL) -+ gB_nkl = cute.local_tile( -+ mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None) -+ ) -+ # (bM, bK, RestM, RestK, RestL) -+ gSFA_mkl = cute.local_tile( -+ mSFA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) -+ ) -+ # (bN, bK, RestN, RestK, RestL) -+ gSFB_nkl = cute.local_tile( -+ mSFB_nkl, -+ cute.slice_(self.mma_tiler_sfb, (0, None, None)), -+ (None, None, None), -+ ) -+ # (bM, bN, RestM, RestN, RestL) -+ gC_mnl = cute.local_tile( -+ mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) -+ ) -+ total_k_block_cnt = cutlass.Int32(cute.size(gA_mkl, mode=[3])) -+ k_block_cnt = total_k_block_cnt -+ if cutlass.const_expr(self.split_k_slices > 1): -+ k_block_cnt = total_k_block_cnt // self.split_k_slices -+ -+ # -+ # Partition global tensor for TiledMMA_A/B/C -+ # -+ thr_mma = tiled_mma.get_slice(mma_tile_coord_v) -+ thr_mma_sfb = tiled_mma_sfb.get_slice(mma_tile_coord_v) -+ # (MMA, MMA_M, MMA_K, RestM, RestK, RestL) -+ tCgA = thr_mma.partition_A(gA_mkl) -+ # (MMA, MMA_N, MMA_K, RestN, RestK, RestL) -+ tCgB = thr_mma.partition_B(gB_nkl) -+ # (MMA, MMA_M, MMA_K, RestM, RestK, RestL) -+ tCgSFA = thr_mma.partition_A(gSFA_mkl) -+ # (MMA, MMA_N, MMA_K, RestN, RestK, RestL) -+ tCgSFB = thr_mma_sfb.partition_B(gSFB_nkl) -+ # (MMA, MMA_M, MMA_N, RestM, RestN, RestL) -+ tCgC = thr_mma.partition_C(gC_mnl) -+ -+ # -+ # Partition global/shared tensor for TMA load A/B -+ # -+ # TMA load A partition_S/D -+ a_cta_layout = cute.make_layout( -+ cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape -+ ) -+ # ((atom_v, rest_v), STAGE) -+ # ((atom_v, rest_v), RestM, RestK, RestL) -+ tAsA, tAgA = cpasync.tma_partition( -+ tma_atom_a, -+ block_in_cluster_coord_vmnk[2], -+ a_cta_layout, -+ cute.group_modes(sA, 0, 3), -+ cute.group_modes(tCgA, 0, 3), -+ ) -+ # TMA load B partition_S/D -+ b_cta_layout = cute.make_layout( -+ cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape -+ ) -+ # ((atom_v, rest_v), STAGE) -+ # ((atom_v, rest_v), RestN, RestK, RestL) -+ tBsB, tBgB = cpasync.tma_partition( -+ tma_atom_b, -+ block_in_cluster_coord_vmnk[1], -+ b_cta_layout, -+ cute.group_modes(sB, 0, 3), -+ cute.group_modes(tCgB, 0, 3), -+ ) -+ -+ # TMA load SFA partition_S/D -+ sfa_cta_layout = a_cta_layout -+ # ((atom_v, rest_v), STAGE) -+ # ((atom_v, rest_v), RestM, RestK, RestL) -+ tAsSFA, tAgSFA = cute.nvgpu.cpasync.tma_partition( -+ tma_atom_sfa, -+ block_in_cluster_coord_vmnk[2], -+ sfa_cta_layout, -+ cute.group_modes(sSFA, 0, 3), -+ cute.group_modes(tCgSFA, 0, 3), -+ ) -+ tAsSFA = cute.filter_zeros(tAsSFA) -+ tAgSFA = cute.filter_zeros(tAgSFA) -+ -+ # TMA load SFB partition_S/D -+ sfb_cta_layout = cute.make_layout( -+ cute.slice_(cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape -+ ) -+ # ((atom_v, rest_v), STAGE) -+ # ((atom_v, rest_v), RestN, RestK, RestL) -+ tBsSFB, tBgSFB = cute.nvgpu.cpasync.tma_partition( -+ tma_atom_sfb, -+ block_in_cluster_coord_sfb_vmnk[1], -+ sfb_cta_layout, -+ cute.group_modes(sSFB, 0, 3), -+ cute.group_modes(tCgSFB, 0, 3), -+ ) -+ tBsSFB = cute.filter_zeros(tBsSFB) -+ tBgSFB = cute.filter_zeros(tBgSFB) -+ -+ # -+ # Partition shared/tensor memory tensor for TiledMMA_A/B/C -+ # -+ # (MMA, MMA_M, MMA_K, STAGE) -+ tCrA = tiled_mma.make_fragment_A(sA) -+ # (MMA, MMA_N, MMA_K, STAGE) -+ tCrB = tiled_mma.make_fragment_B(sB) -+ # (MMA, MMA_M, MMA_N) -+ acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) -+ # (MMA, MMA_M, MMA_N, STAGE) -+ if cutlass.const_expr(self.overlapping_accum): -+ num_acc_stage_overlapped = 2 -+ tCtAcc_fake = tiled_mma.make_fragment_C( -+ cute.append(acc_shape, num_acc_stage_overlapped) -+ ) -+ # (MMA, MMA_M, MMA_N, STAGE) -+ tCtAcc_fake = cute.make_tensor( -+ tCtAcc_fake.iterator, -+ cute.make_layout( -+ tCtAcc_fake.shape, -+ stride=( -+ tCtAcc_fake.stride[0], -+ tCtAcc_fake.stride[1], -+ tCtAcc_fake.stride[2], -+ (256 - self.num_sf_tmem_cols) * tCtAcc_fake.stride[0][1], -+ ), -+ ), -+ ) -+ else: -+ # (MMA, MMA_M, MMA_N, STAGE) -+ tCtAcc_fake = tiled_mma.make_fragment_C( -+ cute.append(acc_shape, self.num_acc_stage) -+ ) -+ -+ # -+ # Cluster wait before tensor memory alloc -+ # -+ if cute.size(self.cluster_shape_mnk) > 1: -+ cute.arch.cluster_wait() -+ else: -+ cute.arch.barrier( -+ barrier_id=self.cta_sync_bar_id, number_of_threads=self.threads_per_cta -+ ) -+ -+ # PDL bookend: griddepcontrol_wait is always emitted; the actual PDL -+ # behavior is controlled by use_pdl= in .launch(). The griddepcontrol -+ # instructions are effectively no-ops when PDL is not enabled at -+ # the driver level. -+ griddepcontrol_wait() -+ -+ # -+ # Specialized TMA load warp -+ # -+ if warp_idx == self.tma_warp_id: -+ # -+ # Persistent tile scheduling loop -+ # -+ tile_sched = _SplitKOneTileScheduler.create( -+ cute.arch.block_idx(), cta_rank_in_cluster -+ ) -+ work_tile = tile_sched.initial_work_tile_info() -+ -+ ab_producer_state = pipeline.make_pipeline_state( -+ pipeline.PipelineUserType.Producer, self.num_ab_stage -+ ) -+ -+ while work_tile.is_valid_tile: -+ # Get tile coord from tile scheduler -+ cur_tile_coord = work_tile.tile_idx -+ mma_tile_coord_mnl = ( -+ cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), -+ cur_tile_coord[1], -+ cur_tile_coord[2], -+ ) -+ -+ input_l = mma_tile_coord_mnl[2] -+ k_block_start = cutlass.Int32(0) -+ if cutlass.const_expr(self.split_k_slices > 1): -+ # C L is the split-K index; dense inputs always use L=0. -+ input_l = cutlass.Int32(0) -+ k_block_start = cutlass.Int32(mma_tile_coord_mnl[2]) * k_block_cnt -+ -+ # -+ # Slice to per mma tile index -+ # -+ # ((atom_v, rest_v), RestK) -+ tAgA_slice = tAgA[(None, mma_tile_coord_mnl[0], None, input_l)] -+ # ((atom_v, rest_v), RestK) -+ tBgB_slice = tBgB[(None, mma_tile_coord_mnl[1], None, input_l)] -+ -+ # ((atom_v, rest_v), RestK) -+ tAgSFA_slice = tAgSFA[(None, mma_tile_coord_mnl[0], None, input_l)] -+ slice_n = mma_tile_coord_mnl[1] -+ if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64): -+ slice_n = mma_tile_coord_mnl[1] // 2 -+ # ((atom_v, rest_v), RestK) -+ tBgSFB_slice = tBgSFB[(None, slice_n, None, input_l)] -+ -+ # Peek (try_wait) AB buffer empty for k_block = prefetch_k_block_cnt -+ ab_producer_state.reset_count() -+ peek_ab_empty_status = cutlass.Boolean(1) -+ if ab_producer_state.count < k_block_cnt: -+ peek_ab_empty_status = ab_pipeline.producer_try_acquire( -+ ab_producer_state -+ ) -+ # -+ # Tma load loop -+ # -+ for _k_block in cutlass.range(0, k_block_cnt, 1, unroll=1): -+ # Conditionally wait for AB buffer empty -+ ab_pipeline.producer_acquire( -+ ab_producer_state, peek_ab_empty_status -+ ) -+ -+ # TMA load A/B/SFA/SFB -+ cute.copy( -+ tma_atom_a, -+ tAgA_slice[(None, k_block_start + ab_producer_state.count)], -+ tAsA[(None, ab_producer_state.index)], -+ tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), -+ mcast_mask=a_full_mcast_mask, -+ ) -+ cute.copy( -+ tma_atom_b, -+ tBgB_slice[(None, k_block_start + ab_producer_state.count)], -+ tBsB[(None, ab_producer_state.index)], -+ tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), -+ mcast_mask=b_full_mcast_mask, -+ ) -+ cute.copy( -+ tma_atom_sfa, -+ tAgSFA_slice[(None, k_block_start + ab_producer_state.count)], -+ tAsSFA[(None, ab_producer_state.index)], -+ tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), -+ mcast_mask=sfa_full_mcast_mask, -+ ) -+ cute.copy( -+ tma_atom_sfb, -+ tBgSFB_slice[(None, k_block_start + ab_producer_state.count)], -+ tBsSFB[(None, ab_producer_state.index)], -+ tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), -+ mcast_mask=sfb_full_mcast_mask, -+ ) -+ -+ # Peek (try_wait) AB buffer empty for k_block = prefetch_k_block_cnt + k_block + 1 -+ ab_producer_state.advance() -+ peek_ab_empty_status = cutlass.Boolean(1) -+ if ab_producer_state.count < k_block_cnt: -+ peek_ab_empty_status = ab_pipeline.producer_try_acquire( -+ ab_producer_state -+ ) -+ -+ # -+ # Advance to next tile -+ # -+ tile_sched.advance_to_next_work() -+ work_tile = tile_sched.get_current_work() -+ -+ # -+ # Wait A/B buffer empty -+ # -+ ab_pipeline.producer_tail(ab_producer_state) -+ -+ # -+ # Specialized MMA warp -+ # -+ if warp_idx == self.mma_warp_id: -+ # -+ # Bar sync for retrieve tensor memory ptr from shared mem -+ # -+ tmem_ptr_read_threads = 32 * len((self.mma_warp_id, *self.epilog_warp_id)) -+ cute.arch.barrier( -+ barrier_id=self.tmem_ptr_sync_bar_id, -+ number_of_threads=tmem_ptr_read_threads, -+ ) -+ -+ # -+ # Retrieving tensor memory ptr and make accumulator/SFA/SFB tensor -+ # -+ # Make accumulator tmem tensor -+ acc_tmem_ptr = cute.arch.retrieve_tmem_ptr( -+ self.acc_dtype, -+ alignment=16, -+ ptr_to_buffer_holding_addr=tmem_holding_buf, -+ ) -+ # (MMA, MMA_M, MMA_N, STAGE) -+ tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) -+ -+ # Make SFA tmem tensor -+ sfa_tmem_ptr = cute.recast_ptr( -+ acc_tmem_ptr + self.num_accumulator_tmem_cols, -+ dtype=self.sf_dtype, -+ ) -+ # (MMA, MMA_M, MMA_K) -+ tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa( -+ tiled_mma, -+ self.mma_tiler, -+ self.sf_vec_size, -+ cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)), -+ ) -+ tCtSFA = cute.make_tensor(sfa_tmem_ptr, tCtSFA_layout) -+ -+ # Make SFB tmem tensor -+ sfb_tmem_ptr = cute.recast_ptr( -+ acc_tmem_ptr + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols, -+ dtype=self.sf_dtype, -+ ) -+ # (MMA, MMA_N, MMA_K) -+ tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb( -+ tiled_mma, -+ self.mma_tiler, -+ self.sf_vec_size, -+ cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)), -+ ) -+ tCtSFB = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout) -+ # -+ # Partition for S2T copy of SFA/SFB -+ # -+ tiled_copy_s2t_sfa, tCsSFA_compact_s2t, tCtSFA_compact_s2t = ( -+ self.mainloop_s2t_copy_and_partition(sSFA, tCtSFA) -+ ) -+ tiled_copy_s2t_sfb, tCsSFB_compact_s2t, tCtSFB_compact_s2t = ( -+ self.mainloop_s2t_copy_and_partition(sSFB, tCtSFB) -+ ) -+ -+ # -+ # Persistent tile scheduling loop -+ # -+ tile_sched = _SplitKOneTileScheduler.create( -+ cute.arch.block_idx(), cta_rank_in_cluster -+ ) -+ work_tile = tile_sched.initial_work_tile_info() -+ -+ ab_consumer_state = pipeline.make_pipeline_state( -+ pipeline.PipelineUserType.Consumer, self.num_ab_stage -+ ) -+ acc_producer_state = pipeline.make_pipeline_state( -+ pipeline.PipelineUserType.Producer, self.num_acc_stage -+ ) -+ -+ while work_tile.is_valid_tile: -+ # Get tile coord from tile scheduler -+ cur_tile_coord = work_tile.tile_idx -+ mma_tile_coord_mnl = ( -+ cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), -+ cur_tile_coord[1], -+ cur_tile_coord[2], -+ ) -+ -+ if cutlass.const_expr(self.overlapping_accum): -+ acc_stage_index = acc_producer_state.phase ^ 1 -+ else: -+ acc_stage_index = acc_producer_state.index -+ -+ # Set tensor memory buffer for current tile -+ # (MMA, MMA_M, MMA_N) -+ tCtAcc = tCtAcc_base[(None, None, None, acc_stage_index)] -+ -+ # Peek (try_wait) AB buffer full for k_block = 0 -+ ab_consumer_state.reset_count() -+ peek_ab_full_status = cutlass.Boolean(1) -+ if ab_consumer_state.count < k_block_cnt and is_leader_cta: -+ peek_ab_full_status = ab_pipeline.consumer_try_wait( -+ ab_consumer_state -+ ) -+ -+ # -+ # Wait for accumulator buffer empty -+ # -+ if is_leader_cta: -+ acc_pipeline.producer_acquire(acc_producer_state) -+ -+ tCtSFB_mma = tCtSFB -+ if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 192): -+ # If this is an ODD tile, shift the TMEM start address for cta_tile_shape_n=192 case by two words (ignores first 64 columns of SFB) -+ offset = ( -+ cutlass.Int32(2) -+ if mma_tile_coord_mnl[1] % 2 == 1 -+ else cutlass.Int32(0) -+ ) -+ shifted_ptr = cute.recast_ptr( -+ acc_tmem_ptr -+ + self.num_accumulator_tmem_cols -+ + self.num_sfa_tmem_cols -+ + offset, -+ dtype=self.sf_dtype, -+ ) -+ tCtSFB_mma = cute.make_tensor(shifted_ptr, tCtSFB_layout) -+ elif cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64): -+ # Move in increments of 64 columns of SFB -+ offset = cutlass.Int32((mma_tile_coord_mnl[1] % 2) * 2) -+ shifted_ptr = cute.recast_ptr( -+ acc_tmem_ptr -+ + self.num_accumulator_tmem_cols -+ + self.num_sfa_tmem_cols -+ + offset, -+ dtype=self.sf_dtype, -+ ) -+ tCtSFB_mma = cute.make_tensor(shifted_ptr, tCtSFB_layout) -+ -+ # -+ # Reset the ACCUMULATE field for each tile -+ # -+ tiled_mma.set(tcgen05.Field.ACCUMULATE, False) -+ -+ # -+ # Mma mainloop -+ # -+ for _k_block in range(k_block_cnt): -+ if is_leader_cta: -+ # Conditionally wait for AB buffer full -+ ab_pipeline.consumer_wait( -+ ab_consumer_state, peek_ab_full_status -+ ) -+ -+ # Copy SFA/SFB from smem to tmem -+ s2t_stage_coord = ( -+ None, -+ None, -+ None, -+ None, -+ ab_consumer_state.index, -+ ) -+ tCsSFA_compact_s2t_staged = tCsSFA_compact_s2t[s2t_stage_coord] -+ tCsSFB_compact_s2t_staged = tCsSFB_compact_s2t[s2t_stage_coord] -+ cute.copy( -+ tiled_copy_s2t_sfa, -+ tCsSFA_compact_s2t_staged, -+ tCtSFA_compact_s2t, -+ ) -+ cute.copy( -+ tiled_copy_s2t_sfb, -+ tCsSFB_compact_s2t_staged, -+ tCtSFB_compact_s2t, -+ ) -+ -+ # tCtAcc += tCrA * tCrSFA * tCrB * tCrSFB -+ num_kphases = cute.size(tCrA, mode=[2]) -+ for kphase_idx in cutlass.range(num_kphases, unroll_full=True): -+ kphase_coord = ( -+ None, -+ None, -+ kphase_idx, -+ ab_consumer_state.index, -+ ) -+ -+ # Set SFA/SFB tensor to tiled_mma -+ sf_kphase_coord = (None, None, kphase_idx) -+ tiled_mma.set( -+ tcgen05.Field.SFA, -+ tCtSFA[sf_kphase_coord].iterator, -+ ) -+ tiled_mma.set( -+ tcgen05.Field.SFB, -+ tCtSFB_mma[sf_kphase_coord].iterator, -+ ) -+ -+ cute.gemm( -+ tiled_mma, -+ tCtAcc, -+ tCrA[kphase_coord], -+ tCrB[kphase_coord], -+ tCtAcc, -+ ) -+ -+ # Enable accumulate on tCtAcc after first kphase -+ tiled_mma.set(tcgen05.Field.ACCUMULATE, True) -+ -+ # Async arrive AB buffer empty -+ ab_pipeline.consumer_release(ab_consumer_state) -+ -+ # Peek (try_wait) AB buffer full for k_block = k_block + 1 -+ ab_consumer_state.advance() -+ peek_ab_full_status = cutlass.Boolean(1) -+ if ab_consumer_state.count < k_block_cnt: -+ if is_leader_cta: -+ peek_ab_full_status = ab_pipeline.consumer_try_wait( -+ ab_consumer_state -+ ) -+ -+ # -+ # Async arrive accumulator buffer full -+ # -+ if is_leader_cta: -+ acc_pipeline.producer_commit(acc_producer_state) -+ acc_producer_state.advance() -+ -+ # -+ # Advance to next tile -+ # -+ tile_sched.advance_to_next_work() -+ work_tile = tile_sched.get_current_work() -+ -+ # -+ # Wait for accumulator buffer empty -+ # -+ acc_pipeline.producer_tail(acc_producer_state) -+ # -+ # Specialized epilogue warps -+ # -+ if warp_idx < self.mma_warp_id: -+ # -+ # Alloc tensor memory buffer -+ # -+ if warp_idx == self.epilog_warp_id[0]: -+ cute.arch.alloc_tmem( -+ self.num_tmem_alloc_cols, -+ tmem_holding_buf, -+ is_two_cta=False, -+ ) -+ -+ # -+ # Bar sync for retrieve tensor memory ptr from shared memory -+ # -+ tmem_ptr_read_threads = 32 * len((self.mma_warp_id, *self.epilog_warp_id)) -+ cute.arch.barrier( -+ barrier_id=self.tmem_ptr_sync_bar_id, -+ number_of_threads=tmem_ptr_read_threads, -+ ) -+ -+ # -+ # Retrieving tensor memory ptr and make accumulator tensor -+ # -+ acc_tmem_ptr = cute.arch.retrieve_tmem_ptr( -+ self.acc_dtype, -+ alignment=16, -+ ptr_to_buffer_holding_addr=tmem_holding_buf, -+ ) -+ # (MMA, MMA_M, MMA_N, STAGE) -+ tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) -+ -+ # -+ # Partition for epilogue -+ # -+ epi_tidx = tidx -+ tiled_copy_t2r, tTR_tAcc_base, tTR_rAcc = ( -+ self.epilog_tmem_copy_and_partition( -+ epi_tidx, tCtAcc_base, tCgC, epi_tile, False -+ ) -+ ) -+ -+ tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype) -+ tiled_copy_r2s, tRS_rC, tRS_sC = self.epilog_smem_copy_and_partition( -+ tiled_copy_t2r, tTR_rC, epi_tidx, sC -+ ) -+ tma_atom_c, bSG_sC, bSG_gC_partitioned = ( -+ self.epilog_gmem_copy_and_partition( -+ epi_tidx, tma_atom_c, tCgC, epi_tile, sC -+ ) -+ ) -+ -+ # -+ # Persistent tile scheduling loop -+ # -+ tile_sched = _SplitKOneTileScheduler.create( -+ cute.arch.block_idx(), cta_rank_in_cluster -+ ) -+ work_tile = tile_sched.initial_work_tile_info() -+ -+ acc_consumer_state = pipeline.make_pipeline_state( -+ pipeline.PipelineUserType.Consumer, self.num_acc_stage -+ ) -+ -+ # Threads/warps participating in tma store pipeline -+ c_producer_group = pipeline.CooperativeGroup( -+ pipeline.Agent.Thread, -+ 32 * len(self.epilog_warp_id), -+ ) -+ c_pipeline = pipeline.PipelineTmaStore.create( -+ num_stages=self.num_c_stage, -+ producer_group=c_producer_group, -+ ) -+ -+ while work_tile.is_valid_tile: -+ # Get tile coord from tile scheduler -+ cur_tile_coord = work_tile.tile_idx -+ mma_tile_coord_mnl = ( -+ cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), -+ cur_tile_coord[1], -+ cur_tile_coord[2], -+ ) -+ -+ # -+ # Slice to per mma tile index -+ # -+ # ((ATOM_V, REST_V), EPI_M, EPI_N) -+ bSG_gC = bSG_gC_partitioned[ -+ ( -+ None, -+ None, -+ None, -+ mma_tile_coord_mnl[0], -+ mma_tile_coord_mnl[1], -+ 0, -+ ) -+ ] -+ -+ if cutlass.const_expr(self.overlapping_accum): -+ acc_stage_index = acc_consumer_state.phase -+ reverse_subtile = ( -+ cutlass.Boolean(True) -+ if acc_stage_index == 0 -+ else cutlass.Boolean(False) -+ ) -+ else: -+ acc_stage_index = acc_consumer_state.index -+ -+ # Set tensor memory buffer for current tile -+ # (T2R, T2R_M, T2R_N, EPI_M, EPI_M) -+ tTR_tAcc = tTR_tAcc_base[ -+ (None, None, None, None, None, acc_stage_index) -+ ] -+ -+ # -+ # Wait for accumulator buffer full -+ # -+ acc_pipeline.consumer_wait(acc_consumer_state) -+ -+ tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) -+ bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC)) -+ -+ # -+ # Store accumulator to global memory in sub-tiles -+ # -+ subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) -+ -+ for subtile_idx in cutlass.range(subtile_cnt): -+ real_subtile_idx = subtile_idx -+ if cutlass.const_expr(self.overlapping_accum): -+ if reverse_subtile: -+ real_subtile_idx = ( -+ self.cta_tile_shape_mnk[1] // self.epi_tile_n -+ - 1 -+ - subtile_idx -+ ) -+ # -+ # Load accumulator from tensor memory buffer to register -+ # -+ tTR_tAcc_mn = tTR_tAcc[(None, None, None, real_subtile_idx)] -+ cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) -+ -+ # -+ # Async arrive accumulator buffer empty earlier when overlapping_accum is enabled -+ # -+ if cutlass.const_expr(self.overlapping_accum): -+ if subtile_idx == self.iter_acc_early_release_in_epilogue: -+ # Fence for TMEM load -+ cute.arch.fence_view_async_tmem_load() -+ with cute.arch.elect_one(): -+ acc_pipeline.consumer_release(acc_consumer_state) -+ acc_consumer_state.advance() -+ # Retile the FP32 accumulator fragment into the same register -+ # order used by the existing shared-memory epilogue. -+ tRS_rAcc = tiled_copy_r2s.retile(tTR_rAcc) -+ owner_rank = cutlass.Int32(real_subtile_idx % self.split_k_slices) -+ peer_count = self.split_k_slices - 1 -+ values_per_thread = self.reduce_values_per_thread -+ values_per_peer = self.num_epilogue_threads * values_per_thread -+ subtile_mailbox_base = ( -+ real_subtile_idx * peer_count * values_per_peer -+ ) -+ reduce_mbar_ptr = reduce_full_mbar_ptr + real_subtile_idx -+ -+ if cta_rank_in_cluster != owner_rank: -+ peer_slot = cta_rank_in_cluster -+ if cta_rank_in_cluster > owner_rank: -+ peer_slot = cta_rank_in_cluster - cutlass.Int32(1) -+ remote_thread_base = ( -+ subtile_mailbox_base -+ + peer_slot * values_per_peer -+ + epi_tidx * values_per_thread -+ ) -+ for value_idx in cutlass.range_constexpr( -+ 0, values_per_thread, 4 -+ ): -+ _store_shared_remote_v4( -+ tRS_rAcc[value_idx], -+ tRS_rAcc[value_idx + 1], -+ tRS_rAcc[value_idx + 2], -+ tRS_rAcc[value_idx + 3], -+ reduce_mailbox.iterator -+ + remote_thread_base -+ + value_idx, -+ reduce_mbar_ptr, -+ owner_rank, -+ ) -+ else: -+ if warp_idx == self.epilog_warp_id[0]: -+ with cute.arch.elect_one(): -+ cute.arch.mbarrier_arrive_and_expect_tx( -+ reduce_mbar_ptr, -+ peer_count -+ * (cute.size(epi_tile[0]) * cute.size(epi_tile[1])) -+ * 4, -+ ) -+ cute.arch.mbarrier_wait(reduce_mbar_ptr, phase=0) -+ -+ local_thread_base = ( -+ subtile_mailbox_base + epi_tidx * values_per_thread -+ ) -+ for peer_slot in cutlass.range_constexpr(peer_count): -+ peer_thread_base = ( -+ local_thread_base + peer_slot * values_per_peer -+ ) -+ for value_idx in cutlass.range_constexpr(values_per_thread): -+ tRS_rAcc[value_idx] = ( -+ tRS_rAcc[value_idx] -+ + reduce_mailbox[peer_thread_base + value_idx] -+ ) -+ -+ # Apply alpha and cast only after the FP32 cluster sum. -+ acc_vec = tRS_rAcc.load() -+ acc_vec = epilogue_op((alpha_value * acc_vec).to(self.c_dtype)) -+ tRS_rC.store(acc_vec) -+ -+ c_buffer = subtile_idx % self.num_c_stage -+ cute.copy( -+ tiled_copy_r2s, -+ tRS_rC, -+ tRS_sC[(None, None, None, c_buffer)], -+ ) -+ cute.arch.fence_proxy( -+ "async.shared", -+ space="cta", -+ ) -+ epilog_threads = 32 * len(self.epilog_warp_id) -+ cute.arch.barrier( -+ barrier_id=self.epilog_sync_bar_id, -+ number_of_threads=epilog_threads, -+ ) -+ -+ if warp_idx == self.epilog_warp_id[0]: -+ cute.copy( -+ tma_atom_c, -+ bSG_sC[(None, c_buffer)], -+ bSG_gC[(None, real_subtile_idx)], -+ ) -+ c_pipeline.producer_commit() -+ c_pipeline.producer_acquire() -+ cute.arch.barrier( -+ barrier_id=self.epilog_sync_bar_id, -+ number_of_threads=epilog_threads, -+ ) -+ -+ # -+ # Async arrive accumulator buffer empty -+ # -+ if cutlass.const_expr(not self.overlapping_accum): -+ with cute.arch.elect_one(): -+ acc_pipeline.consumer_release(acc_consumer_state) -+ acc_consumer_state.advance() -+ -+ # -+ # Advance to next tile -+ # -+ tile_sched.advance_to_next_work() -+ work_tile = tile_sched.get_current_work() -+ -+ # -+ # Dealloc the tensor memory buffer -+ # -+ if warp_idx == self.epilog_warp_id[0]: -+ cute.arch.relinquish_tmem_alloc_permit(is_two_cta=False) -+ epilog_threads = 32 * len(self.epilog_warp_id) -+ cute.arch.barrier( -+ barrier_id=self.epilog_sync_bar_id, number_of_threads=epilog_threads -+ ) -+ if warp_idx == self.epilog_warp_id[0]: -+ cute.arch.dealloc_tmem( -+ acc_tmem_ptr, self.num_tmem_alloc_cols, is_two_cta=False -+ ) -+ # -+ # Wait for C store complete -+ # -+ c_pipeline.producer_tail() -+ -+ griddepcontrol_launch_dependents() -+ -+ def _compute_grid( -+ self, -+ c: cute.Tensor, -+ cta_tile_shape_mnk: Tuple[int, int, int], -+ ) -> Tuple[int, int, int]: -+ """Return one physical K-cluster for every output tile.""" -+ c_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0)) -+ gc = cute.zipped_divide(c, tiler=c_shape) -+ num_ctas_mnl = gc[(0, (None, None, None))].shape -+ return (num_ctas_mnl[0], num_ctas_mnl[1], self.split_k_slices) From da732ecf49dc99e5d1b7aa49f6894cb5abfe7695 Mon Sep 17 00:00:00 2001 From: jiahanc <173873397+jiahanc@users.noreply.github.com> Date: Fri, 10 Jul 2026 09:28:08 -0700 Subject: [PATCH 09/10] apply autotuner fix MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 中文：应用 AutoTuner non-Tensor 输入保护修复。 --- .../fixed_seq_len/minimaxm3_fp4_b300.sh | 19 +++++++++- ...lashinfer-autotuner-non-tensor-guard.patch | 38 +++++++++++++++++++ 2 files changed, 55 insertions(+), 2 deletions(-) create mode 100644 benchmarks/single_node/fixed_seq_len/patches/flashinfer-autotuner-non-tensor-guard.patch diff --git a/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh b/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh index 5b732b7f5e..bfdcdfd398 100755 --- a/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh +++ b/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh @@ -7,7 +7,8 @@ # # At runtime the recipe swaps the image's FlashInfer for the first pinned # nightly containing the upstream SM100 low-M MXFP8 split-K kernel -# (flashinfer-ai/flashinfer#3847). +# (flashinfer-ai/flashinfer#3847), then backports the AutoTuner non-Tensor guard +# fix from flashinfer-ai/flashinfer#3918. source "$(dirname "$0")/../../benchmark_lib.sh" @@ -23,7 +24,7 @@ check_env_vars \ RANDOM_RANGE_RATIO \ RESULT_FILENAME -# --- FlashInfer nightly ------------------------------------------------------ +# --- FlashInfer nightly + AutoTuner non-Tensor guard patch ------------------ FLASHINFER_VERSION=0.6.15.dev20260710 FLASHINFER_NIGHTLY_TAG=nightly-v0.6.15-20260710 FLASHINFER_RELEASE_URL="https://github.com/flashinfer-ai/flashinfer/releases/download/${FLASHINFER_NIGHTLY_TAG}" @@ -36,6 +37,20 @@ python3 -m pip install \ "${FLASHINFER_RELEASE_URL}/flashinfer_jit_cache-${FLASHINFER_VERSION}+cu130-cp39-abi3-manylinux_2_28_$(uname -m).whl" \ || { echo "FlashInfer nightly install failed" >&2; exit 1; } +# The pinned nightly predates flashinfer-ai/flashinfer#3918. Apply only its +# runtime fix; the upstream test change is intentionally not backported. +FLASHINFER_PATCH="$(dirname "$0")/patches/flashinfer-autotuner-non-tensor-guard.patch" +if ! command -v patch >/dev/null 2>&1; then + apt-get update -y && apt-get install -y --no-install-recommends patch \ + || { echo "Failed to install patch(1)" >&2; exit 1; } +fi +SITE_PACKAGES=$(dirname "$(python3 -c "import importlib.util; print(importlib.util.find_spec('flashinfer').submodule_search_locations[0])")") \ + || { echo "Could not locate the installed flashinfer package" >&2; exit 1; } +patch --dry-run -p1 -d "${SITE_PACKAGES}" < "${FLASHINFER_PATCH}" >/dev/null \ + || { echo "FlashInfer AutoTuner non-Tensor guard patch does not apply" >&2; exit 1; } +patch -p1 -d "${SITE_PACKAGES}" < "${FLASHINFER_PATCH}" \ + || { echo "FlashInfer AutoTuner non-Tensor guard patch failed" >&2; exit 1; } + # ----------------------------------------------------------------------------- if [[ -n "${MODEL_PATH:-}" ]]; then diff --git a/benchmarks/single_node/fixed_seq_len/patches/flashinfer-autotuner-non-tensor-guard.patch b/benchmarks/single_node/fixed_seq_len/patches/flashinfer-autotuner-non-tensor-guard.patch new file mode 100644 index 0000000000..711af8b849 --- /dev/null +++ b/benchmarks/single_node/fixed_seq_len/patches/flashinfer-autotuner-non-tensor-guard.patch @@ -0,0 +1,38 @@ +diff --git a/flashinfer/autotuner/autotuner.py b/flashinfer/autotuner/autotuner.py +index 54b90a0..76443ca 100644 +--- a/flashinfer/autotuner/autotuner.py ++++ b/flashinfer/autotuner/autotuner.py +@@ -2074,14 +2074,19 @@ class AutoTuner: + def _prepare_input_tensors_with_batches( + self, +- inputs: list[torch.Tensor], ++ inputs: list[Any], + tuning_config: TuningConfig, +- ) -> list[list[torch.Tensor]]: ++ ) -> list[list[Any]]: + """Create multiple input copies to flush the L2 cache between profiling iterations.""" + if not tuning_config.use_cold_l2_cache: + return [inputs] + +- one_buffer_bytes = sum(input.numel() * input.element_size() for input in inputs) ++ one_buffer_bytes = sum( ++ input.numel() * input.element_size() ++ if isinstance(input, torch.Tensor) ++ else 0 ++ for input in inputs ++ ) + if one_buffer_bytes <= 0: + logger.debug( + "[Autotuner] No tensor inputs or zero-sized tensors; falling back to single-batch profiling." + ) +@@ -2093,7 +2098,9 @@ class AutoTuner: + + inputs_list = [inputs] + for _ in range(num_buffers - 1): +- inputs_list.append(list(t.clone() for t in inputs)) ++ inputs_list.append( ++ [t.clone() if isinstance(t, torch.Tensor) else t for t in inputs] ++ ) + + logger.debug( + f"[Autotuner] use_cold_l2_cache={tuning_config.use_cold_l2_cache}, use {num_buffers} different tensors for profiling" From 0f5015faee311b44d1a39c5df0bcc3bc51e7773d Mon Sep 17 00:00:00 2001 From: jiahanc <173873397+jiahanc@users.noreply.github.com> Date: Fri, 10 Jul 2026 17:35:32 -0700 Subject: [PATCH 10/10] use trtllm ar MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 中文：恢复显式使用 FlashInfer TRT-LLM all-reduce backend。 --- benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh | 1 + perf-changelog.yaml | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh b/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh index bfdcdfd398..f9df77df72 100755 --- a/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh +++ b/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh @@ -72,6 +72,7 @@ SERVER_LOG=/workspace/server.log export VLLM_ENGINE_READY_TIMEOUT_S=3600 export VLLM_FLOAT32_MATMUL_PRECISION=high +export VLLM_FLASHINFER_ALLREDUCE_BACKEND=trtllm export VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=1800 if [ "${DP_ATTENTION}" = "true" ]; then diff --git a/perf-changelog.yaml b/perf-changelog.yaml index 2cd5c553ea..e3c2f626bc 100644 --- a/perf-changelog.yaml +++ b/perf-changelog.yaml @@ -4722,5 +4722,5 @@ description: - "Update the vLLM image to nightly-2afa3f7e950264bb179d030c23a1ed1f46558fd9" - "Install FlashInfer 0.6.15.dev20260710 with upstream SM100 low-M MXFP8 split-K support" - - "Remove the explicit FlashInfer TRT-LLM all-reduce backend override" + - "Use the explicit FlashInfer TRT-LLM all-reduce backend override" pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2124