wehub-resource-sync
59a0a3844c
PR Test AMD / cancel-on-close (push) Has been skipped
PR Test NVIDIA ARM / scan (push) Has been skipped
PR Test NVIDIA / cancel-on-close (push) Has been skipped
PR Test AMD / scan (push) Has been skipped
PR Test NVIDIA ARM / cancel-on-close (push) Has been skipped
PR Test NVIDIA / scan (push) Has been skipped
Release Docker Images / build (cu129-torch-2.11.0) (push) Has been skipped
Release Docker Images / build (cu130-torch-2.11.0) (push) Has been skipped
Release PyPI / publish (push) Has been skipped
Scheduler Python Test / test (push) Successful in 27m19s
Docs / build (push) Successful in 28m8s
Scheduler C++ Test / test (push) Successful in 28m19s
Scheduler C++ Test / test-flat (push) Successful in 28m18s
Docs / deploy (push) Has been cancelled
PR Test AMD / finish (push) Has been cancelled
PR Test NVIDIA / finish (push) Has been cancelled
PR Test NVIDIA ARM / finish (push) Has been cancelled
PR Test NVIDIA ARM / ${{ matrix.name }} (${{ matrix.runner }}) (push) Has been cancelled
PR Test AMD / ${{ matrix.name }} (${{ matrix.runner }}) (push) Has been cancelled
PR Test NVIDIA / ${{ matrix.name }} (${{ matrix.runner }}) (push) Has been cancelled
672 lines
24 KiB
Python
672 lines
24 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2025, Tri Dao.
|
|
# SPDX-FileCopyrightText: Copyright (c) 2025, Wentao Guo, Ted Zadouri, Tri Dao.
|
|
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
#
|
|
# This file contains code derived from the Quack library:
|
|
# https://github.com/Dao-AILab/quack
|
|
# Originally integrated into TensorRT-LLM:
|
|
# https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/cute_dsl_kernels/argmax.py
|
|
#
|
|
# Argmax kernel using CuTe DSL.
|
|
#
|
|
# This module pulls in ``cuda.bindings.driver`` and ``cutlass``, which are
|
|
# NVIDIA-only Python packages. **Do not import this module directly from the
|
|
# runtime.** Use ``tokenspeed_kernel.ops.sampling.cute_dsl`` instead — that
|
|
# wrapper gates on ``current_platform().is_nvidia`` before importing this file
|
|
# and falls back to ``torch.argmax`` on every other platform.
|
|
|
|
from typing import Optional, Tuple, Type
|
|
|
|
import cuda.bindings.driver as cuda
|
|
import cutlass
|
|
import cutlass.cute as cute
|
|
import torch
|
|
from cutlass._mlir.dialects import llvm
|
|
from cutlass.cute.arch.nvvm_wrappers import FULL_MASK
|
|
from cutlass.cute.runtime import from_dlpack
|
|
from cutlass.cute.typing import Float32, Int, Int32
|
|
from cutlass.cutlass_dsl import T, dsl_user_op
|
|
|
|
# ============================================================================
|
|
# Torch to CuTE dtype mapping
|
|
# ============================================================================
|
|
torch2cute_dtype_map = {
|
|
torch.float16: cutlass.Float16,
|
|
torch.bfloat16: cutlass.BFloat16,
|
|
torch.float32: cutlass.Float32,
|
|
}
|
|
|
|
|
|
# ============================================================================
|
|
# CUDA Graph compatibility wrapper
|
|
# ============================================================================
|
|
class CUDAGraphCompatibleWrapper:
|
|
"""Wrapper to make tensors compatible with CUDA graph capture for DLPack export."""
|
|
|
|
def __init__(self, tensor):
|
|
self._tensor = tensor
|
|
|
|
def __dlpack__(self, stream=None):
|
|
return self._tensor.__dlpack__(stream=-1)
|
|
|
|
def __dlpack_device__(self):
|
|
return self._tensor.__dlpack_device__()
|
|
|
|
|
|
# ============================================================================
|
|
# Utility functions from quack/utils.py
|
|
# ============================================================================
|
|
@dsl_user_op
|
|
def elem_pointer(
|
|
x: cute.Tensor, coord: cute.Coord, *, loc=None, ip=None
|
|
) -> cute.Pointer:
|
|
return x.iterator + cute.crd2idx(coord, x.layout, loc=loc, ip=ip)
|
|
|
|
|
|
@dsl_user_op
|
|
def set_block_rank(
|
|
smem_ptr: cute.Pointer, peer_cta_rank_in_cluster: cute.Int32, *, loc=None, ip=None
|
|
) -> cutlass.Int32:
|
|
smem_ptr_i32 = smem_ptr.toint(loc=loc, ip=ip).ir_value()
|
|
return cutlass.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(
|
|
val: float | Float32 | cutlass.Int64,
|
|
smem_ptr: cute.Pointer,
|
|
mbar_ptr: cute.Pointer,
|
|
peer_cta_rank_in_cluster: cute.typing.Int,
|
|
*,
|
|
loc=None,
|
|
ip=None,
|
|
) -> None:
|
|
remote_smem_ptr_i32 = set_block_rank(
|
|
smem_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip
|
|
).ir_value()
|
|
remote_mbar_ptr_i32 = set_block_rank(
|
|
mbar_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip
|
|
).ir_value()
|
|
if cutlass.const_expr(isinstance(val, float)):
|
|
val = Float32(val)
|
|
assert isinstance(val, (Float32, Int32, cutlass.Int64))
|
|
suffix = {Float32: "f32", Int32: "s32", cutlass.Int64: "s64"}[type(val)]
|
|
constraint = {Float32: "f", Int32: "r", cutlass.Int64: "l"}[type(val)]
|
|
llvm.inline_asm(
|
|
None,
|
|
[remote_smem_ptr_i32, val.ir_value(loc=loc, ip=ip), remote_mbar_ptr_i32],
|
|
f"st.async.shared::cluster.mbarrier::complete_tx::bytes.{suffix} [$0], $1, [$2];",
|
|
f"r,{constraint},r",
|
|
has_side_effects=True,
|
|
is_align_stack=False,
|
|
asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
)
|
|
|
|
|
|
@cute.jit
|
|
def predicate_k(tAcA: cute.Tensor, limit: cutlass.Int32) -> cute.Tensor:
|
|
tApA = cute.make_rmem_tensor(
|
|
cute.make_layout(
|
|
(
|
|
cute.size(tAcA, mode=[0, 1]),
|
|
cute.size(tAcA, mode=[1]),
|
|
cute.size(tAcA, mode=[2]),
|
|
),
|
|
stride=(cute.size(tAcA, mode=[2]), 0, 1),
|
|
),
|
|
cutlass.Boolean,
|
|
)
|
|
for rest_v in cutlass.range_constexpr(tApA.shape[0]):
|
|
for rest_k in cutlass.range_constexpr(tApA.shape[2]):
|
|
tApA[rest_v, 0, rest_k] = cute.elem_less(
|
|
tAcA[(0, rest_v), 0, rest_k][1], limit
|
|
)
|
|
return tApA
|
|
|
|
|
|
@cute.jit
|
|
def fill_oob(
|
|
tXsX: cute.Tensor, tXpX: Optional[cute.Tensor], fill_value: cute.Numeric
|
|
) -> None:
|
|
tXrX_fill = cute.make_fragment_like(tXsX[(None, 0), None, 0])
|
|
tXrX_fill.fill(fill_value)
|
|
for rest_v in cutlass.range_constexpr(tXsX.shape[0][1]):
|
|
for rest_k in cutlass.range_constexpr(tXsX.shape[2]):
|
|
if cutlass.const_expr(tXpX is not None):
|
|
if not tXpX[rest_v, 0, rest_k]:
|
|
cute.autovec_copy(tXrX_fill, tXsX[(None, rest_v), None, rest_k])
|
|
else:
|
|
cute.autovec_copy(tXrX_fill, tXsX[(None, rest_v), None, rest_k])
|
|
|
|
|
|
@dsl_user_op
|
|
def domain_offset_i64(
|
|
coord: cute.Coord, tensor: cute.Tensor, *, loc=None, ip=None
|
|
) -> cute.Tensor:
|
|
flat_coord_i64 = tuple(cutlass.Int64(c) for c in cute.flatten(coord))
|
|
flat_stride = cute.flatten_to_tuple(tensor.stride)
|
|
offset = sum(c * s for c, s in zip(flat_coord_i64, flat_stride))
|
|
new_ptr = cute.make_ptr(
|
|
tensor.element_type,
|
|
tensor.iterator.toint() + offset * tensor.element_type.width // 8,
|
|
tensor.memspace,
|
|
assumed_align=tensor.iterator.max_alignment,
|
|
)
|
|
return cute.make_tensor(new_ptr, tensor.layout)
|
|
|
|
|
|
# ============================================================================
|
|
# Inline PTX for redux.sync operations
|
|
# ============================================================================
|
|
@dsl_user_op
|
|
def ptx_redux_sync_max_f32(
|
|
value: Float32, mask: Int = FULL_MASK, *, loc=None, ip=None
|
|
) -> Float32:
|
|
return Float32(
|
|
llvm.inline_asm(
|
|
T.f32(),
|
|
[
|
|
Float32(value).ir_value(loc=loc, ip=ip),
|
|
Int32(mask).ir_value(loc=loc, ip=ip),
|
|
],
|
|
"""redux.sync.max.f32 $0, $1, $2;""",
|
|
"=f,f,i",
|
|
has_side_effects=True,
|
|
is_align_stack=False,
|
|
asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
)
|
|
)
|
|
|
|
|
|
@dsl_user_op
|
|
def ptx_redux_sync_min_u32(
|
|
value: Int32, mask: Int = FULL_MASK, *, loc=None, ip=None
|
|
) -> Int32:
|
|
return Int32(
|
|
llvm.inline_asm(
|
|
T.i32(),
|
|
[
|
|
Int32(value).ir_value(loc=loc, ip=ip),
|
|
Int32(mask).ir_value(loc=loc, ip=ip),
|
|
],
|
|
"""redux.sync.min.u32 $0, $1, $2;""",
|
|
"=r,r,i",
|
|
has_side_effects=True,
|
|
is_align_stack=False,
|
|
asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
)
|
|
)
|
|
|
|
|
|
@dsl_user_op
|
|
def ptx_select_argmax_candidate(
|
|
current_max: Float32, warp_max: Float32, current_argmax: Int32, *, loc=None, ip=None
|
|
) -> Int32:
|
|
return Int32(
|
|
llvm.inline_asm(
|
|
T.i32(),
|
|
[
|
|
Float32(current_max).ir_value(loc=loc, ip=ip),
|
|
Float32(warp_max).ir_value(loc=loc, ip=ip),
|
|
Int32(current_argmax).ir_value(loc=loc, ip=ip),
|
|
],
|
|
"""{
|
|
.reg .pred p;
|
|
setp.eq.f32 p, $1, $2;
|
|
selp.s32 $0, $3, 0xffffffff, p;
|
|
}""",
|
|
"=r,f,f,r",
|
|
has_side_effects=False,
|
|
is_align_stack=False,
|
|
asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
)
|
|
)
|
|
|
|
|
|
@cute.jit
|
|
def warp_argmax_redux(current_max: Float32, current_argmax: Int32):
|
|
"""Redux-based warp argmax - only works on sm_100f (Blackwell)."""
|
|
warp_max = ptx_redux_sync_max_f32(current_max)
|
|
candidate_idx = ptx_select_argmax_candidate(current_max, warp_max, current_argmax)
|
|
winning_idx = ptx_redux_sync_min_u32(candidate_idx)
|
|
return warp_max, winning_idx
|
|
|
|
|
|
@cute.jit
|
|
def warp_reduce_argmax(current_max: Float32, current_argmax: Int32):
|
|
"""Shuffle-based warp argmax - works on all architectures (Hopper+).
|
|
|
|
Ties are broken to the lowest index, matching ``torch.argmax``. The redux
|
|
path (``warp_argmax_redux``) already does this via ``redux.sync.min.u32``;
|
|
this path used to keep whichever side the strict ``>`` happened to retain,
|
|
so the tie-break-to-lowest invariant could fail on Hopper. The extra
|
|
``==`` arm makes the two paths agree on Hopper and Blackwell.
|
|
"""
|
|
warp_max = current_max
|
|
warp_argmax = current_argmax
|
|
|
|
# Butterfly shuffle reduction.
|
|
for i in cutlass.range_constexpr(int(5)): # log2(32) = 5 iterations
|
|
other_max = cute.arch.shuffle_sync_bfly(warp_max, offset=1 << i)
|
|
other_argmax = cute.arch.shuffle_sync_bfly(warp_argmax, offset=1 << i)
|
|
|
|
if other_max > warp_max:
|
|
warp_max = other_max
|
|
warp_argmax = other_argmax
|
|
elif other_max == warp_max:
|
|
if other_argmax < warp_argmax:
|
|
warp_argmax = other_argmax
|
|
|
|
return warp_max, warp_argmax
|
|
|
|
|
|
# ============================================================================
|
|
# Reduction Base class
|
|
# ============================================================================
|
|
class ReductionBase:
|
|
def __init__(
|
|
self,
|
|
dtype: Type[cutlass.Numeric],
|
|
N: int,
|
|
stage: int,
|
|
reduction_dtype=cutlass.Float32,
|
|
):
|
|
self.dtype = dtype
|
|
self.N = N
|
|
self.stage = stage
|
|
self.reduction_dtype = reduction_dtype
|
|
|
|
def _calculate_threads_per_row(self):
|
|
raise NotImplementedError()
|
|
|
|
def _set_cluster_n(self):
|
|
self.cluster_n = 1
|
|
|
|
def _get_num_threads(self):
|
|
return 128 if self.N <= 16384 else 256
|
|
|
|
def _get_tv_layout(self, num_copy_bits=128):
|
|
vecsize = num_copy_bits // self.dtype.width
|
|
num_threads = self._get_num_threads()
|
|
threads_per_row = self._calculate_threads_per_row()
|
|
num_blocks_N = cute.ceil_div(
|
|
self.N // vecsize, threads_per_row * self.cluster_n
|
|
)
|
|
cols_per_block = num_threads // threads_per_row
|
|
tiler_mn = (cols_per_block, vecsize * num_blocks_N * threads_per_row)
|
|
tv_layout = cute.make_layout(
|
|
((threads_per_row, cols_per_block), (vecsize, num_blocks_N)),
|
|
stride=(
|
|
(vecsize * cols_per_block, 1),
|
|
(cols_per_block, cols_per_block * vecsize * threads_per_row),
|
|
),
|
|
)
|
|
return tiler_mn, tv_layout
|
|
|
|
def _smem_size_in_bytes(self, tiler_mn, num_warps):
|
|
return (
|
|
cute.size_in_bytes(self.dtype, cute.make_layout(tiler_mn))
|
|
+ self.stage
|
|
* num_warps
|
|
* self.cluster_n
|
|
* (self.reduction_dtype.width // 8)
|
|
+ self.stage * (cutlass.Int64.width // 8)
|
|
)
|
|
|
|
def _get_reduction_buffer_layout(self, tv_layout: cute.Layout, cluster_n: int):
|
|
num_warps = cute.size(tv_layout, mode=[0]) // cute.arch.WARP_SIZE
|
|
warps_per_row = max(tv_layout.shape[0][0] // cute.arch.WARP_SIZE, 1)
|
|
return cute.make_ordered_layout(
|
|
(num_warps // warps_per_row, (warps_per_row, cluster_n), self.stage),
|
|
order=(1, 0, 2),
|
|
)
|
|
|
|
def _allocate_reduction_buffer_and_mbar(
|
|
self, smem: cutlass.utils.SmemAllocator, tv_layout: cute.Layout
|
|
) -> Tuple[cute.Tensor, Optional[cute.Pointer]]:
|
|
reduction_buffer = smem.allocate_tensor(
|
|
self.reduction_dtype,
|
|
self._get_reduction_buffer_layout(tv_layout, self.cluster_n),
|
|
byte_alignment=4,
|
|
)
|
|
if cutlass.const_expr(self.cluster_n > 1):
|
|
mbar_ptr = smem.allocate_array(cutlass.Int64, num_elems=self.stage)
|
|
else:
|
|
mbar_ptr = None
|
|
return reduction_buffer, mbar_ptr
|
|
|
|
@cute.jit
|
|
def _initialize_cluster(
|
|
self, tidx: cutlass.Int32, mbar_ptr: cute.Pointer, num_warps: int
|
|
):
|
|
if cutlass.const_expr(self.cluster_n > 1):
|
|
if tidx < self.stage:
|
|
cute.arch.mbarrier_init(mbar_ptr + tidx, 1)
|
|
cute.arch.mbarrier_init_fence()
|
|
cute.arch.cluster_arrive_relaxed()
|
|
|
|
|
|
# ============================================================================
|
|
# Argmax Kernel class
|
|
# ============================================================================
|
|
class ArgmaxKernel(ReductionBase):
|
|
def __init__(self, dtype: Type[cutlass.Numeric], N: int, use_redux: bool = False):
|
|
super().__init__(dtype, N, stage=1, reduction_dtype=cutlass.Float32)
|
|
# use_redux=True for Blackwell (sm_100f), False for Hopper (sm_90)
|
|
self.use_redux = use_redux
|
|
|
|
def _calculate_threads_per_row(self):
|
|
N = self.N
|
|
return (
|
|
8
|
|
if N <= 64
|
|
else (
|
|
16
|
|
if N <= 128
|
|
else (
|
|
32
|
|
if N <= 3072
|
|
else (64 if N <= 6144 else (128 if N <= 16384 else 256))
|
|
)
|
|
)
|
|
)
|
|
|
|
def _set_cluster_n(self):
|
|
N = self.N
|
|
if cutlass.const_expr(self.dtype.width == 16):
|
|
self.cluster_n = (
|
|
1
|
|
if N <= 16 * 1024
|
|
else (
|
|
2
|
|
if N <= 32 * 1024
|
|
else (4 if N <= 64 * 1024 else (8 if N <= 128 * 1024 else 16))
|
|
)
|
|
)
|
|
else:
|
|
self.cluster_n = (
|
|
1
|
|
if N <= 32 * 1024
|
|
else (
|
|
2
|
|
if N <= 64 * 1024
|
|
else (4 if N <= 128 * 1024 else (8 if N <= 256 * 1024 else 16))
|
|
)
|
|
)
|
|
|
|
def _get_reduction_buffer_layout(self, tv_layout: cute.Layout, cluster_n: int):
|
|
num_warps = cute.size(tv_layout, mode=[0]) // cute.arch.WARP_SIZE
|
|
warps_per_row = max(tv_layout.shape[0][0] // cute.arch.WARP_SIZE, 1)
|
|
return cute.make_ordered_layout(
|
|
(num_warps // warps_per_row, (warps_per_row, cluster_n), self.stage, 2),
|
|
order=(1, 0, 2, 3),
|
|
)
|
|
|
|
def _smem_size_in_bytes(self, tiler_mn, num_warps):
|
|
return (
|
|
cute.size_in_bytes(self.dtype, cute.make_layout(tiler_mn))
|
|
+ 2
|
|
* self.stage
|
|
* num_warps
|
|
* self.cluster_n
|
|
* (self.reduction_dtype.width // 8)
|
|
+ self.stage * (cutlass.Int64.width // 8)
|
|
)
|
|
|
|
@cute.jit
|
|
def __call__(
|
|
self,
|
|
mX: cute.Tensor,
|
|
mO_max: cute.Tensor,
|
|
mO_idx: cute.Tensor,
|
|
stream: cuda.CUstream,
|
|
):
|
|
self._set_cluster_n()
|
|
tiler_mn, tv_layout = self._get_tv_layout()
|
|
num_threads = cute.size(tv_layout, mode=[0])
|
|
num_warps = num_threads // cute.arch.WARP_SIZE
|
|
|
|
self.kernel(mX, mO_max, mO_idx, tv_layout, tiler_mn).launch(
|
|
grid=[cute.ceil_div(mX.shape[0], tiler_mn[0]), self.cluster_n, 1],
|
|
block=[num_threads, 1, 1],
|
|
cluster=(
|
|
[1, self.cluster_n, 1]
|
|
if cutlass.const_expr(self.cluster_n > 1)
|
|
else None
|
|
),
|
|
smem=self._smem_size_in_bytes(tiler_mn, num_warps),
|
|
stream=stream,
|
|
)
|
|
|
|
@cute.kernel
|
|
def kernel(
|
|
self,
|
|
mX: cute.Tensor,
|
|
mO_max: cute.Tensor,
|
|
mO_idx: cute.Tensor,
|
|
tv_layout: cute.Layout,
|
|
tiler_mn: cute.Shape,
|
|
):
|
|
tidx, _, _ = cute.arch.thread_idx()
|
|
bidx, bidy, bidz = cute.arch.block_idx()
|
|
|
|
if cutlass.const_expr(self.cluster_n > 1):
|
|
cluster_y = cute.arch.block_idx()[1]
|
|
else:
|
|
cluster_y = cutlass.const_expr(0)
|
|
|
|
shape = mX.shape
|
|
idX = cute.make_identity_tensor(shape)
|
|
|
|
mX = domain_offset_i64((bidx * tiler_mn[0], 0), mX)
|
|
gX = cute.local_tile(mX, tiler_mn, (0, cluster_y))
|
|
# Each output is 1D (M,); only a row-axis offset is needed.
|
|
mO_max = domain_offset_i64((bidx * tiler_mn[0],), mO_max)
|
|
mO_idx = domain_offset_i64((bidx * tiler_mn[0],), mO_idx)
|
|
cX = cute.local_tile(idX, tiler_mn, (bidx, cluster_y))
|
|
|
|
smem = cutlass.utils.SmemAllocator()
|
|
sX = smem.allocate_tensor(
|
|
mX.element_type,
|
|
cute.make_ordered_layout(tiler_mn, order=(1, 0)),
|
|
byte_alignment=16,
|
|
)
|
|
reduction_buffer, mbar_ptr = self._allocate_reduction_buffer_and_mbar(
|
|
smem, tv_layout
|
|
)
|
|
|
|
copy_atom_load_X = cute.make_copy_atom(
|
|
cute.nvgpu.cpasync.CopyG2SOp(), mX.element_type, num_bits_per_copy=128
|
|
)
|
|
thr_copy_X = cute.make_tiled_copy(
|
|
copy_atom_load_X, tv_layout, tiler_mn
|
|
).get_slice(tidx)
|
|
|
|
tXgX = thr_copy_X.partition_S(gX)
|
|
tXsX = thr_copy_X.partition_D(sX)
|
|
tXcX = thr_copy_X.partition_S(cX)[(0, None), None, None]
|
|
|
|
tvlayout_cX = cute.composition(cX, tv_layout)
|
|
thr_coord = (tidx, (None, None))
|
|
thr_cX = tvlayout_cX[thr_coord]
|
|
|
|
tXrX = cute.make_fragment_like(tXgX)
|
|
num_warps = cute.size(tv_layout, mode=[0]) // cute.arch.WARP_SIZE
|
|
self._initialize_cluster(tidx, mbar_ptr, num_warps)
|
|
|
|
is_even_N = cutlass.const_expr(shape[1] == tiler_mn[1] * self.cluster_n)
|
|
tXpX = (
|
|
predicate_k(thr_copy_X.partition_S(cX), limit=shape[1])
|
|
if not is_even_N
|
|
else None
|
|
)
|
|
|
|
if tXcX[0][0] < shape[0]:
|
|
cute.copy(copy_atom_load_X, tXgX, tXsX, pred=tXpX)
|
|
cute.arch.cp_async_commit_group()
|
|
cute.arch.cp_async_wait_group(0)
|
|
|
|
if cutlass.const_expr(not is_even_N):
|
|
fill_oob(tXsX, tXpX, -tXsX.element_type.inf)
|
|
|
|
cute.autovec_copy(tXsX, tXrX)
|
|
x = tXrX.load().to(cute.Float32)
|
|
|
|
current_max = -tXsX.element_type.inf
|
|
current_argmax = Int32(0xFFFFFFFF)
|
|
|
|
for i in cutlass.range_constexpr(thr_cX.shape[0]):
|
|
for j in cutlass.range_constexpr(thr_cX.shape[1]):
|
|
col_idx = thr_cX[i, j][1]
|
|
linear_idx = i + j * thr_cX.shape[0]
|
|
element_value1 = x[linear_idx]
|
|
if element_value1 > current_max:
|
|
current_max = element_value1
|
|
current_argmax = Int32(col_idx)
|
|
|
|
lane_idx, warp_idx = cute.arch.lane_idx(), cute.arch.warp_idx()
|
|
if cutlass.const_expr(self.use_redux):
|
|
warp_max, warp_argmax = warp_argmax_redux(current_max, current_argmax)
|
|
else:
|
|
warp_max, warp_argmax = warp_reduce_argmax(current_max, current_argmax)
|
|
|
|
if cutlass.const_expr(self.cluster_n == 1):
|
|
warps_per_row = cute.size(reduction_buffer.shape[1])
|
|
row_idx, col_idx = warp_idx // warps_per_row, warp_idx % warps_per_row
|
|
|
|
if lane_idx == 0:
|
|
reduction_buffer[row_idx, col_idx, 0, 0] = warp_max
|
|
reduction_buffer[row_idx, col_idx, 0, 1] = warp_argmax.to(
|
|
cutlass.Float32
|
|
)
|
|
|
|
cute.arch.barrier()
|
|
block_reduce_max = -tXsX.element_type.inf
|
|
block_reduce_argmax = Int32(0xFFFFFFFF)
|
|
|
|
if lane_idx < warps_per_row:
|
|
block_reduce_max = reduction_buffer[row_idx, lane_idx, 0, 0]
|
|
block_reduce_argmax = reduction_buffer[row_idx, lane_idx, 0, 1].to(
|
|
cutlass.Int32
|
|
)
|
|
|
|
if cutlass.const_expr(self.use_redux):
|
|
warp_max, warp_argmax = warp_argmax_redux(
|
|
block_reduce_max, block_reduce_argmax
|
|
)
|
|
else:
|
|
warp_max, warp_argmax = warp_reduce_argmax(
|
|
block_reduce_max, block_reduce_argmax
|
|
)
|
|
else:
|
|
cute.arch.cluster_wait()
|
|
warps_per_row, cluster_n = reduction_buffer.shape[1]
|
|
cta_rank_in_cluster = cute.arch.block_idx_in_cluster()
|
|
rows_per_block, (warps_per_row, cluster_n), _, _ = reduction_buffer.shape
|
|
row_idx, col_idx = warp_idx // warps_per_row, warp_idx % warps_per_row
|
|
|
|
if warp_idx == 0:
|
|
with cute.arch.elect_one():
|
|
num_warps = rows_per_block * warps_per_row
|
|
cute.arch.mbarrier_arrive_and_expect_tx(
|
|
mbar_ptr,
|
|
num_warps
|
|
* cluster_n
|
|
* 2
|
|
* reduction_buffer.element_type.width
|
|
// 8,
|
|
)
|
|
|
|
if lane_idx < cluster_n:
|
|
store_shared_remote(
|
|
warp_max,
|
|
elem_pointer(
|
|
reduction_buffer,
|
|
(row_idx, (col_idx, cta_rank_in_cluster), 0, 0),
|
|
),
|
|
mbar_ptr,
|
|
peer_cta_rank_in_cluster=lane_idx,
|
|
)
|
|
store_shared_remote(
|
|
warp_argmax.to(cutlass.Float32),
|
|
elem_pointer(
|
|
reduction_buffer,
|
|
(row_idx, (col_idx, cta_rank_in_cluster), 0, 1),
|
|
),
|
|
mbar_ptr,
|
|
peer_cta_rank_in_cluster=lane_idx,
|
|
)
|
|
|
|
cute.arch.mbarrier_wait(mbar_ptr, phase=0)
|
|
block_reduce_val = -tXsX.element_type.inf
|
|
block_reduce_argmax = Int32(0xFFFFFFFF)
|
|
num_iter = cute.ceil_div(warps_per_row * cluster_n, cute.arch.WARP_SIZE)
|
|
|
|
for i in cutlass.range_constexpr(num_iter):
|
|
idx = lane_idx + i * cute.arch.WARP_SIZE
|
|
if idx < cute.size(reduction_buffer, mode=[1]):
|
|
element_max = reduction_buffer[row_idx, idx, 0, 0]
|
|
element_argmax = reduction_buffer[row_idx, idx, 0, 1].to(
|
|
cutlass.Int32
|
|
)
|
|
# Tie-break to the lowest index so the cluster-side
|
|
# reduction agrees with torch.argmax / the redux path.
|
|
if element_max > block_reduce_val:
|
|
block_reduce_val = element_max
|
|
block_reduce_argmax = element_argmax
|
|
elif element_max == block_reduce_val:
|
|
if element_argmax < block_reduce_argmax:
|
|
block_reduce_argmax = element_argmax
|
|
|
|
if cutlass.const_expr(self.use_redux):
|
|
warp_max, warp_argmax = warp_argmax_redux(
|
|
block_reduce_val, block_reduce_argmax
|
|
)
|
|
else:
|
|
warp_max, warp_argmax = warp_reduce_argmax(
|
|
block_reduce_val, block_reduce_argmax
|
|
)
|
|
|
|
row_idx = tXcX[0][0]
|
|
warps_per_row = tv_layout.shape[0][0] // cute.arch.WARP_SIZE
|
|
local_row_idx = row_idx - (bidx * tiler_mn[0])
|
|
first_warp_for_row = local_row_idx * warps_per_row
|
|
first_thread_for_row = first_warp_for_row * cute.arch.WARP_SIZE
|
|
|
|
if (
|
|
tidx == first_thread_for_row
|
|
and row_idx < shape[0]
|
|
and local_row_idx >= 0
|
|
and local_row_idx < tiler_mn[0]
|
|
and (self.cluster_n == 1 or bidy == 0)
|
|
):
|
|
# A row whose elements never beat -inf (all-NaN / all -inf) leaves
|
|
# the argmax at its 0xFFFFFFFF sentinel; emit the in-range index 0.
|
|
if warp_argmax == Int32(0xFFFFFFFF):
|
|
warp_argmax = Int32(0)
|
|
mO_max[local_row_idx] = warp_max.to(mO_max.element_type)
|
|
mO_idx[local_row_idx] = warp_argmax.to(mO_idx.element_type)
|