401 lines
15 KiB
Python
401 lines
15 KiB
Python
# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: BSD-3-Clause
|
|
|
|
# Redistribution and use in source and binary forms, with or without
|
|
# modification, are permitted provided that the following conditions are met:
|
|
|
|
# 1. Redistributions of source code must retain the above copyright notice, this
|
|
# list of conditions and the following disclaimer.
|
|
|
|
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
|
# this list of conditions and the following disclaimer in the documentation
|
|
# and/or other materials provided with the distribution.
|
|
|
|
# 3. Neither the name of the copyright holder nor the names of its
|
|
# contributors may be used to endorse or promote products derived from
|
|
# this software without specific prior written permission.
|
|
|
|
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
|
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
|
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
|
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
|
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
|
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
|
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
|
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
|
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
|
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
|
|
|
from typing import Tuple, Optional
|
|
|
|
import cutlass
|
|
import cutlass.cute as cute
|
|
import cutlass.cute.nvgpu.tcgen05 as tcgen05
|
|
import cutlass.pipeline as pipeline
|
|
|
|
|
|
@cute.jit
|
|
def load_qk(
|
|
iterations: int,
|
|
kv_step: cutlass.Int32,
|
|
k_args: Tuple,
|
|
scale_k_args: Optional[Tuple] = None,
|
|
q_args: Optional[Tuple] = None,
|
|
) -> Tuple[pipeline.PipelineProducer, pipeline.PipelineProducer]:
|
|
if cutlass.const_expr(q_args is not None):
|
|
tQgQ, tQsQ, tma_atom_q, load_q_producer = q_args
|
|
else:
|
|
tQgQ, tQsQ, tma_atom_q, load_q_producer = None, None, None, None
|
|
tKgK, tKsK, tma_atom_k, load_k_producer = k_args
|
|
tKgScaleK, tKsScaleK, tma_atom_scale_k, load_scale_k_producer = scale_k_args
|
|
|
|
scale_k_handle = load_scale_k_producer.acquire_and_advance()
|
|
cute.copy(
|
|
tma_atom_scale_k,
|
|
tKgScaleK[None, kv_step],
|
|
tKsScaleK[None, scale_k_handle.index],
|
|
tma_bar_ptr=scale_k_handle.barrier,
|
|
)
|
|
for iter in cutlass.range(iterations, unroll=1):
|
|
if cutlass.const_expr(q_args is not None):
|
|
q_handle = load_q_producer.acquire_and_advance()
|
|
cute.copy(
|
|
tma_atom_q,
|
|
tQgQ[None, iter],
|
|
tQsQ[None, q_handle.index],
|
|
tma_bar_ptr=q_handle.barrier,
|
|
)
|
|
k_handle = load_k_producer.acquire_and_advance()
|
|
cute.copy(
|
|
tma_atom_k,
|
|
tKgK[None, kv_step, iter],
|
|
tKsK[None, k_handle.index],
|
|
tma_bar_ptr=k_handle.barrier,
|
|
)
|
|
if cutlass.const_expr(q_args is not None):
|
|
return load_k_producer, load_scale_k_producer, load_q_producer
|
|
else:
|
|
return load_k_producer, load_scale_k_producer
|
|
|
|
|
|
@cute.jit
|
|
def load_v(
|
|
iterations: int,
|
|
kv_step: cutlass.Int32,
|
|
v_args: Tuple,
|
|
scale_v_args: Tuple,
|
|
) -> pipeline.PipelineProducer:
|
|
tVgV, tVsV, tma_atom_v, load_v_producer = v_args
|
|
tScaleVgV, tScaleVsV, tma_atom_scale_v, load_scale_v_producer = scale_v_args
|
|
scale_v_handle = load_scale_v_producer.acquire_and_advance()
|
|
cute.copy(
|
|
tma_atom_scale_v,
|
|
tScaleVgV[None, kv_step],
|
|
tScaleVsV[None, scale_v_handle.index],
|
|
tma_bar_ptr=scale_v_handle.barrier,
|
|
)
|
|
for iter in cutlass.range(iterations, unroll=1):
|
|
v_handle = load_v_producer.acquire_and_advance()
|
|
cute.copy(
|
|
tma_atom_v,
|
|
tVgV[None, iter, kv_step],
|
|
tVsV[None, v_handle.index],
|
|
tma_bar_ptr=v_handle.barrier,
|
|
)
|
|
return load_v_producer, load_scale_v_producer
|
|
|
|
|
|
@cute.jit
|
|
def get_scale_smem_layout(
|
|
scale_granularity: int,
|
|
d_r: int,
|
|
mma_tiler: cute.Tile,
|
|
major_mode: tcgen05.OperandMajorMode,
|
|
) -> Tuple[cute.Layout, cute.Tile]:
|
|
size_mn = mma_tiler[1] // 2 # 2cta by default
|
|
if cutlass.const_expr(major_mode == tcgen05.OperandMajorMode.MN): # v
|
|
scale_tiler = (mma_tiler[2] * d_r,)
|
|
tma_view_layout = cute.make_layout(
|
|
(mma_tiler[2] * d_r),
|
|
)
|
|
assert scale_granularity % mma_tiler[1] == 0, (
|
|
"scale_granularity must be divisible by mma_tiler[1]"
|
|
)
|
|
rest_l = scale_granularity // mma_tiler[1]
|
|
s2r_view_layout = cute.make_layout(
|
|
(size_mn, mma_tiler[2], (rest_l, d_r)),
|
|
stride=(0, d_r, (0, 1)),
|
|
)
|
|
else: # k
|
|
scale_tiler = (mma_tiler[1] * d_r,)
|
|
tma_view_layout = cute.make_layout((size_mn * d_r))
|
|
assert scale_granularity % mma_tiler[2] == 0, (
|
|
"scale_granularity must be divisible by mma_tiler[2]"
|
|
)
|
|
rest_l = scale_granularity // mma_tiler[2]
|
|
s2r_view_layout = cute.make_layout(
|
|
(size_mn, mma_tiler[2], (rest_l, d_r)),
|
|
stride=(d_r, 0, (0, 1)),
|
|
)
|
|
# Apply a trivial swizzle to make it a composed layout, which could be used to construct TMA atom
|
|
tma_view_smem_layout = cute.make_composed_layout(
|
|
cute.make_swizzle(0, 4, 3), 0, tma_view_layout
|
|
)
|
|
return tma_view_smem_layout, scale_tiler, s2r_view_layout
|
|
|
|
|
|
@cute.jit
|
|
def mma_qk(
|
|
iterations: int,
|
|
qk_tiled_mma: cute.TiledMma,
|
|
tensor_args: Tuple,
|
|
pipeline_args: Tuple,
|
|
):
|
|
tStS, tSrQ, tSrK_trans = tensor_args
|
|
mma_s_producer, load_q_consumer, load_q_releaser, dequant_kv_consumer = (
|
|
pipeline_args
|
|
)
|
|
cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster())
|
|
is_leader_cta = cta_rank_in_cluster % 2 == 0
|
|
if is_leader_cta:
|
|
s_handle = mma_s_producer.acquire_and_advance()
|
|
tStS_slice = tStS[None, None, None, s_handle.index]
|
|
qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
|
|
for iter in cutlass.range(iterations, unroll=1):
|
|
if cutlass.const_expr(load_q_consumer is not None):
|
|
load_q_consumer.wait_and_advance()
|
|
tSrQ_slice = tSrQ[None, None, None, iter]
|
|
k_trans_handle = dequant_kv_consumer.wait_and_advance()
|
|
tSrK_trans_slice = tSrK_trans[None, None, None, k_trans_handle.index]
|
|
num_kphases = cute.size(tSrQ_slice, mode=[2])
|
|
for kphase_idx in cutlass.range(num_kphases, unroll_full=True):
|
|
kphase_coord = (None, None, kphase_idx)
|
|
cute.gemm(
|
|
qk_tiled_mma,
|
|
tStS_slice,
|
|
tSrQ_slice[kphase_coord],
|
|
tSrK_trans_slice[kphase_coord],
|
|
tStS_slice,
|
|
)
|
|
qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
|
|
k_trans_handle.release()
|
|
if cutlass.const_expr(load_q_releaser is not None):
|
|
load_q_releaser.release()
|
|
load_q_releaser.advance()
|
|
s_handle.commit()
|
|
return mma_s_producer, load_q_consumer, dequant_kv_consumer
|
|
|
|
|
|
@cute.jit
|
|
def dequant_k(
|
|
iterations: int,
|
|
transform_warp_ids: Tuple,
|
|
dtype_args: Tuple,
|
|
tensor_args: Tuple,
|
|
pipeline_args: Tuple,
|
|
):
|
|
(k_dtype, q_dtype) = dtype_args
|
|
(sOrig, sScale, sTrans) = tensor_args
|
|
(load_kv_consumer, load_scale_consumer, dequant_kv_producer) = pipeline_args
|
|
tidx, _, _ = cute.arch.thread_idx()
|
|
THREADS_PER_WARP = 32
|
|
thread_idx = tidx % (THREADS_PER_WARP * len(transform_warp_ids))
|
|
r2s_copy_atom = cute.make_copy_atom(
|
|
cute.nvgpu.CopyUniversalOp(), k_dtype, num_bits_per_copy=32
|
|
)
|
|
# Construct tiled_copy satisfying 16 contiguous elts per copy atom
|
|
r2s_tiled_copy = cute.make_cotiled_copy(
|
|
r2s_copy_atom,
|
|
cute.make_layout((256, 16), stride=(16, 1)),
|
|
sTrans[(None, None, None, 0)].layout,
|
|
)
|
|
thr_r2s_tiled_copy = r2s_tiled_copy.get_slice(thread_idx)
|
|
tOsOrig = thr_r2s_tiled_copy.partition_S(sOrig)
|
|
tTsTrans = thr_r2s_tiled_copy.partition_D(sTrans)
|
|
tOrOrig = cute.make_rmem_tensor_like(
|
|
cute.append(
|
|
tOsOrig[None, None, None, None, 0].layout,
|
|
cute.make_layout(
|
|
2, stride=cute.cosize(tOsOrig[None, None, None, None, 0].layout)
|
|
),
|
|
),
|
|
k_dtype,
|
|
)
|
|
tTrTrans = cute.make_rmem_tensor_like(
|
|
cute.append(
|
|
tTsTrans[None, None, None, None, 0].layout,
|
|
cute.make_layout(
|
|
2, stride=cute.cosize(tTsTrans[None, None, None, None, 0].layout)
|
|
),
|
|
),
|
|
q_dtype,
|
|
)
|
|
tSsScale = thr_r2s_tiled_copy.partition_S(sScale)
|
|
tSrScale = cute.make_rmem_tensor_like(tSsScale[None, None, None, None, None, 0])
|
|
scale_handle = load_scale_consumer.wait_and_advance()
|
|
cute.autovec_copy(
|
|
tSsScale[None, None, None, None, None, scale_handle.index], tSrScale
|
|
)
|
|
cute.arch.fence_view_async_shared()
|
|
scale_handle.release()
|
|
# prefetch iter = 0
|
|
kv_handle = load_kv_consumer.wait_and_advance()
|
|
cute.autovec_copy(
|
|
tOsOrig[None, None, None, None, kv_handle.index],
|
|
tOrOrig[None, None, None, None, 0],
|
|
)
|
|
transformed_tensor = tOrOrig[None, None, None, None, 0].load().to(q_dtype)
|
|
scale = cute.TensorSSA(
|
|
tSrScale[None, None, None, None, 0].load(),
|
|
transformed_tensor.shape,
|
|
q_dtype,
|
|
)
|
|
transformed_tensor = transformed_tensor * scale
|
|
tTrTrans[None, None, None, None, 0].store(transformed_tensor)
|
|
cute.arch.fence_view_async_shared()
|
|
kv_handle.release()
|
|
for iter in cutlass.range(1, iterations, unroll_full=True):
|
|
kv_trans_handle = dequant_kv_producer.acquire_and_advance()
|
|
cute.autovec_copy(
|
|
tTrTrans[None, None, None, None, (iter - 1) % 2],
|
|
tTsTrans[None, None, None, None, kv_trans_handle.index],
|
|
)
|
|
cute.arch.fence_view_async_shared()
|
|
kv_trans_handle.commit()
|
|
kv_handle = load_kv_consumer.wait_and_advance()
|
|
cute.autovec_copy(
|
|
tOsOrig[None, None, None, None, kv_handle.index],
|
|
tOrOrig[None, None, None, None, iter % 2],
|
|
)
|
|
transformed_tensor = (
|
|
tOrOrig[None, None, None, None, iter % 2].load().to(q_dtype)
|
|
)
|
|
scale = cute.TensorSSA(
|
|
tSrScale[None, None, None, None, iter].load(),
|
|
transformed_tensor.shape,
|
|
q_dtype,
|
|
)
|
|
transformed_tensor = transformed_tensor * scale
|
|
tTrTrans[None, None, None, None, iter % 2].store(transformed_tensor)
|
|
cute.arch.fence_view_async_shared()
|
|
kv_handle.release()
|
|
kv_trans_handle = dequant_kv_producer.acquire_and_advance()
|
|
cute.autovec_copy(
|
|
tTrTrans[None, None, None, None, (iterations - 1) % 2],
|
|
tTsTrans[None, None, None, None, kv_trans_handle.index],
|
|
)
|
|
cute.arch.fence_view_async_shared()
|
|
kv_trans_handle.commit()
|
|
return load_kv_consumer, load_scale_consumer, dequant_kv_producer
|
|
|
|
|
|
@cute.jit
|
|
def dequant_v(
|
|
iterations: int,
|
|
transform_warp_ids: Tuple,
|
|
dtype_args: Tuple,
|
|
tensor_args: Tuple,
|
|
pipeline_args: Tuple,
|
|
):
|
|
(v_dtype, q_dtype) = dtype_args
|
|
(sOrig, sScale, sTrans) = tensor_args
|
|
(load_kv_consumer, load_scale_consumer, dequant_kv_producer) = pipeline_args
|
|
tidx, _, _ = cute.arch.thread_idx()
|
|
THREADS_PER_WARP = 32
|
|
thread_idx = tidx % (THREADS_PER_WARP * len(transform_warp_ids))
|
|
r2s_copy_atom = cute.make_copy_atom(
|
|
cute.nvgpu.CopyUniversalOp(), v_dtype, num_bits_per_copy=32
|
|
)
|
|
# Construct tiled_copy satisfying 16 contiguous elts per copy atom
|
|
r2s_tiled_copy = cute.make_cotiled_copy(
|
|
r2s_copy_atom,
|
|
cute.make_layout((256, 16), stride=(16, 1)),
|
|
sTrans[(None, None, None, 0)].layout,
|
|
)
|
|
thr_r2s_tiled_copy = r2s_tiled_copy.get_slice(thread_idx)
|
|
tOsOrig = thr_r2s_tiled_copy.partition_S(sOrig)
|
|
tTsTrans = thr_r2s_tiled_copy.partition_D(sTrans)
|
|
# double buffer for better perf
|
|
tOrOrig = cute.make_rmem_tensor_like(
|
|
cute.append(
|
|
tOsOrig[None, None, None, None, 0].layout,
|
|
cute.make_layout(
|
|
2, stride=cute.cosize(tOsOrig[None, None, None, None, 0].layout)
|
|
),
|
|
),
|
|
v_dtype,
|
|
)
|
|
tTrTrans = cute.make_rmem_tensor_like(
|
|
cute.append(
|
|
tTsTrans[None, None, None, None, 0].layout,
|
|
cute.make_layout(
|
|
2, stride=cute.cosize(tTsTrans[None, None, None, None, 0].layout)
|
|
),
|
|
),
|
|
q_dtype,
|
|
)
|
|
tSsScale = thr_r2s_tiled_copy.partition_S(sScale)
|
|
tSrScale = cute.make_rmem_tensor_like(tSsScale[None, None, None, None, None, 0])
|
|
scale_v_handle = load_scale_consumer.wait_and_advance()
|
|
cute.autovec_copy(
|
|
tSsScale[None, None, None, None, None, scale_v_handle.index],
|
|
tSrScale,
|
|
)
|
|
cute.arch.fence_view_async_shared()
|
|
scale_v_handle.release()
|
|
# prefetch iter = 0
|
|
kv_handle = load_kv_consumer.wait_and_advance()
|
|
cute.autovec_copy(
|
|
tOsOrig[None, None, None, None, kv_handle.index],
|
|
tOrOrig[None, None, None, None, 0],
|
|
)
|
|
transformed_tensor = tOrOrig[None, None, None, None, 0].load().to(q_dtype)
|
|
scale = cute.TensorSSA(
|
|
tSrScale[None, None, None, None, 0].load(),
|
|
transformed_tensor.shape,
|
|
q_dtype,
|
|
)
|
|
transformed_tensor = transformed_tensor * scale
|
|
tTrTrans[None, None, None, None, 0].store(transformed_tensor)
|
|
cute.arch.fence_view_async_shared()
|
|
kv_handle.release()
|
|
for iter in cutlass.range(1, iterations, unroll_full=True):
|
|
kv_trans_handle = dequant_kv_producer.acquire_and_advance()
|
|
cute.autovec_copy(
|
|
tTrTrans[None, None, None, None, (iter - 1) % 2],
|
|
tTsTrans[None, None, None, None, kv_trans_handle.index],
|
|
)
|
|
cute.arch.fence_view_async_shared()
|
|
kv_trans_handle.commit()
|
|
kv_handle = load_kv_consumer.wait_and_advance()
|
|
cute.autovec_copy(
|
|
tOsOrig[None, None, None, None, kv_handle.index],
|
|
tOrOrig[None, None, None, None, iter % 2],
|
|
)
|
|
transformed_tensor = (
|
|
tOrOrig[None, None, None, None, iter % 2].load().to(q_dtype)
|
|
)
|
|
scale = cute.TensorSSA(
|
|
tSrScale[
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
iter,
|
|
].load(),
|
|
transformed_tensor.shape,
|
|
q_dtype,
|
|
)
|
|
transformed_tensor = transformed_tensor * scale
|
|
tTrTrans[None, None, None, None, iter % 2].store(transformed_tensor)
|
|
cute.arch.fence_view_async_shared()
|
|
kv_handle.release()
|
|
kv_trans_handle = dequant_kv_producer.acquire_and_advance()
|
|
cute.autovec_copy(
|
|
tTrTrans[None, None, None, None, (iterations - 1) % 2],
|
|
tTsTrans[None, None, None, None, kv_trans_handle.index],
|
|
)
|
|
cute.arch.fence_view_async_shared()
|
|
kv_trans_handle.commit()
|
|
return load_kv_consumer, load_scale_consumer, dequant_kv_producer
|