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..58bef9bd51 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,13 @@ # 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). It also patches +# the installed vLLM with the VLLM_MINIMAX_M3_BF16_ROUTER_GEMM knob (bf16 +# router gate weights, cuBLAS bf16 x bf16 -> fp32 logits) and enables it. source "$(dirname "$0")/../../benchmark_lib.sh" @@ -19,6 +25,51 @@ check_env_vars \ RANDOM_RANGE_RATIO \ RESULT_FILENAME +# --- FlashInfer nightly + CuTeDSL split-K gemm patch ------------------------- +# Pinned to nightly release nightly-v0.6.14-20260702. The wheels are not yet on +# the flashinfer.ai/whl/nightly flat index, so install by direct release URL. +FLASHINFER_VERSION=0.6.14.dev20260702 +FLASHINFER_NIGHTLY_TAG=nightly-v0.6.14-20260702 +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; } + +# Patch the installed vLLM with the VLLM_MINIMAX_M3_BF16_ROUTER_GEMM knob +# (bf16 router gate weights + cuBLAS bf16 x bf16 -> fp32 router logits). +# vLLM is not reinstalled above, so guard against re-applying on reruns. +VLLM_ROUTER_PATCH="$(dirname "$0")/patches/minimax-m3-bf16-router-gemm.patch" +if patch -R -p1 -s -f --dry-run -d "${SITE_PACKAGES}" < "$VLLM_ROUTER_PATCH" >/dev/null 2>&1; then + echo "[minimax-m3-bf16-router-gemm] already applied" +else + patch -p1 -d "${SITE_PACKAGES}" < "$VLLM_ROUTER_PATCH" \ + || { echo "MiniMax M3 bf16 router gemm patch failed" >&2; exit 1; } +fi +# ----------------------------------------------------------------------------- + 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" @@ -38,7 +89,9 @@ SERVER_LOG=/workspace/server.log export VLLM_ENGINE_READY_TIMEOUT_S=3600 export VLLM_FLOAT32_MATMUL_PRECISION=high +export VLLM_MINIMAX_M3_BF16_ROUTER_GEMM=1 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 +110,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/benchmarks/single_node/fixed_seq_len/patches/minimax-m3-bf16-router-gemm.patch b/benchmarks/single_node/fixed_seq_len/patches/minimax-m3-bf16-router-gemm.patch new file mode 100644 index 0000000000..3752e12860 --- /dev/null +++ b/benchmarks/single_node/fixed_seq_len/patches/minimax-m3-bf16-router-gemm.patch @@ -0,0 +1,69 @@ +From 9b9d9ce655f671ebf5beb69b7959740c5cb02071 Mon Sep 17 00:00:00 2001 +From: wzhao18 +Date: Wed, 8 Jul 2026 11:44:24 -0700 +Subject: [PATCH] Add knob with bf16 router gemm + +--- + vllm/envs.py | 7 +++++++ + vllm/models/minimax_m3/nvidia/model.py | 9 ++++++++- + 2 files changed, 15 insertions(+), 1 deletion(-) + +diff --git a/vllm/envs.py b/vllm/envs.py +index 13a19b86c0..b978102765 100755 +--- a/vllm/envs.py ++++ b/vllm/envs.py +@@ -189,6 +189,7 @@ if TYPE_CHECKING: + ] = "relax" + VLLM_USE_FUSED_MOE_GROUPED_TOPK: bool = True + VLLM_MOE_SKIP_PADDING: bool = False ++ VLLM_MINIMAX_M3_BF16_ROUTER_GEMM: bool = False + VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER: bool = True + VLLM_USE_FLASHINFER_MOE_INT4: bool = False + VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR: str | None = None +@@ -1497,6 +1498,12 @@ environment_variables: dict[str, Callable[[], Any]] = { + # treats topk_id == -1 as a skip sentinel; off by default because not all + # kernels support it yet. + "VLLM_MOE_SKIP_PADDING": lambda: bool(int(os.getenv("VLLM_MOE_SKIP_PADDING", "0"))), ++ # Store MiniMax-M3 router gate weights in bf16 and compute router logits ++ # with cuBLAS bf16 x bf16 -> fp32 instead of casting activations to fp32. ++ # Faster for batch sizes > 16; the weights are rounded fp32 -> bf16. ++ "VLLM_MINIMAX_M3_BF16_ROUTER_GEMM": lambda: bool( ++ int(os.getenv("VLLM_MINIMAX_M3_BF16_ROUTER_GEMM", "0")) ++ ), + # Allow use of FlashInfer FP8 block-scale GEMM for linear layers. + # This uses TensorRT-LLM kernels and requires SM90+ (Hopper). + "VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER": lambda: bool( +diff --git a/vllm/models/minimax_m3/nvidia/model.py b/vllm/models/minimax_m3/nvidia/model.py +index 31b6f35cb2..0ee7f7e49a 100644 +--- a/vllm/models/minimax_m3/nvidia/model.py ++++ b/vllm/models/minimax_m3/nvidia/model.py +@@ -18,6 +18,7 @@ import torch + from torch import nn + from transformers import PretrainedConfig + ++import vllm.envs as envs + from vllm import _custom_ops as ops + from vllm.compilation.breakable_cudagraph import eager_break_during_capture + from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config +@@ -224,11 +225,17 @@ class MiniMaxM3MoE(nn.Module): + + # Router weights are stored in fp32; GateLinear upcasts the bf16 + # activations and computes the gate in fp32 (fp32 router logits). ++ # With VLLM_MINIMAX_M3_BF16_ROUTER_GEMM, the weights are instead cast ++ # to bf16 and the gate runs cuBLAS bf16 x bf16 -> fp32: no activation ++ # upcast, still fp32 logits, but the weights are rounded to bf16. ++ gate_params_dtype = ( ++ torch.bfloat16 if envs.VLLM_MINIMAX_M3_BF16_ROUTER_GEMM else torch.float32 ++ ) + self.gate = GateLinear( + config.hidden_size, + config.num_local_experts, + bias=False, +- params_dtype=torch.float32, ++ params_dtype=gate_params_dtype, + out_dtype=torch.float32, + prefix=f"{prefix}.gate", + ) +-- +2.43.0 + diff --git a/configs/nvidia-master.yaml b/configs/nvidia-master.yaml index 68a125b5e0..b720e59a6a 100644 --- a/configs/nvidia-master.yaml +++ b/configs/nvidia-master.yaml @@ -13253,7 +13253,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 fc79650307..bf1bed588e 100644 --- a/perf-changelog.yaml +++ b/perf-changelog.yaml @@ -4634,3 +4634,9 @@ - "Add GLM-5.1 NVFP4 GB300 disaggregated dynamo-sglang MTP config covering 8k1k and 1k1k, high-throughput and low-latency topologies" - "Image: lmsysorg/sglang:v0.5.13.post1-cu130" pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2114 + +- config-keys: + - minimaxm3-fp4-b300-vllm + description: + - "update vllm image" + pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2120