diff --git a/python/sglang/kernels/ops/gemm/flashinfer_pr4266_dense_bf16_gemm_sm100_direct.py b/python/sglang/kernels/ops/gemm/flashinfer_pr4266_dense_bf16_gemm_sm100_direct.py deleted file mode 100644 index 5aa69e8f7..000000000 --- a/python/sglang/kernels/ops/gemm/flashinfer_pr4266_dense_bf16_gemm_sm100_direct.py +++ /dev/null @@ -1,367 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. -# Vendored from flashinfer-ai/flashinfer@629147317d4149a12e53bcef27808bac380c283f. -"""Register-prefetch BF16 GEMM for low-M, long-K decode shapes. - -The kernel keeps a complete output dot product inside one CTA and reuses each -prefetched B value across several public-M rows. It is intentionally a -separate autotuner runner from the Blackwell tensor-core split-K kernel: the -two algorithms have different useful shape regions and tactic spaces. -""" - -from __future__ import annotations - -import dataclasses -import functools - -import cuda.bindings.driver as _cuda -import cutlass -import cutlass.cute as cute -import torch as _torch -from cutlass import const_expr -from cutlass.cute import experimental as cute_ext -from cutlass.cute.runtime import from_dlpack - -_VECTOR_WIDTH = 8 -_SUPPORTED_BLOCK_SIZES = (32, 64, 96, 128, 192, 256, 384) -_SUPPORTED_OUTPUTS_PER_BLOCK = (1, 2, 4) -_MAX_M = 32 -_COMPILE_OPTIONS = "--ptxas-options -maxrregcount=64" - - -@dataclasses.dataclass(frozen=True) -class DirectTactic: - """One direct-kernel specialization.""" - - block_size: int - outputs_per_block: int - rows_per_block: int - - -def _default_rows_per_block(m: int) -> int: - if m <= 8: - return m - return next(rows for rows in (8, 4, 2, 1) if m % rows == 0) - - -def validate_tactic(tactic: DirectTactic, m: int, n: int, k: int) -> None: - """Reject a direct tactic that cannot serve ``(m, n, k)``.""" - if tactic.block_size not in _SUPPORTED_BLOCK_SIZES: - raise ValueError(f"unsupported block_size={tactic.block_size}") - if tactic.outputs_per_block not in _SUPPORTED_OUTPUTS_PER_BLOCK: - raise ValueError(f"unsupported outputs_per_block={tactic.outputs_per_block}") - if not 1 <= m <= _MAX_M: - raise ValueError(f"direct GEMM requires 1 <= M <= {_MAX_M}, got {m}") - if not 1 <= tactic.rows_per_block <= m or m % tactic.rows_per_block: - raise ValueError(f"rows_per_block={tactic.rows_per_block} must divide M={m}") - if n <= 0 or n % tactic.outputs_per_block: - raise ValueError( - f"N={n} must be divisible by outputs_per_block={tactic.outputs_per_block}" - ) - k_tile = tactic.block_size * _VECTOR_WIDTH - if k <= 0 or k % k_tile: - raise ValueError(f"K={k} must be divisible by {k_tile}") - - -def default_tactic(m: int, n: int, k: int) -> DirectTactic: - """Choose the measured register-prefetch fallback tactic.""" - block_size = next( - ( - block - for block in (256, 192, 128, 96, 64, 32) - if k % (block * _VECTOR_WIDTH) == 0 - ), - None, - ) - if block_size is None: - raise ValueError("direct GEMM requires a supported 16-byte K tiling") - outputs_per_block = next(outputs for outputs in (2, 1) if n % outputs == 0) - tactic = DirectTactic( - block_size, - outputs_per_block, - _default_rows_per_block(m), - ) - validate_tactic(tactic, m, n, k) - return tactic - - -def autotune_tactics(m: int, n: int, k: int) -> list[DirectTactic]: - """Enumerate the compact tactic space used by FlashInfer autotuning. - - Block sizes cover every configuration exercised in the H100/B200 sweep; - output grouping spans the measured 1/2/4-column choices. Row tiling stays - at the occupancy-oriented default to keep JIT cost bounded. - """ - try: - default = default_tactic(m, n, k) - except ValueError: - return [] - tactics = [default] - for block_size in _SUPPORTED_BLOCK_SIZES: - for outputs_per_block in _SUPPORTED_OUTPUTS_PER_BLOCK: - tactic = DirectTactic( - block_size, - outputs_per_block, - default.rows_per_block, - ) - try: - validate_tactic(tactic, m, n, k) - except ValueError: - continue - tactics.append(tactic) - return list(dict.fromkeys(tactics)) - - -def prefer_direct_bf16_gemm_sm100(m: int, n: int, k: int) -> bool: - """Return the conservative B200 no-autotune crossover heuristic. - - The three bands are a compact fit to a warm/cold sweep over M=1..16,24,32, - 18 N values, and 11 K values. This is deliberately not a blanket rule for - K=8192: direct wins only where public M and N leave the tensor-core path - with too little independent output work. - """ - return k == 8192 and ( - (m == 1 and n <= 4608) or (m <= 4 and n <= 512) or (m <= 8 and n <= 256) - ) - - -class DirectDenseGemmKernel: - """K-specialized direct GEMM with whole-mainloop vector prefetch.""" - - def __init__( - self, - *, - element_type, - num_rows: int, - k_extent: int, - tactic: DirectTactic, - use_pdl: bool, - ) -> None: - validate_tactic(tactic, num_rows, tactic.outputs_per_block, k_extent) - self.element_type = element_type - self.num_rows = num_rows - self.rows_per_block = tactic.rows_per_block - self.k_extent = k_extent - self.block_size = tactic.block_size - self.outputs_per_block = tactic.outputs_per_block - self.vector_width = _VECTOR_WIDTH - self.use_pdl = use_pdl - self.num_warps = tactic.block_size // cute.arch.WARP_SIZE - self.num_k_tiles = k_extent // (tactic.block_size * _VECTOR_WIDTH) - - @cute.jit - def __call__( - self, - gA: cute.Tensor, - gB: cute.Tensor, - gC: cute.Tensor, - stream: _cuda.CUstream, - ) -> None: - n = cute.size(gB, mode=[0]) - copy_a = cute.make_copy_atom( - cute.nvgpu.CopyG2ROp(), - self.element_type, - num_bits_per_copy=self.vector_width * self.element_type.width, - load_cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS, - ) - copy_b = cute.make_copy_atom( - cute.nvgpu.CopyG2ROp(), - self.element_type, - num_bits_per_copy=self.vector_width * self.element_type.width, - load_cache_mode=cute.nvgpu.LoadCacheMode.STREAMING, - ) - self.kernel(gA, gB, gC, copy_a, copy_b).launch( - grid=[ - cute.ceil_div(n, self.outputs_per_block), - self.num_rows // self.rows_per_block, - 1, - ], - block=[self.block_size, 1, 1], - smem=self.rows_per_block * self.outputs_per_block * self.num_warps * 4, - stream=stream, - use_pdl=self.use_pdl, - min_blocks_per_mp=1, - ) - - @cute.kernel - def kernel( - self, - gA: cute.Tensor, - gB: cute.Tensor, - gC: cute.Tensor, - copy_a: cute.CopyAtom, - copy_b: cute.CopyAtom, - ) -> None: - tidx, _, _ = cute.arch.thread_idx() - block_idx, block_m, _ = cute.arch.block_idx() - warp_idx = cute.arch.warp_idx() - - num_rows: cutlass.Constexpr = self.rows_per_block - outputs_per_block: cutlass.Constexpr = self.outputs_per_block - vector_width: cutlass.Constexpr = self.vector_width - block_size: cutlass.Constexpr = self.block_size - num_warps: cutlass.Constexpr = self.num_warps - num_k_tiles: cutlass.Constexpr = self.num_k_tiles - - acc = cute.make_rmem_tensor( - cute.make_layout( - (num_rows, outputs_per_block), stride=(outputs_per_block, 1) - ), - cutlass.Float32, - ) - acc.fill(0.0) - - if const_expr(self.use_pdl): - cute.arch.griddepcontrol_wait() - - n_base = block_idx * outputs_per_block - m_base = block_m * num_rows - gA_vec = cute.logical_divide(gA, (None, vector_width)) - gB_vec = cute.logical_divide(gB, (None, vector_width)) - tA_all = cute.logical_divide(gA_vec, (None, (None, block_size))) - tB_all = cute.logical_divide(gB_vec, (None, (None, block_size))) - tA = tA_all[None, (None, (tidx, None))] - - b_regs = cute.make_rmem_tensor( - cute.make_layout( - (outputs_per_block, num_k_tiles, vector_width), - stride=(num_k_tiles * vector_width, vector_width, 1), - ), - self.element_type, - ) - for ni in cutlass.range_constexpr(outputs_per_block): - tB = tB_all[n_base + ni, (None, (tidx, None))] - for k_tile in cutlass.range_constexpr(num_k_tiles): - cute.copy(copy_b, tB[None, k_tile], b_regs[ni, k_tile, None]) - - a_regs = cute.make_rmem_tensor( - cute.make_layout((num_k_tiles, vector_width), stride=(vector_width, 1)), - self.element_type, - ) - for mi in cutlass.range_constexpr(num_rows): - for k_tile in cutlass.range_constexpr(num_k_tiles): - cute.copy( - copy_a, - tA[m_base + mi, None, k_tile], - a_regs[k_tile, None], - ) - for k_tile in cutlass.range_constexpr(num_k_tiles): - for vi in cutlass.range_constexpr(vector_width): - a_value = a_regs[k_tile, vi].to(cutlass.Float32) - for ni in cutlass.range_constexpr(outputs_per_block): - acc[mi, ni] = acc[mi, ni] + a_value * b_regs[ni, k_tile, vi].to( - cutlass.Float32 - ) - - for mi in cutlass.range_constexpr(num_rows): - for ni in cutlass.range_constexpr(outputs_per_block): - acc[mi, ni] = cute.arch.warp_reduction_sum(acc[mi, ni]) - - smem_layout = cute.make_layout( - (num_rows, outputs_per_block, num_warps), - stride=(outputs_per_block * num_warps, num_warps, 1), - ) - smem = cutlass.utils.SmemAllocator() - partials = smem.allocate_tensor(cutlass.Float32, smem_layout, byte_alignment=16) - with cute.arch.elect_one(): - for mi in cutlass.range_constexpr(num_rows): - for ni in cutlass.range_constexpr(outputs_per_block): - partials[mi, ni, warp_idx] = acc[mi, ni] - - cute.arch.sync_threads() - if tidx == 0: - for mi in cutlass.range_constexpr(num_rows): - for ni in cutlass.range_constexpr(outputs_per_block): - total = cutlass.Float32(0.0) - for warp in cutlass.range_constexpr(num_warps): - total = total + partials[mi, ni, warp] - gC[m_base + mi, n_base + ni] = total.to(self.element_type) - - if const_expr(self.use_pdl): - cute.arch.griddepcontrol_launch_dependents() - - -def _from_dlpack_static(tensor: _torch.Tensor): - # K is specialized and the row stride must retain its 16-byte divisibility - # for the verifier to accept vectorized G2R copies. - return from_dlpack(tensor, assumed_align=32) - - -def _make_compile_repr_tensors(dtype, m: int, n: int, k: int): - return tuple( - _from_dlpack_static(tensor) - for tensor in ( - _torch.empty((m, k), dtype=dtype, device="cuda"), - _torch.empty((n, k), dtype=dtype, device="cuda"), - _torch.empty((m, n), dtype=dtype, device="cuda"), - ) - ) - - -@functools.cache -def _get_compiled_direct_kernel( - dtype, - m: int, - n: int, - k: int, - tactic: DirectTactic, - use_pdl: bool, -): - if dtype != _torch.bfloat16: - raise ValueError(f"direct GEMM supports BF16; got {dtype}") - kernel = DirectDenseGemmKernel( - element_type=cutlass.BFloat16, - num_rows=m, - k_extent=k, - tactic=tactic, - use_pdl=use_pdl, - ) - tensors = _make_compile_repr_tensors(dtype, m, n, k) - stream = _cuda.CUstream(_torch.cuda.current_stream().cuda_stream) - return cute_ext.compile(kernel, *tensors, stream, options=_COMPILE_OPTIONS) - - -def _validate_runtime_tensors(a, b, out, tactic: DirectTactic): - if any(not isinstance(tensor, _torch.Tensor) for tensor in (a, b, out)): - raise ValueError("a, b, and out must be torch tensors") - if a.ndim != 2 or b.ndim != 2 or out.ndim != 2: - raise ValueError("direct GEMM accepts only 2D tensors") - if a.device.type != "cuda" or b.device != a.device or out.device != a.device: - raise ValueError("a, b, and out must be on the same CUDA device") - if a.dtype != _torch.bfloat16 or b.dtype != a.dtype or out.dtype != a.dtype: - raise ValueError("a, b, and out must share BF16 dtype") - if not a.is_contiguous() or not b.T.is_contiguous() or not out.is_contiguous(): - raise ValueError("direct GEMM requires row-major A/out and column-major B") - if any(tensor.data_ptr() % 32 for tensor in (a, b, out)): - raise ValueError("a, b, and out must be 32-byte aligned") - - m, k = a.shape - if b.shape[0] != k: - raise ValueError( - f"incompatible shapes: a is {tuple(a.shape)}, b is {tuple(b.shape)}" - ) - n = b.shape[1] - if out.shape != (m, n): - raise ValueError(f"out must have shape {(m, n)}, got {tuple(out.shape)}") - validate_tactic(tactic, m, n, k) - return m, n, k - - -def run_direct_dense(a, b, out, pdl: bool, tactic: DirectTactic): - """Run direct ``A[M,K] @ B[K,N]`` with the ``mm_bf16`` layouts.""" - m, n, k = _validate_runtime_tensors(a, b, out, tactic) - compiled = _get_compiled_direct_kernel(a.dtype, m, n, k, tactic, pdl) - tensors = tuple(_from_dlpack_static(tensor) for tensor in (a, b.T, out)) - stream = _cuda.CUstream(_torch.cuda.current_stream(a.device).cuda_stream) - compiled(*tensors, stream) - return out - - -__all__ = [ - "DirectTactic", - "autotune_tactics", - "default_tactic", - "prefer_direct_bf16_gemm_sm100", - "run_direct_dense", -] diff --git a/python/sglang/kernels/ops/gemm/flashinfer_pr4266_dense_bf16_gemm_sm100_splitk.py b/python/sglang/kernels/ops/gemm/flashinfer_pr4266_dense_bf16_gemm_sm100_splitk.py deleted file mode 100644 index 92ec82470..000000000 --- a/python/sglang/kernels/ops/gemm/flashinfer_pr4266_dense_bf16_gemm_sm100_splitk.py +++ /dev/null @@ -1,1043 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -# Vendored from flashinfer-ai/flashinfer@629147317d4149a12e53bcef27808bac380c283f. -"""Blackwell low-M BF16/FP16 GEMM with an in-kernel cluster split-K reduction. - -Each cluster rank accumulates an exact K slice in FP32. Peers publish partials -to rank 0 through DSMEM; rank 0 reduces, casts, and stores once. The public -``A[M, K] @ B[K, N]`` problem is swapped internally, so tile dimensions below -use kernel coordinates: kernel-M carries public N and kernel-N carries public M. -""" - -from __future__ import annotations - -import dataclasses - -import cuda.bindings.driver as _cuda -import cutlass -import cutlass.cute as cute -import cutlass.utils as utils -import cutlass.utils.blackwell_helpers as sm100_utils -from cutlass import Int32 -from cutlass._mlir.dialects import llvm -from cutlass.cute import experimental as cute_ext -from cutlass.cute.nvgpu import tcgen05 -from cutlass.cute.runtime import from_dlpack -from cutlass.cutlass_dsl import T, dsl_user_op - -#: Per-CTA SMEM capacity reported by CuTeDSL on SM100/SM103. -_SMEM_CAPACITY_BYTES = 227 * 1024 - -#: K extent of one CTA tile. -_CTA_K = 128 - -#: Kernel-M tiles; 64 increases CTA count for low-M decode shapes. -_SUPPORTED_MMA_M = (64, 128) - -#: Kernel-N carries public M, which is limited to 32. -_SUPPORTED_MMA_N = (8, 16, 32) - -#: Physical cluster-K sizes; split 1 compiles out the DSMEM path. -_SUPPORTED_SPLIT_K = (1, 2, 3, 4) - -#: Largest public M this low-M policy serves. -_MAX_M = 32 - -#: Bytes per FP32 partial exchanged through DSMEM. -_FP32_BYTES = 4 - -#: DSMEM mailbox base alignment, in bytes. -_MAILBOX_ALIGN_BYTES = 128 - -#: Size and alignment of one mbarrier, in bytes. -_MBARRIER_BYTES = 8 - -#: Size and alignment of the TMEM base pointer slot. -_TMEM_POINTER_BYTES = 4 - -#: Bytes per BF16/FP16 element. -_AB_ELEMENT_BYTES = 2 - -#: Alignment of the A/B shared-memory buffers. -_AB_BUFFER_ALIGN_BYTES = 1024 - -#: A/B pipeline stage bounds. -_MIN_AB_STAGES = 2 -_MAX_AB_STAGES = 12 - - -@dataclasses.dataclass(frozen=True, slots=True) -class SplitKTactic: - """One specialization; mma_m carries public N and mma_n carries public M.""" - - mma_m: int - mma_n: int - split_k: int - ab_stages: int - - -def _align_up(value: int, alignment: int) -> int: - return ((value + alignment - 1) // alignment) * alignment - - -def _smem_bytes( - tactic: SplitKTactic, - ab_stages: int, -) -> int: - """Mirror the device allocator's shared-memory layout.""" - cursor = ( - _align_up( - tactic.mma_m * _CTA_K * _AB_ELEMENT_BYTES * ab_stages, - _AB_BUFFER_ALIGN_BYTES, - ) - + tactic.mma_n * _CTA_K * _AB_ELEMENT_BYTES * ab_stages - ) - - cursor = _align_up(cursor, _MBARRIER_BYTES) - cursor += 2 * ab_stages * _MBARRIER_BYTES - cursor += 3 * _MBARRIER_BYTES - cursor = _align_up(cursor, _TMEM_POINTER_BYTES) - cursor += _TMEM_POINTER_BYTES - - if tactic.split_k == 1: - return cursor - - return ( - _align_up( - _align_up(cursor, _MAILBOX_ALIGN_BYTES) - + (tactic.split_k - 1) * tactic.mma_m * tactic.mma_n * _FP32_BYTES, - _MBARRIER_BYTES, - ) - + _MBARRIER_BYTES - ) - - -def _max_ab_stages_for( - tactic: SplitKTactic, - smem_capacity: int, -) -> int: - return next( - ( - stages - for stages in range(_MAX_AB_STAGES, -1, -1) - if _smem_bytes(tactic, stages) <= smem_capacity - ), - 0, - ) - - -def validate_tactic( - tactic: SplitKTactic, - m: int, - n: int, - k: int, - *, - smem_capacity: int = _SMEM_CAPACITY_BYTES, -) -> None: - """Reject a tactic that cannot serve ``(m, n, k)``.""" - if tactic.mma_m not in _SUPPORTED_MMA_M: - raise ValueError(f"unsupported mma_m={tactic.mma_m}") - if tactic.mma_n not in _SUPPORTED_MMA_N: - raise ValueError(f"unsupported mma_n={tactic.mma_n}") - if tactic.split_k not in _SUPPORTED_SPLIT_K: - raise ValueError(f"unsupported split_k={tactic.split_k}") - if not _MIN_AB_STAGES <= tactic.ab_stages <= _MAX_AB_STAGES: - raise ValueError( - f"ab_stages must be in [{_MIN_AB_STAGES}, {_MAX_AB_STAGES}], " - f"got {tactic.ab_stages}" - ) - if not 1 <= m <= _MAX_M: - raise ValueError(f"this low-M policy requires 1 <= M <= {_MAX_M}, got {m}") - if n <= 0: - raise ValueError(f"N must be positive, got {n}") - if k <= 0 or k % _CTA_K or (k // _CTA_K) % tactic.split_k: - raise ValueError( - f"K={k} with CTA_K={_CTA_K} does not divide evenly across " - f"split_k={tactic.split_k}" - ) - smem_bytes = _smem_bytes(tactic, tactic.ab_stages) - if smem_bytes > smem_capacity: - raise ValueError( - f"tactic {tactic} needs {smem_bytes} B of shared memory but only " - f"{smem_capacity} B are available; max ab_stages is " - f"{_max_ab_stages_for(tactic, smem_capacity)}" - ) - - -def autotune_tactics( - m: int, - n: int, - k: int, - *, - smem_capacity: int = _SMEM_CAPACITY_BYTES, -) -> list[SplitKTactic]: - """Return valid tactics in the shape-specific stage window.""" - tactics: list[SplitKTactic] = [] - for mma_m in _SUPPORTED_MMA_M: - for mma_n in _SUPPORTED_MMA_N: - for split_k in _SUPPORTED_SPLIT_K: - base = SplitKTactic(mma_m, mma_n, split_k, _MIN_AB_STAGES) - try: - validate_tactic(base, m, n, k, smem_capacity=smem_capacity) - except ValueError: - continue - max_stages = _max_ab_stages_for(base, smem_capacity) - # Short K favors shallow pipelines; long K stays near the cap. - tactics.extend( - dataclasses.replace(base, ab_stages=ab_stages) - for ab_stages in ( - range(_MIN_AB_STAGES, min(max_stages, 6) + 1) - if k <= 4 * _CTA_K - else range( - min(max(5, max_stages - 2), max_stages), - max_stages + 1, - ) - ) - ) - return tactics - - -def default_tactic(m: int, n: int, k: int) -> SplitKTactic: - """Choose the default occupancy-oriented tactic.""" - if n <= 512: - mma_m = 64 - mma_n = 16 if n == 512 and m > 24 else 8 - requested_split = 4 - else: - mma_n = 8 if m <= 8 else 16 if m <= 16 else 32 - if n <= 3072: - mma_m = 128 if m <= 16 else 64 - requested_split = 4 if m <= 16 else 2 - elif n < 8192: - mma_m = 64 - requested_split = 2 - else: - mma_m = 128 if k <= 1024 and m <= 24 else 64 - requested_split = 1 - - if k <= 4 * _CTA_K: - requested_split = 1 - split_k = next( - split_k - for split_k in reversed(_SUPPORTED_SPLIT_K) - if split_k <= requested_split and (k // _CTA_K) % split_k == 0 - ) - tactic = SplitKTactic(mma_m, mma_n, split_k, _MIN_AB_STAGES) - max_stages = _max_ab_stages_for(tactic, _SMEM_CAPACITY_BYTES) - tactic = dataclasses.replace( - tactic, - ab_stages=( - _MIN_AB_STAGES - if k <= 2 * _CTA_K and m > 8 - else min(max_stages, 6) - if k <= 4 * _CTA_K - else max_stages - ), - ) - validate_tactic(tactic, m, n, k) - return tactic - - -__all__ = [ - "SplitKTactic", - "autotune_tactics", - "default_tactic", - "run_splitk_dense", -] - - -@dsl_user_op -def _map_shared_rank( - smem_ptr: cute.Pointer, - peer_cta_rank_in_cluster: Int32, - *, - loc=None, - ip=None, -) -> Int32: - """Map an SMEM pointer into a peer CTA's address space.""" - return Int32( - llvm.inline_asm( - T.i32(), - [ - smem_ptr.toint(loc=loc, ip=ip).ir_value(), - 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, - value1, - value2, - value3, - smem_ptr: cute.Pointer, - mbar_ptr: cute.Pointer, - peer_cta_rank_in_cluster: Int32, - *, - loc=None, - ip=None, -) -> None: - """Publish four FP32 partials into a peer's SMEM, crediting 16 bytes.""" - llvm.inline_asm( - None, - [ - _map_shared_rank( - smem_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip - ).ir_value(), - 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), - _map_shared_rank( - mbar_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip - ).ir_value(), - ], - "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, - ) - - -#: Rank that gathers partials and stores the output. -OWNER_RANK = 0 - - -class SplitKDenseGemmKernel: - """Standalone BF16/FP16 GEMM with a cluster-local split-K reduction.""" - - def __init__( - self, - *, - tactic: SplitKTactic, - use_pdl: bool, - has_bias: bool, - ) -> None: - self.acc_dtype = cutlass.Float32 - self.cta_m = tactic.mma_m - self.cta_n = tactic.mma_n - self.cta_k = _CTA_K - self.num_ab_stage = tactic.ab_stages - self.split_k = tactic.split_k - self.use_pdl = use_pdl - self.has_bias = has_bias - - self.threads_per_cta = 256 - self.epilog_threads = 128 - self.mma_tiler_mn = (tactic.mma_m, tactic.mma_n) - self.cta_group = tcgen05.CtaGroup.ONE - self.tma_op = cute_ext.OperationTypeEnum.SM90_TMA_LOAD - self.cluster_shape = (1, tactic.split_k, 1) - - values_per_thread = (tactic.mma_m * tactic.mma_n) // self.epilog_threads - if values_per_thread % 4: - raise ValueError( - f"CTA tile ({tactic.mma_m}, {tactic.mma_n}) gives " - f"{values_per_thread} " - "values per epilogue thread; remote stores require a multiple of 4" - ) - self.mailbox_elements = ( - (tactic.split_k - 1) * self.epilog_threads * values_per_thread - ) - self.expected_transaction_bytes = self.mailbox_elements * _FP32_BYTES - - @cute.experimental.jit - def __call__( - self, - a: cute.Tensor, - b: cute.Tensor, - c: cute.Tensor, - bias: cute.Tensor, - stream: _cuda.CUstream, - ): - # Grid-y packs output-N tile and cluster rank. - self.kernel(a, b, c, bias).launch( - grid=( - cute.ceil_div(c.layout.shape[0], self.cta_m), - cute.ceil_div(c.layout.shape[1], self.cta_n) * self.split_k, - c.layout.shape[2], - ), - block=(self.threads_per_cta, 1, 1), - cluster=self.cluster_shape, - smem=cute.Int64(utils.get_smem_capacity_in_bytes("sm_100")), - stream=stream, - use_pdl=self.use_pdl, - ) - - @cute.experimental.kernel - def kernel( - self, - mA: cute.Tensor, # (Gemm_M, Gemm_K, Gemm_L), K-major - mB: cute.Tensor, # (Gemm_N, Gemm_K, Gemm_L), K-major - mC: cute.Tensor, # (Gemm_M, Gemm_N, Gemm_L), M-major - mBias: cute.Tensor, # Broadcast bias; dead when has_bias=False - ): - """Allocate storage and dispatch the specialized warps.""" - stages = self.num_ab_stage - - ab_dtype = mA.element_type - tiled_mma = sm100_utils.make_trivial_tiled_mma( - ab_dtype, - ab_dtype, - utils.LayoutEnum.from_tensor(mA).mma_major_mode(), - utils.LayoutEnum.from_tensor(mB).mma_major_mode(), - self.acc_dtype, - self.cta_group, - self.mma_tiler_mn, - ) - - mnk_tiler = (self.mma_tiler_mn[0], self.mma_tiler_mn[1], self.cta_k) - block_idx = cute.arch.block_idx() - bidx = block_idx[0] - split_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) - n_idx = block_idx[1] // self.split_k - l_idx = block_idx[2] - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - - sA = cute_ext.allocate( - ab_dtype, - cute.AddressSpace.smem, - sm100_utils.make_smem_layout_a(tiled_mma, mnk_tiler, ab_dtype, stages), - alignment=_AB_BUFFER_ALIGN_BYTES, - ) - sB = cute_ext.allocate( - ab_dtype, - cute.AddressSpace.smem, - sm100_utils.make_smem_layout_b(tiled_mma, mnk_tiler, ab_dtype, stages), - alignment=_AB_BUFFER_ALIGN_BYTES, - ) - - acc_layout = cute_ext.make_tmem_layout_acc( - tiled_mma, self.mma_tiler_mn, acc_stage=1 - ) - c_tiler_mn = (self.cta_m, self.cta_n) - - bar_full = cute_ext.allocate( - cutlass.Int64, - cute.AddressSpace.smem, - cute.make_layout(stages), - alignment=_MBARRIER_BYTES, - ).iterator - bar_empty = cute_ext.allocate( - cutlass.Int64, - cute.AddressSpace.smem, - cute.make_layout(stages), - alignment=_MBARRIER_BYTES, - ).iterator - bar_tma_epilog = cute_ext.allocate( - cutlass.Int64, - cute.AddressSpace.smem, - cute.make_layout(1), - alignment=_MBARRIER_BYTES, - ).iterator - bar_mma_epilog = cute_ext.allocate( - cutlass.Int64, - cute.AddressSpace.smem, - cute.make_layout(1), - alignment=_MBARRIER_BYTES, - ).iterator - bar_tmem_alloc = cute_ext.allocate( - cutlass.Int64, - cute.AddressSpace.smem, - cute.make_layout(1), - alignment=_MBARRIER_BYTES, - ).iterator - tmem_base_ptr = cute_ext.allocate( - cutlass.Int32, - cute.AddressSpace.smem, - cute.make_layout(1), - alignment=_TMEM_POINTER_BYTES, - ).iterator - - if cutlass.const_expr(self.split_k > 1): - mailbox = cute_ext.allocate( - cutlass.Float32, - cute.AddressSpace.smem, - cute.make_layout(self.mailbox_elements), - alignment=_MAILBOX_ALIGN_BYTES, - ) - bar_reduce = cute_ext.allocate( - cutlass.Int64, - cute.AddressSpace.smem, - cute.make_layout(1), - alignment=_MBARRIER_BYTES, - ).iterator - else: - # Dummy operands for the compile-time-elided reduction. - mailbox = sA - bar_reduce = bar_mma_epilog - - if warp_idx == 0: - with cute.arch.elect_one(): - for i in range(stages): - cute.arch.mbarrier_init(bar_full + i, 2) - cute.arch.mbarrier_init(bar_empty + i, 1) - cute.arch.mbarrier_init(bar_tma_epilog, 32) - cute.arch.mbarrier_init(bar_mma_epilog, 1) - cute.arch.mbarrier_init(bar_tmem_alloc, 160) - - if cutlass.const_expr(self.split_k > 1): - # Owner arrival plus peer transaction-byte credits. - cute.arch.mbarrier_init(bar_reduce, 1) - - cute.arch.mbarrier_init_fence() - if cutlass.const_expr(self.split_k > 1): - # Publish peer barriers before cross-CTA stores. - cute.arch.cluster_arrive_relaxed() - else: - cute.arch.barrier() - - # Host validation guarantees an equal, tail-free K partition. - k_tile_count = cute.size(mA, mode=[1]) // self.cta_k // self.split_k - k_tile_start = split_rank * k_tile_count - - if cutlass.const_expr(self.split_k > 1): - cute.arch.cluster_wait() - - # Warp 3 is idle; warps 4-7 run the epilogue. - if warp_idx == 0: - self.dma_warp( - bar_full, - bar_empty, - bar_tma_epilog, - cute.local_tile(mA, (self.cta_m, self.cta_k), (bidx, None, l_idx)), - sA, - cute_ext.get_cta_v_map_ab(mA, mnk_tiler, tiled_mma, "A"), - k_tile_start, - k_tile_count, - True, - ) - elif warp_idx == 1: - self.dma_warp( - bar_full, - bar_empty, - bar_tma_epilog, - cute.local_tile(mB, (self.cta_n, self.cta_k), (n_idx, None, l_idx)), - sB, - cute_ext.get_cta_v_map_ab(mB, mnk_tiler, tiled_mma, "B"), - k_tile_start, - k_tile_count, - False, - ) - elif warp_idx == 2: - self.mma_warp( - bar_full, - bar_empty, - bar_mma_epilog, - bar_tmem_alloc, - tiled_mma, - sA, - sB, - tmem_base_ptr, - acc_layout, - self.cta_k // cute.size(tiled_mma.shape_mnk, mode=[2]), - k_tile_count, - ) - elif warp_idx >= 4: - self.epilog_warp( - bar_tma_epilog, - bar_mma_epilog, - bar_tmem_alloc, - tmem_base_ptr, - acc_layout, - cute.local_tile(mC, c_tiler_mn, (bidx, n_idx, l_idx)), - cute.local_tile(mBias, c_tiler_mn, (bidx, n_idx, l_idx)), - cute.arch.thread_idx()[0] - 128, - mC.element_type, - utils.LayoutEnum.from_tensor(mC), - mailbox, - bar_reduce, - split_rank, - ) - - @cute.experimental.jit - def dma_warp( - self, - bar_full, - bar_empty, - bar_tma_epilog, - g_tile: cute.Tensor, - s_tile: cute.Tensor, - cta_v_map: cute.Layout, - k_tile_start: cutlass.Int32, - k_tile_count: cutlass.Int32, - is_a: cutlass.Constexpr, - ): - stages = self.num_ab_stage - if cutlass.const_expr(not is_a and self.use_pdl): - cute.arch.griddepcontrol_wait() - - empty_phase = cutlass.Int32(1) - for k_tile in cutlass.range(k_tile_count, unroll=1): - stage = k_tile % stages - cute.arch.mbarrier_wait(bar_empty + stage, empty_phase) - with cute.arch.elect_one(): - cute.arch.mbarrier_arrive_and_expect_tx( - bar_full + stage, - cute.size_in_bytes( - s_tile.element_type, - cute.slice_(s_tile.layout, (None, None, None, 0)), - ), - ) - cute_ext.tma_load( - g_tile[None, None, k_tile_start + k_tile], - s_tile[None, None, None, stage], - (bar_full + stage).value, - cta_v_map=cta_v_map, - tma_operation_type=self.tma_op, - update_expect_tx=False, - ) - if stage == stages - 1: - empty_phase = empty_phase ^ 1 - - if cutlass.const_expr(is_a and self.use_pdl): - cute.arch.griddepcontrol_launch_dependents() - if cutlass.const_expr(not is_a and self.has_bias): - cute.arch.mbarrier_arrive(bar_tma_epilog) - self._drain_producer(bar_empty, empty_phase, k_tile_count) - - @cute.experimental.jit - def _drain_producer( - self, - bar_empty, - empty_phase: cutlass.Int32, - k_tile_count: cutlass.Int32, - ): - stages = self.num_ab_stage - for tail in cutlass.range(stages, unroll=1): - stage = (tail + k_tile_count) % stages - cute.arch.mbarrier_wait(bar_empty + stage, empty_phase) - if stage == stages - 1: - empty_phase = empty_phase ^ 1 - - @cute.experimental.jit - def mma_warp( - self, - bar_full, - bar_empty, - bar_mma_epilog, - bar_tmem_alloc, - tiled_mma: cute.TiledMma, - sA: cute.Tensor, - sB: cute.Tensor, - tmem_base_ptr, - acc_layout: cutlass.Constexpr, - mma_inst_tile_k: cutlass.Constexpr, - k_tile_count: cutlass.Int32, - ): - num_tmem_cols = 256 - cute.arch.alloc_tmem(num_tmem_cols, tmem_base_ptr, is_two_cta=False) - cute.arch.mbarrier_arrive(bar_tmem_alloc) - cute.arch.relinquish_tmem_alloc_permit(is_two_cta=False) - - tmem_ptr = cute.arch.retrieve_tmem_ptr(self.acc_dtype, 16, tmem_base_ptr) - accumulator = cute.make_tensor(tmem_ptr, acc_layout)[None, None, None, 0] - mma_atom = cute.make_mma_atom(tiled_mma.op) - full_phase = cutlass.Int32(0) - for k_tile in cutlass.range(k_tile_count, unroll=1): - stage = k_tile % self.num_ab_stage - cute.arch.mbarrier_wait(bar_full + stage, full_phase) - for k_block in range(mma_inst_tile_k): - if k_block == 0: - mma_atom.set(tcgen05.Field.ACCUMULATE, k_tile != 0) - else: - mma_atom.set(tcgen05.Field.ACCUMULATE, True) - cute_ext.dot( - mma_atom, - cute.append_ones(sA[None, None, k_block, stage], up_to_rank=3), - cute.append_ones(sB[None, None, k_block, stage], up_to_rank=3), - accumulator, - ) - with cute.arch.elect_one(): - tcgen05.commit(bar_empty + stage, None, self.cta_group) - if stage == self.num_ab_stage - 1: - full_phase = full_phase ^ 1 - - with cute.arch.elect_one(): - tcgen05.commit(bar_mma_epilog, None, self.cta_group) - cute.arch.mbarrier_arrive(bar_tmem_alloc) - cute.arch.mbarrier_wait(bar_tmem_alloc, 1) - cute.arch.dealloc_tmem(tmem_ptr, num_tmem_cols, is_two_cta=False) - - @cute.experimental.jit - def epilog_warp( - self, - bar_tma_epilog, - bar_mma_epilog, - bar_tmem_alloc, - tmem_base_ptr, - acc_layout: cutlass.Constexpr, - gD_tile: cute.Tensor, - gBias_tile: cute.Tensor, - epi_tid: cutlass.Int32, - c_dtype: cutlass.Constexpr, - d_layout: cutlass.Constexpr, - mailbox, - bar_reduce, - split_rank: cutlass.Int32, - ): - # Wait until MMA publishes the TMEM base pointer. - cute.arch.mbarrier_arrive(bar_tmem_alloc) - cute.arch.mbarrier_wait(bar_tmem_alloc, 0) - - acc_view = cute.make_tensor( - cute.arch.retrieve_tmem_ptr(self.acc_dtype, 16, tmem_base_ptr), - acc_layout, - )[((None, None), 0, 0, 0)] - - epi_tile = (self.cta_m, self.cta_n) - tiled_copy_t2r = cute.nvgpu.tcgen05.make_tmem_copy( - sm100_utils.get_tmem_load_op( - (self.cta_m, self.cta_n, self.cta_k), - d_layout, - c_dtype, - self.acc_dtype, - epi_tile, - False, - ), - acc_view, - ) - gD_epi = cute.flat_divide(gD_tile, epi_tile) - - # Match each epilogue thread's TMEM partition in RMEM. - rmem_layout = cute_ext.make_t2r_rmem_layout(tiled_copy_t2r, gD_epi, epi_tid) - rAcc = cute_ext.allocate( - self.acc_dtype, - cute.AddressSpace.rmem, - rmem_layout, - alignment=32, - ) - rD = cute_ext.allocate( - c_dtype, - cute.AddressSpace.rmem, - rmem_layout, - alignment=32, - ) - thr_t2r = tiled_copy_t2r.get_slice(epi_tid) - - if cutlass.const_expr(self.has_bias): - bias_dtype = gBias_tile.element_type - rBias = cute_ext.allocate( - bias_dtype, - cute.AddressSpace.rmem, - rmem_layout, - alignment=32, - ) - rBiasAcc = cute_ext.allocate( - self.acc_dtype, - cute.AddressSpace.rmem, - rmem_layout, - alignment=32, - ) - if split_rank == OWNER_RANK: - cute.arch.mbarrier_wait(bar_tma_epilog, 0) - cute_ext.partition_and_copy( - cute.make_tiled_copy_D( - cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), bias_dtype), - tiled_copy_t2r, - ).get_slice(epi_tid), - cute.flat_divide(gBias_tile, epi_tile)[None, None, 0, 0], - rBias, - ) - rBiasAcc.store(rBias.load().to(self.acc_dtype)) - - cute.arch.mbarrier_wait(bar_mma_epilog, 0) - cute_ext.partition_and_copy(thr_t2r, acc_view, rAcc) - # Make tcgen05.ld visible before TMEM release and RMEM use. - cute.arch.fence_view_async_tmem_load() - cute.arch.mbarrier_arrive(bar_tmem_alloc) - - # Peers publish FP32 partials; only rank 0 reduces and stores. - if cutlass.const_expr(self.split_k > 1): - assert cute.size(rmem_layout) == self.mailbox_elements // ( - (self.split_k - 1) * self.epilog_threads - ) - values_per_thread = cutlass.const_expr(cute.size(rmem_layout)) - values_per_peer = cutlass.const_expr( - self.epilog_threads * values_per_thread - ) - if split_rank != OWNER_RANK: - for value_idx in cutlass.range_constexpr(0, values_per_thread, 4): - _store_shared_remote_v4( - rAcc[value_idx], - rAcc[value_idx + 1], - rAcc[value_idx + 2], - rAcc[value_idx + 3], - mailbox.iterator - + (split_rank - Int32(1)) * values_per_peer - + epi_tid * values_per_thread - + value_idx, - bar_reduce, - Int32(OWNER_RANK), - ) - else: - if epi_tid == 0: - cute.arch.mbarrier_arrive_and_expect_tx( - bar_reduce, self.expected_transaction_bytes - ) - cute.arch.mbarrier_wait(bar_reduce, 0) - for peer in cutlass.range_constexpr(self.split_k - 1): - for value_idx in cutlass.range_constexpr(values_per_thread): - rAcc[value_idx] = ( - rAcc[value_idx] - + mailbox[ - peer * values_per_peer - + epi_tid * values_per_thread - + value_idx - ] - ) - - if split_rank == OWNER_RANK: - if cutlass.const_expr(self.has_bias): - rAcc.store(rAcc.load() + rBiasAcc.load()) - - rD.store(rAcc.load().to(c_dtype)) - # Preserve TMEM coordinates; the copy predicates output tails. - cute_ext.partition_and_copy(thr_t2r, rD, gD_epi[None, None, 0, 0]) - - # The reduction mbarrier covers remote stores; no cluster barrier needed. - - -import torch as _torch - -_SUPPORTED_TORCH_DTYPES = (_torch.bfloat16, _torch.float16) - - -@cute.experimental.jit -def _bmm_no_bias( - gemm_op: cutlass.Constexpr, - a: cute.Tensor, - b: cute.Tensor, - c: cute.Tensor, - stream: _cuda.CUstream, -): - c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) - gemm_op( - cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), - cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), - c, - cute.make_tensor(c.iterator, cute.select(c.layout, mode=[0, 1, 2])), - stream, - ) - - -@cute.experimental.jit -def _bmm_bias( - gemm_op: cutlass.Constexpr, - a: cute.Tensor, - b: cute.Tensor, - c: cute.Tensor, - bias: cute.Tensor, - stream: _cuda.CUstream, -): - gemm_op( - cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), - cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), - cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])), - cute.make_tensor(bias.iterator, cute.select(bias.layout, mode=[1, 2, 0])), - stream, - ) - - -def _from_dlpack_dynamic(tensor, leading_dim: int, assumed_align: int = 32): - return from_dlpack(tensor, assumed_align=assumed_align).mark_layout_dynamic( - leading_dim=leading_dim - ) - - -def _detect_leading_dim(tensor: _torch.Tensor) -> int: - # Ignore synthetic batch stride, including 1x1. - for dim, stride in enumerate(tensor.stride()[1:], start=1): - if stride == 1: - return dim - raise ValueError("tensor has no stride-1 dimension") - - -def _make_layout_tensor( - shape: tuple[int, ...], dtype: _torch.dtype, leading_dim: int -) -> _torch.Tensor: - permutation = [dim for dim in range(len(shape)) if dim != leading_dim] + [ - leading_dim - ] - return _torch.empty( - tuple(shape[dim] for dim in permutation), dtype=dtype, device="cuda" - ).permute([permutation.index(dim) for dim in range(len(shape))]) - - -def _make_compile_repr_tensors( - dtype: _torch.dtype, - has_bias: bool, - a_leading: int, - b_leading: int, - c_leading: int, -): - m, n, k, batch = 64, 8, _CTA_K, 1 - tensors = tuple( - _from_dlpack_dynamic( - _make_layout_tensor(shape, dtype, leading_dim), leading_dim - ) - for shape, leading_dim in zip( - ((batch, n, k), (batch, k, m), (batch, n, m)), - (a_leading, b_leading, c_leading), - strict=True, - ) - ) - if not has_bias: - return (*tensors, None) - return ( - *tensors, - _from_dlpack_dynamic( - _torch.empty((n,), dtype=dtype, device="cuda").as_strided( - size=(batch, n, m), stride=(0, 1, 0) - ), - 1, - 2, - ), - ) - - -def _to_cute_swap(a, b, out, bias): - a_swap = b.unsqueeze(0).transpose(-2, -1) - b_swap = a.unsqueeze(0).transpose(-2, -1) - c_swap = out.unsqueeze(0).transpose(-2, -1) - leading_dims = tuple( - _detect_leading_dim(tensor) for tensor in (a_swap, b_swap, c_swap) - ) - cute_tensors = tuple( - _from_dlpack_dynamic(tensor, leading_dim) - for tensor, leading_dim in zip( - (a_swap, b_swap, c_swap), leading_dims, strict=True - ) - ) - if bias is None: - return (*cute_tensors, None, leading_dims) - return ( - *cute_tensors, - _from_dlpack_dynamic( - bias.as_strided( - size=(1, c_swap.shape[1], c_swap.shape[2]), stride=(0, 1, 0) - ), - 1, - 2, - ), - leading_dims, - ) - - -# Tactic hashes all compile-time tile, split, and stage fields. -_SPLITK_COMPILE_CACHE: dict = {} - - -def _get_compiled_splitk_kernel( - dtype, - tactic: SplitKTactic, - use_pdl: bool, - has_bias: bool, - leading_dims: tuple[int, int, int], -): - key = ( - dtype, - tactic, - use_pdl, - has_bias, - *leading_dims, - ) - cached = _SPLITK_COMPILE_CACHE.get(key) - if cached is not None: - return cached - - if dtype not in _SUPPORTED_TORCH_DTYPES: - raise ValueError( - f"split-K dense GEMM supports {_SUPPORTED_TORCH_DTYPES}; got {dtype}" - ) - - kernel = SplitKDenseGemmKernel( - tactic=tactic, - use_pdl=use_pdl, - has_bias=has_bias, - ) - compile_tensors = _make_compile_repr_tensors(dtype, has_bias, *leading_dims) - stream = _cuda.CUstream(_torch.cuda.current_stream().cuda_stream) - if has_bias: - compiled = cute_ext.compile(_bmm_bias, kernel, *compile_tensors, stream) - else: - compiled = cute_ext.compile(_bmm_no_bias, kernel, *compile_tensors[:3], stream) - _SPLITK_COMPILE_CACHE[key] = compiled - return compiled - - -def _validate_runtime_tensors(a, b, bias, out) -> tuple[int, int, int]: - tensors = (a, b, out) + ((bias,) if bias is not None else ()) - if any(not isinstance(tensor, _torch.Tensor) for tensor in tensors): - raise ValueError("a, b, out, and bias must be torch tensors") - if a.ndim != 2 or b.ndim != 2 or out.ndim != 2: - raise ValueError("split-K dense GEMM accepts only 2D tensors") - if a.device.type != "cuda" or any(tensor.device != a.device for tensor in tensors): - raise ValueError("all tensors must be on the same CUDA device") - if a.dtype not in _SUPPORTED_TORCH_DTYPES or any( - tensor.dtype != a.dtype for tensor in tensors - ): - raise ValueError("a, b, out, and bias must share BF16 or FP16 dtype") - if any( - not (tensor.is_contiguous() or tensor.t().is_contiguous()) - for tensor in (a, b, out) - ): - raise ValueError( - "a, b, and out must be dense row-major or column-major matrices" - ) - if any(tensor.data_ptr() % 32 for tensor in (a, b, out)): - raise ValueError("a, b, and out must be 32-byte aligned") - - m, k = a.shape - if b.shape[0] != k: - raise ValueError( - f"incompatible shapes: a is {tuple(a.shape)}, b is {tuple(b.shape)}" - ) - n = b.shape[1] - if out.shape != (m, n): - raise ValueError(f"out must have shape {(m, n)}, got {tuple(out.shape)}") - if bias is not None and ( - bias.ndim != 1 or bias.shape[0] != n or not bias.is_contiguous() - ): - raise ValueError( - f"bias must be contiguous with shape {(n,)}, " - f"got shape {tuple(bias.shape)} and stride {bias.stride()}" - ) - - return m, n, k - - -def run_splitk_dense( - a, - b, - bias, - out, - pdl: bool, - tactic: SplitKTactic, -): - """Run ``A[M,K] @ B[K,N]`` with the ``mm_bf16`` layouts.""" - validate_tactic(tactic, *_validate_runtime_tensors(a, b, bias, out)) - has_bias = bias is not None - cute_tensors = _to_cute_swap(a, b, out, bias) - compiled = _get_compiled_splitk_kernel( - dtype=a.dtype, - tactic=tactic, - use_pdl=pdl, - has_bias=has_bias, - leading_dims=cute_tensors[4], - ) - stream = _cuda.CUstream(_torch.cuda.current_stream(a.device).cuda_stream) - if has_bias: - compiled(*cute_tensors[:4], stream) - else: - compiled(*cute_tensors[:3], stream) - return out diff --git a/python/sglang/srt/arg_groups/fields/exec_.py b/python/sglang/srt/arg_groups/fields/exec_.py index bb1c3bbed..f80d877c9 100644 --- a/python/sglang/srt/arg_groups/fields/exec_.py +++ b/python/sglang/srt/arg_groups/fields/exec_.py @@ -178,9 +178,9 @@ class ExecKernel: bf16_gemm_backend: A[ str, Arg( - help="Choose the backend for unquantized BF16 GEMM operations. Options: 'auto' (default; selects 'cutedsl' on SM10x GPUs, except deterministic inference selects 'torch'; otherwise uses cuBLAS via torch.nn.functional.linear), 'cutedsl' (SGLang JIT CuTe DSL TGV BF16 GEMM on SM10x; dispatches between the allowlisted low-M Split-K kernel, the CuTe DSL kernel, and cuBLAS; set SGLANG_ENABLE_BF16_SPLITK_GEMM=0 to disable Split-K), 'flashinfer_pr4266' (legacy compatibility alias for the optimized CuTe DSL path), 'gemv', 'torch' (always uses cuBLAS via torch.nn.functional.linear).", + help="Choose the backend for unquantized BF16 GEMM operations. Options: 'auto' (default; selects 'cutedsl' on SM10x GPUs, except deterministic inference selects 'torch'; otherwise uses cuBLAS via torch.nn.functional.linear), 'cutedsl' (SGLang JIT CuTe DSL TGV BF16 GEMM on SM10x; dispatches between the allowlisted low-M Split-K kernel, the CuTe DSL kernel, and cuBLAS; set SGLANG_ENABLE_BF16_SPLITK_GEMM=0 to disable Split-K), 'gemv', 'torch' (always uses cuBLAS via torch.nn.functional.linear).", cli_name="--bf16-gemm-backend", - choices=["auto", "cutedsl", "flashinfer_pr4266", "gemv", "torch"], + choices=["auto", "cutedsl", "gemv", "torch"], ), ] = "auto" dsa_prefill_backend: A[ diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index cc262fee0..1239dd34d 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -1794,8 +1794,6 @@ _DEPRECATED_ENVS: Dict[str, _DeprecatedEnv] = { "SGLANG_OPT_SWA_EVICT_DROP_PAGE_MARGIN": _DeprecatedEnv(), # sconv-family kernels always use the CUDA-JIT ports when supported; no toggle. "SGLANG_OPT_USE_CUDA_SCONV": _DeprecatedEnv(), - # The direct dense BF16 GEMM source is vendored in-tree. - "SGLANG_FLASHINFER_PR4266_SOURCE": _DeprecatedEnv(), # DSV4 compressor V2 is always used. "SGLANG_OPT_USE_COMPRESSOR_V2": _DeprecatedEnv(), "SGLANG_ENABLE_HICACHE_BUFFER_ANCHOR_LOCK": _DeprecatedEnv( diff --git a/python/sglang/srt/layers/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py index d11fe126e..1eeca21f8 100644 --- a/python/sglang/srt/layers/quantization/unquant.py +++ b/python/sglang/srt/layers/quantization/unquant.py @@ -76,7 +76,6 @@ if _use_aiter: class Bf16GemmBackend(Enum): AUTO = "auto" CUTEDSL = "cutedsl" - FLASHINFER_PR4266 = "flashinfer_pr4266" GEMV = "gemv" TORCH = "torch" @@ -89,28 +88,22 @@ class Bf16GemmBackend(Enum): def is_gemv(self) -> bool: return self == Bf16GemmBackend.GEMV - def is_flashinfer_pr4266(self) -> bool: - return self == Bf16GemmBackend.FLASHINFER_PR4266 - - def is_optimized(self) -> bool: - return self.is_cutedsl() or self.is_flashinfer_pr4266() - _BF16_GEMM_BACKEND: Optional[Bf16GemmBackend] = None _cutedsl_bf16_gemm = None _use_cutedsl_bf16_gemm = None _hopper_bf16_gemv = None _use_hopper_bf16_gemv = None -_flashinfer_pr4266_splitk_tactic = None -_flashinfer_pr4266_run_splitk_dense = None -_flashinfer_pr4266_direct_default_tactic = None -_flashinfer_pr4266_prefer_direct = None -_flashinfer_pr4266_run_direct_dense = None +_splitk_tactic = None +_run_splitk_dense = None +_direct_default_tactic = None +_prefer_direct = None +_run_direct_dense = None _enable_bf16_splitk_gemm = False # GB300 TP16 tactics measured under CUDA graph replay with PDL and cold weights. # Unlisted shapes, including M=64, retain the existing TGV/cuBLAS path. -_FLASHINFER_PR4266_TUNED_TACTICS = { +_BF16_SPLITK_TUNED_TACTICS = { (1, 256, 8192): (64, 8, 4, 11), (2, 256, 8192): (64, 8, 4, 11), (4, 256, 8192): (64, 8, 4, 11), @@ -142,23 +135,23 @@ _FLASHINFER_PR4266_TUNED_TACTICS = { } -def use_flashinfer_pr4266_bf16_gemm(m: int, n: int, k: int) -> bool: - return (m, n, k) in _FLASHINFER_PR4266_TUNED_TACTICS +def use_bf16_splitk_gemm(m: int, n: int, k: int) -> bool: + return (m, n, k) in _BF16_SPLITK_TUNED_TACTICS def should_enable_bf16_splitk_gemm(backend: Bf16GemmBackend) -> bool: """Return whether the optional Split-K path should be initialized.""" - return backend.is_optimized() and envs.SGLANG_ENABLE_BF16_SPLITK_GEMM.get() + return backend.is_cutedsl() and envs.SGLANG_ENABLE_BF16_SPLITK_GEMM.get() def initialize_bf16_gemm_config() -> None: global _BF16_GEMM_BACKEND global _cutedsl_bf16_gemm, _use_cutedsl_bf16_gemm - global _flashinfer_pr4266_splitk_tactic - global _flashinfer_pr4266_run_splitk_dense - global _flashinfer_pr4266_direct_default_tactic - global _flashinfer_pr4266_prefer_direct - global _flashinfer_pr4266_run_direct_dense + global _splitk_tactic + global _run_splitk_dense + global _direct_default_tactic + global _prefer_direct + global _run_direct_dense global _enable_bf16_splitk_gemm backend_str = get_exec().kernel.bf16_gemm_backend @@ -183,7 +176,7 @@ def initialize_bf16_gemm_config() -> None: _hopper_bf16_gemv = hopper_bf16_gemv _use_hopper_bf16_gemv = use_hopper_bf16_gemv - elif backend.is_optimized(): + elif backend.is_cutedsl(): if get_exec().deterministic.enable_deterministic_inference: raise ValueError( "--bf16-gemm-backend cutedsl is batch-size dependent and cannot " @@ -204,21 +197,21 @@ def initialize_bf16_gemm_config() -> None: _enable_bf16_splitk_gemm = False if should_enable_bf16_splitk_gemm(backend): - from sglang.kernels.ops.gemm.flashinfer_pr4266_dense_bf16_gemm_sm100_direct import ( + from flashinfer.gemm.kernels.dense_bf16_gemm_direct import ( default_tactic, prefer_direct_bf16_gemm_sm100, run_direct_dense, ) - from sglang.kernels.ops.gemm.flashinfer_pr4266_dense_bf16_gemm_sm100_splitk import ( + from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import ( SplitKTactic, run_splitk_dense, ) - _flashinfer_pr4266_splitk_tactic = SplitKTactic - _flashinfer_pr4266_run_splitk_dense = run_splitk_dense - _flashinfer_pr4266_direct_default_tactic = default_tactic - _flashinfer_pr4266_prefer_direct = prefer_direct_bf16_gemm_sm100 - _flashinfer_pr4266_run_direct_dense = run_direct_dense + _splitk_tactic = SplitKTactic + _run_splitk_dense = run_splitk_dense + _direct_default_tactic = default_tactic + _prefer_direct = prefer_direct_bf16_gemm_sm100 + _run_direct_dense = run_direct_dense _enable_bf16_splitk_gemm = True _BF16_GEMM_BACKEND = backend @@ -230,20 +223,18 @@ def _bf16_gemm_dispatch_fake( return x.new_empty((*x.shape[:-1], weight.shape[0])) -def _flashinfer_pr4266_bf16_gemm( +def _bf16_splitk_gemm( x: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor] ) -> torch.Tensor: x_2d = x.view(-1, x.shape[-1]) out = torch.empty((x_2d.shape[0], weight.shape[0]), dtype=x.dtype, device=x.device) m, n, k = x_2d.shape[0], weight.shape[0], weight.shape[1] - if bias is None and _flashinfer_pr4266_prefer_direct(m, n, k): - tactic = _flashinfer_pr4266_direct_default_tactic(m, n, k) - _flashinfer_pr4266_run_direct_dense(x_2d, weight.T, out, True, tactic) + if bias is None and _prefer_direct(m, n, k): + tactic = _direct_default_tactic(m, n, k) + _run_direct_dense(x_2d, weight.T, out, True, tactic) else: - tactic = _flashinfer_pr4266_splitk_tactic( - *_FLASHINFER_PR4266_TUNED_TACTICS[(m, n, k)] - ) - _flashinfer_pr4266_run_splitk_dense( + tactic = _splitk_tactic(*_BF16_SPLITK_TUNED_TACTICS[(m, n, k)]) + _run_splitk_dense( x_2d, weight.T, bias, @@ -261,10 +252,10 @@ def _bf16_gemm_dispatch_impl( addend: Optional[torch.Tensor] = None, ) -> torch.Tensor: m = x.numel() // x.shape[-1] - if _enable_bf16_splitk_gemm and use_flashinfer_pr4266_bf16_gemm( + if _enable_bf16_splitk_gemm and use_bf16_splitk_gemm( m, weight.shape[0], weight.shape[1] ): - output = _flashinfer_pr4266_bf16_gemm(x, weight, bias) + output = _bf16_splitk_gemm(x, weight, bias) elif ( _use_hopper_bf16_gemv is not None and bias is None @@ -423,7 +414,7 @@ class UnquantizedLinearMethod(LinearMethodBase): return tgemm.mm(x, layer.weight, bias, otype=x.dtype) elif ( - get_bf16_gemm_backend().is_optimized() + get_bf16_gemm_backend().is_cutedsl() and x.is_cuda and x.dtype == torch.bfloat16 and layer.weight.dtype == torch.bfloat16 diff --git a/test/registered/unit/layers/quantization/test_flashinfer_pr4266_bf16_gemm.py b/test/registered/unit/layers/quantization/test_bf16_splitk_gemm.py similarity index 65% rename from test/registered/unit/layers/quantization/test_flashinfer_pr4266_bf16_gemm.py rename to test/registered/unit/layers/quantization/test_bf16_splitk_gemm.py index 569426e56..381656d54 100644 --- a/test/registered/unit/layers/quantization/test_flashinfer_pr4266_bf16_gemm.py +++ b/test/registered/unit/layers/quantization/test_bf16_splitk_gemm.py @@ -2,25 +2,25 @@ import pytest from sglang.srt.environ import envs from sglang.srt.layers.quantization.unquant import ( - _FLASHINFER_PR4266_TUNED_TACTICS, + _BF16_SPLITK_TUNED_TACTICS, Bf16GemmBackend, should_enable_bf16_splitk_gemm, - use_flashinfer_pr4266_bf16_gemm, + use_bf16_splitk_gemm, ) from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=11, suite="base-a-test-cpu") -@pytest.mark.parametrize("m,n,k", _FLASHINFER_PR4266_TUNED_TACTICS) -def test_flashinfer_pr4266_selects_tuned_oakhaven_shape(m: int, n: int, k: int): - assert use_flashinfer_pr4266_bf16_gemm(m, n, k) +@pytest.mark.parametrize("m,n,k", _BF16_SPLITK_TUNED_TACTICS) +def test_splitk_selects_tuned_oakhaven_shape(m: int, n: int, k: int): + assert use_bf16_splitk_gemm(m, n, k) @pytest.mark.parametrize("m", [0, 33, 64]) @pytest.mark.parametrize("n,k", [(256, 8192), (512, 8192), (2304, 8192), (2560, 8192)]) -def test_flashinfer_pr4266_keeps_large_m_on_existing_path(m: int, n: int, k: int): - assert not use_flashinfer_pr4266_bf16_gemm(m, n, k) +def test_splitk_keeps_large_m_on_existing_path(m: int, n: int, k: int): + assert not use_bf16_splitk_gemm(m, n, k) @pytest.mark.parametrize( @@ -32,12 +32,8 @@ def test_flashinfer_pr4266_keeps_large_m_on_existing_path(m: int, n: int, k: int (32, 4096, 8192), ], ) -def test_flashinfer_pr4266_rejects_unmeasured_shapes(shape: tuple[int, int, int]): - assert not use_flashinfer_pr4266_bf16_gemm(*shape) - - -def test_flashinfer_pr4266_backend_is_explicit(): - assert Bf16GemmBackend.FLASHINFER_PR4266.value == "flashinfer_pr4266" +def test_splitk_rejects_unmeasured_shapes(shape: tuple[int, int, int]): + assert not use_bf16_splitk_gemm(*shape) def test_bf16_splitk_is_enabled_by_default(): diff --git a/test/registered/unit/layers/quantization/test_unquant_apply_with_addend.py b/test/registered/unit/layers/quantization/test_unquant_apply_with_addend.py index 85741b4a5..9651d663a 100644 --- a/test/registered/unit/layers/quantization/test_unquant_apply_with_addend.py +++ b/test/registered/unit/layers/quantization/test_unquant_apply_with_addend.py @@ -168,15 +168,11 @@ class TestApplyWithAddend(CustomTestCase): ) elif route == "splitk": enter(patch.object(unquant, "_enable_bf16_splitk_gemm", True)) - enter( - patch.object( - unquant, "use_flashinfer_pr4266_bf16_gemm", lambda *a: True - ) - ) + enter(patch.object(unquant, "use_bf16_splitk_gemm", lambda *a: True)) enter( patch.object( unquant, - "_flashinfer_pr4266_bf16_gemm", + "_bf16_splitk_gemm", _fake_kernel(kernel_calls, route), ) )